Change SchedulerRequirement to SlurmRequirement; add qos and reservat…

…ion (#996)

* Change SchedulerRequirement to SlurmRequirement and add qos and reservation

* Make requirements changes to documentation and continuous integration configuration

* Updated coverage.svg
---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

beeflow/common/config_driver.py

@@ -303,6 +303,10 @@ def filepath_completion_input(*pargs, **kwargs):
                  default='', info='default account time limit (leave blank if none)')
 VALIDATOR.option('job', 'default_partition', validator=lambda val: val.strip(), prompt=True,
                  default='', info='default partition to run jobs on (leave blank if none)')
+VALIDATOR.option('job', 'default_qos', validator=lambda val: val.strip(), prompt=True,
+                 default='', info='default qos to run jobs on (leave blank if none)')
+VALIDATOR.option('job', 'default_reservation', validator=lambda val: val.strip(), prompt=True,
+                 default='', info='default reservation to run jobs on (leave blank if none)')
 
 
 def validate_chrun_opts(opts):

beeflow/common/worker/slurm_worker.py

@@ -25,12 +25,15 @@
 class BaseSlurmWorker(Worker):
     """Base slurm worker code."""
 
-    def __init__(self, default_account='', default_time_limit='', default_partition='', **kwargs):
+    def __init__(self, default_account='', default_time_limit='', default_partition='',
+            default_qos='', default_reservation='', **kwargs):
         """Initialize the base slurm worker."""
         super().__init__(**kwargs)
         self.default_account = default_account
         self.default_time_limit = default_time_limit
         self.default_partition = default_partition
+        self.default_qos = default_qos
+        self.default_reservation = default_reservation
 
     def build_text(self, task):
         """Build text for task script."""
@@ -50,14 +53,17 @@ def build_text(self, task):
         ntasks = task.get_requirement('beeflow:MPIRequirement', 'ntasks', default=nodes)
         # Need to rethink the MPI version parameter
         mpi_version = task.get_requirement('beeflow:MPIRequirement', 'mpiVersion', default='')
-        time_limit = task.get_requirement('beeflow:SchedulerRequirement', 'timeLimit',
+        time_limit = task.get_requirement('beeflow:SlurmRequirement', 'timeLimit',
                                           default=self.default_time_limit)
         time_limit = validation.time_limit(time_limit)
-        account = task.get_requirement('beeflow:SchedulerRequirement', 'account',
+        account = task.get_requirement('beeflow:SlurmRequirement', 'account',
                                        default=self.default_account)
-        partition = task.get_requirement('beeflow:SchedulerRequirement',
-                                         'partition',
+        partition = task.get_requirement('beeflow:SlurmRequirement', 'partition',
                                          default=self.default_partition)
+        qos = task.get_requirement('beeflow:SlurmRequirement', 'qos',
+                                         default=self.default_qos)
+        reservation = task.get_requirement('beeflow:SlurmRequirement', 'reservation',
+                                         default=self.default_reservation)
 
         shell = task.get_requirement('beeflow:ScriptRequirement', 'shell', default="/bin/bash")
         scripts_enabled = task.get_requirement('beeflow:ScriptRequirement', 'enabled',
@@ -81,9 +87,13 @@ def build_text(self, task):
         if time_limit:
             script.append(f'#SBATCH --time={time_limit}')
         if account:
-            script.append(f'#SBATCH -A {account}')
+            script.append(f'#SBATCH --account {account}')
         if partition:
-            script.append(f'#SBATCH -p {partition}')
+            script.append(f'#SBATCH --partition {partition}')
+        if qos:
+            script.append(f'#SBATCH --qos {qos}')
+        if reservation:
+            script.append(f'#SBATCH --reservation {reservation}')
 
         # Return immediately on error
         if shell == "/bin/bash":

beeflow/task_manager/utils.py

@@ -38,7 +38,8 @@ def worker_interface():
         'runner_opts': bc.get('task_manager', 'runner_opts'),
     }
     # Job defaults
-    for default_key in ['default_account', 'default_time_limit', 'default_partition']:
+    for default_key in ['default_account', 'default_time_limit', 'default_partition',
+                        'default_qos', 'default_reservation']:
         worker_kwargs[default_key] = bc.get('job', default_key)
     # Special slurm arguments
     if wls == 'Slurm':

beeflow/wf_manager/resources/wf_update.py

@@ -149,7 +149,7 @@ def handle_state_change(self, state_update, task, wfi, db):
 
         # If the job failed and it doesn't include a checkpoint-restart hint,
         # then fail the entire workflow
-        if state_update.job_state == 'FAILED':
+        if state_update.job_state in ['FAILED', 'SUBMIT_FAIL']:
             set_dependent_tasks_dep_fail(db, wfi, state_update.wf_id, task)
             log.info("Workflow failed")
             wf_id = wfi.workflow_id

ci/bee_config.sh

@@ -39,6 +39,8 @@ setup =
 default_account =
 default_time_limit =
 default_partition =
+default_qos=
+default_reservation=
 
 [graphdb]
 hostname = localhost

docs/sphinx/bee_cwl.rst

@@ -154,25 +154,29 @@ filenames, the ``restart parameter`` will be added to the run command followed
 by the path to the latest checkpoint file, and ``num_tries`` specifies the maximum
 number of times the task will be restarted.
 
-beeflow:SchedulerRequirement
+beeflow:SlurmRequirement
 ----------------------------
 
 This requirement is designed for specifying additional information that will be
-passed to a scheduler such as Slurm on job submission. It currently supports
-the following options:
-
+passed to the Slurm scheduler during job submission. Each of the options can be 
+set in the configuration file bee.conf under the ``job`` section to use for all
+workflows. Setting any beeflow:SlurmRequirement in the CWL file will override the 
+setting in bee.conf. Current options supported are:
+
+* ``account`` - account name to run the job with (often used for charging).
+* ``partition`` - partition to launch job on.
+* ``qos`` - quality of service to use.
+* ``reservation`` - reservation to use to launch job.
 * ``timeLimit`` - time limit for the job in the format that Slurm uses currently.
-* ``account`` - may be useful if running jobs with different accounts (if you
-  want to run all workflows with the same account it's best to set this with
-  the ``default_account`` option under the ``job`` section in the bee.conf file).
-* ``partition`` - partition to launch job with.
 
 An example is shown below::
 
     beeflow:SchedulerRequirement:
       timeLimit: 00:00:10
       account: account12345
-      partition: scaling
+      partition: partition-a
+      qos: long
+      reservation: reservation-a 
 
 beeflow:ScriptRequirement
 -------------------------