Review SLURM configuration and include new configuration option for pre-sbatch commands and gpus

CHANGELOG.md

@@ -1,5 +1,18 @@
 **Note**: Numbers like (\#1234) point to closed Pull Requests on the fractal-server repository.
 
+# 2.3.5
+
+> WARNING: The `pre_submission_commands` SLURM configuration is included as an
+> experimental feature, since it is still not useful for its main intended
+> goal (calling `module load` before running `sbatch`).
+
+* SLURM runners
+    * Expose `gpus` SLURM parameter (\#1678).
+    * For SSH executor, add `pre_submission_commands` (\#1678).
+    * Removed obsolete arguments from `get_slurm_config` function (\#1678).
+* SSH features:
+    * Add `FractalSSH.write_remote_file` method (\#1678).
+
 # 2.3.4
 
 * SSH SLURM runner:

benchmarks/runner/benchmark_runner.py

@@ -59,12 +59,12 @@ def mock_venv(tmp_path: str) -> dict:
             args[
                 "command_non_parallel"
             ] = f"{python} {src_dir / task['executable_non_parallel']}"
-            args["meta_non_paralell"] = task.get("meta_non_paralell")
+            args["meta_non_parallel"] = task.get("meta_non_parallel")
         if task.get("executable_parallel"):
             args[
                 "command_parallel"
             ] = f"{python} {src_dir / task['executable_parallel']}"
-            args["meta_paralell"] = task.get("meta_paralell")
+            args["meta_parallel"] = task.get("meta_parallel")
 
         t = TaskV2Mock(
             id=ind,

benchmarks/runner/mocks.py

@@ -38,8 +38,8 @@ class TaskV2Mock(BaseModel):
 
     command_non_parallel: Optional[str] = None
     command_parallel: Optional[str] = None
-    meta_non_paralell: Optional[dict[str, Any]] = Field(default_factory=dict)
-    meta_paralell: Optional[dict[str, Any]] = Field(default_factory=dict)
+    meta_non_parallel: Optional[dict[str, Any]] = Field(default_factory=dict)
+    meta_parallel: Optional[dict[str, Any]] = Field(default_factory=dict)
     type: Optional[str]
 
     @root_validator(pre=False)

fractal_server/app/runner/executors/slurm/_slurm_config.py

@@ -62,6 +62,8 @@ class _SlurmConfigSet(BaseModel, extra=Extra.forbid):
     time: Optional[str]
     account: Optional[str]
     extra_lines: Optional[list[str]]
+    pre_submission_commands: Optional[list[str]]
+    gpus: Optional[str]
 
 
 class _BatchingConfigSet(BaseModel, extra=Extra.forbid):
@@ -219,6 +221,7 @@ class SlurmConfig(BaseModel, extra=Extra.forbid):
         constraint: Corresponds to SLURM option.
         gres: Corresponds to SLURM option.
         account: Corresponds to SLURM option.
+        gpus: Corresponds to SLURM option.
         time: Corresponds to SLURM option (WARNING: not fully supported).
         prefix: Prefix of configuration lines in SLURM submission scripts.
         shebang_line: Shebang line for SLURM submission scripts.
@@ -240,6 +243,8 @@ class SlurmConfig(BaseModel, extra=Extra.forbid):
             Key-value pairs to be included as `export`-ed variables in SLURM
             submission script, after prepending values with the user's cache
             directory.
+        pre_submission_commands: List of commands to be prepended to the sbatch
+            command.
     """
 
     # Required SLURM parameters (note that the integer attributes are those
@@ -254,6 +259,7 @@ class SlurmConfig(BaseModel, extra=Extra.forbid):
     job_name: Optional[str] = None
     constraint: Optional[str] = None
     gres: Optional[str] = None
+    gpus: Optional[str] = None
     time: Optional[str] = None
     account: Optional[str] = None
 
@@ -274,6 +280,8 @@ class SlurmConfig(BaseModel, extra=Extra.forbid):
     target_num_jobs: int
     max_num_jobs: int
 
+    pre_submission_commands: list[str] = Field(default_factory=list)
+
     def _sorted_extra_lines(self) -> list[str]:
         """
         Return a copy of `self.extra_lines`, where lines starting with
@@ -340,7 +348,14 @@ def to_sbatch_preamble(
             f"{self.prefix} --cpus-per-task={self.cpus_per_task}",
             f"{self.prefix} --mem={mem_per_job_MB}M",
         ]
-        for key in ["job_name", "constraint", "gres", "time", "account"]:
+        for key in [
+            "job_name",
+            "constraint",
+            "gres",
+            "gpus",
+            "time",
+            "account",
+        ]:
             value = getattr(self, key)
             if value is not None:
                 # Handle the `time` parameter

fractal_server/app/runner/executors/slurm/ssh/executor.py

@@ -869,9 +869,22 @@ def _submit_job(self, job: SlurmJob) -> tuple[Future, str]:
 
         # Submit job to SLURM, and get jobid
         sbatch_command = f"sbatch --parsable {job.slurm_script_remote}"
-        sbatch_stdout = self.fractal_ssh.run_command(
-            cmd=sbatch_command,
-        )
+        pre_submission_cmds = job.slurm_config.pre_submission_commands
+        if len(pre_submission_cmds) == 0:
+            sbatch_stdout = self.fractal_ssh.run_command(cmd=sbatch_command)
+        else:
+            logger.debug(f"Now using {pre_submission_cmds=}")
+            script_lines = pre_submission_cmds + [sbatch_command]
+            script_content = "\n".join(script_lines)
+            script_content = f"{script_content}\n"
+            script_path_remote = (
+                f"{job.slurm_script_remote.as_posix()}_wrapper.sh"
+            )
+            self.fractal_ssh.write_remote_file(
+                path=script_path_remote, content=script_content
+            )
+            cmd = f"bash {script_path_remote}"
+            sbatch_stdout = self.fractal_ssh.run_command(cmd=cmd)
 
         # Extract SLURM job ID from stdout
         try:
@@ -881,7 +894,9 @@ def _submit_job(self, job: SlurmJob) -> tuple[Future, str]:
             error_msg = (
                 f"Submit command `{sbatch_command}` returned "
                 f"`{stdout=}` which cannot be cast to an integer "
-                f"SLURM-job ID. Original error:\n{str(e)}"
+                f"SLURM-job ID.\n"
+                f"Note that {pre_submission_cmds=}.\n"
+                f"Original error:\n{str(e)}"
             )
             logger.error(error_msg)
             raise JobExecutionError(info=error_msg)

fractal_server/app/runner/executors/slurm/sudo/executor.py

@@ -1121,6 +1121,12 @@
             slurm_err_path=str(job.slurm_stderr),
         )
 
+        # Print warning for ignored parameter
+        if len(job.slurm_config.pre_submission_commands) > 0:
+            logger.warning(
+                f"Ignoring {job.slurm_config.pre_submission_commands=}."
+            )
+
         # Submit job via sbatch, and retrieve jobid
 
         # Write script content to a job.slurm_script

fractal_server/app/runner/v2/__init__.py

@@ -36,8 +36,8 @@
 from ._local_experimental import (
     process_workflow as local_experimental_process_workflow,
 )
-from ._slurm import process_workflow as slurm_sudo_process_workflow
 from ._slurm_ssh import process_workflow as slurm_ssh_process_workflow
+from ._slurm_sudo import process_workflow as slurm_sudo_process_workflow
 from .handle_failed_job import assemble_filters_failed_job
 from .handle_failed_job import assemble_history_failed_job
 from .handle_failed_job import assemble_images_failed_job

fractal_server/app/runner/v2/_slurm_ssh/get_slurm_config.py → fractal_server/app/runner/v2/_slurm_common/get_slurm_config.py

@@ -18,8 +18,6 @@
 
 def get_slurm_config(
     wftask: WorkflowTaskV2,
-    workflow_dir_local: Path,
-    workflow_dir_remote: Path,
     which_type: Literal["non_parallel", "parallel"],
     config_path: Optional[Path] = None,
 ) -> SlurmConfig:
@@ -43,13 +41,6 @@ def get_slurm_config(
         wftask:
             WorkflowTask for which the SLURM configuration is is to be
             prepared.
-        workflow_dir_local:
-            Server-owned directory to store all task-execution-related relevant
-            files (inputs, outputs, errors, and all meta files related to the
-            job execution). Note: users cannot write directly to this folder.
-        workflow_dir_remote:
-            User-side directory with the same scope as `workflow_dir_local`,
-            and where a user can write.
         config_path:
             Path of a Fractal SLURM configuration file; if `None`, use
             `FRACTAL_SLURM_CONFIG_FILE` variable from settings.
@@ -99,13 +90,13 @@ def get_slurm_config(
     # 1. This block of definitions takes priority over other definitions from
     #    slurm_env which are not under the `needs_gpu` subgroup
     # 2. This block of definitions has lower priority than whatever comes next
-    #    (i.e. from WorkflowTask.meta).
+    #    (i.e. from WorkflowTask.meta_parallel).
     if wftask_meta is not None:
         needs_gpu = wftask_meta.get("needs_gpu", False)
     else:
         needs_gpu = False
     logger.debug(f"[get_slurm_config] {needs_gpu=}")
-    if needs_gpu and slurm_env.gpu_slurm_config is not None:  # FIXME
+    if needs_gpu:
         for key, value in slurm_env.gpu_slurm_config.dict(
             exclude_unset=True, exclude={"mem"}
         ).items():
@@ -143,9 +134,9 @@ def get_slurm_config(
             )
             logger.error(error_msg)
             raise SlurmConfigError(error_msg)
-        for key in ["time", "gres", "constraint"]:
+        for key in ["time", "gres", "gpus", "constraint"]:
             value = wftask_meta.get(key, None)
-            if value:
+            if value is not None:
                 slurm_dict[key] = value
     if wftask_meta is not None:
         extra_lines = wftask_meta.get("extra_lines", [])