diff --git a/maestrowf/abstracts/interfaces/__init__.py b/maestrowf/abstracts/interfaces/__init__.py index e94906a6..0ac2bebd 100644 --- a/maestrowf/abstracts/interfaces/__init__.py +++ b/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") diff --git a/maestrowf/abstracts/interfaces/schedulerscriptadapter.py b/maestrowf/abstracts/interfaces/schedulerscriptadapter.py index 128799ec..1c0c2156 100644 --- a/maestrowf/abstracts/interfaces/schedulerscriptadapter.py +++ b/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): """ diff --git a/maestrowf/datastructures/core/executiongraph.py b/maestrowf/datastructures/core/executiongraph.py index 9ebdbf1e..76e1d3e5 100644 --- a/maestrowf/datastructures/core/executiongraph.py +++ b/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() diff --git a/maestrowf/datastructures/core/study.py b/maestrowf/datastructures/core/study.py index 2aea20df..d19636b3 100644 --- a/maestrowf/datastructures/core/study.py +++ b/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): """ diff --git a/maestrowf/interfaces/__init__.py b/maestrowf/interfaces/__init__.py index 156760a6..f3e0ddc0 100644 --- a/maestrowf/interfaces/__init__.py +++ b/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,8 +106,12 @@ 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}'" \ @@ -79,6 +119,15 @@ def get_adapter(cls, 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 diff --git a/maestrowf/interfaces/script/localparscriptadapter.py b/maestrowf/interfaces/script/localparscriptadapter.py new file mode 100644 index 00000000..50026be6 --- /dev/null +++ b/maestrowf/interfaces/script/localparscriptadapter.py @@ -0,0 +1,597 @@ +############################################################################### +# Copyright (c) 2017, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory +# Written by Francesco Di Natale, dinatale3@llnl.gov. +# +# LLNL-CODE-734340 +# All rights reserved. +# This file is part of MaestroWF, Version: 1.0.0. +# +# For details, see https://github.com/LLNL/maestrowf. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +############################################################################### + +"""Local interface implementation.""" +import logging +import os +import psutil +import re +import signal +import uuid + +from concurrent.futures import ThreadPoolExecutor +from concurrent.futures._base import PENDING as future_PENDING +from functools import partial as func_partial +from threading import Thread, RLock + +from maestrowf.abstracts.enums import JobStatusCode, SubmissionCode, \ + CancelCode, State +from maestrowf.interfaces.script import CancellationRecord, SubmissionRecord +from maestrowf.abstracts.interfaces import SchedulerScriptAdapter +from maestrowf.utils import start_process + + +LOGGER = logging.getLogger(__name__) + +DEFAULT_LOCAL_CMD_FLAGS = { + "cmd": "", + "ntasks": None, + "nodes": None, + "cores per task": None +} + + +def get_default_local_cmd_flags(): + """ + Returns default cmd flags for local adapters + """ + return DEFAULT_LOCAL_CMD_FLAGS + + +class LocalParallelScriptAdapter(SchedulerScriptAdapter): + """A ScriptAdapter class for interfacing for parallel local execution.""" + + key = "local_parallel" + + executor = None + running_steps = {} # needs to be singleton, threadsafe + done_steps = {} # storing completed futures/results to avoid losing them + tasks_lock = RLock() + total_procs = 1 + avail_procs = 1 + + # # The var tag to look for to replace for parallelized commands. + # launcher_var = "$(LAUNCHER)" + # # Allocation regex and compilation + # # Keeping this one here for legacy. + # launcher_regex = re.compile( + # re.escape(launcher_var) + r"\[(?P.*)\]") + + # # We can have multiple requested submission properties. + # # Legacy allocation of nodes and procs. + # legacy_alloc = r"(?P[0-9]+),\s*(?P[0-9]+)" + # # Just allocate based on tasks. + # task_alloc = r"(?P[0-9]+)p" + # # Just allocate based on nodes. + # node_alloc = r"(?P[0-9]+)n" + + def __init__(self, **kwargs): + """ + Initialize an instance of the LocalParallelScriptAdapter. + + The LocalParallelScriptAdapter is the adapter that is used for workflows that + will execute on the user's machine. This adapter constructs shell scripts for + a StudyStep based on user set defaults and local settings present in each step. + + These scripts are submitted to a local scheduler for asynchronous execution + and monitoring. + + :param **kwargs: A dictionary with default settings for the adapter. + """ + LOGGER.debug("kwargs\n--------------------------\n%s", kwargs) + print(f"Self type: {type(self)}") + print(f"Super type: {type(LocalParallelScriptAdapter)}") + print(f"self: {self}\nsuper: {LocalParallelScriptAdapter}") + print(f"LocalParallelScriptAdapter subclass of object: {issubclass(LocalParallelScriptAdapter, object)}") + print(f"self subclass of object: {issubclass(self.__class__, object)}") + print(f"self subclass of LocalParallelScriptAdapter: {issubclass(self.__class__, LocalParallelScriptAdapter)}") + print(f"self bases: {type(self).__bases__}") + print(f"LocalParallelScriptAdapter bases: {LocalParallelScriptAdapter.__bases__}") + print(f"self mro: {type(self).__mro__}") + print(f"LocalParallelScriptAdapter mro: {LocalParallelScriptAdapter.__mro__}") + + print(f"self methods: {sorted(self.__dict__.keys())}") + print(f"base methods: {sorted(LocalParallelScriptAdapter.__dict__.keys())}") + super().__init__(**kwargs) + # super(LocalParallelScriptAdapter, self).__init__(**kwargs) + + # Register keys + self.add_batch_parameter("proc_count", int(kwargs.pop("proc_count", "1"))) + + self._header = { + "procs": "# procs = {procs}", + } + + self.total_procs = self._batch['proc_count'] + self.avail_procs = self.total_procs + self.executor = ThreadPoolExecutor(max_workers=self.total_procs) + + # Setup initial no-op launcher parameters for default operation + # Load cmd flags for launcher token replacement + if "cmd_flags" in kwargs: + cmd_flags = kwargs.pop("cmd_flags") + else: + cmd_flags = get_default_local_cmd_flags() + + self._cmd_flags.update(cmd_flags) + + self._extension = '.sh' # read this from shell key if present? + + def __getstate__(self): + """Helper for excluding threadpool from pickling""" + state = self.__dict__.copy() + del state["executor"] + return state + + def __setstate__(self, state): + self.__dict__.update(state) + # Add baz back since it doesn't exist in the pickle + self.executor = ThreadPoolExecutor(max_workers=self.total_procs) + + def _update_running_steps(self, add_tasks=None, rm_tasks=None): + """Thread safe addition/removal of tasks from running_steps dict""" + if not add_tasks: + add_tasks = [] + + if not rm_tasks: + rm_tasks = [] + + with self.tasks_lock: + for task_id in rm_tasks: + if task_id in self.running_steps: + self.avail_procs += self.running_steps[task_id][2] # give back resources + + self.done_steps[task_id] = self.running_steps.pop(task_id) + + LOGGER.debug("Removing task {}, {} from running_steps".format( + task_id, self.done_steps[task_id])) + else: + LOGGER.error("Tried removing task {} from running_steps, " + "but it was not found.".format(task_id)) + + for task in add_tasks: + for task_id in task: + if task_id in self.running_steps: + LOGGER.error("Tried adding already tracked task {} " + "to running_steps.".format(task_id)) + else: + LOGGER.debug("Adding task {} to running_steps".format(task_id)) + self.running_steps[task_id] = task[task_id] + + # Want any status codes returned here if errors triggered? + + def get_header(self, step): + """ + Generate the header present at the top of execution scripts. + + :param step: A StudyStep instance. + :returns: A string of the header based on internal batch parameters and + the parameter step. + """ + return "" + + def get_parallelize_command(self, procs, nodes, **kwargs): + """ + Generate the parallelization segement of the command line. + + :param procs: Number of processors to allocate to the parallel call. + :param nodes: Number of nodes to allocate to the parallel call + (default = 1). + :returns: A string of the parallelize command configured using nodes + and procs. + :NOTE: this is currently a dummy method -> rework to add user specified + replacements later + :note: need a mechanism to override these/set from outside adapter + """ + if self._parallelize_func: + return self._parallelize_func(procs, nodes, **kwargs) + + args = [ + self._cmd_flags["cmd"], + ] + + return "".join(args) + + def get_scheduler_command(self, step): + """ + Generate the full parallelized command for use in a batch script. + + :param step: A StudyStep instance. + :returns: + 1. A Boolean value - True if command is to be scheduled, False + otherwise. + 2. A string representing the parallelized batch command for the + specified step command. + 3. A string representing the parallelized batch command for the + specified step restart command. + """ + # We should never get a study step that doesn't have a run entry; but + # better to be safe. + if not step.run: + msg = "Malformed StudyStep. A StudyStep requires a run entry." + LOGGER.error(msg) + raise ValueError(msg) + + # If the user is requesting nodes, we need to request the nodes and + # set up the command with scheduling. + _nodes = step.run.get("nodes", 0) + _procs = step.run.get("procs", 0) + + to_be_scheduled = False # Local parallel does not submit batch jobs + + if _nodes or _procs: + cmd = self._substitute_parallel_command( + step.run["cmd"], + **step.run + ) + LOGGER.debug("Running parallel command: %s", cmd) + + # Also check for the restart command and parallelize it too. + restart = "" + if step.run["restart"]: + restart = self._substitute_parallel_command( + step.run["restart"], + **step.run + ) + LOGGER.debug("Restart command: %s", cmd) + LOGGER.info("Running parallel workflow step '%s' locally.", step.name) + # Otherwise, just return the command. It doesn't need scheduling. + else: + LOGGER.info("Running workflow step '%s' locally.", step.name) + to_be_scheduled = False + cmd = step.run["cmd"] + restart = step.run["restart"] + + return to_be_scheduled, cmd, restart + + def _state(self, fut): + """ + Map a scheduler specific job state to a Study.State enum. + + :param fut: Future instance representing a task + :returns: A Study.State enum corresponding to parameter job_state. + """ + # NOTE: should all of this replace what's in check_jobs -> what about error code there? + if fut.running(): + return State.RUNNING + elif fut.done(): + if fut in self.running_steps: + return State.FINISHING + else: + return State.FINISHED + elif fut.cancelled(): + return State.CANCELLED + elif fut._state == future_PENDING: + return State.PENDING + else: + return State.UNKNOWN + + def _write_script(self, ws_path, step): + """ + Write a shell script to the workspace of a workflow step. + + The job_map optional parameter is a map of workflow step names to job + identifiers. This parameter so far is only planned to be used when a + study is configured to be launched in one go (more or less a script + chain using a scheduler's dependency setting). The functionality of + the parameter may change depending on both future intended use. + + :param ws_path: Path to the workspace directory of the step. + :param step: An instance of a StudyStep. + :returns: False (will not be scheduled), the path to the + written script for run["cmd"], and the path to the script written + for run["restart"] (if it exists). + """ + # THIS IS A HACK FOR NOW: make better use of get_scheduler_command later + cmd = step.run["cmd"] + restart = step.run["restart"] + # to_be_scheduled, _, _ = self.get_scheduler_command(step) + to_be_scheduled, cmd, restart = self.get_scheduler_command(step) + + fname = "{}.sh".format(step.name) + script_path = os.path.join(ws_path, fname) + with open(script_path, "w") as script: + script.write("#!{0}\n\n{1}\n".format(self._exec, cmd)) + + if restart: + rname = "{}.restart.sh".format(step.name) + restart_path = os.path.join(ws_path, rname) + + with open(restart_path, "w") as script: + script.write("#!{0}\n\n{1}\n".format(self._exec, restart)) + else: + restart_path = None + + return to_be_scheduled, script_path, restart_path + + def check_jobs(self, joblist): + """ + For the given job list, query execution status. + + This method uses the scontrol show job command and does a + regex search for job information. + + :param joblist: A list of job identifiers to be queried. + :returns: The return code of the status query, and a dictionary of job + identifiers to their status. + """ + + status = {} + status_code = JobStatusCode.OK + LOGGER.debug("Checking jobs {}".format(joblist)) + LOGGER.debug(" Currently running steps: {}".format(self.running_steps)) + for jid in joblist: + LOGGER.debug("Looking for job with id {}".format(jid)) + + if jid in self.running_steps: + LOGGER.debug("Job with id {} is in running steps".format(jid)) + fut, pid, step_procs = self.running_steps[jid] + LOGGER.debug("Job with id {} has future {} with states: done = {}, running = {}, _state = {}".format( + jid, fut, fut.done(), fut.running(), fut._state)) + + elif jid in self.done_steps: + fut, pid, step_procs = self.done_steps.pop(jid) + LOGGER.debug("Job with id {} has future {} with states: done = {}, running = {}, _state = {}".format( + jid, fut, fut.done(), fut.running(), fut._state)) + #status[jid] = State.FINISHED + + else: + LOGGER.debug("Job with id {} is not found in running steps".format(jid)) + # what state reaches this, and what's an appropriate return code? + continue + + # Future states: running, done, cancelled + if fut.done(): + # Check result/process retcode: subprocess errors caught here + # if fut in self.running_steps: + # NOTE: does the avail_procs in update_running_steps always catch ending or + # is there some other check needed here? -> save jid and remove from running steps? + result = fut.result() # this the right place to catch this? + LOGGER.debug("Job {}, has result: {}".format(jid, result)) + status[jid] = State.FINISHED + + elif fut.running(): + status[jid] = State.RUNNING + + elif fut.cancelled(): + # What about cancelled pid's? that would show up under fut.done()... + status[jid] = State.CANCELLED + + elif fut._state == future_PENDING: + LOGGER.debug("Job {}, with Fut {} is pending".format(jid, fut)) + status[jid] = State.PENDING + + else: + print("Job {}, with Fut {} with unknown state".format(jid, fut)) + if jid in self.running_steps: + LOGGER.debug("Job {}, with Fut {} is running?".format(jid, fut)) + LOGGER.debug("Job {}, Fut {} state: running {}, done {}, cancelled {}".format( + jid, fut, fut.done(), fut.running(), fut.cancelled())) + status[jid] = State.UNKNOWN # this ever reached, and if so how + status_code = JobStatusCode.ERROR + + return status_code, status + + def cancel_jobs(self, joblist): + """ + For the given job list, cancel each job. + + :param joblist: A list of job identifiers to be cancelled. + :returns: The return code to indicate if jobs were cancelled. + """ + if not joblist: + return CancellationRecord(CancelCode.OK, 0) + + retcode = 0 + + for jid in joblist: + if jid in self.running_steps: + fut, pid, step_procs = self.running_steps[jid] + if fut.done(): # cleanup + # Really need _submit to return submission record vs retcode/err? + result = fut.result() + LOGGER.debug("Removing job {} from running steps. result = {}".format(jid, result)) + + else: + # Interrupt running subprocesses + try: + self._kill(pid) #process.kill() # better way to kill it? + result=fut.result() # wait on future to exit + LOGGER.debug("Removing job {} from running steps. result = {}".format(jid, result)) + except KeyError: + LOGGER.error("Error, future {}, no longer in running step list".format(fut)) + retcode += 1 + + except: # TODO: catch exceptions from kill(), get_result() + LOGGER.exception("Error, unexpected behavior trying to cancel future {}, " + "and subprocess with id {}".format(fut, pid)) + retcode += 1 + + else: # What state occurs when executiongraph knows of jobs that aren't here? + # if not fut.done() and not fut.running(): + # This shouldn't be hit since execution graph only submits if resources available + LOGGER.error("Error, encounterd job with id {} in unexpected state".format(jid)) + retcode += 1 + + if retcode == 0: + _record = CancellationRecord(CancelCode.OK, retcode) + else: + LOGGER.error("Error code '%s' seen. Unexpected behavior " + "encountered.") # NOTE: what does logger.error inject in %s? + _record = CancellationRecord(CancelCode.ERROR, retcode) + + return _record + + @staticmethod + def _kill(subprocess_pid, sig=signal.SIGTERM, include_parent=True): + """Kill a process tree (including grandchildren) with signal + "sig" and return a (gone, still_alive) tuple. + "on_terminate", if specified, is a callabck function which is + called as soon as a child terminates. + + NOTE: borrowed from recipe in psutil docs + """ + assert subprocess_pid != os.getpid(), "won't kill myself" + parent = psutil.Process(subprocess_pid) + children = parent.children(recursive=True) + if include_parent: + children.append(parent) + + for p in children: + p.send_signal(sig) + LOGGER.debug("Killing process {}".format(p.name())) + + gone, alive = psutil.wait_procs(children, timeout=5, + callback=None) + + LOGGER.debug(" Gone, alive = {}, {}".format(gone, alive)) + # return (gone, alive) + + def _submit(self, p, step, path, cwd, job_map=None, env=None): + """ + Execute the step locally. + + If cwd is specified, the submit method will operate outside of the path + specified by the 'cwd' parameter. + If env is specified, the submit method will set the environment + variables for submission to the specified values. The 'env' parameter + should be a dictionary of environment variables. + + :param p: Subprocess object + :param step: An instance of a StudyStep. + :param path: Path to the script to be executed. + :param cwd: Path to the current working directory. + :param job_map: A map of workflow step names to their job identifiers. + :param env: A dict containing a modified environment for execution. + :returns: The return code of the submission command and job identiifer. + """ + LOGGER.debug("cwd = %s", cwd) + LOGGER.debug("Script to execute: %s", path) + + pid = p.pid + output, err = p.communicate() + retcode = p.wait() + + o_path = os.path.join(cwd, "{}.out".format(step.name)) + e_path = os.path.join(cwd, "{}.err".format(step.name)) + + with open(o_path, "w") as out: + out.write(output) + + with open(e_path, "w") as out: + out.write(err) + + if retcode == 0: + LOGGER.info("Execution returned status OK.") + # REPLACE PID WITH UUID? + return SubmissionRecord(SubmissionCode.OK, retcode, pid) + else: + LOGGER.warning("Execution returned an error: %s", str(err)) + _record = SubmissionRecord(SubmissionCode.ERROR, retcode, pid) + _record.add_info("stderr", str(err)) + return _record + + def submit(self, step, path, cwd, job_map=None, env=None): + """ + Execute the step locally. + + If cwd is specified, the submit method will operate outside of the path + specified by the 'cwd' parameter. + If env is specified, the submit method will set the environment + variables for submission to the specified values. The 'env' parameter + should be a dictionary of environment variables. + + :param step: An instance of a StudyStep. + :param path: Path to the script to be executed. + :param cwd: Path to the current working directory. + :param job_map: A map of workflow step names to their job identifiers. + :param env: A dict containing a modified environment for execution. + :returns: The return code of the submission command and job identiifer. + """ + try: + p = start_process(path, shell=False, cwd=cwd, env=env) + fut = self.executor.submit(self._submit, p, step, path, cwd, job_map=None, env=None) + try: + step_procs = step.run.get("procs") + if not step_procs: + step_procs = 1 + else: + step_procs = int(step_procs) + except ValueError: + step_procs = 1 + LOGGER.error("Starting step {} with no 'procs' attribute with 1 processor.".format(step.name)) + + self.avail_procs -= step_procs + + except: # add some relevant exceptions here.. + LOGGER.warning("Execution returned an error") + + _record = SubmissionRecord(SubmissionCode.ERROR, -1) + #self.avail_procs += step_procs + return _record + + # self.avail_procs -= step._procs + # except: + + # Update running task list. + jid = uuid.uuid4() + self._update_running_steps(add_tasks=[{jid: (fut, p.pid, step_procs)}]) + # fut.add_done_callback(func_partial(self._task_callback, fut, rm_func=self._update_running_steps)) + fut.add_done_callback(self.wrap_callback(self._update_running_steps, jid)) + + LOGGER.info("Execution returned status OK.") + return SubmissionRecord(SubmissionCode.OK, 0, jid) + + @staticmethod + def wrap_callback(rm_func, jid): + def task_callback(fut): + #fut.result() + rm_func(add_tasks=None, rm_tasks=[jid]) + + return task_callback + + @property + def extension(self): + """ + Returns the extension that generated scripts will use. + + :returns: A string of the extension + """ + return self._extension + + + # @property + # def key(self): + # """ + # Return the key name for a ScriptAdapter.. + + # This is used to register the adapter in the ScriptAdapterFactory + # and when writing the workflow specification. + # """ + # return self.key # diff --git a/maestrowf/interfaces/script/slurmscriptadapter.py b/maestrowf/interfaces/script/slurmscriptadapter.py index c15ee5f8..2e95a455 100644 --- a/maestrowf/interfaces/script/slurmscriptadapter.py +++ b/maestrowf/interfaces/script/slurmscriptadapter.py @@ -33,7 +33,8 @@ import os import re -from maestrowf.abstracts.interfaces import SchedulerScriptAdapter +from maestrowf.abstracts.interfaces import SchedulerScriptAdapter, \ + ParallelizeCmd from maestrowf.abstracts.enums import JobStatusCode, State, SubmissionCode, \ CancelCode from maestrowf.interfaces.script import CancellationRecord, SubmissionRecord @@ -41,6 +42,72 @@ LOGGER = logging.getLogger(__name__) +DEFAULT_SLURM_CMD_FLAGS = { + "cmd": "srun", + "depends": "--dependency", + "ntasks": "-n", + "nodes": "-N", + "cores per task": "-c", +} + + +def get_default_slurm_cmd_flags(): + """ + Returns default cmd flags for slurm adapters + """ + return DEFAULT_SLURM_CMD_FLAGS + + +class SlurmParallelizeCmd(ParallelizeCmd): + key = 'slurm' + + def __init__(self, cmd_flags=None, unsupported=None): + self._cmd_flags = {} + self._unsupported = set(["cmd", "depends", "ntasks", "nodes"]) + + self._cmd_flags.update(get_default_slurm_cmd_flags()) + + if cmd_flags: + self._cmd_flags.update(cmd_flags) + + def __call__(self, procs, nodes=None, **kwargs): + """ + Generate the SLURM parallelization segement of the command line. + + :param procs: Number of processors to allocate to the parallel call. + :param nodes: Number of nodes to allocate to the parallel call + (default = 1). + :returns: A string of the parallelize command configured using nodes + and procs. + """ + args = [ + # SLURM srun command + self._cmd_flags["cmd"], + # Processors segment + self._cmd_flags["ntasks"], + str(procs) + ] + + if nodes: + args += [ + self._cmd_flags["nodes"], + str(nodes), + ] + + supported = set(kwargs.keys()) - self._unsupported + for key in supported: + value = kwargs.get(key) + if key not in self._cmd_flags: + LOGGER.warning("'%s' is not supported -- omitted.", key) + continue + if value: + args += [ + self._cmd_flags[key], + "{}".format(str(value)) + ] + + return " ".join(args) + class SlurmScriptAdapter(SchedulerScriptAdapter): """A ScriptAdapter class for interfacing with the SLURM scheduler.""" @@ -99,13 +166,13 @@ def __init__(self, **kwargs): self._exclusive = "#SBATCH --exclusive" self._qos = "#SBATCH --qos={qos}" - self._cmd_flags = { - "cmd": "srun", - "depends": "--dependency", - "ntasks": "-n", - "nodes": "-N", - "cores per task": "-c", - } + # Load cmd flags for launcher token replacement + if "cmd_flags" in kwargs: + cmd_flags = kwargs.pop("cmd_flags") + else: + cmd_flags = get_default_slurm_cmd_flags() + + self._cmd_flags.update(cmd_flags) self._extension = ".slurm.sh" self._unsupported = set(["cmd", "depends", "ntasks", "nodes"]) @@ -174,6 +241,9 @@ def get_parallelize_command(self, procs, nodes=None, **kwargs): :returns: A string of the parallelize command configured using nodes and procs. """ + if self._parallelize_func: + return self._parallelize_func(procs, nodes, **kwargs) + args = [ # SLURM srun command self._cmd_flags["cmd"], diff --git a/tests/interfaces/script/test_slurmscriptadapter.py b/tests/interfaces/script/test_slurmscriptadapter.py index 6a6d3443..186beaf8 100644 --- a/tests/interfaces/script/test_slurmscriptadapter.py +++ b/tests/interfaces/script/test_slurmscriptadapter.py @@ -55,7 +55,12 @@ def test_slurm_adapter_in_factory(): """ saf = ScriptAdapterFactory # Make sure SlurmScriptAdapter is in the facotries object - assert(saf.factories[SlurmScriptAdapter.key] == SlurmScriptAdapter) + print(SlurmScriptAdapter) + test_adapter = SlurmScriptAdapter + # assert(saf.factories[SlurmScriptAdapter.key] is test_adapter) + print(id(saf.factories[SlurmScriptAdapter.key])) + print(id(SlurmScriptAdapter)) + assert(saf.factories[SlurmScriptAdapter.key] is SlurmScriptAdapter) # Make sure the SlurmScriptAdapter key is in the valid adapters assert(SlurmScriptAdapter.key in ScriptAdapterFactory.get_valid_adapters()) # Make sure that get_adapter returns the SlurmScriptAdapter when asking diff --git a/tests/interfaces/test_script_adapter.py b/tests/interfaces/test_script_adapter.py index 259a1878..72e1ba7a 100644 --- a/tests/interfaces/test_script_adapter.py +++ b/tests/interfaces/test_script_adapter.py @@ -34,9 +34,15 @@ as it was converted to dynamically load all ScriptAdapters using a namespace plugin methodology. """ +import os + import pytest from maestrowf.interfaces import ScriptAdapterFactory +from maestrowf.interfaces.script.slurmscriptadapter import SlurmParallelizeCmd +from maestrowf.datastructures.core import StudyStep + +from rich.pretty import pprint def test_factory(): @@ -53,7 +59,7 @@ def test_get_valid_adapters(): the results from get_valid_adapters() """ saf = ScriptAdapterFactory - assert(saf.factories.keys() == ScriptAdapterFactory.get_valid_adapters()) + assert(sorted(saf.factories.keys()) == sorted(ScriptAdapterFactory.get_valid_adapters())) def test_adapter_none_found(): @@ -63,3 +69,34 @@ def test_adapter_none_found(): """ with pytest.raises(Exception): ScriptAdapterFactory.get_adapter('empty-adapter') + +def test_adapter_script_generator(): + """ + Tests script generation and launcher token replacement + """ + test_step = StudyStep() + test_step.name = 'test-step' + test_step.description = 'script writer test' + test_step.run = { + 'cmd': '\n'.join(['$(LAUNCHER) echo "Hello, $(NAME)!" > hello_world.txt', + 'sleep 5']), + 'procs': 1, + 'nodes': '', + 'restart': '\n'.join(['']) + } + pprint(test_step) + + batch_info = { + 'type': 'local_parallel', + 'proc_count': 4, + 'shell': '/bin/bash' + } + adapter = ScriptAdapterFactory.get_adapter(batch_info['type']) + + adapter2 = adapter(**batch_info) + adapter2.register_parallelize_command(SlurmParallelizeCmd) + + print(adapter2) + adapter2.write_script(os.path.abspath('.'), test_step) + + assert(1 == None)