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

parsl/executors/high_throughput/manager_selector.py

@@ -0,0 +1,55 @@
+import random
+from abc import ABCMeta, abstractmethod
+from typing import Dict, List, Set
+
+from parsl.executors.high_throughput.manager_record import ManagerRecord
+
+
+class ManagerSelector(metaclass=ABCMeta):
+
+    @abstractmethod
+    def sort_managers(self, ready_managers: Dict[bytes, ManagerRecord], manager_list: Set[bytes]) -> List[bytes]:
+        """ Sort a given list of managers.
+
+        Any operations pertaining to the sorting and rearrangement of the
+        interesting_managers Set should be performed here.
+        """
+        pass
+
+
+class RandomManagerSelector(ManagerSelector):
+
+    """Returns a shuffled list of interesting_managers
+
+    By default this strategy is used by the interchange. Works well
+    in distributing workloads equally across all availble compute
+    resources. The random workload strategy is not effective in
+    conjunction with elastic scaling behavior as the even task
+    distribution does not allow the scaling down of blocks, leading
+    to wasted resource consumption.
+    """
+
+    def sort_managers(self, ready_managers: Dict[bytes, ManagerRecord], manager_list: Set[bytes]) -> List[bytes]:
+        c_manager_list = list(manager_list)
+        random.shuffle(c_manager_list)
+        return c_manager_list
+
+
+class BlockIdManagerSelector(ManagerSelector):
+
+    """Returns an interesting_managers list sorted by block ID
+
+    Observations:
+    1. BlockID manager selector helps with workloads that see a varying
+    amount of tasks over time. New blocks are prioritized with the
+    blockID manager selector, when used with 'htex_auto_scaling', results
+    in compute cost savings.
+
+    2. Doesn't really work with bag-of-tasks workloads. When all the tasks
+    are put into the queue upfront, all blocks operate at near full
+    utilization for the majority of the workload, which task goes where
+    doesn't really matter.
+    """
+
+    def sort_managers(self, ready_managers: Dict[bytes, ManagerRecord], manager_list: Set[bytes]) -> List[bytes]:
+        return sorted(manager_list, key=lambda x: (ready_managers[x]['block_id'] is not None, ready_managers[x]['block_id']))

parsl/executors/high_throughput/monitoring_info.py

@@ -3,6 +3,7 @@
 # then be acquired by any other code running in
 # a worker context - specifically the monitoring
 # wrapper code.
-from typing import Optional
 from queue import Queue
+from typing import Optional
+
 result_queue: Optional[Queue] = None

parsl/executors/high_throughput/mpi_executor.py

@@ -0,0 +1,113 @@
+"""A simplified interface for HTEx when running in MPI mode"""
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import typeguard
+
+from parsl.data_provider.staging import Staging
+from parsl.executors.high_throughput.executor import (
+    GENERAL_HTEX_PARAM_DOCS,
+    HighThroughputExecutor,
+)
+from parsl.executors.high_throughput.mpi_prefix_composer import (
+    VALID_LAUNCHERS,
+    validate_resource_spec,
+)
+from parsl.executors.status_handling import BlockProviderExecutor
+from parsl.jobs.states import JobStatus
+from parsl.launchers import SimpleLauncher
+from parsl.providers import LocalProvider
+from parsl.providers.base import ExecutionProvider
+
+
+class MPIExecutor(HighThroughputExecutor):
+    __doc__ = f"""A version of :class:`~parsl.HighThroughputExecutor` tuned for executing multi-node (e.g., MPI) tasks.
+
+    The Provider _must_ use the :class:`~parsl.launchers.SimpleLauncher`,
+    which places a single pool of workers on the first node of a block.
+    Each worker can then make system calls which use an MPI launcher (e.g., ``mpirun``, ``srun``)
+    to spawn multi-node tasks.
+
+    Specify the maximum number of multi-node tasks to run at once using ``max_workers_per_block``.
+    The value should be less than or equal to the ``nodes_per_block`` in the Provider.
+
+    Parameters
+    ----------
+    max_workers_per_block: int
+        Maximum number of MPI applications to run at once per block
+
+    mpi_launcher: str
+        Select one from the list of supported MPI launchers:
+        ("srun", "aprun", "mpiexec").
+        default: "mpiexec"
+
+    {GENERAL_HTEX_PARAM_DOCS}
+    """
+
+    @typeguard.typechecked
+    def __init__(self,
+                 label: str = 'MPIExecutor',
+                 provider: ExecutionProvider = LocalProvider(),
+                 launch_cmd: Optional[str] = None,
+                 interchange_launch_cmd: Optional[str] = None,
+                 address: Optional[str] = None,
+                 loopback_address: str = "127.0.0.1",
+                 worker_ports: Optional[Tuple[int, int]] = None,
+                 worker_port_range: Optional[Tuple[int, int]] = (54000, 55000),
+                 interchange_port_range: Optional[Tuple[int, int]] = (55000, 56000),
+                 storage_access: Optional[List[Staging]] = None,
+                 working_dir: Optional[str] = None,
+                 worker_debug: bool = False,
+                 max_workers_per_block: int = 1,
+                 prefetch_capacity: int = 0,
+                 heartbeat_threshold: int = 120,
+                 heartbeat_period: int = 30,
+                 drain_period: Optional[int] = None,
+                 poll_period: int = 10,
+                 address_probe_timeout: Optional[int] = None,
+                 worker_logdir_root: Optional[str] = None,
+                 mpi_launcher: str = "mpiexec",
+                 block_error_handler: Union[bool, Callable[[BlockProviderExecutor, Dict[str, JobStatus]], None]] = True,
+                 encrypted: bool = False):
+        super().__init__(
+            # Hard-coded settings
+            cores_per_worker=1e-9,  # Ensures there will be at least an absurd number of workers
+            max_workers_per_node=max_workers_per_block,
+
+            # Everything else
+            label=label,
+            provider=provider,
+            launch_cmd=launch_cmd,
+            interchange_launch_cmd=interchange_launch_cmd,
+            address=address,
+            loopback_address=loopback_address,
+            worker_ports=worker_ports,
+            worker_port_range=worker_port_range,
+            interchange_port_range=interchange_port_range,
+            storage_access=storage_access,
+            working_dir=working_dir,
+            worker_debug=worker_debug,
+            prefetch_capacity=prefetch_capacity,
+            heartbeat_threshold=heartbeat_threshold,
+            heartbeat_period=heartbeat_period,
+            drain_period=drain_period,
+            poll_period=poll_period,
+            address_probe_timeout=address_probe_timeout,
+            worker_logdir_root=worker_logdir_root,
+            block_error_handler=block_error_handler,
+            encrypted=encrypted
+        )
+        self.enable_mpi_mode = True
+        self.mpi_launcher = mpi_launcher
+
+        self.max_workers_per_block = max_workers_per_block
+
+        if not isinstance(self.provider.launcher, SimpleLauncher):
+            raise TypeError("mpi_mode requires the provider to be configured to use a SimpleLauncher")
+
+        if mpi_launcher not in VALID_LAUNCHERS:
+            raise ValueError(f"mpi_launcher set to:{mpi_launcher} must be set to one of {VALID_LAUNCHERS}")
+
+        self.mpi_launcher = mpi_launcher
+
+    def validate_resource_spec(self, resource_specification: dict):
+        return validate_resource_spec(resource_specification)

parsl/executors/high_throughput/mpi_prefix_composer.py

@@ -1,5 +1,7 @@
 import logging
-from typing import Dict, List, Tuple, Set
+from typing import Dict, List, Tuple
+
+from parsl.executors.errors import InvalidResourceSpecification
 
 logger = logging.getLogger(__name__)
 
@@ -8,16 +10,6 @@ VALID_LAUNCHERS = ('srun',
                    'mpiexec')
 
 
-class InvalidResourceSpecification(Exception):
-    """Exception raised when Invalid keys are supplied via resource specification"""
-
-    def __init__(self, invalid_keys: Set[str]):
-        self.invalid_keys = invalid_keys
-
-    def __str__(self):
-        return f"Invalid resource specification options supplied: {self.invalid_keys}"
-
-
 def validate_resource_spec(resource_spec: Dict[str, str]):
     """Basic validation of keys in the resource_spec
 
@@ -25,6 +17,13 @@ def validate_resource_spec(resource_spec: Dict[str, str]):
         is invalid (e.g, contains invalid keys)
     """
     user_keys = set(resource_spec.keys())
+
+    # empty resource_spec when mpi_mode is set causes parsl to hang
+    # ref issue #3427
+    if len(user_keys) == 0:
+        raise InvalidResourceSpecification(user_keys,
+                                           'MPI mode requires optional parsl_resource_specification keyword argument to be configured')
+
     legal_keys = set(("ranks_per_node",
                       "num_nodes",
                       "num_ranks",

parsl/executors/high_throughput/mpi_resource_management.py

@@ -8,8 +8,7 @@ from enum import Enum
 from typing import Dict, List
 
 from parsl.multiprocessing import SpawnContext
-from parsl.serialize import (pack_res_spec_apply_message,
-                             unpack_res_spec_apply_message)
+from parsl.serialize import pack_res_spec_apply_message, unpack_res_spec_apply_message
 
 logger = logging.getLogger(__name__)
 
@@ -18,7 +17,6 @@ class Scheduler(Enum):
     Unknown = 0
     Slurm = 1
     PBS = 2
-    Cobalt = 3
 
 
 def get_slurm_hosts_list() -> List[str]:
@@ -38,13 +36,6 @@ def get_pbs_hosts_list() -> List[str]:
         return [line.strip() for line in f.readlines()]
 
 
-def get_cobalt_hosts_list() -> List[str]:
-    """Get list of COBALT hosts from envvar: COBALT_NODEFILE"""
-    nodefile_name = os.environ["COBALT_NODEFILE"]
-    with open(nodefile_name) as f:
-        return [line.strip() for line in f.readlines()]
-
-
 def get_nodes_in_batchjob(scheduler: Scheduler) -> List[str]:
     """Get nodelist from all supported schedulers"""
     nodelist = []
@@ -52,8 +43,6 @@ def get_nodes_in_batchjob(scheduler: Scheduler) -> List[str]:
         nodelist = get_slurm_hosts_list()
     elif scheduler == Scheduler.PBS:
         nodelist = get_pbs_hosts_list()
-    elif scheduler == Scheduler.Cobalt:
-        nodelist = get_cobalt_hosts_list()
     else:
         raise RuntimeError(f"mpi_mode does not support scheduler:{scheduler}")
     return nodelist
@@ -65,8 +54,6 @@ def identify_scheduler() -> Scheduler:
         return Scheduler.Slurm
     elif os.environ.get("PBS_NODEFILE"):
         return Scheduler.PBS
-    elif os.environ.get("COBALT_NODEFILE"):
-        return Scheduler.Cobalt
     else:
         return Scheduler.Unknown
 
@@ -173,9 +160,7 @@ class MPITaskScheduler(TaskScheduler):
         """Schedule task if resources are available otherwise backlog the task"""
         user_ns = locals()
         user_ns.update({"__builtins__": __builtins__})
-        _f, _args, _kwargs, resource_spec = unpack_res_spec_apply_message(
-            task_package["buffer"], user_ns, copy=False
-        )
+        _f, _args, _kwargs, resource_spec = unpack_res_spec_apply_message(task_package["buffer"])
 
         nodes_needed = resource_spec.get("num_nodes")
         if nodes_needed:
@@ -190,6 +175,7 @@ class MPITaskScheduler(TaskScheduler):
                 self._map_tasks_to_nodes[task_package["task_id"]] = allocated_nodes
                 buffer = pack_res_spec_apply_message(_f, _args, _kwargs, resource_spec)
                 task_package["buffer"] = buffer
+                task_package["resource_spec"] = resource_spec
 
         self.pending_task_q.put(task_package)
 
@@ -208,8 +194,11 @@ class MPITaskScheduler(TaskScheduler):
         """Return result and relinquish provisioned nodes"""
         result_pkl = self.pending_result_q.get(block, timeout=timeout)
         result_dict = pickle.loads(result_pkl)
+        # TODO (wardlt): If the task did not request nodes, it won't be in `self._map_tasks_to_nodes`.
+        #  Causes Parsl to hang. See Issue #3427
         if result_dict["type"] == "result":
             task_id = result_dict["task_id"]
+            assert task_id in self._map_tasks_to_nodes, "You are about to experience issue #3427"
             nodes_to_reallocate = self._map_tasks_to_nodes[task_id]
             self._return_nodes(nodes_to_reallocate)
             self._schedule_backlog_tasks()

parsl/executors/high_throughput/probe.py

@@ -1,11 +1,13 @@
-import zmq
 import argparse
-import uuid
-import time
 import logging
-from parsl.addresses import get_all_addresses
+import time
+import uuid
+
+import zmq
 from zmq.utils.monitor import recv_monitor_message
 
+from parsl.addresses import get_all_addresses, tcp_url
+
 logger = logging.getLogger(__name__)
 
 
@@ -30,7 +32,8 @@ def probe_addresses(addresses, task_port, timeout=120):
     for addr in addresses:
         socket = context.socket(zmq.DEALER)
         socket.setsockopt(zmq.LINGER, 0)
-        url = "tcp://{}:{}".format(addr, task_port)
+        socket.setsockopt(zmq.IPV6, True)
+        url = tcp_url(addr, task_port)
         logger.debug("Trying to connect back on {}".format(url))
         socket.connect(url)
         addr_map[addr] = {'sock': socket,
@@ -69,8 +72,7 @@ class TestWorker:
 
         address = probe_addresses(addresses, port)
         print("Viable address :", address)
-        self.task_incoming.connect("tcp://{}:{}".format(address, port))
-        print("Here")
+        self.task_incoming.connect(tcp_url(address, port))
 
     def heartbeat(self):
         """ Send heartbeat to the incoming task queue

parsl/executors/high_throughput/process_worker_pool.py

@@ -1,39 +1,44 @@
 #!/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
@@ -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))