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

parsl/executors/radical/executor.py

@@ -1,30 +1,30 @@
 """RadicalPilotExecutor builds on the RADICAL-Pilot/Parsl
 """
+import inspect
+import logging
 import os
+import queue
 import sys
+import threading as mt
 import time
-import parsl
-import queue
-import logging
-import inspect
+from concurrent.futures import Future
+from functools import partial
+from pathlib import PosixPath
+from typing import Dict, Optional
+
 import requests
 import typeguard
-import threading as mt
-
-from functools import partial
-from typing import Optional, Dict
-from pathlib import Path, PosixPath
-from concurrent.futures import Future
 
+import parsl
+from parsl.app.errors import BashExitFailure, RemoteExceptionWrapper
 from parsl.app.python import timeout
-from .rpex_resources import ResourceConfig
 from parsl.data_provider.files import File
-from parsl.utils import RepresentationMixin
-from parsl.app.errors import BashExitFailure
 from parsl.executors.base import ParslExecutor
-from parsl.app.errors import RemoteExceptionWrapper
 from parsl.serialize import deserialize, pack_res_spec_apply_message
-from parsl.serialize.errors import SerializationError, DeserializationError
+from parsl.serialize.errors import DeserializationError, SerializationError
+from parsl.utils import RepresentationMixin
+
+from .rpex_resources import CLIENT, MPI, ResourceConfig
 
 try:
     import radical.pilot as rp
@@ -59,7 +59,7 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
           ``rp.PilotManager`` and ``rp.TaskManager``.
       2. "translate": Unwrap, identify, and parse Parsl ``apps`` into ``rp.TaskDescription``.
       3. "submit": Submit Parsl apps to ``rp.TaskManager``.
-      4. "shut_down": Shut down the RADICAL-Pilot runtime and all associated components.
+      4. "shutdown": Shut down the RADICAL-Pilot runtime and all associated components.
 
     Here is a diagram
 
@@ -133,24 +133,32 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
         self.resource = resource
         self._uid = RPEX.lower()
         self.bulk_mode = bulk_mode
+        self._terminate = mt.Event()
         self.working_dir = working_dir
         self.pilot_kwargs = rpex_pilot_kwargs
         self.future_tasks: Dict[str, Future] = {}
 
         if rpex_cfg:
-            self.rpex_cfg = rpex_cfg
+            self.rpex_cfg = rpex_cfg.get_config()
         elif not rpex_cfg and 'local' in resource:
-            self.rpex_cfg = ResourceConfig()
+            self.rpex_cfg = ResourceConfig().get_config()
         else:
-            raise ValueError('Resource config file must be '
-                             'specified for a non-local execution')
+            raise ValueError('Resource config must be '
+                             'specified for a non-local resources')
 
     def task_state_cb(self, task, state):
         """
         Update the state of Parsl Future apps
         Based on RP task state callbacks.
         """
-        if not task.uid.startswith('master'):
+        # check the Master/Worker state
+        if task.mode in [rp.RAPTOR_MASTER, rp.RAPTOR_WORKER]:
+            if state == rp.FAILED:
+                exception = RuntimeError(f'{task.uid} failed with internal error: {task.stderr}')
+                self._fail_all_tasks(exception)
+
+        # check all other tasks state
+        else:
             parsl_task = self.future_tasks[task.uid]
 
             if state == rp.DONE:
@@ -186,6 +194,23 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
                     else:
                         parsl_task.set_exception('Task failed for an unknown reason')
 
+    def _fail_all_tasks(self, exception):
+        """
+        Fail all outstanding tasks with the given exception.
+
+        This method iterates through all outstanding tasks in the
+        `_future_tasks` dictionary, which have not yet completed,
+        and sets the provided exception as their result, indicating
+        a failure.
+
+        Parameters:
+        - exception: The exception to be set as the result for all
+                     outstanding tasks.
+        """
+        for fut_task in self.future_tasks.values():
+            if not fut_task.done():
+                fut_task.set_exception(exception)
+
     def start(self):
         """Create the Pilot component and pass it.
         """
@@ -202,63 +227,62 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
                    'resource': self.resource}
 
         if not self.resource or 'local' in self.resource:
-            # move the agent sandbox to the working dir mainly
-            # for debugging purposes. This will allow parsl
-            # to include the agent sandbox with the ci artifacts.
-            if os.environ.get("LOCAL_SANDBOX"):
-                pd_init['sandbox'] = self.run_dir
-                os.environ["RADICAL_LOG_LVL"] = "DEBUG"
-
-            logger.info("RPEX will be running in the local mode")
+            os.environ["RADICAL_LOG_LVL"] = "DEBUG"
+            logger.info("RPEX will be running in local mode")
 
         pd = rp.PilotDescription(pd_init)
         pd.verify()
 
-        self.rpex_cfg = self.rpex_cfg._get_cfg_file(path=self.run_dir)
-        cfg = ru.Config(cfg=ru.read_json(self.rpex_cfg))
+        # start RP's main components TMGR, PMGR and Pilot
+        self.tmgr = rp.TaskManager(session=self.session)
+        self.pmgr = rp.PilotManager(session=self.session)
+        self.pilot = self.pmgr.submit_pilots(pd)
 
-        self.master = cfg.master_descr
-        self.n_masters = cfg.n_masters
+        if not self.pilot.description.get('cores') or not self.pilot.description.get('nodes'):
+            logger.warning('no "cores/nodes" per pilot were set, using default resources')
+
+        self.tmgr.add_pilots(self.pilot)
+        self.tmgr.register_callback(self.task_state_cb)
 
-        tds = list()
-        master_path = '{0}/rpex_master.py'.format(PWD)
         worker_path = '{0}/rpex_worker.py'.format(PWD)
 
-        for i in range(self.n_masters):
-            td = rp.TaskDescription(self.master)
-            td.mode = rp.RAPTOR_MASTER
-            td.uid = ru.generate_id('master.%(item_counter)06d', ru.ID_CUSTOM,
+        self.masters = []
+
+        logger.info(f'Starting {self.rpex_cfg.n_masters} masters and {self.rpex_cfg.n_workers} workers for each master')
+
+        # create N masters
+        for _ in range(self.rpex_cfg.n_masters):
+            md = rp.TaskDescription(self.rpex_cfg.master_descr)
+            md.uid = ru.generate_id('rpex.master.%(item_counter)06d', ru.ID_CUSTOM,
                                     ns=self.session.uid)
-            td.ranks = 1
-            td.cores_per_rank = 1
-            td.arguments = [self.rpex_cfg, i]
-            td.input_staging = self._stage_files([File(master_path),
-                                                  File(worker_path),
-                                                  File(self.rpex_cfg)], mode='in')
-            tds.append(td)
 
-        self.pmgr = rp.PilotManager(session=self.session)
-        self.tmgr = rp.TaskManager(session=self.session)
+            # submit the master to the TMGR
+            master = self.tmgr.submit_raptors(md)[0]
+            self.masters.append(master)
+
+            workers = []
+            # create N workers for each master and submit them to the TMGR
+            for _ in range(self.rpex_cfg.n_workers):
+                wd = rp.TaskDescription(self.rpex_cfg.worker_descr)
+                wd.uid = ru.generate_id('rpex.worker.%(item_counter)06d', ru.ID_CUSTOM,
+                                        ns=self.session.uid)
+                wd.raptor_id = master.uid
+                wd.input_staging = self._stage_files([File(worker_path)], mode='in')
+                workers.append(wd)
 
-        # submit pilot(s)
-        pilot = self.pmgr.submit_pilots(pd)
-        if not pilot.description.get('cores'):
-            logger.warning('no "cores" per pilot was set, using default resources {0}'.format(pilot.resources))
+            self.tmgr.submit_workers(workers)
 
-        self.tmgr.submit_tasks(tds)
+        self.select_master = self._cyclic_master_selector()
 
         # prepare or use the current env for the agent/pilot side environment
-        if cfg.pilot_env_mode != 'client':
-            logger.info("creating {0} environment for the executor".format(cfg.pilot_env.name))
-            pilot.prepare_env(env_name=cfg.pilot_env.name,
-                              env_spec=cfg.pilot_env.as_dict())
+        if self.rpex_cfg.pilot_env_mode != CLIENT:
+            logger.info("creating {0} environment for the executor".format(self.rpex_cfg.pilot_env.name))
+            self.pilot.prepare_env(env_name=self.rpex_cfg.pilot_env.name,
+                                   env_spec=self.rpex_cfg.pilot_env.as_dict())
         else:
             client_env = sys.prefix
             logger.info("reusing ({0}) environment for the executor".format(client_env))
 
-        self.tmgr.add_pilots(pilot)
-        self.tmgr.register_callback(self.task_state_cb)
-
         # create a bulking thread to run the actual task submission
         # to RP in bulks
         if self.bulk_mode:
@@ -272,8 +296,21 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
             self._bulk_thread.daemon = True
             self._bulk_thread.start()
 
+            logger.info('bulk mode is on, submitting tasks in bulks')
+
         return True
 
+    def _cyclic_master_selector(self):
+        """
+        Balance tasks submission across N masters and N workers
+        """
+        current_master = 0
+        masters_uids = [m.uid for m in self.masters]
+
+        while True:
+            yield masters_uids[current_master]
+            current_master = (current_master + 1) % len(self.masters)
+
     def unwrap(self, func, args):
         """
         Unwrap a Parsl app and its args for further processing.
@@ -364,22 +401,25 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
 
                 # This is the default mode where the bash_app will be executed as
                 # as a single core process by RP. For cores > 1 the user must use
-                # above or use MPI functions if their code is Python.
+                # task.mode=rp.TASK_EXECUTABLE (above) or use MPI functions if their
+                # code is Python.
                 else:
                     task.mode = rp.TASK_PROC
-                    task.raptor_id = 'master.%06d' % (tid % self.n_masters)
+                    task.raptor_id = next(self.select_master)
                     task.executable = self._pack_and_apply_message(func, args, kwargs)
 
         elif PYTHON in task_type or not task_type:
             task.mode = rp.TASK_FUNCTION
-            task.raptor_id = 'master.%06d' % (tid % self.n_masters)
+            task.raptor_id = next(self.select_master)
             if kwargs.get('walltime'):
                 func = timeout(func, kwargs['walltime'])
 
-            # we process MPI function differently
-            if 'comm' in kwargs:
+            # Check how to serialize the function object
+            if MPI in self.rpex_cfg.worker_type.lower():
+                task.use_mpi = True
                 task.function = rp.PythonTask(func, *args, **kwargs)
             else:
+                task.use_mpi = False
                 task.function = self._pack_and_apply_message(func, args, kwargs)
 
         task.input_staging = self._stage_files(kwargs.get("inputs", []),
@@ -394,7 +434,7 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
         try:
             task.verify()
         except ru.typeddict.TDKeyError as e:
-            raise Exception(f'{e}. Please check Radical.Pilot TaskDescription documentation')
+            raise Exception(f'{e}. Please check: https://radicalpilot.readthedocs.io/en/stable/ documentation')
 
         return task
 
@@ -413,7 +453,11 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
 
     def _unpack_and_set_parsl_exception(self, parsl_task, exception):
         try:
-            s = rp.utils.deserialize_bson(exception)
+            try:
+                s = rp.utils.deserialize_bson(exception)
+            except Exception:
+                s = exception
+
             if isinstance(s, RemoteExceptionWrapper):
                 try:
                     s.reraise()
@@ -421,6 +465,8 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
                     parsl_task.set_exception(e)
             elif isinstance(s, Exception):
                 parsl_task.set_exception(s)
+            elif isinstance(s, str):
+                parsl_task.set_exception(eval(s))
             else:
                 raise ValueError("Unknown exception-like type received: {}".format(type(s)))
         except Exception as e:
@@ -440,16 +486,10 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
                 elif isinstance(k_val, PosixPath):
                     k_val = k_val.__str__()
 
-                # if the stderr/out has no path
-                # then we consider it local and
-                # we just set the path to the cwd
-                if '/' not in k_val:
-                    k_val = CWD + '/' + k_val
-
-                # finally set the stderr/out to
-                # the desired name by the user
+                # set the stderr/out to the desired
+                # name by the user
                 setattr(task, k, k_val)
-                task.sandbox = Path(k_val).parent.__str__()
+                task.sandbox = CWD
 
     def _stage_files(self, files, mode):
         """
@@ -477,7 +517,7 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
                 # this indicates that the user
                 # did not provided a specific
                 # output file and RP will stage out
-                # the task.output from pilot://task_folder
+                # the task.stdout from pilot://task_folder
                 # to the CWD or file.url
                 if '/' not in file.url:
                     f = {'source': file.filename,
@@ -493,7 +533,7 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
 
         bulk = list()
 
-        while True:
+        while not self._terminate.is_set():
 
             now = time.time()  # time of last submission
 
@@ -513,6 +553,9 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
                 if len(bulk) >= self._max_bulk_size:
                     break
 
+                if self._terminate.is_set():
+                    break
+
             if bulk:
                 logger.debug('submit bulk: %d', len(bulk))
                 self.tmgr.submit_tasks(bulk)
@@ -548,7 +591,15 @@ class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
 
     def shutdown(self, hub=True, targets='all', block=False):
         """Shutdown the executor, including all RADICAL-Pilot components."""
-        logger.info("RadicalPilotExecutor shutdown")
+        logger.info("RadicalPilotExecutor is terminating...")
+
+        self._terminate.set()
+
+        # ensure we are in the bulk submssion mode
+        if self.bulk_mode:
+            self._bulk_thread.join()
+
         self.session.close(download=True)
+        logger.info("RadicalPilotExecutor is terminated.")
 
         return True

parsl/executors/radical/rpex_resources.py

@@ -1,17 +1,13 @@
-import sys
 import json
-
+import sys
 from typing import List
 
-_setup_paths: List[str]
+_setup_paths: List[str] = []
 try:
     import radical.pilot as rp
     import radical.utils as ru
 except ImportError:
-    _setup_paths = []
-else:
-    _setup_paths = [rp.sdist_path,
-                    ru.sdist_path]
+    pass
 
 
 MPI = "mpi"
@@ -77,7 +73,7 @@ class ResourceConfig:
 
     pilot_env_setup : list
         List of setup commands/packages for the pilot environment.
-        Default setup includes "parsl", rp.sdist_path, and ru.sdist_path.
+        Default is an empty list.
 
     python_v : str
         The Python version to be used in the pilot environment.
@@ -108,7 +104,7 @@ class ResourceConfig:
     python_v: str = f'{sys.version_info[0]}.{sys.version_info[1]}'
     worker_type: str = DEFAULT_WORKER
 
-    def _get_cfg_file(cls, path=None):
+    def get_config(cls, path=None):
 
         # Default ENV mode for RP is to reuse
         # the client side. If this is not the case,
@@ -126,6 +122,7 @@ class ResourceConfig:
         cfg = {
             'n_masters': cls.masters,
             'n_workers': cls.workers,
+            'worker_type': cls.worker_type,
             'gpus_per_node': cls.worker_gpus_per_node,
             'cores_per_node': cls.worker_cores_per_node,
             'cores_per_master': cls.cores_per_master,
@@ -143,9 +140,10 @@ class ResourceConfig:
             'pilot_env_mode': cls.pilot_env_mode,
 
             'master_descr': {
+                "ranks": 1,
+                "cores_per_rank": 1,
                 "mode": rp.RAPTOR_MASTER,
                 "named_env": cls.pilot_env_name,
-                "executable": "python3 rpex_master.py",
             },
 
             'worker_descr': {
@@ -154,12 +152,16 @@ class ResourceConfig:
                 "raptor_file": "./rpex_worker.py",
                 "raptor_class": cls.worker_type if
                 cls.worker_type.lower() != MPI else MPI_WORKER,
+                "ranks": cls.nodes_per_worker * cls.worker_cores_per_node,
+                "gpus_per_rank": cls.nodes_per_worker * cls.worker_gpus_per_node,
             }}
 
-        # Convert the class instance to a cfg file.
-        config_path = 'rpex.cfg'
+        # Convert the class instance to a Json file or a Config dict.
         if path:
+            config_path = 'rpex.cfg'
             config_path = path + '/' + config_path
-        with open(config_path, 'w') as f:
-            json.dump(cfg, f, indent=4)
-        return config_path
+            with open(config_path, 'w') as f:
+                json.dump(cfg, f, indent=4)
+        else:
+            config_obj = ru.Config(from_dict=cfg)
+            return config_obj

parsl/executors/radical/rpex_worker.py

@@ -1,10 +1,11 @@
 import sys
+
 import radical.pilot as rp
 
 import parsl.app.errors as pe
 from parsl.app.bash import remote_side_bash_executor
-from parsl.serialize import unpack_res_spec_apply_message, serialize
-from parsl.executors.high_throughput.process_worker_pool import execute_task
+from parsl.executors.execute_task import execute_task
+from parsl.serialize import serialize, unpack_res_spec_apply_message
 
 
 class ParslWorker:
@@ -32,7 +33,7 @@ class ParslWorker:
 
         try:
             buffer = rp.utils.deserialize_bson(task['description']['executable'])
-            func, args, kwargs, _resource_spec = unpack_res_spec_apply_message(buffer, {}, copy=False)
+            func, args, kwargs, _resource_spec = unpack_res_spec_apply_message(buffer)
             ret = remote_side_bash_executor(func, *args, **kwargs)
             exc = (None, None)
             val = None