diff --git a/README.md b/README.md
index 9ff81e42..ed763456 100644
--- a/README.md
+++ b/README.md
@@ -9,45 +9,21 @@
[](https://python.org)
[](https://github.com/ROCm/madengine/actions)
[](https://github.com/psf/black)
-[](CHANGELOG.md)
+[](CHANGELOG.md)
[](LICENSE)
> **AI model automation and benchmarking platform for local and distributed execution**
-madengine is a modern CLI tool for running Large Language Models (LLMs) and Deep Learning models across local and distributed environments. Built for the [MAD (Model Automation and Dashboarding)](https://github.com/ROCm/MAD) ecosystem, it provides seamless execution from single GPUs to multi-node clusters.
-
-## π Table of Contents
-
-- [Key Features](#-key-features)
-- [Quick Start](#-quick-start)
-- [Commands](#-commands)
-- [Documentation](#-documentation)
-- [Architecture](#-architecture)
-- [Feature Matrix](#-feature-matrix)
-- [Usage Examples](#-usage-examples)
-- [Model Discovery](#-model-discovery)
-- [Performance Profiling](#-performance-profiling)
-- [Reporting and Database](#-reporting-and-database)
-- [Installation](#-installation)
-- [Tips & Best Practices](#-tips--best-practices)
- - [Log error pattern scan](#log-error-pattern-scan)
- - [Exit codes and CI](#exit-codes-and-ci)
-- [Contributing](#-contributing)
-- [License](#-license)
-- [Links & Resources](#-links--resources)
+madengine is a modern CLI tool for running Large Language Models (LLMs) and Deep Learning models across local and distributed environments. Built for the [MAD (Model Automation and Dashboarding)](https://github.com/ROCm/MAD) ecosystem, it provides seamless execution from single GPUs to multi-node clusters β with the same command working locally, on Kubernetes, and on SLURM.
## β¨ Key Features
-- **π Modern CLI** - Rich terminal output with Typer and Rich
-- **π― Simple Deployment** - Run locally or deploy to Kubernetes/SLURM via configuration
-- **π§ Distributed Launchers** - Full support for torchrun, DeepSpeed, Megatron-LM, TorchTitan, Primus, vLLM, SGLang
-- **π³ Container-Native** - Docker-based execution with GPU support (ROCm, CUDA)
-- **π ROCm Path** - Auto-detect **host** ROCm root (override with top-level `MAD_ROCM_PATH`); in-container `ROCM_PATH` is set independently via `docker_env_vars.MAD_ROCM_PATH` and resolved at Docker run (image OCI env + in-image probe, not host mirroring) β see [Configuration](docs/configuration.md#rocm-path-run-only)
-- **π Performance Tools** - Integrated profiling with rocprof/rocprofv3, [rocm-trace-lite](https://github.com/sunway513/rocm-trace-lite) (RTL), rocblas, MIOpen, RCCL tracing
-- **π― ROCprofv3 Profiles** - 8 pre-configured profiles for compute/memory/communication bottleneck analysis
-- **π Environment Validation** - TheRock ROCm detection and validation tools
-- **βοΈ Intelligent Defaults** - Minimal K8s configs with automatic preset application
-- **π Configurable log scan** - Optional `--additional-context` keys to disable or tune post-run log substring checks (see [Log error pattern scan](#log-error-pattern-scan))
+- **π Modern CLI** β Rich terminal output with Typer and Rich
+- **π― Simple Deployment** β Run locally or deploy to Kubernetes/SLURM by adding a config key; no code changes
+- **π§ Distributed Launchers** β torchrun, DeepSpeed, Megatron-LM, TorchTitan, Primus, vLLM, SGLang
+- **π³ Container-Native** β Docker-based execution with GPU support (ROCm, CUDA)
+- **π Performance Tools** β Integrated profiling with rocprof/rocprofv3, [rocm-trace-lite](https://github.com/sunway513/rocm-trace-lite), rocBLAS/MIOpen/RCCL tracing β see [Profiling](docs/profiling.md)
+- **βοΈ Intelligent Defaults** β Minimal configs auto-merged with presets; host/in-container ROCm path auto-detected β see [Configuration](docs/configuration.md#rocm-path-run-only)
## π Quick Start
@@ -55,504 +31,186 @@ madengine is a modern CLI tool for running Large Language Models (LLMs) and Deep
# Install madengine
pip install git+https://github.com/ROCm/madengine.git
-# Clone MAD package (required for models)
+# Clone the MAD package (required for models)
git clone https://github.com/ROCm/MAD.git && cd MAD
# Discover available models
madengine discover --tags dummy
-# Run locally (full workflow: discover/build/run as configured by the model)
+# Run locally (discover β build β run, as configured by the model)
madengine run --tags dummy
+```
+
+> **Note:** For build operations `gpu_vendor` defaults to `AMD` and `guest_os` to `UBUNTU`. For non-AMD/Ubuntu environments, set them explicitly, e.g. `--additional-context '{"gpu_vendor": "NVIDIA", "guest_os": "CENTOS"}'`.
+
+**Results:** Performance data is written to `perf.csv` (and optionally `perf_entry.csv`), created automatically if missing. Failed runs are recorded with status `FAILURE` so every attempted model appears. See [Exit Codes](docs/cli-reference.md#exit-codes) for CI usage. If ROCm isn't auto-detected, set `MAD_ROCM_PATH` β see [Configuration](docs/configuration.md#rocm-path-run-only).
-# Or with explicit configuration
-madengine run --tags dummy \
- --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'
+## ποΈ Architecture
+
+madengine is organized in layers: the CLI drives orchestrators that discover and build models, then hand off to a local or distributed execution target, which runs the model under the appropriate launcher and emits performance data for reporting.
+
+```mermaid
+flowchart TB
+ subgraph CLI["CLI Layer β Typer + Rich"]
+ C1[discover]
+ C2[build]
+ C3[run]
+ C4[report]
+ C5[database]
+ end
+
+ subgraph ORC["Orchestration Layer"]
+ O1[DiscoverModels]
+ O2[BuildOrchestrator]
+ O3[RunOrchestrator]
+ MAN[(build_manifest.json)]
+ end
+
+ subgraph EXEC["Execution / Deployment Layer"]
+ E1[ContainerRunner
local Docker]
+ E2[DeploymentFactory]
+ K8S[Kubernetes Jobs]
+ SLURM[SLURM Jobs]
+ end
+
+ subgraph LAUNCH["Launcher Layer"]
+ T[Train: torchrun Β· DeepSpeed
Megatron-LM Β· TorchTitan Β· Primus]
+ I[Infer: vLLM Β· SGLang Β· SGLang Disagg]
+ end
+
+ OUT[(perf.csv / JSON)]
+
+ C1 --> O1
+ C2 --> O2
+ C3 --> O3
+ O2 --> MAN --> O3
+ O1 --> O2
+ O3 --> E1
+ O3 --> E2
+ E2 --> K8S
+ E2 --> SLURM
+ E1 --> LAUNCH
+ K8S --> LAUNCH
+ SLURM --> LAUNCH
+ LAUNCH --> OUT
+ OUT --> C4
+ OUT --> C5
```
-> **Note**: For build operations, `gpu_vendor` defaults to `AMD` and `guest_os` defaults to `UBUNTU` if not specified. For production deployments or non-AMD/Ubuntu environments, explicitly specify these values.
+1. **CLI Layer** β five commands: `discover`, `build`, `run`, `report`, `database`
+2. **Orchestration** β `DiscoverModels` finds models; `BuildOrchestrator` builds images and writes `build_manifest.json`; `RunOrchestrator` reads/triggers the build and infers the target
+3. **Execution / Deployment** β local `ContainerRunner`, or `DeploymentFactory` β Kubernetes / SLURM
+4. **Launchers** β distributed training and inference frameworks
+5. **Output & Post-Processing** β `perf.csv`/JSON results β `report` (HTML/email) and `database` (MongoDB)
-If auto-detection does not find your **host** ROCm root, set top-level `MAD_ROCM_PATH` in `--additional-context`. For a different ROCm root **inside the container**, set `docker_env_vars.MAD_ROCM_PATH` in additional context. If you omit it, madengine derives in-container `ROCM_PATH` when running Docker (from the image's baked-in env, then an in-container probe, then `/opt/rocm` β it does **not** copy the host path). You can also set `ROCM_PATH` / `MAD_AUTO_ROCM_PATH=0` for **host** behavior as documented in [docs/configuration.md](docs/configuration.md):
+## π Workflow
-```bash
-# Override host ROCm root:
-madengine run --tags dummy --additional-context '{"MAD_ROCM_PATH": "/path/to/rocm"}'
-# or: export ROCM_PATH=/path/to/rocm && madengine run --tags dummy
-# Override in-container ROCm root independently:
-madengine run --tags dummy --additional-context '{"docker_env_vars": {"MAD_ROCM_PATH": "/path/in/container"}}'
+The core pipeline is the same everywhere: discover models, build images once, run them against a target, then report. Build and run can be separated so images are built once (e.g. in CI) and reused across nodes.
+
+```mermaid
+flowchart LR
+ D[discover
find models by tag] --> B[build
Docker images]
+ B --> M[(build_manifest.json)]
+ M --> R[run
infer target + execute]
+ R --> P[(perf.csv)]
+ P --> RP[report
HTML / email]
+ P --> DB[database
MongoDB]
```
-**Results:** Performance data is written to `perf.csv` (and optionally `perf_entry.csv`). The file is created automatically if missing. Failed runs (including pre-run setup failures) are recorded with status `FAILURE` so every attempted model appears in the table. See [Exit Codes](docs/cli-reference.md#exit-codes) for CI/script usage.
+**Deployment target is inferred from the config** (Convention over Configuration) β no `deploy` flag needed:
-## π Commands
+```mermaid
+flowchart TD
+ A[additional_context] --> Q{which key?}
+ Q -->|k8s / kubernetes| K[Kubernetes deployment]
+ Q -->|slurm| S[SLURM deployment]
+ Q -->|neither| L[Local Docker execution]
+```
-madengine provides five main commands for model automation and benchmarking:
+## π Commands
| Command | Description | Use Case |
|---------|-------------|----------|
-| **[discover](#-model-discovery)** | Find available models | Model exploration and validation |
-| **[build](#building-images)** | Build Docker images | Create containerized models |
-| **[run](#-usage-examples)** | Execute models | Local and distributed execution |
-| **[report](docs/cli-reference.md#report---generate-reports)** | Generate HTML reports | Convert CSV to viewable reports |
-| **[database](docs/cli-reference.md#database---upload-to-mongodb)** | Upload to MongoDB | Store results in database |
-
-**Quick Start:**
+| **[discover](docs/usage.md#model-discovery)** | Find available models | Model exploration and validation |
+| **[build](docs/usage.md#build-workflow)** | Build Docker images | Create containerized models |
+| **[run](docs/usage.md#run-workflow)** | Execute models | Local and distributed execution |
+| **[report](docs/cli-reference.md#report---generate-reports)** | Generate HTML/email reports | Convert CSV to viewable reports |
+| **[database](docs/cli-reference.md#database---upload-to-mongodb)** | Upload to MongoDB | Store results in a database |
```bash
-# Discover models
-madengine discover --tags dummy
+madengine discover --tags dummy # Find models
+madengine build --tags dummy # Build image (AMD/UBUNTU defaults)
+madengine run --tags dummy # Run model
+madengine report to-html --csv-file-path perf_entry.csv # Report
+madengine database --file perf_entry.csv --db mydb --collection results # Upload
+```
-# Build image (uses AMD/UBUNTU defaults)
-madengine build --tags dummy
+For all options and examples, see the **[CLI Reference](docs/cli-reference.md)**.
-# Run model
-madengine run --tags dummy
+## π» Usage Examples
-# For non-AMD/Ubuntu environments, specify explicitly:
-# madengine build --tags dummy --additional-context '{"gpu_vendor": "NVIDIA", "guest_os": "CENTOS"}'
+```bash
+# Local, multi-GPU with torchrun (DDP/FSDP)
+madengine run --tags model \
+ --additional-context '{"docker_gpus": "0,1,2,3",
+ "distributed": {"launcher": "torchrun", "nproc_per_node": 4}}'
-# Generate report
-madengine report to-html --csv-file perf_entry.csv
+# Kubernetes (minimal config, presets auto-applied)
+madengine run --tags model \
+ --additional-context '{"k8s": {"gpu_count": 2}}'
-# Upload results
-madengine database --csv-file perf_entry.csv --db mydb --collection results
+# SLURM (build once, then deploy)
+madengine build --tags model --registry gcr.io/myproject
+madengine run --manifest-file build_manifest.json \
+ --additional-context '{"slurm": {"partition": "gpu", "nodes": 4, "gpus_per_node": 8},
+ "distributed": {"launcher": "torchtitan", "nnodes": 4, "nproc_per_node": 8}}'
```
-For detailed command options, see the **[CLI Command Reference](docs/cli-reference.md)**.
+More local/K8s/SLURM/CI recipes: [Usage Guide](docs/usage.md) Β· [Configuration](docs/configuration.md) Β· [CLI Reference](docs/cli-reference.md).
## π Documentation
| Guide | Description |
|-------|-------------|
| [Installation](docs/installation.md) | Complete installation instructions |
-| [Usage Guide](docs/usage.md) | Commands, workflows, and examples ([`--skip-model-run`](docs/usage.md#skip-model-run-after-build)) |
+| [Usage Guide](docs/usage.md) | Commands, workflows, and examples |
| **[CLI Reference](docs/cli-reference.md)** | **Detailed command options and examples** |
+| [Configuration](docs/configuration.md) | Advanced options, ROCm path, log error scan |
| [Deployment](docs/deployment.md) | Kubernetes and SLURM deployment |
-| [Configuration](docs/configuration.md) | Advanced options; [run log error pattern scan](docs/configuration.md#run-phase-log-error-pattern-scan) |
| [Batch Build](docs/batch-build.md) | Selective builds for CI/CD |
-| [Launchers](docs/launchers.md) | Distributed training frameworks |
+| [Launchers](docs/launchers.md) | Distributed frameworks + capability matrices |
| [Profiling](docs/profiling.md) | Performance analysis tools |
| [Contributing](docs/contributing.md) | How to contribute |
-## ποΈ Architecture
-
-```
- βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
- β madengine CLI v2.0 (Typer + Rich) β
- β discover β build β run β report β database β
- βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
- β β β
- β β βΌ
- β β ββββββββββββββββββββββββ Orchestration Layer ββββββββββββββββββββββββββββ
- β β β Model Discovery (models.json / scripts/ get_models) β
- β β β BuildOrchestrator Β· RunOrchestrator β
- β ββββ |
- ββββββββββββββββ΄ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
- β
- βββββββββββββββββββββββββββββββββββΌββββββββββββββββββββ Infrastructure Layer ββββββββββββββ
- β βΌ βΌ βΌ β
- β ββββββββββββββββ ββββββββββββββββ ββββββββββββββββ β
- β β Local β β Kubernetes β β SLURM β β
- β β Docker β β Jobs β β Jobs β β
- β ββββββββ¬ββββββββ ββββββββ¬ββββββββ ββββββββ¬ββββββββ β
- β ββββββββββββββββββββΌβββββββββββββββββββ β
- βββββββββββββββββββββββββββββββββββΌββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
- βΌ
- βββββββββββββββββββββββββββββββββββββ Launcher Layer (Distribution) βββββββββββββββββββββββ
- β Train: torchrun Β· DeepSpeed Β· Megatron-LM Β· TorchTitan Β· Primus β
- β Infer: vLLM Β· SGLang Β· SGLang Disagg β
- βββββββββββββββββββββββββββββββββββ¬ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
- βΌ
- βββββββββββββββββββββββββββββββ
- β Performance (CSV/JSON) β
- βββββββββββββββ¬ββββββββββββββββ
- β
- βββββββββββββββββββββ΄ββββββββββββββββββββ
- βΌ βΌ
- βββββββββββββββββββββ ββββββββββββββββββββ
- β report β β database β
- β to-html, to-email β β MongoDB upload β
- βββββββββββββββββββββ ββββββββββββββββββββ
-```
-
-**Component Flow:**
-
-1. **CLI Layer** - User interface with 5 commands (discover, build, run, report, database)
-2. **Model Discovery** - Find and validate models from MAD package
-3. **Orchestration** - BuildOrchestrator & RunOrchestrator manage workflows
-4. **Execution Targets** - Local Docker, Kubernetes Jobs, or SLURM Jobs
-5. **Distributed Launchers** - Training (torchrun, DeepSpeed, Megatron-LM, TorchTitan, Primus) and Inference (vLLM, SGLang)
-6. **Performance Output** - CSV/JSON results with metrics
-7. **Post-Processing** - Report generation (HTML/Email) and database upload (MongoDB)
-
-## π― Feature Matrix
-
-### Supported Launchers & Infrastructure
+## π― Supported Launchers
| Launcher | Local | Kubernetes | SLURM | Type | Key Features |
|----------|-------|-----------|-------|------|--------------|
| **torchrun** | β
| β
| β
| Training | PyTorch DDP/FSDP, elastic training |
| **DeepSpeed** | β
| β
| β
| Training | ZeRO optimization, pipeline parallelism |
| **Megatron-LM** | β
| β
| β
| Training | Tensor+Pipeline parallel, large transformers |
-| **TorchTitan** | β
| β
| β
| Training | FSDP2+TP+PP+CP, Llama 3.1 (8B-405B) |
-| **Primus** | β
| β
| β
| Training | Megatron / TorchTitan / MaxText via Primus YAML; `distributed.primus` |
+| **TorchTitan** | β
| β
| β
| Training | FSDP2+TP+PP+CP, Llama 3.1 (8Bβ405B) |
+| **Primus** | β
| β
| β
| Training | Megatron / TorchTitan / MaxText via Primus YAML |
| **vLLM** | β
| β
| β
| Inference | v1 engine, PagedAttention, Ray cluster |
| **SGLang** | β
| β
| β
| Inference | RadixAttention, structured generation |
| **SGLang Disagg** | β | β
| β
| Inference | Disaggregated prefill/decode, Mooncake, 3+ nodes |
-**Note:** All launchers support single-GPU, multi-GPU (single node), and multi-node (where infrastructure allows). See [Launchers Guide](docs/launchers.md) for details.
-
-### Parallelism Capabilities
-
-| Launcher | Tensor Parallel (TP) | Pipeline Parallel (PP) | Data Parallel (DP) | Context Parallel (CP) | FSDP/ZeRO | Expert Parallel (EP) | Primary Use Case |
-|----------|----------------------|------------------------|--------------------|------------------------|-----------|----------------------|------------------|
-| **torchrun** | βManual | βNo | βManual (DDP) | βNo | βManual (FSDP) | βNo | General distributed training |
-| **TorchTitan** | β
Auto | β
Auto | β
Auto (FSDP2) | βManual | β
Auto (FSDP2) | βNo | Large-scale LLM pre-training |
-| **DeepSpeed** | βManual | βManual | β
Auto (ZeRO) | βNo | β
Auto (ZeRO) | βNo | Memory-efficient training |
-| **Megatron-LM** | β
Auto | β
Auto | β
Implicit | β
Auto | βNo | βNo | Large transformer training |
-| **Primus** | βManual | βManual | βManual | βManual | βManual | βNo | Unified pretrain (experiment YAML; backend-specific) |
-| **vLLM** | β
Auto | SLURM: β
Auto (Multi) / K8s: βDisabled | β
Auto (Replicas) | βNo | βNo | βManual | High-throughput inference |
-| **SGLang** | β
Auto | SLURM: β
Auto (Multi) / K8s: βDisabled | βLimited | βNo | βNo | βNo | Inference + structured gen |
-| **SGLang PD Disagg** | β
Auto | βNo | β
Role-based | βNo | βNo | βNo | Optimized prefill/decode |
+All launchers support single-GPU, multi-GPU, and multi-node (where infrastructure allows). See the [Launchers Guide](docs/launchers.md) for the full **parallelism** and **infrastructure** capability matrices.
-**Legend:** β
Auto = supported and configured by madengine; βManual = supported by launcher but requires user configuration; βLimited / βDisabled = launcher or platform limitation. See [Launchers Guide](docs/launchers.md) and [Configuration](docs/configuration.md) for details.
+## π Profiling
-### Infrastructure Capabilities
-
-| Feature | Local | Kubernetes | SLURM |
-|---------|-------|-----------|-------|
-| **Execution** | Docker containers | K8s Jobs | SLURM jobs |
-| **Multi-Node** | β | β
Indexed Jobs | β
Job arrays |
-| **Resource Mgmt** | Manual | Declarative (YAML) | Batch scheduler |
-| **Monitoring** | Docker logs | kubectl/dashboard | squeue/scontrol |
-| **Auto-scaling** | β | β
| β |
-| **Network** | Host | CNI plugin | InfiniBand/Ethernet |
-
-## π» Usage Examples
-
-### Local Execution
-
-```bash
-# Single GPU
-madengine run --tags dummy \
- --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'
-
-# Multi-GPU with torchrun (DDP/FSDP)
-madengine run --tags model \
- --additional-context '{
- "gpu_vendor": "AMD",
- "guest_os": "UBUNTU",
- "docker_gpus": "0,1,2,3",
- "distributed": {
- "launcher": "torchrun",
- "nproc_per_node": 4
- }
- }'
-
-# With DeepSpeed (ZeRO optimization)
-madengine run --tags model \
- --additional-context '{
- "gpu_vendor": "AMD",
- "guest_os": "UBUNTU",
- "docker_gpus": "all",
- "distributed": {
- "launcher": "deepspeed",
- "nproc_per_node": 8
- }
- }'
-```
-
-### Kubernetes Deployment
+madengine ships integrated profiling for AMD ROCm β `rocprof`, eight pre-configured `rocprofv3` profiles (ROCm 7.0+), `rocm-trace-lite`, library tracing (rocBLAS/MIOpen/Tensile/RCCL), and power/VRAM monitors. Tools are stackable via `--additional-context '{"tools": [...]}'`.
```bash
-# Minimal config (auto-defaults applied)
-madengine run --tags model \
- --additional-context '{"k8s": {"gpu_count": 2}}'
-
-# Multi-node inference with vLLM
-madengine run --tags model \
- --additional-context '{
- "k8s": {
- "namespace": "ml-team",
- "gpu_count": 8
- },
- "distributed": {
- "launcher": "vllm",
- "nnodes": 2,
- "nproc_per_node": 4
- }
- }'
-
-# SGLang with structured generation
-madengine run --tags model \
- --additional-context '{
- "k8s": {"gpu_count": 4},
- "distributed": {
- "launcher": "sglang",
- "nproc_per_node": 4
- }
- }'
-```
-
-### SLURM Deployment
-
-```bash
-# Build phase (local or CI)
-madengine build --tags model \
- --registry gcr.io/myproject \
- --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'
-
-# Deploy phase (on SLURM login node)
-madengine run --manifest-file build_manifest.json \
- --additional-context '{
- "slurm": {
- "partition": "gpu",
- "nodes": 4,
- "gpus_per_node": 8,
- "time": "24:00:00"
- },
- "distributed": {
- "launcher": "torchtitan",
- "nnodes": 4,
- "nproc_per_node": 8
- }
- }'
-```
-
-To run on **specific nodes**, set `nodelist` (comma-separated node names). When set, the job is restricted to those nodes and automatic node health preflight is skipped. Example: `"slurm": { "nodelist": "node01,node02", "nodes": 2, ... }`. See [Configuration](docs/configuration.md#slurm-deployment) and [examples/slurm-configs/basic/03-multi-node-basic-nodelist.json](examples/slurm-configs/basic/03-multi-node-basic-nodelist.json).
-
-### Common Workflows
-
-**Development β Testing β Production:**
-
-```bash
-# 1. Develop locally with single GPU
-madengine run --tags model \
- --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'
-
-# 2. Test multi-GPU locally
-madengine run --tags model \
- --additional-context '{
- "gpu_vendor": "AMD",
- "guest_os": "UBUNTU",
- "docker_gpus": "0,1",
- "distributed": {"launcher": "torchrun", "nproc_per_node": 2}
- }'
-
-# 3. Build and push to registry
-madengine build --tags model \
- --registry docker.io/myorg \
- --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'
-
-# 4. Deploy to Kubernetes
-madengine run --manifest-file build_manifest.json
-```
-
-**CI/CD Pipeline:**
-
-```bash
-# Batch build (selective rebuilds)
-madengine build --batch-manifest batch.json \
- --registry docker.io/myorg
-
-# Run tests
-madengine run --manifest-file build_manifest.json \
- --additional-context '{"k8s": {"namespace": "ci-test"}}'
-
-# Generate and email reports
-madengine report to-email --directory ./results --output ci_report.html
-
-# Upload to database
-madengine database --csv-file perf_entry.csv \
- --database-name ci_db --collection-name test_results
-```
-
-See [Usage Guide](docs/usage.md), [Configuration Guide](docs/configuration.md), and [CLI Reference](docs/cli-reference.md) for more examples.
-
-### Building Images
-
-```bash
-# Build single model
-madengine build --tags dummy \
- --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'
-
-# Build with registry (for distributed deployment)
-madengine build --tags model1 model2 \
- --registry localhost:5000 \
- --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'
-
-# Build for multiple GPU architectures
-madengine build --tags model \
- --target-archs gfx908 gfx90a gfx942 \
- --registry gcr.io/myproject
-
-# Batch build mode (selective builds for CI/CD)
-madengine build --batch-manifest examples/build-manifest/batch.json \
- --registry docker.io/myorg
-
-# Clean rebuild (no Docker cache)
-madengine build --tags model --clean-docker-cache \
- --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'
-```
-
-**Output:** Creates `build_manifest.json` with built image names and configurations.
-
-See [Batch Build Guide](docs/batch-build.md) and examples in [`examples/build-manifest/`](examples/build-manifest/).
-
-## π Model Discovery
-
-madengine discovers models from the MAD package using three methods:
-
-```bash
-# Root models (models.json)
-madengine discover --tags pyt_huggingface_bert
-
-# Directory-specific (scripts/{dir}/models.json), scoped tag: {dir}/{model_or_tag}
-madengine discover --tags dummy2/model1
-
-# Dynamic with parameters (scripts/{dir}/get_models_json.py)
-madengine discover --tags dummy3/model3:batch_size=512
+madengine run --tags model --additional-context '{"tools": [{"name": "rocprofv3_compute"}]}'
```
-## π Performance Profiling
-
-madengine includes integrated profiling tools for AMD ROCm:
-
-```bash
-# GPU profiling with rocprof
-madengine run --tags model \
- --additional-context '{
- "gpu_vendor": "AMD",
- "guest_os": "UBUNTU",
- "tools": [{"name": "rocprof"}]
- }'
-
-# ROCprofv3 (ROCm 7.0+) - Advanced profiling with pre-configured profiles
-madengine run --tags model \
- --additional-context '{"tools": [{"name": "rocprofv3_compute"}]}'
-
-# Use configuration files for complex setups
-madengine run --tags model \
- --additional-context-file examples/profiling-configs/rocprofv3_multi_gpu.json
-
-# Library tracing (rocBLAS, MIOpen, Tensile, RCCL)
-madengine run --tags model \
- --additional-context '{"tools": [{"name": "rocblas_trace"}]}'
-
-# rocm-trace-lite β lightweight kernel dispatch trace (SQLite; no rocprofiler-sdk)
-# Requires outbound HTTPS to GitHub on first run unless the wheel is baked into the image
-# (see docs/profiling.md). Do not combine with rocprof / rocprofv3_* on the same run.
-madengine run --tags model \
- --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU", "tools": [{"name": "rocm_trace_lite"}]}'
-
-# Power and VRAM monitoring
-madengine run --tags model \
- --additional-context '{"tools": [
- {"name": "gpu_info_power_profiler"},
- {"name": "gpu_info_vram_profiler"}
- ]}'
-
-# Multiple tools (stackable)
-madengine run --tags model \
- --additional-context '{"tools": [
- {"name": "rocprofv3_memory"},
- {"name": "rocblas_trace"},
- {"name": "gpu_info_power_profiler"}
- ]}'
-```
-
-**Available Tools:**
-
-| Tool | Purpose | Output |
-|------|---------|--------|
-| `rocprof` | GPU kernel profiling | Kernel timings, occupancy |
-| `rocprofv3_compute` | Compute-bound analysis (ROCm 7.0+) | ALU metrics, wave execution |
-| `rocprofv3_memory` | Memory-bound analysis (ROCm 7.0+) | Cache hits, bandwidth |
-| `rocprofv3_communication` | Multi-GPU communication (ROCm 7.0+) | RCCL traces, inter-GPU transfers |
-| `rocprofv3_lightweight` | Minimal overhead profiling (ROCm 7.0+) | HIP and kernel traces |
-| `rocm_trace_lite` | RTL **`lite`** mode β kernel dispatch trace (HSA, SQLite/RPD-style); [`rtl trace --mode lite`](https://sunway513.github.io/rocm-trace-lite/quickstart.html) via `rtl_trace_wrapper.sh` | `rocm_trace_lite_output/trace.db` (and optional `trace.json.gz`, `trace_summary.txt`) |
-| `rocm_trace_lite_default` | RTL **`default`** mode β broader dispatch coverage; higher overhead than `lite` (same outputs paths) | Same as `rocm_trace_lite` |
-| `rocblas_trace` | rocBLAS library calls | Function calls, arguments |
-| `miopen_trace` | MIOpen library calls | Conv/pooling operations |
-| `tensile_trace` | Tensile GEMM library | Matrix multiply details |
-| `rccl_trace` | RCCL collective ops | Communication patterns |
-| `gpu_info_power_profiler` | GPU power consumption | Power usage over time |
-| `gpu_info_vram_profiler` | GPU memory usage | VRAM utilization |
-| `therock_check` | TheRock ROCm validation | Installation detection |
-
-**ROCprofv3 Profiles** (ROCm 7.0+):
-
-madengine provides 8 pre-configured ROCprofv3 profiles for different bottleneck scenarios:
-
-- `rocprofv3_compute` - Compute-bound workloads (transformers, dense ops)
-- `rocprofv3_memory` - Memory-bound workloads (large batches, high-res)
-- `rocprofv3_communication` - Multi-GPU distributed training
-- `rocprofv3_full` - Comprehensive profiling (all metrics, high overhead)
-- `rocprofv3_lightweight` - Minimal overhead (production-friendly)
-- `rocprofv3_perfetto` - Perfetto UI compatible traces
-- `rocprofv3_api_overhead` - API call timing analysis
-- `rocprofv3_pc_sampling` - Kernel hotspot identification
-
-See [`examples/profiling-configs/`](examples/profiling-configs/) for ready-to-use configuration files.
-
-**rocm-trace-lite (`rocm_trace_lite` / `rocm_trace_lite_default`):**
-
-- madengine runs workloads under `scripts/common/tools/rtl_trace_wrapper.sh`, which invokes the `rtl` CLI (or `python3 -m rocm_trace_lite.cli`) with **`RTL_MODE=lite`** or **`RTL_MODE=default`** and writes traces under `rocm_trace_lite_output/`.
-- The trace **pre-script** installs the package from a **[GitHub Release wheel](https://github.com/sunway513/rocm-trace-lite/releases)** (not PyPI). By default it uses a **pinned** `linux_x86_64` wheel for reproducible installs. Set **`ROCM_TRACE_LITE_FOLLOW_LATEST=1`** to resolve the latest wheel via the GitHub API, or **`ROCM_TRACE_LITE_WHEEL_URL`** to a direct `.whl` URL for air-gapped installs or non-x86_64 platforms.
-- Choose **either** `rocm_trace_lite` **or** rocprof / `rocprofv3_*` for a given runβnot both. Details: [Profiling Guide](docs/profiling.md) (section *rocm-trace-lite (RTL)*).
-
-**TheRock Validation:**
-
-```bash
-# Validate TheRock installation (AMD's pip-based ROCm)
-madengine run --tags dummy_therock \
- --additional-context '{"tools": [{"name": "therock_check"}]}'
-```
-
-See [Profiling Guide](docs/profiling.md) for detailed usage and analysis.
-
-## π Reporting and Database
-
-### Generate Reports
-
-Convert performance CSV files to HTML reports:
-
-```bash
-# Single CSV to HTML
-madengine report to-html --csv-file perf_entry.csv
-
-# Consolidated email report (all CSVs in directory)
-madengine report to-email --directory ./results --output summary.html
-```
-
-### Upload to Database
-
-Store performance results in MongoDB:
-
-```bash
-# Set MongoDB connection
-export MONGO_HOST=mongodb.example.com
-export MONGO_PORT=27017
-export MONGO_USER=myuser
-export MONGO_PASSWORD=mypassword
-
-# Upload CSV to MongoDB
-madengine database --csv-file perf_entry.csv \
- --database-name performance_db \
- --collection-name model_runs
-```
-
-**Use Cases:**
-- Track performance over time
-- Compare results across different configurations
-- Build performance dashboards
-- Automated CI/CD reporting
-
-See [CLI Reference](docs/cli-reference.md) for complete options.
+See the [Profiling Guide](docs/profiling.md) and ready-to-use configs in [`examples/profiling-configs/`](examples/profiling-configs/).
## π¦ Installation
```bash
-# Install madengine (all dependencies, including Kubernetes support, are included)
+# Install (all dependencies, including Kubernetes support, included)
pip install git+https://github.com/ROCm/madengine.git
# Development installation
@@ -560,147 +218,41 @@ git clone https://github.com/ROCm/madengine.git
cd madengine && pip install -e .
```
-See [Installation Guide](docs/installation.md) for detailed instructions.
-
-## π‘ Tips & Best Practices
-
-### General Usage
-
-- **Use configuration files** for complex setups instead of long command lines
-- **Test locally first** with single GPU before scaling to multi-node
-- **Enable verbose logging** (`--verbose`) when debugging issues
-- **Use `--live-output`** for real-time monitoring of long-running operations
-
-### Log error pattern scan
-
-After a local Docker run, madengine can scan the captured **run log** for common failure substrings (for example `RuntimeError:`, `CUDA out of memory`, `Traceback`). That helps catch hard failures when exit codes are ambiguous, but some workloads log benign `RuntimeError:` text while tests still pass.
-
-- **Disable** the scan when another signal is authoritative (e.g. pytest/JUnit inside the image): set `"log_error_pattern_scan": false` in `--additional-context` or in the model entry in `models.json`. See [Configuration β Run phase: log error pattern scan](docs/configuration.md#run-phase-log-error-pattern-scan).
-- **Extend exclusions** with `log_error_benign_patterns` (list of strings), or **replace** the default pattern list with `log_error_patterns` (non-empty list of strings) for advanced cases.
-
-### CI / Jenkins
-
-- **Exit codes:** The CLI uses fixed exit codes (`ExitCode` in `madengine.cli.constants`, e.g. `SUCCESS=0`, `RUN_FAILURE=3`, `INVALID_ARGS=4`). Pipelines should treat **non-zero** as failure; no log scraping is required for pass/fail.
-- **Streaming:** In Jenkins, avoid redirecting stdout only to a file (`> file`) without `tee` if you want the console to update during the run. Prefer `... 2>&1 | tee madengine.run.log` with `bash -o pipefail` so the step exit code is still from `madengine`.
-- **Unbuffered Python:** If output still appears in chunks, set `PYTHONUNBUFFERED=1` (or `python -u`) for the `madengine` process.
+See the [Installation Guide](docs/installation.md) for details.
-### Build & Deployment
+## π‘ Tips & Troubleshooting
-- **Separate build and run phases** for distributed deployments
-- **Skip model script:** `madengine run --tags β¦ --skip-model-run` starts the container and runs `pre_scripts`, but skips the model script. Combine with `--keep-alive` for a live container ready for manual exec. Ignored with a warning on SLURM/K8s. See [Usage β Skip model run after build](docs/usage.md#skip-model-run-after-build).
-- **Use registries** for multi-node execution (K8s/SLURM)
-- **Use batch build mode** for CI/CD to optimize build times
-- **Specify `--target-archs`** when building for multiple GPU architectures
+- **Test locally first** with a single GPU before scaling to multi-node; use config files for complex setups.
+- **Debugging:** add `--verbose --live-output`; keep a container for inspection with `--keep-alive`.
+- **CI:** the CLI uses fixed [exit codes](docs/cli-reference.md#exit-codes) (`0` success, `2` build failure, `3` run failure, `4` invalid args) β no log scraping needed.
+- **Log error scan** can flag benign `RuntimeError:` text; disable or tune it via `log_error_pattern_scan` / `log_error_benign_patterns` β see [Configuration](docs/configuration.md#run-phase-log-error-pattern-scan).
-### Performance
-
-- **Start with small timeouts** and increase as needed
-- **Use profiling tools** to identify bottlenecks
-- **Monitor GPU utilization** with `gpu_info_power_profiler`
-- **Profile library calls** with rocBLAS/MIOpen tracing
-
-### Exit codes and CI
-
-madengine uses consistent exit codes for scripts and CI (e.g. Jenkins): `0` = success, `1` = general failure, `2` = build failure, `3` = one or more run failures, `4` = invalid arguments. Failed runs are still written to `perf.csv` with status `FAILURE`. See [CLI Reference β Exit Codes](docs/cli-reference.md#exit-codes) for the full table and examples.
-
-### Troubleshooting
-
-```bash
-# Check model is available
-madengine discover --tags your_model
-
-# Verbose output for debugging
-madengine run --tags model --verbose --live-output
-
-# Keep container alive for inspection
-madengine run --tags model --keep-alive
-
-# Clean rebuild if build fails
-madengine build --tags model --clean-docker-cache --verbose
-```
-
-**ROCm not in /opt/rocm:** Set top-level `MAD_ROCM_PATH` in `--additional-context` for the **host**; for **in-container** paths, set `docker_env_vars.MAD_ROCM_PATH`, or let madengine resolve `ROCM_PATH` at run from the image and probe (see [Configuration](docs/configuration.md#rocm-path-run-only)).
-
-**Common Issues:**
-- **False failures with profiling**: If models show FAILURE but have performance metrics, see [Profiling Troubleshooting](docs/profiling.md#false-failure-detection-with-rocprof)
-- **False failures from `RuntimeError:` in logs**: If the workload logs expected exception text but tests pass, disable or tune the scan with `log_error_pattern_scan` / `log_error_benign_patterns` β see [Configuration](docs/configuration.md#run-phase-log-error-pattern-scan)
-- **ROCProf log errors**: Messages like `E20251230` are informational logs, not errors (fixed in v2.0+)
-- **Configuration errors**: Validate JSON with `python -m json.tool your-config.json`
+More: [Usage β Troubleshooting](docs/usage.md#troubleshooting) Β· [Profiling β False failures](docs/profiling.md#false-failure-detection-with-rocprof).
## π€ Contributing
-We welcome contributions! See [Contributing Guide](docs/contributing.md) for details.
+Contributions are welcome! See the [Contributing Guide](docs/contributing.md).
```bash
git clone https://github.com/ROCm/madengine.git
-cd madengine
-python3 -m venv venv && source venv/bin/activate
+cd madengine && python3 -m venv venv && source venv/bin/activate
pip install -e .
-
-# Run all tests
pytest
-
-# Run specific test module
-pytest tests/unit/test_error_handling.py -v
-
-# Run error pattern tests
-pytest tests/unit/test_error_handling.py::TestErrorPatternMatching -v
```
## π License
-MIT License - see [LICENSE](LICENSE) file for details.
+MIT License β see [LICENSE](LICENSE).
## π Links & Resources
-### Documentation
-- **[CLI Reference](docs/cli-reference.md)** - Complete command options
-- **[Usage Guide](docs/usage.md)** - Workflows and examples
-- **[Deployment Guide](docs/deployment.md)** - Kubernetes/SLURM deployment
-- **[Configuration Guide](docs/configuration.md)** - Advanced configuration
-- **[All Docs](docs/)** - Complete documentation index
-
-### External Resources
-- **MAD Package**: https://github.com/ROCm/MAD
-- **Issues & Support**: https://github.com/ROCm/madengine/issues
-- **ROCm Documentation**: https://rocm.docs.amd.com/
-
-### Getting Help
-
-**Command Help:**
-```bash
-madengine --help # Main help
-madengine --help # Command-specific help
-madengine report --help # Sub-app help
-madengine report to-html --help # Sub-command help
-```
-
-**Quick Checks:**
-```bash
-# Verify installation
-madengine --version
-
-# Discover available models
-madengine discover
-
-# Check specific model
-madengine discover --tags your_model --verbose
-```
-
-**Troubleshooting:**
-- Check [CLI Reference](docs/cli-reference.md) for all command options
-- Enable `--verbose` flag for detailed error messages
-- See [Usage Guide](docs/usage.md) troubleshooting section
-- Report issues: https://github.com/ROCm/madengine/issues
+- **MAD Package:** https://github.com/ROCm/MAD
+- **Issues & Support:** https://github.com/ROCm/madengine/issues
+- **ROCm Documentation:** https://rocm.docs.amd.com/
+- **Command help:** `madengine --help` Β· `madengine --help`
---
## β οΈ Migration Notice (v2.0.0+)
-The CLI has been unified! Starting from v2.0.0:
-- β
Use `madengine` (unified modern CLI with K8s, SLURM, distributed support)
-- β Legacy v1.x CLI has been removed
-
----
-
-**Code Quality**: Clean codebase with no dead code, comprehensive test coverage, and following Python best practices.
+The CLI has been unified. Starting from v2.0.0, use `madengine` (with K8s, SLURM, and distributed support); the legacy v1.x CLI has been removed.
diff --git a/docs/README.md b/docs/README.md
index db31f762..1dd64548 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -35,7 +35,55 @@ Complete documentation for madengine - AI model automation and distributed bench
## ποΈ Architecture
-The architecture diagram (Orchestration, Infrastructure, and Launcher layers) is in the [main README](../README.md#-architecture). Summary:
+The CLI drives orchestrators that discover and build models, then hand off to a local or distributed execution target, which runs the model under the appropriate launcher and emits performance data for reporting. (Same diagram as the [main README](../README.md#-architecture).)
+
+```mermaid
+flowchart TB
+ subgraph CLI["CLI Layer β Typer + Rich"]
+ C1[discover]
+ C2[build]
+ C3[run]
+ C4[report]
+ C5[database]
+ end
+
+ subgraph ORC["Orchestration Layer"]
+ O1[DiscoverModels]
+ O2[BuildOrchestrator]
+ O3[RunOrchestrator]
+ MAN[(build_manifest.json)]
+ end
+
+ subgraph EXEC["Execution / Deployment Layer"]
+ E1[ContainerRunner
local Docker]
+ E2[DeploymentFactory]
+ K8S[Kubernetes Jobs]
+ SLURM[SLURM Jobs]
+ end
+
+ subgraph LAUNCH["Launcher Layer"]
+ T[Train: torchrun Β· DeepSpeed
Megatron-LM Β· TorchTitan Β· Primus]
+ I[Infer: vLLM Β· SGLang Β· SGLang Disagg]
+ end
+
+ OUT[(perf.csv / JSON)]
+
+ C1 --> O1
+ C2 --> O2
+ C3 --> O3
+ O2 --> MAN --> O3
+ O1 --> O2
+ O3 --> E1
+ O3 --> E2
+ E2 --> K8S
+ E2 --> SLURM
+ E1 --> LAUNCH
+ K8S --> LAUNCH
+ SLURM --> LAUNCH
+ LAUNCH --> OUT
+ OUT --> C4
+ OUT --> C5
+```
1. **CLI Layer** - User interface with 5 commands (discover, build, run, report, database)
2. **Model Discovery** - Find and validate models from MAD package
@@ -106,7 +154,7 @@ madengine operates within the MAD (Model Automation and Dashboarding) ecosystem.
- **torchrun** - PyTorch DDP/FSDP
- **deepspeed** - ZeRO optimization
-- **megatron** - Large transformers (K8s + SLURM)
+- **megatron-lm** - Large transformers (K8s + SLURM)
- **torchtitan** - LLM pre-training
- **vllm** - LLM inference
- **sglang** - Structured generation
@@ -119,7 +167,7 @@ This documentation follows these principles:
2. **Progressive disclosure** - Start simple, add complexity as needed
3. **Examples first** - Show working examples before explaining details
4. **Consistent naming** - Files follow simple naming pattern (no prefixes)
-5. **Up-to-date** - Reflects current implementation (v2.0)
+5. **Up-to-date** - Tracks the current implementation; see [CHANGELOG.md](../CHANGELOG.md) for release history
## π€ Contributing to Documentation
diff --git a/docs/batch-build.md b/docs/batch-build.md
index 24983d98..7bac329f 100644
--- a/docs/batch-build.md
+++ b/docs/batch-build.md
@@ -224,11 +224,15 @@ Creates `build_manifest.json` with:
}
},
"built_models": {...},
+ "context": {...},
+ "credentials_required": {...},
"deployment_config": {...},
"summary": {...}
}
```
+> `deployment_config` is only written when `--additional-context` resolves to a non-local deployment (e.g. `slurm`, `k8s`/`kubernetes`, `distributed`, `vllm`, or non-empty `env_vars`). Plain local builds omit this key entirely.
+
## Best Practices
1. **Version Control**: Keep batch manifests in version control for reproducibility
diff --git a/docs/cli-reference.md b/docs/cli-reference.md
index 842e0fdf..adffab03 100644
--- a/docs/cli-reference.md
+++ b/docs/cli-reference.md
@@ -230,13 +230,14 @@ madengine run [OPTIONS]
| `--keep-model-dir` | | FLAG | `False` | Keep model directory after run (local Docker only; ignored with a warning on SLURM/K8s) |
| `--clean-docker-cache` | | FLAG | `False` | Rebuild images without using cache (full workflow) |
| `--skip-model-run` | | FLAG | `False` | Skip the model script inside each container. The container still starts and `pre_scripts` still run; only the model script invocation is skipped (status reported as `SKIPPED`, exit code `0`). Combine with `--keep-alive` to leave a live container for manual exec. Ignored with a warning on SLURM/K8s targets. See [Usage β Skip model run](usage.md#skip-model-run-after-build). |
+| `--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). |
| `--manifest-output` | | TEXT | `build_manifest.json` | Output file for build manifest (full workflow) |
| `--summary-output` | `-s` | TEXT | `None` | Output file for summary JSON |
| `--live-output` | `-l` | FLAG | `False` | Print output in real-time |
-| `--output` | `-o` | TEXT | `perf_entry.csv` | Performance output file |
+| `--output` | `-o` | TEXT | `perf.csv` | Performance output file |
| `--ignore-deprecated` | | FLAG | `False` | Force run deprecated models |
| `--data-config` | | TEXT | `data.json` | Custom data configuration file |
-| `--tools-config` | | TEXT | `tools.json` | Custom tools JSON configuration |
+| `--tools-config` | | TEXT | `./scripts/common/tools.json` | Custom tools JSON configuration |
| `--sys-env-details` | | FLAG | `True` | Generate system config env details |
| `--force-mirror-local` | | TEXT | `None` | Path to force local data mirroring |
| `--disable-skip-gpu-arch` | | FLAG | `False` | Disable skipping models based on GPU architecture |
@@ -356,7 +357,7 @@ madengine run --tags model \
**Performance Output:**
-Results are saved to CSV file (default: `perf_entry.csv`) with metrics including:
+Results are saved to CSV file (default: `perf.csv`) with metrics including:
- Execution time
- GPU utilization
- Memory usage
@@ -384,20 +385,20 @@ madengine report to-html [OPTIONS]
| Option | Short | Type | Required | Description |
|--------|-------|------|----------|-------------|
-| `--csv-file` | | TEXT | **Yes** | Path to the CSV file to convert |
+| `--csv-file-path` | | TEXT | **Yes** | Path to the CSV file to convert |
| `--verbose` | `-v` | FLAG | No | Enable verbose logging |
**Examples:**
```bash
# Convert CSV to HTML
-madengine report to-html --csv-file perf_entry.csv
+madengine report to-html --csv-file-path perf_entry.csv
# With custom CSV file
-madengine report to-html --csv-file results/perf_mi300.csv
+madengine report to-html --csv-file-path results/perf_mi300.csv
# Verbose output
-madengine report to-html --csv-file perf.csv --verbose
+madengine report to-html --csv-file-path perf.csv --verbose
```
**Output:** Creates `{filename}.html` in the same directory as the CSV file.
@@ -444,7 +445,7 @@ madengine report to-email --directory ./results --verbose
### `database` - Upload to MongoDB
-Upload CSV performance data to MongoDB database.
+Upload CSV or JSON performance data to MongoDB (format is auto-detected).
**Usage:**
@@ -456,32 +457,30 @@ madengine database [OPTIONS]
| Option | Short | Type | Default | Required | Description |
|--------|-------|------|---------|----------|-------------|
-| `--csv-file` | | TEXT | `perf_entry.csv` | No | Path to the CSV file to upload |
-| `--database-name` | `--db` | TEXT | `None` | **Yes** | Name of the MongoDB database |
-| `--collection-name` | `--collection` | TEXT | `None` | **Yes** | Name of the MongoDB collection |
-| `--verbose` | `-v` | FLAG | `False` | No | Enable verbose logging |
+| `--file` | `-f` | TEXT | `None` | **Yes** | Path to file (CSV or JSON, auto-detected) |
+| `--database` | `--db` | TEXT | `None` | **Yes** | MongoDB database name |
+| `--collection` | `-c` | TEXT | `None` | **Yes** | MongoDB collection name |
+| `--unique-key` | `-k` | TEXT | `None` | No | Unique field(s) for deduplication (comma-separated, auto-detected if not specified) |
+| `--batch-size` | | INT | `1000` | No | Batch size for bulk operations |
+| `--no-upsert` | | FLAG | `False` | No | Insert only (don't update existing documents) |
+| `--no-index` | | FLAG | `False` | No | Skip automatic index creation |
+| `--dry-run` | | FLAG | `False` | No | Validate without uploading |
+| `--verbose` | `-v` | FLAG | `False` | No | Verbose output |
**Examples:**
```bash
-# Upload to MongoDB
-madengine database \
- --csv-file perf_entry.csv \
- --database-name mydb \
- --collection-name results
+# Upload JSON with auto-detection
+madengine database -f perf_entry_super.json --db mydb -c perf_super
-# Short option names
-madengine database \
- --csv-file perf.csv \
- --db test \
- --collection perf_data
+# Upload CSV with custom unique key
+madengine database -f perf.csv --db test -c results -k model,timestamp
+
+# Dry run to validate
+madengine database -f data.json --db test -c data --dry-run
# With verbose output
-madengine database \
- --csv-file perf.csv \
- --db mydb \
- --collection results \
- --verbose
+madengine database -f perf.csv --db mydb -c results --verbose
```
**Environment Variables:**
@@ -490,10 +489,12 @@ MongoDB connection details are read from environment variables:
| Variable | Description | Example |
|----------|-------------|---------|
-| `MONGO_HOST` | MongoDB host address | `localhost` or `mongodb.example.com` |
-| `MONGO_PORT` | MongoDB port | `27017` |
+| `MONGO_HOST` | MongoDB host address (default: `localhost`) | `localhost` or `mongodb.example.com` |
+| `MONGO_PORT` | MongoDB port (default: `27017`) | `27017` |
| `MONGO_USER` | MongoDB username | `admin` |
| `MONGO_PASSWORD` | MongoDB password | `secretpassword` |
+| `MONGO_AUTH_SOURCE` | MongoDB authentication database (default: `admin`) | `admin` |
+| `MONGO_TIMEOUT_MS` | Server selection timeout in milliseconds (default: `5000`) | `5000` |
**Example Setup:**
@@ -504,7 +505,7 @@ export MONGO_USER=myuser
export MONGO_PASSWORD=mypassword
madengine database \
- --csv-file perf_entry.csv \
+ --file perf_entry.csv \
--db performance_db \
--collection model_runs
```
@@ -637,12 +638,17 @@ madengine recognizes these environment variables:
| `MAD_DOCKERHUB_USER` | Docker Hub username | None |
| `MAD_DOCKERHUB_PASSWORD` | Docker Hub password/token | None |
| `MAD_DOCKERHUB_REPO` | Docker Hub repository | None |
-| `MAD_CONTAINER_IMAGE` | Pre-built container image to use | None |
+| `DOCKER_CONFIG` | Directory holding the Docker `config.json` whose existing login madengine reuses | `~/.docker` |
+| `MAD_SKIP_DOCKER_LOGIN` | Set to `1` to never run `docker login`; always defer to the machine's existing credentials | Unset |
| `MONGO_HOST` | MongoDB host for database command | `localhost` |
| `MONGO_PORT` | MongoDB port for database command | `27017` |
| `MONGO_USER` | MongoDB username | None |
| `MONGO_PASSWORD` | MongoDB password | None |
+> `MAD_CONTAINER_IMAGE` is **not** an environment variable. It is an
+> `--additional-context` key that selects a pre-built image and skips the build
+> phase β see [Configuration](configuration.md#pre-built-container-images).
+
---
## Best Practices
@@ -669,6 +675,6 @@ madengine recognizes these environment variables:
---
-**Version:** 2.1.0
-**Last Updated:** May 2026
+Run `madengine --version` for the installed version (derived from git tags via
+versioningit). Release history is in [CHANGELOG.md](../CHANGELOG.md).
diff --git a/docs/configuration.md b/docs/configuration.md
index 4831cc4f..6147c85c 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -21,8 +21,7 @@ madengine run --tags model --additional-context-file config.json
```json
{
"gpu_vendor": "AMD",
- "guest_os": "UBUNTU",
- "timeout_multiplier": 2.0
+ "guest_os": "UBUNTU"
}
```
@@ -121,6 +120,45 @@ madengine run --tags my_unit_test_suite \
Disabling the scan does **not** change performance metric extraction from the log; it only affects the post-hoc grep used to set `has_errors` for status.
+## 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 |
+| Image reference is already `repo@sha256:...` | Used as-is | Used as-is (already pinned) |
+| 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.
+
## System environment collection (rocEnvTool)
Before each container run, madengine appends `scripts/common/pre_scripts/run_rocenv_tool.sh` to the model's pre-scripts. It captures a CSV snapshot of the container's environment (OS, CPU, GPU, ROCm/CUDA, packages, env vars, NUMA). The output ends up at `_env.csv` in the workspace root.
@@ -179,13 +217,13 @@ Unknown `rocenv_mode` values fall back to `lite` with a warning.
**Overrides** (recommended for CI):
- **Additional context (host):** top-level `"MAD_ROCM_PATH": "/path/to/host/rocm"` β controls where madengine looks for host GPU tools (`rocminfo`, `amd-smi`, etc.).
-- **Additional context (container):** `"docker_env_vars": { "MAD_ROCM_PATH": "/path/inside/image" }` β sets the in-container `ROCM_PATH` for Docker runs. If omitted, at `run` time madengine uses the image OCI `Env` (`ROCM_PATH` / `ROCM_HOME`) if present, then an in-container probe, then defaults to `/opt/rocm`. The host-resolved path is **not** mirrored into the container.
+- **Additional context (container):** `"docker_env_vars": { "ROCM_PATH": "/path/inside/image" }` β sets the in-container `ROCM_PATH` for Docker runs. If omitted, at `run` time madengine uses the image OCI `Env` (`ROCM_PATH` / `ROCM_HOME`) if present, then an in-container probe, then defaults to `/opt/rocm`. The host-resolved path is **not** mirrored into the container.
These two keys are independent, allowing host and container to use different ROCm installations without confusion.
Precedence (host): top-level `MAD_ROCM_PATH` β auto-detect (unless disabled) β `ROCM_PATH` β `/opt/rocm`.
-Precedence (container, **local Docker `run`**, **AMD**): `docker_env_vars.MAD_ROCM_PATH` (maps to `ROCM_PATH` for the workload) or explicit `ROCM_PATH` in `docker_env_vars` β image OCI `Env` (`ROCM_PATH` / `ROCM_HOME`) β in-image probe β default `/opt/rocm` with a warning. Implemented in `ContainerRunner.run_container` after the run image is resolved.
+Precedence (container, **local Docker `run`**, **AMD**): explicit `ROCM_PATH` in `docker_env_vars` β image OCI `Env` (`ROCM_PATH` / `ROCM_HOME`) β in-image probe β default `/opt/rocm` with a warning. Implemented in `ContainerRunner.run_container` after the run image is resolved.
This applies to the run phase; build uses build-only context (no GPU detection) but still honors `MAD_ROCM_PATH` in context when set.
@@ -264,17 +302,25 @@ Pass environment variables to containers:
}
```
-### Custom Base Image
+### Pre-Built Container Images
-Override Docker base image:
+`MAD_CONTAINER_IMAGE` runs the model in an image you already have, **skipping the
+build phase entirely**. madengine validates the image exists locally (pulling it
+if not) and writes a synthetic `build_manifest.json` for it:
-```json
-{
- "MAD_CONTAINER_IMAGE": "rocm/pytorch:custom-tag"
-}
+```bash
+madengine run --tags my_model \
+ --additional-context "{'MAD_CONTAINER_IMAGE': 'rocm/pytorch:custom-tag'}"
```
-Or override BASE_DOCKER in FROM line:
+`--tags` is required in this mode β without it madengine has no models to map
+onto the image and fails with a configuration error. Note this is an
+`--additional-context` key, not an environment variable.
+
+### Custom Base Image
+
+To keep the normal build but change the image the Dockerfile's `FROM` line
+resolves to, override `BASE_DOCKER` as a build argument:
```json
{
@@ -328,13 +374,15 @@ Format: Comma-separated list with hyphen ranges.
### Timeout Settings
+Set a per-model timeout (seconds) in `models.json`:
+
```json
{
- "timeout_multiplier": 2.0
+ "timeout": 7200
}
```
-Or use command-line option:
+Or use the command-line option, which overrides the model's timeout:
```bash
madengine run --tags model --timeout 7200
@@ -388,7 +436,6 @@ Automatically applies (see presets under `src/madengine/deployment/presets/k8s/`
"memory_limit": "64Gi",
"cpu": "16",
"cpu_limit": "32",
- "service_account": "madengine-sa",
"image_pull_policy": "Always",
"ttl_seconds_after_finished": null,
"allow_privileged_profiling": null,
@@ -409,7 +456,6 @@ Automatically applies (see presets under `src/madengine/deployment/presets/k8s/`
- `memory_limit` - Memory limit (default: 2Γ memory request)
- `cpu` - CPU cores request (default: auto-scaled by GPU count)
- `cpu_limit` - CPU cores limit (default: 2Γ CPU request)
-- `service_account` - Service account name
- `image_pull_policy` - `Always`, `IfNotPresent`, or `Never`
- `ttl_seconds_after_finished` - Optional Job TTL in seconds (auto-delete finished Job); `null` to omit
- `allow_privileged_profiling` - `null` means enable elevated `securityContext` when tools/profiling are configured; `true`/`false` to force
@@ -417,6 +463,26 @@ Automatically applies (see presets under `src/madengine/deployment/presets/k8s/`
- `secrets.image_pull_secret_names` - Extra pull secret names (strings) merged with any created from `credential.json` when using `from_local_credentials`
- `secrets.runtime_secret_name` - Required for `existing` (pre-created opaque Secret with key `credential.json`); optional for `omit` if you still mount a runtime Secret
+**Cluster and scheduling keys:**
+- `kubeconfig` - Path to kubeconfig (default: `~/.kube/config`)
+- `gpu_resource_name` - GPU resource name (default: `amd.com/gpu`; use `nvidia.com/gpu` for NVIDIA)
+- `node_selector` - Label selectors for pod placement (default: `{}`)
+- `tolerations` - Tolerations for tainted nodes (default: `[]`)
+- `backoff_limit` - Job retry attempts before marking failed (default: `3`)
+- `output_dir` - Directory for generated manifests (default: `./k8s_manifests`)
+
+**Storage keys:**
+- `data_pvc` - Name of an existing PVC to use for data, skipping auto-creation
+- `storage_class` - Broad fallback StorageClass for both the shared-data PVC and the single-node results PVC
+- `nfs_storage_class` / `data_storage_class` - RWX class for shared data and multi-node results
+- `single_node_results_storage_class` / `multi_node_results_storage_class` - Fine-grained results-PVC overrides (`local_path_storage_class` is the legacy single-node fallback)
+- `results_storage_size` / `data_storage_size` - PVC sizes (defaults: `10Gi` / `100Gi`)
+- `recreate_shared_data_pvc` - Delete and recreate `madengine-shared-data` before use. **Destroys existing data** β back up first; intended for migrating an RWO PVC to RWX
+
+Multi-node jobs require an RWX results StorageClass; madengine warns when one is
+not set. The full key reference, PVC access-mode matrix, and preset defaults are
+in [examples/k8s-configs/README.md](../examples/k8s-configs/README.md).
+
### Multi-Node Kubernetes
```json
@@ -458,9 +524,12 @@ Automatically applies (see presets under `src/madengine/deployment/presets/k8s/`
"nodes": 2,
"nodelist": "node01,node02",
"time": "24:00:00",
- "mem": "64G",
- "mail_user": "user@example.com",
- "mail_type": "ALL"
+ "exclusive": true,
+ "constraint": "mi300x",
+ "exclude": "node07,node09",
+ "modules": ["rocm/6.2"],
+ "network_interface": "ib0",
+ "output_dir": "./slurm_results"
}
}
```
@@ -471,15 +540,23 @@ Automatically applies (see presets under `src/madengine/deployment/presets/k8s/`
- `partition` - SLURM partition name (required)
- `account` - Billing account
- `qos` - Quality of Service
-- `gpus_per_node` - GPUs per node (default: 1)
+- `gpus_per_node` - GPUs per node (default: 8)
- `nodes` - Number of nodes (default: 1)
- `nodelist` - Comma-separated node names to run on (e.g. `"node01,node02"`); when set, job is restricted to these nodes and automatic node health preflight is skipped
- `reservation` - SLURM reservation name; forwarded to srun health/cleanup commands and SBATCH directives
- `exclusive` - Exclusive node access (default: `true`)
-- `time` - Wall time limit HH:MM:SS (required)
-- `mem` - Memory per node (e.g., "64G")
-- `mail_user` - Email for notifications
-- `mail_type` - Notification types (BEGIN, END, FAIL, ALL)
+- `time` - Wall time limit HH:MM:SS (default: `24:00:00`)
+- `constraint` - SBATCH `--constraint` feature expression
+- `exclude` - Comma-separated nodes to exclude; node health preflight appends to this list
+- `modules` - Array of environment modules to `module load` in the job (default: `[]`)
+- `network_interface` - Interface exported as `NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME` (e.g. `ib0`)
+- `output_dir` - Directory for SLURM `.out`/`.err` files (default: `./slurm_results`)
+- `skip_gpus_directive` - Omit the `#SBATCH --gpus-per-node` directive (default: `false`). Set `true` on clusters that expose no GPU GRES and reject any job script carrying it; allocation then relies on `exclusive` / `nproc_per_node`.
+
+Node health preflight, shared-storage, and results-collection keys (`enable_node_check`,
+`auto_cleanup_nodes`, `allow_submit_without_clean_nodes`, `verbose_node_check`,
+`shared_workspace`, `results_dir`) are documented in
+[examples/slurm-configs/README.md](../examples/slurm-configs/README.md).
### Multi-Node SLURM
@@ -523,7 +600,7 @@ Automatically applies (see presets under `src/madengine/deployment/presets/k8s/`
**Supported Launchers:**
- `torchrun` - PyTorch DDP/FSDP
- `deepspeed` - ZeRO optimization
-- `megatron` - Large transformers (K8s + SLURM)
+- `megatron-lm` - Large transformers (K8s + SLURM)
- `torchtitan` - LLM pre-training
- `primus` - Primus unified pretrain
- `vllm` - LLM inference
@@ -558,14 +635,12 @@ See [Launchers Guide](launchers.md) for details.
"launcher": "vllm",
"nnodes": 2,
"nproc_per_node": 4
- },
- "vllm": {
- "tensor_parallel_size": 4,
- "pipeline_parallel_size": 1
}
}
```
+`nproc_per_node` is exported into the container as `VLLM_TENSOR_PARALLEL_SIZE`; there is no separate `vllm.*` config block for tensor/pipeline parallel sizing.
+
## Profiling Configuration
### Basic Profiling
@@ -656,14 +731,11 @@ Configure in `data.json` (MAD package root):
```json
{
- "data_sources": {
- "model_data": {
- "nas": {"path": "/home/datum"},
- "minio": {"path": "s3://datasets/datum"},
- "aws": {"path": "s3://datasets/datum"}
- }
- },
- "mirrorlocal": "/tmp/local_mirror"
+ "model_data": {
+ "nas": {"path": "/home/datum", "mirrorlocal": "/tmp/local_mirror"},
+ "minio": {"path": "s3://datasets/datum"},
+ "aws": {"path": "s3://datasets/datum"}
+ }
}
```
@@ -678,13 +750,13 @@ Configure in `credential.json` (MAD package root):
"password": "your_token",
"repository": "myorg"
},
- "AMD_GITHUB": {
+ "PUBLIC_GITHUB_ROCM_KEY": {
"username": "github_username",
- "password": "github_token"
+ "token": "github_token"
},
"MAD_AWS_S3": {
- "username": "aws_access_key",
- "password": "aws_secret_key"
+ "USERNAME": "aws_access_key",
+ "PASSWORD": "aws_secret_key"
}
}
```
@@ -697,6 +769,42 @@ export MAD_DOCKERHUB_PASSWORD=mytoken
export MAD_DOCKERHUB_REPO=myorg
```
+### Registry Authentication
+
+madengine reuses an existing `docker login` β including an organization access
+token (OAT) β rather than requiring credentials to be duplicated into
+`credential.json`. Ambient credentials are read from
+`${DOCKER_CONFIG:-~/.docker}/config.json` exactly as the Docker CLI reads them,
+covering `auths` entries, `credHelpers`, and `credsStore`.
+
+Blank values are treated as **not configured**, not as credentials. A placeholder
+entry such as `{"username": "", "password": ""}` will never override or break a
+working `docker login`.
+
+| Existing `docker login` | Credentials in `credential.json` / env | Behavior |
+|---|---|---|
+| no | yes (non-blank) | `docker login` with the configured credentials |
+| yes | yes (non-blank) | `docker login` with the configured credentials (explicit wins) |
+| yes | absent or blank | Reuse the existing login; no `docker login` is run |
+| no | absent or blank | Push fails with an actionable error; pull warns and continues |
+
+Before `docker build`, madengine logs in to the base image's registry only when
+that registry has no existing login and usable credentials are configured, so a
+node authenticated with an OAT is never re-authenticated.
+
+Relevant environment variables:
+
+- `DOCKER_CONFIG` β directory holding `config.json` (default `~/.docker`)
+- `MAD_SKIP_DOCKER_LOGIN=1` β never run `docker login`; always defer to the
+ credentials the machine already has
+
+If a base image pull is denied, madengine distinguishes the two causes:
+credentials were rejected (authentication β supply credentials or run
+`docker login`), versus credentials were accepted but the registry granted no
+pull scope for that repository (`insufficient_scope` β an authorization problem,
+where the access token needs to be scoped to the repository and re-running
+`docker login` will not help).
+
## Configuration Priority
For Kubernetes/SLURM deployments:
diff --git a/docs/contributing.md b/docs/contributing.md
index 5e53a300..078b23bc 100644
--- a/docs/contributing.md
+++ b/docs/contributing.md
@@ -61,7 +61,7 @@ pytest
pytest --cov=src/madengine --cov-report=html
# Run specific test file
-pytest tests/test_cli.py
+pytest tests/unit/test_cli.py
# Run tests matching pattern
pytest -k "test_build"
diff --git a/docs/deployment.md b/docs/deployment.md
index fa03e7f5..c913b117 100644
--- a/docs/deployment.md
+++ b/docs/deployment.md
@@ -13,24 +13,28 @@ Deployment is configured via `--additional-context` and happens automatically du
## Deployment Workflow
+Build once, then deploy the resulting manifest to any target:
+
+```mermaid
+flowchart LR
+ B["1. Build Phase
(local or CI/CD)
madengine build --tags model"] --> M[(build_manifest.json
+ image in registry)]
+ M --> D["2. Deploy Phase
madengine run --manifest-file build_manifest.json
--additional-context '{...}'"]
+ D --> T{detect target}
+ T -->|k8s / kubernetes| K[K8s Job]
+ T -->|slurm| S[SLURM script]
+ T -->|neither| L[Local Docker]
```
-βββββββββββββββββββββββββββββββββββββββββββββββ
-β 1. Build Phase (Local or CI/CD) β
-β madengine build --tags model β
-β β Creates Docker image β
-β β Pushes to registry β
-β β Generates build_manifest.json β
-βββββββββββββββββββββββββββββββββββββββββββββββ
- β
-βββββββββββββββββββββββββββββββββββββββββββββββ
-β 2. Deploy Phase (Run with Context) β
-β madengine run β
-β --manifest-file build_manifest.json β
-β --additional-context '{"deploy":...}' β
-β β Detects deployment target β
-β β Creates K8s Job or SLURM script β
-β β Submits and monitors execution β
-βββββββββββββββββββββββββββββββββββββββββββββββ
+
+### Deployment Target Inference
+
+No explicit `deploy` field is required β the target is inferred from the config structure (Convention over Configuration):
+
+```mermaid
+flowchart TD
+ A[additional_context] --> Q{which key?}
+ Q -->|k8s / kubernetes| K[Kubernetes deployment]
+ Q -->|slurm| S[SLURM deployment]
+ Q -->|neither| L[Local Docker execution]
```
## Kubernetes Deployment
@@ -90,11 +94,12 @@ The deployment target is automatically detected from the `k8s` key in the config
}
```
-**Configuration Priority:**
-1. User config (`--additional-context-file`)
-2. Profile presets (single-gpu/multi-gpu)
-3. GPU vendor presets (AMD/NVIDIA)
-4. Base defaults
+**Configuration Priority (lowest to highest precedence):**
+1. Base defaults
+2. GPU vendor preset (AMD/NVIDIA)
+3. GPU vendor multi-GPU preset (AMD only, for multi-GPU/multi-node)
+4. Profile preset (single-gpu/multi-gpu/multi-node)
+5. User config (`--additional-context-file`) β highest precedence, applied last
See [examples/k8s-configs/](../examples/k8s-configs/) for complete examples.
@@ -229,9 +234,7 @@ The deployment target is automatically detected from the `slurm` key in the conf
"qos": "normal",
"gpus_per_node": 8,
"nodes": 1,
- "time": "24:00:00",
- "mail_user": "user@example.com",
- "mail_type": "ALL"
+ "time": "24:00:00"
}
}
```
@@ -245,10 +248,7 @@ The deployment target is automatically detected from the `slurm` key in the conf
- `nodelist`: Comma-separated node names to run on (e.g. `"node01,node02"`); when set, job runs only on these nodes and node health preflight is skipped
- `reservation`: SLURM reservation name; forwarded to srun health/cleanup commands
- `time`: Wall time limit (HH:MM:SS)
-- `mem`: Memory per node (e.g., "64G")
- `exclusive`: Exclusive node access (default: `true`)
-- `mail_user`: Email for job notifications
-- `mail_type`: Notification types (BEGIN, END, FAIL, ALL)
See [examples/slurm-configs/](../examples/slurm-configs/) for complete examples.
diff --git a/docs/img/architecture_overview.png b/docs/img/architecture_overview.png
deleted file mode 100755
index 7bf972b3..00000000
Binary files a/docs/img/architecture_overview.png and /dev/null differ
diff --git a/docs/img/distributed_workflow.png b/docs/img/distributed_workflow.png
deleted file mode 100755
index a6723b44..00000000
Binary files a/docs/img/distributed_workflow.png and /dev/null differ
diff --git a/docs/installation.md b/docs/installation.md
index c8c9eb4c..e4cbd658 100644
--- a/docs/installation.md
+++ b/docs/installation.md
@@ -91,7 +91,6 @@ madengine run --tags dummy \
```bash
# Check installation
madengine --version
-madengine --version
# Test basic functionality (requires MAD package)
cd /path/to/MAD
diff --git a/docs/launchers.md b/docs/launchers.md
index 227557aa..658b252a 100644
--- a/docs/launchers.md
+++ b/docs/launchers.md
@@ -19,7 +19,7 @@ madengine provides unified support for multiple distributed frameworks, enabling
| **Primus** | Training | Megatron/TorchTitan/Jax via Primus config | β
| β
| β
|
| **vLLM** | Inference | High-throughput LLM serving | β
| β
| β
|
| **SGLang** | Inference | Fast LLM inference | β
| β
| β
|
-| **SGLang Disaggregated** | Inference | Large-scale disaggregated inference | β
| β
| β
(min 3) |
+| **SGLang Disaggregated** | Inference | Large-scale disaggregated inference | β
| β
| β
(min 2 on SLURM, 3 on K8s) |
| **slurm_multi** | Escape hatch | Self-managed multi-container topologies | β | β
| β
|
---
@@ -69,11 +69,13 @@ madengine run --manifest-file build_manifest.json
"launcher": "torchrun",
"nnodes": 2,
"nproc_per_node": 8,
- "master_port": 29500
+ "port": 29500
}
}
```
+Note: `distributed.port` sets the master port on **SLURM** (`distributed.get("port", 29500)`). On **Kubernetes**, the master port is instead read from a separate top-level `"launcher"` object (`{"launcher": {"master_port": 29500}}`), not from `distributed`.
+
**Features**:
- Automatic rank assignment
- NCCL backend for GPU communication
@@ -136,7 +138,7 @@ madengine run --manifest-file build_manifest.json
```json
{
"distributed": {
- "launcher": "megatron",
+ "launcher": "megatron-lm",
"nnodes": 4,
"nproc_per_node": 8
}
@@ -276,8 +278,6 @@ Optional **`primus.backend`** (e.g. `MaxText`, `megatron`) emits `export BACKEND
**Container image**: Prefer `docker/primus.ubuntu.amd.Dockerfile` with `COPY scripts/Primus/ /workspace/Primus/` and `PRIMUS_ROOT=/workspace/Primus`. On **Kubernetes**, the Jobβs emptyDir hides image files under `/workspace`; madengine bundles `scripts/Primus/examples/...` into the ConfigMap as `Primus/examples/...` so the init container recreates `/workspace/Primus`. `run.sh` resolves `PRIMUS_ROOT` in that order (see script comments).
**Examples**:
-- SLURM: `examples/slurm-configs/minimal/primus-minimal.json`
-- K8s: `examples/k8s-configs/minimal/primus-minimal.json`
- K8s (Primus vs upstream workload API, MaxText caveats, TorchTitan/Megatron/MaxText sample JSON): `examples/k8s-configs/README.md` section **Primus on Kubernetes**
---
@@ -314,13 +314,13 @@ Optional **`primus.backend`** (e.g. `MaxText`, `megatron`) emits `export BACKEND
**Architecture**:
- Single-node: TP across GPUs, no Ray
- Multi-node (K8s): Data Parallelism with independent replicas per pod
-- Multi-node (SLURM): TP + PP with Ray cluster
+- Multi-node (SLURM): Data Parallelism (one vLLM serve per node, TP only on that node, no shared Ray cluster)
**Environment Variables**:
```bash
VLLM_TENSOR_PARALLEL_SIZE=4
VLLM_PIPELINE_PARALLEL_SIZE=1
-VLLM_DISTRIBUTED_BACKEND="auto" # or "ray" for multi-node
+VLLM_DISTRIBUTED_BACKEND="auto" # or "none" for multi-node SLURM (data parallel)
```
**Examples**:
@@ -419,25 +419,31 @@ SGLang Disaggregated separates inference into specialized node pools:
```
**Minimum Requirements**:
-- **Nodes**: Minimum 3 nodes (1 proxy + 1 prefill + 1 decode)
+- **Nodes**: **SLURM** β minimum 2 nodes (co-located proxy on the first prefill node + 1 prefill + 1 decode). **Kubernetes** β minimum 3 nodes (dedicated proxy + 1 prefill + 1 decode); a 2-node co-located topology is rejected.
- **GPUs**: Minimum 1 GPU per node (for tensor parallelism)
- **Network**: High-speed interconnect (InfiniBand recommended for production)
+The proxy/router is started by the model's `run.sh`, not by the launcher, so on
+SLURM both layouts are valid: a dedicated proxy node (`1 + xP + yD == nnodes`) or
+a proxy co-located on the first prefill node (`xP + yD == nnodes`).
+
**Node Roles**:
1. **Proxy Node (Rank 0)**: Load balancer, request router (mini_lb)
2. **Prefill Nodes**: Process input prompts, generate KV cache
3. **Decode Nodes**: Receive KV cache, generate output tokens
**Automatic Split (Default)**:
-- Uses 40/60 golden ratio for prefill/decode
-- Formula: `prefill = max(1, (nnodes - 1) * 2 // 5)`
+- Uses 40/60 golden ratio for prefill/decode across the worker nodes
+- Formula: `prefill = max(1, (nnodes - 1) * 2 // 5)`, `decode = nnodes - 1 - prefill`
+- `nnodes == 2` is special-cased on SLURM to 1 prefill + 1 decode with a co-located proxy (the general formula would yield 0 decode nodes)
| Total Nodes | Proxy | Prefill | Decode |
|-------------|-------|---------|--------|
-| 3 | 1 | 1 (33%) | 1 (33%) |
-| 5 | 1 | 2 (40%) | 2 (40%) |
-| 7 | 1 | 2 (29%) | 4 (57%) |
-| 11 | 1 | 4 (40%) | 6 (60%) |
+| 2 (SLURM only) | co-located | 1 | 1 |
+| 3 | 1 | 1 | 1 |
+| 5 | 1 | 1 | 3 |
+| 7 | 1 | 2 | 4 |
+| 11 | 1 | 4 | 6 |
**Custom Split (NEW Feature!)**:
@@ -469,7 +475,8 @@ Override automatic split based on workload characteristics:
**Validation Rules**:
- `prefill_nodes >= 1`
- `decode_nodes >= 1`
-- `prefill_nodes + decode_nodes + 1 == nnodes`
+- `prefill_nodes + decode_nodes + 1 == nnodes` (dedicated proxy node), **or**
+ `prefill_nodes + decode_nodes == nnodes` (proxy co-located on the first prefill node, SLURM)
**Features**:
- Disaggregated prefill/decode architecture
@@ -529,14 +536,14 @@ SGLANG_NODE_IPS="10.0.0.1,10.0.0.2,..."
**Performance Tuning**:
```bash
# Start with automatic split
-madengine run --tags model --config minimal-config.json
+madengine run --tags model --additional-context-file minimal-config.json
# Monitor bottleneck (prefill latency vs decode throughput)
# If prefill is bottleneck β increase prefill nodes
# If decode is bottleneck β increase decode nodes
# Apply custom split
-madengine run --tags model --config custom-split-config.json
+madengine run --tags model --additional-context-file custom-split-config.json
```
**Troubleshooting**:
@@ -694,6 +701,34 @@ madengine run --manifest-file build_manifest.json
---
+## Parallelism Capabilities
+
+How each launcher handles the various parallelism strategies. `β
Auto` = supported and configured by madengine; `βManual` = supported by the launcher but requires user configuration; `βLimited` / `βDisabled` = launcher or platform limitation.
+
+| Launcher | Tensor Parallel (TP) | Pipeline Parallel (PP) | Data Parallel (DP) | Context Parallel (CP) | FSDP/ZeRO | Expert Parallel (EP) | Primary Use Case |
+|----------|----------------------|------------------------|--------------------|------------------------|-----------|----------------------|------------------|
+| **torchrun** | βManual | βNo | βManual (DDP) | βNo | βManual (FSDP) | βNo | General distributed training |
+| **TorchTitan** | β
Auto | β
Auto | β
Auto (FSDP2) | βManual | β
Auto (FSDP2) | βNo | Large-scale LLM pre-training |
+| **DeepSpeed** | βManual | βManual | β
Auto (ZeRO) | βNo | β
Auto (ZeRO) | βNo | Memory-efficient training |
+| **Megatron-LM** | β
Auto | β
Auto | β
Implicit | β
Auto | βNo | βNo | Large transformer training |
+| **Primus** | βManual | βManual | βManual | βManual | βManual | βNo | Unified pretrain (experiment YAML; backend-specific) |
+| **vLLM** | β
Auto | SLURM: β
Auto (Multi) / K8s: βDisabled | β
Auto (Replicas) | βNo | βNo | βManual | High-throughput inference |
+| **SGLang** | β
Auto | SLURM: β
Auto (Multi) / K8s: βDisabled | βLimited | βNo | βNo | βNo | Inference + structured gen |
+| **SGLang PD Disagg** | β
Auto | βNo | β
Role-based | βNo | βNo | βNo | Optimized prefill/decode |
+
+## Infrastructure Capabilities
+
+| Feature | Local | Kubernetes | SLURM |
+|---------|-------|-----------|-------|
+| **Execution** | Docker containers | K8s Jobs | SLURM jobs |
+| **Multi-Node** | β | β
Indexed Jobs | β
Job arrays |
+| **Resource Mgmt** | Manual | Declarative (YAML) | Batch scheduler |
+| **Monitoring** | Docker logs | kubectl/dashboard | squeue/scontrol |
+| **Auto-scaling** | β | β
| β |
+| **Network** | Host | CNI plugin | InfiniBand/Ethernet |
+
+---
+
## Configuration Best Practices
### 1. Launcher Selection
@@ -835,7 +870,7 @@ SGLANG_NODE_RANK=${SLURM_PROCID}
```bash
Error: Unknown launcher type 'xyz'
```
-Solution: Use one of: `torchrun`, `deepspeed`, `megatron`, `torchtitan`, `primus`, `vllm`, `sglang`, `sglang-disagg`, `slurm_multi` (or `slurm-multi`)
+Solution: Use one of: `torchrun`, `deepspeed`, `megatron-lm`, `torchtitan`, `primus`, `vllm`, `sglang`, `sglang-disagg`, `slurm_multi` (or `slurm-multi`)
**2. Multi-Node Communication Fails**
```bash
@@ -878,7 +913,7 @@ madengine provides `$MAD_MULTI_NODE_RUNNER` for frameworks that use torchrun:
#!/bin/bash
# Your model script
-# For torchrun/deepspeed/megatron/torchtitan
+# For torchrun/deepspeed/megatron-lm/torchtitan
$MAD_MULTI_NODE_RUNNER your_training_script.py --args
# For primus (no MAD_MULTI_NODE_RUNNER; use run.sh β run_pretrain.sh; PRIMUS_* / BACKEND set by madengine)
diff --git a/docs/profiling.md b/docs/profiling.md
index 5e421ec8..367313ef 100644
--- a/docs/profiling.md
+++ b/docs/profiling.md
@@ -93,6 +93,8 @@ madengine uses `rocprof_wrapper.sh` to automatically handle the transition betwe
```
3. **Backward Compatibility:** The `--` works with both rocprof and rocprofv3, ensuring your configurations work across ROCm versions
+**Related presets:** `rocprof_hip_only` (HIP trace only) and `rocprof_sys` (system trace only) wrap the same `rocprof_wrapper.sh` with fixed flags, for when you don't need a custom `cmd`.
+
**Example - Custom Command with Wrapper:**
```json
{
@@ -258,6 +260,8 @@ ROCprofv3 is the next-generation profiler for ROCm 7.0+ with enhanced features a
- **Metrics**: Program counter sampling at 1000 Hz
- **Output Format**: Perfetto trace with PC samples
+**Other presets:** `rocprofv3` (bare `--runtime-trace` invocation, no wrapper), `rocprofv3_agent` (CSV stats/kernel trace via `rocprof_wrapper.sh`), and `rocprofv3_agent_counter` (same as `rocprofv3_agent` plus hardware counters from `counters/instruction_mix.txt`) are also available as `{"name": "..."}` tool entries.
+
#### Using Pre-Configured Profiles
madengine provides ready-to-use configuration files in `examples/profiling-configs/`:
@@ -321,6 +325,7 @@ Custom counter files are in `scripts/common/tools/counters/`:
- `memory_bound.txt` - Cache and memory metrics
- `communication_bound.txt` - PCIe and synchronization metrics
- `full_profile.txt` - Comprehensive metrics
+- `instruction_mix.txt` - Instruction mix metrics (used by `rocprofv3_agent_counter`)
Create your own counter file:
```text
@@ -360,6 +365,8 @@ Trace rocBLAS API calls and configurations:
**Use Case:** Analyze BLAS operations, identify optimization opportunities
+**Related:** `hipblaslt_trace` traces hipBLASLt calls the same way (sets `HIPBLASLT_TRACE=1`).
+
### miopen_trace - MIOpen Library Tracing
Trace MIOpen API calls for deep learning operations:
@@ -509,9 +516,6 @@ Profile real-time GPU memory consumption:
}
```
This will generate both `gpu_info_power_profiler_output.csv` and `gpu_info_vram_profiler_output.csv`.
-- `SAMPLING_RATE` - Sampling interval in seconds
-- `MODE` - Must be `"vram"` for this tool
-- `DUAL-GCD` - Enable dual-GCD mode
**Supported Platforms:** ROCm and CUDA
@@ -666,8 +670,8 @@ madengine run --tags model \
{
"name": "gpu_info_power_profiler",
"env_vars": {
- "DEVICE": "all",
- "SAMPLING_RATE": "0.1"
+ "POWER_DEVICE": "all",
+ "POWER_SAMPLING_RATE": "0.1"
}
},
{"name": "rccl_trace"}
@@ -777,7 +781,7 @@ Balance detail vs. overhead:
{
"name": "gpu_info_power_profiler",
"env_vars": {
- "SAMPLING_RATE": "1.0" // Less overhead, less detail
+ "POWER_SAMPLING_RATE": "1.0" // Less overhead, less detail
}
}
]
@@ -923,15 +927,19 @@ Tool defaults are defined in `scripts/common/tools.json`:
```json
{
"rocprof": {
- "cmd": "rocprof --hip-trace",
- "env_vars": {}
+ "cmd": "bash ../scripts/common/tools/rocprof_wrapper.sh --runtime-trace --",
+ "env_vars": {},
+ "post_scripts": [
+ {"path": "scripts/common/post_scripts/trace.sh", "args": "rocprof"}
+ ]
},
"gpu_info_power_profiler": {
"env_vars": {
- "DEVICE": "0",
- "SAMPLING_RATE": "0.1",
- "MODE": "power",
- "DUAL-GCD": "false"
+ "POWER_DEVICE": "all",
+ "POWER_SAMPLING_RATE": "0.1",
+ "POWER_MODE": "power",
+ "POWER_DUAL_GCD": "false",
+ "POWER_OUTPUT_FILE": "gpu_info_power_profiler_output.csv"
}
}
}
@@ -956,8 +964,7 @@ Validate [TheRock](https://github.com/ROCm/TheRock) ROCm installations before ru
```bash
madengine run --tags dummy_therock \
- --tools therock_check \
- --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'
+ --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU", "tools": [{"name": "therock_check"}]}'
```
**Standalone detection:**
diff --git a/docs/superpowers/plans/2026-08-27-pinned-image-digest.md b/docs/superpowers/plans/2026-08-27-pinned-image-digest.md
new file mode 100644
index 00000000..dad4b9d3
--- /dev/null
+++ b/docs/superpowers/plans/2026-08-27-pinned-image-digest.md
@@ -0,0 +1,1623 @@
+# 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
new file mode 100644
index 00000000..08d69aea
--- /dev/null
+++ b/docs/superpowers/specs/2026-08-27-pinned-image-digest-design.md
@@ -0,0 +1,156 @@
+# 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/docs/usage.md b/docs/usage.md
index f06557bb..a79d495b 100644
--- a/docs/usage.md
+++ b/docs/usage.md
@@ -46,7 +46,7 @@ madengine provides five main commands:
| `build` | Build Docker images | `--tags`, `--registry`, `--batch-manifest` |
| `run` | Execute models | `--tags`, `--manifest-file`, `--timeout` |
| `report` | Generate HTML reports | `to-html`, `to-email` |
-| `database` | Upload to MongoDB | `--csv-file`, `--database-name` |
+| `database` | Upload to MongoDB | `--file`, `--db` |
For complete command options and detailed examples, see **[CLI Command Reference](cli-reference.md)**.
@@ -66,11 +66,11 @@ madengine run --tags model
# madengine build --tags model --additional-context '{"gpu_vendor": "NVIDIA", "guest_os": "CENTOS"}'
# Generate HTML report
-madengine report to-html --csv-file perf_entry.csv
+madengine report to-html --csv-file-path perf_entry.csv
# Upload to MongoDB
-madengine database --csv-file perf_entry.csv \
- --database-name mydb --collection-name results
+madengine database --file perf_entry.csv \
+ --database mydb --collection results
```
## Model Discovery
@@ -82,7 +82,7 @@ madengine supports three discovery methods:
Central model definitions in MAD package root:
```bash
-madengine discover --tags dummy pyt_huggingface_bert
+madengine discover --tags dummy --tags pyt_huggingface_bert
```
### 2. Directory-Specific Models
@@ -117,18 +117,32 @@ Creates `build_manifest.json`:
```json
{
- "models": [
- {
- "model_name": "my_model",
- "image": "localhost:5000/my_model:20240115_123456",
- "tag": "my_model"
+ "built_images": {
+ "ci-my_model_ubuntu": {
+ "model": "my_model",
+ "docker_image": "ci-my_model_ubuntu",
+ "dockerfile": "docker/my_model.ubuntu.amd.Dockerfile",
+ "build_duration": 42.3,
+ "registry": "localhost:5000"
}
- ],
- "registry": "localhost:5000",
- "build_timestamp": "2024-01-15T12:34:56Z"
+ },
+ "built_models": {
+ "ci-my_model_ubuntu": {
+ "name": "my_model",
+ "dockerfile": "my_model",
+ "n_gpus": "1"
+ }
+ },
+ "context": {
+ "gpu_vendor": "AMD",
+ "guest_os": "UBUNTU"
+ },
+ "credentials_required": []
}
```
+`built_images` and `built_models` are both keyed by the built Docker image name. Depending on the build, the manifest may also include `deployment_config` and `summary` keys.
+
### Build with Deployment Config
Include deployment configuration:
@@ -493,7 +507,7 @@ Convert performance CSV files to viewable HTML reports:
```bash
# Single CSV to HTML
-madengine report to-html --csv-file perf_entry.csv
+madengine report to-html --csv-file-path perf_entry.csv
# Result: Creates perf_entry.html in same directory
```
@@ -532,13 +546,13 @@ export MONGO_PASSWORD=secretpassword
# Upload results
madengine database \
- --csv-file perf_entry.csv \
- --database-name performance_tracking \
- --collection-name model_runs
+ --file perf_entry.csv \
+ --database performance_tracking \
+ --collection model_runs
# Upload specific results
madengine database \
- --csv-file results/perf_mi300.csv \
+ --file results/perf_mi300.csv \
--db benchmarks \
--collection mi300_results
```
@@ -547,15 +561,15 @@ madengine database \
```bash
# 1. Run benchmarks
-madengine run --tags model1 model2 model3 \
+madengine run --tags model1 --tags model2 --tags model3 \
--output perf_entry.csv
# 2. Generate HTML report
-madengine report to-html --csv-file perf_entry.csv
+madengine report to-html --csv-file-path perf_entry.csv
# 3. Upload to database
madengine database \
- --csv-file perf_entry.csv \
+ --file perf_entry.csv \
--db benchmarks \
--collection daily_runs
@@ -586,7 +600,7 @@ Configure distributed training:
**Supported Launchers:**
- `torchrun` - PyTorch DDP/FSDP
- `deepspeed` - ZeRO optimization
-- `megatron` - Large transformers (K8s + SLURM)
+- `megatron-lm` - Large transformers (K8s + SLURM)
- `torchtitan` - LLM pre-training
- `vllm` - LLM inference
- `sglang` - Structured generation
@@ -671,6 +685,8 @@ madengine build --tags model --clean-docker-cache --verbose
| `MAD_DOCKERHUB_USER` | Docker Hub username | `"myusername"` |
| `MAD_DOCKERHUB_PASSWORD` | Docker Hub password | `"mytoken"` |
| `MAD_DOCKERHUB_REPO` | Docker Hub repository | `"myorg"` |
+| `DOCKER_CONFIG` | Directory holding the Docker `config.json` whose existing login is reused | `/etc/docker-oat` |
+| `MAD_SKIP_DOCKER_LOGIN` | Set to `1` to never run `docker login` and always defer to the machine's existing credentials | `"1"` |
## Best Practices
@@ -764,7 +780,7 @@ if [ $? -eq 0 ]; then
# Generate and upload results
madengine report to-email --output ci_results.html
madengine database \
- --csv-file perf.csv \
+ --file perf.csv \
--db ci_results \
--collection ${CI_BUILD_ID}
else
diff --git a/examples/k8s-configs/README.md b/examples/k8s-configs/README.md
index 40843ab3..fbd327dd 100644
--- a/examples/k8s-configs/README.md
+++ b/examples/k8s-configs/README.md
@@ -103,16 +103,16 @@ MODEL_DIR=tests/fixtures/dummy madengine run \
```bash
# For single GPU testing
-cp examples/k8s-configs/01-single-node-single-gpu.json my-config.json
+cp examples/k8s-configs/basic/01-native-single-node-single-gpu.json my-config.json
# For multi-GPU (2 GPUs)
-cp examples/k8s-configs/02-single-node-multi-gpu.json my-config.json
+cp examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu.json my-config.json
# For multi-node distributed (2 nodes Γ 2 GPUs)
-cp examples/k8s-configs/03-multi-node-basic.json my-config.json
+cp examples/k8s-configs/basic/03-torchrun-multi-node-basic.json my-config.json
# For data provider with auto-PVC
-cp examples/k8s-configs/06-data-provider-with-pvc.json my-config.json
+cp examples/k8s-configs/basic/06-data-provider-with-pvc.json my-config.json
```
#### 2. Customize for Your Cluster (Optional)
@@ -158,10 +158,10 @@ Located in [`minimal/`](minimal/) directory:
| File | Description | GPU Count |
|------|-------------|-----------|
-| [`minimal/single-gpu-minimal.json`](minimal/single-gpu-minimal.json) | Single GPU with auto-defaults | 1 |
-| [`minimal/multi-gpu-minimal.json`](minimal/multi-gpu-minimal.json) | Multi-GPU with auto-defaults | 2 |
-| [`minimal/multi-node-minimal.json`](minimal/multi-node-minimal.json) | Multi-node with auto-defaults | 2Γ2 |
-| [`minimal/nvidia-gpu-minimal.json`](minimal/nvidia-gpu-minimal.json) | NVIDIA GPUs with auto-defaults | 4 |
+| [`minimal/torchrun-single-gpu-minimal.json`](minimal/torchrun-single-gpu-minimal.json) | Single GPU with auto-defaults | 1 |
+| [`minimal/torchrun-multi-gpu-minimal.json`](minimal/torchrun-multi-gpu-minimal.json) | Multi-GPU with auto-defaults | 2 |
+| [`minimal/torchrun-multi-node-minimal.json`](minimal/torchrun-multi-node-minimal.json) | Multi-node with auto-defaults | 2Γ2 |
+| [`minimal/torchrun-nvidia-gpu-minimal.json`](minimal/torchrun-nvidia-gpu-minimal.json) | NVIDIA GPUs with auto-defaults | 4 |
| [`minimal/custom-namespace-minimal.json`](minimal/custom-namespace-minimal.json) | Shows override examples | 1 |
**Distributed Launchers:**
@@ -169,7 +169,6 @@ Located in [`minimal/`](minimal/) directory:
| File | Launcher | Description | GPUs |
|------|----------|-------------|------|
| [`minimal/torchtitan-single-node-minimal.json`](minimal/torchtitan-single-node-minimal.json) | TorchTitan | LLM pre-training (single-node) | 8 |
-| [`minimal/primus-minimal.json`](minimal/primus-minimal.json) | primus | Primus pretrain (edit `distributed.primus.config_path`) | 2 |
| [`minimal/vllm-single-node-minimal.json`](minimal/vllm-single-node-minimal.json) | vLLM | LLM inference (single-node) | 4 |
| [`minimal/sglang-single-node-minimal.json`](minimal/sglang-single-node-minimal.json) | SGLang | LLM inference (single-node) | 4 |
@@ -211,24 +210,24 @@ Set `distributed.primus.config_path` to your YAML under the Primus repo layout.
**MaxText caveat:** For **multi-node** MaxText, Primus `run_pretrain.sh` may run **in-container `apt` installs** (InfiniBand-related packages). Many clusters disallow that unless the image is pre-baked or policy allows it. madengine logs a **warning** when MaxText is detected (`backend` or path) and `nnodes > 1`.
-Primus examples under [`basic/`](basic/): [`primus-single-node-multi-gpu.json`](basic/primus-single-node-multi-gpu.json) (one pod, multi-GPU) and [`primus-multi-node.json`](basic/primus-multi-node.json) (Indexed Job). The same files work for TorchTitan, Megatron, and MaxText: set `distributed.primus.config_path` to your experiment YAML, and rely on madengineβs `BACKEND` inference from the model name (`primus_pretrain/__...`) or set `distributed.primus.backend` only when you need an explicit override. Use `docker_env_vars.HF_TOKEN` as a placeholder or runtime secrets β do not commit real tokens.
+To use Primus, start from any `distributed.launcher: "torchrun"` example under [`basic/`](basic/) or [`minimal/`](minimal/) and change `distributed.launcher` to `"primus"`, adding a nested `distributed.primus.config_path` pointing at your experiment YAML (see [Distributed Execution Fields](#distributed-execution-fields) below). This works for TorchTitan, Megatron, and MaxText: set `distributed.primus.config_path` to your experiment YAML, and rely on madengine's `BACKEND` inference from the model name (`primus_pretrain/__...`) or set `distributed.primus.backend` only when you need an explicit override. Use `docker_env_vars.HF_TOKEN` as a placeholder or runtime secrets β do not commit real tokens.
### Full Configs (Reference Examples)
Complete configurations showing all available fields:
-**Training Configs:**
+**Training Configs (basic/):**
| File | GPUs | Nodes | Launcher | Use Case |
|------|------|-------|----------|----------|
-| [`01-single-node-single-gpu.json`](01-single-node-single-gpu.json) | 1 | 1 | None | Basic testing, small models |
-| [`01-single-node-single-gpu-tools.json`](01-single-node-single-gpu-tools.json) | 1 | 1 | None | Single GPU + monitoring |
-| [`02-single-node-multi-gpu.json`](02-single-node-multi-gpu.json) | 2 | 1 | torchrun | Multi-GPU training |
-| [`02-single-node-multi-gpu-tools.json`](02-single-node-multi-gpu-tools.json) | 2 | 1 | torchrun | Multi-GPU + monitoring |
-| [`03-multi-node-basic.json`](03-multi-node-basic.json) | 2/node | 2 | torchrun | Multi-node basics (4 GPUs total) |
-| [`04-multi-node-advanced.json`](04-multi-node-advanced.json) | 2/node | 4 | torchrun | Production multi-node (8 GPUs) |
-| [`05-nvidia-gpu-example.json`](05-nvidia-gpu-example.json) | 4 | 1 | torchrun | NVIDIA GPUs (A100, H100) |
-| [`06-data-provider-with-pvc.json`](06-data-provider-with-pvc.json) | 2 | 1+ | torchrun | **Data provider with auto-PVC** |
+| [`basic/01-native-single-node-single-gpu.json`](basic/01-native-single-node-single-gpu.json) | 1 | 1 | None | Basic testing, small models |
+| [`basic/01-native-single-node-single-gpu-tools.json`](basic/01-native-single-node-single-gpu-tools.json) | 1 | 1 | None | Single GPU + monitoring |
+| [`basic/02-torchrun-single-node-multi-gpu.json`](basic/02-torchrun-single-node-multi-gpu.json) | 2 | 1 | torchrun | Multi-GPU training |
+| [`basic/02-torchrun-single-node-multi-gpu-tools.json`](basic/02-torchrun-single-node-multi-gpu-tools.json) | 2 | 1 | torchrun | Multi-GPU + monitoring |
+| [`basic/03-torchrun-multi-node-basic.json`](basic/03-torchrun-multi-node-basic.json) | 2/node | 2 | torchrun | Multi-node basics (4 GPUs total) |
+| [`basic/04-torchrun-multi-node-advanced.json`](basic/04-torchrun-multi-node-advanced.json) | 2/node | 4 | torchrun | Production multi-node (8 GPUs) |
+| [`basic/05-torchrun-nvidia-gpu-example.json`](basic/05-torchrun-nvidia-gpu-example.json) | 4 | 1 | torchrun | NVIDIA GPUs (A100, H100) |
+| [`basic/06-data-provider-with-pvc.json`](basic/06-data-provider-with-pvc.json) | 2 | 1+ | torchrun | **Data provider with auto-PVC** |
**Distributed Launcher Configs (basic/):**
@@ -237,8 +236,6 @@ Complete configurations showing all available fields:
| [`basic/torchtitan-multi-node-basic.json`](basic/torchtitan-multi-node-basic.json) | 8/node | 4 | TorchTitan | Llama 3.1 70B+ training |
| [`basic/vllm-multi-node-basic.json`](basic/vllm-multi-node-basic.json) | 4/node | 2 | vLLM | High-throughput inference |
| [`basic/sglang-multi-node-basic.json`](basic/sglang-multi-node-basic.json) | 4/node | 2 | SGLang | Distributed inference |
-| [`basic/primus-single-node-multi-gpu.json`](basic/primus-single-node-multi-gpu.json) | 8 | 1 | primus | Primus pretrain (single pod; edit `primus.config_path`) |
-| [`basic/primus-multi-node.json`](basic/primus-multi-node.json) | 8/node | 2+ | primus | Primus pretrain (multi-pod; edit `nnodes`, `primus.config_path`) |
---
@@ -248,26 +245,26 @@ Complete configurations showing all available fields:
| Scenario | Config File | GPUs | Nodes |
|----------|-------------|------|-------|
-| **Quick test** | `01-single-node-single-gpu.json` | 1 | 1 |
-| **Single GPU benchmark** | `01-single-node-single-gpu-tools.json` | 1 | 1 |
-| **Multi-GPU (2 GPUs)** | `02-single-node-multi-gpu.json` | 2 | 1 |
-| **Multi-GPU + monitoring** | `02-single-node-multi-gpu-tools.json` | 2 | 1 |
-| **Multi-node (4 GPUs)** | `03-multi-node-basic.json` | 2Γ2 | 2 |
-| **Multi-node (8 GPUs)** | `04-multi-node-advanced.json` | 2Γ4 | 4 |
-| **NVIDIA GPUs** | `05-nvidia-gpu-example.json` | 4 | 1 |
-| **With data download** | `06-data-provider-with-pvc.json` | 2 | 1+ |
+| **Quick test** | `basic/01-native-single-node-single-gpu.json` | 1 | 1 |
+| **Single GPU benchmark** | `basic/01-native-single-node-single-gpu-tools.json` | 1 | 1 |
+| **Multi-GPU (2 GPUs)** | `basic/02-torchrun-single-node-multi-gpu.json` | 2 | 1 |
+| **Multi-GPU + monitoring** | `basic/02-torchrun-single-node-multi-gpu-tools.json` | 2 | 1 |
+| **Multi-node (4 GPUs)** | `basic/03-torchrun-multi-node-basic.json` | 2Γ2 | 2 |
+| **Multi-node (8 GPUs)** | `basic/04-torchrun-multi-node-advanced.json` | 2Γ4 | 4 |
+| **NVIDIA GPUs** | `basic/05-torchrun-nvidia-gpu-example.json` | 4 | 1 |
+| **With data download** | `basic/06-data-provider-with-pvc.json` | 2 | 1+ |
### By Use Case
| Use Case | Recommended Config |
|----------|-------------------|
-| **Development/Testing** | `01-single-node-single-gpu.json` |
-| **Small models (BERT, ResNet)** | `01-single-node-single-gpu.json` |
-| **Medium models (GPT-2, Stable Diffusion)** | `02-single-node-multi-gpu.json` |
-| **Large models (LLaMA-13B)** | `03-multi-node-basic.json` |
-| **Very large models (LLaMA-70B+)** | `04-multi-node-advanced.json` |
-| **Models requiring datasets** | `06-data-provider-with-pvc.json` |
-| **Busy/shared clusters** | `02-single-node-multi-gpu.json` (2 GPUs) |
+| **Development/Testing** | `basic/01-native-single-node-single-gpu.json` |
+| **Small models (BERT, ResNet)** | `basic/01-native-single-node-single-gpu.json` |
+| **Medium models (GPT-2, Stable Diffusion)** | `basic/02-torchrun-single-node-multi-gpu.json` |
+| **Large models (LLaMA-13B)** | `basic/03-torchrun-multi-node-basic.json` |
+| **Very large models (LLaMA-70B+)** | `basic/04-torchrun-multi-node-advanced.json` |
+| **Models requiring datasets** | `basic/06-data-provider-with-pvc.json` |
+| **Busy/shared clusters** | `basic/02-torchrun-single-node-multi-gpu.json` (2 GPUs) |
---
@@ -278,7 +275,7 @@ Complete configurations showing all available fields:
```bash
MODEL_DIR=tests/fixtures/dummy madengine build \
--tags dummy \
- --additional-context-file examples/k8s-configs/01-single-node-single-gpu.json \
+ --additional-context-file examples/k8s-configs/basic/01-native-single-node-single-gpu.json \
--registry dockerhub
MODEL_DIR=tests/fixtures/dummy madengine run \
@@ -291,7 +288,7 @@ MODEL_DIR=tests/fixtures/dummy madengine run \
```bash
MODEL_DIR=tests/fixtures/dummy madengine build \
--tags dummy_torchrun \
- --additional-context-file examples/k8s-configs/02-single-node-multi-gpu.json \
+ --additional-context-file examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu.json \
--registry dockerhub
MODEL_DIR=tests/fixtures/dummy madengine run \
@@ -304,7 +301,7 @@ MODEL_DIR=tests/fixtures/dummy madengine run \
```bash
MODEL_DIR=tests/fixtures/dummy madengine build \
--tags dummy_torchrun \
- --additional-context-file examples/k8s-configs/03-multi-node-basic.json \
+ --additional-context-file examples/k8s-configs/basic/03-torchrun-multi-node-basic.json \
--registry dockerhub
MODEL_DIR=tests/fixtures/dummy madengine run \
@@ -317,7 +314,7 @@ MODEL_DIR=tests/fixtures/dummy madengine run \
```bash
MODEL_DIR=tests/fixtures/dummy madengine build \
--tags dummy_torchrun_data_minio \
- --additional-context-file examples/k8s-configs/06-data-provider-with-pvc.json \
+ --additional-context-file examples/k8s-configs/basic/06-data-provider-with-pvc.json \
--registry dockerhub
MODEL_DIR=tests/fixtures/dummy madengine run \
@@ -347,7 +344,7 @@ kubectl get pvc madengine-shared-data
**Step 1: Use data provider config**
```bash
madengine build --tags dummy_torchrun_data_minio \
- --additional-context-file examples/k8s-configs/06-data-provider-with-pvc.json \
+ --additional-context-file examples/k8s-configs/basic/06-data-provider-with-pvc.json \
--registry dockerhub
```
@@ -465,7 +462,6 @@ To use an existing PVC instead of auto-creation:
"_comment": "Description of this configuration",
"gpu_vendor": "AMD|NVIDIA",
"guest_os": "UBUNTU",
- "deploy": "k8s",
"k8s": {
"kubeconfig": "~/.kube/config",
@@ -483,8 +479,7 @@ To use an existing PVC instead of auto-creation:
"node_selector": {},
"tolerations": [],
- "data_pvc": null, // Optional: for data providers
- "results_pvc": null // Optional: custom results storage
+ "data_pvc": null // Optional: for data providers
},
"distributed": {
@@ -512,10 +507,10 @@ To use an existing PVC instead of auto-creation:
| Field | Type | Required | Description |
|-------|------|----------|-------------|
-| `gpu_vendor` | string | **Yes** | `"AMD"` or `"NVIDIA"` |
-| `guest_os` | string | **Yes** | `"UBUNTU"`, `"RHEL"`, etc. |
-| `deploy` | string | **Yes** | Must be `"k8s"` |
-| `k8s` | object | **Yes** | Kubernetes configuration |
+| `gpu_vendor` | string | No (default `"AMD"`) | `"AMD"` or `"NVIDIA"` |
+| `guest_os` | string | No (default `"UBUNTU"`) | `"UBUNTU"` or `"CENTOS"` |
+| `deploy` | string | No | Not needed β deploy target is auto-inferred from presence of the `k8s`/`kubernetes` key (or `slurm`); if set, it must agree with the inferred target |
+| `k8s` | object | **Yes** | Kubernetes configuration (presence of this key is what triggers K8s deployment) |
| `distributed` | object | No | Distributed training (for torchrun) |
| `env_vars` | object | No | Custom environment variables |
| `debug` | boolean | No | Enable debug mode (saves manifests) |
@@ -551,7 +546,9 @@ To use an existing PVC instead of auto-creation:
|-------|------|---------|-------------|
| `image_pull_policy` | string | `"Always"` | `"Always"`, `"IfNotPresent"`, or `"Never"` |
| `backoff_limit` | integer | `3` | Retry attempts before marking failed |
-| `host_ipc` | boolean | `false` | Enable shared memory (required for multi-node) |
+| `allow_privileged_profiling` | boolean | `null` | Whether the pod runs with the elevated privileges profiling tools need. If unset (`null`), it's auto-enabled when a `tools` config is present; set explicitly to `true`/`false` to override |
+
+**Note:** `host_ipc` is **not** user-configurable β madengine always sets it automatically to `nnodes > 1` (enabled for multi-node jobs, disabled otherwise). Setting `k8s.host_ipc` in your config has no effect; it is not read anywhere.
**Optional - Node Selection:**
@@ -565,15 +562,18 @@ To use an existing PVC instead of auto-creation:
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `data_pvc` | string | `null` | Data PVC name (auto-created if using data provider) |
-| `results_pvc` | string | `null` | Results PVC name (auto-created by default) |
| `storage_class` | string | **`nfs-banff`** (preset, since 2.0.3) | Generic broad fallback for both the data PVC and the single-node results PVC when no more-specific key is set |
| `nfs_storage_class` | string | **`nfs-banff`** (preset) | RWX class for shared-data / multi-node results |
| `local_path_storage_class` | string | `null` (not in preset since 2.0.3; was **`local-path`** in β€ 2.0.2) | Optional RWO class for single-node `{job}-results`. Still honoured for backward compatibility |
| `data_storage_class` | string | **`nfs-banff`** (preset) | Overrides SC for shared-data only |
| `single_node_results_storage_class` | string | `null` | Overrides single-node results SC (falls back to `local_path_storage_class`, then `storage_class`) |
| `multi_node_results_storage_class` | string | `null` | Overrides multi-node results SC (`nfs_storage_class` if unset) |
+| `results_storage_size` | string | `"10Gi"` | Size of the per-job `{job}-results` PVC |
+| `data_storage_size` | string | `"100Gi"` | Size of the shared `madengine-shared-data` PVC |
| `recreate_shared_data_pvc` | boolean | **`false`** (preset) | If `true`, delete `madengine-shared-data` before create (data loss) |
+**Note:** `results_pvc` (the per-job results PVC name) is **not** user-configurable β madengine always creates and uses `{job_name}-results`. There is no config key that overrides this name.
+
#### Distributed Execution Fields
Configuration for distributed workloads (training and inference):
@@ -650,7 +650,7 @@ CPU: 16 (request), 32 (limit)
GPUs: 2 per node (4 total)
Memory: 64Gi per node
CPU: 16 per node
-host_ipc: true (required!)
+host_ipc: enabled automatically (nnodes > 1)
```
**Multi-Node Advanced (4 nodes Γ 2 GPUs):**
@@ -658,7 +658,7 @@ host_ipc: true (required!)
GPUs: 2 per node (8 total)
Memory: 128Gi per node
CPU: 24 per node
-host_ipc: true
+host_ipc: enabled automatically (nnodes > 1)
PVCs: Recommended for data and results
```
@@ -714,7 +714,7 @@ Write durable outputs under `/results//` in the container so each re
- `HSA_ENABLE_SDMA=0` - Disable SDMA for better P2P
**For multi-node:**
-- `host_ipc: true` - Required for shared memory
+- Shared memory (`host_ipc`) is enabled automatically when `distributed.nnodes > 1` β no config needed
- `HSA_FORCE_FINE_GRAIN_PCIE=1` - Cross-node communication
- `TORCH_NCCL_ASYNC_ERROR_HANDLING=1` - Better error handling
@@ -833,14 +833,7 @@ NCCL WARN ... Unable to find NCCL communicator
**Solutions:**
-1. **Enable host_ipc:**
-```json
-{
- "k8s": {
- "host_ipc": true // Required for multi-node!
- }
-}
-```
+1. **Host IPC is automatic β no config needed:** madengine always enables shared memory (`hostIPC`) for the pod whenever `distributed.nnodes > 1`; there is no `host_ipc` config key to set, and setting one has no effect. If communication is still failing, confirm `nnodes` in your `distributed` config is actually `> 1` (a value of `1` disables host IPC).
2. **Verify headless service:**
```bash
@@ -1008,7 +1001,7 @@ Use Case: Multi-GPU training, testing on busy clusters
Configuration: 2 nodes Γ 2 GPUs per node
Memory: 64Gi per node
CPU: 16 per node
-host_ipc: true (required!)
+host_ipc: enabled automatically (nnodes > 1)
Use Case: Distributed training development
```
@@ -1017,7 +1010,7 @@ Use Case: Distributed training development
Configuration: 4 nodes Γ 2 GPUs per node
Memory: 128Gi per node
CPU: 24 per node
-host_ipc: true
+host_ipc: enabled automatically (nnodes > 1)
PVCs: Recommended
Use Case: Large-scale production training
```
@@ -1031,7 +1024,7 @@ Use Case: Large-scale production training
```bash
# Use minimal config (defaults for everything)
madengine build --tags dummy \
- --additional-context-file examples/k8s-configs/01-single-node-single-gpu.json \
+ --additional-context-file examples/k8s-configs/basic/01-native-single-node-single-gpu.json \
--registry dockerhub
madengine run --manifest-file build_manifest.json
@@ -1042,7 +1035,7 @@ madengine run --manifest-file build_manifest.json
```bash
# Use 2 GPUs to avoid scheduling conflicts
madengine build --tags resnet50 \
- --additional-context-file examples/k8s-configs/02-single-node-multi-gpu.json \
+ --additional-context-file examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu.json \
--registry dockerhub
madengine run --manifest-file build_manifest.json --live-output
@@ -1053,7 +1046,7 @@ madengine run --manifest-file build_manifest.json --live-output
```bash
# Multi-node for large models
madengine build --tags llama_13b \
- --additional-context-file examples/k8s-configs/03-multi-node-basic.json \
+ --additional-context-file examples/k8s-configs/basic/03-torchrun-multi-node-basic.json \
--registry dockerhub
madengine run --manifest-file build_manifest.json --live-output
@@ -1064,7 +1057,7 @@ madengine run --manifest-file build_manifest.json --live-output
```bash
# Data provider with auto-PVC
madengine build --tags bert_large \
- --additional-context-file examples/k8s-configs/06-data-provider-with-pvc.json \
+ --additional-context-file examples/k8s-configs/basic/06-data-provider-with-pvc.json \
--registry dockerhub
madengine run --manifest-file build_manifest.json --live-output
@@ -1079,7 +1072,7 @@ kubectl exec -- ls -lh /data/
```bash
# Use *-tools.json variant for monitoring
madengine build --tags model \
- --additional-context-file examples/k8s-configs/02-single-node-multi-gpu-tools.json \
+ --additional-context-file examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu-tools.json \
--registry dockerhub
madengine run --manifest-file build_manifest.json --live-output
@@ -1096,7 +1089,7 @@ kubectl cp :/results/gpu_info_*.csv ./
```bash
# Copy closest match
-cp examples/k8s-configs/02-single-node-multi-gpu.json my-custom-config.json
+cp examples/k8s-configs/basic/02-torchrun-single-node-multi-gpu.json my-custom-config.json
# Edit
vim my-custom-config.json
@@ -1216,25 +1209,25 @@ kubectl logs | grep NCCL
## π Learning Path
### Level 1: Beginner
-1. Start with `01-single-node-single-gpu.json`
+1. Start with `basic/01-native-single-node-single-gpu.json`
2. Test on single GPU
3. Understand basic K8s concepts
4. Monitor logs and results
### Level 2: Intermediate
-1. Try `02-single-node-multi-gpu.json`
+1. Try `basic/02-torchrun-single-node-multi-gpu.json`
2. Learn distributed execution with torchrun (training workloads)
3. Understand NCCL configuration
4. Profile GPU utilization
### Level 3: Advanced
-1. Deploy `03-multi-node-basic.json`
+1. Deploy `basic/03-torchrun-multi-node-basic.json`
2. Master multi-node networking
3. Optimize NCCL parameters
4. Use PVCs for data and results
### Level 4: Expert
-1. Customize `04-multi-node-advanced.json`
+1. Customize `basic/04-torchrun-multi-node-advanced.json`
2. Fine-tune for your cluster
3. Implement node affinity and tolerations
4. Scale to 8+ nodes
@@ -1250,7 +1243,7 @@ Before deploying to production:
- [ ] Set appropriate memory and CPU limits
- [ ] Configured node selectors (if needed)
- [ ] Set NCCL environment variables
-- [ ] Enabled `host_ipc` for multi-node
+- [ ] Confirmed `distributed.nnodes > 1` for multi-node (host IPC is enabled automatically)
- [ ] Tested with small batch size first
- [ ] Configured PVCs for data (if using data providers)
- [ ] Set up monitoring and logging
@@ -1271,15 +1264,17 @@ Before deploying to production:
```
examples/k8s-configs/
-βββ README.md # This file
-βββ 01-single-node-single-gpu.json # 1 GPU, basic
-βββ 01-single-node-single-gpu-tools.json # 1 GPU + monitoring
-βββ 02-single-node-multi-gpu.json # 2 GPUs, distributed
-βββ 02-single-node-multi-gpu-tools.json # 2 GPUs + monitoring
-βββ 03-multi-node-basic.json # 2 nodes Γ 2 GPUs
-βββ 04-multi-node-advanced.json # 4 nodes Γ 2 GPUs
-βββ 05-nvidia-gpu-example.json # NVIDIA GPUs
-βββ 06-data-provider-with-pvc.json # Data provider + auto-PVC
+βββ README.md # This file
+βββ basic/
+β βββ 01-native-single-node-single-gpu.json # 1 GPU, basic
+β βββ 01-native-single-node-single-gpu-tools.json # 1 GPU + monitoring
+β βββ 02-torchrun-single-node-multi-gpu.json # 2 GPUs, distributed
+β βββ 02-torchrun-single-node-multi-gpu-tools.json # 2 GPUs + monitoring
+β βββ 03-torchrun-multi-node-basic.json # 2 nodes Γ 2 GPUs
+β βββ 04-torchrun-multi-node-advanced.json # 4 nodes Γ 2 GPUs
+β βββ 05-torchrun-nvidia-gpu-example.json # NVIDIA GPUs
+β βββ 06-data-provider-with-pvc.json # Data provider + auto-PVC
+βββ minimal/ # Minimal configs (see minimal/README.md)
```
---
diff --git a/examples/slurm-configs/README.md b/examples/slurm-configs/README.md
index e89f016f..0100fe54 100644
--- a/examples/slurm-configs/README.md
+++ b/examples/slurm-configs/README.md
@@ -42,7 +42,7 @@ The deployment type is **inferred** from the configuration structure:
| File | Description | Nodes | GPUs | Use Case |
|------|-------------|-------|------|----------|
-| `01-torchrun-single-node-single-gpu.json` | Single GPU training | 1 | 1 | Quick tests, small models |
+| `01-single-node-single-gpu.json` | Single GPU training | 1 | 1 | Quick tests, small models |
| `02-single-node-multi-gpu.json` | Single node, 8 GPUs | 1 | 8 | Single-node distributed workload |
| `03-multi-node-basic.json` | 2 nodes, 8 GPUs each | 2 | 16 | Multi-node distributed workload |
| `03-multi-node-basic-nodelist.json` | Same as 03 with `nodelist` | 2 | 16 | Pin job to specific nodes (e.g. node01,node02) |
@@ -58,15 +58,13 @@ The deployment type is **inferred** from the configuration structure:
### Minimal Examples (`minimal/`)
Stripped-down configurations showing only essential fields:
-- `single-gpu-minimal.json` - Minimal single GPU config
-- `multi-gpu-minimal.json` - Minimal 8 GPU config
-- `multi-node-minimal.json` - Minimal 2-node config
-- `primus-minimal.json` - Minimal Primus pretrain (`distributed.launcher: "primus"`; edit `primus.config_path`)
+- `torchrun-single-gpu-minimal.json` - Minimal single GPU config
+- `torchrun-multi-gpu-minimal.json` - Minimal 8 GPU config
+- `torchrun-multi-node-minimal.json` - Minimal 2-node config
- `vllm-single-node-minimal.json` - Minimal vLLM single-node
- `vllm-multi-node-minimal.json` - Minimal vLLM multi-node
- `slurm-multi-minimal.json` - Minimal slurm_multi self-managed launcher (3 nodes)
-For Primus options and environment variables, see [Launchers Guide](../../docs/launchers.md#5-primus).
For slurm_multi escape-hatch launcher, see [Launchers Guide](../../docs/launchers.md#9-slurm_multi-self-managed-escape-hatch).
## π Configuration Workflow
@@ -127,7 +125,7 @@ ssh user@hpc-cluster.example.com
# Phase 1: Build with configuration
MODEL_DIR=models/my-model madengine build \
--tags model_tag \
- --additional-context-file examples/slurm-configs/03-multi-node-basic.json \
+ --additional-context-file examples/slurm-configs/basic/03-multi-node-basic.json \
--manifest-output build_manifest.json
# Phase 2: Run from manifest
@@ -146,7 +144,7 @@ For quick tests without custom `env_vars`:
```bash
madengine run --tags model_tag \
- --additional-context-file examples/slurm-configs/minimal/single-gpu-minimal.json
+ --additional-context-file examples/slurm-configs/minimal/torchrun-single-gpu-minimal.json
```
### 3. CLI Override
@@ -168,7 +166,7 @@ madengine run --tags model_tag \
```bash
# Use base config, override specific fields
madengine run --tags model_tag \
- --additional-context-file examples/slurm-configs/03-multi-node-basic.json \
+ --additional-context-file examples/slurm-configs/basic/03-multi-node-basic.json \
--additional-context '{"slurm": {"nodes": 4, "time": "48:00:00"}}'
```
@@ -392,13 +390,23 @@ madengine uses intelligent multi-layer configuration merging:
"qos": "high", // Quality of Service
"account": "project-name", // SLURM account
"network_interface": "ib0", // Network interface (ib0/eth0)
- "modules": ["rocm/5.7.0"] // Environment modules to load
+ "modules": ["rocm/5.7.0"], // Environment modules to load
+ "enable_node_check": true, // Run node health preflight before submission (default: true)
+ "auto_cleanup_nodes": false, // Automatically clean up unhealthy nodes found during preflight (default: false)
+ "allow_submit_without_clean_nodes": false, // Submit even if fewer clean nodes than requested are found (default: false)
+ "verbose_node_check": false // Print detailed node health-check output (default: false)
}
}
```
**nodelist**: When set to a comma-separated list of node names (e.g. `"node01,node02"`), the job runs only on those nodes. Automatic node health preflight is skipped when `nodelist` is set.
+**Node Health Preflight**: Before submitting multi-node jobs, madengine runs a health check across candidate nodes and pins the job to the healthy ones (skipped if `nodelist` is set).
+- **`enable_node_check`** (default `true`): Enables/disables the preflight health check.
+- **`auto_cleanup_nodes`** (default `false`): Automatically attempts to clean up unhealthy nodes found during preflight.
+- **`allow_submit_without_clean_nodes`** (default `false`): If fewer clean nodes are found than requested, submit anyway instead of failing.
+- **`verbose_node_check`** (default `false`): Print detailed health-check output during preflight.
+
### Distributed Execution Section
```json
@@ -452,7 +460,7 @@ For slurm_multi, the model's `.slurm` script runs on baremetal and manages Docke
```bash
madengine run --tags my_model \
- --additional-context-file examples/slurm-configs/minimal/single-gpu-minimal.json
+ --additional-context-file examples/slurm-configs/minimal/torchrun-single-gpu-minimal.json
```
### Multi-Node Training
@@ -461,7 +469,7 @@ madengine run --tags my_model \
# Build with config
MODEL_DIR=models/my-model madengine build \
--tags training \
- --additional-context-file examples/slurm-configs/03-multi-node-basic.json
+ --additional-context-file examples/slurm-configs/basic/03-multi-node-basic.json
# Run from manifest
MODEL_DIR=models/my-model madengine run \
@@ -498,7 +506,7 @@ MODEL_DIR=models/llama2-70b madengine run \
```bash
madengine build --tags my_model \
- --additional-context-file examples/slurm-configs/04-multi-node-advanced.json
+ --additional-context-file examples/slurm-configs/basic/04-multi-node-advanced.json
madengine run --manifest-file build_manifest.json
```
@@ -759,7 +767,7 @@ module load python/3.9
# 3. Build with configuration
MODEL_DIR=models/my-model madengine build \
--tags llama2_training \
- --additional-context-file examples/slurm-configs/03-multi-node-basic.json \
+ --additional-context-file examples/slurm-configs/basic/03-multi-node-basic.json \
--manifest-output build_manifest.json
# 4. Run from manifest
diff --git a/src/madengine/cli/commands/run.py b/src/madengine/cli/commands/run.py
index c961cbe5..83266ff8 100644
--- a/src/madengine/cli/commands/run.py
+++ b/src/madengine/cli/commands/run.py
@@ -106,6 +106,18 @@ def run(
),
),
] = False,
+ 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,
manifest_output: Annotated[
str,
typer.Option(
@@ -234,6 +246,7 @@ def run(
verbose=verbose,
cleanup_perf=cleanup_perf,
skip_model_run=skip_model_run,
+ require_pinned_image=require_pinned_image,
_separate_phases=True,
)
@@ -338,6 +351,7 @@ def run(
verbose=verbose,
cleanup_perf=cleanup_perf,
skip_model_run=skip_model_run,
+ require_pinned_image=require_pinned_image,
_separate_phases=False, # Full workflow uses .live.log (not .run.live.log)
)
diff --git a/src/madengine/core/auth.py b/src/madengine/core/auth.py
index 15f0a0a6..021aa6dc 100644
--- a/src/madengine/core/auth.py
+++ b/src/madengine/core/auth.py
@@ -11,7 +11,8 @@
import json
import os
import shlex
-from typing import Dict, Optional
+from pathlib import Path
+from typing import Dict, Optional, Tuple
from madengine.core.errors import (
ConfigurationError,
@@ -19,6 +20,19 @@
handle_error,
)
+# Keys under which the Docker CLI stores Docker Hub credentials in config.json.
+# Docker has used several spellings over the years and any of them means
+# "this machine is authenticated to Docker Hub".
+_DOCKERHUB_CONFIG_KEYS: Tuple[str, ...] = (
+ "https://index.docker.io/v1/",
+ "index.docker.io",
+ "registry-1.docker.io",
+ "docker.io",
+)
+
+# Registry values that madengine treats as "Docker Hub" rather than a host.
+_DOCKERHUB_ALIASES = ("docker.io", "dockerhub")
+
def load_credentials() -> Optional[Dict]:
"""Load credentials from credential.json and environment variables.
@@ -81,6 +95,174 @@ def load_credentials() -> Optional[Dict]:
return credentials
+def _registry_config_keys(registry: Optional[str]) -> Tuple[str, ...]:
+ """Map a madengine registry value to the keys Docker uses in config.json.
+
+ Args:
+ registry: Registry URL (e.g. ``"localhost:5000"``, ``"docker.io/rocm"``),
+ or ``None``/empty string for Docker Hub.
+
+ Returns:
+ The config.json ``auths`` keys that would hold credentials for it.
+ """
+ if not registry or registry.lower() in _DOCKERHUB_ALIASES:
+ return _DOCKERHUB_CONFIG_KEYS
+ # Downstream code derives the registry host the same way (docker login ).
+ host = registry.split("/")[0]
+ if host.lower() in _DOCKERHUB_ALIASES:
+ return _DOCKERHUB_CONFIG_KEYS
+ return (host,)
+
+
+def has_ambient_docker_auth(registry: Optional[str]) -> bool:
+ """Report whether the local Docker CLI is already authenticated to ``registry``.
+
+ Reads ``${DOCKER_CONFIG:-~/.docker}/config.json`` the same way the Docker CLI
+ does, so an existing ``docker login`` (e.g. an organisation access token) is
+ honoured instead of being overridden or reported as "no credentials".
+
+ Args:
+ registry: Registry URL, or ``None``/empty string for Docker Hub.
+
+ Returns:
+ ``True`` if a usable credential entry exists for the registry. Any read
+ or parse problem yields ``False``; this function never raises and never
+ logs credential material.
+ """
+ config_dir = os.environ.get("DOCKER_CONFIG") or os.path.join(
+ os.path.expanduser("~"), ".docker"
+ )
+ try:
+ config = json.loads(Path(config_dir, "config.json").read_text(encoding="utf-8"))
+ except (OSError, ValueError):
+ return False
+ if not isinstance(config, dict):
+ return False
+
+ keys = _registry_config_keys(registry)
+ auths = config.get("auths") or {}
+ cred_helpers = config.get("credHelpers") or {}
+ creds_store = config.get("credsStore")
+
+ if isinstance(cred_helpers, dict) and any(key in cred_helpers for key in keys):
+ return True
+ if not isinstance(auths, dict):
+ return False
+ for key in keys:
+ entry = auths.get(key)
+ if not isinstance(entry, dict):
+ continue
+ if entry.get("auth") or entry.get("identitytoken") or entry.get("username"):
+ return True
+ # Credential-store-managed entries are persisted as an empty object;
+ # the secret itself lives in the external store.
+ if creds_store:
+ return True
+ return False
+
+
+def _usable_credentials(creds: object) -> bool:
+ """Report whether a credential entry carries a non-blank username and password.
+
+ Placeholder entries such as ``{"username": "", "password": ""}`` are treated
+ as "not configured" rather than as credentials, matching
+ :func:`madengine.deployment.k8s_secrets.build_registry_secret_data`.
+ """
+ if not isinstance(creds, dict):
+ return False
+ return bool(str(creds.get("username") or "").strip()) and bool(
+ str(creds.get("password") or "").strip()
+ )
+
+
+def _registry_from_image(image: Optional[str]) -> Optional[str]:
+ """Extract the registry host from an image reference.
+
+ Args:
+ image: Image reference (e.g. ``"ghcr.io/org/app:tag"``), or ``None``.
+
+ Returns:
+ The registry host, or ``None`` when the reference targets Docker Hub,
+ is malformed, or was not supplied.
+ """
+ if not image or not image.strip():
+ return None
+ ref = image.strip()
+ if "/" not in ref:
+ return None
+ # Docker treats the first component as a registry only when it looks like a
+ # host; otherwise it is a Docker Hub namespace (e.g. "rocm/private").
+ host = ref.split("/")[0]
+ if not ("." in host or ":" in host or host == "localhost"):
+ return None
+ if host.lower() in _DOCKERHUB_ALIASES:
+ return None
+ return host
+
+
+def explain_registry_denial(
+ log_text: str, image: Optional[str] = None
+) -> Optional[str]:
+ """Turn a registry denial in Docker output into an actionable explanation.
+
+ Distinguishes "authenticated but not authorized for this repository" from
+ "not authenticated at all", because the two need completely different fixes.
+
+ Args:
+ log_text: Docker build/pull output to inspect.
+ image: Optional image reference the denial refers to, for the message.
+
+ Returns:
+ A multi-line hint, or ``None`` if the output shows no registry denial.
+ """
+ lowered = (log_text or "").lower()
+ subject = image.strip() if image and image.strip() else "the base image"
+
+ if "insufficient_scope" in lowered or "authorization failed" in lowered:
+ return (
+ f"Base image pull was denied: {subject}\n"
+ " The registry ACCEPTED the credentials but granted no pull scope "
+ "for this repository.\n"
+ " This is an authorization problem, not a login problem:\n"
+ " - the access token is not scoped to this repository, or\n"
+ " - the repository/tag does not exist under that namespace.\n"
+ " Re-running `docker login` will not fix it; widen the token's "
+ "repository scope instead."
+ )
+
+ denied = (
+ "pull access denied" in lowered
+ or "requested access to the resource is denied" in lowered
+ or "authentication required" in lowered
+ or "unauthorized" in lowered
+ )
+ if not denied:
+ return None
+
+ registry = _registry_from_image(image)
+ if registry is None:
+ fixes = (
+ " - `docker login` on this machine (madengine reuses an existing "
+ "login)\n"
+ ' - add {"dockerhub": {"username": "...", "password": "..."}} to '
+ "credential.json\n"
+ " - export MAD_DOCKERHUB_USER and MAD_DOCKERHUB_PASSWORD"
+ )
+ else:
+ fixes = (
+ f" - `docker login {registry}` on this machine (madengine reuses "
+ "an existing login)\n"
+ f' - add {{"{registry}": {{"username": "...", "password": "..."}}}} '
+ "to credential.json"
+ )
+
+ return (
+ f"Base image pull was denied: {subject}\n"
+ " No usable credentials were presented to the registry.\n"
+ " Fix with any one of:\n" + fixes
+ )
+
+
def login_to_registry(
registry: Optional[str],
credentials: Optional[Dict],
@@ -103,10 +285,17 @@ def login_to_registry(
failure (missing key, invalid format, or docker login error).
Set to ``False`` to log and return instead, allowing the caller
to fall back to pulling public images.
+
+ Precedence: explicit credentials (``credential.json`` / ``MAD_DOCKERHUB_*``)
+ win when they carry a non-blank username and password. Otherwise an existing
+ ``docker login`` on this machine is reused and no login is attempted, so a
+ placeholder credential entry never overrides or breaks working ambient auth.
+ Set ``MAD_SKIP_DOCKER_LOGIN=1`` to always defer to ambient credentials.
"""
- if not credentials:
+ if os.environ.get("MAD_SKIP_DOCKER_LOGIN") == "1":
rich_console.print(
- "[yellow]No credentials provided for registry login[/yellow]"
+ "[yellow]MAD_SKIP_DOCKER_LOGIN=1 - using existing docker login for "
+ f"{registry or 'DockerHub'}[/yellow]"
)
return
@@ -116,8 +305,35 @@ def login_to_registry(
if registry and registry.lower() == "docker.io":
registry_key = "dockerhub"
- if registry_key not in credentials:
- error_msg = f"No credentials found for registry: {registry_key}"
+ entry = (credentials or {}).get(registry_key)
+ creds: Dict = entry if isinstance(entry, dict) else {}
+
+ if not _usable_credentials(creds):
+ # No explicit credentials configured for this registry. If the machine
+ # is already logged in (e.g. an organisation access token), reuse that
+ # instead of failing or clobbering it.
+ if has_ambient_docker_auth(registry):
+ rich_console.print(
+ f"[green]Using existing docker login for "
+ f"{registry or 'DockerHub'} (no explicit credentials "
+ f"configured)[/green]"
+ )
+ return
+
+ if not credentials:
+ rich_console.print(
+ "[yellow]No credentials provided for registry login[/yellow]"
+ )
+ return
+
+ if registry_key not in credentials:
+ error_msg = f"No credentials found for registry: {registry_key}"
+ else:
+ error_msg = (
+ f"Invalid credentials format for registry: {registry_key}"
+ f"\nCredentials must contain non-empty 'username' and "
+ f"'password' fields"
+ )
if registry_key == "dockerhub":
error_msg += (
f"\nPlease add dockerhub credentials to credential.json:\n"
@@ -140,18 +356,7 @@ def login_to_registry(
" }\n"
"}"
)
- rich_console.print(f"[red]{error_msg}[/red]")
- if raise_on_failure:
- raise RuntimeError(error_msg)
- return
-
- creds = credentials[registry_key]
-
- if "username" not in creds or "password" not in creds:
- error_msg = (
- f"Invalid credentials format for registry: {registry_key}"
- f"\nCredentials must contain 'username' and 'password' fields"
- )
+ error_msg += "\nAlternatively, run `docker login` on this machine."
rich_console.print(f"[red]{error_msg}[/red]")
if raise_on_failure:
raise RuntimeError(error_msg)
diff --git a/src/madengine/core/console.py b/src/madengine/core/console.py
index 71105c2f..36e1e161 100644
--- a/src/madengine/core/console.py
+++ b/src/madengine/core/console.py
@@ -229,6 +229,12 @@ def sh(
# Check for failure
success = proc.returncode == 0
+ # When output is captured rather than streamed it is discarded on
+ # failure, and the RuntimeError below carries only the command and the
+ # exit code. Echo it so the log records why the command actually failed.
+ if not success and not canFail and not secret and not self.live_output and outs:
+ print(redact_secrets(outs), flush=True)
+
# Show docker operation completion status
if not secret:
self._show_docker_completion(command, success)
diff --git a/src/madengine/core/image_digest.py b/src/madengine/core/image_digest.py
new file mode 100644
index 00000000..53c21c00
--- /dev/null
+++ b/src/madengine/core/image_digest.py
@@ -0,0 +1,128 @@
+#!/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
+
+from madengine.core.errors import ConfigurationError, create_error_context
+
+# `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}"
+
+
+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:
+ # An already-pinned reference (e.g. MAD_CONTAINER_IMAGE given as
+ # repo@sha256:...) satisfies the guarantee on its own; re-pinning it is
+ # a no-op and rejecting it would be wrong.
+ if parse_repo_digest(registry_image):
+ return registry_image
+
+ 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)
diff --git a/src/madengine/database/README.md b/src/madengine/database/README.md
index 2c8e5f9f..eec76ccb 100644
--- a/src/madengine/database/README.md
+++ b/src/madengine/database/README.md
@@ -1,81 +1,152 @@
-# Database Layer (Future MongoDB Ingestion)
+# Database Layer
-**Status**: Planned for future development
-**Purpose**: Modern data ingestion API for local and distributed deployments
+**Status**: Active
+**Purpose**: Upload CSV/JSON performance results to MongoDB
---
-## π― Objective
+## π― Responsibility
-This directory is reserved for a future unified database ingestion layer that will support:
-- MongoDB data persistence
-- Local result storage
-- Distributed data collection from build and run phases
-- Unified API for performance metrics ingestion
+This module implements MongoDB ingestion for madengine results (e.g. `perf.csv`,
+`perf_entry.csv`, or arbitrary JSON documents). It is a single, self-contained
+file β `mongodb.py` β with no sub-packages. It handles:
----
+- Auto-detecting file format (CSV vs JSON)
+- Loading files with native type preservation (numbers, bools, nested
+ JSON-in-CSV-cell strings)
+- Transforming/normalizing documents (metadata stamping, numpy/pandas type
+ cleanup)
+- Auto-detecting unique fields for deduplication when not specified
+- Bulk upload to MongoDB with batching, upsert, and automatic indexing
-## π Current State
+It is wired directly into the CLI via `madengine database` (see
+`src/madengine/cli/commands/database.py`).
-β οΈ **Not yet implemented**. This directory is a placeholder for future development.
+---
-For current database operations, use the existing `db/` package which handles MySQL operations via SSH.
+## π¦ Components (`mongodb.py`)
+
+### Configuration
+
+- **`MongoDBConfig`** β dataclass holding `host`, `port`, `username`,
+ `password`, `auth_source`, `timeout_ms`. `MongoDBConfig.from_env()` builds
+ one from `MONGO_HOST`, `MONGO_PORT`, `MONGO_USER`, `MONGO_PASSWORD`,
+ `MONGO_AUTH_SOURCE`, `MONGO_TIMEOUT_MS`. The `.uri` property builds the
+ `mongodb://` connection string.
+- **`UploadOptions`** β dataclass controlling upload behavior:
+ `unique_fields`, `upsert`, `batch_size`, `ordered`, `create_indexes`,
+ `index_fields`, `add_metadata`, `metadata_prefix`, `validate_schema`
+ (reserved field, currently unused by the implementation), `dry_run`.
+- **`UploadResult`** β dataclass returned by every upload: `status`
+ (`"success"` / `"partial"` / `"failed"`), `documents_read`,
+ `documents_processed`, `documents_inserted`, `documents_updated`,
+ `documents_failed`, `errors`, `duration_seconds`. Has a
+ `print_summary()` method that renders a formatted Rich summary.
+
+### File loading (Strategy pattern)
+
+- **`DocumentLoader`** (ABC) β defines `load(file_path)` and
+ `infer_schema(documents)`.
+- **`JSONLoader`** β loads a JSON object or array of objects, preserving
+ native types.
+- **`CSVLoader`** β loads via `pandas.read_csv`, preserving native types and
+ attempting to parse cell values that look like JSON (`{...}` / `[...]`)
+ back into dicts/lists.
+- **`detect_file_format(file_path)`** β picks `FileFormat.CSV` or
+ `FileFormat.JSON` from the extension, falling back to content sniffing.
+- **`get_loader(file_format)`** β returns the right loader instance.
+
+### Transformation
+
+- **`DocumentTransformer`** β takes `UploadOptions` and:
+ - `transform(documents)` adds metadata (`_meta_uploaded_at`,
+ `created_date`) and normalizes types (numpy scalars β Python, pandas
+ `Timestamp` β `datetime`, `NaN` β `None`).
+ - `infer_unique_fields(documents)` guesses a dedup key by checking
+ candidate fields (`model`, `name`, `id`, `timestamp`, `date`,
+ `pipeline`) for uniqueness across a sample of documents.
+
+### Upload
+
+- **`MongoDBUploader`** β connection + bulk-write class, usable as a context
+ manager (`with MongoDBUploader(config) as uploader:`).
+ - `connect()` / `disconnect()`
+ - `upload(documents, database_name, collection_name, options)` β
+ `UploadResult`. Creates indexes (if `options.create_indexes`) then does
+ either a plain `insert_many` (when no `unique_fields`/`upsert`) or a
+ batched `bulk_write` of `UpdateOne(..., upsert=True)` operations keyed
+ on `unique_fields`.
+
+### Entry points
+
+- **`upload_file_to_mongodb(file_path, database_name, collection_name, config=None, options=None) -> UploadResult`**
+ β the main entry point. Detects format, loads, auto-infers unique fields
+ if not given, transforms, honors `dry_run` (returns without connecting to
+ MongoDB), then uploads.
+- **`upload_csv_to_mongodb(csv_file_path, database_name, collection_name, mongo_config=None) -> Dict[str, Any]`**
+ β deprecated wrapper around `upload_file_to_mongodb` that returns a legacy
+ dict shape instead of `UploadResult`.
+- **`MongoDBHandler`** β deprecated class-based wrapper (`MongoDBHandler(args).run() -> bool`)
+ kept for backward compatibility with old argparse-style call sites.
---
-## ποΈ Legacy MySQL Tools (Removed)
+## π CLI mapping (`madengine database`)
-**MySQL support has been removed from madengine**. The following tools are no longer available:
+`src/madengine/cli/commands/database.py` is a thin Typer wrapper around
+`upload_file_to_mongodb`:
-| File | Purpose | Status |
-|------|---------|--------|
-| ~~`tools/create_table_db.py`~~ | MySQL table creation | **REMOVED** |
-| ~~`tools/update_table_db.py`~~ | MySQL table updates | **REMOVED** |
-| ~~`db/` package~~ | MySQL operations via SSH | **REMOVED** |
+| CLI flag | Maps to |
+|---|---|
+| `--file` / `-f` | `file_path` |
+| `--database` / `--db` | `database_name` |
+| `--collection` / `-c` | `collection_name` |
+| `--unique-key` / `-k` (comma-separated) | `UploadOptions.unique_fields` |
+| `--batch-size` | `UploadOptions.batch_size` |
+| `--no-upsert` | `UploadOptions.upsert = False` |
+| `--no-index` | `UploadOptions.create_indexes = False` |
+| `--dry-run` | `UploadOptions.dry_run` |
-For database operations, use MongoDB via the `database` command in the new CLI or legacy `mad.py`.
+Connection config always comes from `MongoDBConfig.from_env()` (the CLI does
+not expose host/port/credential flags β set `MONGO_HOST`, `MONGO_PORT`,
+`MONGO_USER`, `MONGO_PASSWORD` instead). See `madengine database --help` or
+`docs/cli-reference.md` for the full flag reference.
---
-## π Future Implementation Plan
-
-When implemented, this layer will provide:
-
-### **1. MongoDB Client** (`mongodb_client.py`)
-```python
-from madengine.database.mongodb_client import MongoDBClient
+## π Usage
-# Connect to local or remote MongoDB
-client = MongoDBClient(connection_string="mongodb://localhost:27017")
+**Via the CLI:**
-# Ingest build results
-client.ingest_build_results(build_manifest)
+```bash
+export MONGO_HOST=localhost
+export MONGO_USER=admin
+export MONGO_PASSWORD=secret
-# Ingest run results
-client.ingest_run_results(run_summary)
+madengine database -f perf.csv --db madengine --collection results -k model,timestamp
+madengine database -f perf_entry.json --db madengine --collection results --dry-run
```
-### **2. Local Storage** (`local_storage.py`)
-```python
-from madengine.database.local_storage import LocalStorage
-
-# Store results locally (JSON, Parquet, etc.)
-storage = LocalStorage(base_path="./madengine_results")
-storage.save_results(results_dict)
-```
+**Via the Python API:**
-### **3. Unified API** (`api.py`)
```python
-from madengine.database import ingest_results
-
-# Works with both local and distributed deployments
-ingest_results(
- results=run_summary,
- target="mongodb", # or "local", "mysql"
- config={"connection": "mongodb://..."}
+from madengine.database import upload_file_to_mongodb, MongoDBConfig, UploadOptions
+
+result = upload_file_to_mongodb(
+ file_path="perf.csv",
+ database_name="madengine",
+ collection_name="results",
+ config=MongoDBConfig.from_env(),
+ options=UploadOptions(unique_fields=["model", "timestamp"], batch_size=500),
)
+
+result.print_summary()
+print(result.status, result.documents_inserted, result.documents_updated)
```
+`config` and `options` are both optional β omit them to use
+`MongoDBConfig.from_env()` and default `UploadOptions()`.
+
---
## π¦ Difference from `db/` Package (Removed)
@@ -87,28 +158,23 @@ ingest_results(
| **Transport** | SSH tunnel | Direct connection |
| **Status** | **REMOVED** | Active |
----
-
-## π Migration Status
-
MySQL support has been fully removed from madengine:
-1. β
**Phase 1**: Removed `db/` package (MySQL operations)
-2. β
**Phase 2**: Removed `tools/create_table_db.py` and `tools/update_table_db.py`
-3. β
**Phase 3**: Removed `utils/ssh_to_db.py` (SSH to MySQL host)
-4. β
**Phase 4**: Removed MySQL dependencies (`mysql-connector-python`, `pymysql`)
+1. β
Removed `db/` package (MySQL operations)
+2. β
Removed `tools/create_table_db.py` and `tools/update_table_db.py`
+3. β
Removed `utils/ssh_to_db.py` (SSH to MySQL host)
+4. β
Removed MySQL dependencies (`mysql-connector-python`, `pymysql`)
-**Current state**: Only MongoDB support remains via the `database/` package.
+**Current state**: Only MongoDB support remains, via this `database/` package.
---
## π References
-- **MongoDB package**: `src/madengine/database/mongodb.py`
-- **CLI database command**: `madengine database --help`
+- **Implementation**: `src/madengine/database/mongodb.py`
+- **CLI command**: `src/madengine/cli/commands/database.py`, `madengine database --help`
---
-**Last Updated**: November 30, 2025
+**Last Updated**: 2026-08-05
**Maintainer**: madengine Team
-
diff --git a/src/madengine/deployment/k8s_template_context.py b/src/madengine/deployment/k8s_template_context.py
index e38b251a..498e2659 100644
--- a/src/madengine/deployment/k8s_template_context.py
+++ b/src/madengine/deployment/k8s_template_context.py
@@ -31,6 +31,7 @@
from madengine.core.additional_context_defaults import DEFAULT_GUEST_OS
from madengine.core.dataprovider import Data
from madengine.core.errors import ConfigurationError
+from madengine.core.image_digest import resolve_pinned_image
from madengine.utils.gpu_config import resolve_runtime_gpus
from madengine.utils.path_utils import get_madengine_root
@@ -483,6 +484,15 @@ def _prepare_template_context(
else None
)
+ # 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,
+ )
+
# Build complete context
context = {
# Job metadata
@@ -515,7 +525,7 @@ def _prepare_template_context(
"data_provider_script": data_provider_script,
"data_provider_script_content": data_provider_script_content,
# Image
- "image": image_info["registry_image"],
+ "image": resolved_image,
"image_pull_policy": self.k8s_config.get("image_pull_policy", "Always"),
# Resources
"gpu_resource_name": self.gpu_resource_name,
diff --git a/src/madengine/deployment/slurm.py b/src/madengine/deployment/slurm.py
index af99c08d..c9cdd21b 100644
--- a/src/madengine/deployment/slurm.py
+++ b/src/madengine/deployment/slurm.py
@@ -13,6 +13,7 @@
import os
import shlex
+import shutil
import subprocess
import time
from pathlib import Path
@@ -28,6 +29,8 @@
)
from .config_loader import ConfigLoader, apply_deployment_config
from .slurm_node_selector import SlurmNodeSelector
+from madengine.core.errors import ConfigurationError
+from madengine.core.image_digest import resolve_pinned_image
from madengine.utils.gpu_config import resolve_runtime_gpus
from madengine.utils.run_details import get_build_number, get_pipeline
from madengine.utils.path_utils import scripts_base_dir_from
@@ -77,6 +80,8 @@ def __init__(self, config: DeploymentConfig):
self.time_limit = self.slurm_config.get("time", "24:00:00")
self.output_dir = Path(self.slurm_config.get("output_dir", "./slurm_results"))
self.reservation = self.slurm_config.get("reservation", None)
+ # Some clusters expose no GPU GRES, so sbatch rejects --gpus-per-node.
+ self.skip_gpus_directive = self.slurm_config.get("skip_gpus_directive", False)
# Setup Jinja2 template engine
template_dir = Path(__file__).parent / "templates" / "slurm"
@@ -226,6 +231,22 @@ def validate(self) -> bool:
self.console.print("[green]β SLURM environment validated[/green]")
return True
+ @staticmethod
+ def _submission_bin_dir() -> Optional[str]:
+ """
+ Directory the madengine console script was resolved from at submission time.
+
+ A batch job is not guaranteed to inherit the submitter's PATH: a site can
+ default sbatch to --export=NONE, and `module load` can rewrite PATH before
+ the job body runs. Passing the directory into the job script lets it put
+ the same madengine back on PATH instead of relying on inheritance.
+
+ Returns:
+ Optional[str]: absolute directory, or None if madengine is not on PATH
+ """
+ cli_path = shutil.which("madengine")
+ return str(Path(cli_path).resolve().parent) if cli_path else None
+
def _validate_cli_availability(self) -> bool:
"""
Validate madengine is available before job submission.
@@ -241,7 +262,9 @@ def _validate_cli_availability(self) -> bool:
["madengine", "--version"],
capture_output=True,
text=True,
- timeout=5,
+ # A cold import off shared/NFS storage can take far longer than a
+ # local one, so this only guards against a hung interpreter.
+ timeout=600,
check=False
)
if result.returncode == 0:
@@ -315,6 +338,11 @@ def prepare(self) -> bool:
return self._prepare_slurm_multi_script(
model_info_peek, docker_image_name=model_keys_peek[0]
)
+ 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
@@ -432,7 +460,29 @@ def _prepare_slurm_multi_script(self, model_info: Dict, docker_image_name: str =
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]"
+ )
+
# Get model args. The wrapper script below is executed by bash, so the
# script name and free-form args string must be shell-quoted to prevent
# embedded metacharacters ($(), backticks, ;, etc.) from being evaluated
@@ -456,7 +506,10 @@ def _prepare_slurm_multi_script(self, model_info: Dict, docker_image_name: str =
f"#SBATCH --partition={self.partition}",
f"#SBATCH --nodes={self.nodes}",
f"#SBATCH --ntasks={self.nodes}",
- f"#SBATCH --gpus-per-node={self.gpus_per_node}",
+ ]
+ if not self.skip_gpus_directive:
+ script_lines.append(f"#SBATCH --gpus-per-node={self.gpus_per_node}")
+ script_lines += [
f"#SBATCH --time={self.time_limit}",
]
# Honour user-configured exclusivity (defaults to True to match the standard SLURM template).
@@ -669,6 +722,7 @@ def debug(self, msg):
"partition": self.partition,
"nodes": self.nodes,
"gpus_per_node": resolved_gpus_per_node, # Use resolved GPU count
+ "skip_gpus_directive": self.skip_gpus_directive,
"time_limit": self.time_limit,
"output_dir": str(self.output_dir),
"master_port": master_port,
@@ -682,6 +736,7 @@ def debug(self, msg):
"qos": self.slurm_config.get("qos"),
"account": self.slurm_config.get("account"),
"modules": self.slurm_config.get("modules", []),
+ "submission_bin_dir": self._submission_bin_dir(),
"env_vars": self.config.additional_context.get("env_vars", {}),
"shared_workspace": self.slurm_config.get("shared_workspace"),
"shared_data": self.config.additional_context.get("shared_data"),
diff --git a/src/madengine/deployment/templates/slurm/job.sh.j2 b/src/madengine/deployment/templates/slurm/job.sh.j2
index 3b236d9f..4cdddba0 100644
--- a/src/madengine/deployment/templates/slurm/job.sh.j2
+++ b/src/madengine/deployment/templates/slurm/job.sh.j2
@@ -6,7 +6,8 @@
#SBATCH --nodes={{ nodes }}
#SBATCH --ntasks={{ nodes }}
#SBATCH --ntasks-per-node=1
-#SBATCH --gpus-per-node={{ gpus_per_node }}
+{% if not skip_gpus_directive %}#SBATCH --gpus-per-node={{ gpus_per_node }}
+{% endif %}
#SBATCH --time={{ time_limit }}
{% if reservation %}
#SBATCH --reservation={{ reservation }}
@@ -41,6 +42,19 @@
module load {{ module }}
{% endfor %}
+# =============================================================================
+# PATH
+# =============================================================================
+# A batch job does not reliably inherit the submitter's PATH: a site can default
+# sbatch to --export=NONE, and the module loads above can rewrite it. Without
+# this the madengine console script is missing on the compute node even though
+# the pre-submission check found it. Re-add the per-user bin directory and the
+# directory madengine itself was resolved from at submission time.
+export PATH="$HOME/.local/bin:$PATH"
+{% if submission_bin_dir %}
+export PATH="{{ submission_bin_dir }}:$PATH"
+{% endif %}
+
# =============================================================================
# Environment Setup (Standard ML Environment Variables)
# =============================================================================
@@ -141,7 +155,15 @@ echo "Submission directory: {{ manifest_file | dirname }}"
# Single-node: Prefer shared storage (submission dir), with local fallback if needed
# Check if submission directory is on shared filesystem
SUBMIT_DIR={{ manifest_file | dirname }}
-if df -T "$SUBMIT_DIR" 2>/dev/null | grep -qE '\bnfs\b|\blustre\b|\bgpfs\b|\bceph\b'; then
+# Ask df for the filesystem type alone. Matching the whole df -T line reads the mount
+# point too, so a local disk under a path such as /mnt/nfs-scratch answered yes and the
+# job then trusted node-local storage to be shared. --output=fstype is GNU coreutils
+# 8.21 and up; the awk form is the fallback for anything older.
+SUBMIT_FSTYPE=$(df --output=fstype "$SUBMIT_DIR" 2>/dev/null | tail -n 1)
+if [ -z "$SUBMIT_FSTYPE" ]; then
+ SUBMIT_FSTYPE=$(df -T "$SUBMIT_DIR" 2>/dev/null | awk 'NR > 1 { print $2; exit }')
+fi
+if printf '%s' "$SUBMIT_FSTYPE" | grep -qE '^(nfs[0-9]*|lustre|gpfs|ceph|beegfs|panfs)$'; then
# Submission directory is on shared storage - use it directly (best practice)
WORKSPACE=$SUBMIT_DIR
WORKSPACE_TYPE="shared-nfs"
diff --git a/src/madengine/execution/README.md b/src/madengine/execution/README.md
index 1277cfa3..6a51f274 100644
--- a/src/madengine/execution/README.md
+++ b/src/madengine/execution/README.md
@@ -44,7 +44,7 @@ result = builder.build_image(
# Build all models
results = builder.build_all_models(
- models_list=[model1, model2, model3],
+ models=[model1, model2, model3],
target_archs=["gfx90a", "gfx942"]
)
@@ -52,6 +52,10 @@ results = builder.build_all_models(
builder.export_build_manifest(output_file="build_manifest.json")
```
+### **`dockerfile_utils.py`**
+
+Helper functions for multi-architecture Dockerfile parsing (e.g., `parse_dockerfile_gpu_variables`, `normalize_architecture_name`, `is_target_arch_compatible_with_variable`, `is_compilation_arch_compatible`) used by `docker_builder.py`.
+
### **`container_runner.py`**
Runs Docker containers locally for model execution.
@@ -73,16 +77,19 @@ runner = ContainerRunner(context, data, console)
# Run model in container
result = runner.run_container(
model_info=model_dict,
- model_docker=docker_client,
- gpu_ids="0,1",
+ docker_image="model1:latest",
timeout=3600
)
# Result includes status, metrics, logs
-print(result["status"]) # "successful", "failed", "timeout"
-print(result["duration"])
+print(result["status"]) # "SUCCESS", "FAILURE", "SKIPPED"
+print(result["test_duration"])
```
+### **`container_runner_helpers.py`**
+
+Backs `container_runner.py`'s timeout management and error-detection features (e.g., `resolve_log_error_scan_config`, `log_text_has_error_pattern`, `resolve_run_status`, `resolve_run_timeout`, `make_run_log_file_path`).
+
---
## ποΈ Architecture
diff --git a/src/madengine/execution/container_runner.py b/src/madengine/execution/container_runner.py
index eab4af7d..4fe29816 100644
--- a/src/madengine/execution/container_runner.py
+++ b/src/madengine/execution/container_runner.py
@@ -23,6 +23,7 @@
from madengine.core.console import Console, redact_secrets
from madengine.core.context import Context
from madengine.core.docker import Docker
+from madengine.core.image_digest import resolve_pinned_image
from madengine.core.timeout import Timeout
from madengine.core.dataprovider import Data
from madengine.utils.ops import PythonicTee, file_print
@@ -2836,7 +2837,21 @@ def run_models_from_manifest(
# Local image mode (MAD_CONTAINER_IMAGE): Use the provided image directly
run_image = build_info.get("docker_image")
self.rich_console.print(f"[yellow]π Using local image: {run_image}[/yellow]")
-
+
+ # This branch also covers build-on-compute-node manifests,
+ # whose docker_image is a registry reference. Enforce here
+ # too, otherwise those manifests would silently bypass the
+ # flag by never reaching the registry branch below.
+ run_image = resolve_pinned_image(
+ run_image,
+ build_info.get("image_digest"),
+ bool(
+ (self.additional_context or {}).get("require_pinned_image")
+ ),
+ model_name=model_info.get("name", ""),
+ )
+
+
# Ensure the local image is available on this node. In a
# multi-node SLURM run only the primary may have the
# locally-built image; the shared-tar cache
@@ -2849,11 +2864,20 @@ def run_models_from_manifest(
)
elif build_info.get("registry_image"):
- # Registry image: Pull from registry
+ # 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(build_info["registry_image"])
+ self.pull_image(pull_target)
# Update docker_image to use registry image
- run_image = build_info["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
diff --git a/src/madengine/execution/docker_builder.py b/src/madengine/execution/docker_builder.py
index bc44aa44..0e0d608c 100644
--- a/src/madengine/execution/docker_builder.py
+++ b/src/madengine/execution/docker_builder.py
@@ -9,6 +9,7 @@
import os
import shlex
+import sys
from pathlib import Path
import time
import json
@@ -16,9 +17,14 @@
import typing
from contextlib import redirect_stdout, redirect_stderr
from rich.console import Console as RichConsole
-from madengine.core.auth import login_to_registry
+from madengine.core.auth import (
+ explain_registry_denial,
+ has_ambient_docker_auth,
+ login_to_registry,
+)
from madengine.core.console import Console
from madengine.core.context import Context
+from madengine.core.image_digest import parse_push_digest, parse_repo_digest
from madengine.utils.ops import PythonicTee
from madengine.execution.dockerfile_utils import (
is_target_arch_compatible_with_variable,
@@ -44,6 +50,7 @@ def __init__(
self.live_output = live_output
self.rich_console = RichConsole()
self.built_images = {} # Track built images
+ self.pushed_digests = {} # registry_image -> digest recorded at push time
self.built_models = {} # Track built models
def get_context_path(self, info: typing.Dict) -> str:
@@ -98,6 +105,72 @@ def get_build_arg(self, run_build_arg: typing.Optional[typing.Dict] = None) -> s
return build_args
+ def _resolve_base_docker(self, dockerfile: str) -> str:
+ """Resolve the base image the Dockerfile builds ``FROM``.
+
+ Prefers a ``BASE_DOCKER`` override from ``docker_build_arg`` context,
+ otherwise reads ``ARG BASE_DOCKER=`` from the Dockerfile.
+
+ Args:
+ dockerfile: Path to the Dockerfile.
+
+ Returns:
+ str: The base image reference, or ``""`` if it cannot be determined.
+ """
+ if (
+ "docker_build_arg" in self.context.ctx
+ and "BASE_DOCKER" in self.context.ctx["docker_build_arg"]
+ ):
+ return str(self.context.ctx["docker_build_arg"]["BASE_DOCKER"])
+ try:
+ return str(
+ self.console.sh(
+ f"grep '^ARG BASE_DOCKER=' {shlex.quote(dockerfile)} | sed -E 's/ARG BASE_DOCKER=//g'"
+ )
+ )
+ except Exception:
+ return ""
+
+ @staticmethod
+ def _registry_of(image: str) -> str:
+ """Return the registry an image reference points at.
+
+ Args:
+ image: An image reference such as ``rocm/pytorch:latest`` or
+ ``myhost:5000/team/img:tag``.
+
+ Returns:
+ str: The registry host, or ``"docker.io"`` for Docker Hub references.
+ """
+ first_segment = (image or "").strip().split("/")[0]
+ if "." in first_segment or ":" in first_segment or first_segment == "localhost":
+ return first_segment
+ return "docker.io"
+
+ def _report_registry_denial(self, log_file_path: str, base_docker: str) -> None:
+ """Print an actionable hint if a build log shows a registry denial.
+
+ Called from inside the build's stdout redirection, so the hint is written
+ to the real stdout as well to make sure it reaches the terminal and not
+ only the build log.
+
+ Args:
+ log_file_path: Path to the build log written by this build.
+ base_docker: The base image the build was pulling.
+ """
+ try:
+ with open(log_file_path, encoding="utf-8", errors="replace") as log:
+ log_text = log.read()
+ except OSError:
+ return
+ hint = explain_registry_denial(log_text, base_docker)
+ if not hint:
+ return
+ message = f"[bold red]β {hint}[/bold red]"
+ self.rich_console.print(f"\n{message}")
+ if not self.live_output and sys.__stdout__ is not None:
+ RichConsole(file=sys.__stdout__).print(f"\n{message}")
+
def build_image(
self,
model_info: typing.Dict,
@@ -195,8 +268,28 @@ def build_image(
with redirect_stdout(
PythonicTee(outlog, self.live_output)
), redirect_stderr(PythonicTee(outlog, self.live_output)):
+ # `docker build --pull` resolves the base image itself, so it only
+ # works if this machine can authenticate to the base image's
+ # registry. Log in when β and only when β there is no existing
+ # login to reuse, so an ambient `docker login` (e.g. an
+ # organisation access token) is left alone.
+ base_docker = self._resolve_base_docker(dockerfile)
+ base_registry = self._registry_of(base_docker)
+ if credentials and not has_ambient_docker_auth(base_registry):
+ login_to_registry(
+ base_registry,
+ credentials,
+ console=self.console,
+ rich_console=self.rich_console,
+ raise_on_failure=False,
+ )
+
print(f"π¨ Executing build command...")
- self.console.sh(build_command, timeout=None)
+ try:
+ self.console.sh(build_command, timeout=None)
+ except Exception:
+ self._report_registry_denial(log_file_path, base_docker)
+ raise
build_duration = time.time() - build_start_time
@@ -206,17 +299,6 @@ def build_image(
self.rich_console.print(f"[dim]{'='*80}[/dim]")
# Get base docker info
- base_docker = ""
- if (
- "docker_build_arg" in self.context.ctx
- and "BASE_DOCKER" in self.context.ctx["docker_build_arg"]
- ):
- base_docker = self.context.ctx["docker_build_arg"]["BASE_DOCKER"]
- else:
- base_docker = self.console.sh(
- f"grep '^ARG BASE_DOCKER=' {shlex.quote(dockerfile)} | sed -E 's/ARG BASE_DOCKER=//g'"
- )
-
print(f"BASE DOCKER is {base_docker}")
# Get docker SHA
@@ -316,7 +398,9 @@ def push_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}")
- self.console.sh(push_command)
+ push_output = self.console.sh(push_command)
+
+ self._record_pushed_digest(registry_image, push_output)
self.rich_console.print(f"[bold green]β
Successfully pushed image to registry:[/bold green] [cyan]{registry_image}[/cyan]")
self.rich_console.print(f"[dim]{'='*80}[/dim]")
@@ -326,6 +410,44 @@ def push_image(
self.rich_console.print(f"[red]β Failed to push image {docker_image} to registry {registry}: {e}[/red]")
raise
+ 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]"
+ )
+
def export_build_manifest(
self,
output_file: str = "build_manifest.json",
@@ -688,6 +810,10 @@ def _build_model_single_arch(
)
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
except Exception as e:
build_info["push_error"] = str(e)
@@ -896,6 +1022,10 @@ def _build_model_for_arch(
try:
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
except Exception as e:
build_info["push_error"] = str(e)
diff --git a/src/madengine/orchestration/build_orchestrator.py b/src/madengine/orchestration/build_orchestrator.py
index 246701cf..17e836e8 100644
--- a/src/madengine/orchestration/build_orchestrator.py
+++ b/src/madengine/orchestration/build_orchestrator.py
@@ -20,7 +20,7 @@
from madengine.core.console import Console
from madengine.core.context import Context
from madengine.core.additional_context_defaults import apply_build_context_defaults
-from madengine.core.auth import load_credentials
+from madengine.core.auth import has_ambient_docker_auth, load_credentials
from madengine.core.errors import (
BuildError,
ConfigurationError,
@@ -975,16 +975,33 @@ def _execute_build_on_compute(
)
if any(pub_reg in registry_lower for pub_reg in public_registries):
if not dockerhub_user or not dockerhub_password:
- raise ConfigurationError(
- f"Registry credentials required for pushing to {registry}",
- context=create_error_context(
- operation="build_on_compute",
- component="BuildOrchestrator",
- additional_info={"registry": registry},
- ),
- suggestions=_matched_hints,
- )
- self.rich_console.print(f" Auth: Will login to registry before push")
+ # An existing `docker login` (e.g. an organisation access token)
+ # is a valid substitute for explicit credentials, so do not hard
+ # fail on it. The generated sbatch script already falls back to
+ # "assume pre-authenticated" when no credentials are supplied.
+ if has_ambient_docker_auth(registry.split("/")[0]):
+ self.rich_console.print(
+ " [yellow]Auth: No explicit credentials; relying on the "
+ "existing docker login[/yellow]"
+ )
+ self.rich_console.print(
+ " [dim]Note: this login was detected on the submit node. "
+ "The compute node only shares it if $HOME (or $DOCKER_CONFIG) "
+ "is on shared storage.[/dim]"
+ )
+ else:
+ raise ConfigurationError(
+ f"Registry credentials required for pushing to {registry}",
+ context=create_error_context(
+ operation="build_on_compute",
+ component="BuildOrchestrator",
+ additional_info={"registry": registry},
+ ),
+ suggestions=_matched_hints
+ + ["Or run `docker login` on the build node"],
+ )
+ else:
+ self.rich_console.print(f" Auth: Will login to registry before push")
else:
# Private/internal registry - may not need auth
self.rich_console.print(f" Auth: Private registry (auth may not be required)")
diff --git a/src/madengine/orchestration/run_orchestrator.py b/src/madengine/orchestration/run_orchestrator.py
index 296af1ee..aa8fa652 100644
--- a/src/madengine/orchestration/run_orchestrator.py
+++ b/src/madengine/orchestration/run_orchestrator.py
@@ -83,6 +83,13 @@ def __init__(self, args, additional_context: Optional[Dict] = None):
merged_context.update(additional_context)
self.additional_context = merged_context
+
+ # 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
+
keys_str = ", ".join(sorted(self.additional_context.keys())) if self.additional_context else "(none)"
self.rich_console.print(f"[dim]Run additional context (CLI):[/dim] [cyan]{keys_str}[/cyan]")
@@ -459,6 +466,7 @@ def _create_manifest_from_local_image(
"deprecated": model.get("deprecated", False),
"skip_gpu_arch": model.get("skip_gpu_arch", []),
"additional_docker_run_options": model.get("additional_docker_run_options", ""),
+ "multiple_results": model.get("multiple_results", ""),
}
# Write manifest to file
@@ -495,7 +503,15 @@ def _load_and_merge_manifest(self, manifest_file: str) -> str:
if "context" not in manifest:
manifest["context"] = {}
- merge_keys = ["tools", "pre_scripts", "post_scripts", "encapsulate_script"]
+ 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",
+ ]
context_updated = False
for key in merge_keys:
if key in self.additional_context:
@@ -764,14 +780,17 @@ def _show_node_info(self):
self.console.sh("echo 'MAD Run Models'")
host_os = self.context.ctx.get("host_os", "")
+ # This is purely informational, but a package manager can block forever on
+ # an interactive prompt (e.g. yum asking to import a repo GPG key) with no
+ # tty to answer it, so every query is capped.
if "HOST_UBUNTU" in host_os:
- print(self.console.sh("apt show rocm-libs -a", canFail=True))
+ print(self.console.sh("timeout 10 apt show rocm-libs -a", canFail=True))
elif "HOST_CENTOS" in host_os:
- print(self.console.sh("yum info rocm-libs", canFail=True))
+ print(self.console.sh("timeout 10 yum info rocm-libs", canFail=True))
elif "HOST_SLES" in host_os:
- print(self.console.sh("zypper info rocm-libs", canFail=True))
+ print(self.console.sh("timeout 10 zypper info rocm-libs", canFail=True))
elif "HOST_AZURE" in host_os:
- print(self.console.sh("tdnf info rocm-libs", canFail=True))
+ print(self.console.sh("timeout 10 tdnf info rocm-libs", canFail=True))
else:
self.rich_console.print("[yellow]Warning: Unable to detect host OS[/yellow]")
diff --git a/src/madengine/reporting/README.md b/src/madengine/reporting/README.md
index 33d8e5a4..ac3ec49e 100644
--- a/src/madengine/reporting/README.md
+++ b/src/madengine/reporting/README.md
@@ -26,43 +26,41 @@ from madengine.reporting.update_perf_csv import update_perf_csv, flatten_tags
# Update CSV with new results
update_perf_csv(
- perf_json_path="results.json",
- output_csv="performance.csv"
+ perf_csv="performance.csv",
+ single_result="results.json",
)
-# Flatten nested tags for CSV export
-flattened = flatten_tags(perf_entry)
+# Flatten nested tags in place (no return value)
+flatten_tags(perf_entry) # mutates perf_entry in place
```
----
-
-## ποΈ Legacy Reporting Tools
+### **`update_perf_super.py`**
-The following legacy-only reporting tools remain in `tools/`:
+Maintains `perf_super.json`, a cumulative superset performance record that also
+captures matched config data, and provides conversion to `perf_super.csv` /
+`perf_entry_super.json` / `perf_entry_super.csv`.
-| File | Purpose | Used By | Status |
-|------|---------|---------|--------|
-| `tools/csv_to_html.py` | Convert CSV to HTML | `mad.py`, `run_models.py` | Legacy only |
-| `tools/csv_to_email.py` | Email CSV reports | `mad.py` | Legacy only |
+**Used by:**
+- β
`execution/container_runner.py` (modern madengine CLI)
-These tools are **NOT** used by the modern `madengine` CLI.
+### **`csv_to_html.py`**
----
+Converts a single CSV file to an HTML table. Provides the `ConvertCsvToHtml`
+handler class (`ConvertCsvToHtml.__init__(self, args: argparse.Namespace)`,
+`.run(self) -> bool`) used by the CLI.
-## π Architecture Decision
-
-**Why is `update_perf_csv.py` in `reporting/` instead of `tools/`?**
+**Used by:**
+- β
`cli/commands/report.py` (backs `madengine report to-html`)
-1. β
**Shared across architectures**: Used by both legacy and new CLI
-2. β
**Active development**: Not deprecated, actively maintained
-3. β
**Clear responsibility**: Performance data processing
-4. β
**Semantic clarity**: Reporting is a distinct concern
+### **`csv_to_email.py`**
-**Why are other CSV tools still in `tools/`?**
+Converts all CSV files in a directory into a single consolidated HTML report
+suitable for emailing. Provides the `ConvertCsvToEmail` handler class
+(`ConvertCsvToEmail.__init__(self, args: argparse.Namespace)`, `.run(self) -> bool`)
+used by the CLI.
-- They are **not used** by the modern `madengine` CLI
-- Kept for backward compatibility only
-- Will be deprecated when legacy CLI is retired
+**Used by:**
+- β
`cli/commands/report.py` (backs `madengine report to-email`)
---
@@ -74,10 +72,10 @@ These tools are **NOT** used by the modern `madengine` CLI.
from madengine.reporting.update_perf_csv import update_perf_csv
# After model execution completes
+perf_csv = "/path/to/performance.csv"
results_json = "/path/to/results.json"
-output_csv = "/path/to/performance.csv"
-update_perf_csv(results_json, output_csv)
+update_perf_csv(perf_csv, single_result=results_json)
```
### **Legacy madengine** (via `run_models.py` or `mad.py`)
@@ -107,6 +105,10 @@ Performance CSV
(Optional) CSV β Email (legacy only)
```
+**Note:** As a side effect, `update_perf_csv()` (and the `handle_*_result()` helpers it
+calls) also always write/append `perf_entry.csv` and `perf_entry.json` with the
+latest result, regardless of the output file passed in.
+
---
## π§ͺ Testing
diff --git a/src/madengine/scripts/common/post_scripts/trace.sh b/src/madengine/scripts/common/post_scripts/trace.sh
index 1e489861..1dbdaaf5 100644
--- a/src/madengine/scripts/common/post_scripts/trace.sh
+++ b/src/madengine/scripts/common/post_scripts/trace.sh
@@ -33,8 +33,12 @@ rpd)
if [ -f "./rocmProfileData/tools/rpd2tracing.py" ]; then
echo "RPD post-script: rpd2tracing.py found"
if [ -f "trace.rpd" ] && [ -s "trace.rpd" ]; then
- python3 ./rocmProfileData/tools/rpd2tracing.py trace.rpd trace.json
- mv trace.rpd trace.json "$OUTPUT"
+ if python3 ./rocmProfileData/tools/rpd2tracing.py trace.rpd trace.json; then
+ mv trace.rpd trace.json "$OUTPUT"
+ else
+ echo "RPD post-script: rpd2tracing.py failed (likely no captured trace data); saving raw trace.rpd only"
+ mv trace.rpd "$OUTPUT"
+ fi
else
echo "RPD post-script: Skipping rpd2tracing.py because trace.rpd is missing or empty"
# Create empty files so the directory structure exists
diff --git a/src/madengine/scripts/common/pre_scripts/gpu_info_pre.sh b/src/madengine/scripts/common/pre_scripts/gpu_info_pre.sh
index 60bd60a0..7903e378 100644
--- a/src/madengine/scripts/common/pre_scripts/gpu_info_pre.sh
+++ b/src/madengine/scripts/common/pre_scripts/gpu_info_pre.sh
@@ -5,15 +5,19 @@
#
gpu_vendor=""
-if [ -f "/usr/bin/nvidia-smi" ]; then
+if command -v nvidia-smi >/dev/null 2>&1; then
echo "NVIDIA GPU detected."
gpu_vendor="NVIDIA"
gpu_architecture=$(nvidia-smi --query-gpu=name --format=csv,noheader | grep -m 1 -E -o ".{0,1}100"| xargs )
python3 -m pip install nvidia-ml-py
-elif [ -f "/opt/rocm/bin/rocm-smi" ]; then
+elif command -v rocm-smi >/dev/null 2>&1 || command -v amd-smi >/dev/null 2>&1; then
echo "AMD GPU detected."
gpu_vendor="AMD"
- gpu_architecture=$(rocminfo | grep -o -m 1 'gfx.*' | xargs )
+ if command -v rocminfo >/dev/null 2>&1; then
+ gpu_architecture=$(rocminfo | grep -o -m 1 'gfx.*' | xargs )
+ else
+ echo "rocminfo not found; skipping AMD GPU architecture detection."
+ fi
MI200="gfx90a"
MI100="gfx908"
MI50="gfx906"
diff --git a/src/madengine/scripts/common/tools.json b/src/madengine/scripts/common/tools.json
index 82869087..b237c63c 100644
--- a/src/madengine/scripts/common/tools.json
+++ b/src/madengine/scripts/common/tools.json
@@ -7,10 +7,8 @@
"args": "rpd"
}
],
- "cmd": "./rocmProfileData/rpd_tracer/runTracer.sh",
- "env_vars": {
- "LD_LIBRARY_PATH": "./rocmProfileData/rpd_tracer:/opt/rocm/lib"
- },
+ "cmd": "LD_LIBRARY_PATH=\"./rocmProfileData/rpd_tracer:${ROCM_PATH:-/opt/rocm}/lib:${LD_LIBRARY_PATH}\" ./rocmProfileData/rpd_tracer/runTracer.sh",
+ "env_vars": {},
"post_scripts": [
{
"path": "scripts/common/post_scripts/trace.sh",
diff --git a/src/madengine/scripts/common/tools/amd_smi_utils.py b/src/madengine/scripts/common/tools/amd_smi_utils.py
index e0e48096..ab245c8e 100644
--- a/src/madengine/scripts/common/tools/amd_smi_utils.py
+++ b/src/madengine/scripts/common/tools/amd_smi_utils.py
@@ -6,16 +6,18 @@
Copyright (c) Advanced Micro Devices, Inc. All rights reserved.
"""
+import os
import sys
import logging
from typing import List, Optional, Dict, Any
-sys.path.append("/opt/rocm/libexec/amdsmi_cli/")
+_ROCM_PATH = os.environ.get("ROCM_PATH", "/opt/rocm")
+sys.path.append(f"{_ROCM_PATH}/libexec/amdsmi_cli/")
try:
from amdsmi_init import amdsmi_interface
from amdsmi_init import amdsmi_cli_init, amdsmi_cli_shutdown
except ImportError:
- raise ImportError("Could not import /opt/rocm/libexec/amdsmi_cli/amdsmi_init.py")
+ raise ImportError(f"Could not import {_ROCM_PATH}/libexec/amdsmi_cli/amdsmi_init.py")
class ProfUtils:
diff --git a/src/madengine/scripts/common/tools/gpu_info_profiler.py b/src/madengine/scripts/common/tools/gpu_info_profiler.py
index 111f655d..14949954 100644
--- a/src/madengine/scripts/common/tools/gpu_info_profiler.py
+++ b/src/madengine/scripts/common/tools/gpu_info_profiler.py
@@ -14,6 +14,7 @@
import sys
import csv
import os
+import shutil
import logging
import typing
import signal
@@ -27,10 +28,11 @@ def check_amd_smi_available() -> bool:
bool: True if amd-smi is available, False otherwise.
"""
# First check for Python bindings (more reliable for programmatic access)
+ rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
try:
- sys.path.append("/opt/rocm/libexec/amdsmi_cli/")
+ sys.path.append(f"{rocm_path}/libexec/amdsmi_cli/")
from amdsmi_init import amdsmi_interface
- logging.debug("amd-smi Python bindings found at /opt/rocm/libexec/amdsmi_cli/")
+ logging.debug(f"amd-smi Python bindings found at {rocm_path}/libexec/amdsmi_cli/")
return True
except ImportError:
logging.debug("amd-smi Python bindings not found")
@@ -74,9 +76,11 @@ def get_rocm_version() -> Optional[float]:
logging.debug(f"hipconfig check failed: {e}")
try:
- # Fallback to /opt/rocm/.info/version
- if os.path.exists("/opt/rocm/.info/version"):
- result = subprocess.run(['cat', '/opt/rocm/.info/version'],
+ # Fallback to $ROCM_PATH/.info/version
+ rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
+ version_file = f"{rocm_path}/.info/version"
+ if os.path.exists(version_file):
+ result = subprocess.run(['cat', version_file],
capture_output=True, text=True, timeout=10)
if result.returncode == 0:
version_str = result.stdout.strip().split('-')[0] # Remove build suffix
@@ -96,15 +100,21 @@ def detect_gpu_vendor() -> tuple[bool, bool]:
Raises:
ValueError: If no GPU management tools are found.
"""
- if os.path.exists("/usr/bin/nvidia-smi"):
+ rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
+ if shutil.which("nvidia-smi"):
return True, False
- elif os.path.exists("/opt/rocm/bin/rocm-smi") or check_amd_smi_available():
+ elif (
+ shutil.which("rocm-smi")
+ or os.path.exists(f"{rocm_path}/bin/rocm-smi")
+ or check_amd_smi_available()
+ ):
return False, True
else:
error_msg = (
"Unable to detect GPU vendor. No GPU management tools found.\n"
- "For NVIDIA: /usr/bin/nvidia-smi not found\n"
- "For AMD: /opt/rocm/bin/rocm-smi and amd-smi not found\n\n"
+ "For NVIDIA: nvidia-smi not found on PATH\n"
+ f"For AMD: rocm-smi not found on PATH or in {rocm_path}/bin, "
+ "and amd-smi not available\n\n"
"Please ensure:\n"
" 1. GPU drivers are installed\n"
" 2. For AMD GPUs: ROCm is properly installed (https://rocm.docs.amd.com)\n"
diff --git a/src/madengine/scripts/common/tools/rocm_smi_utils.py b/src/madengine/scripts/common/tools/rocm_smi_utils.py
index dd73219b..0ef5772f 100644
--- a/src/madengine/scripts/common/tools/rocm_smi_utils.py
+++ b/src/madengine/scripts/common/tools/rocm_smi_utils.py
@@ -6,16 +6,18 @@
Copyright (c) Advanced Micro Devices, Inc. All rights reserved.
"""
+import os
import sys
import logging
from typing import List
-sys.path.append("/opt/rocm/libexec/rocm_smi/")
+_ROCM_PATH = os.environ.get("ROCM_PATH", "/opt/rocm")
+sys.path.append(f"{_ROCM_PATH}/libexec/rocm_smi/")
try:
import rocm_smi
from rsmiBindings import *
except ImportError:
- raise ImportError("Could not import /opt/rocm/libexec/rocm_smi/rocm_smi.py")
+ raise ImportError(f"Could not import {_ROCM_PATH}/libexec/rocm_smi/rocm_smi.py")
class ProfUtils:
diff --git a/tests/integration/test_orchestrator_workflows.py b/tests/integration/test_orchestrator_workflows.py
index 143e56e5..5e51cc7d 100644
--- a/tests/integration/test_orchestrator_workflows.py
+++ b/tests/integration/test_orchestrator_workflows.py
@@ -323,6 +323,9 @@ def test_run_orchestrator_initialization(self, mock_context):
mock_args = MagicMock()
mock_args.additional_context = None
mock_args.live_output = True
+ # A bare MagicMock auto-vivifies truthy attributes; pin the flags this
+ # assertion depends on so they cannot leak into additional_context.
+ mock_args.require_pinned_image = False
orchestrator = RunOrchestrator(mock_args)
@@ -335,6 +338,7 @@ def test_run_orchestrator_additional_context_parsing(self):
mock_args = MagicMock()
mock_args.additional_context = '{"deploy": "slurm", "slurm": {"nodes": 4}}'
mock_args.live_output = True
+ mock_args.require_pinned_image = False
orchestrator = RunOrchestrator(mock_args)
diff --git a/tests/unit/test_auth.py b/tests/unit/test_auth.py
index 47d4a6c7..d2bf2fd8 100644
--- a/tests/unit/test_auth.py
+++ b/tests/unit/test_auth.py
@@ -1,10 +1,16 @@
"""Unit tests for madengine.core.auth module."""
+import json
import os
import pytest
from unittest.mock import MagicMock, mock_open, patch
-from madengine.core.auth import load_credentials, login_to_registry
+from madengine.core.auth import (
+ explain_registry_denial,
+ has_ambient_docker_auth,
+ load_credentials,
+ login_to_registry,
+)
class TestLoadCredentials:
@@ -115,21 +121,23 @@ def test_load_credentials_non_dockerhub_registry(self, mock_file, mock_exists):
assert result["custom_registry"]["token"] == "abc123"
+@patch.dict(os.environ, {"MAD_SKIP_DOCKER_LOGIN": ""}, clear=False)
+@patch("madengine.core.auth.has_ambient_docker_auth", return_value=False)
class TestLoginToRegistry:
- """Tests for login_to_registry()."""
+ """Tests for login_to_registry() when the machine has no existing docker login."""
def _mocks(self):
console = MagicMock()
rich_console = MagicMock()
return console, rich_console
- def test_no_credentials_returns_early(self):
+ def test_no_credentials_returns_early(self, mock_ambient):
"""Passing None credentials logs a warning and returns without error."""
console, rich_console = self._mocks()
login_to_registry("docker.io", None, console, rich_console)
console.sh.assert_not_called()
- def test_missing_registry_key_raises_when_raise_on_failure(self):
+ def test_missing_registry_key_raises_when_raise_on_failure(self, mock_ambient):
"""RuntimeError raised when registry key absent and raise_on_failure=True."""
console, rich_console = self._mocks()
credentials = {"other_registry": {"username": "u", "password": "p"}}
@@ -137,14 +145,14 @@ def test_missing_registry_key_raises_when_raise_on_failure(self):
login_to_registry("myregistry.io", credentials, console, rich_console, raise_on_failure=True)
console.sh.assert_not_called()
- def test_missing_registry_key_returns_when_not_raise_on_failure(self):
+ def test_missing_registry_key_returns_when_not_raise_on_failure(self, mock_ambient):
"""Returns silently when registry key absent and raise_on_failure=False."""
console, rich_console = self._mocks()
credentials = {"other_registry": {"username": "u", "password": "p"}}
login_to_registry("myregistry.io", credentials, console, rich_console, raise_on_failure=False)
console.sh.assert_not_called()
- def test_invalid_credentials_format_raises(self):
+ def test_invalid_credentials_format_raises(self, mock_ambient):
"""RuntimeError raised when username/password fields missing."""
console, rich_console = self._mocks()
credentials = {"dockerhub": {"token": "abc"}}
@@ -152,14 +160,22 @@ def test_invalid_credentials_format_raises(self):
login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=True)
console.sh.assert_not_called()
- def test_invalid_credentials_format_returns_when_not_raise_on_failure(self):
+ def test_invalid_credentials_format_returns_when_not_raise_on_failure(self, mock_ambient):
"""Returns silently when credentials format invalid and raise_on_failure=False."""
console, rich_console = self._mocks()
credentials = {"dockerhub": {"token": "abc"}}
login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=False)
console.sh.assert_not_called()
- def test_docker_io_normalised_to_dockerhub(self):
+ def test_blank_credentials_raise_without_ambient_auth(self, mock_ambient):
+ """Placeholder credentials are treated as absent, not as credentials."""
+ console, rich_console = self._mocks()
+ credentials = {"dockerhub": {"repository": "r", "username": "", "password": ""}}
+ with pytest.raises(RuntimeError, match="username|password"):
+ login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=True)
+ console.sh.assert_not_called()
+
+ def test_docker_io_normalised_to_dockerhub(self, mock_ambient):
"""docker.io registry is looked up under the 'dockerhub' key."""
console, rich_console = self._mocks()
credentials = {"dockerhub": {"username": "user", "password": "pass"}}
@@ -169,7 +185,7 @@ def test_docker_io_normalised_to_dockerhub(self):
# docker.io should not appear in the login command (uses default DockerHub endpoint)
assert "docker.io" not in cmd
- def test_custom_registry_included_in_command(self):
+ def test_custom_registry_included_in_command(self, mock_ambient):
"""Non-DockerHub registry URL is included in the login command."""
console, rich_console = self._mocks()
credentials = {"myregistry.io": {"username": "user", "password": "pass"}}
@@ -178,7 +194,7 @@ def test_custom_registry_included_in_command(self):
cmd = console.sh.call_args[0][0]
assert "myregistry.io" in cmd
- def test_login_failure_raises_when_raise_on_failure(self):
+ def test_login_failure_raises_when_raise_on_failure(self, mock_ambient):
"""docker login error is re-raised when raise_on_failure=True."""
console, rich_console = self._mocks()
console.sh.side_effect = RuntimeError("auth failed")
@@ -186,10 +202,182 @@ def test_login_failure_raises_when_raise_on_failure(self):
with pytest.raises(RuntimeError, match="auth failed"):
login_to_registry(None, credentials, console, rich_console, raise_on_failure=True)
- def test_login_failure_suppressed_when_not_raise_on_failure(self):
+ def test_login_failure_suppressed_when_not_raise_on_failure(self, mock_ambient):
"""docker login error is suppressed when raise_on_failure=False."""
console, rich_console = self._mocks()
console.sh.side_effect = RuntimeError("auth failed")
credentials = {"dockerhub": {"username": "user", "password": "pass"}}
login_to_registry(None, credentials, console, rich_console, raise_on_failure=False)
# Should not propagate the exception
+
+
+@patch.dict(os.environ, {"MAD_SKIP_DOCKER_LOGIN": ""}, clear=False)
+class TestLoginToRegistryWithAmbientAuth:
+ """Tests for login_to_registry() when the machine already has a docker login."""
+
+ def _mocks(self):
+ return MagicMock(), MagicMock()
+
+ @patch("madengine.core.auth.has_ambient_docker_auth", return_value=True)
+ def test_blank_credentials_defer_to_ambient_auth(self, mock_ambient):
+ """Blank credentials never override or break an existing docker login."""
+ console, rich_console = self._mocks()
+ credentials = {"dockerhub": {"repository": "r", "username": "", "password": ""}}
+ # No raise even with raise_on_failure=True: the machine is authenticated.
+ login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=True)
+ console.sh.assert_not_called()
+
+ @patch("madengine.core.auth.has_ambient_docker_auth", return_value=True)
+ def test_missing_registry_key_defers_to_ambient_auth(self, mock_ambient):
+ """A registry with no credential.json entry falls back to the existing login."""
+ console, rich_console = self._mocks()
+ credentials = {"other_registry": {"username": "u", "password": "p"}}
+ login_to_registry("myregistry.io", credentials, console, rich_console, raise_on_failure=True)
+ console.sh.assert_not_called()
+
+ @patch("madengine.core.auth.has_ambient_docker_auth", return_value=True)
+ def test_explicit_credentials_win_over_ambient_auth(self, mock_ambient):
+ """Usable explicit credentials still trigger a login (explicit wins)."""
+ console, rich_console = self._mocks()
+ credentials = {"dockerhub": {"username": "user", "password": "pass"}}
+ login_to_registry("docker.io", credentials, console, rich_console)
+ console.sh.assert_called_once()
+ assert "--username user" in console.sh.call_args[0][0]
+
+ @patch("madengine.core.auth.has_ambient_docker_auth", return_value=False)
+ def test_whitespace_only_credentials_are_not_credentials(self, mock_ambient):
+ """Whitespace-only values are treated as blank."""
+ console, rich_console = self._mocks()
+ credentials = {"dockerhub": {"username": " ", "password": "\t"}}
+ with pytest.raises(RuntimeError, match="username|password"):
+ login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=True)
+ console.sh.assert_not_called()
+
+
+class TestSkipDockerLogin:
+ """Tests for the MAD_SKIP_DOCKER_LOGIN escape hatch."""
+
+ @patch.dict(os.environ, {"MAD_SKIP_DOCKER_LOGIN": "1"}, clear=False)
+ def test_skip_env_var_bypasses_login(self):
+ """MAD_SKIP_DOCKER_LOGIN=1 defers to ambient credentials unconditionally."""
+ console, rich_console = MagicMock(), MagicMock()
+ credentials = {"dockerhub": {"username": "user", "password": "pass"}}
+ login_to_registry("docker.io", credentials, console, rich_console, raise_on_failure=True)
+ console.sh.assert_not_called()
+
+
+class TestHasAmbientDockerAuth:
+ """Tests for has_ambient_docker_auth()."""
+
+ def _write_config(self, tmp_path, config):
+ (tmp_path / "config.json").write_text(json.dumps(config), encoding="utf-8")
+ return {"DOCKER_CONFIG": str(tmp_path)}
+
+ def test_dockerhub_auth_entry_detected(self, tmp_path):
+ """A Docker Hub entry with an auth blob counts as authenticated."""
+ env = self._write_config(
+ tmp_path, {"auths": {"https://index.docker.io/v1/": {"auth": "abc123"}}}
+ )
+ with patch.dict(os.environ, env, clear=False):
+ assert has_ambient_docker_auth(None) is True
+ assert has_ambient_docker_auth("docker.io") is True
+ assert has_ambient_docker_auth("docker.io/rocm/mad-private") is True
+ assert has_ambient_docker_auth("myregistry.io") is False
+
+ def test_identity_token_entry_detected(self, tmp_path):
+ """An identitytoken-only entry counts as authenticated."""
+ env = self._write_config(
+ tmp_path, {"auths": {"index.docker.io": {"identitytoken": "tok"}}}
+ )
+ with patch.dict(os.environ, env, clear=False):
+ assert has_ambient_docker_auth("docker.io") is True
+
+ def test_cred_helper_detected(self, tmp_path):
+ """A credHelpers entry counts as authenticated."""
+ env = self._write_config(tmp_path, {"credHelpers": {"myregistry.io": "ecr-login"}})
+ with patch.dict(os.environ, env, clear=False):
+ assert has_ambient_docker_auth("myregistry.io/team/img") is True
+ assert has_ambient_docker_auth("docker.io") is False
+
+ def test_creds_store_with_empty_auth_entry(self, tmp_path):
+ """credsStore-managed entries are stored empty but are still credentials."""
+ env = self._write_config(
+ tmp_path,
+ {"credsStore": "desktop", "auths": {"https://index.docker.io/v1/": {}}},
+ )
+ with patch.dict(os.environ, env, clear=False):
+ assert has_ambient_docker_auth("docker.io") is True
+ assert has_ambient_docker_auth("other.io") is False
+
+ def test_empty_auth_entry_without_creds_store(self, tmp_path):
+ """An empty entry with no credential store is not usable."""
+ env = self._write_config(tmp_path, {"auths": {"https://index.docker.io/v1/": {}}})
+ with patch.dict(os.environ, env, clear=False):
+ assert has_ambient_docker_auth("docker.io") is False
+
+ def test_registry_with_port_matches_host(self, tmp_path):
+ """host:port registries are matched on the host segment."""
+ env = self._write_config(tmp_path, {"auths": {"localhost:5000": {"auth": "x"}}})
+ with patch.dict(os.environ, env, clear=False):
+ assert has_ambient_docker_auth("localhost:5000/team/img") is True
+
+ def test_missing_config_returns_false(self, tmp_path):
+ """A missing config.json yields False rather than raising."""
+ with patch.dict(os.environ, {"DOCKER_CONFIG": str(tmp_path)}, clear=False):
+ assert has_ambient_docker_auth("docker.io") is False
+
+ def test_corrupt_config_returns_false(self, tmp_path):
+ """A corrupt config.json yields False rather than raising."""
+ (tmp_path / "config.json").write_text("not json{{{", encoding="utf-8")
+ with patch.dict(os.environ, {"DOCKER_CONFIG": str(tmp_path)}, clear=False):
+ assert has_ambient_docker_auth("docker.io") is False
+
+
+class TestExplainRegistryDenial:
+ """Tests for explain_registry_denial()."""
+
+ def test_insufficient_scope_is_an_authorization_problem(self):
+ """insufficient_scope is reported as authorization, not authentication."""
+ log = (
+ "#6 ERROR: pull access denied, repository does not exist or may require "
+ "authorization: server message: insufficient_scope: authorization failed"
+ )
+ hint = explain_registry_denial(log, "rocm/triton-inference-server-dev:x")
+ assert hint is not None
+ assert "authorization problem" in hint
+ assert "rocm/triton-inference-server-dev:x" in hint
+ assert "will not fix it" in hint
+
+ def test_plain_denial_points_at_login(self):
+ """A denial without insufficient_scope points at supplying credentials."""
+ log = "Error response from daemon: pull access denied for rocm/private"
+ hint = explain_registry_denial(log, "rocm/private:latest")
+ assert hint is not None
+ assert "docker login" in hint
+ assert "authorization problem" not in hint
+
+ def test_non_dockerhub_denial_names_that_registry(self):
+ """A denial for another registry does not suggest Docker Hub credentials."""
+ log = "Error response from daemon: pull access denied for ghcr.io/org/app"
+ hint = explain_registry_denial(log, "ghcr.io/org/app:latest")
+ assert hint is not None
+ assert "docker login ghcr.io" in hint
+ assert '"ghcr.io": {"username"' in hint
+ assert "dockerhub" not in hint
+ assert "MAD_DOCKERHUB" not in hint
+
+ def test_dockerhub_namespace_still_suggests_dockerhub(self):
+ """A bare namespace/repo reference is Docker Hub, not a registry host."""
+ log = "Error response from daemon: pull access denied for rocm/private"
+ hint = explain_registry_denial(log, "rocm/private:latest")
+ assert hint is not None
+ assert '"dockerhub"' in hint
+ assert "MAD_DOCKERHUB_USER" in hint
+
+ def test_unrelated_failure_returns_none(self):
+ """Non-registry build failures produce no hint."""
+ assert explain_registry_denial("RUN apt-get install failed: exit code 100") is None
+
+ def test_empty_log_returns_none(self):
+ """Empty output produces no hint."""
+ assert explain_registry_denial("") is None
diff --git a/tests/unit/test_container_runner.py b/tests/unit/test_container_runner.py
index aae79321..1899cf5c 100644
--- a/tests/unit/test_container_runner.py
+++ b/tests/unit/test_container_runner.py
@@ -523,3 +523,175 @@ def test_skip_model_run_false_does_exec_script(self):
assert any(
"run.sh" in c and "cd " in c for c in docker_sh_calls
), f"Model script was not executed: {docker_sh_calls}"
+
+
+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_local_image_entry_with_registry_ref_is_not_a_bypass(self, _mock_csv):
+ """build_info["local_image"] must not smuggle an unpinned registry tag through.
+
+ The build-on-compute-node path writes entries carrying BOTH a truthy
+ local_image and a registry reference in docker_image. That branch runs
+ before the registry branch, so without enforcement here the flag is
+ silently a no-op for those manifests.
+ """
+ with tempfile.TemporaryDirectory() as tmpdir:
+ manifest_path = self._manifest(
+ tmpdir,
+ {
+ "local_image": "ci-m_ubuntu",
+ "docker_image": "myorg/ci:m",
+ "built_on_compute": True,
+ },
+ )
+ runner = self._runner()
+ runner.perf_csv_path = os.path.join(tmpdir, "perf.csv")
+ runner.additional_context = {"require_pinned_image": True}
+
+ with patch.object(
+ runner, "_ensure_local_image_available"
+ ) as mock_ensure, patch.object(runner, "run_container") as mock_run:
+ result = runner.run_models_from_manifest(
+ manifest_file=manifest_path, timeout=60
+ )
+
+ mock_ensure.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_local_image_already_pinned_is_accepted(self, _mock_csv):
+ """An explicitly digest-pinned MAD_CONTAINER_IMAGE already meets the guarantee."""
+ pinned = f"myorg/ci@{DIGEST}"
+ with tempfile.TemporaryDirectory() as tmpdir:
+ manifest_path = self._manifest(
+ tmpdir, {"local_image": True, "docker_image": pinned}
+ )
+ runner = self._runner()
+ runner.perf_csv_path = os.path.join(tmpdir, "perf.csv")
+ runner.additional_context = {"require_pinned_image": True}
+
+ with patch.object(runner, "_ensure_local_image_available"), patch.object(
+ runner, "run_container", return_value={"status": "SUCCESS"}
+ ) as mock_run:
+ runner.run_models_from_manifest(manifest_file=manifest_path, timeout=60)
+
+ assert mock_run.call_args[1]["docker_image"] == pinned
+
+ @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}")
diff --git a/tests/unit/test_docker_builder.py b/tests/unit/test_docker_builder.py
index 3fe97f9b..c9d63171 100644
--- a/tests/unit/test_docker_builder.py
+++ b/tests/unit/test_docker_builder.py
@@ -3,12 +3,14 @@
Currently covers registry image naming for single-arch builds (credentials β repository:tag).
"""
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
import pytest
from madengine.execution.docker_builder import DockerBuilder
+DIGEST = "sha256:" + "df36ef7e" * 8
+
@pytest.fixture
def docker_builder():
@@ -44,3 +46,165 @@ def test_create_registry_image_name_without_credentials_matches_local_tag(docker
None,
)
assert out == "ci-dummy_dummy.ubuntu.amd"
+
+
+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 == {}
+
+
+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]
diff --git a/tests/unit/test_execution.py b/tests/unit/test_execution.py
index dc18121e..b6c844a3 100644
--- a/tests/unit/test_execution.py
+++ b/tests/unit/test_execution.py
@@ -88,6 +88,19 @@ def test_non_ci_tag_sanitizes_full_ref(self):
def test_short_ci_tag_unchanged(self):
assert _docker_image_ref_for_log_naming("ci-model_ubuntu") == "ci-model_ubuntu"
+ 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"
+ )
+
class TestMakeRunLogFilePath:
"""make_run_log_file_path behavior."""
diff --git a/tests/unit/test_image_digest.py b/tests/unit/test_image_digest.py
new file mode 100644
index 00000000..58e24f18
--- /dev/null
+++ b/tests/unit/test_image_digest.py
@@ -0,0 +1,152 @@
+"""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.errors import ConfigurationError
+from madengine.core.image_digest import (
+ build_pinned_reference,
+ parse_push_digest,
+ parse_repo_digest,
+ resolve_pinned_image,
+)
+
+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}"
+
+
+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_enabled_accepts_already_pinned_reference_without_manifest_digest(self):
+ # A reference the user pinned themselves already provides the guarantee,
+ # so there is nothing to look up and nothing to reject.
+ pinned = f"myorg/ci@{DIGEST}"
+ assert resolve_pinned_image(pinned, None, True) == pinned
+
+ 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
diff --git a/tests/unit/test_k8s.py b/tests/unit/test_k8s.py
index 4a1f84a4..b4674b77 100644
--- a/tests/unit/test_k8s.py
+++ b/tests/unit/test_k8s.py
@@ -5,6 +5,8 @@
Integration/e2e tests stay in their own modules.
"""
+import json
+
import pytest
from madengine.deployment.k8s_names import (
@@ -24,10 +26,13 @@
build_registry_secret_data,
)
from madengine.deployment.kubernetes import (
+ KubernetesDeployment,
_pod_job_name_label_selector,
assign_pvc_subdirs_to_pods,
match_pvc_subdir_to_k8s_pod,
)
+from madengine.core.errors import ConfigurationError
+from madengine.deployment.base import DeploymentConfig
def test_merge_secrets_config_defaults():
@@ -274,3 +279,72 @@ def test_invalid_mode_falls_back_to_lite(self):
pre_scripts = []
mixin.gather_system_env_details(pre_scripts, "my_model", rocenv_mode="bogus")
assert pre_scripts[0]["args"] == "my_model_env lite UBUNTU"
+
+
+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
+ )
diff --git a/tests/unit/test_orchestration.py b/tests/unit/test_orchestration.py
index 54a11bd5..f198b9ed 100644
--- a/tests/unit/test_orchestration.py
+++ b/tests/unit/test_orchestration.py
@@ -16,6 +16,7 @@
from madengine.orchestration.build_orchestrator import BuildOrchestrator
from madengine.orchestration.run_orchestrator import RunOrchestrator
from madengine.core.errors import ConfigurationError
+from madengine.cli.utils import create_args_namespace
# ---- image_filtering ----
@@ -160,6 +161,9 @@ def test_initializes_with_args(self, mock_context):
mock_args = MagicMock()
mock_args.additional_context = None
mock_args.live_output = True
+ # A bare MagicMock auto-vivifies truthy attributes; pin the flags this
+ # assertion depends on so they cannot leak into additional_context.
+ mock_args.require_pinned_image = False
orchestrator = RunOrchestrator(mock_args)
@@ -359,3 +363,148 @@ def test_distributed_warns_on_local_only_flags(self, tmp_path):
assert "--keep-alive" in printed
assert "--skip-model-run" in printed
assert "--keep-model-dir" not in printed # was False, must not appear
+
+
+@pytest.mark.unit
+class TestCreateManifestFromLocalImage:
+ """MAD_CONTAINER_IMAGE (local image) mode must carry every models.json field
+ that ContainerRunner relies on -- including multiple_results, whose absence
+ silently drops perf-CSV-based result reporting (falls back to scraping the
+ log for a 'performance: NUMBER METRIC' line and reports FAILURE)."""
+
+ @patch("madengine.orchestration.run_orchestrator.Context")
+ def test_multiple_results_field_is_preserved(self, mock_context, tmp_path):
+ mock_context.return_value.ctx = {}
+ mock_args = MagicMock()
+ mock_args.additional_context = None
+ mock_args.live_output = False
+
+ orchestrator = RunOrchestrator(mock_args)
+ orchestrator.console = MagicMock()
+ orchestrator.rich_console = MagicMock()
+
+ fake_model = {
+ "name": "uber_storefront/v1t1",
+ "tags": ["inference"],
+ "scripts": "run.sh",
+ "n_gpus": "1",
+ "data": "uber_storefront_models",
+ "args": "--model-dir v1t1",
+ "multiple_results": "perf_uber_storefront_v1t1.csv",
+ }
+
+ manifest_output = str(tmp_path / "build_manifest.json")
+
+ with patch(
+ "madengine.utils.discover_models.DiscoverModels.run",
+ return_value=[fake_model],
+ ):
+ orchestrator._create_manifest_from_local_image(
+ image_name="registry.io/org/model:ci-tag",
+ tags=["inference"],
+ manifest_output=manifest_output,
+ )
+
+ with open(manifest_output) as f:
+ manifest = json.load(f)
+
+ built_model = next(iter(manifest["built_models"].values()))
+ assert built_model["multiple_results"] == "perf_uber_storefront_v1t1.csv"
+
+ @patch("madengine.orchestration.run_orchestrator.Context")
+ def test_multiple_results_defaults_to_empty_string(self, mock_context, tmp_path):
+ """Models without multiple_results in models.json still get the key (empty),
+ so ContainerRunner's model_info.get("multiple_results") lookups never KeyError."""
+ mock_context.return_value.ctx = {}
+ mock_args = MagicMock()
+ mock_args.additional_context = None
+ mock_args.live_output = False
+
+ orchestrator = RunOrchestrator(mock_args)
+ orchestrator.console = MagicMock()
+ orchestrator.rich_console = MagicMock()
+
+ fake_model = {
+ "name": "dummy/model",
+ "tags": ["inference"],
+ "scripts": "run.sh",
+ "n_gpus": "1",
+ "data": "",
+ "args": "",
+ }
+
+ manifest_output = str(tmp_path / "build_manifest.json")
+
+ with patch(
+ "madengine.utils.discover_models.DiscoverModels.run",
+ return_value=[fake_model],
+ ):
+ orchestrator._create_manifest_from_local_image(
+ image_name="registry.io/org/model:ci-tag",
+ tags=["inference"],
+ manifest_output=manifest_output,
+ )
+
+ with open(manifest_output) as f:
+ manifest = json.load(f)
+
+ built_model = next(iter(manifest["built_models"].values()))
+ assert built_model["multiple_results"] == ""
+
+
+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
diff --git a/tests/unit/test_slurm_job_template.py b/tests/unit/test_slurm_job_template.py
new file mode 100644
index 00000000..6ce19628
--- /dev/null
+++ b/tests/unit/test_slurm_job_template.py
@@ -0,0 +1,201 @@
+#!/usr/bin/env python3
+"""
+Unit tests for the generated SLURM job script (`job.sh.j2`).
+
+Locks in the portability contract points that clusters keep re-discovering
+downstream (see ROCm/rocm-systems#9055, which patched madengine's source
+rather than filing them):
+
+1. The job script puts madengine back on PATH itself instead of assuming the
+ batch environment inherited the submitter's PATH.
+2. The shared-filesystem probe recognizes `nfs4`, which is what `df -T`
+ reports on most modern NFS mounts.
+3. `slurm.skip_gpus_directive` removes `#SBATCH --gpus-per-node`, which a
+ cluster advertising no GPU GRES rejects outright.
+
+Copyright (c) Advanced Micro Devices, Inc. All rights reserved.
+"""
+
+import json
+import re
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+from madengine.deployment.base import DeploymentConfig
+from madengine.deployment.slurm import SlurmDeployment
+
+
+MODEL_ENTRY = {
+ "name": "dummy_torchrun_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": "",
+}
+
+
+def _build_deployment(
+ tmp_path: Path,
+ slurm_overrides: dict = None,
+ distributed_overrides: dict = None,
+) -> SlurmDeployment:
+ """SlurmDeployment over a minimal torchrun manifest, output_dir under tmp_path."""
+ manifest = {
+ "built_images": {"dummy-image": {"docker_image": "dummy:latest"}},
+ "built_models": {"dummy-image": 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))
+
+ slurm_config = {
+ "partition": "test-partition",
+ "nodes": 2,
+ "gpus_per_node": 8,
+ "time": "01:00:00",
+ "output_dir": str(tmp_path / "slurm_output"),
+ "exclusive": True,
+ }
+ slurm_config.update(slurm_overrides or {})
+
+ distributed_config = {
+ "launcher": "torchrun",
+ "nnodes": 2,
+ "nproc_per_node": 8,
+ "backend": "nccl",
+ "port": 29500,
+ }
+ distributed_config.update(distributed_overrides or {})
+
+ cfg = DeploymentConfig(
+ target="slurm",
+ manifest_file=str(manifest_path),
+ additional_context={
+ "deploy": "slurm",
+ "gpu_vendor": "AMD",
+ "guest_os": "UBUNTU",
+ "slurm": slurm_config,
+ "distributed": distributed_config,
+ },
+ )
+ return SlurmDeployment(cfg)
+
+
+def _render(deployment: SlurmDeployment) -> str:
+ """Render job.sh.j2 exactly as prepare() does, without submitting anything."""
+ context = deployment._prepare_template_context(MODEL_ENTRY)
+ return deployment.jinja_env.get_template("job.sh.j2").render(**context)
+
+
+# ---------------------------------------------------------------------------
+# 1. PATH is re-established inside the job
+
+class TestJobScriptPath:
+ """The job script must not depend on the submitter's PATH being inherited."""
+
+ def test_user_bin_dir_is_prepended(self, tmp_path):
+ script = _render(_build_deployment(tmp_path))
+ assert 'export PATH="$HOME/.local/bin:$PATH"' in script
+
+ def test_submission_bin_dir_is_prepended(self, tmp_path):
+ with patch("madengine.deployment.slurm.shutil.which", return_value="/opt/venv/bin/madengine"):
+ script = _render(_build_deployment(tmp_path))
+ assert 'export PATH="/opt/venv/bin:$PATH"' in script
+
+ def test_no_empty_export_when_cli_not_on_path(self, tmp_path):
+ """madengine missing at submission time must not render an empty PATH entry."""
+ with patch("madengine.deployment.slurm.shutil.which", return_value=None):
+ script = _render(_build_deployment(tmp_path))
+ assert 'export PATH=":$PATH"' not in script
+ assert 'export PATH="$HOME/.local/bin:$PATH"' in script
+
+ def test_path_is_set_before_madengine_is_looked_up(self, tmp_path):
+ """The export is useless if it lands after `command -v madengine`."""
+ with patch("madengine.deployment.slurm.shutil.which", return_value="/opt/venv/bin/madengine"):
+ script = _render(_build_deployment(tmp_path))
+ assert script.index('export PATH="/opt/venv/bin:$PATH"') < script.index("command -v madengine")
+
+
+# ---------------------------------------------------------------------------
+# 2. Shared-filesystem probe
+
+class TestSharedFilesystemProbe:
+ """`df -T` reports nfs4 on modern mounts; the probe must not miss it.
+
+ And it must read the filesystem type, nothing else: the mount point travels on the
+ same `df -T` line, so a local disk under a path such as /mnt/nfs-scratch used to answer
+ yes and the job then trusted node-local storage to be visible from every node.
+ """
+
+ @staticmethod
+ def _probe_pattern(script: str) -> str:
+ match = re.search(r"SUBMIT_FSTYPE\"?\s*\|\s*grep -qE '([^']+)'", script)
+ assert match, "shared-filesystem probe not found in rendered script"
+ return match.group(1)
+
+ @pytest.mark.parametrize("fstype,expected", [
+ ("nfs", True),
+ ("nfs3", True),
+ ("nfs4", True),
+ ("lustre", True),
+ ("gpfs", True),
+ ("ceph", True),
+ ("beegfs", True),
+ ("panfs", True),
+ ("ext4", False),
+ ("xfs", False),
+ ("overlay", False),
+ ("tmpfs", False),
+ ])
+ def test_probe_matches_shared_filesystems(self, tmp_path, fstype, expected):
+ # The probe only exists on the single-node branch of the template.
+ deployment = _build_deployment(tmp_path, {"nodes": 1}, {"nnodes": 1})
+ pattern = self._probe_pattern(_render(deployment))
+ assert bool(re.search(pattern, fstype)) is expected
+
+ def test_the_probe_reads_the_fstype_column_only(self, tmp_path):
+ script = _render(_build_deployment(tmp_path, {"nodes": 1}, {"nnodes": 1}))
+ assert 'df --output=fstype "$SUBMIT_DIR"' in script
+ assert 'df -T "$SUBMIT_DIR" 2>/dev/null | grep' not in script
+
+ def test_a_mount_point_that_says_nfs_does_not_make_a_disk_shared(self, tmp_path):
+ """/mnt/nfs-scratch on ext4 is local, whatever its name suggests."""
+ script = _render(_build_deployment(tmp_path, {"nodes": 1}, {"nnodes": 1}))
+ pattern = self._probe_pattern(script)
+ df_line = "/dev/nvme0n1p2 ext4 104857600 50106368 54751232 48% /mnt/nfs-scratch"
+ assert re.search(pattern, df_line) is None
+ assert re.search(pattern, "ext4") is None
+
+ def test_there_is_a_fallback_for_df_without_output(self, tmp_path):
+ """--output is coreutils 8.21; older df still has to be read correctly."""
+ script = _render(_build_deployment(tmp_path, {"nodes": 1}, {"nnodes": 1}))
+ assert "awk 'NR > 1 { print $2; exit }'" in script
+
+
+# ---------------------------------------------------------------------------
+# 3. GPU GRES directive opt-out
+
+class TestGpusPerNodeDirective:
+ """A cluster with GresTypes=(null) rejects any job carrying --gpus-per-node."""
+
+ def test_directive_present_by_default(self, tmp_path):
+ script = _render(_build_deployment(tmp_path))
+ assert "#SBATCH --gpus-per-node=8" in script
+
+ def test_directive_omitted_when_opted_out(self, tmp_path):
+ script = _render(_build_deployment(tmp_path, {"skip_gpus_directive": True}))
+ assert "--gpus-per-node" not in script
diff --git a/tests/unit/test_slurm_multi.py b/tests/unit/test_slurm_multi.py
index de5c2eae..ae1fac97 100644
--- a/tests/unit/test_slurm_multi.py
+++ b/tests/unit/test_slurm_multi.py
@@ -28,6 +28,7 @@
from madengine.deployment.common import VALID_LAUNCHERS, is_self_managed_launcher, normalize_launcher
from madengine.deployment.base import DeploymentConfig
from madengine.deployment.slurm import SlurmDeployment
+from madengine.core.errors import ConfigurationError
# ---------------------------------------------------------------------------
@@ -595,3 +596,90 @@ def test_below_minimum_nodes_raises(self, deployment_factory):
deployment_factory._generate_sglang_disagg_command(
nnodes=1, nproc_per_node=8, master_port=12345
)
+
+
+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, 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()