Add a local and ssh scheduler definitions

Author: arteymix

bioluigi/config.py

@@ -8,6 +8,11 @@ class bioluigi(luigi.Config):
     scheduler_extra_args = luigi.ListParameter(default=[],
                                                description='List of extra arguments to pass to the scheduler')
     slurm_srun_bin = luigi.Parameter(default='srun')
+    ssh_bin = luigi.Parameter(default='ssh')
+    ssh_remote = luigi.OptionalParameter(default=None)
+    ssh_port = luigi.OptionalIntParameter(default=None)
+    ssh_user = luigi.OptionalParameter(default=None)
+    ssh_identity_file = luigi.OptionalParameter(default=None)
     prefetch_bin = luigi.Parameter(default='prefetch')
     fastqdump_bin = luigi.Parameter(default='fastq-dump')
     cutadapt_bin = luigi.Parameter(default='cutadapt')

bioluigi/scheduled_external_program.py

@@ -2,30 +2,30 @@
 This module provides a flavour of :class:`luigi.contrib.external_program.ExternalProgramTask`
 designed to work with an external scheduler such as Slurm.
 
-It also extend the basic definition of an external task to declare basic
+It also extends the basic definition of an external task to declare basic
 resource consumptions such as the number of CPU and the amount of memory
 required to execute the task.
 """
 
 import datetime
+import os
 from datetime import timedelta
 from typing import Optional
 
 import luigi
-from luigi.contrib.external_program import ExternalProgramTask
 
 from .config import bioluigi
 from .schedulers import get_available_schedulers, get_scheduler, ScheduledTask
 
 cfg = bioluigi()
 
-class ScheduledExternalProgramTask(ExternalProgramTask, ScheduledTask):
+class ScheduledExternalProgramTask(luigi.Task, ScheduledTask):
     """
     Variant of :class:`luigi.contrib.external_program.ExternalProgramTask` that
     executes the task with a :class:`Scheduler`.
     """
     scheduler: str = luigi.ChoiceParameter(default=cfg.scheduler,
-                                           choices=['local'] + [blurb for blurb in get_available_schedulers()],
+                                           choices=[blurb for blurb in get_available_schedulers()],
                                            positional=False, significant=False,
                                            description='Scheduler to use for running the task')
     scheduler_partition: Optional[str] = luigi.OptionalParameter(default=cfg.scheduler_partition, positional=False,
@@ -42,21 +42,17 @@ class ScheduledExternalProgramTask(ExternalProgramTask, ScheduledTask):
     memory: float = luigi.FloatParameter(default=1, positional=False, significant=False,
                                          description='Amount of memory (in gigabyte) to allocate for the task')
 
+    working_directory = luigi.OptionalParameter(default=None, significant=False, positional=False)
+
+    capture_output = luigi.BoolParameter(default=True, significant=False, positional=False)
+
+    def program_environment(self) -> dict[str, str]:
+        env = os.environ.copy()
+        return env
+
     @property
     def resources(self):
-        if self.scheduler == 'local':
-            # local_jobs is actually constrained by the number of workers
-            return {'cpus': self.cpus, 'memory': self.memory}
-        else:
-            return {'{}_jobs'.format(self.scheduler): 1, '{}_cpus'.format(self.scheduler): self.cpus}
+        return get_scheduler(self.scheduler).get_resources_for_task(self)
 
     def run(self):
-        if self.scheduler == 'local':
-            super().run()
-        else:
-            try:
-                _scheduler = get_scheduler(self.scheduler)
-            except KeyError:
-                raise ValueError('Unsupported scheduler {}'.format(self.scheduler))
-            _scheduler.run_task(self)
-
+        get_scheduler(self.scheduler).run_task(self)

bioluigi/schedulers.py

@@ -1,11 +1,11 @@
 import json
 import logging
+import subprocess
 from abc import abstractmethod, ABC
 from datetime import timedelta
-from subprocess import Popen, PIPE
+from subprocess import Popen
 from typing import Optional
 
-import luigi
 from luigi.contrib.external_program import ExternalProgramRunError, ExternalProgramRunContext
 
 from .config import bioluigi
@@ -16,7 +16,7 @@
 
 _schedulers = {}
 
-class ScheduledTask(luigi.Task, ABC):
+class ScheduledTask(ABC):
     """Interface for Luigi tasks that can be scheduled."""
     walltime: timedelta
 
@@ -28,6 +28,8 @@ class ScheduledTask(luigi.Task, ABC):
 
     capture_output: bool
 
+    working_directory: Optional[str]
+
     @abstractmethod
     def program_environment(self) -> dict[str, str]:
         pass
@@ -40,7 +42,10 @@ def get_available_schedulers():
     return _schedulers.keys()
 
 def get_scheduler(blurb):
-    return _schedulers[blurb]
+    try:
+        return _schedulers[blurb]
+    except KeyError:
+        raise ValueError('Unsupported scheduler {}'.format(self.scheduler))
 
 def register_scheduler(blurb):
     """
@@ -56,9 +61,34 @@ def wrapper(cls):
 
     return wrapper
 
+def _run_command(args, env, cwd, capture_output):
+    kwargs = {'env': env, 'universal_newlines': True}
+    if capture_output:
+        kwargs['stdout'] = subprocess.PIPE
+        kwargs['stderr'] = subprocess.PIPE
+    if cwd:
+        kwargs['cwd'] = cwd
+    logger.info('Running command %s%s', ' '.join(args), 'in ' + cwd if cwd else '')
+    proc = Popen(args, **kwargs)
+    with ExternalProgramRunContext(proc):
+        if capture_output:
+            stderr, stdout = proc.communicate()
+            if stdout:
+                logger.info('Program stdout:\n{}'.format(stdout))
+            if stderr:
+                logger.info('Program stderr:\n{}'.format(stderr))
+        else:
+            proc.wait()
+    if proc.returncode != 0:
+        raise ExternalProgramRunError('Program exited with non-zero return code.', tuple(args), env, stdout, stderr)
+
 class Scheduler(ABC):
     @abstractmethod
-    def run_task(self, task):
+    def get_resources_for_task(self, task: ScheduledTask) -> dict[str, int]:
+        pass
+
+    @abstractmethod
+    def run_task(self, task: ScheduledTask):
         pass
 
 @register_scheduler('slurm')
@@ -67,10 +97,13 @@ class SlurmScheduler(Scheduler):
     Scheduler based on Slurm https://slurm.schedmd.com/
     """
 
+    def get_resources_for_task(self, task: ScheduledTask):
+        return {'slurm_jobs': 1, 'slurm_cpus': task.cpus}
+
     def run_task(self, task: ScheduledTask):
         secs = int(task.walltime.total_seconds())
-        srun_args = [cfg.slurm_srun_bin]
-        srun_args.extend([
+        args = [cfg.slurm_srun_bin]
+        args.extend([
             '--verbose',
             '--job-name', task.get_task_family(),
             '--comment', json.dumps({'task_id': task.task_id, 'priority': task.priority}),
@@ -79,19 +112,47 @@ def run_task(self, task: ScheduledTask):
             '--mem', '{}G'.format(int(task.memory)),
             '--cpus-per-task', str(task.cpus)])
         if task.scheduler_partition:
-            srun_args.extend(['--partition', task.scheduler_partition])
+            args.extend(['--partition', task.scheduler_partition])
+        if task.working_directory:
+            args.extend(['--chdir', task.working_directory])
         # FIXME: task.priority is not reliable and does not reflect what the
         # scheduler
         # TODO: srun_args.extend([--priority', str(max(0, cfg.scheduler_priority))])
-        srun_args.extend(map(str, task.scheduler_extra_args))
-        args = list(map(str, task.program_args()))
-        env = task.program_environment()
-        logger.info('Running Slurm command {}'.format(' '.join(srun_args + args)))
-        proc = Popen(srun_args + args, env=env, stdout=PIPE, stderr=PIPE, universal_newlines=True)
-        with ExternalProgramRunContext(proc):
-            stdout, stderr = proc.communicate()
-        if proc.returncode != 0:
-            raise ExternalProgramRunError('Program exited with non-zero return code.', tuple(args), env, stdout, stderr)
-        if task.capture_output:
-            logger.info('Program stdout:\n{}'.format(stdout))
-            logger.info('Program stderr:\n{}'.format(stderr))
+        args.extend(map(str, task.scheduler_extra_args))
+        args.extend(map(str, task.program_args()))
+        _run_command(args, env=(task.program_environment()), cwd=task.working_directory,
+                     capture_output=task.capture_output)
+
+@register_scheduler('local')
+class LocalScheduler(Scheduler):
+    """A scheduler that uses a local subprocess"""
+
+    def get_resources_for_task(self, task: ScheduledTask):
+        return {'cpus': task.cpus, 'memory': task.memory}
+
+    def run_task(self, task: ScheduledTask):
+        _run_command(list(map(str, task.program_args())), env=task.program_environment(), cwd=task.working_directory,
+                     capture_output=task.capture_output)
+
+@register_scheduler('ssh')
+class SshScheduler(Scheduler):
+    """A scheduler that uses SSH to run a task remotely"""
+
+    def get_resources_for_task(self, task: ScheduledTask):
+        return {'ssh_cpus': task.cpus, 'ssh_memory': task.memory}
+
+    def run_task(self, task: ScheduledTask):
+        if not cfg.ssh_remote:
+            raise ValueError('No SSH remote is configured.')
+        args = [cfg.ssh_bin]
+        if cfg.ssh_port:
+            args.extend(['-p', str(cfg.ssh_port)])
+        if cfg.ssh_user:
+            args.extend(['-u', cfg.ssh_user])
+        if cfg.ssh_identity_file:
+            args.extend(['-i', cfg.ssh_identity_file])
+        args.extend(task.scheduler_extra_args)
+        args.append(cfg.ssh_remote)
+        args.extend(map(str, task.program_args()))
+        _run_command(args, env=task.program_environment(), cwd=task.working_directory,
+                     capture_output=task.capture_output)

example.luigi.cfg

@@ -15,6 +15,10 @@ slurm_jobs=1024
 # prevent workers from needlessly waiting on pending jobs if all the nodes are
 # busy.
 slurm_cpus=396
+# For the SSH scheduler, the number of jobs, CPUs and memory are managed per-host
+ssh_cpus=8
+# In gigabytes
+ssh_memory=32
 # Limit the number of I/O-bound jobs
 # The following are considered I/O-bound: prefetch, fastq-dump, fastqc,
 # bamtofastq
@@ -39,11 +43,13 @@ cellranger_bamtofastq_jobs=10
 [bioluigi]
 scheduler=local
 scheduler_extra_args=[]
+# Remote to use for the SSH scheduler
+ssh_remote=
 cutadapt_bin=cutadapt
 star_bin=STAR
 rsem_dir=
 bcftools_bin=bcftools
 vep_bin=vep
 # If this is not specified, it will default to $HOME/.vep
 vep_dir=
-cellranger_bin=cellranger
+cellranger_bin=cellranger

luigi.cfg

@@ -0,0 +1 @@
+example.luigi.cfg

tests/test_scheduled_external_program.py

@@ -5,6 +5,7 @@
 import pytest
 
 from bioluigi.scheduled_external_program import ScheduledExternalProgramTask
+from bioluigi.schedulers import get_available_schedulers
 
 class MyTask(ScheduledExternalProgramTask):
     cpus = 4
@@ -35,19 +36,30 @@ def test_default_scheduler():
     assert task.complete()
 
 def test_local_scheduler():
+    assert 'local' in get_available_schedulers()
     task = MyTask("2", scheduler='local')
+    assert task.resources['cpus'] == 4
+    assert task.resources['memory'] == 1
     assert not task.complete()
     luigi.build([task], local_scheduler=True)
     assert task.complete()
 
 def test_slurm_scheduler():
-    if which('srun') is None:
-        pytest.skip('srun is needed to run Slurm tests.')
+    assert 'slurm' in get_available_schedulers()
     task = MyTask("3", scheduler='slurm')
     assert 'slurm_jobs' in task.resources
     assert 'slurm_cpus' in task.resources
     assert task.resources['slurm_jobs'] == 1
     assert task.resources['slurm_cpus'] == 4
     assert not task.complete()
+    if which('srun') is None:
+        pytest.skip('srun is needed to run Slurm scheduler tests.')
     luigi.build([task], local_scheduler=True)
     assert task.complete()
+
+def test_ssh_scheduler():
+    task = MyTask('5', scheduler='ssh')
+    assert task.resources == {'ssh_cpus': 4, 'ssh_memory': 1}
+    # FIXME:
+    # luigi.build([task], local_scheduler=True)
+    # assert task.complete()