parsl 2024.5.27__py3-none-any.whl → 2024.6.10__py3-none-any.whl

parsl/data_provider/ftp.py

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

parsl/data_provider/globus.py

@@ -1,15 +1,16 @@
-import logging
 import json
-import globus_sdk
+import logging
 import os
-import parsl
-import typeguard
-
 from functools import partial
 from typing import Optional
+
+import globus_sdk
+import typeguard
+
+import parsl
 from parsl.app.app import python_app
-from parsl.utils import RepresentationMixin
 from parsl.data_provider.staging import Staging
+from parsl.utils import RepresentationMixin
 
 logger = logging.getLogger(__name__)
 

parsl/data_provider/http.py

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

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/dflow.py

@@ -1,23 +1,26 @@
 from __future__ import annotations
+
 import atexit
+import concurrent.futures as cf
+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
@@ -31,21 +34,24 @@ 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__)
 
@@ -204,6 +210,8 @@ class DataFlowKernel:
         self.tasks: Dict[int, TaskRecord] = {}
         self.submitter_lock = threading.Lock()
 
+        self.dependency_launch_pool = cf.ThreadPoolExecutor(max_workers=1, thread_name_prefix="Dependency-Launch")
+
         self.dependency_resolver = self.config.dependency_resolver if self.config.dependency_resolver is not None \
             else SHALLOW_DEPENDENCY_RESOLVER
 
@@ -212,9 +220,24 @@ class DataFlowKernel:
     def __enter__(self):
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
-        logger.debug("Exiting the context manager, calling cleanup for DFK")
-        self.cleanup()
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        mode = self.config.exit_mode
+        logger.debug("Exiting context manager, with exit mode '%s'", mode)
+        if mode == "cleanup":
+            logger.info("Calling cleanup for DFK")
+            self.cleanup()
+        elif mode == "skip":
+            logger.info("Skipping all cleanup handling")
+        elif mode == "wait":
+            if exc_type is None:
+                logger.info("Waiting for all tasks to complete")
+                self.wait_for_current_tasks()
+                self.cleanup()
+            else:
+                logger.info("There was an exception - cleaning up without waiting for task completion")
+                self.cleanup()
+        else:
+            raise InternalConsistencyError(f"Exit case for {mode} should be unreachable, validated by typeguard on Config()")
 
     def _send_task_log_info(self, task_record: TaskRecord) -> None:
         if self.monitoring:
@@ -606,9 +629,9 @@ class DataFlowKernel:
         return kwargs.get('_parsl_staging_inhibit', False)
 
     def launch_if_ready(self, task_record: TaskRecord) -> None:
-        """
-        launch_if_ready will launch the specified task, if it is ready
-        to run (for example, without dependencies, and in pending state).
+        """Schedules a task record for re-inspection to see if it is ready
+        for launch and for launch if it is ready. The call will return
+        immediately.
 
         This should be called by any piece of the DataFlowKernel that
         thinks a task may have become ready to run.
@@ -617,13 +640,17 @@ class DataFlowKernel:
         ready to run - launch_if_ready will not incorrectly launch that
         task.
 
-        It is also not an error to call launch_if_ready on a task that has
-        already been launched - launch_if_ready will not re-launch that
-        task.
-
         launch_if_ready is thread safe, so may be called from any thread
         or callback.
         """
+        self.dependency_launch_pool.submit(self._launch_if_ready_async, task_record)
+
+    @wrap_with_logs
+    def _launch_if_ready_async(self, task_record: TaskRecord) -> None:
+        """
+        _launch_if_ready will launch the specified task, if it is ready
+        to run (for example, without dependencies, and in pending state).
+        """
         exec_fu = None
 
         task_id = task_record['id']
@@ -1266,6 +1293,10 @@ class DataFlowKernel:
             self.monitoring.close()
             logger.info("Terminated monitoring")
 
+        logger.info("Terminating dependency launch pool")
+        self.dependency_launch_pool.shutdown()
+        logger.info("Terminated dependency launch pool")
+
         logger.info("Unregistering atexit hook")
         atexit.unregister(self.atexit_cleanup)
         logger.info("Unregistered atexit hook")

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/executor.py

@@ -1,44 +1,39 @@
+import logging
+import math
+import pickle
+import threading
 import typing
+import warnings
 from collections import defaultdict
 from concurrent.futures import Future
-import typeguard
-import logging
-import threading
-import queue
-import pickle
 from dataclasses import dataclass
-from multiprocessing import Process, Queue
-from typing import Dict, Sequence
-from typing import List, Optional, Tuple, Union, Callable
-import math
-import warnings
+from multiprocessing import Process
+from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
+
+import typeguard
 
 import parsl.launchers
-from parsl.usage_tracking.api import UsageInformation
-from parsl.serialize import pack_res_spec_apply_message, deserialize
-from parsl.serialize.errors import SerializationError, DeserializationError
+from parsl import curvezmq
+from parsl.addresses import get_all_addresses
 from parsl.app.errors import RemoteExceptionWrapper
-from parsl.jobs.states import JobStatus, JobState, TERMINAL_STATES
-from parsl.executors.high_throughput import zmq_pipes
-from parsl.executors.high_throughput import interchange
-from parsl.executors.errors import (
-    BadMessage, ScalingFailed,
-)
+from parsl.data_provider.staging import Staging
+from parsl.executors.errors import BadMessage, ScalingFailed
+from parsl.executors.high_throughput import interchange, zmq_pipes
+from parsl.executors.high_throughput.errors import CommandClientTimeoutError
 from parsl.executors.high_throughput.mpi_prefix_composer import (
     VALID_LAUNCHERS,
-    validate_resource_spec
+    validate_resource_spec,
 )
-
-from parsl import curvezmq
 from parsl.executors.status_handling import BlockProviderExecutor
-from parsl.providers.base import ExecutionProvider
-from parsl.data_provider.staging import Staging
-from parsl.addresses import get_all_addresses
-from parsl.process_loggers import wrap_with_logs
-
+from parsl.jobs.states import TERMINAL_STATES, JobState, JobStatus
 from parsl.multiprocessing import ForkProcess
-from parsl.utils import RepresentationMixin
+from parsl.process_loggers import wrap_with_logs
 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 DeserializationError, SerializationError
+from parsl.usage_tracking.api import UsageInformation
+from parsl.utils import RepresentationMixin
 
 logger = logging.getLogger(__name__)
 
@@ -415,13 +410,13 @@ class HighThroughputExecutor(BlockProviderExecutor, RepresentationMixin, UsageIn
             )
 
         self.outgoing_q = zmq_pipes.TasksOutgoing(
-            curvezmq.ClientContext(self.cert_dir), "127.0.0.1", self.interchange_port_range
+            "127.0.0.1", self.interchange_port_range, self.cert_dir
         )
         self.incoming_q = zmq_pipes.ResultsIncoming(
-            curvezmq.ClientContext(self.cert_dir), "127.0.0.1", self.interchange_port_range
+            "127.0.0.1", self.interchange_port_range, self.cert_dir
         )
         self.command_client = zmq_pipes.CommandClient(
-            curvezmq.ClientContext(self.cert_dir), "127.0.0.1", self.interchange_port_range
+            "127.0.0.1", self.interchange_port_range, self.cert_dir
         )
 
         self._queue_management_thread = None
@@ -531,10 +526,9 @@ class HighThroughputExecutor(BlockProviderExecutor, RepresentationMixin, UsageIn
         Starts the interchange process locally and uses an internal command queue to
         get the worker task and result ports that the interchange has bound to.
         """
-        comm_q = Queue(maxsize=10)
         self.interchange_proc = ForkProcess(target=interchange.starter,
-                                            args=(comm_q,),
-                                            kwargs={"client_ports": (self.outgoing_q.port,
+                                            kwargs={"client_address": "127.0.0.1",
+                                                    "client_ports": (self.outgoing_q.port,
                                                                      self.incoming_q.port,
                                                                      self.command_client.port),
                                                     "interchange_address": self.address,
@@ -552,9 +546,10 @@ class HighThroughputExecutor(BlockProviderExecutor, RepresentationMixin, UsageIn
                                             name="HTEX-Interchange"
                                             )
         self.interchange_proc.start()
+
         try:
-            (self.worker_task_port, self.worker_result_port) = comm_q.get(block=True, timeout=120)
-        except queue.Empty:
+            (self.worker_task_port, self.worker_result_port) = self.command_client.run("WORKER_PORTS", timeout_s=120)
+        except CommandClientTimeoutError:
             logger.error("Interchange has not completed initialization in 120s. Aborting")
             raise Exception("Interchange failed to start")