Generalise submission and cancellation arguments

dask_jobqueue/core.py

@@ -53,6 +53,10 @@
         Additional arguments to pass to `dask-worker`
     env_extra : list
         Deprecated: use ``job_script_prologue`` instead. This parameter will be removed in a future version.
+    submit_command_extra : list
+        Extra arguments to pass to the job scheduler submit command
+    cancel_command_extra : list
+        Extra arguments to pass to the job scheduler cancel command
     job_script_prologue : list
         Other commands to add to script before launching worker.
     header_skip : list
@@ -172,6 +176,8 @@ def __init__(
         job_extra=None,
         job_extra_directives=None,
         env_extra=None,
+        submit_command_extra=None,
+        cancel_command_extra=None,
         job_script_prologue=None,
         header_skip=None,
         job_directives_skip=None,
@@ -270,6 +276,29 @@ def __init__(
 
         if env_extra is None:
             env_extra = dask.config.get("jobqueue.%s.env-extra" % self.config_name)
+
+        if submit_command_extra is None:
+            submit_command_extra = dask.config.get(
+                "jobqueue.%s.submit-command-extra" % self.config_name, []
+            )
+        self.submit_command_extra = submit_command_extra
+
+        if self.submit_command is not None:
+            self.submit_command += " " + " ".join(
+                shlex.quote(arg) for arg in self.submit_command_extra
+            )
+
+        if cancel_command_extra is None:
+            cancel_command_extra = dask.config.get(
+                "jobqueue.%s.cancel-command-extra" % self.config_name, []
+            )
+        self.cancel_command_extra = cancel_command_extra
+
+        if self.cancel_command is not None:
+            self.cancel_command += " " + " ".join(
+                shlex.quote(arg) for arg in self.cancel_command_extra
+            )
+
         if job_script_prologue is None:
             job_script_prologue = dask.config.get(
                 "jobqueue.%s.job-script-prologue" % self.config_name

dask_jobqueue/htcondor.py

@@ -1,6 +1,5 @@
 import logging
 import re
-import shlex
 
 import dask
 from dask.utils import parse_bytes
@@ -38,8 +37,6 @@ def __init__(
         name=None,
         disk=None,
         config_name=None,
-        submit_command_extra=None,
-        cancel_command_extra=None,
         **base_class_kwargs
     ):
         super().__init__(
@@ -95,28 +92,6 @@ def __init__(
         if self.job_extra_directives:
             self.job_header_dict.update(self.job_extra_directives)
 
-        if submit_command_extra is None:
-            submit_command_extra = dask.config.get(
-                "jobqueue.%s.submit-command-extra" % self.config_name, []
-            )
-
-        self.submit_command = (
-            HTCondorJob.submit_command
-            + " "
-            + " ".join(shlex.quote(arg) for arg in submit_command_extra)
-        )
-
-        if cancel_command_extra is None:
-            cancel_command_extra = dask.config.get(
-                "jobqueue.%s.cancel-command-extra" % self.config_name, []
-            )
-
-        self.cancel_command = (
-            HTCondorJob.cancel_command
-            + " "
-            + " ".join(shlex.quote(arg) for arg in cancel_command_extra)
-        )
-
     def job_script(self):
         """Construct a job submission script"""
         quoted_arguments = quote_arguments(["-c", self._command_template])
@@ -227,10 +202,6 @@ class HTCondorCluster(JobQueueCluster):
     job_extra_directives : dict
         Extra submit file attributes for the job as key-value pairs.
         They will be inserted as ``key = value``.
-    submit_command_extra : list of str
-        Extra arguments to pass to condor_submit
-    cancel_command_extra : list of str
-        Extra arguments to pass to condor_rm
     {job}
     {cluster}
 

dask_jobqueue/jobqueue.yaml

@@ -122,6 +122,8 @@ jobqueue:
     account: null
     walltime: '00:30:00'
     env-extra: null
+    submit-command-extra: []    # Extra sbatch arguments
+    cancel-command-extra: ["--signal=SIGTERM"]    # Extra scancel arguments
     job-script-prologue: []
     job-cpu: null
     job-mem: null

dask_jobqueue/slurm.py

@@ -26,10 +26,15 @@ def __init__(
         job_cpu=None,
         job_mem=None,
         config_name=None,
+        cancel_command_extra=["--signal=SIGTERM"],
         **base_class_kwargs
     ):
         super().__init__(
-            scheduler=scheduler, name=name, config_name=config_name, **base_class_kwargs
+            scheduler=scheduler,
+            name=name,
+            config_name=config_name,
+            cancel_command_extra=cancel_command_extra,
+            **base_class_kwargs
         )
 
         if queue is None: