diff --git a/docs/deployment.md b/docs/deployment.md index fa03e7f5..2881cee4 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -138,7 +138,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/examples/k8s-configs/README.md b/examples/k8s-configs/README.md index 40843ab3..55fb1632 100644 --- a/examples/k8s-configs/README.md +++ b/examples/k8s-configs/README.md @@ -186,7 +186,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 e89f016f..7559aa10 100644 --- a/examples/slurm-configs/README.md +++ b/examples/slurm-configs/README.md @@ -404,7 +404,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) @@ -418,12 +418,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 69e1367e..44b55c7a 100644 --- a/src/madengine/deployment/base.py +++ b/src/madengine/deployment/base.py @@ -124,6 +124,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..4434668e 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,88 @@ "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(), never typed by a user, and must pass validation +# unchanged so a round-trip through a report does not raise. +_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. + + Empty values mean "no launcher configured" and return None. 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. + + 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. -def canonicalize_distributed_launcher(launcher: Optional[str]) -> Optional[str]: - """Normalize alternate launcher spellings to their canonical form. + Returns: + The canonical launcher name, or None when nothing was configured. - Resolves aliases (e.g. ``sglang_disagg`` → ``sglang-disagg``). Unknown or - empty values are returned unchanged; callers do their own validation. + Raises: + ConfigurationError: If the value is not a recognized launcher. """ if not launcher: - return launcher - return _LAUNCHER_ALIASES.get(launcher, launcher) + 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 normalized in _LAUNCHER_SENTINELS: + return normalized + normalized = _DOCUMENTED_ALIASES.get(normalized, normalized) + if normalized in VALID_LAUNCHERS: + return normalized + + suggestions = [] + 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,40 +145,6 @@ 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. - - 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" - - _SELF_MANAGED_LAUNCHERS: frozenset = frozenset({"slurm_multi"}) @@ -119,10 +155,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..6bf66f1b 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 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,7 +129,7 @@ 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"] @@ -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 e38b251a..17daa8a4 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, @@ -238,7 +238,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: @@ -339,7 +339,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 " @@ -359,7 +359,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]") @@ -475,7 +475,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 088a3fb2..13528be3 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 @@ -631,18 +630,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) @@ -761,7 +751,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) @@ -1583,7 +1573,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", "")), @@ -1642,7 +1632,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 320d049f..8829bc8d 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 4cdddba0..f18674a8 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 eab4af7d..9c10224a 100644 --- a/src/madengine/execution/container_runner.py +++ b/src/madengine/execution/container_runner.py @@ -33,7 +33,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 @@ -329,32 +328,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: @@ -662,6 +637,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. @@ -678,7 +692,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}" @@ -687,9 +701,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( @@ -697,28 +711,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}'; " @@ -1376,14 +1382,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 aae79321..e27c812f 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..8d05936e 100644 --- a/tests/unit/test_deployment.py +++ b/tests/unit/test_deployment.py @@ -6,14 +6,15 @@ 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_rocprofv3_available, - normalize_launcher, + launcher_for_reporting, tools_include_rocprof_family, + validate_launcher, ) @@ -53,40 +54,75 @@ 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" + 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 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" +class TestValidateLauncher: + """validate_launcher accepts one spelling per launcher and rejects the rest.""" - def test_canonical_form_passthrough(self): - assert canonicalize_distributed_launcher("sglang-disagg") == "sglang-disagg" - assert canonicalize_distributed_launcher("torchrun") == "torchrun" + @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_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_empty_means_no_launcher_configured(self, value): + assert validate_launcher(value, source="test") is None + + 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_survive_a_round_trip(self, launcher): + assert validate_launcher(launcher, source="test") == launcher + + @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 TestToolsIncludeRocprofFamily: @@ -226,6 +262,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 de5c2eae..d57355f0 100644 --- a/tests/unit/test_slurm_multi.py +++ b/tests/unit/test_slurm_multi.py @@ -25,7 +25,12 @@ 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 @@ -41,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), @@ -60,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..48c03c14 100644 --- a/tests/unit/test_validators.py +++ b/tests/unit/test_validators.py @@ -311,3 +311,59 @@ 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 + + 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