diff --git a/.gitignore b/.gitignore index c824efdf..9dde930c 100644 --- a/.gitignore +++ b/.gitignore @@ -101,6 +101,9 @@ instance/ # Sphinx documentation docs/_build/ +# Local agent planning artifacts +docs/superpowers/ + # Jupyter Notebook .ipynb_checkpoints diff --git a/docs/deployment.md b/docs/deployment.md index c913b117..24d1b4d6 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -143,7 +143,7 @@ This creates: **Supported Launchers:** - `torchrun` - PyTorch DDP/FSDP - `deepspeed` - ZeRO optimization -- `megatron` - Megatron-LM training +- `megatron-lm` - Megatron-LM training - `torchtitan` - LLM pre-training - `primus` - Primus unified pretrain (Megatron / TorchTitan / MaxText YAML) - `vllm` - LLM inference diff --git a/docs/superpowers/plans/2026-08-27-pinned-image-digest.md b/docs/superpowers/plans/2026-08-27-pinned-image-digest.md deleted file mode 100644 index dad4b9d3..00000000 --- a/docs/superpowers/plans/2026-08-27-pinned-image-digest.md +++ /dev/null @@ -1,1623 +0,0 @@ -# Pinned Image Digest Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Record the real digest of every image madengine pushes, and let users opt into pinning run-time pulls to that digest so a moved registry tag fails loudly instead of silently running the wrong image. - -**Architecture:** A new leaf module `madengine/core/image_digest.py` holds four small pure functions (two parsers, one reference builder, one policy resolver that raises when enforcement is on but a digest is missing). `DockerBuilder.push_image()` records digests into a `self.pushed_digests` dict — mirroring the existing `self.built_images` / `self.built_models` pattern — and the two push call sites copy the digest into `build_info["image_digest"]`, which rides into `build_manifest.json` unchanged. At run time a `--require-pinned-image` flag (mirrored as the `require_pinned_image` additional-context key) flows through `RunOrchestrator.additional_context` into all three execution paths, each calling the same `resolve_pinned_image()` helper. - -**Tech Stack:** Python 3, Typer (CLI), pytest + `unittest.mock`, Jinja2 (SLURM/K8s templates), Docker CLI. - ---- - -## Background you need before starting - -Read the design spec first: `docs/superpowers/specs/2026-08-27-pinned-image-digest-design.md`. - -Facts about this codebase that the tasks below depend on: - -- `Console.sh(command)` (`src/madengine/core/console.py:138`) runs a shell command and **returns its stdout as a stripped string**, raising `RuntimeError` on non-zero exit unless `canFail=True`. This is how we capture `docker push` output. -- `build_info` dicts are created in `DockerBuilder.build_image()` (`src/madengine/execution/docker_builder.py:311-322`) and serialized wholesale into `build_manifest.json` under `built_images` by `export_build_manifest()`. **Any new key added to `build_info` appears in the manifest automatically** — no serializer changes needed. -- There are exactly two places that call `push_image` and set `build_info["registry_image"]`: `docker_builder.py:769` (single-arch) and `docker_builder.py:977` (per-GPU-arch). Both must record the digest. -- `ContainerRunner.run_models_from_manifest()` merges the manifest's `context` dict over its own `additional_context` (`src/madengine/execution/container_runner.py:2799-2800`). `RunOrchestrator._load_and_merge_manifest()` writes selected runtime keys into `manifest["context"]` (`src/madengine/orchestration/run_orchestrator.py:499-503`). Adding our key to that merge list is what makes the flag survive into the nested `madengine run` that the standard SLURM job script executes on each compute node. -- `_docker_image_ref_for_log_naming()` (`src/madengine/execution/container_runner_helpers.py:232`) already strips `@sha256:...`, so pinned references produce the same log filenames as tags. Task 9 locks that in with a test. -- Test style in this repo: `pytest` classes named `TestX` with plain `assert`, `MagicMock` for `Context`/`Console`, `tmp_path` for manifests. Follow the surrounding file's style in each test file you touch. - -### Deliberate deviations from the spec (do not "fix" these) - -1. **SLURM enforcement point.** The spec's table says the pinned reference goes into the generated `srun docker pull` line in `slurm.py:544`. Implementing only that would pull the pinned image and then `docker run` the *tag* — the exact race the feature exists to close. Instead we set `env_vars["DOCKER_IMAGE_NAME"]` to the pinned reference (Task 8); the pull line interpolates that same variable, so both pull and run are pinned by one change. -2. **Digest capture is not added to the build-on-compute-node path** (`build_orchestrator._execute_build_on_compute` at `build_orchestrator.py:748`, which pushes from a generated bash script inside an sbatch job and writes `built_images` entries at `:1224`/`:1239` with no `registry_image` and no digest). Manifests from that path have no `image_digest`, so `--require-pinned-image` runs against them fail fast with the Task 2 error. This is documented in Task 10, not implemented. - ---- - -## File Structure - -**Create:** - -| File | Responsibility | -|---|---| -| `src/madengine/core/image_digest.py` | Pure helpers: parse a digest out of `docker push` / `docker image inspect` output, build a `repo@sha256:...` reference, and apply the enforcement policy (return-as-is / pin / raise). No I/O, no Docker calls. | -| `tests/unit/test_image_digest.py` | Unit tests for all four functions in the module above. | - -**Modify:** - -| File | Change | -|---|---| -| `src/madengine/execution/docker_builder.py` | Capture the pushed digest in `push_image()`; copy it into `build_info["image_digest"]` at both push call sites. | -| `src/madengine/cli/commands/run.py` | Add the `--require-pinned-image` flag; pass it into both `create_args_namespace(...)` calls. | -| `src/madengine/orchestration/run_orchestrator.py` | Fold the flag into `additional_context`; persist it into `manifest["context"]`. | -| `src/madengine/execution/container_runner.py` | Resolve the pinned reference before the registry pull. | -| `src/madengine/deployment/k8s_template_context.py` | Resolve the pinned reference for the pod spec `image` field. | -| `src/madengine/deployment/slurm.py` | Resolve the pinned reference for `DOCKER_IMAGE_NAME` in the slurm_multi wrapper. | -| `docs/cli-reference.md`, `docs/configuration.md` | Document the flag and the context key. | - -**Test files touched:** `tests/unit/test_image_digest.py` (new), `tests/unit/test_docker_builder.py`, `tests/unit/test_orchestration.py`, `tests/unit/test_container_runner.py`, `tests/unit/test_k8s.py`, `tests/unit/test_slurm_multi.py`, `tests/unit/test_execution.py`. - ---- - -## Task 1: Digest parsing and pinned-reference construction - -**Files:** -- Create: `src/madengine/core/image_digest.py` -- Test: `tests/unit/test_image_digest.py` - -- [ ] **Step 1: Write the failing tests** - -Create `tests/unit/test_image_digest.py`: - -```python -"""Unit tests for madengine.core.image_digest. - -Covers digest extraction from `docker push` / `docker image inspect` output and -construction of pinned `repo@sha256:...` references. - -Copyright (c) Advanced Micro Devices, Inc. All rights reserved. -""" - -import pytest - -from madengine.core.image_digest import ( - build_pinned_reference, - parse_push_digest, - parse_repo_digest, -) - - -DIGEST = "sha256:" + "df36ef7e" * 8 # 64 hex chars -OTHER_DIGEST = "sha256:" + "cbb7e5ed" * 8 - - -class TestParsePushDigest: - """parse_push_digest extracts the digest line emitted by `docker push`.""" - - def test_typical_push_output(self): - output = ( - "The push refers to repository [docker.io/myorg/ci]\n" - "a1b2c3d4e5f6: Pushed\n" - "9f8e7d6c5b4a: Layer already exists\n" - f"mymodel: digest: {DIGEST} size: 4738\n" - ) - assert parse_push_digest(output) == DIGEST - - def test_no_space_after_colon(self): - assert parse_push_digest(f"mymodel: digest:{DIGEST} size: 12") == DIGEST - - def test_last_digest_wins_when_multiple(self): - output = ( - f"tag-a: digest: {OTHER_DIGEST} size: 10\n" - f"tag-b: digest: {DIGEST} size: 10\n" - ) - assert parse_push_digest(output) == DIGEST - - def test_uppercase_hex_is_not_matched(self): - assert parse_push_digest("mymodel: digest: sha256:ABC size: 1") is None - - def test_output_without_digest_line(self): - assert parse_push_digest("The push refers to repository [x]\nlayer: Pushed\n") is None - - def test_empty_output(self): - assert parse_push_digest("") is None - - def test_none_output(self): - assert parse_push_digest(None) is None - - -class TestParseRepoDigest: - """parse_repo_digest extracts the digest from a `repo@sha256:...` reference.""" - - def test_repo_digest_reference(self): - assert parse_repo_digest(f"myorg/ci@{DIGEST}") == DIGEST - - def test_registry_with_port(self): - assert parse_repo_digest(f"localhost:5000/myorg/ci@{DIGEST}") == DIGEST - - def test_surrounding_whitespace_and_quotes(self): - assert parse_repo_digest(f" 'myorg/ci@{DIGEST}' \n") == DIGEST - - def test_empty_repodigests_placeholder(self): - # `docker image inspect` prints this when RepoDigests is empty. - assert parse_repo_digest("") is None - - def test_none_output(self): - assert parse_repo_digest(None) is None - - -class TestBuildPinnedReference: - """build_pinned_reference produces repo@sha256:... with any tag stripped.""" - - def test_strips_tag(self): - assert build_pinned_reference(f"myorg/ci:mymodel", DIGEST) == f"myorg/ci@{DIGEST}" - - def test_no_tag(self): - assert build_pinned_reference("myorg/ci", DIGEST) == f"myorg/ci@{DIGEST}" - - def test_registry_port_is_not_mistaken_for_a_tag(self): - out = build_pinned_reference("localhost:5000/myorg/ci:latest", DIGEST) - assert out == f"localhost:5000/myorg/ci@{DIGEST}" - - def test_registry_port_without_tag(self): - out = build_pinned_reference("localhost:5000/myorg/ci", DIGEST) - assert out == f"localhost:5000/myorg/ci@{DIGEST}" - - def test_bare_name_with_tag(self): - assert build_pinned_reference("ci-dummy:latest", DIGEST) == f"ci-dummy@{DIGEST}" - - def test_existing_digest_is_replaced(self): - out = build_pinned_reference(f"myorg/ci@{OTHER_DIGEST}", DIGEST) - assert out == f"myorg/ci@{DIGEST}" -``` - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/unit/test_image_digest.py -v` -Expected: FAIL — `ModuleNotFoundError: No module named 'madengine.core.image_digest'` - -- [ ] **Step 3: Write the implementation** - -Create `src/madengine/core/image_digest.py`: - -```python -#!/usr/bin/env python3 -""" -Registry image digest helpers for madengine. - -Build pushes record the digest of the image they push; runs can optionally be -pinned to that digest so a registry tag that moved between build and run fails -loudly instead of silently resolving to a different image. - -Copyright (c) Advanced Micro Devices, Inc. All rights reserved. -""" - -import re -import typing - - -# `docker push` prints e.g. "mytag: digest: sha256:<64 hex> size: 4738". -_PUSH_DIGEST_RE = re.compile(r"digest:\s*(sha256:[0-9a-f]{64})") - -# `docker image inspect --format '{{index .RepoDigests 0}}'` prints "repo@sha256:<64 hex>". -_REPO_DIGEST_RE = re.compile(r"@(sha256:[0-9a-f]{64})") - - -def parse_push_digest(push_output: typing.Optional[str]) -> typing.Optional[str]: - """Extract the pushed image digest from `docker push` output. - - Args: - push_output: Combined stdout/stderr of the push command. - - Returns: - The digest (``sha256:...``), or None if no digest line was present. - When several digest lines appear the last one is returned, which is the - digest of the reference the push command was invoked with. - """ - if not push_output: - return None - matches = _PUSH_DIGEST_RE.findall(push_output) - return matches[-1] if matches else None - - -def parse_repo_digest(inspect_output: typing.Optional[str]) -> typing.Optional[str]: - """Extract the digest from a ``repo@sha256:...`` reference. - - Args: - inspect_output: Output of - ``docker image inspect --format '{{index .RepoDigests 0}}' ``. - - Returns: - The digest (``sha256:...``), or None if the output held no digest - (e.g. ```` when the image has no RepoDigests entry). - """ - if not inspect_output: - return None - match = _REPO_DIGEST_RE.search(inspect_output) - return match.group(1) if match else None - - -def build_pinned_reference(registry_image: str, digest: str) -> str: - """Build a digest-pinned image reference. - - Any existing tag or digest on ``registry_image`` is dropped. A port in the - registry host (``localhost:5000/org/img``) is not mistaken for a tag because - only the final path segment is inspected for ``:``. - - Args: - registry_image: Image reference, with or without tag/digest. - digest: Digest to pin to (``sha256:...``). - - Returns: - A reference of the form ``repo@sha256:...``. - """ - repo = registry_image.split("@", 1)[0] - last_slash = repo.rfind("/") - tail = repo[last_slash + 1 :] - if ":" in tail: - repo = repo[: last_slash + 1] + tail.split(":", 1)[0] - return f"{repo}@{digest}" -``` - -- [ ] **Step 4: Run tests to verify they pass** - -Run: `pytest tests/unit/test_image_digest.py -v` -Expected: PASS (18 tests) - -- [ ] **Step 5: Commit** - -```bash -git add src/madengine/core/image_digest.py tests/unit/test_image_digest.py -git commit -m "feat(image-digest): add digest parsing and pinned reference helpers" -``` - ---- - -## Task 2: Enforcement policy resolver - -**Files:** -- Modify: `src/madengine/core/image_digest.py` -- Test: `tests/unit/test_image_digest.py` - -`resolve_pinned_image()` is the single place the "pin, pass through, or fail" decision is made. All three execution paths (local Docker, K8s, SLURM) call it so they cannot drift. - -- [ ] **Step 1: Write the failing tests** - -Append to `tests/unit/test_image_digest.py`: - -```python -class TestResolvePinnedImage: - """resolve_pinned_image applies the --require-pinned-image policy.""" - - def test_disabled_returns_image_unchanged_without_digest(self): - assert resolve_pinned_image("myorg/ci:mymodel", None, False) == "myorg/ci:mymodel" - - def test_disabled_returns_image_unchanged_even_with_digest(self): - # Default behaviour must not change for existing users: the digest rides - # along in the manifest but is never used unless enforcement is on. - assert resolve_pinned_image("myorg/ci:mymodel", DIGEST, False) == "myorg/ci:mymodel" - - def test_enabled_with_digest_returns_pinned_reference(self): - out = resolve_pinned_image("myorg/ci:mymodel", DIGEST, True) - assert out == f"myorg/ci@{DIGEST}" - - def test_enabled_without_digest_raises(self): - with pytest.raises(ConfigurationError): - resolve_pinned_image("myorg/ci:mymodel", None, True, model_name="my_model") - - def test_enabled_with_empty_digest_raises(self): - with pytest.raises(ConfigurationError): - resolve_pinned_image("myorg/ci:mymodel", "", True, model_name="my_model") - - def test_error_message_names_model_and_image(self): - with pytest.raises(ConfigurationError) as excinfo: - resolve_pinned_image("myorg/ci:mymodel", None, True, model_name="my_model") - message = str(excinfo.value) - assert "my_model" in message - assert "myorg/ci:mymodel" in message - - def test_error_carries_actionable_suggestions(self): - with pytest.raises(ConfigurationError) as excinfo: - resolve_pinned_image("myorg/ci:mymodel", None, True, model_name="my_model") - assert excinfo.value.suggestions -``` - -And extend the imports at the top of the file: - -```python -from madengine.core.errors import ConfigurationError -from madengine.core.image_digest import ( - build_pinned_reference, - parse_push_digest, - parse_repo_digest, - resolve_pinned_image, -) -``` - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/unit/test_image_digest.py -k ResolvePinned -v` -Expected: FAIL — `ImportError: cannot import name 'resolve_pinned_image'` - -- [ ] **Step 3: Write the implementation** - -Add the import near the top of `src/madengine/core/image_digest.py`, below `import typing`: - -```python -from madengine.core.errors import ConfigurationError, create_error_context -``` - -Append to the same file: - -```python -def resolve_pinned_image( - registry_image: str, - image_digest: typing.Optional[str], - require_pinned: bool, - model_name: str = "", -) -> str: - """Resolve the image reference to use for a run. - - Args: - registry_image: Tagged registry reference from the build manifest. - image_digest: Digest recorded at build time, if any. - require_pinned: True when --require-pinned-image / require_pinned_image is set. - model_name: Model name, used only to make the error message actionable. - - Returns: - ``registry_image`` unchanged when enforcement is off, otherwise a - digest-pinned reference. - - Raises: - ConfigurationError: When enforcement is on but the manifest recorded no - digest. Falling back to the tag is deliberately not done: silent - degradation would defeat the guarantee the caller asked for. - """ - if not require_pinned: - return registry_image - - if not image_digest: - raise ConfigurationError( - f"--require-pinned-image is set but the build manifest records no " - f"image digest for model '{model_name}' (image: {registry_image}). " - f"Refusing to pull by tag.", - context=create_error_context( - operation="resolve_pinned_image", - component="image_digest", - additional_info={"model": model_name, "image": registry_image}, - ), - suggestions=[ - "Rebuild with this version of madengine so the push digest is recorded", - "Drop --require-pinned-image / require_pinned_image to pull by tag", - ], - ) - - return build_pinned_reference(registry_image, image_digest) -``` - -- [ ] **Step 4: Run tests to verify they pass** - -Run: `pytest tests/unit/test_image_digest.py -v` -Expected: PASS (25 tests) - -- [ ] **Step 5: Commit** - -```bash -git add src/madengine/core/image_digest.py tests/unit/test_image_digest.py -git commit -m "feat(image-digest): add resolve_pinned_image enforcement policy" -``` - ---- - -## Task 3: Capture the pushed digest in `DockerBuilder.push_image()` - -**Files:** -- Modify: `src/madengine/execution/docker_builder.py:51`, `:394-403` -- Test: `tests/unit/test_docker_builder.py` - -Digest capture is **always on** and never fails a build. If neither the push output nor `docker image inspect` yields a digest, we log a dim note and move on — the push already succeeded. - -- [ ] **Step 1: Write the failing tests** - -Append to `tests/unit/test_docker_builder.py`: - -```python -DIGEST = "sha256:" + "df36ef7e" * 8 - - -class TestPushImageRecordsDigest: - """push_image records the pushed image digest in builder.pushed_digests.""" - - def _builder(self, sh_side_effect): - ctx = MagicMock() - ctx.ctx = {} - console = MagicMock() - console.sh = MagicMock(side_effect=sh_side_effect) - builder = DockerBuilder(ctx, console) - builder.rich_console = MagicMock() - return builder - - def test_digest_parsed_from_push_output(self): - def sh(command, *args, **kwargs): - if "docker push" in command: - return f"mymodel: digest: {DIGEST} size: 4738" - return "" - - builder = self._builder(sh) - result = builder.push_image("ci-dummy", "localhost:5000", None, "localhost:5000/ci-dummy") - - assert result == "localhost:5000/ci-dummy" - assert builder.pushed_digests["localhost:5000/ci-dummy"] == DIGEST - - def test_falls_back_to_image_inspect_when_push_output_has_no_digest(self): - def sh(command, *args, **kwargs): - if "docker push" in command: - return "The push refers to repository [localhost:5000/ci-dummy]\nlayer: Pushed" - if "docker image inspect" in command: - return f"localhost:5000/ci-dummy@{DIGEST}" - return "" - - builder = self._builder(sh) - builder.push_image("ci-dummy", "localhost:5000", None, "localhost:5000/ci-dummy") - - assert builder.pushed_digests["localhost:5000/ci-dummy"] == DIGEST - inspect_calls = [ - c for c in builder.console.sh.call_args_list if "docker image inspect" in c.args[0] - ] - assert len(inspect_calls) == 1 - assert "RepoDigests" in inspect_calls[0].args[0] - - def test_no_digest_anywhere_leaves_entry_absent_and_push_still_succeeds(self): - def sh(command, *args, **kwargs): - if "docker push" in command: - return "layer: Pushed" - if "docker image inspect" in command: - return "" - return "" - - builder = self._builder(sh) - result = builder.push_image("ci-dummy", "localhost:5000", None, "localhost:5000/ci-dummy") - - assert result == "localhost:5000/ci-dummy" - assert "localhost:5000/ci-dummy" not in builder.pushed_digests - # The gap is noted at dim level, not as a user-facing warning. - printed = " ".join(str(c) for c in builder.rich_console.print.call_args_list) - assert "[dim]" in printed - assert "no pushed digest" in printed.lower() - - def test_inspect_failure_is_swallowed(self): - def sh(command, *args, **kwargs): - if "docker push" in command: - return "layer: Pushed" - if "docker image inspect" in command: - raise RuntimeError("no such image") - return "" - - builder = self._builder(sh) - result = builder.push_image("ci-dummy", "localhost:5000", None, "localhost:5000/ci-dummy") - - assert result == "localhost:5000/ci-dummy" - assert builder.pushed_digests == {} - - def test_no_registry_records_nothing(self): - builder = self._builder(lambda command, *a, **k: "") - result = builder.push_image("ci-dummy") - - assert result == "ci-dummy" - assert builder.pushed_digests == {} -``` - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/unit/test_docker_builder.py -k PushImageRecordsDigest -v` -Expected: FAIL — `AttributeError: 'DockerBuilder' object has no attribute 'pushed_digests'` - -- [ ] **Step 3: Write the implementation** - -Add the import to `src/madengine/execution/docker_builder.py`, after the `from madengine.core.context import Context` line: - -```python -from madengine.core.image_digest import parse_push_digest, parse_repo_digest -``` - -In `DockerBuilder.__init__`, immediately after `self.built_images = {} # Track built images` (line 51), add: - -```python - self.pushed_digests = {} # registry_image -> digest recorded at push time -``` - -Replace the push block in `push_image()` (lines 394-399) — from `# Push the image` through `self.console.sh(push_command)` — with: - -```python - # Push the image - push_command = f"docker push {shlex.quote(registry_image)}" - self.rich_console.print(f"\n[bold blue]🚀 Starting docker push to registry...[/bold blue]") - print(f"📤 Registry: {registry}") - print(f"🏷️ Image: {registry_image}") - push_output = self.console.sh(push_command) - - self._record_pushed_digest(registry_image, push_output) -``` - -Add the new method immediately after `push_image()` (i.e. after its `raise` at line 408, before `def export_build_manifest`): - -```python - def _record_pushed_digest(self, registry_image: str, push_output: str) -> None: - """Record the digest of the image just pushed, for digest-pinned runs. - - Best-effort by design: the push has already succeeded by the time this - runs, so a missing digest is a manifest-completeness gap (noted at dim - level), never a build failure. Runs only consult the recorded digest - when --require-pinned-image is set. - - Args: - registry_image: The reference that was pushed. - push_output: stdout/stderr captured from the push command. - """ - digest = parse_push_digest(push_output) - - if not digest: - # Some registries/mirrors do not print the digest line; ask the - # daemon for the RepoDigests entry it recorded for this push. - try: - inspect_output = self.console.sh( - "docker image inspect --format '{{index .RepoDigests 0}}' " - + shlex.quote(registry_image) - ) - digest = parse_repo_digest(inspect_output) - except Exception: - digest = None - - if not digest: - self.rich_console.print( - f"[dim]No pushed digest recorded for {registry_image}; " - f"--require-pinned-image runs will reject this manifest entry[/dim]" - ) - return - - self.pushed_digests[registry_image] = digest - self.rich_console.print(f"[dim]Pushed digest: {registry_image} -> {digest}[/dim]") -``` - -- [ ] **Step 4: Run tests to verify they pass** - -Run: `pytest tests/unit/test_docker_builder.py -v` -Expected: PASS (all, including the 2 pre-existing naming tests) - -- [ ] **Step 5: Verify no existing push tests regressed** - -Run: `pytest tests/integration/test_docker_integration.py -k push_image -v` -Expected: PASS (6 tests). These mock `Console.sh` with `return_value="Success"`, so the digest parse returns None, the inspect fallback also returns `"Success"` (no digest), and `push_image` still returns the tag unchanged. - -- [ ] **Step 6: Commit** - -```bash -git add src/madengine/execution/docker_builder.py tests/unit/test_docker_builder.py -git commit -m "feat(build): capture pushed image digest during docker push" -``` - ---- - -## Task 4: Write `image_digest` into `build_info` at both push sites - -**Files:** -- Modify: `src/madengine/execution/docker_builder.py:762-772`, `:971-980` -- Test: `tests/unit/test_docker_builder.py` - -`build_info` is serialized wholesale into `build_manifest.json`, so setting the key here is all that is needed to get it into the manifest. - -- [ ] **Step 1: Write the failing tests** - -Append to `tests/unit/test_docker_builder.py`: - -```python -class TestBuildInfoCarriesImageDigest: - """Both push call sites copy the recorded digest into build_info.""" - - def _builder(self): - ctx = MagicMock() - ctx.ctx = {} - builder = DockerBuilder(ctx, MagicMock()) - builder.rich_console = MagicMock() - return builder - - def _run_single_arch(self, builder): - """Drive _build_model_single_arch with everything below push_image stubbed.""" - return builder._build_model_single_arch( - model_info={"name": "dummy", "dockerfile": "docker/dummy"}, - credentials={}, - clean_cache=False, - registry="localhost:5000", - phase_suffix="", - batch_build_metadata=None, - ) - - def test_single_arch_push_sets_image_digest(self): - builder = self._builder() - - def fake_push(docker_image, registry, credentials, explicit_registry_image): - builder.pushed_digests[explicit_registry_image] = DIGEST - return explicit_registry_image - - with patch.object( - builder, "_get_dockerfiles_for_model", return_value=["docker/dummy.ubuntu"] - ), patch.object( - builder, "build_image", return_value={"docker_image": "ci-dummy", "model": "dummy"} - ), patch.object( - builder, "_get_effective_gpu_architecture", return_value="" - ), patch.object( - builder, "_create_registry_image_name", return_value="localhost:5000/ci-dummy" - ), patch.object( - builder, "push_image", side_effect=fake_push - ): - results = self._run_single_arch(builder) - - assert results[0]["registry_image"] == "localhost:5000/ci-dummy" - assert results[0]["image_digest"] == DIGEST - # A push failure must still be recorded the way it is today. - assert "push_error" not in results[0] - - def test_single_arch_push_without_digest_omits_key(self): - builder = self._builder() - - with patch.object( - builder, "_get_dockerfiles_for_model", return_value=["docker/dummy.ubuntu"] - ), patch.object( - builder, "build_image", return_value={"docker_image": "ci-dummy", "model": "dummy"} - ), patch.object( - builder, "_get_effective_gpu_architecture", return_value="" - ), patch.object( - builder, "_create_registry_image_name", return_value="localhost:5000/ci-dummy" - ), patch.object( - builder, "push_image", return_value="localhost:5000/ci-dummy" - ): - results = self._run_single_arch(builder) - - assert results[0]["registry_image"] == "localhost:5000/ci-dummy" - assert "image_digest" not in results[0] -``` - -Extend the import at the top of `tests/unit/test_docker_builder.py`: - -```python -from unittest.mock import MagicMock, patch -``` - -> `_build_model_single_arch(self, model_info, credentials, clean_cache, registry, phase_suffix, batch_build_metadata)` is at `docker_builder.py:729`; the per-arch sibling is `_build_model_for_arch`, which takes an extra `arch` argument. Both are reached from `build_all_models` (`docker_builder.py:581`, `:617`). - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/unit/test_docker_builder.py -k BuildInfoCarriesImageDigest -v` -Expected: FAIL — `KeyError: 'image_digest'` - -- [ ] **Step 3: Write the implementation** - -At the single-arch site (`docker_builder.py:769-770`), replace: - -```python - self.push_image(build_info["docker_image"], registry, credentials, registry_image) - build_info["registry_image"] = registry_image -``` - -with: - -```python - self.push_image(build_info["docker_image"], registry, credentials, registry_image) - build_info["registry_image"] = registry_image - # Recorded at push time; consumed only by --require-pinned-image runs. - pushed_digest = self.pushed_digests.get(registry_image) - if pushed_digest: - build_info["image_digest"] = pushed_digest -``` - -At the per-arch site (`docker_builder.py:977-978`), replace: - -```python - self.push_image(arch_image_name, registry, credentials, registry_image) - build_info["registry_image"] = registry_image -``` - -with: - -```python - self.push_image(arch_image_name, registry, credentials, registry_image) - build_info["registry_image"] = registry_image - # Recorded at push time; consumed only by --require-pinned-image runs. - pushed_digest = self.pushed_digests.get(registry_image) - if pushed_digest: - build_info["image_digest"] = pushed_digest -``` - -- [ ] **Step 4: Run tests to verify they pass** - -Run: `pytest tests/unit/test_docker_builder.py -v` -Expected: PASS - -- [ ] **Step 5: Confirm the manifest structure is unchanged for existing consumers** - -Run: `pytest tests/unit/test_orchestration.py tests/unit/test_slurm_multi.py -v` -Expected: PASS — `image_digest` is purely additive. - -- [ ] **Step 6: Commit** - -```bash -git add src/madengine/execution/docker_builder.py tests/unit/test_docker_builder.py -git commit -m "feat(build): record image_digest in build manifest entries" -``` - ---- - -## Task 5: `--require-pinned-image` flag and context propagation - -**Files:** -- Modify: `src/madengine/cli/commands/run.py:162-168` (flag), `:230-247` and `:318-341` (both `create_args_namespace` calls) -- Modify: `src/madengine/orchestration/run_orchestrator.py:85` (context merge), `:502` (manifest merge keys) -- Test: `tests/unit/test_orchestration.py` - -Two entry points must both work: the CLI flag, and the `require_pinned_image` key inside `--additional-context` (which is how CI pipelines drive madengine). Persisting the key into `manifest["context"]` is what carries the setting to the nested `madengine run` that the SLURM job script executes on each compute node. - -- [ ] **Step 1: Write the failing tests** - -Append to `tests/unit/test_orchestration.py`: - -```python -class TestRequirePinnedImageContext: - """--require-pinned-image and require_pinned_image both reach additional_context.""" - - @patch("madengine.orchestration.run_orchestrator.Context") - def test_cli_flag_sets_context_key(self, mock_context): - args = create_args_namespace( - additional_context=None, - require_pinned_image=True, - live_output=False, - ) - orch = RunOrchestrator(args) - assert orch.additional_context["require_pinned_image"] is True - - @patch("madengine.orchestration.run_orchestrator.Context") - def test_flag_absent_leaves_key_unset(self, mock_context): - args = create_args_namespace( - additional_context=None, - require_pinned_image=False, - live_output=False, - ) - orch = RunOrchestrator(args) - assert "require_pinned_image" not in orch.additional_context - - @patch("madengine.orchestration.run_orchestrator.Context") - def test_additional_context_key_alone_is_honoured(self, mock_context): - args = create_args_namespace( - additional_context="{'require_pinned_image': True}", - live_output=False, - ) - orch = RunOrchestrator(args) - assert orch.additional_context["require_pinned_image"] is True - - @patch("madengine.orchestration.run_orchestrator.Context") - def test_key_is_persisted_into_manifest_context(self, mock_context, tmp_path): - manifest_path = tmp_path / "build_manifest.json" - manifest_path.write_text(json.dumps({ - "built_images": {"img1": {"registry_image": "myorg/ci:m"}}, - "built_models": {"img1": {"name": "m"}}, - "context": {}, - "deployment_config": {}, - })) - - args = create_args_namespace( - additional_context=None, - require_pinned_image=True, - live_output=False, - ) - orch = RunOrchestrator(args) - orch._load_and_merge_manifest(str(manifest_path)) - - written = json.loads(manifest_path.read_text()) - assert written["context"]["require_pinned_image"] is True -``` - -Make sure `json`, `patch`, `create_args_namespace` and `RunOrchestrator` are imported at the top of `tests/unit/test_orchestration.py` — the existing `TestRunOrchestratorInit` and `TestCreateManifestFromLocalImage` classes already import most of these; add only what is missing: - -```python -from madengine.cli.utils import create_args_namespace -``` - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/unit/test_orchestration.py -k RequirePinnedImage -v` -Expected: FAIL — `KeyError: 'require_pinned_image'` - -- [ ] **Step 3: Write the implementation — orchestrator** - -In `src/madengine/orchestration/run_orchestrator.py`, immediately after `self.additional_context = merged_context` (line 85), add: - -```python - # The CLI flag and the require_pinned_image context key are equivalent; - # the key lets CI pipelines that drive madengine through - # --additional-context opt in the same way as for k8s/slurm/tools. - if getattr(args, "require_pinned_image", False): - self.additional_context["require_pinned_image"] = True -``` - -In `_load_and_merge_manifest`, extend the merge key list (line 502) from: - -```python - merge_keys = ["tools", "pre_scripts", "post_scripts", "encapsulate_script"] -``` - -to: - -```python - merge_keys = [ - "tools", - "pre_scripts", - "post_scripts", - "encapsulate_script", - # Persisted so nested runs on SLURM compute nodes (which re-enter - # `madengine run --manifest-file`) inherit the enforcement setting. - "require_pinned_image", - ] -``` - -- [ ] **Step 4: Write the implementation — CLI** - -In `src/madengine/cli/commands/run.py`, add a new option immediately after the `skip_model_run` option block (which ends at line 108 with `] = False,`): - -```python - require_pinned_image: Annotated[ - bool, - typer.Option( - "--require-pinned-image", - help=( - "Pull registry images by the digest recorded in the build manifest " - "instead of by tag. Fails immediately if the manifest has no digest " - "for an image. Equivalent to the 'require_pinned_image' " - "additional-context key." - ), - ), - ] = False, -``` - -Then add `require_pinned_image=require_pinned_image,` to **both** `create_args_namespace(...)` calls — in the manifest-exists branch (next to `skip_model_run=skip_model_run,` around line 245) and in the full-workflow branch (next to `skip_model_run=skip_model_run,` around line 339). - -- [ ] **Step 5: Run tests to verify they pass** - -Run: `pytest tests/unit/test_orchestration.py -v` -Expected: PASS - -- [ ] **Step 6: Verify the flag is wired into the CLI** - -Run: `madengine run --help | grep -A 3 require-pinned-image` -Expected: the option and its help text are listed. - -- [ ] **Step 7: Commit** - -```bash -git add src/madengine/cli/commands/run.py src/madengine/orchestration/run_orchestrator.py tests/unit/test_orchestration.py -git commit -m "feat(run): add --require-pinned-image flag and context propagation" -``` - ---- - -## Task 6: Enforce pinned pulls in local Docker execution - -**Files:** -- Modify: `src/madengine/execution/container_runner.py:2851-2859` -- Test: `tests/unit/test_container_runner.py` - -- [ ] **Step 1: Write the failing tests** - -Append to `tests/unit/test_container_runner.py`: - -```python -DIGEST = "sha256:" + "df36ef7e" * 8 - - -class TestRequirePinnedImageLocalRun: - """run_models_from_manifest honours require_pinned_image for registry pulls.""" - - def _manifest(self, tmpdir, build_info): - manifest_path = os.path.join(tmpdir, "build_manifest.json") - with open(manifest_path, "w") as f: - json.dump( - { - "built_images": {"img1": build_info}, - "built_models": { - "img1": {"name": "m", "tags": "t", "n_gpus": "1", "args": ""} - }, - }, - f, - ) - return manifest_path - - def _runner(self): - ctx = MagicMock() - ctx.ctx = {"docker_env_vars": {"MAD_SYSTEM_GPU_ARCHITECTURE": "gfx90a"}} - ctx.ensure_runtime_context = MagicMock() - console = MagicMock() - console.sh.return_value = "testhost" - runner = ContainerRunner(context=ctx, console=console) - runner.set_credentials({}) - return runner - - @patch("madengine.execution.container_runner.update_perf_csv") - def test_default_pulls_by_tag_even_when_digest_present(self, _mock_csv): - with tempfile.TemporaryDirectory() as tmpdir: - manifest_path = self._manifest( - tmpdir, - {"registry_image": "myorg/ci:m", "image_digest": DIGEST}, - ) - runner = self._runner() - runner.perf_csv_path = os.path.join(tmpdir, "perf.csv") - - with patch.object(runner, "pull_image") as mock_pull, patch.object( - runner, "run_container", return_value={"status": "SUCCESS"} - ): - runner.run_models_from_manifest(manifest_file=manifest_path, timeout=60) - - mock_pull.assert_called_once_with("myorg/ci:m") - - @patch("madengine.execution.container_runner.update_perf_csv") - def test_enabled_pulls_pinned_reference(self, _mock_csv): - with tempfile.TemporaryDirectory() as tmpdir: - manifest_path = self._manifest( - tmpdir, - {"registry_image": "myorg/ci:m", "image_digest": DIGEST}, - ) - runner = self._runner() - runner.perf_csv_path = os.path.join(tmpdir, "perf.csv") - runner.additional_context = {"require_pinned_image": True} - - with patch.object(runner, "pull_image") as mock_pull, patch.object( - runner, "run_container", return_value={"status": "SUCCESS"} - ) as mock_run: - runner.run_models_from_manifest(manifest_file=manifest_path, timeout=60) - - mock_pull.assert_called_once_with(f"myorg/ci@{DIGEST}") - # The container must run the same pinned reference that was pulled. - assert mock_run.call_args[1]["docker_image"] == f"myorg/ci@{DIGEST}" - - @patch("madengine.execution.container_runner.update_perf_csv") - def test_enabled_without_digest_fails_before_pulling(self, _mock_csv): - with tempfile.TemporaryDirectory() as tmpdir: - manifest_path = self._manifest(tmpdir, {"registry_image": "myorg/ci:m"}) - runner = self._runner() - runner.perf_csv_path = os.path.join(tmpdir, "perf.csv") - runner.additional_context = {"require_pinned_image": True} - - with patch.object(runner, "pull_image") as mock_pull, patch.object( - runner, "run_container" - ) as mock_run: - result = runner.run_models_from_manifest( - manifest_file=manifest_path, timeout=60 - ) - - mock_pull.assert_not_called() - mock_run.assert_not_called() - assert len(result["failed_runs"]) == 1 - assert "require-pinned-image" in result["failed_runs"][0]["error"] - - @patch("madengine.execution.container_runner.update_perf_csv") - def test_manifest_context_key_enables_enforcement(self, _mock_csv): - """A nested run on a SLURM compute node inherits the setting via manifest context.""" - with tempfile.TemporaryDirectory() as tmpdir: - manifest_path = os.path.join(tmpdir, "build_manifest.json") - with open(manifest_path, "w") as f: - json.dump( - { - "built_images": { - "img1": {"registry_image": "myorg/ci:m", "image_digest": DIGEST} - }, - "built_models": { - "img1": {"name": "m", "tags": "t", "n_gpus": "1", "args": ""} - }, - "context": {"require_pinned_image": True}, - }, - f, - ) - runner = self._runner() - runner.perf_csv_path = os.path.join(tmpdir, "perf.csv") - - with patch.object(runner, "pull_image") as mock_pull, patch.object( - runner, "run_container", return_value={"status": "SUCCESS"} - ): - runner.run_models_from_manifest(manifest_file=manifest_path, timeout=60) - - mock_pull.assert_called_once_with(f"myorg/ci@{DIGEST}") -``` - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/unit/test_container_runner.py -k RequirePinnedImageLocalRun -v` -Expected: FAIL — `test_enabled_pulls_pinned_reference` asserts the pinned ref but the tag is pulled. - -- [ ] **Step 3: Write the implementation** - -Add the import to `src/madengine/execution/container_runner.py`, after `from madengine.core.docker import Docker`: - -```python -from madengine.core.image_digest import resolve_pinned_image -``` - -Replace the registry branch at `container_runner.py:2851-2859`: - -```python - elif build_info.get("registry_image"): - # Registry image: Pull from registry - try: - self.pull_image(build_info["registry_image"]) - # Update docker_image to use registry image - run_image = build_info["registry_image"] - except Exception as pull_error: - self.rich_console.print(f"[yellow]Warning: Could not pull from registry, using local image[/yellow]") - run_image = image_name -``` - -with: - -```python - elif build_info.get("registry_image"): - # Registry image: Pull from registry. Under - # require_pinned_image this resolves to repo@sha256:... and - # raises (outside the pull try/except, so there is no tag - # fallback) when the manifest recorded no digest. - pull_target = resolve_pinned_image( - build_info["registry_image"], - build_info.get("image_digest"), - bool((self.additional_context or {}).get("require_pinned_image")), - model_name=model_info.get("name", ""), - ) - try: - self.pull_image(pull_target) - # Update docker_image to use registry image - run_image = pull_target - except Exception as pull_error: - self.rich_console.print(f"[yellow]Warning: Could not pull from registry, using local image[/yellow]") - run_image = image_name -``` - -- [ ] **Step 4: Run tests to verify they pass** - -Run: `pytest tests/unit/test_container_runner.py -v` -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add src/madengine/execution/container_runner.py tests/unit/test_container_runner.py -git commit -m "feat(run): pin local docker pulls to manifest digest when required" -``` - ---- - -## Task 7: Enforce pinned images in the Kubernetes pod spec - -**Files:** -- Modify: `src/madengine/deployment/k8s_template_context.py:518` -- Test: `tests/unit/test_k8s.py` - -- [ ] **Step 1: Write the failing tests** - -Append to `tests/unit/test_k8s.py`: - -```python -class TestK8sRequirePinnedImage: - """The generated pod spec image field honours require_pinned_image.""" - - DIGEST = "sha256:" + "df36ef7e" * 8 - - def _template_context(self, tmp_path, monkeypatch, require_pinned, image_digest): - """Build a real template context, the way prepare() does. - - _prepare_template_context reads the manifest and the model's scripts - directory from the current working directory, so the test runs inside - tmp_path with a minimal model tree. - """ - monkeypatch.chdir(tmp_path) - (tmp_path / "scripts" / "dummy").mkdir(parents=True) - (tmp_path / "scripts" / "dummy" / "run.sh").write_text("#!/bin/bash\necho hi\n") - - image_info = {"registry_image": "myorg/ci:m"} - if image_digest: - image_info["image_digest"] = image_digest - model_info = { - "name": "m", - "tags": ["t"], - "n_gpus": "1", - "args": "", - "scripts": "scripts/dummy/run.sh", - "dockerfile": "docker/dummy", - } - manifest = { - "built_images": {"img1": image_info}, - "built_models": {"img1": model_info}, - "context": {}, - } - (tmp_path / "build_manifest.json").write_text(json.dumps(manifest)) - - additional_context = { - "k8s": {"namespace": "default"}, - "gpu_vendor": "AMD", - "guest_os": "UBUNTU", - } - if require_pinned: - additional_context["require_pinned_image"] = True - - cfg = DeploymentConfig( - target="k8s", - manifest_file="build_manifest.json", - additional_context=additional_context, - ) - deployment = KubernetesDeployment(cfg) - return deployment._prepare_template_context(model_info, image_info) - - def test_default_uses_tag(self, tmp_path, monkeypatch): - ctx = self._template_context( - tmp_path, monkeypatch, require_pinned=False, image_digest=self.DIGEST - ) - assert ctx["image"] == "myorg/ci:m" - - def test_enabled_uses_pinned_reference(self, tmp_path, monkeypatch): - ctx = self._template_context( - tmp_path, monkeypatch, require_pinned=True, image_digest=self.DIGEST - ) - assert ctx["image"] == f"myorg/ci@{self.DIGEST}" - - def test_enabled_without_digest_raises(self, tmp_path, monkeypatch): - with pytest.raises(ConfigurationError): - self._template_context( - tmp_path, monkeypatch, require_pinned=True, image_digest=None - ) -``` - -Add whatever of these imports `tests/unit/test_k8s.py` is missing at the top (`pytest` is already there): - -```python -import json - -from madengine.core.errors import ConfigurationError -from madengine.deployment.base import DeploymentConfig -from madengine.deployment.kubernetes import KubernetesDeployment -``` - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/unit/test_k8s.py -k K8sRequirePinnedImage -v` -Expected: FAIL — `test_template_context_wires_the_resolver` fails; the resolver-behaviour tests pass already (they exercise Task 2 code directly, which is intentional: they document the K8s-facing contract). - -- [ ] **Step 3: Write the implementation** - -Add the import to `src/madengine/deployment/k8s_template_context.py`, next to the existing `from madengine.core.errors import ConfigurationError` (line 33): - -```python -from madengine.core.image_digest import resolve_pinned_image -``` - -In `_prepare_template_context`, immediately before the `return {` statement that begins the context dict, add: - -```python - # Under require_pinned_image the pod pulls repo@sha256:... so a moved tag - # surfaces as an ImagePullBackOff rather than a silent wrong-image run. - resolved_image = resolve_pinned_image( - image_info["registry_image"], - image_info.get("image_digest"), - bool(additional_context.get("require_pinned_image")), - model_name=model_name, - ) -``` - -Then change the image entry (line 518) from: - -```python - "image": image_info["registry_image"], -``` - -to: - -```python - "image": resolved_image, -``` - -> `additional_context` and `model_name` are both already local variables in this method (`additional_context = self.config.additional_context.copy()` and `model_name = model_info["name"]` near the top). - -- [ ] **Step 4: Run tests to verify they pass** - -Run: `pytest tests/unit/test_k8s.py -v` -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add src/madengine/deployment/k8s_template_context.py tests/unit/test_k8s.py -git commit -m "feat(k8s): pin pod image to manifest digest when required" -``` - ---- - -## Task 8: Enforce pinned images in the SLURM slurm_multi wrapper - -**Files:** -- Modify: `src/madengine/deployment/slurm.py:440-457` -- Test: `tests/unit/test_slurm_multi.py` - -The standard SLURM template path needs no change — it re-enters `madengine run --manifest-file` on each compute node, which goes through Task 6's local-Docker enforcement using the `require_pinned_image` key that Task 5 persisted into `manifest["context"]`. Only the self-managed `slurm_multi` wrapper, which bypasses that nested run, needs its own resolution. - -Setting `DOCKER_IMAGE_NAME` to the pinned reference pins both the parallel `srun docker pull` (which interpolates this variable) and the `docker run` inside the model's own script. See "Deliberate deviations" above. - -- [ ] **Step 1: Write the failing tests** - -Append to `tests/unit/test_slurm_multi.py`: - -```python -DIGEST = "sha256:" + "df36ef7e" * 8 - - -class TestSlurmMultiRequirePinnedImage: - """slurm_multi wrapper pins DOCKER_IMAGE_NAME (and thus the pull) when required.""" - - IMAGE_KEY = "rocm/pytorch-private:sglang_disagg_mori_20260502" - - def _deployment(self, tmp_path, require_pinned, image_digest): - 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# placeholder\n") - - image_entry = { - "image_name": self.IMAGE_KEY, - "docker_image": self.IMAGE_KEY, - "registry_image": self.IMAGE_KEY, - } - if image_digest: - image_entry["image_digest"] = image_digest - - manifest = { - "built_images": {self.IMAGE_KEY: image_entry}, - "built_models": {self.IMAGE_KEY: PR186_MODEL_ENTRY}, - "context": { - "docker_env_vars": {}, - "docker_mounts": {}, - "docker_build_arg": {}, - "gpu_vendor": "AMD", - "guest_os": "UBUNTU", - "docker_gpus": "all", - }, - } - manifest_path = tmp_path / "build_manifest.json" - manifest_path.write_text(json.dumps(manifest)) - - additional_context = { - "deploy": "slurm", - "gpu_vendor": "AMD", - "guest_os": "UBUNTU", - "slurm": dict( - PR186_MODEL_ENTRY["slurm"], output_dir=str(tmp_path / "slurm_results") - ), - "distributed": PR186_MODEL_ENTRY["distributed"], - } - if require_pinned: - additional_context["require_pinned_image"] = True - - cfg = DeploymentConfig( - target="slurm", - manifest_file=str(manifest_path), - additional_context=additional_context, - ) - return SlurmDeployment(cfg) - - def test_default_exports_tag(self, tmp_path): - dep = self._deployment(tmp_path, require_pinned=False, image_digest=DIGEST) - assert dep.prepare() is True - script_text = Path(dep.script_path).read_text() - assert f"export DOCKER_IMAGE_NAME={shlex.quote(self.IMAGE_KEY)}" in script_text - assert DIGEST not in script_text - - def test_enabled_exports_pinned_reference(self, tmp_path): - dep = self._deployment(tmp_path, require_pinned=True, image_digest=DIGEST) - assert dep.prepare() is True - script_text = Path(dep.script_path).read_text() - - pinned = f"rocm/pytorch-private@{DIGEST}" - assert f"export DOCKER_IMAGE_NAME={shlex.quote(pinned)}" in script_text - # The parallel pull interpolates the same value, so it is pinned too. - assert f"docker pull {pinned}" in script_text - - def test_enabled_without_digest_does_not_silently_fall_through(self, tmp_path): - """A missing digest must abort, not quietly take the standard template path. - - prepare()'s launcher peek wraps the slurm_multi dispatch in a bare - `except Exception: pass`. Without the re-raise added in Step 3b, a - ConfigurationError here would be swallowed and prepare() would generate - an ordinary (unpinned) sbatch script instead — the exact silent - degradation the flag exists to prevent. - """ - dep = self._deployment(tmp_path, require_pinned=True, image_digest=None) - with pytest.raises(ConfigurationError): - dep.prepare() -``` - -Add whatever of these imports `tests/unit/test_slurm_multi.py` is missing at the top: - -```python -import pytest - -from madengine.core.errors import ConfigurationError -``` - -- [ ] **Step 2: Run tests to verify they fail** - -Run: `pytest tests/unit/test_slurm_multi.py -k RequirePinnedImage -v` -Expected: FAIL — `test_enabled_exports_pinned_reference` finds the tag, not the pinned reference. - -- [ ] **Step 3: Write the implementation** - -Add the import to `src/madengine/deployment/slurm.py`, next to the other `madengine.*` imports (after `from madengine.utils.gpu_config import resolve_runtime_gpus`): - -```python -from madengine.core.image_digest import resolve_pinned_image -``` - -In `_prepare_slurm_multi_script`, replace the `DOCKER_IMAGE_NAME` resolution block (lines 440-457): - -```python - # Override DOCKER_IMAGE_NAME with the built image from manifest - # This ensures the run uses the freshly built image, not the base image - # Priority: docker_image_name param > model_info.docker_image > env_vars.DOCKER_IMAGE_NAME - if docker_image_name and docker_image_name.startswith("ci-"): - # The manifest key IS the built image name for madengine-built images - self.console.print(f"[cyan]Using built Docker image: {docker_image_name}[/cyan]") - env_vars["DOCKER_IMAGE_NAME"] = docker_image_name - elif "docker_image" in model_info: - built_image = model_info["docker_image"] - self.console.print(f"[cyan]Using Docker image: {built_image}[/cyan]") - env_vars["DOCKER_IMAGE_NAME"] = built_image - elif "image" in model_info: - # Fallback to 'image' field - built_image = model_info["image"] - self.console.print(f"[cyan]Using Docker image: {built_image}[/cyan]") - env_vars["DOCKER_IMAGE_NAME"] = built_image -``` - -with: - -```python - # Override DOCKER_IMAGE_NAME with the built image from manifest - # This ensures the run uses the freshly built image, not the base image - # Priority: docker_image_name param > model_info.docker_image > env_vars.DOCKER_IMAGE_NAME - if docker_image_name and docker_image_name.startswith("ci-"): - # The manifest key IS the built image name for madengine-built images - self.console.print(f"[cyan]Using built Docker image: {docker_image_name}[/cyan]") - env_vars["DOCKER_IMAGE_NAME"] = docker_image_name - elif "docker_image" in model_info: - built_image = model_info["docker_image"] - self.console.print(f"[cyan]Using Docker image: {built_image}[/cyan]") - env_vars["DOCKER_IMAGE_NAME"] = built_image - elif "image" in model_info: - # Fallback to 'image' field - built_image = model_info["image"] - self.console.print(f"[cyan]Using Docker image: {built_image}[/cyan]") - env_vars["DOCKER_IMAGE_NAME"] = built_image - - # Under require_pinned_image, pin DOCKER_IMAGE_NAME to the digest recorded - # at build time. slurm_multi runs the model's own script (no nested - # `madengine run` on the compute nodes), so enforcement has to happen here. - # Pinning the variable covers both the parallel `srun docker pull` below, - # which interpolates it, and the `docker run` inside the model script. - require_pinned = bool( - self.config.additional_context.get("require_pinned_image") - ) - if require_pinned and env_vars.get("DOCKER_IMAGE_NAME"): - image_entry = (self.manifest.get("built_images") or {}).get( - docker_image_name, {} - ) - env_vars["DOCKER_IMAGE_NAME"] = resolve_pinned_image( - env_vars["DOCKER_IMAGE_NAME"], - image_entry.get("image_digest"), - True, - model_name=model_info.get("name", ""), - ) - self.console.print( - f"[cyan]Pinned Docker image: {env_vars['DOCKER_IMAGE_NAME']}[/cyan]" - ) -``` - -- [ ] **Step 3b: Stop `prepare()` from swallowing the enforcement error** - -`prepare()` (`slurm.py:315-341`) wraps the whole slurm_multi dispatch — including the `_prepare_slurm_multi_script` call itself — in `except Exception: pass`, then falls through to the standard template path. Left as-is, a `ConfigurationError` from Step 3 would be silently discarded and an ordinary *unpinned* sbatch script would be generated instead. - -Narrow the handler so enforcement errors propagate. Replace the `except` clause at `slurm.py:339-341`: - -```python - except Exception: - # Fall through to develop's standard flow on any peek error - pass -``` - -with: - -```python - except ConfigurationError: - # Enforcement failures (e.g. --require-pinned-image with no recorded - # digest) are deliberate aborts, not peek errors. Falling through to - # the standard path here would silently generate an unpinned script. - raise - except Exception: - # Fall through to develop's standard flow on any peek error - pass -``` - -Add the import alongside the other `madengine.*` imports in `slurm.py`: - -```python -from madengine.core.errors import ConfigurationError -``` - -- [ ] **Step 4: Run tests to verify they pass** - -Run: `pytest tests/unit/test_slurm_multi.py -v` -Expected: PASS - -- [ ] **Step 5: Commit** - -```bash -git add src/madengine/deployment/slurm.py tests/unit/test_slurm_multi.py -git commit -m "feat(slurm): pin slurm_multi image to manifest digest when required" -``` - ---- - -## Task 9: Lock in log-filename compatibility for pinned references - -**Files:** -- Test: `tests/unit/test_execution.py` - -`_docker_image_ref_for_log_naming()` already strips `@sha256:...`; this test prevents a future refactor from breaking it and silently changing log/tar filenames when pinning is on. - -- [ ] **Step 1: Write the test** - -Append to the existing `_docker_image_ref_for_log_naming` test class in `tests/unit/test_execution.py`: - -```python - def test_pinned_reference_names_same_as_untagged_reference(self): - digest = "sha256:" + "df36ef7e" * 8 - assert ( - _docker_image_ref_for_log_naming(f"registry/ns/myimg@{digest}") - == _docker_image_ref_for_log_naming("registry/ns/myimg") - ) - - def test_pinned_ci_reference_still_yields_tag(self): - digest = "sha256:" + "df36ef7e" * 8 - assert ( - _docker_image_ref_for_log_naming(f"rocm/ns/img:ci-m_model_df@{digest}") - == "ci-m_model_df" - ) -``` - -- [ ] **Step 2: Run the test** - -Run: `pytest tests/unit/test_execution.py -k log_naming -v` -Expected: PASS immediately — this is a characterization test of behaviour that already exists. If either assertion fails, stop and report it; that would mean pinned references change log filenames and the spec's compatibility claim is wrong. - -- [ ] **Step 3: Commit** - -```bash -git add tests/unit/test_execution.py -git commit -m "test(execution): cover log naming for digest-pinned image references" -``` - ---- - -## Task 10: Documentation and full-suite verification - -**Files:** -- Modify: `docs/cli-reference.md:232` (run options table) -- Modify: `docs/configuration.md` - -- [ ] **Step 1: Add the CLI reference row** - -In `docs/cli-reference.md`, insert a new row into the `run` options table immediately after the `--skip-model-run` row (line 232): - -```markdown -| `--require-pinned-image` | | FLAG | `False` | Pull registry images by the `sha256` digest recorded in the build manifest (`repo@sha256:...`) instead of by tag, so a tag that moved between build and run fails loudly instead of silently running a different image. Fails immediately — with no tag fallback — if the manifest has no digest for an image. Equivalent to the `require_pinned_image` additional-context key. See [Configuration — Pinned image digests](configuration.md#pinned-image-digests). | -``` - -- [ ] **Step 2: Add the configuration section** - -In `docs/configuration.md`, add a new section after the "Run phase: log error pattern scan" section (which ends before "## System environment collection (rocEnvTool)"): - -```markdown -## Pinned image digests - -Every build records the digest of the image it pushes as `image_digest` on each -`built_images` entry in `build_manifest.json`. This capture is always on and -costs nothing: by default the digest is carried along and never used. - -Pass `--require-pinned-image` (or set `"require_pinned_image": true` in -`--additional-context`) to make the run phase pull `repo@sha256:...` instead of -the tag: - -```bash -madengine run --manifest-file build_manifest.json --require-pinned-image - -# Equivalent, for pipelines that drive madengine through additional context -madengine run --manifest-file build_manifest.json \ - --additional-context "{'require_pinned_image': True}" -``` - -| Behaviour | Flag absent (default) | Flag set | -|---|---|---| -| Registry pull | By tag | By digest (`repo@sha256:...`) | -| Manifest has no `image_digest` | Pull by tag | **Fails immediately**, no tag fallback | -| Tag moved since the build | Silently runs the newer image | Registry rejects the pull (`manifest unknown`) | - -Applies to all three execution paths: local Docker, Kubernetes (the pod spec -`image` field), and SLURM. On SLURM the setting is written into the manifest's -`context` block so the nested `madengine run` on each compute node inherits it. - -**Limitations** - -- This does not prevent two concurrent builds from racing to push the same - mutable tag. It converts the resulting silent wrong-image run into a fast, - clear failure. Eliminating the race requires unique tags per build in the - calling CI pipeline. -- Manifests produced by the build-on-compute-node path (SLURM batch builds, - which push from inside a generated sbatch script) carry no `image_digest`. - Runs against those manifests fail fast when the flag is set. -``` - -- [ ] **Step 3: Run the full unit suite** - -Run: `pytest tests/unit -v` -Expected: PASS, no regressions. - -- [ ] **Step 4: Run the integration suite** - -Run: `pytest tests/integration -v -m "not slow"` -Expected: PASS. Pay particular attention to `tests/integration/test_docker_integration.py -k push_image`, which asserts the exact `docker tag` / `docker push` call shapes. - -- [ ] **Step 5: Format and lint the changed files** - -```bash -black src/madengine/core/image_digest.py src/madengine/execution/docker_builder.py \ - src/madengine/execution/container_runner.py src/madengine/deployment/slurm.py \ - src/madengine/deployment/k8s_template_context.py \ - src/madengine/orchestration/run_orchestrator.py src/madengine/cli/commands/run.py \ - tests/unit/test_image_digest.py -isort src/madengine/core/image_digest.py tests/unit/test_image_digest.py -mypy src/madengine/core/image_digest.py -``` - -Expected: `black`/`isort` reformat or report no changes; `mypy` reports no errors in the new module. - -> Only run `black`/`isort` on the files you actually touched. Reformatting untouched files would bloat the diff. - -- [ ] **Step 6: Commit** - -```bash -git add docs/cli-reference.md docs/configuration.md -git commit -m "docs: document --require-pinned-image and image digest capture" -``` - -- [ ] **Step 7: Manual end-to-end smoke check (optional, requires a registry)** - -```bash -# Build and push, then confirm the digest landed in the manifest -madengine build --tags dummy --registry localhost:5000 -python -c "import json; m=json.load(open('build_manifest.json')); print({k: v.get('image_digest') for k, v in m['built_images'].items()})" - -# Run with enforcement and confirm the pull is by digest -madengine run --manifest-file build_manifest.json --require-pinned-image --live-output 2>&1 | grep "docker pull" -``` - -Expected: the manifest prints a `sha256:...` per image, and the pull line contains `@sha256:`. - ---- - -## Verification checklist against the spec's testing plan - -| Spec test | Covered by | -|---|---| -| 1. Push output digest → `image_digest` | Task 3 `test_digest_parsed_from_push_output` + Task 4 `test_single_arch_push_sets_image_digest` | -| 2. Fallback to `docker image inspect` | Task 3 `test_falls_back_to_image_inspect_when_push_output_has_no_digest` | -| 3. Both paths fail → absent key, debug log, build unaffected | Task 3 `test_no_digest_anywhere_leaves_entry_absent_and_push_still_succeeds`, `test_inspect_failure_is_swallowed` | -| 4. Flag absent → pull by tag, no new output | Task 6 `test_default_pulls_by_tag_even_when_digest_present`, Task 7 `test_default_uses_tag`, Task 8 `test_default_exports_tag` | -| 5. Flag present + digest → pinned reference | Task 6 `test_enabled_pulls_pinned_reference`, Task 7 `test_enabled_uses_pinned_reference`, Task 8 `test_enabled_exports_pinned_reference` | -| 6. Flag present, no digest → fail before pull | Task 2 `test_enabled_without_digest_raises`, Task 6 `test_enabled_without_digest_fails_before_pulling` | -| 7. K8s pod spec and SLURM script carry pinned/tag reference | Task 7 + Task 8 | -| 8. Log filename derivation unchanged | Task 9 | -| 9. Existing `docker_sha` / manifest tests unmodified | Task 4 Step 5, Task 10 Steps 3-4 | diff --git a/docs/superpowers/specs/2026-08-27-pinned-image-digest-design.md b/docs/superpowers/specs/2026-08-27-pinned-image-digest-design.md deleted file mode 100644 index 08d69aea..00000000 --- a/docs/superpowers/specs/2026-08-27-pinned-image-digest-design.md +++ /dev/null @@ -1,156 +0,0 @@ -# Design: pin registry pulls to the digest recorded at build time - -## Problem - -A client-perf-hub accuracy run failed because the image pulled at run time -did not match the image pushed at build time: - -- build pushed `sha256:df36ef7e...` -- the CI runner later pulled `sha256:cbb7e5ed...` - -Both shared the same registry tag. A concurrent build of the same model -pushed a newer image to that tag between the two events, so the run phase -silently executed a different image than the one the build phase produced. - -It was assumed madengine already guards against this ("we use it to confirm -if the pulled image is identical to the one in the build manifest"). It does -not. Tracing the code: - -- `build_info["docker_sha"]` (`docker_builder.py:303-308`) is the digest of - the Dockerfile's `FROM` (base) image, not the image madengine builds and - pushes. It is consumed only as a reporting column (`update_perf_csv.py`, - `k8s_results.py`, `slurm.py`). -- `push_image()` (`docker_builder.py:395`) discards `docker push` stdout, - which is where the pushed digest (`digest: sha256:...`) is printed. It is - never captured anywhere. -- The run phase pulls by tag only (`container_runner.py:2854` → - `pull_image()` at `container_runner.py:572`, a plain `docker pull `) - and performs no comparison against anything in the manifest. The only - image-identity checks in the codebase (`_local_image_id`, - `BUILD_FINGERPRINT_LABEL`, `container_runner.py:2464-2503`) exist solely - for cross-node consistency in multi-node SLURM *local-image* mode; they - never touch registry digests. - -This design adds the missing capability: capture the real pushed digest at -build time, and optionally enforce it at run time. - -## Non-goals - -- This does not prevent the tag race itself. Two builds pushing the same - mutable tag (`f"{registry}:{model_name}"`, `build_orchestrator.py:1055`) - will still clobber each other; the loser under the new flag fails fast - with a clear error instead of silently running the wrong image. Actually - eliminating the race (e.g. unique tags per build) is a client-perf-hub / - upstream CI change, out of scope here. -- This does not rename or repurpose `build_info["docker_sha"]`. It is wired - into four existing reporting sinks and column orders; changing its - meaning is a separate, riskier change and is called out only as a - possible future cleanup. -- This does not change default pull behavior for any existing user. See - "Opt-in enforcement" below. - -## Design - -### 1. Capture the pushed digest at build time (always on, default mode included) - -In `DockerBuilder.push_image()` (`docker_builder.py`), after `docker push` -succeeds, parse the digest from its output (`digest: sha256:...`, the same -line format already parsed for base-image SHA in `docker_builder.py:305`). - -If the push output doesn't contain a parseable digest line (registry output -format variance, e.g. some mirrors), fall back to: -``` -docker image inspect --format '{{index .RepoDigests 0}}' -``` -and extract the digest from that. - -If both fail to produce a digest, log a **debug/dim-level** note (not a -user-facing warning) and continue — this is a manifest-completeness gap to -leave a trail for later, not a failure. Push already succeeded; we don't -block on this. - -Store the result as a new field: `build_info["image_digest"]` (e.g. -`"sha256:df36ef7e..."`), separate from `docker_sha`. This flows into -`build_manifest.json` through the existing `built_images` / -`export_build_manifest()` path with no other changes needed. - -This capture step runs unconditionally — no flag gates it. It's the -recording half of the fix, and it's inert (never read by pull logic) unless -strict mode is on. - -### 2. Opt-in enforcement at run time - -New flag: `--require-pinned-image` on the `run` command, mirrored as an -`additional_context` key (e.g. `"require_pinned_image": true`) so CI -pipelines that drive madengine through `additional_context` rather than raw -CLI flags can set it the same way as other behavior-affecting keys -(`k8s`, `slurm`, `tools`, ...). - -**Default (flag absent):** no behavior change whatsoever. Pull by tag, -exactly as today. `image_digest` rides along in the manifest unused. - -**With the flag set**, for every `build_info` entry with a -`registry_image`: -- If `image_digest` is present, construct a pinned reference - `repo@sha256:...` (stripping any existing tag) and pull that image - instead of the tag. The registry itself now enforces identity: a moved - tag surfaces as a normal "manifest unknown" pull failure rather than a - silent wrong-image success. -- If `image_digest` is absent (older manifest, or capture failed at build - time), fail immediately with a clear error naming the model/image and - explaining that the manifest has no recorded digest — do not fall back to - pulling by tag. Silent degradation defeats the purpose of asking for the - guarantee. - -One shared helper (e.g. `build_pinned_reference(registry_image, digest)`) -builds the `repo@sha256:...` string, used identically by all three -enforcement call sites so they can't drift: - -| Path | Location | Change under the flag | -|---|---|---| -| Local Docker | `container_runner.py:2854` | `pull_image(pinned_ref)` | -| Kubernetes | `k8s_template_context.py:518` (`"image": image_info["registry_image"]`) | `"image": pinned_ref` | -| SLURM | `slurm.py:544` (parallel `srun` pull) | pinned ref substituted into the generated pull command | - -### 3. Compatibility check: log filenames - -`container_runner_helpers.py:256` already strips `@sha256:...` before -deriving log/tar filenames from an image reference (`ref_without_digest = -s.split("@", 1)[0]`). Pinned references flow through this unchanged — no -new filename collisions. Confirmed by reading the function; will be -covered by a test regardless. - -## Testing plan - -TDD; each behavior below is a separate test: - -1. `docker push` output containing a `digest: sha256:...` line → - `build_info["image_digest"]` set to that value. -2. `docker push` output without a parseable digest line → falls back to - `docker image inspect --format '{{index .RepoDigests 0}}'`. -3. Both parse paths fail → `image_digest` absent, a debug-level log line is - emitted, and the build/push otherwise succeeds unaffected. -4. Flag absent (default) → run pulls by tag regardless of whether - `image_digest` is present in the manifest; no new log output. -5. Flag present, `image_digest` present → the pull command (or k8s image - field / SLURM pull line) contains `repo@sha256:...`. -6. Flag present, `image_digest` absent → run fails immediately with a clear - error, before any pull is attempted. -7. K8s pod spec and SLURM generated script both carry the pinned reference - when the flag is set, and the untouched tag reference when it isn't. -8. Log/tar filename derivation given a pinned (`@sha256:...`) reference - produces the same filename as today's digest-free reference. -9. Existing tests touching `docker_sha` / `build_manifest.json` structure - continue to pass unmodified (the new field is additive). - -## Rollout note (for the reply to Tej/Rahul) - -- The verification described in the thread ("we use it to confirm...") - does not currently exist in the code; this design builds it, gated - behind `--require-pinned-image` / `require_pinned_image` context key. -- Turning the flag on stops the *symptom* (silently running the wrong - image) by failing fast instead. It does not stop the *cause* (two builds - racing to push the same mutable tag) — that needs a change on the - client-perf-hub / CI side (e.g. unique tags per build). -- client-perf-hub must explicitly opt in for this guarantee to apply to its - runs; it is not automatic on upgrade. diff --git a/examples/k8s-configs/README.md b/examples/k8s-configs/README.md index fbd327dd..3bfbd23a 100644 --- a/examples/k8s-configs/README.md +++ b/examples/k8s-configs/README.md @@ -185,7 +185,7 @@ To validate rendered YAML after a debug run, install [kubeconform](https://githu ### Multi-node DNS (PyTorch vs Ray) -For **PyTorch-native** launchers (`torchrun`, `deepspeed`, `torchtitan`, `megatron`, `primus`), multi-node Jobs use a **headless Service** whose name matches `pod.spec.subdomain`, per Kubernetes DNS rules, so pods get stable per-pod DNS names for rendezvous. +For **PyTorch-native** launchers (`torchrun`, `deepspeed`, `torchtitan`, `megatron-lm`, `primus`), multi-node Jobs use a **headless Service** whose name matches `pod.spec.subdomain`, per Kubernetes DNS rules, so pods get stable per-pod DNS names for rendezvous. For **Ray-based** multi-node (`vllm`, `sglang`), a headless Service may still be created for networking, but **per-pod DNS via `subdomain` is not applied** the same way as for PyTorch; production multi-node Ray on Kubernetes often uses **KubeRay** (see upstream vLLM / Ray docs). Treat Job-based multi-node Ray as a best-effort path. @@ -580,7 +580,7 @@ Configuration for distributed workloads (training and inference): | Field | Type | Default | Description | |-------|------|---------|-------------| -| `launcher` | string | - | Launcher type: `torchrun`, `deepspeed`, `torchtitan`, `megatron`, `primus`, `vllm`, `sglang` | +| `launcher` | string | - | Launcher type: `torchrun`, `deepspeed`, `torchtitan`, `megatron-lm`, `primus`, `vllm`, `sglang` | | `enabled` | boolean | `false` | Enable distributed execution (legacy, prefer `launcher`) | | `backend` | string | `"nccl"` | `"nccl"`, `"gloo"`, or `"mpi"` | | `nnodes` | integer | `1` | Number of nodes | @@ -679,7 +679,7 @@ Write durable outputs under `/results//` in the container so each re **Training Launchers:** - **torchrun**: Standard PyTorch DDP/FSDP training - **deepspeed**: ZeRO optimization for memory efficiency -- **megatron**: Megatron-LM tensor and pipeline parallelism +- **megatron-lm**: Megatron-LM tensor and pipeline parallelism - **torchtitan**: LLM pre-training with multi-dimensional parallelism (FSDP2+TP+PP) - **primus**: Unified Primus pretrain (Megatron / TorchTitan / MaxText experiment YAML; see [Primus on Kubernetes](#primus-on-kubernetes)) diff --git a/examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu-tools.json b/examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu-tools.json index 3c5f80ae..a231ab14 100644 --- a/examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu-tools.json +++ b/examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu-tools.json @@ -28,7 +28,7 @@ "launcher": "torchrun", "nnodes": 1, "nproc_per_node": 2, - "master_port": 29500 + "port": 29500 }, "env_vars": { diff --git a/examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu.json b/examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu.json index be0d7c5e..3fdf6d26 100644 --- a/examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu.json +++ b/examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu.json @@ -27,7 +27,7 @@ "launcher": "torchrun", "nnodes": 1, "nproc_per_node": 2, - "master_port": 29500 + "port": 29500 }, "env_vars": { diff --git a/examples/k8s-configs/basic/03-torchrun-multi-node-basic.json b/examples/k8s-configs/basic/03-torchrun-multi-node-basic.json index 0c2205f9..6db11b10 100644 --- a/examples/k8s-configs/basic/03-torchrun-multi-node-basic.json +++ b/examples/k8s-configs/basic/03-torchrun-multi-node-basic.json @@ -28,7 +28,7 @@ "launcher": "torchrun", "nnodes": 2, "nproc_per_node": 2, - "master_port": 29500 + "port": 29500 }, "env_vars": { diff --git a/examples/k8s-configs/basic/04-torchrun-multi-node-advanced.json b/examples/k8s-configs/basic/04-torchrun-multi-node-advanced.json index 5560ffab..bee06a54 100644 --- a/examples/k8s-configs/basic/04-torchrun-multi-node-advanced.json +++ b/examples/k8s-configs/basic/04-torchrun-multi-node-advanced.json @@ -54,7 +54,7 @@ "launcher": "torchrun", "nnodes": 4, "nproc_per_node": 2, - "master_port": 29500 + "port": 29500 }, "env_vars": { diff --git a/examples/k8s-configs/basic/05-torchrun-nvidia-gpu-example.json b/examples/k8s-configs/basic/05-torchrun-nvidia-gpu-example.json index 7c087acc..6d2f37d0 100644 --- a/examples/k8s-configs/basic/05-torchrun-nvidia-gpu-example.json +++ b/examples/k8s-configs/basic/05-torchrun-nvidia-gpu-example.json @@ -31,7 +31,7 @@ "launcher": "torchrun", "nnodes": 1, "nproc_per_node": 4, - "master_port": 29500 + "port": 29500 }, "env_vars": { diff --git a/examples/k8s-configs/basic/06-data-provider-with-pvc.json b/examples/k8s-configs/basic/06-data-provider-with-pvc.json index 9bd2e47f..4c52f17e 100644 --- a/examples/k8s-configs/basic/06-data-provider-with-pvc.json +++ b/examples/k8s-configs/basic/06-data-provider-with-pvc.json @@ -36,7 +36,7 @@ "nnodes": 1, "nproc_per_node": 2, - "master_port": 29500 + "port": 29500 }, "env_vars": { diff --git a/examples/k8s-configs/basic/megatron-lm-multi-node-basic.json b/examples/k8s-configs/basic/megatron-lm-multi-node-basic.json index e059ba08..b0c0efd5 100644 --- a/examples/k8s-configs/basic/megatron-lm-multi-node-basic.json +++ b/examples/k8s-configs/basic/megatron-lm-multi-node-basic.json @@ -18,10 +18,10 @@ }, "distributed": { - "launcher": "megatron", + "launcher": "megatron-lm", "nnodes": 4, "nproc_per_node": 8, - "master_port": 29500 + "port": 29500 }, "env_vars": { diff --git a/examples/k8s-configs/basic/sglang-disagg-custom-split.json b/examples/k8s-configs/basic/sglang-disagg-custom-split.json index 49aeecb1..97cd7509 100644 --- a/examples/k8s-configs/basic/sglang-disagg-custom-split.json +++ b/examples/k8s-configs/basic/sglang-disagg-custom-split.json @@ -29,7 +29,7 @@ "launcher": "sglang-disagg", "nnodes": 7, "nproc_per_node": 8, - "master_port": 29500, + "port": 29500, "sglang_disagg": { "prefill_nodes": 4, "decode_nodes": 2 diff --git a/examples/k8s-configs/basic/sglang-disagg-multi-node-basic.json b/examples/k8s-configs/basic/sglang-disagg-multi-node-basic.json index c16fd342..25096a9e 100644 --- a/examples/k8s-configs/basic/sglang-disagg-multi-node-basic.json +++ b/examples/k8s-configs/basic/sglang-disagg-multi-node-basic.json @@ -28,7 +28,7 @@ "launcher": "sglang-disagg", "nnodes": 5, "nproc_per_node": 8, - "master_port": 29500 + "port": 29500 }, "context": { diff --git a/examples/k8s-configs/basic/sglang-multi-node-basic.json b/examples/k8s-configs/basic/sglang-multi-node-basic.json index b693260e..59dab1af 100644 --- a/examples/k8s-configs/basic/sglang-multi-node-basic.json +++ b/examples/k8s-configs/basic/sglang-multi-node-basic.json @@ -22,7 +22,7 @@ "launcher": "sglang", "nnodes": 2, "nproc_per_node": 4, - "master_port": 29500 + "port": 29500 }, "context": { diff --git a/examples/k8s-configs/basic/torchtitan-multi-node-basic.json b/examples/k8s-configs/basic/torchtitan-multi-node-basic.json index e350605d..db4d1190 100644 --- a/examples/k8s-configs/basic/torchtitan-multi-node-basic.json +++ b/examples/k8s-configs/basic/torchtitan-multi-node-basic.json @@ -22,7 +22,7 @@ "launcher": "torchtitan", "nnodes": 4, "nproc_per_node": 8, - "master_port": 29500 + "port": 29500 }, "context": { diff --git a/examples/k8s-configs/basic/vllm-multi-node-basic.json b/examples/k8s-configs/basic/vllm-multi-node-basic.json index 4c1b61c9..eb9f4cae 100644 --- a/examples/k8s-configs/basic/vllm-multi-node-basic.json +++ b/examples/k8s-configs/basic/vllm-multi-node-basic.json @@ -22,7 +22,7 @@ "launcher": "vllm", "nnodes": 2, "nproc_per_node": 4, - "master_port": 29500 + "port": 29500 }, "context": { diff --git a/examples/k8s-configs/minimal/megatron-lm-exclude-node.json b/examples/k8s-configs/minimal/megatron-lm-exclude-node.json index 793431a2..09e69e8b 100644 --- a/examples/k8s-configs/minimal/megatron-lm-exclude-node.json +++ b/examples/k8s-configs/minimal/megatron-lm-exclude-node.json @@ -24,7 +24,7 @@ }, "distributed": { - "launcher": "megatron", + "launcher": "megatron-lm", "nnodes": 1, "nproc_per_node": 2 }, diff --git a/examples/k8s-configs/minimal/megatron-lm-minimal.json b/examples/k8s-configs/minimal/megatron-lm-minimal.json index 43266e01..d033a26d 100644 --- a/examples/k8s-configs/minimal/megatron-lm-minimal.json +++ b/examples/k8s-configs/minimal/megatron-lm-minimal.json @@ -14,7 +14,7 @@ }, "distributed": { - "launcher": "megatron", + "launcher": "megatron-lm", "nnodes": 1, "nproc_per_node": 2 }, diff --git a/examples/k8s-configs/minimal/megatron-lm-optimized.json b/examples/k8s-configs/minimal/megatron-lm-optimized.json index 29559308..bfeba1b4 100644 --- a/examples/k8s-configs/minimal/megatron-lm-optimized.json +++ b/examples/k8s-configs/minimal/megatron-lm-optimized.json @@ -29,10 +29,10 @@ "distributed": { "enabled": true, "backend": "nccl", - "launcher": "megatron", + "launcher": "megatron-lm", "nnodes": 1, "nproc_per_node": 2, - "master_port": 29500 + "port": 29500 }, "env_vars": { diff --git a/examples/slurm-configs/README.md b/examples/slurm-configs/README.md index 0100fe54..c75c0356 100644 --- a/examples/slurm-configs/README.md +++ b/examples/slurm-configs/README.md @@ -412,7 +412,7 @@ madengine uses intelligent multi-layer configuration merging: ```json { "distributed": { - "launcher": "torchrun", // Launcher type: torchrun, vllm, sglang, deepspeed, megatron, slurm_multi + "launcher": "torchrun", // Launcher type: torchrun, vllm, sglang, deepspeed, megatron-lm, slurm_multi "backend": "nccl", // Communication backend (nccl/gloo) "port": 29500, // Master node port "nnodes": 2, // Number of nodes (overrides slurm.nodes if set) @@ -426,12 +426,12 @@ madengine uses intelligent multi-layer configuration merging: - `vllm`: vLLM inference engine (TP/PP parallelism) - `sglang`: SGLang inference engine - `deepspeed`: DeepSpeed training framework -- `megatron`: Megatron-LM large model training +- `megatron-lm`: Megatron-LM large model training - `slurm_multi` / `slurm-multi`: Self-managed multi-container topologies (escape hatch) - Custom: Set environment variables, model script handles launcher **Note**: For vLLM and SGLang, the model script handles process spawning directly. -For torchrun/deepspeed/megatron, use `$MAD_MULTI_NODE_RUNNER` in your model script. +For torchrun/deepspeed/megatron-lm, use `$MAD_MULTI_NODE_RUNNER` in your model script. For slurm_multi, the model's `.slurm` script runs on baremetal and manages Docker containers via `srun` internally. ### Environment Variables diff --git a/examples/slurm-configs/basic/09-megatron-lm-multi-node.json b/examples/slurm-configs/basic/09-megatron-lm-multi-node.json index 84e3c3f6..bb718a20 100644 --- a/examples/slurm-configs/basic/09-megatron-lm-multi-node.json +++ b/examples/slurm-configs/basic/09-megatron-lm-multi-node.json @@ -17,10 +17,10 @@ }, "distributed": { - "launcher": "megatron", + "launcher": "megatron-lm", "nnodes": 4, "nproc_per_node": 8, - "master_port": 29500 + "port": 29500 }, "env_vars": { diff --git a/examples/slurm-configs/minimal/megatron-lm-minimal.json b/examples/slurm-configs/minimal/megatron-lm-minimal.json index 9480359e..3461f71a 100644 --- a/examples/slurm-configs/minimal/megatron-lm-minimal.json +++ b/examples/slurm-configs/minimal/megatron-lm-minimal.json @@ -14,7 +14,7 @@ }, "distributed": { - "launcher": "megatron", + "launcher": "megatron-lm", "nnodes": 1, "nproc_per_node": 2 }, diff --git a/src/madengine/cli/validators.py b/src/madengine/cli/validators.py index 68f45856..5ed5ce69 100644 --- a/src/madengine/cli/validators.py +++ b/src/madengine/cli/validators.py @@ -250,6 +250,47 @@ def _validate_gpu_vendor_guest_after_defaults(context: Dict[str, Any]) -> None: ) +def _validate_launcher_after_defaults(context: Dict[str, Any]) -> None: + """Validate any launcher in the context and rewrite it to its canonical spelling. + + Fails here rather than at deploy time: a launcher madengine does not recognize + used to run the model as a plain single-process job and still report SUCCESS, + so the benchmark number was wrong with nothing to indicate it. + """ + from madengine.core.errors import ConfigurationError + from madengine.deployment.common import validate_launcher + + launcher_cfg = context.get("launcher") + if launcher_cfg is not None and not isinstance(launcher_cfg, dict): + # A bare string is a natural mistake, since distributed.launcher *is* a + # string. Name both valid shapes rather than only rejecting this one. + console.print(f"❌ Invalid launcher: [red]{launcher_cfg!r}[/red]") + console.print("💡 'launcher' must be an object. Use one of:") + console.print( + ' [green]{"launcher": {"type": "torchrun", "nnodes": 2}}[/green]' + ) + console.print( + ' [green]{"distributed": {"launcher": "torchrun", "nnodes": 2}}[/green]' + ) + raise typer.Exit(ExitCode.INVALID_ARGS) + + targets = [] + distributed = context.get("distributed") + if isinstance(distributed, dict) and "launcher" in distributed: + targets.append((distributed, "launcher")) + if isinstance(launcher_cfg, dict) and "type" in launcher_cfg: + targets.append((launcher_cfg, "type")) + + for holder, key in targets: + try: + holder[key] = validate_launcher(holder[key], source="additional_context") + except ConfigurationError as exc: + console.print(f"❌ Invalid launcher: [red]{holder[key]!r}[/red]") + for suggestion in exc.suggestions or []: + console.print(f"💡 [green]{suggestion}[/green]") + raise typer.Exit(ExitCode.INVALID_ARGS) + + def finalize_additional_context_dict( context: Dict[str, Any], *, @@ -280,6 +321,7 @@ def finalize_additional_context_dict( validate_additional_context_structure(context) _normalize_docker_build_arg_values(context) _validate_gpu_vendor_guest_after_defaults(context) + _validate_launcher_after_defaults(context) return context diff --git a/src/madengine/deployment/base.py b/src/madengine/deployment/base.py index 53261482..de9fc744 100644 --- a/src/madengine/deployment/base.py +++ b/src/madengine/deployment/base.py @@ -134,6 +134,76 @@ def __init__(self, config: DeploymentConfig): self.config = config self.manifest = self._load_manifest(config.manifest_file) self.console = Console() + self._validate_launchers() + + def _validate_launchers(self) -> None: + """Validate every launcher this deployment will read, and canonicalize in place. + + Deliberately in ``__init__`` rather than ``validate()``: ``execute()`` catches + bare ``Exception`` and returns a FAILED result without re-raising, so a + ConfigurationError raised any later never reaches the handler in + ``cli/commands/run.py``. ``__init__`` runs under DeploymentFactory.create(), + which re-raises ConfigurationError, so the user gets INVALID_ARGS and a message + naming the correct spelling. + + Raises: + ConfigurationError: If any configured launcher is not a valid launcher. + """ + # Imported here: common.py imports from core.errors, and a module-level import + # would make base.py part of that chain for every deployment consumer. + from madengine.core.errors import ConfigurationError, create_error_context + + from .common import validate_launcher + + context = self.config.additional_context or {} + + distributed = context.get("distributed") + if isinstance(distributed, dict) and "launcher" in distributed: + distributed["launcher"] = validate_launcher( + distributed["launcher"], source="additional_context.distributed.launcher" + ) + + launcher_cfg = context.get("launcher") + if launcher_cfg is not None and not isinstance(launcher_cfg, dict): + # A bare string here is a natural mistake, since distributed.launcher *is* + # a string. Left alone it surfaces as AttributeError deep in the K8s + # template context, so name both valid shapes now. + raise ConfigurationError( + f"'launcher' in additional_context must be an object, got " + f"{type(launcher_cfg).__name__} ({launcher_cfg!r})", + context=create_error_context( + operation="validate_launchers", + component="deployment.base", + additional_info={"launcher": launcher_cfg}, + ), + suggestions=[ + 'Use {"launcher": {"type": "torchrun", "nnodes": 2}}', + 'Or {"distributed": {"launcher": "torchrun", "nnodes": 2}}', + ], + ) + if isinstance(launcher_cfg, dict) and "type" in launcher_cfg: + launcher_cfg["type"] = validate_launcher( + launcher_cfg["type"], source="additional_context.launcher.type" + ) + + deployment_config = self.manifest.get("deployment_config") + if isinstance(deployment_config, dict): + manifest_distributed = deployment_config.get("distributed") + if isinstance(manifest_distributed, dict) and "launcher" in manifest_distributed: + manifest_distributed["launcher"] = validate_launcher( + manifest_distributed["launcher"], + source="build_manifest.json deployment_config.distributed.launcher", + ) + + for model_name, model_info in (self.manifest.get("built_models") or {}).items(): + if not isinstance(model_info, dict): + continue + model_distributed = model_info.get("distributed") + if isinstance(model_distributed, dict) and "launcher" in model_distributed: + model_distributed["launcher"] = validate_launcher( + model_distributed["launcher"], + source=f"model '{model_name}' distributed.launcher", + ) def _load_manifest(self, manifest_file: str) -> Dict: """ diff --git a/src/madengine/deployment/common.py b/src/madengine/deployment/common.py index 13657246..88b52bcd 100644 --- a/src/madengine/deployment/common.py +++ b/src/madengine/deployment/common.py @@ -8,11 +8,15 @@ Copyright (c) Advanced Micro Devices, Inc. All rights reserved. """ +import difflib import functools import subprocess from typing import Any, Dict, List, Optional -# Valid distributed launchers (used by normalize_launcher) +from madengine.core.errors import ConfigurationError, create_error_context + +# Valid distributed launchers. Each has exactly one accepted spelling: a value +# that is not in this list is an error, not something to guess at. VALID_LAUNCHERS = [ "torchrun", "torchtitan", @@ -25,22 +29,99 @@ "slurm_multi", ] -# Alternate spellings for distributed launcher values → canonical form. -# Add new aliases here only; do not branch on alternate spellings at dispatch sites. -_LAUNCHER_ALIASES: Dict[str, str] = { - "sglang_disagg": "sglang-disagg", +# The one accepted alternate spelling, kept because docs/launchers.md advertises +# it. Deliberately a named entry rather than a blanket "-" → "_" rewrite, which +# would silently accept hyphen variants of every other launcher too. +_DOCUMENTED_ALIASES: Dict[str, str] = { + "slurm-multi": "slurm_multi", } +# Deployment-mode sentinels meaning "no distributed launcher". They are produced +# by launcher_for_reporting() for the perf.csv ``launcher`` column and are never +# read back as config, so validate_launcher rejects them: a user who writes +# ``launcher: docker`` means "no launcher" and should say so by omitting the key, +# not by naming a value that has no dispatch arm. +_LAUNCHER_SENTINELS = frozenset({"docker", "native"}) + + +def validate_launcher(launcher: Optional[str], *, source: str) -> Optional[str]: + """Validate a user-supplied launcher and return its canonical spelling. -def canonicalize_distributed_launcher(launcher: Optional[str]) -> Optional[str]: - """Normalize alternate launcher spellings to their canonical form. + ``None`` and blank strings mean "no launcher configured" and return None. + Every other value must be a recognized launcher: unknown values raise rather + than falling back, because a silently-defaulted launcher runs the model as a + plain single-process job and still reports SUCCESS. Other falsy values (0, + False, []) are misconfiguration, not absence, and raise. - Resolves aliases (e.g. ``sglang_disagg`` → ``sglang-disagg``). Unknown or - empty values are returned unchanged; callers do their own validation. + Args: + launcher: Raw launcher value from config, model card, or environment. + source: Where the value came from, e.g. ``additional_context.distributed``. + Included in the error so the user knows which file to edit. + + Returns: + The canonical launcher name, or None when nothing was configured. + + Raises: + ConfigurationError: If the value is not a recognized launcher. """ - if not launcher: - return launcher - return _LAUNCHER_ALIASES.get(launcher, launcher) + if launcher is None: + return None + if not isinstance(launcher, str): + raise ConfigurationError( + f"Invalid launcher in {source}: expected a string, got " + f"{type(launcher).__name__} ({launcher!r})", + context=create_error_context( + operation="validate_launcher", + component="deployment.common", + additional_info={"launcher": launcher, "source": source}, + ), + suggestions=[f"Supported launchers: {', '.join(VALID_LAUNCHERS)}"], + ) + + normalized = launcher.strip().lower() + if not normalized: + return None + normalized = _DOCUMENTED_ALIASES.get(normalized, normalized) + if normalized in VALID_LAUNCHERS: + return normalized + + suggestions = [] + if normalized in _LAUNCHER_SENTINELS: + # Emitted into perf.csv by launcher_for_reporting(); has no dispatch arm. + suggestions.append( + f"'{normalized}' is a reporting value for runs with no distributed " + "launcher — omit the launcher key instead of setting it" + ) + else: + close = difflib.get_close_matches(normalized, VALID_LAUNCHERS, n=1, cutoff=0.6) + if close: + suggestions.append(f"Did you mean '{close[0]}'?") + suggestions.append(f"Supported launchers: {', '.join(VALID_LAUNCHERS)}") + raise ConfigurationError( + f"Unknown launcher '{launcher}' in {source}", + context=create_error_context( + operation="validate_launcher", + component="deployment.common", + additional_info={"launcher": launcher, "source": source}, + ), + suggestions=suggestions, + ) + + +def launcher_for_reporting( + launcher_type: Optional[str], deployment_type: str +) -> str: + """Return the launcher to record in perf results. Never raises. + + Values reaching here have already passed validate_launcher at the config + boundary, so this only supplies the sentinel for the "no launcher + configured" case: ``native`` on Kubernetes (the pod is the container), + ``docker`` elsewhere. Reporting must not raise — BaseDeployment.execute() + drops metrics on exception, which would lose results from a successful run. + """ + if launcher_type: + return launcher_type + return "native" if deployment_type == "kubernetes" else "docker" # Tool names that use rocprof / rocprofv3 wrapping and need MPI-aware rocprofv3 on multi-node. @@ -75,38 +156,22 @@ def tools_include_rocprof_family(tools_config: List[Dict]) -> bool: return False -def normalize_launcher(launcher_type: Optional[str], deployment_type: str) -> str: - """ - Normalize launcher field based on deployment type and launcher value. +# Launchers that run one independent replica per node, so a node-local metric is +# a per-replica figure and job throughput is that figure scaled by nnodes. +# +# Deliberately NOT the same set as the Ray-based launchers guarded in the job +# templates: sglang-disagg is Ray-based but its nodes are heterogeneous roles +# (proxy + prefill + decode) cooperating on a single endpoint, not replicas. +# Its throughput is reported once, by the proxy — scaling that by nnodes would +# multiply a whole-cluster number. +PER_REPLICA_LAUNCHERS: frozenset = frozenset({"vllm", "sglang"}) - Logic: - - If launcher is in VALID_LAUNCHERS: keep as-is - - If launcher's hyphen/underscore variant is in VALID_LAUNCHERS: normalize - (e.g. "slurm-multi" -> "slurm_multi") - - If launcher is None/empty/invalid: - * local → "docker" (runs in Docker container) - * slurm → "docker" (typically uses containers on compute nodes) - * kubernetes → "native" (pod itself is the container) - - Args: - launcher_type: Raw launcher type from config (may be None) - deployment_type: "local", "slurm", or "kubernetes" - Returns: - Normalized launcher string - """ - if launcher_type and launcher_type in VALID_LAUNCHERS: - return launcher_type - # Normalize hyphen variant: slurm-multi -> slurm_multi - if launcher_type and launcher_type.replace("-", "_") in VALID_LAUNCHERS: - return launcher_type.replace("-", "_") - if deployment_type == "local": - return "docker" - if deployment_type == "slurm": - return "docker" - if deployment_type == "kubernetes": - return "native" - return "docker" +def is_per_replica_launcher(launcher_type: Optional[str]) -> bool: + """Return True if each node reports its own replica's metric, not the job's.""" + if not launcher_type or not isinstance(launcher_type, str): + return False + return launcher_type.strip().lower() in PER_REPLICA_LAUNCHERS _SELF_MANAGED_LAUNCHERS: frozenset = frozenset({"slurm_multi"}) @@ -119,10 +184,16 @@ def is_self_managed_launcher(launcher_type: Optional[str]) -> bool: directly on the head node and orchestrate Docker containers via srun internally. They bypass the standard sbatch template entirely and are an escape hatch — not peers of the templated launchers (torchrun, vllm, sglang, etc.). + + Callers pass raw model-card values here during early "should we take the + self-managed path?" peeks, before validation has run, so this normalizes + inline and never raises. """ - if not launcher_type: + if not launcher_type or not isinstance(launcher_type, str): return False - return normalize_launcher(launcher_type, "slurm") in _SELF_MANAGED_LAUNCHERS + normalized = launcher_type.strip().lower() + normalized = _DOCUMENTED_ALIASES.get(normalized, normalized) + return normalized in _SELF_MANAGED_LAUNCHERS @functools.lru_cache(maxsize=None) diff --git a/src/madengine/deployment/k8s_results.py b/src/madengine/deployment/k8s_results.py index 6da189b5..35ecb0bf 100644 --- a/src/madengine/deployment/k8s_results.py +++ b/src/madengine/deployment/k8s_results.py @@ -15,7 +15,7 @@ from pathlib import Path from typing import Any, Dict, List, Optional -from .common import normalize_launcher +from .common import is_per_replica_launcher, launcher_for_reporting from madengine.utils.path_utils import scripts_base_dir_from from madengine.utils.run_details import flatten_tags_in_place, get_build_number, get_pipeline @@ -129,9 +129,9 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: ) # Normalize launcher based on deployment type and validity - launcher_type = normalize_launcher(launcher_type, "kubernetes") + launcher_type = launcher_for_reporting(launcher_type, "kubernetes") - is_ray_launcher = launcher_type in ["vllm", "sglang"] + is_per_replica = is_per_replica_launcher(launcher_type) # Sort pods by name to ensure consistent ordering (pod-0 is master) sorted_pods = sorted(pods.items, key=lambda p: p.metadata.name) @@ -146,9 +146,9 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: per_node_metrics = [] # Store performance from each node results["nodes"] = [] # Store per-node details for display - # Special handling for Ray-based launchers (vLLM, SGLang) + # Special handling for data-parallel launchers (vLLM, SGLang) # These report per-replica metrics, need scaling - if is_multinode and is_ray_launcher: + if is_multinode and is_per_replica: self.console.print( f"[cyan]Multi-node Ray deployment: {nnodes} nodes (Data Parallel mode)[/cyan]" ) @@ -207,7 +207,7 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: if perf_data: # For Ray launchers, this is per-replica metric - if is_multinode and is_ray_launcher: + if is_multinode and is_per_replica: perf_data["is_per_replica"] = True per_node_metrics.append(perf_data) self.console.print( @@ -259,7 +259,7 @@ def collect_results(self, deployment_id: str) -> Dict[str, Any]: # ======================================================================== if per_node_metrics: # Special handling for Ray launchers - multiply by nnodes - if is_multinode and is_ray_launcher: + if is_multinode and is_per_replica: original_perf = per_node_metrics[0]["performance"] aggregated_perf = original_perf * nnodes self.console.print( @@ -920,7 +920,7 @@ def _create_failure_record(self, model_info: Dict, build_info: Dict, pod_name: s if nproc_per_node is None: nproc_per_node = int(model_info.get("n_gpus", 1)) # Launcher: use distributed.launcher when set, otherwise "native" for k8s - launcher = normalize_launcher(distributed_config.get("launcher"), "kubernetes") + launcher = launcher_for_reporting(distributed_config.get("launcher"), "kubernetes") # Create a record with the same structure as successful runs # but with performance=0, metric="", and status="FAILED" @@ -996,7 +996,7 @@ def _build_perf_entry_from_aggregated( nproc_per_node = distributed_config.get("nproc_per_node") if nproc_per_node is None: nproc_per_node = int(model_info.get("n_gpus", 1)) - launcher = normalize_launcher(distributed_config.get("launcher"), "kubernetes") + launcher = launcher_for_reporting(distributed_config.get("launcher"), "kubernetes") test_duration = aggregated_record.get("test_duration") or aggregated_record.get("duration", "") run_details = { "model": model_info.get("name", aggregated_record.get("model", "")), @@ -1063,7 +1063,7 @@ def _build_common_info_dict( gpus_per_node = str(nproc_per_node) nnodes_str = str(nnodes) # Launcher: use distributed.launcher when set, otherwise "native" for k8s - launcher = normalize_launcher(distributed_config.get("launcher"), "kubernetes") + launcher = launcher_for_reporting(distributed_config.get("launcher"), "kubernetes") result = { "n_gpus": str(total_gpus), "nnodes": nnodes_str, @@ -1113,7 +1113,7 @@ def _create_multiple_result_row_record( nproc_per_node = int(model_info.get("n_gpus", 1)) # Launcher: use distributed.launcher when set, otherwise "native" for k8s - launcher = normalize_launcher(distributed_config.get("launcher"), "kubernetes") + launcher = launcher_for_reporting(distributed_config.get("launcher"), "kubernetes") result = { "model": item.get("model", model_info.get("name", "")), "n_gpus": str(nnodes * nproc_per_node), diff --git a/src/madengine/deployment/k8s_template_context.py b/src/madengine/deployment/k8s_template_context.py index d7467207..e8528a9d 100644 --- a/src/madengine/deployment/k8s_template_context.py +++ b/src/madengine/deployment/k8s_template_context.py @@ -13,7 +13,7 @@ from pathlib import Path from typing import Any, Dict, List, Optional -from .common import canonicalize_distributed_launcher, configure_multi_node_profiling +from .common import configure_multi_node_profiling from .k8s_names import sanitize_k8s_container_name, sanitize_k8s_label_value from .k8s_secrets import ( CONFIGMAP_MAX_BYTES, @@ -240,7 +240,7 @@ def _prepare_template_context( self.console.print(f"[cyan]Configuring SGLang: {nnodes} nodes × {nproc_per_node} GPUs/node[/cyan]") - elif launcher_type == "megatron": + elif launcher_type == "megatron-lm": if not isinstance(nnodes, int) or nnodes < 1: raise ValueError(f"Invalid nnodes: {nnodes}. Must be positive integer >= 1") if not isinstance(nproc_per_node, int) or nproc_per_node < 1: @@ -341,7 +341,7 @@ def _prepare_template_context( model_args=model_info.get("args", ""), ) - elif canonicalize_distributed_launcher(launcher_type) == "sglang-disagg": + elif launcher_type == "sglang-disagg": if nnodes < 3: raise ValueError( f"SGLang Disaggregated requires minimum 3 nodes " @@ -361,7 +361,7 @@ def _prepare_template_context( model_script=model_info.get("scripts", "run.sh") ) - elif launcher_type == "megatron": + elif launcher_type == "megatron-lm": if nnodes > 1: create_headless_service = True self.console.print(f"[dim]Multi-node Megatron-LM: Creating headless service for pod discovery[/dim]") @@ -477,7 +477,7 @@ def _prepare_template_context( privileged_profiling = bool(ap_prof) _pytorch_native = frozenset( - {"torchrun", "deepspeed", "torchtitan", "megatron", "primus"} + {"torchrun", "deepspeed", "torchtitan", "megatron-lm", "primus"} ) subdomain_val = ( self.service_name diff --git a/src/madengine/deployment/slurm.py b/src/madengine/deployment/slurm.py index a83ebcb1..c630e418 100644 --- a/src/madengine/deployment/slurm.py +++ b/src/madengine/deployment/slurm.py @@ -22,10 +22,9 @@ from .base import BaseDeployment, DeploymentConfig, DeploymentResult, DeploymentStatus, create_jinja_env from .primus_backend import infer_primus_backend_from_model_name, merged_primus_config from .common import ( - canonicalize_distributed_launcher, configure_multi_node_profiling, is_self_managed_launcher, - normalize_launcher, + launcher_for_reporting, ) from .config_loader import ConfigLoader, apply_deployment_config from .slurm_node_selector import SlurmNodeSelector @@ -661,18 +660,9 @@ 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 - - # 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 - # Normalize launcher based on deployment type and validity - launcher_type = normalize_launcher(launcher_type, "slurm") - # Persist the resolved launcher so downstream readers (reporting paths, - # later normalize_launcher calls) see the same value the template used, - # rather than re-deriving from the raw alias and mapping it to "docker". - self.distributed_config["launcher"] = launcher_type + # Extract launcher configuration. Already validated in BaseDeployment.__init__, + # so this is a canonical name or absent. + launcher_type = self.distributed_config.get("launcher") or "torchrun" nnodes = self.distributed_config.get("nnodes", self.nodes) nproc_per_node = self.distributed_config.get("nproc_per_node", resolved_gpus_per_node) @@ -793,7 +783,7 @@ def _generate_launcher_command( return self._generate_sglang_disagg_command(nnodes, nproc_per_node, master_port) elif launcher_type == "deepspeed": return self._generate_deepspeed_command(nnodes, nproc_per_node, master_port) - elif launcher_type == "megatron": + elif launcher_type == "megatron-lm": return self._generate_megatron_command(nnodes, nproc_per_node, master_port) elif launcher_type == "torchtitan": return self._generate_torchtitan_command(nnodes, nproc_per_node, master_port) @@ -1615,7 +1605,7 @@ def _build_perf_entry_from_aggregated( from madengine.utils.config_parser import ConfigParser launcher_type = self.distributed_config.get("launcher", "torchrun") - launcher = normalize_launcher(launcher_type, "slurm") + launcher = launcher_for_reporting(launcher_type, "slurm") run_details = { "model": model_info.get("name", aggregated_record.get("model", "")), @@ -1674,7 +1664,7 @@ def _build_common_info_dict( from madengine.reporting.update_perf_csv import flatten_tags launcher_type = self.distributed_config.get("launcher", "torchrun") - launcher = normalize_launcher(launcher_type, "slurm") + launcher = launcher_for_reporting(launcher_type, "slurm") total_gpus = self.nodes * self.gpus_per_node result = { "n_gpus": str(total_gpus), diff --git a/src/madengine/deployment/templates/kubernetes/job.yaml.j2 b/src/madengine/deployment/templates/kubernetes/job.yaml.j2 index c7599dee..da5e8c92 100644 --- a/src/madengine/deployment/templates/kubernetes/job.yaml.j2 +++ b/src/madengine/deployment/templates/kubernetes/job.yaml.j2 @@ -149,7 +149,7 @@ spec: # - NVIDIA GPUs: Use ONLY CUDA_VISIBLE_DEVICES # Setting both HIP_VISIBLE_DEVICES and CUDA_VISIBLE_DEVICES simultaneously # causes Ray error: "Inconsistent values found" - {% if launcher_type == "vllm" or launcher_type == "sglang" %} + {% if launcher_type in ['vllm', 'sglang', 'sglang-disagg'] %} # Ray-based launchers: Detect GPU vendor and set appropriate variable if command -v rocm-smi &> /dev/null || command -v rocminfo &> /dev/null; then # AMD GPU detected - use HIP_VISIBLE_DEVICES ONLY @@ -181,7 +181,7 @@ spec: export MAD_K8S_JOB=true export MAD_DEPLOYMENT_TYPE=kubernetes - {% if launcher_type == "torchrun" or launcher_type == "deepspeed" or launcher_type == "megatron" or launcher_type == "primus" or launcher_type == "torchtitan" %} + {% if launcher_type in ['torchrun', 'deepspeed', 'megatron-lm', 'primus', 'torchtitan'] %} # {{ launcher_type }} distributed environment (auto-configured from K8s) {% if nnodes > 1 %} # Multi-node {{ launcher_type }} (Indexed Job) diff --git a/src/madengine/deployment/templates/slurm/job.sh.j2 b/src/madengine/deployment/templates/slurm/job.sh.j2 index c692acf3..f3d3f4cf 100644 --- a/src/madengine/deployment/templates/slurm/job.sh.j2 +++ b/src/madengine/deployment/templates/slurm/job.sh.j2 @@ -73,7 +73,7 @@ export GPUS_PER_NODE={{ gpus_per_node }} # IMPORTANT: Ray (vLLM, SGLang) requires HIP_VISIBLE_DEVICES for AMD GPUs # Do NOT set both HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES together GPU_LIST=$(seq -s, 0 $(({{ gpus_per_node }}-1))) -{% if launcher_type == "vllm" or launcher_type == "sglang" %} +{% if launcher_type in ['vllm', 'sglang', 'sglang-disagg'] %} # Ray-based launchers: Detect GPU vendor and set appropriate variable # CRITICAL: Do NOT set both HIP_VISIBLE_DEVICES and CUDA_VISIBLE_DEVICES together if command -v rocm-smi &> /dev/null || command -v rocminfo &> /dev/null; then @@ -444,7 +444,7 @@ echo "==========================================" # Generate launcher-specific command {{ launcher_command }} -{% if launcher_type in ['torchrun', 'deepspeed', 'megatron', 'torchtitan'] %} +{% if launcher_type in ['torchrun', 'deepspeed', 'megatron-lm', 'torchtitan'] %} echo " MAD_MULTI_NODE_RUNNER: ${MAD_MULTI_NODE_RUNNER}" {% endif %} echo "==========================================" @@ -453,7 +453,7 @@ echo "==========================================" echo "Single-node {{ launcher_type|default('torchrun') }} setup" {{ launcher_command }} -{% if launcher_type in ['torchrun', 'deepspeed', 'megatron', 'torchtitan'] %} +{% if launcher_type in ['torchrun', 'deepspeed', 'megatron-lm', 'torchtitan'] %} echo " MAD_MULTI_NODE_RUNNER: ${MAD_MULTI_NODE_RUNNER}" {% endif %} {% endif %} @@ -642,7 +642,7 @@ echo " RANK (node rank): ${RANK}" echo " NODE_RANK: ${NODE_RANK}" echo " NNODES: ${NNODES}" echo " NPROC_PER_NODE: ${GPUS_PER_NODE}" -{% if launcher_type in ['torchrun', 'deepspeed', 'megatron', 'torchtitan'] %} +{% if launcher_type in ['torchrun', 'deepspeed', 'megatron-lm', 'torchtitan'] %} echo " MAD_MULTI_NODE_RUNNER: ${MAD_MULTI_NODE_RUNNER}" {% endif %} if [ "${SLURM_PROCID}" = "0" ]; then diff --git a/src/madengine/execution/container_runner.py b/src/madengine/execution/container_runner.py index d04d23c8..7fcae376 100644 --- a/src/madengine/execution/container_runner.py +++ b/src/madengine/execution/container_runner.py @@ -34,7 +34,6 @@ ) from madengine.reporting.update_perf_super import update_perf_super_json, update_perf_super_csv from madengine.utils.gpu_config import resolve_runtime_gpus -from madengine.deployment.common import canonicalize_distributed_launcher from madengine.utils.config_parser import ConfigParser from madengine.utils.path_utils import scripts_base_dir_from from madengine.utils.run_details import get_build_number, get_pipeline @@ -331,32 +330,8 @@ def create_run_details_dict( except (ValueError, TypeError): total_gpus = resolved_gpu_count - # Extract launcher from multiple sources in priority order: - # 1. additional_context (passed via --additional-context CLI arg) - # 2. model_info distributed config (in models.json) - # 3. MAD_LAUNCHER environment variable - # 4. Default to 'docker' for local deployments - launcher = "" - - # Check additional_context first (highest priority) - if self.additional_context: - distributed_config = self.additional_context.get("distributed", {}) - launcher = distributed_config.get("launcher", "") - if launcher: - print(f"🚀 Launcher from additional_context: {launcher}") - - # Check model_info distributed config - if not launcher and model_info.get("distributed"): - launcher = model_info["distributed"].get("launcher", "") - if launcher: - print(f"🚀 Launcher from model_info: {launcher}") - - # Fallback to environment variable - if not launcher: - launcher = os.environ.get("MAD_LAUNCHER", "") - if launcher: - print(f"🚀 Launcher from MAD_LAUNCHER env: {launcher}") - + launcher = self._resolve_launcher(model_info, announce=True) + # Apply deployment-specific defaults if no launcher specified deployment_type = os.environ.get("MAD_DEPLOYMENT_TYPE", "local") if not launcher: @@ -664,6 +639,45 @@ def get_cpu_arg(self) -> str: cpus = self.context.ctx["docker_cpus"].replace(" ", "") return f"--cpuset-cpus {cpus} " + def _resolve_launcher( + self, model_info: typing.Dict, announce: bool = False + ) -> str: + """Resolve the configured launcher for this run, in source priority order. + + 1. ``additional_context.distributed.launcher`` (the --additional-context CLI arg) + 2. ``model_info.distributed.launcher`` (models.json) + 3. ``MAD_LAUNCHER_TYPE`` / ``MAD_LAUNCHER`` (exported by the SLURM job script) + + Does not validate. Launchers are validated at the config boundary — in + ``cli/validators.py`` for the CLI and ``BaseDeployment.__init__`` for + SLURM/K8s — so anything arriving here is already canonical. Callers on the + local Docker path warn and default rather than fail, which is why this + returns the value as given instead of raising on it. + + Args: + model_info: The model card being run. + announce: Print which source supplied the launcher. + + Returns: + The launcher name, or ``""`` if none is configured. + """ + sources = [ + ( + (self.additional_context or {}).get("distributed", {}).get("launcher", ""), + "additional_context", + ), + ((model_info.get("distributed") or {}).get("launcher", ""), "model_info"), + (os.environ.get("MAD_LAUNCHER_TYPE", ""), "MAD_LAUNCHER_TYPE env"), + (os.environ.get("MAD_LAUNCHER", ""), "MAD_LAUNCHER env"), + ] + for value, label in sources: + if not value: + continue + if announce: + print(f"🚀 Launcher from {label}: {value}") + return value + return "" + def _generate_local_launcher_command(self, launcher_type: str, nproc_per_node: int) -> str: """Generate distributed process launcher command for Docker local deployment. @@ -680,7 +694,7 @@ def _generate_local_launcher_command(self, launcher_type: str, nproc_per_node: i Launcher command string, or empty string for launchers that manage their own process spawning (vllm, sglang). """ - if launcher_type in ("torchrun", "megatron", "megatron-lm", "torchtitan"): + if launcher_type in ("torchrun", "megatron-lm", "torchtitan"): return f"torchrun --standalone --nproc_per_node={nproc_per_node}" elif launcher_type == "deepspeed": return f"deepspeed --num_gpus={nproc_per_node}" @@ -689,9 +703,9 @@ def _generate_local_launcher_command(self, launcher_type: str, nproc_per_node: i else: return f"torchrun --standalone --nproc_per_node={nproc_per_node}" - # Deployment-mode sentinels that normalize_launcher emits for "no real - # launcher". Users may pass these explicitly; defaulting them to torchrun is - # expected, not an error, so they should not trigger an unrecognized warning. + # Deployment-mode sentinels meaning "no real launcher". Users may pass these + # explicitly; defaulting them to torchrun is expected, not an error, so they + # should not trigger an unrecognized warning. _NON_LAUNCHER_SENTINELS = ("docker", "native") def _resolve_local_multi_node_runner_env( @@ -699,28 +713,20 @@ def _resolve_local_multi_node_runner_env( ) -> None: """Set ``docker_env_vars["MAD_MULTI_NODE_RUNNER"]`` for local Docker runs. - No-op if the env var is already set. Resolves launcher from - ``additional_context.distributed.launcher``, then ``model_info.distributed.launcher``, - then ``MAD_LAUNCHER``; falls back to ``torchrun`` for unknown values. - Self-managing launchers (vllm/sglang/sglang-disagg/primus) set the var - to an empty string so downstream scripts under ``set -u`` don't fail. + No-op if the env var is already set. Resolves the launcher via + :meth:`_resolve_launcher`, falling back to ``torchrun``. Self-managing + launchers (vllm/sglang/sglang-disagg/primus) set the var to an empty + string so downstream scripts under ``set -u`` don't fail. """ if "MAD_MULTI_NODE_RUNNER" in self.context.ctx["docker_env_vars"]: return - launcher = "" - if self.additional_context: - launcher = self.additional_context.get("distributed", {}).get("launcher", "") - if not launcher and model_info.get("distributed"): - launcher = model_info["distributed"].get("launcher", "") - if not launcher: - launcher = os.environ.get("MAD_LAUNCHER", "") - canonical_launcher = canonicalize_distributed_launcher(launcher) + launcher = self._resolve_launcher(model_info) valid_local_launchers = ( - "torchrun", "megatron", "megatron-lm", "torchtitan", + "torchrun", "megatron-lm", "torchtitan", "deepspeed", "vllm", "sglang", "sglang-disagg", "primus", ) - if canonical_launcher in valid_local_launchers: - dist_launcher = canonical_launcher + if launcher in valid_local_launchers: + dist_launcher = launcher else: if launcher and launcher not in self._NON_LAUNCHER_SENTINELS: print(f"⚠️ Unrecognized launcher '{launcher}'; " @@ -1379,14 +1385,7 @@ def run_container( # ========== CHECK FOR SELF-MANAGED LAUNCHERS ========== # slurm_multi launchers run scripts directly on the host, # not inside a madengine-managed Docker. The script manages its own containers via srun. - launcher = "" - if self.additional_context: - distributed_config = self.additional_context.get("distributed", {}) - launcher = distributed_config.get("launcher", "") - if not launcher and model_info.get("distributed"): - launcher = model_info["distributed"].get("launcher", "") - if not launcher: - launcher = os.environ.get("MAD_LAUNCHER_TYPE", "") + launcher = self._resolve_launcher(model_info) if is_self_managed_launcher(launcher): self.rich_console.print( f"\n[bold cyan]🖥️ Self-managed launcher (launcher: {launcher})[/bold cyan]" diff --git a/tests/unit/test_container_runner.py b/tests/unit/test_container_runner.py index 702f4e41..8b643c99 100644 --- a/tests/unit/test_container_runner.py +++ b/tests/unit/test_container_runner.py @@ -357,13 +357,11 @@ def test_does_not_override_user_provided_value(self): @pytest.mark.parametrize( "launcher", - ["vllm", "sglang", "sglang-disagg", "sglang_disagg", "primus"], + ["vllm", "sglang", "sglang-disagg", "primus"], ) def test_self_managed_launchers_set_empty_string(self, launcher): """Self-managing launchers set the var to "" (defined but empty), - so downstream scripts under set -u don't fail referencing it. - Covers the ``sglang_disagg`` underscore alias to lock in the - canonicalize_distributed_launcher() routing.""" + so downstream scripts under set -u don't fail referencing it.""" runner = self._runner( additional_context={"distributed": {"launcher": launcher}}, ) diff --git a/tests/unit/test_deployment.py b/tests/unit/test_deployment.py index e054554e..66d52309 100644 --- a/tests/unit/test_deployment.py +++ b/tests/unit/test_deployment.py @@ -6,14 +6,16 @@ import pytest +from madengine.core.errors import ConfigurationError from madengine.deployment.base import BaseDeployment, DeploymentConfig, create_jinja_env from madengine.deployment.common import ( VALID_LAUNCHERS, - canonicalize_distributed_launcher, configure_multi_node_profiling, + is_per_replica_launcher, is_rocprofv3_available, - normalize_launcher, + launcher_for_reporting, tools_include_rocprof_family, + validate_launcher, ) @@ -53,40 +55,102 @@ def test_contains_expected_launchers(self): assert "sglang-disagg" in VALID_LAUNCHERS -class TestNormalizeLauncher: - """normalize_launcher behavior.""" +class TestLauncherForReporting: + """launcher_for_reporting supplies a sentinel only when nothing is configured.""" def test_valid_launcher_passthrough(self): for lt in VALID_LAUNCHERS: - assert normalize_launcher(lt, "kubernetes") == lt - assert normalize_launcher(lt, "slurm") == lt - assert normalize_launcher(lt, "local") == lt + assert launcher_for_reporting(lt, "kubernetes") == lt + assert launcher_for_reporting(lt, "slurm") == lt + assert launcher_for_reporting(lt, "local") == lt - @pytest.mark.parametrize("launcher", [None, "", "invalid"]) - def test_invalid_or_missing_launcher_kubernetes_returns_native(self, launcher): - assert normalize_launcher(launcher, "kubernetes") == "native" + @pytest.mark.parametrize("launcher", [None, ""]) + def test_missing_launcher_kubernetes_returns_native(self, launcher): + assert launcher_for_reporting(launcher, "kubernetes") == "native" @pytest.mark.parametrize("deployment", ["slurm", "local", "unknown"]) - def test_invalid_or_missing_launcher_non_k8s_returns_docker(self, deployment): - assert normalize_launcher(None, deployment) == "docker" - - -class TestCanonicalizeDistributedLauncher: - """canonicalize_distributed_launcher resolves alternate spellings.""" - - def test_underscore_form_maps_to_canonical_hyphen_form(self): - assert canonicalize_distributed_launcher("sglang_disagg") == "sglang-disagg" - - def test_canonical_form_passthrough(self): - assert canonicalize_distributed_launcher("sglang-disagg") == "sglang-disagg" - assert canonicalize_distributed_launcher("torchrun") == "torchrun" - - @pytest.mark.parametrize("value", [None, ""]) - def test_empty_returned_unchanged(self, value): - assert canonicalize_distributed_launcher(value) == value - - def test_unknown_value_returned_unchanged(self): - assert canonicalize_distributed_launcher("bogus_launcher") == "bogus_launcher" + def test_missing_launcher_non_k8s_returns_docker(self, deployment): + assert launcher_for_reporting(None, deployment) == "docker" + + def test_never_raises_on_a_bogus_value(self): + """Reporting runs after a successful job; raising here would drop its metrics.""" + assert launcher_for_reporting("bogus", "slurm") == "bogus" + + +class TestValidateLauncher: + """validate_launcher accepts one spelling per launcher and rejects the rest.""" + + @pytest.mark.parametrize("launcher", VALID_LAUNCHERS) + def test_every_valid_launcher_is_accepted(self, launcher): + assert validate_launcher(launcher, source="test") == launcher + + @pytest.mark.parametrize("value", [None, "", " "]) + def test_empty_means_no_launcher_configured(self, value): + assert validate_launcher(value, source="test") is None + + @pytest.mark.parametrize("value", [0, False, []]) + def test_other_falsy_values_are_misconfiguration_not_absence(self, value): + """`if not launcher` would have waved these through as 'nothing configured'.""" + with pytest.raises(ConfigurationError): + validate_launcher(value, source="test") + + def test_documented_hyphen_alias_for_slurm_multi(self): + """docs/launchers.md advertises slurm-multi; that promise is kept.""" + assert validate_launcher("slurm-multi", source="test") == "slurm_multi" + + @pytest.mark.parametrize("value,canonical", [ + ("Torchrun", "torchrun"), + (" torchrun ", "torchrun"), + ("MEGATRON-LM", "megatron-lm"), + ]) + def test_case_and_whitespace_are_normalized(self, value, canonical): + assert validate_launcher(value, source="test") == canonical + + @pytest.mark.parametrize("launcher", ["docker", "native"]) + def test_reporting_sentinels_are_rejected_as_config(self, launcher): + """They are perf.csv output values with no dispatch arm; accepting one at a + config boundary would reinstate the silent single-process run.""" + with pytest.raises(ConfigurationError) as exc_info: + validate_launcher(launcher, source="additional_context") + assert "omit the launcher key" in " ".join(exc_info.value.suggestions) + + @pytest.mark.parametrize("bad,canonical", [ + ("megatron", "megatron-lm"), + ("megatron_lm", "megatron-lm"), + ("sglang_disagg", "sglang-disagg"), + ]) + def test_rejected_spellings_name_the_canonical_one(self, bad, canonical): + """The did-you-mean is the contract: it fixes the user's config.""" + with pytest.raises(ConfigurationError) as exc_info: + validate_launcher(bad, source="additional_context") + rendered = str(exc_info.value) + " ".join(exc_info.value.suggestions) + assert canonical in rendered + + def test_error_names_the_source(self): + with pytest.raises(ConfigurationError) as exc_info: + validate_launcher("nonsense", source="models.json") + assert "models.json" in str(exc_info.value) + + def test_non_string_is_rejected(self): + with pytest.raises(ConfigurationError): + validate_launcher(["torchrun"], source="test") + + +class TestIsPerReplicaLauncher: + """Which launchers report a node-local metric per replica.""" + + @pytest.mark.parametrize("launcher", ["vllm", "sglang"]) + def test_data_parallel_launchers_are_per_replica(self, launcher): + assert is_per_replica_launcher(launcher) is True + + def test_sglang_disagg_is_not_per_replica(self): + """Ray-based, but its nodes are proxy/prefill/decode roles serving one + endpoint — scaling the proxy's number by nnodes would inflate it.""" + assert is_per_replica_launcher("sglang-disagg") is False + + @pytest.mark.parametrize("launcher", ["torchrun", "megatron-lm", "docker", None, ""]) + def test_everything_else_aggregates_normally(self, launcher): + assert is_per_replica_launcher(launcher) is False class TestToolsIncludeRocprofFamily: @@ -226,6 +290,7 @@ def cleanup(self, deployment_id): pass def _make_deployment(): cfg = MagicMock(spec=DeploymentConfig) cfg.manifest_file = None + cfg.additional_context = {} with patch.object(BaseDeployment, "_load_manifest", return_value={}): return _ConcreteDeployment(cfg) diff --git a/tests/unit/test_launcher_dispatch.py b/tests/unit/test_launcher_dispatch.py new file mode 100644 index 00000000..d9bef809 --- /dev/null +++ b/tests/unit/test_launcher_dispatch.py @@ -0,0 +1,333 @@ +#!/usr/bin/env python3 +""" +Unit tests locking in that every launcher in ``VALID_LAUNCHERS`` actually reaches +a dispatch arm on the backends that claim to support it. + +The defect these exist to prevent: ``megatron-lm`` — the documented spelling, used +in all four shipped example configs — reached no dispatch arm on *either* backend. +Both chains compared against the bare literal ``"megatron"``, so the documented name +fell through to the "unknown launcher" default. SLURM printed a warning and ran with +no distributed setup; Kubernetes produced no launcher command, no headless service, +and no rank env vars. Both reported SUCCESS with a wrong benchmark number. + +Nothing caught it because no test asserted "backend X has an arm for launcher Y". +That assertion is what this file is. + +Copyright (c) Advanced Micro Devices, Inc. All rights reserved. +""" + +import json +from pathlib import Path +from unittest.mock import patch + +import pytest + +from madengine.deployment.base import DeploymentConfig +from madengine.deployment.common import VALID_LAUNCHERS +from madengine.deployment.slurm import SlurmDeployment + + +# slurm_multi bypasses the templated launcher chain entirely: it runs the model's own +# .slurm script on the head node and orchestrates containers via srun itself. It is an +# escape hatch, not a peer of the templated launchers, so it has no dispatch arm on +# either backend. Writing the exception down is the point — an *unexplained* absence +# is exactly what the megatron bug looked like. +SELF_MANAGED = {"slurm_multi"} + +TEMPLATED_LAUNCHERS = [lt for lt in VALID_LAUNCHERS if lt not in SELF_MANAGED] + +MODEL_ENTRY = { + "name": "dummy_multinode", + "url": "", + "dockerfile": "docker/dummy", + "scripts": "scripts/dummy/run.sh", + "n_gpus": "8", + "owner": "mad.support@amd.com", + "training_precision": "", + "tags": ["pyt", "training"], + "timeout": -1, + "args": "", +} + +MANIFEST_CONTEXT = { + "docker_env_vars": {}, + "docker_mounts": {}, + "docker_build_arg": {}, + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "docker_gpus": "all", +} + + +def _write_manifest(tmp_path: Path, distributed: dict) -> Path: + manifest = { + "built_images": {"dummy-image": {"docker_image": "dummy:latest"}}, + "built_models": {"dummy-image": MODEL_ENTRY}, + "context": MANIFEST_CONTEXT, + "deployment_config": {"distributed": distributed}, + } + manifest_path = tmp_path / "build_manifest.json" + manifest_path.write_text(json.dumps(manifest)) + return manifest_path + + +# --------------------------------------------------------------------------- +# SLURM + +def _slurm_deployment(tmp_path: Path, launcher: str, nnodes: int = 3) -> SlurmDeployment: + distributed = { + "launcher": launcher, + "nnodes": nnodes, + "nproc_per_node": 8, + "backend": "nccl", + "port": 29500, + } + cfg = DeploymentConfig( + target="slurm", + manifest_file=str(_write_manifest(tmp_path, distributed)), + additional_context={ + "deploy": "slurm", + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "slurm": { + "partition": "test-partition", + "nodes": nnodes, + "gpus_per_node": 8, + "time": "01:00:00", + "output_dir": str(tmp_path / "slurm_output"), + }, + "distributed": distributed, + }, + ) + return SlurmDeployment(cfg) + + +class TestSlurmLauncherDispatch: + """Every templated launcher must reach its own arm, not the unknown-launcher default.""" + + @pytest.mark.parametrize("launcher", TEMPLATED_LAUNCHERS) + def test_every_valid_launcher_has_an_arm(self, tmp_path, launcher): + deployment = _slurm_deployment(tmp_path, launcher) + with patch.object( + SlurmDeployment, "_generate_basic_env_command", side_effect=AssertionError( + f"launcher '{launcher}' fell through to the unknown-launcher default" + ) + ): + command = deployment._generate_launcher_command( + launcher_type=launcher, + nnodes=3, + nproc_per_node=8, + master_port=29500, + model_name=MODEL_ENTRY["name"], + ) + assert command is not None + + def test_the_documented_megatron_spelling_reaches_megatron(self, tmp_path): + """The original bug: this returned the basic-env fallback, silently.""" + deployment = _slurm_deployment(tmp_path, "megatron-lm") + command = deployment._generate_launcher_command( + launcher_type="megatron-lm", + nnodes=3, + nproc_per_node=8, + master_port=29500, + ) + assert "MAD_MULTI_NODE_RUNNER" in command + + def test_an_unknown_launcher_still_reaches_the_default(self, tmp_path): + """Sanity: the fallback arm is real, so the test above proves something.""" + deployment = _slurm_deployment(tmp_path, "torchrun") + with patch.object( + SlurmDeployment, "_generate_basic_env_command", return_value="basic" + ) as basic: + deployment._generate_launcher_command( + launcher_type="not-a-launcher", + nnodes=3, + nproc_per_node=8, + master_port=29500, + ) + basic.assert_called_once() + + +class TestSlurmRayGpuVisibility: + """Ray-based launchers must not export HIP+ROCR+CUDA together. + + Ray fails with "Inconsistent values found" when more than one visibility + variable is set. sglang-disagg was missing from the guard, so it took the + else-branch that exports all three. + """ + + @staticmethod + def _render(deployment: SlurmDeployment) -> str: + context = deployment._prepare_template_context(MODEL_ENTRY) + return deployment.jinja_env.get_template("job.sh.j2").render(**context) + + @pytest.mark.parametrize("launcher", ["vllm", "sglang", "sglang-disagg"]) + def test_ray_launchers_take_the_single_variable_branch(self, tmp_path, launcher): + script = self._render(_slurm_deployment(tmp_path, launcher)) + assert "unset RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES" in script + assert 'unset CUDA_VISIBLE_DEVICES # Unset to avoid "Inconsistent values" error' in script + + def test_non_ray_launchers_do_not(self, tmp_path): + script = self._render(_slurm_deployment(tmp_path, "torchrun")) + assert "unset RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES" not in script + + +# --------------------------------------------------------------------------- +# Kubernetes + +def _k8s_deployment(tmp_path: Path, launcher: str, nnodes: int = 3): + """KubernetesDeployment over a minimal manifest, with cluster access stubbed. + + ``__init__`` loads kubeconfig and builds API clients, both of which raise + without a cluster. Patch the real class rather than a mixin stub, so the + dispatch chain under test is the one that actually ships. + """ + from madengine.deployment import kubernetes as k8s_module + + distributed = { + "launcher": launcher, + "nnodes": nnodes, + "nproc_per_node": 8, + "backend": "nccl", + } + cfg = DeploymentConfig( + target="kubernetes", + manifest_file=str(_write_manifest(tmp_path, distributed)), + additional_context={ + "deploy": "kubernetes", + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "k8s": {"namespace": "default"}, + "distributed": distributed, + }, + ) + with patch.object(k8s_module, "k8s_config"), patch.object(k8s_module, "client"): + deployment = k8s_module.KubernetesDeployment(cfg) + # Set in deploy(), which the render path below bypasses. + deployment.job_name = "test-job" + deployment.job_label = "test-job" + deployment.service_name = "test-svc" + deployment.main_container_name = "test-svc" + deployment.configmap_name = "test-cm" + return deployment + + +IMAGE_INFO = { + "docker_image": "dummy:latest", + "registry_image": "registry.local/dummy:latest", +} + + +class TestKubernetesLauncherDispatch: + """The K8s dispatch chain must produce a launcher command for every valid launcher.""" + + @pytest.mark.parametrize("launcher", TEMPLATED_LAUNCHERS) + def test_every_valid_launcher_produces_a_launcher_command(self, tmp_path, launcher): + deployment = _k8s_deployment(tmp_path, launcher) + context = deployment._prepare_template_context(MODEL_ENTRY, IMAGE_INFO) + assert context["launcher_type"] == launcher + assert context["launcher_command"], ( + f"launcher '{launcher}' produced no launcher command — it reached no " + f"dispatch arm in the K8s chain" + ) + + def test_the_documented_megatron_spelling_reaches_megatron(self, tmp_path): + """The original bug: this produced no launcher command and no headless service.""" + deployment = _k8s_deployment(tmp_path, "megatron-lm") + context = deployment._prepare_template_context(MODEL_ENTRY, IMAGE_INFO) + assert context["launcher_command"] + assert context["create_headless_service"] is True + + @pytest.mark.parametrize( + "launcher", ["torchrun", "deepspeed", "torchtitan", "megatron-lm", "primus"] + ) + def test_pytorch_native_launchers_get_a_pod_subdomain(self, tmp_path, launcher): + """Multi-node PyTorch-native launchers need DNS for rank discovery.""" + deployment = _k8s_deployment(tmp_path, launcher) + context = deployment._prepare_template_context(MODEL_ENTRY, IMAGE_INFO) + assert context["subdomain"] == deployment.service_name + + +# --------------------------------------------------------------------------- +# Cross-backend parity + +def test_slurm_and_kubernetes_support_the_same_templated_launchers(tmp_path): + """A launcher valid on one backend must be valid on the other, or be a known exception.""" + for launcher in TEMPLATED_LAUNCHERS: + slurm = _slurm_deployment(tmp_path, launcher) + with patch.object( + SlurmDeployment, "_generate_basic_env_command", + side_effect=AssertionError(f"SLURM has no arm for '{launcher}'"), + ): + slurm._generate_launcher_command( + launcher_type=launcher, nnodes=3, nproc_per_node=8, master_port=29500, + model_name=MODEL_ENTRY["name"], + ) + k8s = _k8s_deployment(tmp_path, launcher) + context = k8s._prepare_template_context(MODEL_ENTRY, IMAGE_INFO) + assert context["launcher_command"], f"K8s has no arm for '{launcher}'" + + +# --------------------------------------------------------------------------- +# The config-boundary chokepoint + +class TestBaseDeploymentValidatesLaunchers: + """Launchers are validated in ``BaseDeployment.__init__``, deliberately. + + ``execute()`` catches bare ``Exception`` and returns a FAILED result without + re-raising, so a ConfigurationError raised any later never reaches the CLI's + handler. ``__init__`` runs under DeploymentFactory.create(), which re-raises. + """ + + def test_bad_launcher_in_additional_context_raises(self, tmp_path): + from madengine.core.errors import ConfigurationError + + with pytest.raises(ConfigurationError) as exc_info: + _slurm_deployment(tmp_path, "megatron") + assert "megatron-lm" in " ".join(exc_info.value.suggestions) + + def test_bad_launcher_in_the_manifest_raises(self, tmp_path): + """The manifest is a separate source; the CLI validator never sees it.""" + from madengine.core.errors import ConfigurationError + + manifest_path = _write_manifest(tmp_path, {"launcher": "megatron"}) + cfg = DeploymentConfig( + target="slurm", + manifest_file=str(manifest_path), + additional_context={ + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "slurm": {"partition": "p", "output_dir": str(tmp_path)}, + }, + ) + with pytest.raises(ConfigurationError): + SlurmDeployment(cfg) + + def test_bad_launcher_on_a_model_card_raises(self, tmp_path): + from madengine.core.errors import ConfigurationError + + manifest = { + "built_images": {"dummy-image": {"docker_image": "dummy:latest"}}, + "built_models": { + "dummy-image": {**MODEL_ENTRY, "distributed": {"launcher": "megatron"}} + }, + "context": MANIFEST_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={ + "gpu_vendor": "AMD", + "guest_os": "UBUNTU", + "slurm": {"partition": "p", "output_dir": str(tmp_path)}, + }, + ) + with pytest.raises(ConfigurationError) as exc_info: + SlurmDeployment(cfg) + assert "dummy-image" in str(exc_info.value) + + def test_the_documented_alias_is_canonicalized_in_place(self, tmp_path): + deployment = _slurm_deployment(tmp_path, "slurm-multi") + assert deployment.config.additional_context["distributed"]["launcher"] == "slurm_multi" diff --git a/tests/unit/test_slurm_multi.py b/tests/unit/test_slurm_multi.py index ae1fac97..7ddd7149 100644 --- a/tests/unit/test_slurm_multi.py +++ b/tests/unit/test_slurm_multi.py @@ -25,10 +25,14 @@ import pytest -from madengine.deployment.common import VALID_LAUNCHERS, is_self_managed_launcher, normalize_launcher +from madengine.core.errors import ConfigurationError +from madengine.deployment.common import ( + VALID_LAUNCHERS, + is_self_managed_launcher, + validate_launcher, +) from madengine.deployment.base import DeploymentConfig from madengine.deployment.slurm import SlurmDeployment -from madengine.core.errors import ConfigurationError # --------------------------------------------------------------------------- @@ -42,7 +46,7 @@ def test_slurm_multi_in_valid_launchers(self): @pytest.mark.parametrize("launcher,expected", [ ("slurm_multi", True), - ("slurm-multi", True), # hyphen alias normalized via normalize_launcher + ("slurm-multi", True), # documented hyphen alias (docs/launchers.md) ("torchrun", False), ("vllm", False), ("sglang-disagg", False), @@ -61,14 +65,16 @@ class TestNormalizeSlurmMultiAliases: """slurm-multi (hyphen) normalizes to slurm_multi.""" def test_canonical(self): - assert normalize_launcher("slurm_multi", "slurm") == "slurm_multi" + assert validate_launcher("slurm_multi", source="test") == "slurm_multi" def test_hyphen_alias(self): - assert normalize_launcher("slurm-multi", "slurm") == "slurm_multi" + assert validate_launcher("slurm-multi", source="test") == "slurm_multi" - def test_unknown_falls_through_to_default(self): - # Sanity: unrelated value still returns docker for slurm - assert normalize_launcher("totally-bogus", "slurm") == "docker" + def test_unknown_is_rejected_not_defaulted(self): + # A silently-defaulted launcher runs the model as a plain single-process + # job and still reports SUCCESS, so this must fail loudly instead. + with pytest.raises(ConfigurationError): + validate_launcher("totally-bogus", source="test") # --------------------------------------------------------------------------- diff --git a/tests/unit/test_validators.py b/tests/unit/test_validators.py index 6a78c52d..83beb773 100644 --- a/tests/unit/test_validators.py +++ b/tests/unit/test_validators.py @@ -311,3 +311,71 @@ def test_validate_additional_context_log_error_patterns_rejects_non_string_eleme with pytest.raises(typer.Exit) as exc_info: validate_additional_context(additional_context=bad) assert exc_info.value.exit_code == ExitCode.INVALID_ARGS + + +class TestValidateLauncherContext: + """Launcher validation at the CLI boundary. + + A launcher madengine did not recognize used to run the model as a plain + single-process job and still report SUCCESS, so the benchmark number was + wrong with nothing to indicate it. These lock in loud failure instead. + """ + + @staticmethod + def _context(**extra): + return json.dumps({"gpu_vendor": "AMD", "guest_os": "UBUNTU", **extra}) + + @pytest.mark.parametrize("launcher", ["torchrun", "megatron-lm", "sglang-disagg"]) + def test_valid_launcher_is_accepted(self, launcher): + result = validate_additional_context( + additional_context=self._context(distributed={"launcher": launcher}) + ) + assert result["distributed"]["launcher"] == launcher + + @pytest.mark.parametrize("bad", ["megatron", "megatron_lm", "sglang_disagg"]) + def test_rejected_spellings_exit_with_invalid_args(self, bad, capsys): + with pytest.raises(typer.Exit) as exc_info: + validate_additional_context( + additional_context=self._context(distributed={"launcher": bad}) + ) + assert exc_info.value.exit_code == ExitCode.INVALID_ARGS + + @pytest.mark.parametrize("sentinel", ["docker", "native"]) + def test_reporting_sentinels_are_rejected_at_the_cli_boundary(self, sentinel): + """These are perf.csv output values with no dispatch arm. Accepting one + here would let it fall through to the unknown-launcher default at + dispatch time — the silent single-process run this validation exists + to prevent.""" + with pytest.raises(typer.Exit) as exc_info: + validate_additional_context( + additional_context=self._context(distributed={"launcher": sentinel}) + ) + assert exc_info.value.exit_code == ExitCode.INVALID_ARGS + + def test_documented_slurm_multi_alias_is_accepted_and_canonicalized(self): + result = validate_additional_context( + additional_context=self._context(distributed={"launcher": "slurm-multi"}) + ) + assert result["distributed"]["launcher"] == "slurm_multi" + + def test_case_is_folded(self): + result = validate_additional_context( + additional_context=self._context(distributed={"launcher": "Torchrun"}) + ) + assert result["distributed"]["launcher"] == "torchrun" + + def test_launcher_type_key_is_validated_too(self): + with pytest.raises(typer.Exit) as exc_info: + validate_additional_context( + additional_context=self._context(launcher={"type": "megatron"}) + ) + assert exc_info.value.exit_code == ExitCode.INVALID_ARGS + + def test_bare_string_launcher_is_a_clean_error_not_an_attributeerror(self): + """{"launcher": "torchrun"} is a natural mistake; it used to crash deep in + the K8s template context with AttributeError.""" + with pytest.raises(typer.Exit) as exc_info: + validate_additional_context( + additional_context=self._context(launcher="torchrun") + ) + assert exc_info.value.exit_code == ExitCode.INVALID_ARGS