Repository: dask/dask-jobqueue (python)

Issue:
Provide parameter for "dask-worker"
Issue https://github.com/rapidsai/dask-cuda/issues/3 proposes creating a `dask-cuda-worker` alternative to `dask-worker` which starts one process per GPU with appropriate environment variables.  If we go ahead with that then it would be convenient to compose that approach with dask-jobqueue.  This would be particularly easy if the following line could be made configurable to allow other executables

https://github.com/dask/dask-jobqueue/blob/5a5585dffea0237e452792c7ed0b13d5d1e3f8bb/dask_jobqueue/core.py#L254

Thoughts?

Pull request description:
Make cli worker parameter flexible
Fixes #229 
Depends https://github.com/rapidsai/dask-cuda/pull/1181

Hi all! 

Basic implementation of calling CLI's other than the `distributed` CLI. 

Will require https://github.com/rapidsai/dask-cuda/pull/1181 to be merged first, adding a `main` entrypoint to `dask-cuda-worker`

I have had to add a bit of a shim to filter out some CLI args that are not shared between the `dask-worker` and `dask-worker-cuda` CLIs. 

Very basic test to see that it works:

```python 
import dask
from dask import distributed
from dask_jobqueue.local import LocalCluster

lc_gpu = LocalCluster(worker_command="dask_cuda.cli", cores=2, memory="2GB")
client = distributed.Client(lc_gpu)
lc_gpu.scale(2)
print(lc_gpu.job_script())
> /home/hmacdope/anaconda3/envs/dask_dev/bin/python -m dask_cuda.cli tcp://192.168.1.5:33677 --name dummy-name --nthreads 1 --memory-limit 0.93GiB --death-timeout 60
```

I am a new contributor so please let me know if I am missing anything obvious. 



Patch:
diff --git a/dask_jobqueue/core.py b/dask_jobqueue/core.py
index 01c1756..2883c08 100644
--- a/dask_jobqueue/core.py
+++ b/dask_jobqueue/core.py
@@ -48,6 +48,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 +168,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 +225,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 +339,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]
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/docs/source/changelog.rst b/docs/source/changelog.rst
index 6df2f64..40ccb4c 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -4,6 +4,9 @@ Changelog
 Development version
 -------------------
 
+0.8.2 / 2023-06-15
+------------------
+
 - Extend OARCluster implementation to let OAR take into account the memory parameter (:pr:`598`, :pr:`595`)
 
 0.8.1 / 2022-10-04


Test patch:
diff --git a/dask_jobqueue/tests/test_htcondor.py b/dask_jobqueue/tests/test_htcondor.py
index c5c5bf6..653b72d 100644
--- a/dask_jobqueue/tests/test_htcondor.py
+++ b/dask_jobqueue/tests/test_htcondor.py
@@ -141,6 +141,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_jobqueue_core.py b/dask_jobqueue/tests/test_jobqueue_core.py
index 8f743f7..23d2059 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:
diff --git a/dask_jobqueue/tests/test_lsf.py b/dask_jobqueue/tests/test_lsf.py
index 40c826c..134599b 100644
--- a/dask_jobqueue/tests/test_lsf.py
+++ b/dask_jobqueue/tests/test_lsf.py
@@ -101,9 +101,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",
@@ -130,9 +130,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",
@@ -322,6 +322,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..565d955 100644
--- a/dask_jobqueue/tests/test_pbs.py
+++ b/dask_jobqueue/tests/test_pbs.py
@@ -76,9 +76,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",
@@ -102,9 +102,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")
@@ -361,6 +361,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..d1b5d2f 100644
--- a/dask_jobqueue/tests/test_sge.py
+++ b/dask_jobqueue/tests/test_sge.py
@@ -58,6 +58,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..c2aaed3 100644
--- a/dask_jobqueue/tests/test_slurm.py
+++ b/dask_jobqueue/tests/test_slurm.py
@@ -77,9 +77,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 +111,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")
@@ -193,6 +193,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": [],