executorlib 1.7.2__tar.gz → 1.7.4__tar.gz

{executorlib-1.7.2 → executorlib-1.7.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: executorlib
-Version: 1.7.2
+Version: 1.7.4
 Summary: Up-scale python functions for high performance computing (HPC) with executorlib.
 Project-URL: Homepage, https://github.com/pyiron/executorlib
 Project-URL: Documentation, https://executorlib.readthedocs.io
@@ -47,26 +47,26 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Physics
 Requires-Python: <3.14,>3.9
-Requires-Dist: cloudpickle<=3.1.1,>=2.0.0
+Requires-Dist: cloudpickle<=3.1.2,>=2.0.0
 Requires-Dist: pyzmq<=27.1.0,>=25.0.0
 Provides-Extra: all
-Requires-Dist: h5py<=3.15.0,>=3.6.0; extra == 'all'
-Requires-Dist: ipython<=9.0.2,>=7.33.0; extra == 'all'
+Requires-Dist: h5py<=3.15.1,>=3.6.0; extra == 'all'
+Requires-Dist: ipython<=9.9.0,>=7.33.0; extra == 'all'
 Requires-Dist: mpi4py<=4.1.1,>=3.1.4; extra == 'all'
-Requires-Dist: networkx<=3.4.2,>=2.8.8; extra == 'all'
+Requires-Dist: networkx<=3.6.1,>=2.8.8; extra == 'all'
 Requires-Dist: pygraphviz<=1.14,>=1.10; extra == 'all'
-Requires-Dist: pysqa==0.3.2; extra == 'all'
+Requires-Dist: pysqa==0.3.3; extra == 'all'
 Provides-Extra: cache
-Requires-Dist: h5py<=3.15.0,>=3.6.0; extra == 'cache'
+Requires-Dist: h5py<=3.15.1,>=3.6.0; extra == 'cache'
 Provides-Extra: cluster
-Requires-Dist: h5py<=3.15.0,>=3.6.0; extra == 'cluster'
-Requires-Dist: pysqa==0.3.2; extra == 'cluster'
+Requires-Dist: h5py<=3.15.1,>=3.6.0; extra == 'cluster'
+Requires-Dist: pysqa==0.3.3; extra == 'cluster'
 Provides-Extra: graph
-Requires-Dist: networkx<=3.4.2,>=2.8.8; extra == 'graph'
+Requires-Dist: networkx<=3.6.1,>=2.8.8; extra == 'graph'
 Requires-Dist: pygraphviz<=1.14,>=1.10; extra == 'graph'
 Provides-Extra: graphnotebook
-Requires-Dist: ipython<=9.0.2,>=7.33.0; extra == 'graphnotebook'
-Requires-Dist: networkx<=3.4.2,>=2.8.8; extra == 'graphnotebook'
+Requires-Dist: ipython<=9.9.0,>=7.33.0; extra == 'graphnotebook'
+Requires-Dist: networkx<=3.6.1,>=2.8.8; extra == 'graphnotebook'
 Requires-Dist: pygraphviz<=1.14,>=1.10; extra == 'graphnotebook'
 Provides-Extra: mpi
 Requires-Dist: mpi4py<=4.1.1,>=3.1.4; extra == 'mpi'
@@ -208,6 +208,7 @@ as hierarchical job scheduler within the allocations.
   * [Basic Functionality](https://executorlib.readthedocs.io/en/latest/1-single-node.html#basic-functionality)
   * [Parallel Functions](https://executorlib.readthedocs.io/en/latest/1-single-node.html#parallel-functions)
   * [Performance Optimization](https://executorlib.readthedocs.io/en/latest/1-single-node.html#performance-optimization)
+  * [Advanced Scheduling](https://executorlib.readthedocs.io/en/latest/1-single-node.html#advanced-scheduling)
   * [Testing and Debugging](https://executorlib.readthedocs.io/en/latest/1-single-node.html#testing-and-debugging)
 * [HPC Cluster Executor](https://executorlib.readthedocs.io/en/latest/2-hpc-cluster.html)
   * [SLURM](https://executorlib.readthedocs.io/en/latest/2-hpc-cluster.html#slurm)

{executorlib-1.7.2 → executorlib-1.7.4}/README.md

@@ -134,6 +134,7 @@ as hierarchical job scheduler within the allocations.
   * [Basic Functionality](https://executorlib.readthedocs.io/en/latest/1-single-node.html#basic-functionality)
   * [Parallel Functions](https://executorlib.readthedocs.io/en/latest/1-single-node.html#parallel-functions)
   * [Performance Optimization](https://executorlib.readthedocs.io/en/latest/1-single-node.html#performance-optimization)
+  * [Advanced Scheduling](https://executorlib.readthedocs.io/en/latest/1-single-node.html#advanced-scheduling)
   * [Testing and Debugging](https://executorlib.readthedocs.io/en/latest/1-single-node.html#testing-and-debugging)
 * [HPC Cluster Executor](https://executorlib.readthedocs.io/en/latest/2-hpc-cluster.html)
   * [SLURM](https://executorlib.readthedocs.io/en/latest/2-hpc-cluster.html#slurm)

{executorlib-1.7.2 → executorlib-1.7.4}/pyproject.toml

@@ -1,8 +1,8 @@
 [build-system]
 requires = [
-    "hatchling==1.27.0",
+    "hatchling>=1.27.0,<=1.28.0",
     "hatch-vcs==0.5.0",
-    "cloudpickle>=2.0.0,<=3.1.1",
+    "cloudpickle>=2.0.0,<=3.1.2",
     "pyzmq>=25.0.0,<=27.1.0",
 ]
 build-backend = "hatchling.build"
@@ -29,7 +29,7 @@ classifiers = [
     "Programming Language :: Python :: 3.13",
 ]
 dependencies = [
-    "cloudpickle>=2.0.0,<=3.1.1",
+    "cloudpickle>=2.0.0,<=3.1.2",
     "pyzmq>=25.0.0,<=27.1.0",
 ]
 dynamic = ["version"]
@@ -40,28 +40,28 @@ Documentation = "https://executorlib.readthedocs.io"
 Repository = "https://github.com/pyiron/executorlib"
 
 [project.optional-dependencies]
-cache = ["h5py>=3.6.0,<=3.15.0"]
+cache = ["h5py>=3.6.0,<=3.15.1"]
 graph = [
     "pygraphviz>=1.10,<=1.14",
-    "networkx>=2.8.8,<=3.4.2",
+    "networkx>=2.8.8,<=3.6.1",
 ]
 graphnotebook = [
     "pygraphviz>=1.10,<=1.14",
-    "networkx>=2.8.8,<=3.4.2",
-    "ipython>=7.33.0,<=9.0.2",
+    "networkx>=2.8.8,<=3.6.1",
+    "ipython>=7.33.0,<=9.9.0",
 ]
 mpi = ["mpi4py>=3.1.4,<=4.1.1"]
 cluster = [
-    "pysqa==0.3.2",
-    "h5py>=3.6.0,<=3.15.0",
+    "pysqa==0.3.3",
+    "h5py>=3.6.0,<=3.15.1",
 ]
 all = [
     "mpi4py>=3.1.4,<=4.1.1",
-    "pysqa==0.3.2",
-    "h5py>=3.6.0,<=3.15.0",
+    "pysqa==0.3.3",
+    "h5py>=3.6.0,<=3.15.1",
     "pygraphviz>=1.10,<=1.14",
-    "networkx>=2.8.8,<=3.4.2",
-    "ipython>=7.33.0,<=9.0.2",
+    "networkx>=2.8.8,<=3.6.1",
+    "ipython>=7.33.0,<=9.9.0",
 ]
 
 [tool.ruff]

{executorlib-1.7.2 → executorlib-1.7.4}/src/executorlib/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.7.2'
-__version_tuple__ = version_tuple = (1, 7, 2)
+__version__ = version = '1.7.4'
+__version_tuple__ = version_tuple = (1, 7, 4)
 
 __commit_id__ = commit_id = None

{executorlib-1.7.2 → executorlib-1.7.4}/src/executorlib/executor/base.py

@@ -21,6 +21,7 @@ class BaseExecutor(FutureExecutor, ABC):
 
     def __init__(self, executor: TaskSchedulerBase):
         self._task_scheduler = executor
+        self._is_active = True
 
     @property
     def max_workers(self) -> Optional[int]:
@@ -99,9 +100,49 @@ class BaseExecutor(FutureExecutor, ABC):
         Returns:
             Future: A Future representing the given call.
         """
-        return self._task_scheduler.submit(
-            *([fn] + list(args)), resource_dict=resource_dict, **kwargs
-        )
+        if self._is_active:
+            return self._task_scheduler.submit(
+                *([fn] + list(args)), resource_dict=resource_dict, **kwargs
+            )
+        else:
+            raise RuntimeError("cannot schedule new futures after shutdown")
+
+    def map(
+        self,
+        fn: Callable,
+        *iterables,
+        timeout: Optional[float] = None,
+        chunksize: int = 1,
+    ):
+        """Returns an iterator equivalent to map(fn, iter).
+
+        Args:
+            fn: A callable that will take as many arguments as there are
+                passed iterables.
+            timeout: The maximum number of seconds to wait. If None, then there
+                is no limit on the wait time.
+            chunksize: The size of the chunks the iterable will be broken into
+                before being passed to a child process. This argument is only
+                used by ProcessPoolExecutor; it is ignored by
+                ThreadPoolExecutor.
+
+        Returns:
+            An iterator equivalent to: map(func, *iterables) but the calls may
+            be evaluated out-of-order.
+
+        Raises:
+            TimeoutError: If the entire result iterator could not be generated
+                before the given timeout.
+            Exception: If fn(*args) raises for any values.
+        """
+        if self._is_active:
+            return self._task_scheduler.map(
+                *([fn] + list(iterables)),
+                timeout=timeout,
+                chunksize=chunksize,
+            )
+        else:
+            raise RuntimeError("cannot schedule new futures after shutdown")
 
     def shutdown(self, wait: bool = True, *, cancel_futures: bool = False):
         """
@@ -119,6 +160,7 @@ class BaseExecutor(FutureExecutor, ABC):
                 cancelled.
         """
         self._task_scheduler.shutdown(wait=wait, cancel_futures=cancel_futures)
+        self._is_active = False
 
     def __len__(self) -> int:
         """
@@ -143,3 +185,4 @@ class BaseExecutor(FutureExecutor, ABC):
         Exit method called when exiting the context manager.
         """
         self._task_scheduler.__exit__(*args, **kwargs)
+        self._is_active = False

{executorlib-1.7.2 → executorlib-1.7.4}/src/executorlib/executor/flux.py

@@ -65,6 +65,7 @@ class FluxJobExecutor(BaseExecutor):
         plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+        export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
     Examples:
@@ -105,6 +106,7 @@ class FluxJobExecutor(BaseExecutor):
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
     ):
         """
@@ -152,6 +154,7 @@ class FluxJobExecutor(BaseExecutor):
             plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+            export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
         """
@@ -189,6 +192,7 @@ class FluxJobExecutor(BaseExecutor):
                     refresh_rate=refresh_rate,
                     plot_dependency_graph=plot_dependency_graph,
                     plot_dependency_graph_filename=plot_dependency_graph_filename,
+                    export_workflow_filename=export_workflow_filename,
                 )
             )
         else:
@@ -255,6 +259,7 @@ class FluxClusterExecutor(BaseExecutor):
         plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+        export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
     Examples:
@@ -293,6 +298,7 @@ class FluxClusterExecutor(BaseExecutor):
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
     ):
         """
@@ -338,6 +344,7 @@ class FluxClusterExecutor(BaseExecutor):
             plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+            export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
         """
@@ -420,6 +427,7 @@ class FluxClusterExecutor(BaseExecutor):
                     refresh_rate=refresh_rate,
                     plot_dependency_graph=plot_dependency_graph,
                     plot_dependency_graph_filename=plot_dependency_graph_filename,
+                    export_workflow_filename=export_workflow_filename,
                 )
             )
 

{executorlib-1.7.2 → executorlib-1.7.4}/src/executorlib/executor/single.py

@@ -58,6 +58,7 @@ class SingleNodeExecutor(BaseExecutor):
         plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+        export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
     Examples:
@@ -94,6 +95,7 @@ class SingleNodeExecutor(BaseExecutor):
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
     ):
         """
@@ -138,6 +140,7 @@ class SingleNodeExecutor(BaseExecutor):
             plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+            export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
         """
@@ -171,6 +174,7 @@ class SingleNodeExecutor(BaseExecutor):
                     refresh_rate=refresh_rate,
                     plot_dependency_graph=plot_dependency_graph,
                     plot_dependency_graph_filename=plot_dependency_graph_filename,
+                    export_workflow_filename=export_workflow_filename,
                 )
             )
         else:
@@ -226,6 +230,7 @@ class TestClusterExecutor(BaseExecutor):
         plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+        export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
     Examples:
@@ -262,6 +267,7 @@ class TestClusterExecutor(BaseExecutor):
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
     ):
         """
@@ -299,6 +305,7 @@ class TestClusterExecutor(BaseExecutor):
             plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+            export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
         """
@@ -358,6 +365,7 @@ class TestClusterExecutor(BaseExecutor):
                     refresh_rate=refresh_rate,
                     plot_dependency_graph=plot_dependency_graph,
                     plot_dependency_graph_filename=plot_dependency_graph_filename,
+                    export_workflow_filename=export_workflow_filename,
                 )
             )
 

{executorlib-1.7.2 → executorlib-1.7.4}/src/executorlib/executor/slurm.py

@@ -63,6 +63,7 @@ class SlurmClusterExecutor(BaseExecutor):
         plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+        export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
     Examples:
@@ -101,6 +102,7 @@ class SlurmClusterExecutor(BaseExecutor):
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
     ):
         """
@@ -146,6 +148,7 @@ class SlurmClusterExecutor(BaseExecutor):
             plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+            export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
         """
@@ -225,6 +228,7 @@ class SlurmClusterExecutor(BaseExecutor):
                     refresh_rate=refresh_rate,
                     plot_dependency_graph=plot_dependency_graph,
                     plot_dependency_graph_filename=plot_dependency_graph_filename,
+                    export_workflow_filename=export_workflow_filename,
                 )
             )
 
@@ -275,6 +279,7 @@ class SlurmJobExecutor(BaseExecutor):
         plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+        export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
     Examples:
@@ -312,6 +317,7 @@ class SlurmJobExecutor(BaseExecutor):
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
     ):
         """
@@ -360,6 +366,7 @@ class SlurmJobExecutor(BaseExecutor):
             plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+            export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
 
         """
@@ -394,6 +401,7 @@ class SlurmJobExecutor(BaseExecutor):
                     refresh_rate=refresh_rate,
                     plot_dependency_graph=plot_dependency_graph,
                     plot_dependency_graph_filename=plot_dependency_graph_filename,
+                    export_workflow_filename=export_workflow_filename,
                 )
             )
         else:

{executorlib-1.7.2 → executorlib-1.7.4}/src/executorlib/standalone/serialize.py

@@ -75,7 +75,7 @@ def serialize_funct(
                 "kwargs": fn_kwargs,
             }
         )
-        task_key = fn.__name__ + _get_hash(binary=binary_all)
+        task_key = _get_function_name(fn=fn) + _get_hash(binary=binary_all)
     data = {
         "fn": fn,
         "args": fn_args,
@@ -99,3 +99,10 @@ def _get_hash(binary: bytes) -> str:
     # Remove specification of jupyter kernel from hash to be deterministic
     binary_no_ipykernel = re.sub(b"(?<=/ipykernel_)(.*)(?=/)", b"", binary)
     return str(hashlib.md5(binary_no_ipykernel).hexdigest())
+
+
+def _get_function_name(fn: Callable) -> str:
+    if hasattr(fn, "__name__"):
+        return fn.__name__
+    else:
+        return str(fn.__class__).split("'")[-2].split(".")[-1]

{executorlib-1.7.2 → executorlib-1.7.4}/src/executorlib/task_scheduler/base.py

@@ -143,6 +143,43 @@ class TaskSchedulerBase(FutureExecutor):
             )
         return f
 
+    def map(
+        self,
+        fn: Callable,
+        *iterables,
+        timeout: Optional[float] = None,
+        chunksize: int = 1,
+    ):
+        """Returns an iterator equivalent to map(fn, iter).
+
+        Args:
+            fn: A callable that will take as many arguments as there are
+                passed iterables.
+            timeout: The maximum number of seconds to wait. If None, then there
+                is no limit on the wait time.
+            chunksize: The size of the chunks the iterable will be broken into
+                before being passed to a child process. This argument is only
+                used by ProcessPoolExecutor; it is ignored by
+                ThreadPoolExecutor.
+
+        Returns:
+            An iterator equivalent to: map(func, *iterables) but the calls may
+            be evaluated out-of-order.
+
+        Raises:
+            TimeoutError: If the entire result iterator could not be generated
+                before the given timeout.
+            Exception: If fn(*args) raises for any values.
+        """
+        if isinstance(iterables, (list, tuple)) and any(
+            isinstance(i, Future) for i in iterables
+        ):
+            iterables = tuple(
+                i.result() if isinstance(i, Future) else i for i in iterables
+            )
+
+        return super().map(fn, *iterables, timeout=timeout, chunksize=chunksize)
+
     def shutdown(self, wait: bool = True, *, cancel_futures: bool = False):
         """
         Clean-up the resources associated with the Executor.

{executorlib-1.7.2 → executorlib-1.7.4}/src/executorlib/task_scheduler/interactive/dependency.py

@@ -13,6 +13,7 @@ from executorlib.standalone.interactive.arguments import (
 )
 from executorlib.task_scheduler.base import TaskSchedulerBase
 from executorlib.task_scheduler.interactive.dependency_plot import (
+    export_dependency_graph_function,
     generate_nodes_and_edges_for_plotting,
     generate_task_hash_for_plotting,
     plot_dependency_graph_function,
@@ -28,6 +29,7 @@ class DependencyTaskScheduler(TaskSchedulerBase):
         refresh_rate (float, optional): The refresh rate for updating the executor queue. Defaults to 0.01.
         plot_dependency_graph (bool, optional): Whether to generate and plot the dependency graph. Defaults to False.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+        export_workflow_filename (str): Name of the file to store the exported workflow graph in.
 
     Attributes:
         _future_hash_dict (Dict[str, Future]): A dictionary mapping task hash to future object.
@@ -44,6 +46,7 @@ class DependencyTaskScheduler(TaskSchedulerBase):
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        export_workflow_filename: Optional[str] = None,
     ) -> None:
         super().__init__(max_cores=max_cores)
         self._process_kwargs = {
@@ -61,7 +64,8 @@ class DependencyTaskScheduler(TaskSchedulerBase):
         self._future_hash_dict: dict = {}
         self._task_hash_dict: dict = {}
         self._plot_dependency_graph_filename = plot_dependency_graph_filename
-        if plot_dependency_graph_filename is None:
+        self._export_workflow_filename = export_workflow_filename
+        if plot_dependency_graph_filename is None and export_workflow_filename is None:
             self._generate_dependency_graph = plot_dependency_graph
         else:
             self._generate_dependency_graph = True
@@ -209,11 +213,18 @@ class DependencyTaskScheduler(TaskSchedulerBase):
                     v: k for k, v in self._future_hash_dict.items()
                 },
             )
-            return plot_dependency_graph_function(
-                node_lst=node_lst,
-                edge_lst=edge_lst,
-                filename=self._plot_dependency_graph_filename,
-            )
+            if self._export_workflow_filename is not None:
+                return export_dependency_graph_function(
+                    node_lst=node_lst,
+                    edge_lst=edge_lst,
+                    file_name=self._export_workflow_filename,
+                )
+            else:
+                return plot_dependency_graph_function(
+                    node_lst=node_lst,
+                    edge_lst=edge_lst,
+                    filename=self._plot_dependency_graph_filename,
+                )
         else:
             return None
 

{executorlib-1.7.2 → executorlib-1.7.4}/src/executorlib/task_scheduler/interactive/dependency_plot.py

@@ -1,8 +1,11 @@
+import inspect
+import json
 import os.path
 from concurrent.futures import Future
 from typing import Optional
 
 import cloudpickle
+import numpy as np
 
 from executorlib.standalone.select import FutureSelector
 
@@ -24,6 +27,12 @@ def generate_nodes_and_edges_for_plotting(
     edge_lst: list = []
     hash_id_dict: dict = {}
 
+    def extend_args(funct_dict):
+        sig = inspect.signature(funct_dict["fn"])
+        args = sig.bind(*funct_dict["args"], **funct_dict["kwargs"])
+        funct_dict["signature"] = args.arguments
+        return funct_dict
+
     def add_element(arg, link_to, label=""):
         """
         Add element to the node and edge lists.
@@ -39,6 +48,8 @@ def generate_nodes_and_edges_for_plotting(
                     "start": hash_id_dict[future_hash_inverse_dict[arg._future]],
                     "end": link_to,
                     "label": label + str(arg._selector),
+                    "end_label": label,
+                    "start_label": str(arg._selector),
                 }
             )
         elif isinstance(arg, Future):
@@ -53,39 +64,71 @@ def generate_nodes_and_edges_for_plotting(
             lst_no_future = [a if not isinstance(a, Future) else "$" for a in arg]
             node_id = len(node_lst)
             node_lst.append(
-                {"name": str(lst_no_future), "id": node_id, "shape": "circle"}
+                {
+                    "name": str(lst_no_future),
+                    "value": "python_workflow_definition.shared.get_list",
+                    "id": node_id,
+                    "type": "function",
+                    "shape": "box",
+                }
             )
             edge_lst.append({"start": node_id, "end": link_to, "label": label})
             for i, a in enumerate(arg):
                 if isinstance(a, Future):
-                    add_element(arg=a, link_to=node_id, label="ind: " + str(i))
+                    add_element(arg=a, link_to=node_id, label=str(i))
         elif isinstance(arg, dict) and any(isinstance(a, Future) for a in arg.values()):
             dict_no_future = {
                 kt: vt if not isinstance(vt, Future) else "$" for kt, vt in arg.items()
             }
             node_id = len(node_lst)
             node_lst.append(
-                {"name": str(dict_no_future), "id": node_id, "shape": "circle"}
+                {
+                    "name": str(dict_no_future),
+                    "value": "python_workflow_definition.shared.get_dict",
+                    "id": node_id,
+                    "type": "function",
+                    "shape": "box",
+                }
             )
             edge_lst.append({"start": node_id, "end": link_to, "label": label})
             for kt, vt in arg.items():
-                if isinstance(vt, Future):
-                    add_element(arg=vt, link_to=node_id, label="key: " + kt)
+                add_element(arg=vt, link_to=node_id, label=kt)
         else:
-            node_id = len(node_lst)
-            node_lst.append({"name": str(arg), "id": node_id, "shape": "circle"})
+            value_dict = {
+                str(n["value"]): n["id"] for n in node_lst if n["type"] == "input"
+            }
+            if str(arg) not in value_dict:
+                node_id = len(node_lst)
+                node_lst.append(
+                    {
+                        "name": label,
+                        "value": arg,
+                        "id": node_id,
+                        "type": "input",
+                        "shape": "circle",
+                    }
+                )
+            else:
+                node_id = value_dict[str(arg)]
             edge_lst.append({"start": node_id, "end": link_to, "label": label})
 
-    for k, v in task_hash_dict.items():
+    task_hash_modified_dict = {
+        k: extend_args(funct_dict=v) for k, v in task_hash_dict.items()
+    }
+
+    for k, v in task_hash_modified_dict.items():
         hash_id_dict[k] = len(node_lst)
         node_lst.append(
-            {"name": v["fn"].__name__, "id": hash_id_dict[k], "shape": "box"}
+            {
+                "name": v["fn"].__name__,
+                "type": "function",
+                "value": v["fn"].__module__ + "." + v["fn"].__name__,
+                "id": hash_id_dict[k],
+                "shape": "box",
+            }
         )
-    for k, task_dict in task_hash_dict.items():
-        for arg in task_dict["args"]:
-            add_element(arg=arg, link_to=hash_id_dict[k], label="")
-
-        for kw, v in task_dict["kwargs"].items():
+    for k, task_dict in task_hash_modified_dict.items():
+        for kw, v in task_dict["signature"].items():
             add_element(arg=v, link_to=hash_id_dict[k], label=str(kw))
 
     return node_lst, edge_lst
@@ -175,7 +218,10 @@ def plot_dependency_graph_function(
 
     graph = nx.DiGraph()
     for node in node_lst:
-        graph.add_node(node["id"], label=node["name"], shape=node["shape"])
+        if node["type"] == "input":
+            graph.add_node(node["id"], label=str(node["value"]), shape=node["shape"])
+        else:
+            graph.add_node(node["id"], label=str(node["name"]), shape=node["shape"])
     for edge in edge_lst:
         graph.add_edge(edge["start"], edge["end"], label=edge["label"])
     if filename is not None:
@@ -186,3 +232,76 @@ def plot_dependency_graph_function(
         from IPython.display import SVG, display  # noqa
 
         display(SVG(nx.nx_agraph.to_agraph(graph).draw(prog="dot", format="svg")))
+
+
+def export_dependency_graph_function(
+    node_lst: list, edge_lst: list, file_name: str = "workflow.json"
+):
+    """
+    Export the graph visualization of nodes and edges as a JSON dictionary.
+
+    Args:
+        node_lst (list): List of nodes.
+        edge_lst (list): List of edges.
+        file_name (str): Name of the file to store the exported graph in.
+    """
+    pwd_nodes_lst = []
+    for n in node_lst:
+        if n["type"] == "function":
+            pwd_nodes_lst.append(
+                {"id": n["id"], "type": n["type"], "value": n["value"]}
+            )
+        elif n["type"] == "input" and isinstance(n["value"], np.ndarray):
+            pwd_nodes_lst.append(
+                {
+                    "id": n["id"],
+                    "type": n["type"],
+                    "value": n["value"].tolist(),
+                    "name": n["name"],
+                }
+            )
+        else:
+            pwd_nodes_lst.append(
+                {
+                    "id": n["id"],
+                    "type": n["type"],
+                    "value": n["value"],
+                    "name": n["name"],
+                }
+            )
+
+    final_node = {"id": len(pwd_nodes_lst), "type": "output", "name": "result"}
+    pwd_nodes_lst.append(final_node)
+    pwd_edges_lst = [
+        (
+            {
+                "target": e["end"],
+                "targetPort": e["label"],
+                "source": e["start"],
+                "sourcePort": None,
+            }
+            if "start_label" not in e
+            else {
+                "target": e["end"],
+                "targetPort": e["end_label"],
+                "source": e["start"],
+                "sourcePort": e["start_label"],
+            }
+        )
+        for e in edge_lst
+    ]
+    pwd_edges_lst.append(
+        {
+            "target": final_node["id"],
+            "targetPort": None,
+            "source": max([e["target"] for e in pwd_edges_lst]),
+            "sourcePort": None,
+        }
+    )
+    pwd_dict = {
+        "version": "0.1.0",
+        "nodes": pwd_nodes_lst,
+        "edges": pwd_edges_lst,
+    }
+    with open(file_name, "w") as f:
+        json.dump(pwd_dict, f, indent=4)