parsl 2024.3.11__py3-none-any.whl → 2025.1.13__py3-none-any.whl

{parsl-2024.3.11.data → parsl-2025.1.13.data}/scripts/process_worker_pool.py

@@ -1,41 +1,47 @@
 #!python
 
 import argparse
+import json
 import logging
+import math
+import multiprocessing
 import os
-import sys
+import pickle
 import platform
+import queue
+import subprocess
+import sys
 import threading
-import pickle
 import time
-import queue
 import uuid
-from typing import Sequence, Optional, Dict, List
-
-import zmq
-import math
-import json
-import psutil
-import multiprocessing
 from multiprocessing.managers import DictProxy
 from multiprocessing.sharedctypes import Synchronized
+from typing import Dict, List, Optional, Sequence
+
+import psutil
+import zmq
 
 from parsl import curvezmq
-from parsl.process_loggers import wrap_with_logs
-from parsl.version import VERSION as PARSL_VERSION
+from parsl.addresses import tcp_url
 from parsl.app.errors import RemoteExceptionWrapper
+from parsl.executors.execute_task import execute_task
 from parsl.executors.high_throughput.errors import WorkerLost
-from parsl.executors.high_throughput.probe import probe_addresses
-from parsl.multiprocessing import SpawnContext
-from parsl.serialize import unpack_res_spec_apply_message, serialize
+from parsl.executors.high_throughput.mpi_prefix_composer import (
+    VALID_LAUNCHERS,
+    compose_all,
+)
 from parsl.executors.high_throughput.mpi_resource_management import (
+    MPITaskScheduler,
     TaskScheduler,
-    MPITaskScheduler
 )
-
-from parsl.executors.high_throughput.mpi_prefix_composer import compose_all, VALID_LAUNCHERS
+from parsl.executors.high_throughput.probe import probe_addresses
+from parsl.multiprocessing import SpawnContext
+from parsl.process_loggers import wrap_with_logs
+from parsl.serialize import serialize
+from parsl.version import VERSION as PARSL_VERSION
 
 HEARTBEAT_CODE = (2 ** 32) - 1
+DRAINED_CODE = (2 ** 32) - 2
 
 
 class Manager:
@@ -73,7 +79,8 @@ class Manager:
                  enable_mpi_mode: bool = False,
                  mpi_launcher: str = "mpiexec",
                  available_accelerators: Sequence[str],
-                 cert_dir: Optional[str]):
+                 cert_dir: Optional[str],
+                 drain_period: Optional[int]):
         """
         Parameters
         ----------
@@ -138,6 +145,9 @@ class Manager:
 
         cert_dir : str | None
             Path to the certificate directory.
+
+        drain_period: int | None
+            Number of seconds to drain after  TODO: could be a nicer timespec involving m,s,h qualifiers for user friendliness?
         """
 
         logger.info("Manager initializing")
@@ -150,8 +160,8 @@ class Manager:
                 raise Exception("No viable address found")
             else:
                 logger.info("Connection to Interchange successful on {}".format(ix_address))
-                task_q_url = "tcp://{}:{}".format(ix_address, task_port)
-                result_q_url = "tcp://{}:{}".format(ix_address, result_port)
+                task_q_url = tcp_url(ix_address, task_port)
+                result_q_url = tcp_url(ix_address, result_port)
                 logger.info("Task url : {}".format(task_q_url))
                 logger.info("Result url : {}".format(result_q_url))
         except Exception:
@@ -176,6 +186,7 @@ class Manager:
 
         self.uid = uid
         self.block_id = block_id
+        self.start_time = time.time()
 
         self.enable_mpi_mode = enable_mpi_mode
         self.mpi_launcher = mpi_launcher
@@ -227,6 +238,14 @@ class Manager:
         self.heartbeat_period = heartbeat_period
         self.heartbeat_threshold = heartbeat_threshold
         self.poll_period = poll_period
+
+        self.drain_time: float
+        if drain_period:
+            self.drain_time = self._start_time + drain_period
+            logger.info(f"Will request drain at {self.drain_time}")
+        else:
+            self.drain_time = float('inf')
+
         self.cpu_affinity = cpu_affinity
 
         # Define accelerator available, adjust worker count accordingly
@@ -247,6 +266,7 @@ class Manager:
                'worker_count': self.worker_count,
                'uid': self.uid,
                'block_id': self.block_id,
+               'start_time': self.start_time,
                'prefetch_capacity': self.prefetch_capacity,
                'max_capacity': self.worker_count + self.prefetch_capacity,
                'os': platform.system(),
@@ -262,10 +282,19 @@ class Manager:
         """ Send heartbeat to the incoming task queue
         """
         msg = {'type': 'heartbeat'}
+        # don't need to dumps and encode this every time - could do as a global on import?
         b_msg = json.dumps(msg).encode('utf-8')
         self.task_incoming.send(b_msg)
         logger.debug("Sent heartbeat")
 
+    def drain_to_incoming(self):
+        """ Send heartbeat to the incoming task queue
+        """
+        msg = {'type': 'drain'}
+        b_msg = json.dumps(msg).encode('utf-8')
+        self.task_incoming.send(b_msg)
+        logger.debug("Sent drain")
+
     @wrap_with_logs
     def pull_tasks(self, kill_event):
         """ Pull tasks from the incoming tasks zmq pipe onto the internal
@@ -298,6 +327,7 @@ class Manager:
             # time here are correctly copy-pasted from the relevant if
             # statements.
             next_interesting_event_time = min(last_beat + self.heartbeat_period,
+                                              self.drain_time,
                                               last_interchange_contact + self.heartbeat_threshold)
             try:
                 pending_task_count = self.pending_task_queue.qsize()
@@ -312,6 +342,17 @@ class Manager:
                 self.heartbeat_to_incoming()
                 last_beat = time.time()
 
+            if time.time() > self.drain_time:
+                logger.info("Requesting drain")
+                self.drain_to_incoming()
+                # This will start the pool draining...
+                # Drained exit behaviour does not happen here. It will be
+                # driven by the interchange sending a DRAINED_CODE message.
+
+                # now set drain time to the far future so we don't send a drain
+                # message every iteration.
+                self.drain_time = float('inf')
+
             poll_duration_s = max(0, next_interesting_event_time - time.time())
             socks = dict(poller.poll(timeout=poll_duration_s * 1000))
 
@@ -322,10 +363,14 @@ class Manager:
 
                 if tasks == HEARTBEAT_CODE:
                     logger.debug("Got heartbeat from interchange")
-
+                elif tasks == DRAINED_CODE:
+                    logger.info("Got fully drained message from interchange - setting kill flag")
+                    kill_event.set()
                 else:
                     task_recv_counter += len(tasks)
-                    logger.debug("Got executor tasks: {}, cumulative count of tasks: {}".format([t['task_id'] for t in tasks], task_recv_counter))
+                    logger.debug("Got executor tasks: {}, cumulative count of tasks: {}".format(
+                        [t['task_id'] for t in tasks], task_recv_counter
+                    ))
 
                     for task in tasks:
                         self.task_scheduler.put_task(task)
@@ -490,9 +535,8 @@ class Manager:
         self._worker_watchdog_thread.start()
         self._monitoring_handler_thread.start()
 
-        logger.info("Loop start")
+        logger.info("Manager threads started")
 
-        # TODO : Add mechanism in this loop to stop the worker pool
         # This might need a multiprocessing event to signal back.
         self._kill_event.wait()
         logger.critical("Received kill event, terminating worker processes")
@@ -548,45 +592,13 @@ def update_resource_spec_env_vars(mpi_launcher: str, resource_spec: Dict, node_i
         os.environ[key] = prefix_table[key]
 
 
-def execute_task(bufs, mpi_launcher: Optional[str] = None):
-    """Deserialize the buffer and execute the task.
-
-    Returns the result or throws exception.
-    """
-    user_ns = locals()
-    user_ns.update({'__builtins__': __builtins__})
-
-    f, args, kwargs, resource_spec = unpack_res_spec_apply_message(bufs, user_ns, copy=False)
-
-    for varname in resource_spec:
-        envname = "PARSL_" + str(varname).upper()
-        os.environ[envname] = str(resource_spec[varname])
-
-    if resource_spec.get("MPI_NODELIST"):
-        worker_id = os.environ['PARSL_WORKER_RANK']
-        nodes_for_task = resource_spec["MPI_NODELIST"].split(',')
-        logger.info(f"Launching task on provisioned nodes: {nodes_for_task}")
-        assert mpi_launcher
-        update_resource_spec_env_vars(mpi_launcher,
-                                      resource_spec=resource_spec,
-                                      node_info=nodes_for_task)
-    # We might need to look into callability of the function from itself
-    # since we change it's name in the new namespace
-    prefix = "parsl_"
-    fname = prefix + "f"
-    argname = prefix + "args"
-    kwargname = prefix + "kwargs"
-    resultname = prefix + "result"
-
-    user_ns.update({fname: f,
-                    argname: args,
-                    kwargname: kwargs,
-                    resultname: resultname})
-
-    code = "{0} = {1}(*{2}, **{3})".format(resultname, fname,
-                                           argname, kwargname)
-    exec(code, user_ns, user_ns)
-    return user_ns.get(resultname)
+def _init_mpi_env(mpi_launcher: str, resource_spec: Dict):
+    node_list = resource_spec.get("MPI_NODELIST")
+    if node_list is None:
+        return
+    nodes_for_task = node_list.split(',')
+    logger.info(f"Launching task on provisioned nodes: {nodes_for_task}")
+    update_resource_spec_env_vars(mpi_launcher=mpi_launcher, resource_spec=resource_spec, node_info=nodes_for_task)
 
 
 @wrap_with_logs(target="worker_log")
@@ -608,14 +620,6 @@ def worker(
     debug: bool,
     mpi_launcher: str,
 ):
-    """
-
-    Put request token into queue
-    Get task from task_queue
-    Pop request from queue
-    Put result into result_queue
-    """
-
     # override the global logger inherited from the __main__ process (which
     # usually logs to manager.log) with one specific to this worker.
     global logger
@@ -640,7 +644,8 @@ def worker(
     # If desired, set process affinity
     if cpu_affinity != "none":
         # Count the number of cores per worker
-        avail_cores = sorted(os.sched_getaffinity(0))  # Get the available threads
+        # OSX does not implement os.sched_getaffinity
+        avail_cores = sorted(os.sched_getaffinity(0))  # type: ignore[attr-defined, unused-ignore]
         cores_per_worker = len(avail_cores) // pool_size
         assert cores_per_worker > 0, "Affinity does not work if there are more workers than cores"
 
@@ -680,12 +685,39 @@ def worker(
         os.environ["KMP_AFFINITY"] = f"explicit,proclist=[{proc_list}]"  # For Intel OpenMP
 
         # Set the affinity for this worker
-        os.sched_setaffinity(0, my_cores)
+        # OSX does not implement os.sched_setaffinity so type checking
+        # is ignored here in two ways:
+        # On a platform without sched_setaffinity, that attribute will not
+        # be defined, so ignore[attr-defined] will tell mypy to ignore this
+        # incorrect-for-OS X attribute access.
+        # On a platform with sched_setaffinity, that type: ignore message
+        # will be redundant, and ignore[unused-ignore] tells mypy to ignore
+        # that this ignore is unneeded.
+        os.sched_setaffinity(0, my_cores)  # type: ignore[attr-defined, unused-ignore]
         logger.info("Set worker CPU affinity to {}".format(my_cores))
 
     # If desired, pin to accelerator
     if accelerator is not None:
-        os.environ["CUDA_VISIBLE_DEVICES"] = accelerator
+
+        # If CUDA devices, find total number of devices to allow for MPS
+        # See: https://developer.nvidia.com/system-management-interface
+        nvidia_smi_cmd = "nvidia-smi -L > /dev/null && nvidia-smi -L | wc -l"
+        nvidia_smi_ret = subprocess.run(nvidia_smi_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        if nvidia_smi_ret.returncode == 0:
+            num_cuda_devices = int(nvidia_smi_ret.stdout.split()[0])
+        else:
+            num_cuda_devices = None
+
+        try:
+            if num_cuda_devices is not None:
+                procs_per_cuda_device = pool_size // num_cuda_devices
+                partitioned_accelerator = str(int(accelerator) // procs_per_cuda_device)  # multiple workers will share a GPU
+                os.environ["CUDA_VISIBLE_DEVICES"] = partitioned_accelerator
+                logger.info(f'Pinned worker to partitioned cuda device: {partitioned_accelerator}')
+            else:
+                os.environ["CUDA_VISIBLE_DEVICES"] = accelerator
+        except (TypeError, ValueError, ZeroDivisionError):
+            os.environ["CUDA_VISIBLE_DEVICES"] = accelerator
         os.environ["ROCR_VISIBLE_DEVICES"] = accelerator
         os.environ["ZE_AFFINITY_MASK"] = accelerator
         os.environ["ZE_ENABLE_PCI_ID_DEVICE_ORDER"] = '1'
@@ -724,8 +756,10 @@ def worker(
             ready_worker_count.value -= 1
         worker_enqueued = False
 
+        _init_mpi_env(mpi_launcher=mpi_launcher, resource_spec=req["resource_spec"])
+
         try:
-            result = execute_task(req['buffer'], mpi_launcher=mpi_launcher)
+            result = execute_task(req['buffer'])
             serialized_result = serialize(result, buffer_threshold=1000000)
         except Exception as e:
             logger.info('Caught an exception: {}'.format(e))
@@ -804,6 +838,8 @@ if __name__ == "__main__":
                         help="Heartbeat period in seconds. Uses manager default unless set")
     parser.add_argument("--hb_threshold", default=120,
                         help="Heartbeat threshold in seconds. Uses manager default unless set")
+    parser.add_argument("--drain_period", default=None,
+                        help="Drain this pool after specified number of seconds. By default, does not drain.")
     parser.add_argument("--address_probe_timeout", default=30,
                         help="Timeout to probe for viable address to interchange. Default: 30s")
     parser.add_argument("--poll", default=10,
@@ -824,7 +860,7 @@ if __name__ == "__main__":
                         required=True,
                         help="Whether/how workers should control CPU affinity.")
     parser.add_argument("--available-accelerators", type=str, nargs="*",
-                        help="Names of available accelerators")
+                        help="Names of available accelerators, if not given assumed to be zero accelerators available", default=[])
     parser.add_argument("--enable_mpi_mode", action='store_true',
                         help="Enable MPI mode")
     parser.add_argument("--mpi-launcher", type=str, choices=VALID_LAUNCHERS,
@@ -856,6 +892,7 @@ if __name__ == "__main__":
         logger.info("Prefetch capacity: {}".format(args.prefetch_capacity))
         logger.info("Heartbeat threshold: {}".format(args.hb_threshold))
         logger.info("Heartbeat period: {}".format(args.hb_period))
+        logger.info("Drain period: {}".format(args.drain_period))
         logger.info("CPU affinity: {}".format(args.cpu_affinity))
         logger.info("Accelerators: {}".format(" ".join(args.available_accelerators)))
         logger.info("enable_mpi_mode: {}".format(args.enable_mpi_mode))
@@ -876,6 +913,7 @@ if __name__ == "__main__":
                           prefetch_capacity=int(args.prefetch_capacity),
                           heartbeat_threshold=int(args.hb_threshold),
                           heartbeat_period=int(args.hb_period),
+                          drain_period=None if args.drain_period == "None" else int(args.drain_period),
                           poll_period=int(args.poll),
                           cpu_affinity=args.cpu_affinity,
                           enable_mpi_mode=args.enable_mpi_mode,

parsl-2025.1.13.dist-info/METADATA

@@ -0,0 +1,96 @@
+Metadata-Version: 2.1
+Name: parsl
+Version: 2025.1.13
+Summary: Simple data dependent workflows in Python
+Home-page: https://github.com/Parsl/parsl
+Download-URL: https://github.com/Parsl/parsl/archive/2025.01.13.tar.gz
+Author: The Parsl Team
+Author-email: parsl@googlegroups.com
+License: Apache 2.0
+Keywords: Workflows,Scientific computing
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.9.0
+License-File: LICENSE
+Requires-Dist: pyzmq>=17.1.2
+Requires-Dist: typeguard!=3.*,<5,>=2.10
+Requires-Dist: typing-extensions<5,>=4.6
+Requires-Dist: globus-sdk
+Requires-Dist: dill
+Requires-Dist: tblib
+Requires-Dist: requests
+Requires-Dist: psutil>=5.5.1
+Requires-Dist: setproctitle
+Requires-Dist: filelock<4,>=3.13
+Provides-Extra: all
+Requires-Dist: sqlalchemy<2.1,>=2; extra == "all"
+Requires-Dist: pydot>=1.4.2; extra == "all"
+Requires-Dist: networkx<3.3,>=3.2; extra == "all"
+Requires-Dist: Flask>=1.0.2; extra == "all"
+Requires-Dist: flask-sqlalchemy; extra == "all"
+Requires-Dist: pandas<3,>=2.2; extra == "all"
+Requires-Dist: plotly; extra == "all"
+Requires-Dist: python-daemon; extra == "all"
+Requires-Dist: boto3; extra == "all"
+Requires-Dist: kubernetes; extra == "all"
+Requires-Dist: ipython<=8.6.0; extra == "all"
+Requires-Dist: nbsphinx; extra == "all"
+Requires-Dist: sphinx<7.2,>=7.1; extra == "all"
+Requires-Dist: sphinx-rtd-theme; extra == "all"
+Requires-Dist: google-auth; extra == "all"
+Requires-Dist: google-api-python-client; extra == "all"
+Requires-Dist: python-gssapi; extra == "all"
+Requires-Dist: azure<=4; extra == "all"
+Requires-Dist: msrestazure; extra == "all"
+Requires-Dist: work-queue; extra == "all"
+Requires-Dist: pyyaml; extra == "all"
+Requires-Dist: cffi; extra == "all"
+Requires-Dist: jsonschema; extra == "all"
+Requires-Dist: proxystore; extra == "all"
+Requires-Dist: radical.pilot==1.90; extra == "all"
+Requires-Dist: radical.utils==1.90; extra == "all"
+Provides-Extra: aws
+Requires-Dist: boto3; extra == "aws"
+Provides-Extra: azure
+Requires-Dist: azure<=4; extra == "azure"
+Requires-Dist: msrestazure; extra == "azure"
+Provides-Extra: docs
+Requires-Dist: ipython<=8.6.0; extra == "docs"
+Requires-Dist: nbsphinx; extra == "docs"
+Requires-Dist: sphinx<7.2,>=7.1; extra == "docs"
+Requires-Dist: sphinx-rtd-theme; extra == "docs"
+Provides-Extra: flux
+Requires-Dist: pyyaml; extra == "flux"
+Requires-Dist: cffi; extra == "flux"
+Requires-Dist: jsonschema; extra == "flux"
+Provides-Extra: google_cloud
+Requires-Dist: google-auth; extra == "google-cloud"
+Requires-Dist: google-api-python-client; extra == "google-cloud"
+Provides-Extra: gssapi
+Requires-Dist: python-gssapi; extra == "gssapi"
+Provides-Extra: kubernetes
+Requires-Dist: kubernetes; extra == "kubernetes"
+Provides-Extra: monitoring
+Requires-Dist: sqlalchemy<2.1,>=2; extra == "monitoring"
+Provides-Extra: proxystore
+Requires-Dist: proxystore; extra == "proxystore"
+Provides-Extra: radical-pilot
+Requires-Dist: radical.pilot==1.90; extra == "radical-pilot"
+Requires-Dist: radical.utils==1.90; extra == "radical-pilot"
+Provides-Extra: visualization
+Requires-Dist: pydot>=1.4.2; extra == "visualization"
+Requires-Dist: networkx<3.3,>=3.2; extra == "visualization"
+Requires-Dist: Flask>=1.0.2; extra == "visualization"
+Requires-Dist: flask-sqlalchemy; extra == "visualization"
+Requires-Dist: pandas<3,>=2.2; extra == "visualization"
+Requires-Dist: plotly; extra == "visualization"
+Requires-Dist: python-daemon; extra == "visualization"
+Provides-Extra: workqueue
+Requires-Dist: work-queue; extra == "workqueue"
+
+Simple parallel workflows system for Python