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

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,7 @@ class _SlurmConfigSet(BaseModel, extra=Extra.forbid):
     time: Optional[str]
     account: Optional[str]
     extra_lines: Optional[list[str]]
+    pre_submission_commands: Optional[list[str]]
 
 
 class _BatchingConfigSet(BaseModel, extra=Extra.forbid):
@@ -240,6 +241,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
@@ -274,6 +277,9 @@ class SlurmConfig(BaseModel, extra=Extra.forbid):
     target_num_jobs: int
     max_num_jobs: int
 
+    # FIXME
+    pre_submission_cmds: list[str] = Field(default_factory=list)
+
     def _sorted_extra_lines(self) -> list[str]:
         """
         Return a copy of `self.extra_lines`, where lines starting with

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

@@ -869,15 +869,30 @@
 
         # 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_cmds
+        if len(pre_submission_cmds) == 0:
+            sbatch_stdout = self.fractal_ssh.run_command(cmd=sbatch_command)
+        else:
+            logger.warning(f"Now using {pre_submission_cmds=}")
+            script_lines = pre_submission_cmds + [sbatch_command]
+            script_content = "\n".join(script_lines)
+            script_path_remote = (
+                f"{job.slurm_script_remote.as_posix()}_wrapper.sh"
+            )
+            with self.fractal_ssh.sftp().open(
+                filename=script_path_remote, mode="w"
+            ) as f:
+                f.write(script_content)
+            cmd = f"bash {script_path_remote}"
+            sbatch_stdout = self.fractal_ssh.run_command(cmd=cmd)
 
         # Extract SLURM job ID from stdout
         try:
             stdout = sbatch_stdout.strip("\n")
             jobid = int(stdout)
         except ValueError as e:
+            if len(pre_submission_cmds) > 0:
+                logger.warning(f"I used {pre_submission_cmds=}")
             error_msg = (
                 f"Submit command `{sbatch_command}` returned "
                 f"`{stdout=}` which cannot be cast to an integer "

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 is not None:
         for key, value in slurm_env.gpu_slurm_config.dict(
             exclude_unset=True, exclude={"mem"}
         ).items():

fractal_server/app/runner/v2/_slurm_ssh/_submit_setup.py

@@ -17,8 +17,10 @@
 from typing import Literal
 
 from ...task_files import get_task_file_paths
-from .get_slurm_config import get_slurm_config
 from fractal_server.app.models.v2 import WorkflowTaskV2
+from fractal_server.app.runner.v2._slurm_common.get_slurm_config import (
+    get_slurm_config,
+)
 
 
 def _slurm_submit_setup(
@@ -62,8 +64,6 @@ def _slurm_submit_setup(
     # Get SlurmConfig object
     slurm_config = get_slurm_config(
         wftask=wftask,
-        workflow_dir_local=workflow_dir_local,
-        workflow_dir_remote=workflow_dir_remote,
         which_type=which_type,
     )
 

fractal_server/app/runner/v2/_slurm/_submit_setup.py → fractal_server/app/runner/v2/_slurm_sudo/_submit_setup.py

@@ -17,8 +17,10 @@
 from typing import Literal
 
 from ...task_files import get_task_file_paths
-from .get_slurm_config import get_slurm_config
 from fractal_server.app.models.v2 import WorkflowTaskV2
+from fractal_server.app.runner.v2._slurm_common.get_slurm_config import (
+    get_slurm_config,
+)
 
 
 def _slurm_submit_setup(
@@ -62,8 +64,6 @@ def _slurm_submit_setup(
     # Get SlurmConfig object
     slurm_config = get_slurm_config(
         wftask=wftask,
-        workflow_dir_local=workflow_dir_local,
-        workflow_dir_remote=workflow_dir_remote,
         which_type=which_type,
     )
 

tests/v2/00_ssh/test_FractalSSH.py

@@ -258,3 +258,23 @@ def test_remove_folder_input_validation():
             safe_root="/actual_root",
         )
     print(e.value)
+
+
+def test_run_command_session_persistence(
+    fractal_ssh: FractalSSH, tmp777_path: Path
+):
+    """
+    FIXME add explanation
+    """
+
+    script1 = tmp777_path / "set_var.sh"
+    with script1.open("w") as f:
+        f.write("export MYVAR=123\necho $MYVAR")
+    script2 = tmp777_path / "read_var.sh"
+    with script2.open("w") as f:
+        f.write("echo $MYVAR")
+
+    stdout = fractal_ssh.run_command(cmd=f"bash {script1.as_posix()}")
+    assert stdout.strip("\n") == "123"
+    stdout = fractal_ssh.run_command(cmd=f"bash {script2.as_posix()}")
+    assert stdout.strip("\n") == ""