Repository: acorg/slurm-pipeline

Issue:
Add a command line option to have job launching sleep between steps
To give a distributed filesystem time to settle. Without this it appears that a job may be launched even when `sbatch` is run with a dependency if the dependent job is not yet apparent (due to filesystem consistency slowness).

Discussion:

Patch:
diff --git a/bin/slurm-pipeline.py b/bin/slurm-pipeline.py
index 971d493..ea55018 100755
--- a/bin/slurm-pipeline.py
+++ b/bin/slurm-pipeline.py
@@ -33,11 +33,19 @@ parser.add_argument(
           'help text for --first-step for how this affects the environment '
           'in which step scripts are executed.'))
 
+parser.add_argument(
+    '--sleep', type=float, default=0.0,
+    help=('Specify the (floating point) number of seconds to sleep for '
+          'between running step scripts. This can be used to allow a '
+          'distributed file system to settle, so that jobs that have been '
+          'scheduled can be seen when used as dependencies in later '
+          'invocations of sbatch.'))
+
 args, scriptArgs = parser.parse_known_args()
 
 sp = SlurmPipeline(args.specification, force=args.force,
                    firstStep=args.firstStep, lastStep=args.lastStep,
-                   scriptArgs=scriptArgs)
+                   sleep=args.sleep, scriptArgs=scriptArgs)
 sp.schedule()
 
 print(sp.toJSON())
diff --git a/setup.py b/setup.py
index a0032f6..c5558ca 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 from setuptools import setup
 
 setup(name='slurm-pipeline',
-      version='1.0.8',
+      version='1.0.9',
       packages=['slurm_pipeline'],
       include_package_data=True,
       url='https://github.com/acorg/slurm-pipeline',
diff --git a/slurm_pipeline/pipeline.py b/slurm_pipeline/pipeline.py
index 196c42f..5353c2f 100644
--- a/slurm_pipeline/pipeline.py
+++ b/slurm_pipeline/pipeline.py
@@ -1,7 +1,7 @@
 import os
 from os import path, environ
 import re
-from time import time
+import time
 from six import string_types
 from json import load, dumps
 import subprocess
@@ -45,6 +45,11 @@ class SlurmPipeline(object):
     @param lastStep: If not C{None}, the name of the last specification step
           to execute. See above docs for C{firstStep} for how this affects the
           calling of step scripts.
+    @param sleep: Gives the C{float} number of seconds to sleep for between
+          running step scripts. This can be used to allow a distributed file
+          system to settle, so that jobs that have been scheduled can be seen
+          when used as dependencies in later invocations of sbatch. Pass 0.0
+          for no sleep.
     @param scriptArgs: A C{list} of C{str} arguments that should be put on the
         command line of all steps that have no dependencies.
     """
@@ -56,7 +61,7 @@ class SlurmPipeline(object):
     TASK_NAME_LINE = re.compile('^TASK:\s+(\S+)\s*')
 
     def __init__(self, specification, force=False, firstStep=None,
-                 lastStep=None, scriptArgs=None):
+                 lastStep=None, sleep=0.0, scriptArgs=None):
         if isinstance(specification, string_types):
             specification = self._loadSpecification(specification)
         # self.steps will be keyed by step name, with values that are
@@ -68,6 +73,7 @@ class SlurmPipeline(object):
         self.lastStep = lastStep
         self._checkSpecification(specification)
         self.specification = specification
+        self._sleep = sleep
         self._scriptArgs = scriptArgs
         self._scriptArgsStr = (
             ' '.join(map(str, scriptArgs)) if scriptArgs else '')
@@ -83,11 +89,12 @@ class SlurmPipeline(object):
         if 'scheduledAt' in self.specification:
             raise SchedulingError('Specification has already been scheduled')
         else:
-            self.specification['scheduledAt'] = time()
+            self.specification['scheduledAt'] = time.time()
 
             firstStepFound = lastStepFound = False
+            nSteps = len(self.specification['steps'])
 
-            for step in self.specification['steps']:
+            for stepIndex, step in enumerate(self.specification['steps']):
                 if self.firstStep is not None:
                     if firstStepFound:
                         if self.lastStep is not None:
@@ -112,6 +119,11 @@ class SlurmPipeline(object):
 
                 self._scheduleStep(step, simulate)
 
+                # If we're supposed to pause between scheduling steps and
+                # this is not the last step, then sleep.
+                if self._sleep > 0.0 and stepIndex < nSteps - 1:
+                    time.sleep(self._sleep)
+
     def _scheduleStep(self, step, simulate):
         """
         Schedule a single execution step.
@@ -184,7 +196,7 @@ class SlurmPipeline(object):
             env.pop('SP_DEPENDENCY_ARG', None)
             self._runStepScript(step, args, env)
 
-        step['scheduledAt'] = time()
+        step['scheduledAt'] = time.time()
 
     def _runStepScript(self, step, args, env):
         """


Test patch:
diff --git a/test/test_pipeline.py b/test/test_pipeline.py
index 5a99a3e..b4e200c 100644
--- a/test/test_pipeline.py
+++ b/test/test_pipeline.py
@@ -421,14 +421,14 @@ class TestRunner(TestCase):
         with patch.object(subprocess, 'check_output',
                           return_value='Submitted batch job 4\n'):
             runner = SlurmPipeline(
-                    {
-                        'steps': [
-                            {
-                                'name': 'name1',
-                                'script': 'script1',
-                            },
-                        ],
-                    })
+                {
+                    'steps': [
+                        {
+                            'name': 'name1',
+                            'script': 'script1',
+                        },
+                    ],
+                })
             runner.schedule()
             error = '^Specification has already been scheduled$'
             assertRaisesRegex(self, SchedulingError, error,
@@ -1157,3 +1157,98 @@ class TestRunner(TestCase):
 
             env3 = mockMethod.mock_calls[2][2]['env']
             self.assertEqual('1', env3['SP_SIMULATE'])
+
+    @patch('subprocess.check_output')
+    @patch('time.sleep')
+    @patch('os.access')
+    @patch('os.path.exists')
+    def testSleep(self, existsMock, accessMock, sleepMock, subprocessMock):
+        """
+        If a sleep argument is given to SlurmPipeline, sleep must be called
+        between steps with the expected number of seconds.
+        """
+        runner = SlurmPipeline(
+            {
+                'steps': [
+                    {
+                        'name': 'name1',
+                        'script': 'script1',
+                    },
+                    {
+                        'name': 'name2',
+                        'script': 'script2',
+                    },
+                    {
+                        'dependencies': ['name1'],
+                        'name': 'name3',
+                        'script': 'script3',
+                    },
+                ],
+            }, sleep=1.0)
+        runner.schedule()
+
+        sleepMock.assert_has_calls([call(1.0), call(1.0)])
+
+    @patch('subprocess.check_output')
+    @patch('time.sleep')
+    @patch('os.access')
+    @patch('os.path.exists')
+    def testSleepNotCalledByDefault(self, existsMock, accessMock, sleepMock,
+                                    subprocessMock):
+        """
+        If no sleep argument is given to SlurmPipeline, sleep must not be
+        called.
+        """
+        runner = SlurmPipeline(
+            {
+                'steps': [
+                    {
+                        'name': 'name1',
+                        'script': 'script1',
+                    },
+                    {
+                        'name': 'name2',
+                        'script': 'script2',
+                    },
+                    {
+                        'dependencies': ['name1'],
+                        'name': 'name3',
+                        'script': 'script3',
+                    },
+                ],
+            })
+        runner.schedule()
+
+        self.assertFalse(sleepMock.called)
+
+    @patch('subprocess.check_output')
+    @patch('time.sleep')
+    @patch('os.access')
+    @patch('os.path.exists')
+    def testSleepNotCalledWhenZero(self, existsMock, accessMock, sleepMock,
+                                   subprocessMock):
+        """
+        If a sleep argument of 0.0 is given to SlurmPipeline, sleep must not be
+        called.
+        """
+        runner = SlurmPipeline(
+            {
+                'steps': [
+                    {
+                        'name': 'name1',
+                        'script': 'script1',
+                    },
+                    {
+                        'name': 'name2',
+                        'script': 'script2',
+                    },
+                    {
+                        'dependencies': ['name1'],
+                        'name': 'name3',
+                        'script': 'script3',
+                    },
+                ],
+            }, sleep=0.0)
+        runner.schedule()
+
+        self.assertFalse(sleepMock.called)