Repository: dask/dask-jobqueue
Pull request #610 (2023-07-12 12:37:32)

Asynchronous job submission and removal
This solution is using asyncio.create_subprocess_exec and similar to what was done in PR #473.
However, this other PR lacked an update job removal, which is included here.

The `test_jobqueue_job_call` required an overhaul to make it compatible with the async code. Its form is basically the same as the other async tests now.

In issue #470 another solution using `run_in_executor` was proposed. I consider the solution with `create_subprocess_exec` more rigorous though, because with `create_subprocess_exec` the code is fully using the async API, instead of wrapping non-async functionality in a `run_in_executor` call.

jrueb: I have removed the `weakref.finalize` call here, which was one way to remove jobs at program exit. This is because `finalize` interferes with the fact that `_close_job` is a coroutine now. At program exit, Distributed already takes care of job removal here
https://github.com/dask/distributed/blob/8e0c79bb82b5d9e3f4f9538935520da960588df6/distributed/deploy/spec.py#L692C1-L692C1
In fact, before this PR, jobs were removed twice, which made the program exit unnecessary lengthy.

There is one issue left however: As you can see in the linked Distributed code, there is a timeout of 10 seconds. This is not enough to remove much more than 350 jobs (at least in my tests with HTCondor). One may want to ask to increase this timeout.

Patch:
diff --git a/dask_jobqueue/core.py b/dask_jobqueue/core.py
index 01c1756..b265ac8 100644
--- a/dask_jobqueue/core.py
+++ b/dask_jobqueue/core.py
@@ -4,9 +4,8 @@ import math
 import os
 import re
 import shlex
-import subprocess
+import asyncio
 import sys
-import weakref
 import abc
 import tempfile
 import copy
@@ -48,6 +47,8 @@ job_parameters = """
         Seconds to wait for a scheduler before closing workers
     extra : list
         Deprecated: use ``worker_extra_args`` instead. This parameter will be removed in a future version.
+    worker_command : list
+        Command to run when launching a worker.  Defaults to "distributed.cli.dask_worker"
     worker_extra_args : list
         Additional arguments to pass to `dask-worker`
     env_extra : list
@@ -166,6 +167,7 @@ class Job(ProcessInterface, abc.ABC):
         death_timeout=None,
         local_directory=None,
         extra=None,
+        worker_command=None,
         worker_extra_args=None,
         job_extra=None,
         job_extra_directives=None,
@@ -222,6 +224,10 @@ class Job(ProcessInterface, abc.ABC):
             )
         if extra is None:
             extra = dask.config.get("jobqueue.%s.extra" % self.config_name)
+        if worker_command is None:
+            worker_command = dask.config.get(
+                "jobqueue.%s.worker-command" % self.config_name
+            )
         if worker_extra_args is None:
             worker_extra_args = dask.config.get(
                 "jobqueue.%s.worker-extra-args" % self.config_name
@@ -332,17 +338,23 @@ class Job(ProcessInterface, abc.ABC):
         self._job_script_prologue = job_script_prologue
 
         # dask-worker command line build
-        dask_worker_command = "%(python)s -m distributed.cli.dask_worker" % dict(
-            python=python
+        dask_worker_command = "%(python)s -m %(worker_command)s" % dict(
+            python=python,
+            worker_command=worker_command
         )
+
         command_args = [dask_worker_command, self.scheduler]
-        command_args += ["--nthreads", self.worker_process_threads]
-        if processes is not None and processes > 1:
-            command_args += ["--nworkers", processes]
 
-        command_args += ["--memory-limit", self.worker_process_memory]
+        # common
         command_args += ["--name", str(name)]
-        command_args += ["--nanny" if nanny else "--no-nanny"]
+        command_args += ["--nthreads", self.worker_process_threads]
+        command_args += ["--memory-limit", self.worker_process_memory]
+
+        #  distributed.cli.dask_worker specific
+        if worker_command == "distributed.cli.dask_worker":
+            if processes is not None and processes > 1:
+                command_args += ["--nworkers", processes]
+            command_args += ["--nanny" if nanny else "--no-nanny"]
 
         if death_timeout is not None:
             command_args += ["--death-timeout", death_timeout]
@@ -390,8 +402,7 @@ class Job(ProcessInterface, abc.ABC):
             yield fn
 
     async def _submit_job(self, script_filename):
-        # Should we make this async friendly?
-        return self._call(shlex.split(self.submit_command) + [script_filename])
+        return await self._call(shlex.split(self.submit_command) + [script_filename])
 
     @property
     def worker_process_threads(self):
@@ -411,8 +422,6 @@ class Job(ProcessInterface, abc.ABC):
             out = await self._submit_job(fn)
             self.job_id = self._job_id_from_submit_output(out)
 
-        weakref.finalize(self, self._close_job, self.job_id, self.cancel_command)
-
         logger.debug("Starting job: %s", self.job_id)
         await super().start()
 
@@ -439,18 +448,18 @@ class Job(ProcessInterface, abc.ABC):
 
     async def close(self):
         logger.debug("Stopping worker: %s job: %s", self.name, self.job_id)
-        self._close_job(self.job_id, self.cancel_command)
+        await self._close_job(self.job_id, self.cancel_command)
 
     @classmethod
-    def _close_job(cls, job_id, cancel_command):
+    async def _close_job(cls, job_id, cancel_command):
         if job_id:
             with suppress(RuntimeError):  # deleting job when job already gone
-                cls._call(shlex.split(cancel_command) + [job_id])
+                await cls._call(shlex.split(cancel_command) + [job_id])
             logger.debug("Closed job %s", job_id)
 
     @staticmethod
-    def _call(cmd, **kwargs):
-        """Call a command using subprocess.Popen.
+    async def _call(cmd, **kwargs):
+        """Call a command using asyncio.create_subprocess_exec.
 
         This centralizes calls out to the command line, providing consistent
         outputs, logging, and an opportunity to go asynchronous in the future.
@@ -459,7 +468,7 @@ class Job(ProcessInterface, abc.ABC):
         ----------
         cmd: List(str))
             A command, each of which is a list of strings to hand to
-            subprocess.Popen
+            asyncio.create_subprocess_exec
 
         Examples
         --------
@@ -478,11 +487,14 @@ class Job(ProcessInterface, abc.ABC):
             "Executing the following command to command line\n{}".format(cmd_str)
         )
 
-        proc = subprocess.Popen(
-            cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, **kwargs
+        proc = await asyncio.create_subprocess_exec(
+            *cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            **kwargs
         )
 
-        out, err = proc.communicate()
+        out, err = await proc.communicate()
         out, err = out.decode(), err.decode()
 
         if proc.returncode != 0:
@@ -745,7 +757,6 @@ or by setting this value in the config file found in `~/.config/dask/jobqueue.ya
         for key, value in worker_security_dict.items():
             # dump worker in-memory keys for use in job_script
             if value is not None and "\n" in value:
-
                 try:
                     f = tempfile.NamedTemporaryFile(
                         mode="wt",
diff --git a/dask_jobqueue/jobqueue.yaml b/dask_jobqueue/jobqueue.yaml
index bd7b9c5..3bcb8c5 100644
--- a/dask_jobqueue/jobqueue.yaml
+++ b/dask_jobqueue/jobqueue.yaml
@@ -12,6 +12,7 @@ jobqueue:
     local-directory: null       # Location of fast local storage like /scratch or $TMPDIR
     shared-temp-directory: null       # Shared directory currently used to dump temporary security objects for workers
     extra: null                 # deprecated: use worker-extra-args
+    worker-command: "distributed.cli.dask_worker" # Command to launch a worker
     worker-extra-args: []       # Additional arguments to pass to `dask-worker`
 
     # OAR resource manager options
@@ -44,6 +45,7 @@ jobqueue:
     local-directory: null       # Location of fast local storage like /scratch or $TMPDIR
     shared-temp-directory: null       # Shared directory currently used to dump temporary security objects for workers
     extra: null                 # deprecated: use worker-extra-args
+    worker-command: "distributed.cli.dask_worker" # Command to launch a worker
     worker-extra-args: []       # Additional arguments to pass to `dask-worker`
 
     # PBS resource manager options
@@ -75,6 +77,7 @@ jobqueue:
     local-directory: null       # Location of fast local storage like /scratch or $TMPDIR
     shared-temp-directory: null       # Shared directory currently used to dump temporary security objects for workers
     extra: null                 # deprecated: use worker-extra-args
+    worker-command: "distributed.cli.dask_worker" # Command to launch a worker
     worker-extra-args: []       # Additional arguments to pass to `dask-worker`
 
     # SGE resource manager options
@@ -106,6 +109,7 @@ jobqueue:
     local-directory: null       # Location of fast local storage like /scratch or $TMPDIR
     shared-temp-directory: null       # Shared directory currently used to dump temporary security objects for workers
     extra: null                 # deprecated: use worker-extra-args
+    worker-command: "distributed.cli.dask_worker" # Command to launch a worker
     worker-extra-args: []       # Additional arguments to pass to `dask-worker`
 
     # SLURM resource manager options
@@ -138,6 +142,7 @@ jobqueue:
     local-directory: null       # Location of fast local storage like /scratch or $TMPDIR
     shared-temp-directory: null       # Shared directory currently used to dump temporary security objects for workers
     extra: null                 # deprecated: use worker-extra-args
+    worker-command: "distributed.cli.dask_worker" # Command to launch a worker
     worker-extra-args: []       # Additional arguments to pass to `dask-worker`
 
     # PBS resource manager options
@@ -169,6 +174,7 @@ jobqueue:
     local-directory: null       # Location of fast local storage like /scratch or $TMPDIR
     shared-temp-directory: null       # Shared directory currently used to dump temporary security objects for workers
     extra: null                 # deprecated: use worker-extra-args
+    worker-command: "distributed.cli.dask_worker" # Command to launch a worker
     worker-extra-args: []       # Additional arguments to pass to `dask-worker`
 
     # LSF resource manager options
@@ -203,6 +209,7 @@ jobqueue:
     local-directory: null       # Location of fast local storage like /scratch or $TMPDIR
     shared-temp-directory: null       # Shared directory currently used to dump temporary security objects for workers
     extra: null                 # deprecated: use worker-extra-args
+    worker-command: "distributed.cli.dask_worker" # Command to launch a worker
     worker-extra-args: []       # Additional arguments to pass to `dask-worker`
 
     # HTCondor Resource Manager options
@@ -232,6 +239,7 @@ jobqueue:
     local-directory: null       # Location of fast local storage like /scratch or $TMPDIR
     shared-temp-directory: null       # Shared directory currently used to dump temporary security objects for workers
     extra: null                 # deprecated: use worker-extra-args
+    worker-command: "distributed.cli.dask_worker" # Command to launch a worker
     worker-extra-args: []       # Additional arguments to pass to `dask-worker`
 
     env-extra: null
diff --git a/dask_jobqueue/local.py b/dask_jobqueue/local.py
index e7b0aa6..68e98c7 100644
--- a/dask_jobqueue/local.py
+++ b/dask_jobqueue/local.py
@@ -71,7 +71,7 @@ class LocalJob(Job):
         return str(self.process.pid)
 
     @classmethod
-    def _close_job(self, job_id, cancel_command):
+    async def _close_job(self, job_id, cancel_command):
         os.kill(int(job_id), 9)
         # from distributed.utils_test import terminate_process
         # terminate_process(self.process)
diff --git a/dask_jobqueue/oar.py b/dask_jobqueue/oar.py
index 8dba580..0a23f4b 100644
--- a/dask_jobqueue/oar.py
+++ b/dask_jobqueue/oar.py
@@ -11,7 +11,6 @@ logger = logging.getLogger(__name__)
 
 
 class OARJob(Job):
-
     # Override class variables
     submit_command = "oarsub"
     cancel_command = "oardel"


Test patch:
diff --git a/dask_jobqueue/tests/test_htcondor.py b/dask_jobqueue/tests/test_htcondor.py
index c5c5bf6..ea57e00 100644
--- a/dask_jobqueue/tests/test_htcondor.py
+++ b/dask_jobqueue/tests/test_htcondor.py
@@ -106,7 +106,6 @@ def test_extra_args_broken_cancel(loop):
         cancel_command_extra=["-name", "wrong.docker"],
     ) as cluster:
         with Client(cluster) as client:
-
             cluster.scale(2)
 
             client.wait_for_workers(2, timeout=QUEUE_WAIT)
@@ -141,6 +140,7 @@ def test_config_name_htcondor_takes_custom_config():
         "interface": None,
         "death-timeout": None,
         "extra": None,
+        "worker-command": None,
         "worker-extra-args": [],
         "env-extra": None,
         "job-script-prologue": [],
diff --git a/dask_jobqueue/tests/test_job.py b/dask_jobqueue/tests/test_job.py
index f24e114..497ceeb 100644
--- a/dask_jobqueue/tests/test_job.py
+++ b/dask_jobqueue/tests/test_job.py
@@ -432,25 +432,25 @@ def test_deprecation_job_extra(Cluster):
     assert "old_param" in job_script
 
 
-def test_jobqueue_job_call(tmpdir, Cluster):
-    cluster = Cluster(cores=1, memory="1GB")
-
-    path = tmpdir.join("test.py")
-    path.write('print("this is the stdout")')
-
-    out = cluster.job_cls._call([sys.executable, path.strpath])
-    assert out == "this is the stdout\n"
-
-    path_with_error = tmpdir.join("non-zero-exit-code.py")
-    path_with_error.write('print("this is the stdout")\n1/0')
-
-    match = (
-        "Command exited with non-zero exit code.+"
-        "Exit code: 1.+"
-        "stdout:\nthis is the stdout.+"
-        "stderr:.+ZeroDivisionError"
-    )
+@pytest.mark.asyncio
+async def test_jobqueue_job_call(tmpdir, Cluster):
+    async with Cluster(cores=1, memory="1GB", asynchronous=True) as cluster:
+        path = tmpdir.join("test.py")
+        path.write('print("this is the stdout")')
+
+        out = await cluster.job_cls._call([sys.executable, path.strpath])
+        assert out == "this is the stdout\n"
+
+        path_with_error = tmpdir.join("non-zero-exit-code.py")
+        path_with_error.write('print("this is the stdout")\n1/0')
+
+        match = (
+            "Command exited with non-zero exit code.+"
+            "Exit code: 1.+"
+            "stdout:\nthis is the stdout.+"
+            "stderr:.+ZeroDivisionError"
+        )
 
-    match = re.compile(match, re.DOTALL)
-    with pytest.raises(RuntimeError, match=match):
-        cluster.job_cls._call([sys.executable, path_with_error.strpath])
+        match = re.compile(match, re.DOTALL)
+        with pytest.raises(RuntimeError, match=match):
+            await cluster.job_cls._call([sys.executable, path_with_error.strpath])
diff --git a/dask_jobqueue/tests/test_jobqueue_core.py b/dask_jobqueue/tests/test_jobqueue_core.py
index 8f743f7..ceadd12 100644
--- a/dask_jobqueue/tests/test_jobqueue_core.py
+++ b/dask_jobqueue/tests/test_jobqueue_core.py
@@ -49,6 +49,9 @@ def test_command_template(Cluster):
         assert " --local-directory /scratch" in cluster._dummy_job._command_template
         assert " --preload mymodule" in cluster._dummy_job._command_template
 
+    with Cluster(cores=2, memory="4GB", worker_command="dask_cuda.cli") as cluster:
+        assert "dask_cuda.cli" in cluster._dummy_job._command_template
+
 
 def test_shebang_settings(Cluster, request):
     if Cluster is HTCondorCluster or Cluster is LocalCluster:
@@ -244,7 +247,6 @@ def get_interface_and_port(index=0):
 
 
 def test_scheduler_options(Cluster):
-
     interface, port = get_interface_and_port()
 
     with Cluster(
@@ -321,7 +323,6 @@ def test_import_scheduler_options_from_config(Cluster):
     with dask.config.set(
         {"jobqueue.%s.scheduler-options" % default_config_name: scheduler_options}
     ):
-
         with Cluster(cores=2, memory="2GB") as cluster:
             scheduler_options = cluster.scheduler_spec["options"]
             assert scheduler_options.get("interface") == config_scheduler_interface
diff --git a/dask_jobqueue/tests/test_lsf.py b/dask_jobqueue/tests/test_lsf.py
index 40c826c..87e2cd6 100644
--- a/dask_jobqueue/tests/test_lsf.py
+++ b/dask_jobqueue/tests/test_lsf.py
@@ -17,7 +17,6 @@ from . import QUEUE_WAIT
 
 def test_header():
     with LSFCluster(walltime="00:02", processes=4, cores=8, memory="8GB") as cluster:
-
         assert "#BSUB" in cluster.job_header
         assert "#BSUB -J dask-worker" in cluster.job_header
         assert "#BSUB -n 8" in cluster.job_header
@@ -35,7 +34,6 @@ def test_header():
         ncpus=24,
         mem=100000000000,
     ) as cluster:
-
         assert "#BSUB -q general" in cluster.job_header
         assert "#BSUB -J dask-worker" in cluster.job_header
         assert "#BSUB -n 24" in cluster.job_header
@@ -54,7 +52,6 @@ def test_header():
         ncpus=24,
         mem=100000000000,
     ) as cluster:
-
         assert "#BSUB -q general" in cluster.job_header
         assert "#BSUB -J dask-worker" in cluster.job_header
         assert "#BSUB -n 24" in cluster.job_header
@@ -65,7 +62,6 @@ def test_header():
         assert '#BSUB -P "Dask On LSF"' in cluster.job_header
 
     with LSFCluster(cores=4, memory="8GB") as cluster:
-
         assert "#BSUB -n" in cluster.job_header
         assert "#BSUB -W" in cluster.job_header
         assert "#BSUB -M" in cluster.job_header
@@ -75,7 +71,6 @@ def test_header():
     with LSFCluster(
         cores=4, memory="8GB", job_extra_directives=["-u email@domain.com"]
     ) as cluster:
-
         assert "#BSUB -u email@domain.com" in cluster.job_header
         assert "#BSUB -n" in cluster.job_header
         assert "#BSUB -W" in cluster.job_header
@@ -86,7 +81,6 @@ def test_header():
 
 def test_job_script():
     with LSFCluster(walltime="00:02", processes=4, cores=8, memory="28GB") as cluster:
-
         job_script = cluster.job_script()
         assert "#BSUB" in job_script
         assert "#BSUB -J dask-worker" in job_script
@@ -101,9 +95,9 @@ def test_job_script():
             in job_script
         )
         formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
-        assert (
-            f"--nthreads 2 --nworkers 4 --memory-limit {formatted_bytes}" in job_script
-        )
+        assert ("--nthreads 2" in job_script)
+        assert ("--nworkers 4" in job_script)
+        assert (f"--memory-limit {formatted_bytes}" in job_script)
 
     with LSFCluster(
         queue="general",
@@ -114,7 +108,6 @@ def test_job_script():
         ncpus=24,
         mem=100000000000,
     ) as cluster:
-
         job_script = cluster.job_script()
         assert "#BSUB -q general" in cluster.job_header
         assert "#BSUB -J dask-worker" in cluster.job_header
@@ -130,9 +123,9 @@ def test_job_script():
             in job_script
         )
         formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
-        assert (
-            f"--nthreads 2 --nworkers 4 --memory-limit {formatted_bytes}" in job_script
-        )
+        assert ("--nthreads 2" in job_script)
+        assert ("--nworkers 4" in job_script)
+        assert (f"--memory-limit {formatted_bytes}" in job_script)
 
     with LSFCluster(
         walltime="1:00",
@@ -141,7 +134,6 @@ def test_job_script():
         project="Dask On LSF",
         job_extra_directives=["-R rusage[mem=16GB]"],
     ) as cluster:
-
         job_script = cluster.job_script()
 
         assert "#BSUB -J dask-worker" in cluster.job_header
@@ -322,6 +314,7 @@ def test_config_name_lsf_takes_custom_config():
         "local-directory": "/foo",
         "shared-temp-directory": None,
         "extra": None,
+        "worker-command": None,
         "worker-extra-args": [],
         "env-extra": None,
         "job-script-prologue": [],
diff --git a/dask_jobqueue/tests/test_oar.py b/dask_jobqueue/tests/test_oar.py
index 52fd51b..c938488 100644
--- a/dask_jobqueue/tests/test_oar.py
+++ b/dask_jobqueue/tests/test_oar.py
@@ -82,9 +82,9 @@ def test_job_script():
             in job_script
         )
         formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
-        assert (
-            f"--nthreads 2 --nworkers 4 --memory-limit {formatted_bytes}" in job_script
-        )
+        assert ("--nthreads 2" in job_script)
+        assert ("--nworkers 4" in job_script)
+        assert (f"--memory-limit {formatted_bytes}" in job_script)
 
     with OARCluster(
         walltime="00:02:00",
@@ -115,9 +115,9 @@ def test_job_script():
             in job_script
         )
         formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
-        assert (
-            f"--nthreads 2 --nworkers 4 --memory-limit {formatted_bytes}" in job_script
-        )
+        assert ("--nthreads 2" in job_script)
+        assert ("--nworkers 4" in job_script)
+        assert (f"--memory-limit {formatted_bytes}" in job_script)
 
 
 def test_config_name_oar_takes_custom_config():
@@ -137,6 +137,7 @@ def test_config_name_oar_takes_custom_config():
         "local-directory": "/foo",
         "shared-temp-directory": None,
         "extra": None,
+        "worker-command": None,
         "worker-extra-args": [],
         "env-extra": None,
         "job-script-prologue": [],
diff --git a/dask_jobqueue/tests/test_pbs.py b/dask_jobqueue/tests/test_pbs.py
index 3959c86..8ec3a01 100644
--- a/dask_jobqueue/tests/test_pbs.py
+++ b/dask_jobqueue/tests/test_pbs.py
@@ -18,7 +18,6 @@ def test_header(Cluster):
     with Cluster(
         walltime="00:02:00", processes=4, cores=8, memory="28GB", name="dask-worker"
     ) as cluster:
-
         assert "#PBS" in cluster.job_header
         assert "#PBS -N dask-worker" in cluster.job_header
         assert "#PBS -l select=1:ncpus=8:mem=27GB" in cluster.job_header
@@ -34,7 +33,6 @@ def test_header(Cluster):
         resource_spec="select=1:ncpus=24:mem=100GB",
         memory="28GB",
     ) as cluster:
-
         assert "#PBS -q regular" in cluster.job_header
         assert "#PBS -N dask-worker" in cluster.job_header
         assert "#PBS -l select=1:ncpus=24:mem=100GB" in cluster.job_header
@@ -43,7 +41,6 @@ def test_header(Cluster):
         assert "#PBS -A DaskOnPBS" in cluster.job_header
 
     with Cluster(cores=4, memory="8GB") as cluster:
-
         assert "#PBS -j oe" not in cluster.job_header
         assert "#PBS -N" in cluster.job_header
         assert "#PBS -l walltime=" in cluster.job_header
@@ -51,7 +48,6 @@ def test_header(Cluster):
         assert "#PBS -q" not in cluster.job_header
 
     with Cluster(cores=4, memory="8GB", job_extra_directives=["-j oe"]) as cluster:
-
         assert "#PBS -j oe" in cluster.job_header
         assert "#PBS -N" in cluster.job_header
         assert "#PBS -l walltime=" in cluster.job_header
@@ -62,7 +58,6 @@ def test_header(Cluster):
 @pytest.mark.parametrize("Cluster", [PBSCluster, MoabCluster])
 def test_job_script(Cluster):
     with Cluster(walltime="00:02:00", processes=4, cores=8, memory="28GB") as cluster:
-
         job_script = cluster.job_script()
         assert "#PBS" in job_script
         assert "#PBS -N dask-worker" in job_script
@@ -76,9 +71,9 @@ def test_job_script(Cluster):
             in job_script
         )
         formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
-        assert (
-            f"--nthreads 2 --nworkers 4 --memory-limit {formatted_bytes}" in job_script
-        )
+        assert ("--nthreads 2" in job_script)
+        assert ("--nworkers 4" in job_script)
+        assert (f"--memory-limit {formatted_bytes}" in job_script)
 
     with Cluster(
         queue="regular",
@@ -88,7 +83,6 @@ def test_job_script(Cluster):
         resource_spec="select=1:ncpus=24:mem=100GB",
         memory="28GB",
     ) as cluster:
-
         job_script = cluster.job_script()
         assert "#PBS -q regular" in job_script
         assert "#PBS -N dask-worker" in job_script
@@ -102,9 +96,9 @@ def test_job_script(Cluster):
             in job_script
         )
         formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
-        assert (
-            f"--nthreads 2 --nworkers 4 --memory-limit {formatted_bytes}" in job_script
-        )
+        assert ("--nthreads 2" in job_script)
+        assert ("--nworkers 4" in job_script)
+        assert (f"--memory-limit {formatted_bytes}" in job_script)
 
 
 @pytest.mark.env("pbs")
@@ -119,7 +113,6 @@ def test_basic(loop):
         loop=loop,
     ) as cluster:
         with Client(cluster) as client:
-
             cluster.scale(2)
             client.wait_for_workers(2, timeout=QUEUE_WAIT)
 
@@ -154,7 +147,6 @@ def test_scale_cores_memory(loop):
         loop=loop,
     ) as cluster:
         with Client(cluster) as client:
-
             cluster.scale(cores=2)
             client.wait_for_workers(1, timeout=QUEUE_WAIT)
 
@@ -188,7 +180,6 @@ def test_basic_scale_edge_cases(loop):
         job_extra_directives=["-V"],
         loop=loop,
     ) as cluster:
-
         cluster.scale(2)
         cluster.scale(0)
 
@@ -299,7 +290,6 @@ def test_scale_grouped(loop):
         loop=loop,
     ) as cluster:
         with Client(cluster) as client:
-
             cluster.scale(4)  # Start 2 jobs
 
             start = time()
@@ -361,6 +351,7 @@ def test_config_name_pbs_takes_custom_config():
         "local-directory": "/foo",
         "shared-temp-directory": None,
         "extra": None,
+        "worker-command": None,
         "worker-extra-args": [],
         "env-extra": None,
         "job-script-prologue": [],
diff --git a/dask_jobqueue/tests/test_sge.py b/dask_jobqueue/tests/test_sge.py
index 10e2da2..ab76122 100644
--- a/dask_jobqueue/tests/test_sge.py
+++ b/dask_jobqueue/tests/test_sge.py
@@ -16,7 +16,6 @@ def test_basic(loop):
         walltime="00:02:00", cores=8, processes=4, memory="2GiB", loop=loop
     ) as cluster:
         with Client(cluster, loop=loop) as client:
-
             cluster.scale(2)
 
             start = time()
@@ -58,6 +57,7 @@ def test_config_name_sge_takes_custom_config():
         "local-directory": "/foo",
         "shared-temp-directory": None,
         "extra": None,
+        "worker-command": None,
         "worker-extra-args": [],
         "env-extra": None,
         "job-script-prologue": [],
diff --git a/dask_jobqueue/tests/test_slurm.py b/dask_jobqueue/tests/test_slurm.py
index 80769d5..51dbbfb 100644
--- a/dask_jobqueue/tests/test_slurm.py
+++ b/dask_jobqueue/tests/test_slurm.py
@@ -16,7 +16,6 @@ def test_header():
     with SLURMCluster(
         walltime="00:02:00", processes=4, cores=8, memory="28GB"
     ) as cluster:
-
         assert "#SBATCH" in cluster.job_header
         assert "#SBATCH -J dask-worker" in cluster.job_header
         assert "#SBATCH -n 1" in cluster.job_header
@@ -35,7 +34,6 @@ def test_header():
         job_cpu=16,
         job_mem="100G",
     ) as cluster:
-
         assert "#SBATCH --cpus-per-task=16" in cluster.job_header
         assert "#SBATCH --cpus-per-task=8" not in cluster.job_header
         assert "#SBATCH --mem=100G" in cluster.job_header
@@ -44,7 +42,6 @@ def test_header():
         assert "#SBATCH -p regular" in cluster.job_header
 
     with SLURMCluster(cores=4, memory="8GB") as cluster:
-
         assert "#SBATCH" in cluster.job_header
         assert "#SBATCH -J " in cluster.job_header
         assert "#SBATCH -n 1" in cluster.job_header
@@ -57,7 +54,6 @@ def test_job_script():
     with SLURMCluster(
         walltime="00:02:00", processes=4, cores=8, memory="28GB"
     ) as cluster:
-
         job_script = cluster.job_script()
         assert "#SBATCH" in job_script
         assert "#SBATCH -J dask-worker" in job_script
@@ -77,9 +73,9 @@ def test_job_script():
             in job_script
         )
         formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
-        assert (
-            f"--nthreads 2 --nworkers 4 --memory-limit {formatted_bytes}" in job_script
-        )
+        assert ("--nthreads 2" in job_script)
+        assert ("--nworkers 4" in job_script)
+        assert (f"--memory-limit {formatted_bytes}" in job_script)
 
     with SLURMCluster(
         walltime="00:02:00",
@@ -111,9 +107,9 @@ def test_job_script():
             in job_script
         )
         formatted_bytes = format_bytes(parse_bytes("7GB")).replace(" ", "")
-        assert (
-            f"--nthreads 2 --nworkers 4 --memory-limit {formatted_bytes}" in job_script
-        )
+        assert ("--nthreads 2" in job_script)
+        assert ("--nworkers 4" in job_script)
+        assert (f"--memory-limit {formatted_bytes}" in job_script)
 
 
 @pytest.mark.env("slurm")
@@ -127,7 +123,6 @@ def test_basic(loop):
         loop=loop,
     ) as cluster:
         with Client(cluster) as client:
-
             cluster.scale(2)
 
             start = time()
@@ -193,6 +188,7 @@ def test_config_name_slurm_takes_custom_config():
         "local-directory": "/foo",
         "shared-temp-directory": None,
         "extra": None,
+        "worker-command": None,
         "worker-extra-args": [],
         "env-extra": None,
         "job-script-prologue": [],