Parallel local

maestrowf/abstracts/interfaces/__init__.py

@@ -32,8 +32,8 @@
 """
 
 from maestrowf.abstracts.interfaces.schedulerscriptadapter import \
-    SchedulerScriptAdapter
+    SchedulerScriptAdapter, ParallelizeCmd
 from maestrowf.abstracts.interfaces.scriptadapter import ScriptAdapter
 
 
-__all__ = ("SchedulerScriptAdapter", "ScriptAdapter")
+__all__ = ("SchedulerScriptAdapter", "ScriptAdapter", "ParallelizeCmd")

maestrowf/abstracts/interfaces/schedulerscriptadapter.py

@@ -39,6 +39,21 @@
 LOGGER = logging.getLogger(__name__)
 
 
+@six.add_metaclass(ABCMeta)
+class ParallelizeCmd():
+    """
+    Callable object which generates the concrete scheduler task
+    launch commands.
+    """
+    def __init__(self, cmd_flags=None, unsupported=None):
+        self._cmd_flags = {}
+        self._unsupported = {}
+
+    @abstractmethod
+    def __call__(self, procs, nodes, **kwargs):
+        pass
+
+
 @six.add_metaclass(ABCMeta)
 class SchedulerScriptAdapter(ScriptAdapter):
     """
@@ -82,6 +97,9 @@ def __init__(self, **kwargs):
         super(SchedulerScriptAdapter, self).__init__(**kwargs)
         self._batch = {}
 
+        # Store launcher -> concrete parallel command mappings
+        self._cmd_flags = {}
+
     def add_batch_parameter(self, name, value):
         """
         Add a parameter to the ScriptAdapter instance.
@@ -92,6 +110,17 @@ def add_batch_parameter(self, name, value):
         """
         self._batch[name] = value
 
+    def add_cmd_flags(self, name, value):
+        """
+        Add command flags to the ScriptAdapter instance.  Command flags
+        map the launcher tokens to specific parallel invocations
+        """
+        if name in self._cmd_flags:
+            LOGGER.info(f"Overwriting '{name}' value ('{self._cmd_flags[name]}"
+                        f"') in cmd_flags with '{value}'")
+
+        self._cmd_flags[name] = value
+
     @abstractmethod
     def get_header(self, step):
         """
@@ -103,6 +132,10 @@ def get_header(self, step):
         """
         pass
 
+    def register_parallelize_command(self, parallelize_func):
+        # Note do some validation here -> callable types only
+        self._parallelize_func = parallelize_func()
+
     @abstractmethod
     def get_parallelize_command(self, procs, nodes, **kwargs):
         """

maestrowf/datastructures/core/executiongraph.py

@@ -9,6 +9,7 @@
 import tempfile
 from time import sleep
 from filelock import FileLock, Timeout
+from concurrent import futures
 
 from maestrowf.abstracts import PickleInterface
 from maestrowf.abstracts.enums import JobStatusCode, State, SubmissionCode, \
@@ -362,6 +363,7 @@ def __init__(self, submission_attempts=1, submission_throttle=0,
         self._submission_throttle = submission_throttle
         self.dry_run = dry_run
 
+
         # A map that tracks the dependencies of a step.
         # NOTE: I don't know how performant the Python dict structure is, but
         # we'll use it for now. I think this may want to be changed to an AVL
@@ -453,6 +455,7 @@ def set_adapter(self, adapter):
             msg = "'{}' adapter must be specfied in ScriptAdapterFactory." \
                   .format(adapter)
             LOGGER.error(msg)
+            LOGGER.error("Valid adapters: {}".format(ScriptAdapterFactory.get_valid_adapters()))
             raise TypeError(msg)
 
         self._adapter = adapter
@@ -534,7 +537,9 @@ def generate_scripts(self):
 
         # Set up the adapter.
         LOGGER.info("Generating scripts...")
-        adapter = ScriptAdapterFactory.get_adapter(self._adapter["type"])
+        # adapter = ScriptAdapterFactory.get_adapter(self._adapter["type"],
+        #                                            self._adapter.get('parallel_type')
+        adapter = ScriptAdapterFactory.get_adapter(self._adapter["type"])                                                   
         adapter = adapter(**self._adapter)
 
         self._check_tmp_dir()
@@ -893,9 +898,19 @@ def execute_ready_steps(self):
 
         # We now have a collection of ready steps. Execute.
         # If we don't have a submission limit, go ahead and submit all.
+        # Check requested resources -> nprocs
+        # nthreads = 1
+        # if self._local_procs > 0:
+        #     nthreads = self._local_procs
+
+        # if adapter.total_procs != nthreads:
+        #     adapter.total_procs = nthreads
+        #     adapter.avail_procs = nthreads
+
         if self._submission_throttle == 0:
             LOGGER.info("Launching all ready steps...")
             _available = len(self.ready_steps)
+            LOGGER.info("Ready steps: {}".format(self.ready_steps))
         # Else, we have a limit -- adhere to it.
         else:
             # Compute the number of available slots we have for execution.
@@ -909,9 +924,23 @@ def execute_ready_steps(self):
             _available = min(_available, len(self.ready_steps))
             LOGGER.info("Found %d available slots...", _available)
 
+        # for i in range(0, _available):
+        #     # Pop the record and execute using the helper method.
+        #     _record = self.values[self.ready_steps.popleft()]
+
+        #     # If we get to this point and we've cancelled, cancel the record.
+        #     if self.is_canceled:
+        #         logger.info("Cancelling '%s' -- continuing.", _record.name)
+        #         _record.mark_end(State.CANCELLED)
+        #         self.cancelled_steps.add(_record.name)
+        #         continue
+
+        #     logger.debug("Launching job %d -- %s", i, _record.name)
+        #     self._execute_record(_record, adapter)
+
         for i in range(0, _available):
             # Pop the record and execute using the helper method.
-            _record = self.values[self.ready_steps.popleft()]
+            _record = self.values[self.ready_steps[0]]
 
             # If we get to this point and we've cancelled, cancel the record.
             if self.is_canceled:
@@ -920,8 +949,32 @@ def execute_ready_steps(self):
                 self.cancelled_steps.add(_record.name)
                 continue
 
-            LOGGER.debug("Launching job %d -- %s", i, _record.name)
-            self._execute_record(_record, adapter)
+            # NOTE: verify this actually updates avail_procs on the fly, thus allowing the
+            # available tasks to be fully consumed before going back to sleep
+            LOGGER.debug("Attempting to submit step {} with total procs = {}, available procs = {}".format(_record.step.name, adapter.total_procs, adapter.avail_procs))
+            avail_procs = adapter.avail_procs
+            LOGGER.debug("avail_procs from the adapter = %d", avail_procs)
+            LOGGER.debug("total_procs from the adapter = %d", adapter.total_procs)
+            step_procs = _record.step.run.get("procs")
+            if not step_procs:
+                step_procs = 1
+            else:
+                try:
+                    step_procs = int(step_procs)
+                except ValueError:
+                    step_procs = 1
+                    LOGGER.error("Setting step {} with no 'procs' attribute to use 1 processor.".format(_record.step.name))
+
+            # NOTE: better place to set this default, and do type conversions?
+            if step_procs <= avail_procs:
+                LOGGER.debug("Launching job %d -- %s", i, _record.name)
+                self.ready_steps.popleft()  # remove key now that it's sure to run
+                self._execute_record(_record, adapter)
+            else:
+                LOGGER.debug("step_procs thought > avail_procs: %d : %d", step_procs, avail_procs)
+                if avail_procs == 0:
+                    break       # exit loop early to avoid needless churn
+
 
         # check the status of the study upon finishing this round of execution
         completion_status = self._check_study_completion()

maestrowf/datastructures/core/study.py

@@ -152,6 +152,19 @@ def __ne__(self, other):
         """
         return not self.__eq__(other)
 
+    def __rich_repr__(self):
+        """
+        Unique representation of this StudyStep object structure for 
+        pretty printing with the Rich library
+
+        NOTE: does this account for parameter expanded steps?
+        """
+        yield 'name', self.name
+        yield 'description', self.description
+        yield 'nickname', self.nickname
+        yield 'run', self.run
+
+
 
 class Study(DAG, PickleInterface):
     """

maestrowf/interfaces/__init__.py

@@ -31,8 +31,7 @@
 import logging
 import pkgutil
 import inspect
-from maestrowf.abstracts.interfaces import ScriptAdapter
-
+from maestrowf.abstracts.interfaces import ScriptAdapter, ParallelizeCmd
 
 __all__ = ("ScriptAdapterFactory",)
 LOGGER = logging.getLogger(__name__)
@@ -52,6 +51,7 @@ def iter_adapters():
     mods = [(name, ispkg) for finder, name, ispkg in pkgutil.iter_modules(
             loader.load_module('maestrowf.interfaces.script').__path__,
             loader.load_module('maestrowf.interfaces.script').__name__ + ".")]
+
     cs = []
     for name, _ in mods:
         # get loader for every module
@@ -61,6 +61,42 @@ def iter_adapters():
             if isinstance(cls, type) and issubclass(cls, ScriptAdapter) and \
                     not inspect.isabstract(cls):
                 cs.append(cls)
+                print("FOUND CLASS '{}' with key '{}'".format(cls.__name__, cls.key))
+            if isinstance(cls, type) and issubclass(cls, ScriptAdapter):
+                print("FOUND CLASS '{}'".format(cls.__name__))
+                LOGGER.debug("Found class '{}'".format(cls.__name__))
+
+    return cs
+
+
+def iter_parallel_cmds():
+    """
+    Based off of packaging.python.org loop over a namespace and find the
+    modules. This has been adapted for this particular use case of loading
+    all classes implementing ParallelizeCmd loaded from all modules in
+    maestrowf.interfaces.script.
+    :return: an iterable of the classes existing in the namespace
+    """
+    # get loader for the script adapter package
+    loader = pkgutil.get_loader('maestrowf.interfaces.script')
+    # get all of the modules in the package
+    mods = [(name, ispkg) for finder, name, ispkg in pkgutil.iter_modules(
+            loader.load_module('maestrowf.interfaces.script').__path__,
+            loader.load_module('maestrowf.interfaces.script').__name__ + ".")]
+
+    cs = []
+    for name, _ in mods:
+        # get loader for every module
+        m = pkgutil.get_loader(name).load_module(name)
+        # get all classes that implement ParallelizeCmd and are not abstract
+        for n, cls in m.__dict__.items():
+            if isinstance(cls, type) and issubclass(cls, ParallelizeCmd) and \
+                    not inspect.isabstract(cls):
+                cs.append(cls)
+                print("FOUND CLASS '{}' with key '{}'".format(cls.__name__, cls.key))
+            if isinstance(cls, type) and issubclass(cls, ParallelizeCmd):
+                print("FOUND CLASS '{}'".format(cls.__name__))
+                LOGGER.debug("Found class '{}'".format(cls.__name__))
 
     return cs
 
@@ -70,15 +106,28 @@ class ScriptAdapterFactory(object):
        adapter.key: adapter for adapter in iter_adapters()
     }
 
+    parallel_cmd_factories = {
+        parallel_cmd.key: parallel_cmd for parallel_cmd in iter_parallel_cmds()
+    }
+
     @classmethod
-    def get_adapter(cls, adapter_id):
+    def get_adapter(cls, adapter_id, parallel_type=None):
         if adapter_id.lower() not in cls.factories:
             msg = "Adapter '{0}' not found. Specify an adapter that exists " \
                   "or implement a new one mapping to the '{0}'" \
                   .format(str(adapter_id))
             LOGGER.error(msg)
             raise Exception(msg)
 
+        # adapter = cls.factories[adapter_id]
+
+        # # Attach any specified parallelize cmd generators
+        # if parallel_type and adapter_id in cls.parallel_cmd_factories:
+        #     adapter.register_parallelize_command(
+        #         cls.parallel_cmd_factories[adapter_id]
+        #     )
+
+        print(f"Returning: {cls.factories[adapter_id]}")
         return cls.factories[adapter_id]
 
     @classmethod