Repository: acorg/slurm-pipeline

Issue:
Add --nice command line option
To allow setting of priority.


Discussion:

Patch:
diff --git a/Makefile b/Makefile
index 99c405f..08eef5b 100644
--- a/Makefile
+++ b/Makefile
@@ -1,4 +1,4 @@
-.PHONY: check, tcheck, pep8, pyflakes, lint, wc, clean, clobber, upload
+.PHONY: check, tcheck, examples-test, pep8, pyflakes, lint, wc, clean, clobber, upload
 
 check:
 	python -m discover -v
@@ -6,6 +6,12 @@ check:
 tcheck:
 	trial --rterrors test
 
+examples-test:
+	make -C examples/word-count run clean
+	make -C examples/word-count-with-skipping run clean
+	make -C examples/blast run clean
+	make -C examples/blast-with-force-and-simulate run clean
+
 pep8:
 	find . -path './.tox' -prune -o -path './build' -prune -o -path './dist' -prune -o -name '*.py' -print0 | xargs -0 pep8
 
diff --git a/README.md b/README.md
index 8e85e06..230930e 100644
--- a/README.md
+++ b/README.md
@@ -248,7 +248,7 @@ the full list:
 ## Step script environment variables
 
 The following environment variables are set when a step's script is
-exectued:
+executed:
 
 * `SP_ORIGINAL_ARGS` contains the (space-separated) list of arguments
   originally passed to `slurm-pipeline.py`. Most scripts will not need to
@@ -275,26 +275,42 @@ exectued:
   invokes `sbatch` to guarantee that the execution of the script does not
   begin until after the tasks from all dependent steps have finished
   successfully.
+* `SP_NICE_ARG` contains a string that should be put on the command line when
+  calling `sbatch`. This sets the priority level of the SLURM jobs. The numeric
+  nice value can be set using the `--nice` option when running `slurm-pipeline.py`.
+  See `man sbatch` for details on nice values. Note that calling `sbatch` with
+  this value is not enforced. The `--nice` option simply provides a simple way
+  to specify a priority value on the command line and to pass it to scripts.
+  Scripts can always ignore it or use their own value. If no value is given,
+  `SP_NICE_ARG` will contain just the string `--nice`, which will tell SLURM
+  to use a default nice value. It is useful to use a default nice value as it
+  allows a regular user to later submit jobs with a higher priority (lower nice
+  value). A regular user cannot use a negative nice value, so if a default
+  nice value was not used, all jobs get nice value `0` which prevents the user
+  from submitting higher priority jobs later on.
+
+## Calling sbatch with SP_DEPENDENCY_ARG and SP_NICE_ARG
+
+The canonical way to use `SP_DEPENDENCY_ARG` and `SP_NICE_ARG` when calling
+`sbatch` in a step shell script is as follows:
 
-    The canonical way to use `SP_DEPENDENCY_ARG` in a step script is as
-    follows:
+```sh
+jobid=`sbatch $SP_DEPENDENCY_ARG $SP_NICE_ARG script.sh | cut -f4 -d' '`
+echo "TASK: $task $jobid"
+```
 
-    ```sh
-    jobid=`sbatch -n 1 $SP_DEPENDENCY_ARG submit.sh $task | cut -f4 -d' '`
-    echo "TASK: $task $jobid"
-    ```
+This calls `sbatch` with the dependency and nice arguments (if any) and
+gets the job id from the `sbatch` output (`sbatch` prints a line like
+`Submitted batch job 3779695`) and the `cut` in the above pulls out just
+the job id. The task name is then emitted, along with the job id.
 
-    This calls `sbatch` with the dependency argument (if any) and
-    simultaneously gets the job id from the `sbatch` output (`sbatch` prints a
-    line like `Submitted batch job 3779695`) and the `cut` in the above pulls
-    out just the job id. The task name is then emitted, along with the job id.
 
 ## Separation of concerns
 
 `slurm-pipeline.py` doesn't interact with SLURM at all. Actually, the
-*only* thing it knows about SLURM is how to construct a dependency argument
-for `sbatch` (so it could in theory be generalized to support other
-workload managers).
+*only* things it knows about SLURM is how to construct dependency and nice
+arguments for `sbatch` (so it could in theory be generalized to support
+other workload managers).
 
 To use `slurm-pipeline.py` you need to make a specification file such as
 the one above to indicate the steps in your pipeline, their scripts, and
@@ -716,19 +732,20 @@ the `examples` directory.
 SLURM allows users to submit scripts for later execution. Thus there are
 two distinct phases of operation: the time of scheduling and the later
 time(s) of script excecution. When using `slurm-pipeline.py` it is
-important to understand this distinction.
+very important to keep this distinction in mind.
 
 The reason is that `slurm-pipeline.py` only examines the output of
-scheduling scripts for task names and job ids. If a scheduling script calls
-`sbatch` to execute a later script, the output of that later script cannot
-be checked for `TASK: xxx 97322` style output because `slurm-pipeline.py`
-is completely unaware of the existence of that script. Thus all tasks and
-job dependencies must be established at the time of scheduling.
+*scheduling* scripts for task names and job ids. If a scheduling script
+calls `sbatch` to execute a later script, the output of that later script
+(when it is finally run by SLURM) cannot be checked for `TASK: xxx 97322`
+style output because `slurm-pipeline.py` is completely unaware of the
+existence of that script. In other words, *all tasks and job dependencies
+must be established at the time of scheduling.*
 
 Normally this is not an issue, as many pipelines fall nicely into the model
 used by `slurm-pipeline.py`. But sometimes it is necessary to write a step
-script that performs a slow synchronous operation in order to emit
-tasks. For example, you might have a very large input file that you want to
+script that performs a slow synchronous operation in order to emit tasks.
+For example, you might have a very large input file that you want to
 process in smaller pieces. You can use `split` to break the file into
 pieces and emit task names such as `xaa`, `xab` etc, but you must do this
 synchronously (i.e., in the step script, not in a script submitted to
diff --git a/bin/slurm-pipeline.py b/bin/slurm-pipeline.py
index d041de9..93fef8a 100755
--- a/bin/slurm-pipeline.py
+++ b/bin/slurm-pipeline.py
@@ -64,6 +64,12 @@ parser.add_argument(
           'i.e., those with no dependencies, in the current specification may '
           'begin.'))
 
+parser.add_argument(
+    '--nice', type=int,
+    help=('A numeric nice (priority) value, in the range -10000 (highest '
+          'priority) to 10000 (lowest priority). Note that only privileged '
+          'users can specify a negative adjustment.'))
+
 args, scriptArgs = parser.parse_known_args()
 
 sp = SlurmPipeline(args.specification)
@@ -73,7 +79,7 @@ startAfter = list(map(int, args.startAfter)) if args.startAfter else None
 status = sp.schedule(
     force=args.force, firstStep=args.firstStep, lastStep=args.lastStep,
     sleep=args.sleep, scriptArgs=scriptArgs, skip=args.skip,
-    startAfter=startAfter)
+    startAfter=startAfter, nice=args.nice)
 
 statusAsJSON = sp.specificationToJSON(status)
 
diff --git a/examples/blast-with-force-and-simulate/README.md b/examples/blast-with-force-and-simulate/README.md
new file mode 100644
index 0000000..bf81acc
--- /dev/null
+++ b/examples/blast-with-force-and-simulate/README.md
@@ -0,0 +1,1 @@
+A description of this example can be found in the [top-level README](../../README.md) file.
diff --git a/examples/blast/2-run-blast.sh b/examples/blast/2-run-blast.sh
index 816264b..75a6760 100755
--- a/examples/blast/2-run-blast.sh
+++ b/examples/blast/2-run-blast.sh
@@ -2,7 +2,7 @@
 
 ./blast.sh -query $1 -outfmt "6 bitscore sseqid qseqid"
 
-# Clean up (remove the split FASTA file made by split-fasta.sh).
+# Clean up (remove the split FASTA file made by 1-split-fasta.sh).
 rm $1
 
 echo "TASK: $1"
diff --git a/examples/blast/README.md b/examples/blast/README.md
new file mode 100644
index 0000000..bf81acc
--- /dev/null
+++ b/examples/blast/README.md
@@ -0,0 +1,1 @@
+A description of this example can be found in the [top-level README](../../README.md) file.
diff --git a/examples/blast/blast.sh b/examples/blast/blast.sh
index 4943ed0..a44c43f 100755
--- a/examples/blast/blast.sh
+++ b/examples/blast/blast.sh
@@ -1,10 +1,9 @@
 #!/bin/bash
 
-
 # $1 = "--query" (to simulate running BLAST), which we just ignore.
 # $3 = "--outfmt" (ditto).
 
-# $2 is given to us by run-blast.sh (it's one of the x?? FASTA files). Pull
+# $2 is given to us by 2-run-blast.sh (it's one of the x?? FASTA files). Pull
 # out the query id so we can make fake BLAST output for it.
 queryId=`head -n 1 $2 | cut -c2-`
 
diff --git a/examples/double-collect/README.md b/examples/double-collect/README.md
new file mode 100644
index 0000000..bf81acc
--- /dev/null
+++ b/examples/double-collect/README.md
@@ -0,0 +1,1 @@
+A description of this example can be found in the [top-level README](../../README.md) file.
diff --git a/examples/word-count-with-skipping/README.md b/examples/word-count-with-skipping/README.md
new file mode 100644
index 0000000..bf81acc
--- /dev/null
+++ b/examples/word-count-with-skipping/README.md
@@ -0,0 +1,1 @@
+A description of this example can be found in the [top-level README](../../README.md) file.
diff --git a/examples/word-count/README.md b/examples/word-count/README.md
new file mode 100644
index 0000000..bf81acc
--- /dev/null
+++ b/examples/word-count/README.md
@@ -0,0 +1,1 @@
+A description of this example can be found in the [top-level README](../../README.md) file.
diff --git a/setup.py b/setup.py
index da5b6b6..209d169 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 from setuptools import setup
 
 setup(name='slurm-pipeline',
-      version='1.1.2',
+      version='1.1.3',
       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 8ee6f8d..817c9ae 100644
--- a/slurm_pipeline/pipeline.py
+++ b/slurm_pipeline/pipeline.py
@@ -26,6 +26,11 @@ class SlurmPipeline(SlurmPipelineBase):
     # job ids. The following regex just matches the first part of that.
     TASK_NAME_LINE = re.compile('^TASK:\s+(\S+)\s*')
 
+    # Limits on the --nice argument to sbatch. In later SLURM versions the
+    # limits are +/-2147483645. See https://slurm.schedmd.com/sbatch.html
+    NICE_HIGHEST = -10000
+    NICE_LOWEST = 10000
+
     @staticmethod
     def checkSpecification(specification):
         """
@@ -54,7 +59,7 @@ class SlurmPipeline(SlurmPipelineBase):
         SlurmPipelineBase.checkSpecification(specification)
 
     def schedule(self, force=False, firstStep=None, lastStep=None, sleep=0.0,
-                 scriptArgs=None, skip=None, startAfter=None):
+                 scriptArgs=None, skip=None, startAfter=None, nice=None):
         """
         Schedule the running of our execution specification.
 
@@ -86,8 +91,12 @@ class SlurmPipeline(SlurmPipelineBase):
             (either successully or unsuccessully, it doesn't matter) before
             steps in the current specification may start. If C{None}, steps in
             the current specification may start immediately.
+        @param nice: An C{int} nice (priority) value, in the range
+            self.NICE_HIGHEST to self.NICE_LOWEST. Note that only
+            privileged users can specify a negative adjustment.
         @raise SchedulingError: If there is a problem with the first, last, or
-            skipped steps, as determined by self._checkRuntime.
+            skipped steps, as determined by self._checkRuntime. ValueError if
+            C{nice} is not numeric or is out of its allowed range.
         @return: A specification C{dict}. This is a copy of the original
             specification, updated with information about this scheduling.
         """
@@ -97,11 +106,12 @@ class SlurmPipeline(SlurmPipelineBase):
         if nSteps and lastStep is not None and firstStep is None:
             firstStep = list(specification['steps'])[0]
         skip = set(skip or ())
-        self._checkRuntime(steps, firstStep, lastStep, skip)
+        self._checkRuntime(steps, firstStep, lastStep, skip, nice)
         specification.update({
             'force': force,
             'firstStep': firstStep,
             'lastStep': lastStep,
+            'nice': nice,
             'scheduledAt': time.time(),
             'scriptArgs': scriptArgs,
             'skip': skip,
@@ -110,6 +120,8 @@ class SlurmPipeline(SlurmPipelineBase):
         })
 
         environ['SP_FORCE'] = str(int(force))
+        environ['SP_NICE_ARG'] = (
+            '--nice' if nice is None else '--nice %d' % nice)
         firstStepFound = lastStepFound = False
 
         for stepIndex, stepName in enumerate(steps):
@@ -297,8 +309,8 @@ class SlurmPipeline(SlurmPipelineBase):
                         (taskName, jobIds, script, step['name']))
                 tasks[taskName].update(jobIds)
 
-    @staticmethod
-    def _checkRuntime(steps, firstStep=None, lastStep=None, skip=None):
+    def _checkRuntime(self, steps, firstStep=None, lastStep=None, skip=None,
+                      nice=None):
         """
         Check that a proposed scheduling makes sense.
 
@@ -309,9 +321,13 @@ class SlurmPipeline(SlurmPipelineBase):
         @param lastStep: If not C{None}, the name of the last specification
             step to execute.
         @param skip: A C{set} of C{str} step names that should be skipped.
+        @param nice: An C{int} nice (priority) value, in the range
+            self.NICE_HIGHEST to self.NICE_LOWEST. Note that only privileged
+            users can specify a negative adjustment.
         @raise SchedulingError: if the last step occurs before the first, if
-            the last or first step are unknown, or if asked to skip a
-            non-existent step.
+            the last or first step are unknown, if asked to skip a
+            non-existent step, or if C{nice} is not numeric or is out of its
+            allowed range (see above).
         @return: An C{OrderedDict} keyed by specification step name,
             with values that are step C{dict}s. This provides convenient /
             direct access to steps by name.
@@ -326,6 +342,23 @@ class SlurmPipeline(SlurmPipelineBase):
             raise SchedulingError(
                 'Last step %r not found in specification' % lastStep)
 
+        if nice is not None:
+            try:
+                nice = int(nice)
+            except ValueError:
+                raise SchedulingError(
+                    'Nice (priority) value %r is not numeric' % nice)
+            else:
+                if nice < self.NICE_HIGHEST or nice > self.NICE_LOWEST:
+                    raise SchedulingError(
+                        'Nice (priority) value %r is outside the allowed '
+                        '[%d, %d] range' %
+                        (nice, self.NICE_HIGHEST, self.NICE_LOWEST))
+
+        if lastStep is not None and lastStep not in steps:
+            raise SchedulingError(
+                'Last step %r not found in specification' % lastStep)
+
         for step in steps.values():
             if step['name'] == firstStep:
                 firstStepFound = True


Test patch:
diff --git a/test/test_pipeline.py b/test/test_pipeline.py
index 560145c..a665d01 100644
--- a/test/test_pipeline.py
+++ b/test/test_pipeline.py
@@ -871,6 +871,64 @@ class TestSlurmPipeline(TestCase):
         env = subprocessMock.mock_calls[0][2]['env']
         self.assertEqual('1', env['SP_FORCE'])
 
+    @patch('subprocess.check_output')
+    @patch('os.access')
+    @patch('os.path.exists')
+    def testDefaultNice(self, existsMock, accessMock, subprocessMock):
+        """
+        If no nice value is given to schedule, SP_NICE_ARG must be set to
+        '--nice' in the step execution environment.
+        """
+        subprocessMock.return_value = ''
+
+        sp = SlurmPipeline(
+            {
+                'steps': [
+                    {
+                        'name': 'name1',
+                        'script': 'script1',
+                    },
+                ],
+            })
+        sp.schedule()
+
+        subprocessMock.assert_has_calls([
+            call(['script1'], cwd='.', universal_newlines=True,
+                 stdin=DEVNULL, env=ANY),
+        ])
+
+        env = subprocessMock.mock_calls[0][2]['env']
+        self.assertEqual('--nice', env['SP_NICE_ARG'])
+
+    @patch('subprocess.check_output')
+    @patch('os.access')
+    @patch('os.path.exists')
+    def testSpecificNice(self, existsMock, accessMock, subprocessMock):
+        """
+        If a specific nice value is given to schedule, SP_NICE_ARG must be set
+        to the expected value in the step execution environment.
+        """
+        subprocessMock.return_value = ''
+
+        sp = SlurmPipeline(
+            {
+                'steps': [
+                    {
+                        'name': 'name1',
+                        'script': 'script1',
+                    },
+                ],
+            })
+        sp.schedule(nice=40)
+
+        subprocessMock.assert_has_calls([
+            call(['script1'], cwd='.', universal_newlines=True,
+                 stdin=DEVNULL, env=ANY),
+        ])
+
+        env = subprocessMock.mock_calls[0][2]['env']
+        self.assertEqual('--nice 40', env['SP_NICE_ARG'])
+
     @patch('subprocess.check_output')
     @patch('os.access')
     @patch('os.path.exists')
@@ -1378,6 +1436,7 @@ class TestSlurmPipeline(TestCase):
             {
                 'firstStep': 'name2',
                 'lastStep': None,
+                'nice': None,
                 'force': True,
                 'scheduledAt': 10.0,
                 'scriptArgs': None,
@@ -1495,3 +1554,70 @@ class TestSlurmPipeline(TestCase):
         env3 = subprocessMock.mock_calls[2][2]['env']
         self.assertEqual('--dependency=afternotok:238?afternotok:560',
                          env3['SP_DEPENDENCY_ARG'])
+
+    @patch('subprocess.check_output')
+    @patch('os.access')
+    @patch('os.path.exists')
+    def testStringNice(self, existsMock, accessMock, subprocessMock):
+        """
+        If a string nice value is passed to schedule, a SchedulingError
+        must be raised.
+        """
+        subprocessMock.return_value = 'TASK: xxx 123\n'
+        sp = SlurmPipeline(
+            {
+                'steps': [
+                    {
+                        'name': 'name1',
+                        'script': 'script1',
+                    },
+                ],
+            })
+        error = "^Nice \(priority\) value 'x' is not numeric$"
+        assertRaisesRegex(self, SchedulingError, error, sp.schedule, nice='x')
+
+    @patch('subprocess.check_output')
+    @patch('os.access')
+    @patch('os.path.exists')
+    def testNiceTooBig(self, existsMock, accessMock, subprocessMock):
+        """
+        If a nice value that is too big (> 10000) is passed to schedule, a
+        SchedulingError must be raised.
+        """
+        subprocessMock.return_value = 'TASK: xxx 123\n'
+        sp = SlurmPipeline(
+            {
+                'steps': [
+                    {
+                        'name': 'name1',
+                        'script': 'script1',
+                    },
+                ],
+            })
+        error = ("^Nice \(priority\) value 10001 is outside the allowed "
+                 "\[-10000, 10000\] range$")
+        assertRaisesRegex(self, SchedulingError, error, sp.schedule,
+                          nice=10001)
+
+    @patch('subprocess.check_output')
+    @patch('os.access')
+    @patch('os.path.exists')
+    def testNiceTooSmall(self, existsMock, accessMock, subprocessMock):
+        """
+        If a nice value that is too small (< -10000) is passed to schedule, a
+        SchedulingError must be raised.
+        """
+        subprocessMock.return_value = 'TASK: xxx 123\n'
+        sp = SlurmPipeline(
+            {
+                'steps': [
+                    {
+                        'name': 'name1',
+                        'script': 'script1',
+                    },
+                ],
+            })
+        error = ("^Nice \(priority\) value -10001 is outside the allowed "
+                 "\[-10000, 10000\] range$")
+        assertRaisesRegex(self, SchedulingError, error, sp.schedule,
+                          nice=-10001)