diff --git a/README.md b/README.md index a7c5e8e..cb8ea7a 100644 --- a/README.md +++ b/README.md @@ -214,7 +214,8 @@ Then open `http://localhost:8501` in your browser. Override default values to match your cluster: ```yaml -slurm_partition: "gpu" # which partition/queue to submit to +slurm_partition: "gpu" # partition(s) to submit inference to; one name, + # "gpu-el8,gpu-training" or a YAML list for several slurm_qos: "normal" # optional QoS if your site uses it structure_inference_gpus_per_task: 1 # number of GPUs each inference job needs structure_inference_gpu_model: "" # "" lets SLURM pick any GPU in the partition; set a model to pin @@ -260,8 +261,15 @@ you hit these. When set this drives `--exclude` per job and **overrides** `structure_inference_gpu_model` (the two would conflict). It's the practical "fit to GPU" lever: requested host RAM is a separate pool and does not size GPU VRAM, but excluding too-small GPUs by length does. Use explicit comma node lists - (bracket ranges may be glob-expanded by the shell). Multi-partition routing (e.g. EMBL's bigger - `gpu-training` cards) is out of scope — keep one partition and let unified memory spill the tail. + (bracket ranges may be glob-expanded by the shell). VRAM-tier routing works *within* the listed + partition(s); it excludes nodes by name, so if you span partitions make sure the tier node lists + cover every partition you submit to. +- **Span several partitions** by giving `slurm_partition` more than one name — a comma-separated + string (`"gpu-el8,gpu-training"`) or a YAML list. The plugin passes them straight to `sbatch -p`, + and SLURM starts each inference job on whichever partition frees up first, so jobs aren't stuck + behind one busy queue (e.g. spill onto EMBL's bigger `gpu-training` cards). Every listed partition + must accept the job's GPUs, `--mem` and walltime (`structure_inference_max_runtime` ≤ each + partition's `MaxTime`); a partition the job doesn't fit is simply skipped by SLURM. - **Exclude specific nodes** with `slurm_exclude_nodes` → passed verbatim to `sbatch --exclude` (e.g. `"gpu50,gpu51"`). Use it as a fallback for nodes whose GPU the container can't use — e.g. a CUDA compute capability newer than the container's bundled `ptxas` (fails `ptxas too old` / diff --git a/config/config.yaml b/config/config.yaml index d3c1ed8..80886a9 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -140,6 +140,19 @@ batch_max_tokens: 0 alphafold_inference_threads: 8 # SLURM resources +# Partition(s) for the GPU structure_inference jobs. Give ONE partition, or SEVERAL +# to let SLURM start each job on whichever one frees up first (it schedules onto the +# soonest-available partition). Accepts a single name, a comma/space-separated string, +# or a YAML list -- all equivalent: +# slurm_partition: "gpu-el8" # single partition +# slurm_partition: "gpu-el8,gpu-training" # comma-separated +# slurm_partition: # YAML list +# - gpu-el8 +# - gpu-training +# All listed partitions must accept the job's resources (GPUs, --mem, and a walltime +# within each partition's MaxTime -- see structure_inference_max_runtime); a partition +# the job cannot fit is simply skipped by SLURM. Only structure_inference uses this; +# the other (CPU) rules run on the cluster's default partition. slurm_partition: "gpu-el8" slurm_qos: "normal" structure_inference_gpus_per_task: 1 diff --git a/test/test_normalize_partitions.py b/test/test_normalize_partitions.py new file mode 100644 index 0000000..caa237b --- /dev/null +++ b/test/test_normalize_partitions.py @@ -0,0 +1,77 @@ +"""Unit tests for ``normalize_partitions`` (multi-partition support). + +``common.smk`` is plain Python (stdlib-only imports), so it is loaded by path and +its functions tested directly. Run with ``python test/test_normalize_partitions.py`` +or ``pytest test/test_normalize_partitions.py``. +""" + +import importlib.machinery +import importlib.util +import shlex +from pathlib import Path + +_COMMON = Path(__file__).resolve().parents[1] / "workflow" / "rules" / "common.smk" +_spec = importlib.util.spec_from_loader( + "aps_common", importlib.machinery.SourceFileLoader("aps_common", str(_COMMON)) +) +common = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(common) +normalize_partitions = common.normalize_partitions + + +def test_single_partition_unchanged(): + assert normalize_partitions("gpu-el8") == "gpu-el8" + + +def test_none_and_empty_return_none(): + assert normalize_partitions(None) is None + assert normalize_partitions("") is None + assert normalize_partitions(" ") is None + assert normalize_partitions([]) is None + assert normalize_partitions([" ", ""]) is None + + +def test_comma_separated_string_is_normalized(): + assert normalize_partitions("gpu-el8,transform") == "gpu-el8,transform" + # surrounding whitespace around commas is stripped + assert normalize_partitions("gpu-el8, transform ,training") == ( + "gpu-el8,transform,training" + ) + + +def test_whitespace_separated_string(): + assert normalize_partitions("gpu-el8 transform") == "gpu-el8,transform" + + +def test_yaml_list_is_joined_with_commas(): + assert normalize_partitions(["gpu-el8", "transform"]) == "gpu-el8,transform" + assert normalize_partitions(("gpu-el8", "gpu-training")) == "gpu-el8,gpu-training" + + +def test_duplicates_removed_order_preserved(): + assert normalize_partitions(["gpu-el8", "transform", "gpu-el8"]) == ( + "gpu-el8,transform" + ) + assert normalize_partitions("transform,gpu-el8,transform") == ( + "transform,gpu-el8" + ) + + +def test_list_entries_are_stripped(): + assert normalize_partitions([" gpu-el8 ", " transform"]) == "gpu-el8,transform" + + +def test_result_survives_shlex_quote_unquoted(): + # The SLURM plugin runs the partition through shlex.quote before passing it to + # `sbatch -p`. A comma-joined list must survive that untouched (comma is a + # shell-safe character) so sbatch receives the full partition list. + result = normalize_partitions("gpu-el8,transform") + assert shlex.quote(result) == result == "gpu-el8,transform" + + +if __name__ == "__main__": + for name, fn in sorted(globals().items()): + if name.startswith("test_") and callable(fn): + fn() + print(f"ok {name}") + print("all normalize_partitions tests passed") diff --git a/workflow/Snakefile b/workflow/Snakefile index 0955030..8b58afb 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -63,7 +63,11 @@ input_files = config.get("input_files", []) if isinstance(input_files, (str, Path)): input_files = [input_files] -DEFAULT_SLURM_PARTITION = config.get("slurm_partition") +# slurm_partition may be a single name, a comma/space-separated string or a YAML +# list; normalize to one comma-separated string (e.g. "gpu-el8,transform") that +# sbatch -p understands natively (SLURM runs the job on whichever partition frees +# up first). See config.yaml for usage. +DEFAULT_SLURM_PARTITION = normalize_partitions(config.get("slurm_partition")) DEFAULT_SLURM_GRES = config.get("slurm_gres") DEFAULT_SLURM_QOS = config.get("slurm_qos") DEFAULT_STRUCTURE_INFERENCE_GPUS = config.get("structure_inference_gpus_per_task", 1) diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 410076c..80bf4aa 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -506,6 +506,41 @@ def prepare_container_binds( os.environ.setdefault(var, "1") +def normalize_partitions(value: Any) -> str | None: + """Normalise a ``slurm_partition`` config value to a comma-separated string. + + SLURM's ``sbatch -p`` natively accepts several partitions as a comma list + (``-p gpu-el8,transform``) and schedules the job onto whichever one lets it + start soonest. This lets a user list every GPU partition they may run on so + inference jobs are not stuck behind one busy queue. + + Accepts any of: + + * a YAML list/tuple: ``[gpu-el8, transform]`` + * a comma- and/or whitespace-separated string: ``"gpu-el8, transform"`` + * a single partition string: ``"gpu-el8"`` (unchanged) + * ``None`` / empty -> ``None`` (caller supplies its own fallback) + + Returns a de-duplicated, order-preserving comma-joined string (no spaces, so + it survives ``shlex.quote`` unquoted and reaches ``sbatch`` verbatim), or + ``None`` when no partition is given. + """ + if value is None: + return None + if isinstance(value, (list, tuple, set)): + items = list(value) + else: + # A single scalar; split on commas and any surrounding whitespace so both + # "a,b", "a, b" and "a b" are accepted. + items = str(value).replace(",", " ").split() + names: list[str] = [] + for item in items: + name = str(item).strip() + if name and name not in names: + names.append(name) + return ",".join(names) if names else None + + def linear_resources( *, mem: int = 800,