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

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

@@ -1,39 +1,44 @@
 #!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
@@ -155,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:
@@ -181,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
@@ -260,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(),
@@ -335,14 +342,17 @@ class Manager:
                 self.heartbeat_to_incoming()
                 last_beat = time.time()
 
-            if self.drain_time and time.time() > self.drain_time:
+            if time.time() > self.drain_time:
                 logger.info("Requesting drain")
                 self.drain_to_incoming()
-                self.drain_time = None
                 # 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))
 
@@ -354,11 +364,13 @@ class Manager:
                 if tasks == HEARTBEAT_CODE:
                     logger.debug("Got heartbeat from interchange")
                 elif tasks == DRAINED_CODE:
-                    logger.info("Got fulled drained message from interchange - setting kill flag")
+                    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)
@@ -580,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")
@@ -640,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
@@ -672,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"
 
@@ -712,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'
@@ -756,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))

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