Repository: DataBiosphere/toil

Issue:
SLURM integration
Hi,

I was wondering if there is any plan to adapt a SLURM batch system into toil. I have noticed that you evaluated this a while ago in #18 and #19 , and he conclusion was that SLURM does not meet the use case for Toil, that works with scalable systems like amazon EC2.

I've seen though that you guys provide an abstraction of batch system class [here](https://github.com/BD2KGenomics/toil/blob/master/src/toil/batchSystems/abstractBatchSystem.py). And there is for example [an implementation for LSF](https://github.com/BD2KGenomics/toil/blob/master/src/toil/batchSystems/lsf.py).

Unfortunately the documentation on how to implement this class is [uncomplete](http://toil.readthedocs.org/en/latest/batchSystem.html), but having those examples should be possible to do it for SLURM.

Do you guys think its possible and see a fit for this?

Thanks a lot!



Discussion:
I've also heard from @chapmanb that SLURM support would be helpful for running [bcbio-nextgen](https://github.com/chapmanb/bcbio-nextgen) on Toil.

My experience with SLURM has been similar to [this comment](https://github.com/BD2KGenomics/toil/issues/19#issuecomment-89384559) though, where a static cluster is brought up on EC2 with SLURM, Docker, and NFS to run [Nextflow](http://www.nextflow.io/) pipelines.  If I remember correctly, SLURM handles worker nodes dropping out after initialization, but you can't dynamically add new worker nodes.

Relates to #19 

It looks to me as if SLURM has two responsibilities: 1) provisioning of nodes and 2) scheduling of ~~nodes~~ tasks. The abstract batch system API would cover 2) while 1) falls into the realm of the upcoming auto-scaling and provisioning support (issues #349, #600 and PR #541).

We don't have the bandwidth to take this on right now but a PR is always welcome. I would target 2) first. And there would have to be a change to 

https://github.com/BD2KGenomics/cgcloud/blob/master/jenkins/src/cgcloud/jenkins/toil_jenkins_slave.py

for deploying SLURM on a single node for testing such that we can test the SLURM support on continuous integration. 

s/scheduling of nodes/scheduling of tasks/ above

Hannes and Michael;
Thanks so much for the discussion. Our use case is on existing local cluster HPC deployments running SLURM, so we wouldn't need node provisioning or in-flight addition of worker nodes. The goal was to have a framework that can run CWL jobs on existing resources and then have the option to scale out to more elastic environments through your AWS/Azure support. If we were to plug SLURM in through the Batch system API, would that be a good fit?

Yes. When implementing this, the best batch system implementation to use as a model are singleMachine and mesos. If the task submission to SLURM happens via a command line utility as opposed to a library/API, the GridEngine batch system is also worth looking at since it also uses a CLI to submit tasks.

SLURM is very similar in operation to GridEngine.  An alternative to using the command-line approach is to use DRMAA which would, in a sense, align work across batch systems. That said, DRMAA support and availability varies by HPC cluster system.

It would be preferable if there were a batch system implementation that generically works on every DRMAA-compatible backend. Not only would that open Toil up to more backends but using a full-fledged API might be easier and more efficient than using a CLI and parsing program output. I have never heard of DRMAA until today but if you think it is easier to write a batch system that uses it, go for it.

Other systems like Galaxy use DRMAA alongside other cluster support systems:

https://github.com/galaxyproject/galaxy/blob/9264d90ea471214b09666b82781acd49d84f19ec/lib/galaxy/jobs/runners/drmaa.py

Another (perhaps cleaner and newer) abstraction is SAGA:

https://github.com/radical-cybertools/saga-python

In any case, I wouldn't wrap the bare cmdline tools, so it's a matter of choosing the best abstraction/dependency.

Maybe for context, this is all part of the same request. We are using bcbio for the Scottish Genome Partnership (and other projects) on several SLURM HPC installations and want to push the CWL functionality a bit further. A quick poll here shows that none of the sites have DRMAA or SAGA deployed, and I'm _hoping_ we can avoid additional dependencies just as @chapmanb managed to wrap pretty much everything on our side in Conda.

That said, if it's the only reasonable way forward we'll cope ;-)

Fair point, countering my own argument with: [ipython-cluster-helper SLURM code](https://github.com/roryk/ipython-cluster-helper/blob/master/cluster_helper/slurm.py).

That pointer may be very helpful @brainstorm , thanks!

Thanks for the references to `GridEngine` and `singleMachine` as well @hannes-ucsc . I'll try to implement basic SLURM support by looking at these classes and will submit a Pull Request when I'm done so that you can have a look and give me feedback.

Thanks all!


Resolving patch:
diff --git a/src/toil/batchSystems/slurm.py b/src/toil/batchSystems/slurm.py
new file mode 100644
--- /dev/null
+++ b/src/toil/batchSystems/slurm.py
@@ -0,0 +1,397 @@
+# Copyright (c) 2016 Duke Center for Genomic and Computational Biology
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+import logging
+import os
+from pipes import quote
+import subprocess
+import time
+import math
+from Queue import Queue, Empty
+from threading import Thread
+
+from toil.batchSystems.abstractBatchSystem import AbstractBatchSystem
+
+logger = logging.getLogger(__name__)
+
+sleepSeconds = 1
+
+
+class MemoryString:
+    def __init__(self, string):
+        if string[-1] == 'K' or string[-1] == 'M' or string[-1] == 'G' or string[-1] == 'T':
+            self.unit = string[-1]
+            self.val = float(string[:-1])
+        else:
+            self.unit = 'B'
+            self.val = float(string)
+        self.bytes = self.byteVal()
+
+    def __str__(self):
+        if self.unit != 'B':
+            return str(self.val) + self.unit
+        else:
+            return str(self.val)
+
+    def byteVal(self):
+        if self.unit == 'B':
+            return self.val
+        elif self.unit == 'K':
+            return self.val * 1024
+        elif self.unit == 'M':
+            return self.val * 1048576
+        elif self.unit == 'G':
+            return self.val * 1073741824
+        elif self.unit == 'T':
+            return self.val * 1099511627776
+
+    def __cmp__(self, other):
+        return cmp(self.bytes, other.bytes)
+
+
+class Worker(Thread):
+    def __init__(self, newJobsQueue, updatedJobsQueue, killQueue, killedJobsQueue, boss):
+        Thread.__init__(self)
+        self.newJobsQueue = newJobsQueue
+        self.updatedJobsQueue = updatedJobsQueue
+        self.killQueue = killQueue
+        self.killedJobsQueue = killedJobsQueue
+        self.waitingJobs = list()
+        self.runningJobs = set()
+        self.boss = boss
+        self.allocatedCpus = dict()
+        self.slurmJobIDs = dict()
+
+    def parse_elapsed(self, elapsed):
+        # slurm returns elapsed time in days-hours:minutes:seconds format
+        # Sometimes it will only return minutes:seconds, so days may be omitted
+        # For ease of calculating, we'll make sure all the delimeters are ':'
+        # Then reverse the list so that we're always counting up from seconds -> minutes -> hours -> days
+        total_seconds = 0
+        try:
+            elapsed = elapsed.replace('-',':').split(':')
+            elapsed.reverse()
+            seconds_per_unit = [1, 60, 3600, 86400]
+            for index, multiplier in enumerate(seconds_per_unit):
+                if index < len(elapsed):
+                    total_seconds += multiplier * int(elapsed[index])
+        except ValueError:
+            pass # slurm may return INVALID instead of a time
+        return total_seconds
+
+
+    def getRunningJobIDs(self):
+        # Should return a dictionary of Job IDs and number of seconds
+        times = {}
+        currentjobs = dict((str(self.slurmJobIDs[x]), x) for x in self.runningJobs)
+        # currentjobs is a dictionary that maps a slurm job id (string) to our own internal job id
+        # squeue arguments:
+        # -h for no header
+        # --format to get jobid i, state %t and time days-hours:minutes:seconds
+
+        lines = subprocess.check_output(['squeue', '-h', '--format', '%i %t %M']).split('\n')
+        for line in lines:
+            values = line.split()
+            if len(values) < 3:
+                continue
+            slurm_jobid, state, elapsed_time = values
+            if slurm_jobid in currentjobs and state == 'R':
+                seconds_running = self.parse_elapsed(elapsed_time)
+                times[currentjobs[slurm_jobid]] = seconds_running
+
+        return times
+
+    def getSlurmID(self, jobID):
+        if not jobID in self.slurmJobIDs:
+            RuntimeError("Unknown jobID, could not be converted")
+
+        job = self.slurmJobIDs[jobID]
+        return str(job)
+
+    def forgetJob(self, jobID):
+        self.runningJobs.remove(jobID)
+        del self.allocatedCpus[jobID]
+        del self.slurmJobIDs[jobID]
+
+    def killJobs(self):
+        # Load hit list:
+        killList = list()
+        while True:
+            try:
+                jobId = self.killQueue.get(block=False)
+            except Empty:
+                break
+            else:
+                killList.append(jobId)
+
+        if not killList:
+            return False
+
+        # Do the dirty job
+        for jobID in list(killList):
+            if jobID in self.runningJobs:
+                logger.debug('Killing job: %s', jobID)
+                subprocess.check_call(['scancel', self.getSlurmID(jobID)])
+            else:
+                if jobID in self.waitingJobs:
+                    self.waitingJobs.remove(jobID)
+                self.killedJobsQueue.put(jobID)
+                killList.remove(jobID)
+
+        # Wait to confirm the kill
+        while killList:
+            for jobID in list(killList):
+                if self.getJobExitCode(self.slurmJobIDs[jobID]) is not None:
+                    logger.debug('Adding jobID %s to killedJobsQueue', jobID)
+                    self.killedJobsQueue.put(jobID)
+                    killList.remove(jobID)
+                    self.forgetJob(jobID)
+            if len(killList) > 0:
+                logger.warn("Some jobs weren't killed, trying again in %is.", sleepSeconds)
+                time.sleep(sleepSeconds)
+
+        return True
+
+    def createJobs(self, newJob):
+        activity = False
+        # Load new job id if present:
+        if newJob is not None:
+            self.waitingJobs.append(newJob)
+        # Launch jobs as necessary:
+        while len(self.waitingJobs) > 0 and sum(self.allocatedCpus.values()) < int(
+                self.boss.maxCores):
+            activity = True
+            jobID, cpu, memory, command = self.waitingJobs.pop(0)
+            sbatch_line = self.prepareSbatch(cpu, memory, jobID) + ['--wrap={}'.format(command)]
+            slurmJobID = self.sbatch(sbatch_line)
+            self.slurmJobIDs[jobID] = slurmJobID
+            self.runningJobs.add(jobID)
+            self.allocatedCpus[jobID] = cpu
+        return activity
+
+    def checkOnJobs(self):
+        activity = False
+        logger.debug('List of running jobs: %r', self.runningJobs)
+        for jobID in list(self.runningJobs):
+            logger.debug("Checking status of internal job id %d", jobID)
+            status = self.getJobExitCode(self.slurmJobIDs[jobID])
+            if status is not None:
+                activity = True
+                self.updatedJobsQueue.put((jobID, status))
+                self.forgetJob(jobID)
+        return activity
+
+    def run(self):
+        while True:
+            activity = False
+            newJob = None
+            if not self.newJobsQueue.empty():
+                activity = True
+                newJob = self.newJobsQueue.get()
+                if newJob is None:
+                    logger.debug('Received queue sentinel.')
+                    break
+            activity |= self.killJobs()
+            activity |= self.createJobs(newJob)
+            activity |= self.checkOnJobs()
+            if not activity:
+                logger.debug('No activity, sleeping for %is', sleepSeconds)
+                time.sleep(sleepSeconds)
+
+    def prepareSbatch(self, cpu, mem, jobID):
+        #  Returns the sbatch command line before the script to run
+        sbatch_line = ['sbatch', '-Q', '-J', 'toil_job_{}'.format(jobID)]
+
+        if self.boss.environment:
+            for k, v in self.boss.environment.iteritems():
+                quoted_value = quote(os.environ[k] if v is None else v)
+                sbatch_line.append('--export={}={}'.format(k, quoted_value))
+
+        if mem is not None:
+            # memory passed in is in bytes, but slurm expects megabytes
+            sbatch_line.append('--mem={}'.format(int(mem)/2**20))
+        if cpu is not None:
+            sbatch_line.append('--cpus-per-task={}'.format(int(math.ceil(cpu))))
+
+        return sbatch_line
+
+    def sbatch(self, sbatch_line):
+        logger.debug("Running %r", sbatch_line)
+        try:
+            output = subprocess.check_output(sbatch_line, stderr=subprocess.STDOUT)
+            # sbatch prints a line like 'Submitted batch job 2954103'
+            result = int(output.strip().split()[-1])
+            logger.debug("sbatch submitted job %d", result)
+            return result
+        except subprocess.CalledProcessError as e:
+            logger.error("sbatch command failed with code %d: %s", e.returncode, e.output)
+            raise e
+        except OSError as e:
+            logger.error("sbatch command failed")
+            raise e
+
+    def getJobExitCode(self, slurmJobID):
+        logger.debug("Getting exit code for slurm job %d", slurmJobID)
+        # SLURM job exit codes are obtained by running sacct.
+        # sacct returns
+        # -n : no header
+        # -j : job
+        # --format : specify output columns
+        # -P : separate columns with pipes
+        # -S 1970-01-01 override start time limit
+        args = ['sacct', '-n', '-j', str(slurmJobID), '--format','State,ExitCode', '-P', '-S', '1970-01-01']
+        process = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+        for line in process.stdout:
+            values = line.strip().split('|')
+            if len(values) < 2:
+                continue
+            state, exitcode = values
+            logger.debug("sacct Job state is %s", state)
+            # If Job is in a running state, return None to indicate we don't have an update
+            if state in ('PENDING', 'RUNNING','CONFIGURING','COMPLETING','RESIZING','SUSPENDED'):
+                return None
+            status, _ = exitcode.split(':')
+            logger.debug("sacct exitcode is %s, returning status %s", exitcode, status )
+            return int(status)
+        logger.debug("Did not find exit code for job in sacct output")
+        return None
+
+
+class SlurmBatchSystem(AbstractBatchSystem):
+    """
+    The interface for SLURM
+    """
+
+    def __init__(self, config, maxCores, maxMemory, maxDisk):
+        AbstractBatchSystem.__init__(self, config, maxCores, maxMemory, maxDisk)
+        self.slurmResultsFile = self._getResultsFileName(config.jobStore)
+        # Reset the job queue and results (initially, we do this again once we've killed the jobs)
+        self.slurmResultsFileHandle = open(self.slurmResultsFile, 'w')
+        # We lose any previous state in this file, and ensure the files existence
+        self.slurmResultsFileHandle.close()
+        self.currentJobs = set()
+        self.maxCPU, self.maxMEM = self.obtainSystemConstants()
+        self.nextJobID = 0
+        self.newJobsQueue = Queue()
+        self.updatedJobsQueue = Queue()
+        self.killQueue = Queue()
+        self.killedJobsQueue = Queue()
+        self.worker = Worker(self.newJobsQueue, self.updatedJobsQueue, self.killQueue,
+                             self.killedJobsQueue, self)
+        self.worker.start()
+
+    def __des__(self):
+        # Closes the file handle associated with the results file.
+        self.slurmResultsFileHandle.close()
+
+    def issueBatchJob(self, command, memory, cores, disk):
+        self.checkResourceRequest(memory, cores, disk)
+        jobID = self.nextJobID
+        self.nextJobID += 1
+        self.currentJobs.add(jobID)
+        self.newJobsQueue.put((jobID, cores, memory, command))
+        logger.debug("Issued the job command: %s with job id: %s ", command, str(jobID))
+        return jobID
+
+    def killBatchJobs(self, jobIDs):
+        """
+        Kills the given jobs, represented as Job ids, then checks they are dead by checking
+        they are not in the list of issued jobs.
+        """
+        jobIDs = set(jobIDs)
+        logger.debug('Jobs to be killed: %r', jobIDs)
+        for jobID in jobIDs:
+            self.killQueue.put(jobID)
+        while jobIDs:
+            killedJobId = self.killedJobsQueue.get()
+            if killedJobId is None:
+                break
+            jobIDs.remove(killedJobId)
+            if killedJobId in self.currentJobs:
+                self.currentJobs.remove(killedJobId)
+            if jobIDs:
+                logger.debug('Some kills (%s) still pending, sleeping %is', len(jobIDs),
+                             sleepSeconds)
+                time.sleep(sleepSeconds)
+
+    def getIssuedBatchJobIDs(self):
+        """
+        Gets the list of jobs issued to SLURM.
+        """
+        return list(self.currentJobs)
+
+    def getRunningBatchJobIDs(self):
+        return self.worker.getRunningJobIDs()
+
+    def getUpdatedBatchJob(self, maxWait):
+        try:
+            i = self.updatedJobsQueue.get(timeout=maxWait)
+        except Empty:
+            return None
+        logger.debug('UpdatedJobsQueue Item: %s', i)
+        jobID, retcode = i
+        self.updatedJobsQueue.task_done()
+        self.currentJobs.remove(jobID)
+        return i
+
+    def shutdown(self):
+        """
+        Signals worker to shutdown (via sentinel) then cleanly joins the thread
+        """
+        newJobsQueue = self.newJobsQueue
+        self.newJobsQueue = None
+
+        newJobsQueue.put(None)
+        self.worker.join()
+
+    def getWaitDuration(self):
+        return 1.0
+
+    @classmethod
+    def getRescueBatchJobFrequency(cls):
+        return 1800  # Half an hour
+
+    @staticmethod
+    def obtainSystemConstants():
+        # sinfo -Ne --format '%m,%c'
+        # sinfo arguments:
+        # -N for node-oriented
+        # -h for no header
+        # -e for exact values (e.g. don't return 32+)
+        # --format to get memory, cpu
+        max_cpu = 0
+        max_mem = MemoryString('0')
+        lines = subprocess.check_output(['sinfo', '-Nhe', '--format', '%m %c']).split('\n')
+        for line in lines:
+            values = line.split()
+            if len(values) < 2:
+                continue
+            mem, cpu = values
+            max_cpu = max(max_cpu, int(cpu))
+            max_mem = max(max_mem, MemoryString(mem + 'M'))
+        if max_cpu == 0 or max_mem.byteVal() ==  0:
+            RuntimeError('sinfo did not return memory or cpu info')
+        return max_cpu, max_mem
+
+
+    def setEnv(self, name, value=None):
+        # if value and ',' in value:
+        #     raise ValueError("GridEngine does not support commata in environment variable values")
+        return AbstractBatchSystem.setEnv(self, name, value)
+
+    @staticmethod
+    def supportsWorkerCleanup():
+        return False
+
diff --git a/src/toil/common.py b/src/toil/common.py
--- a/src/toil/common.py
+++ b/src/toil/common.py
@@ -547,6 +547,10 @@ def createBatchSystem(config, jobStore=None, userScript=None):
 
             kwargs['masterAddress'] = config.mesosMasterAddress
 
+        elif config.batchSystem == 'slurm' or config.batchSystem == 'Slurm':
+            from toil.batchSystems.slurm import SlurmBatchSystem
+            batchSystemClass = SlurmBatchSystem
+
         else:
             raise RuntimeError('Unrecognised batch system: %s' % config.batchSystem)
 


Test patch:
diff --git a/src/toil/test/__init__.py b/src/toil/test/__init__.py
--- a/src/toil/test/__init__.py
+++ b/src/toil/test/__init__.py
@@ -227,6 +227,22 @@ def needs_parasol(test_item):
         return test_item
 
 
+def needs_slurm(test_item):
+    """
+    Use as a decorator before test classes or methods to only run them if Slurm is installed.
+    """
+    test_item = _mark_test('slurm', test_item)
+    try:
+        with open(os.devnull, 'r+') as devnull:
+            subprocess.Popen('squeue', stdout=devnull, stderr=devnull, stdin=devnull)
+    except OSError:
+        return unittest.skip("Skipping test. Install Slurm to include this test.")(test_item)
+    except:
+        raise
+    else:
+        return test_item
+
+
 def needs_encryption(test_item):
     """
     Use as a decorator before test classes or methods to only run them if PyNaCl is installed
diff --git a/src/toil/test/batchSystems/batchSystemTest.py b/src/toil/test/batchSystems/batchSystemTest.py
--- a/src/toil/test/batchSystems/batchSystemTest.py
+++ b/src/toil/test/batchSystems/batchSystemTest.py
@@ -34,7 +34,7 @@
 from toil.batchSystems.singleMachine import SingleMachineBatchSystem
 from toil.batchSystems.abstractBatchSystem import InsufficientSystemResources
 from toil.job import Job
-from toil.test import ToilTest, needs_mesos, needs_parasol, needs_gridengine
+from toil.test import ToilTest, needs_mesos, needs_parasol, needs_gridengine, needs_slurm
 
 log = logging.getLogger(__name__)
 
@@ -468,6 +468,36 @@ def setUpClass(cls):
         logging.basicConfig(level=logging.DEBUG)
 
 
+@needs_slurm
+class SlurmBatchSystemTest(hidden.AbstractBatchSystemTest):
+    """
+    Tests against the Slurm batch system
+    """
+
+    def _createDummyConfig(self):
+        config = super(SlurmBatchSystemTest, self)._createDummyConfig()
+        # can't use _getTestJobStorePath since that method removes the directory
+        config.jobStore = self._createTempDir('jobStore')
+        return config
+
+    def createBatchSystem(self):
+        from toil.batchSystems.slurm import SlurmBatchSystem
+        return SlurmBatchSystem(config=self.config, maxCores=numCores, maxMemory=1000e9,
+                                maxDisk=1e9)
+
+    @classmethod
+    def setUpClass(cls):
+        super(SlurmBatchSystemTest, cls).setUpClass()
+        logging.basicConfig(level=logging.DEBUG)
+
+    def tearDown(self):
+        super(SlurmBatchSystemTest, self).tearDown()
+        # Cleanup 'slurm-%j.out' produced by sbatch
+        from glob import glob
+        for f in glob('slurm-*.out'):
+            os.unlink(f)
+
+
 @contextmanager
 def tempFileContaining(content):
     fd, path = tempfile.mkstemp(suffix='.py')