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

parsl/executors/high_throughput/process_worker_pool.py

@@ -1,41 +1,47 @@
 #!/usr/bin/env python3
 
 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/executors/high_throughput/zmq_pipes.py

@@ -1,10 +1,19 @@
 #!/usr/bin/env python3
 
-import zmq
 import logging
 import threading
+import time
+from typing import Optional
+
+import zmq
 
 from parsl import curvezmq
+from parsl.addresses import tcp_url
+from parsl.errors import InternalConsistencyError
+from parsl.executors.high_throughput.errors import (
+    CommandClientBadError,
+    CommandClientTimeoutError,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -12,25 +21,29 @@ logger = logging.getLogger(__name__)
 class CommandClient:
     """ CommandClient
     """
-    def __init__(self, zmq_context: curvezmq.ClientContext, ip_address, port_range):
+    def __init__(self, ip_address, port_range, cert_dir: Optional[str] = None):
         """
         Parameters
         ----------
 
-        zmq_context: curvezmq.ClientContext
-            CurveZMQ client context used to create secure sockets
         ip_address: str
            IP address of the client (where Parsl runs)
+
         port_range: tuple(int, int)
            Port range for the comms between client and interchange
 
+        cert_dir: str | None
+            Path to the certificate directory. Setting this to None will disable encryption.
+            default: None
+
         """
-        self.zmq_context = zmq_context
+        self.zmq_context = curvezmq.ClientContext(cert_dir)
         self.ip_address = ip_address
         self.port_range = port_range
         self.port = None
         self.create_socket_and_bind()
         self._lock = threading.Lock()
+        self.ok = True
 
     def create_socket_and_bind(self):
         """ Creates socket and binds to a port.
@@ -40,13 +53,13 @@ class CommandClient:
         self.zmq_socket = self.zmq_context.socket(zmq.REQ)
         self.zmq_socket.setsockopt(zmq.LINGER, 0)
         if self.port is None:
-            self.port = self.zmq_socket.bind_to_random_port("tcp://{}".format(self.ip_address),
+            self.port = self.zmq_socket.bind_to_random_port(tcp_url(self.ip_address),
                                                             min_port=self.port_range[0],
                                                             max_port=self.port_range[1])
         else:
-            self.zmq_socket.bind("tcp://{}:{}".format(self.ip_address, self.port))
+            self.zmq_socket.bind(tcp_url(self.ip_address, self.port))
 
-    def run(self, message, max_retries=3):
+    def run(self, message, max_retries=3, timeout_s=None):
         """ This function needs to be fast at the same time aware of the possibility of
         ZMQ pipes overflowing.
 
@@ -54,13 +67,43 @@ class CommandClient:
         in ZMQ sockets reaching a broken state once there are ~10k tasks in flight.
         This issue can be magnified if each the serialized buffer itself is larger.
         """
+        if not self.ok:
+            raise CommandClientBadError()
+
+        start_time_s = time.monotonic()
+
         reply = '__PARSL_ZMQ_PIPES_MAGIC__'
         with self._lock:
             for _ in range(max_retries):
                 try:
                     logger.debug("Sending command client command")
+
+                    if timeout_s is not None:
+                        remaining_time_s = start_time_s + timeout_s - time.monotonic()
+                        poll_result = self.zmq_socket.poll(timeout=remaining_time_s * 1000, flags=zmq.POLLOUT)
+                        if poll_result == zmq.POLLOUT:
+                            pass  # this is OK, so continue
+                        elif poll_result == 0:
+                            raise CommandClientTimeoutError("Waiting for command channel to be ready for a command")
+                        else:
+                            raise InternalConsistencyError(f"ZMQ poll returned unexpected value: {poll_result}")
+
                     self.zmq_socket.send_pyobj(message, copy=True)
-                    logger.debug("Waiting for command client response")
+
+                    if timeout_s is not None:
+                        logger.debug("Polling for command client response or timeout")
+                        remaining_time_s = start_time_s + timeout_s - time.monotonic()
+                        poll_result = self.zmq_socket.poll(timeout=remaining_time_s * 1000, flags=zmq.POLLIN)
+                        if poll_result == zmq.POLLIN:
+                            pass  # this is OK, so continue
+                        elif poll_result == 0:
+                            logger.error("Command timed-out - command client is now bad forever")
+                            self.ok = False
+                            raise CommandClientTimeoutError("Waiting for a reply from command channel")
+                        else:
+                            raise InternalConsistencyError(f"ZMQ poll returned unexpected value: {poll_result}")
+
+                    logger.debug("Receiving command client response")
                     reply = self.zmq_socket.recv_pyobj()
                     logger.debug("Received command client response")
                 except zmq.ZMQError:
@@ -85,23 +128,26 @@ class CommandClient:
 class TasksOutgoing:
     """ Outgoing task queue from the executor to the Interchange
     """
-    def __init__(self, zmq_context: curvezmq.ClientContext, ip_address, port_range):
+    def __init__(self, ip_address, port_range, cert_dir: Optional[str] = None):
         """
         Parameters
         ----------
 
-        zmq_context: curvezmq.ClientContext
-            CurveZMQ client context used to create secure sockets
         ip_address: str
            IP address of the client (where Parsl runs)
+
         port_range: tuple(int, int)
            Port range for the comms between client and interchange
 
+        cert_dir: str | None
+            Path to the certificate directory. Setting this to None will disable encryption.
+            default: None
+
         """
-        self.zmq_context = zmq_context
+        self.zmq_context = curvezmq.ClientContext(cert_dir)
         self.zmq_socket = self.zmq_context.socket(zmq.DEALER)
         self.zmq_socket.set_hwm(0)
-        self.port = self.zmq_socket.bind_to_random_port("tcp://{}".format(ip_address),
+        self.port = self.zmq_socket.bind_to_random_port(tcp_url(ip_address),
                                                         min_port=port_range[0],
                                                         max_port=port_range[1])
         self.poller = zmq.Poller()
@@ -138,31 +184,43 @@ class ResultsIncoming:
     """ Incoming results queue from the Interchange to the executor
     """
 
-    def __init__(self, zmq_context: curvezmq.ClientContext, ip_address, port_range):
+    def __init__(self, ip_address, port_range, cert_dir: Optional[str] = None):
         """
         Parameters
         ----------
 
-        zmq_context: curvezmq.ClientContext
-            CurveZMQ client context used to create secure sockets
         ip_address: str
            IP address of the client (where Parsl runs)
+
         port_range: tuple(int, int)
            Port range for the comms between client and interchange
 
+        cert_dir: str | None
+            Path to the certificate directory. Setting this to None will disable encryption.
+            default: None
+
         """
-        self.zmq_context = zmq_context
+        self.zmq_context = curvezmq.ClientContext(cert_dir)
         self.results_receiver = self.zmq_context.socket(zmq.DEALER)
         self.results_receiver.set_hwm(0)
-        self.port = self.results_receiver.bind_to_random_port("tcp://{}".format(ip_address),
+        self.port = self.results_receiver.bind_to_random_port(tcp_url(ip_address),
                                                               min_port=port_range[0],
                                                               max_port=port_range[1])
+        self.poller = zmq.Poller()
+        self.poller.register(self.results_receiver, zmq.POLLIN)
 
-    def get(self):
+    def get(self, timeout_ms=None):
+        """Get a message from the queue, returning None if timeout expires
+        without a message. timeout is measured in milliseconds.
+        """
         logger.debug("Waiting for ResultsIncoming message")
-        m = self.results_receiver.recv_multipart()
-        logger.debug("Received ResultsIncoming message")
-        return m
+        socks = dict(self.poller.poll(timeout=timeout_ms))
+        if self.results_receiver in socks and socks[self.results_receiver] == zmq.POLLIN:
+            m = self.results_receiver.recv_multipart()
+            logger.debug("Received ResultsIncoming message")
+            return m
+        else:
+            return None
 
     def close(self):
         self.results_receiver.close()