Repository: dask/dask-jobqueue
Pull request #575 (2022-08-11 13:02:33)

rename `env_extra`
Renamed `env_extra` to `job_script_prologue` and added a FutureWarning for `env_extra`.

If `job_script_prologue` is empty,  `env_extra` will still be respected (see new tests).

Part of resolving #323


Patch:
diff --git a/dask_jobqueue/core.py b/dask_jobqueue/core.py
index dbf7a5a..3094c7e 100644
--- a/dask_jobqueue/core.py
+++ b/dask_jobqueue/core.py
@@ -49,6 +49,8 @@ job_parameters = """
     extra : list
         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.
+    job_script_prologue : list
         Other commands to add to script before launching worker.
     header_skip : list
         Lines to skip in the header.
@@ -129,7 +131,7 @@ class Job(ProcessInterface, abc.ABC):
 %(shebang)s
 
 %(job_header)s
-%(env_header)s
+%(job_script_prologue)s
 %(worker_command)s
 """.lstrip()
 
@@ -155,6 +157,7 @@ class Job(ProcessInterface, abc.ABC):
         local_directory=None,
         extra=None,
         env_extra=None,
+        job_script_prologue=None,
         header_skip=None,
         log_directory=None,
         shebang=None,
@@ -207,6 +210,21 @@ class Job(ProcessInterface, abc.ABC):
             extra = dask.config.get("jobqueue.%s.extra" % self.config_name)
         if env_extra is None:
             env_extra = dask.config.get("jobqueue.%s.env-extra" % self.config_name)
+        if job_script_prologue is None:
+            job_script_prologue = dask.config.get(
+                "jobqueue.%s.job-script-prologue" % self.config_name
+            )
+        if env_extra is not None:
+            warn = (
+                "env_extra has been renamed to job_script_prologue. "
+                "You are still using it (even if only set to []; please also check config files)."
+                "If you did not set job_script_prologue yet, env_extra will be respected for now, "
+                "but it will be removed in a future release. "
+                "If you already set job_script_prologue, env_extra is ignored and you can remove it."
+            )
+            warnings.warn(warn, FutureWarning)
+            if not job_script_prologue:
+                job_script_prologue = env_extra
         if header_skip is None:
             header_skip = dask.config.get(
                 "jobqueue.%s.header-skip" % self.config_name, ()
@@ -245,7 +263,7 @@ class Job(ProcessInterface, abc.ABC):
 
         self.shebang = shebang
 
-        self._env_extra = env_extra
+        self._job_script_prologue = job_script_prologue
         self.header_skip = set(header_skip)
 
         # dask-worker command line build
@@ -299,7 +317,7 @@ class Job(ProcessInterface, abc.ABC):
         pieces = {
             "shebang": self.shebang,
             "job_header": header,
-            "env_header": "\n".join(filter(None, self._env_extra)),
+            "job_script_prologue": "\n".join(filter(None, self._job_script_prologue)),
             "worker_command": self._command_template,
         }
         return self._script_template % pieces
diff --git a/dask_jobqueue/htcondor.py b/dask_jobqueue/htcondor.py
index b3dcfd0..f3d5a96 100644
--- a/dask_jobqueue/htcondor.py
+++ b/dask_jobqueue/htcondor.py
@@ -61,12 +61,12 @@ Queue
         else:
             self.job_extra = job_extra
 
-        if self._env_extra is not None:
-            # Overwrite command template: prepend commands from env_extra separated by semicolon.
+        if self._job_script_prologue is not None:
+            # Overwrite command template: prepend commands from job_script_prologue separated by semicolon.
             # This is special for HTCondor, because lines to execute on the worker node cannot be
             # simply added to the submit script like for other batch systems.
             self._command_template = "; ".join(
-                self._env_extra + [self._command_template]
+                self._job_script_prologue + [self._command_template]
             )
 
         self.job_header_dict = {
@@ -246,12 +246,12 @@ class HTCondorCluster(JobQueueCluster):
 
     >>> cluster.adapt(maximum_jobs=20)
 
-    If setup commands need to be run before starting the worker on the worker node, ``env_extra`` can be used,
+    If setup commands need to be run before starting the worker on the worker node, ``job_script_prologue`` can be used,
     e.g., to activate a virtual environment:
 
     >>> from dask_jobqueue.htcondor import HTCondorCluster
     >>> cluster = HTCondorCluster(cores=1, memory="2GB", disk="4GB",
-                                  env_extra=['cd /some/path/', 'source venv/bin/activate'])
+                                  job_script_prologue=['cd /some/path/', 'source venv/bin/activate'])
 
     Note that environment variables are no longer passed via the ``Environment`` parameter in the submit
     description file. If you explictly want to set that, you need to use ``job_extra``.
diff --git a/dask_jobqueue/jobqueue.yaml b/dask_jobqueue/jobqueue.yaml
index 85da804..568345b 100644
--- a/dask_jobqueue/jobqueue.yaml
+++ b/dask_jobqueue/jobqueue.yaml
@@ -18,7 +18,8 @@ jobqueue:
     queue: null
     project: null
     walltime: '00:30:00'
-    env-extra: []
+    env-extra: null
+    job-script-prologue: []
     resource-spec: null
     job-extra: []
     log-directory: null
@@ -45,7 +46,8 @@ jobqueue:
     queue: null
     project: null
     walltime: '00:30:00'
-    env-extra: []
+    env-extra: null
+    job-script-prologue: []
     resource-spec: null
     job-extra: []
     log-directory: null
@@ -72,7 +74,8 @@ jobqueue:
     queue: null
     project: null
     walltime: '00:30:00'
-    env-extra: []
+    env-extra: null
+    job-script-prologue: []
     job-extra: []
     log-directory: null
     resource-spec: null
@@ -99,7 +102,8 @@ jobqueue:
     queue: null
     project: null
     walltime: '00:30:00'
-    env-extra: []
+    env-extra: null
+    job-script-prologue: []
     job-cpu: null
     job-mem: null
     job-extra: []
@@ -127,7 +131,8 @@ jobqueue:
     queue: null
     project: null
     walltime: '00:30:00'
-    env-extra: []
+    env-extra: null
+    job-script-prologue: []
     resource-spec: null
     job-extra: []
     log-directory: null
@@ -154,7 +159,8 @@ jobqueue:
     queue: null
     project: null
     walltime: '00:30'
-    env-extra: []
+    env-extra: null
+    job-script-prologue: []
     ncpus: null
     mem: null
     job-extra: []
@@ -181,7 +187,8 @@ jobqueue:
 
     # HTCondor Resource Manager options
     disk: null                  # Total amount of disk per job
-    env-extra: []
+    env-extra: null
+    job-script-prologue: []
     job-extra: {}               # Extra submit attributes
     submit-command-extra: []    # Extra condor_submit arguments
     cancel-command-extra: []    # Extra condor_rm arguments
@@ -204,7 +211,8 @@ jobqueue:
     shared-temp-directory: null       # Shared directory currently used to dump temporary security objects for workers
     extra: []
 
-    env-extra: []
+    env-extra: null
+    job-script-prologue: []
     job-extra: []
     log-directory: null
     
diff --git a/docs/source/advanced-tips-and-tricks.rst b/docs/source/advanced-tips-and-tricks.rst
index 237adc6..be635a1 100644
--- a/docs/source/advanced-tips-and-tricks.rst
+++ b/docs/source/advanced-tips-and-tricks.rst
@@ -68,21 +68,24 @@ accepted option on some SLURM clusters. The error was something like this:
   sbatch: error: Memory specification can not be satisfied
   sbatch: error: Batch job submission failed: Requested node configuration is not available
 
-Run setup commands before starting the worker with ``env_extra``
-----------------------------------------------------------------
+Run setup commands before starting the worker with ``job_script_prologue``
+--------------------------------------------------------------------------
+
+*Note: the parameter* ``job_script_prologue`` *was named* ``env_extra`` *until version 0.7.4.* ``env_extra`` *can still
+be used, but is considered deprecated and will be removed in a future version.*
 
 Sometimes you need to run some setup commands before the actual worker can be started. This includes
 setting environment variables, loading environment modules, sourcing/activating a virtual environment,
 or activating conda/mamba environments.
 
-This can be achieved using the ``env_extra`` parameter. Example for setting up a virtual environment:
+This can be achieved using the ``job_script_prologue`` parameter. Example for setting up a virtual environment:
 
 .. code-block:: python
 
    from dask_jobqueue.htcondor import HTCondorCluster
-   env_extra = ['cd /some/path', 'source venv/bin/activate']
+   job_script_prologue = ['cd /some/path', 'source venv/bin/activate']
    cluster = HTCondorCluster(cores=1, memory="2GB", disk="4GB", log_directory = 'logs', python='python3',
-                             env_extra=env_extra)
+                             job_script_prologue=job_script_prologue)
    print(cluster.job_script())
 
 For ``HTCondorCluster``, the commands will be prepended to the actual python call in the ``Arguments``
diff --git a/docs/source/debug.rst b/docs/source/debug.rst
index 7e22bb7..b48056c 100644
--- a/docs/source/debug.rst
+++ b/docs/source/debug.rst
@@ -27,7 +27,7 @@ a real job script file, and submit it using ``qsub``, ``sbatch``, ``bsub`` or
 what is appropriate for you job queuing system.
 
 To correct any problem detected at this point, you could try to use
-``job_extra`` or ``env_extra`` kwargs when initializing your cluster object.
+``job_extra`` or ``job_script_prologue`` kwargs when initializing your cluster object.
 
 In particular, pay attention to the python executable used to launch the
 workers, which by default is the one used to launch the scheduler (this makes
diff --git a/docs/source/examples.rst b/docs/source/examples.rst
index 4f9a382..5552787 100644
--- a/docs/source/examples.rst
+++ b/docs/source/examples.rst
@@ -129,9 +129,11 @@ SLURM Deployment: Low-priority node usage
                            memory="16GB",
                            project="co_laika",
                            queue='savio2_bigmem',
-                           env_extra=['export LANG="en_US.utf8"',
-                                      'export LANGUAGE="en_US.utf8"',
-                                      'export LC_ALL="en_US.utf8"'],
+                           job_script_prologue=[
+                               'export LANG="en_US.utf8"',
+                               'export LANGUAGE="en_US.utf8"',
+                               'export LC_ALL="en_US.utf8"'
+                           ],
                            job_extra=['--qos="savio_lowprio"'])
 
 


Test patch:
diff --git a/dask_jobqueue/tests/test_htcondor.py b/dask_jobqueue/tests/test_htcondor.py
index 2147e21..41aa838 100644
--- a/dask_jobqueue/tests/test_htcondor.py
+++ b/dask_jobqueue/tests/test_htcondor.py
@@ -27,7 +27,7 @@ def test_job_script():
         processes=2,
         memory="100MB",
         disk="100MB",
-        env_extra=[
+        job_script_prologue=[
             'export LANG="en_US.utf8"',
             'export LC_ALL="en_US.utf8"',
             "cd /some/path/",
@@ -140,7 +140,8 @@ def test_config_name_htcondor_takes_custom_config():
         "interface": None,
         "death-timeout": None,
         "extra": [],
-        "env-extra": [],
+        "env-extra": None,
+        "job-script-prologue": [],
         "log-directory": None,
         "shebang": "#!/usr/bin/env condor_submit",
         "local-directory": "/tmp",
diff --git a/dask_jobqueue/tests/test_job.py b/dask_jobqueue/tests/test_job.py
index 5458b78..c20b62b 100644
--- a/dask_jobqueue/tests/test_job.py
+++ b/dask_jobqueue/tests/test_job.py
@@ -143,3 +143,67 @@ async def test_nworkers_scale():
 def test_docstring_cluster(Cluster):
     assert "cores :" in Cluster.__doc__
     assert Cluster.__name__[: -len("Cluster")] in Cluster.__doc__
+
+
+def test_deprecation(Cluster):
+    import warnings
+
+    # test issuing of warning
+    warnings.simplefilter("always")
+
+    job_cls = Cluster.job_cls
+    with warnings.catch_warnings(record=True) as w:
+        # should give a warning
+        job = job_cls(cores=1, memory="1 GB", env_extra=["env_extra is used"])
+        assert len(w) == 1
+        assert issubclass(w[0].category, FutureWarning)
+        assert "env_extra has been renamed" in str(w[0].message)
+    with warnings.catch_warnings(record=True) as w:
+        # should give a warning
+        job = job_cls(
+            cores=1,
+            memory="1 GB",
+            env_extra=["env_extra is used"],
+            job_script_prologue=["job_script_prologue is used"],
+        )
+        assert len(w) == 1
+        assert issubclass(w[0].category, FutureWarning)
+        assert "env_extra has been renamed" in str(w[0].message)
+    with warnings.catch_warnings(record=True) as w:
+        # should not give a warning
+        job = job_cls(
+            cores=1,
+            memory="1 GB",
+            job_script_prologue=["job_script_prologue is used, env_extra not"],
+        )
+        assert len(w) == 0
+
+    # the rest is not about the warning but about behaviour: if job_script_prologue is not
+    # set, env_extra should still be used if provided
+    warnings.simplefilter("ignore")
+    job = job_cls(
+        cores=1,
+        memory="1 GB",
+        env_extra=["env_extra"],
+        job_script_prologue=["job_script_prologue"],
+    )
+    job_script = job.job_script()
+    assert "env_extra" not in job_script
+    assert "job_script_prologue" in job_script
+
+    job = job_cls(
+        cores=1,
+        memory="1 GB",
+        env_extra=["env_extra"],
+    )
+    job_script = job.job_script()
+    assert "env_extra" in job_script
+
+    job = job_cls(
+        cores=1,
+        memory="1 GB",
+        env_extra=["env_extra"],
+        job_script_prologue=[],
+    )
+    job_script = job.job_script()
+    assert "env_extra" in job_script
diff --git a/dask_jobqueue/tests/test_lsf.py b/dask_jobqueue/tests/test_lsf.py
index 20b47c8..f236ac4 100644
--- a/dask_jobqueue/tests/test_lsf.py
+++ b/dask_jobqueue/tests/test_lsf.py
@@ -321,7 +321,8 @@ def test_config_name_lsf_takes_custom_config():
         "local-directory": "/foo",
         "shared-temp-directory": None,
         "extra": [],
-        "env-extra": [],
+        "env-extra": None,
+        "job-script-prologue": [],
         "log-directory": None,
         "shebang": "#!/usr/bin/env bash",
         "use-stdin": None,
diff --git a/dask_jobqueue/tests/test_oar.py b/dask_jobqueue/tests/test_oar.py
index 0e156e8..a57b3dd 100644
--- a/dask_jobqueue/tests/test_oar.py
+++ b/dask_jobqueue/tests/test_oar.py
@@ -64,7 +64,7 @@ def test_job_script():
         processes=4,
         cores=8,
         memory="28GB",
-        env_extra=[
+        job_script_prologue=[
             'export LANG="en_US.utf8"',
             'export LANGUAGE="en_US.utf8"',
             'export LC_ALL="en_US.utf8"',
@@ -109,7 +109,8 @@ def test_config_name_oar_takes_custom_config():
         "local-directory": "/foo",
         "shared-temp-directory": None,
         "extra": [],
-        "env-extra": [],
+        "env-extra": None,
+        "job-script-prologue": [],
         "log-directory": None,
         "shebang": "#!/usr/bin/env bash",
         "job-cpu": None,
diff --git a/dask_jobqueue/tests/test_pbs.py b/dask_jobqueue/tests/test_pbs.py
index cb6973e..b62de96 100644
--- a/dask_jobqueue/tests/test_pbs.py
+++ b/dask_jobqueue/tests/test_pbs.py
@@ -360,7 +360,8 @@ def test_config_name_pbs_takes_custom_config():
         "local-directory": "/foo",
         "shared-temp-directory": None,
         "extra": [],
-        "env-extra": [],
+        "env-extra": None,
+        "job-script-prologue": [],
         "log-directory": None,
         "shebang": "#!/usr/bin/env bash",
         "job-cpu": None,
diff --git a/dask_jobqueue/tests/test_sge.py b/dask_jobqueue/tests/test_sge.py
index d24502c..4977705 100644
--- a/dask_jobqueue/tests/test_sge.py
+++ b/dask_jobqueue/tests/test_sge.py
@@ -57,7 +57,8 @@ def test_config_name_sge_takes_custom_config():
         "local-directory": "/foo",
         "shared-temp-directory": None,
         "extra": [],
-        "env-extra": [],
+        "env-extra": None,
+        "job-script-prologue": [],
         "log-directory": None,
         "shebang": "#!/usr/bin/env bash",
         "job-cpu": None,
@@ -79,7 +80,7 @@ def test_job_script(tmpdir):
         queue="my-queue",
         project="my-project",
         walltime="02:00:00",
-        env_extra=["export MY_VAR=my_var"],
+        job_script_prologue=["export MY_VAR=my_var"],
         job_extra=["-w e", "-m e"],
         log_directory=log_directory,
         resource_spec="h_vmem=12G,mem_req=12G",
diff --git a/dask_jobqueue/tests/test_slurm.py b/dask_jobqueue/tests/test_slurm.py
index 74fca2b..1f8701f 100644
--- a/dask_jobqueue/tests/test_slurm.py
+++ b/dask_jobqueue/tests/test_slurm.py
@@ -86,7 +86,7 @@ def test_job_script():
         processes=4,
         cores=8,
         memory="28GB",
-        env_extra=[
+        job_script_prologue=[
             'export LANG="en_US.utf8"',
             'export LANGUAGE="en_US.utf8"',
             'export LC_ALL="en_US.utf8"',
@@ -192,7 +192,8 @@ def test_config_name_slurm_takes_custom_config():
         "local-directory": "/foo",
         "shared-temp-directory": None,
         "extra": [],
-        "env-extra": [],
+        "env-extra": None,
+        "job-script-prologue": [],
         "log-directory": None,
         "shebang": "#!/usr/bin/env bash",
         "job-cpu": None,
@@ -232,7 +233,7 @@ def test_worker_name_uses_cluster_name(loop):
         cores=1,
         memory="2GB",
         name="test-$MY_ENV_VARIABLE",
-        env_extra=["MY_ENV_VARIABLE=my-env-variable-value"],
+        job_script_prologue=["MY_ENV_VARIABLE=my-env-variable-value"],
         loop=loop,
     ) as cluster:
         with Client(cluster) as client: