From 89b4bb41b877fa337a3cb31ddc24579ec5a349b1 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 29 Jul 2026 03:48:07 -0700 Subject: [PATCH 1/7] Make Puzzletron examples portable Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 8 ++- .../nemotron3/nano_30b_a3b_bf16/model.yaml | 2 +- .../orchestration/qwen_moe/runner.slurm.yaml | 15 ++--- .../orchestration/runner.slurm.example.yaml | 15 ++--- nv-internal/puzzletron_defaults.example.yaml | 12 ++++ nv-internal/sepehr_defaults.yaml | 15 ----- .../torch/puzzletron/test_portable_configs.py | 65 +++++++++++++++++++ 7 files changed, 96 insertions(+), 36 deletions(-) create mode 100644 nv-internal/puzzletron_defaults.example.yaml delete mode 100644 nv-internal/sepehr_defaults.yaml create mode 100644 tests/unit/torch/puzzletron/test_portable_configs.py diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index b7be3aa0c5a..87effb28c5f 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -22,7 +22,7 @@ campaign report. The configs below are the exact current-code entry points for the completed campaigns. Each verified report is a self-contained HTML file that embeds all -sanity-check outputs, stage manifests, and evaluation results — they can be +sanity-check outputs, stage manifests, and evaluation results, they can be 100s of MB. Download them to disk and open locally rather than previewing in a browser tab. @@ -80,10 +80,12 @@ local defaults-versus-customize decisions at every section: ```bash python examples/puzzletron/puzzletron_setup_v2.py \ - --defaults nv-internal/sepehr_defaults.yaml + --defaults nv-internal/puzzletron_defaults.example.yaml ``` -The defaults file is loaded only when passed explicitly. Selection prompts have +The example defaults use only repository-relative values. Copy the file and add +site-specific data, scheduler, and container settings before selecting it. +Defaults are loaded only when passed explicitly. Selection prompts have a visible **← Back** action; text and 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: diff --git a/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml index 672dd813da8..5ab35f7b6ed 100644 --- a/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml +++ b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml @@ -1,7 +1,7 @@ # @package _global_ display_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 -input_hf_model_path: /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/ssameni/hf/hub/models--nvidia--NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/snapshots/cbd3fa9f933d55ef16a84236559f4ee2a0526848 +input_hf_model_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 model_info: hf_repo: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 diff --git a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml b/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml index b9a4e840e6d..643e5f6845d 100644 --- a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml +++ b/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml @@ -11,14 +11,11 @@ runner: interactive_max_nodes: 2 max_nodes: 20 time_limit: "4:00:00" - log_dir: /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/ssameni/engineering/puzzle_runs/qwen-moe/logs + log_dir: puzzle_runs/qwen-moe/logs execution_contract: - repository: /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/ssameni/engineering/modelopt_qwen - venv: /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/ssameni/engineering/modelopt_qwen/.venv_new - container: /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/ssameni/enroot/pytorch_25p05.sqsh - container_mounts: /lustre:/lustre - prerun_commands: - - source /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/ssameni/setup-envs.sh - - export VLLM_ROOT=/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/ssameni/engineering/vllm_new - - export AUTOMODEL_ROOT=/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/ssameni/engineering/Automodel_new + repository: . + venv: .venv + container: + container_mounts: + prerun_commands: [] postrun_commands: [] diff --git a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml b/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml index 49eec6b49f3..589fe713e04 100644 --- a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml +++ b/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Example Slurm runner config for the v2 orchestrator.""" +# Example Slurm runner config for the v2 orchestrator. runner: kind: slurm @@ -11,15 +11,14 @@ runner: partition_interactive: interactive partition_batch: batch # Optional. CPU/IO stages fall back to one GPU when omitted. - partition_cpu: null + partition_cpu: interactive_max_nodes: 2 time_limit: "4:00:00" log_dir: puzzle_runs/logs execution_contract: - repository: /shared/modelopt - venv: /shared/modelopt/.venv - container: /path/to/container.sqsh - container_mounts: /shared:/shared - prerun_commands: - - source /path/to/site-setup.sh + repository: . + venv: .venv + container: + container_mounts: + prerun_commands: [] postrun_commands: [] diff --git a/nv-internal/puzzletron_defaults.example.yaml b/nv-internal/puzzletron_defaults.example.yaml new file mode 100644 index 00000000000..8525134203f --- /dev/null +++ b/nv-internal/puzzletron_defaults.example.yaml @@ -0,0 +1,12 @@ +schema_version: 1 +infrastructure: + execution_contract: + repository: . + venv: .venv + container: + container_mounts: + prerun_commands: [] + runner: + kind: slurm + slurm: + partition_cpu: diff --git a/nv-internal/sepehr_defaults.yaml b/nv-internal/sepehr_defaults.yaml deleted file mode 100644 index 80eab3bbe44..00000000000 --- a/nv-internal/sepehr_defaults.yaml +++ /dev/null @@ -1,15 +0,0 @@ -schema_version: 1 -data: - source: /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/ssameni/engineering/Puzzle-KD-Nemotron-Post-Training-Dataset-v2/ -infrastructure: - execution_contract: - venv: .venv_new - container: /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/ssameni/enroot/pytorch_25p05.sqsh - container_mounts: /lustre:/lustre - prerun_commands: - - source /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/ssameni/setup-envs.sh - runner: - kind: slurm - slurm: - account: coreai_dlalgo_llm - partition_cpu: cpu_interactive diff --git a/tests/unit/torch/puzzletron/test_portable_configs.py b/tests/unit/torch/puzzletron/test_portable_configs.py new file mode 100644 index 00000000000..5c6ff638b01 --- /dev/null +++ b/tests/unit/torch/puzzletron/test_portable_configs.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. + +"""Regression tests for portable Puzzletron example configuration.""" + +from pathlib import Path + +import yaml + +from puzzletron_orchestrator.compiler import load_runner_config +from puzzletron_setup.v2.defaults import load_defaults + +REPOSITORY_ROOT = Path(__file__).parents[4] + + +def test_runner_examples_use_repository_relative_defaults() -> None: + runner_paths = ( + "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", + "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", + ) + + for relative_path in runner_paths: + runner = load_runner_config(REPOSITORY_ROOT / relative_path) + assert runner.contract.repository == "." + assert runner.contract.venv == ".venv" + assert runner.contract.container is None + assert runner.contract.container_mounts is None + assert not runner.contract.prerun_commands + + +def test_setup_defaults_are_portable_and_valid() -> None: + path = REPOSITORY_ROOT / "nv-internal/puzzletron_defaults.example.yaml" + + defaults = load_defaults(path) + + assert defaults["infrastructure"]["execution_contract"] == { + "repository": ".", + "venv": ".venv", + "container": None, + "container_mounts": None, + "prerun_commands": [], + } + + +def test_nemotron_model_uses_hugging_face_identity() -> None: + path = ( + REPOSITORY_ROOT + / "examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml" + ) + + config = yaml.safe_load(path.read_text()) + + assert config["input_hf_model_path"] == config["model_info"]["hf_repo"] From 940e5f1211835c62add3fcdbdd83d35ff3970952 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 29 Jul 2026 10:02:26 -0700 Subject: [PATCH 2/7] Refine portable Puzzletron examples Signed-off-by: Johannes Rausch --- .../orchestration/execution.example.yaml | 2 +- .../orchestration/qwen_moe/runner.slurm.yaml | 12 +++- .../runner.baremetal.example.yaml | 10 +-- .../orchestration/runner.slurm.example.yaml | 12 ++-- nv-internal/puzzletron_defaults.example.yaml | 11 +++- puzzletron_setup/wizard.py | 2 +- .../torch/puzzletron/test_portable_configs.py | 62 ++++++++++++++++++- 7 files changed, 92 insertions(+), 19 deletions(-) diff --git a/examples/puzzletron/configs/orchestration/execution.example.yaml b/examples/puzzletron/configs/orchestration/execution.example.yaml index 0e764993bc2..562d0cd5b84 100644 --- a/examples/puzzletron/configs/orchestration/execution.example.yaml +++ b/examples/puzzletron/configs/orchestration/execution.example.yaml @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Example runner and execution configs for the v2 orchestrator.""" +# Example execution config for the v2 orchestrator. execution: defaults: diff --git a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml b/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml index 643e5f6845d..8083e27cc6d 100644 --- a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml +++ b/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml @@ -4,10 +4,14 @@ runner: kind: slurm slurm: - account: coreai_dlalgo_llm + # Required. Replace with the Slurm account for your site. + account: your_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: null interactive_max_nodes: 2 max_nodes: 20 time_limit: "4:00:00" @@ -15,7 +19,9 @@ runner: execution_contract: repository: . venv: .venv - container: - container_mounts: + # Optional. Leave null to execute directly in the worker environment. + container: null + # Optional and used only when a container is configured. + container_mounts: null prerun_commands: [] postrun_commands: [] diff --git a/examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml b/examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml index fa020ad36b3..890ca577a70 100644 --- a/examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml +++ b/examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml @@ -1,11 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Example bare-metal runner config for the v2 orchestrator.""" +# Example bare-metal runner config for the v2 orchestrator. runner: kind: baremetal inventory: + # Replace these generic names with hosts reachable from the controller. rendezvous_host: node-a rendezvous_port_base: 29500 hosts: @@ -14,6 +15,7 @@ runner: - hostname: node-b gpus: 8 execution_contract: - repository: /shared/modelopt - venv: /shared/modelopt/.venv - setup_env: "" + repository: . + venv: .venv + # Optional. Set to a worker-side environment setup script when needed. + setup_env: null diff --git a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml b/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml index 589fe713e04..869108fda15 100644 --- a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml +++ b/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml @@ -6,19 +6,23 @@ runner: kind: slurm slurm: + # Required. Replace with the Slurm account for your site. account: your_account + # Replace these generic partition names if your site uses different names. partition: batch partition_interactive: interactive partition_batch: batch - # Optional. CPU/IO stages fall back to one GPU when omitted. - partition_cpu: + # Optional. CPU/IO stages use the regular one-node partition when unset. + partition_cpu: null interactive_max_nodes: 2 time_limit: "4:00:00" log_dir: puzzle_runs/logs execution_contract: repository: . venv: .venv - container: - container_mounts: + # Optional. Leave null to execute directly in the worker environment. + container: null + # Optional and used only when a container is configured. + container_mounts: null prerun_commands: [] postrun_commands: [] diff --git a/nv-internal/puzzletron_defaults.example.yaml b/nv-internal/puzzletron_defaults.example.yaml index 8525134203f..29355cc23dd 100644 --- a/nv-internal/puzzletron_defaults.example.yaml +++ b/nv-internal/puzzletron_defaults.example.yaml @@ -3,10 +3,15 @@ infrastructure: execution_contract: repository: . venv: .venv - container: - container_mounts: + # Optional. Leave null to execute directly in the worker environment. + container: null + # Optional and used only when a container is configured. + container_mounts: null prerun_commands: [] runner: kind: slurm slurm: - partition_cpu: + # Add the required Slurm account for your site before using this file. + # account: your_account + # Optional. CPU/IO stages use the regular partition when unset. + partition_cpu: null diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index 52a9e9374f5..c658888bd19 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -1219,7 +1219,7 @@ def _ask_infrastructure( default="", description=( "Use /host/path:/container/path; separate multiple mounts with commas, for " - "example /lustre:/lustre,/datasets:/datasets." + "example /data:/data,/models:/models." ), ).strip() prerun = prompts.text("Pre-run commands separated by ';;' (blank for none):", default="") diff --git a/tests/unit/torch/puzzletron/test_portable_configs.py b/tests/unit/torch/puzzletron/test_portable_configs.py index 5c6ff638b01..61856f71aa4 100644 --- a/tests/unit/torch/puzzletron/test_portable_configs.py +++ b/tests/unit/torch/puzzletron/test_portable_configs.py @@ -18,26 +18,44 @@ from pathlib import Path import yaml - from puzzletron_orchestrator.compiler import load_runner_config + from puzzletron_setup.v2.defaults import load_defaults REPOSITORY_ROOT = Path(__file__).parents[4] def test_runner_examples_use_repository_relative_defaults() -> None: - runner_paths = ( + slurm_runner_paths = ( "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", ) - for relative_path in runner_paths: + for relative_path in slurm_runner_paths: runner = load_runner_config(REPOSITORY_ROOT / relative_path) assert runner.contract.repository == "." assert runner.contract.venv == ".venv" assert runner.contract.container is None assert runner.contract.container_mounts is None assert not runner.contract.prerun_commands + assert runner.slurm is not None + assert runner.slurm.account == "your_account" + assert runner.slurm.partition_cpu is None + + baremetal = load_runner_config( + REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml" + ) + assert baremetal.contract.repository == "." + assert baremetal.contract.venv == ".venv" + assert baremetal.contract.setup_env is None + + +def test_execution_example_is_valid_yaml() -> None: + path = REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/execution.example.yaml" + + config = yaml.safe_load(path.read_text()) + + assert config["execution"]["defaults"]["gpus_per_node"] == 8 def test_setup_defaults_are_portable_and_valid() -> None: @@ -63,3 +81,41 @@ def test_nemotron_model_uses_hugging_face_identity() -> None: config = yaml.safe_load(path.read_text()) assert config["input_hf_model_path"] == config["model_info"]["hf_repo"] + + +def test_active_examples_use_portable_value_shapes() -> None: + slurm_paths = ( + "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", + "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", + ) + + for relative_path in slurm_paths: + runner = load_runner_config(REPOSITORY_ROOT / relative_path) + assert not Path(runner.contract.repository).is_absolute() + assert not Path(runner.contract.venv).is_absolute() + assert runner.slurm is not None + assert runner.slurm.account.startswith("REPLACE_WITH_") + + baremetal = load_runner_config( + REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml" + ) + assert baremetal.baremetal is not None + assert baremetal.baremetal.rendezvous_host.startswith("REPLACE_WITH_") + assert all(host.hostname.startswith("REPLACE_WITH_") for host in baremetal.baremetal.hosts) + + defaults_paths = sorted((REPOSITORY_ROOT / "nv-internal").glob("*_defaults*.yaml")) + assert [path.name for path in defaults_paths] == ["puzzletron_defaults.example.yaml"] + + +def test_optional_yaml_values_are_explicit() -> None: + paths = ( + "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", + "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", + "nv-internal/puzzletron_defaults.example.yaml", + ) + + for relative_path in paths: + content = (REPOSITORY_ROOT / relative_path).read_text() + assert "partition_cpu: null" in content + assert "container: null" in content + assert "container_mounts: null" in content From 848f7c5bc58c7433c8ed798a0f1e6c237a441fb0 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 29 Jul 2026 15:06:20 -0700 Subject: [PATCH 3/7] Clarify Puzzletron example inputs Signed-off-by: Johannes Rausch --- .../nemotron3/nano_30b_a3b_bf16/model.yaml | 3 + .../families/qwen3_5/qwen3p5_9b/model.yaml | 3 + .../qwen3_5/qwen3p6_35b_a3b/model.yaml | 3 + .../orchestration/qwen_moe/runner.slurm.yaml | 17 +++-- .../runner.baremetal.example.yaml | 16 +++-- .../orchestration/runner.slurm.example.yaml | 17 +++-- nv-internal/puzzletron_defaults.example.yaml | 14 ++++- puzzletron_setup/wizard.py | 32 ++++++++-- .../torch/puzzletron/test_portable_configs.py | 62 ++++++++++++++++--- 9 files changed, 140 insertions(+), 27 deletions(-) diff --git a/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml index 5ab35f7b6ed..0020cb6acbe 100644 --- a/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml +++ b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml @@ -1,9 +1,12 @@ # @package _global_ display_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 +# Runnable public Hugging Face repository ID; this is not a placeholder. +# Setup-generated bundles replace it with the model source selected by the user. input_hf_model_path: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 model_info: + # Canonical public identity used for model metadata and revision checks. hf_repo: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 hf_revision: cbd3fa9f933d55ef16a84236559f4ee2a0526848 model_type: nemotron_h diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_9b/model.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_9b/model.yaml index b81ac8c7d98..e162df8cc9c 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_9b/model.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_9b/model.yaml @@ -1,8 +1,11 @@ # @package _global_ +# Runnable public Hugging Face repository ID; this is not a placeholder. +# Setup-generated bundles replace it with the model source selected by the user. input_hf_model_path: Qwen/Qwen3.5-9B model_info: + # Canonical public identity used for model metadata and revision checks. hf_repo: Qwen/Qwen3.5-9B model_type: qwen3_5_text architectures: [Qwen3_5ForConditionalGeneration] diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p6_35b_a3b/model.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p6_35b_a3b/model.yaml index b3c25d14704..737ff41e219 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p6_35b_a3b/model.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p6_35b_a3b/model.yaml @@ -1,12 +1,15 @@ # @package _global_ display_name: Qwen/Qwen3.6-35B-A3B +# Runnable public Hugging Face repository ID; this is not a placeholder. +# Setup-generated bundles replace it with the model source selected by the user. input_hf_model_path: Qwen/Qwen3.6-35B-A3B # Required by depth RPC / scoring stages; unlike dense Qwen configs this is not # inherited from a family-level descriptor key. descriptor: qwen3_5_moe model_info: + # Canonical public identity used for model metadata and revision checks. hf_repo: Qwen/Qwen3.6-35B-A3B hf_revision: 995ad96eacd98c81ed38be0c5b274b04031597b0 model_type: qwen3_5_moe diff --git a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml b/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml index 8083e27cc6d..3e7cf0d35b1 100644 --- a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml +++ b/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml @@ -4,8 +4,8 @@ runner: kind: slurm slurm: - # Required. Replace with the Slurm account for your site. - account: your_account + # 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 @@ -17,11 +17,20 @@ runner: time_limit: "4:00:00" log_dir: puzzle_runs/qwen-moe/logs execution_contract: + # Runnable when orchestration starts from a checkout visible on every worker. + # Otherwise, replace "." with the worker-visible repository path. repository: . + # Sourced as /bin/activate after changing into the repository. + # Use a repository-relative path or an absolute path visible on every worker. venv: .venv - # Optional. Leave null to execute directly in the worker environment. + # Optional. Set to an image/path accepted by the cluster's srun container plugin. + # Leave null to execute directly in the worker environment. container: null - # Optional and used only when a container is configured. + # Optional and used only with a container. Use comma-separated + # /host/path:/container/path entries, for example /data:/data,/models:/models. container_mounts: null + # Optional shell commands run in order before virtualenv activation. + # Example: ["module load cuda", "export HF_HOME=/data/huggingface"] prerun_commands: [] + # Optional shell commands run when the stage payload exits. postrun_commands: [] diff --git a/examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml b/examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml index 890ca577a70..03f91ffac8f 100644 --- a/examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml +++ b/examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml @@ -6,16 +6,22 @@ runner: kind: baremetal inventory: - # Replace these generic names with hosts reachable from the controller. - rendezvous_host: node-a + # Required placeholders. Replace both host values with SSH hosts reachable + # from the controller; the rendezvous host must match one inventory hostname. + rendezvous_host: REPLACE_WITH_PRIMARY_HOST rendezvous_port_base: 29500 hosts: - - hostname: node-a + - hostname: REPLACE_WITH_PRIMARY_HOST gpus: 8 - - hostname: node-b + - hostname: REPLACE_WITH_SECONDARY_HOST gpus: 8 execution_contract: + # Runnable when the same checkout is visible on every SSH host. + # Otherwise, replace "." with the repository path shared by those hosts. repository: . + # Sourced as /bin/activate after changing into the repository. + # Use a repository-relative path or an absolute path visible on every host. venv: .venv - # Optional. Set to a worker-side environment setup script when needed. + # Optional script sourced before virtualenv activation on every host. + # Leave null when the worker login environment needs no additional setup. setup_env: null diff --git a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml b/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml index 869108fda15..84942376cf6 100644 --- a/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml +++ b/examples/puzzletron/configs/orchestration/runner.slurm.example.yaml @@ -6,8 +6,8 @@ runner: kind: slurm slurm: - # Required. Replace with the Slurm account for your site. - account: your_account + # 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 @@ -18,11 +18,20 @@ runner: time_limit: "4:00:00" log_dir: puzzle_runs/logs execution_contract: + # Runnable when orchestration starts from a checkout visible on every worker. + # Otherwise, replace "." with the worker-visible repository path. repository: . + # Sourced as /bin/activate after changing into the repository. + # Use a repository-relative path or an absolute path visible on every worker. venv: .venv - # Optional. Leave null to execute directly in the worker environment. + # Optional. Set to an image/path accepted by the cluster's srun container plugin. + # Leave null to execute directly in the worker environment. container: null - # Optional and used only when a container is configured. + # Optional and used only with a container. Use comma-separated + # /host/path:/container/path entries, for example /data:/data,/models:/models. container_mounts: null + # Optional shell commands run in order before virtualenv activation. + # Example: ["module load cuda", "export HF_HOME=/data/huggingface"] prerun_commands: [] + # Optional shell commands run when the stage payload exits. postrun_commands: [] diff --git a/nv-internal/puzzletron_defaults.example.yaml b/nv-internal/puzzletron_defaults.example.yaml index 29355cc23dd..e1fb42f4203 100644 --- a/nv-internal/puzzletron_defaults.example.yaml +++ b/nv-internal/puzzletron_defaults.example.yaml @@ -1,17 +1,25 @@ schema_version: 1 infrastructure: execution_contract: + # Runnable when setup and orchestration start from a checkout visible on every worker. + # Otherwise, replace "." with the worker-visible repository path. repository: . + # Sourced as /bin/activate after changing into the repository. + # Use a repository-relative path or an absolute path visible on every worker. venv: .venv - # Optional. Leave null to execute directly in the worker environment. + # Optional. Set to an image/path accepted by the cluster's srun container plugin. + # Leave null to execute directly in the worker environment. container: null - # Optional and used only when a container is configured. + # Optional and used only with a container. Use comma-separated + # /host/path:/container/path entries, for example /data:/data,/models:/models. container_mounts: null + # Optional shell commands run in order before virtualenv activation. + # Example: ["module load cuda", "export HF_HOME=/data/huggingface"] prerun_commands: [] runner: kind: slurm slurm: # Add the required Slurm account for your site before using this file. - # account: your_account + # account: REPLACE_WITH_SLURM_ACCOUNT # Optional. CPU/IO stages use the regular partition when unset. partition_cpu: null diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index c658888bd19..65ead3422f1 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -1211,9 +1211,29 @@ def _ask_infrastructure( [("Slurm", "slurm"), ("SSH bare metal", "baremetal")], default="slurm", ) - repository = prompts.text("Repository path on workers:", default=str(Path.cwd())) - venv = prompts.text("Python virtual environment on workers:", default=".venv") - container = prompts.text("Container image/path (blank for none):", default="").strip() + repository = prompts.text( + "Repository path on workers:", + default=str(Path.cwd()), + description=( + "Use the checkout path visible on every worker and inside the container, if used." + ), + ) + venv = prompts.text( + "Python virtual environment on workers:", + default=".venv", + description=( + "Use a repository-relative path or an absolute worker-visible path containing " + "bin/activate." + ), + ) + container = prompts.text( + "Container image/path (blank for none):", + default="", + description=( + "Use an image/path accepted by the cluster's srun container plugin, or leave " + "blank to run directly in the worker environment." + ), + ).strip() mounts = prompts.text( "Container mounts (blank for none):", default="", @@ -1222,7 +1242,11 @@ def _ask_infrastructure( "example /data:/data,/models:/models." ), ).strip() - prerun = prompts.text("Pre-run commands separated by ';;' (blank for none):", default="") + prerun = prompts.text( + "Pre-run commands separated by ';;' (blank for none):", + default="", + description="Commands run in order before virtualenv activation.", + ) gpus_per_node = prompts.integer("GPUs per node:", default=8) _print_mesh_guidance(model) mesh_checkpoint = prompts.checkpoint() diff --git a/tests/unit/torch/puzzletron/test_portable_configs.py b/tests/unit/torch/puzzletron/test_portable_configs.py index 61856f71aa4..f9e4c67c391 100644 --- a/tests/unit/torch/puzzletron/test_portable_configs.py +++ b/tests/unit/torch/puzzletron/test_portable_configs.py @@ -39,7 +39,7 @@ def test_runner_examples_use_repository_relative_defaults() -> None: assert runner.contract.container_mounts is None assert not runner.contract.prerun_commands assert runner.slurm is not None - assert runner.slurm.account == "your_account" + assert runner.slurm.account == "REPLACE_WITH_SLURM_ACCOUNT" assert runner.slurm.partition_cpu is None baremetal = load_runner_config( @@ -48,6 +48,12 @@ def test_runner_examples_use_repository_relative_defaults() -> None: assert baremetal.contract.repository == "." assert baremetal.contract.venv == ".venv" assert baremetal.contract.setup_env is None + assert baremetal.baremetal is not None + assert baremetal.baremetal.rendezvous_host == "REPLACE_WITH_PRIMARY_HOST" + assert [host.hostname for host in baremetal.baremetal.hosts] == [ + "REPLACE_WITH_PRIMARY_HOST", + "REPLACE_WITH_SECONDARY_HOST", + ] def test_execution_example_is_valid_yaml() -> None: @@ -72,15 +78,21 @@ def test_setup_defaults_are_portable_and_valid() -> None: } -def test_nemotron_model_uses_hugging_face_identity() -> None: - path = ( - REPOSITORY_ROOT - / "examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml" +def test_model_examples_label_runnable_hugging_face_identities() -> None: + paths = ( + "examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml", + "examples/puzzletron/configs/families/qwen3_5/qwen3p5_9b/model.yaml", + "examples/puzzletron/configs/families/qwen3_5/qwen3p6_35b_a3b/model.yaml", ) - config = yaml.safe_load(path.read_text()) + for relative_path in paths: + content = (REPOSITORY_ROOT / relative_path).read_text() + config = yaml.safe_load(content) - assert config["input_hf_model_path"] == config["model_info"]["hf_repo"] + assert config["input_hf_model_path"] == config["model_info"]["hf_repo"] + assert not config["input_hf_model_path"].startswith("REPLACE_WITH_") + assert "Runnable public Hugging Face repository ID" in content + assert "Setup-generated bundles replace it" in content def test_active_examples_use_portable_value_shapes() -> None: @@ -119,3 +131,39 @@ def test_optional_yaml_values_are_explicit() -> None: assert "partition_cpu: null" in content assert "container: null" in content assert "container_mounts: null" in content + + +def test_required_example_values_use_visible_placeholders() -> None: + slurm_paths = ( + "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", + "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", + ) + + for relative_path in slurm_paths: + content = (REPOSITORY_ROOT / relative_path).read_text() + assert "account: REPLACE_WITH_SLURM_ACCOUNT" in content + + defaults = (REPOSITORY_ROOT / "nv-internal/puzzletron_defaults.example.yaml").read_text() + assert "# account: REPLACE_WITH_SLURM_ACCOUNT" in defaults + + baremetal = ( + REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml" + ).read_text() + assert baremetal.count("REPLACE_WITH_PRIMARY_HOST") == 2 + assert baremetal.count("REPLACE_WITH_SECONDARY_HOST") == 1 + + +def test_execution_contract_examples_explain_runnable_and_optional_values() -> None: + slurm_paths = ( + "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", + "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", + "nv-internal/puzzletron_defaults.example.yaml", + ) + + for relative_path in slurm_paths: + content = (REPOSITORY_ROOT / relative_path).read_text() + assert "checkout visible on every worker" in content + assert "Sourced as /bin/activate" in content + assert "accepted by the cluster's srun container plugin" in content + assert "/data:/data,/models:/models" in content + assert "before virtualenv activation" in content From 3c8ef29d290db4fcd26c57787f4cd08d1d98e0d1 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Fri, 31 Jul 2026 13:56:32 +0200 Subject: [PATCH 4/7] Restore Qwen Slurm execution placeholders Keep the Qwen MoE runner's repository, environment, container, mount, and setup hooks visible without retaining site-specific paths. Add a regression test for the portable placeholder contract. Signed-off-by: Johannes Rausch --- .../orchestration/qwen_moe/runner.slurm.yaml | 28 +++++------ .../torch/puzzletron/test_portable_configs.py | 49 +++++++++++++------ 2 files changed, 46 insertions(+), 31 deletions(-) diff --git a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml b/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml index 3e7cf0d35b1..4bfd30b80bf 100644 --- a/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml +++ b/examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml @@ -17,20 +17,18 @@ runner: time_limit: "4:00:00" log_dir: puzzle_runs/qwen-moe/logs execution_contract: - # Runnable when orchestration starts from a checkout visible on every worker. - # Otherwise, replace "." with the worker-visible repository path. - repository: . - # Sourced as /bin/activate after changing into the repository. - # Use a repository-relative path or an absolute path visible on every worker. - venv: .venv - # Optional. Set to an image/path accepted by the cluster's srun container plugin. - # Leave null to execute directly in the worker environment. - container: null - # Optional and used only with a container. Use comma-separated - # /host/path:/container/path entries, for example /data:/data,/models:/models. - container_mounts: null - # Optional shell commands run in order before virtualenv activation. - # Example: ["module load cuda", "export HF_HOME=/data/huggingface"] - prerun_commands: [] + # Replace with the ModelOpt checkout path visible on every worker and in the container. + repository: REPLACE_WITH_WORKER_VISIBLE_MODELOPT_CHECKOUT + # Replace with the virtual environment path to source on every worker. + venv: REPLACE_WITH_WORKER_VISIBLE_MODELOPT_VENV + # Replace with an image or path accepted by the site's srun container plugin. + container: REPLACE_WITH_SLURM_CONTAINER_IMAGE + # Replace with the host and container paths required by the campaign. + container_mounts: "REPLACE_WITH_HOST_PATH:REPLACE_WITH_CONTAINER_PATH" + # Replace these values with the site's setup script and source checkout paths. + prerun_commands: + - source REPLACE_WITH_SITE_SETUP_SCRIPT + - export VLLM_ROOT=REPLACE_WITH_WORKER_VISIBLE_VLLM_CHECKOUT + - export AUTOMODEL_ROOT=REPLACE_WITH_WORKER_VISIBLE_AUTOMODEL_CHECKOUT # Optional shell commands run when the stage payload exits. postrun_commands: [] diff --git a/tests/unit/torch/puzzletron/test_portable_configs.py b/tests/unit/torch/puzzletron/test_portable_configs.py index f9e4c67c391..5907428a38e 100644 --- a/tests/unit/torch/puzzletron/test_portable_configs.py +++ b/tests/unit/torch/puzzletron/test_portable_configs.py @@ -26,21 +26,17 @@ def test_runner_examples_use_repository_relative_defaults() -> None: - slurm_runner_paths = ( - "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", - "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", + slurm = load_runner_config( + REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml" ) - - for relative_path in slurm_runner_paths: - runner = load_runner_config(REPOSITORY_ROOT / relative_path) - assert runner.contract.repository == "." - assert runner.contract.venv == ".venv" - assert runner.contract.container is None - assert runner.contract.container_mounts is None - assert not runner.contract.prerun_commands - assert runner.slurm is not None - assert runner.slurm.account == "REPLACE_WITH_SLURM_ACCOUNT" - assert runner.slurm.partition_cpu is None + assert slurm.contract.repository == "." + assert slurm.contract.venv == ".venv" + assert slurm.contract.container is None + assert slurm.contract.container_mounts is None + assert not slurm.contract.prerun_commands + assert slurm.slurm is not None + assert slurm.slurm.account == "REPLACE_WITH_SLURM_ACCOUNT" + assert slurm.slurm.partition_cpu is None baremetal = load_runner_config( REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml" @@ -56,6 +52,29 @@ def test_runner_examples_use_repository_relative_defaults() -> 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" + ) + + assert runner.contract.repository == "REPLACE_WITH_WORKER_VISIBLE_MODELOPT_CHECKOUT" + assert runner.contract.venv == "REPLACE_WITH_WORKER_VISIBLE_MODELOPT_VENV" + assert runner.contract.container == "REPLACE_WITH_SLURM_CONTAINER_IMAGE" + assert ( + runner.contract.container_mounts + == "REPLACE_WITH_HOST_PATH:REPLACE_WITH_CONTAINER_PATH" + ) + assert runner.contract.prerun_commands == ( + "source REPLACE_WITH_SITE_SETUP_SCRIPT", + "export VLLM_ROOT=REPLACE_WITH_WORKER_VISIBLE_VLLM_CHECKOUT", + "export AUTOMODEL_ROOT=REPLACE_WITH_WORKER_VISIBLE_AUTOMODEL_CHECKOUT", + ) + assert runner.slurm is not None + assert runner.slurm.account == "REPLACE_WITH_SLURM_ACCOUNT" + assert runner.slurm.partition_cpu is None + + def test_execution_example_is_valid_yaml() -> None: path = REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/execution.example.yaml" @@ -122,7 +141,6 @@ def test_active_examples_use_portable_value_shapes() -> None: def test_optional_yaml_values_are_explicit() -> None: paths = ( "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", - "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", "nv-internal/puzzletron_defaults.example.yaml", ) @@ -156,7 +174,6 @@ def test_required_example_values_use_visible_placeholders() -> None: def test_execution_contract_examples_explain_runnable_and_optional_values() -> None: slurm_paths = ( "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", - "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", "nv-internal/puzzletron_defaults.example.yaml", ) From 63a89ffa25d55a082c6e2bf65d1d4fa89ca7468a Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Sat, 1 Aug 2026 13:28:40 +0200 Subject: [PATCH 5/7] Tighten portable configuration tests Signed-off-by: Johannes Rausch --- .../torch/puzzletron/test_portable_configs.py | 94 +++---------------- 1 file changed, 14 insertions(+), 80 deletions(-) diff --git a/tests/unit/torch/puzzletron/test_portable_configs.py b/tests/unit/torch/puzzletron/test_portable_configs.py index 5907428a38e..de372b674f8 100644 --- a/tests/unit/torch/puzzletron/test_portable_configs.py +++ b/tests/unit/torch/puzzletron/test_portable_configs.py @@ -18,14 +18,14 @@ from pathlib import Path import yaml -from puzzletron_orchestrator.compiler import load_runner_config +from puzzletron_orchestrator.compiler import load_runner_config from puzzletron_setup.v2.defaults import load_defaults -REPOSITORY_ROOT = Path(__file__).parents[4] +REPOSITORY_ROOT = Path(__file__).resolve().parents[4] -def test_runner_examples_use_repository_relative_defaults() -> None: +def test_slurm_runner_example_uses_portable_defaults() -> None: slurm = load_runner_config( REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml" ) @@ -38,6 +38,8 @@ def test_runner_examples_use_repository_relative_defaults() -> None: assert slurm.slurm.account == "REPLACE_WITH_SLURM_ACCOUNT" assert slurm.slurm.partition_cpu is None + +def test_baremetal_runner_example_uses_portable_defaults() -> None: baremetal = load_runner_config( REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml" ) @@ -54,17 +56,13 @@ def test_runner_examples_use_repository_relative_defaults() -> 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" + REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml" ) assert runner.contract.repository == "REPLACE_WITH_WORKER_VISIBLE_MODELOPT_CHECKOUT" assert runner.contract.venv == "REPLACE_WITH_WORKER_VISIBLE_MODELOPT_VENV" assert runner.contract.container == "REPLACE_WITH_SLURM_CONTAINER_IMAGE" - assert ( - runner.contract.container_mounts - == "REPLACE_WITH_HOST_PATH:REPLACE_WITH_CONTAINER_PATH" - ) + assert runner.contract.container_mounts == "REPLACE_WITH_HOST_PATH:REPLACE_WITH_CONTAINER_PATH" assert runner.contract.prerun_commands == ( "source REPLACE_WITH_SITE_SETUP_SCRIPT", "export VLLM_ROOT=REPLACE_WITH_WORKER_VISIBLE_VLLM_CHECKOUT", @@ -95,9 +93,13 @@ def test_setup_defaults_are_portable_and_valid() -> None: "container_mounts": None, "prerun_commands": [], } + assert defaults["infrastructure"]["runner"]["slurm"]["partition_cpu"] is None + + defaults_paths = sorted((REPOSITORY_ROOT / "nv-internal").glob("*_defaults*.yaml")) + assert [path.name for path in defaults_paths] == ["puzzletron_defaults.example.yaml"] -def test_model_examples_label_runnable_hugging_face_identities() -> None: +def test_model_examples_use_public_hugging_face_identities() -> None: paths = ( "examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/model.yaml", "examples/puzzletron/configs/families/qwen3_5/qwen3p5_9b/model.yaml", @@ -110,77 +112,9 @@ def test_model_examples_label_runnable_hugging_face_identities() -> None: assert config["input_hf_model_path"] == config["model_info"]["hf_repo"] assert not config["input_hf_model_path"].startswith("REPLACE_WITH_") - assert "Runnable public Hugging Face repository ID" in content - assert "Setup-generated bundles replace it" in content - - -def test_active_examples_use_portable_value_shapes() -> None: - slurm_paths = ( - "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", - "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", - ) - - for relative_path in slurm_paths: - runner = load_runner_config(REPOSITORY_ROOT / relative_path) - assert not Path(runner.contract.repository).is_absolute() - assert not Path(runner.contract.venv).is_absolute() - assert runner.slurm is not None - assert runner.slurm.account.startswith("REPLACE_WITH_") - - baremetal = load_runner_config( - REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml" - ) - assert baremetal.baremetal is not None - assert baremetal.baremetal.rendezvous_host.startswith("REPLACE_WITH_") - assert all(host.hostname.startswith("REPLACE_WITH_") for host in baremetal.baremetal.hosts) - - defaults_paths = sorted((REPOSITORY_ROOT / "nv-internal").glob("*_defaults*.yaml")) - assert [path.name for path in defaults_paths] == ["puzzletron_defaults.example.yaml"] -def test_optional_yaml_values_are_explicit() -> None: - paths = ( - "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", - "nv-internal/puzzletron_defaults.example.yaml", - ) - - for relative_path in paths: - content = (REPOSITORY_ROOT / relative_path).read_text() - assert "partition_cpu: null" in content - assert "container: null" in content - assert "container_mounts: null" in content - - -def test_required_example_values_use_visible_placeholders() -> None: - slurm_paths = ( - "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", - "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml", - ) - - for relative_path in slurm_paths: - content = (REPOSITORY_ROOT / relative_path).read_text() - assert "account: REPLACE_WITH_SLURM_ACCOUNT" in content - +def test_setup_defaults_example_identifies_required_slurm_account() -> None: defaults = (REPOSITORY_ROOT / "nv-internal/puzzletron_defaults.example.yaml").read_text() - assert "# account: REPLACE_WITH_SLURM_ACCOUNT" in defaults - - baremetal = ( - REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml" - ).read_text() - assert baremetal.count("REPLACE_WITH_PRIMARY_HOST") == 2 - assert baremetal.count("REPLACE_WITH_SECONDARY_HOST") == 1 - -def test_execution_contract_examples_explain_runnable_and_optional_values() -> None: - slurm_paths = ( - "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml", - "nv-internal/puzzletron_defaults.example.yaml", - ) - - for relative_path in slurm_paths: - content = (REPOSITORY_ROOT / relative_path).read_text() - assert "checkout visible on every worker" in content - assert "Sourced as /bin/activate" in content - assert "accepted by the cluster's srun container plugin" in content - assert "/data:/data,/models:/models" in content - assert "before virtualenv activation" in content + assert "# account: REPLACE_WITH_SLURM_ACCOUNT" in defaults From 99326c5adfefb0dd98644bad3c24dae660430b15 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Sat, 1 Aug 2026 13:59:54 +0200 Subject: [PATCH 6/7] Move defaults example and tighten portability tests Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 2 +- .../configs/setup/defaults.example.yaml | 0 .../torch/puzzletron/test_portable_configs.py | 81 +++++++++---------- 3 files changed, 37 insertions(+), 46 deletions(-) rename nv-internal/puzzletron_defaults.example.yaml => examples/puzzletron/configs/setup/defaults.example.yaml (100%) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 87effb28c5f..53b5d23ea01 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -80,7 +80,7 @@ local defaults-versus-customize decisions at every section: ```bash python examples/puzzletron/puzzletron_setup_v2.py \ - --defaults nv-internal/puzzletron_defaults.example.yaml + --defaults examples/puzzletron/configs/setup/defaults.example.yaml ``` The example defaults use only repository-relative values. Copy the file and add diff --git a/nv-internal/puzzletron_defaults.example.yaml b/examples/puzzletron/configs/setup/defaults.example.yaml similarity index 100% rename from nv-internal/puzzletron_defaults.example.yaml rename to examples/puzzletron/configs/setup/defaults.example.yaml diff --git a/tests/unit/torch/puzzletron/test_portable_configs.py b/tests/unit/torch/puzzletron/test_portable_configs.py index de372b674f8..25f414a3a20 100644 --- a/tests/unit/torch/puzzletron/test_portable_configs.py +++ b/tests/unit/torch/puzzletron/test_portable_configs.py @@ -19,39 +19,38 @@ import yaml -from puzzletron_orchestrator.compiler import load_runner_config +from puzzletron_orchestrator.compiler import load_execution_config, load_runner_config from puzzletron_setup.v2.defaults import load_defaults REPOSITORY_ROOT = Path(__file__).resolve().parents[4] -def test_slurm_runner_example_uses_portable_defaults() -> None: +def test_slurm_runner_example_is_portable() -> None: slurm = load_runner_config( REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.slurm.example.yaml" ) - assert slurm.contract.repository == "." - assert slurm.contract.venv == ".venv" + assert not Path(slurm.contract.repository).is_absolute() + assert not Path(slurm.contract.venv).is_absolute() assert slurm.contract.container is None assert slurm.contract.container_mounts is None assert not slurm.contract.prerun_commands assert slurm.slurm is not None - assert slurm.slurm.account == "REPLACE_WITH_SLURM_ACCOUNT" + assert slurm.slurm.account.startswith("REPLACE_WITH_") assert slurm.slurm.partition_cpu is None -def test_baremetal_runner_example_uses_portable_defaults() -> None: +def test_baremetal_runner_example_is_portable() -> None: baremetal = load_runner_config( REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/runner.baremetal.example.yaml" ) - assert baremetal.contract.repository == "." - assert baremetal.contract.venv == ".venv" + assert not Path(baremetal.contract.repository).is_absolute() + assert not Path(baremetal.contract.venv).is_absolute() assert baremetal.contract.setup_env is None assert baremetal.baremetal is not None - assert baremetal.baremetal.rendezvous_host == "REPLACE_WITH_PRIMARY_HOST" - assert [host.hostname for host in baremetal.baremetal.hosts] == [ - "REPLACE_WITH_PRIMARY_HOST", - "REPLACE_WITH_SECONDARY_HOST", - ] + hostnames = [host.hostname for host in baremetal.baremetal.hosts] + assert hostnames + assert baremetal.baremetal.rendezvous_host in hostnames + assert all(hostname.startswith("REPLACE_WITH_") for hostname in hostnames) def test_qwen_slurm_runner_preserves_portable_environment_contract() -> None: @@ -59,44 +58,43 @@ def test_qwen_slurm_runner_preserves_portable_environment_contract() -> None: REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/qwen_moe/runner.slurm.yaml" ) - assert runner.contract.repository == "REPLACE_WITH_WORKER_VISIBLE_MODELOPT_CHECKOUT" - assert runner.contract.venv == "REPLACE_WITH_WORKER_VISIBLE_MODELOPT_VENV" - assert runner.contract.container == "REPLACE_WITH_SLURM_CONTAINER_IMAGE" - assert runner.contract.container_mounts == "REPLACE_WITH_HOST_PATH:REPLACE_WITH_CONTAINER_PATH" - assert runner.contract.prerun_commands == ( - "source REPLACE_WITH_SITE_SETUP_SCRIPT", - "export VLLM_ROOT=REPLACE_WITH_WORKER_VISIBLE_VLLM_CHECKOUT", - "export AUTOMODEL_ROOT=REPLACE_WITH_WORKER_VISIBLE_AUTOMODEL_CHECKOUT", + contract_values = ( + runner.contract.repository, + runner.contract.venv, + runner.contract.container, + runner.contract.container_mounts, ) + assert all(value and value.startswith("REPLACE_WITH_") for value in contract_values) + assert runner.contract.prerun_commands + assert all("REPLACE_WITH_" in command for command in runner.contract.prerun_commands) assert runner.slurm is not None - assert runner.slurm.account == "REPLACE_WITH_SLURM_ACCOUNT" + assert runner.slurm.account.startswith("REPLACE_WITH_") assert runner.slurm.partition_cpu is None -def test_execution_example_is_valid_yaml() -> None: +def test_execution_example_is_loadable() -> None: path = REPOSITORY_ROOT / "examples/puzzletron/configs/orchestration/execution.example.yaml" - config = yaml.safe_load(path.read_text()) + execution = load_execution_config(path) - assert config["execution"]["defaults"]["gpus_per_node"] == 8 + assert set(execution) >= {"defaults", "stages"} -def test_setup_defaults_are_portable_and_valid() -> None: - path = REPOSITORY_ROOT / "nv-internal/puzzletron_defaults.example.yaml" +def test_setup_defaults_example_is_portable() -> None: + path = REPOSITORY_ROOT / "examples/puzzletron/configs/setup/defaults.example.yaml" defaults = load_defaults(path) - assert defaults["infrastructure"]["execution_contract"] == { - "repository": ".", - "venv": ".venv", - "container": None, - "container_mounts": None, - "prerun_commands": [], - } - assert defaults["infrastructure"]["runner"]["slurm"]["partition_cpu"] is None + contract = defaults["infrastructure"]["execution_contract"] + assert not Path(contract["repository"]).is_absolute() + assert not Path(contract["venv"]).is_absolute() + assert contract["container"] is None + assert contract["container_mounts"] is None + assert not contract["prerun_commands"] - defaults_paths = sorted((REPOSITORY_ROOT / "nv-internal").glob("*_defaults*.yaml")) - assert [path.name for path in defaults_paths] == ["puzzletron_defaults.example.yaml"] + slurm = defaults["infrastructure"]["runner"]["slurm"] + assert "account" not in slurm + assert slurm["partition_cpu"] is None def test_model_examples_use_public_hugging_face_identities() -> None: @@ -107,14 +105,7 @@ def test_model_examples_use_public_hugging_face_identities() -> None: ) for relative_path in paths: - content = (REPOSITORY_ROOT / relative_path).read_text() - config = yaml.safe_load(content) + config = yaml.safe_load((REPOSITORY_ROOT / relative_path).read_text()) assert config["input_hf_model_path"] == config["model_info"]["hf_repo"] assert not config["input_hf_model_path"].startswith("REPLACE_WITH_") - - -def test_setup_defaults_example_identifies_required_slurm_account() -> None: - defaults = (REPOSITORY_ROOT / "nv-internal/puzzletron_defaults.example.yaml").read_text() - - assert "# account: REPLACE_WITH_SLURM_ACCOUNT" in defaults From 66d59ff7d152d565ee32f8317860e8e89d6a66dc Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Sat, 1 Aug 2026 20:09:58 +0200 Subject: [PATCH 7/7] Address Puzzletron portability review feedback Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 2 +- .../puzzletron/configs/families/qwen3_5/qwen3p5_9b/model.yaml | 1 + puzzletron_setup/wizard.py | 2 +- tests/unit/torch/puzzletron/test_portable_configs.py | 2 ++ 4 files changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 53b5d23ea01..4d9aab4e058 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -22,7 +22,7 @@ campaign report. The configs below are the exact current-code entry points for the completed campaigns. Each verified report is a self-contained HTML file that embeds all -sanity-check outputs, stage manifests, and evaluation results, they can be +sanity-check outputs, stage manifests, and evaluation results; they can be 100s of MB. Download them to disk and open locally rather than previewing in a browser tab. diff --git a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_9b/model.yaml b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_9b/model.yaml index e162df8cc9c..07f53ca4091 100644 --- a/examples/puzzletron/configs/families/qwen3_5/qwen3p5_9b/model.yaml +++ b/examples/puzzletron/configs/families/qwen3_5/qwen3p5_9b/model.yaml @@ -7,6 +7,7 @@ input_hf_model_path: Qwen/Qwen3.5-9B model_info: # Canonical public identity used for model metadata and revision checks. hf_repo: Qwen/Qwen3.5-9B + hf_revision: 21eca8a083a2121a92fba681f4a7c72cf20ff1a7 model_type: qwen3_5_text architectures: [Qwen3_5ForConditionalGeneration] num_hidden_layers: 32 diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index 65ead3422f1..6c5f9f59b76 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -1213,7 +1213,7 @@ def _ask_infrastructure( ) repository = prompts.text( "Repository path on workers:", - default=str(Path.cwd()), + default=".", description=( "Use the checkout path visible on every worker and inside the container, if used." ), diff --git a/tests/unit/torch/puzzletron/test_portable_configs.py b/tests/unit/torch/puzzletron/test_portable_configs.py index 25f414a3a20..084e65087fc 100644 --- a/tests/unit/torch/puzzletron/test_portable_configs.py +++ b/tests/unit/torch/puzzletron/test_portable_configs.py @@ -15,6 +15,7 @@ """Regression tests for portable Puzzletron example configuration.""" +import re from pathlib import Path import yaml @@ -109,3 +110,4 @@ def test_model_examples_use_public_hugging_face_identities() -> None: assert config["input_hf_model_path"] == config["model_info"]["hf_repo"] assert not config["input_hf_model_path"].startswith("REPLACE_WITH_") + assert re.fullmatch(r"[0-9a-f]{40}", config["model_info"]["hf_revision"])