executors 0.6.0__py2.py3-none-any.whl

executors/__init__.py

@@ -0,0 +1,57 @@
+import os
+import sys
+import logging
+import executors.lsf as lsf
+import executors.slurm as slurm
+import executors.local as local
+import executors.pbsubmit as pbsubmit
+import executors.torque as torque
+
+logger = logging.getLogger(__name__)
+
+def get(name, partition='default', **kwargs):
+    if name == 'slurm':
+        return slurm.Executor(partition, **kwargs)
+    if name == 'pbsubmit':
+        return pbsubmit.Executor(partition, **kwargs)
+    if name == 'torque':
+        return torque.Executor(partition, **kwargs)
+    if name == 'lsf':
+        return lsf.Executor(partition, **kwargs)
+    if name == 'local':
+        return local.Executor(**kwargs)
+    raise SchedulerNotFound(name)
+
+def probe(partition, **kwargs):
+    if slurm.Executor.available():
+        logger.debug('detected slurm job scheduler')
+        return slurm.Executor(partition, **kwargs)
+    if torque.Executor.available():
+        logger.debug('detected torque job scheduler')
+        return torque.Executor(partition, **kwargs)
+    if pbsubmit.Executor.available():
+        logger.debug('detected pbsubmit job scheduler')
+        return pbsubmit.Executor(partition, **kwargs)
+    if lsf.Executor.available():
+        return lsf.Executor(partition, **kwargs)
+    logger.debug('no schedulers detected, so returning a local executor')
+    return local.Executor(**kwargs)
+
+def which(x):
+    for p in os.environ.get('PATH').split(os.pathsep):
+        p = os.path.join(p, x)
+        if os.path.exists(p):
+            return os.path.abspath(p)
+    return None
+
+class CalledProcessError(Exception):
+    def __init__(self, message, returncode, cmd, stdout, stderr):
+        super(CalledProcessError, self).__init__(message)
+        self.returncode = returncode
+        self.cmd = cmd
+        self.stdout = stdout
+        self.stderr = stderr
+
+class SchedulerNotFound(Exception):
+    pass
+

executors/__version__.py

@@ -0,0 +1,6 @@
+__title__ = 'executors'
+__description__ = 'Job scheduling abstractions'
+__url__ = 'https://github.com/harvard-nrg/executors'
+__version__ = '0.6.0'
+__author__ = 'Neuroinformatics Research Group'
+__author_email__ = 'info@neuroinfo.org'

executors/commons/__init__.py

@@ -0,0 +1,21 @@
+import os
+import re
+from itertools import zip_longest
+
+def which(x):
+    for p in os.environ.get('PATH').split(os.pathsep):
+        p = os.path.join(p, x)
+        if os.path.exists(p):
+            return os.path.abspath(p)
+    return None
+
+def match(s, expressions):
+    for expr in expressions:
+        if re.match(s, expr):
+            return True
+    return False
+
+def grouper(iterable, n, fillvalue=None):
+    args = [iter(iterable)] * n
+    return zip_longest(*args, fillvalue=fillvalue)
+

executors/exceptions/__init__.py

@@ -0,0 +1,8 @@
+class ExecutorNotFound(Exception):
+    pass
+
+class CommandNotFound(Exception):
+    pass
+
+class TimeoutError(Exception):
+    pass

executors/local/__init__.py

@@ -0,0 +1,107 @@
+import os
+import io
+import six
+import csv
+import time
+import logging
+import subprocess as sp
+import executors.commons as commons
+from executors.models import AbstractExecutor
+from executors.exceptions import ExecutorNotFound, CommandNotFound, TimeoutError
+from ratelimit import limits, sleep_and_retry
+
+logger = logging.getLogger('local')
+
+class Executor(AbstractExecutor):
+    def __init__(self, **kwargs):
+        self.ptab = dict()
+        self.poll_timeout = 60
+        self.poll_delay = 1
+
+    @staticmethod
+    def available():
+        return True
+
+    def submit(self, job, **kwargs):
+        '''
+        Submit a job locally. 
+
+        :param job: Job object
+        :type job: :mod:`executors.models.Job`
+        '''
+        cmd = job.command
+        if isinstance(cmd, six.string_types):
+            cmd = shlex.split(command)
+        logger.debug(cmd)
+        if job.error is None:
+            job.error = job.output
+        p = sp.Popen(
+            cmd,
+            stdout=open(job.output, 'w'),
+            stderr=open(job.error, 'w')
+        )
+        job.pid = p.pid
+        os.rename(job.output, job.output.replace('%j', str(job.pid)))
+        os.rename(job.error, job.error.replace('%j', str(job.pid)))
+        self.ptab[job.pid] = p
+
+    def cancel(self, job, wait=False):
+        '''
+        Kill the process.
+
+        Since killing a job is inherently asyncronous, pass 
+        wait=True to wait until the job state is updated and 
+        the process is confirmed dead.
+
+        :param job: Job object
+        :type job: :mod:`executors.models.Job`
+        :param wait: Wait until the job state is updated
+        :type wait: bool
+        '''
+        if not wait:
+            self._cancel_async(job.pid)
+            return
+        # cancel job and wait for the job to be confirmed dead
+        self._cancel_async(job.pid)
+        logger.debug('waiting for job %s to be killed', job_id)
+        tic = time.time()
+        while 1:
+            self.update(job)
+            if not job.active:
+                break
+            if time.time() - tic > self.poll_timeout:
+                raise TimeoutError('exceeded wait time {0}s for job {1}'.format(self.poll_timeout, job_id))
+            time.sleep(self.poll_delay)
+
+    def _cancel_async(self, pid):
+        logger.debug('killing pid %s', pid)
+        self.ptab[pid].kill()
+
+    def update(self, job, wait=False):
+        '''
+        Update a single job state.
+
+        :param job: Job object
+        :type job: :mod:`executors.models.Job`
+        :param wait: Wait for job state to be updated
+        :type wait: bool
+        '''
+        p = self.ptab[job.pid]
+        p.poll()
+        job.active = True
+        if p.returncode is not None:
+            job.active = False
+            job.returncode = p.returncode
+
+    def update_many(self, jobs, wait=False):
+        '''
+        Update multiple job states.
+
+        :param jobs: List of :mod:`executors.models.Job` objects
+        :type jobs: list
+        :param wait: Wait for job state to be updated
+        :type wait: bool
+        '''
+        # this should be implemented as one call
+        for job in jobs:
+            self.update(job, wait=wait)

executors/lsf/__init__.py

@@ -0,0 +1,178 @@
+import os
+import re
+import logging
+import subprocess as sp
+from executors.commons import which
+from executors.models import AbstractExecutor
+from executors.exceptions import ExecutorNotFound, CommandNotFound
+
+logger = logging.getLogger(__name__)
+
+class Executor(AbstractExecutor):
+    ACTIVE = (
+        'PEND',    # pending
+        'WAIT',    # waiting
+        'RUN',     # running
+        'PSUSP',   # suspended while pending
+        'USUSP',   # suspended while running
+        'SSUSP'    # suspended for other reason
+    )
+    INACTIVE = (
+        'DONE',    # completed successfully
+        'EXIT',    # completed unsuccessfully
+        'ZOMBI'    # zombie state
+    )
+
+    def __init__(self, partition, **kwargs):
+        if not self.available():
+            raise BsubNotFound()
+        self.partition = partition
+        self.polling_interval = 5
+        self.timeout = 60
+        self._default_args = self.default_args(**kwargs)
+
+    def default_args(self, **kwargs):
+        args = list()
+        for k,v in iter(kwargs.items()):
+            logger.warn('unrecognized executor argument "%s"', k)
+        return args
+
+    @staticmethod
+    def available():
+        if which('bsub'):
+            return True
+        return False
+
+    def submit(self, job):
+        prefix = '{0}-%j'.format(job.name) if job.name else '%j'
+        if not job.output:
+            job.output = os.path.expanduser('~/{0}.out'.format(prefix))
+        if not job.error:
+            job.error = os.path.expanduser('~/{0}.err'.format(prefix))
+        command = job.command
+        if isinstance(command, list):
+            command = sp.list2cmdline(command)
+        if not which('bsub'):
+            raise CommandNotFound('bsub')
+        cmd = [
+            'bsub',
+            '-q', self.partition
+        ]
+        cmd.extend(self._default_args)
+        cmd.extend(self._arguments(job))
+        cmd.extend([
+            command
+        ])
+        logger.debug(sp.list2cmdline(cmd))
+        output = sp.check_output(cmd, stderr=sp.STDOUT).strip().decode()
+        pid = re.search(r'^Job <(\d+)>', output).group(1)
+        logger.debug('parsed job id %s', pid)
+        job.pid = pid
+ 
+    def update(self, job, wait=False):
+        status = self.bjobs(job)
+        job_state = status['job_state']
+        exit_status = status['exit_status']
+        output_path = status['output_path']
+        error_path = status['error_path']
+        logger.debug('job {0} is in {1} state'.format(job.pid, job_state))
+        if job_state in Executor.ACTIVE:
+            job.active = True
+        elif job_state in Executor.INACTIVE:
+            job.active = False
+            job.returncode = int(exit_status)
+
+    def update_many(self, jobs, wait=False):
+        for job in jobs:
+            self.update(job, wait=wait)
+
+    def cancel(self, job, wait=False):
+        if not which('bkill'):
+            raise CommandNotFound('bkill')
+        cmd = [
+            'bkill',
+            job.pid
+        ]
+        try:
+            logger.debug(cmd)
+            sp.check_output(cmd, stderr=sp.PIPE)
+        except sp.CalledProcessError as e:
+            # qdel will return a 255 exit status if it tries to query the 
+            # state of a Job ID that is already in a 'C' state or if the 
+            # Job ID is unknown. We should pass on either of these states
+            if e.returncode == 255:
+                logger.debug('job %s is in a completed state or unknown and cannot be cancelled', job.pid)
+                pass
+            raise e
+
+    def bjobs(self, job):
+        if not which('bjobs'):
+            raise CommandNotFound('bjobs')
+        cmd = [
+            'bjobs',
+            '-l',
+            job.pid
+        ]
+        logger.debug(cmd)
+        try:
+            output = sp.check_output(cmd).strip().decode()
+        except sp.CalledProcessError as e:
+            raise e
+        pid = re.match(r'Job <(\d+)>', output).group(1)
+        job_state = re.search(r'Status <(\w+)>', output).group(1)
+        exit_status = None
+        if job_state in Executor.INACTIVE:
+            exit_status = 0
+            if job_state == 'EXIT':
+                exit_status = re.search(r'Exited with exit code (\d+).', output).group(1)
+        return {
+            'pid': pid,
+            'job_state': job_state,
+            'exit_status': exit_status,
+            'output_path': job.output,
+            'error_path': job.error
+        }
+
+    def _parse_mem_value(self, s):
+        try:
+            match = re.match(r'^(\d+)(K|KB|M|MB|G|GB|T|TB)$', s)
+            size,unit = match.group(1),match.group(2)
+        except:
+            raise IndecipherableMemoryArgument(m)
+        if unit in ('K', 'KB'):
+            memarg = .001 * int(size)
+        elif unit in ('M', 'MB'):
+            memarg = int(size)
+        elif unit in ('G', 'GB'):
+            memarg = int(size) * 1000
+        elif unit in ('T', 'TB'):
+            memarg = int(size) * 1000000
+        if memarg < 1:
+            memarg = 1
+        logger.debug('translated memory argument %s', memarg)
+        return str(int(memarg))
+
+    def _arguments(self, job):
+        arguments = list()
+        qsub_opts = dict()
+        if hasattr(job, 'output') and job.output:
+            o = job.output.replace('%j', '%J')
+            arguments.extend(['-o', os.path.expanduser(o)])
+        if hasattr(job, 'error') and job.error:
+            e = job.error.replace('%j', '%J')
+            arguments.extend(['-e', os.path.expanduser(e)])
+        if hasattr(job, 'name') and job.name:
+            arguments.extend(['-J', job.name])
+        if hasattr(job, 'memory') and job.memory:
+            arguments.extend(['-M', self._parse_mem_value(job.memory)])
+        if hasattr(job, 'parent') and job.parent:
+            arguments.extend(['-w', 'done({0})'.format(job.parent.pid)])
+        if hasattr(job, 'processors') and job.processors:
+            arguments.extend(['-n', job.processors])
+        return arguments
+
+class BsubNotFound(ExecutorNotFound):
+    pass
+
+class IndecipherableMemoryArgument(Exception):
+    pass

executors/models/__init__.py

@@ -0,0 +1,172 @@
+import abc
+import time
+import logging
+import executors.commons as commons
+
+logger = logging.getLogger(__name__)
+
+class Job(object):
+    def __init__(self, command, memory, time, cpus=1, nodes=1, gpus=None, name=None, output=None, error=None, parent=None):
+        self.name = name
+        self.command = command
+        self.memory = memory
+        self.time = time
+        self.parent = parent
+        self.output = output
+        self.error = error
+        self.cpus = cpus
+        self.gpus = gpus
+        self.nodes = nodes
+        # these fields are managed by an Executor
+        self.pid = None
+        self.returncode = None
+        self.active = None
+
+class JobArray(object):
+    def __init__(self, executor, cancel_on_fail=False):
+        '''
+        :param executor: Executor instance
+        :type executor: :mod:`executors.models.AbstractExecutor`
+        :param cancel_on_fail: Cancel remaining jobs if any job fails
+        :type cancel_on_fail: bool
+        '''
+        self.E = executor
+        self.cancel_on_fail = cancel_on_fail
+        self.array = list()
+        self.running = dict() 
+        self.complete = dict()
+        self.failed = dict()
+        self.active = 0
+        self.inactive = 0
+    
+    def add(self, job):
+        '''
+        Add a job object to this job array.
+
+        :param job: Job object
+        :type job: :mod:`executors.models.Job`
+        '''
+        self.array.append(job)
+
+    def submit(self, limit=None, delay=None):
+        '''
+        Submit the job array. To rate limit the number of jobs running 
+        concurrently, use the limit parameter. To inject an artificial 
+        delay between job submissions, use the delay parameter.
+
+        Setting limit=None will submit all jobs and this method will return 
+        immediately. You are responsible for calling JobArray.wait if you want 
+        a blocking call. Providing a limit parameter will turn this method into 
+        a blocking call since it needs to continuously monitor running jobs.
+
+        :param limit: Limit the number of jobs running concurrently.
+        :type limit: int
+        :param delay: Wait N seconds between each job submission.
+        :type delay: int
+        '''
+        if not limit:
+            limit = len(self.array)
+        cancel = False
+        for job in self.array:
+            submitted = False
+            while not submitted:
+                if len(self.running) < limit:
+                    self.E.submit(job)
+                    self.running[job.pid] = job
+                    logger.debug('%s was submitted with pid %s', job.name, job.pid)
+                    submitted = True
+                    if delay:
+                        time.sleep(delay)
+                else:
+                    self.update()
+                    if len(self.failed) and self.cancel_on_fail:
+                        logger.debug('cancelling gradual submission')
+                        cancel = True
+                        break
+            if cancel:
+                break
+        if limit:
+            self.wait()
+    
+    def update(self):
+        '''
+        Update all jobs in the job array.
+        '''
+        self.E.update_many(self.running.values())
+        for pid in list(self.running):
+            job = self.running[pid]
+            if job.returncode == None:
+                continue
+            elif job.returncode == 0:
+                logger.debug('job %s (%s) returncode is %s', job.pid, job.name, job.returncode)
+                self.complete[pid] = job
+                del self.running[pid]
+            elif job.returncode > 0:
+                logger.debug('job %s (%s) returncode is %s', job.pid, job.name, job.returncode)
+                self.failed[pid] = job
+                del self.running[pid]
+
+    def wait(self, confirm_cancel=False):
+        '''
+        Wait for all jobs in the job array to complete. 
+
+        If the job array is configured with cancel_on_fail=True, you can 
+        pass confirm_cancel=True to wait for the executor to confirm that 
+        all jobs have been cancelled. Without this, jobs will only be 
+        signaled to cancel.
+
+        :param confirm_cancel: Wait until jobs have been cancelled.
+        :type confirm_cancel: bool
+        '''
+        while 1:
+            self.update()
+            # if any jobs have failed and self.cancel_on_fail=True, cancel remaining jobs
+            if len(self.failed) > 0 and self.cancel_on_fail:
+                self.cancel(confirm=confirm_cancel)
+                return
+            if len(self.running) == 0:
+                return
+
+    def cancel(self, confirm=False):
+        '''
+        Cancel all jobs in the job array. Pass confirm=True to wait for the 
+        executor to confirm that all jobs have indeed been cancelled. Without 
+        this parameter, jobs will only be signaled to cancel.
+
+        :param confirm: Confirm that jobs have been cancelled
+        :type confirm: bool
+        '''
+        for job in self.running.values():
+            self.E.cancel(job, wait=confirm)
+
+class AbstractExecutor(object):
+    __metaclass__ = abc.ABCMeta
+
+    @abc.abstractmethod
+    def __init__(self, partition):
+        pass
+
+    @abc.abstractproperty
+    def ACTIVE(self):
+        raise NotImplementedError
+
+    @abc.abstractproperty
+    def INACTIVE(self):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def submit(self, job):
+        pass
+
+    @abc.abstractmethod
+    def cancel(self, job):
+        pass
+
+    @abc.abstractmethod
+    def update(self, job):
+        pass
+
+    @abc.abstractmethod
+    def update_many(self, job):
+        pass
+

executors/pbsubmit/__init__.py

@@ -0,0 +1,234 @@
+import os
+import re
+import logging
+import paramiko
+import getpass as gp
+import subprocess as sp
+import xml.etree.ElementTree as et
+from executors.commons import which
+from executors.models import AbstractExecutor
+from executors.exceptions import ExecutorNotFound, CommandNotFound, TimeoutError
+
+logger = logging.getLogger(__name__)
+
+class Executor(AbstractExecutor):
+    ACTIVE = (
+        'Q',    # queued
+        'R',    # running
+        'H',    # held
+        'E'     # exited
+    )
+    INACTIVE = (
+        'C',    # complete
+    )
+
+    def __init__(self, partition, **kwargs):
+        if not self.available():
+            raise PBSubmitNotFound()
+        self.partition = partition
+        self.polling_interval = 5
+        self.timeout = 60
+        self._default_args = self.default_args(**kwargs)
+
+    def default_args(self, **kwargs):
+        args = list()
+        for k,v in iter(kwargs.items()):
+            if k == 'nodes':
+                args.extend([
+                    '-l', '+'.join(v)
+                ])
+            else:
+                logger.warn('unrecognized Executor argument "%s"', k)
+        return args
+
+    @staticmethod
+    def available():
+        if which('pbsubmit'):
+            return True
+        return False
+
+    def submit(self, job):
+        command = job.command
+        if isinstance(command, list):
+            command = sp.list2cmdline(command)
+        if not which('pbsubmit'):
+            raise CommandNotFound('pbsubmit')
+        cmd = [
+            'pbsubmit',
+            '-q', self.partition
+        ]
+        cmd.extend(self._arguments(job))
+        cmd.extend([
+            '-c', command
+        ])
+        logger.debug(sp.list2cmdline(cmd))
+        output = sp.check_output(cmd, stderr=sp.STDOUT).decode('utf-8')
+        output = output.strip().split('\n')
+        pid = output[-1]
+        job.pid = pid
+        self._alter_logs(job) # insert pid into stdout and stderr files
+        pbsjob = re.match(r'^Opening pbsjob_(\d+)', output[0]).groups(0)[0]
+        job.pbsjob = pbsjob
+
+    def _alter_logs(self, job):
+        match = re.match(r'^(\d+)\.', job.pid)
+        pid = match.group(1)        
+        qalter_args = list()
+        if job.output and '%j' in job.output:
+            output = job.output.replace('%j', pid)
+            qalter_args.extend(['-o', os.path.expanduser(output)])
+        if job.error and '%j' in job.error:
+            error = job.error.replace('%j', pid)
+            qalter_args.extend(['-e', os.path.expanduser(error)])
+        if qalter_args:
+            if not which('qalter'):
+                raise CommandNotFound('qalter')
+            cmd = ['qalter'] + qalter_args + [pid]
+            sp.check_output(cmd)
+
+    def update(self, job, wait=False):
+        xml = self.qstat(job)
+        job_state = xml.findtext('.//job_state')
+        exit_status = xml.findtext('.//exit_status')
+        output_path = re.sub(r'^.*:', '', xml.findtext('.//Output_Path'))
+        error_path = re.sub(r'^.*:', '', xml.findtext('.//Error_Path'))
+        logger.debug('job {0} is in {1} state'.format(job.pid, job_state))
+        if job_state in Executor.ACTIVE:
+            job.active = True
+        elif job_state in Executor.INACTIVE:
+            job.active = False
+            job.returncode = int(exit_status)
+
+    def update_many(self, jobs, wait=False):
+        for job in jobs:
+            self.update(job, wait=wait)
+
+    def cancel(self, job, wait=False):
+        if not which('qdel'):
+            raise CommandNotFound('qdel')
+        cmd = [
+            'qdel',
+            job.pid
+        ]
+        try:
+            logger.debug(cmd)
+            sp.check_output(cmd, stderr=sp.PIPE)
+        except sp.CalledProcessError as e:
+            # qdel will return a 153 exit status if it tries to query the 
+            # state of a Job ID that is already in a 'C' state, or a 170 
+            # exit status if the Job ID is unknown. We should pass on either
+            # of these states. A Job ID can become unknown only minutes after 
+            # a job has entered the 'C' state.
+            if e.returncode == 170:
+                logger.debug('job %s is in a completed state and cannot be cancelled', job.pid)
+                pass
+            elif e.returncode == 153:
+                logger.debug('job %s is unknown and cannot be cancelled', job.pid)
+                pass
+            else:
+                raise e
+
+    def qstat(self, job):
+        if not which('qstat'):
+            raise CommandNotFound('qstat')
+        cmd = [
+            'qstat',
+            '-x',
+            '-f',
+            job.pid
+        ]
+        logger.debug(cmd)
+        try:
+            output = sp.check_output(cmd)
+        except sp.CalledProcessError as e:
+            if e.returncode == 170:
+                logger.debug('job %s already in completed state, falling back to jobinfo', job.pid)
+                output = self.jobinfo(job)
+            elif e.returncode == 153:
+                logger.debug('job %s unknown to the scheduler, falling back to jobinfo', job.pid)
+                output = self.jobinfo(job)
+            else:
+                raise e
+        return et.fromstring(output.strip())
+
+    def jobinfo(self, job, node='launchpad'):
+        cmd = [
+            'jobinfo',
+            job.pid
+        ]
+        username = gp.getuser()
+        # ssh into head node to get job info
+        logging.getLogger('paramiko').setLevel(logging.INFO)
+        client = paramiko.SSHClient()
+        client.load_system_host_keys()
+        client.connect(node)
+        logger.debug('jobinfo command %s', cmd)
+        _,stdout,_ = client.exec_command(sp.list2cmdline(cmd))
+        stdout = stdout.read()
+        client.close()
+        # get job pid without domain  
+        match = re.match(r'^(\d+)\.', job.pid)
+        pid = match.group(1)
+        # parse exit status
+        stdout = stdout.decode('utf-8').strip().split('\n')
+        match = re.match(r'^\s+Exit status: (-?\d+)$', stdout[-1])
+        exit_status = match.group(1)
+        logger.debug('discovered exit status: %s', exit_status)
+        # build XML output similar to qstat -x
+        root = et.Element('jobinfo')
+        et.SubElement(root, 'job_state').text = 'C'
+        et.SubElement(root, 'exit_status').text = exit_status
+        et.SubElement(root, 'Output_Path').text = '/pbs/{0}/pbsjob_{1}.o{2}'.format(username, job.pbsjob, pid)
+        et.SubElement(root, 'Error_Path').text = '/pbs/{0}/pbsjob_{1}.e{2}'.format(username, job.pbsjob, pid)
+        return et.tostring(root)
+
+    def _parse_mem_value(self, s):
+        try:
+            match = re.match(r'^(\d+)(K|KB|M|MB|G|GB|T|TB)$', s)
+            size,unit = match.group(1),match.group(2)
+        except:
+            raise IndecipherableMemoryArgument(m)
+        if unit in ('K', 'KB'):
+            unit = 'kb'
+        elif unit in ('M', 'MB'):
+            unit = 'mb'
+        elif unit in ('G', 'GB'):
+            unit = 'gb'
+        elif unit in ('T', 'TB'):
+            unit = 'tb'
+        memarg = size + unit
+        logger.debug('translated memory argument %s', memarg)
+        return size + unit
+
+    def _arguments(self, job):
+        arguments = list()
+        qsub_opts = dict()
+        if hasattr(job, 'output') and job.output:
+            arguments.extend(['-O', os.path.expanduser(job.output)])
+        if hasattr(job, 'error') and job.error:
+            arguments.extend(['-E', os.path.expanduser(job.error)])
+        if hasattr(job, 'parent') and job.parent:
+            arguments.extend(['-W', 'depend=afterok:{0}'.format(job.parent.pid)])
+        if hasattr(job, 'name') and job.name:
+            arguments.extend(['-o', '-N {0}'.format(job.name)])
+        if hasattr(job, 'memory') and job.memory:
+            qsub_opts['vmem'] = self._parse_mem_value(job.memory)
+        if hasattr(job, 'processors') and job.processors:
+            qsub_opts['ppn'] = job.processors
+        # build and append pass-through qsub options
+        qsub_opts = 'nodes={NODES}:ppn={PPN},vmem={VMEM}'.format(
+            NODES=qsub_opts.get('nodes', 1),
+            PPN=qsub_opts.get('ppn', 1),
+            VMEM=qsub_opts.get('vmem', '1gb')
+        )
+        arguments.extend(['-l', qsub_opts])
+        return arguments
+
+class PBSubmitNotFound(ExecutorNotFound):
+    pass
+
+class IndecipherableMemoryArgument(Exception):
+    pass
+
+class QstatError(Exception):
+    pass

executors/slurm/__init__.py

@@ -0,0 +1,293 @@
+import os
+import io
+import six
+import csv
+import time
+import logging
+import subprocess as sp
+import executors.commons as commons
+from executors.models import AbstractExecutor
+from executors.exceptions import ExecutorNotFound, CommandNotFound, TimeoutError
+from ratelimit import limits, sleep_and_retry
+
+logger = logging.getLogger('slurm')
+
+class Executor(AbstractExecutor):
+    ACTIVE = (
+        r'PENDING',
+        r'CONFIGURING',
+        r'RUNNING',
+        r'RESIZING',
+        r'SUSPENDED',
+        r'COMPLETING'
+    )
+    INACTIVE = (
+        r'COMPLETED',
+        r'CANCELLED',
+        r'CANCELLED by \d+',
+        r'FAILED',
+        r'OUT_OF_MEMORY', 
+        r'NODE_FAIL',
+        r'PREEMPTED', 
+        r'TIMEOUT'
+    )
+
+    def __init__(self, partition, **kwargs):
+        '''
+        :param partition: Slurm partition
+        :type partition: str
+        :param nodelist: List of nodes to include (optional)
+        :type nodelist: list
+        :param exclude: Nodes to exclude (optional)
+        :type nodelist: list
+        :param reservation: Slurm reservation (optional)
+        :type reservation: str
+        '''
+        if not self.available():
+            raise SlurmNotFound()
+        self.partition = partition
+        self.poll_timeout = 60
+        self.poll_delay = 1
+        self.args = self._translate(kwargs)
+
+    def _translate(self, kwargs):
+        args = list()
+        for k,v in iter(kwargs.items()):
+            if k == 'nodelist':
+                if not isinstance(v, list):
+                    raise SbatchError('nodelist argument must be a list')
+                args.extend([
+                    '--nodelist', ','.join(v)
+                ])
+            elif k == 'exclude':
+                if not isinstance(v, list):
+                    raise SbatchError('nodelist argument must be a list')
+                args.extend([
+                    '--exclude', ','.join(v)
+                ])
+            elif k == 'reservation':
+                args.extend([
+                    '--reservation', v
+                ])
+            else:
+                logger.warn('unrecognized argument "%s"', k)
+        return args
+
+    @staticmethod
+    def available():
+        '''
+        Check if Slurm is available on the system.
+        '''
+        if commons.which('sbatch'):
+            return True
+        return False
+
+    def submit(self, job, **kwargs):
+        '''
+        Submit a job with sbatch. Pass wrap=False to disable wrapping the 
+        command.
+
+        :param job: Job object
+        :type job: :mod:`executors.models.Job`
+        :param wrap: Disable wrapping
+        :type wrap: bool
+        '''
+        if not commons.which('sbatch'):
+            raise CommandNotFound('sbatch')
+        if not self.partition:
+            raise SbatchError('no slurm partition was defined')
+        cmd = [
+            'sbatch',
+            '--parsable',
+            '--partition', self.partition
+        ]
+        cmd.extend(self.args)
+        cmd.extend(self._arguments(job))
+        wrap = kwargs.get('wrap', True)
+        command = job.command
+        if wrap:
+            if isinstance(command, list):
+                command = sp.list2cmdline(command)
+            cmd.extend([
+                '--wrap', command
+            ])
+        else:
+            if isinstance(command, six.string_types):
+                command = shlex.split(command)
+            cmd.extend(command)
+        logger.debug(cmd)
+        pid =  sp.check_output(cmd).strip().decode()
+        job.pid = pid
+
+    def cancel(self, job, wait=False):
+        '''
+        Send scancel to a Slurm job. 
+
+        Since cancelling a job in Slurm is inherently asyncronous, pass 
+        wait=True to wait until the job state is updated.
+
+        :param job: Job object
+        :type job: :mod:`executors.models.Job`
+        :param wait: Wait until the job state is updated
+        :type wait: bool
+        '''
+        job_id = job.pid + '.batch'
+        if not wait:
+            self._cancel_async(job.pid)
+            return
+        # cancel job and wait for the job to be registered as cancelled
+        self._cancel_async(job.pid)
+        logger.debug('waiting for job %s to be registered as CANCELLED', job_id)
+        tic = time.time()
+        while 1:
+            self.update(job)
+            if not job.active:
+                break
+            if time.time() - tic > self.poll_timeout:
+                raise TimeoutError('exceeded wait time {0}s for job {1}'.format(self.poll_timeout, job_id))
+            time.sleep(self.poll_delay)
+
+    def _cancel_async(self, job_id):
+        if not commons.which('scancel'):
+            raise CommandNotFound('scancel')
+        cmd = [
+            'scancel',
+            job_id
+        ]
+        logger.debug(cmd)
+        sp.check_output(cmd)
+
+    def update(self, job, wait=False):
+        '''
+        Update a single job state.
+
+        Since querying job state in Slurm is inherently asyncronous, you must 
+        pass wait=True to wait until the job state is updated.
+
+        :param job: Job object
+        :type job: :mod:`executors.models.Job`
+        :param wait: Wait for job state to be updated
+        :type wait: bool
+        '''
+        # run sacct
+        rows = self.sacct(job, wait=wait)
+        # we should only have one row at this point
+        if len(rows) == 0:
+            return rows
+        elif len(rows) > 1:
+            raise SacctError('more rows than expected after parsing sacct output: {0}'.format(rows))
+        row = rows.pop()
+        job_state = row['State']
+        job_code,sbatch_code = row['ExitCode'].split(':')
+        # assume the return code is the greater of the two
+        exit_status = max([int(sbatch_code), int(job_code)])
+        logger.debug('pid {0} is in "{1}" state'.format(job.pid, job_state))
+        if commons.match(job_state, Executor.ACTIVE):
+            job.active = True
+        elif commons.match(job_state, Executor.INACTIVE):
+            job.active = False
+            job.returncode = int(exit_status)
+
+    def update_many(self, jobs, wait=False):
+        '''
+        Update multiple job states.
+
+        Since querying job state in Slurm is inherently asyncronous, you must 
+        pass wait=True to wait until the job state is updated.
+
+        :param jobs: List of :mod:`executors.models.Job` objects
+        :type jobs: list
+        :param wait: Wait for job state to be updated
+        :type wait: bool
+        '''
+        # this should be implemented as one call
+        for job in jobs:
+            self.update(job, wait=wait)
+
+    def sacct(self, job, wait=False):
+        '''
+        Run the sacct and return all rows. This method is rate limited to 
+        5 calls every 20 seconds.
+
+        Since sacct is inherently asyncronous, pass wait=True to wait until 
+        the job state is updated.
+
+        :param job: Job object
+        :type job: :mod:`executors.models.Job`
+        :param wait: Wait for job state to be updated
+        :type wait: bool
+        '''
+        job_id = job.pid + '.batch'
+        # return whatever sacct output is immediately
+        if not wait:
+            return self._sacct_async(job_id)
+        # wait for the job to appear in sacct or timeout
+        logger.debug('waiting for job %s to appear in sacct', job_id)
+        tic = time.time()
+        while 1:
+            rows = self._sacct_async(job_id)
+            if rows:
+                return rows
+            if time.time() - tic > self.poll_timeout:
+                raise TimeoutError('exceeded wait time {0}s for job {1}'.format(self.poll_timeout, job_id))
+            time.sleep(self.poll_delay)
+
+    @sleep_and_retry
+    @limits(calls=5, period=20)
+    def _sacct_async(self, job_id):
+        '''
+        Run sacct command on a job and serialize output. This method is rate 
+        limited to 5 calls every 20 seconds.
+
+        :param job_id: Slurm job ID
+        :type job_id: str
+        :returns: List of sacct rows
+        :rtype: list
+        '''
+        # build the sacct command
+        if not commons.which('sacct'):
+            raise CommandNotFound('sacct') 
+        cmd = [
+            'sacct',
+            '--parsable2',
+            '--delimiter', ',',
+            '--brief',
+            '--jobs',
+            job_id
+        ]
+        # execute the sacct command, serialize, and return the result
+        logger.debug(cmd)
+        output = sp.check_output(cmd, universal_newlines=True).strip()
+        output = csv.DictReader(io.StringIO(six.u(output)))
+        return [row for row in output]
+
+    def _arguments(self, job):
+        arguments = list()
+        if hasattr(job, 'name') and job.name:
+            arguments.extend(['--job-name', job.name])
+        if hasattr(job, 'memory') and job.memory:
+            arguments.extend(['--mem', job.memory])
+        if hasattr(job, 'cpus') and job.cpus:
+            arguments.extend(['--cpus-per-task', str(job.cpus)])
+        if hasattr(job, 'gpus') and job.gpus:
+            arguments.extend(['--gres', f'gpu:{job.gpus}'])
+        if hasattr(job, 'nodes') and job.nodes:
+            arguments.extend(['--nodes', str(job.nodes)])
+        if hasattr(job, 'output') and job.output:
+            arguments.extend(['--output', os.path.expanduser(job.output)])
+        if hasattr(job, 'error') and job.error:
+            arguments.extend(['--error', os.path.expanduser(job.error)])
+        if hasattr(job, 'time') and job.time:
+            arguments.extend(['--time', str(job.time)])
+        if hasattr(job, 'parent') and job.parent:
+            arguments.extend(['--dependency', 'afterok:{0}'.format(job.parent.pid)])
+        return arguments
+
+class SacctError(Exception):
+    pass
+
+class SlurmNotFound(ExecutorNotFound):
+    pass
+
+class SbatchError(Exception):
+    pass

executors/torque/__init__.py

@@ -0,0 +1,218 @@
+import os
+import re
+import logging
+import subprocess as sp
+from executors.commons import which
+from executors.models import AbstractExecutor
+from executors.exceptions import ExecutorNotFound, CommandNotFound, TimeoutError
+from ratelimit import limits, sleep_and_retry
+
+logger = logging.getLogger(__name__)
+
+class Executor(AbstractExecutor):
+    ACTIVE = (
+        'Q',    # queued
+        'R',    # running
+        'H',    # held
+        'E'     # exited
+    )
+    INACTIVE = (
+        'C',    # complete
+    )
+
+    def __init__(self, partition, **kwargs):
+        if not self.available():
+            raise PBSubmitNotFound()
+        self.partition = partition
+        self.polling_interval = 5
+        self.timeout = 60
+        self._default_args = self.default_args(**kwargs)
+
+    def default_args(self, **kwargs):
+        args = list()
+        for k,v in iter(kwargs.items()):
+            if k == 'nodes':
+                args.extend([
+                    '-l', '+'.join(v)
+                ])
+            else:
+                logger.warn('unrecognized Executor argument "%s"', k)
+        return args
+
+    @staticmethod
+    def available():
+        if which('qsub'):
+            return True
+        return False
+
+    def submit(self, job):
+        prefix = '{0}-%j'.format(job.name) if job.name else '%j'
+        if not job.output:
+            job.output = os.path.expanduser('~/{0}.out'.format(prefix))
+        if not job.error:
+            job.error = os.path.expanduser('~/{0}.err'.format(prefix))
+        command = job.command
+        if isinstance(command, list):
+            command = sp.list2cmdline(command)
+        if not which('qsub'):
+            raise CommandNotFound('qsub')
+
+        cmd = [
+            'qsub',
+            '-q', self.partition,
+            '-d', os.getcwd()
+        ]
+        cmd.extend(self._default_args)
+        cmd.extend(self._arguments(job))
+        job_script = f'#!/bin/bash\n{command}'
+        logger.debug(sp.list2cmdline(cmd))
+        output = sp.check_output(
+            cmd,
+            stderr=sp.STDOUT,
+            input=job_script.encode(),
+        ).decode().strip()
+        pid = re.search(r'^(\d+)', output).group(0)
+        job.pid = pid
+        self._alter_logs(job)
+        logger.debug('parsed job id %s', pid)
+
+    def _alter_logs(self, job):
+        pid = job.pid
+        qalter_args = list()
+        if job.output and '%j' in job.output:
+            output = job.output.replace('%j', pid)
+            qalter_args.extend(['-o', os.path.expanduser(output)])
+        if job.error and '%j' in job.error:
+            error = job.error.replace('%j', pid)
+            qalter_args.extend(['-e', os.path.expanduser(error)])
+        if qalter_args:
+            if not which('qalter'):
+                raise CommandNotFound('qalter')
+            cmd = ['qalter'] + qalter_args + [pid]
+            sp.check_output(cmd)
+
+    def update(self, job, wait=False):
+        try:
+            output = self.qstat(job)
+        except QstatUnknownJobError as e:
+            job.active = False
+            job.returncode = 1
+            return
+        job_state = re.search(r'job_state = (.*)', output).group(1)
+        exit_status = re.search(r'exit_status = (.*)', output)
+        if not exit_status:
+            exit_status = -1
+        else:
+            exit_status = exit_status.group(1)
+        output_path = re.search(r'Output_Path = (.*)', output).group(1)
+        error_path = re.search(r'Error_Path = (.*)', output).group(1)
+        logger.debug('job {0} is in {1} state'.format(job.pid, job_state))
+        if job_state in Executor.ACTIVE:
+            job.active = True
+        elif job_state in Executor.INACTIVE:
+            job.active = False
+            job.returncode = int(exit_status)
+
+    def update_many(self, jobs, wait=False):
+        for job in jobs:
+            self.update(job, wait=wait)
+
+    def cancel(self, job, wait=False):
+        if not which('qdel'):
+            raise CommandNotFound('qdel')
+        cmd = [
+            'qdel',
+            job.pid
+        ]
+        try:
+            logger.debug(cmd)
+            sp.check_output(cmd, stderr=sp.PIPE)
+        except sp.CalledProcessError as e:
+            # qdel will return a 153 exit status if it tries to query the 
+            # state of a Job ID that is already in a 'C' state, or a 170 
+            # exit status if the Job ID is unknown. We should pass on either
+            # of these states. A Job ID can become unknown only minutes after 
+            # a job has entered the 'C' state.
+            if e.returncode == 170:
+                logger.debug('job %s is in a completed state and cannot be cancelled', job.pid)
+                pass
+            elif e.returncode == 153:
+                logger.debug('job %s is unknown and cannot be cancelled', job.pid)
+                pass
+            else:
+                raise e
+
+    @sleep_and_retry
+    @limits(calls=5, period=20)
+    def qstat(self, job):
+        if not which('qstat'):
+            raise CommandNotFound('qstat')
+        cmd = [
+            'qstat',
+            '-f',
+            job.pid
+        ]
+        logger.debug(cmd)
+        try:
+            output = sp.check_output(cmd)
+        except sp.CalledProcessError as e:
+            if e.returncode == 153:
+                logger.debug('job %s is unknown to the scheduler', job.pid)
+                raise QstatUnknownJobError(job)
+            else:
+                raise e
+        return output.decode()
+
+    def _parse_mem_value(self, s):
+        try:
+            match = re.match(r'^(\d+)(K|KB|M|MB|G|GB|T|TB)$', s)
+            size,unit = match.group(1),match.group(2)
+        except:
+            raise IndecipherableMemoryArgument(m)
+        if unit in ('K', 'KB'):
+            unit = 'kb'
+        elif unit in ('M', 'MB'):
+            unit = 'mb'
+        elif unit in ('G', 'GB'):
+            unit = 'gb'
+        elif unit in ('T', 'TB'):
+            unit = 'tb'
+        memarg = size + unit
+        logger.debug('translated memory argument %s', memarg)
+        return size + unit
+
+    def _arguments(self, job):
+        arguments = list()
+        qsub_opts = dict()
+        if hasattr(job, 'output') and job.output:
+            arguments.extend(['-o', os.path.expanduser(job.output)])
+        if hasattr(job, 'error') and job.error:
+            arguments.extend(['-e', os.path.expanduser(job.error)])
+        if hasattr(job, 'parent') and job.parent:
+            arguments.extend(['-W', 'depend=afterok:{0}'.format(job.parent.pid)])
+        if hasattr(job, 'name') and job.name:
+            arguments.extend(['-N', job.name])
+        if hasattr(job, 'memory') and job.memory:
+            qsub_opts['vmem'] = self._parse_mem_value(job.memory)
+        if hasattr(job, 'processors') and job.processors:
+            qsub_opts['ppn'] = job.processors
+        # build and append pass-through qsub options
+        qsub_opts = 'nodes={NODES}:ppn={PPN},vmem={VMEM}'.format(
+            NODES=qsub_opts.get('nodes', 1),
+            PPN=qsub_opts.get('ppn', 1),
+            VMEM=qsub_opts.get('vmem', '1gb')
+        )
+        arguments.extend(['-l', qsub_opts])
+        return arguments
+
+class PBSubmitNotFound(ExecutorNotFound):
+    pass
+
+class IndecipherableMemoryArgument(Exception):
+    pass
+
+class QstatError(Exception):
+    pass
+
+class QstatUnknownJobError(QstatError):
+    pass

executors-0.6.0.dist-info/METADATA

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: executors
+Version: 0.6.0
+Summary: Job scheduling abstractions
+Home-page: https://github.com/harvard-nrg/executors
+Author: Neuroinformatics Research Group
+Author-email: info@neuroinfo.org
+Requires-Dist: paramiko
+Requires-Dist: ratelimit
+Requires-Dist: six
+Dynamic: author
+Dynamic: author-email
+Dynamic: home-page
+Dynamic: requires-dist
+Dynamic: summary

executors-0.6.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+executors/__init__.py,sha256=QhvhwwmmvxPCHepGIebKuzXlYCpd85lujDeQmARTpOk,1853
+executors/__version__.py,sha256=t_YmKUbJUjUJc4cCZ2ICdWjqUyVFGtq--LHZGhr9SJ4,234
+executors/commons/__init__.py,sha256=Cfw3ZiWf8k3VEEAaPMzP0_JdyOzs7ylndXzNp0sGx5E,494
+executors/exceptions/__init__.py,sha256=DmXrC2plEO1sX67p6JQamtuLDO0inNqMbR8ZfJzRP80,129
+executors/local/__init__.py,sha256=g4ZWFAZtt3OrTHBXXNzU2uS_lFpr7nDdoGXErHGjOUY,3176
+executors/lsf/__init__.py,sha256=DkafJvJyB1fTpUvQ-0GfblzO_adJWU8F_uyM2nGMkX4,6039
+executors/models/__init__.py,sha256=QKX4RmtMQ43xIyBp2PwAFatldAf2ASdk8KCtBdIGszo,5596
+executors/pbsubmit/__init__.py,sha256=nBEpCaQfJUhgVNLt8GvX7pAoWof3tBe3c8HeYgm9uNA,8357
+executors/slurm/__init__.py,sha256=nyVFuo48mh9fI8MdtnpHyWgMauXGRsgprIh6DxrbbDo,9736
+executors/torque/__init__.py,sha256=e6d-IyPKVgZWTnomB8XFLXVvuSWHfjem8oh4TNzaX7c,7326
+executors-0.6.0.dist-info/METADATA,sha256=nPc7S8JJKRRjWh3ebOOVXiomy5ZICcm1EP4Xj-V14pc,380
+executors-0.6.0.dist-info/WHEEL,sha256=Mk1ST5gDzEO5il5kYREiBnzzM469m5sI8ESPl7TRhJY,110
+executors-0.6.0.dist-info/top_level.txt,sha256=XjxPatJhppvboRZNbeIzCodjlqjE8Z-FRTeTcmTTiZc,10
+executors-0.6.0.dist-info/RECORD,,

executors-0.6.0.dist-info/WHEEL

@@ -0,0 +1,6 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.10.2)
+Root-Is-Purelib: true
+Tag: py2-none-any
+Tag: py3-none-any
+

executors-0.6.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+executors