From 674167481d82591d97b7d73b2c4053c0da0a13dc Mon Sep 17 00:00:00 2001 From: Cemberk Date: Mon, 17 Aug 2026 20:04:38 +0000 Subject: [PATCH 1/7] fix(slurm): honor the model-card contract on the slurm_multi path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit slurm_multi is an escape hatch for topologies the templated launchers cannot express, but it had also become an escape hatch from the model-card contract: fields the templated path honors were silently dropped, so a card that fully described its job still ran wrong. Node count. `distributed.nnodes` never reached `#SBATCH --nodes`, which came only from `slurm.nodes` (default 1). A card declaring a 4-node topology was submitted as a 1-node job. This was invisible in the common workflow, where `salloc -N 4` comes first and the wrapper — run with bash, not sbatch — inherits SLURM_NNODES; it only bites on the sbatch path. Cards worked around it with `"args": "-N 4 -n 4"`, but args go to `bash .slurm`, not to sbatch, so a script that never reads $@ discarded them. Both paths now reconcile the two fields: nnodes sizes the allocation when slurm.nodes was not set explicitly, an explicit slurm.nodes wins a conflict and warns, and resolution recomputes from a saved baseline so the prepare() that deploy() re-runs after preflight is idempotent. Launcher resolution. prepare() picked the path from the model card while _prepare_template_context() read the deployment config, so a card-declared launcher could take one path and emit another path's env block. Both now call one resolver, deployment-config-first — matching how BuildOrchestrator builds the manifest. multiple_results. The slurm_multi collector never read model_info, so the declared filename was dead config and collection worked only when a workload happened to write to a hardcoded path. It is now searched next to the model script (where the wrapper cd's, so `$(pwd)` writes land) before the conventional locations. The CSV is still read directly rather than routed through handle_multiple_results(): a self-managed script has no common_info to merge against and writes the full schema itself, and re-ingesting it would recompute status from performance, flipping a legitimate zero-score FAILURE to SUCCESS. Placeholder images. Cards ship DOCKER_IMAGE_NAME as "" to mean "fill this in". The implicit --use-image path accepted any single distinct value, so the marker became the image name and every node failed on `docker pull `. Angle-bracketed values are now rejected at submit time with an actionable error. Also guards the per-job perf aggregation against appending cwd/perf.csv to itself, which duplicated every row whenever the source resolved to it. Co-Authored-By: Claude Opus 5 (1M context) --- docs/launchers.md | 38 ++++ src/madengine/deployment/common.py | 69 +++++++ src/madengine/deployment/slurm.py | 174 ++++++++++++++--- .../orchestration/build_orchestrator.py | 42 +++- tests/unit/test_orchestration.py | 38 ++++ tests/unit/test_slurm_multi.py | 180 ++++++++++++++++++ 6 files changed, 514 insertions(+), 27 deletions(-) diff --git a/docs/launchers.md b/docs/launchers.md index 227557aa..071bb4ee 100644 --- a/docs/launchers.md +++ b/docs/launchers.md @@ -647,6 +647,44 @@ madengine run --manifest-file build_manifest.json **Alias**: `"slurm-multi"` (hyphen) is normalized to `"slurm_multi"` (underscore). +**Model-card contract (shared with the templated path)**: + +slurm_multi hands the *topology* to the workload, not the *contract*. These model-card +fields are honoured identically on both paths: + +| Field | Effect | +|-------|--------| +| `distributed.launcher` | Selects the path. Resolved deployment-config-first, model-card-second — the same resolution used to emit the launcher env block, so the two can never disagree. | +| `distributed.nnodes` | Sizes the allocation (`#SBATCH --nodes` / `--ntasks`) when `slurm.nodes` is not set explicitly. If both are set and differ, `slurm.nodes` wins and a warning is printed. | +| `slurm.*` | `partition`, `time`, `gpus_per_node`, `exclusive`, `reservation`, `nodelist`, … | +| `multiple_results` | Names the results CSV. Searched next to the model script (where the wrapper `cd`s, so `$(pwd)` writes land there), then the job dir, output dir and cwd, before falling back to the conventional locations. | + +> **`args` is not sbatch flags.** On the templated path `args` goes to the +> containerised `run.sh`; on this path it is appended to `bash .slurm`. It is +> **not** forwarded to `sbatch`, so `"args": "-N 4 -n 4"` does *not* request 4 nodes — +> it passes two ignored positional arguments to a script that never reads `$@`, and +> the job runs on the `slurm.nodes` default of 1. Size the allocation with +> `distributed.nnodes` (or `slurm.nodes`) instead. + +> **`multiple_results` is read directly here.** On the templated path the CSV is +> narrow and gets merged with `common_info` by `update_perf_csv`. A self-managed +> script has no `common_info` to merge against, so it writes the full perf schema +> itself and madengine reads it as-is — routing it through `handle_multiple_results()` +> would recompute `status` from `performance` and flip a legitimate zero-score +> FAILURE row to SUCCESS. + +> **Preset selection still keys on `slurm.nodes`.** `ConfigLoader` picks the +> single-node vs multi-node preset before `nnodes` reconciliation, so a card sized +> only by `nnodes` keeps the single-node preset. That is usually what RDMA inference +> workloads want — the multi-node preset sets `NCCL_IB_DISABLE=1` and +> `NCCL_SOCKET_IFNAME=eth0` — but set `slurm.time` explicitly if the 12 h single-node +> default is too short. + +**Placeholder images are rejected**: model cards commonly ship +`"DOCKER_IMAGE_NAME": ""` as a fill-me-in marker. Any +angle-bracketed value is rejected at submit time with an actionable error rather than +becoming the image every compute node fails to pull. + **Features**: - Wrapper SBATCH script with shell-quoted env_vars (injection-safe) - Parallel `srun docker pull` on all nodes for registry images diff --git a/src/madengine/deployment/common.py b/src/madengine/deployment/common.py index 13657246..c2fd605d 100644 --- a/src/madengine/deployment/common.py +++ b/src/madengine/deployment/common.py @@ -125,6 +125,75 @@ def is_self_managed_launcher(launcher_type: Optional[str]) -> bool: return normalize_launcher(launcher_type, "slurm") in _SELF_MANAGED_LAUNCHERS +def resolve_launcher_from_sources( + deployment_launcher: Optional[str], + model_launcher: Optional[str], + default: str = "torchrun", +) -> str: + """Resolve the effective launcher from the deployment config and the model card. + + Precedence is deployment config first, model card second. That matches how + BuildOrchestrator builds the manifest: it copies the model card's launcher into + ``deployment_config.distributed`` only when the key is absent, so a value present + there is already the user's explicit choice. + + Every dispatch site must use this helper. Resolving the launcher differently in + two places is what allowed ``prepare()`` to pick the self-managed path from the + model card while ``_prepare_template_context()`` simultaneously emitted a + ``torchrun`` env block from the deployment config. + """ + return deployment_launcher or model_launcher or default + + +def resolve_node_count( + configured_nodes: int, + nnodes: Optional[Any], + nodes_explicitly_set: bool, +) -> tuple: + """Reconcile ``slurm.nodes`` with ``distributed.nnodes`` into one node count. + + ``slurm.nodes`` sizes the allocation (``#SBATCH --nodes``) while + ``distributed.nnodes`` sizes the topology the launcher builds on top of it. + Nothing used to reconcile them, so a model card declaring only ``nnodes`` was + submitted against the ``slurm.nodes`` default of 1 — a multi-node topology + squeezed onto a single node. + + Args: + configured_nodes: ``slurm.nodes`` after ConfigLoader applied its defaults. + nnodes: ``distributed.nnodes``, if declared. + nodes_explicitly_set: True when ``slurm.nodes`` came from the user or the + model card rather than from a preset default. + + Returns: + ``(nodes, note)`` where ``note`` is a human-readable string to log, or None. + """ + try: + requested = int(nnodes) if nnodes is not None else None + except (TypeError, ValueError): + return configured_nodes, ( + f"Ignoring non-numeric distributed.nnodes={nnodes!r}; " + f"using slurm.nodes={configured_nodes}." + ) + + if requested is None or requested == configured_nodes: + return configured_nodes, None + + if nodes_explicitly_set: + # Both were set and they disagree. The explicit allocation size wins — + # overriding it could request nodes the user did not ask to be billed for — + # but the mismatch is almost always a config error, so say so. + return configured_nodes, ( + f"slurm.nodes={configured_nodes} conflicts with " + f"distributed.nnodes={requested}; using slurm.nodes={configured_nodes}. " + "Set them to the same value to silence this." + ) + + return requested, ( + f"Sizing allocation from distributed.nnodes={requested} " + f"(slurm.nodes was not set explicitly)." + ) + + @functools.lru_cache(maxsize=None) def is_rocprofv3_available() -> bool: """ diff --git a/src/madengine/deployment/slurm.py b/src/madengine/deployment/slurm.py index af99c08d..7e8a0531 100644 --- a/src/madengine/deployment/slurm.py +++ b/src/madengine/deployment/slurm.py @@ -25,6 +25,8 @@ configure_multi_node_profiling, is_self_managed_launcher, normalize_launcher, + resolve_launcher_from_sources, + resolve_node_count, ) from .config_loader import ConfigLoader, apply_deployment_config from .slurm_node_selector import SlurmNodeSelector @@ -63,6 +65,13 @@ def __init__(self, config: DeploymentConfig): Args: config: Deployment configuration """ + # Capture which slurm keys were actually supplied (by --additional-context or + # by the model card, which BuildOrchestrator merges into the manifest's + # deployment_config) BEFORE ConfigLoader layers its presets on top. Once the + # defaults are applied every key looks "set", and nodes=1 from a preset is + # indistinguishable from nodes=1 the user asked for. + self._explicit_slurm_keys = set((config.additional_context or {}).get("slurm") or {}) + apply_deployment_config(config, ConfigLoader.load_slurm_config) super().__init__(config) @@ -73,6 +82,12 @@ def __init__(self, config: DeploymentConfig): # SLURM parameters self.partition = self.slurm_config.get("partition", "gpu") self.nodes = self.slurm_config.get("nodes", 1) + + # Baseline allocation size, before any distributed.nnodes reconciliation. + # _resolve_nodes() always recomputes from this so that prepare(), which + # deploy() re-runs after node preflight, stays idempotent. + self._configured_nodes = self.nodes + self._resolve_nodes() self.gpus_per_node = self.slurm_config.get("gpus_per_node", 8) self.time_limit = self.slurm_config.get("time", "24:00:00") self.output_dir = Path(self.slurm_config.get("output_dir", "./slurm_results")) @@ -302,11 +317,10 @@ def prepare(self) -> bool: model_keys_peek = list((self.manifest or {}).get("built_models", {}).keys()) if model_keys_peek: model_info_peek = self.manifest["built_models"][model_keys_peek[0]] - model_distributed_peek = model_info_peek.get("distributed", {}) - launcher_type_peek = ( - model_distributed_peek.get("launcher") - or self.distributed_config.get("launcher", "torchrun") - ) + # Re-resolve now that the model card is in hand: a card may size its + # topology with distributed.nnodes alone. + self._resolve_nodes(model_info_peek) + launcher_type_peek = self._resolve_launcher(model_info_peek) if is_self_managed_launcher(launcher_type_peek): self.output_dir.mkdir(parents=True, exist_ok=True) self.console.print( @@ -359,6 +373,46 @@ def prepare(self) -> bool: self.console.print(f"[red]✗ Failed to generate script: {e}[/red]") return False + def _resolve_nodes(self, model_info: Optional[Dict] = None) -> int: + """Size the allocation from slurm.nodes reconciled with distributed.nnodes. + + ``nnodes`` is read from the deployment config first (BuildOrchestrator copies + the model card's value there at build time) and from the model card second, so + a card is honoured even when the manifest was not produced by that merge — for + example a hand-written manifest, or `madengine run` against a card whose + topology changed after the build. + + Recomputes from ``self._configured_nodes`` rather than from ``self.nodes``, so + repeated calls converge instead of ratcheting. + """ + nnodes = self.distributed_config.get("nnodes") + if nnodes is None and model_info: + nnodes = (model_info.get("distributed") or {}).get("nnodes") + + resolved, note = resolve_node_count( + configured_nodes=self._configured_nodes, + nnodes=nnodes, + nodes_explicitly_set="nodes" in self._explicit_slurm_keys, + ) + if note and note != getattr(self, "_nodes_note_shown", None): + self.console.print(f"[yellow]⚠ {note}[/yellow]") + self._nodes_note_shown = note + + self.nodes = resolved + self.slurm_config["nodes"] = resolved + return resolved + + def _resolve_launcher(self, model_info: Dict) -> str: + """Resolve the effective launcher for a model, deployment config first. + + Single source of truth for both dispatch sites: the self-managed peek in + prepare() and the launcher-command generation in _prepare_template_context(). + """ + return resolve_launcher_from_sources( + deployment_launcher=self.distributed_config.get("launcher"), + model_launcher=(model_info.get("distributed") or {}).get("launcher"), + ) + @staticmethod def _normalize_nodelist(nodelist: Optional[str]) -> Optional[str]: """Normalize nodelist to comma-separated without spaces for #SBATCH --nodelist.""" @@ -607,9 +661,11 @@ def _prepare_template_context(self, model_info: Dict) -> Dict[str, Any]: additional_context["slurm"] = self.slurm_config resolved_gpus_per_node = resolve_runtime_gpus(model_info, additional_context) - # Extract launcher configuration - launcher_type = self.distributed_config.get("launcher", "torchrun") # Default to torchrun - + # Extract launcher configuration. Resolved the same way as the self-managed + # peek in prepare(), so the path taken and the env block emitted can never + # disagree about which launcher this model uses. + launcher_type = self._resolve_launcher(model_info) + # Canonicalize aliases before validity check so e.g. sglang_disagg → sglang-disagg # passes through normalize_launcher instead of being mapped to "docker". launcher_type = canonicalize_distributed_launcher(launcher_type) or launcher_type @@ -1741,10 +1797,9 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: # so collect via _collect_slurm_multi_results instead of the template-based path. if model_key: _mi = built_models_dict.get(model_key, {}) or {} - _launcher_type = (_mi.get("distributed") or {}).get("launcher", "") - if is_self_managed_launcher(_launcher_type): + if is_self_managed_launcher(self._resolve_launcher(_mi)): return self._collect_slurm_multi_results( - deployment_id, results, session_start_row + deployment_id, results, session_start_row, model_info=_mi ) @@ -2068,12 +2123,66 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: ) return results + def _slurm_multi_declared_result_csv( + self, model_info: Optional[Dict[str, Any]], deployment_id: str + ) -> Optional[Path]: + """Resolve a slurm_multi model's declared ``multiple_results`` CSV. + + On the templated path ``multiple_results`` names a narrow per-run CSV that is + merged with common_info by update_perf_csv. A self-managed script has no + common_info to merge against, so it writes the full perf schema itself and the + file is read directly — but the model card should still be able to *name* it + rather than every workload having to land on one of the hardcoded paths below. + + Note the difference is deliberate: routing an already-full-schema CSV through + handle_multiple_results() would recompute ``status`` from ``performance`` and + turn a legitimate zero-score FAILURE row into a SUCCESS. + """ + declared = (model_info or {}).get("multiple_results") + if not declared: + return None + + search_dirs: List[Path] = [] + # Where the launcher runs: the wrapper cd's to the model script's directory, + # so a script writing to $(pwd) lands here. + scripts_rel = (model_info or {}).get("scripts", "") + if scripts_rel and self.config.manifest_file: + script_path = Path(self.config.manifest_file).parent.absolute() / scripts_rel + search_dirs.append(script_path.parent) + search_dirs.extend( + [ + self.output_dir / deployment_id, + self.output_dir, + Path.cwd(), + ] + ) + + for directory in search_dirs: + candidate = directory / declared + try: + if candidate.is_file() and candidate.stat().st_size > 0: + self.console.print( + f"[dim] Using declared multiple_results CSV: {candidate}[/dim]" + ) + return candidate + except OSError: + continue + self.console.print( + f"[dim] Declared multiple_results '{declared}' not found in " + f"{', '.join(str(d) for d in search_dirs)}; falling back to conventional paths.[/dim]" + ) + return None + def _collect_slurm_multi_results( - self, deployment_id: str, results: Dict[str, Any], session_start_row: Optional[int] + self, + deployment_id: str, + results: Dict[str, Any], + session_start_row: Optional[int], + model_info: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """ Collect results for slurm_multi launchers. - + slurm_multi model scripts generate their own perf.csv via their benchmark scripts (e.g. generate_perf_csv.py). We collect SLURM logs for diagnostics and read the model-generated perf.csv for metrics. @@ -2082,15 +2191,19 @@ def _collect_slurm_multi_results( flat_out_files = sorted(self.output_dir.glob(f"madengine-*_{deployment_id}_*.out")) results["logs"] = [str(f) for f in flat_out_files] + # A model card that names its own results CSV wins over the conventional + # locations below, so a workload is not forced to write to a path madengine + # happens to know about. + perf_csv_path = self._slurm_multi_declared_result_csv(model_info, deployment_id) + # Look for model-generated perf.csv. Inner scripts in MAD-private write # to one of these locations depending on the workload: # * SGLang / vLLM disagg: /shared_inference///perf.csv # * Large EP / KV cache: /slurm_output/perf_csv/**.csv # Plus the legacy /perf.csv path some flows still use. - # Priority: results_dir config > shared_inference NFS > slurm_output/perf_csv - # > /perf.csv (with NFS-propagation retry). - perf_csv_path = None - if self.slurm_config.get("results_dir"): + # Priority: declared multiple_results > results_dir config > shared_inference + # NFS > slurm_output/perf_csv > /perf.csv (with NFS-propagation retry). + if not perf_csv_path and self.slurm_config.get("results_dir"): results_dir = Path(self.slurm_config["results_dir"]) candidates = list(results_dir.glob("perf*.csv")) if candidates: @@ -2139,16 +2252,25 @@ def _collect_slurm_multi_results( import shutil cwd_perf = Path("perf.csv") try: - if cwd_perf.exists(): - with open(perf_csv_path, "r") as src, open(cwd_perf, "a") as dst: - next(src, None) # skip per-job header so cwd CSV stays single-headed - for line in src: - dst.write(line) + # The source can legitimately resolve to the cwd perf.csv itself — + # via the /perf.csv fallback below, or a model card declaring + # multiple_results: "perf.csv". Appending a file to itself would + # duplicate every row, so there is nothing to aggregate. + if cwd_perf.exists() and perf_csv_path.resolve() == cwd_perf.resolve(): + self.console.print( + "[dim]Per-job perf is already the cwd perf.csv; nothing to aggregate[/dim]" + ) else: - shutil.copy(str(perf_csv_path), str(cwd_perf)) - self.console.print( - f"[green]✓ Aggregated per-job perf into {cwd_perf}[/green]" - ) + if cwd_perf.exists(): + with open(perf_csv_path, "r") as src, open(cwd_perf, "a") as dst: + next(src, None) # skip per-job header so cwd CSV stays single-headed + for line in src: + dst.write(line) + else: + shutil.copy(str(perf_csv_path), str(cwd_perf)) + self.console.print( + f"[green]✓ Aggregated per-job perf into {cwd_perf}[/green]" + ) except Exception as e: self.console.print( f"[yellow]⚠ Could not aggregate per-job perf into cwd perf.csv: {e}[/yellow]" diff --git a/src/madengine/orchestration/build_orchestrator.py b/src/madengine/orchestration/build_orchestrator.py index 17e836e8..0a00d037 100644 --- a/src/madengine/orchestration/build_orchestrator.py +++ b/src/madengine/orchestration/build_orchestrator.py @@ -282,6 +282,13 @@ def execute( } if len(card_images) == 1: implicit_image = next(iter(card_images)) + # Reject ""-style markers here rather than + # letting them become the image name every compute node tries + # to pull. + self._reject_placeholder_image( + implicit_image, + [m.get("name", "unknown") for m in slurm_multi_models], + ) self.rich_console.print( f"[dim]slurm_multi: no --registry/--use-image given; " f"using DOCKER_IMAGE_NAME from model card -> {implicit_image}[/dim]" @@ -665,6 +672,37 @@ def _execute_with_prebuilt_image( ), ) from e + # Model cards ship DOCKER_IMAGE_NAME as an angle-bracketed placeholder + # (e.g. "") to mark "you must supply this". A placeholder is + # not a usable image reference, and silently accepting one produces a confusing + # `docker pull ` failure on every compute node instead of an + # actionable message at submit time. + @staticmethod + def _is_placeholder_image(image: Optional[str]) -> bool: + """Return True if the value is a fill-me-in marker rather than an image ref.""" + if not image: + return True + candidate = image.strip() + return candidate.startswith("<") or candidate.endswith(">") + + def _reject_placeholder_image(self, image: str, model_names: List[str]) -> None: + """Raise a ConfigurationError if the resolved image is a placeholder.""" + if not self._is_placeholder_image(image): + return + raise ConfigurationError( + f"Model card DOCKER_IMAGE_NAME is a placeholder, not an image: {image!r}", + context=create_error_context( + operation="resolve_image", + component="BuildOrchestrator", + additional_info={"image": image, "model_names": model_names}, + ), + suggestions=[ + "Pass the real image explicitly: --use-image /:", + "Or build and push from the model's dockerfile: --registry ", + "Or replace DOCKER_IMAGE_NAME in the model card env_vars with a real image", + ], + ) + def _resolve_image_from_model_card(self) -> str: """ Resolve Docker image name from model card's DOCKER_IMAGE_NAME env var. @@ -742,7 +780,9 @@ def _resolve_image_from_model_card(self) -> str: ) else: self.rich_console.print(f"[green]✓ Auto-detected image: {resolved_image}[/green]\n") - + + self._reject_placeholder_image(resolved_image, list(images_found)) + return resolved_image def _execute_build_on_compute( diff --git a/tests/unit/test_orchestration.py b/tests/unit/test_orchestration.py index aaf24fcc..82079e46 100644 --- a/tests/unit/test_orchestration.py +++ b/tests/unit/test_orchestration.py @@ -446,3 +446,41 @@ def test_multiple_results_defaults_to_empty_string(self, mock_context, tmp_path) built_model = next(iter(manifest["built_models"].values())) assert built_model["multiple_results"] == "" + + +class TestPlaceholderImageRejection: + """Model cards ship DOCKER_IMAGE_NAME as a "" marker. + + The implicit --use-image path used to accept any single distinct card value, so + the placeholder became the image name and every compute node failed on + `docker pull ` instead of the user getting told at submit time. + """ + + @pytest.mark.parametrize("value", [ + "", + "", + " ", + "", + None, + ]) + def test_placeholders_detected(self, value): + assert BuildOrchestrator._is_placeholder_image(value) is True + + @pytest.mark.parametrize("value", [ + "rocm/vllm:latest", + "docker.io/myorg/img:tag", + "ci-pyt_vllm_kimi_k3_mi300x", + "localhost:5000/img", + ]) + def test_real_images_accepted(self, value): + assert BuildOrchestrator._is_placeholder_image(value) is False + + def test_reject_raises_configuration_error(self): + orchestrator = BuildOrchestrator.__new__(BuildOrchestrator) + with pytest.raises(ConfigurationError) as exc: + orchestrator._reject_placeholder_image("", ["m1"]) + assert "placeholder" in str(exc.value).lower() + + def test_reject_passes_through_real_image(self): + orchestrator = BuildOrchestrator.__new__(BuildOrchestrator) + orchestrator._reject_placeholder_image("rocm/vllm:latest", ["m1"]) diff --git a/tests/unit/test_slurm_multi.py b/tests/unit/test_slurm_multi.py index de5c2eae..cc286c85 100644 --- a/tests/unit/test_slurm_multi.py +++ b/tests/unit/test_slurm_multi.py @@ -595,3 +595,183 @@ def test_below_minimum_nodes_raises(self, deployment_factory): deployment_factory._generate_sglang_disagg_command( nnodes=1, nproc_per_node=8, master_port=12345 ) + + +# --------------------------------------------------------------------------- +# 6. Path-parity contract: the model card drives the allocation and the launcher +# +# slurm_multi is an escape hatch, not a second product. These lock in the pieces +# of the model-card contract that the templated path already honoured and the +# self-managed path silently dropped. + +from madengine.deployment.common import ( # noqa: E402 + resolve_launcher_from_sources, + resolve_node_count, +) + + +class TestResolveNodeCount: + """slurm.nodes and distributed.nnodes reconcile into one allocation size.""" + + def test_nnodes_absent_keeps_configured(self): + assert resolve_node_count(1, None, False) == (1, None) + + def test_nnodes_sizes_allocation_when_nodes_defaulted(self): + """The bug this fixes: a card declaring only nnodes ran on the nodes=1 default.""" + nodes, note = resolve_node_count(1, 4, nodes_explicitly_set=False) + assert nodes == 4 + assert note and "distributed.nnodes=4" in note + + def test_agreement_is_silent(self): + assert resolve_node_count(4, 4, nodes_explicitly_set=True) == (4, None) + + def test_explicit_nodes_wins_conflict_but_warns(self): + """Never silently allocate more nodes than the user asked to be billed for.""" + nodes, note = resolve_node_count(6, 4, nodes_explicitly_set=True) + assert nodes == 6 + assert note and "conflicts" in note + + @pytest.mark.parametrize("bad", ["bad", "", [], {}]) + def test_non_numeric_nnodes_falls_back(self, bad): + nodes, note = resolve_node_count(2, bad, nodes_explicitly_set=False) + assert nodes == 2 + if bad != "": + assert note is not None + + def test_idempotent_across_repeated_resolution(self): + """deploy() re-runs prepare() after preflight; resolution must not ratchet.""" + nodes, _ = resolve_node_count(1, 4, nodes_explicitly_set=False) + again, _ = resolve_node_count(1, 4, nodes_explicitly_set=False) + assert nodes == again == 4 + + +class TestResolveLauncherFromSources: + """One resolver for both dispatch sites, deployment config first.""" + + def test_deployment_config_wins(self): + assert resolve_launcher_from_sources("vllm", "slurm_multi") == "vllm" + + def test_falls_back_to_model_card(self): + assert resolve_launcher_from_sources(None, "slurm_multi") == "slurm_multi" + + def test_default_when_neither_declared(self): + assert resolve_launcher_from_sources(None, None) == "torchrun" + + def test_dispatch_and_env_block_cannot_disagree(self): + """ + prepare() used to read the card while _prepare_template_context() read the + deployment config, so a card-declared launcher could pick one path and emit + another path's env block. Both now call this, so they agree by construction. + """ + for deployment, card in [(None, "vllm"), ("sglang", "vllm"), (None, None)]: + assert resolve_launcher_from_sources(deployment, card) == \ + resolve_launcher_from_sources(deployment, card) + + +class TestSlurmMultiDeclaredResultsCsv: + """A slurm_multi card's `multiple_results` names its own results CSV.""" + + @pytest.fixture + def deployment(self, tmp_path: Path) -> SlurmDeployment: + script_rel = "scripts/wl/run.slurm" + script_abs = tmp_path / script_rel + script_abs.parent.mkdir(parents=True, exist_ok=True) + script_abs.write_text("#!/bin/bash\n") + + model = { + "name": "wl", + "scripts": script_rel, + "multiple_results": "perf_WL.csv", + "distributed": {"launcher": "slurm_multi"}, + } + manifest = { + "built_images": {}, + "built_models": {"img:tag": model}, + "context": {}, + } + manifest_path = tmp_path / "build_manifest.json" + manifest_path.write_text(json.dumps(manifest)) + cfg = DeploymentConfig( + target="slurm", + manifest_file=str(manifest_path), + additional_context={ + "slurm": {"output_dir": str(tmp_path / "slurm_results")}, + "distributed": {"launcher": "slurm_multi"}, + }, + ) + d = SlurmDeployment(cfg) + d._model_for_test = model + d._script_dir_for_test = script_abs.parent + return d + + def test_finds_csv_next_to_model_script(self, deployment): + """The wrapper cd's to the script dir, so $(pwd) writes land there.""" + target = deployment._script_dir_for_test / "perf_WL.csv" + target.write_text("model,performance,metric\nwl,1.0,tok/s\n") + found = deployment._slurm_multi_declared_result_csv( + deployment._model_for_test, "12345" + ) + assert found == target + + def test_ignores_empty_file(self, deployment): + (deployment._script_dir_for_test / "perf_WL.csv").write_text("") + assert deployment._slurm_multi_declared_result_csv( + deployment._model_for_test, "12345" + ) is None + + def test_returns_none_without_declaration(self, deployment): + assert deployment._slurm_multi_declared_result_csv({"name": "wl"}, "12345") is None + + def test_returns_none_when_model_info_missing(self, deployment): + assert deployment._slurm_multi_declared_result_csv(None, "12345") is None + + +class TestSlurmMultiPerfAggregation: + """Aggregating the per-job CSV into cwd/perf.csv must not append it to itself.""" + + @pytest.fixture + def deployment(self, tmp_path: Path) -> SlurmDeployment: + manifest = {"built_images": {}, "built_models": {}, "context": {}} + manifest_path = tmp_path / "build_manifest.json" + manifest_path.write_text(json.dumps(manifest)) + cfg = DeploymentConfig( + target="slurm", + manifest_file=str(manifest_path), + additional_context={"slurm": {"output_dir": str(tmp_path / "slurm_results")}}, + ) + d = SlurmDeployment(cfg) + d.output_dir.mkdir(parents=True, exist_ok=True) + return d + + def test_cwd_perf_source_is_not_duplicated(self, deployment, tmp_path, monkeypatch): + """ + The /perf.csv fallback (and a card declaring multiple_results: + "perf.csv") makes source and destination the same file. Appending it to + itself would double every row on each collection. + """ + monkeypatch.chdir(tmp_path) + rows = "model,performance,metric,status\nwl,1.0,tok/s,SUCCESS\n" + Path("perf.csv").write_text(rows) + + deployment._collect_slurm_multi_results( + "12345", {"perf_files": [], "logs": [], "successful_runs": [], "failed_runs": []}, None + ) + + assert Path("perf.csv").read_text() == rows + + def test_distinct_source_is_appended(self, deployment, tmp_path, monkeypatch): + """A genuinely separate per-job CSV still aggregates, minus its header.""" + monkeypatch.chdir(tmp_path) + Path("perf.csv").write_text("model,performance,metric,status\nold,1.0,tok/s,SUCCESS\n") + job_csv = tmp_path / "job" / "perf.csv" + job_csv.parent.mkdir() + job_csv.write_text("model,performance,metric,status\nnew,2.0,tok/s,SUCCESS\n") + deployment.slurm_config["results_dir"] = str(job_csv.parent) + + deployment._collect_slurm_multi_results( + "12345", {"perf_files": [], "logs": [], "successful_runs": [], "failed_runs": []}, None + ) + + text = Path("perf.csv").read_text() + assert "old,1.0" in text and "new,2.0" in text + assert text.count("model,performance") == 1 From def98bf0754c89149bb6519a46ed7988a6aa86d7 Mon Sep 17 00:00:00 2001 From: Cemberk Date: Mon, 17 Aug 2026 20:08:20 +0000 Subject: [PATCH 2/7] fix(run): a self-managed SLURM launcher implies a SLURM deployment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Target inference is convention-over-configuration: it keys purely on the presence of a `slurm` or `k8s` block. But model cards routinely declare only `distributed.launcher: slurm_multi` and no `slurm` block — all 37 slurm_multi entries in ROCm/MAD are shaped that way. Those inferred "local" and were handed to the container runner, so the slurm_multi path was never reached and the model's .slurm script would be run as an ordinary local container workload. slurm_multi drives sbatch/srun directly, so it is a SLURM deployment by construction. Infer it as one when no explicit block says otherwise; an explicit k8s or slurm block still wins, since that is the user's stated intent. Co-Authored-By: Claude Opus 5 (1M context) --- .../orchestration/run_orchestrator.py | 19 +++++++++----- tests/unit/test_orchestration.py | 26 +++++++++++++++++++ 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/src/madengine/orchestration/run_orchestrator.py b/src/madengine/orchestration/run_orchestrator.py index b0338d8c..c0d10d3f 100644 --- a/src/madengine/orchestration/run_orchestrator.py +++ b/src/madengine/orchestration/run_orchestrator.py @@ -31,6 +31,7 @@ ExecutionError, create_error_context, ) +from madengine.deployment.common import is_self_managed_launcher from madengine.utils.session_tracker import SessionTracker from madengine.orchestration.image_filtering import ( filter_images_by_gpu_compatibility as _filter_by_gpu_compat, @@ -1110,19 +1111,25 @@ def _infer_deployment_target(self, config: Dict) -> str: Convention over Configuration: - Presence of "k8s" or "kubernetes" field → k8s deployment - Presence of "slurm" field → slurm deployment - - Neither present → local execution - + - A self-managed SLURM launcher (slurm_multi) → slurm deployment + - None of the above → local execution + Args: config: Configuration dictionary - + Returns: Deployment target: "k8s", "slurm", or "local" """ if "k8s" in config or "kubernetes" in config: return "k8s" - elif "slurm" in config: + if "slurm" in config: return "slurm" - else: - return "local" + # slurm_multi runs the model's own .slurm script through sbatch/srun, so it is + # a SLURM deployment by construction. Without this a model card that declared + # the launcher but no `slurm` block inferred "local" and was handed to the + # container runner, which never reaches the slurm_multi path at all. + if is_self_managed_launcher((config.get("distributed") or {}).get("launcher")): + return "slurm" + return "local" diff --git a/tests/unit/test_orchestration.py b/tests/unit/test_orchestration.py index 82079e46..b20b1d04 100644 --- a/tests/unit/test_orchestration.py +++ b/tests/unit/test_orchestration.py @@ -484,3 +484,29 @@ def test_reject_raises_configuration_error(self): def test_reject_passes_through_real_image(self): orchestrator = BuildOrchestrator.__new__(BuildOrchestrator) orchestrator._reject_placeholder_image("rocm/vllm:latest", ["m1"]) + + +class TestSelfManagedLauncherImpliesSlurm: + """A slurm_multi model card without a `slurm` block still deploys to SLURM. + + Target inference keys on the presence of a `slurm`/`k8s` block. Model cards + routinely declare only `distributed.launcher: slurm_multi`, which inferred + "local" and handed the job to the container runner — so the slurm_multi path + was never reached and the model's .slurm script was run as a local workload. + """ + + @pytest.mark.parametrize("config,expected", [ + ({}, "local"), + ({"slurm": {}}, "slurm"), + ({"k8s": {}}, "k8s"), + ({"distributed": {"launcher": "slurm_multi"}}, "slurm"), + ({"distributed": {"launcher": "slurm-multi"}}, "slurm"), + ({"distributed": {"launcher": "torchrun"}}, "local"), + ({"distributed": {"launcher": "vllm"}}, "local"), + ({"distributed": {}}, "local"), + # An explicit k8s block still wins; slurm_multi is SLURM-only by construction + # but the explicit target is the user's stated intent. + ({"k8s": {}, "distributed": {"launcher": "slurm_multi"}}, "k8s"), + ]) + def test_inference(self, config, expected): + assert RunOrchestrator._infer_deployment_target(None, config) == expected From b18b264516efdbfcea780fab45f5cbe169106f8d Mon Sep 17 00:00:00 2001 From: Stephen Shao Date: Wed, 23 Sep 2026 21:27:04 +0000 Subject: [PATCH 3/7] fix(slurm): reject non-positive distributed.nnodes; drop a tautological test Two review findings from Copilot on #176: resolve_node_count() accepted nnodes values of 0 or -1 (and their string forms) and propagated them into slurm.nodes, so a typo in a model card became '#SBATCH --nodes=0' and failed at submit time with a scheduler error that named nothing useful. Non-positive values now fall back to the configured slurm.nodes with a note, matching the non-numeric path. test_dispatch_and_env_block_cannot_disagree compared the function's result to itself, so it could never fail. Replaced with an explicit expectation matrix, including a falsy deployment launcher to pin down that it does not override the model card. Co-Authored-By: Claude Opus 5 --- docs/launchers.md | 2 +- src/madengine/deployment/common.py | 9 +++++++++ tests/unit/test_slurm_multi.py | 30 ++++++++++++++++++++++++++---- 3 files changed, 36 insertions(+), 5 deletions(-) diff --git a/docs/launchers.md b/docs/launchers.md index 676e56b2..f2bdfcfc 100644 --- a/docs/launchers.md +++ b/docs/launchers.md @@ -662,7 +662,7 @@ fields are honoured identically on both paths: | Field | Effect | |-------|--------| | `distributed.launcher` | Selects the path. Resolved deployment-config-first, model-card-second — the same resolution used to emit the launcher env block, so the two can never disagree. | -| `distributed.nnodes` | Sizes the allocation (`#SBATCH --nodes` / `--ntasks`) when `slurm.nodes` is not set explicitly. If both are set and differ, `slurm.nodes` wins and a warning is printed. | +| `distributed.nnodes` | Sizes the allocation (`#SBATCH --nodes` / `--ntasks`) when `slurm.nodes` is not set explicitly. If both are set and differ, `slurm.nodes` wins and a warning is printed. Non-numeric or non-positive values are ignored with a warning — `slurm.nodes` is used instead. | | `slurm.*` | `partition`, `time`, `gpus_per_node`, `exclusive`, `reservation`, `nodelist`, … | | `multiple_results` | Names the results CSV. Searched next to the model script (where the wrapper `cd`s, so `$(pwd)` writes land there), then the job dir, output dir and cwd, before falling back to the conventional locations. | diff --git a/src/madengine/deployment/common.py b/src/madengine/deployment/common.py index c2fd605d..63842b01 100644 --- a/src/madengine/deployment/common.py +++ b/src/madengine/deployment/common.py @@ -175,6 +175,15 @@ def resolve_node_count( f"using slurm.nodes={configured_nodes}." ) + if requested is not None and requested < 1: + # An allocation of zero or fewer nodes is not a thing sbatch can schedule; + # passing it through would fail at submit time with a scheduler error that + # says nothing about the model card that caused it. + return configured_nodes, ( + f"Ignoring distributed.nnodes={requested} (must be >= 1); " + f"using slurm.nodes={configured_nodes}." + ) + if requested is None or requested == configured_nodes: return configured_nodes, None diff --git a/tests/unit/test_slurm_multi.py b/tests/unit/test_slurm_multi.py index 22903262..6d0b3f86 100644 --- a/tests/unit/test_slurm_multi.py +++ b/tests/unit/test_slurm_multi.py @@ -639,6 +639,18 @@ def test_non_numeric_nnodes_falls_back(self, bad): if bad != "": assert note is not None + @pytest.mark.parametrize("bad", [0, -1, "0", "-1"]) + def test_non_positive_nnodes_falls_back(self, bad): + """sbatch cannot schedule --nodes=0; reject it here, not at submit time.""" + nodes, note = resolve_node_count(2, bad, nodes_explicitly_set=False) + assert nodes == 2 + assert note and ">= 1" in note + + def test_non_positive_nnodes_does_not_override_explicit_nodes(self): + nodes, note = resolve_node_count(4, 0, nodes_explicitly_set=True) + assert nodes == 4 + assert note is not None + def test_idempotent_across_repeated_resolution(self): """deploy() re-runs prepare() after preflight; resolution must not ratchet.""" nodes, _ = resolve_node_count(1, 4, nodes_explicitly_set=False) @@ -658,15 +670,25 @@ def test_falls_back_to_model_card(self): def test_default_when_neither_declared(self): assert resolve_launcher_from_sources(None, None) == "torchrun" - def test_dispatch_and_env_block_cannot_disagree(self): + @pytest.mark.parametrize("deployment,card,expected", [ + (None, "vllm", "vllm"), + ("sglang", "vllm", "sglang"), + (None, None, "torchrun"), + ("sglang", None, "sglang"), + ("slurm_multi", "torchrun", "slurm_multi"), + # A falsy deployment value is "not declared", not "declared as empty": + # ConfigLoader materializes distributed.launcher as "" when the preset + # carries the key but neither the user nor the card filled it in. + ("", "slurm_multi", "slurm_multi"), + ("", None, "torchrun"), + ]) + def test_dispatch_and_env_block_cannot_disagree(self, deployment, card, expected): """ prepare() used to read the card while _prepare_template_context() read the deployment config, so a card-declared launcher could pick one path and emit another path's env block. Both now call this, so they agree by construction. """ - for deployment, card in [(None, "vllm"), ("sglang", "vllm"), (None, None)]: - assert resolve_launcher_from_sources(deployment, card) == \ - resolve_launcher_from_sources(deployment, card) + assert resolve_launcher_from_sources(deployment, card) == expected class TestSlurmMultiDeclaredResultsCsv: From bc790220b89c7fe24f55cdcadc1ebd974f696cbe Mon Sep 17 00:00:00 2001 From: Stephen Shao Date: Wed, 23 Sep 2026 22:05:48 +0000 Subject: [PATCH 4/7] fix(slurm): persist launcher/slurm.* for all build paths, track explicit-key provenance Addresses remaining PR #176 review comments: - Extract the model-card distributed/slurm merge (previously only run by _execute_with_prebuilt_image) into BuildOrchestrator._merge_model_config_into_manifest and call it from the normal Docker-build path too, so a slurm_multi model card's distributed.launcher reaches deployment_config regardless of build path. Without this, run --manifest-file inferred "local" for models built without --use-image. - Record which slurm.* keys were actually explicit (--additional-context or model card) as deployment_config._explicit_slurm_keys in the manifest, and have SlurmDeployment prefer that provenance over inferring explicitness from which keys are present. A manifest's slurm dict is serialized after ConfigLoader applies preset defaults, so a persisted nodes=1 default was indistinguishable from a real user/model-card setting and could silently override distributed.nnodes. - Expand the slurm.* merge whitelist to include account, qos, modules, skip_gpus_directive, and results_dir, matching what SlurmDeployment actually reads and docs/configuration.md documents. Update docs/launchers.md's model-card contract table to match. Co-Authored-By: Claude Sonnet 5 --- docs/launchers.md | 2 +- src/madengine/deployment/slurm.py | 17 +- .../orchestration/build_orchestrator.py | 163 +++++++++++------- .../orchestration/run_orchestrator.py | 10 +- tests/unit/test_orchestration.py | 91 ++++++++++ tests/unit/test_slurm_multi.py | 53 ++++++ 6 files changed, 268 insertions(+), 68 deletions(-) diff --git a/docs/launchers.md b/docs/launchers.md index f2bdfcfc..ffc89a0a 100644 --- a/docs/launchers.md +++ b/docs/launchers.md @@ -663,7 +663,7 @@ fields are honoured identically on both paths: |-------|--------| | `distributed.launcher` | Selects the path. Resolved deployment-config-first, model-card-second — the same resolution used to emit the launcher env block, so the two can never disagree. | | `distributed.nnodes` | Sizes the allocation (`#SBATCH --nodes` / `--ntasks`) when `slurm.nodes` is not set explicitly. If both are set and differ, `slurm.nodes` wins and a warning is printed. Non-numeric or non-positive values are ignored with a warning — `slurm.nodes` is used instead. | -| `slurm.*` | `partition`, `time`, `gpus_per_node`, `exclusive`, `reservation`, `nodelist`, … | +| `slurm.*` | `partition`, `nodes`, `gpus_per_node`, `time`, `exclusive`, `reservation`, `output_dir`, `nodelist`, `account`, `qos`, `modules`, `skip_gpus_directive`, `results_dir` — the full set the prebuilt-image and normal build manifest merges promote into `deployment_config.slurm`. Other `slurm.*` fields documented in [configuration.md](configuration.md) are read at deploy time but are not copied from the model card into a manifest's `deployment_config`. | | `multiple_results` | Names the results CSV. Searched next to the model script (where the wrapper `cd`s, so `$(pwd)` writes land there), then the job dir, output dir and cwd, before falling back to the conventional locations. | > **`args` is not sbatch flags.** On the templated path `args` goes to the diff --git a/src/madengine/deployment/slurm.py b/src/madengine/deployment/slurm.py index 992dff1e..21b5ad5e 100644 --- a/src/madengine/deployment/slurm.py +++ b/src/madengine/deployment/slurm.py @@ -74,7 +74,22 @@ def __init__(self, config: DeploymentConfig): # deployment_config) BEFORE ConfigLoader layers its presets on top. Once the # defaults are applied every key looks "set", and nodes=1 from a preset is # indistinguishable from nodes=1 the user asked for. - self._explicit_slurm_keys = set((config.additional_context or {}).get("slurm") or {}) + # + # A build_manifest.json persists the slurm dict *after* BuildOrchestrator's + # own ConfigLoader defaulting, so inferring explicitness from present keys + # here would treat a persisted preset default (e.g. nodes=1) as explicit. + # Prefer the "_explicit_slurm_keys" provenance BuildOrchestrator recorded at + # build time when present; only fall back to inferring from dict keys for a + # config that never went through a manifest (e.g. direct --additional-context). + explicit_keys_from_manifest = (config.additional_context or {}).get( + "_explicit_slurm_keys" + ) + if explicit_keys_from_manifest is not None: + self._explicit_slurm_keys = set(explicit_keys_from_manifest) + else: + self._explicit_slurm_keys = set( + (config.additional_context or {}).get("slurm") or {} + ) apply_deployment_config(config, ConfigLoader.load_slurm_config) super().__init__(config) diff --git a/src/madengine/orchestration/build_orchestrator.py b/src/madengine/orchestration/build_orchestrator.py index 8cbbd4e5..bbcf4684 100644 --- a/src/madengine/orchestration/build_orchestrator.py +++ b/src/madengine/orchestration/build_orchestrator.py @@ -427,6 +427,11 @@ def execute( # Step 6: Save deployment_config to manifest self._save_deployment_config(manifest_output) + # Merge model card distributed/slurm config into deployment_config so + # a slurm_multi launcher (or slurm.* fields) is visible at run time + # even without --additional-context. See _merge_model_config_into_manifest. + self._merge_model_config_into_manifest(manifest_output, models) + self.rich_console.print(f"[green]✓ Build complete: {manifest_output}[/green]") self.rich_console.print(f"[dim]{'=' * 60}[/dim]\n") @@ -591,70 +596,11 @@ def _execute_with_prebuilt_image( # Save deployment config self._save_deployment_config(manifest_output) - - # Merge model's distributed and slurm config into deployment_config - # This ensures launcher and slurm settings are in deployment_config even if not in additional-context - if models: - with open(manifest_output, "r") as f: - saved_manifest = json.load(f) - - if "deployment_config" not in saved_manifest: - saved_manifest["deployment_config"] = {} - - # Merge model's distributed config from the first model. - # If multiple models have differing distributed configs, warn — only the first wins here. - # Use json.dumps for the hash key so nested dicts (e.g. sglang_disagg / vllm_disagg) - # don't trigger TypeError: unhashable type: 'dict' from `tuple(sorted(items()))`. - if len(models) > 1: - distinct_distributed = { - json.dumps(m.get("distributed") or {}, sort_keys=True, default=str) - for m in models - } - if len(distinct_distributed) > 1: - self.rich_console.print( - "[yellow]Warning: discovered models have differing distributed configs; " - f"using {models[0].get('name', '')}'s config.[/yellow]" - ) - model_distributed = models[0].get("distributed", {}) - if model_distributed: - if "distributed" not in saved_manifest["deployment_config"]: - saved_manifest["deployment_config"]["distributed"] = {} - - # Copy launcher and other critical fields from model config - for key in ["launcher", "nnodes", "nproc_per_node", "backend", "port", "sglang_disagg", "vllm_disagg"]: - if key in model_distributed and key not in saved_manifest["deployment_config"]["distributed"]: - saved_manifest["deployment_config"]["distributed"][key] = model_distributed[key] - - # Merge model's slurm config into deployment_config.slurm from the first model. - # This enables run phase to auto-detect SLURM deployment without --additional-context. - # Warn when multiple models have differing slurm configs (only the first wins here). - # json.dumps key for the same unhashable-nested-dict reason as above. - if len(models) > 1: - distinct_slurm = { - json.dumps(m.get("slurm") or {}, sort_keys=True, default=str) - for m in models - } - if len(distinct_slurm) > 1: - self.rich_console.print( - "[yellow]Warning: discovered models have differing slurm configs; " - f"using {models[0].get('name', '')}'s config.[/yellow]" - ) - model_slurm = models[0].get("slurm", {}) - if model_slurm: - if "slurm" not in saved_manifest["deployment_config"]: - saved_manifest["deployment_config"]["slurm"] = {} - - # Copy slurm settings from model config (model card fills in - # values not explicitly set by --additional-context). - # Use _original_user_slurm_keys (captured before ConfigLoader - # applies defaults) so model card values override defaults - # but user's explicit CLI values still win. - for key in ["partition", "nodes", "gpus_per_node", "time", "exclusive", "reservation", "output_dir", "nodelist"]: - if key in model_slurm and key not in self._original_user_slurm_keys: - saved_manifest["deployment_config"]["slurm"][key] = model_slurm[key] - - with open(manifest_output, "w") as f: - json.dump(saved_manifest, f, indent=2) + + # Merge model's distributed and slurm config into deployment_config. + # This ensures launcher and slurm settings are in deployment_config even + # if not in additional-context (see _merge_model_config_into_manifest). + self._merge_model_config_into_manifest(manifest_output, models) self.rich_console.print(f"[green]✓ Generated manifest: {manifest_output}[/green]") self.rich_console.print(f" Pre-built image: {use_image}") @@ -1343,6 +1289,95 @@ def _save_build_summary(self, manifest_file: str, build_summary: Dict): except Exception as e: self.rich_console.print(f"[yellow]Warning: Could not save build summary: {e}[/yellow]") + def _merge_model_config_into_manifest(self, manifest_output: str, models: list): + """Merge model card distributed/slurm config into manifest deployment_config. + + Used by every build path (normal Docker build, prebuilt-image, implicit + DOCKER_IMAGE_NAME) so a model card's ``distributed.launcher`` and ``slurm.*`` + fields land in ``deployment_config`` regardless of how the image was built. + Without this, only the prebuilt-image path populated ``deployment_config``, + so a model card declaring ``distributed.launcher: slurm_multi`` was invisible + to ``run --manifest-file`` target inference for a normal build. + """ + if not models: + return + + with open(manifest_output, "r") as f: + saved_manifest = json.load(f) + + if "deployment_config" not in saved_manifest: + saved_manifest["deployment_config"] = {} + + # Merge model's distributed config from the first model. + # If multiple models have differing distributed configs, warn — only the first wins here. + # Use json.dumps for the hash key so nested dicts (e.g. sglang_disagg / vllm_disagg) + # don't trigger TypeError: unhashable type: 'dict' from `tuple(sorted(items()))`. + if len(models) > 1: + distinct_distributed = { + json.dumps(m.get("distributed") or {}, sort_keys=True, default=str) + for m in models + } + if len(distinct_distributed) > 1: + self.rich_console.print( + "[yellow]Warning: discovered models have differing distributed configs; " + f"using {models[0].get('name', '')}'s config.[/yellow]" + ) + model_distributed = models[0].get("distributed", {}) + if model_distributed: + if "distributed" not in saved_manifest["deployment_config"]: + saved_manifest["deployment_config"]["distributed"] = {} + + # Copy launcher and other critical fields from model config + for key in ["launcher", "nnodes", "nproc_per_node", "backend", "port", "sglang_disagg", "vllm_disagg"]: + if key in model_distributed and key not in saved_manifest["deployment_config"]["distributed"]: + saved_manifest["deployment_config"]["distributed"][key] = model_distributed[key] + + # Merge model's slurm config into deployment_config.slurm from the first model. + # This enables run phase to auto-detect SLURM deployment without --additional-context. + # Warn when multiple models have differing slurm configs (only the first wins here). + # json.dumps key for the same unhashable-nested-dict reason as above. + if len(models) > 1: + distinct_slurm = { + json.dumps(m.get("slurm") or {}, sort_keys=True, default=str) + for m in models + } + if len(distinct_slurm) > 1: + self.rich_console.print( + "[yellow]Warning: discovered models have differing slurm configs; " + f"using {models[0].get('name', '')}'s config.[/yellow]" + ) + model_slurm = models[0].get("slurm", {}) + explicit_slurm_keys = set(self._original_user_slurm_keys) + if model_slurm: + if "slurm" not in saved_manifest["deployment_config"]: + saved_manifest["deployment_config"]["slurm"] = {} + + # Copy slurm settings from model config (model card fills in + # values not explicitly set by --additional-context). + # Use _original_user_slurm_keys (captured before ConfigLoader + # applies defaults) so model card values override defaults + # but user's explicit CLI values still win. + for key in [ + "partition", "nodes", "gpus_per_node", "time", "exclusive", + "reservation", "output_dir", "nodelist", "account", "qos", + "modules", "skip_gpus_directive", "results_dir", + ]: + if key in model_slurm and key not in self._original_user_slurm_keys: + saved_manifest["deployment_config"]["slurm"][key] = model_slurm[key] + explicit_slurm_keys.add(key) + + # Record which slurm keys are explicit (user --additional-context or model + # card) as opposed to ConfigLoader preset defaults, so a later `run + # --manifest-file` can tell a real distributed.nnodes-vs-slurm.nodes + # conflict apart from a preset default silently outranking the model card. + if explicit_slurm_keys: + saved_manifest["deployment_config"]["_explicit_slurm_keys"] = sorted( + explicit_slurm_keys + ) + + with open(manifest_output, "w") as f: + json.dump(saved_manifest, f, indent=2) + def _save_deployment_config(self, manifest_file: str): """Save deployment_config from --additional-context to manifest.""" if not self.additional_context: diff --git a/src/madengine/orchestration/run_orchestrator.py b/src/madengine/orchestration/run_orchestrator.py index 3f9b58d3..30bb3027 100644 --- a/src/madengine/orchestration/run_orchestrator.py +++ b/src/madengine/orchestration/run_orchestrator.py @@ -238,8 +238,14 @@ def execute( if not self.additional_context: self.additional_context = {} - # Merge deployment_config into additional_context (for deployment layer to use) - for key in ["slurm", "k8s", "kubernetes", "distributed", "vllm", "env_vars", "debug"]: + # Merge deployment_config into additional_context (for deployment layer to use). + # "_explicit_slurm_keys" carries provenance for slurm.* fields (user/model-card + # explicit vs. ConfigLoader preset default) captured at build time — see + # BuildOrchestrator._merge_model_config_into_manifest. + for key in [ + "slurm", "k8s", "kubernetes", "distributed", "vllm", "env_vars", + "debug", "_explicit_slurm_keys", + ]: if key in deployment_config and key not in self.additional_context: self.additional_context[key] = deployment_config[key] diff --git a/tests/unit/test_orchestration.py b/tests/unit/test_orchestration.py index daacf07a..2740585b 100644 --- a/tests/unit/test_orchestration.py +++ b/tests/unit/test_orchestration.py @@ -638,6 +638,97 @@ def test_inference(self, config, expected): assert RunOrchestrator._infer_deployment_target(None, config) == expected +@pytest.mark.unit +class TestMergeModelConfigIntoManifest: + """BuildOrchestrator._merge_model_config_into_manifest is shared by every build + path (normal Docker build, prebuilt-image, implicit DOCKER_IMAGE_NAME), so a model + card's distributed/slurm declarations reach deployment_config no matter how the + image was built. + """ + + @patch("madengine.orchestration.build_orchestrator.Context") + @patch("os.path.exists", return_value=False) + def _make_orchestrator(self, mock_exists, mock_context, additional_context=None): + mock_args = MagicMock() + mock_args.additional_context = additional_context + mock_args.additional_context_file = None + mock_args.live_output = True + return BuildOrchestrator(mock_args) + + def test_normal_build_path_persists_model_card_launcher(self, tmp_path): + """Regression: only the prebuilt-image path used to copy distributed.launcher + into deployment_config, so a normal Docker build of a slurm_multi model card + left deployment_config without a launcher and `run` inferred "local".""" + orchestrator = self._make_orchestrator() + manifest_file = tmp_path / "build_manifest.json" + manifest_file.write_text(json.dumps({"built_models": {}})) + + models = [{"name": "m1", "distributed": {"launcher": "slurm_multi"}}] + orchestrator._merge_model_config_into_manifest(str(manifest_file), models) + + saved = json.loads(manifest_file.read_text()) + assert saved["deployment_config"]["distributed"]["launcher"] == "slurm_multi" + + def test_expanded_slurm_whitelist_promotes_all_documented_fields(self, tmp_path): + """docs/launchers.md documents these slurm.* fields as part of the model-card + contract; the merge whitelist must actually promote every one of them.""" + orchestrator = self._make_orchestrator() + manifest_file = tmp_path / "build_manifest.json" + manifest_file.write_text(json.dumps({"built_models": {}})) + + model_slurm = { + "partition": "gpu", + "nodes": 4, + "gpus_per_node": 8, + "time": "12:00:00", + "exclusive": True, + "reservation": "myres", + "output_dir": "./out", + "nodelist": "node[1-4]", + "account": "myaccount", + "qos": "high", + "modules": ["rocm/6.0"], + "skip_gpus_directive": True, + "results_dir": "./results", + } + models = [{"name": "m1", "slurm": model_slurm}] + orchestrator._merge_model_config_into_manifest(str(manifest_file), models) + + saved_slurm = json.loads(manifest_file.read_text())["deployment_config"]["slurm"] + for key, value in model_slurm.items(): + assert saved_slurm[key] == value + + def test_explicit_slurm_keys_provenance_distinguishes_model_card_from_default( + self, tmp_path + ): + """Regression: a manifest's slurm dict is written *after* ConfigLoader applies + preset defaults, so a persisted nodes=1 default was indistinguishable from a + real user/model-card setting on a later `run --manifest-file`. The merge must + record which keys were actually explicit (user additional-context or model + card) as opposed to a preset default.""" + orchestrator = self._make_orchestrator( + additional_context='{"slurm": {"partition": "gpu"}}' + ) + manifest_file = tmp_path / "build_manifest.json" + manifest_file.write_text(json.dumps({"built_models": {}})) + + # Model card declares nnodes-sizing via slurm.nodes; "time" here simulates a + # ConfigLoader-applied preset default already present in additional_context, + # which must NOT be treated as explicit. + models = [{"name": "m1", "slurm": {"nodes": 4}}] + orchestrator._merge_model_config_into_manifest(str(manifest_file), models) + + explicit_keys = set( + json.loads(manifest_file.read_text())["deployment_config"][ + "_explicit_slurm_keys" + ] + ) + assert "partition" in explicit_keys # from --additional-context + assert "nodes" in explicit_keys # copied from the model card + assert "time" not in explicit_keys # never declared by user or model card + + +@pytest.mark.unit class TestRequirePinnedImageContext: """--require-pinned-image and require_pinned_image both reach additional_context.""" diff --git a/tests/unit/test_slurm_multi.py b/tests/unit/test_slurm_multi.py index 6d0b3f86..021a625b 100644 --- a/tests/unit/test_slurm_multi.py +++ b/tests/unit/test_slurm_multi.py @@ -691,6 +691,59 @@ def test_dispatch_and_env_block_cannot_disagree(self, deployment, card, expected assert resolve_launcher_from_sources(deployment, card) == expected +class TestExplicitSlurmKeysManifestProvenance: + """A manifest's slurm dict is written *after* ConfigLoader applies preset + defaults, so inferring "explicit" from which keys are present in the manifest + would treat a persisted default (e.g. nodes=1) as if the user or model card had + set it. BuildOrchestrator records real provenance as "_explicit_slurm_keys" in + deployment_config; SlurmDeployment must prefer that over inferring from keys. + """ + + def _deployment(self, tmp_path: Path, additional_context: dict) -> SlurmDeployment: + manifest = {"built_images": {}, "built_models": {}, "context": {}} + manifest_path = tmp_path / "build_manifest.json" + manifest_path.write_text(json.dumps(manifest)) + cfg = DeploymentConfig( + target="slurm", + manifest_file=str(manifest_path), + additional_context=additional_context, + ) + return SlurmDeployment(cfg) + + def test_persisted_default_is_not_treated_as_explicit(self, tmp_path): + """Regression: slurm.nodes=1 came from a ConfigLoader preset default at build + time, not from the user or model card, so distributed.nnodes must still win.""" + d = self._deployment(tmp_path, { + "slurm": {"nodes": 1, "output_dir": str(tmp_path / "slurm_results")}, + "distributed": {"nnodes": 4}, + "_explicit_slurm_keys": ["output_dir"], + }) + assert "nodes" not in d._explicit_slurm_keys + assert d.nodes == 4 + + def test_manifest_provenance_respects_real_explicit_conflict(self, tmp_path): + """When the manifest says slurm.nodes really was explicit, an explicit + conflict with distributed.nnodes must still win (with a warning).""" + d = self._deployment(tmp_path, { + "slurm": {"nodes": 6, "output_dir": str(tmp_path / "slurm_results")}, + "distributed": {"nnodes": 4}, + "_explicit_slurm_keys": ["nodes", "output_dir"], + }) + assert "nodes" in d._explicit_slurm_keys + assert d.nodes == 6 + + def test_falls_back_to_key_inference_without_manifest_provenance(self, tmp_path): + """Direct --additional-context (no manifest round trip) has no + "_explicit_slurm_keys" key; behavior falls back to inferring from the + slurm dict's own keys, as before this fix.""" + d = self._deployment(tmp_path, { + "slurm": {"nodes": 6, "output_dir": str(tmp_path / "slurm_results")}, + "distributed": {"nnodes": 4}, + }) + assert d._explicit_slurm_keys == {"nodes", "output_dir"} + assert d.nodes == 6 + + class TestSlurmMultiDeclaredResultsCsv: """A slurm_multi card's `multiple_results` names its own results CSV.""" From 5f67ee72b384ad86c52c8e487ebf6fbe389536fd Mon Sep 17 00:00:00 2001 From: Stephen Shao Date: Wed, 23 Sep 2026 23:09:20 +0000 Subject: [PATCH 5/7] fix(slurm): close the three open Copilot findings on the parity path --build-on-compute wrote its own manifest and returned before _merge_model_config_into_manifest, so it skipped the shared model-field and provenance handling: the card's multiple_results was dropped, no _explicit_slurm_keys was recorded, and because the merge read self.additional_context (post-ConfigLoader) the preset defaults outranked the card. Route it through the same helper and apply the same explicit-key precedence rule when deriving its slurm config. _prepare_slurm_multi_script promoted account/qos/modules into deployment_config and then dropped them: unlike job.sh.j2 the wrapper emitted no #SBATCH --account/--qos and ran no module load, so a card declaring them was silently ignored on the self-managed path. Emit them, keeping the directives inside the header block sbatch actually parses. _load_and_merge_manifest replaced the manifest's slurm block with the runtime one but left _explicit_slurm_keys describing the old block, so a build-time explicit nodes marked the runtime default nodes=1 as deliberate and suppressed distributed.nnodes. Recompute the provenance from the runtime keys whenever runtime slurm is supplied. Also stop _merge_model_config_into_manifest from turning a missing manifest into a BuildError. It runs after _save_build_summary and _save_deployment_config, which both degrade to a warning when export_build_manifest() produced nothing; the new step did not, which failed test_build_then_run_workflow. Co-Authored-By: Claude Opus 5 --- docs/launchers.md | 2 +- src/madengine/deployment/slurm.py | 26 +- .../orchestration/build_orchestrator.py | 48 ++- .../orchestration/run_orchestrator.py | 12 + tests/unit/test_slurm_multi.py | 284 ++++++++++++++++++ 5 files changed, 363 insertions(+), 9 deletions(-) diff --git a/docs/launchers.md b/docs/launchers.md index ffc89a0a..f2a24c0c 100644 --- a/docs/launchers.md +++ b/docs/launchers.md @@ -663,7 +663,7 @@ fields are honoured identically on both paths: |-------|--------| | `distributed.launcher` | Selects the path. Resolved deployment-config-first, model-card-second — the same resolution used to emit the launcher env block, so the two can never disagree. | | `distributed.nnodes` | Sizes the allocation (`#SBATCH --nodes` / `--ntasks`) when `slurm.nodes` is not set explicitly. If both are set and differ, `slurm.nodes` wins and a warning is printed. Non-numeric or non-positive values are ignored with a warning — `slurm.nodes` is used instead. | -| `slurm.*` | `partition`, `nodes`, `gpus_per_node`, `time`, `exclusive`, `reservation`, `output_dir`, `nodelist`, `account`, `qos`, `modules`, `skip_gpus_directive`, `results_dir` — the full set the prebuilt-image and normal build manifest merges promote into `deployment_config.slurm`. Other `slurm.*` fields documented in [configuration.md](configuration.md) are read at deploy time but are not copied from the model card into a manifest's `deployment_config`. | +| `slurm.*` | `partition`, `nodes`, `gpus_per_node`, `time`, `exclusive`, `reservation`, `output_dir`, `nodelist`, `account`, `qos`, `modules`, `skip_gpus_directive`, `results_dir` — the full set every build path (normal build, prebuilt image, `--build-on-compute`) promotes into `deployment_config.slurm`. Other `slurm.*` fields documented in [configuration.md](configuration.md) are read at deploy time but are not copied from the model card into a manifest's `deployment_config`. | | `multiple_results` | Names the results CSV. Searched next to the model script (where the wrapper `cd`s, so `$(pwd)` writes land there), then the job dir, output dir and cwd, before falling back to the conventional locations. | > **`args` is not sbatch flags.** On the templated path `args` goes to the diff --git a/src/madengine/deployment/slurm.py b/src/madengine/deployment/slurm.py index 21b5ad5e..b83a81f0 100644 --- a/src/madengine/deployment/slurm.py +++ b/src/madengine/deployment/slurm.py @@ -594,7 +594,17 @@ def _prepare_slurm_multi_script(self, model_info: Dict, docker_image_name: str = nodelist = self._normalize_nodelist(self.slurm_config.get("nodelist")) if nodelist: script_lines.append(f"#SBATCH --nodelist={nodelist}") - + + # Accounting directives, same as job.sh.j2. Sites that bill jobs reject a + # submission without them, so a model card declaring qos/account has to be + # honoured here too or the self-managed path silently drops it. + qos = self.slurm_config.get("qos") + if qos: + script_lines.append(f"#SBATCH --qos={qos}") + account = self.slurm_config.get("account") + if account: + script_lines.append(f"#SBATCH --account={account}") + script_lines.extend([ "", f"# slurm_multi launcher script for {model_info['name']}", @@ -602,9 +612,19 @@ def _prepare_slurm_multi_script(self, model_info: Dict, docker_image_name: str = "", "set -e", "", - "# Environment variables", ]) - + + # `module load` before anything else runs, matching job.sh.j2: these bring + # the site's docker/rocm/mpi into PATH for the model's own srun calls. + modules = self.slurm_config.get("modules", []) or [] + if modules: + script_lines.append("# Load required modules") + for module in modules: + script_lines.append(f"module load {shlex.quote(str(module))}") + script_lines.append("") + + script_lines.append("# Environment variables") + for key, value in env_vars.items(): script_lines.append(f"export {key}={shlex.quote(str(value))}") diff --git a/src/madengine/orchestration/build_orchestrator.py b/src/madengine/orchestration/build_orchestrator.py index bbcf4684..b04af35c 100644 --- a/src/madengine/orchestration/build_orchestrator.py +++ b/src/madengine/orchestration/build_orchestrator.py @@ -865,15 +865,30 @@ def _execute_build_on_compute( # All models are built in a single sbatch job, so one SLURM config applies to all. first_model = models[0] model_slurm_config = first_model.get("slurm", {}) + # self.additional_context is post-ConfigLoader, so it carries preset defaults + # (e.g. nodes=1) alongside anything the user actually passed. Only the keys + # captured before defaulting may outrank the model card — otherwise a preset + # default silently beats a declared slurm.* field. Same precedence rule as + # _merge_model_config_into_manifest. context_slurm_config = self.additional_context.get("slurm", {}) - slurm_config = {**model_slurm_config, **context_slurm_config} + slurm_config = { + **context_slurm_config, + **{ + key: value + for key, value in model_slurm_config.items() + if key not in self._original_user_slurm_keys + }, + } self.rich_console.print(f"[green]✓ Found {len(models)} model(s)[/green]\n") self.rich_console.print("[bold cyan]📋 SLURM Configuration (merged):[/bold cyan]") if model_slurm_config: self.rich_console.print(f" [dim]From model card:[/dim] {list(model_slurm_config.keys())}") - if context_slurm_config: - self.rich_console.print(f" [dim]From --additional-context (overrides):[/dim] {list(context_slurm_config.keys())}") + if self._original_user_slurm_keys: + self.rich_console.print( + f" [dim]From --additional-context (overrides):[/dim] " + f"{sorted(self._original_user_slurm_keys)}" + ) # Validate required fields partition = slurm_config.get("partition") @@ -1221,6 +1236,10 @@ def _execute_build_on_compute( "data": m.get("data", ""), "n_gpus": m.get("n_gpus", "8"), "tags": m.get("tags", []), + # Same field set as _execute_with_prebuilt_image: without this a + # model card's declared results CSV is lost, and the run falls + # back to log scraping with nothing to scrape. + "multiple_results": m.get("multiple_results", ""), "slurm": slurm_config, "distributed": m.get("distributed", {}), "env_vars": {**m.get("env_vars", {}), "DOCKER_IMAGE_NAME": rim}, @@ -1248,6 +1267,14 @@ def _execute_build_on_compute( with open(manifest_output, "w") as f: json.dump(manifest, f, indent=2) + # Route this path through the same model-field/provenance handling as the + # other build paths, so a model card's distributed.* lands in + # deployment_config and "_explicit_slurm_keys" records which slurm.* keys + # were really set (vs. a ConfigLoader preset default). Without it, run + # would infer explicitness from the persisted keys and let a default + # nodes=1 outrank the card's distributed.nnodes. + self._merge_model_config_into_manifest(manifest_output, models) + self.rich_console.print(f"[green]✓ Build completed on compute node[/green]") for pmd in per_model_data: self.rich_console.print(f"[green]✓ Image pushed: {pmd['registry_image_name']}[/green]") @@ -1302,8 +1329,19 @@ def _merge_model_config_into_manifest(self, manifest_output: str, models: list): if not models: return - with open(manifest_output, "r") as f: - saved_manifest = json.load(f) + try: + with open(manifest_output, "r") as f: + saved_manifest = json.load(f) + except Exception as e: + # Best-effort, same as the _save_build_summary/_save_deployment_config + # steps that run just before this one: export_build_manifest() may have + # written nothing (no manifest, or an unreadable one). That is already + # surfaced by those steps' warnings, so don't escalate it to a BuildError + # and lose the build results. + self.rich_console.print( + f"[yellow]Warning: Could not merge model config into manifest: {e}[/yellow]" + ) + return if "deployment_config" not in saved_manifest: saved_manifest["deployment_config"] = {} diff --git a/src/madengine/orchestration/run_orchestrator.py b/src/madengine/orchestration/run_orchestrator.py index 30bb3027..227a32c4 100644 --- a/src/madengine/orchestration/run_orchestrator.py +++ b/src/madengine/orchestration/run_orchestrator.py @@ -508,6 +508,18 @@ def _load_and_merge_manifest(self, manifest_file: str) -> str: for key in ["deploy", "slurm", "k8s", "kubernetes", "distributed", "vllm", "env_vars", "debug"]: if key in self.additional_context: stored_config[key] = self.additional_context[key] + # "_explicit_slurm_keys" describes the slurm block it was recorded + # against. Replacing that block above makes the build-time + # provenance stale: a manifest built with an explicit slurm.nodes + # would keep "nodes" marked explicit even though the runtime slurm + # block never sets it, so SlurmDeployment would treat its own + # default nodes=1 as deliberate and ignore distributed.nnodes. + # RunOrchestrator does not run --additional-context through + # ConfigLoader, so the runtime keys are exactly what the user set. + if "slurm" in self.additional_context: + stored_config["_explicit_slurm_keys"] = sorted( + self.additional_context.get("slurm") or {} + ) manifest["deployment_config"] = stored_config # Merge context (tools, pre_scripts, post_scripts, encapsulate_script) diff --git a/tests/unit/test_slurm_multi.py b/tests/unit/test_slurm_multi.py index 021a625b..c6c1918e 100644 --- a/tests/unit/test_slurm_multi.py +++ b/tests/unit/test_slurm_multi.py @@ -941,3 +941,287 @@ def test_enabled_without_digest_does_not_silently_fall_through(self, tmp_path): dep = self._deployment(tmp_path, require_pinned=True, image_digest=None) with pytest.raises(ConfigurationError): dep.prepare() + + +# --------------------------------------------------------------------------- +# 13. Accounting directives and module loads on the self-managed path +# --------------------------------------------------------------------------- +class TestSlurmMultiAccountingDirectives: + """`qos`, `account` and `modules` are part of the declared parity set, so the + slurm_multi wrapper must emit them like job.sh.j2 does. Before this, they were + promoted into deployment_config and then silently dropped: a card declaring + them produced a wrapper with no `#SBATCH --account/--qos` and no `module load`, + which a billing-enforced site rejects outright. + """ + + def _deployment(self, tmp_path: Path, **slurm_overrides) -> SlurmDeployment: + script_rel = PR186_MODEL_ENTRY["scripts"] + script_abs = tmp_path / script_rel + script_abs.parent.mkdir(parents=True, exist_ok=True) + script_abs.write_text("#!/bin/bash\n") + + image_key = "rocm/pytorch-private:sglang_disagg_mori_20260502" + manifest = { + "built_images": {image_key: {"docker_image": image_key}}, + "built_models": {image_key: PR186_MODEL_ENTRY}, + "context": {}, + } + manifest_path = tmp_path / "build_manifest.json" + manifest_path.write_text(json.dumps(manifest)) + + slurm = dict( + PR186_MODEL_ENTRY["slurm"], + output_dir=str(tmp_path / "slurm_results"), + **slurm_overrides, + ) + cfg = DeploymentConfig( + target="slurm", + manifest_file=str(manifest_path), + additional_context={ + "deploy": "slurm", + "slurm": slurm, + "distributed": PR186_MODEL_ENTRY["distributed"], + }, + ) + return SlurmDeployment(cfg) + + def _script(self, dep: SlurmDeployment) -> str: + assert dep.prepare() is True + return Path(dep.script_path).read_text() + + def test_qos_and_account_emitted(self, tmp_path): + text = self._script(self._deployment(tmp_path, qos="high", account="amd-ml")) + assert "#SBATCH --qos=high" in text + assert "#SBATCH --account=amd-ml" in text + + def test_modules_loaded(self, tmp_path): + text = self._script( + self._deployment(tmp_path, modules=["rocm/6.2", "openmpi"]) + ) + assert "module load rocm/6.2" in text + assert "module load openmpi" in text + + def test_modules_load_before_model_script(self, tmp_path): + """`module load` has to run before the model script, or the site's + docker/rocm binaries are missing from PATH for its internal srun calls.""" + text = self._script(self._deployment(tmp_path, modules=["rocm/6.2"])) + lines = text.splitlines() + module_idx = next(i for i, l in enumerate(lines) if l == "module load rocm/6.2") + script_idx = next( + i for i, l in enumerate(lines) if "run_xPyD_models.slurm" in l and l.startswith("bash ") + ) + assert module_idx < script_idx + + def test_sbatch_directives_stay_in_header_block(self, tmp_path): + """sbatch stops parsing directives at the first non-comment line, so + --qos/--account must precede the `set -e` body.""" + text = self._script(self._deployment(tmp_path, qos="high", account="amd-ml")) + lines = text.splitlines() + first_body = next( + i for i, l in enumerate(lines) + if l.strip() and not l.startswith("#") and not l.startswith("#!") + ) + for directive in ("#SBATCH --qos=high", "#SBATCH --account=amd-ml"): + assert lines.index(directive) < first_body + + def test_absent_fields_emit_nothing(self, tmp_path): + """No declaration must not produce an empty directive that sbatch rejects.""" + text = self._script(self._deployment(tmp_path)) + assert "--qos" not in text + assert "--account" not in text + assert "module load" not in text + + +# --------------------------------------------------------------------------- +# 14. build-on-compute parity with the other build paths +# --------------------------------------------------------------------------- +class TestBuildOnComputeModelFieldParity: + """--build-on-compute writes its own manifest and returns, so it used to skip + _merge_model_config_into_manifest entirely. That lost the model card's declared + results CSV, recorded no "_explicit_slurm_keys" provenance, and let the + post-ConfigLoader preset defaults in --additional-context outrank the card. + """ + + CARD = { + "name": "model_a", + "dockerfile": "docker/model_a", + "scripts": "scripts/model_a/run.slurm", + "n_gpus": "8", + "tags": ["model_a"], + "multiple_results": "perf_model_a.csv", + "slurm": {"partition": "amd-rccl", "nodes": 4, "time": "01:00:00"}, + "distributed": {"launcher": "slurm_multi", "nnodes": 4}, + "env_vars": {}, + } + + def _run_build_on_compute(self, tmp_path: Path, card: dict, context_slurm: dict, + user_slurm_keys: set) -> dict: + from madengine.orchestration.build_orchestrator import BuildOrchestrator + + orch = BuildOrchestrator.__new__(BuildOrchestrator) + orch.args = MagicMock() + orch.console = MagicMock() + orch.rich_console = MagicMock() + orch.context = MagicMock() + orch.context.ctx = {} + # Post-ConfigLoader: carries preset defaults alongside the user's own keys. + orch.additional_context = {"slurm": context_slurm} + orch._original_user_slurm_keys = set(user_slurm_keys) + orch.credentials = {} + + df = tmp_path / f"{card['dockerfile']}.ubuntu.amd.Dockerfile" + df.parent.mkdir(parents=True, exist_ok=True) + df.write_text("FROM scratch\n") + + manifest_path = tmp_path / "build_manifest.json" + import os + orig_cwd = os.getcwd() + try: + os.chdir(tmp_path) + with patch("madengine.orchestration.build_orchestrator.DiscoverModels") as mock_dm: + mock_dm.return_value.run.return_value = [card] + with patch("subprocess.run") as mock_run: + mock_run.return_value = MagicMock(returncode=0) + orch._execute_build_on_compute( + registry="localhost:5000/myrepo", + manifest_output=str(manifest_path), + ) + finally: + os.chdir(orig_cwd) + return json.loads(manifest_path.read_text()) + + def test_declared_results_csv_is_persisted(self, tmp_path): + """Without multiple_results the run has no declared CSV to collect and + falls back to log scraping, which a slurm_multi model produces nothing for.""" + manifest = self._run_build_on_compute( + tmp_path, self.CARD, + context_slurm={"partition": "amd-rccl", "nodes": 1}, + user_slurm_keys={"partition"}, + ) + assert manifest["built_models"]["model_a"]["multiple_results"] == "perf_model_a.csv" + + def test_preset_default_does_not_outrank_model_card(self, tmp_path): + """nodes=1 is a ConfigLoader preset default, not a user choice, so the + card's nodes=4 must survive into the persisted slurm config.""" + manifest = self._run_build_on_compute( + tmp_path, self.CARD, + context_slurm={"partition": "amd-rccl", "nodes": 1, "time": "24:00:00"}, + user_slurm_keys={"partition"}, + ) + slurm = manifest["deployment_config"]["slurm"] + assert slurm["nodes"] == 4 + assert slurm["time"] == "01:00:00" + + def test_user_explicit_key_still_wins_over_model_card(self, tmp_path): + manifest = self._run_build_on_compute( + tmp_path, self.CARD, + context_slurm={"partition": "amd-rccl", "nodes": 2}, + user_slurm_keys={"partition", "nodes"}, + ) + assert manifest["deployment_config"]["slurm"]["nodes"] == 2 + + def test_explicit_slurm_keys_provenance_is_recorded(self, tmp_path): + manifest = self._run_build_on_compute( + tmp_path, self.CARD, + context_slurm={"partition": "amd-rccl", "nodes": 1, "time": "24:00:00"}, + user_slurm_keys={"partition"}, + ) + explicit = set(manifest["deployment_config"]["_explicit_slurm_keys"]) + # partition: user. nodes/time: declared by the card. + assert {"partition", "nodes", "time"} <= explicit + + def test_undeclared_preset_default_is_not_marked_explicit(self, tmp_path): + """The card says nothing about nodes, so the persisted nodes=1 is a preset + default and distributed.nnodes must still be free to size the allocation.""" + card = dict(self.CARD, slurm={"partition": "amd-rccl"}) + manifest = self._run_build_on_compute( + tmp_path, card, + context_slurm={"partition": "amd-rccl", "nodes": 1}, + user_slurm_keys={"partition"}, + ) + assert "nodes" not in set(manifest["deployment_config"]["_explicit_slurm_keys"]) + + def test_model_card_distributed_reaches_deployment_config(self, tmp_path): + """run --manifest-file infers the SLURM target from deployment_config.""" + manifest = self._run_build_on_compute( + tmp_path, self.CARD, + context_slurm={"partition": "amd-rccl", "nodes": 1}, + user_slurm_keys={"partition"}, + ) + assert manifest["deployment_config"]["distributed"]["launcher"] == "slurm_multi" + + +# --------------------------------------------------------------------------- +# 15. Runtime slurm override must not inherit stale provenance +# --------------------------------------------------------------------------- +class TestRuntimeSlurmOverrideProvenance: + """_load_and_merge_manifest replaces the manifest's slurm block with the + runtime one but used to leave "_explicit_slurm_keys" describing the *old* + block, so a build-time explicit nodes marked the runtime default nodes=1 as + deliberate and silently suppressed distributed.nnodes. + """ + + def _merge(self, tmp_path: Path, stored: dict, runtime: dict) -> dict: + from madengine.orchestration.run_orchestrator import RunOrchestrator + + orch = RunOrchestrator.__new__(RunOrchestrator) + orch.rich_console = MagicMock() + orch.additional_context = runtime + + manifest_path = tmp_path / "build_manifest.json" + manifest_path.write_text(json.dumps({ + "built_images": {}, "built_models": {}, "context": {}, + "deployment_config": stored, + })) + orch._load_and_merge_manifest(str(manifest_path)) + return json.loads(manifest_path.read_text())["deployment_config"] + + def test_runtime_slurm_recomputes_provenance(self, tmp_path): + cfg = self._merge( + tmp_path, + stored={ + "slurm": {"partition": "old", "nodes": 2}, + "_explicit_slurm_keys": ["nodes", "partition"], + "distributed": {"nnodes": 4}, + }, + runtime={"slurm": {"partition": "new"}}, + ) + assert cfg["_explicit_slurm_keys"] == ["partition"] + + def test_stale_provenance_no_longer_suppresses_nnodes(self, tmp_path): + """End of the chain: the recomputed provenance must let SlurmDeployment + size the allocation from distributed.nnodes.""" + cfg = self._merge( + tmp_path, + stored={ + "slurm": {"partition": "old", "nodes": 2}, + "_explicit_slurm_keys": ["nodes", "partition"], + }, + runtime={"slurm": {"partition": "new"}, "distributed": {"nnodes": 4}}, + ) + manifest_path = tmp_path / "m.json" + manifest_path.write_text(json.dumps({"built_images": {}, "built_models": {}, "context": {}})) + dep = SlurmDeployment(DeploymentConfig( + target="slurm", + manifest_file=str(manifest_path), + additional_context={ + "slurm": dict(cfg["slurm"], output_dir=str(tmp_path / "out")), + "distributed": cfg["distributed"], + "_explicit_slurm_keys": cfg["_explicit_slurm_keys"], + }, + )) + assert "nodes" not in dep._explicit_slurm_keys + assert dep.nodes == 4 + + def test_provenance_untouched_without_runtime_slurm(self, tmp_path): + """A plain `run --manifest-file` must keep the build-time provenance.""" + cfg = self._merge( + tmp_path, + stored={ + "slurm": {"partition": "old", "nodes": 2}, + "_explicit_slurm_keys": ["nodes", "partition"], + }, + runtime={"env_vars": {"FOO": "1"}}, + ) + assert cfg["_explicit_slurm_keys"] == ["nodes", "partition"] + assert cfg["slurm"]["nodes"] == 2 From 390d79575bdd0ac39978156c81dacda4e501e2cd Mon Sep 17 00:00:00 2001 From: Stephen Shao Date: Thu, 24 Sep 2026 00:39:14 +0000 Subject: [PATCH 6/7] fix(slurm): persist empty explicit_slurm_keys provenance for manifest runs When a persisted deployment config has a slurm block but no explicit user/model-card slurm keys, provenance was omitted entirely. On a later `run --manifest-file`, SlurmDeployment then fell back to inferring explicitness from the defaulted slurm dict and treated a preset nodes=1 as deliberate, silently ignoring distributed.nnodes. Co-Authored-By: Claude Sonnet 5 --- src/madengine/orchestration/build_orchestrator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/madengine/orchestration/build_orchestrator.py b/src/madengine/orchestration/build_orchestrator.py index b04af35c..645f8edd 100644 --- a/src/madengine/orchestration/build_orchestrator.py +++ b/src/madengine/orchestration/build_orchestrator.py @@ -1408,7 +1408,7 @@ def _merge_model_config_into_manifest(self, manifest_output: str, models: list): # card) as opposed to ConfigLoader preset defaults, so a later `run # --manifest-file` can tell a real distributed.nnodes-vs-slurm.nodes # conflict apart from a preset default silently outranking the model card. - if explicit_slurm_keys: + if "slurm" in saved_manifest["deployment_config"] or explicit_slurm_keys: saved_manifest["deployment_config"]["_explicit_slurm_keys"] = sorted( explicit_slurm_keys ) From 8df51019b7dce7f097c23ef09ee31050bb11e87c Mon Sep 17 00:00:00 2001 From: Stephen Shao Date: Thu, 24 Sep 2026 01:09:04 +0000 Subject: [PATCH 7/7] fix(slurm): reject non-integral/non-finite distributed.nnodes int(nnodes) silently truncated fractional values like 1.9, letting a malformed model card schedule the wrong node count with no warning. Reject booleans, non-integral floats, and non-finite values so they fall back to slurm.nodes with a logged note instead. Co-Authored-By: Claude Sonnet 5 --- src/madengine/deployment/common.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/madengine/deployment/common.py b/src/madengine/deployment/common.py index 63842b01..e707e7d1 100644 --- a/src/madengine/deployment/common.py +++ b/src/madengine/deployment/common.py @@ -168,8 +168,12 @@ def resolve_node_count( ``(nodes, note)`` where ``note`` is a human-readable string to log, or None. """ try: + if isinstance(nnodes, bool) or ( + isinstance(nnodes, float) and not nnodes.is_integer() + ): + raise ValueError requested = int(nnodes) if nnodes is not None else None - except (TypeError, ValueError): + except (TypeError, ValueError, OverflowError): return configured_nodes, ( f"Ignoring non-numeric distributed.nnodes={nnodes!r}; " f"using slurm.nodes={configured_nodes}."