feat: remote stage in

.github/workflows/ci.yml

@@ -24,8 +24,14 @@ jobs:
         run: |
           poetry install --sync
 
+      - name: Debug black environment
+        continue-on-error: true
+        run: |
+          poetry run black --version
+          poetry run black --check --diff --target-version py311 .
+
       - name: Check formatting
-        run: poetry run black --check .
+        run: poetry run black --check --target-version py311 .
 
   linting:
     runs-on: ubuntu-latest

docs/further.md

@@ -190,6 +190,30 @@ The scoring algorithm calculates a score by summing the ratios of requested reso
 
 If no suitable partition is found based on the job's resource requirements, the plugin falls back to the default SLURM behavior, which typically uses the cluster's default partition or any partition specified explicitly in the job's resources.
 
+#### Automatic Stage-In
+
+In HPC we stage-in files from global parallel file system onto node local storage. This serves a number of purposes:
+
+- avoiding random I/O which throttles performance tremendously and may occur, for instance, when several SLURM jobs with the same program attempt to use the same reference files (e.g. with alignment programs)
+- minimizing latency for repeated file access, since local storage (e.g., SSD or NVMe) typically offers much lower access times than network-attached parallel filesystems.
+- improving fault tolerance and resilience during job execution — if the parallel filesystem experiences transient issues, locally staged files remain accessible.
+
+to name a few.
+
+To enable automatic stage-in of files a workflow rule needs to designate either the [`random` access file pattern or the `multi` access pattern](https://snakemake.readthedocs.io/en/stable/snakefiles/storage.html#access-pattern-annotation). Furthermore a flag `--slurm-node-local-prefix` can be used to indicate a node local directory available for usage within a SLURM job. As always with Snakemake, we may use the flag on the command line and in a configuration file, e.g.:
+
+```yaml
+slurm-node-local-prefix: /localscratch/$SLURM_JOB_ID
+```
+
+Environment variables may just be written as literals - even if the environment variable in only set within a cluster job.
+
+If the automatic stage-in process is not desirable, `--slurm-suppress-auto-stagein` will turn it off.
+
+Stage-in processes will be performed on all nodes of a job. Using `sbcast` for files <= 4GB and `scp` for files > 4GB. Note that stage-in processes with `scp` cross-node `ssh` needs to be working on your cluster.
+
+.. note:: Currently, there is no distinction between different node local files systems, e.g. `/tmpfs`, `nvme`s on SSD or dedicated ramdisks. Only one node local directory prefix is allowed. 
+
 
 #### Standard Resources
 

pyproject.toml

@@ -39,6 +39,7 @@ omit = [".*", "*/site-packages/*", "Snakefile"]
 
 [tool.black]
 line-length = 88
+target-version = ["py311"]
 extend-exclude = 'tests/test_scontrol_parsing\.py'
 
 [tool.pytest.ini_options]

snakemake_executor_plugin_slurm/__init__.py

@@ -34,6 +34,10 @@
 from snakemake_interface_executor_plugins.jobs import (
     JobExecutorInterface,
 )
+from snakemake.io.flags.access_patterns import (
+    AccessPattern,
+    STORE_KEY,
+)
 from snakemake_interface_common.exceptions import WorkflowError
 
 from .accounts import (
@@ -44,10 +48,12 @@
 from .utils import (
     get_max_array_size,
     get_job_wildcards,
+    encode_deferred_envvars,
     pending_jobs_for_rule,
     delete_slurm_environment,
     delete_empty_dirs,
     set_gres_string,
+    get_file_size,
 )
 from .job_status_query import (
     get_min_job_age,
@@ -229,6 +235,32 @@ class ExecutorSettings(ExecutorSettingsBase):
         },
     )
 
+    node_local_prefix: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "A cluster specific node local local global directory. "
+            "Similar to '.remote_job_local_storage_prefix' of the fs-storage "
+            "plugin. Used to ensure stage-in when input is flagged for random "
+            "or mixed access patterns - see https://jgu.to/kwenq ."
+            "Other than the fs-storage plugin, this function only works to "
+            "stage-in input files (no stage-out) but supports transferring "
+            "onto all nodes of the SLURM job. For big files (> 2GB)'ssh' from "
+            "node to node must be enabled."
+        },
+    )
+    suppress_auto_stagein: bool = field(
+        default=False,
+        metadata={
+            "help": "By default, when node_local_prefix is set, Snakemake will "
+            "automatically stage in input files with random or mixed access "
+            "patterns to the node local directory. Setting this flag will "
+            "suppress this behavior, so that no automatic staging will be "
+            "performed.",
+            "env_var": False,
+            "required": False,
+        },
+    )
+
     init_seconds_before_status_checks: Optional[int] = field(
         default=40,
         metadata={
@@ -636,9 +668,17 @@ def additional_general_args(self):
         passed to `exec_job`.
         """
         general_args = "--executor slurm-jobstep --jobs 1"
-        # need to pass
+        # switch for passing a command as script vs using `sbatch --wrap`
         if self.workflow.executor_settings.pass_command_as_script:
             general_args += " --slurm-jobstep-pass-command-as-script"
+        # attempt auto stage-in if not suppressed and node_local_prefix is set
+        if (
+            self.workflow.executor_settings.node_local_prefix
+            and not self.workflow.executor_settings.suppress_auto_stagein
+        ):
+            ld = self.workflow.executor_settings.node_local_prefix
+            ld = encode_deferred_envvars(ld)
+            general_args += f" --slurm-jobstep-node-local-prefix={shlex.quote(ld)}"
         return general_args
 
     def run_jobs(self, jobs: List[JobExecutorInterface]):
@@ -648,6 +688,33 @@ def run_jobs(self, jobs: List[JobExecutorInterface]):
         - `run_array_jobs` for array job submission, or
         - `run_pool_jobs` for pool job submission (to be implemented in the future).
         """
+        # check whether current jobs are using random or multi access patters
+        # and node_local_prefix is set
+        if (
+            not self.workflow.executor_settings.suppress_auto_stagein
+            and self.workflow.executor_settings.node_local_prefix
+        ):
+            for job in jobs:
+                for inp in job.input:
+                    has_random_or_mixed = any(
+                        isinstance(inp.flags.get(STORE_KEY), AccessPattern)
+                        and inp.flags[STORE_KEY]
+                        in {AccessPattern.RANDOM, AccessPattern.MIXED}
+                        for inp in job.input
+                    )
+                if has_random_or_mixed:
+                    size = get_file_size(inp.path)
+                    if size is not None and size > 100:
+                        self.logger.warning(
+                            f"Job '{job.name}' has input file '{inp.path}' with "
+                            f"random or mixed access pattern and size {size} "
+                            "bytes. Snakemake will attempt to stage in this file "
+                            "to the node local directory specified by "
+                            f"{self.workflow.executor_settings.node_local_prefix}. "
+                            "However, for files of this size the process might be "
+                            "slow."
+                        )
+
         if self._main_event_loop is None:
             try:
                 self._main_event_loop = asyncio.get_running_loop()

snakemake_executor_plugin_slurm/utils.py

@@ -71,6 +71,33 @@ def get_job_wildcards(job: JobExecutorInterface) -> str:
     return wildcard_str
 
 
+_DEFERRED_ENVVAR_PATTERN = re.compile(
+    r"""
+    (?<!\\)\$
+    (?:
+        \{(?P<braced>[A-Za-z_][A-Za-z0-9_]*)\}
+        |(?P<bare>[A-Za-z_][A-Za-z0-9_]*)
+    )
+    """,
+    re.VERBOSE,
+)
+
+
+def encode_deferred_envvars(value: str) -> str:
+    """Replace shell-style env variables with neutral markers.
+
+    The marker keeps submit-time shell quoting safe. The matching jobstep-side
+    code can expand ``__ENV_NAME__`` back from ``os.environ`` once the job is
+    running in its real job context.
+    """
+
+    def replace(match: re.Match[str]) -> str:
+        name = match.group("braced") or match.group("bare")
+        return f"__ENV_{name}__"
+
+    return _DEFERRED_ENVVAR_PATTERN.sub(replace, value)
+
+
 def pending_jobs_for_rule(dag: DAGExecutorInterface, rule_name: str) -> int:
     """Count jobs of a rule that are currently eligible for scheduling.
 
@@ -406,3 +433,29 @@ def set_gres_string(job: JobExecutorInterface) -> str:
         gres += f" --gpus={gpu_string}"
 
     return gres
+
+
+def get_file_size(job: JobExecutorInterface) -> int:
+    """
+    Function to get the file size of the input files for a job. This is used
+    to determine the resources needed for the job.
+
+    Args:
+        job: The JobExecutorInterface instance representing the job
+    Returns:
+        The total file size of the input files for the job, in GB.
+    """
+    total_size_bytes = 0
+    for input_file in job.input_files:
+        if os.path.isfile(input_file):
+            total_size_bytes += os.path.getsize(input_file)
+        elif os.path.isdir(input_file):
+            for dirpath, _, filenames in os.walk(input_file):
+                for filename in filenames:
+                    file_path = os.path.join(dirpath, filename)
+                    if os.path.isfile(file_path):
+                        total_size_bytes += os.path.getsize(file_path)
+
+    # Convert bytes to GB
+    total_size_gb = total_size_bytes / (1024**3)
+    return round_half_up(total_size_gb)  # implicit conversion to int

tests/test_node_local_prefix.py

@@ -0,0 +1,20 @@
+from snakemake_executor_plugin_slurm.utils import encode_deferred_envvars
+
+
+def test_encode_deferred_envvars_rewrites_simple_env_vars():
+    assert (
+        encode_deferred_envvars("/localscratch/$SLURM_JOB_ID")
+        == "/localscratch/__ENV_SLURM_JOB_ID__"
+    )
+    assert (
+        encode_deferred_envvars("/localscratch/${SLURM_JOB_ID}/run")
+        == "/localscratch/__ENV_SLURM_JOB_ID__/run"
+    )
+
+
+def test_encode_deferred_envvars_leaves_literal_text_alone():
+    assert encode_deferred_envvars("/localscratch/job-1") == "/localscratch/job-1"
+    assert (
+        encode_deferred_envvars(r"/localscratch/\$SLURM_JOB_ID")
+        == r"/localscratch/\$SLURM_JOB_ID"
+    )

tests/test_stagein.py

@@ -0,0 +1,123 @@
+from pathlib import Path
+from typing import Optional
+
+import snakemake.common.tests
+import snakemake.settings.types as settings
+from snakemake import api
+from snakemake.settings.types import ConfigSettings
+from snakemake_interface_executor_plugins.settings import ExecutorSettingsBase
+
+from snakemake_executor_plugin_slurm import ExecutorSettings
+
+
+class LocalStageinTestcasesBase(snakemake.common.tests.TestWorkflowsLocalStorageBase):
+    """Resolve stage-in testcases from this plugin's tests/testcases directory."""
+
+    def get_config_settings(self) -> Optional[ConfigSettings]:
+        return None
+
+    def run_workflow(self, test_name, tmp_path, deployment_method=frozenset()):
+        test_path = Path(__file__).parent / "testcases" / test_name
+        if not test_path.exists():
+            return super().run_workflow(test_name, tmp_path, deployment_method)
+
+        if self.omit_tmp:
+            tmp_path = test_path
+        else:
+            tmp_path = Path(tmp_path) / test_name
+            self._copy_test_files(test_path, tmp_path)
+
+        resource_settings = self.get_resource_settings()
+
+        if self._common_settings().local_exec:
+            resource_settings.cores = 3
+            resource_settings.nodes = None
+        else:
+            resource_settings.cores = 1
+            resource_settings.nodes = 3
+
+        with api.SnakemakeApi(
+            settings.OutputSettings(
+                verbose=True,
+                show_failed_logs=True,
+            ),
+        ) as snakemake_api:
+            workflow_api = snakemake_api.workflow(
+                config_settings=self.get_config_settings(),
+                resource_settings=resource_settings,
+                storage_settings=settings.StorageSettings(
+                    default_storage_provider=self.get_default_storage_provider(),
+                    default_storage_prefix=self.get_default_storage_prefix(),
+                    shared_fs_usage=(
+                        settings.SharedFSUsage.all()
+                        if self.get_assume_shared_fs()
+                        else frozenset()
+                    ),
+                ),
+                deployment_settings=self.get_deployment_settings(deployment_method),
+                storage_provider_settings=self.get_default_storage_provider_settings(),
+                workdir=Path(tmp_path),
+                snakefile=tmp_path / "Snakefile",
+            )
+
+            dag_api = workflow_api.dag()
+            dag_api.execute_workflow(
+                executor=self.get_executor(),
+                executor_settings=self.get_executor_settings(),
+                execution_settings=settings.ExecutionSettings(
+                    latency_wait=self.latency_wait,
+                ),
+                remote_execution_settings=self.get_remote_execution_settings(),
+            )
+
+
+class TestStageInSbcast(LocalStageinTestcasesBase):
+    """Integration test for sbcast stage-in with a small synthetic input."""
+
+    __test__ = True
+
+    def get_executor(self) -> str:
+        return "slurm"
+
+    def get_executor_settings(self) -> Optional[ExecutorSettingsBase]:
+        return ExecutorSettings(
+            init_seconds_before_status_checks=2,
+            node_local_prefix="/tmp/snakemake-stagein",
+        )
+
+    def test_stagein_sbcast(self, tmp_path):
+        self.run_workflow("stagein_sbcast", tmp_path)
+
+        run_dir = Path(tmp_path) / "stagein_sbcast"
+        count_file = run_dir / "counted" / "count.txt"
+        sbcast_log = run_dir / "log" / "sbcast.log"
+
+        assert count_file.exists()
+        assert sbcast_log.exists()
+        assert count_file.read_text().strip().startswith("3 ")
+
+
+class TestStageInSSH(LocalStageinTestcasesBase):
+    """Optional integration test for SSH stage-in with a sparse large input."""
+
+    __test__ = True
+
+    def get_executor(self) -> str:
+        return "slurm"
+
+    def get_executor_settings(self) -> Optional[ExecutorSettingsBase]:
+        return ExecutorSettings(
+            init_seconds_before_status_checks=2,
+            node_local_prefix="/tmp/snakemake-stagein",
+        )
+
+    def test_stagein_ssh(self, tmp_path):
+        self.run_workflow("stagein_ssh", tmp_path)
+
+        run_dir = Path(tmp_path) / "stagein_ssh"
+        size_file = run_dir / "measured" / "size.txt"
+        ssh_log = run_dir / "log" / "ssh.log"
+
+        assert size_file.exists()
+        assert ssh_log.exists()
+        assert "data/big" in size_file.read_text()

tests/testcases/stagein_sbcast/Snakefile

@@ -0,0 +1,25 @@
+inputflags:
+    access.sequential
+
+
+rule all:
+    input:
+        "counted/count.txt"
+
+
+rule make_words:
+    output:
+        "data/words"
+    shell:
+        "printf 'alpha\nbeta\ngamma\n' > {output}"
+
+
+rule use_sbcast:
+    input:
+        access.random("data/words")
+    output:
+        "counted/count.txt"
+    log:
+        "log/sbcast.log"
+    shell:
+        "wc -l {input} > {output}"