diff --git a/dask_jobqueue/core.py b/dask_jobqueue/core.py index 62af14aa..080c17b3 100644 --- a/dask_jobqueue/core.py +++ b/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 diff --git a/dask_jobqueue/htcondor.py b/dask_jobqueue/htcondor.py index 4ab5c07a..196ffb3b 100644 --- a/dask_jobqueue/htcondor.py +++ b/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} diff --git a/dask_jobqueue/jobqueue.yaml b/dask_jobqueue/jobqueue.yaml index f9424158..06fb6e9d 100644 --- a/dask_jobqueue/jobqueue.yaml +++ b/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 diff --git a/dask_jobqueue/slurm.py b/dask_jobqueue/slurm.py index 8e6c1a07..8b8942b4 100644 --- a/dask_jobqueue/slurm.py +++ b/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: