WIP: Adding config to nx-parallel

_nx_parallel/__init__.py

@@ -9,6 +9,20 @@ def get_info():
         "package": "nx_parallel",
         "url": "https://github.com/networkx/nx-parallel",
         "short_summary": "Parallel backend for NetworkX algorithms",
+        "default_config": {
+            "backend": None,
+            "n_jobs": None,
+            "verbose": 0,
+            "temp_folder": None,
+            "max_nbytes": "1M",
+            "mmap_mode": "r",
+            "prefer": None,
+            "require": None,
+            "return_as": "list",
+            "timeout": None,
+            "pre_dispatch": "2 * n_jobs",
+            "batch_size": "auto",
+        },
         "functions": {
             "number_of_isolates": {
                 "url": "https://github.com/networkx/nx-parallel/blob/main/nx_parallel/algorithms/isolate.py#L8",
@@ -150,18 +164,23 @@ def get_info():
                     'get_chunks : str, function (default = "chunks")': "A function that takes in an iterable of all the nodes as input and returns an iterable `node_chunks`. The default chunking is done by slicing the `G.nodes` into `n` chunks, where `n` is the number of CPU cores."
                 },
             },
+            "get_curr_configs": {
+                "url": "https://github.com/networkx/nx-parallel/blob/main/nx_parallel/utils/config.py#L44",
+                "additional_docs": "Returns the current configuration settings for nx_parallel.",
+                "additional_parameters": None,
+            },
             "chunks": {
-                "url": "https://github.com/networkx/nx-parallel/blob/main/nx_parallel/utils/chunk.py#L8",
+                "url": "https://github.com/networkx/nx-parallel/blob/main/nx_parallel/utils/chunk.py#L9",
                 "additional_docs": "Divides an iterable into chunks of size n",
                 "additional_parameters": None,
             },
             "cpu_count": {
-                "url": "https://github.com/networkx/nx-parallel/blob/main/nx_parallel/utils/chunk.py#L18",
+                "url": "https://github.com/networkx/nx-parallel/blob/main/nx_parallel/utils/chunk.py#L19",
                 "additional_docs": "Returns the number of logical CPUs or cores",
                 "additional_parameters": None,
             },
             "create_iterables": {
-                "url": "https://github.com/networkx/nx-parallel/blob/main/nx_parallel/utils/chunk.py#L26",
+                "url": "https://github.com/networkx/nx-parallel/blob/main/nx_parallel/utils/chunk.py#L28",
                 "additional_docs": "Creates an iterable of function inputs for parallel computation based on the provided iterator type.",
                 "additional_parameters": {
                     "G : NetworkX graph": "iterator : str Type of iterator. Valid values are 'node', 'edge', 'isolate'",

_nx_parallel/update_get_info.py

@@ -1,5 +1,6 @@
 import os
 import ast
+from nx_parallel.utils.config import _configs
 
 __all__ = [
     "get_funcs_info",
@@ -134,7 +135,8 @@ def get_url(file_path, function_name):
 
 # Creating a temp__init__.py file
 
-string = '''# This file was automatically generated by update_get_info.py
+string = (
+    '''# This file was automatically generated by update_get_info.py
 
 
 def get_info():
@@ -145,7 +147,11 @@ def get_info():
         "package": "nx_parallel",
         "url": "https://github.com/networkx/nx-parallel",
         "short_summary": "Parallel backend for NetworkX algorithms",
-        "functions": '''
+        "default_config": '''
+    + str(_configs.get_config_dict())
+    + """,
+        "functions": """
+)
 
 with open("_nx_parallel/temp__init__.py", "w") as f:
     f.write(string + str(get_funcs_info()) + "}\n")

nx_parallel/algorithms/cluster.py

@@ -47,15 +47,18 @@ def _compute_clustering_chunk(node_iter_chunk):
     else:
         node_iter = list(G.nbunch_iter(nodes))
 
-    total_cores = nxp.cpu_count()
+    configs = nxp.get_curr_configs()
+    cores = nxp.cpu_count()
+    total_cores = cores if cores is not None else -1
+    configs["n_jobs"] = total_cores
 
     if get_chunks == "chunks":
         num_in_chunk = max(len(node_iter) // total_cores, 1)
         node_iter_chunks = nxp.chunks(node_iter, num_in_chunk)
     else:
         node_iter_chunks = get_chunks(node_iter)
 
-    result = Parallel(n_jobs=total_cores)(
+    result = Parallel(**configs)(
         delayed(_compute_clustering_chunk)(node_iter_chunk)
         for node_iter_chunk in node_iter_chunks
     )

nx_parallel/utils/__init__.py

@@ -1 +1,2 @@
 from .chunk import *
+from .config import *

nx_parallel/utils/chunk.py

@@ -1,6 +1,7 @@
 import itertools
 import os
 import networkx as nx
+import nx_parallel as nxp
 
 __all__ = ["chunks", "cpu_count", "create_iterables"]
 
@@ -20,7 +21,8 @@ def cpu_count():
     # Check if we are running under pytest
     if "PYTEST_CURRENT_TEST" in os.environ:
         return 2
-    return os.cpu_count()
+    else:
+        return nxp.get_curr_configs("n_jobs")
 
 
 def create_iterables(G, iterator, n_cores, list_of_iterator=None):

nx_parallel/utils/config.py

@@ -0,0 +1,55 @@
+from networkx.utils.configs import Config
+from typing import Union
+from dataclasses import asdict
+import networkx as nx
+
+__all__ = [
+    "NxpConfig",
+    "_configs",
+    "get_curr_configs",
+]
+
+
+class NxpConfig(Config):
+    backend: str = None
+    n_jobs: int = None
+    verbose: int = 0
+    temp_folder: str = None
+    max_nbytes: Union[int, str] = "1M"
+    mmap_mode: str = "r"
+    prefer: str = None
+    require: str = None
+    return_as: str = "list"
+    timeout: float = None
+    pre_dispatch: str = "2 * n_jobs"
+    batch_size: int = "auto"
+
+    def get_config_dict(self, config=None):
+        """Return the default configuration as a dictionary."""
+        config_dict = asdict(self)
+        if config is None:
+            return config_dict
+        elif isinstance(config, list):
+            new_config = {k: config_dict[k] for k in config if k in config_dict}
+            return new_config
+        elif config in config_dict:
+            return config_dict[config]
+        else:
+            raise KeyError(f"Invalid config: {config}")
+
+
+_configs = NxpConfig()
+
+
+def get_curr_configs(config=None):
+    """Returns the current configuration settings for nx_parallel."""
+    config_dict = dict(nx.config.backends.parallel)
+    if config is None:
+        return config_dict
+    elif isinstance(config, list):
+        new_config = {k: config_dict[k] for k in config if k in config_dict}
+        return new_config
+    elif config in config_dict:
+        return config_dict[config]
+    else:
+        raise KeyError(f"Invalid config: {config}")