Pass configured environment variables to singularity

cluster_utils/job.py

@@ -4,9 +4,8 @@
 import os
 import pathlib
 import time
-import typing
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import pandas as pd
 
@@ -103,15 +102,19 @@ def generate_execution_cmd(self, paths, cmd_prefix: Optional[str] = None):
         set_cwd = "cd {}".format(paths["main_path"])
 
         if "variables" in paths:
-            if not isinstance(paths, dict):
+            if not isinstance(paths["variables"], dict):
                 raise ValueError(
                     'Expected type dict for "variables", but got type'
                     f' {type(paths["variables"])} instead'
                 )
+            env_variables = {
+                str(name): str(value) for name, value in paths["variables"].items()
+            }
             set_env_variables = "\n".join(
-                f"export {name}={value}" for name, value in paths["variables"].items()
+                f'export {name}="{value}"' for name, value in env_variables.items()
             )
         else:
+            env_variables = None
             set_env_variables = ""
 
         if "pre_job_script" in paths:
@@ -181,6 +184,7 @@ def generate_execution_cmd(self, paths, cmd_prefix: Optional[str] = None):
                 self.singularity_settings,
                 paths["main_path"],
                 current_setting["working_dir"],
+                env_variables,
             )
 
         if cmd_prefix:
@@ -189,10 +193,10 @@ def generate_execution_cmd(self, paths, cmd_prefix: Optional[str] = None):
         res = "\n".join(
             [
                 set_cwd,
-                pre_job_script,
                 virtual_env_activate,
                 conda_env_activate,
                 set_env_variables,
+                pre_job_script,
                 exec_cmd,
             ]
         )
@@ -202,8 +206,9 @@ def singularity_wrap(
         self,
         exec_cmd: str,
         singularity_settings: SingularitySettings,
-        exec_dir: typing.Union[str, os.PathLike],
-        working_dir: typing.Union[str, os.PathLike],
+        exec_dir: Union[str, os.PathLike],
+        working_dir: Union[str, os.PathLike],
+        env_variables: Optional[Dict[str, str]],
     ) -> str:
         """Wrap the given command to execute it in a Singularity container.
 
@@ -223,6 +228,9 @@ def singularity_wrap(
         # create model directory (so it can be bound into the container)
         working_dir.mkdir(exist_ok=True)
 
+        if env_variables is None:
+            env_variables = {}
+
         # construct singularity command
         cwd = os.fspath(exec_dir)
         bind_dirs = ["/tmp", os.fspath(working_dir), cwd]
@@ -231,6 +239,9 @@ def singularity_wrap(
             "run" if singularity_settings.use_run else "exec",
             "--bind=%s" % ",".join(bind_dirs),
             "--pwd=%s" % cwd,
+            " ".join(
+                f'--env {name}="{value}"' for name, value in env_variables.items()
+            ),
             *singularity_settings.args,
             os.fspath(singularity_image),
         ]

docs/configuration.rst

@@ -86,7 +86,7 @@
     that when running on the cluster this directory also contains the stdout/stderr of
     the jobs (but not when running locally).
 
 .. confval:: remove_working_dirs: bool = {grid_search: false, hp_optimization: true}
 
     Remove the working directories of the jobs (including the parameters used for that
     job, saved metrics and potentially other output files like checkpoints) once they
@@ -134,14 +134,17 @@
 
     .. confval:: environment_setup.variables: dict[str]:
 
-        Environment variables to set. Variables are set after a virtual/conda
-        environment is activated, thus override environment variables set before.
+        Environment variables to set. Variables are set *after* a virtual/conda environment 
+        is activated, thus override environment variables set before. They are also set 
+        *before* the :confval:`environment_setup.pre_job_script`: this can be useful to pass 
+        parameters to the script, e.g. to setup a generic script that changes its behavior based 
+        on the values defined in the cluster_utils config file.
 
     .. confval:: environment_setup.is_python_script: bool, default=true:
 
         Whether the target to run is a Python script.
 
     .. confval:: environment_setup.run_as_module: bool, default=false:
 
         Whether to run the script as a Python module
         (``python -m my_package.my_module``) or as a script
@@ -319,7 +322,7 @@
     ``--signal`` of ``sbatch``).  If not set, no signal is sent.
 
     See example :doc:`examples/slurm_timeout_signal`.
 .. confval:: cluster_requirements.extra_submission_options: list[str]
 
     List of additional options for ``sbatch``.  Can be used if a specific
     setting is needed which is not already covered by the options above.
@@ -408,7 +411,7 @@
     **Required.**
 
     The optimisation method that is used to find good hyperparameters.
-    Supported methods are 
+    Supported methods are
 
     - cem_metaoptimizer
     - nevergrad \*
@@ -452,7 +455,7 @@
          * - IntNormal
            - Normal distribution using integer values.
          * - IntLogNormal
-           - Log-normal distribution using integer values. 
+           - Log-normal distribution using integer values.
          * - Discrete
            - Discrete list of values.
     - ``bounds``:  List ``[min_value, max_value]``