parsl 2024.5.20__py3-none-any.whl → 2024.6.3__py3-none-any.whl

parsl/data_provider/rsync.py

@@ -1,8 +1,8 @@
 import logging
 import os
 
-from parsl.utils import RepresentationMixin
 from parsl.data_provider.staging import Staging
+from parsl.utils import RepresentationMixin
 
 logger = logging.getLogger(__name__)
 

parsl/data_provider/staging.py

@@ -1,9 +1,9 @@
 from concurrent.futures import Future
-from typing import Optional, Callable
+from typing import TYPE_CHECKING, Callable, Optional
+
 from parsl.app.futures import DataFuture
 from parsl.data_provider.files import File
 
-from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from parsl.data_provider.data_manager import DataManager
 

parsl/data_provider/zip.py

@@ -1,16 +1,15 @@
-import filelock
 import logging
 import os
-import parsl
 import zipfile
-
 from typing import Tuple
 
-from parsl.data_provider.staging import Staging
+import filelock
+
+import parsl
 from parsl.data_provider.files import File
+from parsl.data_provider.staging import Staging
 from parsl.errors import ParslError
 
-
 logger = logging.getLogger(__name__)
 
 

parsl/dataflow/dependency_resolvers.py

@@ -0,0 +1,115 @@
+from concurrent.futures import Future
+from dataclasses import dataclass
+from functools import singledispatch
+from typing import Callable, Sequence
+
+
+@dataclass
+class DependencyResolver:
+    """A DependencyResolver describes how app dependencies can be resolved.
+    It is specified as two functions: `traverse_to_gather` which turns an
+    app parameter into a sequence of futures which must be waited for before
+    the task can be executed (for example, in the case of
+    `DEEP_DEPENDENCY_RESOLVER` this traverses structures such as lists to
+    find every contained ``Future``), and `traverse_to_unwrap` which turns an
+    app parameter into its value to be passed to the app on execution
+    (for example in the case of `DEEP_DEPENDENCY_RESOLVER` this replaces a
+    list containing futures with a new list containing the values of those
+    resolved futures).
+
+    By default, Parsl will use `SHALLOW_DEPENDENCY_RESOLVER` which only
+    resolves Futures passed directly as arguments.
+    """
+    traverse_to_gather: Callable[[object], Sequence[Future]]
+    traverse_to_unwrap: Callable[[object], object]
+
+
+@singledispatch
+def shallow_traverse_to_gather(o):
+    # objects in general do not expose futures that we can see
+    return []
+
+
+@singledispatch
+def shallow_traverse_to_unwrap(o):
+    # objects in general unwrap to themselves
+    return o
+
+
+@shallow_traverse_to_gather.register
+def _(fut: Future):
+    return [fut]
+
+
+@shallow_traverse_to_unwrap.register
+@singledispatch
+def _(fut: Future):
+    assert fut.done()
+    return fut.result()
+
+
+@singledispatch
+def deep_traverse_to_gather(o):
+    # objects in general do not expose futures that we can see
+    return []
+
+
+@singledispatch
+def deep_traverse_to_unwrap(o):
+    # objects in general unwrap to themselves
+    return o
+
+
+@deep_traverse_to_gather.register
+def _(fut: Future):
+    return [fut]
+
+
+@deep_traverse_to_unwrap.register
+@singledispatch
+def _(fut: Future):
+    assert fut.done()
+    return fut.result()
+
+
+@deep_traverse_to_gather.register(tuple)
+@deep_traverse_to_gather.register(list)
+@deep_traverse_to_gather.register(set)
+def _(iterable):
+    return [e for v in iterable for e in deep_traverse_to_gather(v)]
+
+
+@deep_traverse_to_unwrap.register(tuple)
+@deep_traverse_to_unwrap.register(list)
+@deep_traverse_to_unwrap.register(set)
+@singledispatch
+def _(iterable):
+
+    type_ = type(iterable)
+    return type_(map(deep_traverse_to_unwrap, iterable))
+
+
+@deep_traverse_to_gather.register(dict)
+def _(dictionary):
+    futures = []
+    for key, value in dictionary.items():
+        futures.extend(deep_traverse_to_gather(key))
+        futures.extend(deep_traverse_to_gather(value))
+    return futures
+
+
+@deep_traverse_to_unwrap.register(dict)
+def _(dictionary):
+    unwrapped_dict = {}
+    for key, value in dictionary.items():
+        key = deep_traverse_to_unwrap(key)
+        value = deep_traverse_to_unwrap(value)
+        unwrapped_dict[key] = value
+    return unwrapped_dict
+
+
+DEEP_DEPENDENCY_RESOLVER = DependencyResolver(traverse_to_gather=deep_traverse_to_gather,
+                                              traverse_to_unwrap=deep_traverse_to_unwrap)
+
+SHALLOW_DEPENDENCY_RESOLVER = DependencyResolver(traverse_to_gather=shallow_traverse_to_gather,
+                                                 traverse_to_unwrap=shallow_traverse_to_unwrap)

parsl/dataflow/dflow.py

@@ -1,23 +1,25 @@
 from __future__ import annotations
+
 import atexit
+import datetime
+import inspect
 import logging
 import os
 import pathlib
 import pickle
 import random
-import time
-import typeguard
-import inspect
-import threading
 import sys
-import datetime
+import threading
+import time
+from concurrent.futures import Future
+from functools import partial
 from getpass import getuser
-from typeguard import typechecked
+from socket import gethostname
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 from uuid import uuid4
-from socket import gethostname
-from concurrent.futures import Future
-from functools import partial
+
+import typeguard
+from typeguard import typechecked
 
 import parsl
 from parsl.app.errors import RemoteExceptionWrapper
@@ -26,25 +28,29 @@ from parsl.channels import Channel
 from parsl.config import Config
 from parsl.data_provider.data_manager import DataManager
 from parsl.data_provider.files import File
+from parsl.dataflow.dependency_resolvers import SHALLOW_DEPENDENCY_RESOLVER
 from parsl.dataflow.errors import BadCheckpoint, DependencyError, JoinError
 from parsl.dataflow.futures import AppFuture
 from parsl.dataflow.memoization import Memoizer
 from parsl.dataflow.rundirs import make_rundir
-from parsl.dataflow.states import States, FINAL_STATES, FINAL_FAILURE_STATES
+from parsl.dataflow.states import FINAL_FAILURE_STATES, FINAL_STATES, States
 from parsl.dataflow.taskrecord import TaskRecord
-from parsl.errors import ConfigurationError, InternalConsistencyError, NoDataFlowKernelError
-from parsl.jobs.job_status_poller import JobStatusPoller
-from parsl.usage_tracking.usage import UsageTracker
+from parsl.errors import (
+    ConfigurationError,
+    InternalConsistencyError,
+    NoDataFlowKernelError,
+)
 from parsl.executors.base import ParslExecutor
 from parsl.executors.status_handling import BlockProviderExecutor
 from parsl.executors.threads import ThreadPoolExecutor
+from parsl.jobs.job_status_poller import JobStatusPoller
 from parsl.monitoring import MonitoringHub
+from parsl.monitoring.message_type import MessageType
 from parsl.monitoring.remote import monitor_wrapper
 from parsl.process_loggers import wrap_with_logs
 from parsl.providers.base import ExecutionProvider
-from parsl.utils import get_version, get_std_fname_mode, get_all_checkpoints, Timer
-
-from parsl.monitoring.message_type import MessageType
+from parsl.usage_tracking.usage import UsageTracker
+from parsl.utils import Timer, get_all_checkpoints, get_std_fname_mode, get_version
 
 logger = logging.getLogger(__name__)
 
@@ -203,6 +209,9 @@ class DataFlowKernel:
         self.tasks: Dict[int, TaskRecord] = {}
         self.submitter_lock = threading.Lock()
 
+        self.dependency_resolver = self.config.dependency_resolver if self.config.dependency_resolver is not None \
+            else SHALLOW_DEPENDENCY_RESOLVER
+
         atexit.register(self.atexit_cleanup)
 
     def __enter__(self):
@@ -852,8 +861,11 @@ class DataFlowKernel:
         depends: List[Future] = []
 
         def check_dep(d: Any) -> None:
-            if isinstance(d, Future):
-                depends.extend([d])
+            try:
+                depends.extend(self.dependency_resolver.traverse_to_gather(d))
+            except Exception:
+                logger.exception("Exception in dependency_resolver.traverse_to_gather")
+                raise
 
         # Check the positional args
         for dep in args:
@@ -870,7 +882,8 @@ class DataFlowKernel:
 
         return depends
 
-    def _unwrap_futures(self, args, kwargs):
+    def _unwrap_futures(self, args: Sequence[Any], kwargs: Dict[str, Any]) \
+            -> Tuple[Sequence[Any], Dict[str, Any], Sequence[Tuple[Exception, str]]]:
         """This function should be called when all dependencies have completed.
 
         It will rewrite the arguments for that task, replacing each Future
@@ -891,53 +904,40 @@ class DataFlowKernel:
         """
         dep_failures = []
 
+        def append_failure(e: Exception, dep: Future) -> None:
+            # If this Future is associated with a task inside this DFK,
+            # then refer to the task ID.
+            # Otherwise make a repr of the Future object.
+            if hasattr(dep, 'task_record') and dep.task_record['dfk'] == self:
+                tid = "task " + repr(dep.task_record['id'])
+            else:
+                tid = repr(dep)
+            dep_failures.extend([(e, tid)])
+
         # Replace item in args
         new_args = []
         for dep in args:
-            if isinstance(dep, Future):
-                try:
-                    new_args.extend([dep.result()])
-                except Exception as e:
-                    # If this Future is associated with a task inside this DFK,
-                    # then refer to the task ID.
-                    # Otherwise make a repr of the Future object.
-                    if hasattr(dep, 'task_record') and dep.task_record['dfk'] == self:
-                        tid = "task " + repr(dep.task_record['id'])
-                    else:
-                        tid = repr(dep)
-                    dep_failures.extend([(e, tid)])
-            else:
-                new_args.extend([dep])
+            try:
+                new_args.extend([self.dependency_resolver.traverse_to_unwrap(dep)])
+            except Exception as e:
+                append_failure(e, dep)
 
         # Check for explicit kwargs ex, fu_1=<fut>
         for key in kwargs:
             dep = kwargs[key]
-            if isinstance(dep, Future):
-                try:
-                    kwargs[key] = dep.result()
-                except Exception as e:
-                    if hasattr(dep, 'task_record'):
-                        tid = dep.task_record['id']
-                    else:
-                        tid = None
-                    dep_failures.extend([(e, tid)])
+            try:
+                kwargs[key] = self.dependency_resolver.traverse_to_unwrap(dep)
+            except Exception as e:
+                append_failure(e, dep)
 
         # Check for futures in inputs=[<fut>...]
         if 'inputs' in kwargs:
             new_inputs = []
             for dep in kwargs['inputs']:
-                if isinstance(dep, Future):
-                    try:
-                        new_inputs.extend([dep.result()])
-                    except Exception as e:
-                        if hasattr(dep, 'task_record'):
-                            tid = dep.task_record['id']
-                        else:
-                            tid = None
-                        dep_failures.extend([(e, tid)])
-
-                else:
-                    new_inputs.extend([dep])
+                try:
+                    new_inputs.extend([self.dependency_resolver.traverse_to_unwrap(dep)])
+                except Exception as e:
+                    append_failure(e, dep)
             kwargs['inputs'] = new_inputs
 
         return new_args, kwargs, dep_failures
@@ -1042,6 +1042,8 @@ class DataFlowKernel:
 
         func = self._add_output_deps(executor, app_args, app_kwargs, app_fu, func)
 
+        logger.debug("Added output dependencies")
+
         # Replace the function invocation in the TaskRecord with whatever file-staging
         # substitutions have been made.
         task_record.update({
@@ -1053,8 +1055,10 @@ class DataFlowKernel:
 
         self.tasks[task_id] = task_record
 
+        logger.debug("Gathering dependencies")
         # Get the list of dependencies for the task
         depends = self._gather_all_deps(app_args, app_kwargs)
+        logger.debug("Gathered dependencies")
         task_record['depends'] = depends
 
         depend_descs = []
@@ -1271,6 +1275,13 @@ class DataFlowKernel:
         atexit.unregister(self.atexit_cleanup)
         logger.info("Unregistered atexit hook")
 
+        if DataFlowKernelLoader._dfk is self:
+            logger.info("Unregistering default DFK")
+            parsl.clear()
+            logger.info("Unregistered default DFK")
+        else:
+            logger.debug("Cleaning up non-default DFK - not unregistering")
+
         logger.info("DFK cleanup complete")
 
     def checkpoint(self, tasks: Optional[Sequence[TaskRecord]] = None) -> str:

parsl/dataflow/errors.py

@@ -1,6 +1,7 @@
-from parsl.errors import ParslError
 from typing import Optional, Sequence, Tuple
 
+from parsl.errors import ParslError
+
 
 class DataFlowException(ParslError):
     """Base class for all exceptions.

parsl/dataflow/futures.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
 
-from concurrent.futures import Future
 import logging
 import threading
+from concurrent.futures import Future
 from typing import Any, Optional, Sequence, Union
 
 import parsl.app.app as app
-
 from parsl.app.futures import DataFuture
 from parsl.dataflow.taskrecord import TaskRecord
 

parsl/dataflow/memoization.py

@@ -1,18 +1,18 @@
 from __future__ import annotations
+
 import hashlib
-from functools import lru_cache, singledispatch
 import logging
 import pickle
-from parsl.dataflow.taskrecord import TaskRecord
+from functools import lru_cache, singledispatch
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
-from typing import Dict, Any, List, Optional, TYPE_CHECKING
+from parsl.dataflow.taskrecord import TaskRecord
 
 if TYPE_CHECKING:
     from parsl import DataFlowKernel  # import loop at runtime - needed for typechecking - TODO turn into "if typing:"
 
-from concurrent.futures import Future
-
 import types
+from concurrent.futures import Future
 
 logger = logging.getLogger(__name__)
 

parsl/dataflow/rundirs.py

@@ -1,6 +1,6 @@
+import logging
 import os
 from glob import glob
-import logging
 
 logger = logging.getLogger(__name__)
 

parsl/dataflow/taskrecord.py

@@ -1,19 +1,18 @@
 from __future__ import annotations
 
-import threading
 import datetime
-from typing_extensions import TypedDict
+import threading
 from concurrent.futures import Future
 
-
 # only for type checking:
-from typing import Any, Callable, Dict, Optional, List, Sequence, TYPE_CHECKING, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Union
+
+from typing_extensions import TypedDict
 
 if TYPE_CHECKING:
     from parsl.dataflow.futures import AppFuture
 
 import parsl.dataflow.dflow as dflow
-
 from parsl.dataflow.states import States
 
 

parsl/executors/__init__.py

@@ -1,8 +1,8 @@
-from parsl.executors.threads import ThreadPoolExecutor
-from parsl.executors.workqueue.executor import WorkQueueExecutor
+from parsl.executors.flux.executor import FluxExecutor
 from parsl.executors.high_throughput.executor import HighThroughputExecutor
 from parsl.executors.high_throughput.mpi_executor import MPIExecutor
-from parsl.executors.flux.executor import FluxExecutor
+from parsl.executors.threads import ThreadPoolExecutor
+from parsl.executors.workqueue.executor import WorkQueueExecutor
 
 __all__ = ['ThreadPoolExecutor',
            'HighThroughputExecutor',

parsl/executors/base.py

@@ -2,6 +2,7 @@ import os
 from abc import ABCMeta, abstractmethod
 from concurrent.futures import Future
 from typing import Any, Callable, Dict, Optional
+
 from typing_extensions import Literal, Self
 
 from parsl.monitoring.radios import MonitoringRadio

parsl/executors/flux/execute_parsl_task.py

@@ -1,12 +1,12 @@
 """Script for executing tasks inside of Flux jobs."""
 
 import argparse
-import os
 import logging
+import os
 
+from parsl.executors.flux import TaskResult
 from parsl.executors.high_throughput.process_worker_pool import execute_task
 from parsl.serialize import serialize
-from parsl.executors.flux import TaskResult
 
 
 def main():

parsl/executors/flux/executor.py

@@ -1,33 +1,32 @@
 """Defines the FluxExecutor class."""
 
+import collections
 import concurrent.futures as cf
 import functools
+import itertools
 import os
+import queue
+import shutil
 import sys
-import uuid
 import threading
-import itertools
-import shutil
-import queue
-from socket import gethostname
-import collections
-from collections.abc import Mapping, Callable
-from typing import Optional, Any, Dict
+import uuid
 import weakref
+from collections.abc import Callable, Mapping
+from socket import gethostname
+from typing import Any, Dict, Optional
 
 import zmq
 
-from parsl.utils import RepresentationMixin
+from parsl.app.errors import AppException
 from parsl.executors.base import ParslExecutor
+from parsl.executors.errors import ScalingFailed
 from parsl.executors.flux.execute_parsl_task import __file__ as _WORKER_PATH
 from parsl.executors.flux.flux_instance_manager import __file__ as _MANAGER_PATH
-from parsl.executors.errors import ScalingFailed
 from parsl.providers import LocalProvider
 from parsl.providers.base import ExecutionProvider
 from parsl.serialize import deserialize, pack_res_spec_apply_message
 from parsl.serialize.errors import SerializationError
-from parsl.app.errors import AppException
-
+from parsl.utils import RepresentationMixin
 
 _WORKER_PATH = os.path.realpath(_WORKER_PATH)
 _MANAGER_PATH = os.path.realpath(_MANAGER_PATH)

parsl/executors/flux/flux_instance_manager.py

@@ -1,10 +1,10 @@
 """Script meant to be the initial program of a Flux instance."""
 
 import argparse
+import logging
 import os
 from os.path import dirname
-import logging
-from socket import gethostname, gethostbyname
+from socket import gethostbyname, gethostname
 
 import zmq
 
@@ -16,8 +16,8 @@ def main():
     encapsulating Flux instance.
     """
     # flux imports only available when launched under Flux instance
-    import flux.job
     import flux
+    import flux.job
 
     logging.basicConfig(
         level=logging.DEBUG, format="%(asctime)s [%(levelname)s] %(message)s"

parsl/executors/high_throughput/errors.py

@@ -10,3 +10,13 @@ class WorkerLost(Exception):
 
     def __str__(self):
         return self.__repr__()
+
+
+class CommandClientTimeoutError(Exception):
+    """Raised when the command client times out waiting for a response.
+    """
+
+
+class CommandClientBadError(Exception):
+    """Raised when the command client is bad from an earlier timeout.
+    """