From f71a23350fd73fc7cb3be450c425b0b45a822420 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 25 Aug 2026 18:48:40 +0200 Subject: [PATCH 1/4] Validate Puzzletron configuration before launch Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 720 +++--------------- .../orchestration/execution.example.yaml | 12 + .../qwen3p5_0p8b/runner.slurm.yaml | 5 - .../qwen_moe/execution.production.yaml | 20 +- .../orchestration/qwen_moe/runner.slurm.yaml | 10 +- .../orchestration/runner.slurm.example.yaml | 12 +- .../configs/setup/defaults.example.yaml | 5 +- examples/puzzletron/docs/campaign_reports.md | 24 +- .../docs/configuration_overrides.md | 20 + examples/puzzletron/docs/environment_setup.md | 226 ++++++ .../puzzletron/docs/legacy_nano_campaign.md | 117 +++ .../docs/orchestration_operations.md | 46 ++ .../puzzletron/docs/qwen3p5_0p8b_smoke.md | 44 ++ examples/puzzletron/docs/setup_wizard.md | 74 ++ .../puzzletron/docs/slurm_configuration.md | 67 ++ examples/puzzletron/orchestrate.py | 58 +- modelopt/torch/puzzletron/_config_aliases.py | 65 ++ .../puzzletron/orchestration/adapters/pool.py | 2 +- .../orchestration/adapters/post_mip.py | 2 +- .../orchestration/adapters/sharded.py | 31 +- .../orchestration/adapters/stage_compat.py | 2 +- .../puzzletron/orchestration/compiler.py | 353 +++++++-- .../torch/puzzletron/orchestration/config.py | 53 +- .../puzzletron/orchestration/controller.py | 2 +- .../orchestration/executors/slurm.py | 22 +- .../puzzletron/orchestration/identity.py | 25 +- .../torch/puzzletron/orchestration/process.py | 64 ++ .../puzzletron/orchestration/reporting.py | 24 +- .../torch/puzzletron/orchestration/schema.py | 83 +- modelopt/torch/puzzletron/pipeline_config.py | 2 + puzzletron_setup/bundle.py | 22 +- puzzletron_setup/state.py | 12 + puzzletron_setup/v2/bundle.py | 56 +- puzzletron_setup/v2/defaults.py | 16 +- puzzletron_setup/v2/resolved.py | 18 +- puzzletron_setup/v2/wizard.py | 72 +- puzzletron_setup/v2/wizard_common.py | 6 +- puzzletron_setup/wizard.py | 26 +- .../unit/torch/puzzletron/test_data_config.py | 96 ++- .../torch/puzzletron/test_example_runner.py | 2 +- .../puzzletron/test_orchestration_compiler.py | 323 +++++++- .../test_orchestration_executors.py | 248 +++++- .../test_orchestration_lightweight.py | 133 +++- .../test_orchestration_reporting.py | 54 +- .../test_orchestration_shutdown_progress.py | 4 +- .../torch/puzzletron/test_portable_configs.py | 43 +- .../test_qwen3p5_0p8b_smoke_plan.py | 1 - .../torch/puzzletron/test_setup_bundle.py | 63 +- .../torch/puzzletron/test_setup_v2_quick.py | 61 ++ .../test_setup_v2_resolved_config.py | 97 ++- 50 files changed, 2579 insertions(+), 964 deletions(-) create mode 100644 examples/puzzletron/docs/configuration_overrides.md create mode 100644 examples/puzzletron/docs/environment_setup.md create mode 100644 examples/puzzletron/docs/legacy_nano_campaign.md create mode 100644 examples/puzzletron/docs/orchestration_operations.md create mode 100644 examples/puzzletron/docs/qwen3p5_0p8b_smoke.md create mode 100644 examples/puzzletron/docs/setup_wizard.md create mode 100644 examples/puzzletron/docs/slurm_configuration.md create mode 100644 modelopt/torch/puzzletron/_config_aliases.py create mode 100644 modelopt/torch/puzzletron/orchestration/process.py diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 86af42abd85..d27a0ee6bf3 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -8,19 +8,28 @@ distill the selected model. ## Table of Contents - [Start here](#start-here) -- [Setup wizard](#setup-wizard) - [Installation](#installation) +- [Setup wizard](#setup-wizard) - [Evaluate a checkpoint](#evaluate-a-checkpoint) - [Run with an agent](#run-with-an-agent) - [Configuration](#configuration) +- [Experiment overrides](#experiment-overrides) +- [Slurm configuration](#slurm-configuration) +- [Qwen 3.5 smoke test](#qwen-35-smoke-test) - [Run a campaign](#run-a-campaign) -- [Campaign stages](#campaign-stages) +- [Controller operations](#controller-operations) +- [MIP runs](#mip-runs) +- [Post-MIP pipelines](#post-mip-pipelines) +- [Sanity validation](#sanity-validation) - [Reports](#reports) +- [Legacy Nano campaign](#legacy-nano-campaign) +- [Architecture](#architecture) ## Start here -- **New campaign:** use the [setup wizard](#setup-wizard) to generate validated - smoke and production bundles. +- **New campaign:** complete the [installation](#installation), then use the + [setup wizard](#setup-wizard) to generate validated smoke and production + bundles. - **Generated campaign:** complete the [installation](#installation), then [run the campaign](#run-a-campaign) with its generated bundle. - **Checkpoint evaluation:** use [Evaluate a checkpoint](#evaluate-a-checkpoint) @@ -30,47 +39,34 @@ distill the selected model. - **Existing results:** see [Reports](#reports) to regenerate a campaign report or inspect the retained examples. -## Setup wizard - -The schema-driven Puzzletron v2 setup wizard inspects a local checkpoint or -Hugging Face model configuration and generates self-contained smoke and -production experiment, runner, and execution bundles. Its setup environment -does not require PyTorch or model weights. +## Installation -The repository-root [`puzzletron_setup`](../../puzzletron_setup/) package keeps -this configuration-only flow outside `modelopt.torch`, so starting the wizard -does not initialize the ModelOpt or PyTorch runtime. At the **Model** prompt, -provide an existing local checkpoint/config path or a Hugging Face model URL or -repository ID; the wizard reads configuration metadata, not model weights. At -the **Dataset** prompt, provide an existing local dataset path or a Hugging Face -dataset URL or repository ID. +See [environment setup](docs/environment_setup.md) for worker containers, +pinned CUDA and PyTorch packages, patched dependencies, model-specific kernels, +bare-metal environments, and verification. -Install the setup dependencies: +Use a lightweight environment for the setup wizard and controller: ```bash -python -m pip install -r examples/puzzletron/requirements-setup.txt +python3 -m venv .venv-puzzletron-control +source .venv-puzzletron-control/bin/activate +python -m pip install \ + -r examples/puzzletron/requirements-setup.txt \ + -r examples/puzzletron/requirements-orchestrator.txt ``` -The guided flow offers three profiles: +GPU workers use the environment or container declared in the generated runner +file. Prepare that environment before launch, then run the smoke bundle first. -- **Quick smoke** is the fastest way to verify the campaign shape. -- **Balanced pruning** is recommended for a first real campaign. -- **High-confidence search** spends more runtime on scoring and sanity checks. +## Setup wizard -The selected profile supplies pruning and search defaults from the detected -model family's `setup_v2_defaults.yaml`, including geometry-specific refinements -when available. Setup then asks for the model and dataset and requires explicit -acceptance or customization of infrastructure-specific worker and cluster -defaults. +See the [setup wizard guide](docs/setup_wizard.md) for profiles, hosted dataset +handling, full configuration mode, generated files, and resuming an interrupted +setup. -For a first-class hosted dataset, setup records a worker-visible local output -path. The explicitly selected campaign directory contains one generated -`README.md` runbook beside its smoke and production bundles; resuming setup -updates that same runbook rather than creating one per launch. Under **Prepare -dataset**, it contains the exact acquisition command. The wizard inspects -dataset metadata but does not download or materialize rows. Run that command -from the full worker environment before launching the campaign. A custom local -dataset is treated as already prepared and is referenced directly. +The setup wizard reads a local checkpoint or Hugging Face model configuration +and generates validated smoke and production bundles. It does not load model +weights. Start the wizard with the repository's example defaults file: @@ -79,250 +75,17 @@ python examples/puzzletron/puzzletron_setup_v2.py \ --defaults examples/puzzletron/configs/setup/defaults.example.yaml ``` -The example defaults use only repository-relative values. Copy the file and add -site-specific data, scheduler, and container settings before selecting it. -The defaults file is loaded only when passed explicitly and takes precedence -over the selected profile. To expose every per-section and nested setting, use -the advanced flow explicitly: - -```bash -python examples/puzzletron/puzzletron_setup_v2.py --full -``` - -Press **Esc** to go back from any prompt. Selection prompts show a visible -**← Back** action, and text or numeric prompts accept `:back`. -Every accepted answer and the exact navigation frame are saved in -`answers_v2.yaml`, so an interrupted session can resume with: - -```bash -python examples/puzzletron/puzzletron_setup_v2.py --resume /path/to/campaign -``` - -The wizard supports reusable execution profiles, multiple deployment -measurements, independent optimization goals, and editable downstream flows. -The defaults keep the common path concise while preserving detailed controls -for advanced campaigns. See [Configuration](#configuration) for the generated -bundle structure and extension points. - -The final review writes `resolved_defaults.yaml`, `README.md`, and validated -`smoke/` and `production/` bundles transactionally. The wizard validates both -bundles and writes a `dry-run-plan.txt` file in each, but neither bundle is -submitted and production is not gated on smoke. The wizard never launches the -orchestrator. - -## Installation - -Puzzletron uses one Python environment for ModelOpt, the patched vLLM fork, -AutoModel, and official AIPerf. Install PyTorch first and build every CUDA -extension against that same installation; mixing PyTorch or CUDA builds can -cause import failures or incorrect GPU execution. AIPerf uses the official PyPI -package; no custom AIPerf fork is required. - -### 1. Choose a compatible worker environment - -For a reproducible public bootstrap, start with: - -```text -nvcr.io/nvidia/cuda:12.9.2-cudnn-devel-ubuntu24.04 -``` - -This image is a bootstrap example, not a required runner image. For a Slurm -campaign, set `runner.execution_contract.container` to an image or path accepted -at your site, or leave it unset to execute directly in the worker environment. -Bare-metal runners use the host environment selected by -`runner.execution_contract.venv`. In either case, keep the CUDA and PyTorch -combination compatible with the pinned `cu129` packages below and run the -environment checks before launch. - -The bootstrap commands below assume a container. For bare-metal runners, skip -the Docker example and the `/workspace` and `apt-get` steps. Install equivalent -Python and build dependencies through your site's host-environment tooling, -then create or select the worker virtual environment referenced by -`runner.execution_contract.venv` and adapt the remaining paths accordingly. - -For example: - -```bash -export PUZZLETRON_WORKSPACE=/absolute/path/to/workspace -docker run --gpus all --ipc=host --rm -it \ - -v "${PUZZLETRON_WORKSPACE}:/workspace" \ - -w /workspace \ - nvcr.io/nvidia/cuda:12.9.2-cudnn-devel-ubuntu24.04 bash -``` - -Inside the container, install Python 3.12 and the build tools used by editable -packages and optional CUDA extensions: - -```bash -apt-get update -DEBIAN_FRONTEND=noninteractive apt-get install -y \ - build-essential cmake git ninja-build \ - python3 python3-dev python3-pip python3-venv -``` - -### 2. Clone the tracked forks - -Keep ModelOpt and the two Puzzletron forks as siblings: - -The core compatibility pins used by the CPU CI lane are recorded once in the -machine-readable [CI environment](ci_environment.json). Nox reads that file -directly. The full GPU setup below uses the same core package versions and adds -the CUDA-specific builds, patched vLLM runtime, and AIPerf. - -```bash -export MODEL_OPT_ROOT=/workspace/modelopt -export VLLM_ROOT=/workspace/vllm -export AUTOMODEL_ROOT=/workspace/Automodel -export PUZZLETRON_CI_ENVIRONMENT="${MODEL_OPT_ROOT}/examples/puzzletron/ci_environment.json" -export AUTOMODEL_REF="$(python3 -c \ - 'import json, sys; print(json.load(open(sys.argv[1]))["nemo_automodel"]["commit"])' \ - "${PUZZLETRON_CI_ENVIRONMENT}")" - -git clone --branch feature/add_anymodel_to_vllm --single-branch \ - https://github.com/Separius/vllm.git "${VLLM_ROOT}" -git clone --branch puzzletron --single-branch \ - https://github.com/Separius/Automodel.git "${AUTOMODEL_ROOT}" -git -C "${AUTOMODEL_ROOT}" checkout --detach "${AUTOMODEL_REF}" -``` - -```text -/workspace/ -├── modelopt/ -├── vllm/ -└── Automodel/ -``` - -### 3. Create the environment and install runtime packages - -The patched vLLM branch uses the PyTorch version recorded in the CI environment -with CUDA 12.9. Install that combination before anything that compiles CUDA -code: - -```bash -python3 -m venv /workspace/.venv -source /workspace/.venv/bin/activate - -export PUZZLETRON_TORCH_VERSION="$(python -c \ - 'import json, sys; print(json.load(open(sys.argv[1]))["torch"])' \ - "${PUZZLETRON_CI_ENVIRONMENT}")" -export PUZZLETRON_TORCHVISION_VERSION="$(python -c \ - 'import json, sys; print(json.load(open(sys.argv[1]))["torchvision"])' \ - "${PUZZLETRON_CI_ENVIRONMENT}")" -export PUZZLETRON_TRANSFORMERS_VERSION="$(python -c \ - 'import json, sys; print(json.load(open(sys.argv[1]))["transformers"])' \ - "${PUZZLETRON_CI_ENVIRONMENT}")" - -python -m pip install --upgrade \ - pip "setuptools>=80,<81" "setuptools-scm>=8" setuptools-rust \ - wheel "packaging>=24.2" "cmake>=3.26.1" ninja jinja2 - -python -m pip install \ - "torch==${PUZZLETRON_TORCH_VERSION}" \ - "torchvision==${PUZZLETRON_TORCHVISION_VERSION}" \ - "torchaudio==${PUZZLETRON_TORCH_VERSION}" \ - --index-url https://download.pytorch.org/whl/cu129 - -VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 \ - python -m pip install --no-build-isolation -e "${VLLM_ROOT}" - -python -m pip install -e "${AUTOMODEL_ROOT}" -python -m pip install aiperf -python -m pip install -e "${MODEL_OPT_ROOT}[hf,puzzletron]" -python -m pip install "transformers==${PUZZLETRON_TRANSFORMERS_VERSION}" -python -m pip install -r "${MODEL_OPT_ROOT}/examples/puzzletron/requirements.txt" -``` - -Do not add `--no-deps`: the packages need their declared Python dependencies. -`--no-build-isolation` ensures compiled extensions use the active PyTorch -installation; it does not disable dependency installation. - -Install only the model-specific kernels required by the target architecture: - -```bash -# Mixture of experts -python -m pip install --no-build-isolation \ - "git+https://github.com/fanshiqing/grouped_gemm@v1.1.4" - -# Mamba -python -m pip install "mamba-ssm[causal-conv1d]" --no-build-isolation - -# Linear attention -python -m pip install "flash-linear-attention[cuda]" -``` - -### 4. Verify the exact environment - -Run these checks inside the same container and venv used for Puzzletron jobs: - -```bash -test "$(git -C "${VLLM_ROOT}" remote get-url origin)" = \ - "https://github.com/Separius/vllm.git" -test "$(git -C "${VLLM_ROOT}" branch --show-current)" = \ - "feature/add_anymodel_to_vllm" -test "$(git -C "${AUTOMODEL_ROOT}" remote get-url origin)" = \ - "https://github.com/Separius/Automodel.git" -test "$(git -C "${AUTOMODEL_ROOT}" rev-parse HEAD)" = "${AUTOMODEL_REF}" - -git -C "${MODEL_OPT_ROOT}" rev-parse HEAD -git -C "${VLLM_ROOT}" rev-parse HEAD -git -C "${AUTOMODEL_ROOT}" rev-parse HEAD -``` - -```bash -python - <<'PY' -import importlib.metadata as metadata -import json -import os - -from packaging.version import Version - -import aiperf -import lmms_eval -import modelopt -import nemo_automodel -import torch -import transformers -import vllm - -with open(os.environ["PUZZLETRON_CI_ENVIRONMENT"], encoding="utf-8") as stream: - ci_environment = json.load(stream) - -for package in ( - "torch", - "vllm", - "nemo-automodel", - "aiperf", - "lmms-eval", - "nvidia-modelopt", -): - print(package, metadata.version(package)) - -print("torch CUDA", torch.version.cuda) -print("CUDA available", torch.cuda.is_available()) -print("modelopt", modelopt.__file__) -print("vllm", vllm.__file__) - -assert Version(torch.__version__).release == Version(ci_environment["torch"]).release -assert Version(metadata.version("torchvision")).release == Version( - ci_environment["torchvision"] -).release -assert transformers.__version__ == ci_environment["transformers"] -assert metadata.version("lmms-eval") == ci_environment["lmms_eval"] -assert Version(metadata.version("nemo-automodel")).base_version == ( - ci_environment["nemo_automodel"]["base_version"] -) -assert torch.version.cuda == "12.9" -assert torch.cuda.is_available() -PY - -python -m pip check -``` - -Record the three source revisions and verification output with the campaign. -Re-run verification after pulling either fork or rebuilding a CUDA extension. +Choose **Balanced pruning** for a first campaign, review the detected model and +infrastructure settings, and select an output directory. The generated +`README.md` contains any dataset preparation command and the exact paths for +the smoke and production bundles. The wizard prepares files but does not submit +jobs. ## Evaluate a checkpoint +See [checkpoint evaluation](docs/checkpoint_evaluation.md) for task selection, +full evaluation, result locations, and model-detection overrides. + Basic evaluation is independent of MIP and the campaign DAG. In the Puzzletron worker environment, run any compatible local Hugging Face checkpoint directly: @@ -333,11 +96,9 @@ python examples/puzzletron/evaluate_lmms_checkpoint.py \ ``` The default one-GPU smoke evaluates eight samples each from IFEval and GSM8K. -Qwen 3.5 checkpoints are configured automatically. See -[checkpoint evaluation](docs/checkpoint_evaluation.md) to choose tasks, run a -full evaluation, find results, or override model detection. For options not -covered by the convenience command, append `--lmms-eval-args` followed by the -native lmms-eval options. +Qwen 3.5 checkpoints are configured automatically. For options not covered by +the convenience command, append `--lmms-eval-args` followed by the native +lmms-eval options. ## Run with an agent @@ -382,96 +143,41 @@ export PUZZLETRON_RUN_ROOT=/shared/puzzle_runs/my_campaign `PUZZLETRON_RUN_ROOT` is a convenience used by the checked-in experiment YAMLs to resolve `puzzle_dir`. Generated bundles write their chosen `puzzle_dir` directly. In both cases, `puzzle_dir` is the canonical location for artifacts, -logs, manifests, and controller state. - -Independent runs, variants, solution pools, resource constraints, and -homogeneous search are documented in [MIP runs](docs/mip_profiles.md). Configure -candidate evaluation, filtering, materialization, and distillation with -[post-MIP pipelines](docs/post_mip_pipeline.md). - -### Qwen 3.5 0.8B MIP smoke - -The focused Qwen 3.5 0.8B example pins the public checkpoint revision. Its -default model config follows the tracked 0.8B runtime campaign and searches -only the FFN intermediate sizes `[3072, 2048]`. The `mip_smoke.yaml` run -composes that config directly. It enables the composite scenario route required -by named-profile MIP while allowing only the teacher embedding width; depth, -attention, and GDN axes also remain at their teacher values. It is the first -runtime-validation target. - -The experimental `advanced.yaml` overlay remains an explicit follow-up. Its -broader axis structure is adapted from the Qwen 3.5 9B config and its target -values are derived from the pinned 0.8B geometry. Those advanced targets were -not selected from a completed 0.8B campaign and have not been fully -runtime-validated. In particular, its `gdn_key_head_dim` 128 to 96 target still -lacks physical-runtime equivalence evidence; that blocker does not apply to the -FFN-only MIP smoke. - -The checked-in one-GPU execution plan ends at `mip`. Bypass, vLLM serving -statistics, evaluation, AIPerf, and distillation are deliberately outside this -smoke boundary. CPU plan tests validate composition and scheduling only; the -opt-in GPU test must pass before treating the MIP smoke route as runtime-validated. -Review and replace every site placeholder in the runner before submission, -then inspect the complete plan without launching work: +manifests, controller state, and logs unless `runner.slurm.log_dir` relocates +attempt logs. -```bash -python examples/puzzletron/orchestrate.py \ - --experiment examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/mip_smoke.yaml \ - --runner examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml \ - --execution examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.smoke.yaml \ - --stage full --dry-run -``` +## Experiment overrides -The real-checkpoint acceptance test is an explicit manual gate and is not run -by generic GPU CI. From a reviewed worker-visible checkout and environment on -one H100 80GB GPU, with model access configured, run: +See [experiment overrides](docs/configuration_overrides.md) for temporary +changes without editing the checked-in YAML. -```bash -python -m pytest -v -s --run-manual \ - tests/gpu/torch/puzzletron/test_qwen3p5_0p8b_smoke.py -``` +Overrides can select another run root, adjust a campaign value, or change one +stage while preserving the source configuration. Validate the resolved config +before launch so misspelled or misplaced fields fail at the command boundary. -Retain the source revision, resolved environment or container, GPU model, -command, and complete pytest log as the result record. Treat the route as -runtime-validated only when the test passes and confirms a successful MIP -manifest, the exact `params-90` active profile, and at least one feasible MIP -scenario. The test runs the orchestrator locally with isolated temporary data -and cache roots; it does not submit scheduler work or consume the runner -placeholders above. +## Slurm configuration -## Run a campaign +See [Slurm configuration](docs/slurm_configuration.md) for partition lists, +CPU-only stages, log directories, and accepted compatibility fields. -The v2 orchestrator lives in -[`modelopt/torch/puzzletron/orchestration/`](../../modelopt/torch/puzzletron/orchestration/) -and is launched through -[`examples/puzzletron/orchestrate.py`](orchestrate.py). It separates: +Use the checked-in runner and execution examples as templates, replace their +site placeholders, and inspect the plan with `--dry-run` before launch. Runner +files own infrastructure; execution files own per-stage strategy and resource +selection. -- experiment semantics (`--experiment`, the existing Puzzletron YAML); -- runner infrastructure (`--runner`, Slurm or bare-metal inventory plus container/venv); -- execution semantics (`--execution`, per-stage strategy, `instances`, and optional mesh overrides). +## Qwen 3.5 smoke test -The repository-root -[`puzzletron_orchestrator`](../../puzzletron_orchestrator/) package is a -dependency-light facade over that canonical implementation. The CLI imports -through the facade instead of `modelopt.torch`, avoiding ModelOpt's eager -PyTorch initialization. It requires only Python 3.10+, PyYAML, and Rich. This -lets it run on a Slurm login node while GPU jobs use the full environment -declared by the runner: +See the [Qwen 3.5 0.8B smoke guide](docs/qwen3p5_0p8b_smoke.md) for the +one-GPU route, dry run, and manual GPU acceptance test. -```bash -python3 -m venv .venv-orchestrator -source .venv-orchestrator/bin/activate -python -m pip install -r examples/puzzletron/requirements-orchestrator.txt -``` +This focused campaign checks the MIP path on a small public checkpoint before +larger model or cluster runs. -The login-node environment must expose `sbatch`, `squeue`, and `sacct`. It does -not need PyTorch, Hydra, ModelOpt installation, CUDA, or the worker container. - -### Generated v2 campaign +## Run a campaign -Setup-v2-generated bundles encode evaluation, filtering, materialization, -AIPerf, and distillation in one campaign DAG. Run the generated smoke bundle -first to validate the environment and campaign wiring: +Activate the control environment and run the generated smoke bundle first. The +smoke run checks the worker environment and campaign wiring before the larger +production run: ```bash PUZZLETRON_BUNDLE=/path/to/generated/campaign/smoke @@ -483,241 +189,52 @@ python examples/puzzletron/orchestrate.py \ --stage full ``` -After the smoke campaign succeeds, replace `smoke` with `production` and run -the same command for the full campaign. Add `--dry-run` to inspect either plan -without submitting work, or select one stage while iterating, for example -`--stage mip --dry-run`. - -The setup wizard can also add downstream evaluation for materialized campaign -candidates. See [post-MIP pipelines](docs/post_mip_pipeline.md) to configure it -or add it to an existing campaign. - -### Legacy checked-in Nano campaign - -The checked-in Nano experiment uses the legacy `zero_shot_evaluation`, -`aiperf`, and global distillation stages. Its online-solution path still needs -two explicit preparation steps because the orchestrator runs and aggregates -the shards but does not create the online evaluation plan or materialize the -selected finalists. Use site-specific runner and execution configs, then run -the prerequisite DAG through MIP by temporarily disabling the downstream -stages: - -The repository already provides the bounded materializer, and setup emits the -command that invokes it. The Nano experiment pins its public model source and -revision but inherits a repository-relative `dataset_path` from `base.yaml`, so -a real run must override that value with a compatible, materialized Hugging Face -dataset directory visible at the same path on every worker. Dataset -materialization is outside the dependency-light setup and controller -environments. If you do not already have a compatible `datasets.save_to_disk` -directory, prepare Puzzle-KD from the full worker environment before starting -the controller: +After the smoke campaign succeeds, change `smoke` to `production` and run the +same command. Add `--dry-run` before either launch to inspect the plan without +submitting jobs. -```bash -export PUZZLETRON_DATASET=/shared/datasets/puzzle-kd-v2 +## Controller operations -python examples/puzzletron/materialize_dataset.py puzzle_kd_v2 \ - --output "$PUZZLETRON_DATASET" \ - --train-samples 8192 \ - --validation-samples 1024 \ - --seed 408 -``` +See [controller operations](docs/orchestration_operations.md) for individual +stages, non-interactive behavior, logging options, execution strategies, +recovery, and controller records. -Pass the same `dataset_path` override to every controller invocation so plan -compilation and worker commands resolve the same input: +The controller shows live progress and keeps resume state under +`${puzzle_dir}/orchestration/`. Press `q` or Ctrl-C in an interactive terminal +to cancel, detach, or continue. Run the same command again to recover a detached +campaign. -```bash -PUZZLETRON_EXPERIMENT=examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/default.yaml -PUZZLETRON_RUNNER=/path/to/runner.yaml -PUZZLETRON_EXECUTION=/path/to/execution.yaml -export PUZZLETRON_RUN_ROOT=/shared/puzzle_runs/my_campaign -export PUZZLETRON_DATASET=/shared/datasets/puzzle-kd-v2 +## MIP runs -python examples/puzzletron/orchestrate.py \ - --experiment "$PUZZLETRON_EXPERIMENT" \ - --runner "$PUZZLETRON_RUNNER" \ - --execution "$PUZZLETRON_EXECUTION" \ - --stage full \ - --override "dataset_path=$PUZZLETRON_DATASET" \ - --override zero_shot_evaluation.enabled=false \ - --override aiperf.enabled=false \ - --override global_distillation_sanity.enabled=false \ - --override global_distillation.enabled=false \ - --override post_distillation_evaluation.enabled=false -``` +See [MIP runs](docs/mip_profiles.md) for variants, solution pools, objectives, +resource constraints, workload measurements, and homogeneous search. -Prepare the online evaluation plan, then run and aggregate its shards. The -profile IDs below match the Nano example; for another legacy experiment using -this path, pass every entry from its `zero_shot_evaluation.profile_ids` list. +Named profiles let one campaign compare candidate architectures against +different parameter, runtime, or memory goals without duplicating the earlier +importance and scoring stages. -```bash -python examples/puzzletron/run_profile_online_evaluation.py \ - --puzzle-dir "$PUZZLETRON_RUN_ROOT" \ - --profile-id params-075 \ - --profile-id runtime-075 \ - --profile-id memory-075 \ - --profile-id params-075-num-experts-only \ - --profile-id params-075-expert-dim-only \ - --profile-id params-075-num-experts-and-expert-dim \ - --prepare +## Post-MIP pipelines -python examples/puzzletron/orchestrate.py \ - --experiment "$PUZZLETRON_EXPERIMENT" \ - --runner "$PUZZLETRON_RUNNER" \ - --execution "$PUZZLETRON_EXECUTION" \ - --stage zero_shot_evaluation \ - --override "dataset_path=$PUZZLETRON_DATASET" -``` +See [post-MIP pipelines](docs/post_mip_pipeline.md) for candidate evaluation, +filtering, materialization, AIPerf, and distillation. -Materialize the evaluated finalists for the Nano example's configured AIPerf -profile before running AIPerf and the remaining enabled stages. This helper -loads ModelOpt and Safetensors, so run it in the full worker environment from -the installation steps above, not in the dependency-light controller -environment. If the runner uses a container, enter it with the same mounts -before activating the worker venv. For another legacy experiment using this -path, use its `aiperf.profile_id` value. +These downstream nodes turn selected MIP solutions into evaluated or +materialized checkpoints and can continue through serving measurements and +global distillation. -```bash -# On the worker host or in the worker container: -cd /path/to/modelopt -source /path/to/full-modelopt-venv/bin/activate -export PUZZLETRON_RUN_ROOT=/shared/puzzle_runs/my_campaign - -python examples/puzzletron/prepare_online_profile_finalists.py \ - --puzzle-dir "$PUZZLETRON_RUN_ROOT" \ - --config examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/default.yaml \ - --profile-id runtime-075 \ - --count 1 -``` - -Return to the login node before launching the controller. The login node must -provide the scheduler commands listed above. - -```bash -cd /path/to/modelopt -source .venv-orchestrator/bin/activate -PUZZLETRON_EXPERIMENT=examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/default.yaml -PUZZLETRON_RUNNER=/path/to/runner.yaml -PUZZLETRON_EXECUTION=/path/to/execution.yaml -export PUZZLETRON_RUN_ROOT=/shared/puzzle_runs/my_campaign -export PUZZLETRON_DATASET=/shared/datasets/puzzle-kd-v2 - -python examples/puzzletron/orchestrate.py \ - --experiment "$PUZZLETRON_EXPERIMENT" \ - --runner "$PUZZLETRON_RUNNER" \ - --execution "$PUZZLETRON_EXECUTION" \ - --stage full \ - --override "dataset_path=$PUZZLETRON_DATASET" -``` - -For the legacy Nano experiment, the final `--stage full` resumes from the -verified completed stages. Do not use it as the initial command for legacy -configs with `mode: online_solutions`; setup-v2-generated bundles use their -dynamic `post.*` DAG instead. - -### Monitor and resume - -The launch command is a blocking foreground controller: it submits every -dependency-ready branch concurrently, polls scheduler state, and exits when the -selected plan completes or fails. Progress is colorized automatically on a TTY. -Interactive terminals show a live stage table with status, nodes/tasks/GPUs, -elapsed time, and best-effort ETA when a stage exposes current/total progress. -Completed stages remain visible; dependency waits, failed stages, and descendants -blocked by failures are labeled explicitly. Redirected output falls back to -timestamped one-line progress updates. -Press `q` or Ctrl-C in an interactive terminal to choose between cancelling all -active jobs and quitting, leaving jobs running and detaching the controller, or -resuming the campaign. Non-interactive Ctrl-C and SIGTERM retain the safe -cancel-and-quit behavior. A detached controller preserves durable handles, so the -same command recovers the running jobs. -Use `--color always` when piping through `tee`, `--color never` for plain logs, -and `--poll-interval SECONDS` to change the default five-second scheduler poll. - -Durable controller state is written under `${puzzle_dir}/orchestration/`. The -controller supports `single`, `sharded`, and `persistent_pool` strategies, -stdlib-first Slurm and SSH executors, attempt recovery, and semantic stage -validation through WorkAdapters. See -[`configs/orchestration/`](configs/orchestration/) for starter runner and -execution files. - -Accepted rank-zero stage results also write immutable, checksum-validated -execution records under `/manifests/executions/`. Puzzletron -validates these records when resuming a stage; the records identify existing -outputs but do not copy or make those outputs immutable. - -## Campaign stages - -Select one stage with the same v2 experiment, runner, and execution configs: - -```bash -PUZZLETRON_BUNDLE=/path/to/generated/campaign/production - -python examples/puzzletron/orchestrate.py \ - --experiment "$PUZZLETRON_BUNDLE/experiment.yaml" \ - --runner "$PUZZLETRON_BUNDLE/runner.yaml" \ - --execution "$PUZZLETRON_BUNDLE/execution.yaml" \ - --stage width_importance -``` - -The dependency-free [`StageSpec` registry](../../modelopt/torch/puzzletron/stages/graph.py) -is the authoritative contract for every public stage's identity, dependencies, -enablement, semantic config sections, and static completion artifacts. Add or -change those properties there. Default execution strategies remain -scheduler-specific and live in the -[orchestration compiler](../../modelopt/torch/puzzletron/orchestration/compiler.py); -handlers, scheduler adapters, mesh resolution, and heavyweight artifact -validators remain separate runtime concerns. See the -[v2 architecture](docs/v2_architecture.md) for orchestration internals and -maintainer guidance. - -| Stage | Purpose | -|---|---| -| `convert` | Convert the immutable Hugging Face teacher into the configured backend format. | -| `tokenize_data` | Build deterministic train and validation token caches. | -| `vllm_stats` | Measure exact runtime and memory costs for candidate subblocks. | -| `depth_importance` | Rank cumulative block or subblock removals. | -| `width_importance` | Collect activation-based rankings for every enabled width axis. | -| `sort` | Reorder the teacher so nested prefixes implement ranked width choices. | -| `sort_sanity` | Check that sorting preserves teacher outputs. | -| `width_sanity` | Compare ranked, original-order, and reverse slices on representative layers. | -| `slicing_sanity` | Verify dynamic slicing against physical materialization. | -| `bypass_sanity` | Overfit small local-distillation cases before production bypass. | -| `bypass` | Train nested replacement blocks across the configured search space. | -| `build_library` | Assemble sorted, bypassed, and no-op replacement candidates. | -| `replacement_scoring` | Score replacing one block or subblock at a time. | -| `mip` | Solve heterogeneous and homogeneous architecture searches under named constraints. | -| `zero_shot_evaluation` | Evaluate selected MIP recipes online without materializing every checkpoint. | -| `aiperf` | Materialize selected finalists and benchmark serving performance. | -| `global_distillation_sanity` | Overfit the selected global student as a correctness check. | -| `global_distillation` | Distill the selected architecture at the configured production scale. | -| `post_distillation_evaluation` | Evaluate the final distilled checkpoint. | - -## Interpret width sanity results - -Puzzletron separates implementation correctness from ranking quality: - -- Sort and slicing equivalence failures are correctness errors and always fail - their stages. -- A width-ranking miss means the activation-sorted candidate underperformed an - original or reverse control. It is a quality warning by default. - -To also fail a sanity stage on ranking-quality warnings, enable strict warning -handling: - -```yaml -sanity: - fail_on_warnings: true -``` +## Sanity validation -See [Sorting, width ranking, and slicing sanity](docs/sanity_validation.md) for -the slicing mental model, measured metrics, comparison controls, tolerances, -worked example, and qualification guidance. +See [sanity validation](docs/sanity_validation.md) for correctness checks, +ranking warnings, comparison controls, tolerances, and qualification guidance. -Independent DAG branches may run concurrently when they have disjoint writers. -Long-running stages should resume their durable checkpoints or immutable shards -rather than restarting completed work. +Sorting and slicing equivalence failures are correctness errors. Ranking +quality misses are warnings unless strict warning handling is enabled. ## Reports +See [campaign reports](docs/campaign_reports.md) for cache controls and the +evidence status of retained example reports. + After the selected plan completes cleanly, the v2 orchestrator generates the final campaign report through the configured runner. Reporting is nonfatal to the completed campaign, but a failed report attempt is recorded in the @@ -732,25 +249,22 @@ python examples/puzzletron/generate_campaign_progress_report.py \ --model-name 'My model' ``` -The output is -`/artifacts/campaign_report/campaign_report.html`. It is a -self-contained file suitable for sharing. Section source and configuration -fingerprints are cached under -`/artifacts/campaign_report/section_cache`, so unchanged sections -are reused and only affected sections rebuild. Use `--rebuild-section aiperf` -(repeatable) for selected sections, or `--no-cache` for an intentional full -rebuild. - -### Retained campaign reports - -The [campaign report catalog](docs/campaign_reports.md) records each retained -report's producer state, reproduction and support status, metadata origin, -relationship to current configuration files, and known limitations. An entry -marked as not reproduced is not a current model-support claim. Detailed run -facts remain in the reports. - -Each retained report is a self-contained HTML file that embeds sanity-check -outputs, stage manifests, and evaluation results; it can be hundreds of MB. -Download it to disk and open it locally rather than previewing it in a browser -tab. Interpret its evaluation results together with the reproduction status -and unresolved findings in the catalog. +Open +`/artifacts/campaign_report/campaign_report.html` locally. + +## Legacy Nano campaign + +See the [legacy Nano campaign](docs/legacy_nano_campaign.md) for the separate +online evaluation and finalist-materialization workflow used by the checked-in +Nano configuration. + +That configuration uses `mode: online_solutions`; it does not use the generated +campaign DAG's integrated evaluation and materialization nodes. + +## Architecture + +See the [v2 architecture](docs/v2_architecture.md) for the stage registry, +campaign DAG, scheduler-neutral control plane, and maintainer guidance. + +The experiment config owns model and algorithm semantics, the runner owns the +worker environment, and the execution config owns per-stage orchestration. diff --git a/examples/puzzletron/configs/orchestration/execution.example.yaml b/examples/puzzletron/configs/orchestration/execution.example.yaml index 2bfd3f675cb..413036096c9 100644 --- a/examples/puzzletron/configs/orchestration/execution.example.yaml +++ b/examples/puzzletron/configs/orchestration/execution.example.yaml @@ -10,12 +10,24 @@ execution: artifact_settling_timeout_seconds: 300 gpus_per_node: 8 stages: + # CPU-only stages request no GPUs and can override the runner partition. + # A list lets Slurm select any eligible CPU partition. convert: strategy: single instances: 1 + resource: cpu + partition: + - REPLACE_WITH_PRIMARY_CPU_SLURM_PARTITION + - REPLACE_WITH_ALTERNATE_CPU_SLURM_PARTITION tokenize_data: strategy: single instances: 1 + resource: cpu + partition: REPLACE_WITH_PRIMARY_CPU_SLURM_PARTITION + # Final report generation is always CPU-only and accepts a partition + # override without resource or strategy fields. + final_report: + partition: REPLACE_WITH_PRIMARY_CPU_SLURM_PARTITION vllm_stats: strategy: sharded instances: 16 diff --git a/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml b/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml index 84fd11dc8a9..e67c1e54fc9 100644 --- a/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml +++ b/examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml @@ -8,13 +8,8 @@ runner: slurm: account: REPLACE_WITH_SLURM_ACCOUNT partition: REPLACE_WITH_SLURM_PARTITION - partition_batch: REPLACE_WITH_SLURM_PARTITION - partition_interactive: - partition_cpu: - interactive_max_nodes: 1 max_nodes: 1 time_limit: "1:00:00" - log_dir: puzzle_runs/qwen3p5_0p8b_smoke/logs execution_contract: repository: REPLACE_WITH_WORKER_VISIBLE_MODELOPT_CHECKOUT venv: REPLACE_WITH_WORKER_VISIBLE_MODELOPT_VENV diff --git a/examples/puzzletron/configs/orchestration/qwen_moe/execution.production.yaml b/examples/puzzletron/configs/orchestration/qwen_moe/execution.production.yaml index 10d7c78a73e..a563b7d368a 100644 --- a/examples/puzzletron/configs/orchestration/qwen_moe/execution.production.yaml +++ b/examples/puzzletron/configs/orchestration/qwen_moe/execution.production.yaml @@ -9,16 +9,16 @@ execution: stages: convert: {strategy: single, instances: 1} tokenize_data: {strategy: single, instances: 1} - # One interactive node (8 GPUs): pack 8 vLLM shards onto that single node. - vllm_stats: {strategy: sharded, instances: 8, partition: interactive} - width_importance: {strategy: single, instances: 1, partition: interactive} + # One node (8 GPUs): pack 8 vLLM shards onto that single node. + vllm_stats: {strategy: sharded, instances: 8} + width_importance: {strategy: single, instances: 1} # One gang-scheduled allocation: four 8-GPU workers, coordinator on node 0. - depth_importance: {strategy: persistent_pool, instances: 4, partition: batch} - sort: {strategy: single, instances: 1, partition: interactive} - sort_sanity: {strategy: single, instances: 1, partition: interactive} - width_sanity: {strategy: single, instances: 1, partition: interactive} - slicing_sanity: {strategy: single, instances: 1, partition: interactive} - bypass_sanity: {strategy: single, instances: 1, partition: interactive} + depth_importance: {strategy: persistent_pool, instances: 4} + sort: {strategy: single, instances: 1} + sort_sanity: {strategy: single, instances: 1} + width_sanity: {strategy: single, instances: 1} + slicing_sanity: {strategy: single, instances: 1} + bypass_sanity: {strategy: single, instances: 1} bypass: {strategy: single, instances: 1} build_library: strategy: single @@ -32,7 +32,7 @@ execution: dp_shard: 1 dp_replicate: 1 # One gang-scheduled allocation: two 8-GPU workers, coordinator on node 0. - replacement_scoring: {strategy: persistent_pool, instances: 2, partition: interactive} + replacement_scoring: {strategy: persistent_pool, instances: 2} mip: {strategy: single, instances: 1} zero_shot_evaluation: {strategy: sharded, instances: 8} aiperf: {strategy: sharded, instances: 8} diff --git a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml b/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml index f1b4e8817c6..b3b2489c682 100644 --- a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml +++ b/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml @@ -6,16 +6,10 @@ runner: slurm: # Required placeholder. Replace with the Slurm account for your site. account: REPLACE_WITH_SLURM_ACCOUNT - # Replace these generic partition names if your site uses different names. - partition: batch - partition_interactive: interactive - partition_batch: batch - # Optional. CPU/IO stages use the regular one-node partition when unset. - partition_cpu: - interactive_max_nodes: 2 + # Replace with one partition name or a list of eligible names for this site. + partition: REPLACE_WITH_SLURM_PARTITION max_nodes: 20 time_limit: "4:00:00" - log_dir: puzzle_runs/qwen-moe/logs execution_contract: # Replace with the ModelOpt checkout path visible on every worker and in the container. repository: REPLACE_WITH_WORKER_VISIBLE_MODELOPT_CHECKOUT diff --git a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml b/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml index 3ca048d9393..3b1a47fd6c3 100644 --- a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml +++ b/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml @@ -8,13 +8,11 @@ runner: slurm: # Required placeholder. Replace with the Slurm account for your site. account: REPLACE_WITH_SLURM_ACCOUNT - # Replace these generic partition names if your site uses different names. - partition: batch - partition_interactive: interactive - partition_batch: batch - # Optional. CPU/IO stages use the regular one-node partition when unset. - partition_cpu: - interactive_max_nodes: 2 + # Optional default for stages without a partition override. Use one name or + # a list of eligible names. Slurm uses the site default when this is omitted. + partition: + - REPLACE_WITH_PRIMARY_SLURM_PARTITION + - REPLACE_WITH_ALTERNATE_SLURM_PARTITION time_limit: "4:00:00" log_dir: puzzle_runs/logs execution_contract: diff --git a/examples/puzzletron/configs/setup/defaults.example.yaml b/examples/puzzletron/configs/setup/defaults.example.yaml index ade0d88bd5a..cb42c45625e 100644 --- a/examples/puzzletron/configs/setup/defaults.example.yaml +++ b/examples/puzzletron/configs/setup/defaults.example.yaml @@ -19,5 +19,6 @@ infrastructure: slurm: # Add the required Slurm account for your site before using this file. # account: REPLACE_WITH_SLURM_ACCOUNT - # Optional. CPU/IO stages use the regular partition when unset. - partition_cpu: + # Optional. Use one partition name or a list of eligible names. When + # omitted, Slurm uses the site's default partition. + partition: diff --git a/examples/puzzletron/docs/campaign_reports.md b/examples/puzzletron/docs/campaign_reports.md index 12b1272a224..95a7eeb2f54 100644 --- a/examples/puzzletron/docs/campaign_reports.md +++ b/examples/puzzletron/docs/campaign_reports.md @@ -1,11 +1,31 @@ # Puzzletron Campaign Reports -This page catalogs retained Puzzletron campaign reports and the status of their -evidence. The compact [campaign report index](../reports/campaign_report_index.yaml) +The orchestrator generates a cumulative HTML report after a campaign. Regenerate +it without rerunning model work: + +```bash +python examples/puzzletron/generate_campaign_progress_report.py \ + --puzzle-dir /shared/puzzle_runs/my_campaign \ + --model-name 'My model' +``` + +The output is +`/artifacts/campaign_report/campaign_report.html`. Section inputs +and configuration fingerprints are cached under +`/artifacts/campaign_report/section_cache`. Use +`--rebuild-section aiperf` to rebuild one section, or `--no-cache` to rebuild +the whole report. + +This page also catalogs retained Puzzletron campaign reports and the status of +their evidence. The compact [campaign report index](../reports/campaign_report_index.yaml) records each report's producer state, reproduction and support status, metadata origin, current-configuration relationship, and known limitations. Detailed run facts remain in the reports. +Retained reports are self-contained HTML files and may be hundreds of MB. +Download them and open them locally. Interpret their results together with the +reproduction status and unresolved findings below. + ## Report status | Model | Report | Producer state | Reproduction | Support | Current configuration relationship | diff --git a/examples/puzzletron/docs/configuration_overrides.md b/examples/puzzletron/docs/configuration_overrides.md new file mode 100644 index 00000000000..35eea2dd7ae --- /dev/null +++ b/examples/puzzletron/docs/configuration_overrides.md @@ -0,0 +1,20 @@ +# Experiment overrides + +Use command-line overrides for temporary experiment value changes. Append a +repeatable `--override KEY=VALUE` to the orchestrator command and inspect the +result with `--dry-run` before launch: + +```bash +--override mip.runs.params-90.solver.num_solutions=4 \ +--override ++runtime_annotations.reason=capacity-check \ +--dry-run +``` + +Plain `KEY=VALUE` and explicit `++KEY=VALUE` both add or replace experiment +values. The controller and GPU workers interpret these forms identically. +Single-plus add (`+KEY=VALUE`) and delete (`~KEY`) operators are not +supported. Put structural changes in a copied run config so they remain easy to +review. + +Overrides apply only to the experiment config. Edit or copy the runner and +execution files when changing site or scheduler settings. diff --git a/examples/puzzletron/docs/environment_setup.md b/examples/puzzletron/docs/environment_setup.md new file mode 100644 index 00000000000..9eaa8e137c1 --- /dev/null +++ b/examples/puzzletron/docs/environment_setup.md @@ -0,0 +1,226 @@ +# Environment setup + +Puzzletron uses two environments: + +- a lightweight control environment for the setup wizard and orchestrator; +- a GPU worker environment for ModelOpt, the patched vLLM fork, AutoModel, and + AIPerf. + +The runner file connects them. `runner.execution_contract.venv` selects the +worker virtual environment, and `runner.execution_contract.container` selects +an optional Slurm container. + +## Control environment + +The setup wizard and orchestrator do not import PyTorch or initialize CUDA. +Create one environment for both: + +```bash +python3 -m venv .venv-puzzletron-control +source .venv-puzzletron-control/bin/activate +python -m pip install \ + -r examples/puzzletron/requirements-setup.txt \ + -r examples/puzzletron/requirements-orchestrator.txt +``` + +A Slurm login node also needs `sbatch`, `squeue`, and `sacct`. It does not need +ModelOpt, CUDA, the worker container, or the worker virtual environment. + +## Worker environment + +Use one Python environment for ModelOpt, patched vLLM, AutoModel, and official +AIPerf. Install PyTorch first and build every CUDA extension against that +installation. Mixing PyTorch or CUDA builds can cause import failures or +incorrect GPU execution. + +### Choose a container or host environment + +This CUDA image provides a reproducible bootstrap: + +```text +nvcr.io/nvidia/cuda:12.9.2-cudnn-devel-ubuntu24.04 +``` + +The image is an example, not a required runner image. Slurm campaigns can set +`runner.execution_contract.container` to an image or path accepted by the +site, or omit it to execute directly in the worker environment. Bare-metal +runners use the host environment selected by `runner.execution_contract.venv`. + +The commands below assume a container. For bare metal, skip the Docker, +`/workspace`, and `apt-get` steps. Install equivalent Python and build tools +through the host-environment tooling and adapt the paths. + +```bash +export PUZZLETRON_WORKSPACE=/absolute/path/to/workspace +docker run --gpus all --ipc=host --rm -it \ + -v "${PUZZLETRON_WORKSPACE}:/workspace" \ + -w /workspace \ + nvcr.io/nvidia/cuda:12.9.2-cudnn-devel-ubuntu24.04 bash +``` + +Inside the container, install Python and the build tools used by editable +packages and optional CUDA extensions: + +```bash +apt-get update +DEBIAN_FRONTEND=noninteractive apt-get install -y \ + build-essential cmake git ninja-build \ + python3 python3-dev python3-pip python3-venv +``` + +### Clone the tracked forks + +Keep ModelOpt and the two Puzzletron forks as siblings. The machine-readable +[CI environment](../ci_environment.json) records the shared compatibility pins. + +```bash +export MODEL_OPT_ROOT=/workspace/modelopt +export VLLM_ROOT=/workspace/vllm +export AUTOMODEL_ROOT=/workspace/Automodel +export PUZZLETRON_CI_ENVIRONMENT="${MODEL_OPT_ROOT}/examples/puzzletron/ci_environment.json" +export AUTOMODEL_REF="$(python3 -c \ + 'import json, sys; print(json.load(open(sys.argv[1]))["nemo_automodel"]["commit"])' \ + "${PUZZLETRON_CI_ENVIRONMENT}")" + +git clone --branch feature/add_anymodel_to_vllm --single-branch \ + https://github.com/Separius/vllm.git "${VLLM_ROOT}" +git clone --branch puzzletron --single-branch \ + https://github.com/Separius/Automodel.git "${AUTOMODEL_ROOT}" +git -C "${AUTOMODEL_ROOT}" checkout --detach "${AUTOMODEL_REF}" +``` + +```text +/workspace/ +├── modelopt/ +├── vllm/ +└── Automodel/ +``` + +### Install runtime packages + +The patched vLLM branch uses the PyTorch version recorded in the CI environment +with CUDA 12.9. Install that combination before compiling CUDA code: + +```bash +python3 -m venv /workspace/.venv +source /workspace/.venv/bin/activate + +export PUZZLETRON_TORCH_VERSION="$(python -c \ + 'import json, sys; print(json.load(open(sys.argv[1]))["torch"])' \ + "${PUZZLETRON_CI_ENVIRONMENT}")" +export PUZZLETRON_TORCHVISION_VERSION="$(python -c \ + 'import json, sys; print(json.load(open(sys.argv[1]))["torchvision"])' \ + "${PUZZLETRON_CI_ENVIRONMENT}")" +export PUZZLETRON_TRANSFORMERS_VERSION="$(python -c \ + 'import json, sys; print(json.load(open(sys.argv[1]))["transformers"])' \ + "${PUZZLETRON_CI_ENVIRONMENT}")" + +python -m pip install --upgrade \ + pip "setuptools>=80,<81" "setuptools-scm>=8" setuptools-rust \ + wheel "packaging>=24.2" "cmake>=3.26.1" ninja jinja2 + +python -m pip install \ + "torch==${PUZZLETRON_TORCH_VERSION}" \ + "torchvision==${PUZZLETRON_TORCHVISION_VERSION}" \ + "torchaudio==${PUZZLETRON_TORCH_VERSION}" \ + --index-url https://download.pytorch.org/whl/cu129 + +VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 \ + python -m pip install --no-build-isolation -e "${VLLM_ROOT}" + +python -m pip install -e "${AUTOMODEL_ROOT}" +python -m pip install aiperf +python -m pip install -e "${MODEL_OPT_ROOT}[hf,puzzletron]" +python -m pip install "transformers==${PUZZLETRON_TRANSFORMERS_VERSION}" +python -m pip install -r "${MODEL_OPT_ROOT}/examples/puzzletron/requirements.txt" +``` + +Do not add `--no-deps`; these packages need their declared Python dependencies. +`--no-build-isolation` makes compiled extensions use the active PyTorch +installation. It does not disable dependency installation. + +Install only the kernels required by the target architecture: + +```bash +# Mixture of experts +python -m pip install --no-build-isolation \ + "git+https://github.com/fanshiqing/grouped_gemm@v1.1.4" + +# Mamba +python -m pip install "mamba-ssm[causal-conv1d]" --no-build-isolation + +# Linear attention +python -m pip install "flash-linear-attention[cuda]" +``` + +## Verify the worker environment + +Run these checks inside the same container and virtual environment used by +Puzzletron jobs: + +```bash +test "$(git -C "${VLLM_ROOT}" remote get-url origin)" = \ + "https://github.com/Separius/vllm.git" +test "$(git -C "${VLLM_ROOT}" branch --show-current)" = \ + "feature/add_anymodel_to_vllm" +test "$(git -C "${AUTOMODEL_ROOT}" remote get-url origin)" = \ + "https://github.com/Separius/Automodel.git" +test "$(git -C "${AUTOMODEL_ROOT}" rev-parse HEAD)" = "${AUTOMODEL_REF}" + +git -C "${MODEL_OPT_ROOT}" rev-parse HEAD +git -C "${VLLM_ROOT}" rev-parse HEAD +git -C "${AUTOMODEL_ROOT}" rev-parse HEAD +``` + +```bash +python - <<'PY' +import importlib.metadata as metadata +import json +import os + +from packaging.version import Version + +import aiperf +import lmms_eval +import modelopt +import nemo_automodel +import torch +import transformers +import vllm + +with open(os.environ["PUZZLETRON_CI_ENVIRONMENT"], encoding="utf-8") as stream: + ci_environment = json.load(stream) + +for package in ( + "torch", + "vllm", + "nemo-automodel", + "aiperf", + "lmms-eval", + "nvidia-modelopt", +): + print(package, metadata.version(package)) + +print("torch CUDA", torch.version.cuda) +print("CUDA available", torch.cuda.is_available()) +print("modelopt", modelopt.__file__) +print("vllm", vllm.__file__) + +assert Version(torch.__version__).release == Version(ci_environment["torch"]).release +assert Version(metadata.version("torchvision")).release == Version( + ci_environment["torchvision"] +).release +assert transformers.__version__ == ci_environment["transformers"] +assert metadata.version("lmms-eval") == ci_environment["lmms_eval"] +assert Version(metadata.version("nemo-automodel")).base_version == ( + ci_environment["nemo_automodel"]["base_version"] +) +assert torch.version.cuda == "12.9" +assert torch.cuda.is_available() +PY + +python -m pip check +``` + +Record the three source revisions and verification output with the campaign. +Repeat verification after pulling either fork or rebuilding a CUDA extension. diff --git a/examples/puzzletron/docs/legacy_nano_campaign.md b/examples/puzzletron/docs/legacy_nano_campaign.md new file mode 100644 index 00000000000..6b4927bb77a --- /dev/null +++ b/examples/puzzletron/docs/legacy_nano_campaign.md @@ -0,0 +1,117 @@ +# Legacy Nano campaign + +The checked-in Nano experiment uses `mode: online_solutions` with the +`zero_shot_evaluation`, `aiperf`, and global distillation stages. This config +requires explicit commands to create the evaluation plan and materialize +selected finalists. Setup v2 represents those operations as campaign-DAG +nodes; its recommended flow includes evaluation and materialization. + +## Prepare the dataset and MIP results + +Use site-specific runner and execution configs. The Nano experiment pins its +public model source and revision but inherits a repository-relative +`dataset_path` from `base.yaml`. Override it with a materialized Hugging Face +dataset directory that is visible at the same path on every worker. + +If needed, prepare Puzzle-KD from the full worker environment before starting +the controller: + +```bash +export PUZZLETRON_DATASET=/shared/datasets/puzzle-kd-v2 + +python examples/puzzletron/materialize_dataset.py puzzle_kd_v2 \ + --output "$PUZZLETRON_DATASET" \ + --train-samples 8192 \ + --validation-samples 1024 \ + --seed 408 +``` + +Pass the same dataset override to every controller invocation. First run the +campaign through MIP with the downstream stages disabled: + +```bash +PUZZLETRON_EXPERIMENT=examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/default.yaml +PUZZLETRON_RUNNER=/path/to/runner.yaml +PUZZLETRON_EXECUTION=/path/to/execution.yaml +export PUZZLETRON_RUN_ROOT=/shared/puzzle_runs/my_campaign +export PUZZLETRON_DATASET=/shared/datasets/puzzle-kd-v2 + +python examples/puzzletron/orchestrate.py \ + --experiment "$PUZZLETRON_EXPERIMENT" \ + --runner "$PUZZLETRON_RUNNER" \ + --execution "$PUZZLETRON_EXECUTION" \ + --stage full \ + --override "dataset_path=$PUZZLETRON_DATASET" \ + --override zero_shot_evaluation.enabled=false \ + --override aiperf.enabled=false \ + --override global_distillation_sanity.enabled=false \ + --override global_distillation.enabled=false \ + --override post_distillation_evaluation.enabled=false +``` + +## Evaluate candidate profiles + +Prepare the online evaluation plan, then run and aggregate its shards. These +profile IDs match the Nano example. For another legacy experiment, pass every +entry from its `zero_shot_evaluation.profile_ids` list. + +```bash +python examples/puzzletron/run_profile_online_evaluation.py \ + --puzzle-dir "$PUZZLETRON_RUN_ROOT" \ + --profile-id params-075 \ + --profile-id runtime-075 \ + --profile-id memory-075 \ + --profile-id params-075-num-experts-only \ + --profile-id params-075-expert-dim-only \ + --profile-id params-075-num-experts-and-expert-dim \ + --prepare + +python examples/puzzletron/orchestrate.py \ + --experiment "$PUZZLETRON_EXPERIMENT" \ + --runner "$PUZZLETRON_RUNNER" \ + --execution "$PUZZLETRON_EXECUTION" \ + --stage zero_shot_evaluation \ + --override "dataset_path=$PUZZLETRON_DATASET" +``` + +## Materialize finalists and resume + +Materialize the evaluated finalists for the configured AIPerf profile from the +full worker environment. If the runner uses a container, enter it with the same +mounts before activating the worker virtual environment. For another legacy +experiment, use its `aiperf.profile_id` value. + +```bash +# On the worker host or in the worker container: +cd /path/to/modelopt +source /path/to/full-modelopt-venv/bin/activate +export PUZZLETRON_RUN_ROOT=/shared/puzzle_runs/my_campaign + +python examples/puzzletron/prepare_online_profile_finalists.py \ + --puzzle-dir "$PUZZLETRON_RUN_ROOT" \ + --config examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/default.yaml \ + --profile-id runtime-075 \ + --count 1 +``` + +Return to the login node and resume the remaining enabled stages: + +```bash +cd /path/to/modelopt +source .venv-puzzletron-control/bin/activate +PUZZLETRON_EXPERIMENT=examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/default.yaml +PUZZLETRON_RUNNER=/path/to/runner.yaml +PUZZLETRON_EXECUTION=/path/to/execution.yaml +export PUZZLETRON_RUN_ROOT=/shared/puzzle_runs/my_campaign +export PUZZLETRON_DATASET=/shared/datasets/puzzle-kd-v2 + +python examples/puzzletron/orchestrate.py \ + --experiment "$PUZZLETRON_EXPERIMENT" \ + --runner "$PUZZLETRON_RUNNER" \ + --execution "$PUZZLETRON_EXECUTION" \ + --stage full \ + --override "dataset_path=$PUZZLETRON_DATASET" +``` + +The final command resumes from verified completed stages. Do not use it as the +initial command for legacy configs with `mode: online_solutions`. diff --git a/examples/puzzletron/docs/orchestration_operations.md b/examples/puzzletron/docs/orchestration_operations.md new file mode 100644 index 00000000000..a982cf02a6c --- /dev/null +++ b/examples/puzzletron/docs/orchestration_operations.md @@ -0,0 +1,46 @@ +# Controller operation and recovery + +Run one stage with the same experiment, runner, and execution files used for a +full campaign: + +```bash +PUZZLETRON_BUNDLE=/path/to/generated/campaign/production + +python examples/puzzletron/orchestrate.py \ + --experiment "$PUZZLETRON_BUNDLE/experiment.yaml" \ + --runner "$PUZZLETRON_BUNDLE/runner.yaml" \ + --execution "$PUZZLETRON_BUNDLE/execution.yaml" \ + --stage width_importance +``` + +The launch command runs a foreground controller. It submits every +dependency-ready branch concurrently, polls scheduler state, and exits when the +selected plan completes or fails. + +## Progress and interruption + +Interactive terminals show a live stage table with status, resources, elapsed +time, and a best-effort ETA when a stage reports progress. Completed stages, +dependency waits, failures, and descendants blocked by failures remain visible. +Redirected output uses timestamped one-line updates instead. + +Press `q` or Ctrl-C in an interactive terminal to cancel active jobs and quit, +detach while leaving jobs running, or continue. Non-interactive Ctrl-C and +SIGTERM cancel active work and quit. A detached controller preserves durable +handles, so running the same command recovers the active jobs. + +Use `--color always` when piping through `tee`, `--color never` for plain logs, +and `--poll-interval SECONDS` to change the default five-second poll interval. + +## State and execution records + +Durable controller state is written under `${puzzle_dir}/orchestration/`. The +controller supports `single`, `sharded`, and `persistent_pool` strategies, +Slurm and SSH executors, attempt recovery, and semantic stage validation. See +the [`configs/orchestration/`](../configs/orchestration/) directory for starter +runner and execution files. + +Accepted rank-zero stage results also write checksum-validated execution +records under `/manifests/executions/`. Puzzletron validates these +records when resuming a stage. They identify existing outputs but do not copy +or make those outputs immutable. diff --git a/examples/puzzletron/docs/qwen3p5_0p8b_smoke.md b/examples/puzzletron/docs/qwen3p5_0p8b_smoke.md new file mode 100644 index 00000000000..2cb30263ac0 --- /dev/null +++ b/examples/puzzletron/docs/qwen3p5_0p8b_smoke.md @@ -0,0 +1,44 @@ +# Qwen 3.5 0.8B MIP smoke + +The focused Qwen 3.5 0.8B example pins the public checkpoint revision. Its +default model config searches only the FFN intermediate sizes `[3072, 2048]`. +The `mip_smoke.yaml` run enables the composite scenario route required by +named-profile MIP while keeping depth, attention, GDN, and embedding width at +their teacher values. + +The experimental `advanced.yaml` overlay covers more axes. Its target values +were derived from the pinned 0.8B geometry rather than selected by a completed +campaign, and the overlay has not been fully runtime-validated. In particular, +the `gdn_key_head_dim` reduction from 128 to 96 does not yet have physical +runtime equivalence evidence. This does not affect the FFN-only smoke route. + +## Inspect the plan + +The checked-in one-GPU execution plan ends at `mip`. It does not run bypass, +vLLM serving statistics, evaluation, AIPerf, or distillation. Replace every +site placeholder in the runner, then inspect the plan without launching work: + +```bash +python examples/puzzletron/orchestrate.py \ + --experiment examples/puzzletron/configs/families/qwen3_5/qwen3p5_0p8b/runs/mip_smoke.yaml \ + --runner examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml \ + --execution examples/puzzletron/configs/orchestration/qwen3p5_0p8b/execution.smoke.yaml \ + --stage full --dry-run +``` + +## Run the GPU acceptance test + +CPU plan tests cover composition and scheduling only. The real-checkpoint test +is a manual gate and is not part of generic GPU CI. Run it from a reviewed, +worker-visible checkout on one H100 80GB GPU with model access configured: + +```bash +python -m pytest -v -s --run-manual \ + tests/gpu/torch/puzzletron/test_qwen3p5_0p8b_smoke.py +``` + +Treat the route as runtime-validated only when the test passes and confirms a +successful MIP manifest, the `params-90` active profile, and at least one +feasible MIP scenario. Retain the source revision, environment or container, +GPU model, command, and complete pytest log with the result. The test uses +isolated temporary data and cache roots and does not submit scheduler work. diff --git a/examples/puzzletron/docs/setup_wizard.md b/examples/puzzletron/docs/setup_wizard.md new file mode 100644 index 00000000000..8d7999c9536 --- /dev/null +++ b/examples/puzzletron/docs/setup_wizard.md @@ -0,0 +1,74 @@ +# Setup wizard + +The Puzzletron setup wizard inspects a local checkpoint configuration or a +Hugging Face model configuration and generates self-contained smoke and +production bundles. It reads configuration metadata, not model weights, and +does not submit jobs. + +## Profiles + +The guided flow offers three profiles: + +- **Quick smoke** creates the smallest campaign for checking campaign shape. +- **Balanced pruning** provides the recommended defaults for a first campaign. +- **High-confidence search** spends more runtime on scoring and sanity checks. + +The selected profile supplies pruning and search defaults from the detected +model family's `setup_v2_defaults.yaml`. The wizard then asks for the model, +dataset, worker environment, and cluster settings. + +## Models and datasets + +At the **Model** prompt, provide an existing local checkpoint or configuration +path, or a Hugging Face model URL or repository ID. + +At the **Dataset** prompt, provide an existing local dataset path, a Hugging +Face dataset URL, or a repository ID. For a hosted dataset, setup records a +worker-visible output path. The generated campaign `README.md` contains the +exact acquisition command. Run that command from the worker environment before +launching the campaign. A local dataset is referenced directly. + +## Defaults and advanced mode + +Start the guided flow with the example defaults: + +```bash +python examples/puzzletron/puzzletron_setup_v2.py \ + --defaults examples/puzzletron/configs/setup/defaults.example.yaml +``` + +The example uses repository-relative values. Copy it and add site-specific +data, scheduler, and container settings before selecting it. The defaults file +is loaded only when passed explicitly and takes precedence over the selected +profile. + +Use the full flow to expose every section and nested setting: + +```bash +python examples/puzzletron/puzzletron_setup_v2.py --full +``` + +## Navigation and resume + +Press **Esc** to return from any prompt. Selection prompts include a visible +**← Back** action, and text or numeric prompts accept `:back`. + +The wizard saves accepted answers and navigation state in `answers_v2.yaml`. +Resume an interrupted setup with: + +```bash +python examples/puzzletron/puzzletron_setup_v2.py --resume /path/to/campaign +``` + +## Generated files + +The final review writes `resolved_defaults.yaml`, one campaign `README.md`, and +validated `smoke/` and `production/` bundles. Each bundle contains experiment, +runner, and execution YAML plus a `dry-run-plan.txt`. The wizard does not submit +either bundle, and the production bundle is not automatically gated on smoke. + +The generated configuration can include reusable execution profiles, multiple +deployment measurements, independent optimization goals, and editable +downstream flows. See [experiment overrides](configuration_overrides.md), +[Slurm configuration](slurm_configuration.md), and +[post-MIP pipelines](post_mip_pipeline.md) for those controls. diff --git a/examples/puzzletron/docs/slurm_configuration.md b/examples/puzzletron/docs/slurm_configuration.md new file mode 100644 index 00000000000..278a6965f66 --- /dev/null +++ b/examples/puzzletron/docs/slurm_configuration.md @@ -0,0 +1,67 @@ +# Slurm configuration + +Use the runner file for site-wide Slurm settings and the execution file for +stage-specific choices. + +## Partitions and logs + +`runner.slurm.partition` sets the default for stages without a partition +override. It accepts one partition name or a list of eligible names. Omit it to +use the site's Slurm default. A stage can set +`execution.stages..partition` to one name or its own eligible list. + +`runner.slurm.log_dir` sets the directory used for every attempt log, including +the final-report attempt. When omitted, logs are written below +`/logs`. + +The runner loader accepts `partition_interactive`, `partition_batch`, +`partition_cpu`, and `interactive_max_nodes` as compatibility fields. They +infer stage routing from role names and node count, which assumes a particular +site layout and duplicates execution-stage settings. Maintained configs and +examples use `runner.slurm.partition` with stage overrides instead. + +The production examples also avoid literal `interactive` and `batch` stage +overrides because those partition names are not portable between Slurm sites. + +## CPU-only stages + +Set `resource: cpu` with a CPU partition override for work that does not need a +GPU. Other stages continue to use the runner default: + +```yaml +runner: + kind: slurm + slurm: + partition: + - gpu-general + - gpu-overflow + +execution: + stages: + convert: + strategy: single + resource: cpu + partition: + - cpu-general + - cpu-overflow + width_importance: + strategy: single +``` + +Slurm selects one partition from each eligible list. The +[`runner.slurm.example.yaml`](../configs/orchestration/runner.slurm.example.yaml) +and [`execution.example.yaml`](../configs/orchestration/execution.example.yaml) +files show the runner default and per-stage CPU routing together. The CPU-only +`final_report` task accepts only a `partition` override. + +## Scheduler settings and model settings + +Do not put `sequence_parallel` under +`execution.stages..parallel`. That mapping controls scheduler allocation +and accepts mesh dimensions such as `tp`, `pp`, and `dp_replicate`. +`sequence_parallel` changes model execution and belongs in the experiment's +model-parallel profile. Setup-generated execution files omit it for this +reason. + +Runner and execution files reject unknown fields and suggest the closest valid +name when possible. diff --git a/examples/puzzletron/orchestrate.py b/examples/puzzletron/orchestrate.py index c81d551c9e0..4ceee633714 100644 --- a/examples/puzzletron/orchestrate.py +++ b/examples/puzzletron/orchestrate.py @@ -1,6 +1,18 @@ #!/usr/bin/env python3 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """CLI entrypoint for the Puzzletron v2 campaign orchestrator.""" @@ -11,6 +23,8 @@ import sys from pathlib import Path +import yaml + REPOSITORY_ROOT = Path(__file__).resolve().parents[2] if str(REPOSITORY_ROOT) not in sys.path: sys.path.insert(0, str(REPOSITORY_ROOT)) @@ -29,9 +43,19 @@ def _build_parser() -> argparse.ArgumentParser: parser.add_argument("--experiment", required=True, help="Path to the experiment YAML.") parser.add_argument("--runner", required=True, help="Path to the runner environment YAML.") parser.add_argument("--execution", required=True, help="Path to the execution semantics YAML.") - parser.add_argument("--stage", default="full", help="Stage id or 'full' for all enabled stages.") - parser.add_argument("--override", action="append", default=[], help="Hydra-style config override.") - parser.add_argument("--dry-run", action="store_true", help="Print packed submissions without submitting.") + parser.add_argument( + "--stage", default="full", help="Stage id or 'full' for all enabled stages." + ) + parser.add_argument( + "--override", + action="append", + default=[], + metavar="KEY=VALUE", + help="Repeatable config override; KEY=VALUE and ++KEY=VALUE are supported.", + ) + parser.add_argument( + "--dry-run", action="store_true", help="Print packed submissions without submitting." + ) parser.add_argument("--local", action="store_true", help="Use the local subprocess executor.") parser.add_argument("--once", action="store_true", help="Run one controller iteration.") parser.add_argument("--max-iterations", type=int, default=None) @@ -53,21 +77,25 @@ def _build_parser() -> argparse.ArgumentParser: def main(argv: list[str] | None = None) -> int: args = _build_parser().parse_args(argv) logger = OrchestratorLogger(color=args.color) - runner = load_runner_config(args.runner) - execution = load_execution_config(args.execution) - plan = compile_campaign_plan( - experiment_config_path=args.experiment, - runner=runner, - execution=execution, - overrides=args.override, - stage_filter=args.stage, - ) + try: + runner = load_runner_config(args.runner) + execution = load_execution_config(args.execution) + plan = compile_campaign_plan( + experiment_config_path=args.experiment, + runner=runner, + execution=execution, + overrides=args.override, + stage_filter=args.stage, + ) + submissions = dry_run_plan(plan, overrides=args.override) if args.dry_run else None + except (KeyError, OSError, TypeError, ValueError, yaml.YAMLError) as error: + logger.error(f"cannot build campaign plan: {error}") + return 2 if args.dry_run: - submissions = dry_run_plan(plan, overrides=args.override) + assert submissions is not None logger.banner("dry-run only; no jobs will be submitted") logger.plan( - f"{len(plan.stages)} stage(s), {len(submissions)} submission(s), " - f"root={plan.puzzle_dir}" + f"{len(plan.stages)} stage(s), {len(submissions)} submission(s), root={plan.puzzle_dir}" ) for node in plan.stages: count = sum(item.stage_id == node.stage_id for item in submissions) diff --git a/modelopt/torch/puzzletron/_config_aliases.py b/modelopt/torch/puzzletron/_config_aliases.py new file mode 100644 index 00000000000..03f4cf0a7b6 --- /dev/null +++ b/modelopt/torch/puzzletron/_config_aliases.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared compatibility checks for Puzzletron experiment configuration.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, Mapping + +__all__ = [] + +_COMPATIBILITY_ALIASES = ( + ("puzzle_dir", "experiment.dir", "puzzle_dir", True), + ("input_hf_model_path", "model.source", "input_hf_model_path", False), + ("teacher_dir", "convert.teacher_dir", "teacher_dir", True), + ("dataset_path", "data.path", "dataset_path", True), + ("trust_remote_code", "model.trust_remote_code", "model.trust_remote_code", False), +) + + +def _lookup(config: Mapping[str, Any], dotted_path: str) -> Any: + value: Any = config + for key in dotted_path.split("."): + if not isinstance(value, Mapping) or key not in value: + raise KeyError(dotted_path) + value = value[key] + return value + + +def _compatible_values(left: Any, right: Any, *, path_like: bool) -> bool: + if path_like: + return ( + Path(os.path.normpath(str(left))).expanduser() + == Path(os.path.normpath(str(right))).expanduser() + ) + return left == right + + +def _validate_compatibility_aliases(config: Mapping[str, Any]) -> None: + for legacy_path, canonical_path, preferred_override, path_like in _COMPATIBILITY_ALIASES: + try: + legacy_value = _lookup(config, legacy_path) + canonical_value = _lookup(config, canonical_path) + except KeyError: + continue + if not _compatible_values(legacy_value, canonical_value, path_like=path_like): + raise ValueError( + f"Experiment aliases {legacy_path!r} and {canonical_path!r} disagree; " + f"keep them identical or override {preferred_override!r} so composed " + "references stay synchronized" + ) diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index 1502a503a21..23675deb090 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/pool.py @@ -258,7 +258,7 @@ def command( ) -> AttemptSpec: repo = Path(runner.contract.repository) role = item.metadata.get("role", "worker") - log_dir = plan.puzzle_dir / "logs" + log_dir = plan.log_dir replacement_puzzle_dir = ( _replacement_puzzle_dir(plan, item.metadata.get("width")) if node.stage_id == "replacement_scoring" diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index 4d05c519f3e..88295e9d4b7 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -244,7 +244,7 @@ def command( ) for override in overrides or (): argv.extend(["--override", override]) - log_path = plan.puzzle_dir / "logs" / f"{node.stage_id}_{item.shard_index}_{attempt_id}.log" + log_path = plan.log_dir / f"{node.stage_id}_{item.shard_index}_{attempt_id}.log" # evaluation/global_kd always call torch.distributed, so even 1-GPU # workers need torchrun to export RANK/WORLD_SIZE. distributed_worker = node_type in {"evaluation", "global_kd"} diff --git a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py index cce0df7fc18..56afcc60ba8 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py @@ -1,17 +1,29 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """Sharded stage adapter for independent worker instances.""" from __future__ import annotations -import subprocess import time import uuid from dataclasses import replace from pathlib import Path from ..executors.slurm import SlurmExecutor +from ..process import run_argv from ..schema import ( AttemptSpec, CampaignPlan, @@ -32,7 +44,7 @@ __all__ = ["ShardedStageAdapter"] -_SHARDED_ENTRYPOINTS = { +_SHARDED_ENTRYPOINTS: dict[str, tuple[str, list[str]]] = { "vllm_stats": ("examples/puzzletron/run_runtime_stats_shard.py", []), "zero_shot_evaluation": ( "examples/puzzletron/run_profile_online_evaluation.py", @@ -63,7 +75,8 @@ def _run_slurm_aggregate( if slurm is None: raise ValueError("Slurm aggregation requires runner.slurm") attempt_id = str(uuid.uuid4()) - log_path = plan.puzzle_dir / "logs" / f"{node.stage_id}_merge_{attempt_id}.log" + partition = node.partition or slurm.partition_cpu + log_path = plan.log_dir / f"{node.stage_id}_merge_{attempt_id}.log" attempt = AttemptSpec( attempt_id=attempt_id, work_id=f"{node.stage_id}:aggregate", @@ -78,7 +91,7 @@ def _run_slurm_aggregate( contract_hash=plan.contract_hash, metadata={ "gpus_per_node": 0, - "partition": slurm.partition_cpu or slurm.partition_for_nodes(1), + **({"partition": partition} if partition else {}), }, task_topology=TaskTopology(task_count=1, gpus_per_task=0), ) @@ -207,7 +220,7 @@ def command( overrides: list[str] | None = None, ) -> AttemptSpec: repo = Path(runner.contract.repository) - log_dir = plan.puzzle_dir / "logs" + log_dir = plan.log_dir logical_count = int(item.metadata.get("logical_shard_count", node.instances)) script, extra_args = _SHARDED_ENTRYPOINTS.get( node.stage_id, @@ -217,8 +230,7 @@ def command( measurement_id = item.metadata.get("measurement_id") name_suffix = f"_{measurement_id}" if measurement_id else "" log_path = str( - log_dir - / f"{node.stage_id}{name_suffix}_shard{item.shard_index}_{attempt_id}.log" + log_dir / f"{node.stage_id}{name_suffix}_shard{item.shard_index}_{attempt_id}.log" ) if node.stage_id == "aiperf": aiperf = plan.experiment_config.get("aiperf") or {} @@ -353,12 +365,9 @@ def aggregate( command=command, ) else: - result = subprocess.run( + result = run_argv( command, cwd=plan.runner.contract.repository, - capture_output=True, - text=True, - check=False, ) if result.returncode: raise RuntimeError( diff --git a/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py b/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py index 8940bd91678..95f1f79aa24 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py @@ -671,7 +671,7 @@ def command( ) -> AttemptSpec: repo = Path(runner.contract.repository) main_py = repo / "examples" / "puzzletron" / "main.py" - log_dir = plan.puzzle_dir / "logs" + log_dir = plan.log_dir log_path = str(log_dir / f"{node.stage_id}_{attempt_id}.log") argv: list[str] = [ "python", diff --git a/modelopt/torch/puzzletron/orchestration/compiler.py b/modelopt/torch/puzzletron/orchestration/compiler.py index b41d5e203a2..142c489da9a 100644 --- a/modelopt/torch/puzzletron/orchestration/compiler.py +++ b/modelopt/torch/puzzletron/orchestration/compiler.py @@ -18,7 +18,9 @@ from __future__ import annotations import math +from collections.abc import Sequence from dataclasses import asdict +from difflib import get_close_matches from pathlib import Path from typing import Any, Mapping @@ -40,16 +42,19 @@ ExecutionContract, ExecutionStrategy, FailurePolicy, + HaltPolicy, ParallelMeshOverride, RunnerEnvironment, SlurmRunnerConfig, StageExecutionSpec, StagePlanNode, + normalize_slurm_partition, ) from .stages import ( configured_parent_stage_ids, configured_stage_ids, distributed_stage_ids, + stage_ids, topological_mapping_items, ) from .vllm_measurements import normalize_vllm_measurements @@ -65,6 +70,63 @@ _CONTROLLER_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] _DEFAULT_ARTIFACT_SETTLING_TIMEOUT_SECONDS = 300.0 +_RUNNER_FIELDS = {"kind", "execution_contract", "slurm", "inventory"} +_EXECUTION_CONTRACT_FIELDS = { + "repository", + "venv", + "container", + "container_mounts", + "mounts", + "setup_env", + "prerun_commands", + "prerun", + "postrun_commands", + "postrun", +} +_SLURM_FIELDS = { + "account", + "partition", + "partition_interactive", + "partition_batch", + "partition_cpu", + "interactive_max_nodes", + "max_nodes", + "time_limit", + "qos", + "log_dir", +} +_INVENTORY_FIELDS = {"hosts", "rendezvous_host", "rendezvous_port_base"} +_HOST_FIELDS = {"hostname", "gpus"} +_EXECUTION_FIELDS = {"defaults", "stages"} +_EXECUTION_DEFAULT_FIELDS = { + "artifact_settling_timeout_seconds", + "failure_policy", + "halt_policy", + "gpus_per_node", + "partition", + "resource", +} +_STAGE_EXECUTION_FIELDS = { + "strategy", + "instances", + "num_jobs", + "failure_policy", + "gpus_per_node", + "partition", + "resource", + "parallel", +} +_FINAL_REPORT_FIELDS = {"partition"} +_PARALLEL_FIELDS = { + "tp", + "cp", + "pp", + "ep", + "dp", + "dp_shard", + "dp_replicate", +} + def _mapping(value: Any) -> dict[str, Any]: return dict(value) if isinstance(value, Mapping) else {} @@ -162,14 +224,14 @@ def _post_mip_stage_metadata(config: Mapping[str, Any]) -> tuple[dict[str, Any], ) global_node_ids.add(str(node_id)) node_type = str(node_value.get("type") or "") - metadata = _POST_MIP_NODE_METADATA.get(node_type) - if metadata is None: + node_metadata = _POST_MIP_NODE_METADATA.get(node_type) + if node_metadata is None: raise ValueError(f"unknown post-MIP node type {node_type!r}") - if not metadata.get("implemented", True): + if not node_metadata.get("implemented", True): raise NotImplementedError( f"post-MIP node type {node_type!r} is declared but not implemented" ) - prepared_nodes[str(node_id)] = (dict(node_value), metadata) + prepared_nodes[str(node_id)] = (dict(node_value), dict(node_metadata)) def dependency_ids( _node_id: str, @@ -265,6 +327,121 @@ def dependency_ids( return tuple(compiled) +def _required_mapping(value: Any, *, path: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"{path} must be a mapping") + return dict(value) + + +def _reject_unknown_fields( + payload: Mapping[str, Any], + allowed: set[str], + *, + path: str, +) -> None: + for field in payload: + if not isinstance(field, str): + raise TypeError(f"{path} field names must be strings; got {field!r}") + if field in allowed: + continue + suggestion = get_close_matches(field, sorted(allowed), n=1) + suffix = f"; did you mean {suggestion[0]!r}?" if suggestion else "" + raise ValueError(f"Unknown config field {path}.{field}{suffix}") + + +def _positive_int(value: Any, *, path: str) -> int: + if isinstance(value, bool): + raise TypeError(f"{path} must be a positive integer") + if isinstance(value, int): + parsed = value + elif isinstance(value, str) and value.strip().isdigit(): + parsed = int(value) + else: + raise TypeError(f"{path} must be a positive integer") + if parsed < 1: + raise ValueError(f"{path} must be at least 1") + return parsed + + +def _command_sequence(value: Any, *, path: str) -> tuple[str, ...]: + if value is None: + return () + if isinstance(value, str): + return (value,) + if not isinstance(value, Sequence) or any(not isinstance(item, str) for item in value): + raise TypeError(f"{path} must be a string or a sequence of strings") + return tuple(value) + + +def _validate_execution_payload(execution: Mapping[str, Any]) -> None: + _reject_unknown_fields(execution, _EXECUTION_FIELDS, path="execution") + defaults = _required_mapping(execution.get("defaults", {}), path="execution.defaults") + _reject_unknown_fields(defaults, _EXECUTION_DEFAULT_FIELDS, path="execution.defaults") + if "failure_policy" in defaults: + FailurePolicy(str(defaults["failure_policy"])) + if "halt_policy" in defaults: + HaltPolicy(str(defaults["halt_policy"])) + if "gpus_per_node" in defaults: + _positive_int(defaults["gpus_per_node"], path="execution.defaults.gpus_per_node") + if "partition" in defaults: + normalize_slurm_partition(defaults["partition"], path="execution.defaults.partition") + if "resource" in defaults and str(defaults["resource"]) not in {"cpu", "gpu"}: + raise ValueError("execution.defaults.resource must be 'cpu' or 'gpu'") + + stages = _required_mapping(execution.get("stages", {}), path="execution.stages") + for stage_id, raw_stage in stages.items(): + if not isinstance(stage_id, str) or not stage_id: + raise TypeError(f"execution.stages keys must be non-empty strings; got {stage_id!r}") + stage_path = f"execution.stages.{stage_id}" + stage = _required_mapping(raw_stage, path=stage_path) + allowed_fields = ( + _FINAL_REPORT_FIELDS if stage_id == "final_report" else _STAGE_EXECUTION_FIELDS + ) + _reject_unknown_fields(stage, allowed_fields, path=stage_path) + if stage_id == "final_report": + if "partition" in stage: + normalize_slurm_partition(stage["partition"], path=f"{stage_path}.partition") + continue + if "instances" in stage and "num_jobs" in stage: + raise ValueError(f"{stage_path} cannot set both instances and legacy num_jobs") + if "strategy" in stage: + ExecutionStrategy(str(stage["strategy"])) + if "failure_policy" in stage: + FailurePolicy(str(stage["failure_policy"])) + if "instances" in stage: + _positive_int(stage["instances"], path=f"{stage_path}.instances") + if "num_jobs" in stage: + _positive_int(stage["num_jobs"], path=f"{stage_path}.num_jobs") + if "gpus_per_node" in stage: + _positive_int(stage["gpus_per_node"], path=f"{stage_path}.gpus_per_node") + if "partition" in stage: + normalize_slurm_partition(stage["partition"], path=f"{stage_path}.partition") + if "resource" in stage and str(stage["resource"]) not in {"cpu", "gpu"}: + raise ValueError(f"{stage_path}.resource must be 'cpu' or 'gpu'") + if "parallel" in stage: + parallel = _required_mapping(stage["parallel"], path=f"{stage_path}.parallel") + if "sequence_parallel" in parallel: + raise ValueError( + f"{stage_path}.parallel.sequence_parallel belongs in the experiment " + "model-parallel profile; it does not affect scheduler allocation" + ) + _reject_unknown_fields(parallel, _PARALLEL_FIELDS, path=f"{stage_path}.parallel") + if "dp" in parallel and "dp_replicate" in parallel: + raise ValueError(f"{stage_path}.parallel cannot set both dp and dp_replicate") + for field, value in parallel.items(): + _positive_int(value, path=f"{stage_path}.parallel.{field}") + + +def _validate_execution_stage_ids( + execution: Mapping[str, Any], + *, + dynamic_stage_ids: Sequence[str], +) -> None: + stages = _required_mapping(execution.get("stages", {}), path="execution.stages") + allowed = {*stage_ids(), *dynamic_stage_ids, "final_report"} + _reject_unknown_fields(stages, allowed, path="execution.stages") + + def _load_yaml(path: str | Path) -> dict[str, Any]: payload = yaml.safe_load(Path(path).read_text()) if payload is None: @@ -291,59 +468,109 @@ def load_runner_config(path: str | Path) -> RunnerEnvironment: """Load a runner environment YAML file.""" payload = _load_yaml(path) - runner = _mapping(payload.get("runner")) + _reject_unknown_fields(payload, {"runner"}, path="config") + runner = _required_mapping(payload.get("runner"), path="runner") + _reject_unknown_fields(runner, _RUNNER_FIELDS, path="runner") kind = str(runner.get("kind", "slurm")) - contract_payload = _mapping(runner.get("execution_contract")) - prerun = contract_payload.get("prerun_commands") or contract_payload.get("prerun") or () - postrun = contract_payload.get("postrun_commands") or contract_payload.get("postrun") or () - if isinstance(prerun, str): - prerun = (prerun,) - if isinstance(postrun, str): - postrun = (postrun,) + contract_payload = _required_mapping( + runner.get("execution_contract", {}), path="runner.execution_contract" + ) + _reject_unknown_fields( + contract_payload, + _EXECUTION_CONTRACT_FIELDS, + path="runner.execution_contract", + ) + for canonical, alias in ( + ("container_mounts", "mounts"), + ("prerun_commands", "prerun"), + ("postrun_commands", "postrun"), + ): + if canonical in contract_payload and alias in contract_payload: + raise ValueError( + f"runner.execution_contract cannot set both {canonical} and legacy {alias}" + ) + prerun = _command_sequence( + contract_payload.get("prerun_commands", contract_payload.get("prerun")), + path="runner.execution_contract.prerun_commands", + ) + postrun = _command_sequence( + contract_payload.get("postrun_commands", contract_payload.get("postrun")), + path="runner.execution_contract.postrun_commands", + ) contract = ExecutionContract( repository=str(contract_payload.get("repository", ".")), venv=str(contract_payload.get("venv", ".venv")), container=contract_payload.get("container"), container_mounts=contract_payload.get("container_mounts") or contract_payload.get("mounts"), setup_env=contract_payload.get("setup_env"), - prerun_commands=tuple(str(item) for item in prerun), - postrun_commands=tuple(str(item) for item in postrun), + prerun_commands=prerun, + postrun_commands=postrun, ) slurm = None baremetal = None if kind == "slurm": - slurm_payload = _mapping(runner.get("slurm")) + if "inventory" in runner: + raise ValueError("runner.inventory is only valid when runner.kind is 'baremetal'") + slurm_payload = _required_mapping(runner.get("slurm", {}), path="runner.slurm") + _reject_unknown_fields(slurm_payload, _SLURM_FIELDS, path="runner.slurm") + max_nodes = ( + _positive_int(slurm_payload["max_nodes"], path="runner.slurm.max_nodes") + if slurm_payload.get("max_nodes") is not None + else None + ) slurm = SlurmRunnerConfig( account=str(slurm_payload.get("account", "")), - partition=str( - slurm_payload.get( - "partition", - slurm_payload.get("partition_batch", "batch"), - ) - ), + partition=slurm_payload.get("partition"), partition_interactive=slurm_payload.get("partition_interactive"), partition_batch=slurm_payload.get("partition_batch"), partition_cpu=slurm_payload.get("partition_cpu"), - interactive_max_nodes=int(slurm_payload.get("interactive_max_nodes", 2)), - max_nodes=( - int(slurm_payload["max_nodes"]) - if slurm_payload.get("max_nodes") is not None - else None + interactive_max_nodes=_positive_int( + slurm_payload.get("interactive_max_nodes", 2), + path="runner.slurm.interactive_max_nodes", ), + max_nodes=max_nodes, time_limit=str(slurm_payload.get("time_limit", "4:00:00")), qos=slurm_payload.get("qos"), log_dir=slurm_payload.get("log_dir"), ) elif kind == "baremetal": - inventory = _mapping(runner.get("inventory")) - hosts = tuple( - BareMetalHost(hostname=str(item["hostname"]), gpus=int(item.get("gpus", 8))) - for item in inventory.get("hosts", []) - ) + if "slurm" in runner: + raise ValueError("runner.slurm is only valid when runner.kind is 'slurm'") + inventory = _required_mapping(runner.get("inventory", {}), path="runner.inventory") + _reject_unknown_fields(inventory, _INVENTORY_FIELDS, path="runner.inventory") + raw_hosts = inventory.get("hosts", ()) + if isinstance(raw_hosts, (str, bytes)) or not isinstance(raw_hosts, Sequence): + raise TypeError("runner.inventory.hosts must be a sequence of host mappings") + hosts_list = [] + for index, raw_host in enumerate(raw_hosts): + host_path = f"runner.inventory.hosts[{index}]" + host = _required_mapping(raw_host, path=host_path) + _reject_unknown_fields(host, _HOST_FIELDS, path=host_path) + hostname = str(host.get("hostname", "")).strip() + if not hostname: + raise ValueError(f"{host_path}.hostname must be non-empty") + hosts_list.append( + BareMetalHost( + hostname=hostname, + gpus=_positive_int(host.get("gpus", 8), path=f"{host_path}.gpus"), + ) + ) + hosts = tuple(hosts_list) + if not hosts: + raise ValueError("runner.inventory.hosts must contain at least one host") + hostnames = [host.hostname for host in hosts] + if len(hostnames) != len(set(hostnames)): + raise ValueError("runner.inventory.hosts contains duplicate hostnames") + rendezvous_host = inventory.get("rendezvous_host") + if rendezvous_host is not None and str(rendezvous_host) not in hostnames: + raise ValueError("runner.inventory.rendezvous_host must name an inventory host") baremetal = BareMetalRunnerConfig( hosts=hosts, - rendezvous_host=inventory.get("rendezvous_host"), - rendezvous_port_base=int(inventory.get("rendezvous_port_base", 29500)), + rendezvous_host=str(rendezvous_host) if rendezvous_host is not None else None, + rendezvous_port_base=_positive_int( + inventory.get("rendezvous_port_base", 29500), + path="runner.inventory.rendezvous_port_base", + ), ) else: raise ValueError(f"Unsupported runner kind: {kind}") @@ -353,7 +580,6 @@ def load_runner_config(path: str | Path) -> RunnerEnvironment: contract=contract, slurm=slurm, baremetal=baremetal, - defaults=_mapping(runner.get("defaults")), ) updated_contract = with_contract_hash(environment) return RunnerEnvironment( @@ -361,7 +587,6 @@ def load_runner_config(path: str | Path) -> RunnerEnvironment: contract=updated_contract, slurm=environment.slurm, baremetal=environment.baremetal, - defaults=environment.defaults, ) @@ -369,7 +594,10 @@ def load_execution_config(path: str | Path) -> dict[str, Any]: """Load execution semantics YAML.""" payload = _load_yaml(path) - return _mapping(payload.get("execution")) + _reject_unknown_fields(payload, {"execution"}, path="config") + execution = _required_mapping(payload.get("execution"), path="execution") + _validate_execution_payload(execution) + return execution def _resolve_artifact_settling_timeout_seconds( @@ -434,7 +662,9 @@ def resolve_stage_execution_specs( defaults = _mapping(execution.get("defaults")) _resolve_artifact_settling_timeout_seconds(defaults) - default_gpus_per_node = int(defaults.get("gpus_per_node", 8)) + default_gpus_per_node = _positive_int( + defaults.get("gpus_per_node", 8), path="execution.defaults.gpus_per_node" + ) default_policy = FailurePolicy(str(defaults.get("failure_policy", FailurePolicy.STRICT.value))) stage_payload = _mapping(execution.get("stages")) dynamic_defaults = dict(dynamic_defaults or {}) @@ -450,12 +680,24 @@ def resolve_stage_execution_specs( ) else: strategy = ExecutionStrategy(str(strategy_name)) - instances = int(payload.get("instances", 1)) - if strategy is ExecutionStrategy.SHARDED and instances == 1: - instances = int(payload.get("instances", payload.get("num_jobs", 1))) + instances = _positive_int( + payload.get("instances", payload.get("num_jobs", 1)), + path=f"execution.stages.{stage_id}.instances", + ) + if strategy is ExecutionStrategy.SINGLE and instances != 1: + raise ValueError( + f"execution.stages.{stage_id}.instances must be 1 for strategy 'single'" + ) policy = FailurePolicy(str(payload.get("failure_policy", default_policy.value))) gpus_per_node = payload.get("gpus_per_node", defaults.get("gpus_per_node")) - partition = payload.get("partition", defaults.get("partition")) + partition_path = ( + f"execution.stages.{stage_id}.partition" + if "partition" in payload + else "execution.defaults.partition" + ) + partition = normalize_slurm_partition( + payload.get("partition", defaults.get("partition")), path=partition_path + ) resource = str(payload.get("resource", defaults.get("resource", "gpu"))) if resource not in {"cpu", "gpu"}: raise ValueError( @@ -464,13 +706,13 @@ def resolve_stage_execution_specs( resolved[stage_id] = StageExecutionSpec( stage_id=stage_id, strategy=strategy, - instances=max(1, instances), + instances=instances, failure_policy=policy, mesh_override=_parse_mesh_override(payload.get("parallel")), gpus_per_node=( int(gpus_per_node) if gpus_per_node is not None else default_gpus_per_node ), - partition=str(partition) if partition is not None else None, + partition=partition, resource=resource, ) return resolved @@ -486,6 +728,7 @@ def compile_campaign_plan( ) -> CampaignPlan: """Compile one campaign plan from experiment + runner + execution configs.""" + _validate_execution_payload(execution) experiment_path = Path(experiment_config_path) experiment_config = load_experiment_config(experiment_path, overrides=overrides or []) puzzle_dir = Path( @@ -494,6 +737,10 @@ def compile_campaign_plan( or "." ) post_mip_stages = _post_mip_stage_metadata(experiment_config) + _validate_execution_stage_ids( + execution, + dynamic_stage_ids=tuple(row["stage_id"] for row in post_mip_stages), + ) enabled = configured_stage_ids( experiment_config, dynamic_post_mip_stage_ids=(row["stage_id"] for row in post_mip_stages), @@ -510,6 +757,19 @@ def compile_campaign_plan( enabled, dynamic_defaults=dynamic_execution_defaults, ) + execution_defaults = _mapping(execution.get("defaults")) + final_report = _mapping(_mapping(execution.get("stages")).get("final_report")) + final_report_partition_path = ( + "execution.stages.final_report.partition" + if "partition" in final_report + else "execution.defaults.partition" + ) + final_report_partition = normalize_slurm_partition( + final_report.get("partition", execution_defaults.get("partition")), + path=final_report_partition_path, + ) + if final_report_partition is None and runner.slurm is not None: + final_report_partition = runner.slurm.partition_cpu distributed = set(distributed_stage_ids()) nodes: list[StagePlanNode] = [] post_mip_by_stage = {row["stage_id"]: row for row in post_mip_stages} @@ -619,10 +879,11 @@ def compile_campaign_plan( puzzle_dir=puzzle_dir, experiment_config=experiment_config, runner=runner, - execution_defaults=_mapping(execution.get("defaults")), + execution_defaults=execution_defaults, stages=tuple(nodes), contract_hash=contract_hash, overrides=tuple(overrides or ()), + final_report_partition=final_report_partition, ) @@ -636,6 +897,10 @@ def plan_to_dict(plan: CampaignPlan) -> dict[str, Any]: "overrides": list(plan.overrides), "runner_kind": plan.runner.kind, "execution_defaults": dict(plan.execution_defaults), + "final_report": { + "resource": "cpu", + "partition": plan.final_report_partition, + }, "stages": [ { "stage_id": node.stage_id, diff --git a/modelopt/torch/puzzletron/orchestration/config.py b/modelopt/torch/puzzletron/orchestration/config.py index f5265eb8e36..5fefd3610f9 100644 --- a/modelopt/torch/puzzletron/orchestration/config.py +++ b/modelopt/torch/puzzletron/orchestration/config.py @@ -25,6 +25,11 @@ import yaml +if __package__.startswith("puzzletron_orchestrator"): + from puzzletron_orchestrator._config_aliases import _validate_compatibility_aliases +else: + from .._config_aliases import _validate_compatibility_aliases + __all__ = ["load_experiment_config"] _INTERPOLATION = re.compile(r"\$\{([^${}]*)\}") @@ -139,8 +144,8 @@ def _resolve_expression(expression: str, config: Mapping[str, Any]) -> Any: return {"__type__": expression.removeprefix("get_object:")} try: return deepcopy(_lookup(config, expression)) - except KeyError: - return "${" + expression + "}" + except KeyError as error: + raise ValueError(f"Unknown config interpolation {expression!r}") from error def _resolve_string(value: str, config: Mapping[str, Any]) -> Any: @@ -177,36 +182,36 @@ def _resolve(value: Any, config: Mapping[str, Any]) -> Any: def _apply_override(config: dict[str, Any], override: str) -> None: if override.startswith("~"): - raise ValueError(f"Deletion overrides are not supported: {override!r}") + raise ValueError( + f"Deletion overrides are not supported by the dependency-light controller: {override!r}" + ) key, separator, raw_value = override.partition("=") + key = key.strip() + if key.startswith("~"): + raise ValueError( + f"Deletion overrides are not supported by the dependency-light controller: {override!r}" + ) if not separator: raise ValueError(f"Override must have KEY=VALUE form: {override!r}") - addition_only = False - allow_missing = False - if key.startswith("++"): - key = key[2:] - allow_missing = True - elif key.startswith("+"): - key = key[1:] - addition_only = True - allow_missing = True - if not key or key.startswith(("+", "~")): - raise ValueError(f"Unsupported Hydra override form: {override!r}") + if key.startswith("+") and not key.startswith("++"): + raise ValueError( + "Single-plus Hydra overrides are not supported by the dependency-light " + f"controller; use KEY=VALUE or ++KEY=VALUE: {override!r}" + ) + key = key.removeprefix("++") keys = key.split(".") + if ( + not key + or any(not part for part in keys) + or any(part.startswith(("+", "~")) for part in keys) + ): + raise ValueError(f"Override has an invalid dotted key: {override!r}") target = config for part in keys[:-1]: - if part not in target: - if not allow_missing: - raise ValueError(f"Override path does not exist: {override!r}") - target[part] = {} - child = target[part] + child = target.setdefault(part, {}) if not isinstance(child, dict): raise ValueError(f"Override path crosses a scalar: {override!r}") target = child - if addition_only and keys[-1] in target: - raise ValueError(f"Addition override already exists: {override!r}") - if not allow_missing and keys[-1] not in target: - raise ValueError(f"Override key does not exist: {override!r}") target[keys[-1]] = _load_yaml(raw_value) @@ -235,6 +240,8 @@ def load_experiment_config( else: raise ValueError(f"Config interpolation did not converge: {config_path}") + _validate_compatibility_aliases(config) + config["_runtime"] = { "config_path": str(config_path), "overrides": list(overrides or ()), diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index d98a18b7def..e524fdc80b3 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.py @@ -200,7 +200,7 @@ def __init__( ) -> None: self.plan = plan self.store = CampaignStateStore(plan.puzzle_dir) - (plan.puzzle_dir / "logs").mkdir(parents=True, exist_ok=True) + plan.log_dir.mkdir(parents=True, exist_ok=True) self.executor = executor or create_executor(plan, local=local) self.poll_interval_seconds = poll_interval_seconds self.logger = logger or OrchestratorLogger() diff --git a/modelopt/torch/puzzletron/orchestration/executors/slurm.py b/modelopt/torch/puzzletron/orchestration/executors/slurm.py index 4460bb96fd2..140734a2322 100644 --- a/modelopt/torch/puzzletron/orchestration/executors/slurm.py +++ b/modelopt/torch/puzzletron/orchestration/executors/slurm.py @@ -20,11 +20,11 @@ import math import os import shlex -import subprocess import time from pathlib import Path from typing import Sequence +from ..process import ProcessResult, run_argv from ..schema import AttemptSpec, JobHandle, JobState, JobStatus, RunnerEnvironment from ..task_launcher import TASK_IDENTITY_ENV_KEYS from ..task_topology import resolve_task_topology @@ -33,8 +33,8 @@ __all__ = ["SlurmExecutor", "render_hook_lines", "render_sbatch_script"] -def _run_command(argv: Sequence[str]) -> subprocess.CompletedProcess[str]: - return subprocess.run(list(argv), capture_output=True, text=True, check=False) +def _run_command(argv: Sequence[str]) -> ProcessResult: + return run_argv(argv) def _slurm_job_id(handle: JobHandle) -> str | None: @@ -61,7 +61,7 @@ def _slurm_job_id(handle: JobHandle) -> str | None: _CANCEL_POLL_SECONDS = 1.0 -def _is_transient_submit_error(result: subprocess.CompletedProcess[str]) -> bool: +def _is_transient_submit_error(result: ProcessResult) -> bool: detail = f"{result.stderr}\n{result.stdout}".lower() return any(marker in detail for marker in _TRANSIENT_SUBMIT_ERRORS) @@ -106,7 +106,7 @@ def render_sbatch_script( *, attempt: AttemptSpec, runner: RunnerEnvironment, - partition: str, + partition: str | None, account: str, time_limit: str, qos: str | None, @@ -173,10 +173,11 @@ def render_sbatch_script( header_lines = [ "#!/bin/bash", f"#SBATCH --job-name={job_name}", - f"#SBATCH --partition={partition}", f"#SBATCH --account={account}", f"#SBATCH --time={time_limit}", ] + if partition: + header_lines.insert(2, f"#SBATCH --partition={partition}") if qos: header_lines.append(f"#SBATCH --qos={qos}") header_lines.extend( @@ -258,10 +259,13 @@ def __init__(self, runner: RunnerEnvironment, *, scripts_dir: Path | None = None def submit(self, attempt: AttemptSpec) -> JobHandle: slurm = self.runner.slurm - assert slurm is not None - partition = str( - attempt.metadata.get("partition") or slurm.partition_for_nodes(attempt.allocation_nodes) + if slurm is None: + raise RuntimeError("Slurm executor lost its runner configuration") + partition = attempt.metadata.get("partition") or slurm.partition_for_nodes( + attempt.allocation_nodes ) + if partition is not None: + partition = str(partition) job_name = f"pt-{attempt.stage_id[:18]}-{attempt.attempt_id[:8]}" script_path = self.scripts_dir / f"{attempt.stage_id}_{attempt.attempt_id}.sh" script_path.write_text( diff --git a/modelopt/torch/puzzletron/orchestration/identity.py b/modelopt/torch/puzzletron/orchestration/identity.py index 3cc415fa6d2..8fab4141d1e 100644 --- a/modelopt/torch/puzzletron/orchestration/identity.py +++ b/modelopt/torch/puzzletron/orchestration/identity.py @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """Execution-contract identity hashing for orchestration attempts.""" @@ -72,13 +84,9 @@ def mip_input_artifact_paths( { f"width/{width}/stats": base / "subblock_stats.json", f"width/{width}/scores": base / score_name, - f"width/{width}/canonical": base - / "single_sequence_replacement_solutions.json", + f"width/{width}/canonical": base / "single_sequence_replacement_solutions.json", f"width/{width}/library": base / "replacement_library.json", - f"width/{width}/teacher_config": base - / "ckpts" - / "sorted_teacher" - / "config.json", + f"width/{width}/teacher_config": base / "ckpts" / "sorted_teacher" / "config.json", f"width/{width}/teacher_index": base / "ckpts" / "sorted_teacher" @@ -95,9 +103,7 @@ def artifact_snapshot_identity(paths: Mapping[str, str | Path]) -> str: for label, raw_path in sorted(paths.items()): path = Path(raw_path) members = ( - sorted(item for item in path.rglob("*") if item.is_file()) - if path.is_dir() - else [path] + sorted(item for item in path.rglob("*") if item.is_file()) if path.is_dir() else [path] ) if not path.exists(): rows.append({"label": label, "missing": True}) @@ -142,6 +148,7 @@ def execution_contract_hash(runner: RunnerEnvironment) -> str: "max_nodes": runner.slurm.max_nodes, "time_limit": runner.slurm.time_limit, "qos": runner.slurm.qos, + "log_dir": runner.slurm.log_dir, } if runner.baremetal is not None: payload["baremetal"] = { diff --git a/modelopt/torch/puzzletron/orchestration/process.py b/modelopt/torch/puzzletron/orchestration/process.py new file mode 100644 index 00000000000..397d0c11d18 --- /dev/null +++ b/modelopt/torch/puzzletron/orchestration/process.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shell-free process execution for dependency-light orchestration.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Sequence + from pathlib import Path + +__all__ = ["ProcessResult", "run_argv"] + + +@dataclass(frozen=True) +class ProcessResult: + """Captured result from one argv-only child process.""" + + args: tuple[str, ...] + returncode: int + stdout: str + stderr: str + + +async def _run_argv_async(argv: tuple[str, ...], *, cwd: str | Path | None) -> ProcessResult: + # create_subprocess_exec passes argv directly to the child without a command shell. + process = await asyncio.create_subprocess_exec( + *argv, + cwd=str(cwd) if cwd is not None else None, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout, stderr = await process.communicate() + returncode = process.returncode + if returncode is None: + raise RuntimeError("child process did not terminate after communicate()") + return ProcessResult( + args=argv, + returncode=returncode, + stdout=stdout.decode(errors="replace"), + stderr=stderr.decode(errors="replace"), + ) + + +def run_argv(argv: Sequence[str], *, cwd: str | Path | None = None) -> ProcessResult: + """Run an explicit argv sequence without a shell and capture text output.""" + + return asyncio.run(_run_argv_async(tuple(str(part) for part in argv), cwd=cwd)) diff --git a/modelopt/torch/puzzletron/orchestration/reporting.py b/modelopt/torch/puzzletron/orchestration/reporting.py index aaec8ea98df..c89749cf61b 100644 --- a/modelopt/torch/puzzletron/orchestration/reporting.py +++ b/modelopt/torch/puzzletron/orchestration/reporting.py @@ -1,13 +1,27 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """Runner-backed final campaign report contracts.""" from __future__ import annotations from dataclasses import dataclass -from pathlib import Path -from typing import Any, Mapping +from typing import TYPE_CHECKING, Any, Mapping + +if TYPE_CHECKING: + from pathlib import Path from .schema import AttemptSpec, CampaignPlan, CommandSpec, TaskLauncher, TaskTopology @@ -68,9 +82,9 @@ def build_final_report_attempt(plan: CampaignPlan, *, attempt_id: str) -> Attemp """Build one direct, CPU-only attempt in the campaign runner environment.""" metadata: dict[str, Any] = {"gpus_per_node": 0} - if plan.runner.slurm is not None and plan.runner.slurm.partition_cpu: - metadata["partition"] = plan.runner.slurm.partition_cpu - log_path = plan.puzzle_dir / "logs" / f"final_report_{attempt_id}.log" + if plan.final_report_partition is not None: + metadata["partition"] = plan.final_report_partition + log_path = plan.log_dir / f"final_report_{attempt_id}.log" return AttemptSpec( attempt_id=attempt_id, work_id="final_report:0", diff --git a/modelopt/torch/puzzletron/orchestration/schema.py b/modelopt/torch/puzzletron/orchestration/schema.py index 0853abe88de..4d24c631d11 100644 --- a/modelopt/torch/puzzletron/orchestration/schema.py +++ b/modelopt/torch/puzzletron/orchestration/schema.py @@ -17,12 +17,31 @@ from __future__ import annotations +from collections.abc import Sequence from dataclasses import dataclass, field from enum import Enum -from typing import TYPE_CHECKING, Any, Mapping - -if TYPE_CHECKING: - from pathlib import Path +from pathlib import Path +from typing import Any, Mapping, cast + + +def normalize_slurm_partition(value: Any, *, path: str) -> str | None: + """Normalize one or more eligible Slurm partitions for ``--partition``.""" + + if value is None: + return None + values = value.split(",") if isinstance(value, str) else value + if isinstance(values, (str, bytes)) or not isinstance(values, Sequence): + raise TypeError(f"{path} must be a partition name or a sequence of names") + if any(not isinstance(item, str) for item in values): + raise TypeError(f"{path} must contain only partition names") + partitions = [item.strip() for item in values] + if not partitions or any(not item for item in partitions): + raise ValueError(f"{path} must contain at least one non-empty partition name") + if any("," in item or any(character.isspace() for character in item) for item in partitions): + raise ValueError(f"{path} contains an invalid partition name") + if len(set(partitions)) != len(partitions): + raise ValueError(f"{path} contains duplicate partition names") + return ",".join(partitions) class ExecutionStrategy(str, Enum): @@ -135,28 +154,46 @@ class SlurmRunnerConfig: """Slurm-specific runner facts.""" account: str - partition: str = "batch" - partition_interactive: str | None = None - partition_batch: str | None = None - partition_cpu: str | None = None + partition: str | Sequence[str] | None = None + partition_interactive: str | Sequence[str] | None = None + partition_batch: str | Sequence[str] | None = None + partition_cpu: str | Sequence[str] | None = None interactive_max_nodes: int = 2 max_nodes: int | None = None time_limit: str = "4:00:00" qos: str | None = None log_dir: str | None = None - def partition_for_nodes(self, nodes: int) -> str: - """Pick interactive for short/small jobs and batch otherwise.""" - - interactive = self.partition_interactive or ( - self.partition if self.partition == "interactive" else None - ) - batch = self.partition_batch or ( - self.partition if self.partition != "interactive" else "batch" - ) + def __post_init__(self) -> None: + for field_name in ( + "partition", + "partition_interactive", + "partition_batch", + "partition_cpu", + ): + object.__setattr__( + self, + field_name, + normalize_slurm_partition( + getattr(self, field_name), path=f"runner.slurm.{field_name}" + ), + ) + if self.interactive_max_nodes <= 0: + raise ValueError("runner.slurm.interactive_max_nodes must be positive") + + def partition_for_nodes(self, nodes: int) -> str | None: + """Resolve the canonical partition or a deprecated role-based fallback.""" + + partition = cast("str | None", self.partition) + partition_interactive = cast("str | None", self.partition_interactive) + partition_batch = cast("str | None", self.partition_batch) + if self.partition_interactive is None and self.partition_batch is None: + return partition + interactive = partition_interactive or (partition if partition == "interactive" else None) + batch = partition_batch or (partition if partition != "interactive" else "batch") if interactive and nodes <= self.interactive_max_nodes: return interactive - return batch or self.partition + return batch or partition @dataclass(frozen=True) @@ -184,7 +221,6 @@ class RunnerEnvironment: contract: ExecutionContract slurm: SlurmRunnerConfig | None = None baremetal: BareMetalRunnerConfig | None = None - defaults: Mapping[str, Any] = field(default_factory=dict) @dataclass(frozen=True) @@ -219,6 +255,15 @@ class CampaignPlan: stages: tuple[StagePlanNode, ...] contract_hash: str overrides: tuple[str, ...] = () + final_report_partition: str | None = None + + @property + def log_dir(self) -> Path: + """Return the configured shared log directory for every campaign attempt.""" + + if self.runner.slurm is not None and self.runner.slurm.log_dir: + return Path(self.runner.slurm.log_dir).expanduser() + return self.puzzle_dir / "logs" @dataclass(frozen=True) diff --git a/modelopt/torch/puzzletron/pipeline_config.py b/modelopt/torch/puzzletron/pipeline_config.py index e2c956d9c2a..0b43a04b455 100644 --- a/modelopt/torch/puzzletron/pipeline_config.py +++ b/modelopt/torch/puzzletron/pipeline_config.py @@ -22,6 +22,7 @@ import hydra from omegaconf import DictConfig, OmegaConf +from ._config_aliases import _validate_compatibility_aliases from .dataset.config import PuzzletronDataSpec from .granularity import resolve_granularity from .tools.hydra_utils import initialize_hydra_config_for_dir, register_hydra_resolvers @@ -174,6 +175,7 @@ def normalize_pipeline_config(config: DictConfig | dict[str, Any]) -> dict[str, dictionary. """ cfg = _to_plain(config) + _validate_compatibility_aliases(cfg) if "parallel" in cfg: raise ValueError( "top-level parallel was removed; configure automodel.parallel on each " diff --git a/puzzletron_setup/bundle.py b/puzzletron_setup/bundle.py index 3ee3e8f0936..42e922d819c 100644 --- a/puzzletron_setup/bundle.py +++ b/puzzletron_setup/bundle.py @@ -145,7 +145,6 @@ def _parallel(mesh: Mapping[str, Any]) -> dict[str, Any]: "ep": int(mesh.get("ep", 1)), "dp_shard": int(mesh.get("dp_shard", 1)), "dp_replicate": int(mesh.get("dp_replicate", 1)), - "sequence_parallel": False, } @@ -155,7 +154,7 @@ def _serving_parallel(topology: Mapping[str, Any]) -> dict[str, Any]: mesh = vllm_topology_to_mesh(topology) except (TypeError, ValueError) as error: raise SetupError(str(error)) from error - return {**mesh.as_dict(), "sequence_parallel": False} + return mesh.as_dict() def _aligned_batch_size(mesh: Mapping[str, Any], requested: int = 1) -> int: @@ -722,9 +721,7 @@ def render_runner(state: Mapping[str, Any], budget: str) -> dict[str, Any]: "execution_contract": deepcopy(infrastructure["execution_contract"]), } if runner["kind"] == "slurm": - slurm = deepcopy(_mapping(runner_answers.get("slurm"))) - slurm["partition"] = slurm.get("partition_batch", "batch") - runner["slurm"] = slurm + runner["slurm"] = deepcopy(_mapping(runner_answers.get("slurm"))) elif runner["kind"] == "baremetal": runner["inventory"] = deepcopy(_mapping(runner_answers.get("inventory"))) else: @@ -819,8 +816,10 @@ def _dynamic_stage_entries( "instances": 1 if cpu_stage else max(1, instances), "gpus_per_node": gpus_per_node, } - if cpu_stage and cpu_partition: - entry.update(resource="cpu", partition=cpu_partition) + if cpu_stage: + entry["resource"] = "cpu" + if cpu_partition: + entry["partition"] = cpu_partition if node_type == "evaluation": entry["parallel"] = dict(common) elif node_type in {"aiperf", "downstream_evaluation"}: @@ -854,7 +853,7 @@ def render_execution( pool_workers = int(workers.get("pool", 1)) sharded_workers = int(workers.get("sharded", 1)) embedding_widths = list(_mapping(experiment.get("embedding_pruning")).get("widths") or ()) - stages = { + stages: dict[str, dict[str, Any]] = { "convert": {"strategy": "single", "instances": 1, "parallel": single_gpu}, "tokenize_data": {"strategy": "single", "instances": 1}, "vllm_stats": { @@ -902,9 +901,10 @@ def render_execution( ) ) cpu_stage_ids = {"convert", "tokenize_data", "build_library", "mip"} - if cpu_partition: - for stage_id in cpu_stage_ids: - stages[stage_id].update(resource="cpu", partition=cpu_partition) + for stage_id in cpu_stage_ids: + stages[stage_id]["resource"] = "cpu" + if cpu_partition: + stages[stage_id]["partition"] = cpu_partition for stage_id, stage in stages.items(): stage.setdefault("gpus_per_node", gpus_per_node) if "short_kd" in stage_id or "global_kd" in stage_id: diff --git a/puzzletron_setup/state.py b/puzzletron_setup/state.py index 904d7bb9a02..ddb6cfe696e 100644 --- a/puzzletron_setup/state.py +++ b/puzzletron_setup/state.py @@ -1,5 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """Versioned, interruption-safe answer persistence for Puzzletron setup.""" diff --git a/puzzletron_setup/v2/bundle.py b/puzzletron_setup/v2/bundle.py index d4c47ee28e5..a34771868a8 100644 --- a/puzzletron_setup/v2/bundle.py +++ b/puzzletron_setup/v2/bundle.py @@ -276,7 +276,8 @@ def render_experiment_v2(state: WizardState, budget: str) -> dict[str, Any]: def _render_runner_v2(config: ResolvedCampaignConfig, budget: str) -> dict[str, Any]: rendered = render_runner(_legacy_state_from_resolved(config), budget) - return _deep_merge(rendered, _plain(config.compatibility.runner)) + compatibility = _plain(config.compatibility.runner) + return _deep_merge(rendered, compatibility) def render_runner_v2(state: WizardState, budget: str) -> dict[str, Any]: @@ -307,9 +308,16 @@ def _render_execution_v2( entry["partition"] = resource.partition if resource.profile_name: profile = config.parallel_profiles.get(resource.profile_name) - entry["parallel"] = profile._parallel() if profile is not None else {} + parallel = profile._parallel() if profile is not None else {} + entry["parallel"] = { + key: value for key, value in parallel.items() if key != "sequence_parallel" + } elif resource.parallel is not None: - entry["parallel"] = _plain(resource.parallel) + entry["parallel"] = { + key: value + for key, value in _plain(resource.parallel).items() + if key != "sequence_parallel" + } stages[str(stage_id)] = entry return rendered @@ -386,25 +394,41 @@ def _bundle_readme( "The command is idempotent only when the existing manifest matches these answers.", "", "```bash", - " ".join(shlex.quote(part) for part in acquisition_command), + shlex.join(acquisition_command), "```", ) ) for budget in ("smoke", "production"): bundle = campaign_dir / budget - orchestrator_command = ( - f"python {orchestrator} --experiment {bundle / 'experiment.yaml'} " - f"--runner {bundle / 'runner.yaml'} --execution {bundle / 'execution.yaml'} " - "--stage full" - ) + orchestrator_args = [ + "python", + str(orchestrator), + "--experiment", + str(bundle / "experiment.yaml"), + "--runner", + str(bundle / "runner.yaml"), + "--execution", + str(bundle / "execution.yaml"), + "--stage", + "full", + ] + inspect_command = shlex.join([*orchestrator_args, "--dry-run"]) + launch_command = shlex.join(orchestrator_args) lines.extend( [ "", f"## {budget.title()}", "", + "Inspect the complete plan without submitting jobs:", + "", "```bash", - f"{orchestrator_command} --dry-run", - orchestrator_command, + inspect_command, + "```", + "", + "After reviewing the plan and worker paths, launch the campaign:", + "", + "```bash", + launch_command, "```", ] ) @@ -414,8 +438,14 @@ def _bundle_readme( "## Resume setup", "", "```bash", - f"python {Path(repository) / 'examples/puzzletron/puzzletron_setup_v2.py'} " - f"--resume {campaign_dir}", + shlex.join( + [ + "python", + str(Path(repository) / "examples/puzzletron/puzzletron_setup_v2.py"), + "--resume", + str(campaign_dir), + ] + ), "```", "", ] diff --git a/puzzletron_setup/v2/defaults.py b/puzzletron_setup/v2/defaults.py index 7718e2ff2f2..d88e187fad0 100644 --- a/puzzletron_setup/v2/defaults.py +++ b/puzzletron_setup/v2/defaults.py @@ -51,9 +51,7 @@ "kind": "slurm", "slurm": { "account": "", - "partition_interactive": "interactive", - "partition_batch": "batch", - "partition_cpu": None, + "partition": None, "time_limit": "4:00:00", "qos": None, "max_nodes": 64, @@ -121,7 +119,6 @@ class _AnyMapping: "data.acquisition.num_samples": 1, "data.acquisition.max_shards_per_subset": 1, "infrastructure.gpus_per_node": 1, - "infrastructure.runner.slurm.interactive_max_nodes": 1, "infrastructure.runner.slurm.max_nodes": 1, "pruning.depth_remove": 0, "pruning.depth_importance_samples": 1, @@ -165,7 +162,11 @@ class _AnyMapping: "infrastructure.execution_contract.prerun_commands", "infrastructure.execution_contract.postrun_commands", } -_STRING_OR_SEQUENCE_PATHS = {"data.subsets", "data.acquisition.subsets"} +_STRING_OR_SEQUENCE_PATHS = { + "data.subsets", + "data.acquisition.subsets", + "infrastructure.runner.slurm.partition", +} _SCHEMA = { "schema_version": None, "campaign": { @@ -208,10 +209,7 @@ class _AnyMapping: "kind": None, "slurm": { "account": None, - "partition_interactive": None, - "partition_batch": None, - "partition_cpu": None, - "interactive_max_nodes": None, + "partition": None, "max_nodes": None, "time_limit": None, "qos": None, diff --git a/puzzletron_setup/v2/resolved.py b/puzzletron_setup/v2/resolved.py index e42fb8d7aca..fb6c494c743 100644 --- a/puzzletron_setup/v2/resolved.py +++ b/puzzletron_setup/v2/resolved.py @@ -25,6 +25,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any +from puzzletron_orchestrator.schema import normalize_slurm_partition from puzzletron_setup import SetupError from .defaults import BUILTIN_DEFAULTS @@ -591,7 +592,9 @@ def _stage_resource(stage_id: str, raw: Mapping[str, Any]) -> ResolvedStageResou instances=int(raw.get("instances", 1)), resource=str(raw.get("resource", "gpu")), gpus_per_node=(int(raw["gpus_per_node"]) if raw.get("gpus_per_node") is not None else None), - partition=(str(raw["partition"]) if raw.get("partition") else None), + partition=normalize_slurm_partition( + raw.get("partition"), path=f"stages.{stage_id}.partition" + ), profile_name=(str(raw["profile_name"]) if raw.get("profile_name") else None), parallel=_mapping(parallel) if isinstance(parallel, Mapping) else None, extra={key: value for key, value in raw.items() if key not in known}, @@ -638,12 +641,19 @@ def effective(path: str, default: Any = _USE_BUILTIN_DEFAULT) -> Any: runner_kind=str(effective("infrastructure.runner.kind")), slurm={ "account": effective("infrastructure.runner.slurm.account"), - "partition_interactive": effective("infrastructure.runner.slurm.partition_interactive"), - "partition_batch": effective("infrastructure.runner.slurm.partition_batch"), - "partition_cpu": effective("infrastructure.runner.slurm.partition_cpu"), + "partition": effective("infrastructure.runner.slurm.partition"), + "partition_interactive": effective( + "infrastructure.runner.slurm.partition_interactive", None + ), + "partition_batch": effective("infrastructure.runner.slurm.partition_batch", None), + "partition_cpu": effective("infrastructure.runner.slurm.partition_cpu", None), + "interactive_max_nodes": effective( + "infrastructure.runner.slurm.interactive_max_nodes", 2 + ), "time_limit": effective("infrastructure.runner.slurm.time_limit"), "qos": effective("infrastructure.runner.slurm.qos"), "max_nodes": effective("infrastructure.runner.slurm.max_nodes"), + "log_dir": effective("infrastructure.runner.slurm.log_dir", None), }, execution_contract={ "repository": effective("infrastructure.execution_contract.repository"), diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index 3c0a1935178..d68962f3402 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -19,7 +19,7 @@ from __future__ import annotations from collections import OrderedDict -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from pathlib import Path from typing import Any @@ -52,7 +52,8 @@ from .parallel_validation import validate_automodel_parallelism, validate_vllm_parallelism from .post_mip import FlowDraft, NodeDraft, PostMIPFlowEditor, recommended_flow from .presets import QUICK_SETUP_PRESETS, get_setup_preset -from .prompts import BACK, InteractiveBackend, PromptBackend, PromptChoice +from .prompts import BACK as _BACK +from .prompts import InteractiveBackend, PromptBackend, PromptChoice from .resources import ( ParallelProfile, ResourceProfileRegistry, @@ -82,6 +83,8 @@ from .wizard_common import _text_field as _text_field from .wizard_common import _vllm_granularity_choices as _vllm_granularity_choices +BACK: Any = _BACK + __all__ = ["SECTION_BUILDERS", "run_wizard_v2"] _CUSTOM_MODEL_SOURCE = "__custom_model_source__" @@ -95,6 +98,18 @@ _NEMOTRON_VLM_ADAPTER = "nemotron_vlm_v2" _NEMOTRON_VLM_DEFAULT_SUBSETS = ("sparsetables", "plotqa_cot", "wiki_en") + +def _render_partition_default(value: Any) -> str: + """Render a scalar-or-list partition default for the text prompt.""" + if value is None: + return "" + if isinstance(value, str): + return value + if isinstance(value, Sequence): + return ",".join(str(item) for item in value) + return str(value) + + SUPPORTED_MODEL_GROUPS = ( ( "Nemotron 3", @@ -564,6 +579,7 @@ def data_section( if selection is BACK: return False catalog, selected_subsets, weights = selection + assert selected_subsets is not None by_name = {item.name: item for item in catalog.subsets} subset_selection = { "source": catalog.source, @@ -686,9 +702,7 @@ def infrastructure_section( ("infrastructure.execution_contract.container_mounts", None), ("infrastructure.runner.kind", "slurm"), ("infrastructure.runner.slurm.account", ""), - ("infrastructure.runner.slurm.partition_interactive", "interactive"), - ("infrastructure.runner.slurm.partition_batch", "batch"), - ("infrastructure.runner.slurm.partition_cpu", None), + ("infrastructure.runner.slurm.partition", None), ("infrastructure.runner.slurm.time_limit", "4:00:00"), ("infrastructure.runner.slurm.qos", None), ("infrastructure.runner.slurm.max_nodes", 64), @@ -781,12 +795,10 @@ def infrastructure_section( ), ("infrastructure.runner.slurm.account", "Slurm account:", ""), ( - "infrastructure.runner.slurm.partition_interactive", - "Interactive partition:", - "interactive", + "infrastructure.runner.slurm.partition", + "Eligible Slurm partitions (comma-separated; blank for site default):", + "", ), - ("infrastructure.runner.slurm.partition_batch", "Batch partition:", "batch"), - ("infrastructure.runner.slurm.partition_cpu", "CPU partition:", ""), ("infrastructure.runner.slurm.time_limit", "Default time limit:", "4:00:00"), ): value = _text_field( @@ -796,6 +808,11 @@ def infrastructure_section( label, fallback, validate=validate_worker_path if path in worker_path_fields else None, + render_default=( + _render_partition_default + if path == "infrastructure.runner.slurm.partition" + else None + ), ) if value is BACK: return False @@ -1512,6 +1529,21 @@ def _axis_selection_validation( return True +def _axis_selection_validator( + axis: Any, + *, + require_reduced: bool, +) -> Callable[[Any], bool | str]: + def validate(raw_values: Any) -> bool | str: + return _axis_selection_validation( + axis, + raw_values, + require_reduced=require_reduced, + ) + + return validate + + def width_axes_section(session: WizardSession, resolver: DefaultsResolver, context: dict) -> bool: pruning = _pruning_payload(session.state) inventory = context["model"].inventory @@ -1564,12 +1596,9 @@ def width_axes_section(session: WizardSession, resolver: DefaultsResolver, conte f"Values for {axis.label}:", [(str(value), value) for value in axis.values], defaults=selected, - validate=lambda values, axis=axis, require_reduced=require_reduced: ( - _axis_selection_validation( - axis, - values, - require_reduced=require_reduced, - ) + validate=_axis_selection_validator( + axis, + require_reduced=require_reduced, ), ) if selected is BACK: @@ -4033,15 +4062,13 @@ def _configure_dynamic_resources( ) resources = _mapping_copy(session.state.collection("stage_resources")) gpus_per_node = int(session.state.get_field("infrastructure.gpus_per_node", 8)) - cpu_partition = session.state.get_field("infrastructure.runner.slurm.partition_cpu", None) for node_id, node in tuple(editor.flow(flow_id).nodes.items()): stage_id = f"post.{flow_id}.{node_id}" if node.node_type in {"filter", "manual_filter", "materialize"}: resources[stage_id] = { "strategy": "single", "instances": 1, - "resource": "cpu" if cpu_partition else "gpu", - "partition": cpu_partition, + "resource": "cpu", "gpus_per_node": gpus_per_node, } continue @@ -4385,12 +4412,7 @@ def output_review_section( "venv": session.state.get_field("infrastructure.execution_contract.venv"), "container": session.state.get_field("infrastructure.execution_contract.container"), "slurm_account": session.state.get_field("infrastructure.runner.slurm.account"), - "interactive_partition": session.state.get_field( - "infrastructure.runner.slurm.partition_interactive" - ), - "batch_partition": session.state.get_field( - "infrastructure.runner.slurm.partition_batch" - ), + "partition": session.state.get_field("infrastructure.runner.slurm.partition"), "gpus_per_node": session.state.get_field("infrastructure.gpus_per_node"), }, "results": session.state.get_field("output.result_root"), diff --git a/puzzletron_setup/v2/wizard_common.py b/puzzletron_setup/v2/wizard_common.py index 479f2439dd6..10ec36b22a7 100644 --- a/puzzletron_setup/v2/wizard_common.py +++ b/puzzletron_setup/v2/wizard_common.py @@ -259,12 +259,16 @@ def _text_field( fallback: str = "", *, validate: Callable[[Any], bool | str] | None = None, + render_default: Callable[[Any], str] | None = None, ) -> Any: resolved = _resolved(session.state, resolver, path, fallback) + default = ( + render_default(resolved.value) if render_default is not None else str(resolved.value or "") + ) value = session.text( path, label, - default=str(resolved.value or ""), + default=default, validate=validate, ) if value is not BACK: diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index da8e631ded3..0110bd79fff 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -44,7 +44,7 @@ from .state import AnswerState if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import Mapping, Sequence __all__ = ["run_wizard"] @@ -349,7 +349,7 @@ def _ask_pruning(prompts: PromptSession, state: AnswerState, model: InspectedMod "and additional GPU cost." ), ) - bypass = {"enabled": bypass_enabled, "sanity": bypass_sanity} + bypass: dict[str, Any] = {"enabled": bypass_enabled, "sanity": bypass_sanity} data = state.section("data") if bypass_enabled: bypass.update( @@ -859,8 +859,10 @@ def _downstream_evaluation_metric_suggestions(node_id: str, config: Mapping[str, tasks = [item.strip() for item in tasks.split(",") if item.strip()] for task in tasks: task_name = str(task).strip() - for metric in _DOWNSTREAM_EVALUATION_METRICS_BY_TASK.get(task_name, ()): - suggestions.append(f"{node_id}.{task_name}.{metric}") + suggestions.extend( + f"{node_id}.{task_name}.{metric}" + for metric in _DOWNSTREAM_EVALUATION_METRICS_BY_TASK.get(task_name, ()) + ) return suggestions @@ -891,6 +893,7 @@ def node_id(name: str) -> str: best = node_id("best") if objective is None: objective = next(iter(run.get("objectives") or ()), {}) + assert objective is not None objective_metric = str(objective.get("metric", "metrics.lm_loss")) objective_direction = str(objective.get("direction", "minimize")) nodes: OrderedDict[str, Any] = OrderedDict() @@ -1338,7 +1341,7 @@ def post_mip_gpus_per_instance(node_type: str, default: int) -> int: return rows -def _print_resource_rows(rows: list[Mapping[str, Any]]) -> None: +def _print_resource_rows(rows: Sequence[Mapping[str, Any]]) -> None: print("\nDerived resource plan:") print(f"{'Stage':30} {'Instances':>10} {'GPU/instance':>14} {'Nodes':>8}") for row in rows: @@ -1432,20 +1435,17 @@ def _ask_infrastructure( } runner: dict[str, Any] = {"kind": runner_kind} if runner_kind == "slurm": - cpu_partition = prompts.text( - "CPU partition (blank to use one GPU node for CPU/IO stages):", + partition = prompts.text( + "Eligible Slurm partitions (comma-separated; blank for site default):", default="", description=( - "Used for conversion, tokenization, sorting, block-library construction, " - "MIP, filters, and materialization. Leave blank when the cluster has no " - "CPU-only partition." + "Slurm may start the job in any listed partition. Configure a stage-level " + "partition only when that stage needs a different eligible set." ), ).strip() runner["slurm"] = { "account": prompts.text("Slurm account:", default=""), - "partition_interactive": prompts.text("Interactive partition:", default="interactive"), - "partition_batch": prompts.text("Batch partition:", default="batch"), - "partition_cpu": cpu_partition or None, + "partition": partition or None, "time_limit": prompts.text("Default time limit:", default="4:00:00"), "qos": prompts.text("QoS (blank for none):", default="").strip() or None, "max_nodes": prompts.integer("Maximum simultaneous nodes:", default=64), diff --git a/tests/unit/torch/puzzletron/test_data_config.py b/tests/unit/torch/puzzletron/test_data_config.py index 3b7a13d654d..0df8642c9f8 100644 --- a/tests/unit/torch/puzzletron/test_data_config.py +++ b/tests/unit/torch/puzzletron/test_data_config.py @@ -1,5 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for Puzzletron data and experiment configuration normalization.""" + +from pathlib import Path import pytest from omegaconf import OmegaConf @@ -8,7 +24,9 @@ from modelopt.torch.puzzletron.pipeline_config import ( adapt_runtime_hydra_config, normalize_pipeline_config, + pipeline_config_from_path, ) +from puzzletron_orchestrator.config import load_experiment_config def test_canonical_packed_multimodal_data_spec(): @@ -80,13 +98,85 @@ def test_pipeline_defaults_sort_sanity_to_include_reverse(): def test_pipeline_preserves_explicit_reverse_sort_opt_out(): - canonical = normalize_pipeline_config( - {"sort_sanity": {"include_reverse": False}} - ) + canonical = normalize_pipeline_config({"sort_sanity": {"include_reverse": False}}) assert canonical["sort_sanity"]["include_reverse"] is False +@pytest.mark.parametrize( + "loader", + [pipeline_config_from_path, load_experiment_config], + ids=["pipeline", "controller"], +) +@pytest.mark.parametrize( + ("legacy_path", "canonical_section", "canonical_key", "legacy", "canonical", "preferred"), + [ + ("puzzle_dir", "experiment", "dir", "first", "second", "puzzle_dir"), + ( + "input_hf_model_path", + "model", + "source", + "model-a", + "model-b", + "input_hf_model_path", + ), + ("teacher_dir", "convert", "teacher_dir", "first", "second", "teacher_dir"), + ("dataset_path", "data", "path", "first", "second", "dataset_path"), + ( + "trust_remote_code", + "model", + "trust_remote_code", + False, + True, + "model.trust_remote_code", + ), + ], +) +def test_pipeline_and_controller_loaders_reject_conflicting_compatibility_aliases( + tmp_path: Path, + loader, + legacy_path: str, + canonical_section: str, + canonical_key: str, + legacy, + canonical, + preferred: str, +) -> None: + experiment = tmp_path / "experiment.yaml" + experiment.write_text( + OmegaConf.to_yaml( + OmegaConf.create( + { + legacy_path: legacy, + canonical_section: {canonical_key: canonical}, + } + ) + ) + ) + + with pytest.raises(ValueError, match=rf"override '{preferred}'.*stay synchronized"): + loader(experiment) + + +@pytest.mark.parametrize( + "override", + ["runtime_annotations.reason=ad-hoc", "++runtime_annotations.reason=ad-hoc"], +) +def test_pipeline_and_controller_loaders_apply_overrides_with_parity( + tmp_path: Path, + override: str, +) -> None: + experiment = tmp_path / "experiment.yaml" + experiment.write_text("experiment:\n dir: run\n") + + pipeline = pipeline_config_from_path(experiment, overrides=[override]) + controller = load_experiment_config(experiment, overrides=[override]) + + assert ( + pipeline["runtime_annotations"] == controller["runtime_annotations"] == {"reason": "ad-hoc"} + ) + + def test_runtime_adapter_derives_legacy_loader_fields_from_canonical_data(): canonical = normalize_pipeline_config( { diff --git a/tests/unit/torch/puzzletron/test_example_runner.py b/tests/unit/torch/puzzletron/test_example_runner.py index 6ea9dead9e1..a797cc843ce 100644 --- a/tests/unit/torch/puzzletron/test_example_runner.py +++ b/tests/unit/torch/puzzletron/test_example_runner.py @@ -181,7 +181,7 @@ def test_loaded_stage_run_publishes_distinct_authored_and_effective_config(tmp_p ) + "\n" ) - override = "+slicing_sanity.tolerance=0.25" + override = "++slicing_sanity.tolerance=0.25" config = pipeline_config_from_path(config_path, overrides=[override]) manifest_path = tmp_path / "manifests" / "slicing_sanity.json" diff --git a/tests/unit/torch/puzzletron/test_orchestration_compiler.py b/tests/unit/torch/puzzletron/test_orchestration_compiler.py index a79335c3917..3ed7b10f77b 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_compiler.py +++ b/tests/unit/torch/puzzletron/test_orchestration_compiler.py @@ -25,10 +25,18 @@ compile_campaign_plan, load_execution_config, load_runner_config, + plan_to_dict, resolve_stage_execution_specs, ) from modelopt.torch.puzzletron.orchestration.controller import CampaignController -from modelopt.torch.puzzletron.orchestration.schema import ExecutionStrategy, HaltPolicy +from modelopt.torch.puzzletron.orchestration.identity import execution_contract_hash, hash_payload +from modelopt.torch.puzzletron.orchestration.schema import ( + ExecutionContract, + ExecutionStrategy, + HaltPolicy, + RunnerEnvironment, + SlurmRunnerConfig, +) @pytest.fixture @@ -122,6 +130,269 @@ def test_resolve_stage_execution_specs_assigns_default_strategies(tmp_configs): assert configured["vllm_stats"].instances == 16 +def test_runner_config_rejects_unknown_field_with_suggestion(tmp_configs) -> None: + _, runner_path, _ = tmp_configs + payload = yaml.safe_load(runner_path.read_text()) + payload["runner"]["slurm"]["partition_name"] = "gpu" + runner_path.write_text(yaml.safe_dump(payload)) + + with pytest.raises( + ValueError, + match="runner.slurm.partition_name; did you mean 'partition'", + ): + load_runner_config(runner_path) + + +def test_runner_config_rejects_unused_defaults_mapping(tmp_configs) -> None: + _, runner_path, _ = tmp_configs + payload = yaml.safe_load(runner_path.read_text()) + payload["runner"]["defaults"] = {"partition": "ignored"} + runner_path.write_text(yaml.safe_dump(payload)) + + with pytest.raises(ValueError, match=r"Unknown config field runner\.defaults"): + load_runner_config(runner_path) + + +def test_runner_config_preserves_legacy_partition_routing(tmp_configs) -> None: + _, runner_path, _ = tmp_configs + payload = yaml.safe_load(runner_path.read_text()) + payload["runner"]["slurm"].update( + { + "partition_interactive": "interactive", + "partition_batch": "batch", + "partition_cpu": "cpu", + "interactive_max_nodes": 2, + } + ) + runner_path.write_text(yaml.safe_dump(payload)) + + runner = load_runner_config(runner_path) + + assert runner.slurm is not None + assert runner.slurm.partition_for_nodes(1) == "interactive" + assert runner.slurm.partition_for_nodes(3) == "batch" + assert runner.slurm.partition_cpu == "cpu" + + +def test_runner_config_normalizes_multiple_eligible_partitions(tmp_configs) -> None: + _, runner_path, _ = tmp_configs + payload = yaml.safe_load(runner_path.read_text()) + payload["runner"]["slurm"]["partition"] = ["gpu-a", "gpu-b"] + runner_path.write_text(yaml.safe_dump(payload)) + + runner = load_runner_config(runner_path) + + assert runner.slurm is not None + assert runner.slurm.partition == "gpu-a,gpu-b" + + +def test_partition_set_changes_slurm_execution_contract_identity() -> None: + first = RunnerEnvironment( + kind="slurm", + contract=ExecutionContract(repository="/repo", venv="/venv"), + slurm=SlurmRunnerConfig(account="acct", partition=["gpu-a", "gpu-b"]), + ) + second = RunnerEnvironment( + kind="slurm", + contract=ExecutionContract(repository="/repo", venv="/venv"), + slurm=SlurmRunnerConfig(account="acct", partition=["gpu-a", "gpu-c"]), + ) + + assert execution_contract_hash(first) != execution_contract_hash(second) + + +def test_partition_schema_migration_changes_slurm_execution_contract_identity() -> None: + runner = RunnerEnvironment( + kind="slurm", + contract=ExecutionContract(repository="/repo", venv="/venv"), + slurm=SlurmRunnerConfig(account="acct", partition="batch"), + ) + legacy_identity = hash_payload( + { + "repository": "/repo", + "venv": "/venv", + "container": None, + "container_mounts": None, + "setup_env": None, + "prerun_commands": [], + "postrun_commands": [], + "runner_kind": "slurm", + "task_topology_contract": 1, + "slurm": { + "account": "acct", + "partition_interactive": None, + "partition_batch": "batch", + "partition_cpu": None, + "interactive_max_nodes": 2, + "max_nodes": None, + "time_limit": "4:00:00", + "qos": None, + }, + } + ) + + assert execution_contract_hash(runner) != legacy_identity + + +def test_runner_config_rejects_duplicate_partitions(tmp_configs) -> None: + _, runner_path, _ = tmp_configs + payload = yaml.safe_load(runner_path.read_text()) + payload["runner"]["slurm"]["partition"] = ["gpu", "gpu"] + runner_path.write_text(yaml.safe_dump(payload)) + + with pytest.raises( + ValueError, + match=r"runner\.slurm\.partition contains duplicate partition names", + ): + load_runner_config(runner_path) + + +def test_runner_config_rejects_partition_directive_injection(tmp_configs) -> None: + _, runner_path, _ = tmp_configs + payload = yaml.safe_load(runner_path.read_text()) + payload["runner"]["slurm"]["partition"] = "gpu\n#SBATCH --qos=unexpected" + runner_path.write_text(yaml.safe_dump(payload)) + + with pytest.raises( + ValueError, + match=r"runner\.slurm\.partition contains an invalid partition name", + ): + load_runner_config(runner_path) + + +@pytest.mark.parametrize("scope", ["defaults", "stage"]) +def test_execution_config_rejects_partition_directive_injection(tmp_configs, scope: str) -> None: + _, _, execution_path = tmp_configs + payload = yaml.safe_load(execution_path.read_text()) + if scope == "defaults": + payload["execution"]["defaults"]["partition"] = "gpu\n#SBATCH --qos=unexpected" + error_path = r"execution\.defaults\.partition" + else: + payload["execution"]["stages"]["vllm_stats"]["partition"] = "gpu\n#SBATCH --qos=unexpected" + error_path = r"execution\.stages\.vllm_stats\.partition" + execution_path.write_text(yaml.safe_dump(payload)) + + with pytest.raises(ValueError, match=rf"{error_path} contains an invalid partition name"): + load_execution_config(execution_path) + + +def test_runner_config_rejects_invalid_command_sequence(tmp_configs) -> None: + _, runner_path, _ = tmp_configs + payload = yaml.safe_load(runner_path.read_text()) + payload["runner"]["execution_contract"]["prerun_commands"] = ["module load cuda", 7] + runner_path.write_text(yaml.safe_dump(payload)) + + with pytest.raises(TypeError, match="prerun_commands must be a string or a sequence"): + load_runner_config(runner_path) + + +@pytest.mark.parametrize( + ("canonical", "legacy", "value"), + [ + ("container_mounts", "mounts", ["/host:/container"]), + ("prerun_commands", "prerun", ["module load cuda"]), + ("postrun_commands", "postrun", ["echo done"]), + ], +) +def test_runner_config_rejects_canonical_and_legacy_contract_fields( + tmp_configs, canonical: str, legacy: str, value: list[str] +) -> None: + _, runner_path, _ = tmp_configs + payload = yaml.safe_load(runner_path.read_text()) + contract = payload["runner"]["execution_contract"] + contract[canonical] = value + contract[legacy] = value + runner_path.write_text(yaml.safe_dump(payload)) + + with pytest.raises( + ValueError, + match=rf"cannot set both {canonical} and legacy {legacy}", + ): + load_runner_config(runner_path) + + +def test_execution_config_rejects_unknown_nested_field(tmp_configs) -> None: + _, _, execution_path = tmp_configs + payload = yaml.safe_load(execution_path.read_text()) + payload["execution"]["stages"]["width_importance"]["instance_count"] = 2 + execution_path.write_text(yaml.safe_dump(payload)) + + with pytest.raises( + ValueError, + match="width_importance.instance_count; did you mean 'instances'", + ): + load_execution_config(execution_path) + + +def test_execution_config_rejects_non_partition_final_report_fields(tmp_configs) -> None: + _, _, execution_path = tmp_configs + payload = yaml.safe_load(execution_path.read_text()) + payload["execution"]["stages"]["final_report"] = {"resource": "cpu"} + execution_path.write_text(yaml.safe_dump(payload)) + + with pytest.raises( + ValueError, + match=r"Unknown config field execution\.stages\.final_report\.resource", + ): + load_execution_config(execution_path) + + +def test_execution_config_rejects_fractional_instance_count(tmp_configs) -> None: + _, _, execution_path = tmp_configs + payload = yaml.safe_load(execution_path.read_text()) + payload["execution"]["stages"]["width_importance"]["instances"] = 1.5 + execution_path.write_text(yaml.safe_dump(payload)) + + with pytest.raises(TypeError, match="width_importance.instances must be a positive integer"): + load_execution_config(execution_path) + + +def test_execution_config_rejects_model_runtime_field_in_allocation_mesh( + tmp_configs, +) -> None: + _, _, execution_path = tmp_configs + payload = yaml.safe_load(execution_path.read_text()) + payload["execution"]["stages"]["width_importance"]["parallel"] = { + "tp": 1, + "sequence_parallel": True, + } + execution_path.write_text(yaml.safe_dump(payload)) + + with pytest.raises( + ValueError, + match="sequence_parallel belongs in the experiment model-parallel profile", + ): + load_execution_config(execution_path) + + +def test_compile_campaign_plan_rejects_unknown_execution_stage(tmp_configs) -> None: + experiment_path, runner_path, execution_path = tmp_configs + execution = load_execution_config(execution_path) + execution["stages"]["width_importnace"] = {"strategy": "single", "instances": 1} + + with pytest.raises(ValueError, match="width_importnace; did you mean 'width_importance'"): + compile_campaign_plan( + experiment_config_path=experiment_path, + runner=load_runner_config(runner_path), + execution=execution, + ) + + +def test_execution_config_rejects_single_strategy_with_multiple_instances( + tmp_configs, +) -> None: + experiment_path, runner_path, execution_path = tmp_configs + execution = load_execution_config(execution_path) + execution["stages"]["width_importance"]["instances"] = 2 + + with pytest.raises(ValueError, match="must be 1 for strategy 'single'"): + compile_campaign_plan( + experiment_config_path=experiment_path, + runner=load_runner_config(runner_path), + execution=execution, + ) + + def test_compile_campaign_plan_packs_vllm_stats_instances(tmp_configs): experiment_path, runner_path, execution_path = tmp_configs runner = load_runner_config(runner_path) @@ -212,15 +483,13 @@ def test_compile_campaign_plan_rejects_invalid_artifact_settling_timeout( ) -def test_compile_campaign_plan_uses_cpu_partition_without_gpus(tmp_configs): +def test_compile_campaign_plan_uses_stage_partition_list_without_gpus(tmp_configs): experiment_path, runner_path, execution_path = tmp_configs - runner_payload = yaml.safe_load(runner_path.read_text()) - runner_payload["runner"]["slurm"]["partition_cpu"] = "cpu" - runner_path.write_text(yaml.safe_dump(runner_payload)) execution_payload = yaml.safe_load(execution_path.read_text()) execution_payload["execution"]["stages"]["convert"] = { "strategy": "single", "resource": "cpu", + "partition": ["cpu-a", "cpu-b"], } execution_path.write_text(yaml.safe_dump(execution_payload)) @@ -232,12 +501,54 @@ def test_compile_campaign_plan_uses_cpu_partition_without_gpus(tmp_configs): convert = next(node for node in plan.stages if node.stage_id == "convert") assert convert.resource == "cpu" - assert convert.partition == "cpu" + assert convert.partition == "cpu-a,cpu-b" assert convert.gpus_per_instance == 0 assert convert.total_gpus == 0 assert convert.nodes == 1 +def test_compile_campaign_plan_migrates_legacy_cpu_partition(tmp_configs): + experiment_path, runner_path, execution_path = tmp_configs + runner_payload = yaml.safe_load(runner_path.read_text()) + runner_payload["runner"]["slurm"]["partition_cpu"] = "cpu" + runner_path.write_text(yaml.safe_dump(runner_payload)) + execution_payload = yaml.safe_load(execution_path.read_text()) + execution_payload["execution"]["stages"]["convert"] = { + "strategy": "single", + "resource": "cpu", + } + execution_path.write_text(yaml.safe_dump(execution_payload)) + + plan = compile_campaign_plan( + experiment_config_path=experiment_path, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + ) + convert = next(node for node in plan.stages if node.stage_id == "convert") + + assert convert.partition == "cpu" + assert plan.final_report_partition == "cpu" + + +def test_compile_campaign_plan_routes_final_report_to_eligible_cpu_partitions(tmp_configs): + experiment_path, runner_path, execution_path = tmp_configs + execution_payload = yaml.safe_load(execution_path.read_text()) + execution_payload["execution"]["stages"]["final_report"] = {"partition": ["cpu-a", "cpu-b"]} + execution_path.write_text(yaml.safe_dump(execution_payload)) + + plan = compile_campaign_plan( + experiment_config_path=experiment_path, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + ) + + assert plan.final_report_partition == "cpu-a,cpu-b" + assert plan_to_dict(plan)["final_report"] == { + "resource": "cpu", + "partition": "cpu-a,cpu-b", + } + + def test_post_mip_compiler_topologically_orders_serialized_nodes() -> None: config = { "mip": {"runs": {"memory": {}}}, diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index 617b53fb66f..5fa6fad0ca6 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -16,9 +16,11 @@ """Tests for orchestration executors.""" import subprocess +import sys import time from pathlib import Path +import pytest import yaml import puzzletron_orchestrator.adapters.sharded as sharded_module @@ -31,6 +33,7 @@ from puzzletron_orchestrator.executors.baremetal import BareMetalSSHExecutor from puzzletron_orchestrator.executors.local import LocalExecutor from puzzletron_orchestrator.executors.slurm import SlurmExecutor, render_sbatch_script +from puzzletron_orchestrator.process import run_argv from puzzletron_orchestrator.schema import ( AttemptSpec, BareMetalHost, @@ -47,10 +50,22 @@ SlurmRunnerConfig, StagePlanNode, TaskTopology, + WorkItem, WorkPlan, ) +def test_run_argv_captures_text_without_a_shell(tmp_path: Path): + result = run_argv( + (sys.executable, "-c", "print('captured')"), + cwd=tmp_path, + ) + + assert result.returncode == 0 + assert result.stdout == "captured\n" + assert result.stderr == "" + + def test_local_executor_runs_successful_command(tmp_path: Path): executor = LocalExecutor() log_path = tmp_path / "ok.log" @@ -141,9 +156,7 @@ def test_render_sbatch_script_requests_gpus_per_node(): ), slurm=SlurmRunnerConfig( account="acct", - partition="batch", - partition_interactive="interactive", - partition_batch="batch", + partition=["gpu-a", "gpu-b"], ), ) attempt = AttemptSpec( @@ -159,7 +172,7 @@ def test_render_sbatch_script_requests_gpus_per_node(): script = render_sbatch_script( attempt=attempt, runner=runner, - partition=runner.slurm.partition_for_nodes(attempt.allocation_nodes), + partition=runner.slurm.partition, account="acct", time_limit="4:00:00", qos=None, @@ -167,7 +180,7 @@ def test_render_sbatch_script_requests_gpus_per_node(): ) assert "#SBATCH --gpus-per-node=8" in script assert "#SBATCH --nodes=2" in script - assert "#SBATCH --partition=interactive" in script + assert "#SBATCH --partition=gpu-a,gpu-b" in script assert "source /site/setup-envs.sh" in script assert script.startswith("#!/bin/bash\n") @@ -176,7 +189,7 @@ def test_render_sbatch_script_omits_gpu_requests_for_cpu_stage(): runner = RunnerEnvironment( kind="slurm", contract=ExecutionContract(repository="/repo", venv="/repo/.venv"), - slurm=SlurmRunnerConfig(account="acct", partition_cpu="cpu"), + slurm=SlurmRunnerConfig(account="acct", partition="cpu"), ) attempt = AttemptSpec( attempt_id="a1", @@ -206,9 +219,137 @@ def test_render_sbatch_script_omits_gpu_requests_for_cpu_stage(): assert "--gpu-bind" not in srun -def test_vllm_aggregation_uses_slurm_execution_contract(tmp_path: Path, monkeypatch): +@pytest.mark.parametrize("configured_log_dir", [False, True], ids=("fallback", "configured")) +def test_work_adapters_use_effective_campaign_log_dir(tmp_path: Path, configured_log_dir: bool): + """Every adapter must use the plan's effective shared log directory.""" + + custom_log_dir = tmp_path / "shared-logs" if configured_log_dir else None + runner = RunnerEnvironment( + kind="slurm", + contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), + slurm=SlurmRunnerConfig( + account="acct", + partition="batch", + log_dir=str(custom_log_dir) if custom_log_dir else None, + ), + ) + + def stage(stage_id: str, strategy: ExecutionStrategy) -> StagePlanNode: + return StagePlanNode( + stage_id=stage_id, + strategy=strategy, + instances=1, + failure_policy=FailurePolicy.STRICT, + mesh={}, + gpus_per_instance=1, + gpus_per_node=8, + nodes=1, + total_gpus=1, + exclusive=False, + parents=(), + distributed=False, + partition="batch", + ) + + nodes = ( + stage("convert", ExecutionStrategy.SINGLE), + stage("depth_importance", ExecutionStrategy.PERSISTENT_POOL), + stage("vllm_stats", ExecutionStrategy.SHARDED), + stage("post.profile.online_eval", ExecutionStrategy.SHARDED), + ) + plan = CampaignPlan( + experiment_config_path=str(tmp_path / "experiment.yaml"), + puzzle_dir=tmp_path / "run", + experiment_config={ + "puzzle_dir": str(tmp_path / "run"), + "depth_importance": {"output_dir": str(tmp_path / "depth")}, + "post_mip": {"flows": {"profile": {"nodes": {"online_eval": {"type": "evaluation"}}}}}, + }, + runner=runner, + execution_defaults={"gpus_per_node": 8}, + stages=nodes, + contract_hash="contract", + ) + items = ( + WorkItem("convert:0", "convert", 0, 1, 1), + WorkItem( + "depth_importance:gang", + "depth_importance", + 0, + 1, + 1, + metadata={"role": "gang", "worker_count": 1}, + ), + WorkItem("vllm_stats:0", "vllm_stats", 0, 1, 1), + WorkItem( + "post.profile.online_eval:0", + "post.profile.online_eval", + 0, + 1, + 1, + metadata={"logical_shard_count": 1}, + ), + ) + expected_log_dir = custom_log_dir or plan.puzzle_dir / "logs" + + for node, item in zip(nodes, items, strict=True): + attempt = adapter_for_stage(node).command( + plan=plan, + node=node, + item=item, + attempt_id="a1", + runner=runner, + ) + assert attempt.command.log_path is not None + assert Path(attempt.command.log_path).parent == expected_log_dir + + +@pytest.mark.parametrize( + ("slurm_kwargs", "node_partition", "expected_partition", "configured_log_dir"), + [ + ({"partition": "runner-a"}, "reserved-a,reserved-b", "reserved-a,reserved-b", False), + ({"partition": "runner-a"}, "reserved-a,reserved-b", "reserved-a,reserved-b", True), + ( + { + "partition_interactive": "interactive", + "partition_batch": "batch", + "partition_cpu": "cpu", + }, + None, + "cpu", + False, + ), + ( + { + "partition_interactive": "interactive", + "partition_batch": "batch", + "partition_cpu": "cpu", + }, + None, + "cpu", + True, + ), + ], + ids=( + "stage-override-fallback-log", + "stage-override-configured-log", + "legacy-cpu-fallback-log", + "legacy-cpu-configured-log", + ), +) +def test_vllm_aggregation_uses_slurm_execution_contract( + tmp_path: Path, + monkeypatch, + slurm_kwargs, + node_partition, + expected_partition, + configured_log_dir, +): """Controller-side merges must run in the same container/venv as workers.""" + slurm_kwargs = dict(slurm_kwargs) + if configured_log_dir: + slurm_kwargs["log_dir"] = str(tmp_path / "shared-logs") runner = RunnerEnvironment( kind="slurm", contract=ExecutionContract( @@ -216,10 +357,7 @@ def test_vllm_aggregation_uses_slurm_execution_contract(tmp_path: Path, monkeypa venv=".venv-worker", container="/images/pytorch.sqsh", ), - slurm=SlurmRunnerConfig( - account="acct", - partition_cpu="cpu", - ), + slurm=SlurmRunnerConfig(account="acct", **slurm_kwargs), ) node = StagePlanNode( stage_id="vllm_stats", @@ -234,6 +372,7 @@ def test_vllm_aggregation_uses_slurm_execution_contract(tmp_path: Path, monkeypa exclusive=False, parents=("convert",), distributed=False, + partition=node_partition, ) plan = CampaignPlan( experiment_config_path=str(tmp_path / "experiment.yaml"), @@ -275,7 +414,7 @@ def poll(self, handles): def fail_local_subprocess(*_args, **_kwargs): raise AssertionError("aggregation escaped the Slurm execution contract") - monkeypatch.setattr(sharded_module.subprocess, "run", fail_local_subprocess) + monkeypatch.setattr(sharded_module, "run_argv", fail_local_subprocess) adapter = adapter_for_stage(node) result = adapter.aggregate( plan=plan, @@ -293,7 +432,10 @@ def fail_local_subprocess(*_args, **_kwargs): attempt = submitted[0] assert attempt.allocation_gpus == 0 assert attempt.metadata["gpus_per_node"] == 0 - assert attempt.metadata["partition"] == "cpu" + assert attempt.metadata["partition"] == expected_partition + expected_log_dir = tmp_path / ("shared-logs" if configured_log_dir else "run/logs") + assert attempt.command.log_path is not None + assert Path(attempt.command.log_path).parent == expected_log_dir assert attempt.command.argv[-1] == "--merge" assert result.summary["merge_handle"] == "slurm-123" @@ -309,9 +451,7 @@ def test_render_sbatch_script_never_requests_exclusive(): ), slurm=SlurmRunnerConfig( account="acct", - partition="batch", - partition_interactive="interactive", - partition_batch="batch", + partition="gpu", ), ) attempt = AttemptSpec( @@ -342,15 +482,45 @@ def test_render_sbatch_script_never_requests_exclusive(): assert "tee -a" not in script -def test_partition_for_nodes_prefers_batch_above_interactive_cap(): - slurm = SlurmRunnerConfig( - account="acct", - partition_interactive="interactive", - partition_batch="batch", - interactive_max_nodes=2, +@pytest.mark.parametrize( + ("runner_partition", "metadata", "expected_directive"), + [ + (["runner-a", "runner-b"], {}, "#SBATCH --partition=runner-a,runner-b"), + (["runner-a"], {"partition": "stage-a,stage-b"}, "#SBATCH --partition=stage-a,stage-b"), + (None, {}, None), + ], +) +def test_slurm_submit_resolves_partition_precedence( + tmp_path: Path, + monkeypatch, + runner_partition, + metadata, + expected_directive, +): + monkeypatch.setattr( + "puzzletron_orchestrator.executors.slurm._run_command", + lambda argv: subprocess.CompletedProcess(argv, 0, stdout="12345\n", stderr=""), + ) + runner = RunnerEnvironment( + kind="slurm", + contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), + slurm=SlurmRunnerConfig(account="acct", partition=runner_partition), ) - assert slurm.partition_for_nodes(2) == "interactive" - assert slurm.partition_for_nodes(3) == "batch" + attempt = AttemptSpec( + attempt_id="a1", + work_id="mip:0", + stage_id="mip", + command=CommandSpec(argv=("python", "worker.py")), + metadata=metadata, + ) + + SlurmExecutor(runner, scripts_dir=tmp_path / "sbatch").submit(attempt) + + script = (tmp_path / "sbatch" / "mip_a1.sh").read_text() + if expected_directive is None: + assert "#SBATCH --partition=" not in script + else: + assert expected_directive in script def test_slurm_submit_retries_transient_controller_timeout(tmp_path: Path, monkeypatch): @@ -385,7 +555,7 @@ class _Result: runner = RunnerEnvironment( kind="slurm", contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), - slurm=SlurmRunnerConfig(account="acct", partition="interactive"), + slurm=SlurmRunnerConfig(account="acct", partition="gpu"), ) attempt = AttemptSpec( attempt_id="abcdef12-3456", @@ -441,7 +611,7 @@ class _Result: runner = RunnerEnvironment( kind="slurm", contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), - slurm=SlurmRunnerConfig(account="acct", partition="interactive"), + slurm=SlurmRunnerConfig(account="acct", partition="gpu"), ) attempt = AttemptSpec( attempt_id="abcdef12-3456", @@ -461,7 +631,7 @@ def _slurm_executor(tmp_path: Path) -> SlurmExecutor: RunnerEnvironment( kind="slurm", contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), - slurm=SlurmRunnerConfig(account="acct", partition="interactive"), + slurm=SlurmRunnerConfig(account="acct", partition="gpu"), ), scripts_dir=tmp_path / "sbatch", ) @@ -561,7 +731,7 @@ def test_depth_pool_uses_one_four_node_gang_allocation(tmp_path: Path): runner = RunnerEnvironment( kind="slurm", contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), - slurm=SlurmRunnerConfig(account="acct", partition_batch="batch"), + slurm=SlurmRunnerConfig(account="acct", partition="batch"), ) node = StagePlanNode( stage_id="depth_importance", @@ -623,7 +793,7 @@ def test_depth_pool_packs_four_two_gpu_workers_per_node(tmp_path: Path): runner = RunnerEnvironment( kind="slurm", contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), - slurm=SlurmRunnerConfig(account="acct", partition_batch="batch"), + slurm=SlurmRunnerConfig(account="acct", partition="batch"), ) node = StagePlanNode( stage_id="depth_importance", @@ -684,7 +854,7 @@ def test_post_mip_workers_share_one_packed_allocation(tmp_path: Path): runner = RunnerEnvironment( kind="slurm", contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), - slurm=SlurmRunnerConfig(account="acct", partition_batch="batch"), + slurm=SlurmRunnerConfig(account="acct", partition="batch"), ) node = StagePlanNode( stage_id="post.profile.online_eval", @@ -737,7 +907,7 @@ def test_replacement_pool_uses_one_four_node_gang_allocation(tmp_path: Path): runner = RunnerEnvironment( kind="slurm", contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), - slurm=SlurmRunnerConfig(account="acct", partition_batch="batch"), + slurm=SlurmRunnerConfig(account="acct", partition="batch"), ) node = StagePlanNode( stage_id="replacement_scoring", @@ -788,7 +958,7 @@ def _replacement_width_attempts( runner = RunnerEnvironment( kind="slurm", contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), - slurm=SlurmRunnerConfig(account="acct", partition_batch="batch"), + slurm=SlurmRunnerConfig(account="acct", partition="batch"), ) node = StagePlanNode( stage_id="replacement_scoring", @@ -928,7 +1098,7 @@ def test_replacement_pool_completion_identity_changes_with_embedding_widths(tmp_ ] == ["4", "4", "4", "4"] -def test_stage_partition_override_forces_batch(tmp_path: Path): +def test_stage_partition_override_uses_eligible_partition_list(tmp_path: Path): experiment = tmp_path / "experiment.yaml" experiment.write_text( yaml.safe_dump( @@ -956,9 +1126,7 @@ def test_stage_partition_override_forces_batch(tmp_path: Path): "kind": "slurm", "slurm": { "account": "test", - "partition_interactive": "interactive", - "partition_batch": "batch", - "interactive_max_nodes": 2, + "partition": ["gpu-a", "gpu-b"], }, "execution_contract": { "repository": str(tmp_path), @@ -978,7 +1146,7 @@ def test_stage_partition_override_forces_batch(tmp_path: Path): "vllm_stats": { "strategy": "sharded", "instances": 4, - "partition": "batch", + "partition": ["reserved-a", "reserved-b"], }, }, } @@ -991,8 +1159,8 @@ def test_stage_partition_override_forces_batch(tmp_path: Path): execution=load_execution_config(execution), ) node = next(item for item in plan.stages if item.stage_id == "vllm_stats") - assert node.partition == "batch" - assert plan.runner.slurm.partition_for_nodes(1) == "interactive" + assert node.partition == "reserved-a,reserved-b" + assert plan.runner.slurm.partition == "gpu-a,gpu-b" adapter = adapter_for_stage(node) attempt = adapter.command( plan=plan, @@ -1001,7 +1169,7 @@ def test_stage_partition_override_forces_batch(tmp_path: Path): attempt_id="a1", runner=plan.runner, ) - assert attempt.metadata["partition"] == "batch" + assert attempt.metadata["partition"] == "reserved-a,reserved-b" def _baremetal_runner() -> RunnerEnvironment: diff --git a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py index edf9bd78192..8d757d93c5e 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py +++ b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py @@ -235,7 +235,7 @@ def test_load_experiment_config_matches_hydra_scientific_number_semantics( """ ) - config = load_experiment_config(experiment, overrides=["+threshold=1e-4"]) + config = load_experiment_config(experiment, overrides=["++threshold=1e-4"]) assert config["bypass"]["best_val_loss"] == 1e9 assert config["bypass"]["training"] == { @@ -274,13 +274,7 @@ def test_load_experiment_config_distinguishes_hydra_addition_modes( experiment = tmp_path / "experiment.yaml" experiment.write_text("value: 1\n") - added = load_experiment_config(experiment, overrides=["+added.value=2"]) - with pytest.raises(ValueError, match="^Addition override already exists"): - load_experiment_config(experiment, overrides=["+value=2"]) - with pytest.raises(ValueError, match="^Override key does not exist"): - load_experiment_config(experiment, overrides=["missing=2"]) - with pytest.raises(ValueError, match="^Override path does not exist"): - load_experiment_config(experiment, overrides=["missing.value=2"]) + added = load_experiment_config(experiment, overrides=["added.value=2"]) replaced = load_experiment_config(experiment, overrides=["++value=2"]) created = load_experiment_config(experiment, overrides=["++created.value=3"]) @@ -289,6 +283,129 @@ def test_load_experiment_config_distinguishes_hydra_addition_modes( assert created["created"] == {"value": 3} +@pytest.mark.parametrize("override", ["+experiment.dir=other", "~experiment.dir"]) +def test_load_experiment_config_rejects_unsupported_hydra_operators( + tmp_path: Path, + override: str, +) -> None: + experiment = tmp_path / "experiment.yaml" + experiment.write_text("experiment:\n dir: run\n") + + with pytest.raises(ValueError, match="not supported by the dependency-light controller"): + load_experiment_config(experiment, overrides=[override]) + + +def test_load_experiment_config_rejects_unknown_interpolation(tmp_path: Path) -> None: + experiment = tmp_path / "experiment.yaml" + experiment.write_text("experiment:\n dir: ${missing.value}\n") + + with pytest.raises(ValueError, match="Unknown config interpolation 'missing.value'"): + load_experiment_config(experiment) + + +def test_orchestrator_cli_reports_config_errors_without_traceback(tmp_path: Path) -> None: + runner = tmp_path / "runner.yaml" + runner.write_text( + yaml.safe_dump( + { + "runner": { + "kind": "slurm", + "slurm": {"account": "test", "partition_name": "gpu"}, + } + } + ) + ) + + result = subprocess.run( + [ + sys.executable, + "examples/puzzletron/orchestrate.py", + "--experiment", + str(tmp_path / "experiment.yaml"), + "--runner", + str(runner), + "--execution", + str(tmp_path / "execution.yaml"), + "--dry-run", + ], + cwd=REPOSITORY_ROOT, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 2 + assert "cannot build campaign plan" in result.stderr + assert "partition" in result.stderr + assert "Traceback" not in result.stderr + + +def test_orchestrator_cli_reports_dry_run_adapter_errors_without_traceback( + tmp_path: Path, +) -> None: + experiment = tmp_path / "experiment.yaml" + runner = tmp_path / "runner.yaml" + execution = tmp_path / "execution.yaml" + experiment.write_text( + yaml.safe_dump( + { + "experiment": {"dir": str(tmp_path / "run")}, + "embedding_pruning": {"enabled": True, "widths": []}, + "replacement_scoring": {"enabled": True}, + } + ) + ) + runner.write_text( + yaml.safe_dump( + { + "runner": { + "kind": "slurm", + "slurm": {"account": "test", "partition": "gpu"}, + } + } + ) + ) + execution.write_text( + yaml.safe_dump( + { + "execution": { + "defaults": {"gpus_per_node": 1}, + "stages": { + "replacement_scoring": { + "strategy": "persistent_pool", + "instances": 1, + } + }, + } + } + ) + ) + + result = subprocess.run( + [ + sys.executable, + "examples/puzzletron/orchestrate.py", + "--experiment", + str(experiment), + "--runner", + str(runner), + "--execution", + str(execution), + "--stage", + "replacement_scoring", + "--dry-run", + ], + cwd=REPOSITORY_ROOT, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 2 + assert "embedding replacement scoring requires at least one width" in result.stderr + assert "Traceback" not in result.stderr + + def test_convert_completeness_requires_runtime_subblock_library( tmp_path: Path, write_terminal_manifest ) -> None: diff --git a/tests/unit/torch/puzzletron/test_orchestration_reporting.py b/tests/unit/torch/puzzletron/test_orchestration_reporting.py index 5b5246067fb..d413b6a9072 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_reporting.py +++ b/tests/unit/torch/puzzletron/test_orchestration_reporting.py @@ -1,12 +1,27 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. """Tests for the runner-backed orchestration report finalizer.""" from __future__ import annotations -from collections.abc import Sequence from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Sequence from puzzletron_orchestrator.controller import CampaignController from puzzletron_orchestrator.executors.base import Executor @@ -27,7 +42,13 @@ ) -def _plan(tmp_path: Path, *, partition_cpu: str | None = "cpu") -> CampaignPlan: +def _plan( + tmp_path: Path, + *, + partition: str | None = "shared", + final_report_partition: str | None = None, + log_dir: Path | None = None, +) -> CampaignPlan: return CampaignPlan( experiment_config_path=str(tmp_path / "experiment.yaml"), puzzle_dir=tmp_path / "run", @@ -43,19 +64,19 @@ def _plan(tmp_path: Path, *, partition_cpu: str | None = "cpu") -> CampaignPlan: ), slurm=SlurmRunnerConfig( account="test", - partition_interactive="interactive", - partition_batch="batch", - partition_cpu=partition_cpu, + partition=partition, + log_dir=str(log_dir) if log_dir is not None else None, ), ), execution_defaults={}, stages=(), contract_hash="contract", + final_report_partition=final_report_partition, ) -def test_build_final_report_attempt_uses_runner_cpu_task(tmp_path: Path): - plan = _plan(tmp_path) +def test_build_final_report_attempt_uses_configured_cpu_partition(tmp_path: Path): + plan = _plan(tmp_path, final_report_partition="cpu-a,cpu-b") attempt = build_final_report_attempt(plan, attempt_id="report-attempt") @@ -77,20 +98,29 @@ def test_build_final_report_attempt_uses_runner_cpu_task(tmp_path: Path): assert attempt.allocation_nodes == 1 assert attempt.allocation_gpus == 0 assert attempt.exclusive is False - assert attempt.metadata == {"gpus_per_node": 0, "partition": "cpu"} + assert attempt.metadata == {"gpus_per_node": 0, "partition": "cpu-a,cpu-b"} assert attempt.task_topology.task_count == 1 assert attempt.task_topology.gpus_per_task == 0 assert attempt.task_topology.launcher is TaskLauncher.DIRECT -def test_build_final_report_attempt_uses_normal_partition_fallback(tmp_path: Path): - plan = _plan(tmp_path, partition_cpu=None) - - attempt = build_final_report_attempt(plan, attempt_id="report-attempt") +def test_build_final_report_attempt_uses_runner_default_when_unconfigured(tmp_path: Path): + attempt = build_final_report_attempt(_plan(tmp_path), attempt_id="report-attempt") assert attempt.metadata == {"gpus_per_node": 0} +def test_build_final_report_attempt_uses_configured_log_directory(tmp_path: Path): + log_dir = tmp_path / "shared-logs" + + attempt = build_final_report_attempt( + _plan(tmp_path, log_dir=log_dir), + attempt_id="report-attempt", + ) + + assert attempt.command.log_path == str(log_dir / "final_report_report-attempt.log") + + class _ReportExecutor(Executor): backend = "fake" diff --git a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py index 17704f2ad23..1cc8f633ed3 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py +++ b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py @@ -660,7 +660,7 @@ def poll(self, handles): def test_controller_rejects_overrides_that_differ_from_compiled_plan(tmp_path: Path): experiment, runner_path, execution_path = _write_configs(tmp_path) - compiled_overrides = ["+convert.model_path=/models/compiled"] + compiled_overrides = ["++convert.model_path=/models/compiled"] plan = compile_campaign_plan( experiment_config_path=experiment, runner=load_runner_config(runner_path), @@ -678,7 +678,7 @@ def test_controller_rejects_overrides_that_differ_from_compiled_plan(tmp_path: P ) with pytest.raises(ValueError, match="must match the overrides compiled"): - controller.run(overrides=["+convert.model_path=/models/runtime"], once=True) + controller.run(overrides=["++convert.model_path=/models/runtime"], once=True) assert executor.submitted_stage_ids == [] assert controller.store.list_attempts("convert") == [] diff --git a/tests/unit/torch/puzzletron/test_portable_configs.py b/tests/unit/torch/puzzletron/test_portable_configs.py index 6fc6552dba4..e0cb8f01baa 100644 --- a/tests/unit/torch/puzzletron/test_portable_configs.py +++ b/tests/unit/torch/puzzletron/test_portable_configs.py @@ -18,6 +18,7 @@ import re from pathlib import Path +import pytest import yaml from puzzletron_orchestrator.compiler import load_execution_config, load_runner_config @@ -25,6 +26,12 @@ from puzzletron_setup.v2.defaults import load_defaults REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +SLURM_RUNNER_CONFIGS = ( + "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", + "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", + "examples/puzzletron/configs/orchestration/qwen3p5_0p8b/runner.slurm.yaml", +) +NAMED_SLURM_RUNNER_CONFIGS = SLURM_RUNNER_CONFIGS[1:] NEMOTRON3_NANO_30B_MODEL_CONFIG = ( "examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml" ) @@ -36,9 +43,8 @@ def test_slurm_runner_example_is_portable() -> None: - slurm = load_runner_config( - REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml" - ) + path = REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml" + slurm = load_runner_config(path) assert slurm.contract.repository == WORKER_REPOSITORY_PLACEHOLDER assert slurm.contract.venv == WORKER_VENV_PLACEHOLDER assert slurm.contract.container is None @@ -46,7 +52,10 @@ def test_slurm_runner_example_is_portable() -> None: assert not slurm.contract.prerun_commands assert slurm.slurm is not None assert slurm.slurm.account.startswith("REPLACE_WITH_") - assert slurm.slurm.partition_cpu is None + assert slurm.slurm.partition == ( + "REPLACE_WITH_PRIMARY_SLURM_PARTITION,REPLACE_WITH_ALTERNATE_SLURM_PARTITION" + ) + assert slurm.slurm.log_dir == "puzzle_runs/logs" def test_baremetal_runner_example_is_portable() -> None: @@ -64,9 +73,8 @@ def test_baremetal_runner_example_is_portable() -> None: def test_qwen_slurm_runner_preserves_portable_environment_contract() -> None: - runner = load_runner_config( - REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml" - ) + path = REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml" + runner = load_runner_config(path) contract_values = ( runner.contract.repository, @@ -83,7 +91,24 @@ def test_qwen_slurm_runner_preserves_portable_environment_contract() -> None: assert all("REPLACE_WITH_" in command for command in runner.contract.prerun_commands) assert runner.slurm is not None assert runner.slurm.account.startswith("REPLACE_WITH_") - assert runner.slurm.partition_cpu is None + assert runner.slurm.partition.startswith("REPLACE_WITH_") + + +@pytest.mark.parametrize("relative_path", SLURM_RUNNER_CONFIGS) +def test_checked_in_slurm_runners_only_emit_generic_partition( + relative_path: str, +) -> None: + payload = yaml.safe_load((REPOSITORY_ROOT / relative_path).read_text()) + + assert "partition" in payload["runner"]["slurm"] + assert not any(key.startswith("partition_") for key in payload["runner"]["slurm"]) + + +@pytest.mark.parametrize("relative_path", NAMED_SLURM_RUNNER_CONFIGS) +def test_named_slurm_runners_keep_logs_below_the_campaign_root(relative_path: str) -> None: + payload = yaml.safe_load((REPOSITORY_ROOT / relative_path).read_text()) + + assert "log_dir" not in payload["runner"]["slurm"] def test_execution_example_is_loadable() -> None: @@ -108,7 +133,7 @@ def test_setup_defaults_example_is_portable() -> None: slurm = defaults["infrastructure"]["runner"]["slurm"] assert "account" not in slurm - assert slurm["partition_cpu"] is None + assert slurm["partition"] is None def test_model_examples_use_public_hugging_face_identities() -> None: diff --git a/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_smoke_plan.py b/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_smoke_plan.py index 937eedbb312..be42243aa70 100644 --- a/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_smoke_plan.py +++ b/tests/unit/torch/puzzletron/test_qwen3p5_0p8b_smoke_plan.py @@ -151,7 +151,6 @@ def test_qwen3p5_0p8b_runner_requires_an_explicit_site_contract() -> None: assert runner.slurm.time_limit == "1:00:00" assert runner.slurm.account.startswith("REPLACE_WITH_") assert runner.slurm.partition.startswith("REPLACE_WITH_") - assert runner.slurm.partition_batch.startswith("REPLACE_WITH_") assert runner.contract.repository.startswith("REPLACE_WITH_") assert runner.contract.venv.startswith("REPLACE_WITH_") assert runner.contract.container is not None diff --git a/tests/unit/torch/puzzletron/test_setup_bundle.py b/tests/unit/torch/puzzletron/test_setup_bundle.py index 5f4f264ea65..219cc031cb5 100644 --- a/tests/unit/torch/puzzletron/test_setup_bundle.py +++ b/tests/unit/torch/puzzletron/test_setup_bundle.py @@ -13,9 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Tests for scheduler-neutral configs emitted by the Puzzletron setup wizard.""" import pytest @@ -39,6 +36,36 @@ ) +def test_resume_preserves_legacy_partition_routing(tmp_path) -> None: + state = AnswerState.start(tmp_path / "campaign", detailed=False) + state.record_many( + "infrastructure", + { + "runner": { + "kind": "slurm", + "slurm": { + "partition_interactive": "interactive", + "partition_batch": "batch", + "partition_cpu": "cpu", + "interactive_max_nodes": 2, + }, + }, + "execution_contract": {"repository": "/repo", "venv": "/repo/.venv"}, + "gpus_per_node": 8, + "workers": {"pool": 1, "sharded": 1}, + "meshes": {"common": {}, "bypass": {}, "global_kd": {}}, + }, + ) + + resumed = AnswerState.resume(state.path) + runner = render_runner(resumed.payload, "production") + execution = render_execution(resumed.payload, {}, "production") + + assert runner["runner"]["slurm"]["partition_interactive"] == "interactive" + assert runner["runner"]["slurm"]["partition_batch"] == "batch" + assert execution["execution"]["stages"]["convert"]["partition"] == "cpu" + + class _NormalMipPrompts: def __init__(self) -> None: self.messages: list[str] = [] @@ -114,7 +141,6 @@ def test_serving_parallel_treats_vllm_expert_parallelism_as_boolean_mode() -> No "ep": 1, "dp_shard": 1, "dp_replicate": 4, - "sequence_parallel": False, } topology["expert_parallel_size"] = 4 @@ -460,7 +486,6 @@ def test_render_execution_uses_vllm_runtime_topology() -> None: "ep": 1, "dp_shard": 1, "dp_replicate": 2, - "sequence_parallel": False, } assert execution["execution"]["stages"]["bypass"]["strategy"] == "single" assert execution["execution"]["stages"]["bypass"]["instances"] == 1 @@ -472,7 +497,6 @@ def test_render_execution_uses_vllm_runtime_topology() -> None: "ep": 1, "dp_shard": 1, "dp_replicate": 1, - "sequence_parallel": False, } @@ -525,10 +549,7 @@ def test_render_execution_uses_common_mesh_for_post_mip_evaluation_only() -> Non execution = render_execution(state, experiment, "production")["execution"]["stages"] - assert execution["post.run.eval"]["parallel"] == { - **common, - "sequence_parallel": False, - } + assert execution["post.run.eval"]["parallel"] == common assert execution["post.run.serve"]["parallel"] == { "tp": 2, "cp": 1, @@ -536,7 +557,6 @@ def test_render_execution_uses_common_mesh_for_post_mip_evaluation_only() -> Non "ep": 1, "dp_shard": 1, "dp_replicate": 1, - "sequence_parallel": False, } assert execution["post.run.materialized"]["parallel"] == { "tp": 1, @@ -545,7 +565,6 @@ def test_render_execution_uses_common_mesh_for_post_mip_evaluation_only() -> Non "ep": 1, "dp_shard": 1, "dp_replicate": 1, - "sequence_parallel": False, } assert execution["post.run.materialized"]["instances"] == 1 @@ -622,7 +641,6 @@ def test_render_execution_uses_vllm_mesh_for_post_mip_downstream_evaluation() -> "ep": 1, "dp_shard": 1, "dp_replicate": 1, - "sequence_parallel": False, } @@ -778,7 +796,7 @@ def test_render_execution_ignores_legacy_aiperf_worker_override() -> None: assert stages["post.run.serving"]["instances"] == 4 -def test_render_execution_uses_cpu_partition_for_io_bound_stages() -> None: +def test_render_execution_marks_io_bound_stages_as_cpu_resources() -> None: state = { "answers": { "infrastructure": { @@ -786,8 +804,7 @@ def test_render_execution_uses_cpu_partition_for_io_bound_stages() -> None: "kind": "slurm", "slurm": { "account": "acct", - "partition_batch": "batch", - "partition_cpu": "cpu", + "partition": ["cluster-a", "cluster-b"], }, }, "execution_contract": {"repository": "/repo", "venv": "/repo/.venv"}, @@ -816,7 +833,7 @@ def test_render_execution_uses_cpu_partition_for_io_bound_stages() -> None: runner = render_runner(state, "production") stages = render_execution(state, experiment, "production")["execution"]["stages"] - assert runner["runner"]["slurm"]["partition_cpu"] == "cpu" + assert runner["runner"]["slurm"]["partition"] == ["cluster-a", "cluster-b"] for stage_id in ( "convert", "tokenize_data", @@ -826,7 +843,7 @@ def test_render_execution_uses_cpu_partition_for_io_bound_stages() -> None: "post.run.materialized", ): assert stages[stage_id]["resource"] == "cpu" - assert stages[stage_id]["partition"] == "cpu" + assert "partition" not in stages[stage_id] assert stages[stage_id]["instances"] == 1 assert "resource" not in stages["sort"] assert "partition" not in stages["sort"] @@ -838,14 +855,13 @@ def test_render_execution_uses_cpu_partition_for_io_bound_stages() -> None: "ep": 1, "dp_shard": 1, "dp_replicate": 1, - "sequence_parallel": False, } assert stages["post.run.eval"]["instances"] == 8 assert stages["post.run.serving"]["instances"] == 8 assert stages["post.run.short_kd"]["instances"] == 8 -def test_render_execution_falls_back_to_one_gpu_for_cpu_stages() -> None: +def test_render_execution_uses_cpu_resources_without_partition_taxonomy() -> None: state = { "answers": { "infrastructure": { @@ -853,8 +869,7 @@ def test_render_execution_falls_back_to_one_gpu_for_cpu_stages() -> None: "kind": "slurm", "slurm": { "account": "acct", - "partition_batch": "batch", - "partition_cpu": None, + "partition": None, }, }, "gpus_per_node": 8, @@ -877,8 +892,8 @@ def test_render_execution_falls_back_to_one_gpu_for_cpu_stages() -> None: stages = render_execution(state, experiment, "production")["execution"]["stages"] - assert stages["mip"].get("resource", "gpu") == "gpu" - assert stages["post.run.materialized"].get("resource", "gpu") == "gpu" + assert stages["mip"]["resource"] == "cpu" + assert stages["post.run.materialized"]["resource"] == "cpu" assert stages["post.run.materialized"]["instances"] == 1 diff --git a/tests/unit/torch/puzzletron/test_setup_v2_quick.py b/tests/unit/torch/puzzletron/test_setup_v2_quick.py index 55819c319eb..9551d8673d8 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_quick.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_quick.py @@ -979,6 +979,67 @@ def test_guided_infrastructure_prompts_for_unresolved_worker_paths(tmp_path): assert backend.remaining == 0 +def test_customize_partition_prompt_renders_list_default_as_comma_separated(tmp_path): + state = WizardState.start(tmp_path / "campaign", defaults_path=None) + resolver = DefaultsResolver( + file_defaults={ + "infrastructure": { + "runner": {"slurm": {"partition": ["gpu-a", "gpu-b"]}}, + } + } + ) + + class PartitionDefaultBackend(ScriptedBackend): + partition_default = None + + def text(self, message: str, default: str) -> str: + if message.startswith("Eligible Slurm partitions"): + self.partition_default = default + return default + return super().text(message, default) + + backend = PartitionDefaultBackend( + [ + "customize", + "/worker/modelopt", + "/worker/venv", + "", + "", + "acct", + "4:00:00", + "8", + "", + ] + ) + + assert infrastructure_section( + WizardSession(state, backend), + resolver, + {}, + ) + + assert backend.partition_default == "gpu-a,gpu-b" + assert state.get_field("infrastructure.runner.slurm.partition") == "gpu-a,gpu-b" + assert backend.remaining == 0 + + +def test_resume_preserves_legacy_partition_fields(tmp_path): + state = WizardState.start(tmp_path / "campaign", defaults_path=None) + legacy = { + "partition_batch": "batch", + "partition_interactive": "interactive", + "partition_cpu": "cpu", + "interactive_max_nodes": 2, + } + for field, value in legacy.items(): + state.set_field(f"infrastructure.runner.slurm.{field}", value, source="user") + + resumed = WizardState.resume(state.path) + + for field, value in legacy.items(): + assert resumed.get_field(f"infrastructure.runner.slurm.{field}") == value + + def test_full_section_keeps_the_existing_customize_prompt(tmp_path): state = WizardState.start(tmp_path / "campaign", defaults_path=None) backend = ScriptedBackend(["customize"]) diff --git a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py b/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py index 8d7f95de990..256e489c1cc 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py @@ -18,6 +18,7 @@ from __future__ import annotations +import shlex from copy import deepcopy from dataclasses import FrozenInstanceError from pathlib import Path @@ -87,8 +88,7 @@ def _campaign_state(tmp_path: Path) -> WizardState: "infrastructure.execution_contract.venv": ".venv", "infrastructure.runner.kind": "slurm", "infrastructure.runner.slurm.account": "account", - "infrastructure.runner.slurm.partition_batch": "batch", - "infrastructure.runner.slurm.partition_cpu": "cpu", + "infrastructure.runner.slurm.partition": "cluster-default", "infrastructure.gpus_per_node": 8, "output.result_root": "/results", } @@ -190,7 +190,7 @@ def _campaign_state(tmp_path: Path) -> WizardState: "automodel": {"parallel": {"tp": 99}}, }, }, - "runner_overrides": {"runner": {"slurm": {"partition": "late-override"}}}, + "runner_overrides": {"runner": {"slurm": {"partition": ["late-a", "late-b"]}}}, "default_resolutions": { "pruning.depth_remove": {"value": 0, "source": "preset"}, "mip.num_solutions": {"value": 8, "source": "defaults_file"}, @@ -370,10 +370,69 @@ def test_runner_compatibility_override_is_applied_to_resolved_runner(tmp_path: P runner = render_runner_v2(state, "production") - assert runner["runner"]["slurm"]["partition"] == "late-override" + assert runner["runner"]["slurm"]["partition"] == ["late-a", "late-b"] assert runner["runner"]["slurm"]["account"] == "account" +def test_runner_compatibility_override_accepts_scalar_partition(tmp_path: Path) -> None: + state = _campaign_state(tmp_path) + state.set_collection( + "runner_overrides", + { + "runner": { + "slurm": { + "partition": "cluster-c", + } + } + }, + ) + + runner = render_runner_v2(state, "production") + + assert runner["runner"]["slurm"]["partition"] == "cluster-c" + + +def test_stage_resource_accepts_multiple_eligible_partitions(tmp_path: Path) -> None: + state = _campaign_state(tmp_path) + resources = deepcopy(state.collection("stage_resources")) + resources["width_importance"]["partition"] = ["gpu-a", "gpu-b"] + state.set_collection("stage_resources", resources) + + execution = render_execution_v2(state, "production") + + assert execution["execution"]["stages"]["width_importance"]["partition"] == "gpu-a,gpu-b" + + +def test_generated_readme_separates_plan_inspection_from_launch(tmp_path: Path) -> None: + campaign_dir = tmp_path / "campaign with spaces" + repository = "/worker checkout" + readme = bundle_module._bundle_readme(campaign_dir, repository) + + commands = [shlex.split(line) for line in readme.splitlines() if line.startswith("python ")] + orchestrator_commands = [ + command for command in commands if command[1].endswith("/orchestrate.py") + ] + inspection_commands = [command for command in orchestrator_commands if "--dry-run" in command] + launch_commands = [command for command in orchestrator_commands if "--dry-run" not in command] + + assert len(inspection_commands) == len(launch_commands) == 2 + assert {command[command.index("--experiment") + 1] for command in orchestrator_commands} == { + str(campaign_dir / "smoke" / "experiment.yaml"), + str(campaign_dir / "production" / "experiment.yaml"), + } + assert all( + command[1] == f"{repository}/examples/puzzletron/orchestrate.py" + for command in orchestrator_commands + ) + resume_command = next(command for command in commands if "--resume" in command) + assert resume_command == [ + "python", + f"{repository}/examples/puzzletron/puzzletron_setup_v2.py", + "--resume", + str(campaign_dir), + ] + + def test_execution_uses_resolved_stage_resource_and_parallel_profile(tmp_path: Path) -> None: state = _campaign_state(tmp_path) @@ -391,11 +450,37 @@ def test_execution_uses_resolved_stage_resource_and_parallel_profile(tmp_path: P "ep": 1, "dp_shard": 1, "dp_replicate": 1, - "sequence_parallel": True, }, } +def test_execution_strips_model_runtime_fields_from_inline_parallel_mesh(tmp_path: Path) -> None: + state = _campaign_state(tmp_path) + resources = deepcopy(state.collection("stage_resources")) + resources["width_importance"].pop("profile_name") + resources["width_importance"]["parallel"] = { + "tp": 2, + "cp": 1, + "pp": 1, + "ep": 1, + "dp_shard": 1, + "dp_replicate": 1, + "sequence_parallel": True, + } + state.set_collection("stage_resources", resources) + + execution = render_execution_v2(state, "production") + + assert execution["execution"]["stages"]["width_importance"]["parallel"] == { + "tp": 2, + "cp": 1, + "pp": 1, + "ep": 1, + "dp_shard": 1, + "dp_replicate": 1, + } + + def test_legacy_meshes_follow_bound_consumer_profiles(tmp_path: Path) -> None: state = _campaign_state(tmp_path) state.set_collection( @@ -510,7 +595,7 @@ def mutate_state_after_smoke(path: Path) -> BundleValidation: assert "mutated" not in experiment assert execution["execution"]["stages"]["width_importance"]["instances"] == 1 assert execution["execution"]["stages"]["width_importance"]["parallel"]["tp"] == 2 - assert runner["runner"]["slurm"]["partition"] == "late-override" + assert runner["runner"]["slurm"]["partition"] == ["late-a", "late-b"] provenance = yaml.safe_load((state.campaign_dir / "resolved_defaults.yaml").read_text()) assert provenance["stages.width_importance.batch"] == { "value": 4, From c39006989e867f7f535bc6174894eae86d8f0176 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 26 Aug 2026 09:50:40 +0200 Subject: [PATCH 2/4] Fix Puzzletron orchestration edge cases Signed-off-by: Johannes Rausch --- .../orchestration/runner.slurm.example.yaml | 2 +- .../docs/orchestration_operations.md | 6 ++++-- .../puzzletron/docs/slurm_configuration.md | 3 ++- .../orchestration/adapters/post_mip.py | 4 +--- .../torch/puzzletron/orchestration/schema.py | 3 ++- puzzletron_setup/v2/resolved.py | 10 +++++++--- .../test_orchestration_executors.py | 19 +++++++++++++++++++ .../test_orchestration_lightweight.py | 2 ++ .../torch/puzzletron/test_portable_configs.py | 2 +- .../torch/puzzletron/test_post_mip_adapter.py | 9 +++++++++ .../test_setup_v2_resolved_config.py | 11 +++++++++++ 11 files changed, 59 insertions(+), 12 deletions(-) diff --git a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml b/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml index 3b1a47fd6c3..cce5898216b 100644 --- a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml +++ b/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml @@ -14,7 +14,7 @@ runner: - REPLACE_WITH_PRIMARY_SLURM_PARTITION - REPLACE_WITH_ALTERNATE_SLURM_PARTITION time_limit: "4:00:00" - log_dir: puzzle_runs/logs + log_dir: logs execution_contract: # Required. Use the checkout path visible on every worker and in the container, if used. repository: REPLACE_WITH_WORKER_VISIBLE_MODELOPT_CHECKOUT diff --git a/examples/puzzletron/docs/orchestration_operations.md b/examples/puzzletron/docs/orchestration_operations.md index a982cf02a6c..4fec6ec5b87 100644 --- a/examples/puzzletron/docs/orchestration_operations.md +++ b/examples/puzzletron/docs/orchestration_operations.md @@ -29,8 +29,10 @@ detach while leaving jobs running, or continue. Non-interactive Ctrl-C and SIGTERM cancel active work and quit. A detached controller preserves durable handles, so running the same command recovers the active jobs. -Use `--color always` when piping through `tee`, `--color never` for plain logs, -and `--poll-interval SECONDS` to change the default five-second poll interval. +Redirect stderr before piping through `tee` (for example, append +`2>&1 | tee run.log`) so progress output is captured. Use `--color always` for +colored output, `--color never` for plain logs, and `--poll-interval SECONDS` +to change the default five-second poll interval. ## State and execution records diff --git a/examples/puzzletron/docs/slurm_configuration.md b/examples/puzzletron/docs/slurm_configuration.md index 278a6965f66..81039c5844d 100644 --- a/examples/puzzletron/docs/slurm_configuration.md +++ b/examples/puzzletron/docs/slurm_configuration.md @@ -12,7 +12,8 @@ use the site's Slurm default. A stage can set `runner.slurm.log_dir` sets the directory used for every attempt log, including the final-report attempt. When omitted, logs are written below -`/logs`. +`/logs`. Relative values are resolved from `puzzle_dir`; absolute +paths are used as written. The runner loader accepts `partition_interactive`, `partition_batch`, `partition_cpu`, and `interactive_max_nodes` as compatibility fields. They diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index 88295e9d4b7..7ec057bffbb 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -179,9 +179,7 @@ def plan(self, plan: CampaignPlan, node: StagePlanNode) -> WorkPlan: available = _available_evaluation_candidates(plan, node.stage_id) if available is not None: if available < 1: - raise RuntimeError( - f"{node.stage_id} has no candidate architectures to evaluate" - ) + raise ValueError(f"{node.stage_id} has no candidate architectures to evaluate") count = min(count, available) count = _full_node_instance_count(node, count) if count == 1: diff --git a/modelopt/torch/puzzletron/orchestration/schema.py b/modelopt/torch/puzzletron/orchestration/schema.py index 4d24c631d11..5c229e7ffd9 100644 --- a/modelopt/torch/puzzletron/orchestration/schema.py +++ b/modelopt/torch/puzzletron/orchestration/schema.py @@ -262,7 +262,8 @@ def log_dir(self) -> Path: """Return the configured shared log directory for every campaign attempt.""" if self.runner.slurm is not None and self.runner.slurm.log_dir: - return Path(self.runner.slurm.log_dir).expanduser() + configured = Path(self.runner.slurm.log_dir).expanduser() + return configured if configured.is_absolute() else self.puzzle_dir / configured return self.puzzle_dir / "logs" diff --git a/puzzletron_setup/v2/resolved.py b/puzzletron_setup/v2/resolved.py index fb6c494c743..420edeb2614 100644 --- a/puzzletron_setup/v2/resolved.py +++ b/puzzletron_setup/v2/resolved.py @@ -586,15 +586,19 @@ def _stage_resource(stage_id: str, raw: Mapping[str, Any]) -> ResolvedStageResou "parallel", } parallel = raw.get("parallel") + try: + partition = normalize_slurm_partition( + raw.get("partition"), path=f"stages.{stage_id}.partition" + ) + except (TypeError, ValueError) as error: + raise SetupError(str(error)) from error return ResolvedStageResource( stage_id=stage_id, strategy=str(raw.get("strategy", "single")), instances=int(raw.get("instances", 1)), resource=str(raw.get("resource", "gpu")), gpus_per_node=(int(raw["gpus_per_node"]) if raw.get("gpus_per_node") is not None else None), - partition=normalize_slurm_partition( - raw.get("partition"), path=f"stages.{stage_id}.partition" - ), + partition=partition, profile_name=(str(raw["profile_name"]) if raw.get("profile_name") else None), parallel=_mapping(parallel) if isinstance(parallel, Mapping) else None, extra={key: value for key, value in raw.items() if key not in known}, diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index 5fa6fad0ca6..e406fde3f5e 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -304,6 +304,25 @@ def stage(stage_id: str, strategy: ExecutionStrategy) -> StagePlanNode: assert Path(attempt.command.log_path).parent == expected_log_dir +def test_campaign_plan_anchors_relative_log_dir_to_puzzle_dir(tmp_path: Path) -> None: + runner = RunnerEnvironment( + kind="slurm", + contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), + slurm=SlurmRunnerConfig(account="acct", log_dir="shared-logs"), + ) + plan = CampaignPlan( + experiment_config_path=str(tmp_path / "experiment.yaml"), + puzzle_dir=tmp_path / "run", + experiment_config={}, + runner=runner, + execution_defaults={}, + stages=(), + contract_hash="contract", + ) + + assert plan.log_dir == tmp_path / "run" / "shared-logs" + + @pytest.mark.parametrize( ("slurm_kwargs", "node_partition", "expected_partition", "configured_log_dir"), [ diff --git a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py index 8d757d93c5e..fe2d02777df 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py +++ b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py @@ -332,6 +332,7 @@ def test_orchestrator_cli_reports_config_errors_without_traceback(tmp_path: Path capture_output=True, text=True, check=False, + timeout=30, ) assert result.returncode == 2 @@ -399,6 +400,7 @@ def test_orchestrator_cli_reports_dry_run_adapter_errors_without_traceback( capture_output=True, text=True, check=False, + timeout=30, ) assert result.returncode == 2 diff --git a/tests/unit/torch/puzzletron/test_portable_configs.py b/tests/unit/torch/puzzletron/test_portable_configs.py index e0cb8f01baa..d00fd56282f 100644 --- a/tests/unit/torch/puzzletron/test_portable_configs.py +++ b/tests/unit/torch/puzzletron/test_portable_configs.py @@ -55,7 +55,7 @@ def test_slurm_runner_example_is_portable() -> None: assert slurm.slurm.partition == ( "REPLACE_WITH_PRIMARY_SLURM_PARTITION,REPLACE_WITH_ALTERNATE_SLURM_PARTITION" ) - assert slurm.slurm.log_dir == "puzzle_runs/logs" + assert slurm.slurm.log_dir == "logs" def test_baremetal_runner_example_is_portable() -> None: diff --git a/tests/unit/torch/puzzletron/test_post_mip_adapter.py b/tests/unit/torch/puzzletron/test_post_mip_adapter.py index cadbc951e94..a9a0fb88153 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_adapter.py +++ b/tests/unit/torch/puzzletron/test_post_mip_adapter.py @@ -135,6 +135,15 @@ def test_post_mip_evaluation_clamps_workers_to_available_candidates(tmp_path: Pa assert [item.work_id for item in work_plan.items] == [f"{node.stage_id}:0"] +def test_post_mip_evaluation_rejects_empty_candidate_set(tmp_path: Path, monkeypatch): + plan, node = _plan(tmp_path, stage_id="post.params.online_eval", node_type="evaluation") + identity_api = _candidate_count_api(0) + monkeypatch.setattr(post_mip_adapter, "_post_mip_identity_api", lambda: identity_api) + + with pytest.raises(ValueError, match="has no candidate architectures to evaluate"): + PostMIPAdapter().plan(plan, node) + + def test_post_mip_evaluation_uses_torchrun_for_single_gpu_workers(tmp_path: Path): plan, node = _plan(tmp_path, stage_id="post.params.online_eval", node_type="evaluation") attempt = PostMIPAdapter().command( diff --git a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py b/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py index 256e489c1cc..f67781e5550 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_resolved_config.py @@ -28,6 +28,7 @@ import yaml import puzzletron_setup.v2.bundle as bundle_module +from puzzletron_setup import SetupError from puzzletron_setup.v2.bundle import ( build_bundles_v2, render_execution_v2, @@ -403,6 +404,16 @@ def test_stage_resource_accepts_multiple_eligible_partitions(tmp_path: Path) -> assert execution["execution"]["stages"]["width_importance"]["partition"] == "gpu-a,gpu-b" +def test_stage_resource_reports_invalid_partition_as_setup_error(tmp_path: Path) -> None: + state = _campaign_state(tmp_path) + resources = deepcopy(state.collection("stage_resources")) + resources["width_importance"]["partition"] = [] + state.set_collection("stage_resources", resources) + + with pytest.raises(SetupError, match=r"stages\.width_importance\.partition"): + resolve_campaign_config(state) + + def test_generated_readme_separates_plan_inspection_from_launch(tmp_path: Path) -> None: campaign_dir = tmp_path / "campaign with spaces" repository = "/worker checkout" From 5444bd91f7eb9d6400147ec363f39f5e22b36a3c Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 26 Aug 2026 11:57:22 +0200 Subject: [PATCH 3/4] Reuse Puzzletron command runner Share the existing argv-only subprocess path so Slurm remains shell-free without security-check suppressions. Signed-off-by: Johannes Rausch --- .../puzzletron/orchestration/executors/slurm.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/puzzletron/orchestration/executors/slurm.py b/modelopt/torch/puzzletron/orchestration/executors/slurm.py index 049faae7a8c..05a03ef0a35 100644 --- a/modelopt/torch/puzzletron/orchestration/executors/slurm.py +++ b/modelopt/torch/puzzletron/orchestration/executors/slurm.py @@ -20,23 +20,22 @@ import math import os import shlex -import subprocess # nosec B404 import time from pathlib import Path -from typing import Sequence +from typing import Protocol, Sequence from ..schema import AttemptSpec, JobHandle, JobState, JobStatus, RunnerEnvironment from ..task_launcher import TASK_IDENTITY_ENV_KEYS from ..task_topology import resolve_task_topology +from .baremetal import _run_command from .base import Executor __all__ = ["SlurmExecutor", "render_hook_lines", "render_sbatch_script"] -def _run_command(argv: Sequence[str]) -> subprocess.CompletedProcess[str]: - return subprocess.run( # nosec B603 - explicit argv is executed without a shell - list(argv), capture_output=True, text=True, check=False - ) +class _CapturedStreams(Protocol): + stdout: str + stderr: str def _slurm_job_id(handle: JobHandle) -> str | None: @@ -63,7 +62,7 @@ def _slurm_job_id(handle: JobHandle) -> str | None: _CANCEL_POLL_SECONDS = 1.0 -def _is_transient_submit_error(result: subprocess.CompletedProcess[str]) -> bool: +def _is_transient_submit_error(result: _CapturedStreams) -> bool: detail = f"{result.stderr}\n{result.stdout}".lower() return any(marker in detail for marker in _TRANSIENT_SUBMIT_ERRORS) From 09484398824e47c23daf6383c1c91f37e723b073 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 26 Aug 2026 13:13:00 +0200 Subject: [PATCH 4/4] Honor configured aggregation log directory Signed-off-by: Johannes Rausch --- modelopt/torch/puzzletron/orchestration/adapters/sharded.py | 2 +- tests/unit/torch/puzzletron/test_orchestration_executors.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py index a99044cd392..125cd4c0ea8 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py @@ -139,7 +139,7 @@ def _run_local_aggregate( """Run a controller-side merge through the reviewed local executor.""" attempt_id = str(uuid.uuid4()) - log_path = plan.puzzle_dir / "logs" / f"{node.stage_id}_merge_{attempt_id}.log" + log_path = plan.log_dir / f"{node.stage_id}_merge_{attempt_id}.log" attempt = AttemptSpec( attempt_id=attempt_id, work_id=f"{node.stage_id}:aggregate", diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index 1a8d9dc1299..4ba9b9ec23d 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -498,6 +498,7 @@ def test_aiperf_aggregation_uses_reviewed_local_executor(tmp_path: Path, monkeyp runner = RunnerEnvironment( kind="local", contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), + slurm=SlurmRunnerConfig(account="test", log_dir="shared-logs"), ) node = StagePlanNode( stage_id="aiperf", @@ -565,6 +566,8 @@ def poll(self, handles): attempt = submitted[0] assert attempt.allocation_gpus == 0 assert attempt.command.cwd == str(tmp_path) + assert attempt.command.log_path is not None + assert Path(attempt.command.log_path).parent == tmp_path / "run" / "shared-logs" assert attempt.command.argv[-1] == "--merge"