From 096944100b3b2bcc4d7ce63de669940457f2e1e1 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 11 Aug 2026 21:17:07 +0200 Subject: [PATCH 01/24] Add Puzzletron v2 GPU quality baseline Replace the legacy model matrix with one hermetic current-route campaign, and fix the configuration and orchestration contracts it exposes. Keep the dedicated GPU target separate from generic GPU coverage. Signed-off-by: Johannes Rausch --- .../distributed_eval/run_coordinator.sh | 13 +- examples/puzzletron/embedding_pipeline.py | 1 + .../finalize_replacement_scoring.py | 130 ++++- .../distributed_eval/automodel_executor.py | 15 + .../puzzletron/orchestration/adapters/pool.py | 26 +- .../torch/puzzletron/orchestration/config.py | 38 +- .../puzzletron/orchestration/controller.py | 393 ++++++++++++-- .../puzzletron/orchestration/task_launcher.py | 15 +- modelopt/torch/puzzletron/stages/graph.py | 8 +- noxfile.py | 105 ++-- puzzletron_setup/bundle.py | 4 +- tests/gpu/torch/puzzletron/test_puzzletron.py | 418 +++++++++++++++ .../test_automodel_solution_scoring.py | 4 + .../test_orchestration_executors.py | 40 ++ .../test_orchestration_lightweight.py | 77 +++ .../test_orchestration_shutdown_progress.py | 486 +++++++++++++++++- .../test_orchestration_task_topology.py | 46 +- .../torch/puzzletron/test_setup_bundle.py | 10 +- .../unit/torch/puzzletron/test_stage_graph.py | 25 + .../torch/puzzletron/test_width_scenarios.py | 111 ++++ .../test_width_slice_equivalence.py | 5 +- 21 files changed, 1841 insertions(+), 129 deletions(-) create mode 100644 tests/gpu/torch/puzzletron/test_puzzletron.py diff --git a/examples/puzzletron/distributed_eval/run_coordinator.sh b/examples/puzzletron/distributed_eval/run_coordinator.sh index ac22b6774fd..62f5c2edf27 100755 --- a/examples/puzzletron/distributed_eval/run_coordinator.sh +++ b/examples/puzzletron/distributed_eval/run_coordinator.sh @@ -88,6 +88,11 @@ from pathlib import Path import subprocess import sys +from examples.puzzletron.finalize_replacement_scoring import ( + finalization_marker_is_current, + write_finalization_marker, +) + ( completion_dir_text, marker_name, @@ -104,8 +109,12 @@ completion_dir.mkdir(parents=True, exist_ok=True) with (completion_dir / ".finalize.lock").open("a+") as lock: fcntl.flock(lock, fcntl.LOCK_EX) finalized = completion_dir / "finalized" - if finalized.is_file(): + root = Path(puzzle_dir) + root_summary = root / "artifacts" / "replacement_scoring" / "summary.json" + root_manifest = root / "manifests" / "replacement_scoring.json" + if finalization_marker_is_current(finalized, root_manifest, root_summary): raise SystemExit(0) + finalized.unlink(missing_ok=True) completed = tuple(completion_dir.glob("*.done")) expected = int(expected_text) if len(completed) < expected: @@ -126,7 +135,7 @@ with (completion_dir / ".finalize.lock").open("a+") as lock: ], check=True, ) - finalized.touch() + write_finalization_marker(finalized, root_manifest) PY else "${PYTHON_BIN}" "${SCRIPT_DIR}/../finalize_replacement_scoring.py" \ diff --git a/examples/puzzletron/embedding_pipeline.py b/examples/puzzletron/embedding_pipeline.py index 1b6346563d0..92d5f9340ea 100644 --- a/examples/puzzletron/embedding_pipeline.py +++ b/examples/puzzletron/embedding_pipeline.py @@ -148,6 +148,7 @@ def _scenario_overrides(config: dict, scenario: Path) -> tuple[str, ...]: f"teacher_dir={teacher}", f"convert.teacher_dir={teacher}", "bypass.enabled=false", + "embedding_pruning.enabled=false", f"replacement_library_path={scenario / 'replacement_library.json'}", f"build_replacement_library.source_checkpoint_dir={teacher}", "calc_subblock_stats.runtime_stats.execution=inline", diff --git a/examples/puzzletron/finalize_replacement_scoring.py b/examples/puzzletron/finalize_replacement_scoring.py index d8d3244eb6f..e3514f0ea91 100644 --- a/examples/puzzletron/finalize_replacement_scoring.py +++ b/examples/puzzletron/finalize_replacement_scoring.py @@ -7,47 +7,121 @@ from __future__ import annotations import argparse +import json +import os from pathlib import Path -from embedding_pipeline import finalize_replacement_scoring_diagnostics +if __package__: + from .embedding_pipeline import finalize_replacement_scoring_diagnostics +else: + from embedding_pipeline import finalize_replacement_scoring_diagnostics from modelopt.torch.puzzletron.diagnostics import generate_replace_block_report +from modelopt.torch.puzzletron.manifest import StageManifest, write_stage_manifest from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path +def _successful_manifest_identity(manifest_path: str | Path) -> str | None: + try: + manifest = json.loads(Path(manifest_path).read_text()) + except (OSError, ValueError): + return None + if manifest.get("stage") != "replacement_scoring" or manifest.get("status") != "success": + return None + identity = manifest.get("semantic_identity") + return str(identity) if identity else None + + +def finalization_marker_is_current( + marker_path: str | Path, + manifest_path: str | Path, + summary_path: str | Path, +) -> bool: + """Return whether a pool marker names the currently published result.""" + + try: + marker_identity = Path(marker_path).read_text().strip() + summary = json.loads(Path(summary_path).read_text()) + manifest = json.loads(Path(manifest_path).read_text()) + except OSError: + return False + except ValueError: + return False + return bool( + marker_identity + and marker_identity == _successful_manifest_identity(manifest_path) + and summary == (manifest.get("outputs") or {}).get("report") + ) + + +def write_finalization_marker(marker_path: str | Path, manifest_path: str | Path) -> None: + """Atomically bind a pool marker to the published manifest identity.""" + + identity = _successful_manifest_identity(manifest_path) + if identity is None: + raise RuntimeError(f"replacement-scoring manifest is not successful: {manifest_path}") + marker = Path(marker_path) + temporary = marker.with_suffix(marker.suffix + ".tmp") + temporary.write_text(identity + "\n") + temporary.replace(marker) + + +def finalize_replacement_scoring( + config_path: str | Path, + puzzle_dir: str | Path, + *, + overrides: list[str] | None = None, +) -> dict: + """Publish replacement reports and their canonical terminal manifest.""" + + config = pipeline_config_from_path(config_path, overrides=overrides) + config["puzzle_dir"] = str(puzzle_dir) + embedding = config.get("embedding_pruning") or {} + if bool(embedding.get("enabled", False)): + report = finalize_replacement_scoring_diagnostics(config) + else: + puzzle_dir = Path(puzzle_dir) + scoring = config.get("replacement_scoring") or {} + granularity = str(scoring.get("granularity", "block")) + stem = ( + "single_subblock_replacement_solutions" + if granularity == "subblock" + else "single_sequence_replacement_solutions" + ) + report = generate_replace_block_report( + puzzle_dir, + scores_dir=puzzle_dir / f"{stem}--validation", + output_dir=puzzle_dir / "artifacts" / "replacement_scoring", + granularity=granularity, + default_metric=str( + scoring.get("default_metric", "normalized_mse_loss_hidden_states") + ), + default_layer_count=int(scoring.get("default_layer_count", 5)), + anchor_count=int(scoring.get("anchor_count", 3)), + trend_relative_tolerance=float(scoring.get("trend_relative_tolerance", 0.02)), + ) + + manifest = StageManifest(stage="replacement_scoring", inputs={"config": config}, config=config) + manifest.complete(outputs={"report": report}) + write_stage_manifest( + Path(puzzle_dir) / "manifests" / "replacement_scoring.json", + manifest, + ) + return report + + def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", required=True) parser.add_argument("--puzzle-dir", required=True) args = parser.parse_args() - config = pipeline_config_from_path(args.config) - config["puzzle_dir"] = args.puzzle_dir - embedding = config.get("embedding_pruning") or {} - if bool(embedding.get("enabled", False)): - finalize_replacement_scoring_diagnostics(config) - return - - puzzle_dir = Path(args.puzzle_dir) - scoring = config.get("replacement_scoring") or {} - granularity = str(scoring.get("granularity", "block")) - stem = ( - "single_subblock_replacement_solutions" - if granularity == "subblock" - else "single_sequence_replacement_solutions" - ) - generate_replace_block_report( - puzzle_dir, - scores_dir=puzzle_dir / f"{stem}--validation", - output_dir=puzzle_dir / "artifacts" / "replacement_scoring", - granularity=granularity, - default_metric=str( - scoring.get("default_metric", "normalized_mse_loss_hidden_states") - ), - default_layer_count=int(scoring.get("default_layer_count", 5)), - anchor_count=int(scoring.get("anchor_count", 3)), - trend_relative_tolerance=float(scoring.get("trend_relative_tolerance", 0.02)), - ) + overrides = [ + override + for override in os.environ.get("FINALIZE_OVERRIDES", "").splitlines() + if override + ] + finalize_replacement_scoring(args.config, args.puzzle_dir, overrides=overrides) if __name__ == "__main__": diff --git a/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py b/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py index 531070d2620..f4db0f1ae27 100644 --- a/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py +++ b/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py @@ -45,6 +45,8 @@ def __init__(self, hydra_cfg): self.source_hidden_width = None self.sliced_teacher_baseline = None self.latest_observability = None + self.latest_score_device_type = None + self.visible_cuda_device_count = None self._setup_complete = False def capabilities(self) -> dict: @@ -247,10 +249,15 @@ def evaluate(self, request: EvaluationRequest) -> EvaluationResult | None: "hidden_width": self.source_hidden_width, "sliced_teacher_baseline": self.sliced_teacher_baseline, "observability": self.latest_observability, + "score_device_type": getattr(self, "latest_score_device_type", None), + "visible_cuda_device_count": getattr( + self, "visible_cuda_device_count", None + ), }, ) def _score(self, prune_target: dict | list[dict] | None) -> dict | None: + import torch import torch.distributed as torch_dist import modelopt.torch.utils.distributed as dist @@ -264,6 +271,7 @@ def _score(self, prune_target: dict | list[dict] | None) -> dict | None: recipe = self.recipe cache = self.cache params = self.params + self.visible_cuda_device_count = torch.cuda.device_count() per_batch = [] tp_group = recipe.tensor_parallel_group() candidate_lm_head = recipe.lm_head_weight() if recipe.has_outputs else None @@ -298,6 +306,13 @@ def _score(self, prune_target: dict | list[dict] | None) -> dict | None: for batch_idx, (hidden, targets) in enumerate(recipe.iterate_captures()): if hidden is None: continue + device_type = hidden.device.type + if self.latest_score_device_type not in (None, device_type): + raise RuntimeError( + "AutoModel scoring tensors changed device type within one executor: " + f"{self.latest_score_device_type!r} -> {device_type!r}" + ) + self.latest_score_device_type = device_type teacher_hidden = cache.hidden( batch_idx, device=hidden.device, diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index 73a489fe9d3..8c506b23614 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/pool.py @@ -7,6 +7,7 @@ from pathlib import Path +from ..identity import stable_hash from ..schema import ( AttemptSpec, CampaignPlan, @@ -19,6 +20,7 @@ WorkItem, WorkPlan, ) +from ..stages import semantic_stage_config from .base import WorkAdapter from .packing import packed_allocation from .stage_compat import stage_is_complete, stage_output_patterns @@ -130,6 +132,23 @@ def _replacement_overrides(plan: CampaignPlan, puzzle_dir: Path) -> tuple[str, . return tuple(overrides) +def _replacement_completion_identity( + plan: CampaignPlan, + root_overrides: list[str], +) -> str: + return stable_hash( + { + "contract_hash": plan.contract_hash, + "semantic_config": semantic_stage_config( + plan.experiment_config, + "replacement_scoring", + ), + "root_overrides": root_overrides, + }, + prefix="replacement_scoring_completion", + ) + + class PersistentPoolAdapter(WorkAdapter): """Launch one coordinator plus resident worker pool.""" @@ -233,7 +252,8 @@ def command( else plan.puzzle_dir ) campaign_dir = replacement_puzzle_dir / "distributed_eval" / node.stage_id - effective_overrides = list(overrides or []) + root_overrides = list(overrides or []) + effective_overrides = list(root_overrides) if node.stage_id == "replacement_scoring": effective_overrides.extend(_replacement_overrides(plan, replacement_puzzle_dir)) if role == "gang": @@ -254,6 +274,7 @@ def command( script = repo / "examples/puzzletron/distributed_eval/run_depth_pool.sh" else: env.update(_replacement_environment(plan, replacement_puzzle_dir)) + env["FINALIZE_OVERRIDES"] = "\n".join(root_overrides) replacement_widths = _replacement_widths(plan) if len(replacement_widths) > 1: width = int(item.metadata["width"]) @@ -264,7 +285,7 @@ def command( / "artifacts" / "replacement_scoring" / ".pool_completion" - / plan.contract_hash + / _replacement_completion_identity(plan, root_overrides) ), "FINALIZE_COMPLETION_MARKER": f"width-{width}", "FINALIZE_EXPECTED_COMPLETIONS": str( @@ -326,6 +347,7 @@ def command( env["DISTRIBUTED_EVAL_OVERRIDES"] = f"{existing}\n{override}".strip() if node.stage_id == "replacement_scoring": env.update(_replacement_environment(plan, replacement_puzzle_dir)) + env["FINALIZE_OVERRIDES"] = "\n".join(root_overrides) elif node.stage_id == "depth_importance": depth = plan.experiment_config.get("depth_importance") or {} env["OUTPUT_DIR"] = str( diff --git a/modelopt/torch/puzzletron/orchestration/config.py b/modelopt/torch/puzzletron/orchestration/config.py index 0e1c1594216..190ddc50372 100644 --- a/modelopt/torch/puzzletron/orchestration/config.py +++ b/modelopt/torch/puzzletron/orchestration/config.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. """Lightweight experiment-config composition for the Puzzletron controller.""" @@ -16,6 +28,26 @@ __all__ = ["load_experiment_config"] _INTERPOLATION = re.compile(r"\$\{([^${}]*)\}") +_SCIENTIFIC_FLOAT = re.compile(r"^[+-]?[0-9][0-9_]*[eE][+-]?[0-9]+$") + + +class _HydraSafeLoader(yaml.SafeLoader): + """Parse plain scientific notation with Hydra-compatible numeric semantics.""" + + +_HydraSafeLoader.add_implicit_resolver( + "tag:yaml.org,2002:float", + _SCIENTIFIC_FLOAT, + list("-+0123456789"), +) + + +def _load_yaml(value: str) -> Any: + loader = _HydraSafeLoader(value) + try: + return loader.get_single_data() + finally: + loader.dispose() def _mapping(value: Any, *, source: Path) -> dict[str, Any]: @@ -60,7 +92,7 @@ def _compose(path: Path, *, root: Path, stack: tuple[Path, ...]) -> dict[str, An if path in stack: chain = " -> ".join(str(item) for item in (*stack, path)) raise ValueError(f"Config defaults cycle: {chain}") - payload = _mapping(yaml.safe_load(path.read_text()), source=path) + payload = _mapping(_load_yaml(path.read_text()), source=path) defaults = payload.pop("defaults", []) if not isinstance(defaults, list): raise ValueError(f"defaults must be a list: {path}") @@ -104,7 +136,7 @@ def _resolve_expression(expression: str, config: Mapping[str, Any]) -> Any: if expression.startswith("to_path:"): return expression.removeprefix("to_path:") if expression.startswith("get_object:"): - return "${" + expression + "}" + return {"__type__": expression.removeprefix("get_object:")} try: return deepcopy(_lookup(config, expression)) except KeyError: @@ -154,7 +186,7 @@ def _apply_override(config: dict[str, Any], override: str) -> None: if not isinstance(child, dict): raise ValueError(f"Override path crosses a scalar: {override!r}") target = child - target[keys[-1]] = yaml.safe_load(raw_value) + target[keys[-1]] = _load_yaml(raw_value) def load_experiment_config( diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index f36c1105c48..48183aa0249 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.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 - """Durable campaign controller loop.""" from __future__ import annotations @@ -25,9 +22,9 @@ import time import uuid from collections.abc import Mapping -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path -from typing import Any +from typing import Any, Literal from .adapters.post_mip import ManualInputRequired from .adapters.registry import adapter_for_stage @@ -35,10 +32,12 @@ from .compiler import plan_to_dict from .dashboard import StageView, format_duration, progress_eta, progress_fraction from .executors import BareMetalSSHExecutor, Executor, LocalExecutor, SlurmExecutor +from .identity import stable_hash from .logging import OrchestratorLogger from .progress import summarize_stage_artifacts from .reporting import FinalReportResult, build_final_report_attempt, final_report_paths from .schema import ( + AttemptSpec, CampaignPlan, FailureClass, FailurePolicy, @@ -47,8 +46,10 @@ JobState, JobStatus, StagePlanNode, + ValidatedResult, + WorkPlan, ) -from .stages import stage_display_name +from .stages import semantic_stage_config, stage_display_name from .state import ( CampaignStateStore, PersistedAttempt, @@ -62,6 +63,9 @@ __all__ = ["CampaignController", "create_executor", "dry_run_plan"] +_ARTIFACT_SETTLING_TIMEOUT_SECONDS = 300.0 + + def create_executor(plan: CampaignPlan, *, local: bool = False) -> Executor: if local: return LocalExecutor(plan.runner) @@ -116,6 +120,14 @@ class DryRunSubmission: argv: tuple[str, ...] +@dataclass(frozen=True) +class _FinalizationFailure: + phase: Literal["aggregation", "validation"] + reason: str + artifacts: tuple[str, ...] = () + exception_type: str | None = None + + def dry_run_plan( plan: CampaignPlan, *, @@ -189,6 +201,7 @@ def __init__( self._shutting_down = False self._interactive_ready = False self._failed_stages: set[str] = set() + self._finalization_failures: dict[str, _FinalizationFailure] = {} self._manual_waiting: ManualInputRequired | None = None defaults = dict(plan.execution_defaults or {}) self._halt_policy = HaltPolicy(str(defaults.get("halt_policy", HaltPolicy.DRAIN.value))) @@ -311,22 +324,129 @@ def _log_completed_stages(self) -> None: ): self.logger.skip(f"{node.stage_id}: completion artifacts validated") - @staticmethod - def _completed_work_ids(attempts: list[dict[str, Any]]) -> set[str]: - return { - str(attempt["work_id"]) - for attempt in attempts - if attempt.get("status") == JobState.COMPLETED.value - } + def _required_completed_attempts( + self, + node: StagePlanNode, + attempts: list[dict[str, Any]], + ) -> list[dict[str, Any]] | None: + work_plan = adapter_for_stage(node).plan(self.plan, node) + stage_execution_identity = self._stage_execution_identity(node, work_plan) + completed: list[dict[str, Any]] = [] + for item in work_plan.items: + matches = [ + attempt + for attempt in attempts + if attempt.get("work_id") == item.work_id + and attempt.get("status") == JobState.COMPLETED.value + and attempt.get("contract_hash") == self.plan.contract_hash + and isinstance(attempt.get("metadata"), Mapping) + and attempt["metadata"].get("stage_execution_identity") == stage_execution_identity + ] + if not matches: + return None + + def _completion_time(attempt: dict[str, Any]) -> float: + completed_at = attempt.get("completed_at") + if isinstance(completed_at, (int, float)): + return float(completed_at) + submitted_at = attempt.get("submitted_at") + return float(submitted_at) if isinstance(submitted_at, (int, float)) else 0.0 + + completed.append( + max( + matches, + key=_completion_time, + ) + ) + return completed + + def _stage_execution_identity( + self, + node: StagePlanNode, + work_plan: WorkPlan | None = None, + ) -> str: + work_plan = work_plan or adapter_for_stage(node).plan(self.plan, node) + compiled_node = next( + stage + for stage in plan_to_dict(self.plan)["stages"] + if stage["stage_id"] == node.stage_id + ) + return stable_hash( + { + "execution_contract_hash": self.plan.contract_hash, + "semantic_config": semantic_stage_config( + self.plan.experiment_config, node.stage_id + ), + "compiled_node": compiled_node, + "work_items": [ + { + "work_id": item.work_id, + "shard_index": item.shard_index, + "shard_count": item.shard_count, + "gpus_per_instance": item.gpus_per_instance, + "local_gpu_ids": list(item.local_gpu_ids), + "metadata": dict(item.metadata), + } + for item in work_plan.items + ], + }, + prefix=f"{node.stage_id}_execution", + ) + + def _bind_attempt_to_stage_execution( + self, + node: StagePlanNode, + work_plan: WorkPlan, + attempt: AttemptSpec, + ) -> AttemptSpec: + return replace( + attempt, + metadata={ + **dict(attempt.metadata), + "stage_execution_identity": self._stage_execution_identity(node, work_plan), + }, + ) def _required_work_is_completed( self, node: StagePlanNode, attempts: list[dict[str, Any]], ) -> bool: - work_plan = adapter_for_stage(node).plan(self.plan, node) - required = {item.work_id for item in work_plan.items} - return required.issubset(self._completed_work_ids(attempts)) + return self._required_completed_attempts(node, attempts) is not None + + def _legacy_completed_attempts( + self, + node: StagePlanNode, + attempts: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + work_ids = {item.work_id for item in adapter_for_stage(node).plan(self.plan, node).items} + return [ + attempt + for attempt in attempts + if attempt.get("work_id") in work_ids + and attempt.get("status") == JobState.COMPLETED.value + and attempt.get("contract_hash") == self.plan.contract_hash + and ( + not isinstance(attempt.get("metadata"), Mapping) + or "stage_execution_identity" not in attempt["metadata"] + ) + ] + + def _completed_work_artifact_settling_elapsed( + self, + node: StagePlanNode, + attempts: list[dict[str, Any]], + ) -> float | None: + completed = self._required_completed_attempts(node, attempts) + if completed is None: + return None + completed_at: list[float] = [] + for attempt in completed: + value = attempt.get("completed_at") + if not isinstance(value, (int, float)): + return _ARTIFACT_SETTLING_TIMEOUT_SECONDS + completed_at.append(float(value)) + return max(0.0, time.time() - max(completed_at)) def _policy_allows_retry(self, node: StagePlanNode, failure: FailureClass) -> bool: if failure in {FailureClass.SUCCESS, FailureClass.CANCELLED}: @@ -362,14 +482,94 @@ def _stage_has_active_or_completed_work(self, node: StagePlanNode) -> bool: return True if self._required_work_is_completed(node, attempts): # Aggregation is attempted before submission in the controller loop. - # If its outputs are still incomplete, rerun the work instead of - # remaining permanently blocked by historical completed attempts. - return False - return False + # A scheduler-successful attempt must never overlap with a duplicate + # while distributed filesystems are still publishing its artifacts. + # The controller loop either validates those outputs or records a + # bounded settling failure. Historical records without a completion + # timestamp fail validation immediately because their settling age + # cannot be established safely. + return self._completed_work_artifact_settling_elapsed(node, attempts) is not None + return bool(self._legacy_completed_attempts(node, attempts)) def _stage_is_active(self, stage_id: str) -> bool: return any(work_id.startswith(f"{stage_id}:") for _, work_id, _ in self._active.values()) + def _recover_failed_stages(self) -> None: + for node in self.plan.stages: + record = self.store.load_stage_record(node.stage_id) + stage_execution_identity = self._stage_execution_identity(node) + current_failure = bool(record and record.attempts) and all( + attempt.contract_hash == self.plan.contract_hash + and (attempt.metadata or {}).get("stage_execution_identity") + == stage_execution_identity + for attempt in record.attempts + ) + legacy_incompatibility = bool(record and record.attempts) and all( + attempt.contract_hash == self.plan.contract_hash + and (attempt.metadata or {}).get("stage_execution_identity_incompatible") is True + for attempt in record.attempts + ) + if ( + record is None + or record.status != JobState.FAILED.value + or not record.attempts + or not (current_failure or legacy_incompatibility) + or stage_is_complete(self.plan.experiment_config, node.stage_id) + ): + continue + self._failed_stages.add(node.stage_id) + self.logger.error(f"{node.stage_id}: recovered terminal stage validation failure") + + def _fail_legacy_completed_attempts( + self, + node: StagePlanNode, + attempts: list[dict[str, Any]], + ) -> bool: + legacy_attempts = self._legacy_completed_attempts(node, attempts) + if not legacy_attempts: + return False + persisted_attempts = [ + PersistedAttempt( + attempt_id=attempt["attempt_id"], + work_id=attempt["work_id"], + stage_id=node.stage_id, + status=attempt.get("status", JobState.COMPLETED.value), + contract_hash=attempt["contract_hash"], + metadata={ + **dict(attempt.get("metadata") or {}), + "stage_execution_identity_incompatible": True, + }, + ) + for attempt in legacy_attempts + ] + reason = ( + "completed attempt metadata predates stage execution identities; " + "refusing automatic resubmission" + ) + self._failed_stages.add(node.stage_id) + self.store.write_stage_record( + StageRunRecord( + stage_id=node.stage_id, + status=JobState.FAILED.value, + attempts=persisted_attempts, + aggregated=False, + ) + ) + self.store.append_event( + "stage_execution_identity_incompatible", + { + "stage_id": node.stage_id, + "failure_class": FailureClass.CONFIG.value, + "contract_hash": self.plan.contract_hash, + "stage_execution_identity": self._stage_execution_identity(node), + "attempt_ids": [attempt.attempt_id for attempt in persisted_attempts], + "incompatibility": "missing_stage_execution_identity", + "reason": reason, + }, + ) + self.logger.error(f"{node.stage_id}: {reason}") + return True + def _available_nodes(self) -> int | None: slurm = self.plan.runner.slurm if slurm is None or slurm.max_nodes is None: @@ -409,13 +609,17 @@ def _submit_stage(self, node: StagePlanNode, *, overrides: list[str] | None = No self.logger.success(f"{item.work_id}: already complete, skipping") continue attempt_id = str(uuid.uuid4()) - attempt = adapter.command( - plan=self.plan, - node=node, - item=item, - attempt_id=attempt_id, - runner=self.plan.runner, - overrides=overrides, + attempt = self._bind_attempt_to_stage_execution( + node, + work_plan, + adapter.command( + plan=self.plan, + node=node, + item=item, + attempt_id=attempt_id, + runner=self.plan.runner, + overrides=overrides, + ), ) if available_nodes is not None and attempt.allocation_nodes > available_nodes: self.logger.wait( @@ -491,32 +695,127 @@ def _finalize_stage(self, node: StagePlanNode) -> bool: + "\n" ) temporary.replace(decision_path) - aggregate = adapter.aggregate(plan=self.plan, node=node, work_plan=work_plan) + try: + aggregate = adapter.aggregate(plan=self.plan, node=node, work_plan=work_plan) + except (OSError, ValueError, RuntimeError) as error: + return self._record_stage_aggregation_failure( + node, + error, + ) + except (OSError, ValueError, RuntimeError) as error: + return self._record_stage_aggregation_failure( + node, + error, + ) validation = adapter.validate(plan=self.plan, node=node) if not validation.valid: - self.logger.warning(f"{node.stage_id}: {validation.reason}") - return False - attempts = [ + return self._record_stage_validation_failure(node, validation) + self._finalization_failures.pop(node.stage_id, None) + attempts = self._persisted_stage_attempts(node) + self.store.write_stage_record( + StageRunRecord( + stage_id=node.stage_id, + status=JobState.COMPLETED.value, + attempts=attempts, + aggregated=aggregate is not None, + ) + ) + self.store.append_event("stage_completed", {"stage_id": node.stage_id}) + self.logger.success( + f"{node.stage_id} complete; artifacts={', '.join(validation.artifacts) or 'validated'}" + ) + return True + + def _record_stage_validation_failure( + self, + node: StagePlanNode, + validation: ValidatedResult, + ) -> bool: + self._finalization_failures[node.stage_id] = _FinalizationFailure( + phase="validation", + reason=validation.reason, + artifacts=validation.artifacts, + ) + self.logger.warning(f"{node.stage_id}: {validation.reason}") + return False + + def _record_stage_aggregation_failure( + self, + node: StagePlanNode, + error: OSError | ValueError | RuntimeError, + ) -> bool: + reason = f"stage aggregation failed: {type(error).__name__}: {error}" + self._finalization_failures[node.stage_id] = _FinalizationFailure( + phase="aggregation", + reason=reason, + exception_type=type(error).__name__, + ) + self.logger.warning(f"{node.stage_id}: {reason}") + return False + + def _persisted_stage_attempts(self, node: StagePlanNode) -> list[PersistedAttempt]: + stage_execution_identity = self._stage_execution_identity(node) + return [ PersistedAttempt( attempt_id=attempt["attempt_id"], work_id=attempt["work_id"], stage_id=node.stage_id, status=attempt.get("status", JobState.COMPLETED.value), - contract_hash=attempt.get("contract_hash", self.plan.contract_hash), + contract_hash=attempt["contract_hash"], + metadata=dict(attempt["metadata"]), ) for attempt in self.store.list_attempts(node.stage_id) + if attempt.get("contract_hash") == self.plan.contract_hash + and isinstance(attempt.get("metadata"), Mapping) + and attempt["metadata"].get("stage_execution_identity") == stage_execution_identity ] + + def _fail_stage_if_artifacts_did_not_settle( + self, + node: StagePlanNode, + attempts: list[dict[str, Any]], + ) -> bool: + if node.stage_id in self._failed_stages: + return True + elapsed = self._completed_work_artifact_settling_elapsed(node, attempts) + if elapsed is None or elapsed < _ARTIFACT_SETTLING_TIMEOUT_SECONDS: + return False + failure = self._finalization_failures.get(node.stage_id) + reason = failure.reason if failure is not None else "required artifacts are incomplete" + expected_artifacts = list(failure.artifacts) if failure is not None else [] + phase = failure.phase if failure is not None else "validation" + persisted_attempts = self._persisted_stage_attempts(node) + self._failed_stages.add(node.stage_id) self.store.write_stage_record( StageRunRecord( stage_id=node.stage_id, - status=JobState.COMPLETED.value, - attempts=attempts, - aggregated=aggregate is not None, + status=JobState.FAILED.value, + attempts=persisted_attempts, + aggregated=False, ) ) - self.store.append_event("stage_completed", {"stage_id": node.stage_id}) - self.logger.success( - f"{node.stage_id} complete; artifacts={', '.join(validation.artifacts) or 'validated'}" + event_type = ( + "stage_aggregation_failed" if phase == "aggregation" else "stage_validation_failed" + ) + self.store.append_event( + event_type, + { + "stage_id": node.stage_id, + "phase": phase, + "exception_type": failure.exception_type if failure is not None else None, + "failure_class": FailureClass.TIMEOUT_FATAL.value, + "contract_hash": self.plan.contract_hash, + "stage_execution_identity": self._stage_execution_identity(node), + "attempt_ids": [attempt.attempt_id for attempt in persisted_attempts], + "elapsed_seconds": elapsed, + "timeout_seconds": _ARTIFACT_SETTLING_TIMEOUT_SECONDS, + "reason": reason, + "expected_artifacts": expected_artifacts, + }, + ) + self.logger.error( + f"{node.stage_id}: completed work outputs did not settle within " + f"{_ARTIFACT_SETTLING_TIMEOUT_SECONDS:g}s: {reason}" ) return True @@ -1029,6 +1328,7 @@ def _on_signal(signum: int, _frame: object | None) -> None: f"instances={node.instances}, total_gpus={node.total_gpus}" ) self._recover_active_attempts() + self._recover_failed_stages() self._log_completed_stages() self._refresh_dashboard() self.terminal_controls.start() @@ -1052,14 +1352,27 @@ def _on_signal(signum: int, _frame: object | None) -> None: for node in self.plan.stages: if stage_is_complete(self.plan.experiment_config, node.stage_id): continue + if node.stage_id in self._failed_stages: + continue stage_attempts = self.store.list_attempts(node.stage_id) if stage_attempts and not self._stage_is_active(node.stage_id): if self._required_work_is_completed(node, stage_attempts): - self._finalize_stage(node) + finalized = self._finalize_stage(node) + if not finalized and self._manual_waiting is None: + self._fail_stage_if_artifacts_did_not_settle( + node, stage_attempts + ) if self._manual_waiting is not None: break + else: + self._fail_legacy_completed_attempts(node, stage_attempts) if self._manual_waiting is not None: break + if self._failed_stages and self._should_fail_fast(): + halted = True + self._refresh_dashboard() + self.shutdown(reason="fatal stage validation failure") + break for node in self._ready_nodes(): if self._shutdown_requested: break @@ -1160,7 +1473,7 @@ def _on_signal(signum: int, _frame: object | None) -> None: elif cancelled: self.logger.shutdown("campaign stopped by user; rerun the same command to resume") elif halted: - self.logger.error("campaign halted after a failed attempt") + self.logger.error("campaign halted after a stage failure") elif self._manual_waiting is not None: self.logger.wait("campaign paused for a durable manual-filter decision") else: diff --git a/modelopt/torch/puzzletron/orchestration/task_launcher.py b/modelopt/torch/puzzletron/orchestration/task_launcher.py index 4d4b21c414a..3f4a69ea2fd 100644 --- a/modelopt/torch/puzzletron/orchestration/task_launcher.py +++ b/modelopt/torch/puzzletron/orchestration/task_launcher.py @@ -22,6 +22,7 @@ "TaskBinding", "build_task_command", "main", + "rendezvous_endpoint", "rendezvous_port", "resolve_task_binding", ] @@ -76,6 +77,14 @@ def rendezvous_port(attempt_id: str, group_index: int, group_count: int) -> int: return port_start + (seed % (port_span - group_count + 1)) + group_index +def rendezvous_endpoint(binding: TaskBinding) -> str: + """Return a collision-free local endpoint or the shared multi-node endpoint.""" + + if binding.group_size == 1: + return "localhost:0" + return f"{binding.master_addr}:{binding.master_port}" + + def resolve_task_binding( *, attempt_id: str, @@ -131,7 +140,7 @@ def build_task_command( command = tuple(str(part) for part in payload) if launcher is TaskLauncher.DIRECT: return command - rendezvous_host = "localhost" if binding.group_size == 1 else binding.master_addr + endpoint = rendezvous_endpoint(binding) return ( "python", "-m", @@ -139,7 +148,7 @@ def build_task_command( f"--nnodes={binding.group_size}", f"--nproc-per-node={gpus_per_task}", "--rdzv-backend=c10d", - f"--rdzv-endpoint={rendezvous_host}:{binding.master_port}", + f"--rdzv-endpoint={endpoint}", f"--rdzv-id={binding.rendezvous_id}", "--no-python", *command, @@ -241,7 +250,7 @@ def main(argv: Sequence[str] | None = None) -> int: f"host={binding.hostname} task={binding.task_index} " f"local={binding.local_task_index} gpus={','.join(visible_gpus)} " f"group={binding.group_index} rank={binding.group_rank}/{binding.group_size} " - f"endpoint={binding.master_addr}:{binding.master_port} " + f"endpoint={rendezvous_endpoint(binding)} " f"rdzv_id={binding.rendezvous_id}", flush=True, ) diff --git a/modelopt/torch/puzzletron/stages/graph.py b/modelopt/torch/puzzletron/stages/graph.py index d1e2785d084..779a2b9e385 100644 --- a/modelopt/torch/puzzletron/stages/graph.py +++ b/modelopt/torch/puzzletron/stages/graph.py @@ -525,7 +525,13 @@ def semantic_stage_config(config: Mapping[str, Any], stage_id: str) -> dict[str, spec = STAGE_REGISTRY.get(stage_id) stage_sections = (stage_id,) if spec is None else spec.semantic_config_sections sections = dict.fromkeys((*SHARED_SEMANTIC_CONFIG_SECTIONS, *stage_sections)) - return {key: config[key] for key in sections if key in config} + return { + key: config[key] + for key in sections + if key in config + and config[key] is not None + and not (isinstance(config[key], Mapping) and not config[key]) + } def stage_display_name(stage_id: str, *, granularity: str | None = None) -> str: diff --git a/noxfile.py b/noxfile.py index ab87c28e1da..25b1a2cc4c2 100644 --- a/noxfile.py +++ b/noxfile.py @@ -65,6 +65,40 @@ ) +def _verify_puzzletron_v2_environment(session): + """Fail before collection when the dedicated Puzzletron runtime drifts.""" + expected_versions = { + "python": PUZZLETRON_V2_CI_ENVIRONMENT["python"], + "torch": PUZZLETRON_V2_CI_ENVIRONMENT["torch"], + "torchvision": PUZZLETRON_V2_CI_ENVIRONMENT["torchvision"], + "transformers": PUZZLETRON_V2_CI_ENVIRONMENT["transformers"], + "lmms-eval": PUZZLETRON_V2_CI_ENVIRONMENT["lmms_eval"], + "nemo-automodel": PUZZLETRON_V2_AUTOMODEL_SOURCE["base_version"], + } + session.run( + "python", + "-c", + ( + "import sys; " + "from importlib.metadata import version; " + "from packaging.version import Version; " + f"expected = {expected_versions!r}; " + "actual = {" + "'python': f'{sys.version_info.major}.{sys.version_info.minor}', " + "'torch': Version(version('torch')).base_version, " + "'torchvision': Version(version('torchvision')).base_version, " + "'transformers': Version(version('transformers')).base_version, " + "'lmms-eval': Version(version('lmms-eval')).base_version, " + "'nemo-automodel': Version(version('nemo-automodel')).base_version}; " + "mismatches = {name: (actual[name], expected_version) " + "for name, expected_version in expected.items() " + "if actual[name] != expected_version}; " + "assert not mismatches, " + "f'Pinned Puzzletron CI environment mismatch: {mismatches}'" + ), + ) + + def _cov_args(): """Return --cov when COVERAGE_PROCESS_START is set (CI only).""" return ["--cov"] if os.environ.get("COVERAGE_PROCESS_START") else [] @@ -109,36 +143,7 @@ def puzzletron_v2(session): PUZZLETRON_V2_AUTOMODEL, ) session.run("uv", "pip", "check") - expected_versions = { - "python": PUZZLETRON_V2_CI_ENVIRONMENT["python"], - "torch": PUZZLETRON_V2_CI_ENVIRONMENT["torch"], - "torchvision": PUZZLETRON_V2_CI_ENVIRONMENT["torchvision"], - "transformers": PUZZLETRON_V2_CI_ENVIRONMENT["transformers"], - "lmms-eval": PUZZLETRON_V2_CI_ENVIRONMENT["lmms_eval"], - "nemo-automodel": PUZZLETRON_V2_AUTOMODEL_SOURCE["base_version"], - } - session.run( - "python", - "-c", - ( - "import sys; " - "from importlib.metadata import version; " - "from packaging.version import Version; " - f"expected = {expected_versions!r}; " - "actual = {" - "'python': f'{sys.version_info.major}.{sys.version_info.minor}', " - "'torch': Version(version('torch')).base_version, " - "'torchvision': Version(version('torchvision')).base_version, " - "'transformers': Version(version('transformers')).base_version, " - "'lmms-eval': Version(version('lmms-eval')).base_version, " - "'nemo-automodel': Version(version('nemo-automodel')).base_version}; " - "mismatches = {name: (actual[name], expected_version) " - "for name, expected_version in expected.items() " - "if actual[name] != expected_version}; " - "assert not mismatches, " - "f'Pinned Puzzletron CI environment mismatch: {mismatches}'" - ), - ) + _verify_puzzletron_v2_environment(session) session.run( "python", "-m", @@ -171,6 +176,7 @@ def partial_unit(session, subset): # ─── GPU sessions (run inside containers — no new venv) ────────────────────── +# The generic and Puzzletron sessions use their dedicated container environments directly. # `venv_backend="none"` skips creating a new venv so the session runs directly in the container's # existing Python environment (e.g. /opt/venv in NeMo) instead of an isolated one. # Use `python -m pip/pytest` to ensure the container's active venv Python is used, @@ -200,7 +206,44 @@ def gpu(session): "git+https://github.com/state-spaces/mamba.git", "git+https://github.com/Dao-AILab/causal-conv1d.git", ) - session.run("python", "-m", "pytest", "tests/gpu", *_cov_args()) + session.run( + "python", + "-m", + "pytest", + "tests/gpu", + "--ignore=tests/gpu/torch/puzzletron/test_puzzletron.py", + *_cov_args(), + ) + + +# Container: dedicated Puzzletron v2 GPU image with the pinned ci_environment.json runtime. +@nox.session(venv_backend="none") +def gpu_puzzletron(session): + """Run the focused Puzzletron suite in its pinned one-GPU image.""" + _verify_puzzletron_v2_environment(session) + session.run( + "python", + "-c", + ( + "import torch; " + "assert torch.cuda.is_available(), 'Puzzletron GPU CI requires CUDA'; " + "assert torch.cuda.device_count() == 1, " + "f'Puzzletron GPU CI requires exactly one visible GPU, got {torch.cuda.device_count()}'; " + "assert torch.version.cuda == '12.9', " + "f'Puzzletron GPU CI requires CUDA 12.9, got {torch.version.cuda}'" + ), + ) + session.run( + "python", + "-m", + "pytest", + "-o", + "addopts=", + ( + "tests/gpu/torch/puzzletron/test_puzzletron.py::" + "test_tiny_qwen_campaign_uses_current_public_route" + ), + ) # Container: nvcr.io/nvidia/nemo:26.04 or later diff --git a/puzzletron_setup/bundle.py b/puzzletron_setup/bundle.py index 3ee3e8f0936..4713bfb8e9b 100644 --- a/puzzletron_setup/bundle.py +++ b/puzzletron_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 - """Render and validate portable Puzzletron setup bundles.""" from __future__ import annotations @@ -470,6 +467,7 @@ def render_experiment(state: Mapping[str, Any], budget: str) -> dict[str, Any]: "modality": data["modality"], "layout": data["layout"], "max_sample_length": sequence_length, + "sequence_length": sequence_length, "path": data["source"], "acquisition": deepcopy(data_acquisition) if data_acquisition else None, "packing": ( diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py new file mode 100644 index 00000000000..61f85a6f743 --- /dev/null +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -0,0 +1,418 @@ +# 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. + +"""Hermetic GPU coverage for the public Puzzletron campaign route.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from hashlib import sha256 +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import pytest +import torch +import yaml +from _test_utils.torch.transformers_models import create_tiny_qwen3_5_dir +from datasets import Dataset, DatasetDict +from transformers import AutoModelForCausalLM + +from modelopt.torch.puzzletron.orchestration.adapters.stage_compat import stage_is_complete +from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path +from modelopt.torch.puzzletron.stages.graph import enabled_stage_ids +from puzzletron_orchestrator.adapters.registry import adapter_for_stage +from puzzletron_orchestrator.compiler import ( + compile_campaign_plan, + load_execution_config, + load_runner_config, +) +from puzzletron_setup.v2.wizard import _DEFAULT_DATA_SOURCE, _DEFAULT_MODEL_SOURCE, run_wizard_v2 + +if TYPE_CHECKING: + from collections.abc import Sequence + + from puzzletron_setup.v2.prompts import PromptChoice + + +class _DefaultsBackend: + """Select resolved defaults while supplying the test campaign directory.""" + + def __init__(self, campaign_dir: Path) -> None: + self.campaign_dir = campaign_dir + + def text(self, message: str, default: str) -> Any: + if message == "Campaign directory:": + return str(self.campaign_dir) + return default + + def select( + self, + message: str, + choices: Sequence[PromptChoice], + default: Any, + ) -> Any: + if message == "Model:": + return _DEFAULT_MODEL_SOURCE + if message == "Dataset:": + return _DEFAULT_DATA_SOURCE + if message == "Post Mip:": + return "customize" + if message.startswith("Post-MIP flow for "): + return "none" + if default is not None: + return default + return next(choice.value for choice in choices if choice.disabled is None) + + def checkbox( + self, + message: str, + choices: Sequence[PromptChoice], + defaults: Sequence[Any], + ) -> Any: + del message, choices + return list(defaults) + + +def _save_messages_dataset(path: Path) -> None: + response = ( + "Compression removes redundant parameters while preserving useful model behavior. " * 16 + ).strip() + messages = [ + {"role": "user", "content": "What is model compression?"}, + {"role": "assistant", "content": response}, + ] + rows = [{"messages": messages}] * 8 + DatasetDict( + { + "train": Dataset.from_list(rows), + "validation": Dataset.from_list(rows), + } + ).save_to_disk(str(path)) + + +def _artifact_digests(paths: Sequence[Path]) -> dict[str, str]: + return {str(path): sha256(path.read_bytes()).hexdigest() for path in paths} + + +def _nested_values(value: Any, key: str) -> list[Any]: + values = [] + if isinstance(value, dict): + values.extend(item for name, item in value.items() if name == key) + for item in value.values(): + values.extend(_nested_values(item, key)) + elif isinstance(value, list): + for item in value: + values.extend(_nested_values(item, key)) + return values + + +def _tensor_values(value: Any): + if torch.is_tensor(value): + yield value + elif isinstance(value, dict): + for item in value.values(): + yield from _tensor_values(item) + elif isinstance(value, (list, tuple)): + for item in value: + yield from _tensor_values(item) + + +def _run_campaign( + project_root_path: Path, + smoke_bundle: Path, + environment: dict[str, str], +) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [ + sys.executable, + str(project_root_path / "examples/puzzletron/orchestrate.py"), + "--experiment", + str(smoke_bundle / "experiment.yaml"), + "--runner", + str(smoke_bundle / "runner.yaml"), + "--execution", + str(smoke_bundle / "execution.yaml"), + "--stage", + "full", + "--local", + "--poll-interval", + "0.05", + "--color", + "never", + "--override", + "tokenize_data.workers=1", + "--override", + "+replacement_scoring.automodel.lm_head_backend=streaming", + ], + cwd=project_root_path, + env=environment, + capture_output=True, + text=True, + timeout=900, + check=False, + ) + + +def _assert_campaign_succeeded( + completed: subprocess.CompletedProcess[str], + smoke_root: Path, +) -> None: + if completed.returncode == 0: + return + logs = sorted( + smoke_root.glob("logs/**/*.log"), + key=lambda path: path.stat().st_mtime_ns, + ) + log_tail = logs[-1].read_text(errors="replace")[-12000:] if logs else "no task log found" + pytest.fail( + "Tiny Qwen Puzzletron campaign failed.\n" + f"stdout tail:\n{completed.stdout[-12000:]}\n" + f"stderr tail:\n{completed.stderr[-12000:]}\n" + f"latest task-log tail:\n{log_tail}" + ) + + +@pytest.mark.timeout(1200) +def test_tiny_qwen_campaign_uses_current_public_route( + project_root_path: Path, + tmp_path: Path, +) -> None: + """Run every required stage through the public orchestrator; requires one CUDA GPU.""" + + assert torch.cuda.is_available(), "Puzzletron GPU CI requires CUDA" + assert torch.cuda.device_count() == 1, "Puzzletron GPU CI requires one visible GPU" + assert torch.equal(torch.arange(4, device="cuda").cpu(), torch.arange(4)) + model_dir = create_tiny_qwen3_5_dir( + tmp_path / "model", + with_tokenizer=True, + hidden_size=512, + intermediate_size=768, + max_position_embeddings=128, + num_hidden_layers=2, + layer_types=["full_attention"] * 2, + ) + dataset_dir = tmp_path / "dataset" + campaign_dir = tmp_path / "campaign" + result_root = tmp_path / "results" + cache_dir = tmp_path / "cache" + defaults_path = tmp_path / "defaults.yaml" + _save_messages_dataset(dataset_dir) + defaults_path.write_text( + yaml.safe_dump( + { + "schema_version": 1, + "model": { + "source": str(model_dir), + "trust_remote_code": False, + "force_hf": False, + }, + "data": { + "source": str(dataset_dir), + "modality": "text", + "layout": "fixed", + "sequence_length": 32, + }, + "pruning": { + "depth_remove": 0, + "depth_importance_samples": 2, + "width_importance_samples": 2, + "replacement_samples": 2, + "sort_sanity": False, + "width_sanity": False, + "slicing_sanity": False, + "replacement_granularity": "block", + "axes": { + "hidden_width": {"values": [256]}, + "kv_groups": {"values": [2]}, + "q_heads_per_group": {"values": [2]}, + "ffn_intermediate": {"values": [768, 512, 256]}, + "gdn_key_groups": {"values": [2]}, + "gdn_value_heads_per_group": {"values": [2]}, + "gdn_key_head_dim": {"values": [8]}, + "gdn_value_head_dim": {"values": [8]}, + }, + "bypass": {"enabled": False}, + }, + "vllm": { + "enabled": False, + "prefill_seq_len": 32, + "generation_seq_len": 8, + "batch_size": 1, + "max_num_seqs": 1, + }, + "mip": { + "goal_metric": "params", + "goal_value": "90%", + "num_solutions": 1, + }, + "stages": { + "width_importance": {"batch": 1}, + "replacement_scoring": {"batch": 1, "instances": 1}, + }, + "output": {"result_root": str(result_root)}, + "infrastructure": { + "gpus_per_node": 1, + "execution_contract": { + "repository": str(project_root_path), + "venv": sys.prefix, + "container": None, + "container_mounts": None, + "prerun_commands": [], + "postrun_commands": [], + }, + }, + }, + sort_keys=False, + ) + ) + + generated = run_wizard_v2( + resume=None, + defaults_path=defaults_path, + backend=_DefaultsBackend(campaign_dir), + ) + smoke_bundle = generated / "smoke" + for name in ("experiment.yaml", "runner.yaml", "execution.yaml"): + assert (smoke_bundle / name).is_file() + + overrides = [ + "tokenize_data.workers=1", + "+replacement_scoring.automodel.lm_head_backend=streaming", + ] + config = pipeline_config_from_path(smoke_bundle / "experiment.yaml", overrides=overrides) + assert config["embedding_pruning"]["widths"] == [256] + post_mip_flows = config["post_mip"]["flows"] + assert len(post_mip_flows) == 1 + serving_config = next(iter(post_mip_flows.values()))["nodes"]["serving"]["config"] + assert serving_config["input_tokens"] == 32 + assert serving_config["output_tokens"] == 8 + compiled_plan = compile_campaign_plan( + experiment_config_path=smoke_bundle / "experiment.yaml", + runner=load_runner_config(smoke_bundle / "runner.yaml"), + execution=load_execution_config(smoke_bundle / "execution.yaml"), + overrides=overrides, + stage_filter="full", + ) + replacement_node = next( + node for node in compiled_plan.stages if node.stage_id == "replacement_scoring" + ) + replacement_work = adapter_for_stage(replacement_node).plan( + compiled_plan, replacement_node + ) + assert replacement_node.instances == 1 + assert replacement_node.gpus_per_instance == 1 + assert replacement_node.total_gpus == 1 + assert replacement_node.nodes == 1 + assert len(replacement_work.items) == 1 + assert replacement_work.items[0].metadata["worker_count"] == 1 + + environment = os.environ.copy() + environment.update( + { + "CUDA_VISIBLE_DEVICES": os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0], + "HF_DATASETS_OFFLINE": "1", + "HF_HOME": str(cache_dir / "huggingface"), + "HF_HUB_OFFLINE": "1", + "HF_DATASETS_CACHE": str(cache_dir / "datasets"), + "TORCH_HOME": str(cache_dir / "torch"), + "TRANSFORMERS_OFFLINE": "1", + "XDG_CACHE_HOME": str(cache_dir / "xdg"), + } + ) + smoke_root = result_root / "smoke" + completed = _run_campaign(project_root_path, smoke_bundle, environment) + _assert_campaign_succeeded(completed, smoke_root) + + stages = enabled_stage_ids(config) + assert config["model"]["descriptor_override"] == "qwen3_5_text" + assert all(stage_is_complete(config, stage) for stage in stages) + + manifests = [smoke_root / "manifests" / f"{stage}.json" for stage in stages] + assert all(json.loads(path.read_text())["status"] == "success" for path in manifests) + + pass_manifests = list( + smoke_root.glob("pruning/pruning_scores/automodel/*/activation_passes_manifest.json") + ) + assert len(pass_manifests) == 1 + score_files = list(pass_manifests[0].parent.glob("*/rank_*.pth")) + assert score_files + score_tensors = [ + tensor + for score_file in score_files + for tensor in _tensor_values(torch.load(score_file, map_location="cpu", weights_only=False)) + ] + assert score_tensors + assert all(tensor.numel() and torch.isfinite(tensor).all() for tensor in score_tensors) + + replacement_summary = json.loads( + (smoke_root / "artifacts/replacement_scoring/summary.json").read_text() + ) + assert replacement_summary["widths"] == [256] + assert replacement_summary["scenario_count"] == 1 + replacement_results = [ + json.loads(path.read_text()) + for path in smoke_root.glob( + "scenarios/*/distributed_eval/replacement_scoring/results/**/*.json" + ) + ] + assert replacement_results + assert all( + result["provenance"]["score_device_type"] == "cuda" + and result["provenance"]["visible_cuda_device_count"] == 1 + for result in replacement_results + ) + + for checkpoint in (smoke_root / "ckpts/teacher", smoke_root / "ckpts/sorted_teacher"): + assert (checkpoint / "config.json").is_file() + assert list(checkpoint.glob("*.safetensors")) + candidate_library = json.loads((smoke_root / "candidate_library.json").read_text()) + assert 256 in _nested_values(candidate_library, "intermediate_size") + assert 512 in _nested_values(candidate_library, "intermediate_size") + + active_profiles = json.loads((smoke_root / "mip/active_profiles.json").read_text()) + assert active_profiles["status"] == "success" + grids = [ + json.loads((smoke_root / "mip" / "profiles" / profile_id / "mip_grid.json").read_text()) + for profile_id in active_profiles["profile_ids"] + ] + assert all(grid["status"] == "success" for grid in grids) + solution_paths = [ + Path(scenario["solution_path"]) + for grid in grids + for scenario in grid["scenarios"] + if scenario["status"] == "feasible" + ] + solutions = [solution for path in solution_paths for solution in json.loads(path.read_text())] + assert solutions + assert 256 in _nested_values(solutions, "intermediate_size") + + model = AutoModelForCausalLM.from_pretrained( + smoke_root / "ckpts/sorted_teacher", + dtype=torch.bfloat16, + local_files_only=True, + ).cuda() + with torch.no_grad(): + logits = model(torch.tensor([[1, 2, 3, 4]], device="cuda")).logits + assert torch.isfinite(logits).all() + + durable_paths = [*manifests, *pass_manifests] + durable_paths.extend(smoke_root.glob("mip/profiles/*/mip_grid.json")) + before_resume = _artifact_digests(durable_paths) + resumed = _run_campaign(project_root_path, smoke_bundle, environment) + _assert_campaign_succeeded(resumed, smoke_root) + assert _artifact_digests(durable_paths) == before_resume diff --git a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py index d1918c2c7e4..637ffa61589 100644 --- a/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py +++ b/tests/unit/torch/puzzletron/test_automodel_solution_scoring.py @@ -289,6 +289,8 @@ def test_rpc_executor_scores_cumulative_depth_removals(monkeypatch): executor.params = {"micro_batch_size": 4} executor.sliced_teacher_baseline = {"lm_loss": {"avg": 0.0}} executor.latest_observability = None + executor.latest_score_device_type = "cpu" + executor.visible_cuda_device_count = 0 monkeypatch.setattr( executor, "_score", @@ -309,6 +311,8 @@ def test_rpc_executor_scores_cumulative_depth_removals(monkeypatch): assert result.metrics["lm_loss"]["avg"] == 1.0 assert [target["layer_idx"] for target in captured[0]] == [0, 1] assert result.provenance["micro_batch_size"] == 4 + assert result.provenance["score_device_type"] == "cpu" + assert result.provenance["visible_cuda_device_count"] == 0 def test_runtime_fingerprint_ignores_distributed_compute_dtype_transition(): diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index 5e77a624586..fc60799a0b3 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -830,6 +830,7 @@ def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path) item=item, attempt_id=f"a{index}", runner=runner, + overrides=["+replacement_scoring.automodel.lm_head_backend=streaming"], ) for index, item in enumerate(work_plan.items) ] @@ -838,6 +839,17 @@ def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path) assert [attempt.task_topology.task_count for attempt in attempts] == [4, 4] assert [attempt.task_topology.gpus_per_task for attempt in attempts] == [4, 4] assert [attempt.command.env["WORKER_COUNT"] for attempt in attempts] == ["4", "4"] + assert [attempt.command.env["FINALIZE_OVERRIDES"] for attempt in attempts] == [ + "+replacement_scoring.automodel.lm_head_backend=streaming", + "+replacement_scoring.automodel.lm_head_backend=streaming", + ] + assert all( + "puzzle_dir=" not in attempt.command.env["FINALIZE_OVERRIDES"] for attempt in attempts + ) + assert all( + "puzzle_dir=" in attempt.command.env["DISTRIBUTED_EVAL_OVERRIDES"] + for attempt in attempts + ) assert [attempt.command.env["FINALIZE_EXPECTED_COMPLETIONS"] for attempt in attempts] == [ "2", "2", @@ -850,6 +862,34 @@ def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path) attempts[0].command.env["FINALIZE_COMPLETION_DIR"] == attempts[1].command.env["FINALIZE_COMPLETION_DIR"] ) + changed_plan = CampaignPlan( + experiment_config_path=plan.experiment_config_path, + puzzle_dir=plan.puzzle_dir, + experiment_config={ + **plan.experiment_config, + "replacement_scoring": { + "granularity": "subblock", + "default_metric": "mse_loss_hidden_states", + }, + }, + runner=runner, + execution_defaults=plan.execution_defaults, + stages=(node,), + contract_hash=plan.contract_hash, + ) + changed_work_plan = adapter.plan(changed_plan, node) + changed_attempt = adapter.command( + plan=changed_plan, + node=node, + item=changed_work_plan.items[0], + attempt_id="changed", + runner=runner, + overrides=["+replacement_scoring.automodel.lm_head_backend=streaming"], + ) + assert ( + changed_attempt.command.env["FINALIZE_COMPLETION_DIR"] + != attempts[0].command.env["FINALIZE_COMPLETION_DIR"] + ) assert attempts[0].command.env["PUZZLE_DIR"].endswith("scenarios/width-2048/depth-00") assert attempts[1].command.env["PUZZLE_DIR"].endswith("scenarios/width-1792/depth-00") diff --git a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py index 5ba0266162c..7c2a85cfa58 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py +++ b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py @@ -22,6 +22,7 @@ import os import subprocess import sys +from copy import deepcopy from pathlib import Path import yaml @@ -164,6 +165,7 @@ def test_load_experiment_config_composes_defaults_and_interpolation( "puzzle_dir": "${oc.env:RUN_ROOT,unused}", "pruning": {"automodel": {"parallel": {"pp": 2, "dp_shard": 4}}}, "teacher_dir": "${puzzle_dir}/ckpts/teacher", + "hook_class": "${get_object:package.module.Hook}", } ) ) @@ -184,6 +186,7 @@ def test_load_experiment_config_composes_defaults_and_interpolation( assert config["puzzle_dir"] == str(tmp_path / "run") assert config["teacher_dir"] == str(tmp_path / "run" / "ckpts" / "teacher") assert config["copy"] == config["teacher_dir"] + assert config["hook_class"] == {"__type__": "package.module.Hook"} assert config["pruning"]["automodel"]["parallel"] == { "pp": 1, "dp_shard": 4, @@ -192,6 +195,44 @@ def test_load_experiment_config_composes_defaults_and_interpolation( assert config["_runtime"]["config_path"] == str(experiment) +def test_load_experiment_config_matches_hydra_scientific_number_semantics( + tmp_path: Path, +) -> None: + experiment = tmp_path / "experiment.yaml" + experiment.write_text( + """\ +defaults: [_self_] +bypass: + best_val_loss: 1e+9 + training: + learning_rate: 1e-4 + min_lr_factor: 1e-5 + schedule: [1e-4, 1e-5, \"1e-4\"] +quoted: \"1e-4\" +""" + ) + + config = load_experiment_config(experiment, overrides=["threshold=1e-4"]) + + assert config["bypass"]["best_val_loss"] == 1e9 + assert config["bypass"]["training"] == { + "learning_rate": 1e-4, + "min_lr_factor": 1e-5, + } + assert config["bypass"]["schedule"] == [1e-4, 1e-5, "1e-4"] + assert config["quoted"] == "1e-4" + assert config["threshold"] == 1e-4 + assert all( + isinstance(value, float) + for value in ( + config["bypass"]["best_val_loss"], + config["bypass"]["training"]["learning_rate"], + config["bypass"]["training"]["min_lr_factor"], + config["threshold"], + ) + ) + + def test_convert_completeness_requires_runtime_subblock_library( tmp_path: Path, write_terminal_manifest ) -> None: @@ -326,6 +367,42 @@ def test_build_library_requires_its_own_complete_outputs( assert stage_is_complete(config, "build_library") +def test_build_library_completion_accepts_equivalent_loader_and_worker_configs( + tmp_path: Path, write_terminal_manifest +) -> None: + from puzzletron_orchestrator.adapters.stage_compat import stage_is_complete + + experiment = tmp_path / "experiment.yaml" + experiment.write_text( + f"""\ +defaults: [_self_] +puzzle_dir: {tmp_path} +build_library: + enabled: true +bypass: + best_val_loss: 1e+9 + training: + learning_rate: 1e-4 + min_lr_factor: 1e-5 +""" + ) + controller_config = load_experiment_config(experiment) + worker_config = deepcopy(controller_config) + worker_config["library"] = {} + + write_terminal_manifest(tmp_path, "build_library", config=worker_config) + for name in ( + "replacement_library.json", + "candidate_library.json", + "subblock_stats.json", + ): + (tmp_path / name).write_text("{}") + + assert stage_is_complete(controller_config, "build_library") + controller_config["bypass"]["best_val_loss"] = 2e9 + assert not stage_is_complete(controller_config, "build_library") + + def test_embedding_build_library_requires_every_width_scenario( tmp_path: Path, write_terminal_manifest ) -> None: diff --git a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py index a6c2412d0ab..aa1653ed054 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py +++ b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py @@ -41,8 +41,15 @@ from puzzletron_orchestrator.controller import CampaignController from puzzletron_orchestrator.executors.base import Executor from puzzletron_orchestrator.progress import summarize_active_progress, summarize_stage_artifacts -from puzzletron_orchestrator.schema import AttemptSpec, CommandSpec, JobHandle, JobState, JobStatus -from puzzletron_orchestrator.state import StageRunRecord +from puzzletron_orchestrator.schema import ( + AttemptSpec, + CommandSpec, + JobHandle, + JobState, + JobStatus, + ValidatedResult, +) +from puzzletron_orchestrator.state import PersistedAttempt, StageRunRecord from puzzletron_orchestrator.terminal import ShutdownAction if TYPE_CHECKING: @@ -260,6 +267,7 @@ def submit(self, attempt: AttemptSpec) -> JobHandle: }, ) self._handles[handle.handle_id] = handle + self._attempts[handle.handle_id] = attempt return handle @staticmethod @@ -469,7 +477,7 @@ def test_controller_waits_for_parent_job_after_artifact_appears(tmp_path: Path, assert controller._parents_ready(child) -def test_controller_completed_retry_satisfies_work_plan(tmp_path: Path, monkeypatch): +def test_controller_current_contract_completion_satisfies_work_plan(tmp_path: Path, monkeypatch): experiment, runner_path, execution_path = _write_configs(tmp_path) plan = compile_campaign_plan( experiment_config_path=experiment, @@ -478,14 +486,479 @@ def test_controller_completed_retry_satisfies_work_plan(tmp_path: Path, monkeypa stage_filter="convert", ) controller = CampaignController(plan, executor=_FakeExecutor()) + node = plan.stages[0] + stage_execution_identity = controller._stage_execution_identity(node) attempts = [ - {"work_id": "convert:0", "status": JobState.CANCELLED.value}, - {"work_id": "convert:0", "status": JobState.COMPLETED.value}, + { + "work_id": "convert:0", + "status": JobState.CANCELLED.value, + "contract_hash": plan.contract_hash, + "metadata": {"stage_execution_identity": stage_execution_identity}, + }, + { + "work_id": "convert:0", + "status": JobState.COMPLETED.value, + "contract_hash": plan.contract_hash, + "metadata": {"stage_execution_identity": stage_execution_identity}, + }, ] assert controller._required_work_is_completed(plan.stages[0], attempts) monkeypatch.setattr(controller.store, "list_attempts", lambda _stage_id=None: attempts) - assert not controller._stage_has_active_or_completed_work(plan.stages[0]) + assert controller._stage_has_active_or_completed_work(plan.stages[0]) + assert not controller._submit_stage(plan.stages[0]) + + +def test_controller_fails_closed_for_legacy_completed_attempt(tmp_path: Path): + experiment, runner_path, execution_path = _write_configs(tmp_path) + plan = compile_campaign_plan( + experiment_config_path=experiment, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + stage_filter="convert", + ) + node = plan.stages[0] + delegate = adapter_for_stage(node) + item = delegate.plan(plan, node).items[0] + attempt = delegate.command( + plan=plan, + node=node, + item=item, + attempt_id="legacy-completed-attempt", + runner=plan.runner, + ) + executor = _TrackingFakeExecutor() + controller = CampaignController(plan, executor=executor) + controller.store.save_attempt(attempt, None, JobState.COMPLETED.value) + + result = controller.run(once=True) + + assert result["halted"] is True + assert result["failed_stages"] == ["convert"] + assert executor.submitted_stage_ids == [] + events = list(controller.store.events_root.glob("*_stage_execution_identity_incompatible.json")) + assert len(events) == 1 + event = json.loads(events[0].read_text()) + assert event["payload"]["failure_class"] == "config" + assert event["payload"]["incompatibility"] == "missing_stage_execution_identity" + assert event["payload"]["attempt_ids"] == ["legacy-completed-attempt"] + + recovered_executor = _TrackingFakeExecutor() + recovered = CampaignController(plan, executor=recovered_executor) + recovered_result = recovered.run(once=True) + + assert recovered_result["halted"] is True + assert recovered_result["failed_stages"] == ["convert"] + assert recovered_executor.submitted_stage_ids == [] + assert ( + len(list(controller.store.events_root.glob("*_stage_execution_identity_incompatible.json"))) + == 1 + ) + + +def test_controller_resubmits_completed_work_when_stage_semantics_change( + tmp_path: Path, monkeypatch +): + experiment, runner_path, execution_path = _write_configs(tmp_path) + runner = load_runner_config(runner_path) + execution = load_execution_config(execution_path) + plan_a = compile_campaign_plan( + experiment_config_path=experiment, + runner=runner, + execution=execution, + stage_filter="convert", + ) + config_b = yaml.safe_load(experiment.read_text()) + config_b["convert"]["model_path"] = "/models/replacement" + experiment.write_text(yaml.safe_dump(config_b)) + plan_b = compile_campaign_plan( + experiment_config_path=experiment, + runner=runner, + execution=execution, + stage_filter="convert", + ) + executor = _TrackingFakeExecutor() + controller_a = CampaignController(plan_a, executor=_FakeExecutor()) + controller_b = CampaignController(plan_b, executor=executor) + identity_a = controller_a._stage_execution_identity(plan_a.stages[0]) + identity_b = controller_b._stage_execution_identity(plan_b.stages[0]) + attempts = [ + { + "work_id": "convert:0", + "status": JobState.COMPLETED.value, + "contract_hash": plan_a.contract_hash, + "metadata": {"stage_execution_identity": identity_a}, + } + ] + monkeypatch.setattr(controller_b.store, "list_attempts", lambda _stage_id=None: attempts) + + assert plan_a.contract_hash == plan_b.contract_hash + assert identity_a != identity_b + assert not controller_b._required_work_is_completed(plan_b.stages[0], attempts) + assert controller_b._submit_stage(plan_b.stages[0]) + assert executor.submitted_stage_ids == ["convert"] + submitted_attempt = next(iter(executor._attempts.values())) + assert submitted_attempt.metadata["stage_execution_identity"] == identity_b + persisted_attempts = CampaignController(plan_b, executor=_FakeExecutor()).store.list_attempts( + "convert" + ) + assert persisted_attempts[-1]["metadata"]["stage_execution_identity"] == identity_b + + +def test_stage_execution_identity_ignores_unrelated_config(tmp_path: Path): + experiment, runner_path, execution_path = _write_configs(tmp_path) + runner = load_runner_config(runner_path) + execution = load_execution_config(execution_path) + plan_a = compile_campaign_plan( + experiment_config_path=experiment, + runner=runner, + execution=execution, + stage_filter="convert", + ) + config_b = yaml.safe_load(experiment.read_text()) + config_b["report"] = {"title": "unrelated presentation change"} + experiment.write_text(yaml.safe_dump(config_b)) + plan_b = compile_campaign_plan( + experiment_config_path=experiment, + runner=runner, + execution=execution, + stage_filter="convert", + ) + + identity_a = CampaignController(plan_a, executor=_FakeExecutor())._stage_execution_identity( + plan_a.stages[0] + ) + identity_b = CampaignController(plan_b, executor=_FakeExecutor())._stage_execution_identity( + plan_b.stages[0] + ) + + assert identity_a == identity_b + + +def test_controller_revalidates_recent_completed_work_before_resubmitting( + tmp_path: Path, monkeypatch +): + experiment, runner_path, execution_path = _write_configs(tmp_path) + plan = compile_campaign_plan( + experiment_config_path=experiment, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + stage_filter="convert", + ) + node = plan.stages[0] + delegate = adapter_for_stage(node) + item = delegate.plan(plan, node).items[0] + attempt = delegate.command( + plan=plan, + node=node, + item=item, + attempt_id="completed-attempt", + runner=plan.runner, + ) + executor = _TrackingFakeExecutor() + controller = CampaignController(plan, executor=executor, poll_interval_seconds=45.0) + attempt = controller._bind_attempt_to_stage_execution(node, delegate.plan(plan, node), attempt) + now = [1000.0] + monkeypatch.setattr("puzzletron_orchestrator.controller.time.time", lambda: now[0]) + monkeypatch.setattr( + controller, + "_interruptible_sleep", + lambda seconds: now.__setitem__(0, now[0] + seconds), + ) + controller.store.save_attempt(attempt, None, JobState.RUNNING.value) + controller.store.update_attempt_status( + item.work_id, + attempt.attempt_id, + JobStatus( + handle=JobHandle( + backend="fake", + handle_id="completed-handle", + attempt_id=attempt.attempt_id, + ), + state=JobState.COMPLETED, + ), + ) + + class _DelayedVisibilityAdapter: + def __init__(self) -> None: + self.validation_count = 0 + + def __getattr__(self, name): + return getattr(delegate, name) + + def aggregate(self, *, plan, node, work_plan): + return None + + def validate(self, *, plan, node): + self.validation_count += 1 + return ValidatedResult( + valid=self.validation_count >= 4, + reason="stage outputs missing", + artifacts=("ckpts/teacher/config.json",), + ) + + delayed = _DelayedVisibilityAdapter() + monkeypatch.setattr( + "puzzletron_orchestrator.controller.adapter_for_stage", lambda _node: delayed + ) + monkeypatch.setattr( + "puzzletron_orchestrator.controller.stage_is_complete", + lambda _config, stage_id: controller.store.stage_is_complete(stage_id), + ) + + result = controller.run(max_iterations=4) + + assert result["halted"] is False + assert delayed.validation_count == 4 + assert "convert" not in executor.submitted_stage_ids + assert len(controller.store.list_attempts("convert")) == 1 + assert controller.store.stage_is_complete("convert") + + +def test_controller_retries_aggregation_until_outputs_are_visible(tmp_path: Path, monkeypatch): + experiment, runner_path, execution_path = _write_configs(tmp_path) + plan = compile_campaign_plan( + experiment_config_path=experiment, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + stage_filter="convert", + ) + node = plan.stages[0] + delegate = adapter_for_stage(node) + controller = CampaignController(plan, executor=_FakeExecutor()) + + class _DelayedAggregationAdapter: + def __init__(self) -> None: + self.aggregate_count = 0 + self.validation_count = 0 + + def __getattr__(self, name): + return getattr(delegate, name) + + def aggregate(self, *, plan, node, work_plan): + self.aggregate_count += 1 + if self.aggregate_count < 3: + raise FileNotFoundError("shards are still publishing") + return {"status": "complete"} + + def validate(self, *, plan, node): + self.validation_count += 1 + return ValidatedResult(valid=True, reason="stage outputs present") + + delayed = _DelayedAggregationAdapter() + monkeypatch.setattr( + "puzzletron_orchestrator.controller.adapter_for_stage", lambda _node: delayed + ) + + assert [controller._finalize_stage(node) for _ in range(3)] == [False, False, True] + assert delayed.aggregate_count == 3 + assert delayed.validation_count == 1 + assert controller.store.stage_is_complete("convert") + + +def test_controller_propagates_aggregation_programming_errors(tmp_path: Path, monkeypatch): + experiment, runner_path, execution_path = _write_configs(tmp_path) + plan = compile_campaign_plan( + experiment_config_path=experiment, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + stage_filter="convert", + ) + node = plan.stages[0] + delegate = adapter_for_stage(node) + controller = CampaignController(plan, executor=_FakeExecutor()) + + class _BrokenAggregationAdapter: + def __getattr__(self, name): + return getattr(delegate, name) + + def aggregate(self, *, plan, node, work_plan): + raise TypeError("aggregation programming error") + + monkeypatch.setattr( + "puzzletron_orchestrator.controller.adapter_for_stage", + lambda _node: _BrokenAggregationAdapter(), + ) + + with pytest.raises(TypeError, match="aggregation programming error"): + controller._finalize_stage(node) + + +@pytest.mark.parametrize("aggregation_failure", [False, True]) +def test_controller_fails_when_completed_work_artifacts_do_not_settle( + tmp_path: Path, monkeypatch, aggregation_failure: bool +): + experiment, runner_path, execution_path = _write_configs(tmp_path) + plan = compile_campaign_plan( + experiment_config_path=experiment, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + stage_filter="convert", + ) + node = plan.stages[0] + executor = _TrackingFakeExecutor() + controller = CampaignController(plan, executor=executor, poll_interval_seconds=150.0) + now = [1000.0] + delegate = adapter_for_stage(node) + item = delegate.plan(plan, node).items[0] + attempt = delegate.command( + plan=plan, + node=node, + item=item, + attempt_id="completed-attempt", + runner=plan.runner, + ) + attempt = controller._bind_attempt_to_stage_execution(node, delegate.plan(plan, node), attempt) + monkeypatch.setattr("puzzletron_orchestrator.controller.time.time", lambda: now[0]) + monkeypatch.setattr( + controller, + "_interruptible_sleep", + lambda seconds: now.__setitem__(0, now[0] + seconds), + ) + controller.store.save_attempt(attempt, None, JobState.RUNNING.value) + controller.store.update_attempt_status( + item.work_id, + attempt.attempt_id, + JobStatus( + handle=JobHandle( + backend="fake", + handle_id="completed-handle", + attempt_id=attempt.attempt_id, + ), + state=JobState.COMPLETED, + ), + ) + + class _MissingArtifactsAdapter: + def __getattr__(self, name): + return getattr(delegate, name) + + def aggregate(self, *, plan, node, work_plan): + if aggregation_failure: + raise FileNotFoundError("shards are still publishing") + + def validate(self, *, plan, node): + if aggregation_failure: + pytest.fail("validation must wait until aggregation succeeds") + return ValidatedResult( + valid=False, + reason="stage outputs missing", + artifacts=("ckpts/teacher/config.json",), + ) + + missing = _MissingArtifactsAdapter() + monkeypatch.setattr( + "puzzletron_orchestrator.controller.adapter_for_stage", lambda _node: missing + ) + monkeypatch.setattr( + "puzzletron_orchestrator.controller.stage_is_complete", + lambda _config, stage_id: controller.store.stage_is_complete(stage_id), + ) + + result = controller.run(max_iterations=3) + + assert result["halted"] is True + assert result["failed_stages"] == ["convert"] + assert executor.submitted_stage_ids == [] + assert len(controller.store.list_attempts("convert")) == 1 + phase = "aggregation" if aggregation_failure else "validation" + event_name = f"stage_{phase}_failed" + event_paths = list(controller.store.events_root.glob(f"*_{event_name}.json")) + assert len(event_paths) == 1 + event = json.loads(event_paths[0].read_text()) + assert event["payload"]["stage_id"] == "convert" + assert event["payload"]["failure_class"] == "timeout_fatal" + assert event["payload"]["contract_hash"] == plan.contract_hash + assert event["payload"]["stage_execution_identity"] == controller._stage_execution_identity( + node + ) + assert event["payload"]["phase"] == phase + assert event["payload"]["exception_type"] == ( + "FileNotFoundError" if aggregation_failure else None + ) + assert event["payload"]["attempt_ids"] == ["completed-attempt"] + assert event["payload"]["elapsed_seconds"] == 300.0 + expected_reason = ( + "stage aggregation failed: FileNotFoundError: shards are still publishing" + if aggregation_failure + else "stage outputs missing" + ) + assert event["payload"]["reason"] == expected_reason + assert event["payload"]["expected_artifacts"] == ( + [] if aggregation_failure else ["ckpts/teacher/config.json"] + ) + stage_record = controller.store.load_stage_record("convert") + assert stage_record is not None + assert stage_record.status == JobState.FAILED.value + assert stage_record.attempts[0].status == JobState.COMPLETED.value + + recovered_executor = _TrackingFakeExecutor() + recovered = CampaignController(plan, executor=recovered_executor) + recovered_result = recovered.run(once=True) + + assert recovered_result["halted"] is True + assert recovered_result["failed_stages"] == ["convert"] + assert recovered_executor.submitted_stage_ids == [] + assert len(list(controller.store.events_root.glob(f"*_{event_name}.json"))) == 1 + + +@pytest.mark.parametrize("stale_kind", ["contract", "stage_execution"]) +def test_controller_ignores_stale_contract_or_execution_stage_failure( + tmp_path: Path, stale_kind: str +): + experiment, runner_path, execution_path = _write_configs(tmp_path) + runner = load_runner_config(runner_path) + execution = load_execution_config(execution_path) + plan_a = compile_campaign_plan( + experiment_config_path=experiment, + runner=runner, + execution=execution, + stage_filter="convert", + ) + plan = plan_a + if stale_kind == "stage_execution": + config_b = yaml.safe_load(experiment.read_text()) + config_b["convert"]["model_path"] = "/models/replacement" + experiment.write_text(yaml.safe_dump(config_b)) + plan = compile_campaign_plan( + experiment_config_path=experiment, + runner=runner, + execution=execution, + stage_filter="convert", + ) + executor = _TrackingFakeExecutor() + controller = CampaignController(plan, executor=executor) + original_execution_identity = CampaignController( + plan_a, executor=_FakeExecutor() + )._stage_execution_identity(plan_a.stages[0]) + current_execution_identity = controller._stage_execution_identity(plan.stages[0]) + if stale_kind == "stage_execution": + assert plan_a.contract_hash == plan.contract_hash + assert original_execution_identity != current_execution_identity + controller.store.write_stage_record( + StageRunRecord( + stage_id="convert", + status=JobState.FAILED.value, + attempts=[ + PersistedAttempt( + attempt_id="stale-attempt", + work_id="convert:0", + stage_id="convert", + status=JobState.COMPLETED.value, + contract_hash=( + "stale-contract" if stale_kind == "contract" else plan.contract_hash + ), + metadata={"stage_execution_identity": original_execution_identity}, + ) + ], + ) + ) + + result = controller.run(once=True) + + assert result["halted"] is False + assert result["failed_stages"] == [] + assert executor.submitted_stage_ids == ["convert"] def test_controller_completed_summary_excludes_store_only_completion(tmp_path: Path, monkeypatch): @@ -536,6 +1009,7 @@ def test_controller_aggregates_completed_work_before_resubmitting( ) executor = _FakeExecutor() controller = CampaignController(plan, executor=executor, poll_interval_seconds=0.01) + attempt = controller._bind_attempt_to_stage_execution(node, delegate.plan(plan, node), attempt) controller.store.save_attempt(attempt, None, JobState.COMPLETED.value) class _AggregateAdapter: diff --git a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py index 22791749112..c62f40c5794 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py +++ b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py @@ -172,7 +172,23 @@ def _task_binding(*, group_size: int) -> task_launcher.TaskBinding: ) -def test_single_node_torchrun_uses_localhost_for_rendezvous() -> None: +def test_rendezvous_ports_are_stable_distinct_and_in_range() -> None: + first = task_launcher.rendezvous_port("attempt-a", 0, 2) + second = task_launcher.rendezvous_port("attempt-a", 1, 2) + + assert first == task_launcher.rendezvous_port("attempt-a", 0, 2) + assert first != second + assert 20000 <= first < 50000 + assert 20000 <= second < 50000 + + +@pytest.mark.parametrize("group_index", [-1, 2]) +def test_rendezvous_port_rejects_invalid_group_index(group_index: int) -> None: + with pytest.raises(ValueError, match="must be between"): + task_launcher.rendezvous_port("attempt-a", group_index, 2) + + +def test_single_node_torchrun_lets_c10d_choose_a_free_local_port() -> None: command = task_launcher.build_task_command( payload=("python", "worker.py"), launcher=TaskLauncher.TORCHRUN, @@ -180,7 +196,19 @@ def test_single_node_torchrun_uses_localhost_for_rendezvous() -> None: gpus_per_task=4, ) - assert "--rdzv-endpoint=localhost:23456" in command + assert command == ( + "python", + "-m", + "torch.distributed.run", + "--nnodes=1", + "--nproc-per-node=4", + "--rdzv-backend=c10d", + "--rdzv-endpoint=localhost:0", + "--rdzv-id=attempt-a-group-0", + "--no-python", + "python", + "worker.py", + ) def test_multi_node_torchrun_uses_master_hostname_for_rendezvous() -> None: @@ -191,7 +219,19 @@ def test_multi_node_torchrun_uses_master_hostname_for_rendezvous() -> None: gpus_per_task=4, ) - assert "--rdzv-endpoint=node-a:23456" in command + assert command == ( + "python", + "-m", + "torch.distributed.run", + "--nnodes=2", + "--nproc-per-node=4", + "--rdzv-backend=c10d", + "--rdzv-endpoint=node-a:23456", + "--rdzv-id=attempt-a-group-0", + "--no-python", + "python", + "worker.py", + ) def test_direct_launcher_does_not_wrap_payload() -> None: diff --git a/tests/unit/torch/puzzletron/test_setup_bundle.py b/tests/unit/torch/puzzletron/test_setup_bundle.py index 5f4f264ea65..828f5deb855 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 @@ -258,6 +255,13 @@ def test_custom_dataset_rendering_does_not_add_acquisition_fields() -> None: assert "acquisition" not in experiment["data"] +def test_rendered_data_keeps_controller_and_worker_sequence_length_in_sync() -> None: + experiment = render_experiment(_nemotron_render_state(latent_moe=False), "production") + + assert experiment["data"]["sequence_length"] == 2048 + assert experiment["data"]["sequence_length"] == experiment["data"]["max_sample_length"] + + def test_packed_text_uses_native_automodel_data_instead_of_fixed_token_memmaps() -> None: state = _nemotron_render_state(latent_moe=False) state["answers"]["data"]["layout"] = "packed_varlen" diff --git a/tests/unit/torch/puzzletron/test_stage_graph.py b/tests/unit/torch/puzzletron/test_stage_graph.py index d9736ee7bd7..a1941176dc9 100644 --- a/tests/unit/torch/puzzletron/test_stage_graph.py +++ b/tests/unit/torch/puzzletron/test_stage_graph.py @@ -143,6 +143,31 @@ def test_dynamic_stage_semantic_projection_keeps_stage_id_fallback() -> None: assert semantic_stage_config(config, "post.custom") == config +def test_semantic_projection_normalizes_empty_optional_sections() -> None: + baseline = semantic_stage_config({"build_library": {"enabled": True}}, "build_library") + + assert ( + semantic_stage_config( + {"build_library": {"enabled": True}, "library": None}, "build_library" + ) + == baseline + ) + assert ( + semantic_stage_config({"build_library": {"enabled": True}, "library": {}}, "build_library") + == baseline + ) + assert ( + semantic_stage_config( + { + "build_library": {"enabled": True}, + "library": {"vllm": {"enabled": True}}, + }, + "build_library", + ) + != baseline + ) + + def test_registry_uses_the_approved_fixed_dependencies(): assert selected_parent_stage_ids("tokenize_data", {}) == ("convert",) assert selected_parent_stage_ids("vllm_stats", {}) == ("convert",) diff --git a/tests/unit/torch/puzzletron/test_width_scenarios.py b/tests/unit/torch/puzzletron/test_width_scenarios.py index daa9be9a0dd..bb579e53aab 100644 --- a/tests/unit/torch/puzzletron/test_width_scenarios.py +++ b/tests/unit/torch/puzzletron/test_width_scenarios.py @@ -22,6 +22,7 @@ import pytest +import examples.puzzletron.finalize_replacement_scoring as replacement_finalizer from examples.puzzletron.embedding_pipeline import ( _project_vllm_stats_to_scenarios, _visible_gpu_count, @@ -42,6 +43,114 @@ ) from modelopt.torch.puzzletron.replacement_library.library import ReplacementLibrary from modelopt.torch.puzzletron.scenarios import ScenarioKey +from puzzletron_orchestrator.adapters.stage_compat import stage_is_complete + + +def test_replacement_scoring_finalizer_publishes_current_terminal_manifest( + tmp_path, monkeypatch +): + config_path = tmp_path / "experiment.yaml" + config_path.touch() + root_override = "+replacement_scoring.automodel.lm_head_backend=streaming" + loaded_overrides = [] + + def load_config(path, *, overrides=None): + assert path == config_path + loaded_overrides.extend(overrides or ()) + return { + "model": {"path": "tiny-qwen"}, + "embedding_pruning": {"enabled": True, "widths": [256]}, + "replacement_scoring": { + "granularity": "subblock", + "automodel": {"lm_head_backend": "streaming"}, + }, + } + + monkeypatch.setattr(replacement_finalizer, "pipeline_config_from_path", load_config) + report = {"scenario_count": 1, "widths": [256]} + + def publish_report(config): + summary = tmp_path / "artifacts" / "replacement_scoring" / "summary.json" + summary.parent.mkdir(parents=True) + summary.write_text(json.dumps(report)) + return report + + monkeypatch.setattr( + replacement_finalizer, + "finalize_replacement_scoring_diagnostics", + publish_report, + ) + + published_report = replacement_finalizer.finalize_replacement_scoring( + config_path, + tmp_path, + overrides=[root_override], + ) + + manifest_path = tmp_path / "manifests" / "replacement_scoring.json" + manifest = json.loads(manifest_path.read_text()) + assert loaded_overrides == [root_override] + assert published_report == report + assert manifest["stage"] == "replacement_scoring" + assert manifest["status"] == "success" + assert manifest["semantic_config"]["replacement_scoring"]["automodel"] == { + "lm_head_backend": "streaming" + } + assert manifest["outputs"]["report"] == report + assert stage_is_complete( + { + "puzzle_dir": str(tmp_path), + "model": {"path": "tiny-qwen"}, + "embedding_pruning": {"enabled": True, "widths": [256]}, + "replacement_scoring": { + "granularity": "subblock", + "automodel": {"lm_head_backend": "streaming"}, + }, + }, + "replacement_scoring", + ) + + marker_a = tmp_path / "completion-a" / "finalized" + marker_a.parent.mkdir() + replacement_finalizer.write_finalization_marker(marker_a, manifest_path) + summary = tmp_path / "artifacts" / "replacement_scoring" / "summary.json" + assert replacement_finalizer.finalization_marker_is_current(marker_a, manifest_path, summary) + + manifest_a = manifest_path.read_text() + summary_payload = summary.read_text() + summary.unlink() + assert not replacement_finalizer.finalization_marker_is_current( + marker_a, manifest_path, summary + ) + summary.write_text(summary_payload) + summary.write_text(json.dumps({"scenario_count": 2, "widths": [256]})) + assert not replacement_finalizer.finalization_marker_is_current( + marker_a, manifest_path, summary + ) + summary.write_text(summary_payload) + manifest_path.unlink() + assert not replacement_finalizer.finalization_marker_is_current( + marker_a, manifest_path, summary + ) + manifest_path.write_text(manifest_a) + + manifest_b = json.loads(manifest_a) + manifest_b["semantic_identity"] = "replacement_scoring_semantic_b" + manifest_path.write_text(json.dumps(manifest_b)) + assert not replacement_finalizer.finalization_marker_is_current( + marker_a, manifest_path, summary + ) + + marker_b = tmp_path / "completion-b" / "finalized" + marker_b.parent.mkdir() + replacement_finalizer.write_finalization_marker(marker_b, manifest_path) + assert replacement_finalizer.finalization_marker_is_current(marker_b, manifest_path, summary) + + manifest_path.write_text(manifest_a) + assert replacement_finalizer.finalization_marker_is_current(marker_a, manifest_path, summary) + assert not replacement_finalizer.finalization_marker_is_current( + marker_b, manifest_path, summary + ) def _write_scenario_manifest( @@ -397,6 +506,8 @@ def test_embedding_pipeline_launches_block_library_with_torchrun(tmp_path): assert command[1:4] == ("-m", "torch.distributed.run", "--standalone") assert "--nproc_per_node=1" in command + overrides = [command[index + 1] for index, value in enumerate(command) if value == "--override"] + assert "embedding_pruning.enabled=false" in overrides def test_embedding_pipeline_skips_composite_work_on_nonzero_rank(tmp_path, monkeypatch): diff --git a/tests/unit/torch/puzzletron/test_width_slice_equivalence.py b/tests/unit/torch/puzzletron/test_width_slice_equivalence.py index 4b924ae8ab0..c684e01874b 100644 --- a/tests/unit/torch/puzzletron/test_width_slice_equivalence.py +++ b/tests/unit/torch/puzzletron/test_width_slice_equivalence.py @@ -26,6 +26,7 @@ import pytest import torch import transformers +from _test_utils.torch.transformers_models import create_tiny_llama_dir, create_tiny_qwen3_5_dir from transformers import LlamaForCausalLM from modelopt.torch.puzzletron.anymodel.models.llama.llama_model_descriptor import ( @@ -57,10 +58,6 @@ from modelopt.torch.puzzletron.stages.diagnostics import width_slice_equivalence_stage from modelopt.torch.puzzletron.tools.checkpoint_utils import load_model_config from modelopt.torch.puzzletron.utils.data import dataloaders as dataloader_module -from tests._test_utils.torch.transformers_models import ( - create_tiny_llama_dir, - create_tiny_qwen3_5_dir, -) if TYPE_CHECKING: from pathlib import Path From 7fbe3ee665db81c953297a97dc9fe5840fa6d060 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 11 Aug 2026 22:31:14 +0200 Subject: [PATCH 02/24] Preserve authored Puzzletron stage identity Keep controller-compatible authored configuration separate from normalized worker settings so equivalent stage manifests remain resumable. Signed-off-by: Johannes Rausch --- .../finalize_replacement_scoring.py | 12 +-- examples/puzzletron/main.py | 4 +- .../puzzletron/run_axis_diagnostic_worker.py | 38 ++++------ examples/puzzletron/tokenize_data.py | 4 +- modelopt/torch/puzzletron/manifest.py | 27 +++++++ modelopt/torch/puzzletron/pipeline_config.py | 18 +++++ modelopt/torch/puzzletron/stage_runner.py | 11 ++- .../test_diagnostic_scoring_config.py | 26 ++++++- .../torch/puzzletron/test_tokenize_data.py | 73 +++++++++++++++++++ 9 files changed, 169 insertions(+), 44 deletions(-) diff --git a/examples/puzzletron/finalize_replacement_scoring.py b/examples/puzzletron/finalize_replacement_scoring.py index e3514f0ea91..6953534580f 100644 --- a/examples/puzzletron/finalize_replacement_scoring.py +++ b/examples/puzzletron/finalize_replacement_scoring.py @@ -17,7 +17,7 @@ from embedding_pipeline import finalize_replacement_scoring_diagnostics from modelopt.torch.puzzletron.diagnostics import generate_replace_block_report -from modelopt.torch.puzzletron.manifest import StageManifest, write_stage_manifest +from modelopt.torch.puzzletron.manifest import stage_manifest_from_config, write_stage_manifest from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path @@ -93,15 +93,13 @@ def finalize_replacement_scoring( scores_dir=puzzle_dir / f"{stem}--validation", output_dir=puzzle_dir / "artifacts" / "replacement_scoring", granularity=granularity, - default_metric=str( - scoring.get("default_metric", "normalized_mse_loss_hidden_states") - ), + default_metric=str(scoring.get("default_metric", "normalized_mse_loss_hidden_states")), default_layer_count=int(scoring.get("default_layer_count", 5)), anchor_count=int(scoring.get("anchor_count", 3)), trend_relative_tolerance=float(scoring.get("trend_relative_tolerance", 0.02)), ) - manifest = StageManifest(stage="replacement_scoring", inputs={"config": config}, config=config) + manifest = stage_manifest_from_config("replacement_scoring", config) manifest.complete(outputs={"report": report}) write_stage_manifest( Path(puzzle_dir) / "manifests" / "replacement_scoring.json", @@ -117,9 +115,7 @@ def main() -> None: args = parser.parse_args() overrides = [ - override - for override in os.environ.get("FINALIZE_OVERRIDES", "").splitlines() - if override + override for override in os.environ.get("FINALIZE_OVERRIDES", "").splitlines() if override ] finalize_replacement_scoring(args.config, args.puzzle_dir, overrides=overrides) diff --git a/examples/puzzletron/main.py b/examples/puzzletron/main.py index 2c6343f5601..c8cc04df390 100644 --- a/examples/puzzletron/main.py +++ b/examples/puzzletron/main.py @@ -41,8 +41,8 @@ import modelopt.torch.puzzletron as mtpz from modelopt.torch.puzzletron.manifest import ( - StageManifest, semantic_stage_config, + stage_manifest_from_config, validate_stage_execution_record, write_stage_manifest, ) @@ -470,7 +470,7 @@ def _run_worker(args: argparse.Namespace) -> None: def _complete_composite_stage(config: dict, stage: str, outputs: dict): puzzle_dir = Path(config.get("puzzle_dir") or (config.get("experiment") or {})["dir"]) manifest_path = puzzle_dir / "manifests" / f"{stage}.json" - manifest = StageManifest(stage=stage, inputs={"config": config}, config=config) + manifest = stage_manifest_from_config(stage, config) manifest.complete(outputs=outputs) write_stage_manifest(manifest_path, manifest) return mtpz.stage_runner.StageResult( diff --git a/examples/puzzletron/run_axis_diagnostic_worker.py b/examples/puzzletron/run_axis_diagnostic_worker.py index e2b0d679dd7..97cdf4c66d4 100755 --- a/examples/puzzletron/run_axis_diagnostic_worker.py +++ b/examples/puzzletron/run_axis_diagnostic_worker.py @@ -14,10 +14,11 @@ generate_campaign_progress_report, ) from modelopt.torch.puzzletron.diagnostics.width_sanity import aggregate_width_sanity -from modelopt.torch.puzzletron.manifest import StageManifest, write_stage_manifest +from modelopt.torch.puzzletron.manifest import stage_manifest_from_config, write_stage_manifest from modelopt.torch.puzzletron.pipeline_config import ( load_runtime_hydra_config, pipeline_config_from_path, + rebase_authored_pipeline_config, ) from modelopt.torch.puzzletron.stage_runner import run_stage from modelopt.torch.puzzletron.stages.diagnostics import _PRIMARY_METRICS @@ -26,10 +27,9 @@ def _axes(config: dict) -> list[str]: search_axes = (config.get("search_space") or {}).get("axes") or {} - non_sortable = set( - str(axis) - for axis in (config.get("width_sanity") or {}).get("non_sortable_axes", ()) - ) + non_sortable = { + str(axis) for axis in (config.get("width_sanity") or {}).get("non_sortable_axes", ()) + } enabled = [ str(axis) for axis, axis_cfg in search_axes.items() @@ -118,7 +118,7 @@ def _worker_config(config: dict, axis: str, config_path: Path) -> dict: config["width_sanity"] = diagnostic runtime["overrides"] = runtime_overrides config["_runtime"] = runtime - return config + return rebase_authored_pipeline_config(config) def _validate_worker_topology(config: dict, axis: str) -> None: @@ -138,13 +138,7 @@ def _validate_worker_topology(config: dict, axis: str) -> None: "axis diagnostic dp_shard must be divisible by ep because EP is overlaid " f"on FSDP shards: parallel={parallel}" ) - expected = ( - sizes["tp"] - * sizes["cp"] - * sizes["pp"] - * sizes["dp_shard"] - * sizes["dp_replicate"] - ) + expected = sizes["tp"] * sizes["cp"] * sizes["pp"] * sizes["dp_shard"] * sizes["dp_replicate"] world_size = int(os.environ.get("WORLD_SIZE", "1")) if expected != world_size: raise ValueError( @@ -194,9 +188,7 @@ def _finalize(config_path: Path) -> None: worker_manifests = {} for axis in axes: safe = _safe_axis(axis) - manifest_path = ( - puzzle_dir / ".axis_workers" / safe / "manifests" / "width_sanity.json" - ) + manifest_path = puzzle_dir / ".axis_workers" / safe / "manifests" / "width_sanity.json" artifact_dir = puzzle_dir / "artifacts" / f"activation_diagnostic_axis_{safe}" summary_path = artifact_dir / "activation_diagnostic_summary.json" if manifest_path.is_file(): @@ -244,11 +236,7 @@ def _finalize(config_path: Path) -> None: parallel_execution = { "workers": len(axes), "gpus_per_worker": ( - sizes["tp"] - * sizes["cp"] - * sizes["pp"] - * sizes["dp_shard"] - * sizes["dp_replicate"] + sizes["tp"] * sizes["cp"] * sizes["pp"] * sizes["dp_shard"] * sizes["dp_replicate"] ), **sizes, } @@ -266,10 +254,10 @@ def _finalize(config_path: Path) -> None: artifacts_dir.mkdir(parents=True, exist_ok=True) summary_path = artifacts_dir / "summary.json" summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n") - manifest = StageManifest( - stage=stage, - inputs={"config": config, "worker_manifests": worker_manifests}, - config=config, + manifest = stage_manifest_from_config( + stage, + config, + inputs={"worker_manifests": worker_manifests}, ) manifest.complete( outputs={ diff --git a/examples/puzzletron/tokenize_data.py b/examples/puzzletron/tokenize_data.py index 3600d6d9b6c..3ef1472a686 100644 --- a/examples/puzzletron/tokenize_data.py +++ b/examples/puzzletron/tokenize_data.py @@ -21,7 +21,7 @@ import sys from pathlib import Path -from modelopt.torch.puzzletron.manifest import StageManifest, write_stage_manifest +from modelopt.torch.puzzletron.manifest import stage_manifest_from_config, write_stage_manifest from modelopt.torch.puzzletron.stage_runner import StageResult from modelopt.torch.puzzletron.stages.graph import StageSkipReason, stage_is_enabled from puzzletron_orchestrator.token_caches import resolve_tokenize_caches @@ -35,7 +35,7 @@ def tokenize_data_stage(config: dict) -> StageResult: stage_config = config.get("tokenize_data") or {} puzzle_dir = Path(config.get("puzzle_dir") or (config.get("experiment") or {})["dir"]) manifest_path = puzzle_dir / "manifests" / "tokenize_data.json" - manifest = StageManifest(stage="tokenize_data", inputs={"config": config}, config=config) + manifest = stage_manifest_from_config("tokenize_data", config) if not stage_is_enabled("tokenize_data", config): skip_reason = StageSkipReason.DISABLED manifest.complete( diff --git a/modelopt/torch/puzzletron/manifest.py b/modelopt/torch/puzzletron/manifest.py index 9fa4d439bc3..bbcf70fbeab 100644 --- a/modelopt/torch/puzzletron/manifest.py +++ b/modelopt/torch/puzzletron/manifest.py @@ -23,6 +23,7 @@ import shutil import tempfile from collections.abc import Mapping +from copy import deepcopy from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path @@ -48,6 +49,7 @@ "StageManifest", "read_stage_manifest", "semantic_stage_config", + "stage_manifest_from_config", "validate_stage_execution_record", "write_stage_execution_record", "write_stage_manifest", @@ -355,6 +357,31 @@ def to_dict(self) -> dict[str, Any]: return payload +def stage_manifest_from_config( + stage: str, + config: Mapping[str, Any], + *, + inputs: Mapping[str, Any] | None = None, + effective_config: Mapping[str, Any] | None = None, + **manifest_fields: Any, +) -> StageManifest: + """Build a worker manifest with separate authored and effective config views.""" + + runtime = config.get("_runtime") + authored = runtime.get("authored_config") if isinstance(runtime, Mapping) else None + authored_config = deepcopy(dict(authored if isinstance(authored, Mapping) else config)) + manifest_inputs = deepcopy(dict(inputs or {})) + manifest_inputs["config"] = deepcopy(authored_config) + resolved_config = config if effective_config is None else effective_config + return StageManifest( + stage=stage, + inputs=manifest_inputs, + config=authored_config, + effective_config=deepcopy(dict(resolved_config)), + **manifest_fields, + ) + + def write_stage_manifest(path: str | Path, manifest: StageManifest) -> None: """Atomically write a stage manifest from rank zero. diff --git a/modelopt/torch/puzzletron/pipeline_config.py b/modelopt/torch/puzzletron/pipeline_config.py index e2d49695f54..e4f497300c9 100644 --- a/modelopt/torch/puzzletron/pipeline_config.py +++ b/modelopt/torch/puzzletron/pipeline_config.py @@ -33,6 +33,7 @@ "load_runtime_hydra_config", "normalize_pipeline_config", "pipeline_config_from_path", + "rebase_authored_pipeline_config", ] @@ -276,8 +277,12 @@ def pipeline_config_from_path( node_index: int = 0, ) -> dict[str, Any]: """Load a Hydra YAML and attach runtime metadata for stage handlers.""" + # Defer the orchestration-package import until this module is initialized. + from .orchestration.config import load_experiment_config + register_hydra_resolvers() path = Path(config_path).resolve() + authored_config = load_experiment_config(path, overrides=overrides) config_dir, config_name = _config_root_and_name(path) hydra_cfg = initialize_hydra_config_for_dir( config_dir=str(config_dir), @@ -290,10 +295,23 @@ def pipeline_config_from_path( "overrides": list(overrides or []), "num_nodes": int(num_nodes), "node_index": int(node_index), + "authored_config": authored_config, } return cfg +def rebase_authored_pipeline_config(config: dict[str, Any]) -> dict[str, Any]: + """Bind an intentionally derived worker config as its own authored view.""" + + runtime = deepcopy(dict(config.get("_runtime") or {})) + runtime.pop("authored_config", None) + authored_config = deepcopy(config) + authored_config["_runtime"] = deepcopy(runtime) + runtime["authored_config"] = authored_config + config["_runtime"] = runtime + return config + + def load_runtime_hydra_config(config: dict[str, Any]) -> DictConfig: """Reconstruct the instantiated Hydra config used by GPU-heavy stages.""" runtime = dict(config.get("_runtime") or {}) diff --git a/modelopt/torch/puzzletron/stage_runner.py b/modelopt/torch/puzzletron/stage_runner.py index be8e3cb9d79..8a26c412684 100644 --- a/modelopt/torch/puzzletron/stage_runner.py +++ b/modelopt/torch/puzzletron/stage_runner.py @@ -28,7 +28,7 @@ resolve_descriptor_by_name, resolve_descriptor_from_pretrained, ) -from .manifest import StageManifest, write_stage_manifest +from .manifest import StageManifest, stage_manifest_from_config, write_stage_manifest from .pipeline_config import canonical_stage_name, normalize_pipeline_config __all__ = ["STAGES", "StageResult", "normalize_config", "run_stage"] @@ -210,7 +210,7 @@ def run_stage( return _skip_stage( cfg, stage, - StageManifest(stage=stage, inputs={"config": cfg}, config=cfg), + stage_manifest_from_config(stage, cfg), reason=StageSkipReason.DISABLED, message=f"Stage '{stage}' is disabled by configuration.", ) @@ -232,13 +232,12 @@ def run_stage( descriptor_confidence=resolution.confidence, ) runtime_cfg["_runtime"] = runtime - manifest = StageManifest( - stage=stage, + manifest = stage_manifest_from_config( + stage, + cfg, inputs={ - "config": cfg, "descriptor_resolution": resolution.to_dict() if resolution else None, }, - config=cfg, effective_config=copy.deepcopy(runtime_cfg), capability_snapshot=resolution.capabilities.to_dict() if resolution else None, ) diff --git a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py index 28e45a5a541..6b74760aba8 100644 --- a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py +++ b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py @@ -14,6 +14,7 @@ # limitations under the License. import json +from copy import deepcopy from pathlib import Path from types import SimpleNamespace @@ -27,7 +28,7 @@ _validate_worker_topology, _worker_config, ) -from modelopt.torch.puzzletron.manifest import StageManifest +from modelopt.torch.puzzletron.manifest import StageManifest, stage_manifest_from_config from modelopt.torch.puzzletron.stages import diagnostics from modelopt.torch.puzzletron.stages.diagnostics import _scoring_cfg_for_method @@ -304,6 +305,29 @@ def test_axis_worker_preserves_requested_layers_and_targets_per_axis(tmp_path: P ) +def test_axis_workers_bind_distinct_authored_semantic_configs(tmp_path: Path): + config = { + "experiment": {"dir": str(tmp_path / "run")}, + "width_sanity": {"automodel": {"parallel": {"tp": 1}}}, + "_runtime": { + "config_path": str(tmp_path / "experiment.yaml"), + "authored_config": { + "experiment": {"dir": str(tmp_path / "run")}, + "width_sanity": {"enabled": False}, + }, + }, + } + + kv_config = _worker_config(deepcopy(config), "kv_groups", tmp_path / "experiment.yaml") + ffn_config = _worker_config(deepcopy(config), "ffn_intermediate", tmp_path / "experiment.yaml") + kv_manifest = stage_manifest_from_config("width_sanity", kv_config) + ffn_manifest = stage_manifest_from_config("width_sanity", ffn_config) + + assert kv_manifest.config["width_sanity"]["axes"] == ["kv_groups"] + assert ffn_manifest.config["width_sanity"]["axes"] == ["ffn_intermediate"] + assert kv_manifest.semantic_identity != ffn_manifest.semantic_identity + + def test_axis_worker_excludes_non_sortable_axes_from_width_diagnostics(): config = { "search_space": { diff --git a/tests/unit/torch/puzzletron/test_tokenize_data.py b/tests/unit/torch/puzzletron/test_tokenize_data.py index ae7a14c749d..a9d901015f3 100644 --- a/tests/unit/torch/puzzletron/test_tokenize_data.py +++ b/tests/unit/torch/puzzletron/test_tokenize_data.py @@ -15,12 +15,15 @@ """Tests for tokenize_data cache resolution and stage execution.""" +import json from pathlib import Path import pytest from examples.puzzletron.tokenize_data import tokenize_data_stage from modelopt.torch.puzzletron.orchestration.adapters.stage_compat import stage_is_complete +from modelopt.torch.puzzletron.orchestration.config import load_experiment_config +from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path from puzzletron_orchestrator.token_caches import resolve_tokenize_caches @@ -111,6 +114,76 @@ def _run(command, *, check): assert stage_is_complete(config, "tokenize_data") +def test_tokenize_data_manifest_accepts_equivalent_controller_and_worker_configs( + tmp_path, monkeypatch, write_token_cache +): + output = tmp_path / "dataset_cache" / "train.tokens" + experiment = tmp_path / "experiment.yaml" + experiment.write_text( + f"""\ +defaults: [_self_] +puzzle_dir: {tmp_path} +dataset_path: {tmp_path / "dataset"} +model: + source: {tmp_path / "model"} + trust_remote_code: false +convert: + teacher_dir: {tmp_path / "teacher"} +data: + modality: text + layout: fixed + max_sample_length: 8 +search_space: + axes: + hidden_width: + enabled: true + values: [256] +sort_sanity: + enabled: false +width_sanity: + enabled: false +tokenize_data: + enabled: true + workers: 1 + caches: + - output: {output} + split: train + num_samples: 1 + seq_length: 8 + shuffle_seed: 1 +""" + ) + controller_config = load_experiment_config(experiment) + worker_config = pipeline_config_from_path(experiment) + + assert worker_config["sort_sanity"]["include_reverse"] is True + assert worker_config["width_sanity"]["target_values"] == {"hidden_width": 256} + assert "include_reverse" not in controller_config["sort_sanity"] + assert "target_values" not in controller_config["width_sanity"] + + caches = resolve_tokenize_caches(worker_config) + + def _run(command, *, check): + assert check is True + cache_output = Path(command[command.index("--output") + 1]) + cache = next(cache for cache in caches if Path(cache["output"]) == cache_output) + write_token_cache(worker_config, cache) + + monkeypatch.setattr("examples.puzzletron.tokenize_data.subprocess.run", _run) + tokenize_data_stage(worker_config) + + manifest = json.loads((tmp_path / "manifests" / "tokenize_data.json").read_text()) + assert manifest["semantic_config"]["sort_sanity"] == {"enabled": False} + assert manifest["semantic_config"]["width_sanity"] == {"enabled": False} + resolved_path = tmp_path / manifest["execution_record"]["resolved_config_path"] + resolved_config = json.loads(resolved_path.read_text())["resolved_stage_config"] + assert resolved_config["sort_sanity"]["include_reverse"] is True + assert resolved_config["width_sanity"]["target_values"] == {"hidden_width": 256} + assert stage_is_complete(controller_config, "tokenize_data") + controller_config["sort_sanity"]["enabled"] = True + assert not stage_is_complete(controller_config, "tokenize_data") + + def test_tokenize_data_stage_passes_trust_remote_code_only_when_enabled(tmp_path, monkeypatch): commands = [] monkeypatch.setattr( From 4ed2e744197674ed53dd61d10fbc528ceadcfa04 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Tue, 11 Aug 2026 23:02:02 +0200 Subject: [PATCH 03/24] Fix Puzzletron worker semantic validation Use the authored configuration for semantic compatibility checks while preserving normalized worker settings in execution records and synthesized post-MIP candidates. This keeps worker self-validation aligned with controller resume checks. Signed-off-by: Johannes Rausch --- modelopt/torch/puzzletron/manifest.py | 2 +- modelopt/torch/puzzletron/post_mip/runner.py | 4 +- modelopt/torch/puzzletron/stages/graph.py | 21 +++++-- .../torch/puzzletron/test_post_mip_runner.py | 55 +++++++++++++++++++ .../unit/torch/puzzletron/test_stage_graph.py | 22 ++++++++ .../torch/puzzletron/test_tokenize_data.py | 5 +- 6 files changed, 101 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/puzzletron/manifest.py b/modelopt/torch/puzzletron/manifest.py index bbcf70fbeab..8082c9b4aaf 100644 --- a/modelopt/torch/puzzletron/manifest.py +++ b/modelopt/torch/puzzletron/manifest.py @@ -138,7 +138,7 @@ def write_stage_execution_record( or stable_hash(authored_config, prefix=f"{stage}_cfg") ) resolved_config_content = _resolved_config_content( - semantic_stage_config(dict(effective_config), stage) + semantic_stage_config(dict(effective_config), stage, use_authored=False) if isinstance(effective_config, Mapping) else effective_config ) diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 08aeff370a1..cb054134f16 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -451,7 +451,9 @@ def _evaluate_checkpoint( stage=f"{node.stage_id}.{source.architecture_id}", inputs={"config": candidate}, config=candidate, - semantic_config=semantic_stage_config(candidate, "zero_shot_evaluation"), + semantic_config=semantic_stage_config( + candidate, "zero_shot_evaluation", use_authored=False + ), ) evaluation_stage(candidate, manifest) rows = json.loads((output / "evaluation_summary.json").read_text()) diff --git a/modelopt/torch/puzzletron/stages/graph.py b/modelopt/torch/puzzletron/stages/graph.py index 779a2b9e385..a7e1deaa331 100644 --- a/modelopt/torch/puzzletron/stages/graph.py +++ b/modelopt/torch/puzzletron/stages/graph.py @@ -514,23 +514,34 @@ def stage_spec(stage_id: str) -> StageSpec: raise ValueError(f"Unknown Puzzletron stage {stage_id!r}") from error -def semantic_stage_config(config: Mapping[str, Any], stage_id: str) -> dict[str, Any]: +def semantic_stage_config( + config: Mapping[str, Any], stage_id: str, *, use_authored: bool = True +) -> dict[str, Any]: """Return configuration that can change the semantic result of one stage. Public stages declare their semantic sections alongside their other scheduler-neutral metadata. Dynamic stages that are not in the public registry retain the historical stage-ID section fallback. + + Normalized worker configurations retain the independently loaded authored + configuration under ``_runtime.authored_config``. Semantic compatibility + uses that authored view by default, while execution records may explicitly + request the effective worker view. """ + runtime = config.get("_runtime") + authored = runtime.get("authored_config") if isinstance(runtime, Mapping) else None + selected = authored if use_authored and isinstance(authored, Mapping) else config + spec = STAGE_REGISTRY.get(stage_id) stage_sections = (stage_id,) if spec is None else spec.semantic_config_sections sections = dict.fromkeys((*SHARED_SEMANTIC_CONFIG_SECTIONS, *stage_sections)) return { - key: config[key] + key: selected[key] for key in sections - if key in config - and config[key] is not None - and not (isinstance(config[key], Mapping) and not config[key]) + if key in selected + and selected[key] is not None + and not (isinstance(selected[key], Mapping) and not selected[key]) } diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index 47f6e17012e..82bb5613277 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -20,6 +20,7 @@ from omegaconf import OmegaConf +import modelopt.torch.puzzletron.stages.future as future_stages from modelopt.torch.puzzletron.post_mip import runner from modelopt.torch.puzzletron.post_mip.records import ArtifactKind from modelopt.torch.puzzletron.post_mip.runner import ( @@ -132,6 +133,60 @@ def test_online_eval_injects_resolved_hidden_width_into_solution(monkeypatch): assert work.raw_solution["hidden_width"] == 1792 +def test_checkpoint_evaluation_manifest_uses_candidate_effective_config(monkeypatch, tmp_path): + observed = {} + checkpoint = tmp_path / "checkpoint" + node = SimpleNamespace( + node_id="evaluation", + stage_id="post.params.evaluation", + config={"config": {"tasks": ["candidate-task"]}}, + ) + source = SimpleNamespace( + architecture_id="architecture", + artifact={"checkpoint": str(checkpoint)}, + ) + config = { + "puzzle_dir": str(tmp_path), + "zero_shot_evaluation": {"enabled": False}, + "_runtime": { + "authored_config": { + "puzzle_dir": str(tmp_path), + "zero_shot_evaluation": {"enabled": False}, + } + }, + } + + def _evaluation_stage(candidate, manifest): + observed["semantic_config"] = manifest.semantic_config + output = Path(candidate["zero_shot_evaluation"]["output_dir"]) + output.mkdir(parents=True) + (output / "evaluation_summary.json").write_text( + json.dumps( + [ + { + "checkpoint": str(checkpoint), + "metrics": {"score": 1.0}, + "result_path": str(output / "result.json"), + } + ] + ) + ) + + monkeypatch.setattr(future_stages, "evaluation_stage", _evaluation_stage) + + result = runner._evaluate_checkpoint(config, node, source, "execution") + + assert observed["semantic_config"]["zero_shot_evaluation"] == { + "enabled": True, + "checkpoints": [str(checkpoint)], + "output_dir": str( + tmp_path / "artifacts/post_mip/nodes/evaluation/executions/execution/raw/architecture" + ), + "tasks": ["candidate-task"], + } + assert result["metrics"] == {"score": 1.0} + + def test_aiperf_consumes_request_count_without_forwarding_setup_only_keys( monkeypatch, tmp_path, diff --git a/tests/unit/torch/puzzletron/test_stage_graph.py b/tests/unit/torch/puzzletron/test_stage_graph.py index a1941176dc9..5379bdd5f0e 100644 --- a/tests/unit/torch/puzzletron/test_stage_graph.py +++ b/tests/unit/torch/puzzletron/test_stage_graph.py @@ -168,6 +168,28 @@ def test_semantic_projection_normalizes_empty_optional_sections() -> None: ) +def test_semantic_projection_uses_authored_config_unless_effective_view_is_requested() -> None: + config = { + "model": {"source": "normalized-model"}, + "convert": {"teacher_dir": "normalized-teacher"}, + "_runtime": { + "authored_config": { + "model": {"source": "authored-model"}, + "convert": {"teacher_dir": "authored-teacher"}, + } + }, + } + + assert semantic_stage_config(config, "convert") == { + "model": {"source": "authored-model"}, + "convert": {"teacher_dir": "authored-teacher"}, + } + assert semantic_stage_config(config, "convert", use_authored=False) == { + "model": {"source": "normalized-model"}, + "convert": {"teacher_dir": "normalized-teacher"}, + } + + def test_registry_uses_the_approved_fixed_dependencies(): assert selected_parent_stage_ids("tokenize_data", {}) == ("convert",) assert selected_parent_stage_ids("vllm_stats", {}) == ("convert",) diff --git a/tests/unit/torch/puzzletron/test_tokenize_data.py b/tests/unit/torch/puzzletron/test_tokenize_data.py index a9d901015f3..cc2f0676394 100644 --- a/tests/unit/torch/puzzletron/test_tokenize_data.py +++ b/tests/unit/torch/puzzletron/test_tokenize_data.py @@ -20,6 +20,7 @@ import pytest +from examples.puzzletron.main import _validate_worker_result from examples.puzzletron.tokenize_data import tokenize_data_stage from modelopt.torch.puzzletron.orchestration.adapters.stage_compat import stage_is_complete from modelopt.torch.puzzletron.orchestration.config import load_experiment_config @@ -170,7 +171,7 @@ def _run(command, *, check): write_token_cache(worker_config, cache) monkeypatch.setattr("examples.puzzletron.tokenize_data.subprocess.run", _run) - tokenize_data_stage(worker_config) + result = tokenize_data_stage(worker_config) manifest = json.loads((tmp_path / "manifests" / "tokenize_data.json").read_text()) assert manifest["semantic_config"]["sort_sanity"] == {"enabled": False} @@ -179,6 +180,8 @@ def _run(command, *, check): resolved_config = json.loads(resolved_path.read_text())["resolved_stage_config"] assert resolved_config["sort_sanity"]["include_reverse"] is True assert resolved_config["width_sanity"]["target_values"] == {"hidden_width": 256} + _validate_worker_result(worker_config, result, expected_stage="tokenize_data") + assert stage_is_complete(worker_config, "tokenize_data") assert stage_is_complete(controller_config, "tokenize_data") controller_config["sort_sanity"]["enabled"] = True assert not stage_is_complete(controller_config, "tokenize_data") From dfb1083d5679c4c777d0977f28b93ea960fe1fa0 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 01:50:24 +0200 Subject: [PATCH 04/24] Fix Puzzletron GPU baseline regressions The current-route baseline exposed generic rendezvous, finalization, provenance, and AIPerf integration defects. Preserve compiled execution identity across workers, recover aggregation publication failures, verify exact source dependencies, and keep local tokenizer loading compatible with the pinned AIPerf environment. Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 11 +- examples/puzzletron/ci_environment.json | 8 +- examples/puzzletron/ci_environment.py | 68 ++++++++++ .../distributed_eval/run_depth_pool.sh | 24 +++- .../distributed_eval/run_replacement_pool.sh | 24 +++- .../puzzletron/distributed_eval/run_worker.sh | 33 ++++- .../torch/puzzletron/benchmarks/aiperf.py | 11 +- .../puzzletron/orchestration/adapters/pool.py | 28 ++-- .../puzzletron/orchestration/compiler.py | 2 + .../puzzletron/orchestration/controller.py | 56 +++++++- .../torch/puzzletron/orchestration/schema.py | 13 ++ .../puzzletron/orchestration/task_launcher.py | 18 ++- noxfile.py | 55 +++++--- tests/_test_utils/torch/puzzletron/utils.py | 2 + .../test_aiperf_context_capacity.py | 24 ++++ .../torch/puzzletron/test_ci_environment.py | 62 +++++++++ .../test_diagnostic_scoring_config.py | 2 + .../test_orchestration_executors.py | 51 ++++++- .../test_orchestration_shutdown_progress.py | 60 ++++++++- .../test_orchestration_task_topology.py | 126 ++++++++++++++++++ 20 files changed, 611 insertions(+), 67 deletions(-) create mode 100644 examples/puzzletron/ci_environment.py create mode 100644 tests/unit/torch/puzzletron/test_ci_environment.py diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 86af42abd85..bdcb5175491 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -276,6 +276,8 @@ import os from packaging.version import Version +from examples.puzzletron.ci_environment import verify_installed_vcs_source + import aiperf import lmms_eval import modelopt @@ -307,10 +309,17 @@ 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("lmms-eval")).base_version == ( + ci_environment["lmms_eval"]["base_version"] +) assert Version(metadata.version("nemo-automodel")).base_version == ( ci_environment["nemo_automodel"]["base_version"] ) +for package, source in ( + ("lmms-eval", ci_environment["lmms_eval"]), + ("nemo-automodel", ci_environment["nemo_automodel"]), +): + verify_installed_vcs_source(package, source) assert torch.version.cuda == "12.9" assert torch.cuda.is_available() PY diff --git a/examples/puzzletron/ci_environment.json b/examples/puzzletron/ci_environment.json index 23f43e28623..9e39fd5d6b8 100644 --- a/examples/puzzletron/ci_environment.json +++ b/examples/puzzletron/ci_environment.json @@ -1,11 +1,15 @@ { "schema_version": 1, - "scope": "puzzletron_v2_cpu_ci", + "scope": "puzzletron_v2_ci", "python": "3.12", "torch": "2.11.0", "torchvision": "0.26.0", "transformers": "5.8.1", - "lmms_eval": "0.7.0", + "lmms_eval": { + "base_version": "0.7.0", + "repository": "https://github.com/EvolvingLMMs-Lab/lmms-eval.git", + "commit": "15c32bfec165df13c269ddd3cda03b2ed9137825" + }, "nemo_automodel": { "base_version": "0.5.0", "repository": "https://github.com/Separius/Automodel.git", diff --git a/examples/puzzletron/ci_environment.py b/examples/puzzletron/ci_environment.py new file mode 100644 index 00000000000..fe6924e9841 --- /dev/null +++ b/examples/puzzletron/ci_environment.py @@ -0,0 +1,68 @@ +# 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. + +"""Verification helpers for the pinned Puzzletron CI environment.""" + +from __future__ import annotations + +import json +import subprocess +from importlib import metadata +from typing import Any +from urllib.parse import unquote, urlparse + +__all__ = ["verify_installed_vcs_source"] + + +def _normalized_repository(url: object) -> str: + return str(url or "").removesuffix(".git").rstrip("/") + + +def _installed_vcs_source(package: str) -> tuple[str | None, str | None]: + payload = json.loads(metadata.distribution(package).read_text("direct_url.json") or "{}") + vcs_info = payload.get("vcs_info") or {} + if vcs_info.get("commit_id"): + return payload.get("url"), vcs_info["commit_id"] + if (payload.get("dir_info") or {}).get("editable") and str(payload.get("url", "")).startswith( + "file:" + ): + root = unquote(urlparse(payload["url"]).path) + repository = subprocess.check_output( + ["git", "-C", root, "remote", "get-url", "origin"], text=True + ).strip() + commit = subprocess.check_output( + ["git", "-C", root, "rev-parse", "HEAD"], text=True + ).strip() + dirty = subprocess.check_output( + ["git", "-C", root, "status", "--porcelain", "--untracked-files=all"], + text=True, + ).strip() + if dirty: + raise RuntimeError(f"Pinned Puzzletron dependency {package!r} is dirty: {dirty}") + return repository, commit + return payload.get("url"), vcs_info.get("commit_id") + + +def verify_installed_vcs_source(package: str, expected: dict[str, Any]) -> None: + """Require an installed VCS dependency to match its repository and commit.""" + + repository, commit = _installed_vcs_source(package) + expected_source = (_normalized_repository(expected["repository"]), expected["commit"]) + actual_source = (_normalized_repository(repository), commit) + if actual_source != expected_source: + raise RuntimeError( + f"Pinned Puzzletron dependency {package!r} source mismatch: " + f"actual={actual_source!r}, expected={expected_source!r}" + ) diff --git a/examples/puzzletron/distributed_eval/run_depth_pool.sh b/examples/puzzletron/distributed_eval/run_depth_pool.sh index 3b8b663c9b4..2ab40f3796b 100644 --- a/examples/puzzletron/distributed_eval/run_depth_pool.sh +++ b/examples/puzzletron/distributed_eval/run_depth_pool.sh @@ -1,6 +1,18 @@ #!/usr/bin/env bash # 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. set -Eeuo pipefail @@ -10,10 +22,12 @@ set -Eeuo pipefail : "${WORKER_COUNT:?set WORKER_COUNT to the number of worker groups}" : "${PUZZLETRON_GROUP_INDEX:=${PUZZLETRON_TASK_INDEX:-${SLURM_PROCID:-}}}" : "${PUZZLETRON_GROUP_INDEX:?run this script as one orchestrator worker-group task}" +: "${PUZZLETRON_GROUP_RANK:=0}" PYTHON_BIN="${PYTHON_BIN:-python}" SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" GROUP_INDEX="${PUZZLETRON_GROUP_INDEX}" +GROUP_RANK="${PUZZLETRON_GROUP_RANK}" JOB_ID="${SLURM_JOB_ID:-local}" WORKER_PREFIX="${JOB_ID}-depth-" MANIFEST_PATH="${CAMPAIGN_DIR}/manifest.json" @@ -38,7 +52,7 @@ cleanup() { local rc=$? trap - EXIT INT TERM set +e - if [[ "${GROUP_INDEX}" == "0" ]]; then + if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" ]]; then drain_workers fi if [[ -n "${worker_pid}" ]] && kill -0 "${worker_pid}" 2>/dev/null; then @@ -50,7 +64,7 @@ cleanup() { trap cleanup EXIT INT TERM # Rank 0 creates the shared campaign before any worker attempts to open it. -if [[ "${GROUP_INDEX}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then +if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then CUDA_VISIBLE_DEVICES="" "${PYTHON_BIN}" \ -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ @@ -72,21 +86,17 @@ done # Every scheduler task owns one GPU slice and starts one worker group. Multiple # independent worker groups may share a node. -export NNODES=1 -export NODE_RANK=0 export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORLD_SIZE}}" export WORKER_GROUP_INDEX="${GROUP_INDEX}" export WORKER_ID="${WORKER_PREFIX}${GROUP_INDEX}" export WORKER_HOST="${WORKER_HOST:-$(hostname -f)}" export WORKER_PORT="${WORKER_PORT:-$((5010 + GROUP_INDEX))}" -export RDZV_ENDPOINT="127.0.0.1:$((29500 + GROUP_INDEX))" -export RDZV_ID="depth-${JOB_ID}-${GROUP_INDEX}" bash "${SCRIPT_DIR}/run_worker.sh" & worker_pid=$! coordinator_rc=0 -if [[ "${GROUP_INDEX}" == "0" ]]; then +if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" ]]; then # Do not start depth iteration zero until every resident model is ready. CUDA_VISIBLE_DEVICES="" "${PYTHON_BIN}" - \ "${CAMPAIGN_DIR}" \ diff --git a/examples/puzzletron/distributed_eval/run_replacement_pool.sh b/examples/puzzletron/distributed_eval/run_replacement_pool.sh index 283a9ef11a3..c5ebc9ebc56 100755 --- a/examples/puzzletron/distributed_eval/run_replacement_pool.sh +++ b/examples/puzzletron/distributed_eval/run_replacement_pool.sh @@ -1,6 +1,18 @@ #!/usr/bin/env bash # 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. set -Eeuo pipefail @@ -10,10 +22,12 @@ set -Eeuo pipefail : "${WORKER_COUNT:?set WORKER_COUNT to the number of worker groups}" : "${PUZZLETRON_GROUP_INDEX:=${PUZZLETRON_TASK_INDEX:-${SLURM_PROCID:-}}}" : "${PUZZLETRON_GROUP_INDEX:?run this script as one orchestrator worker-group task}" +: "${PUZZLETRON_GROUP_RANK:=0}" PYTHON_BIN="${PYTHON_BIN:-python}" SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" GROUP_INDEX="${PUZZLETRON_GROUP_INDEX}" +GROUP_RANK="${PUZZLETRON_GROUP_RANK}" JOB_ID="${SLURM_JOB_ID:-local}" WORKER_PREFIX="${JOB_ID}-replacement-" MANIFEST_PATH="${CAMPAIGN_DIR}/manifest.json" @@ -38,7 +52,7 @@ cleanup() { local rc=$? trap - EXIT INT TERM set +e - if [[ "${GROUP_INDEX}" == "0" ]]; then + if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" ]]; then drain_workers fi if [[ -n "${worker_pid}" ]] && kill -0 "${worker_pid}" 2>/dev/null; then @@ -49,7 +63,7 @@ cleanup() { } trap cleanup EXIT INT TERM -if [[ "${GROUP_INDEX}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then +if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" && ! -f "${MANIFEST_PATH}" ]]; then CUDA_VISIBLE_DEVICES="" "${PYTHON_BIN}" \ -m modelopt.torch.puzzletron.distributed_eval.cli init \ --campaign-dir "${CAMPAIGN_DIR}" \ @@ -69,21 +83,17 @@ while [[ ! -f "${MANIFEST_PATH}" ]]; do sleep 1 done -export NNODES=1 -export NODE_RANK=0 export NPROC_PER_NODE="${NPROC_PER_NODE:-${WORLD_SIZE}}" export WORKER_GROUP_INDEX="${GROUP_INDEX}" export WORKER_ID="${WORKER_PREFIX}${GROUP_INDEX}" export WORKER_HOST="${WORKER_HOST:-$(hostname -f)}" export WORKER_PORT="${WORKER_PORT:-$((5010 + GROUP_INDEX))}" -export RDZV_ENDPOINT="127.0.0.1:$((29500 + GROUP_INDEX))" -export RDZV_ID="replacement-${JOB_ID}-${GROUP_INDEX}" bash "${SCRIPT_DIR}/run_worker.sh" & worker_pid=$! coordinator_rc=0 -if [[ "${GROUP_INDEX}" == "0" ]]; then +if [[ "${GROUP_INDEX}" == "0" && "${GROUP_RANK}" == "0" ]]; then CUDA_VISIBLE_DEVICES="" "${PYTHON_BIN}" - \ "${CAMPAIGN_DIR}" \ "${WORKER_COUNT}" \ diff --git a/examples/puzzletron/distributed_eval/run_worker.sh b/examples/puzzletron/distributed_eval/run_worker.sh index b95b77f3119..6cfe637a915 100755 --- a/examples/puzzletron/distributed_eval/run_worker.sh +++ b/examples/puzzletron/distributed_eval/run_worker.sh @@ -1,4 +1,19 @@ #!/usr/bin/env bash +# 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. + set -Eeuo pipefail : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}" @@ -6,11 +21,21 @@ set -Eeuo pipefail PYTHON_BIN="${PYTHON_BIN:-python}" TORCHRUN="${TORCHRUN:-torchrun}" -NNODES="${NNODES:-1}" NPROC_PER_NODE="${NPROC_PER_NODE:-8}" -NODE_RANK="${NODE_RANK:-0}" -RDZV_ID="${RDZV_ID:-distributed-eval-${SLURM_JOB_ID:-local}}" -RDZV_ENDPOINT="${RDZV_ENDPOINT:-127.0.0.1:29500}" +if [[ -n "${PUZZLETRON_GROUP_SIZE:-}" ]]; then + : "${PUZZLETRON_GROUP_RANK:?set PUZZLETRON_GROUP_RANK with task identity}" + : "${PUZZLETRON_RENDEZVOUS_ENDPOINT:?set PUZZLETRON_RENDEZVOUS_ENDPOINT with task identity}" + : "${PUZZLETRON_RENDEZVOUS_ID:?set PUZZLETRON_RENDEZVOUS_ID with task identity}" + NNODES="${PUZZLETRON_GROUP_SIZE}" + NODE_RANK="${PUZZLETRON_GROUP_RANK}" + RDZV_ENDPOINT="${PUZZLETRON_RENDEZVOUS_ENDPOINT}" + RDZV_ID="${PUZZLETRON_RENDEZVOUS_ID}" +else + NNODES="${NNODES:-1}" + NODE_RANK="${NODE_RANK:-0}" + RDZV_ID="${RDZV_ID:-distributed-eval-${SLURM_JOB_ID:-local}}" + RDZV_ENDPOINT="${RDZV_ENDPOINT:-127.0.0.1:29500}" +fi WORKER_HOST="${WORKER_HOST:-$(hostname -f)}" WORKER_PORT="${WORKER_PORT:-5010}" WORKER_ID="${WORKER_ID:-${SLURM_JOB_ID:-local}-group-${WORKER_GROUP_INDEX:-0}}" diff --git a/modelopt/torch/puzzletron/benchmarks/aiperf.py b/modelopt/torch/puzzletron/benchmarks/aiperf.py index 872822777bd..561f9c2e018 100644 --- a/modelopt/torch/puzzletron/benchmarks/aiperf.py +++ b/modelopt/torch/puzzletron/benchmarks/aiperf.py @@ -355,6 +355,14 @@ def _clean_subprocess_environment( return env +def _aiperf_subprocess_environment(env: dict[str, str]) -> dict[str, str]: + """Avoid AIPerf's offline resolver for an explicit local tokenizer path.""" + resolved = dict(env) + resolved.pop("HF_HUB_OFFLINE", None) + resolved.pop("TRANSFORMERS_OFFLINE", None) + return resolved + + def run_aiperf_sweep( checkpoint_dir: str | Path, *, @@ -430,6 +438,7 @@ def run_aiperf_sweep( ) for key, value in (topology.get("env") or {}).items(): env[str(key)] = str(value) + aiperf_env = _aiperf_subprocess_environment(env) cached: dict[int, BenchmarkResult] = {} missing: list[tuple[int, Path, list[str], str]] = [] for concurrency in concurrency_values: @@ -496,7 +505,7 @@ def run_aiperf_sweep( command, check=True, timeout=benchmark_timeout, - env=env, + env=aiperf_env, ) export = run_dir / "profile_export_aiperf.json" if not export.is_file(): diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index 8c506b23614..a1eacd356cc 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/pool.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. """Persistent pool adapter for coordinator/worker stages.""" @@ -258,12 +270,15 @@ def command( effective_overrides.extend(_replacement_overrides(plan, replacement_puzzle_dir)) if role == "gang": worker_count = int(item.metadata.get("worker_count", node.instances)) + allocation_nodes, allocation_gpus, topology = packed_allocation( + node, instances=worker_count + ) env = { "CAMPAIGN_DIR": str(campaign_dir), "CONFIG_PATH": plan.experiment_config_path, "PUZZLE_DIR": str(replacement_puzzle_dir), "WORLD_SIZE": str(node.gpus_per_instance), - "NPROC_PER_NODE": str(node.gpus_per_instance), + "NPROC_PER_NODE": str(topology.gpus_per_task), "WORKER_COUNT": str(worker_count), } if node.stage_id == "depth_importance": @@ -288,9 +303,7 @@ def command( / _replacement_completion_identity(plan, root_overrides) ), "FINALIZE_COMPLETION_MARKER": f"width-{width}", - "FINALIZE_EXPECTED_COMPLETIONS": str( - len(replacement_widths) - ), + "FINALIZE_EXPECTED_COMPLETIONS": str(len(replacement_widths)), } ) script = repo / "examples/puzzletron/distributed_eval/run_replacement_pool.sh" @@ -298,9 +311,6 @@ def command( existing = env.get("DISTRIBUTED_EVAL_OVERRIDES", "") env["DISTRIBUTED_EVAL_OVERRIDES"] = f"{existing}\n{override}".strip() log_path = str(log_dir / f"{node.stage_id}_gang_{attempt_id}.log") - allocation_nodes, allocation_gpus, topology = packed_allocation( - node, instances=worker_count - ) return AttemptSpec( attempt_id=attempt_id, work_id=item.work_id, @@ -356,13 +366,9 @@ def command( if role == "worker": env["CUDA_VISIBLE_DEVICES"] = ",".join(str(gpu) for gpu in item.local_gpu_ids) env["NPROC_PER_NODE"] = str(node.gpus_per_instance) - env["NNODES"] = "1" - env["NODE_RANK"] = "0" worker_id = int(item.metadata.get("worker_id", 0)) env["WORKER_GROUP_INDEX"] = str(worker_id) env["WORKER_PORT"] = str(5010 + worker_id) - env["RDZV_ENDPOINT"] = f"127.0.0.1:{29500 + worker_id}" - env["RDZV_ID"] = f"{node.stage_id}-{attempt_id}" argv = ("bash", str(script)) # GPU partitions reject zero-GPU jobs; coordinators still need one GPU slot. allocation_gpus = 1 if role == "coordinator" else node.gpus_per_instance diff --git a/modelopt/torch/puzzletron/orchestration/compiler.py b/modelopt/torch/puzzletron/orchestration/compiler.py index 1eaf8992e49..58ac3229046 100644 --- a/modelopt/torch/puzzletron/orchestration/compiler.py +++ b/modelopt/torch/puzzletron/orchestration/compiler.py @@ -596,6 +596,7 @@ def compile_campaign_plan( execution_defaults=_mapping(execution.get("defaults")), stages=tuple(nodes), contract_hash=contract_hash, + overrides=tuple(overrides or ()), ) @@ -606,6 +607,7 @@ def plan_to_dict(plan: CampaignPlan) -> dict[str, Any]: "experiment_config_path": plan.experiment_config_path, "puzzle_dir": str(plan.puzzle_dir), "contract_hash": plan.contract_hash, + "overrides": list(plan.overrides), "runner_kind": plan.runner.kind, "execution_defaults": dict(plan.execution_defaults), "stages": [ diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index 48183aa0249..7d478d3ebff 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.py @@ -133,6 +133,10 @@ def dry_run_plan( *, overrides: list[str] | None = None, ) -> list[DryRunSubmission]: + if overrides is not None and tuple(overrides) != plan.overrides: + raise ValueError( + "dry-run overrides must match the overrides compiled into the campaign plan" + ) submissions: list[DryRunSubmission] = [] for node in plan.stages: adapter = adapter_for_stage(node) @@ -145,7 +149,7 @@ def dry_run_plan( item=item, attempt_id=attempt_id, runner=plan.runner, - overrides=overrides, + overrides=list(plan.overrides), ) topology = resolve_task_topology(attempt) submissions.append( @@ -378,6 +382,7 @@ def _stage_execution_identity( self.plan.experiment_config, node.stage_id ), "compiled_node": compiled_node, + "root_overrides": list(self.plan.overrides), "work_items": [ { "work_id": item.work_id, @@ -517,6 +522,29 @@ def _recover_failed_stages(self) -> None: or stage_is_complete(self.plan.experiment_config, node.stage_id) ): continue + finalization_failures = [ + (attempt.metadata or {}).get("stage_finalization_failure") + for attempt in record.attempts + ] + if all( + isinstance(failure, Mapping) and failure.get("phase") == "aggregation" + for failure in finalization_failures + ): + failure = finalization_failures[0] + assert isinstance(failure, Mapping) + self._finalization_failures[node.stage_id] = _FinalizationFailure( + phase="aggregation", + reason=str(failure.get("reason") or "stage aggregation failed"), + exception_type=( + str(failure["exception_type"]) + if failure.get("exception_type") is not None + else None + ), + ) + self.logger.wait( + f"{node.stage_id}: recovered aggregation failure; retrying finalization" + ) + continue self._failed_stages.add(node.stage_id) self.logger.error(f"{node.stage_id}: recovered terminal stage validation failure") @@ -582,7 +610,7 @@ def _available_nodes(self) -> int | None: active_nodes += int((attempt.get("allocation") or {}).get("nodes", 0)) return max(0, slurm.max_nodes - active_nodes) - def _submit_stage(self, node: StagePlanNode, *, overrides: list[str] | None = None) -> bool: + def _submit_stage(self, node: StagePlanNode) -> bool: if self._stage_has_active_or_completed_work(node): return False adapter = adapter_for_stage(node) @@ -618,7 +646,7 @@ def _submit_stage(self, node: StagePlanNode, *, overrides: list[str] | None = No item=item, attempt_id=attempt_id, runner=self.plan.runner, - overrides=overrides, + overrides=list(self.plan.overrides), ), ) if available_nodes is not None and attempt.allocation_nodes > available_nodes: @@ -784,7 +812,20 @@ def _fail_stage_if_artifacts_did_not_settle( reason = failure.reason if failure is not None else "required artifacts are incomplete" expected_artifacts = list(failure.artifacts) if failure is not None else [] phase = failure.phase if failure is not None else "validation" - persisted_attempts = self._persisted_stage_attempts(node) + persisted_attempts = [ + replace( + attempt, + metadata={ + **dict(attempt.metadata or {}), + "stage_finalization_failure": { + "phase": phase, + "reason": reason, + "exception_type": (failure.exception_type if failure is not None else None), + }, + }, + ) + for attempt in self._persisted_stage_attempts(node) + ] self._failed_stages.add(node.stage_id) self.store.write_stage_record( StageRunRecord( @@ -1281,6 +1322,11 @@ def run( ) -> dict[str, Any]: """Run the controller until all stages complete or a fatal failure occurs.""" + if overrides is not None and tuple(overrides) != self.plan.overrides: + raise ValueError( + "runtime overrides must match the overrides compiled into the campaign plan" + ) + iterations = 0 halted = False cancelled = False @@ -1376,7 +1422,7 @@ def _on_signal(signum: int, _frame: object | None) -> None: for node in self._ready_nodes(): if self._shutdown_requested: break - self._submit_stage(node, overrides=overrides) + self._submit_stage(node) self._refresh_dashboard( drain_pending=bool( self._failed_stages and (self._active or self._ready_nodes()) diff --git a/modelopt/torch/puzzletron/orchestration/schema.py b/modelopt/torch/puzzletron/orchestration/schema.py index 5b2e5a20227..0853abe88de 100644 --- a/modelopt/torch/puzzletron/orchestration/schema.py +++ b/modelopt/torch/puzzletron/orchestration/schema.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. """Public contracts for Puzzletron campaign orchestration.""" @@ -206,6 +218,7 @@ class CampaignPlan: execution_defaults: Mapping[str, Any] stages: tuple[StagePlanNode, ...] contract_hash: str + overrides: tuple[str, ...] = () @dataclass(frozen=True) diff --git a/modelopt/torch/puzzletron/orchestration/task_launcher.py b/modelopt/torch/puzzletron/orchestration/task_launcher.py index 3f4a69ea2fd..e61a3ebaa58 100644 --- a/modelopt/torch/puzzletron/orchestration/task_launcher.py +++ b/modelopt/torch/puzzletron/orchestration/task_launcher.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. """Dependency-light launcher for one task in an orchestration attempt.""" @@ -39,6 +51,7 @@ "PUZZLETRON_LOCAL_TASK_INDEX", "PUZZLETRON_MASTER_ADDR", "PUZZLETRON_MASTER_PORT", + "PUZZLETRON_RENDEZVOUS_ENDPOINT", "PUZZLETRON_RENDEZVOUS_ID", "PUZZLETRON_TASK_HOSTS", "PUZZLETRON_TASK_INDEX", @@ -212,9 +225,7 @@ def main(argv: Sequence[str] | None = None) -> int: f"task {task_index} expected {args.gpus_per_task} visible GPUs, got {visible_gpus}" ) tasks_per_node = ( - args.task_count - if args.gpus_per_task == 0 - else args.gpus_per_node // args.gpus_per_task + args.task_count if args.gpus_per_task == 0 else args.gpus_per_node // args.gpus_per_task ) expected_hosts = math.ceil(args.task_count / tasks_per_node) if expected_hosts > args.nodes: @@ -243,6 +254,7 @@ def main(argv: Sequence[str] | None = None) -> int: PUZZLETRON_GROUP_SIZE=str(binding.group_size), PUZZLETRON_MASTER_ADDR=binding.master_addr, PUZZLETRON_MASTER_PORT=str(binding.master_port), + PUZZLETRON_RENDEZVOUS_ENDPOINT=rendezvous_endpoint(binding), PUZZLETRON_RENDEZVOUS_ID=binding.rendezvous_id, ) print( diff --git a/noxfile.py b/noxfile.py index 25b1a2cc4c2..f6ebc70faf2 100644 --- a/noxfile.py +++ b/noxfile.py @@ -58,6 +58,7 @@ with PUZZLETRON_V2_CI_ENVIRONMENT_PATH.open(encoding="utf-8") as environment_file: PUZZLETRON_V2_CI_ENVIRONMENT = json.load(environment_file) PUZZLETRON_V2_AUTOMODEL_SOURCE = PUZZLETRON_V2_CI_ENVIRONMENT["nemo_automodel"] +PUZZLETRON_V2_LMMS_SOURCE = PUZZLETRON_V2_CI_ENVIRONMENT["lmms_eval"] PUZZLETRON_V2_AUTOMODEL = ( "nemo-automodel @ git+" f"{PUZZLETRON_V2_AUTOMODEL_SOURCE['repository']}@" @@ -72,30 +73,46 @@ def _verify_puzzletron_v2_environment(session): "torch": PUZZLETRON_V2_CI_ENVIRONMENT["torch"], "torchvision": PUZZLETRON_V2_CI_ENVIRONMENT["torchvision"], "transformers": PUZZLETRON_V2_CI_ENVIRONMENT["transformers"], - "lmms-eval": PUZZLETRON_V2_CI_ENVIRONMENT["lmms_eval"], + "lmms-eval": PUZZLETRON_V2_LMMS_SOURCE["base_version"], "nemo-automodel": PUZZLETRON_V2_AUTOMODEL_SOURCE["base_version"], } + expected_vcs = { + "lmms-eval": PUZZLETRON_V2_LMMS_SOURCE, + "nemo-automodel": PUZZLETRON_V2_AUTOMODEL_SOURCE, + } session.run( "python", "-c", - ( - "import sys; " - "from importlib.metadata import version; " - "from packaging.version import Version; " - f"expected = {expected_versions!r}; " - "actual = {" - "'python': f'{sys.version_info.major}.{sys.version_info.minor}', " - "'torch': Version(version('torch')).base_version, " - "'torchvision': Version(version('torchvision')).base_version, " - "'transformers': Version(version('transformers')).base_version, " - "'lmms-eval': Version(version('lmms-eval')).base_version, " - "'nemo-automodel': Version(version('nemo-automodel')).base_version}; " - "mismatches = {name: (actual[name], expected_version) " - "for name, expected_version in expected.items() " - "if actual[name] != expected_version}; " - "assert not mismatches, " - "f'Pinned Puzzletron CI environment mismatch: {mismatches}'" - ), + f""" +import json +import sys +from importlib.metadata import distribution, version + +from packaging.version import Version + +from examples.puzzletron.ci_environment import verify_installed_vcs_source + +expected = {expected_versions!r} +expected_vcs = {expected_vcs!r} +actual = {{ + "python": f"{{sys.version_info.major}}.{{sys.version_info.minor}}", + "torch": Version(version("torch")).base_version, + "torchvision": Version(version("torchvision")).base_version, + "transformers": Version(version("transformers")).base_version, + "lmms-eval": Version(version("lmms-eval")).base_version, + "nemo-automodel": Version(version("nemo-automodel")).base_version, +}} +mismatches = {{ + name: (actual[name], expected_version) + for name, expected_version in expected.items() + if actual[name] != expected_version +}} + +for name, source in expected_vcs.items(): + verify_installed_vcs_source(name, source) + +assert not mismatches, f"Pinned Puzzletron CI environment mismatch: {{mismatches}}" +""", ) diff --git a/tests/_test_utils/torch/puzzletron/utils.py b/tests/_test_utils/torch/puzzletron/utils.py index 091047ec955..6738cec7d90 100644 --- a/tests/_test_utils/torch/puzzletron/utils.py +++ b/tests/_test_utils/torch/puzzletron/utils.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Cross-model tiny-checkpoint helpers for tests.""" + import os import torch diff --git a/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py b/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py index 616b4d24d9b..452c3280663 100644 --- a/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py +++ b/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Tests for Puzzletron AIPerf context-capacity handling.""" + import json from pathlib import Path from types import SimpleNamespace @@ -20,6 +22,7 @@ import pytest from modelopt.torch.puzzletron.benchmarks.aiperf import ( + _aiperf_subprocess_environment, _canonical_topology, _clean_subprocess_environment, _exact_length_extra_inputs, @@ -103,6 +106,26 @@ def test_prepare_vllm_checkpoint_leaves_native_teacher_unchanged(tmp_path): assert _prepare_vllm_checkpoint(tmp_path) is False +def test_aiperf_environment_avoids_broken_offline_local_path_resolution(): + expected_source = { + "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + "HF_DATASETS_OFFLINE": "1", + "HF_HOME": "/cache/huggingface", + "UNCHANGED": "value", + } + source = dict(expected_source) + + resolved = _aiperf_subprocess_environment(source) + + assert "HF_HUB_OFFLINE" not in resolved + assert "TRANSFORMERS_OFFLINE" not in resolved + assert resolved["HF_DATASETS_OFFLINE"] == "1" + assert resolved["HF_HOME"] == "/cache/huggingface" + assert resolved["UNCHANGED"] == "value" + assert source == expected_source + + def test_canonical_topology_covers_tp_pp_dp_effective_ep_and_context_parallel(): topology = _canonical_topology( { @@ -179,6 +202,7 @@ def test_profile_command_maps_each_workload_answer_to_aiperf_cli(tmp_path): assert command[command.index("--synthetic-input-tokens-stddev") + 1] == "0" assert command[command.index("--output-tokens-mean") + 1] == "128" assert command[command.index("--output-tokens-stddev") + 1] == "0" + assert command[command.index("--tokenizer") + 1] == str(tmp_path / "tokenizer") assert "--use-server-token-count" in command diff --git a/tests/unit/torch/puzzletron/test_ci_environment.py b/tests/unit/torch/puzzletron/test_ci_environment.py new file mode 100644 index 00000000000..72878bcaf84 --- /dev/null +++ b/tests/unit/torch/puzzletron/test_ci_environment.py @@ -0,0 +1,62 @@ +# 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 CI environment provenance checks.""" + +import json + +import pytest + +from examples.puzzletron import ci_environment + + +class _Distribution: + def __init__(self, payload: dict): + self.payload = payload + + def read_text(self, filename: str) -> str | None: + assert filename == "direct_url.json" + return json.dumps(self.payload) + + +def test_editable_pinned_dependency_must_be_clean(monkeypatch): + monkeypatch.setattr( + ci_environment.metadata, + "distribution", + lambda _package: _Distribution( + {"url": "file:///src/automodel", "dir_info": {"editable": True}} + ), + ) + outputs = iter( + [ + "https://github.com/Separius/Automodel.git\n", + "b22cd029d806197e249f2cc4a42c5de91713b772\n", + " M nemo_automodel/model.py\n", + ] + ) + monkeypatch.setattr( + ci_environment.subprocess, + "check_output", + lambda *_args, **_kwargs: next(outputs), + ) + + with pytest.raises(RuntimeError, match="dependency 'nemo-automodel' is dirty"): + ci_environment.verify_installed_vcs_source( + "nemo-automodel", + { + "repository": "https://github.com/Separius/Automodel.git", + "commit": "b22cd029d806197e249f2cc4a42c5de91713b772", + }, + ) diff --git a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py index 6b74760aba8..21c4bdced44 100644 --- a/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py +++ b/tests/unit/torch/puzzletron/test_diagnostic_scoring_config.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Tests for Puzzletron diagnostic scoring configuration.""" + import json from copy import deepcopy from pathlib import Path diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index fc60799a0b3..4ac48262fa0 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -680,6 +680,54 @@ def test_depth_pool_packs_four_two_gpu_workers_per_node(tmp_path: Path): assert "--gpus-per-task=2" in script +def test_depth_pool_splits_one_sixteen_gpu_worker_across_two_nodes(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"), + ) + node = StagePlanNode( + stage_id="depth_importance", + strategy=ExecutionStrategy.PERSISTENT_POOL, + instances=1, + failure_policy=FailurePolicy.STRICT, + mesh={"tp": 2, "cp": 1, "pp": 2, "ep": 2, "dp_shard": 2, "dp_replicate": 1}, + gpus_per_instance=16, + gpus_per_node=8, + nodes=2, + total_gpus=16, + exclusive=True, + parents=("tokenize_data",), + distributed=True, + partition="batch", + ) + plan = CampaignPlan( + experiment_config_path=str(tmp_path / "experiment.yaml"), + puzzle_dir=tmp_path / "run", + experiment_config={"depth_importance": {"output_dir": str(tmp_path / "depth")}}, + runner=runner, + execution_defaults={"gpus_per_node": 8}, + stages=(node,), + contract_hash="contract", + ) + + adapter = adapter_for_stage(node) + item = adapter.plan(plan, node).items[0] + attempt = adapter.command( + plan=plan, + node=node, + item=item, + attempt_id="a1", + runner=runner, + ) + + assert attempt.allocation_nodes == 2 + assert attempt.task_topology.task_count == 2 + assert attempt.task_topology.tasks_per_group == 2 + assert attempt.task_topology.gpus_per_task == 8 + assert attempt.command.env["NPROC_PER_NODE"] == "8" + + def test_post_mip_workers_share_one_packed_allocation(tmp_path: Path): runner = RunnerEnvironment( kind="slurm", @@ -847,8 +895,7 @@ def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path) "puzzle_dir=" not in attempt.command.env["FINALIZE_OVERRIDES"] for attempt in attempts ) assert all( - "puzzle_dir=" in attempt.command.env["DISTRIBUTED_EVAL_OVERRIDES"] - for attempt in attempts + "puzzle_dir=" in attempt.command.env["DISTRIBUTED_EVAL_OVERRIDES"] for attempt in attempts ) assert [attempt.command.env["FINALIZE_EXPECTED_COMPLETIONS"] for attempt in attempts] == [ "2", diff --git a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py index aa1653ed054..b64d3e4dfe4 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py +++ b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py @@ -635,6 +635,47 @@ def test_stage_execution_identity_ignores_unrelated_config(tmp_path: Path): assert identity_a == identity_b +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"] + plan = compile_campaign_plan( + experiment_config_path=experiment, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + overrides=compiled_overrides, + stage_filter="convert", + ) + executor = _TrackingFakeExecutor() + controller = CampaignController(plan, executor=executor) + baseline_plan = compile_campaign_plan( + experiment_config_path=experiment, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + stage_filter="convert", + ) + + with pytest.raises(ValueError, match="must match the overrides compiled"): + controller.run(overrides=["convert.model_path=/models/runtime"], once=True) + + assert executor.submitted_stage_ids == [] + assert controller.store.list_attempts("convert") == [] + + result = controller.run(overrides=compiled_overrides, once=True) + + assert result["halted"] is False + assert executor.submitted_stage_ids == ["convert"] + submitted_attempt = next(iter(executor._attempts.values())) + override_index = submitted_attempt.command.argv.index("--override") + assert submitted_attempt.command.argv[override_index + 1] == compiled_overrides[0] + assert submitted_attempt.metadata["stage_execution_identity"] == ( + controller._stage_execution_identity(plan.stages[0]) + ) + baseline_identity = CampaignController( + baseline_plan, executor=_FakeExecutor() + )._stage_execution_identity(baseline_plan.stages[0]) + assert submitted_attempt.metadata["stage_execution_identity"] != baseline_identity + + def test_controller_revalidates_recent_completed_work_before_resubmitting( tmp_path: Path, monkeypatch ): @@ -830,16 +871,20 @@ def test_controller_fails_when_completed_work_artifacts_do_not_settle( ) class _MissingArtifactsAdapter: + aggregation_ready = False + def __getattr__(self, name): return getattr(delegate, name) def aggregate(self, *, plan, node, work_plan): - if aggregation_failure: + if aggregation_failure and not self.aggregation_ready: raise FileNotFoundError("shards are still publishing") + return {"status": "complete"} def validate(self, *, plan, node): if aggregation_failure: - pytest.fail("validation must wait until aggregation succeeds") + assert self.aggregation_ready + return ValidatedResult(valid=True, reason="stage outputs present") return ValidatedResult( valid=False, reason="stage outputs missing", @@ -891,15 +936,20 @@ def validate(self, *, plan, node): assert stage_record is not None assert stage_record.status == JobState.FAILED.value assert stage_record.attempts[0].status == JobState.COMPLETED.value + assert stage_record.attempts[0].metadata["stage_finalization_failure"]["phase"] == phase recovered_executor = _TrackingFakeExecutor() + missing.aggregation_ready = aggregation_failure recovered = CampaignController(plan, executor=recovered_executor) recovered_result = recovered.run(once=True) - assert recovered_result["halted"] is True - assert recovered_result["failed_stages"] == ["convert"] - assert recovered_executor.submitted_stage_ids == [] + assert recovered_result["halted"] is (not aggregation_failure) + assert recovered_result["failed_stages"] == ([] if aggregation_failure else ["convert"]) + assert recovered_executor.submitted_stage_ids == ( + ["final_report"] if aggregation_failure else [] + ) assert len(list(controller.store.events_root.glob(f"*_{event_name}.json"))) == 1 + assert recovered.store.stage_is_complete("convert") is aggregation_failure @pytest.mark.parametrize("stale_kind", ["contract", "stage_execution"]) diff --git a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py index c62f40c5794..65be2c36cfe 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py +++ b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py @@ -1,8 +1,24 @@ # 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 explicit orchestration task topology.""" +import os +import subprocess +from pathlib import Path + import pytest from puzzletron_orchestrator import task_launcher @@ -156,6 +172,116 @@ def fake_execvpe(executable, command, env) -> None: assert result == 0 assert captured["env"]["CUDA_VISIBLE_DEVICES"] == expected assert captured["env"]["PUZZLETRON_TASK_LAUNCHER"] == "direct" + assert captured["env"]["PUZZLETRON_RENDEZVOUS_ENDPOINT"] == "localhost:0" + + +def test_task_launcher_exports_shared_multi_node_rendezvous(monkeypatch) -> None: + captured: dict[str, object] = {} + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7") + monkeypatch.setenv("PUZZLETRON_TASK_INDEX", "1") + monkeypatch.setenv("PUZZLETRON_LOCAL_TASK_INDEX", "0") + monkeypatch.setenv("PUZZLETRON_TASK_HOSTS", "node-a,node-b") + + def fake_execvpe(executable, command, env) -> None: + captured.update(executable=executable, command=command, env=env) + + monkeypatch.setattr(task_launcher.os, "execvpe", fake_execvpe) + + assert ( + task_launcher.main( + [ + "--attempt-id", + "attempt-a", + "--nodes", + "2", + "--gpus-per-node", + "8", + "--task-count", + "2", + "--gpus-per-task", + "8", + "--tasks-per-group", + "2", + "--launcher", + "direct", + "--", + "python", + "worker.py", + ] + ) + == 0 + ) + + env = captured["env"] + assert isinstance(env, dict) + assert env["PUZZLETRON_GROUP_SIZE"] == "2" + assert env["PUZZLETRON_GROUP_RANK"] == "1" + assert env["PUZZLETRON_RENDEZVOUS_ENDPOINT"].startswith("node-a:") + assert env["PUZZLETRON_RENDEZVOUS_ID"] == "attempt-a-group-0" + + +@pytest.mark.parametrize( + ("group_size", "group_rank", "endpoint"), + [("1", "0", "localhost:0"), ("2", "1", "node-a:23456")], +) +def test_run_worker_consumes_task_launcher_identity( + tmp_path: Path, + group_size: str, + group_rank: str, + endpoint: str, +) -> None: + script = Path(__file__).parents[4] / "examples/puzzletron/distributed_eval/run_worker.sh" + env = { + **os.environ, + "CAMPAIGN_DIR": str(tmp_path / "campaign"), + "CONFIG_PATH": str(tmp_path / "experiment.yaml"), + "TORCHRUN": "/bin/echo", + "NPROC_PER_NODE": "4", + "PUZZLETRON_GROUP_SIZE": group_size, + "PUZZLETRON_GROUP_RANK": group_rank, + "PUZZLETRON_RENDEZVOUS_ENDPOINT": endpoint, + "PUZZLETRON_RENDEZVOUS_ID": "attempt-a-group-0", + } + + result = subprocess.run( + ["bash", str(script)], + env=env, + check=True, + capture_output=True, + text=True, + ) + + assert f"--nnodes {group_size}" in result.stdout + assert f"--node-rank {group_rank}" in result.stdout + assert f"--rdzv-endpoint {endpoint}" in result.stdout + assert "--rdzv-id attempt-a-group-0" in result.stdout + + +@pytest.mark.parametrize("script_name", ["run_replacement_pool.sh", "run_depth_pool.sh"]) +def test_nonzero_group_rank_does_not_own_pool_control_path( + tmp_path: Path, script_name: str +) -> None: + campaign_dir = tmp_path / "campaign" + campaign_dir.mkdir() + (campaign_dir / "manifest.json").write_text("{}\n") + script = Path(__file__).parents[4] / "examples/puzzletron/distributed_eval" / script_name + env = { + **os.environ, + "CAMPAIGN_DIR": str(campaign_dir), + "CONFIG_PATH": str(tmp_path / "experiment.yaml"), + "WORLD_SIZE": "2", + "WORKER_COUNT": "1", + "NPROC_PER_NODE": "1", + "TORCHRUN": "/bin/true", + "PYTHON_BIN": "/bin/false", + "PUZZLETRON_GROUP_INDEX": "0", + "PUZZLETRON_GROUP_RANK": "1", + "PUZZLETRON_GROUP_SIZE": "2", + "PUZZLETRON_RENDEZVOUS_ENDPOINT": "node-a:23456", + "PUZZLETRON_RENDEZVOUS_ID": "attempt-a-group-0", + } + + subprocess.run(["bash", str(script)], env=env, check=True, timeout=10) def _task_binding(*, group_size: int) -> task_launcher.TaskBinding: From 342942683ef7287cd0fdcdc33ec850f299accbda Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 01:54:53 +0200 Subject: [PATCH 05/24] Fix Puzzletron GPU artifact assertion Replacement scoring artifacts are nested beneath both width and depth scenario directories. Include the depth directory when validating the distributed evaluation results. Signed-off-by: Johannes Rausch --- tests/gpu/torch/puzzletron/test_puzzletron.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index 61f85a6f743..c79682f425c 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -367,7 +367,7 @@ def test_tiny_qwen_campaign_uses_current_public_route( replacement_results = [ json.loads(path.read_text()) for path in smoke_root.glob( - "scenarios/*/distributed_eval/replacement_scoring/results/**/*.json" + "scenarios/width-*/depth-*/distributed_eval/replacement_scoring/results/**/*.json" ) ] assert replacement_results From c51e9897213108f07909a74de7a4d5ad1fde6c25 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 02:02:17 +0200 Subject: [PATCH 06/24] Fix Puzzletron CI style checks Apply the repository formatter and full license headers across the affected files. Normalize parameterized test values and document the fused-kernel compatibility names required by the test double. Signed-off-by: Johannes Rausch --- .../distributed_eval/run_coordinator.sh | 15 ++++++++++++ .../finalize_replacement_scoring.py | 12 ++++++++++ .../puzzletron/run_axis_diagnostic_worker.py | 15 ++++++++++++ .../distributed_eval/automodel_executor.py | 23 ++++++++++++++----- tests/gpu/torch/puzzletron/test_puzzletron.py | 4 +--- .../torch/puzzletron/test_width_scenarios.py | 4 +--- 6 files changed, 61 insertions(+), 12 deletions(-) diff --git a/examples/puzzletron/distributed_eval/run_coordinator.sh b/examples/puzzletron/distributed_eval/run_coordinator.sh index 62f5c2edf27..dc06ba6ca48 100755 --- a/examples/puzzletron/distributed_eval/run_coordinator.sh +++ b/examples/puzzletron/distributed_eval/run_coordinator.sh @@ -1,4 +1,19 @@ #!/usr/bin/env bash +# 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. + set -Eeuo pipefail : "${CAMPAIGN_DIR:?set CAMPAIGN_DIR}" diff --git a/examples/puzzletron/finalize_replacement_scoring.py b/examples/puzzletron/finalize_replacement_scoring.py index 6953534580f..d4e62e595ad 100644 --- a/examples/puzzletron/finalize_replacement_scoring.py +++ b/examples/puzzletron/finalize_replacement_scoring.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. """Publish replacement-scoring reports after distributed evaluation.""" diff --git a/examples/puzzletron/run_axis_diagnostic_worker.py b/examples/puzzletron/run_axis_diagnostic_worker.py index 97cdf4c66d4..6cd40a58359 100755 --- a/examples/puzzletron/run_axis_diagnostic_worker.py +++ b/examples/puzzletron/run_axis_diagnostic_worker.py @@ -1,4 +1,19 @@ #!/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. + """Run or finalize one independently distributed width-diagnostic axis.""" from __future__ import annotations diff --git a/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py b/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py index f4db0f1ae27..f5d8d39345a 100644 --- a/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py +++ b/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py @@ -1,3 +1,18 @@ +# 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. + """Long-lived AutoModel executor for distributed replace-block evaluation.""" from __future__ import annotations @@ -250,9 +265,7 @@ def evaluate(self, request: EvaluationRequest) -> EvaluationResult | None: "sliced_teacher_baseline": self.sliced_teacher_baseline, "observability": self.latest_observability, "score_device_type": getattr(self, "latest_score_device_type", None), - "visible_cuda_device_count": getattr( - self, "visible_cuda_device_count", None - ), + "visible_cuda_device_count": getattr(self, "visible_cuda_device_count", None), }, ) @@ -287,9 +300,7 @@ def _score(self, prune_target: dict | list[dict] | None) -> dict | None: with ExitStack() as stack: for target in prune_targets: layer_idx = int(target["layer_idx"]) - bypass_dir = target.pop( - "bypass_checkpoint_dir", self.bypass_checkpoint_dir - ) + bypass_dir = target.pop("bypass_checkpoint_dir", self.bypass_checkpoint_dir) stack.enter_context( recipe.block_checkpoint_overlay_context(bypass_dir, layer_idx) if bypass_dir is not None diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index c79682f425c..d97e716555e 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -311,9 +311,7 @@ def test_tiny_qwen_campaign_uses_current_public_route( replacement_node = next( node for node in compiled_plan.stages if node.stage_id == "replacement_scoring" ) - replacement_work = adapter_for_stage(replacement_node).plan( - compiled_plan, replacement_node - ) + replacement_work = adapter_for_stage(replacement_node).plan(compiled_plan, replacement_node) assert replacement_node.instances == 1 assert replacement_node.gpus_per_instance == 1 assert replacement_node.total_gpus == 1 diff --git a/tests/unit/torch/puzzletron/test_width_scenarios.py b/tests/unit/torch/puzzletron/test_width_scenarios.py index bb579e53aab..fa2f6afdcf6 100644 --- a/tests/unit/torch/puzzletron/test_width_scenarios.py +++ b/tests/unit/torch/puzzletron/test_width_scenarios.py @@ -46,9 +46,7 @@ from puzzletron_orchestrator.adapters.stage_compat import stage_is_complete -def test_replacement_scoring_finalizer_publishes_current_terminal_manifest( - tmp_path, monkeypatch -): +def test_replacement_scoring_finalizer_publishes_current_terminal_manifest(tmp_path, monkeypatch): config_path = tmp_path / "experiment.yaml" config_path.touch() root_override = "+replacement_scoring.automodel.lm_head_backend=streaming" From 7f082eb611cfb514525b6a7e91e34e3bb952ed3a Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 02:15:15 +0200 Subject: [PATCH 07/24] Fix Puzzletron MIP solution assertion Accept either configured pruned FFN width in the selected architecture. The 90 percent parameter constraint does not guarantee that the optimizer chooses the most aggressive candidate. Signed-off-by: Johannes Rausch --- tests/gpu/torch/puzzletron/test_puzzletron.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index d97e716555e..fca04ba1c87 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -397,7 +397,12 @@ def test_tiny_qwen_campaign_uses_current_public_route( ] solutions = [solution for path in solution_paths for solution in json.loads(path.read_text())] assert solutions - assert 256 in _nested_values(solutions, "intermediate_size") + solution_ffn_widths = { + int(width) + for solution in solutions + for width in _nested_values(solution["chosen_block_configs"], "intermediate_size") + } + assert solution_ffn_widths.intersection({256, 512}) model = AutoModelForCausalLM.from_pretrained( smoke_root / "ckpts/sorted_teacher", From 12de31390c6f8e5cca70bc642ea1a2523fcd1524 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 02:38:11 +0200 Subject: [PATCH 08/24] Fix Puzzletron checkpoint reload probe Disable KV caching only for the synthetic all-full-attention Qwen 3.5 forward. This preserves the CUDA reload check without relying on a linear-attention cache layer that the hermetic fixture intentionally omits. Signed-off-by: Johannes Rausch --- tests/gpu/torch/puzzletron/test_puzzletron.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index fca04ba1c87..c0506bcff61 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -410,7 +410,10 @@ def test_tiny_qwen_campaign_uses_current_public_route( local_files_only=True, ).cuda() with torch.no_grad(): - logits = model(torch.tensor([[1, 2, 3, 4]], device="cuda")).logits + logits = model( + torch.tensor([[1, 2, 3, 4]], device="cuda"), + use_cache=False, + ).logits assert torch.isfinite(logits).all() durable_paths = [*manifests, *pass_manifests] From b2196078348eb236e365a8d4aa1f11cd91da525a Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 03:34:50 +0200 Subject: [PATCH 09/24] Exercise full Puzzletron GPU campaign lifecycle Share the hermetic tiny-Qwen campaign fixture and validate the selected KD checkpoint, post-MIP report, and no-op resume through the public orchestrator. Signed-off-by: Johannes Rausch --- .../torch/puzzletron/tiny_qwen_campaign.py | 329 +++++++++ tests/gpu/torch/puzzletron/test_puzzletron.py | 665 ++++++++++-------- 2 files changed, 698 insertions(+), 296 deletions(-) create mode 100644 tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py diff --git a/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py b/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py new file mode 100644 index 00000000000..bf41d716161 --- /dev/null +++ b/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py @@ -0,0 +1,329 @@ +# 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. + +"""Reusable hermetic tiny-Qwen campaign for Puzzletron integration tests.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import yaml +from _test_utils.torch.transformers_models import create_tiny_qwen3_5_dir +from datasets import Dataset, DatasetDict + +from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path +from puzzletron_orchestrator.compiler import ( + compile_campaign_plan, + load_execution_config, + load_runner_config, +) +from puzzletron_setup.v2.wizard import _DEFAULT_DATA_SOURCE, _DEFAULT_MODEL_SOURCE, run_wizard_v2 + +__all__ = ["TinyQwenCampaign", "build_tiny_qwen_campaign"] + +if TYPE_CHECKING: + from collections.abc import Sequence + from pathlib import Path + + from puzzletron_orchestrator.schema import CampaignPlan + from puzzletron_setup.v2.prompts import PromptChoice + + +class _DefaultsBackend: + """Select resolved guided defaults while supplying the campaign directory.""" + + def __init__(self, campaign_dir: Path) -> None: + self.campaign_dir = campaign_dir + + def text(self, message: str, default: str) -> Any: + if message == "Campaign directory:": + return str(self.campaign_dir) + return default + + def select( + self, + message: str, + choices: Sequence[PromptChoice], + default: Any, + ) -> Any: + if message == "Model:": + return _DEFAULT_MODEL_SOURCE + if message == "Dataset:": + return _DEFAULT_DATA_SOURCE + if default is not None: + return default + return next(choice.value for choice in choices if choice.disabled is None) + + def checkbox( + self, + message: str, + choices: Sequence[PromptChoice], + defaults: Sequence[Any], + ) -> Any: + del message, choices + return list(defaults) + + +@dataclass(frozen=True) +class TinyQwenCampaign: + """Generated campaign bundle plus its exact execution contract.""" + + project_root: Path + smoke_bundle: Path + smoke_root: Path + flow_id: str + overrides: tuple[str, ...] + environment: dict[str, str] + config: dict[str, Any] + compiled_plan: CampaignPlan + + def run(self, *, timeout: int = 2100) -> subprocess.CompletedProcess[str]: + """Run or resume the full campaign through the public local orchestrator.""" + + command = [ + sys.executable, + str(self.project_root / "examples/puzzletron/orchestrate.py"), + "--experiment", + str(self.smoke_bundle / "experiment.yaml"), + "--runner", + str(self.smoke_bundle / "runner.yaml"), + "--execution", + str(self.smoke_bundle / "execution.yaml"), + "--stage", + "full", + "--local", + "--poll-interval", + "0.05", + "--color", + "never", + ] + for override in self.overrides: + command.extend(("--override", override)) + return subprocess.run( + command, + cwd=self.project_root, + env=self.environment, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + + def require_success(self, completed: subprocess.CompletedProcess[str]) -> dict[str, Any]: + """Return the controller result or raise with the most useful task log.""" + + if completed.returncode != 0: + logs = sorted( + self.smoke_root.glob("logs/**/*.log"), + key=lambda path: path.stat().st_mtime_ns, + ) + log_tail = ( + logs[-1].read_text(errors="replace")[-12000:] if logs else "no task log found" + ) + raise AssertionError( + "Tiny Qwen Puzzletron campaign failed.\n" + f"stdout tail:\n{completed.stdout[-12000:]}\n" + f"stderr tail:\n{completed.stderr[-12000:]}\n" + f"latest task-log tail:\n{log_tail}" + ) + try: + payload = json.loads(completed.stdout) + except json.JSONDecodeError as error: + raise AssertionError( + "Puzzletron orchestrator did not emit its JSON result.\n" + f"stdout tail:\n{completed.stdout[-12000:]}\n" + f"stderr tail:\n{completed.stderr[-12000:]}" + ) from error + if not isinstance(payload, dict): + raise AssertionError(f"unexpected Puzzletron result payload: {payload!r}") + return payload + + +def _save_messages_dataset(path: Path) -> None: + response = ( + "Compression removes redundant parameters while preserving useful model behavior. " * 16 + ).strip() + messages = [ + {"role": "user", "content": "What is model compression?"}, + {"role": "assistant", "content": response}, + ] + rows = [{"messages": messages}] * 8 + DatasetDict( + { + "train": Dataset.from_list(rows), + "validation": Dataset.from_list(rows), + } + ).save_to_disk(str(path)) + + +def _post_mip_overrides(flow_id: str) -> tuple[str, ...]: + prefix = f"post_mip.flows.{flow_id}.nodes" + return ( + "tokenize_data.workers=1", + "+replacement_scoring.automodel.lm_head_backend=streaming", + f"{prefix}.online_eval.config.eval_samples=2", + f"{prefix}.best_lm.top_k=3", + f"{prefix}.serving.config.request_count=4", + f"{prefix}.fastest.top_k=2", + f"{prefix}.short_kd.config.max_steps=2", + f"{prefix}.short_kd.config.global_batch_size=1", + f"{prefix}.short_kd.config.local_batch_size=1", + f"+{prefix}.short_kd.config.checkpoint_every_steps=2", + f"{prefix}.final_eval.config.eval_samples=2", + ) + + +def build_tiny_qwen_campaign( + project_root: Path, + tmp_path: Path, +) -> TinyQwenCampaign: + """Generate the sole tiny-Qwen setup-to-resume Puzzletron E2E fixture.""" + + model_dir = create_tiny_qwen3_5_dir( + tmp_path / "model", + with_tokenizer=True, + hidden_size=512, + intermediate_size=768, + max_position_embeddings=128, + num_hidden_layers=2, + layer_types=["full_attention"] * 2, + ) + dataset_dir = tmp_path / "dataset" + campaign_dir = tmp_path / "campaign" + result_root = tmp_path / "results" + cache_dir = tmp_path / "cache" + defaults_path = tmp_path / "defaults.yaml" + _save_messages_dataset(dataset_dir) + defaults_path.write_text( + yaml.safe_dump( + { + "schema_version": 1, + "model": { + "source": str(model_dir), + "trust_remote_code": False, + "force_hf": False, + }, + "data": { + "source": str(dataset_dir), + "modality": "text", + "layout": "fixed", + "sequence_length": 32, + }, + "pruning": { + "depth_remove": 0, + "depth_importance_samples": 2, + "width_importance_samples": 2, + "replacement_samples": 2, + "sort_sanity": False, + "width_sanity": False, + "slicing_sanity": False, + "replacement_granularity": "block", + "axes": { + "hidden_width": {"values": [256]}, + "kv_groups": {"values": [2]}, + "q_heads_per_group": {"values": [2]}, + "ffn_intermediate": {"values": [768, 512, 256]}, + "gdn_key_groups": {"values": [2]}, + "gdn_value_heads_per_group": {"values": [2]}, + "gdn_key_head_dim": {"values": [8]}, + "gdn_value_head_dim": {"values": [8]}, + }, + "bypass": {"enabled": False}, + }, + "vllm": { + "enabled": False, + "prefill_seq_len": 32, + "generation_seq_len": 8, + "batch_size": 1, + "max_num_seqs": 1, + }, + "mip": { + "goal_metric": "params", + "goal_value": "90%", + "num_solutions": 3, + }, + "stages": { + "width_importance": {"batch": 1}, + "replacement_scoring": {"batch": 1, "instances": 1}, + }, + "output": {"result_root": str(result_root)}, + "infrastructure": { + "gpus_per_node": 1, + "execution_contract": { + "repository": str(project_root), + "venv": sys.prefix, + "container": None, + "container_mounts": None, + "prerun_commands": [], + "postrun_commands": [], + }, + }, + }, + sort_keys=False, + ) + ) + + generated = run_wizard_v2( + resume=None, + defaults_path=defaults_path, + backend=_DefaultsBackend(campaign_dir), + ) + smoke_bundle = generated / "smoke" + experiment = yaml.safe_load((smoke_bundle / "experiment.yaml").read_text()) + flows = dict((experiment.get("post_mip") or {}).get("flows") or {}) + if len(flows) != 1: + raise AssertionError(f"expected one recommended post-MIP flow, found {sorted(flows)}") + flow_id = next(iter(flows)) + overrides = _post_mip_overrides(flow_id) + config = pipeline_config_from_path(smoke_bundle / "experiment.yaml", overrides=overrides) + compiled_plan = compile_campaign_plan( + experiment_config_path=smoke_bundle / "experiment.yaml", + runner=load_runner_config(smoke_bundle / "runner.yaml"), + execution=load_execution_config(smoke_bundle / "execution.yaml"), + overrides=overrides, + stage_filter="full", + ) + environment = os.environ.copy() + environment.update( + { + "CUDA_VISIBLE_DEVICES": os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0], + "HF_DATASETS_OFFLINE": "1", + "HF_HOME": str(cache_dir / "huggingface"), + "HF_HUB_OFFLINE": "1", + "HF_DATASETS_CACHE": str(cache_dir / "datasets"), + "AIPERF_TOKENIZER_ALIAS_DIR": str(cache_dir / "aiperf-tokenizers"), + "TOKENIZERS_PARALLELISM": "false", + "TORCH_HOME": str(cache_dir / "torch"), + "TRANSFORMERS_OFFLINE": "1", + "VLLM_CACHE_ROOT": str(cache_dir / "vllm"), + "WANDB_DISABLED": "true", + "XDG_CACHE_HOME": str(cache_dir / "xdg"), + } + ) + return TinyQwenCampaign( + project_root=project_root, + smoke_bundle=smoke_bundle, + smoke_root=result_root / "smoke", + flow_id=flow_id, + overrides=overrides, + environment=environment, + config=config, + compiled_plan=compiled_plan, + ) diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index c0506bcff61..839e945d27a 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -18,95 +18,32 @@ from __future__ import annotations import json -import os -import subprocess -import sys +import math from hashlib import sha256 +from itertools import pairwise from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import Any import pytest import torch -import yaml -from _test_utils.torch.transformers_models import create_tiny_qwen3_5_dir -from datasets import Dataset, DatasetDict +from _test_utils.torch.puzzletron.tiny_qwen_campaign import ( + TinyQwenCampaign, + build_tiny_qwen_campaign, +) from transformers import AutoModelForCausalLM from modelopt.torch.puzzletron.orchestration.adapters.stage_compat import stage_is_complete -from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path -from modelopt.torch.puzzletron.stages.graph import enabled_stage_ids +from modelopt.torch.puzzletron.orchestration.state import CampaignStateStore +from modelopt.torch.puzzletron.post_mip.records import ArtifactKind, CandidateLedger, CandidateSet from puzzletron_orchestrator.adapters.registry import adapter_for_stage -from puzzletron_orchestrator.compiler import ( - compile_campaign_plan, - load_execution_config, - load_runner_config, -) -from puzzletron_setup.v2.wizard import _DEFAULT_DATA_SOURCE, _DEFAULT_MODEL_SOURCE, run_wizard_v2 - -if TYPE_CHECKING: - from collections.abc import Sequence - - from puzzletron_setup.v2.prompts import PromptChoice - - -class _DefaultsBackend: - """Select resolved defaults while supplying the test campaign directory.""" - - def __init__(self, campaign_dir: Path) -> None: - self.campaign_dir = campaign_dir - - def text(self, message: str, default: str) -> Any: - if message == "Campaign directory:": - return str(self.campaign_dir) - return default - - def select( - self, - message: str, - choices: Sequence[PromptChoice], - default: Any, - ) -> Any: - if message == "Model:": - return _DEFAULT_MODEL_SOURCE - if message == "Dataset:": - return _DEFAULT_DATA_SOURCE - if message == "Post Mip:": - return "customize" - if message.startswith("Post-MIP flow for "): - return "none" - if default is not None: - return default - return next(choice.value for choice in choices if choice.disabled is None) - - def checkbox( - self, - message: str, - choices: Sequence[PromptChoice], - defaults: Sequence[Any], - ) -> Any: - del message, choices - return list(defaults) - - -def _save_messages_dataset(path: Path) -> None: - response = ( - "Compression removes redundant parameters while preserving useful model behavior. " * 16 - ).strip() - messages = [ - {"role": "user", "content": "What is model compression?"}, - {"role": "assistant", "content": response}, - ] - rows = [{"messages": messages}] * 8 - DatasetDict( - { - "train": Dataset.from_list(rows), - "validation": Dataset.from_list(rows), - } - ).save_to_disk(str(path)) -def _artifact_digests(paths: Sequence[Path]) -> dict[str, str]: - return {str(path): sha256(path.read_bytes()).hexdigest() for path in paths} +def _json(path: Path) -> Any: + return json.loads(path.read_text()) + + +def _artifact_digests(paths: list[Path]) -> dict[str, str]: + return {str(path): sha256(path.read_bytes()).hexdigest() for path in sorted(set(paths))} def _nested_values(value: Any, key: str) -> list[Any]: @@ -132,186 +69,52 @@ def _tensor_values(value: Any): yield from _tensor_values(item) -def _run_campaign( - project_root_path: Path, - smoke_bundle: Path, - environment: dict[str, str], -) -> subprocess.CompletedProcess[str]: - return subprocess.run( - [ - sys.executable, - str(project_root_path / "examples/puzzletron/orchestrate.py"), - "--experiment", - str(smoke_bundle / "experiment.yaml"), - "--runner", - str(smoke_bundle / "runner.yaml"), - "--execution", - str(smoke_bundle / "execution.yaml"), - "--stage", - "full", - "--local", - "--poll-interval", - "0.05", - "--color", - "never", - "--override", - "tokenize_data.workers=1", - "--override", - "+replacement_scoring.automodel.lm_head_backend=streaming", - ], - cwd=project_root_path, - env=environment, - capture_output=True, - text=True, - timeout=900, - check=False, - ) - - -def _assert_campaign_succeeded( - completed: subprocess.CompletedProcess[str], - smoke_root: Path, -) -> None: - if completed.returncode == 0: - return - logs = sorted( - smoke_root.glob("logs/**/*.log"), - key=lambda path: path.stat().st_mtime_ns, - ) - log_tail = logs[-1].read_text(errors="replace")[-12000:] if logs else "no task log found" - pytest.fail( - "Tiny Qwen Puzzletron campaign failed.\n" - f"stdout tail:\n{completed.stdout[-12000:]}\n" - f"stderr tail:\n{completed.stderr[-12000:]}\n" - f"latest task-log tail:\n{log_tail}" - ) - - -@pytest.mark.timeout(1200) -def test_tiny_qwen_campaign_uses_current_public_route( - project_root_path: Path, - tmp_path: Path, -) -> None: - """Run every required stage through the public orchestrator; requires one CUDA GPU.""" - - assert torch.cuda.is_available(), "Puzzletron GPU CI requires CUDA" - assert torch.cuda.device_count() == 1, "Puzzletron GPU CI requires one visible GPU" - assert torch.equal(torch.arange(4, device="cuda").cpu(), torch.arange(4)) - model_dir = create_tiny_qwen3_5_dir( - tmp_path / "model", - with_tokenizer=True, - hidden_size=512, - intermediate_size=768, - max_position_embeddings=128, - num_hidden_layers=2, - layer_types=["full_attention"] * 2, - ) - dataset_dir = tmp_path / "dataset" - campaign_dir = tmp_path / "campaign" - result_root = tmp_path / "results" - cache_dir = tmp_path / "cache" - defaults_path = tmp_path / "defaults.yaml" - _save_messages_dataset(dataset_dir) - defaults_path.write_text( - yaml.safe_dump( - { - "schema_version": 1, - "model": { - "source": str(model_dir), - "trust_remote_code": False, - "force_hf": False, - }, - "data": { - "source": str(dataset_dir), - "modality": "text", - "layout": "fixed", - "sequence_length": 32, - }, - "pruning": { - "depth_remove": 0, - "depth_importance_samples": 2, - "width_importance_samples": 2, - "replacement_samples": 2, - "sort_sanity": False, - "width_sanity": False, - "slicing_sanity": False, - "replacement_granularity": "block", - "axes": { - "hidden_width": {"values": [256]}, - "kv_groups": {"values": [2]}, - "q_heads_per_group": {"values": [2]}, - "ffn_intermediate": {"values": [768, 512, 256]}, - "gdn_key_groups": {"values": [2]}, - "gdn_value_heads_per_group": {"values": [2]}, - "gdn_key_head_dim": {"values": [8]}, - "gdn_value_head_dim": {"values": [8]}, - }, - "bypass": {"enabled": False}, - }, - "vllm": { - "enabled": False, - "prefill_seq_len": 32, - "generation_seq_len": 8, - "batch_size": 1, - "max_num_seqs": 1, - }, - "mip": { - "goal_metric": "params", - "goal_value": "90%", - "num_solutions": 1, - }, - "stages": { - "width_importance": {"batch": 1}, - "replacement_scoring": {"batch": 1, "instances": 1}, - }, - "output": {"result_root": str(result_root)}, - "infrastructure": { - "gpus_per_node": 1, - "execution_contract": { - "repository": str(project_root_path), - "venv": sys.prefix, - "container": None, - "container_mounts": None, - "prerun_commands": [], - "postrun_commands": [], - }, - }, - }, - sort_keys=False, - ) +def _assert_compiled_route(campaign: TinyQwenCampaign) -> None: + for name in ("experiment.yaml", "runner.yaml", "execution.yaml"): + assert (campaign.smoke_bundle / name).is_file() + assert campaign.config["model"]["descriptor_override"] == "qwen3_5_text" + assert campaign.config["embedding_pruning"]["widths"] == [256] + + expected_nodes = ( + "online_eval", + "best_lm", + "materialized", + "serving", + "fastest", + "short_kd", + "final_eval", + "best", ) - - generated = run_wizard_v2( - resume=None, - defaults_path=defaults_path, - backend=_DefaultsBackend(campaign_dir), + prefix = f"post.{campaign.flow_id}." + post_nodes = tuple( + node for node in campaign.compiled_plan.stages if node.stage_id.startswith(prefix) ) - smoke_bundle = generated / "smoke" - for name in ("experiment.yaml", "runner.yaml", "execution.yaml"): - assert (smoke_bundle / name).is_file() + assert tuple(node.stage_id.removeprefix(prefix) for node in post_nodes) == expected_nodes + assert post_nodes[0].parents == ("mip",) + for parent, node in pairwise(post_nodes): + assert node.parents == (parent.stage_id,) + + flow = campaign.config["post_mip"]["flows"][campaign.flow_id]["nodes"] + assert tuple(flow) == expected_nodes + assert flow["online_eval"]["config"]["eval_samples"] == 2 + assert flow["best_lm"]["top_k"] == 3 + assert flow["serving"]["config"]["input_tokens"] == 32 + assert flow["serving"]["config"]["output_tokens"] == 8 + assert flow["serving"]["config"]["request_count"] == 4 + assert flow["fastest"]["top_k"] == 2 + assert flow["short_kd"]["config"]["max_steps"] == 2 + assert flow["short_kd"]["config"]["global_batch_size"] == 1 + assert flow["short_kd"]["config"]["local_batch_size"] == 1 + assert flow["short_kd"]["config"]["checkpoint_every_steps"] == 2 + assert flow["final_eval"]["config"]["eval_samples"] == 2 + assert flow["best"]["top_k"] == 1 - overrides = [ - "tokenize_data.workers=1", - "+replacement_scoring.automodel.lm_head_backend=streaming", - ] - config = pipeline_config_from_path(smoke_bundle / "experiment.yaml", overrides=overrides) - assert config["embedding_pruning"]["widths"] == [256] - post_mip_flows = config["post_mip"]["flows"] - assert len(post_mip_flows) == 1 - serving_config = next(iter(post_mip_flows.values()))["nodes"]["serving"]["config"] - assert serving_config["input_tokens"] == 32 - assert serving_config["output_tokens"] == 8 - compiled_plan = compile_campaign_plan( - experiment_config_path=smoke_bundle / "experiment.yaml", - runner=load_runner_config(smoke_bundle / "runner.yaml"), - execution=load_execution_config(smoke_bundle / "execution.yaml"), - overrides=overrides, - stage_filter="full", - ) replacement_node = next( - node for node in compiled_plan.stages if node.stage_id == "replacement_scoring" + node for node in campaign.compiled_plan.stages if node.stage_id == "replacement_scoring" + ) + replacement_work = adapter_for_stage(replacement_node).plan( + campaign.compiled_plan, replacement_node ) - replacement_work = adapter_for_stage(replacement_node).plan(compiled_plan, replacement_node) assert replacement_node.instances == 1 assert replacement_node.gpus_per_instance == 1 assert replacement_node.total_gpus == 1 @@ -319,32 +122,11 @@ def test_tiny_qwen_campaign_uses_current_public_route( assert len(replacement_work.items) == 1 assert replacement_work.items[0].metadata["worker_count"] == 1 - environment = os.environ.copy() - environment.update( - { - "CUDA_VISIBLE_DEVICES": os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0], - "HF_DATASETS_OFFLINE": "1", - "HF_HOME": str(cache_dir / "huggingface"), - "HF_HUB_OFFLINE": "1", - "HF_DATASETS_CACHE": str(cache_dir / "datasets"), - "TORCH_HOME": str(cache_dir / "torch"), - "TRANSFORMERS_OFFLINE": "1", - "XDG_CACHE_HOME": str(cache_dir / "xdg"), - } - ) - smoke_root = result_root / "smoke" - completed = _run_campaign(project_root_path, smoke_bundle, environment) - _assert_campaign_succeeded(completed, smoke_root) - - stages = enabled_stage_ids(config) - assert config["model"]["descriptor_override"] == "qwen3_5_text" - assert all(stage_is_complete(config, stage) for stage in stages) - - manifests = [smoke_root / "manifests" / f"{stage}.json" for stage in stages] - assert all(json.loads(path.read_text())["status"] == "success" for path in manifests) +def _assert_pruning_and_mip_artifacts(campaign: TinyQwenCampaign) -> list[Path]: + root = campaign.smoke_root pass_manifests = list( - smoke_root.glob("pruning/pruning_scores/automodel/*/activation_passes_manifest.json") + root.glob("pruning/pruning_scores/automodel/*/activation_passes_manifest.json") ) assert len(pass_manifests) == 1 score_files = list(pass_manifests[0].parent.glob("*/rank_*.pth")) @@ -357,14 +139,12 @@ def test_tiny_qwen_campaign_uses_current_public_route( assert score_tensors assert all(tensor.numel() and torch.isfinite(tensor).all() for tensor in score_tensors) - replacement_summary = json.loads( - (smoke_root / "artifacts/replacement_scoring/summary.json").read_text() - ) + replacement_summary = _json(root / "artifacts/replacement_scoring/summary.json") assert replacement_summary["widths"] == [256] assert replacement_summary["scenario_count"] == 1 replacement_results = [ - json.loads(path.read_text()) - for path in smoke_root.glob( + _json(path) + for path in root.glob( "scenarios/width-*/depth-*/distributed_eval/replacement_scoring/results/**/*.json" ) ] @@ -375,17 +155,13 @@ def test_tiny_qwen_campaign_uses_current_public_route( for result in replacement_results ) - for checkpoint in (smoke_root / "ckpts/teacher", smoke_root / "ckpts/sorted_teacher"): - assert (checkpoint / "config.json").is_file() - assert list(checkpoint.glob("*.safetensors")) - candidate_library = json.loads((smoke_root / "candidate_library.json").read_text()) + candidate_library = _json(root / "candidate_library.json") assert 256 in _nested_values(candidate_library, "intermediate_size") assert 512 in _nested_values(candidate_library, "intermediate_size") - - active_profiles = json.loads((smoke_root / "mip/active_profiles.json").read_text()) + active_profiles = _json(root / "mip/active_profiles.json") assert active_profiles["status"] == "success" grids = [ - json.loads((smoke_root / "mip" / "profiles" / profile_id / "mip_grid.json").read_text()) + _json(root / "mip" / "profiles" / profile_id / "mip_grid.json") for profile_id in active_profiles["profile_ids"] ] assert all(grid["status"] == "success" for grid in grids) @@ -395,17 +171,298 @@ def test_tiny_qwen_campaign_uses_current_public_route( for scenario in grid["scenarios"] if scenario["status"] == "feasible" ] - solutions = [solution for path in solution_paths for solution in json.loads(path.read_text())] - assert solutions + solutions = [solution for path in solution_paths for solution in _json(path)] + assert len(solutions) >= 3 solution_ffn_widths = { int(width) for solution in solutions for width in _nested_values(solution["chosen_block_configs"], "intermediate_size") } assert solution_ffn_widths.intersection({256, 512}) + return [*pass_manifests, *score_files, *solution_paths] + + +def _assert_node_publication( + campaign: TinyQwenCampaign, + ledger: CandidateLedger, + node_id: str, + predecessor_identity: str | None, +) -> tuple[dict[str, Any], str]: + node_root = campaign.smoke_root / "artifacts/post_mip/nodes" / node_id + summary = _json(node_root / "summary.json") + current = _json(node_root / "current.json") + index = _json(node_root / "index.json") + execution_identity = summary["execution_identity"] + execution_root = node_root / "executions" / execution_identity + candidate_path = Path(summary["candidate_set_path"]) + observations_path = Path(summary["observations_path"]) + candidate_set = ledger.load_candidate_set(node_id) + candidate_payload = _json(candidate_path) + + assert summary["status"] == "success" + assert summary["stage_id"] == f"post.{campaign.flow_id}.{node_id}" + assert summary["node_id"] == node_id + assert current["execution_identity"] == execution_identity + assert index["current"] == execution_identity + assert execution_identity in index["executions"] + assert candidate_path == execution_root / "candidate_set.json" + assert observations_path == execution_root / "observations.json" + assert _json(execution_root / "summary.json") == summary + assert candidate_path.is_file() and observations_path.is_file() + assert candidate_set.flow_id == campaign.flow_id + assert candidate_set.node_id == node_id + assert candidate_set.producer_execution_identity == execution_identity + assert candidate_set == CandidateSet.create( + campaign.flow_id, + node_id, + candidate_set.revision_ids, + producer_execution_identity=execution_identity, + ) + assert candidate_payload == { + "flow_id": candidate_set.flow_id, + "identity": candidate_set.identity, + "node_id": candidate_set.node_id, + "producer_execution_identity": candidate_set.producer_execution_identity, + "revision_ids": list(candidate_set.revision_ids), + } + assert summary["output_count"] == len(candidate_set.revision_ids) + assert summary["input_count"] == len(_json(observations_path)) + if predecessor_identity is not None: + assert summary["execution_contract"]["candidate_set"] == predecessor_identity + return summary, candidate_set.identity + + +def _finite_metric(ledger: CandidateLedger, revision_id: str, reference: str) -> float: + value = ledger.resolve_metric(revision_id, reference) + assert value is not None and math.isfinite(value) + return value + + +def _assert_post_mip_and_final_checkpoint( + campaign: TinyQwenCampaign, +) -> tuple[Path, list[Path]]: + root = campaign.smoke_root + ledger = CandidateLedger(root / "artifacts/post_mip") + node_ids = ( + "online_eval", + "best_lm", + "materialized", + "serving", + "fastest", + "short_kd", + "final_eval", + "best", + ) + predecessor_identity = None + summaries = {} + for node_id in node_ids: + summaries[node_id], predecessor_identity = _assert_node_publication( + campaign, ledger, node_id, predecessor_identity + ) + + online = ledger.load_candidate_set("online_eval") + assert len(online.revision_ids) >= 3 + online_observations = ledger.observations["online_eval"] + assert all(row.status == "success" for row in online_observations.values()) + assert all(Path(row.artifacts["result_path"]).is_file() for row in online_observations.values()) + online_losses = { + revision_id: _finite_metric(ledger, revision_id, "online_eval.lm_loss") + for revision_id in online.revision_ids + } + + best_lm = ledger.load_candidate_set("best_lm") + assert len(best_lm.revision_ids) == 3 + assert ( + best_lm.revision_ids + == tuple( + revision_id + for _loss, revision_id in sorted((loss, rid) for rid, loss in online_losses.items()) + )[:3] + ) + assert { + revision_id + for revision_id, row in ledger.observations["best_lm"].items() + if row.status == "selected" + } == set(best_lm.revision_ids) + + materialized = ledger.load_candidate_set("materialized") + assert len(materialized.revision_ids) == 3 + materialized_summary = summaries["materialized"] + materialized_root = ( + root + / "artifacts/post_mip/nodes/materialized/executions" + / materialized_summary["execution_identity"] + ) + for revision_id in materialized.revision_ids: + revision = ledger.revisions[revision_id] + assert revision.producer_node == "materialized" + assert revision.parent_revision_id in set(best_lm.revision_ids) + assert revision.artifact_kind is ArtifactKind.CHECKPOINT + checkpoint = Path(revision.artifact["checkpoint"]) + assert checkpoint == materialized_root / "checkpoints" / revision.architecture_id + assert (checkpoint / "config.json").is_file() + assert list(checkpoint.glob("*.safetensors")) + assert { + ledger.revisions[revision_id].parent_revision_id + for revision_id in materialized.revision_ids + } == set(best_lm.revision_ids) + + serving = ledger.load_candidate_set("serving") + assert len(serving.revision_ids) >= 2 + assert set(serving.revision_ids) <= set(materialized.revision_ids) + serving_throughputs = { + revision_id: _finite_metric( + ledger, revision_id, "serving.concurrency_1.output_token_throughput" + ) + for revision_id in serving.revision_ids + } + assert all(throughput > 0 for throughput in serving_throughputs.values()) + serving_observations = ledger.observations["serving"] + assert set(serving_observations) == set(materialized.revision_ids) + assert { + revision_id for revision_id, row in serving_observations.items() if row.status == "success" + } == set(serving.revision_ids) + assert all( + row.status in {"failed", "timed_out"} and row.error and row.output_revision_id is None + for row in serving_observations.values() + if row.status != "success" + ) + aiperf_paths = [ + Path(path) + for row in serving_observations.values() + if row.status == "success" + for artifacts in row.artifacts["result_paths"] + for path in artifacts.values() + ] + assert aiperf_paths and all(path.is_file() for path in aiperf_paths) + + fastest = ledger.load_candidate_set("fastest") + assert len(fastest.revision_ids) == 2 + assert fastest.revision_ids == tuple( + revision_id + for _throughput, revision_id in sorted( + (throughput, revision_id) for revision_id, throughput in serving_throughputs.items() + )[-2:][::-1] + ) + short_kd = ledger.load_candidate_set("short_kd") + assert len(short_kd.revision_ids) == 2 + short_kd_summary = summaries["short_kd"] + short_kd_root = ( + root + / "artifacts/post_mip/nodes/short_kd/executions" + / short_kd_summary["execution_identity"] + ) + kd_paths = [] + for revision_id in short_kd.revision_ids: + revision = ledger.revisions[revision_id] + assert revision.producer_node == "short_kd" + assert revision.parent_revision_id in set(materialized.revision_ids) + assert revision.artifact_kind is ArtifactKind.CHECKPOINT + checkpoint = Path(revision.artifact["checkpoint"]) + summary_path = Path(revision.artifact["summary_path"]) + architecture_root = short_kd_root / "checkpoints" / revision.architecture_id + assert summary_path == architecture_root / "global_distillation_summary.json" + assert checkpoint.is_relative_to(architecture_root) + kd_summary = _json(summary_path) + records = kd_summary["records"] + assert kd_summary["max_steps"] == 2 + assert len({int(record.get("step", record.get("global_step"))) for record in records}) >= 2 + assert all( + math.isfinite(float(record.get("loss", record.get("train_loss")))) for record in records + ) + assert Path(kd_summary["post_kd_checkpoint"]) == checkpoint + assert (checkpoint / "config.json").is_file() + assert list(checkpoint.glob("*.safetensors")) + assert (checkpoint.parents[1] / "saving_completed").is_file() + kd_paths.extend( + [ + summary_path, + checkpoint.parents[2] / "training.jsonl", + checkpoint / "config.json", + *checkpoint.glob("*.safetensors"), + checkpoint.parents[1] / "saving_completed", + ] + ) + assert { + ledger.revisions[revision_id].parent_revision_id for revision_id in short_kd.revision_ids + } == set(fastest.revision_ids) + + final_eval = ledger.load_candidate_set("final_eval") + assert set(final_eval.revision_ids) == set(short_kd.revision_ids) + final_losses = { + revision_id: _finite_metric(ledger, revision_id, "final_eval.lm_loss") + for revision_id in final_eval.revision_ids + } + assert all(row.status == "success" for row in ledger.observations["final_eval"].values()) + best = ledger.load_candidate_set("best") + assert len(best.revision_ids) == 1 + selected_id = best.revision_ids[0] + assert final_losses[selected_id] == min(final_losses.values()) + selected = ledger.revisions[selected_id] + final_checkpoint = Path(selected.artifact["checkpoint"]) + assert selected.producer_node == "short_kd" + assert selected.artifact_kind is ArtifactKind.CHECKPOINT + assert summaries["best"]["checkpoints"] == [str(final_checkpoint)] + + post_paths = [root / "artifacts/post_mip/candidate_registry.json"] + post_paths.extend(root.glob("artifacts/post_mip/nodes/*/current.json")) + post_paths.extend(root.glob("artifacts/post_mip/nodes/*/index.json")) + post_paths.extend(root.glob("artifacts/post_mip/nodes/*/summary.json")) + post_paths.extend(root.glob("artifacts/post_mip/nodes/*/executions/*/candidate_set.json")) + post_paths.extend(root.glob("artifacts/post_mip/nodes/*/executions/*/observations.json")) + post_paths.extend(root.glob("artifacts/post_mip/nodes/*/executions/*/summary.json")) + return final_checkpoint, [*post_paths, *aiperf_paths, *kd_paths] + + +def _assert_final_report(campaign: TinyQwenCampaign, result: dict[str, Any]) -> None: + report_root = campaign.smoke_root / "artifacts/campaign_report" + html_path = report_root / "campaign_report.html" + manifest_path = report_root / "report_manifest.json" + assert result["report_status"] == "completed" + assert Path(result["report_path"]) == html_path + assert Path(result["report_manifest_path"]) == manifest_path + manifest = _json(manifest_path) + assert manifest["schema_version"] == 1 + assert manifest["verification"] == "passed" + assert manifest["campaign_identity"] + html = html_path.read_text() + for node_id in ("online_eval", "serving", "short_kd", "final_eval"): + section_id = "post-" + "-".join( + part.replace("_", "-") for part in (campaign.flow_id, node_id) + ) + assert f'id="{section_id}"' in html + + +@pytest.mark.timeout(2400) +def test_tiny_qwen_campaign_uses_current_public_route( + project_root_path: Path, + tmp_path: Path, +) -> None: + """Run the full route in the pinned CUDA 12.9 image with one visible GPU.""" + assert torch.cuda.is_available(), "Puzzletron GPU CI requires CUDA" + assert torch.cuda.device_count() == 1, "Puzzletron GPU CI requires one visible GPU" + assert torch.equal(torch.arange(4, device="cuda").cpu(), torch.arange(4)) + campaign = build_tiny_qwen_campaign(project_root_path, tmp_path) + _assert_compiled_route(campaign) + + completed = campaign.run() + result = campaign.require_success(completed) + stage_ids = tuple(node.stage_id for node in campaign.compiled_plan.stages) + assert tuple(result["completed"]) == stage_ids + assert result["failed_stages"] == [] + assert not result["halted"] + assert not result["cancelled"] + assert not result["detached"] + assert all(stage_is_complete(campaign.config, stage_id) for stage_id in stage_ids) + + manifest_paths = sorted(campaign.smoke_root.glob("manifests/*.json")) + assert manifest_paths + assert all(_json(path)["status"] == "success" for path in manifest_paths) + pruning_paths = _assert_pruning_and_mip_artifacts(campaign) + final_checkpoint, post_paths = _assert_post_mip_and_final_checkpoint(campaign) model = AutoModelForCausalLM.from_pretrained( - smoke_root / "ckpts/sorted_teacher", + final_checkpoint, dtype=torch.bfloat16, local_files_only=True, ).cuda() @@ -415,10 +472,26 @@ def test_tiny_qwen_campaign_uses_current_public_route( use_cache=False, ).logits assert torch.isfinite(logits).all() + _assert_final_report(campaign, result) - durable_paths = [*manifests, *pass_manifests] - durable_paths.extend(smoke_root.glob("mip/profiles/*/mip_grid.json")) + state = CampaignStateStore(campaign.smoke_root) + attempts_before = { + (attempt["work_id"], attempt["attempt_id"]): attempt for attempt in state.list_attempts() + } + durable_paths = [*manifest_paths, *pruning_paths, *post_paths] before_resume = _artifact_digests(durable_paths) - resumed = _run_campaign(project_root_path, smoke_bundle, environment) - _assert_campaign_succeeded(resumed, smoke_root) + resumed = campaign.run() + resumed_result = campaign.require_success(resumed) + assert tuple(resumed_result["completed"]) == stage_ids + assert resumed_result["failed_stages"] == [] + assert not resumed_result["halted"] + assert not resumed_result["cancelled"] + assert not resumed_result["detached"] + assert resumed_result["report_status"] == "completed" + attempts_after = { + (attempt["work_id"], attempt["attempt_id"]): attempt for attempt in state.list_attempts() + } + assert attempts_after == attempts_before assert _artifact_digests(durable_paths) == before_resume + resumed_checkpoint, _ = _assert_post_mip_and_final_checkpoint(campaign) + assert resumed_checkpoint == final_checkpoint From 26e0fa438de3c8ca82966c25492b0e7f19188afc Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 04:36:33 +0200 Subject: [PATCH 10/24] Fix composite stage manifest validation Avoid fingerprinting the canonical stage pointer as one of its own immutable outputs, and keep test manifests aligned with the execution-record schema. Signed-off-by: Johannes Rausch --- examples/puzzletron/main.py | 1 - tests/unit/torch/puzzletron/conftest.py | 8 +- .../torch/puzzletron/test_example_runner.py | 152 +++++++++++++++++- 3 files changed, 157 insertions(+), 4 deletions(-) diff --git a/examples/puzzletron/main.py b/examples/puzzletron/main.py index c8cc04df390..35c08aa599f 100644 --- a/examples/puzzletron/main.py +++ b/examples/puzzletron/main.py @@ -454,7 +454,6 @@ def _run_worker(args: argparse.Namespace) -> None: stage=args.worker_stage, gpus_per_node=gpus_per_node, ) - outputs["base_manifest"] = str(result.manifest_path) result = _complete_composite_stage(cfg, args.worker_stage, outputs) if int(os.environ.get("RANK", "0")) == 0: if result.status == "failed": diff --git a/tests/unit/torch/puzzletron/conftest.py b/tests/unit/torch/puzzletron/conftest.py index b4d6298797a..2fe8bb38437 100644 --- a/tests/unit/torch/puzzletron/conftest.py +++ b/tests/unit/torch/puzzletron/conftest.py @@ -21,7 +21,7 @@ import pytest -from puzzletron_orchestrator.identity import stable_hash +from puzzletron_orchestrator.identity import canonicalize, stable_hash from puzzletron_orchestrator.stages import semantic_stage_config @@ -36,7 +36,9 @@ def write( config: dict[str, object], **extra: object, ) -> None: - semantic_config = semantic_stage_config(config, stage) + authored_config = canonicalize(config) + config_identity = stable_hash(authored_config, prefix=f"{stage}_cfg") + semantic_config = semantic_stage_config(authored_config, stage) semantic_config_identity = stable_hash(semantic_config, prefix=f"{stage}_semantic_cfg") capability_snapshot = extra.get("capability_snapshot") semantic_identity = stable_hash( @@ -54,6 +56,8 @@ def write( { "stage": stage, "status": "success", + "config": authored_config, + "config_identity": config_identity, "semantic_config": semantic_config, "semantic_config_identity": semantic_config_identity, "semantic_identity": semantic_identity, diff --git a/tests/unit/torch/puzzletron/test_example_runner.py b/tests/unit/torch/puzzletron/test_example_runner.py index 7e5c4b17022..818c40bf361 100644 --- a/tests/unit/torch/puzzletron/test_example_runner.py +++ b/tests/unit/torch/puzzletron/test_example_runner.py @@ -1,4 +1,43 @@ -from examples.puzzletron.main import build_worker_command +# 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 public Puzzletron stage runner.""" + +from __future__ import annotations + +import json +from copy import deepcopy +from types import SimpleNamespace +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path + +from examples.puzzletron import main as puzzletron_main +from examples.puzzletron.main import ( + _complete_composite_stage, + _validate_worker_result, + build_worker_command, +) +from modelopt.torch.puzzletron.manifest import ( + stage_manifest_from_config, + validate_stage_execution_record, + write_stage_manifest, +) +from modelopt.torch.puzzletron.orchestration.adapters.stage_compat import stage_is_complete +from modelopt.torch.puzzletron.stage_runner import StageResult def test_worker_command_propagates_gpu_count_to_composite_followups(): @@ -10,3 +49,114 @@ def test_worker_command_propagates_gpu_count_to_composite_followups(): ) assert command[command.index("--gpus-per-node") + 1] == "1" + + +def test_build_library_worker_does_not_forward_mutable_base_manifest( + tmp_path: Path, monkeypatch +) -> None: + config = { + "puzzle_dir": str(tmp_path), + "embedding_pruning": {"enabled": True, "widths": [256]}, + "build_library": {"enabled": True}, + "execution": {"gpus_per_node": 1}, + } + initial = StageResult( + stage="build_library", + status="success", + manifest_path=tmp_path / "manifests" / "build_library.json", + message="initial root build", + ) + captured_outputs = {} + + monkeypatch.setattr( + puzzletron_main.mtpz.pipeline_config, + "pipeline_config_from_path", + lambda *_args, **_kwargs: deepcopy(config), + ) + monkeypatch.setattr( + puzzletron_main.mtpz.stage_runner, + "run_stage", + lambda *_args, **_kwargs: initial, + ) + monkeypatch.setattr( + puzzletron_main, + "_run_embedding_stage", + lambda **_kwargs: { + "stage": "build_library", + "widths": [256], + "scenarios_root": str(tmp_path / "scenarios"), + }, + ) + + def complete_composite(_config, _stage, outputs): + captured_outputs.update(outputs) + return initial + + monkeypatch.setattr(puzzletron_main, "_complete_composite_stage", complete_composite) + monkeypatch.setattr(puzzletron_main, "_validate_worker_result", lambda *_args, **_kwargs: None) + monkeypatch.setattr(puzzletron_main, "refresh_campaign_report", lambda *_args, **_kwargs: None) + monkeypatch.setattr(puzzletron_main.mtpz.tools, "mprint", lambda *_args, **_kwargs: None) + + puzzletron_main._run_worker( + SimpleNamespace( + config=tmp_path / "experiment.yaml", + override=[], + worker_stage="build_library", + scenario_child=False, + gpus_per_node=1, + ) + ) + + assert captured_outputs == { + "stage": "build_library", + "widths": [256], + "scenarios_root": str(tmp_path / "scenarios"), + } + assert "base_manifest" not in captured_outputs + + +def test_build_library_composite_preserves_authored_and_effective_config(tmp_path: Path) -> None: + authored_config = { + "puzzle_dir": str(tmp_path), + "experiment": {"dir": str(tmp_path)}, + "model": {"source": "example/model"}, + "build_library": {"enabled": True}, + "embedding_pruning": {"enabled": False}, + "vllm_stats": {"subblock_stats_filename": "subblock_stats.json"}, + } + worker_config = deepcopy(authored_config) + worker_config["build_library"]["include_noops"] = False + worker_config["_runtime"] = { + "config_path": str(tmp_path / "experiment.yaml"), + "authored_config": deepcopy(authored_config), + } + + outputs = {} + for name in ("replacement_library.json", "candidate_library.json", "subblock_stats.json"): + path = tmp_path / name + path.write_text("{}\n") + outputs[name.removesuffix(".json")] = str(path) + + manifest_path = tmp_path / "manifests" / "build_library.json" + initial = stage_manifest_from_config("build_library", worker_config) + initial.complete(outputs=outputs) + write_stage_manifest(manifest_path, initial) + initial_pointer = json.loads(manifest_path.read_text()) + + result = _complete_composite_stage( + worker_config, + "build_library", + {"stage": "build_library", "widths": [], "scenarios_root": str(tmp_path / "scenarios")}, + ) + + pointer = json.loads(manifest_path.read_text()) + resolved = json.loads( + (tmp_path / pointer["execution_record"]["resolved_config_path"]).read_text() + ) + assert pointer["config"] == authored_config + assert pointer["semantic_config"] == initial_pointer["semantic_config"] + assert resolved["resolved_stage_config"]["build_library"]["include_noops"] is False + validate_stage_execution_record(manifest_path, expected_stage="build_library") + _validate_worker_result(worker_config, result, expected_stage="build_library") + assert stage_is_complete(authored_config, "build_library") + assert stage_is_complete(worker_config, "build_library") From 79145c696683a0ebe84017a625b62b2111ca6ab6 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 05:00:05 +0200 Subject: [PATCH 11/24] Forward overrides to post-MIP aggregation Keep shard workers and aggregation on the same campaign configuration so they derive one execution identity and consume the same results. Signed-off-by: Johannes Rausch --- .../orchestration/adapters/post_mip.py | 21 +++++---- .../torch/puzzletron/test_post_mip_adapter.py | 46 +++++++++++++++++++ 2 files changed, 58 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index 24fc77bc0a9..f67b8e204aa 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -223,16 +223,19 @@ def aggregate( ) -> PublishedOutput | None: repo = Path(plan.runner.contract.repository) script = repo / "examples" / "puzzletron" / "run_post_mip_node.py" + argv = [ + "python", + str(script), + "--config", + plan.experiment_config_path, + "--stage-id", + node.stage_id, + "--aggregate", + ] + for override in plan.overrides: + argv.extend(["--override", override]) result = subprocess.run( - ( - "python", - str(script), - "--config", - plan.experiment_config_path, - "--stage-id", - node.stage_id, - "--aggregate", - ), + argv, cwd=repo, capture_output=True, text=True, diff --git a/tests/unit/torch/puzzletron/test_post_mip_adapter.py b/tests/unit/torch/puzzletron/test_post_mip_adapter.py index cffd6bf639e..9d0244c5489 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_adapter.py +++ b/tests/unit/torch/puzzletron/test_post_mip_adapter.py @@ -15,6 +15,9 @@ """Tests for post-MIP orchestration adapter launch policy.""" +import json +import subprocess +from dataclasses import replace from pathlib import Path from puzzletron_orchestrator.adapters.post_mip import PostMIPAdapter @@ -27,6 +30,7 @@ StagePlanNode, TaskLauncher, WorkItem, + WorkPlan, ) @@ -139,3 +143,45 @@ def test_post_mip_filter_keeps_direct_launcher(tmp_path: Path): ) assert attempt.task_topology.launcher is TaskLauncher.DIRECT + + +def test_post_mip_aggregation_forwards_campaign_overrides(tmp_path: Path, monkeypatch): + plan, node = _plan(tmp_path, stage_id="post.params.online_eval", node_type="evaluation") + plan = replace( + plan, + overrides=( + "post_mip.flows.params.nodes.online_eval.config.eval_samples=2", + "+post_mip.flows.params.nodes.short_kd.config.checkpoint_every_steps=2", + ), + ) + commands = [] + + def run(command, **_kwargs): + commands.append(tuple(command)) + return subprocess.CompletedProcess(command, 0, stdout=json.dumps({"status": "success"})) + + monkeypatch.setattr("puzzletron_orchestrator.adapters.post_mip.subprocess.run", run) + + publication = PostMIPAdapter().aggregate( + plan=plan, + node=node, + work_plan=WorkPlan(stage_id=node.stage_id, strategy=node.strategy, items=()), + ) + + assert commands == [ + ( + "python", + str(tmp_path / "examples" / "puzzletron" / "run_post_mip_node.py"), + "--config", + plan.experiment_config_path, + "--stage-id", + node.stage_id, + "--aggregate", + "--override", + plan.overrides[0], + "--override", + plan.overrides[1], + ) + ] + assert publication is not None + assert publication.summary == {"status": "success"} From 5c9c435490434374bd87fc7c8b6c4b0fa17c921c Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 05:29:12 +0200 Subject: [PATCH 12/24] Preserve heterogeneous KD checkpoint metadata Refresh AnyModel interchange metadata before publishing a consolidated checkpoint, and reload selected heterogeneous artifacts through the descriptor-aware ModelOpt path. Signed-off-by: Johannes Rausch --- .../distillation/global_kd_recipe.py | 8 ++++ tests/gpu/torch/puzzletron/test_puzzletron.py | 40 ++++++++++++++++--- .../puzzletron/test_global_kd_canonical.py | 17 +++++++- 3 files changed, 58 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py index c874ff0e332..7f59d6c4453 100644 --- a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py +++ b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py @@ -13,6 +13,7 @@ import dataclasses import hashlib +import json import os import types from collections import deque @@ -635,6 +636,13 @@ def save_checkpoint( if self.dist_env.is_main: from pathlib import Path + consolidated = Path(checkpoint_path, "model", "consolidated") + config_path = consolidated / "config.json" + config = json.loads(config_path.read_text()) if config_path.is_file() else {} + if config.get("block_configs"): + from ..utils.vllm_adapter import refresh_realized_checkpoint_config + + refresh_realized_checkpoint_config(consolidated) Path(checkpoint_path, "saving_completed").touch() if torch.distributed.is_initialized(): torch.distributed.barrier() diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index 839e945d27a..9f57d0d8796 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -30,10 +30,10 @@ TinyQwenCampaign, build_tiny_qwen_campaign, ) -from transformers import AutoModelForCausalLM - +from modelopt.torch.puzzletron.anymodel.registry import resolve_descriptor_from_pretrained from modelopt.torch.puzzletron.orchestration.adapters.stage_compat import stage_is_complete from modelopt.torch.puzzletron.orchestration.state import CampaignStateStore +from modelopt.torch.puzzletron.plugins.automodel.load import load_anymodel_for_scoring from modelopt.torch.puzzletron.post_mip.records import ArtifactKind, CandidateLedger, CandidateSet from puzzletron_orchestrator.adapters.registry import adapter_for_stage @@ -374,6 +374,12 @@ def _assert_post_mip_and_final_checkpoint( assert (checkpoint / "config.json").is_file() assert list(checkpoint.glob("*.safetensors")) assert (checkpoint.parents[1] / "saving_completed").is_file() + parent = ledger.revisions[revision.parent_revision_id] + block_configs = _json(checkpoint / "config.json")["block_configs"] + assert block_configs + assert block_configs == _json(Path(parent.artifact["checkpoint"]) / "config.json")[ + "block_configs" + ] kd_paths.extend( [ summary_path, @@ -461,11 +467,35 @@ def test_tiny_qwen_campaign_uses_current_public_route( assert all(_json(path)["status"] == "success" for path in manifest_paths) pruning_paths = _assert_pruning_and_mip_artifacts(campaign) final_checkpoint, post_paths = _assert_post_mip_and_final_checkpoint(campaign) - model = AutoModelForCausalLM.from_pretrained( - final_checkpoint, - dtype=torch.bfloat16, + resolution = resolve_descriptor_from_pretrained(str(final_checkpoint)) + selected_config = _json(final_checkpoint / "config.json") + assert selected_config["architectures"] == ["AnyModel"] + assert selected_config["base_architecture"] == "Qwen3_5ForCausalLM" + selected_block_configs = selected_config["block_configs"] + selected_ffn_widths = [] + for block_config in selected_block_configs: + widths = _nested_values(block_config, "intermediate_size") + assert len(widths) == 1 + selected_ffn_widths.append(int(widths[0])) + per_layer_config = (selected_config.get("text_config") or selected_config)[ + "per_layer_config" + ] + assert [ + int( + per_layer_config.get(str(index), {}).get( + "intermediate_size", selected_config["intermediate_size"] + ) + ) + for index in range(len(selected_block_configs)) + ] == selected_ffn_widths + model = load_anymodel_for_scoring( + str(final_checkpoint), + anymodel_descriptor=resolution.name, + force_hf=True, + torch_dtype=torch.bfloat16, local_files_only=True, ).cuda() + assert [layer.mlp.down_proj.in_features for layer in model.model.layers] == selected_ffn_widths with torch.no_grad(): logits = model( torch.tensor([[1, 2, 3, 4]], device="cuda"), diff --git a/tests/unit/torch/puzzletron/test_global_kd_canonical.py b/tests/unit/torch/puzzletron/test_global_kd_canonical.py index 7c9e0c12f76..1d7cec1b33d 100644 --- a/tests/unit/torch/puzzletron/test_global_kd_canonical.py +++ b/tests/unit/torch/puzzletron/test_global_kd_canonical.py @@ -740,10 +740,11 @@ def __init__(self): assert all(value.item() > 0 for value in recipe._gradient_squared.values()) -def test_global_kd_checkpoint_forwards_best_metric_key(tmp_path): +def test_global_kd_checkpoint_forwards_best_metric_key(tmp_path, monkeypatch): from modelopt.torch.puzzletron.distillation.global_kd_recipe import _WeightedObjectiveMixin calls = [] + refreshes = [] class BaseRecipe: def save_checkpoint( @@ -755,7 +756,12 @@ def save_checkpoint( best_metric_key="default", ): calls.append((epoch, step, train_loss, val_loss, best_metric_key)) - (tmp_path / f"epoch_{epoch}_step_{step}").mkdir() + checkpoint = tmp_path / f"epoch_{epoch}_step_{step}" + consolidated = checkpoint / "model/consolidated" + consolidated.mkdir(parents=True) + (consolidated / "config.json").write_text( + json.dumps({"block_configs": [{"subblock_configs": []}]}) + ) return "saved" class Recipe(_WeightedObjectiveMixin, BaseRecipe): @@ -768,6 +774,12 @@ class Recipe(_WeightedObjectiveMixin, BaseRecipe): {"config": type("Config", (), {"checkpoint_dir": tmp_path})()}, )() recipe.dist_env = type("DistEnv", (), {"is_main": True})() + monkeypatch.setattr( + "modelopt.torch.puzzletron.utils.vllm_adapter.refresh_realized_checkpoint_config", + lambda path: refreshes.append( + (path, (tmp_path / "epoch_2_step_17/saving_completed").exists()) + ), + ) result = recipe.save_checkpoint( 2, @@ -779,6 +791,7 @@ class Recipe(_WeightedObjectiveMixin, BaseRecipe): assert result == "saved" assert calls == [(2, 17, 0.5, {"lm_loss": 0.25}, "lm_loss")] + assert refreshes == [(tmp_path / "epoch_2_step_17/model/consolidated", False)] assert (tmp_path / "epoch_2_step_17" / "saving_completed").is_file() From baceb70192a13587d878c64435c1f30cda3f9340 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 06:19:02 +0200 Subject: [PATCH 13/24] Format Puzzletron GPU test imports Signed-off-by: Johannes Rausch --- tests/gpu/torch/puzzletron/test_puzzletron.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index 9f57d0d8796..d017c6c8104 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -30,6 +30,7 @@ TinyQwenCampaign, build_tiny_qwen_campaign, ) + from modelopt.torch.puzzletron.anymodel.registry import resolve_descriptor_from_pretrained from modelopt.torch.puzzletron.orchestration.adapters.stage_compat import stage_is_complete from modelopt.torch.puzzletron.orchestration.state import CampaignStateStore From 2c43fbc79692cb73bfb875d6469c8c2f64ef5b57 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 15:12:35 +0200 Subject: [PATCH 14/24] Harden Puzzletron GPU baseline recovery Bind post-MIP attempts to the producer contract and preserve explicit AIPerf security policies so resume resubmits stale work without weakening offline execution. Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 2 +- .../finalize_replacement_scoring.py | 9 + .../puzzletron/run_profile_aiperf_worker.py | 15 + .../torch/puzzletron/benchmarks/aiperf.py | 104 +++- .../distillation/global_kd_recipe.py | 125 +++-- .../distributed_eval/automodel_executor.py | 1 + .../puzzletron/orchestration/adapters/base.py | 30 +- .../orchestration/adapters/post_mip.py | 69 ++- .../orchestration/adapters/sharded.py | 8 +- .../orchestration/adapters/stage_compat.py | 151 +----- .../puzzletron/orchestration/controller.py | 106 ++-- .../torch/puzzletron/post_mip/identity.py | 222 +++++++++ modelopt/torch/puzzletron/post_mip/runner.py | 74 +-- modelopt/torch/puzzletron/stages/future.py | 11 + .../torch/puzzletron/utils/vllm_adapter.py | 29 +- noxfile.py | 4 +- .../torch/puzzletron/tiny_qwen_campaign.py | 28 +- .../test_aiperf_context_capacity.py | 65 ++- .../torch/puzzletron/test_ci_environment.py | 49 +- .../puzzletron/test_global_kd_canonical.py | 134 ++--- .../test_orchestration_executors.py | 47 ++ .../test_orchestration_lightweight.py | 4 - .../test_orchestration_shutdown_progress.py | 56 +++ .../test_orchestration_task_topology.py | 1 + .../test_post_mip_execution_identity.py | 462 ++++++++++++++++++ .../torch/puzzletron/test_post_mip_runner.py | 5 +- .../puzzletron/test_vllm_axis_contract.py | 37 +- .../torch/puzzletron/test_width_scenarios.py | 21 +- 28 files changed, 1389 insertions(+), 480 deletions(-) create mode 100644 modelopt/torch/puzzletron/post_mip/identity.py create mode 100644 tests/unit/torch/puzzletron/test_post_mip_execution_identity.py diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index bdcb5175491..611f9a272bd 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -269,7 +269,7 @@ git -C "${AUTOMODEL_ROOT}" rev-parse HEAD ``` ```bash -python - <<'PY' +PYTHONPATH="${MODEL_OPT_ROOT}" python - <<'PY' import importlib.metadata as metadata import json import os diff --git a/examples/puzzletron/finalize_replacement_scoring.py b/examples/puzzletron/finalize_replacement_scoring.py index d4e62e595ad..71e08142af9 100644 --- a/examples/puzzletron/finalize_replacement_scoring.py +++ b/examples/puzzletron/finalize_replacement_scoring.py @@ -32,12 +32,21 @@ from modelopt.torch.puzzletron.manifest import stage_manifest_from_config, write_stage_manifest from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path +__all__ = [ + "finalization_marker_is_current", + "finalize_replacement_scoring", + "main", + "write_finalization_marker", +] + def _successful_manifest_identity(manifest_path: str | Path) -> str | None: try: manifest = json.loads(Path(manifest_path).read_text()) except (OSError, ValueError): return None + if not isinstance(manifest, dict): + return None if manifest.get("stage") != "replacement_scoring" or manifest.get("status") != "success": return None identity = manifest.get("semantic_identity") diff --git a/examples/puzzletron/run_profile_aiperf_worker.py b/examples/puzzletron/run_profile_aiperf_worker.py index 8fa381ca1ec..e2cc145ef66 100644 --- a/examples/puzzletron/run_profile_aiperf_worker.py +++ b/examples/puzzletron/run_profile_aiperf_worker.py @@ -188,6 +188,8 @@ def run_worker( concurrencies: tuple[int, ...] | None = None, request_count: int | None = None, benchmark_timeout: float = 7200, + trust_remote_code: bool = False, + allow_aiperf_v011_online_tokenizer_resolution: bool = False, ) -> Path: # Worker execution needs the GPU stack; result merging intentionally remains # usable by the dependency-light login-node orchestrator. @@ -241,6 +243,10 @@ def run_worker( seed=42, gpu_telemetry="pynvml", benchmark_timeout=benchmark_timeout, + trust_remote_code=trust_remote_code, + allow_aiperf_v011_online_tokenizer_resolution=( + allow_aiperf_v011_online_tokenizer_resolution + ), ) rows.extend(result.model_dump(mode="json") for result in results) selection_suffix = "" @@ -335,6 +341,11 @@ def main() -> None: parser.add_argument("--concurrency", type=int, action="append", default=[]) parser.add_argument("--request-count", type=int) parser.add_argument("--benchmark-timeout", type=float, default=7200) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument( + "--allow-aiperf-v011-online-tokenizer-resolution", + action="store_true", + ) parser.add_argument("--preflight", action="store_true") parser.add_argument("--merge", action="store_true") args = parser.parse_args() @@ -370,6 +381,10 @@ def main() -> None: concurrencies=tuple(args.concurrency) or None, request_count=args.request_count, benchmark_timeout=args.benchmark_timeout, + trust_remote_code=args.trust_remote_code, + allow_aiperf_v011_online_tokenizer_resolution=( + args.allow_aiperf_v011_online_tokenizer_resolution + ), ) print(output) diff --git a/modelopt/torch/puzzletron/benchmarks/aiperf.py b/modelopt/torch/puzzletron/benchmarks/aiperf.py index 561f9c2e018..0ea6fca9d0b 100644 --- a/modelopt/torch/puzzletron/benchmarks/aiperf.py +++ b/modelopt/torch/puzzletron/benchmarks/aiperf.py @@ -46,7 +46,11 @@ _CHECKPOINT_PREPARE_LOCK = Lock() -def _prepare_vllm_checkpoint(checkpoint_dir: Path) -> bool: +def _prepare_vllm_checkpoint( + checkpoint_dir: Path, + *, + trust_remote_code: bool = False, +) -> bool: """Restore AnyModel metadata lost by generic HF checkpoint consolidation.""" with _CHECKPOINT_PREPARE_LOCK: config = json.loads((checkpoint_dir / "config.json").read_text()) @@ -57,7 +61,10 @@ def _prepare_vllm_checkpoint(checkpoint_dir: Path) -> bool: return False from ..utils.vllm_adapter import refresh_realized_checkpoint_config - refresh_realized_checkpoint_config(checkpoint_dir) + refresh_realized_checkpoint_config( + checkpoint_dir, + trust_remote_code=trust_remote_code, + ) return True @@ -309,6 +316,39 @@ def _profile_command( return command +def _vllm_server_command( + *, + checkpoint_dir: Path, + port: int, + model_name: str, + input_tokens: int, + output_tokens: int, + topology: dict[str, Any], + trust_remote_code: bool, +) -> list[str]: + """Build the vLLM command under the caller's explicit code-trust policy.""" + + command = [ + "vllm", + "serve", + str(checkpoint_dir), + "--host", + "127.0.0.1", + "--port", + str(port), + "--served-model-name", + model_name, + "--max-model-len", + str(_server_max_model_len(input_tokens, output_tokens, topology)), + ] + if trust_remote_code: + command.append("--trust-remote-code") + command.extend(_topology_vllm_args(topology)) + command.extend(_descriptor_vllm_args(checkpoint_dir)) + command.extend(str(arg) for arg in topology.get("extra_vllm_args", ())) + return command + + def _clean_subprocess_environment( gpu_ids: str, *, architecture_id: str, topology_id: str ) -> dict[str, str]: @@ -355,11 +395,17 @@ def _clean_subprocess_environment( return env -def _aiperf_subprocess_environment(env: dict[str, str]) -> dict[str, str]: - """Avoid AIPerf's offline resolver for an explicit local tokenizer path.""" +def _aiperf_subprocess_environment( + env: dict[str, str], + *, + allow_aiperf_v011_online_tokenizer_resolution: bool = False, +) -> dict[str, str]: + """Optionally relax AIPerf v0.11's offline local-tokenizer resolution.""" + resolved = dict(env) - resolved.pop("HF_HUB_OFFLINE", None) - resolved.pop("TRANSFORMERS_OFFLINE", None) + if allow_aiperf_v011_online_tokenizer_resolution: + resolved.pop("HF_HUB_OFFLINE", None) + resolved.pop("TRANSFORMERS_OFFLINE", None) return resolved @@ -384,13 +430,18 @@ def run_aiperf_sweep( readiness_timeout: float = 1200, benchmark_timeout: float = 600, gpu_telemetry: str | None = "pynvml", + trust_remote_code: bool = False, + allow_aiperf_v011_online_tokenizer_resolution: bool = False, ) -> list[BenchmarkResult]: - """Run multiple concurrencies against one persistent vLLM server.""" + """Run a serving sweep while preserving offline and remote-code policy by default.""" checkpoint_dir = Path(checkpoint_dir).resolve() artifact_dir = Path(artifact_dir).resolve() artifact_dir.mkdir(parents=True, exist_ok=True) - _prepare_vllm_checkpoint(checkpoint_dir) + _prepare_vllm_checkpoint( + checkpoint_dir, + trust_remote_code=trust_remote_code, + ) concurrency_values = tuple(int(value) for value in concurrencies) if not concurrency_values or len(set(concurrency_values)) != len(concurrency_values): raise ValueError("AIPerf concurrencies must be non-empty and unique") @@ -413,23 +464,15 @@ def run_aiperf_sweep( model_name = f"puzzletron-{architecture_id[:16]}" tokenizer_dir = _short_tokenizer_alias(checkpoint_dir, artifact_dir) server_log = artifact_dir / "vllm_server.log" - server_cmd = [ - "vllm", - "serve", - str(checkpoint_dir), - "--host", - "127.0.0.1", - "--port", - str(port), - "--served-model-name", - model_name, - "--max-model-len", - str(_server_max_model_len(input_tokens, output_tokens, topology)), - "--trust-remote-code", - ] - server_cmd.extend(_topology_vllm_args(topology)) - server_cmd.extend(_descriptor_vllm_args(checkpoint_dir)) - server_cmd.extend(str(arg) for arg in topology.get("extra_vllm_args", ())) + server_cmd = _vllm_server_command( + checkpoint_dir=checkpoint_dir, + port=port, + model_name=model_name, + input_tokens=input_tokens, + output_tokens=output_tokens, + topology=topology, + trust_remote_code=trust_remote_code, + ) executable = _resolve_executable(executable) env = _clean_subprocess_environment( gpu_ids, @@ -438,7 +481,12 @@ def run_aiperf_sweep( ) for key, value in (topology.get("env") or {}).items(): env[str(key)] = str(value) - aiperf_env = _aiperf_subprocess_environment(env) + aiperf_env = _aiperf_subprocess_environment( + env, + allow_aiperf_v011_online_tokenizer_resolution=( + allow_aiperf_v011_online_tokenizer_resolution + ), + ) cached: dict[int, BenchmarkResult] = {} missing: list[tuple[int, Path, list[str], str]] = [] for concurrency in concurrency_values: @@ -472,6 +520,10 @@ def run_aiperf_sweep( "endpoint_type": endpoint_type, "extra_inputs": _exact_length_extra_inputs(extra_inputs, output_tokens), "use_server_token_count": use_server_token_count, + "trust_remote_code": trust_remote_code, + "allow_aiperf_v011_online_tokenizer_resolution": ( + allow_aiperf_v011_online_tokenizer_resolution + ), "revisions": revisions, }, prefix="aiperf_result", diff --git a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py index 7f59d6c4453..5d76e630dbf 100644 --- a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py +++ b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py @@ -110,9 +110,7 @@ def _global_kd_checkpoint_adapter_context(model_parts, descriptor_name: str | No block_configs = _config_value(text_config, "block_configs") if not block_configs: continue - active_descriptor = descriptor_name or _config_value( - config, "anymodel_descriptor" - ) + active_descriptor = descriptor_name or _config_value(config, "anymodel_descriptor") if not active_descriptor: continue descriptor = AutoModelDescriptorFactory.get(str(active_descriptor)) @@ -160,13 +158,17 @@ def non_strict(*args, _original=original, _name=name, options=None, **kwargs): and len(args) >= 2 ): model, optimizer = args[:2] - model_parameters = {id(parameter): fqn for fqn, parameter in model.named_parameters()} + model_parameters = { + id(parameter): fqn for fqn, parameter in model.named_parameters() + } unmatched = [] for group_index, group in enumerate(optimizer.param_groups): for parameter_index, parameter in enumerate(group["params"]): if id(parameter) in model_parameters: continue - local = parameter.to_local() if isinstance(parameter, DTensor) else parameter + local = ( + parameter.to_local() if isinstance(parameter, DTensor) else parameter + ) unmatched.append( { "group": group_index, @@ -175,7 +177,9 @@ def non_strict(*args, _original=original, _name=name, options=None, **kwargs): "local_shape": tuple(local.shape), "requires_grad": bool(parameter.requires_grad), "has_grad": parameter.grad is not None, - "state": sorted(str(key) for key in optimizer.state.get(parameter, {})), + "state": sorted( + str(key) for key in optimizer.state.get(parameter, {}) + ), } ) if unmatched: @@ -313,9 +317,7 @@ def _project_teacher_hidden_on_reference_mesh(hidden, teacher_head, reference_lo return projection(hidden) if projection is not None else teacher_head(hidden) local_hidden = hidden.to_local() if isinstance(hidden, DTensor) else hidden reference_local = ( - reference_logits.to_local() - if isinstance(reference_logits, DTensor) - else reference_logits + reference_logits.to_local() if isinstance(reference_logits, DTensor) else reference_logits ) expected_vocab = int(reference_local.shape[-1]) @@ -551,9 +553,7 @@ def _set_teacher_mtp_enabled(teacher): def _split_output(output, model): seq_idx = None - main_is_hidden = bool( - getattr(model, "_puzzletron_distillation_hidden_output", False) - ) + main_is_hidden = bool(getattr(model, "_puzzletron_distillation_hidden_output", False)) if isinstance(output, tuple): values = list(output) if values and isinstance(values[-1], torch.Tensor) and values[-1].dtype == torch.int32: @@ -596,8 +596,7 @@ def _configure_objective(self): self._objective_step_cursor = {name: 0 for name in self.objective} self._loss_topology_logged = False self._gradient_squared = { - name: torch.tensor(0.0) - for name in ("vision", "projector", "language", "mtp") + name: torch.tensor(0.0) for name in ("vision", "projector", "language", "mtp") } self._gradient_hook_handles = [] self._vision_monitors = [] @@ -627,11 +626,9 @@ def save_checkpoint( val_loss, best_metric_key=best_metric_key, ) - checkpoint_path = ( - os.path.join( - str(self.checkpointer.config.checkpoint_dir), - f"epoch_{epoch}_step_{step}", - ) + checkpoint_path = os.path.join( + str(self.checkpointer.config.checkpoint_dir), + f"epoch_{epoch}_step_{step}", ) if self.dist_env.is_main: from pathlib import Path @@ -642,7 +639,11 @@ def save_checkpoint( if config.get("block_configs"): from ..utils.vllm_adapter import refresh_realized_checkpoint_config - refresh_realized_checkpoint_config(consolidated) + model_config = _config_value(getattr(self, "cfg", None), "model") + refresh_realized_checkpoint_config( + consolidated, + trust_remote_code=bool(_config_value(model_config, "trust_remote_code")), + ) Path(checkpoint_path, "saving_completed").touch() if torch.distributed.is_initialized(): torch.distributed.barrier() @@ -710,16 +711,10 @@ def observability_metadata(self): for role in sorted(roles) }, "vision_output_checksums": sorted( - checksum - for item in gathered - for checksum in item["vision_output_checksums"] + checksum for item in gathered for checksum in item["vision_output_checksums"] ), "media_input_checksums": sorted( - set( - checksum - for item in gathered - for checksum in item["media_input_checksums"] - ) + set(checksum for item in gathered for checksum in item["media_input_checksums"]) ), } @@ -773,7 +768,9 @@ def _remove_text_inactive_optimizer_parameters(self) -> None: if not inactive_ids: return - optimizers = self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + optimizers = ( + self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + ) removed = 0 for optimizer in optimizers: for group in optimizer.param_groups: @@ -826,11 +823,11 @@ def observe_optimizer_step(optimizer, _args, _kwargs): continue gradient = parameter.grad value = gradient.to_local() if isinstance(gradient, DTensor) else gradient - self._gradient_squared[group].add_( - value.detach().float().square().sum() - ) + self._gradient_squared[group].add_(value.detach().float().square().sum()) - optimizers = self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + optimizers = ( + self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + ) for optimizer in optimizers: self._gradient_hook_handles.append( optimizer.register_step_pre_hook(observe_optimizer_step) @@ -860,12 +857,12 @@ def _rebind_optimizer_to_current_model_parameters(self, model=None) -> None: seen.add(id(parameter)) current_parameters.append(parameter) - optimizers = self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + optimizers = ( + self.optimizer if isinstance(self.optimizer, (list, tuple)) else [self.optimizer] + ) for optimizer in optimizers: optimizer_parameters = [ - parameter - for group in optimizer.param_groups - for parameter in group["params"] + parameter for group in optimizer.param_groups for parameter in group["params"] ] current_ids = {id(parameter) for parameter in current_parameters} if all(id(parameter) in current_ids for parameter in optimizer_parameters): @@ -1002,9 +999,7 @@ def _teacher_loss_model(self): if teacher_pp is None: return self.teacher_model return next( - part - for part, stage in zip(teacher_pp.parts, teacher_pp.info.stages) - if stage.is_last + part for part, stage in zip(teacher_pp.parts, teacher_pp.info.stages) if stage.is_last ) @staticmethod @@ -1167,9 +1162,7 @@ def _mtp_objective_losses( student_head = _get_lm_head_module(student_model) if student_is_hidden else None teacher_head = ( - _get_lm_head_module(teacher_model) - if needs_mtp_kd and teacher_is_hidden - else None + _get_lm_head_module(teacher_model) if needs_mtp_kd and teacher_is_hidden else None ) if student_is_hidden and student_head is None: raise ValueError("MTP losses require an accessible student lm_head") @@ -1195,9 +1188,7 @@ def _mtp_objective_losses( depth_labels = torch.where(rolled == seq_idx, depth_labels, -100) flat_student = self._flatten_tokens(student_value) - flat_teacher = ( - self._flatten_tokens(teacher_values[depth]) if needs_mtp_kd else None - ) + flat_teacher = self._flatten_tokens(teacher_values[depth]) if needs_mtp_kd else None flat_labels = depth_labels.reshape(-1) for start in range(0, flat_student.shape[0], chunk_size): stop = min(start + chunk_size, flat_student.shape[0]) @@ -1281,7 +1272,10 @@ def _objective_loss(self, student_out, teacher_out, labels, model, num_label_tok rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 if rank == 0: placements = ( - tuple(type(item).__name__ + f"({getattr(item, 'dim', '')})" for item in student_logits.placements) + tuple( + type(item).__name__ + f"({getattr(item, 'dim', '')})" + for item in student_logits.placements + ) if isinstance(student_logits, DTensor) else ("replicated_tensor",) ) @@ -1395,9 +1389,7 @@ def loss_wrapper(student_out, target, **_kwargs): if teacher_out is None: raise RuntimeError("Teacher PP output queue is empty") model = next( - part - for part, stage in zip(self.model_parts, self.pp.info.stages) - if stage.is_last + part for part, stage in zip(self.model_parts, self.pp.info.stages) if stage.is_last ) # PP schedules rescale by the optimizer-step label count after # backward, so every microbatch loss must remain an unnormalized sum. @@ -1409,9 +1401,7 @@ def loss_wrapper(student_out, target, **_kwargs): return loss_wrapper -class KnowledgeDistillationRecipeForNextTokenPrediction( - _WeightedObjectiveMixin, _AutoModelLLMKD -): +class KnowledgeDistillationRecipeForNextTokenPrediction(_WeightedObjectiveMixin, _AutoModelLLMKD): """AutoModel LLM KD with independently weighted main/MTP objectives.""" def setup(self): @@ -1484,7 +1474,9 @@ def _forward_backward_step( is_train=is_train, loss_buffer=loss_buffer, ) - batch = {key: value.to(self.dist_env.device, non_blocking=True) for key, value in batch.items()} + batch = { + key: value.to(self.dist_env.device, non_blocking=True) for key, value in batch.items() + } # Current AutoModel CP owns label sharding through the batch mapping. # Keep labels present until CP has padded/sharded every sequence tensor, # then remove the CP-local labels for the weighted objective. @@ -1504,7 +1496,9 @@ def _forward_backward_step( teacher_out = None if self.needs_teacher: with ScopedModuleOffloading(self.teacher_model, enabled=False), torch.no_grad(): - teacher_out = self.teacher_model(**filter_forward_kwargs(self.teacher_model, batch)) + teacher_out = self.teacher_model( + **filter_forward_kwargs(self.teacher_model, batch) + ) student_out = model(**filter_forward_kwargs(model, batch)) total, terms = self._objective_loss( student_out, teacher_out, labels, model, num_label_tokens @@ -1747,7 +1741,9 @@ def _forward_backward_step( with torch.no_grad(): prepared = model(_pre_embed_only=True, **media) if self.needs_teacher and "inputs_embeds" in prepared: - _validate_cp_pre_embed_teacher_compatibility(prepared["inputs_embeds"], self.teacher_model) + _validate_cp_pre_embed_teacher_compatibility( + prepared["inputs_embeds"], self.teacher_model + ) for key in VLM_INPUT_KEYS: batch.pop(key, None) batch.update(prepared) @@ -1804,9 +1800,7 @@ def prepare_cp_inputs(pp, parts, values): batch = prepare_cp_inputs(self.pp, self.model_parts, batch) if self.needs_teacher: - teacher_batch = prepare_cp_inputs( - self.teacher_pp, self.teacher_pp.parts, teacher_batch - ) + teacher_batch = prepare_cp_inputs(self.teacher_pp, self.teacher_pp.parts, teacher_batch) train_ctx, batch = make_cp_batch_and_ctx(self.device_mesh, batch) labels = batch.pop("labels") model_input_key = "inputs_embeds" if "inputs_embeds" in batch else "input_ids" @@ -1832,9 +1826,7 @@ def prepare_cp_inputs(pp, parts, values): if self.needs_teacher: teacher_ctx, teacher_batch = make_cp_batch_and_ctx(self.device_mesh, teacher_batch) teacher_labels = teacher_batch.pop("labels") - teacher_input_key = ( - "inputs_embeds" if "inputs_embeds" in teacher_batch else "input_ids" - ) + teacher_input_key = "inputs_embeds" if "inputs_embeds" in teacher_batch else "input_ids" teacher_input = teacher_batch.pop(teacher_input_key) with teacher_ctx(): teacher_targets = ( @@ -1845,8 +1837,9 @@ def prepare_cp_inputs(pp, parts, values): set_pp_vlm_chunk_specs(self.teacher_pp.info.schedule, teacher_batch) capture = self.teacher_model._teacher_logits_capture capture.clear() - with torch.no_grad(), stage_vlm_media_for_pp( - self.teacher_pp, self.teacher_pp.parts, teacher_batch + with ( + torch.no_grad(), + stage_vlm_media_for_pp(self.teacher_pp, self.teacher_pp.parts, teacher_batch), ): teacher_losses = [] if self.teacher_pp.info.has_last_stage else None if self.teacher_pp.info.has_first_stage: @@ -1888,9 +1881,7 @@ def prepare_cp_inputs(pp, parts, values): ) def _run_train_optim_step(self, batches, max_grad_norm=None): - log_data = FinetuneRecipeForVLM._run_train_optim_step( - self, batches, max_grad_norm - ) + log_data = FinetuneRecipeForVLM._run_train_optim_step(self, batches, max_grad_norm) # The shared publisher normalizes last-stage PP microbatch sums and # forwards every objective term to rank zero. diff --git a/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py b/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py index f5d8d39345a..082952f158e 100644 --- a/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py +++ b/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py @@ -270,6 +270,7 @@ def evaluate(self, request: EvaluationRequest) -> EvaluationResult | None: ) def _score(self, prune_target: dict | list[dict] | None) -> dict | None: + # Keep framework imports lazy so this executor remains dependency-light until setup. import torch import torch.distributed as torch_dist diff --git a/modelopt/torch/puzzletron/orchestration/adapters/base.py b/modelopt/torch/puzzletron/orchestration/adapters/base.py index ebf71dd4de8..bdd172b95e9 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/base.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/base.py @@ -22,7 +22,11 @@ WorkPlan, ) -__all__ = ["WorkAdapter"] +__all__ = ["ExecutionIdentityProjectionUnavailable", "WorkAdapter"] + + +class ExecutionIdentityProjectionUnavailable(RuntimeError): + """The current upstream state does not yet define an adapter identity projection.""" class WorkAdapter(ABC): @@ -71,6 +75,30 @@ def aggregate( ) -> PublishedOutput | None: return None + def execution_identity_projection( + self, + *, + plan: CampaignPlan, + node: StagePlanNode, + work_plan: WorkPlan, + ) -> Mapping[str, Any]: + """Return adapter-owned, currently resolvable execution inputs.""" + + return {} + + def prepare_execution_identity_projection( + self, + *, + plan: CampaignPlan, + node: StagePlanNode, + ) -> None: + """Prepare mutable adapter inputs immediately before a new attempt is bound. + + The default is intentionally empty. Read-only currentness checks call only + :meth:`execution_identity_projection`; adapters that require preparation + must implement it here instead of mutating from their projection method. + """ + def classify_failure( self, *, diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index f67b8e204aa..cd320740888 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -20,8 +20,7 @@ import json import subprocess from pathlib import Path - -from puzzletron_orchestrator.post_mip.records import CandidateLedger +from typing import Any from ..schema import ( AttemptSpec, @@ -39,6 +38,21 @@ from .packing import packed_allocation from .stage_compat import _hf_checkpoint_is_complete, post_mip_summary_is_current +if __package__.startswith("puzzletron_orchestrator."): + from puzzletron_orchestrator.post_mip.identity import ( + PostMIPExecutionContractUnavailable, + expected_post_mip_candidate_count, + expected_post_mip_execution_contract, + prepare_post_mip_candidate_ledger, + ) +else: + from ...post_mip.identity import ( + PostMIPExecutionContractUnavailable, + expected_post_mip_candidate_count, + expected_post_mip_execution_contract, + prepare_post_mip_candidate_ledger, + ) + __all__ = ["ManualInputRequired", "PostMIPAdapter"] @@ -75,23 +89,15 @@ def _node_root(plan: CampaignPlan, stage_id: str) -> Path: return plan.puzzle_dir / "artifacts" / "post_mip" / "nodes" / _node_id(stage_id) -def _available_evaluation_candidates(plan: CampaignPlan, stage_id: str, config: dict) -> int | None: - input_id = str(config.get("input", "source")) - ledger = CandidateLedger(plan.puzzle_dir / "artifacts" / "post_mip") - if input_id == "source": - active_mip = plan.puzzle_dir / "mip" / "active_profiles.json" - if not active_mip.is_file(): - return None - ledger.ingest_mip(plan.puzzle_dir) - _prefix, flow_id, _node_id_value = stage_id.split(".", 2) - flow = plan.experiment_config["post_mip"]["flows"][flow_id] - candidate_set = ledger.root_set(flow_id, flow["source"]) - else: - current = plan.puzzle_dir / "artifacts" / "post_mip" / "nodes" / input_id / "current.json" - if not current.is_file(): - return None - candidate_set = ledger.load_candidate_set(input_id) - return len(candidate_set.revision_ids) +def _identity_config(plan: CampaignPlan) -> dict[str, Any]: + return {**plan.experiment_config, "puzzle_dir": str(plan.puzzle_dir)} + + +def _available_evaluation_candidates(plan: CampaignPlan, stage_id: str) -> int | None: + try: + return expected_post_mip_candidate_count(_identity_config(plan), stage_id) + except PostMIPExecutionContractUnavailable: + return None def _full_node_instance_count(node: StagePlanNode, count: int) -> int: @@ -110,12 +116,35 @@ class PostMIPAdapter(WorkAdapter): strategy = ExecutionStrategy.SHARDED + def prepare_execution_identity_projection( + self, + *, + plan: CampaignPlan, + node: StagePlanNode, + ) -> None: + """Prepare the candidate registry only on the attempt-submission path.""" + + del node + prepare_post_mip_candidate_ledger(_identity_config(plan)) + + def execution_identity_projection( + self, + *, + plan: CampaignPlan, + node: StagePlanNode, + work_plan: WorkPlan, + ) -> dict[str, Any]: + """Bind scheduler attempts to the canonical producer execution contract.""" + + del work_plan + return expected_post_mip_execution_contract(_identity_config(plan), node.stage_id) + def plan(self, plan: CampaignPlan, node: StagePlanNode) -> WorkPlan: config = _node_config(plan, node.stage_id) node_type = str(config.get("type")) count = 1 if node_type in {"filter", "manual_filter"} else node.instances if node_type in {"evaluation", "downstream_evaluation"}: - available = _available_evaluation_candidates(plan, node.stage_id, config) + available = _available_evaluation_candidates(plan, node.stage_id) if available is not None: if available < 1: raise RuntimeError( diff --git a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py index cce0df7fc18..432cf055874 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py @@ -217,11 +217,11 @@ 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 {} + model = plan.experiment_config.get("model") or {} argv = [ "python", str(script_path), @@ -236,6 +236,10 @@ def command( "--output-tokens", str(aiperf.get("output_tokens", 1024)), ] + if bool(aiperf.get("trust_remote_code", model.get("trust_remote_code", False))): + argv.append("--trust-remote-code") + if bool(aiperf.get("allow_aiperf_v011_online_tokenizer_resolution", False)): + argv.append("--allow-aiperf-v011-online-tokenizer-resolution") else: argv = ["python", str(script_path), "--config", plan.experiment_config_path] argv.extend(extra_args) diff --git a/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py b/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py index cc8ef517c42..a63a324af4e 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py @@ -50,11 +50,13 @@ stage_manifest_uses_execution_record, validate_stage_execution_record, ) + from puzzletron_orchestrator.post_mip.identity import expected_post_mip_execution_identity else: from ...execution_record import ( stage_manifest_uses_execution_record, validate_stage_execution_record, ) + from ...post_mip.identity import expected_post_mip_execution_identity __all__ = [ "StageCompatAdapter", @@ -457,155 +459,18 @@ def _patterns_present(puzzle_dir: Path, patterns: tuple[str, ...]) -> bool: ) -def _prefixed_hash(prefix: str, payload: Mapping[str, Any]) -> str: - return f"{prefix}_{hash_payload(payload)[:16]}" - - -def _post_input_candidate_set( - config: Mapping[str, Any], puzzle_dir: Path, stage_id: str -) -> tuple[Mapping[str, Any], Mapping[str, Any]]: - _prefix, flow_id, node_id = stage_id.split(".", 2) - flow = config["post_mip"]["flows"][flow_id] - node = flow["nodes"][node_id] - input_id = str(node.get("input", "source")) - registry = _read_mapping(puzzle_dir / "artifacts" / "post_mip" / "candidate_registry.json") - if registry is None: - raise RuntimeError("post-MIP candidate registry is unavailable") - if input_id != "source": - current = _read_mapping( - puzzle_dir / "artifacts" / "post_mip" / "nodes" / input_id / "current.json" - ) - if current is None: - raise RuntimeError(f"post-MIP input node {input_id!r} has no current execution") - candidate_set = _read_mapping( - puzzle_dir - / "artifacts" - / "post_mip" - / "nodes" - / input_id - / "executions" - / str(current["execution_identity"]) - / "candidate_set.json" - ) - if candidate_set is None: - raise RuntimeError(f"post-MIP input node {input_id!r} has no candidate set") - identity_payload = { - key: candidate_set[key] - for key in ( - "flow_id", - "node_id", - "revision_ids", - "producer_execution_identity", - ) - } - if candidate_set.get("identity") != _prefixed_hash("candidate_set", identity_payload): - raise RuntimeError(f"post-MIP input node {input_id!r} has an invalid candidate set") - return candidate_set, registry - - active = _read_mapping(puzzle_dir / "mip" / "active_profiles.json") - if active is None or active.get("status") != "success": - raise RuntimeError("active MIP profile manifest is unavailable") - active_execution = str(active["execution_identity"]) - active_profiles = {str(value) for value in active.get("profile_ids") or ()} - if ( - registry.get("active_mip_execution_identity") != active_execution - or set(registry.get("active_profile_ids") or ()) != active_profiles - ): - raise RuntimeError("post-MIP registry does not reflect the active MIP execution") - source = flow["source"] - variants = source.get("variants", "all") - objectives = source.get("objectives", "all") - if isinstance(variants, str) and variants != "all": - variants = [variants] - if isinstance(objectives, str) and objectives != "all": - objectives = [objectives] - revision_ids = [] - for architecture in dict(registry.get("architectures") or {}).values(): - origins = [ - origin - for origin in architecture.get("origins") or () - if origin.get("profile_id") in active_profiles - and origin.get("mip_execution_identity") == active_execution - and origin.get("run_id") == source["run"] - and (variants == "all" or origin.get("variant_id") in variants) - and (objectives == "all" or (origin.get("objective") or {}).get("metric") in objectives) - ] - if origins: - origins.sort( - key=lambda origin: ( - str(origin.get("profile_id")), - str(origin.get("kind")), - int(origin.get("rank", 0)), - ) - ) - revision_ids.append(str(origins[0]["revision_id"])) - revision_ids = sorted(dict.fromkeys(revision_ids)) - payload = { - "flow_id": flow_id, - "node_id": "source", - "revision_ids": revision_ids, - "producer_execution_identity": active_execution, - } - return { - **payload, - "identity": _prefixed_hash("candidate_set", payload), - }, registry - - def post_mip_summary_is_current( config: Mapping[str, Any], puzzle_dir: Path, stage_id: str, summary: Mapping[str, Any] ) -> bool: """Validate a node summary without importing the PyTorch-backed worker package.""" try: - _prefix, flow_id, node_id = stage_id.split(".", 2) - node = dict(config["post_mip"]["flows"][flow_id]["nodes"][node_id]) - candidate_set, registry = _post_input_candidate_set(config, puzzle_dir, stage_id) - owners = set() - if node.get("type") == "filter": - if node.get("mode") in {"top_k", "threshold"}: - references = [node["metric"]] - else: - references = [entry["metric"] for entry in node.get("metrics") or ()] - owners.update( - str(reference).partition(".")[0] - for reference in references - if not str(reference).startswith("mip.") - ) - model_source = str(node.get("model_source", "latest")) - if model_source not in {"latest", "origin"}: - owners.add(model_source) - dependency_executions = {} - for owner in sorted(owners): - current = _read_mapping( - puzzle_dir / "artifacts" / "post_mip" / "nodes" / owner / "current.json" - ) - if current is None: - return False - dependency_executions[owner] = current["execution_identity"] - revision_ids = [str(value) for value in candidate_set.get("revision_ids") or ()] - revisions = dict(registry.get("revisions") or {}) - if model_source == "latest": - source_revisions = {value: value for value in revision_ids} - elif model_source == "origin": - source_revisions = {} - for value in revision_ids: - current = value - while revisions[current].get("parent_revision_id") is not None: - current = str(revisions[current]["parent_revision_id"]) - source_revisions[value] = current - else: - recorded = (summary.get("execution_contract") or {}).get("source_revisions") or {} - if set(recorded) != set(revision_ids): - return False - source_revisions = dict(recorded) - contract = { - "candidate_set": candidate_set["identity"], - "node": node, - "dependency_executions": dependency_executions, - "source_revisions": source_revisions, - } - return summary.get("execution_identity") == _prefixed_hash("post_mip_execution", contract) + effective_config = dict(config) + effective_config["puzzle_dir"] = str(puzzle_dir) + return summary.get("execution_identity") == expected_post_mip_execution_identity( + effective_config, + stage_id, + ) except (KeyError, OSError, RuntimeError, TypeError, ValueError): return False diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index 7d478d3ebff..3dae60dd441 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.py @@ -26,6 +26,7 @@ from pathlib import Path from typing import Any, Literal +from .adapters.base import ExecutionIdentityProjectionUnavailable from .adapters.post_mip import ManualInputRequired from .adapters.registry import adapter_for_stage from .adapters.stage_compat import stage_is_complete @@ -333,8 +334,11 @@ def _required_completed_attempts( node: StagePlanNode, attempts: list[dict[str, Any]], ) -> list[dict[str, Any]] | None: - work_plan = adapter_for_stage(node).plan(self.plan, node) - stage_execution_identity = self._stage_execution_identity(node, work_plan) + try: + work_plan = adapter_for_stage(node).plan(self.plan, node) + stage_execution_identity = self._stage_execution_identity(node, work_plan) + except ExecutionIdentityProjectionUnavailable: + return None completed: list[dict[str, Any]] = [] for item in work_plan.items: matches = [ @@ -369,34 +373,38 @@ def _stage_execution_identity( node: StagePlanNode, work_plan: WorkPlan | None = None, ) -> str: - work_plan = work_plan or adapter_for_stage(node).plan(self.plan, node) + adapter = adapter_for_stage(node) + work_plan = work_plan or adapter.plan(self.plan, node) compiled_node = next( stage for stage in plan_to_dict(self.plan)["stages"] if stage["stage_id"] == node.stage_id ) - return stable_hash( - { - "execution_contract_hash": self.plan.contract_hash, - "semantic_config": semantic_stage_config( - self.plan.experiment_config, node.stage_id - ), - "compiled_node": compiled_node, - "root_overrides": list(self.plan.overrides), - "work_items": [ - { - "work_id": item.work_id, - "shard_index": item.shard_index, - "shard_count": item.shard_count, - "gpus_per_instance": item.gpus_per_instance, - "local_gpu_ids": list(item.local_gpu_ids), - "metadata": dict(item.metadata), - } - for item in work_plan.items - ], - }, - prefix=f"{node.stage_id}_execution", + payload = { + "execution_contract_hash": self.plan.contract_hash, + "semantic_config": semantic_stage_config(self.plan.experiment_config, node.stage_id), + "compiled_node": compiled_node, + "root_overrides": list(self.plan.overrides), + "work_items": [ + { + "work_id": item.work_id, + "shard_index": item.shard_index, + "shard_count": item.shard_count, + "gpus_per_instance": item.gpus_per_instance, + "local_gpu_ids": list(item.local_gpu_ids), + "metadata": dict(item.metadata), + } + for item in work_plan.items + ], + } + adapter_projection = adapter.execution_identity_projection( + plan=self.plan, + node=node, + work_plan=work_plan, ) + if adapter_projection: + payload["adapter_projection"] = dict(adapter_projection) + return stable_hash(payload, prefix=f"{node.stage_id}_execution") def _bind_attempt_to_stage_execution( self, @@ -502,7 +510,12 @@ def _stage_is_active(self, stage_id: str) -> bool: def _recover_failed_stages(self) -> None: for node in self.plan.stages: record = self.store.load_stage_record(node.stage_id) - stage_execution_identity = self._stage_execution_identity(node) + if record is None or record.status != JobState.FAILED.value or not record.attempts: + continue + try: + stage_execution_identity = self._stage_execution_identity(node) + except ExecutionIdentityProjectionUnavailable: + continue current_failure = bool(record and record.attempts) and all( attempt.contract_hash == self.plan.contract_hash and (attempt.metadata or {}).get("stage_execution_identity") @@ -514,12 +527,9 @@ def _recover_failed_stages(self) -> None: and (attempt.metadata or {}).get("stage_execution_identity_incompatible") is True for attempt in record.attempts ) - if ( - record is None - or record.status != JobState.FAILED.value - or not record.attempts - or not (current_failure or legacy_incompatibility) - or stage_is_complete(self.plan.experiment_config, node.stage_id) + if not (current_failure or legacy_incompatibility) or stage_is_complete( + self.plan.experiment_config, + node.stage_id, ): continue finalization_failures = [ @@ -556,6 +566,10 @@ def _fail_legacy_completed_attempts( legacy_attempts = self._legacy_completed_attempts(node, attempts) if not legacy_attempts: return False + try: + stage_execution_identity = self._stage_execution_identity(node) + except ExecutionIdentityProjectionUnavailable: + return False persisted_attempts = [ PersistedAttempt( attempt_id=attempt["attempt_id"], @@ -589,7 +603,7 @@ def _fail_legacy_completed_attempts( "stage_id": node.stage_id, "failure_class": FailureClass.CONFIG.value, "contract_hash": self.plan.contract_hash, - "stage_execution_identity": self._stage_execution_identity(node), + "stage_execution_identity": stage_execution_identity, "attempt_ids": [attempt.attempt_id for attempt in persisted_attempts], "incompatibility": "missing_stage_execution_identity", "reason": reason, @@ -614,6 +628,7 @@ def _submit_stage(self, node: StagePlanNode) -> bool: if self._stage_has_active_or_completed_work(node): return False adapter = adapter_for_stage(node) + adapter.prepare_execution_identity_projection(plan=self.plan, node=node) work_plan = adapter.plan(self.plan, node) handles: list[tuple[JobHandle, str, str]] = [] available_nodes = self._available_nodes() @@ -688,6 +703,18 @@ def _submit_stage(self, node: StagePlanNode) -> bool: self._last_states[handle.handle_id] = JobState.RUNNING return bool(handles) + def _wait_for_manual_input( + self, + node: StagePlanNode, + request: ManualInputRequired, + ) -> bool: + self._manual_waiting = request + self.logger.wait( + f"{node.stage_id}: manual review is ready; write manual_decision.json " + "and rerun the controller" + ) + return False + def _finalize_stage(self, node: StagePlanNode) -> bool: adapter = adapter_for_stage(node) work_plan = adapter.plan(self.plan, node) @@ -696,12 +723,7 @@ def _finalize_stage(self, node: StagePlanNode) -> bool: aggregate = adapter.aggregate(plan=self.plan, node=node, work_plan=work_plan) except ManualInputRequired as request: if not self.terminal_controls.enabled: - self._manual_waiting = request - self.logger.wait( - f"{node.stage_id}: manual review is ready; write manual_decision.json " - "and rerun the controller" - ) - return False + return self._wait_for_manual_input(node, request) selected = self.terminal_controls.choose_revisions(request.prompt, request.revision_ids) decision_path = ( self.plan.puzzle_dir @@ -725,6 +747,8 @@ def _finalize_stage(self, node: StagePlanNode) -> bool: temporary.replace(decision_path) try: aggregate = adapter.aggregate(plan=self.plan, node=node, work_plan=work_plan) + except ManualInputRequired as request: + return self._wait_for_manual_input(node, request) except (OSError, ValueError, RuntimeError) as error: return self._record_stage_aggregation_failure( node, @@ -1401,7 +1425,11 @@ def _on_signal(signum: int, _frame: object | None) -> None: if node.stage_id in self._failed_stages: continue stage_attempts = self.store.list_attempts(node.stage_id) - if stage_attempts and not self._stage_is_active(node.stage_id): + if ( + stage_attempts + and not self._stage_is_active(node.stage_id) + and self._parents_ready(node) + ): if self._required_work_is_completed(node, stage_attempts): finalized = self._finalize_stage(node) if not finalized and self._manual_waiting is None: diff --git a/modelopt/torch/puzzletron/post_mip/identity.py b/modelopt/torch/puzzletron/post_mip/identity.py new file mode 100644 index 00000000000..e21ca4fab4b --- /dev/null +++ b/modelopt/torch/puzzletron/post_mip/identity.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Canonical execution identities for post-MIP nodes.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from ..identity import stable_hash +from .base import CompiledPostMIPNode, compile_post_mip_flows +from .records import CandidateLedger, CandidateSet + +if __package__.startswith("puzzletron_orchestrator."): + from puzzletron_orchestrator.adapters.base import ExecutionIdentityProjectionUnavailable +else: + from ..orchestration.adapters.base import ExecutionIdentityProjectionUnavailable + +__all__ = [ + "PostMIPExecutionContractUnavailable", + "expected_post_mip_candidate_count", + "expected_post_mip_execution_contract", + "expected_post_mip_execution_identity", + "post_mip_execution_contract", + "post_mip_execution_contract_identity", + "post_mip_execution_identity", + "prepare_post_mip_candidate_ledger", +] + + +class PostMIPExecutionContractUnavailable(ExecutionIdentityProjectionUnavailable): + """The current upstream artifacts do not yet define a post-MIP execution.""" + + +def _puzzle_dir(config: Mapping[str, Any]) -> Path: + return Path(config.get("puzzle_dir") or (config.get("experiment") or {})["dir"]) + + +def _compiled_node(config: Mapping[str, Any], stage_id: str) -> CompiledPostMIPNode: + matches = [node for node in compile_post_mip_flows(config) if node.stage_id == stage_id] + if len(matches) != 1: + raise ValueError(f"expected one compiled post-MIP node {stage_id!r}, found {len(matches)}") + node = matches[0] + if not node.capabilities.implemented: + raise NotImplementedError(f"post-MIP node type {node.node_type!r} is not implemented") + return node + + +def _input_candidate_set( + ledger: CandidateLedger, + config: Mapping[str, Any], + node: CompiledPostMIPNode, +) -> CandidateSet: + if node.input_id == "source": + flow = (config.get("post_mip") or {})["flows"][node.flow_id] + candidate_set = ledger.root_set(node.flow_id, flow["source"]) + else: + candidate_set = ledger.load_candidate_set(node.input_id) + canonical = CandidateSet.create( + candidate_set.flow_id, + candidate_set.node_id, + candidate_set.revision_ids, + producer_execution_identity=candidate_set.producer_execution_identity, + ) + if canonical != candidate_set: + raise RuntimeError(f"post-MIP input node {node.input_id!r} has an invalid candidate set") + return candidate_set + + +def _active_mip_contract(config: Mapping[str, Any]) -> tuple[str, set[str]]: + active_path = _puzzle_dir(config) / "mip" / "active_profiles.json" + if not active_path.is_file(): + raise PostMIPExecutionContractUnavailable("active MIP profile manifest is unavailable") + active = json.loads(active_path.read_text()) + if not isinstance(active, Mapping): + raise TypeError(f"active MIP profile manifest must contain a mapping: {active_path}") + status = active.get("status") + if not isinstance(status, str) or not status: + raise ValueError(f"active MIP profile manifest has an invalid status: {active_path}") + if status != "success": + raise PostMIPExecutionContractUnavailable( + f"active MIP profile manifest is not complete: {active_path}" + ) + execution_identity = active.get("execution_identity") + if not isinstance(execution_identity, str) or not execution_identity: + raise ValueError(f"active MIP profile manifest has no execution identity: {active_path}") + profile_values = active.get("profile_ids") + if ( + not isinstance(profile_values, list) + or not profile_values + or any(not isinstance(value, str) or not value for value in profile_values) + ): + raise ValueError(f"active MIP profile manifest has invalid profile IDs: {active_path}") + return execution_identity, set(profile_values) + + +def _expected_post_mip_inputs( + config: Mapping[str, Any], stage_id: str +) -> tuple[CompiledPostMIPNode, CandidateSet, CandidateLedger]: + active_execution, active_profiles = _active_mip_contract(config) + node = _compiled_node(config, stage_id) + ledger = CandidateLedger(_puzzle_dir(config) / "artifacts" / "post_mip") + if not ledger.registry_path.is_file(): + raise PostMIPExecutionContractUnavailable("post-MIP candidate registry is unavailable") + if ( + ledger.active_mip_execution_identity != active_execution + or ledger.active_profile_ids != active_profiles + ): + raise PostMIPExecutionContractUnavailable( + "post-MIP candidate registry does not reflect the active MIP execution" + ) + try: + candidate_set = _input_candidate_set(ledger, config, node) + except FileNotFoundError as error: + raise PostMIPExecutionContractUnavailable( + f"post-MIP inputs for {stage_id!r} are unavailable" + ) from error + return node, candidate_set, ledger + + +def post_mip_execution_contract( + config: Mapping[str, Any], + node: CompiledPostMIPNode, + candidate_set: CandidateSet, + ledger: CandidateLedger, +) -> dict[str, Any]: + """Return the exact node, input, dependency, and source-revision contract.""" + + dependency_owners = { + reference.partition(".")[0] + for reference in node.metric_references + if not reference.startswith("mip.") + } + if node.model_source not in {"latest", "origin"}: + dependency_owners.add(node.model_source) + dependency_executions = {} + for owner in sorted(dependency_owners): + current_path = ( + _puzzle_dir(config) / "artifacts" / "post_mip" / "nodes" / owner / "current.json" + ) + dependency_executions[owner] = json.loads(current_path.read_text())["execution_identity"] + source_revisions = { + revision_id: ledger.source_revision(revision_id, node.model_source).revision_id + for revision_id in candidate_set.revision_ids + } + return { + "candidate_set": candidate_set.identity, + "node": node.config, + "dependency_executions": dependency_executions, + "source_revisions": source_revisions, + } + + +def post_mip_execution_contract_identity(contract: Mapping[str, Any]) -> str: + """Hash one already-resolved canonical post-MIP execution contract.""" + + return stable_hash(contract, prefix="post_mip_execution") + + +def post_mip_execution_identity( + config: Mapping[str, Any], + node: CompiledPostMIPNode, + candidate_set: CandidateSet, + ledger: CandidateLedger, +) -> str: + """Return the producer identity for one resolved node execution.""" + + return post_mip_execution_contract_identity( + post_mip_execution_contract(config, node, candidate_set, ledger) + ) + + +def expected_post_mip_execution_contract( + config: Mapping[str, Any], stage_id: str +) -> dict[str, Any]: + """Resolve the currently runnable contract for a compiled post-MIP stage.""" + + node, candidate_set, ledger = _expected_post_mip_inputs(config, stage_id) + try: + return post_mip_execution_contract(config, node, candidate_set, ledger) + except FileNotFoundError as error: + raise PostMIPExecutionContractUnavailable( + f"post-MIP inputs for {stage_id!r} are unavailable" + ) from error + + +def expected_post_mip_candidate_count(config: Mapping[str, Any], stage_id: str) -> int: + """Return the candidate count for the currently runnable node contract.""" + + _node, candidate_set, _ledger = _expected_post_mip_inputs(config, stage_id) + return len(candidate_set.revision_ids) + + +def prepare_post_mip_candidate_ledger(config: Mapping[str, Any]) -> None: + """Publish a candidate ledger for the active successful MIP execution if needed.""" + + active_execution, active_profiles = _active_mip_contract(config) + puzzle_dir = _puzzle_dir(config) + ledger = CandidateLedger(puzzle_dir / "artifacts" / "post_mip") + if ( + ledger.registry_path.is_file() + and ledger.active_mip_execution_identity == active_execution + and ledger.active_profile_ids == active_profiles + ): + return + ledger.ingest_mip(puzzle_dir) + if ( + ledger.active_mip_execution_identity != active_execution + or ledger.active_profile_ids != active_profiles + ): + raise RuntimeError("post-MIP candidate registry preparation produced stale state") + + +def expected_post_mip_execution_identity(config: Mapping[str, Any], stage_id: str) -> str: + """Return the producer identity expected for the current post-MIP inputs.""" + + return post_mip_execution_contract_identity( + expected_post_mip_execution_contract(config, stage_id) + ) diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index cb054134f16..1c3243ffbd2 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -33,6 +33,11 @@ from ..identity import canonicalize, stable_hash from .base import CompiledPostMIPNode, NodeKind, compile_post_mip_flows from .filters import apply_filter +from .identity import ( + expected_post_mip_execution_identity, + post_mip_execution_contract, + post_mip_execution_identity, +) from .records import ArtifactKind, CandidateLedger, CandidateSet, NodeObservation __all__ = [ @@ -125,62 +130,6 @@ def _execution_root( return _node_root(config, node) / "executions" / execution_identity -def _execution_contract( - config: Mapping[str, Any], - node: CompiledPostMIPNode, - candidate_set: CandidateSet, - ledger: CandidateLedger, -) -> dict[str, Any]: - dependency_owners = { - reference.partition(".")[0] - for reference in node.metric_references - if not reference.startswith("mip.") - } - if node.model_source not in {"latest", "origin"}: - dependency_owners.add(node.model_source) - dependency_executions = {} - for owner in sorted(dependency_owners): - current_path = ( - _puzzle_dir(config) / "artifacts" / "post_mip" / "nodes" / owner / "current.json" - ) - dependency_executions[owner] = json.loads(current_path.read_text())["execution_identity"] - source_revisions = { - revision_id: ledger.source_revision(revision_id, node.model_source).revision_id - for revision_id in candidate_set.revision_ids - } - return { - "candidate_set": candidate_set.identity, - "node": node.config, - "dependency_executions": dependency_executions, - "source_revisions": source_revisions, - } - - -def _execution_identity( - config: Mapping[str, Any], - node: CompiledPostMIPNode, - candidate_set: CandidateSet, - ledger: CandidateLedger, -) -> str: - return stable_hash( - _execution_contract(config, node, candidate_set, ledger), - prefix="post_mip_execution", - ) - - -def expected_post_mip_execution_identity(config: Mapping[str, Any], stage_id: str) -> str: - """Return the identity a completed post-MIP stage must have right now.""" - - node = _compiled_node(config, stage_id) - ledger = _ledger(config) - active = json.loads((_puzzle_dir(config) / "mip" / "active_profiles.json").read_text()) - if active.get("status") != "success" or ledger.active_mip_execution_identity != active.get( - "execution_identity" - ): - raise RuntimeError("post-MIP ledger does not reflect the active MIP execution") - return _execution_identity(config, node, _input_set(ledger, config, node), ledger) - - def _raw_solution(source) -> dict[str, Any]: path = Path(str(source.artifact["solution_path"])) rows = json.loads(path.read_text()) @@ -496,6 +445,12 @@ def _aiperf( for concurrency in concurrencies } topology = dict(settings.pop("topology", {}) or {}) + trust_remote_code = bool( + settings.pop( + "trust_remote_code", + (config.get("model") or {}).get("trust_remote_code", False), + ) + ) gpu_ids = os.environ.get("CUDA_VISIBLE_DEVICES", "") if not gpu_ids: gpu_ids = ",".join(str(index) for index in range(int(topology.get("gpu_group_size", 1)))) @@ -512,6 +467,7 @@ def _aiperf( request_counts=request_counts, solution_id=source.architecture_id, profile_id=node.flow_id, + trust_remote_code=trust_remote_code, **settings, ) metrics = {} @@ -650,7 +606,7 @@ def run_post_mip_node_shard( ledger = _ledger(config) ledger.ingest_mip(_puzzle_dir(config)) candidate_set = _input_set(ledger, config, node) - execution_identity = _execution_identity(config, node, candidate_set, ledger) + execution_identity = post_mip_execution_identity(config, node, candidate_set, ledger) revision_ids = candidate_set.revision_ids[shard_index::shard_count] output_path = ( _execution_root(config, node, execution_identity) @@ -752,7 +708,7 @@ def aggregate_post_mip_node(config: dict[str, Any], stage_id: str) -> dict[str, ledger = _ledger(config) ledger.ingest_mip(_puzzle_dir(config)) input_set = _input_set(ledger, config, node) - execution_identity = _execution_identity(config, node, input_set, ledger) + execution_identity = post_mip_execution_identity(config, node, input_set, ledger) timed_out_candidates = [] if node.node_type == "filter": observations, output_set = _aggregate_filter(ledger, node, input_set, execution_identity) @@ -873,7 +829,7 @@ def aggregate_post_mip_node(config: dict[str, Any], stage_id: str) -> dict[str, "observations_path": str(observations_path), "candidate_set_path": str(candidate_set_path), "execution_identity": execution_identity, - "execution_contract": _execution_contract(config, node, input_set, ledger), + "execution_contract": post_mip_execution_contract(config, node, input_set, ledger), "checkpoints": sorted( { str(ledger.revisions[revision_id].artifact["checkpoint"]) diff --git a/modelopt/torch/puzzletron/stages/future.py b/modelopt/torch/puzzletron/stages/future.py index 39253eb58b8..64b7a7986ff 100644 --- a/modelopt/torch/puzzletron/stages/future.py +++ b/modelopt/torch/puzzletron/stages/future.py @@ -382,6 +382,13 @@ def aiperf_stage(config: dict[str, Any], manifest: StageManifest): for index in range(0, len(visible), group_size) ] output_dir = Path(stage_cfg.get("output_dir", puzzle_dir / "artifacts" / "aiperf")) + model_cfg = dict(config.get("model") or {}) + trust_remote_code = bool( + stage_cfg.get("trust_remote_code", model_cfg.get("trust_remote_code", False)) + ) + allow_aiperf_v011_online_tokenizer_resolution = bool( + stage_cfg.get("allow_aiperf_v011_online_tokenizer_resolution", False) + ) work = _aiperf_checkpoint_work(checkpoints, list(stage_cfg.get("concurrency", [1, 2, 4, 8]))) pool: Queue[str] = Queue() for gpu_group in gpu_groups: @@ -414,6 +421,10 @@ def _run_checkpoint(item): extra_inputs=dict(stage_cfg.get("extra_inputs") or {}), use_server_token_count=bool(stage_cfg.get("use_server_token_count", True)), seed=int(stage_cfg.get("seed", 42)), + trust_remote_code=trust_remote_code, + allow_aiperf_v011_online_tokenizer_resolution=( + allow_aiperf_v011_online_tokenizer_resolution + ), ) finally: pool.put(gpu_ids) diff --git a/modelopt/torch/puzzletron/utils/vllm_adapter.py b/modelopt/torch/puzzletron/utils/vllm_adapter.py index 3622fd0a44e..fa8aec78dc2 100644 --- a/modelopt/torch/puzzletron/utils/vllm_adapter.py +++ b/modelopt/torch/puzzletron/utils/vllm_adapter.py @@ -218,14 +218,21 @@ def configure_anymodel_metadata(hf_config: Any, descriptor: Any) -> bool: return True -def refresh_realized_checkpoint_config(checkpoint_dir: str | Path) -> Path: - """Rebuild the vLLM interchange fields of an already-realized checkpoint.""" +def refresh_realized_checkpoint_config( + checkpoint_dir: str | Path, + *, + trust_remote_code: bool = False, +) -> Path: + """Rebuild vLLM interchange fields without trusting checkpoint code by default.""" from transformers import AutoConfig from ..anymodel.registry import resolve_descriptor checkpoint_dir = Path(checkpoint_dir) - config = AutoConfig.from_pretrained(checkpoint_dir, trust_remote_code=True) + config = AutoConfig.from_pretrained( + checkpoint_dir, + trust_remote_code=trust_remote_code, + ) descriptor = resolve_descriptor(config).descriptor text_config = _get_text_config(config) @@ -338,22 +345,16 @@ def _convert_block_entry( is_full = _get(attn, "sliding_window_size") == "full" head_dim_field = "global_head_dim" if (is_full or k_eq_v) else "head_dim" current_head_dim = ( - global_global_head_dim - if head_dim_field == "global_head_dim" - else global_head_dim + global_global_head_dim if head_dim_field == "global_head_dim" else global_head_dim ) if qk_head_dim != current_head_dim: entry[head_dim_field] = qk_head_dim window = _get(attn, "sliding_window_size") if window is not None: - desired_type = ( - "full_attention" if window == "full" else "sliding_attention" - ) + desired_type = "full_attention" if window == "full" else "sliding_attention" current_type = ( - global_layer_types[layer_idx] - if layer_idx < len(global_layer_types) - else None + global_layer_types[layer_idx] if layer_idx < len(global_layer_types) else None ) if global_layer_types and current_type != desired_type: layer_types = list(global_layer_types) @@ -462,7 +463,9 @@ def _derive_per_layer_config( global_moe_latent_size = ( _get(text_config, moe_latent_field) if moe_latent_field is not None else None ) - global_mamba_values = {hf_field: _get(text_config, hf_field) for hf_field in mamba_fields.values()} + global_mamba_values = { + hf_field: _get(text_config, hf_field) for hf_field in mamba_fields.values() + } global_q_lora_rank = _get(text_config, "q_lora_rank") global_kv_lora_rank = _get(text_config, "kv_lora_rank") global_layer_types = list(_get(text_config, "layer_types") or ()) diff --git a/noxfile.py b/noxfile.py index f6ebc70faf2..81771d49ade 100644 --- a/noxfile.py +++ b/noxfile.py @@ -233,10 +233,12 @@ def gpu(session): ) +# Manual one-GPU runner for the focused Puzzletron lifecycle test. # Container: dedicated Puzzletron v2 GPU image with the pinned ci_environment.json runtime. @nox.session(venv_backend="none") def gpu_puzzletron(session): - """Run the focused Puzzletron suite in its pinned one-GPU image.""" + """Verify the pinned runtime, then run the focused Puzzletron lifecycle GPU test.""" + session.env["CUDA_VISIBLE_DEVICES"] = os.environ.get("CUDA_VISIBLE_DEVICES", "0") _verify_puzzletron_v2_environment(session) session.run( "python", diff --git a/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py b/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py index bf41d716161..f7d521e9c9b 100644 --- a/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py +++ b/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Reusable hermetic tiny-Qwen campaign for Puzzletron integration tests.""" +"""Reusable local-data tiny-Qwen campaign for Puzzletron integration tests.""" from __future__ import annotations @@ -34,7 +34,7 @@ load_execution_config, load_runner_config, ) -from puzzletron_setup.v2.wizard import _DEFAULT_DATA_SOURCE, _DEFAULT_MODEL_SOURCE, run_wizard_v2 +from puzzletron_setup.v2.wizard import run_wizard_v2 __all__ = ["TinyQwenCampaign", "build_tiny_qwen_campaign"] @@ -51,9 +51,11 @@ class _DefaultsBackend: def __init__(self, campaign_dir: Path) -> None: self.campaign_dir = campaign_dir + self.answered: set[str] = set() def text(self, message: str, default: str) -> Any: if message == "Campaign directory:": + self.answered.add(message) return str(self.campaign_dir) return default @@ -63,10 +65,14 @@ def select( choices: Sequence[PromptChoice], default: Any, ) -> Any: - if message == "Model:": - return _DEFAULT_MODEL_SOURCE - if message == "Dataset:": - return _DEFAULT_DATA_SOURCE + if message in {"Model:", "Dataset:"}: + defaults = [choice.value for choice in choices if choice.title.startswith("Default —")] + if len(defaults) != 1: + raise AssertionError( + f"expected one resolved default for {message!r}, found {len(defaults)}" + ) + self.answered.add(message) + return defaults[0] if default is not None: return default return next(choice.value for choice in choices if choice.disabled is None) @@ -181,6 +187,7 @@ def _post_mip_overrides(flow_id: str) -> tuple[str, ...]: f"{prefix}.online_eval.config.eval_samples=2", f"{prefix}.best_lm.top_k=3", f"{prefix}.serving.config.request_count=4", + f"+{prefix}.serving.config.allow_aiperf_v011_online_tokenizer_resolution=true", f"{prefix}.fastest.top_k=2", f"{prefix}.short_kd.config.max_steps=2", f"{prefix}.short_kd.config.global_batch_size=1", @@ -280,11 +287,15 @@ def build_tiny_qwen_campaign( ) ) + backend = _DefaultsBackend(campaign_dir) generated = run_wizard_v2( resume=None, defaults_path=defaults_path, - backend=_DefaultsBackend(campaign_dir), + backend=backend, ) + expected_prompts = {"Campaign directory:", "Model:", "Dataset:"} + if backend.answered != expected_prompts: + raise AssertionError(f"wizard prompt contract changed; answered {sorted(backend.answered)}") smoke_bundle = generated / "smoke" experiment = yaml.safe_load((smoke_bundle / "experiment.yaml").read_text()) flows = dict((experiment.get("post_mip") or {}).get("flows") or {}) @@ -293,6 +304,9 @@ def build_tiny_qwen_campaign( flow_id = next(iter(flows)) overrides = _post_mip_overrides(flow_id) config = pipeline_config_from_path(smoke_bundle / "experiment.yaml", overrides=overrides) + serving_config = config["post_mip"]["flows"][flow_id]["nodes"]["serving"]["config"] + if serving_config.get("allow_aiperf_v011_online_tokenizer_resolution") is not True: + raise AssertionError("tiny-Qwen AIPerf compatibility policy override was not composed") compiled_plan = compile_campaign_plan( experiment_config_path=smoke_bundle / "experiment.yaml", runner=load_runner_config(smoke_bundle / "runner.yaml"), diff --git a/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py b/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py index 452c3280663..4935a0dcc0c 100644 --- a/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py +++ b/tests/unit/torch/puzzletron/test_aiperf_context_capacity.py @@ -31,6 +31,7 @@ _profile_command, _server_max_model_len, _topology_vllm_args, + _vllm_server_command, ) @@ -92,11 +93,27 @@ def test_prepare_vllm_checkpoint_refreshes_heterogeneous_metadata(monkeypatch, t observed = [] monkeypatch.setattr( "modelopt.torch.puzzletron.utils.vllm_adapter.refresh_realized_checkpoint_config", - lambda path: observed.append(path), + lambda path, **kwargs: observed.append((path, kwargs)), ) assert _prepare_vllm_checkpoint(tmp_path) is True - assert observed == [tmp_path] + assert observed == [(tmp_path, {"trust_remote_code": False})] + + +def test_prepare_vllm_checkpoint_preserves_explicit_remote_code_trust(monkeypatch, tmp_path): + config = { + "architectures": ["BaseModel"], + "text_config": {"per_layer_config": {"0": {"intermediate_size": 8}}}, + } + (tmp_path / "config.json").write_text(json.dumps(config)) + observed = [] + monkeypatch.setattr( + "modelopt.torch.puzzletron.utils.vllm_adapter.refresh_realized_checkpoint_config", + lambda path, **kwargs: observed.append((path, kwargs)), + ) + + assert _prepare_vllm_checkpoint(tmp_path, trust_remote_code=True) is True + assert observed == [(tmp_path, {"trust_remote_code": True})] def test_prepare_vllm_checkpoint_leaves_native_teacher_unchanged(tmp_path): @@ -106,18 +123,34 @@ def test_prepare_vllm_checkpoint_leaves_native_teacher_unchanged(tmp_path): assert _prepare_vllm_checkpoint(tmp_path) is False -def test_aiperf_environment_avoids_broken_offline_local_path_resolution(): - expected_source = { +def _offline_environment() -> dict[str, str]: + return { "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", "HF_DATASETS_OFFLINE": "1", "HF_HOME": "/cache/huggingface", "UNCHANGED": "value", } - source = dict(expected_source) + + +def test_aiperf_environment_preserves_offline_policy_by_default(): + source = _offline_environment() resolved = _aiperf_subprocess_environment(source) + assert resolved == source + assert resolved is not source + + +def test_aiperf_v011_online_tokenizer_relaxation_is_explicit_and_non_mutating(): + expected_source = _offline_environment() + source = dict(expected_source) + + resolved = _aiperf_subprocess_environment( + source, + allow_aiperf_v011_online_tokenizer_resolution=True, + ) + assert "HF_HUB_OFFLINE" not in resolved assert "TRANSFORMERS_OFFLINE" not in resolved assert resolved["HF_DATASETS_OFFLINE"] == "1" @@ -178,6 +211,28 @@ def test_vllm_topology_args_enable_dp_and_expert_parallel_only_when_requested(): assert "--expert-parallel-size" not in ep_args +@pytest.mark.parametrize("trust_remote_code", [False, True]) +def test_vllm_server_command_applies_explicit_remote_code_policy( + monkeypatch, tmp_path, trust_remote_code +): + monkeypatch.setattr( + "modelopt.torch.puzzletron.benchmarks.aiperf._descriptor_vllm_args", + lambda _checkpoint: [], + ) + + command = _vllm_server_command( + checkpoint_dir=tmp_path, + port=8000, + model_name="served-model", + input_tokens=32, + output_tokens=8, + topology={"gpu_group_size": 1}, + trust_remote_code=trust_remote_code, + ) + + assert ("--trust-remote-code" in command) is trust_remote_code + + def test_profile_command_maps_each_workload_answer_to_aiperf_cli(tmp_path): command = _profile_command( executable=Path("/opt/aiperf/bin/aiperf"), diff --git a/tests/unit/torch/puzzletron/test_ci_environment.py b/tests/unit/torch/puzzletron/test_ci_environment.py index 72878bcaf84..b2bd41cdc63 100644 --- a/tests/unit/torch/puzzletron/test_ci_environment.py +++ b/tests/unit/torch/puzzletron/test_ci_environment.py @@ -31,6 +31,50 @@ def read_text(self, filename: str) -> str | None: return json.dumps(self.payload) +_EXPECTED_SOURCE = { + "repository": "https://github.com/Separius/Automodel.git", + "commit": "b22cd029d806197e249f2cc4a42c5de91713b772", +} + + +def _pep610_source(repository: str, commit: str) -> dict: + return { + "url": repository, + "vcs_info": {"vcs": "git", "commit_id": commit}, + } + + +def test_pep610_exact_source_is_accepted(monkeypatch): + monkeypatch.setattr( + ci_environment.metadata, + "distribution", + lambda _package: _Distribution( + _pep610_source(_EXPECTED_SOURCE["repository"], _EXPECTED_SOURCE["commit"]) + ), + ) + + ci_environment.verify_installed_vcs_source("nemo-automodel", _EXPECTED_SOURCE) + + +@pytest.mark.parametrize( + ("repository", "commit"), + [ + ("https://github.com/example/Automodel.git", _EXPECTED_SOURCE["commit"]), + (_EXPECTED_SOURCE["repository"], "0" * 40), + ], + ids=("repository", "commit"), +) +def test_pep610_vcs_source_mismatch_is_rejected(monkeypatch, repository, commit): + monkeypatch.setattr( + ci_environment.metadata, + "distribution", + lambda _package: _Distribution(_pep610_source(repository, commit)), + ) + + with pytest.raises(RuntimeError, match="source mismatch"): + ci_environment.verify_installed_vcs_source("nemo-automodel", _EXPECTED_SOURCE) + + def test_editable_pinned_dependency_must_be_clean(monkeypatch): monkeypatch.setattr( ci_environment.metadata, @@ -55,8 +99,5 @@ def test_editable_pinned_dependency_must_be_clean(monkeypatch): with pytest.raises(RuntimeError, match="dependency 'nemo-automodel' is dirty"): ci_environment.verify_installed_vcs_source( "nemo-automodel", - { - "repository": "https://github.com/Separius/Automodel.git", - "commit": "b22cd029d806197e249f2cc4a42c5de91713b772", - }, + _EXPECTED_SOURCE, ) diff --git a/tests/unit/torch/puzzletron/test_global_kd_canonical.py b/tests/unit/torch/puzzletron/test_global_kd_canonical.py index 1d7cec1b33d..19c1fe79476 100644 --- a/tests/unit/torch/puzzletron/test_global_kd_canonical.py +++ b/tests/unit/torch/puzzletron/test_global_kd_canonical.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +"""Tests for Puzzletron's canonical global-distillation behavior.""" + import json from contextlib import contextmanager from pathlib import Path @@ -87,9 +89,7 @@ def native_state_dict_adapter_context(cls, block_configs): assert observed[0][0].to_dict() == {"subblock_configs": []} -def test_distillation_overfit_stage_disables_mtp_objectives_by_default( - monkeypatch, tmp_path -): +def test_distillation_overfit_stage_disables_mtp_objectives_by_default(monkeypatch, tmp_path): """Nano-like checkpoints without MTP must not enable MTP loss implicitly.""" from modelopt.torch.puzzletron.manifest import StageManifest from modelopt.torch.puzzletron.stages import future @@ -146,9 +146,7 @@ def fake_run_global_kd(kd_config): }, } - future.distillation_overfit_stage( - config, StageManifest(stage="global_distillation_sanity") - ) + future.distillation_overfit_stage(config, StageManifest(stage="global_distillation_sanity")) assert captured["objective"] == { "main_ce": {"weight": 1.0}, @@ -171,19 +169,10 @@ def test_global_distillation_summary_publishes_canonical_training_records(tmp_pa training_log = output_dir / "checkpoints/training.jsonl" training_log.parent.mkdir(parents=True) training_log.write_text( - json.dumps({"step": 1, "loss": 2.0}) - + "\n" - + json.dumps({"step": 2, "loss": 1.0}) - + "\n" + json.dumps({"step": 1, "loss": 2.0}) + "\n" + json.dumps({"step": 2, "loss": 1.0}) + "\n" ) for step in (1, 2): - checkpoint = ( - output_dir - / "checkpoints" - / f"epoch_0_step_{step}" - / "model" - / "consolidated" - ) + checkpoint = output_dir / "checkpoints" / f"epoch_0_step_{step}" / "model" / "consolidated" checkpoint.mkdir(parents=True) (checkpoint / "config.json").write_text("{}") (checkpoint.parents[1] / "saving_completed").touch() @@ -209,9 +198,7 @@ def test_global_distillation_summary_publishes_canonical_training_records(tmp_pa assert payload["max_steps"] == 256 assert payload["sequence_length"] == 16384 assert payload["records"][-1] == {"step": 2, "loss": 1.0} - assert payload["post_kd_checkpoint"].endswith( - "checkpoints/epoch_0_step_2/model/consolidated" - ) + assert payload["post_kd_checkpoint"].endswith("checkpoints/epoch_0_step_2/model/consolidated") def test_global_kd_metric_logger_flushes_every_optimizer_step(): @@ -227,9 +214,7 @@ def test_global_kd_metric_logger_flushes_every_optimizer_step(): assert logger.flush is True -def test_global_kd_uses_memory_bounded_1f1b_by_default_and_allows_override( - tmp_path, monkeypatch -): +def test_global_kd_uses_memory_bounded_1f1b_by_default_and_allows_override(tmp_path, monkeypatch): monkeypatch.setattr( "modelopt.torch.puzzletron.plugins.automodel.config.inject_descriptor_pipeline_config", lambda *args, **kwargs: None, @@ -262,10 +247,7 @@ def test_global_kd_uses_memory_bounded_1f1b_by_default_and_allows_override( assert default_recipe["distributed"]["pipeline"]["pp_microbatch_size"] == 1 assert default_recipe["distributed"]["pipeline"]["pp_batch_size"] == 8 assert default_recipe["dataloader"]["batch_size"] == 8 - assert ( - override_recipe["distributed"]["pipeline"]["pp_schedule"] - == "interleaved1f1b" - ) + assert override_recipe["distributed"]["pipeline"]["pp_schedule"] == "interleaved1f1b" def test_global_kd_auto_domain_uses_canonical_text_dataset(tmp_path, monkeypatch): @@ -310,10 +292,7 @@ def test_global_kd_auto_domain_uses_canonical_text_dataset(tmp_path, monkeypatch assert recipe["recipe"] == "KnowledgeDistillationRecipeForNextTokenPrediction" assert recipe["dataset"] == { - "_target_": ( - "modelopt.torch.puzzletron.distillation.dataset." - "make_puzzletron_llm_dataset" - ), + "_target_": ("modelopt.torch.puzzletron.distillation.dataset.make_puzzletron_llm_dataset"), "dataset_path": str(tmp_path / "dataset"), "split": "train", "num_samples": 4096, @@ -326,9 +305,7 @@ def test_global_kd_auto_domain_uses_canonical_text_dataset(tmp_path, monkeypatch ) -def test_global_kd_packed_text_uses_native_chat_data_and_canonical_pack_size( - tmp_path, monkeypatch -): +def test_global_kd_packed_text_uses_native_chat_data_and_canonical_pack_size(tmp_path, monkeypatch): monkeypatch.setattr( "modelopt.torch.puzzletron.plugins.automodel.config.inject_descriptor_model_kwargs", lambda *args, **kwargs: None, @@ -361,9 +338,7 @@ def test_global_kd_packed_text_uses_native_chat_data_and_canonical_pack_size( recipe = build_automodel_global_kd_recipe(config) - assert recipe["dataset"]["_target_"].endswith( - "make_puzzletron_chat_dataset" - ) + assert recipe["dataset"]["_target_"].endswith("make_puzzletron_chat_dataset") assert recipe["dataloader"]["collate_fn"].endswith("default_collater") assert recipe["packed_sequence"] == { "packed_sequence_size": 256, @@ -395,12 +370,16 @@ def test_global_kd_recipe_publishes_explicit_resume_policy(tmp_path, monkeypatch "pp": 1, } - assert build_automodel_global_kd_recipe(GlobalKDConfig(**common, resume=True))[ - "puzzletron_resume" - ] is True - assert build_automodel_global_kd_recipe(GlobalKDConfig(**common, resume=False))[ - "puzzletron_resume" - ] is False + assert ( + build_automodel_global_kd_recipe(GlobalKDConfig(**common, resume=True))["puzzletron_resume"] + is True + ) + assert ( + build_automodel_global_kd_recipe(GlobalKDConfig(**common, resume=False))[ + "puzzletron_resume" + ] + is False + ) def test_global_kd_config_preserves_per_model_dtype_overrides(tmp_path): @@ -430,12 +409,12 @@ def test_global_kd_config_preserves_per_model_dtype_overrides(tmp_path): kd = build_global_kd_config(config) - assert global_automodel._model_recipe(kd, teacher=False, domain="llm")[ - "torch_dtype" - ] == "float32" - assert global_automodel._model_recipe(kd, teacher=True, domain="llm")[ - "torch_dtype" - ] == "bfloat16" + assert ( + global_automodel._model_recipe(kd, teacher=False, domain="llm")["torch_dtype"] == "float32" + ) + assert ( + global_automodel._model_recipe(kd, teacher=True, domain="llm")["torch_dtype"] == "bfloat16" + ) def test_global_kd_load_checkpoint_honors_resume_policy(): @@ -491,9 +470,7 @@ def test_global_distillation_stage_promotes_canonical_namespace(tmp_path): assert (kd_config.pp, kd_config.cp, kd_config.dp) == (2, 4, 8) -def test_global_kd_preserves_physical_dp_mesh_when_ep_overlays_shards( - tmp_path, monkeypatch -): +def test_global_kd_preserves_physical_dp_mesh_when_ep_overlays_shards(tmp_path, monkeypatch): monkeypatch.setattr( "modelopt.torch.puzzletron.plugins.automodel.config.inject_descriptor_pipeline_config", lambda *args, **kwargs: None, @@ -563,18 +540,14 @@ def from_local(cls, local, *, device_mesh, placements, run_check): source_mesh = Mesh() head_mesh = Mesh() hidden = FakeDTensor(torch.ones(2, 3), source_mesh) - base_layer = type( - "BaseLayer", (), {"weight": FakeDTensor(torch.ones(4, 3), head_mesh)} - )() + base_layer = type("BaseLayer", (), {"weight": FakeDTensor(torch.ones(4, 3), head_mesh)})() head = type("WrappedHead", (), {"base_layer": base_layer})() aligned = global_kd_recipe._align_dtensor_to_module_mesh(hidden, head) assert aligned.device_mesh is head_mesh assert aligned.placements == hidden.placements - assert FakeDTensor.calls == [ - (hidden.local, head_mesh, hidden.placements, False) - ] + assert FakeDTensor.calls == [(hidden.local, head_mesh, hidden.placements, False)] def test_teacher_mtp_projection_uses_local_head_and_student_logit_mesh(monkeypatch): @@ -603,15 +576,13 @@ def from_local(cls, local, *, device_mesh, placements, run_check): teacher_mesh = Mesh() student_mesh = Mesh() hidden = FakeDTensor(torch.tensor([[1.0, 2.0]]), teacher_mesh, ("replicate",)) - weight = FakeDTensor( - torch.tensor([[1.0, 0.0], [0.0, 2.0]]), teacher_mesh, ("shard0",) - ) - head = type("WrappedHead", (), {"base_layer": type("Base", (), {"weight": weight, "bias": None})()})() + weight = FakeDTensor(torch.tensor([[1.0, 0.0], [0.0, 2.0]]), teacher_mesh, ("shard0",)) + head = type( + "WrappedHead", (), {"base_layer": type("Base", (), {"weight": weight, "bias": None})()} + )() reference = FakeDTensor(torch.empty(1, 2), student_mesh, ("shard_vocab",)) - projected = global_kd_recipe._project_teacher_hidden_on_reference_mesh( - hidden, head, reference - ) + projected = global_kd_recipe._project_teacher_hidden_on_reference_mesh(hidden, head, reference) assert projected.device_mesh is student_mesh assert projected.placements == reference.placements @@ -741,6 +712,7 @@ def __init__(self): def test_global_kd_checkpoint_forwards_best_metric_key(tmp_path, monkeypatch): + # Lazy import keeps the optional NeMo AutoModel runtime out of test collection. from modelopt.torch.puzzletron.distillation.global_kd_recipe import _WeightedObjectiveMixin calls = [] @@ -774,10 +746,11 @@ class Recipe(_WeightedObjectiveMixin, BaseRecipe): {"config": type("Config", (), {"checkpoint_dir": tmp_path})()}, )() recipe.dist_env = type("DistEnv", (), {"is_main": True})() + recipe.cfg = {"model": {"trust_remote_code": True}} monkeypatch.setattr( "modelopt.torch.puzzletron.utils.vllm_adapter.refresh_realized_checkpoint_config", - lambda path: refreshes.append( - (path, (tmp_path / "epoch_2_step_17/saving_completed").exists()) + lambda path, **kwargs: refreshes.append( + (path, kwargs, (tmp_path / "epoch_2_step_17/saving_completed").exists()) ), ) @@ -791,7 +764,13 @@ class Recipe(_WeightedObjectiveMixin, BaseRecipe): assert result == "saved" assert calls == [(2, 17, 0.5, {"lm_loss": 0.25}, "lm_loss")] - assert refreshes == [(tmp_path / "epoch_2_step_17/model/consolidated", False)] + assert refreshes == [ + ( + tmp_path / "epoch_2_step_17/model/consolidated", + {"trust_remote_code": True}, + False, + ) + ] assert (tmp_path / "epoch_2_step_17" / "saving_completed").is_file() @@ -813,9 +792,7 @@ def save_model(self, model, path): def save_optimizer(self, saved_optimizer, model, path, scheduler): del path, scheduler saved_models.append(model) - assert saved_optimizer.param_groups[0]["params"] == list( - original.parameters() - ) + assert saved_optimizer.param_groups[0]["params"] == list(original.parameters()) recipe = object.__new__(_WeightedObjectiveMixin) recipe.model_parts = [original] @@ -867,9 +844,7 @@ def __init__(self): recipe._remove_text_inactive_optimizer_parameters() optimized = { - id(parameter) - for group in recipe.optimizer.param_groups - for parameter in group["params"] + id(parameter) for group in recipe.optimizer.param_groups for parameter in group["params"] } assert all(id(parameter) not in optimized for parameter in model.visual.parameters()) assert all(id(parameter) not in optimized for parameter in model.mm_projector.parameters()) @@ -963,9 +938,7 @@ def test_llm_pp_optimizer_step_publishes_every_weighted_objective(monkeypatch): lambda self, batches, max_grad_norm: log_data, ) - recipe = object.__new__( - global_kd_recipe.KnowledgeDistillationRecipeForNextTokenPrediction - ) + recipe = object.__new__(global_kd_recipe.KnowledgeDistillationRecipeForNextTokenPrediction) recipe.needs_teacher = True recipe.pp_enabled = True recipe.device_mesh = type( @@ -986,8 +959,7 @@ def test_llm_pp_optimizer_step_publishes_every_weighted_objective(monkeypatch): } recipe._objective_step_cursor = dict.fromkeys(recipe._objective_buffers, 0) recipe._gradient_squared = { - name: torch.tensor(0.0) - for name in ("vision", "projector", "language", "mtp") + name: torch.tensor(0.0) for name in ("vision", "projector", "language", "mtp") } recipe._dp_allreduce = lambda value, include_cp: value @@ -1045,9 +1017,7 @@ def test_global_kd_uses_canonical_multimodal_packing_and_train_all(tmp_path, mon assert kd.freeze_policy == "train_all" assert kd.teacher_descriptor == "qwen3_5" assert kd.student_descriptor == "qwen3_5" - assert recipe["dataset"]["_target_"].endswith( - "load_materialized_conversation_dataset" - ) + assert recipe["dataset"]["_target_"].endswith("load_materialized_conversation_dataset") assert recipe["dataset"]["path_or_dataset"] == str(tmp_path / "intersyn") assert recipe["packed_sequence"]["pack_size"] == 2048 assert recipe["packed_sequence"]["max_packs"] == 128 diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index 4ac48262fa0..9bdea6d8dbd 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -619,6 +619,53 @@ def test_depth_pool_uses_one_four_node_gang_allocation(tmp_path: Path): assert "run_depth_pool.sh" in script +def test_legacy_aiperf_worker_receives_explicit_security_policy(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"), + ) + node = StagePlanNode( + stage_id="aiperf", + strategy=ExecutionStrategy.SHARDED, + instances=1, + failure_policy=FailurePolicy.STRICT, + mesh={}, + gpus_per_instance=1, + gpus_per_node=8, + nodes=1, + total_gpus=1, + exclusive=False, + parents=("mip",), + distributed=False, + ) + plan = CampaignPlan( + experiment_config_path=str(tmp_path / "experiment.yaml"), + puzzle_dir=tmp_path / "run", + experiment_config={ + "model": {"trust_remote_code": True}, + "aiperf": {"allow_aiperf_v011_online_tokenizer_resolution": True}, + }, + runner=runner, + execution_defaults={"gpus_per_node": 8}, + stages=(node,), + contract_hash="contract", + ) + adapter = adapter_for_stage(node) + work_plan = adapter.plan(plan, node) + + attempt = adapter.command( + plan=plan, + node=node, + item=work_plan.items[0], + attempt_id="a1", + runner=runner, + ) + + assert "--trust-remote-code" in attempt.command.argv + assert "--allow-aiperf-v011-online-tokenizer-resolution" in attempt.command.argv + + def test_depth_pool_packs_four_two_gpu_workers_per_node(tmp_path: Path): runner = RunnerEnvironment( kind="slurm", diff --git a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py index 7c2a85cfa58..2adc3b957c5 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py +++ b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py @@ -354,8 +354,6 @@ def test_depth_completeness_requires_matching_complete_trajectory( def test_build_library_requires_its_own_complete_outputs( tmp_path: Path, write_terminal_manifest ) -> None: - from puzzletron_orchestrator.adapters.stage_compat import stage_is_complete - config = {"puzzle_dir": str(tmp_path)} write_terminal_manifest(tmp_path, "build_library", config=config) (tmp_path / "subblock_stats.json").write_text("{}") @@ -370,8 +368,6 @@ def test_build_library_requires_its_own_complete_outputs( def test_build_library_completion_accepts_equivalent_loader_and_worker_configs( tmp_path: Path, write_terminal_manifest ) -> None: - from puzzletron_orchestrator.adapters.stage_compat import stage_is_complete - experiment = tmp_path / "experiment.yaml" experiment.write_text( f"""\ diff --git a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py index b64d3e4dfe4..c0b0eb2cf49 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py +++ b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py @@ -29,6 +29,7 @@ import pytest import yaml +from puzzletron_orchestrator.adapters.post_mip import ManualInputRequired from puzzletron_orchestrator.adapters.registry import adapter_for_stage from puzzletron_orchestrator.adapters.stage_compat import ( stage_is_complete as artifacts_are_complete, @@ -825,6 +826,61 @@ def aggregate(self, *, plan, node, work_plan): controller._finalize_stage(node) +def test_controller_preserves_repeated_manual_input_request(tmp_path: Path, monkeypatch): + experiment, runner_path, execution_path = _write_configs(tmp_path) + plan = compile_campaign_plan( + experiment_config_path=experiment, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + stage_filter="convert", + ) + node = plan.stages[0] + delegate = adapter_for_stage(node) + + class _Controls: + enabled = True + + @staticmethod + def choose_revisions(_prompt, revision_ids): + return revision_ids + + controller = CampaignController( + plan, + executor=_FakeExecutor(), + terminal_controls=_Controls(), + ) + decision_dir = plan.puzzle_dir / "artifacts/post_mip/nodes/manual" + decision_dir.mkdir(parents=True) + + class _RepeatedManualAdapter: + aggregate_count = 0 + + def __getattr__(self, name): + return getattr(delegate, name) + + def aggregate(self, *, plan, node, work_plan): + del plan, node, work_plan + self.aggregate_count += 1 + raise ManualInputRequired( + "manual", + ("revision-a",), + "Select a revision", + f"execution-{self.aggregate_count}", + ) + + adapter = _RepeatedManualAdapter() + monkeypatch.setattr( + "puzzletron_orchestrator.controller.adapter_for_stage", + lambda _node: adapter, + ) + + assert controller._finalize_stage(node) is False + assert adapter.aggregate_count == 2 + assert controller._manual_waiting is not None + assert controller._manual_waiting.execution_identity == "execution-2" + assert controller._finalization_failures == {} + + @pytest.mark.parametrize("aggregation_failure", [False, True]) def test_controller_fails_when_completed_work_artifacts_do_not_settle( tmp_path: Path, monkeypatch, aggregation_failure: bool diff --git a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py index 65be2c36cfe..570015efccd 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_task_topology.py +++ b/tests/unit/torch/puzzletron/test_orchestration_task_topology.py @@ -249,6 +249,7 @@ def test_run_worker_consumes_task_launcher_identity( check=True, capture_output=True, text=True, + timeout=10, ) assert f"--nnodes {group_size}" in result.stdout diff --git a/tests/unit/torch/puzzletron/test_post_mip_execution_identity.py b/tests/unit/torch/puzzletron/test_post_mip_execution_identity.py new file mode 100644 index 00000000000..292d9d48051 --- /dev/null +++ b/tests/unit/torch/puzzletron/test_post_mip_execution_identity.py @@ -0,0 +1,462 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for post-MIP controller execution identities.""" + +from __future__ import annotations + +import copy +import json +from pathlib import Path + +import pytest + +from puzzletron_orchestrator.adapters.registry import adapter_for_stage +from puzzletron_orchestrator.adapters.stage_compat import stage_is_complete +from puzzletron_orchestrator.controller import CampaignController +from puzzletron_orchestrator.post_mip.base import compile_post_mip_flows +from puzzletron_orchestrator.post_mip.identity import ( + PostMIPExecutionContractUnavailable, + expected_post_mip_execution_contract, + post_mip_execution_contract_identity, +) +from puzzletron_orchestrator.post_mip.records import ( + ArchitectureCandidate, + ArtifactKind, + CandidateLedger, + CandidateSet, + NodeObservation, +) +from puzzletron_orchestrator.schema import ( + CampaignPlan, + ExecutionContract, + ExecutionStrategy, + FailurePolicy, + JobHandle, + JobState, + RunnerEnvironment, + StagePlanNode, +) +from puzzletron_orchestrator.state import PersistedAttempt, StageRunRecord + + +class _TrackingExecutor: + backend = "fake" + + def __init__(self) -> None: + self.attempts = [] + + def submit(self, attempt): + self.attempts.append(attempt) + return JobHandle( + backend=self.backend, + handle_id=f"fake-{attempt.attempt_id}", + attempt_id=attempt.attempt_id, + ) + + +def _identity_fixture(tmp_path: Path) -> tuple[dict, CandidateLedger, tuple[str, ...]]: + puzzle_dir = tmp_path / "run" + puzzle_dir.mkdir() + config = { + "puzzle_dir": str(puzzle_dir), + "mip": {"runs": {"tiny": {}}}, + "post_mip": { + "flows": { + "params": { + "source": {"run": "tiny"}, + "nodes": { + "score": {"type": "evaluation", "config": {"eval_samples": 2}}, + "select": { + "type": "filter", + "input": "score", + "mode": "top_k", + "metric": "score.loss", + "top_k": 2, + "config": {"label": "baseline"}, + }, + "materialize": {"type": "materialize", "input": "select"}, + "final": { + "type": "evaluation", + "input": "select", + "model_source": "materialize", + }, + }, + } + } + }, + } + (puzzle_dir / "mip").mkdir() + (puzzle_dir / "mip" / "active_profiles.json").write_text( + '{"status":"success","execution_identity":"mip-a","profile_ids":["p0"]}\n' + ) + ledger = CandidateLedger(puzzle_dir / "artifacts" / "post_mip") + ledger.active_mip_execution_identity = "mip-a" + ledger.active_profile_ids = {"p0"} + root_ids = [] + for index in range(3): + architecture_id = f"architecture-{index}" + revision = ledger.add_revision( + architecture_id=architecture_id, + artifact_kind=ArtifactKind.CONFIG, + artifact={"kind": "heterogeneous", "mip_metrics": {"loss": float(index)}}, + parent_revision_id=None, + producer_node="mip", + ) + ledger.architectures[architecture_id] = ArchitectureCandidate( + architecture_id=architecture_id, + block_configs=[], + mip_metrics={"loss": float(index)}, + origins=[ + { + "profile_id": "p0", + "mip_execution_identity": "mip-a", + "run_id": "tiny", + "variant_id": "base", + "objective": {"metric": "params"}, + "kind": "heterogeneous", + "rank": index, + "revision_id": revision.revision_id, + } + ], + origin_revision_id=revision.revision_id, + ) + root_ids.append(revision.revision_id) + roots = tuple(root_ids) + ledger.publish() + score_set = CandidateSet.create( + "params", "score", roots[:2], producer_execution_identity="score-a" + ) + ledger.publish_node( + "score", + [ + NodeObservation( + node_id="score", + input_revision_id=revision_id, + source_revision_id=revision_id, + output_revision_id=revision_id, + status="success", + metrics={"loss": float(index)}, + ) + for index, revision_id in enumerate(roots[:2]) + ], + score_set, + "score-a", + ) + select_set = CandidateSet.create( + "params", "select", roots[:2], producer_execution_identity="select-a" + ) + ledger.publish_node( + "select", + [ + NodeObservation( + node_id="select", + input_revision_id=revision_id, + source_revision_id=revision_id, + output_revision_id=revision_id, + status="selected", + ) + for revision_id in roots[:2] + ], + select_set, + "select-a", + ) + materialized_ids = [] + materialized_observations = [] + for index, revision_id in enumerate(roots[:2]): + revision = ledger.add_revision( + architecture_id=f"architecture-{index}", + artifact_kind=ArtifactKind.CHECKPOINT, + artifact={"checkpoint": f"/checkpoint/{index}"}, + parent_revision_id=revision_id, + producer_node="materialize", + ) + materialized_ids.append(revision.revision_id) + materialized_observations.append( + NodeObservation( + node_id="materialize", + input_revision_id=revision_id, + source_revision_id=revision_id, + output_revision_id=revision.revision_id, + status="success", + ) + ) + ledger.publish_node( + "materialize", + materialized_observations, + CandidateSet.create( + "params", + "materialize", + materialized_ids, + producer_execution_identity="materialize-a", + ), + "materialize-a", + ) + return config, ledger, roots + + +def _plan(config: dict, stage_id: str) -> tuple[CampaignPlan, StagePlanNode]: + compiled = next(node for node in compile_post_mip_flows(config) if node.stage_id == stage_id) + node = StagePlanNode( + stage_id=stage_id, + strategy=ExecutionStrategy.SHARDED, + instances=2, + failure_policy=FailurePolicy.STRICT, + mesh={}, + gpus_per_instance=1, + gpus_per_node=8, + nodes=1, + total_gpus=2, + exclusive=False, + parents=compiled.dependency_stage_ids, + distributed=True, + ) + puzzle_dir = Path(config["puzzle_dir"]) + return ( + CampaignPlan( + experiment_config_path=str(puzzle_dir.parent / "experiment.yaml"), + puzzle_dir=puzzle_dir, + experiment_config=config, + runner=RunnerEnvironment( + kind="slurm", + contract=ExecutionContract( + repository=str(puzzle_dir.parent), + venv=str(puzzle_dir.parent / ".venv"), + ), + ), + execution_defaults={"gpus_per_node": 8}, + stages=(node,), + contract_hash="contract", + ), + node, + ) + + +def _controller_identity(config: dict, stage_id: str) -> str: + plan, node = _plan(config, stage_id) + return CampaignController(plan, executor=object())._stage_execution_identity(node) + + +def test_post_mip_identity_tracks_nested_config_and_same_size_candidate_set(tmp_path: Path): + config, ledger, roots = _identity_fixture(tmp_path) + baseline = _controller_identity(config, "post.params.select") + + changed_config = copy.deepcopy(config) + changed_config["post_mip"]["flows"]["params"]["nodes"]["select"]["config"]["label"] = "changed" + assert _controller_identity(changed_config, "post.params.select") != baseline + + replacement_set = CandidateSet.create( + "params", "score", roots[1:], producer_execution_identity="score-b" + ) + ledger.publish_node( + "score", + [ + NodeObservation( + node_id="score", + input_revision_id=revision_id, + source_revision_id=revision_id, + output_revision_id=revision_id, + status="success", + metrics={"loss": float(index)}, + ) + for index, revision_id in enumerate(roots[1:]) + ], + replacement_set, + "score-b", + ) + assert _controller_identity(config, "post.params.select") != baseline + + +def test_post_mip_currentness_does_not_initialize_the_candidate_registry(tmp_path: Path): + config, _ledger, _roots = _identity_fixture(tmp_path) + config["post_mip"]["flows"]["params"]["nodes"]["root_select"] = { + "type": "filter", + "mode": "top_k", + "metric": "mip.loss", + "top_k": 1, + } + stage_id = "post.params.root_select" + execution_identity = _controller_identity(config, stage_id) + summary_path = Path(config["puzzle_dir"]) / "artifacts/post_mip/nodes/root_select/summary.json" + summary_path.parent.mkdir(parents=True, exist_ok=True) + summary_path.write_text( + json.dumps({"status": "success", "execution_identity": execution_identity}) + "\n" + ) + registry = Path(config["puzzle_dir"]) / "artifacts/post_mip/candidate_registry.json" + registry.unlink() + + assert not stage_is_complete(config, stage_id) + assert not registry.exists() + + +def test_post_mip_submission_prepares_a_stale_candidate_registry(tmp_path: Path): + config, ledger, _roots = _identity_fixture(tmp_path) + config["post_mip"]["flows"]["params"]["nodes"]["root_select"] = { + "type": "filter", + "mode": "top_k", + "metric": "mip.loss", + "top_k": 1, + } + active_path = Path(config["puzzle_dir"]) / "mip/active_profiles.json" + active_path.write_text( + '{"status":"success","execution_identity":"mip-b","profile_ids":["p1"]}\n' + ) + assert ledger.active_mip_execution_identity == "mip-a" + + plan, node = _plan(config, "post.params.root_select") + executor = _TrackingExecutor() + controller = CampaignController(plan, executor=executor) + + assert controller._submit_stage(node) + refreshed = CandidateLedger(Path(config["puzzle_dir"]) / "artifacts/post_mip") + + assert len(executor.attempts) == 1 + assert ( + executor.attempts[0] + .metadata["stage_execution_identity"] + .startswith("post.params.root_select_execution_") + ) + assert refreshed.active_mip_execution_identity == "mip-b" + assert refreshed.active_profile_ids == {"p1"} + + +@pytest.mark.parametrize("status", ["pending", "running"]) +def test_non_success_active_mip_defers_identity_without_mutation(tmp_path: Path, status: str): + config, ledger, _roots = _identity_fixture(tmp_path) + active_path = Path(config["puzzle_dir"]) / "mip/active_profiles.json" + registry_before = ledger.registry_path.read_bytes() + active_path.write_text(json.dumps({"status": status}) + "\n") + + with pytest.raises(PostMIPExecutionContractUnavailable): + expected_post_mip_execution_contract(config, "post.params.select") + + assert ledger.registry_path.read_bytes() == registry_before + + +def test_malformed_active_mip_fails_closed(tmp_path: Path): + config, _ledger, _roots = _identity_fixture(tmp_path) + active_path = Path(config["puzzle_dir"]) / "mip/active_profiles.json" + active_path.write_text('{"status":"success","execution_identity":"mip-a","profile_ids":"p0"}\n') + + with pytest.raises(ValueError, match="invalid profile IDs"): + expected_post_mip_execution_contract(config, "post.params.select") + + +def test_post_mip_identity_tracks_dependency_execution_and_source_mapping(tmp_path: Path): + config, _ledger, roots = _identity_fixture(tmp_path) + contract = expected_post_mip_execution_contract(config, "post.params.final") + dependency_only = copy.deepcopy(contract) + dependency_only["dependency_executions"]["materialize"] = "materialize-b" + source_only = copy.deepcopy(contract) + source_only["source_revisions"][roots[0]] = "replacement-revision" + + assert post_mip_execution_contract_identity(dependency_only) != ( + post_mip_execution_contract_identity(contract) + ) + assert post_mip_execution_contract_identity(source_only) != ( + post_mip_execution_contract_identity(contract) + ) + + +def _assert_changed_identity_resubmits(config: dict, changed_config: dict, mutate=None) -> None: + plan_a, node_a = _plan(config, "post.params.select") + controller_a = CampaignController(plan_a, executor=object()) + adapter = adapter_for_stage(node_a) + work_plan = adapter.plan(plan_a, node_a) + attempt = controller_a._bind_attempt_to_stage_execution( + node_a, + work_plan, + adapter.command( + plan=plan_a, + node=node_a, + item=work_plan.items[0], + attempt_id="attempt-a", + runner=plan_a.runner, + ), + ) + controller_a.store.save_attempt( + attempt, + JobHandle(backend="fake", handle_id="fake-attempt-a", attempt_id=attempt.attempt_id), + JobState.COMPLETED.value, + ) + if mutate is not None: + mutate() + + plan_b, node_b = _plan(changed_config, "post.params.select") + executor = _TrackingExecutor() + controller_b = CampaignController(plan_b, executor=executor) + prior = controller_b.store.list_attempts(node_b.stage_id) + identity_b = controller_b._stage_execution_identity(node_b) + + assert not controller_b._required_work_is_completed(node_b, prior) + assert controller_b._submit_stage(node_b) + assert executor.attempts[0].metadata["stage_execution_identity"] == identity_b + + +def test_changed_post_mip_config_resubmits_completed_work(tmp_path: Path): + config, _ledger, _roots = _identity_fixture(tmp_path) + changed_config = copy.deepcopy(config) + changed_config["post_mip"]["flows"]["params"]["nodes"]["select"]["config"]["label"] = "changed" + + _assert_changed_identity_resubmits(config, changed_config) + + +def test_same_size_candidate_set_change_resubmits_completed_work(tmp_path: Path): + config, ledger, roots = _identity_fixture(tmp_path) + + def replace_candidate_set() -> None: + replacement_set = CandidateSet.create( + "params", "score", roots[1:], producer_execution_identity="score-b" + ) + ledger.publish_node( + "score", + [ + NodeObservation( + node_id="score", + input_revision_id=revision_id, + source_revision_id=revision_id, + output_revision_id=revision_id, + status="success", + metrics={"loss": float(index)}, + ) + for index, revision_id in enumerate(roots[1:]) + ], + replacement_set, + "score-b", + ) + + _assert_changed_identity_resubmits(config, config, replace_candidate_set) + + +@pytest.mark.parametrize( + "attempt_status", + [JobState.COMPLETED.value, JobState.PENDING.value, JobState.RUNNING.value], +) +def test_unresolved_future_post_mip_node_defers_failed_record_recovery( + tmp_path: Path, attempt_status: str +): + config, _ledger, _roots = _identity_fixture(tmp_path) + plan, node = _plan(config, "post.params.final") + controller = CampaignController(plan, executor=object()) + current = plan.puzzle_dir / "artifacts/post_mip/nodes/materialize/current.json" + current.unlink() + controller.store.write_stage_record( + StageRunRecord( + stage_id=node.stage_id, + status=JobState.FAILED.value, + attempts=[ + PersistedAttempt( + attempt_id="old", + work_id=f"{node.stage_id}:gang", + stage_id=node.stage_id, + status=attempt_status, + contract_hash=plan.contract_hash, + metadata={"stage_execution_identity": "old"}, + ) + ], + ) + ) + + controller._recover_failed_stages() + + assert controller._failed_stages == set() diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index 82bb5613277..f98c5379529 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -213,6 +213,7 @@ def fake_run_aiperf_sweep(checkpoint, **settings): "minimum_request_count": 4, "requests_per_concurrency": 2, "best_selection_mode": "individual_best", + "allow_aiperf_v011_online_tokenizer_resolution": True, "input_tokens": 1024, "output_tokens": 128, "topology": {"gpu_group_size": 1}, @@ -225,7 +226,7 @@ def fake_run_aiperf_sweep(checkpoint, **settings): ) result = runner._aiperf( - {"puzzle_dir": str(tmp_path)}, + {"puzzle_dir": str(tmp_path), "model": {"trust_remote_code": True}}, node, source, "execution", @@ -234,6 +235,8 @@ def fake_run_aiperf_sweep(checkpoint, **settings): assert captured["checkpoint"] == str(tmp_path / "checkpoint") assert captured["concurrencies"] == (8,) assert captured["request_counts"] == {8: 23} + assert captured["trust_remote_code"] is True + assert captured["allow_aiperf_v011_online_tokenizer_resolution"] is True assert "request_count" not in captured assert "minimum_request_count" not in captured assert "requests_per_concurrency" not in captured diff --git a/tests/unit/torch/puzzletron/test_vllm_axis_contract.py b/tests/unit/torch/puzzletron/test_vllm_axis_contract.py index 08acb053a15..bd76b62c41f 100644 --- a/tests/unit/torch/puzzletron/test_vllm_axis_contract.py +++ b/tests/unit/torch/puzzletron/test_vllm_axis_contract.py @@ -1,3 +1,5 @@ +"""Tests for Puzzletron vLLM geometry and checkpoint interchange contracts.""" + from types import SimpleNamespace from modelopt.torch.puzzletron.block_config import ( @@ -8,9 +10,38 @@ MoEConfig, ) from modelopt.torch.puzzletron.candidates import build_candidate_library +from modelopt.torch.puzzletron.utils import vllm_adapter from modelopt.torch.puzzletron.utils.vllm_adapter import convert_block_configs_to_per_layer_config +def test_checkpoint_config_refresh_does_not_trust_remote_code_by_default( + tmp_path, monkeypatch +) -> None: + observed = [] + config = SimpleNamespace(to_json_string=lambda **_kwargs: "{}\n") + + def load_config(path, **kwargs): + observed.append((path, kwargs)) + return config + + monkeypatch.setattr("transformers.AutoConfig.from_pretrained", load_config) + monkeypatch.setattr( + "modelopt.torch.puzzletron.anymodel.registry.resolve_descriptor", + lambda _config: SimpleNamespace(descriptor=object()), + ) + monkeypatch.setattr(vllm_adapter, "configure_anymodel_metadata", lambda *_args: True) + monkeypatch.setattr( + vllm_adapter, + "convert_block_configs_to_per_layer_config", + lambda *_args, **_kwargs: True, + ) + + config_path = vllm_adapter.refresh_realized_checkpoint_config(tmp_path) + + assert observed == [(tmp_path, {"trust_remote_code": False})] + assert config_path.read_text() == "{}\n" + + def test_mla_search_axes_create_cartesian_typed_candidates() -> None: teacher = BlockConfig( subblock_configs=(MLAConfig(num_heads=16, q_lora_rank=768, kv_lora_rank=512),) @@ -285,12 +316,10 @@ def test_qwen35_moe_descriptor_exposes_bounded_runtime_benchmark_contract() -> N flat_text_config = SimpleNamespace(hidden_size=2048) nested_vlm_config = SimpleNamespace(text_config=flat_text_config) assert ( - Qwen3P5MoeVLModelDescriptor.get_language_model_config(flat_text_config) - is flat_text_config + Qwen3P5MoeVLModelDescriptor.get_language_model_config(flat_text_config) is flat_text_config ) assert ( - Qwen3P5MoeVLModelDescriptor.get_language_model_config(nested_vlm_config) - is flat_text_config + Qwen3P5MoeVLModelDescriptor.get_language_model_config(nested_vlm_config) is flat_text_config ) base = Qwen3P5MoeTextModelDescriptor.runtime_benchmark_base_block_config(runtime) assert base.require_subblock("attention").num_query_heads == 16 diff --git a/tests/unit/torch/puzzletron/test_width_scenarios.py b/tests/unit/torch/puzzletron/test_width_scenarios.py index fa2f6afdcf6..58307c0c0c5 100644 --- a/tests/unit/torch/puzzletron/test_width_scenarios.py +++ b/tests/unit/torch/puzzletron/test_width_scenarios.py @@ -120,7 +120,6 @@ def publish_report(config): assert not replacement_finalizer.finalization_marker_is_current( marker_a, manifest_path, summary ) - summary.write_text(summary_payload) summary.write_text(json.dumps({"scenario_count": 2, "widths": [256]})) assert not replacement_finalizer.finalization_marker_is_current( marker_a, manifest_path, summary @@ -151,6 +150,26 @@ def publish_report(config): ) +@pytest.mark.parametrize( + "payload", + [[], "manifest", 1, None], + ids=("list", "string", "integer", "null"), +) +def test_replacement_scoring_marker_rejects_non_object_manifests(tmp_path, payload): + marker = tmp_path / "finalized" + marker.write_text("replacement_scoring_semantic_test\n") + manifest = tmp_path / "manifest.json" + manifest.write_text(json.dumps(payload)) + summary = tmp_path / "summary.json" + summary.write_text("{}\n") + + assert not replacement_finalizer.finalization_marker_is_current( + marker, + manifest, + summary, + ) + + def _write_scenario_manifest( puzzle_dir: Path, width: int, From a667d9d30393637abace2a5bfad41ed7fd92640f Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 15:40:00 +0200 Subject: [PATCH 15/24] Avoid Puzzletron post-MIP import cycle Load post-MIP execution identity helpers only when dynamic-stage currentness or planning needs them, keeping normal Puzzletron package initialization acyclic. Signed-off-by: Johannes Rausch --- .../orchestration/adapters/post_mip.py | 36 +++++++++---------- .../orchestration/adapters/stage_compat.py | 9 +++-- .../test_orchestration_lightweight.py | 17 +++++++++ 3 files changed, 41 insertions(+), 21 deletions(-) diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index cd320740888..76d2c5c49aa 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -38,24 +38,19 @@ from .packing import packed_allocation from .stage_compat import _hf_checkpoint_is_complete, post_mip_summary_is_current -if __package__.startswith("puzzletron_orchestrator."): - from puzzletron_orchestrator.post_mip.identity import ( - PostMIPExecutionContractUnavailable, - expected_post_mip_candidate_count, - expected_post_mip_execution_contract, - prepare_post_mip_candidate_ledger, - ) -else: - from ...post_mip.identity import ( - PostMIPExecutionContractUnavailable, - expected_post_mip_candidate_count, - expected_post_mip_execution_contract, - prepare_post_mip_candidate_ledger, - ) - __all__ = ["ManualInputRequired", "PostMIPAdapter"] +def _post_mip_identity_api() -> Any: + """Load the producer identity contract after orchestration initialization.""" + + if __package__.startswith("puzzletron_orchestrator."): + from puzzletron_orchestrator.post_mip import identity as identity_api + else: + from ...post_mip import identity as identity_api + return identity_api + + class ManualInputRequired(RuntimeError): """A durable manual-filter review exists and needs a user decision.""" @@ -94,9 +89,10 @@ def _identity_config(plan: CampaignPlan) -> dict[str, Any]: def _available_evaluation_candidates(plan: CampaignPlan, stage_id: str) -> int | None: + identity_api = _post_mip_identity_api() try: - return expected_post_mip_candidate_count(_identity_config(plan), stage_id) - except PostMIPExecutionContractUnavailable: + return identity_api.expected_post_mip_candidate_count(_identity_config(plan), stage_id) + except identity_api.PostMIPExecutionContractUnavailable: return None @@ -125,7 +121,7 @@ def prepare_execution_identity_projection( """Prepare the candidate registry only on the attempt-submission path.""" del node - prepare_post_mip_candidate_ledger(_identity_config(plan)) + _post_mip_identity_api().prepare_post_mip_candidate_ledger(_identity_config(plan)) def execution_identity_projection( self, @@ -137,7 +133,9 @@ def execution_identity_projection( """Bind scheduler attempts to the canonical producer execution contract.""" del work_plan - return expected_post_mip_execution_contract(_identity_config(plan), node.stage_id) + return _post_mip_identity_api().expected_post_mip_execution_contract( + _identity_config(plan), node.stage_id + ) def plan(self, plan: CampaignPlan, node: StagePlanNode) -> WorkPlan: config = _node_config(plan, node.stage_id) diff --git a/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py b/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py index a63a324af4e..cbadfbd241b 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py @@ -50,13 +50,11 @@ stage_manifest_uses_execution_record, validate_stage_execution_record, ) - from puzzletron_orchestrator.post_mip.identity import expected_post_mip_execution_identity else: from ...execution_record import ( stage_manifest_uses_execution_record, validate_stage_execution_record, ) - from ...post_mip.identity import expected_post_mip_execution_identity __all__ = [ "StageCompatAdapter", @@ -465,6 +463,13 @@ def post_mip_summary_is_current( """Validate a node summary without importing the PyTorch-backed worker package.""" try: + if __package__.startswith("puzzletron_orchestrator."): + from puzzletron_orchestrator.post_mip.identity import ( + expected_post_mip_execution_identity, + ) + else: + from ...post_mip.identity import expected_post_mip_execution_identity + effective_config = dict(config) effective_config["puzzle_dir"] = str(puzzle_dir) return summary.get("execution_identity") == expected_post_mip_execution_identity( diff --git a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py index 2adc3b957c5..eb35846bde4 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py +++ b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py @@ -55,6 +55,23 @@ def test_lightweight_package_does_not_import_torch() -> None: assert result.returncode == 0, result.stderr +def test_pipeline_config_import_does_not_cycle_through_post_mip() -> None: + result = subprocess.run( + [ + sys.executable, + "-c", + "from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path; " + "assert callable(pipeline_config_from_path)", + ], + cwd=REPOSITORY_ROOT, + capture_output=True, + text=True, + check=False, + timeout=30, + ) + assert result.returncode == 0, result.stderr + + def test_named_vllm_measurement_gpu_group_includes_data_parallelism() -> None: from puzzletron_orchestrator.vllm_measurements import normalize_vllm_measurements From c48615940c831e8a71a8cc79c575a61208664ea2 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 17:05:42 +0200 Subject: [PATCH 16/24] Format Puzzletron GPU lifecycle test Signed-off-by: Johannes Rausch --- tests/gpu/torch/puzzletron/test_puzzletron.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index d017c6c8104..6893622b1f2 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -378,9 +378,10 @@ def _assert_post_mip_and_final_checkpoint( parent = ledger.revisions[revision.parent_revision_id] block_configs = _json(checkpoint / "config.json")["block_configs"] assert block_configs - assert block_configs == _json(Path(parent.artifact["checkpoint"]) / "config.json")[ - "block_configs" - ] + assert ( + block_configs + == _json(Path(parent.artifact["checkpoint"]) / "config.json")["block_configs"] + ) kd_paths.extend( [ summary_path, @@ -478,9 +479,7 @@ def test_tiny_qwen_campaign_uses_current_public_route( widths = _nested_values(block_config, "intermediate_size") assert len(widths) == 1 selected_ffn_widths.append(int(widths[0])) - per_layer_config = (selected_config.get("text_config") or selected_config)[ - "per_layer_config" - ] + per_layer_config = (selected_config.get("text_config") or selected_config)["per_layer_config"] assert [ int( per_layer_config.get(str(index), {}).get( From 8baa662ed9511b54553555169750821041c62914 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 16:52:00 +0200 Subject: [PATCH 17/24] Fix Puzzletron quality checks Signed-off-by: Johannes Rausch (cherry picked from commit 9bf94f22640b155d0f71a78a8afca9b255e582e4) --- .../puzzletron/distillation/global_kd_recipe.py | 15 +++++++++++++++ .../puzzletron/orchestration/adapters/base.py | 12 ++++++++++++ .../orchestration/adapters/post_mip.py | 7 +++++-- .../orchestration/adapters/sharded.py | 17 +++++++++++++++-- .../puzzletron/orchestration/task_launcher.py | 3 ++- modelopt/torch/puzzletron/post_mip/identity.py | 12 ++++++++++++ tests/gpu/torch/puzzletron/test_puzzletron.py | 4 ++-- .../puzzletron/test_global_kd_canonical.py | 12 ++++++++++++ .../test_post_mip_execution_identity.py | 12 ++++++++++++ .../torch/puzzletron/test_vllm_axis_contract.py | 15 +++++++++++++++ 10 files changed, 102 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py index 5d76e630dbf..419e3199932 100644 --- a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py +++ b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py @@ -1,3 +1,18 @@ +# 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. + # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/modelopt/torch/puzzletron/orchestration/adapters/base.py b/modelopt/torch/puzzletron/orchestration/adapters/base.py index bdd172b95e9..080aeaee5d8 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/base.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/base.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. """WorkAdapter contract for stage orchestration.""" diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index 76d2c5c49aa..754e0a31bf3 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -18,7 +18,9 @@ from __future__ import annotations import json -import subprocess + +# Aggregation uses an explicit argv without a shell. +import subprocess # nosec B404 from pathlib import Path from typing import Any @@ -261,7 +263,8 @@ def aggregate( ] for override in plan.overrides: argv.extend(["--override", override]) - result = subprocess.run( + # The argv starts with the fixed Python entry point and never uses a shell. + result = subprocess.run( # nosec B603 argv, cwd=repo, capture_output=True, diff --git a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py index 432cf055874..c69ac1eace0 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py @@ -1,11 +1,24 @@ # 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 +# Commands are compiled argv lists and never use a shell. +import subprocess # nosec B404 import time import uuid from dataclasses import replace @@ -357,7 +370,7 @@ def aggregate( command=command, ) else: - result = subprocess.run( + result = subprocess.run( # nosec B603 command, cwd=plan.runner.contract.repository, capture_output=True, diff --git a/modelopt/torch/puzzletron/orchestration/task_launcher.py b/modelopt/torch/puzzletron/orchestration/task_launcher.py index e61a3ebaa58..53bca458027 100644 --- a/modelopt/torch/puzzletron/orchestration/task_launcher.py +++ b/modelopt/torch/puzzletron/orchestration/task_launcher.py @@ -272,7 +272,8 @@ def main(argv: Sequence[str] | None = None) -> int: binding=binding, gpus_per_task=args.gpus_per_task, ) - os.execvpe(command[0], command, env) + # Replace this launcher with the already compiled task argv. + os.execvpe(command[0], command, env) # nosec B606 return 0 diff --git a/modelopt/torch/puzzletron/post_mip/identity.py b/modelopt/torch/puzzletron/post_mip/identity.py index e21ca4fab4b..d9a24cce00e 100644 --- a/modelopt/torch/puzzletron/post_mip/identity.py +++ b/modelopt/torch/puzzletron/post_mip/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. """Canonical execution identities for post-MIP nodes.""" diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index 6893622b1f2..00d0ff6821e 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -135,7 +135,7 @@ def _assert_pruning_and_mip_artifacts(campaign: TinyQwenCampaign) -> list[Path]: score_tensors = [ tensor for score_file in score_files - for tensor in _tensor_values(torch.load(score_file, map_location="cpu", weights_only=False)) + for tensor in _tensor_values(torch.load(score_file, map_location="cpu", weights_only=True)) ] assert score_tensors assert all(tensor.numel() and torch.isfinite(tensor).all() for tensor in score_tensors) @@ -510,7 +510,7 @@ def test_tiny_qwen_campaign_uses_current_public_route( } durable_paths = [*manifest_paths, *pruning_paths, *post_paths] before_resume = _artifact_digests(durable_paths) - resumed = campaign.run() + resumed = campaign.run(timeout=300) resumed_result = campaign.require_success(resumed) assert tuple(resumed_result["completed"]) == stage_ids assert resumed_result["failed_stages"] == [] diff --git a/tests/unit/torch/puzzletron/test_global_kd_canonical.py b/tests/unit/torch/puzzletron/test_global_kd_canonical.py index 19c1fe79476..8897c6961d2 100644 --- a/tests/unit/torch/puzzletron/test_global_kd_canonical.py +++ b/tests/unit/torch/puzzletron/test_global_kd_canonical.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. """Tests for Puzzletron's canonical global-distillation behavior.""" diff --git a/tests/unit/torch/puzzletron/test_post_mip_execution_identity.py b/tests/unit/torch/puzzletron/test_post_mip_execution_identity.py index 292d9d48051..6de5e206c21 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_execution_identity.py +++ b/tests/unit/torch/puzzletron/test_post_mip_execution_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. """Tests for post-MIP controller execution identities.""" diff --git a/tests/unit/torch/puzzletron/test_vllm_axis_contract.py b/tests/unit/torch/puzzletron/test_vllm_axis_contract.py index bd76b62c41f..efa5cc77d3c 100644 --- a/tests/unit/torch/puzzletron/test_vllm_axis_contract.py +++ b/tests/unit/torch/puzzletron/test_vllm_axis_contract.py @@ -1,3 +1,18 @@ +# 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 vLLM geometry and checkpoint interchange contracts.""" from types import SimpleNamespace From 8c2ac0b8217b24acd9affc94d1648183167f5f80 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 17:39:22 +0200 Subject: [PATCH 18/24] Fix Puzzletron changed-file typing Signed-off-by: Johannes Rausch (cherry picked from commit 6065146820736e918069019b8b41945926d889c1) --- .../finalize_replacement_scoring.py | 17 +++-- examples/puzzletron/main.py | 18 +++-- .../distillation/global_kd_recipe.py | 70 +++++++++++++------ .../distributed_eval/automodel_executor.py | 69 +++++++++++------- modelopt/torch/puzzletron/manifest.py | 10 +-- .../puzzletron/orchestration/adapters/pool.py | 14 ++-- .../puzzletron/orchestration/compiler.py | 2 +- modelopt/torch/puzzletron/stages/future.py | 26 ++++--- puzzletron_setup/bundle.py | 2 +- pyproject.toml | 1 + 10 files changed, 145 insertions(+), 84 deletions(-) diff --git a/examples/puzzletron/finalize_replacement_scoring.py b/examples/puzzletron/finalize_replacement_scoring.py index 71e08142af9..44a7a6ece4e 100644 --- a/examples/puzzletron/finalize_replacement_scoring.py +++ b/examples/puzzletron/finalize_replacement_scoring.py @@ -23,15 +23,22 @@ import os from pathlib import Path -if __package__: - from .embedding_pipeline import finalize_replacement_scoring_diagnostics -else: - from embedding_pipeline import finalize_replacement_scoring_diagnostics - from modelopt.torch.puzzletron.diagnostics import generate_replace_block_report from modelopt.torch.puzzletron.manifest import stage_manifest_from_config, write_stage_manifest from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path + +def finalize_replacement_scoring_diagnostics(config: dict): + """Preserve the finalizer seam across package and script entry points.""" + if __package__: + from .embedding_pipeline import finalize_replacement_scoring_diagnostics as package_finalize + + return package_finalize(config) + from embedding_pipeline import finalize_replacement_scoring_diagnostics as script_finalize + + return script_finalize(config) + + __all__ = [ "finalization_marker_is_current", "finalize_replacement_scoring", diff --git a/examples/puzzletron/main.py b/examples/puzzletron/main.py index 35c08aa599f..d7c60c6ab2b 100644 --- a/examples/puzzletron/main.py +++ b/examples/puzzletron/main.py @@ -415,6 +415,17 @@ def _run_embedding_stage( ) +def _run_tokenize_data_stage(config: dict): + """Run tokenization from either the package or standalone entry point.""" + if __package__: + from .tokenize_data import tokenize_data_stage as package_tokenize_data_stage + + return package_tokenize_data_stage(config) + from tokenize_data import tokenize_data_stage as script_tokenize_data_stage + + return script_tokenize_data_stage(config) + + def _run_worker(args: argparse.Namespace) -> None: cfg = mtpz.pipeline_config.pipeline_config_from_path( args.config, @@ -431,12 +442,7 @@ def _run_worker(args: argparse.Namespace) -> None: if not _stage_enabled(cfg, args.worker_stage): result = mtpz.stage_runner.run_stage(cfg, args.worker_stage, handlers={}) elif args.worker_stage == "tokenize_data": - if __package__: - from .tokenize_data import tokenize_data_stage - else: - from tokenize_data import tokenize_data_stage - - result = tokenize_data_stage(cfg) + result = _run_tokenize_data_stage(cfg) elif embedding_root and args.worker_stage in composite_only: outputs = _run_embedding_stage( config_path=args.config, diff --git a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py index 419e3199932..cbd57993690 100644 --- a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py +++ b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py @@ -33,7 +33,7 @@ import types from collections import deque from contextlib import nullcontext -from typing import Any +from typing import Any, Callable import torch from nemo_automodel.components.distributed.config import DistributedSetup @@ -205,7 +205,7 @@ def non_strict(*args, _original=original, _name=name, options=None, **kwargs): ) return _original(*args, options=relaxed(options), **kwargs) - non_strict._puzzletron_pp_non_strict = True + setattr(non_strict, "_puzzletron_pp_non_strict", True) setattr(stateful_wrappers, name, non_strict) @@ -233,23 +233,23 @@ def _attach_global_kd_gdn_traces(parts, *, prefix: str, trace_backward: bool = F def _forward_end(_module, _args, output, *, layer_idx=layer_idx): _trace_global_kd_phase(f"{prefix}_gdn_{layer_idx}_forward_end") if trace_backward and isinstance(output, torch.Tensor) and output.requires_grad: - output.register_hook( - lambda grad, layer_idx=layer_idx: ( - _trace_global_kd_phase(f"{prefix}_gdn_{layer_idx}_backward_begin"), - grad, - )[1] - ) + + def trace_backward_begin(grad, *, layer_idx=layer_idx): + _trace_global_kd_phase(f"{prefix}_gdn_{layer_idx}_backward_begin") + return grad + + output.register_hook(trace_backward_begin) module.register_forward_hook(_forward_end) if trace_backward: parameter = next(module.parameters(), None) if parameter is not None and parameter.requires_grad: - parameter.register_hook( - lambda grad, layer_idx=layer_idx: ( - _trace_global_kd_phase(f"{prefix}_gdn_{layer_idx}_parameter_grad"), - grad, - )[1] - ) + + def trace_parameter_grad(grad, *, layer_idx=layer_idx): + _trace_global_kd_phase(f"{prefix}_gdn_{layer_idx}_parameter_grad") + return grad + + parameter.register_hook(trace_parameter_grad) def _instantiate(node): @@ -593,6 +593,23 @@ def _split_output(output, model): class _WeightedObjectiveMixin: + """Objective behavior shared by the LLM and VLM AutoModel recipe bases.""" + + cfg: Any + checkpointer: Any + device_mesh: Any + dist_env: Any + loss_fn: Any + metric_logger_train: Any + model_parts: Any + optimizer: Any + pp: Any + pp_enabled: bool + teacher_model: Any + _ce_loss_buffer: list[torch.Tensor] + _kd_loss_buffer: list[torch.Tensor] + _dp_allreduce: Callable[..., torch.Tensor] + def _configure_objective(self): objective = self.cfg.get("objective", {}) self.objective = { @@ -634,7 +651,7 @@ def save_checkpoint( ): """Publish a completion marker only after model and optimizer DCP succeed.""" - result = super().save_checkpoint( + result = super().save_checkpoint( # type: ignore[misc] epoch, step, train_loss, @@ -716,20 +733,23 @@ def observability_metadata(self): } if not torch.distributed.is_initialized(): return local - gathered = [None] * torch.distributed.get_world_size() + gathered: list[dict[str, Any] | None] = [None] * torch.distributed.get_world_size() torch.distributed.all_gather_object(gathered, local) - roles = set().union(*(item["vision_by_role"] for item in gathered)) + observations = [item for item in gathered if item is not None] + if len(observations) != len(gathered): + raise RuntimeError("Missing global KD observability metadata from a distributed rank") + roles = set().union(*(item["vision_by_role"] for item in observations)) return { - "vision_forward_count": sum(item["vision_forward_count"] for item in gathered), + "vision_forward_count": sum(item["vision_forward_count"] for item in observations), "vision_by_role": { - role: sum(item["vision_by_role"].get(role, 0) for item in gathered) + role: sum(item["vision_by_role"].get(role, 0) for item in observations) for role in sorted(roles) }, "vision_output_checksums": sorted( - checksum for item in gathered for checksum in item["vision_output_checksums"] + checksum for item in observations for checksum in item["vision_output_checksums"] ), "media_input_checksums": sorted( - set(checksum for item in gathered for checksum in item["media_input_checksums"]) + set(checksum for item in observations for checksum in item["media_input_checksums"]) ), } @@ -812,7 +832,7 @@ def load_checkpoint(self, restore_from=None): return None if getattr(self, "_puzzletron_global_kd_domain", None) == "llm": self._remove_text_inactive_optimizer_parameters() - return super().load_checkpoint(restore_from or "LATEST") + return super().load_checkpoint(restore_from or "LATEST") # type: ignore[misc] def _install_gradient_norm_observers(self): self._gradient_squared = { @@ -1219,8 +1239,12 @@ def _chunk_objectives(s_chunk, t_chunk, chunk_labels): phase = f"mtp_depth_{depth}_chunk_{start}_{stop}" _trace_global_kd_phase(f"{phase}_student_head_begin") if student_is_hidden: + if student_head is None: + raise RuntimeError("MTP hidden-state projection is missing its lm_head") s_chunk = _align_dtensor_to_module_mesh(s_chunk, student_head) - s_logits = student_head(s_chunk) if student_is_hidden else s_chunk + s_logits = student_head(s_chunk) + else: + s_logits = s_chunk _trace_global_kd_phase(f"{phase}_student_head_end") zero = self._local_zero(s_logits) _trace_global_kd_phase(f"{phase}_ce_begin") diff --git a/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py b/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py index 082952f158e..c7083140671 100644 --- a/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py +++ b/modelopt/torch/puzzletron/distributed_eval/automodel_executor.py @@ -20,6 +20,7 @@ import time from contextlib import ExitStack, nullcontext from pathlib import Path +from typing import Any from ..anymodel.model_descriptor import ModelDescriptorFactory from ..anymodel.registry import resolve_descriptor_from_pretrained @@ -49,19 +50,19 @@ class AutoModelReplaceBlockExecutor: def __init__(self, hydra_cfg): self.cfg = hydra_cfg - self.recipe = None - self.cache = None - self.params = None - self.teacher_block_configs = None - self.num_q = None - self.head_dim = None - self.bypass_checkpoint_dir = None + self.recipe: Any | None = None + self.cache: Any | None = None + self.params: dict[str, Any] | None = None + self.teacher_block_configs: Any | None = None + self.num_q: int | None = None + self.head_dim: int | None = None + self.bypass_checkpoint_dir: Path | None = None self.is_output_writer = False - self.source_hidden_width = None - self.sliced_teacher_baseline = None - self.latest_observability = None - self.latest_score_device_type = None - self.visible_cuda_device_count = None + self.source_hidden_width: int | None = None + self.sliced_teacher_baseline: dict[str, Any] | None = None + self.latest_observability: dict[str, Any] | None = None + self.latest_score_device_type: str | None = None + self.visible_cuda_device_count: int | None = None self._setup_complete = False def capabilities(self) -> dict: @@ -91,7 +92,8 @@ def setup(self) -> None: from ..tools.checkpoint_utils import load_model_config scoring = self.cfg.scoring - self.params = solution_scoring_params(self.cfg) + params = solution_scoring_params(self.cfg) + self.params = params apply_patch() teacher_dir = Path( scoring.get("teacher_dir", None) or f"{self.cfg.puzzle_dir}/ckpts/teacher" @@ -132,18 +134,19 @@ def setup(self) -> None: recipe_dict = build_solution_recipe_config(self.cfg, target_dir) distributed = recipe_dict.get("distributed", {}) validate_force_hf_ep( - self.params["force_hf"], + params["force_hf"], int(distributed.get("ep_size", 1) or 1), ) target_recipe = _run_recipe( recipe_dict, scoring, - self.params["eval_iters"], - self.params["use_puzzletron_dataloader"], - self.params["data_cfg"], + params["eval_iters"], + params["use_puzzletron_dataloader"], + params["data_cfg"], ) - self.cache = TeacherTargetCache(device=self.params["teacher_cache_device"]) - _extract_teacher_targets(target_recipe, self.cache, self.params) + cache = TeacherTargetCache(device=params["teacher_cache_device"]) + self.cache = cache + _extract_teacher_targets(target_recipe, cache, params) dist.barrier() if source_dir.resolve() == target_dir.resolve(): @@ -154,9 +157,9 @@ def setup(self) -> None: self.recipe = _run_recipe( build_solution_recipe_config(self.cfg, source_dir), scoring, - self.params["eval_iters"], - self.params["use_puzzletron_dataloader"], - self.params["data_cfg"], + params["eval_iters"], + params["use_puzzletron_dataloader"], + params["data_cfg"], ) # AutoModel PP containers can retain a final norm/LM head on a rank # where the pipeline stage does not actually execute them. Elect from @@ -166,12 +169,15 @@ def setup(self) -> None: import torch.distributed as torch_dist rank = torch_dist.get_rank() if torch_dist.is_initialized() else 0 - observed = bool(len(self.cache)) - observed_by_rank = [(rank, observed)] + observed = bool(len(cache)) + observed_by_rank: list[tuple[int, bool] | None] = [(rank, observed)] if torch_dist.is_initialized(): observed_by_rank = [None] * torch_dist.get_world_size() torch_dist.all_gather_object(observed_by_rank, (rank, observed)) - output_ranks = [item_rank for item_rank, has_capture in observed_by_rank if has_capture] + observations = [item for item in observed_by_rank if item is not None] + if len(observations) != len(observed_by_rank): + raise RuntimeError("Missing AutoModel output-rank observation from a distributed rank") + output_ranks = [item_rank for item_rank, has_capture in observations if has_capture] if not output_ranks: raise RuntimeError("No AutoModel rank captured teacher final hidden states") self.is_output_writer = observed and rank == min(output_ranks) @@ -189,6 +195,9 @@ def evaluate(self, request: EvaluationRequest) -> EvaluationResult | None: raise NotImplementedError(f"Unsupported evaluation handler {request.handler!r}") if not self._setup_complete: raise RuntimeError("AutoModelReplaceBlockExecutor.setup() was not called") + params = self.params + if params is None: + raise RuntimeError("AutoModelReplaceBlockExecutor setup state is incomplete") from ..plugins.automodel.solution_launch import _solution_prune_target from ..replacement_library.replacement_utils import parse_layer_replacement @@ -260,7 +269,7 @@ def evaluate(self, request: EvaluationRequest) -> EvaluationResult | None: provenance={ "handler": request.handler, "evaluator_revision": request.evaluator_revision, - "micro_batch_size": self.params.get("micro_batch_size"), + "micro_batch_size": params.get("micro_batch_size"), "hidden_width": self.source_hidden_width, "sliced_teacher_baseline": self.sliced_teacher_baseline, "observability": self.latest_observability, @@ -285,11 +294,17 @@ def _score(self, prune_target: dict | list[dict] | None) -> dict | None: recipe = self.recipe cache = self.cache params = self.params + if recipe is None or cache is None or params is None: + raise RuntimeError("AutoModelReplaceBlockExecutor setup state is incomplete") self.visible_cuda_device_count = torch.cuda.device_count() per_batch = [] tp_group = recipe.tensor_parallel_group() candidate_lm_head = recipe.lm_head_weight() if recipe.has_outputs else None - raw_targets = prune_target if isinstance(prune_target, list) else [prune_target] + raw_targets: list[dict | None] = [] + if isinstance(prune_target, list): + raw_targets.extend(prune_target) + else: + raw_targets.append(prune_target) prune_targets = [dict(target) for target in raw_targets if target is not None] layer_indices = [int(target["layer_idx"]) for target in prune_targets] owned_layers = [ diff --git a/modelopt/torch/puzzletron/manifest.py b/modelopt/torch/puzzletron/manifest.py index 8082c9b4aaf..0b0bfd0eacb 100644 --- a/modelopt/torch/puzzletron/manifest.py +++ b/modelopt/torch/puzzletron/manifest.py @@ -118,12 +118,12 @@ def write_stage_execution_record( ) -> dict[str, Any]: """Persist immutable resolved config and artifact metadata for one execution.""" - manifest_path = _path_without_symlinks(Path(manifest_path), description="stage manifest path") - if manifest_path.parent.name != "manifests": + manifest_file = _path_without_symlinks(Path(manifest_path), description="stage manifest path") + if manifest_file.parent.name != "manifests": raise ValueError( - f"stage manifest must use the campaign manifests directory: {manifest_path}" + f"stage manifest must use the campaign manifests directory: {manifest_file}" ) - root = manifest_path.parent.parent + root = manifest_file.parent.parent stage = str(manifest_payload.get("stage") or "") if not stage or Path(stage).name != stage or stage in {".", ".."}: raise ValueError(f"invalid stage identifier for execution record: {stage!r}") @@ -235,7 +235,7 @@ def write_stage_execution_record( "sha256": resolved_sha256, }, "stage_manifest": { - "path": _portable_relative_path(manifest_path, root), + "path": _portable_relative_path(manifest_file, root), "semantic_identity": manifest_payload.get("semantic_identity"), }, "artifact_contract": "stage-manifest-output-pointers/v1", diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index a1eacd356cc..8f392a224fd 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/pool.py @@ -67,6 +67,14 @@ def _replacement_puzzle_dir(plan: CampaignPlan, width: int | None) -> Path: return plan.puzzle_dir / "scenarios" / f"width-{int(width):04d}" / "depth-00" +def _replacement_work_id(stage_id: str, width: int | None, width_count: int) -> str: + if width_count == 1: + return f"{stage_id}:gang" + if width is None: + raise RuntimeError("multi-width replacement scoring requires concrete widths") + return f"{stage_id}:width-{width:04d}" + + def _replacement_environment(plan: CampaignPlan, puzzle_dir: Path) -> dict[str, str]: scoring = plan.experiment_config.get("replacement_scoring") or {} granularity = str(scoring.get("granularity", "block")) @@ -177,11 +185,7 @@ def plan(self, plan: CampaignPlan, node: StagePlanNode) -> WorkPlan: workers_per_width, remainder = divmod(node.instances, len(widths)) items = tuple( WorkItem( - work_id=( - f"{node.stage_id}:gang" - if len(widths) == 1 - else f"{node.stage_id}:width-{int(width):04d}" - ), + work_id=_replacement_work_id(node.stage_id, width, len(widths)), stage_id=node.stage_id, shard_index=index, shard_count=len(widths), diff --git a/modelopt/torch/puzzletron/orchestration/compiler.py b/modelopt/torch/puzzletron/orchestration/compiler.py index 58ac3229046..42ff0fb5a96 100644 --- a/modelopt/torch/puzzletron/orchestration/compiler.py +++ b/modelopt/torch/puzzletron/orchestration/compiler.py @@ -76,7 +76,7 @@ def _mapping(value: Any) -> dict[str, Any]: "aiperf": ExecutionStrategy.SHARDED, } -_POST_MIP_NODE_METADATA = { +_POST_MIP_NODE_METADATA: dict[str, dict[str, Any]] = { "filter": {"kind": "selector", "accepts": {"config", "checkpoint"}}, "manual_filter": {"kind": "selector", "accepts": {"config", "checkpoint"}}, "materialize": { diff --git a/modelopt/torch/puzzletron/stages/future.py b/modelopt/torch/puzzletron/stages/future.py index 64b7a7986ff..9e2166bad67 100644 --- a/modelopt/torch/puzzletron/stages/future.py +++ b/modelopt/torch/puzzletron/stages/future.py @@ -23,7 +23,7 @@ from dataclasses import asdict from pathlib import Path from queue import Queue -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Sequence from ..anymodel.model_descriptor import ModelDescriptorFactory from ..anymodel.registry import resolve_descriptor_from_pretrained @@ -197,11 +197,11 @@ def _select_evaluated_candidates( def _with_teacher_checkpoint( - teacher_dir: str | Path | None, candidates: list[tuple[str, str | Path]] + teacher_dir: str | Path | None, candidates: Sequence[tuple[str, str | Path]] ) -> list[tuple[str, str | Path]]: """Prepend the configured teacher while keeping downstream checkpoints unique.""" if teacher_dir is None: - return candidates + return list(candidates) teacher_key = str(teacher_dir) return [("teacher", teacher_dir)] + [ (name, checkpoint) @@ -302,7 +302,10 @@ def aiperf_stage(config: dict[str, Any], manifest: StageManifest): ) from ..benchmarks import run_aiperf_sweep, write_aiperf_report - puzzle_dir = Path((config.get("experiment") or {}).get("dir")) + experiment_dir = (config.get("experiment") or {}).get("dir") + if experiment_dir is None: + raise ValueError("AIPerf requires experiment.dir") + puzzle_dir = Path(experiment_dir) teacher_dir = (config.get("convert") or {}).get("teacher_dir") checkpoint_root = Path( stage_cfg.get( @@ -310,7 +313,7 @@ def aiperf_stage(config: dict[str, Any], manifest: StageManifest): puzzle_dir / "mip" / "puzzle_solutions" / "depth_tournament" / "solutions--checkpoints", ) ) - checkpoints = [] + checkpoints: list[tuple[str, str | Path]] = [] if stage_cfg.get("checkpoint_source") == "global_kd": checkpoints.extend(_scenario_grid_global_kd_checkpoints(puzzle_dir)) elif stage_cfg.get("checkpoint_source") == "scenario_grid": @@ -464,20 +467,21 @@ def evaluation_stage(config: dict[str, Any], manifest: StageManifest): puzzle_dir = Path((config.get("experiment") or {})["dir"]) configured = stage_cfg.get("checkpoints") + raw_checkpoint_entries: list[tuple[str, str | Path]] if configured: - checkpoint_entries = [(Path(path).name, Path(path)) for path in configured] + raw_checkpoint_entries = [(Path(path).name, Path(path)) for path in configured] elif stage_cfg.get("checkpoint_source") == "global_kd": - checkpoint_entries = _scenario_grid_global_kd_checkpoints(puzzle_dir) + raw_checkpoint_entries = _scenario_grid_global_kd_checkpoints(puzzle_dir) else: - checkpoint_entries = [ + raw_checkpoint_entries = [ (name, checkpoint) for name, checkpoint in _profile_solution_checkpoints( puzzle_dir, stage_cfg.get("profile_id") ) if name != "teacher" ] - if not checkpoint_entries: - checkpoint_entries = [ + if not raw_checkpoint_entries: + raw_checkpoint_entries = [ (path.parent.name, path.parent) for path in sorted( (puzzle_dir / "scenarios").glob( @@ -488,7 +492,7 @@ def evaluation_stage(config: dict[str, Any], manifest: StageManifest): teacher_dir = (config.get("convert") or {}).get("teacher_dir") checkpoint_entries = [ (name, Path(checkpoint)) - for name, checkpoint in _with_teacher_checkpoint(teacher_dir, checkpoint_entries) + for name, checkpoint in _with_teacher_checkpoint(teacher_dir, raw_checkpoint_entries) ] if len(checkpoint_entries) == 1 and teacher_dir is None: raise FileNotFoundError("exact evaluation found no scenario checkpoints") diff --git a/puzzletron_setup/bundle.py b/puzzletron_setup/bundle.py index 4713bfb8e9b..e0f33b7dc10 100644 --- a/puzzletron_setup/bundle.py +++ b/puzzletron_setup/bundle.py @@ -852,7 +852,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": { diff --git a/pyproject.toml b/pyproject.toml index a6f55d7e46c..6242d121223 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -101,6 +101,7 @@ dev-lint = [ "pre-commit==4.3.0", "ruff==0.12.11", "types-docutils==0.22.3.20260724", + "types-PyYAML==6.0.12.20260724", ] dev-docs = [ "autodoc_pydantic>=2.1.0", From f12abd17e936cc21c2adb717a05509d7a25bebbd Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 19:02:34 +0200 Subject: [PATCH 19/24] Validate Puzzletron security policy inputs Reject truthy string and numeric values across AIPerf execution routes, validate checkpoint collections, and keep the cold-import regression independent of subprocess coverage. Signed-off-by: Johannes Rausch --- examples/puzzletron/main.py | 1 + .../orchestration/adapters/sharded.py | 20 +++- modelopt/torch/puzzletron/post_mip/runner.py | 11 +- modelopt/torch/puzzletron/security_policy.py | 27 +++++ modelopt/torch/puzzletron/stages/future.py | 21 ++-- .../torch/puzzletron/test_future_stages.py | 83 +++++++++++-- .../test_orchestration_executors.py | 111 +++++++++++++----- .../test_orchestration_lightweight.py | 5 + .../torch/puzzletron/test_post_mip_runner.py | 38 ++++++ 9 files changed, 269 insertions(+), 48 deletions(-) create mode 100644 modelopt/torch/puzzletron/security_policy.py diff --git a/examples/puzzletron/main.py b/examples/puzzletron/main.py index d7c60c6ab2b..485c6d2cd6d 100644 --- a/examples/puzzletron/main.py +++ b/examples/puzzletron/main.py @@ -417,6 +417,7 @@ def _run_embedding_stage( def _run_tokenize_data_stage(config: dict): """Run tokenization from either the package or standalone entry point.""" + # The package and standalone entry points require different import paths. if __package__: from .tokenize_data import tokenize_data_stage as package_tokenize_data_stage diff --git a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py index c69ac1eace0..a6b54d5511b 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py @@ -23,6 +23,7 @@ import uuid from dataclasses import replace from pathlib import Path +from typing import TYPE_CHECKING from ..executors.slurm import SlurmExecutor from ..schema import ( @@ -43,6 +44,13 @@ from .packing import packed_allocation from .stage_compat import stage_is_complete, stage_output_patterns +if TYPE_CHECKING: + from ...security_policy import require_boolean_policy +elif __package__.startswith("puzzletron_orchestrator."): + from puzzletron_orchestrator.security_policy import require_boolean_policy +else: + from ...security_policy import require_boolean_policy + __all__ = ["ShardedStageAdapter"] _SHARDED_ENTRYPOINTS = { @@ -249,9 +257,17 @@ def command( "--output-tokens", str(aiperf.get("output_tokens", 1024)), ] - if bool(aiperf.get("trust_remote_code", model.get("trust_remote_code", False))): + trust_remote_code = require_boolean_policy( + aiperf.get("trust_remote_code", model.get("trust_remote_code", False)), + path="aiperf.trust_remote_code", + ) + allow_online_tokenizer_resolution = require_boolean_policy( + aiperf.get("allow_aiperf_v011_online_tokenizer_resolution", False), + path="aiperf.allow_aiperf_v011_online_tokenizer_resolution", + ) + if trust_remote_code: argv.append("--trust-remote-code") - if bool(aiperf.get("allow_aiperf_v011_online_tokenizer_resolution", False)): + if allow_online_tokenizer_resolution: argv.append("--allow-aiperf-v011-online-tokenizer-resolution") else: argv = ["python", str(script_path), "--config", plan.experiment_config_path] diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 1c3243ffbd2..615a2b92ef2 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -31,6 +31,7 @@ from ..evaluation import DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS, run_lmms_eval_checkpoint from ..identity import canonicalize, stable_hash +from ..security_policy import require_boolean_policy from .base import CompiledPostMIPNode, NodeKind, compile_post_mip_flows from .filters import apply_filter from .identity import ( @@ -445,11 +446,16 @@ def _aiperf( for concurrency in concurrencies } topology = dict(settings.pop("topology", {}) or {}) - trust_remote_code = bool( + trust_remote_code = require_boolean_policy( settings.pop( "trust_remote_code", (config.get("model") or {}).get("trust_remote_code", False), - ) + ), + path="post_mip.aiperf.config.trust_remote_code", + ) + allow_online_tokenizer_resolution = require_boolean_policy( + settings.pop("allow_aiperf_v011_online_tokenizer_resolution", False), + path="post_mip.aiperf.config.allow_aiperf_v011_online_tokenizer_resolution", ) gpu_ids = os.environ.get("CUDA_VISIBLE_DEVICES", "") if not gpu_ids: @@ -468,6 +474,7 @@ def _aiperf( solution_id=source.architecture_id, profile_id=node.flow_id, trust_remote_code=trust_remote_code, + allow_aiperf_v011_online_tokenizer_resolution=allow_online_tokenizer_resolution, **settings, ) metrics = {} diff --git a/modelopt/torch/puzzletron/security_policy.py b/modelopt/torch/puzzletron/security_policy.py new file mode 100644 index 00000000000..148f9548b3d --- /dev/null +++ b/modelopt/torch/puzzletron/security_policy.py @@ -0,0 +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. + +"""Validation helpers for security-sensitive Puzzletron configuration.""" + +from typing import Any + +__all__ = ["require_boolean_policy"] + + +def require_boolean_policy(value: Any, *, path: str) -> bool: + """Return a policy boolean without accepting truthy strings or numbers.""" + if not isinstance(value, bool): + raise ValueError(f"{path} must be a boolean") + return value diff --git a/modelopt/torch/puzzletron/stages/future.py b/modelopt/torch/puzzletron/stages/future.py index 9e2166bad67..6aa7e13c0c6 100644 --- a/modelopt/torch/puzzletron/stages/future.py +++ b/modelopt/torch/puzzletron/stages/future.py @@ -33,6 +33,7 @@ build_global_kd_config, run_global_kd, ) +from ..security_policy import require_boolean_policy from .common import complete_stage from .graph import StageSkipReason @@ -300,6 +301,15 @@ def aiperf_stage(config: dict[str, Any], manifest: StageManifest): skip_reason=StageSkipReason.DISABLED, message="AIPerf is disabled.", ) + model_cfg = dict(config.get("model") or {}) + trust_remote_code = require_boolean_policy( + stage_cfg.get("trust_remote_code", model_cfg.get("trust_remote_code", False)), + path="aiperf.trust_remote_code", + ) + allow_aiperf_v011_online_tokenizer_resolution = require_boolean_policy( + stage_cfg.get("allow_aiperf_v011_online_tokenizer_resolution", False), + path="aiperf.allow_aiperf_v011_online_tokenizer_resolution", + ) from ..benchmarks import run_aiperf_sweep, write_aiperf_report experiment_dir = (config.get("experiment") or {}).get("dir") @@ -385,13 +395,6 @@ def aiperf_stage(config: dict[str, Any], manifest: StageManifest): for index in range(0, len(visible), group_size) ] output_dir = Path(stage_cfg.get("output_dir", puzzle_dir / "artifacts" / "aiperf")) - model_cfg = dict(config.get("model") or {}) - trust_remote_code = bool( - stage_cfg.get("trust_remote_code", model_cfg.get("trust_remote_code", False)) - ) - allow_aiperf_v011_online_tokenizer_resolution = bool( - stage_cfg.get("allow_aiperf_v011_online_tokenizer_resolution", False) - ) work = _aiperf_checkpoint_work(checkpoints, list(stage_cfg.get("concurrency", [1, 2, 4, 8]))) pool: Queue[str] = Queue() for gpu_group in gpu_groups: @@ -453,6 +456,9 @@ def evaluation_stage(config: dict[str, Any], manifest: StageManifest): skip_reason=StageSkipReason.DISABLED, message="Zero-shot evaluation is disabled.", ) + configured = stage_cfg.get("checkpoints") + if configured is not None and not isinstance(configured, (list, tuple)): + raise ValueError("zero_shot_evaluation.checkpoints must be a list or tuple") from omegaconf import OmegaConf import modelopt.torch.utils.distributed as dist @@ -466,7 +472,6 @@ def evaluation_stage(config: dict[str, Any], manifest: StageManifest): from .pipeline import _distributed puzzle_dir = Path((config.get("experiment") or {})["dir"]) - configured = stage_cfg.get("checkpoints") raw_checkpoint_entries: list[tuple[str, str | Path]] if configured: raw_checkpoint_entries = [(Path(path).name, Path(path)) for path in configured] diff --git a/tests/unit/torch/puzzletron/test_future_stages.py b/tests/unit/torch/puzzletron/test_future_stages.py index 28a89e2fd92..1c27c85b3ac 100644 --- a/tests/unit/torch/puzzletron/test_future_stages.py +++ b/tests/unit/torch/puzzletron/test_future_stages.py @@ -1,5 +1,19 @@ # 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. + +"""Future-stage configuration and artifact-selection contracts.""" import json from pathlib import Path @@ -7,6 +21,65 @@ import pytest import torch +from modelopt.torch.puzzletron.security_policy import require_boolean_policy + + +@pytest.mark.parametrize("value", ["false", 0, 1, None], ids=["string", "zero", "one", "none"]) +def test_security_policy_rejects_non_boolean_values(value): + with pytest.raises(ValueError, match="^policy must be a boolean$"): + require_boolean_policy(value, path="policy") + + +@pytest.mark.parametrize( + ("config", "path"), + [ + ( + {"aiperf": {"enabled": True, "trust_remote_code": "false"}}, + "aiperf.trust_remote_code", + ), + ( + { + "aiperf": { + "enabled": True, + "allow_aiperf_v011_online_tokenizer_resolution": "false", + } + }, + "aiperf.allow_aiperf_v011_online_tokenizer_resolution", + ), + ( + {"aiperf": {"enabled": True}, "model": {"trust_remote_code": "false"}}, + "aiperf.trust_remote_code", + ), + ], +) +def test_aiperf_stage_rejects_non_boolean_security_policy(config, path): + from modelopt.torch.puzzletron.stages import future + + with pytest.raises(ValueError) as error: + future.aiperf_stage(config, object()) + assert str(error.value) == f"{path} must be a boolean" + + +@pytest.mark.parametrize( + "configured", + ["/checkpoint", 7, {"student": "/checkpoint"}], + ids=["string", "integer", "mapping"], +) +def test_evaluation_stage_rejects_non_list_or_tuple_checkpoints(configured): + from modelopt.torch.puzzletron.stages import future + + config = { + "zero_shot_evaluation": { + "enabled": True, + "checkpoints": configured, + } + } + with pytest.raises( + ValueError, + match=r"^zero_shot_evaluation\.checkpoints must be a list or tuple$", + ): + future.evaluation_stage(config, object()) + def test_distillation_sanity_accepts_packed_cache_without_raw_dataset(tmp_path): from modelopt.torch.puzzletron.stages.future import _distillation_dataset_source @@ -72,9 +145,7 @@ def test_evaluation_descriptor_honors_explicit_legacy_override(monkeypatch, tmp_ ) -def test_scenario_grid_kd_builds_one_isolated_config_per_realized_checkpoint( - monkeypatch, tmp_path -): +def test_scenario_grid_kd_builds_one_isolated_config_per_realized_checkpoint(monkeypatch, tmp_path): from modelopt.torch.puzzletron.stages import future puzzle_dir = tmp_path / "model" @@ -226,12 +297,8 @@ def fail_first(value): observed.append(value) raise RuntimeError("stop") - try: + with pytest.raises(RuntimeError, match="^stop$"): future._bounded_map(fail_first, range(5), max_workers=1) - except RuntimeError as error: - assert str(error) == "stop" - else: - raise AssertionError("worker failure must propagate") assert observed == [0] diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index 9bdea6d8dbd..176dcc70737 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -19,6 +19,7 @@ import time from pathlib import Path +import pytest import yaml import puzzletron_orchestrator.adapters.sharded as sharded_module @@ -51,6 +52,42 @@ ) +def _legacy_aiperf_plan( + tmp_path: Path, experiment_config: dict +) -> tuple[CampaignPlan, StagePlanNode]: + runner = RunnerEnvironment( + kind="slurm", + contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), + slurm=SlurmRunnerConfig(account="acct", partition_batch="batch"), + ) + node = StagePlanNode( + stage_id="aiperf", + strategy=ExecutionStrategy.SHARDED, + instances=1, + failure_policy=FailurePolicy.STRICT, + mesh={}, + gpus_per_instance=1, + gpus_per_node=8, + nodes=1, + total_gpus=1, + exclusive=False, + parents=("mip",), + distributed=False, + ) + return ( + CampaignPlan( + experiment_config_path=str(tmp_path / "experiment.yaml"), + puzzle_dir=tmp_path / "run", + experiment_config=experiment_config, + runner=runner, + execution_defaults={"gpus_per_node": 8}, + stages=(node,), + contract_hash="contract", + ), + node, + ) + + def test_local_executor_runs_successful_command(tmp_path: Path): executor = LocalExecutor() log_path = tmp_path / "ok.log" @@ -620,36 +657,12 @@ def test_depth_pool_uses_one_four_node_gang_allocation(tmp_path: Path): def test_legacy_aiperf_worker_receives_explicit_security_policy(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"), - ) - node = StagePlanNode( - stage_id="aiperf", - strategy=ExecutionStrategy.SHARDED, - instances=1, - failure_policy=FailurePolicy.STRICT, - mesh={}, - gpus_per_instance=1, - gpus_per_node=8, - nodes=1, - total_gpus=1, - exclusive=False, - parents=("mip",), - distributed=False, - ) - plan = CampaignPlan( - experiment_config_path=str(tmp_path / "experiment.yaml"), - puzzle_dir=tmp_path / "run", - experiment_config={ + plan, node = _legacy_aiperf_plan( + tmp_path, + { "model": {"trust_remote_code": True}, "aiperf": {"allow_aiperf_v011_online_tokenizer_resolution": True}, }, - runner=runner, - execution_defaults={"gpus_per_node": 8}, - stages=(node,), - contract_hash="contract", ) adapter = adapter_for_stage(node) work_plan = adapter.plan(plan, node) @@ -659,13 +672,55 @@ def test_legacy_aiperf_worker_receives_explicit_security_policy(tmp_path: Path): node=node, item=work_plan.items[0], attempt_id="a1", - runner=runner, + runner=plan.runner, ) assert "--trust-remote-code" in attempt.command.argv assert "--allow-aiperf-v011-online-tokenizer-resolution" in attempt.command.argv +def test_legacy_aiperf_worker_keeps_security_policies_disabled_by_default(tmp_path: Path): + plan, node = _legacy_aiperf_plan(tmp_path, {}) + adapter = adapter_for_stage(node) + default_attempt = adapter.command( + plan=plan, + node=node, + item=adapter.plan(plan, node).items[0], + attempt_id="a2", + runner=plan.runner, + ) + assert "--trust-remote-code" not in default_attempt.command.argv + assert "--allow-aiperf-v011-online-tokenizer-resolution" not in default_attempt.command.argv + + +@pytest.mark.parametrize( + ("experiment_config", "path"), + [ + ({"model": {"trust_remote_code": "false"}}, "aiperf.trust_remote_code"), + ({"aiperf": {"trust_remote_code": "false"}}, "aiperf.trust_remote_code"), + ( + {"aiperf": {"allow_aiperf_v011_online_tokenizer_resolution": "false"}}, + "aiperf.allow_aiperf_v011_online_tokenizer_resolution", + ), + ], +) +def test_legacy_aiperf_worker_rejects_non_boolean_security_policy( + tmp_path: Path, experiment_config: dict, path: str +): + plan, node = _legacy_aiperf_plan(tmp_path, experiment_config) + adapter = adapter_for_stage(node) + + with pytest.raises(ValueError) as error: + adapter.command( + plan=plan, + node=node, + item=adapter.plan(plan, node).items[0], + attempt_id="invalid", + runner=plan.runner, + ) + assert str(error.value) == f"{path} must be a boolean" + + def test_depth_pool_packs_four_two_gpu_workers_per_node(tmp_path: Path): runner = RunnerEnvironment( kind="slurm", diff --git a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py index eb35846bde4..d531255184e 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py +++ b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py @@ -56,6 +56,10 @@ def test_lightweight_package_does_not_import_torch() -> None: def test_pipeline_config_import_does_not_cycle_through_post_mip() -> None: + environment = dict(os.environ) + # This subprocess checks cold importability, not subprocess coverage collection. + environment.pop("COVERAGE_PROCESS_START", None) + environment.pop("COVERAGE_FILE", None) result = subprocess.run( [ sys.executable, @@ -64,6 +68,7 @@ def test_pipeline_config_import_does_not_cycle_through_post_mip() -> None: "assert callable(pipeline_config_from_path)", ], cwd=REPOSITORY_ROOT, + env=environment, capture_output=True, text=True, check=False, diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index f98c5379529..ed301ab5bb1 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -18,6 +18,7 @@ from pathlib import Path from types import SimpleNamespace +import pytest from omegaconf import OmegaConf import modelopt.torch.puzzletron.stages.future as future_stages @@ -244,6 +245,43 @@ def fake_run_aiperf_sweep(checkpoint, **settings): assert result["metrics"] == {} +@pytest.mark.parametrize( + ("config", "settings", "path"), + [ + ( + {"model": {"trust_remote_code": "false"}}, + {}, + "post_mip.aiperf.config.trust_remote_code", + ), + ( + {}, + {"trust_remote_code": "false"}, + "post_mip.aiperf.config.trust_remote_code", + ), + ( + {}, + {"allow_aiperf_v011_online_tokenizer_resolution": "false"}, + "post_mip.aiperf.config.allow_aiperf_v011_online_tokenizer_resolution", + ), + ], +) +def test_aiperf_rejects_non_boolean_security_policy(config, settings, path, tmp_path): + node = SimpleNamespace( + node_id="serving", + flow_id="params", + config={"config": settings}, + ) + source = SimpleNamespace( + architecture_id="architecture", + artifact={"checkpoint": str(tmp_path / "checkpoint")}, + ) + config = {"puzzle_dir": str(tmp_path), **config} + + with pytest.raises(ValueError) as error: + runner._aiperf(config, node, source, "execution") + assert str(error.value) == f"{path} must be a boolean" + + def test_downstream_evaluation_delegates_to_generic_checkpoint_evaluator(monkeypatch, tmp_path): checkpoint = tmp_path / "checkpoint" checkpoint.mkdir() From 3e5dc6c26de9c167e6a94f9b4b08ceb9a23cc577 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 19:15:36 +0200 Subject: [PATCH 20/24] Use module-scope future stage imports Keep the new configuration-boundary tests consistent with the repository import convention without changing their behavior. Signed-off-by: Johannes Rausch --- tests/unit/torch/puzzletron/test_future_stages.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/unit/torch/puzzletron/test_future_stages.py b/tests/unit/torch/puzzletron/test_future_stages.py index 1c27c85b3ac..b20150ecdb2 100644 --- a/tests/unit/torch/puzzletron/test_future_stages.py +++ b/tests/unit/torch/puzzletron/test_future_stages.py @@ -22,6 +22,7 @@ import torch from modelopt.torch.puzzletron.security_policy import require_boolean_policy +from modelopt.torch.puzzletron.stages.future import aiperf_stage, evaluation_stage @pytest.mark.parametrize("value", ["false", 0, 1, None], ids=["string", "zero", "one", "none"]) @@ -53,10 +54,8 @@ def test_security_policy_rejects_non_boolean_values(value): ], ) def test_aiperf_stage_rejects_non_boolean_security_policy(config, path): - from modelopt.torch.puzzletron.stages import future - with pytest.raises(ValueError) as error: - future.aiperf_stage(config, object()) + aiperf_stage(config, object()) assert str(error.value) == f"{path} must be a boolean" @@ -66,8 +65,6 @@ def test_aiperf_stage_rejects_non_boolean_security_policy(config, path): ids=["string", "integer", "mapping"], ) def test_evaluation_stage_rejects_non_list_or_tuple_checkpoints(configured): - from modelopt.torch.puzzletron.stages import future - config = { "zero_shot_evaluation": { "enabled": True, @@ -78,7 +75,7 @@ def test_evaluation_stage_rejects_non_list_or_tuple_checkpoints(configured): ValueError, match=r"^zero_shot_evaluation\.checkpoints must be a list or tuple$", ): - future.evaluation_stage(config, object()) + evaluation_stage(config, object()) def test_distillation_sanity_accepts_packed_cache_without_raw_dataset(tmp_path): From 630d69ebfde13788929d9f6f78b0106141c3dfcb Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Thu, 13 Aug 2026 01:45:33 +0200 Subject: [PATCH 21/24] Harden Puzzletron recovery and policy handling Reject ambiguous security-policy values and give resumed controllers a fresh artifact-settling window before recording terminal failure. Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 5 ++++ .../puzzletron/run_profile_aiperf_worker.py | 14 +++++++++- .../distillation/global_kd_recipe.py | 14 +++++++--- .../orchestration/adapters/post_mip.py | 2 +- .../orchestration/adapters/sharded.py | 10 +++++-- .../puzzletron/orchestration/controller.py | 5 +++- .../puzzletron/test_global_kd_canonical.py | 12 +++++++++ .../test_orchestration_executors.py | 2 +- .../test_orchestration_shutdown_progress.py | 26 +++++++++++++++++++ 9 files changed, 81 insertions(+), 9 deletions(-) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 611f9a272bd..fa14b3d7522 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -501,6 +501,11 @@ 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. +Remote model code and AIPerf v0.11 online tokenizer resolution are disabled by +default. Enable remote code only for a trusted model source. The tokenizer +compatibility option permits the AIPerf child process to resolve its tokenizer +online even when the surrounding campaign is configured for offline loading. + ### Legacy checked-in Nano campaign The checked-in Nano experiment uses the legacy `zero_shot_evaluation`, diff --git a/examples/puzzletron/run_profile_aiperf_worker.py b/examples/puzzletron/run_profile_aiperf_worker.py index e2cc145ef66..930813a3b61 100644 --- a/examples/puzzletron/run_profile_aiperf_worker.py +++ b/examples/puzzletron/run_profile_aiperf_worker.py @@ -191,6 +191,13 @@ def run_worker( trust_remote_code: bool = False, allow_aiperf_v011_online_tokenizer_resolution: bool = False, ) -> Path: + """Run one AIPerf shard with security-sensitive behavior disabled by default. + + ``trust_remote_code`` is only appropriate for trusted model sources. The + AIPerf v0.11 compatibility option permits online tokenizer resolution for + the AIPerf child process even when the campaign otherwise runs offline. + """ + # Worker execution needs the GPU stack; result merging intentionally remains # usable by the dependency-light login-node orchestrator. from modelopt.torch.puzzletron.benchmarks import run_aiperf_sweep @@ -341,10 +348,15 @@ def main() -> None: parser.add_argument("--concurrency", type=int, action="append", default=[]) parser.add_argument("--request-count", type=int) parser.add_argument("--benchmark-timeout", type=float, default=7200) - parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Allow remote model code; use only with trusted model sources.", + ) parser.add_argument( "--allow-aiperf-v011-online-tokenizer-resolution", action="store_true", + help="Permit online tokenizer resolution in the AIPerf v0.11 child process.", ) parser.add_argument("--preflight", action="store_true") parser.add_argument("--merge", action="store_true") diff --git a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py index cbd57993690..784cc8d640c 100644 --- a/modelopt/torch/puzzletron/distillation/global_kd_recipe.py +++ b/modelopt/torch/puzzletron/distillation/global_kd_recipe.py @@ -33,6 +33,7 @@ import types from collections import deque from contextlib import nullcontext +from pathlib import Path from typing import Any, Callable import torch @@ -77,6 +78,7 @@ from ..plugins.automodel.batch_adapter import VisionForwardMonitor from ..plugins.automodel.pp_utils import set_pp_vlm_chunk_specs +from ..security_policy import require_boolean_policy from .flash_kld import TrainingFlashKLD @@ -663,8 +665,6 @@ def save_checkpoint( f"epoch_{epoch}_step_{step}", ) if self.dist_env.is_main: - from pathlib import Path - consolidated = Path(checkpoint_path, "model", "consolidated") config_path = consolidated / "config.json" config = json.loads(config_path.read_text()) if config_path.is_file() else {} @@ -672,9 +672,17 @@ def save_checkpoint( from ..utils.vllm_adapter import refresh_realized_checkpoint_config model_config = _config_value(getattr(self, "cfg", None), "model") + configured_trust = _config_value(model_config, "trust_remote_code") refresh_realized_checkpoint_config( consolidated, - trust_remote_code=bool(_config_value(model_config, "trust_remote_code")), + trust_remote_code=( + False + if configured_trust is None + else require_boolean_policy( + configured_trust, + path="model.trust_remote_code", + ) + ), ) Path(checkpoint_path, "saving_completed").touch() if torch.distributed.is_initialized(): diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index 754e0a31bf3..0d8971c19f1 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -46,7 +46,7 @@ def _post_mip_identity_api() -> Any: """Load the producer identity contract after orchestration initialization.""" - if __package__.startswith("puzzletron_orchestrator."): + if (__package__ or "").startswith("puzzletron_orchestrator."): from puzzletron_orchestrator.post_mip import identity as identity_api else: from ...post_mip import identity as identity_api diff --git a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py index a6b54d5511b..568f4ac3b0e 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/sharded.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/sharded.py @@ -257,9 +257,15 @@ def command( "--output-tokens", str(aiperf.get("output_tokens", 1024)), ] + if "trust_remote_code" in aiperf: + trust_remote_code_path = "aiperf.trust_remote_code" + trust_remote_code_value = aiperf["trust_remote_code"] + else: + trust_remote_code_path = "model.trust_remote_code" + trust_remote_code_value = model.get("trust_remote_code", False) trust_remote_code = require_boolean_policy( - aiperf.get("trust_remote_code", model.get("trust_remote_code", False)), - path="aiperf.trust_remote_code", + trust_remote_code_value, + path=trust_remote_code_path, ) allow_online_tokenizer_resolution = require_boolean_policy( aiperf.get("allow_aiperf_v011_online_tokenizer_resolution", False), diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index 3dae60dd441..2cdff944750 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.py @@ -207,6 +207,7 @@ def __init__( self._interactive_ready = False self._failed_stages: set[str] = set() self._finalization_failures: dict[str, _FinalizationFailure] = {} + self._first_completion_observed: dict[str, float] = {} self._manual_waiting: ManualInputRequired | None = None defaults = dict(plan.execution_defaults or {}) self._halt_policy = HaltPolicy(str(defaults.get("halt_policy", HaltPolicy.DRAIN.value))) @@ -459,7 +460,9 @@ def _completed_work_artifact_settling_elapsed( if not isinstance(value, (int, float)): return _ARTIFACT_SETTLING_TIMEOUT_SECONDS completed_at.append(float(value)) - return max(0.0, time.time() - max(completed_at)) + now = time.time() + first_observed = self._first_completion_observed.setdefault(node.stage_id, now) + return max(0.0, now - max(max(completed_at), first_observed)) def _policy_allows_retry(self, node: StagePlanNode, failure: FailureClass) -> bool: if failure in {FailureClass.SUCCESS, FailureClass.CANCELLED}: diff --git a/tests/unit/torch/puzzletron/test_global_kd_canonical.py b/tests/unit/torch/puzzletron/test_global_kd_canonical.py index 8897c6961d2..b959ab97775 100644 --- a/tests/unit/torch/puzzletron/test_global_kd_canonical.py +++ b/tests/unit/torch/puzzletron/test_global_kd_canonical.py @@ -20,6 +20,7 @@ from pathlib import Path from types import SimpleNamespace +import pytest import torch from modelopt.torch.puzzletron.distillation.global_automodel import ( @@ -785,6 +786,17 @@ class Recipe(_WeightedObjectiveMixin, BaseRecipe): ] assert (tmp_path / "epoch_2_step_17" / "saving_completed").is_file() + recipe.cfg = {"model": {"trust_remote_code": "false"}} + with pytest.raises(ValueError, match=r"^model\.trust_remote_code must be a boolean$"): + recipe.save_checkpoint( + 2, + 18, + 0.5, + {"lm_loss": 0.25}, + best_metric_key="lm_loss", + ) + assert not (tmp_path / "epoch_2_step_18" / "saving_completed").exists() + def test_global_kd_optimizer_save_uses_the_actual_pipeline_model_parts(): import torch diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index 176dcc70737..d2e03c9afa7 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -696,7 +696,7 @@ def test_legacy_aiperf_worker_keeps_security_policies_disabled_by_default(tmp_pa @pytest.mark.parametrize( ("experiment_config", "path"), [ - ({"model": {"trust_remote_code": "false"}}, "aiperf.trust_remote_code"), + ({"model": {"trust_remote_code": "false"}}, "model.trust_remote_code"), ({"aiperf": {"trust_remote_code": "false"}}, "aiperf.trust_remote_code"), ( {"aiperf": {"allow_aiperf_v011_online_tokenizer_resolution": "false"}}, diff --git a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py index c0b0eb2cf49..a09497766bc 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py +++ b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py @@ -798,6 +798,32 @@ def validate(self, *, plan, node): assert controller.store.stage_is_complete("convert") +def test_resumed_controller_gets_a_fresh_artifact_settling_window(tmp_path: Path, monkeypatch): + experiment, runner_path, execution_path = _write_configs(tmp_path) + plan = compile_campaign_plan( + experiment_config_path=experiment, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + stage_filter="convert", + ) + controller = CampaignController(plan, executor=_FakeExecutor()) + node = plan.stages[0] + attempts = [{"completed_at": 1000.0}] + now = [2000.0] + monkeypatch.setattr( + controller, + "_required_completed_attempts", + lambda _node, _attempts: attempts, + ) + monkeypatch.setattr("puzzletron_orchestrator.controller.time.time", lambda: now[0]) + + assert controller._completed_work_artifact_settling_elapsed(node, attempts) == 0.0 + now[0] += 299.0 + assert controller._completed_work_artifact_settling_elapsed(node, attempts) == 299.0 + now[0] += 1.0 + assert controller._completed_work_artifact_settling_elapsed(node, attempts) == 300.0 + + def test_controller_propagates_aggregation_programming_errors(tmp_path: Path, monkeypatch): experiment, runner_path, execution_path = _write_configs(tmp_path) plan = compile_campaign_plan( From 10f02ec66a9500ae5da7ab87645db969938189b7 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Thu, 13 Aug 2026 01:52:54 +0200 Subject: [PATCH 22/24] Add non-interactive Puzzletron setup Let automation accept resolved setup defaults through the public CLI so integration tests and unattended campaigns share the same bundle-generation path. Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 15 ++++ puzzletron_setup/v2/cli.py | 32 ++++++- puzzletron_setup/v2/prompts.py | 55 ++++++++++++ puzzletron_setup/v2/wizard.py | 43 +++++++++- .../torch/puzzletron/tiny_qwen_campaign.py | 76 +++++------------ .../torch/puzzletron/test_setup_v2_quick.py | 85 +++++++++++++++---- 6 files changed, 231 insertions(+), 75 deletions(-) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index fa14b3d7522..e8ea98c4d8c 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -85,6 +85,21 @@ 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: +Automation can use the same setup entry point without answering prompts. The +defaults file must provide every required value that has no built-in default: + +```bash +python examples/puzzletron/puzzletron_setup_v2.py \ + --defaults /path/to/setup-v2-defaults.yaml \ + --campaign-dir /path/to/campaign \ + --profile smoke \ + --non-interactive +``` + +Non-interactive setup fails instead of guessing when a required answer has no +resolved default. It generates and validates the same smoke and production +bundles as the interactive wizard. + ```bash python examples/puzzletron/puzzletron_setup_v2.py --full ``` diff --git a/puzzletron_setup/v2/cli.py b/puzzletron_setup/v2/cli.py index bbce654d261..b3d7848739b 100644 --- a/puzzletron_setup/v2/cli.py +++ b/puzzletron_setup/v2/cli.py @@ -23,6 +23,9 @@ from puzzletron_setup import SetupError +from .presets import QUICK_SETUP_PRESETS +from .prompts import NonInteractiveBackend + if TYPE_CHECKING: from collections.abc import Sequence @@ -43,6 +46,22 @@ def _parser() -> argparse.ArgumentParser: type=Path, help="Explicit versioned defaults YAML; never discovered automatically.", ) + parser.add_argument( + "--campaign-dir", + type=Path, + help="Campaign directory for a new setup; bypasses that interactive prompt.", + ) + parser.add_argument( + "--profile", + choices=tuple(preset.name for preset in QUICK_SETUP_PRESETS), + default="balanced", + help="Guided setup profile for a new campaign (default: balanced).", + ) + parser.add_argument( + "--non-interactive", + action="store_true", + help="Accept resolved defaults and fail if any required answer has no default.", + ) parser.add_argument( "--full", action="store_true", @@ -56,7 +75,15 @@ def _parser() -> argparse.ArgumentParser: def main(argv: Sequence[str] | None = None) -> int: """Run the setup-v2 command-line interface.""" - args = _parser().parse_args(argv) + parser = _parser() + args = parser.parse_args(argv) + if args.resume is not None and args.campaign_dir is not None: + parser.error("--campaign-dir cannot be combined with --resume") + if args.non_interactive and args.resume is None: + if args.campaign_dir is None: + parser.error("--non-interactive requires --campaign-dir for a new campaign") + if args.defaults is None: + parser.error("--non-interactive requires --defaults for a new campaign") # Keep heavyweight model inspection out of --help and argument-error paths. from .wizard import run_wizard_v2 @@ -65,6 +92,9 @@ def main(argv: Sequence[str] | None = None) -> int: resume=args.resume, defaults_path=args.defaults, full=args.full, + campaign_dir=args.campaign_dir, + setup_profile=args.profile, + backend=NonInteractiveBackend() if args.non_interactive else None, ) except KeyboardInterrupt: target = args.resume or "" diff --git a/puzzletron_setup/v2/prompts.py b/puzzletron_setup/v2/prompts.py index 93e7386deca..8249efae40c 100644 --- a/puzzletron_setup/v2/prompts.py +++ b/puzzletron_setup/v2/prompts.py @@ -29,6 +29,7 @@ __all__ = [ "BACK", "InteractiveBackend", + "NonInteractiveBackend", "PromptBackend", "PromptChoice", "ScriptedBackend", @@ -202,6 +203,60 @@ def checkbox( return list(values) +class NonInteractiveBackend: + """Accept resolved prompt defaults without depending on prompt ordering or labels.""" + + def text(self, message: str, default: str) -> Any: + """Return a non-empty resolved text default.""" + if not str(default).strip(): + raise SetupError(f"Non-interactive setup requires a default for {message!r}.") + return default + + def select( + self, + message: str, + choices: Sequence[PromptChoice], + default: Any, + ) -> Any: + """Return the enabled resolved choice default.""" + enabled = [choice for choice in choices if choice.disabled is None] + if default is None: + if len(enabled) == 1: + return enabled[0].value + raise SetupError(f"Non-interactive setup requires a default for {message!r}.") + selected = next((choice for choice in choices if choice.value == default), None) + if selected is None: + raise SetupError( + f"Non-interactive default {default!r} is not a choice for {message!r}." + ) + if selected.disabled is not None: + raise SetupError( + f"Non-interactive default {default!r} is unavailable for {message!r}: " + f"{selected.disabled}" + ) + return selected.value + + def checkbox( + self, + message: str, + choices: Sequence[PromptChoice], + defaults: Sequence[Any], + ) -> Any: + """Return the enabled resolved checkbox defaults.""" + for value in defaults: + choice = next((item for item in choices if item.value == value), None) + if choice is None: + raise SetupError( + f"Non-interactive default {value!r} is not a choice for {message!r}." + ) + if choice.disabled is not None: + raise SetupError( + f"Non-interactive default {value!r} is unavailable for {message!r}: " + f"{choice.disabled}" + ) + return list(defaults) + + class ScriptedBackend: """Deterministic non-interactive backend for embedding and automation.""" diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index 3c0a1935178..bffda496dc3 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -321,11 +321,16 @@ def _select_model_source( resolver: DefaultsResolver, ) -> Any: while True: + explicit_default = resolver.file_default("model.source") family = session.select( "model.source_family", "Model:", _model_family_choices(resolver), - default=_CUSTOM_MODEL_SOURCE, + default=( + _DEFAULT_MODEL_SOURCE + if explicit_default is not None and explicit_default.value + else _CUSTOM_MODEL_SOURCE + ), ) if family is BACK: return BACK @@ -4482,7 +4487,23 @@ def _fresh_state( defaults_path: Path | None, *, full: bool, + campaign_dir: Path | None = None, + setup_profile: str = "balanced", ) -> WizardState: + if campaign_dir is not None: + if full: + return WizardState.start( + Path(campaign_dir).expanduser(), + defaults_path=defaults_path, + setup_mode="full", + ) + get_setup_preset(setup_profile) + return WizardState.start( + Path(campaign_dir).expanduser(), + defaults_path=defaults_path, + setup_mode="quick", + preset=setup_profile, + ) if full: while True: value = backend.text("Campaign directory:", "") @@ -4498,7 +4519,7 @@ def _fresh_state( ) while True: - preset = _select_setup_preset(backend) + preset = _select_setup_preset(backend, default=setup_profile) if preset is BACK: continue while True: @@ -4534,8 +4555,16 @@ def run_wizard_v2( defaults_path: Path | None, backend: PromptBackend | None = None, full: bool = False, + campaign_dir: Path | None = None, + setup_profile: str = "balanced", ) -> Path: - """Run setup v2, save every answer, validate bundles, and never launch jobs.""" + """Resolve setup answers, validate both bundles, and never launch jobs. + + ``campaign_dir`` and ``setup_profile`` let automation bypass only the new + campaign prompts while using the same wizard sections and bundle renderer. + """ + if resume is not None and campaign_dir is not None: + raise SetupError("campaign_dir cannot be combined with resume") backend = backend or InteractiveBackend() print("Welcome to Puzzletron setup v2.") if resume is None: @@ -4547,7 +4576,13 @@ def run_wizard_v2( "defaults from a profile." ) print(" Use --full only when you need every advanced control.") - state = _fresh_state(backend, defaults_path, full=full) + state = _fresh_state( + backend, + defaults_path, + full=full, + campaign_dir=campaign_dir, + setup_profile=setup_profile, + ) else: state = WizardState.resume(resume) if full and state.setup_mode != "full": diff --git a/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py b/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py index f7d521e9c9b..385c73532fd 100644 --- a/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py +++ b/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py @@ -34,57 +34,13 @@ load_execution_config, load_runner_config, ) -from puzzletron_setup.v2.wizard import run_wizard_v2 __all__ = ["TinyQwenCampaign", "build_tiny_qwen_campaign"] if TYPE_CHECKING: - from collections.abc import Sequence from pathlib import Path from puzzletron_orchestrator.schema import CampaignPlan - from puzzletron_setup.v2.prompts import PromptChoice - - -class _DefaultsBackend: - """Select resolved guided defaults while supplying the campaign directory.""" - - def __init__(self, campaign_dir: Path) -> None: - self.campaign_dir = campaign_dir - self.answered: set[str] = set() - - def text(self, message: str, default: str) -> Any: - if message == "Campaign directory:": - self.answered.add(message) - return str(self.campaign_dir) - return default - - def select( - self, - message: str, - choices: Sequence[PromptChoice], - default: Any, - ) -> Any: - if message in {"Model:", "Dataset:"}: - defaults = [choice.value for choice in choices if choice.title.startswith("Default —")] - if len(defaults) != 1: - raise AssertionError( - f"expected one resolved default for {message!r}, found {len(defaults)}" - ) - self.answered.add(message) - return defaults[0] - if default is not None: - return default - return next(choice.value for choice in choices if choice.disabled is None) - - def checkbox( - self, - message: str, - choices: Sequence[PromptChoice], - defaults: Sequence[Any], - ) -> Any: - del message, choices - return list(defaults) @dataclass(frozen=True) @@ -287,16 +243,30 @@ def build_tiny_qwen_campaign( ) ) - backend = _DefaultsBackend(campaign_dir) - generated = run_wizard_v2( - resume=None, - defaults_path=defaults_path, - backend=backend, + setup = subprocess.run( + [ + sys.executable, + str(project_root / "examples/puzzletron/puzzletron_setup_v2.py"), + "--defaults", + str(defaults_path), + "--campaign-dir", + str(campaign_dir), + "--profile", + "balanced", + "--non-interactive", + ], + cwd=project_root, + capture_output=True, + text=True, + check=False, ) - expected_prompts = {"Campaign directory:", "Model:", "Dataset:"} - if backend.answered != expected_prompts: - raise AssertionError(f"wizard prompt contract changed; answered {sorted(backend.answered)}") - smoke_bundle = generated / "smoke" + if setup.returncode != 0: + raise AssertionError( + "Tiny Qwen Puzzletron setup failed.\n" + f"stdout:\n{setup.stdout[-12000:]}\n" + f"stderr:\n{setup.stderr[-12000:]}" + ) + smoke_bundle = campaign_dir / "smoke" experiment = yaml.safe_load((smoke_bundle / "experiment.yaml").read_text()) flows = dict((experiment.get("post_mip") or {}).get("flows") or {}) if len(flows) != 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..00331e4c227 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_quick.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_quick.py @@ -36,6 +36,7 @@ from puzzletron_setup.v2.prompts import ( BACK, InteractiveBackend, + NonInteractiveBackend, PromptChoice, ScriptedBackend, _bind_escape_back, @@ -44,7 +45,6 @@ from puzzletron_setup.v2.state import WizardState from puzzletron_setup.v2.wizard import ( _CUSTOM_DATA_SOURCE, - _CUSTOM_MODEL_SOURCE, _PUZZLE_KD_DATA_SOURCE, _acquisition_sample_requirements, _fresh_state, @@ -364,6 +364,17 @@ def test_fresh_guided_state_records_profile_and_cli_full_is_explicit(tmp_path): assert _parser().parse_args(["--full"]).full is True +def test_non_interactive_backend_uses_semantic_defaults() -> None: + backend = NonInteractiveBackend() + choices = [PromptChoice("First", "first"), PromptChoice("Second", "second")] + + assert backend.text("Path:", "/resolved/path") == "/resolved/path" + assert backend.select("Choice:", choices, "second") == "second" + assert backend.checkbox("Choices:", choices, ["first"]) == ["first"] + with pytest.raises(SetupError, match="requires a default"): + backend.text("Path:", "") + + def test_cli_forwards_full_to_the_wizard(monkeypatch, tmp_path): captured = {} @@ -377,6 +388,37 @@ def run_wizard_v2(**kwargs): assert captured["full"] is True +def test_cli_forwards_non_interactive_campaign_contract(monkeypatch, tmp_path): + captured = {} + defaults = tmp_path / "defaults.yaml" + defaults.write_text("schema_version: 1\n") + campaign = tmp_path / "campaign" + + def run_wizard_v2(**kwargs): + captured.update(kwargs) + return campaign + + monkeypatch.setattr(wizard_module, "run_wizard_v2", run_wizard_v2) + + assert ( + cli_module.main( + [ + "--defaults", + str(defaults), + "--campaign-dir", + str(campaign), + "--profile", + "smoke", + "--non-interactive", + ] + ) + == 0 + ) + assert captured["campaign_dir"] == campaign + assert captured["setup_profile"] == "smoke" + assert isinstance(captured["backend"], NonInteractiveBackend) + + @pytest.mark.parametrize("full", [False, True]) def test_fresh_state_reprompts_for_empty_campaign_directory(tmp_path, full, capsys): campaign = tmp_path / ("full" if full else "guided") @@ -872,29 +914,38 @@ def test_guided_wizard_runs_real_sections_and_generates_valid_bundles( "infer_dataset_modality", lambda source: SimpleNamespace(modality="text", evidence="local fixture"), ) - backend = ScriptedBackend( - [ - "smoke", - str(campaign), - _CUSTOM_MODEL_SOURCE, - str(model_path), - _CUSTOM_DATA_SOURCE, - str(dataset), - "defaults", - "/worker/modelopt", - "/worker/venv", - True, - ] + defaults = tmp_path / "defaults.yaml" + defaults.write_text( + yaml.safe_dump( + { + "schema_version": 1, + "model": {"source": str(model_path)}, + "data": { + "source": str(dataset), + "modality": "text", + "layout": "fixed", + "sequence_length": 32, + }, + "infrastructure": { + "execution_contract": { + "repository": "/worker/modelopt", + "venv": "/worker/venv", + } + }, + }, + sort_keys=False, + ) ) result = wizard_module.run_wizard_v2( resume=None, - defaults_path=None, - backend=backend, + defaults_path=defaults, + backend=NonInteractiveBackend(), + campaign_dir=campaign, + setup_profile="smoke", ) assert result == campaign.resolve() - assert backend.remaining == 0 assert (campaign / "smoke" / "experiment.yaml").is_file() assert (campaign / "production" / "experiment.yaml").is_file() assert (campaign / "resolved_defaults.yaml").is_file() From 57c87e11ad7439b124f34a2dce7db24c176d0947 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Thu, 13 Aug 2026 02:49:14 +0200 Subject: [PATCH 23/24] Fix Puzzletron post-MIP test import Signed-off-by: Johannes Rausch --- tests/unit/torch/puzzletron/test_post_mip_runner.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index ed301ab5bb1..bd439197deb 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -15,6 +15,7 @@ """Tests for post-MIP execution, including managed downstream evaluation.""" +import json from pathlib import Path from types import SimpleNamespace From ba510fcb838c35c2c9ea3dad94938982e241dd6b Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Thu, 13 Aug 2026 03:13:06 +0200 Subject: [PATCH 24/24] Fix non-interactive setup defaults Preserve explicit empty defaults while failing closed on missing required values, and bound setup subprocess execution so failures surface within the test timeout. Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 2 +- puzzletron_setup/v2/prompts.py | 14 ++--- puzzletron_setup/v2/wizard.py | 4 +- .../torch/puzzletron/tiny_qwen_campaign.py | 53 +++++++++++++------ .../torch/puzzletron/test_setup_v2_quick.py | 3 +- 5 files changed, 48 insertions(+), 28 deletions(-) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index e8ea98c4d8c..43e2ed0d3c6 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -86,7 +86,7 @@ over the selected profile. To expose every per-section and nested setting, use the advanced flow explicitly: Automation can use the same setup entry point without answering prompts. The -defaults file must provide every required value that has no built-in default: +defaults file must provide every required value that has no resolved default: ```bash python examples/puzzletron/puzzletron_setup_v2.py \ diff --git a/puzzletron_setup/v2/prompts.py b/puzzletron_setup/v2/prompts.py index 8249efae40c..b6bb55a3d06 100644 --- a/puzzletron_setup/v2/prompts.py +++ b/puzzletron_setup/v2/prompts.py @@ -56,7 +56,7 @@ class PromptChoice: class PromptBackend(Protocol): """Minimal backend used by the navigable wizard session.""" - def text(self, message: str, default: str) -> Any: + def text(self, message: str, default: str | None) -> Any: """Request a text value.""" raise NotImplementedError @@ -128,10 +128,10 @@ class InteractiveBackend: _BACK_TITLE = "← Back" - def text(self, message: str, default: str) -> Any: + def text(self, message: str, default: str | None) -> Any: """Request text while supporting semantic Back navigation.""" print(" Press Esc to go back (or type :back).") - question = _bind_escape_back(_questionary().text(message, default=default)) + question = _bind_escape_back(_questionary().text(message, default=default or "")) value = _answer(question) if value is BACK: return BACK @@ -206,9 +206,9 @@ def checkbox( class NonInteractiveBackend: """Accept resolved prompt defaults without depending on prompt ordering or labels.""" - def text(self, message: str, default: str) -> Any: - """Return a non-empty resolved text default.""" - if not str(default).strip(): + def text(self, message: str, default: str | None) -> Any: + """Return a resolved text default, including an explicit empty string.""" + if default is None: raise SetupError(f"Non-interactive setup requires a default for {message!r}.") return default @@ -275,7 +275,7 @@ def _next(self) -> Any: value = self._answers.popleft() return BACK if value == ":back" else value - def text(self, message: str, default: str) -> Any: + def text(self, message: str, default: str | None) -> Any: """Return the next scripted text answer.""" del message, default return self._next() diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index bffda496dc3..e566521e280 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -4506,7 +4506,7 @@ def _fresh_state( ) if full: while True: - value = backend.text("Campaign directory:", "") + value = backend.text("Campaign directory:", None) if value is BACK: continue if not str(value).strip(): @@ -4523,7 +4523,7 @@ def _fresh_state( if preset is BACK: continue while True: - value = backend.text("Campaign directory:", "") + value = backend.text("Campaign directory:", None) if value is BACK: break if not str(value).strip(): diff --git a/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py b/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py index 385c73532fd..72e8d6fc5d6 100644 --- a/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py +++ b/tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py @@ -243,23 +243,42 @@ def build_tiny_qwen_campaign( ) ) - setup = subprocess.run( - [ - sys.executable, - str(project_root / "examples/puzzletron/puzzletron_setup_v2.py"), - "--defaults", - str(defaults_path), - "--campaign-dir", - str(campaign_dir), - "--profile", - "balanced", - "--non-interactive", - ], - cwd=project_root, - capture_output=True, - text=True, - check=False, - ) + setup_command = [ + sys.executable, + str(project_root / "examples/puzzletron/puzzletron_setup_v2.py"), + "--defaults", + str(defaults_path), + "--campaign-dir", + str(campaign_dir), + "--profile", + "balanced", + "--non-interactive", + ] + try: + setup = subprocess.run( + setup_command, + cwd=project_root, + capture_output=True, + text=True, + check=False, + timeout=300, + ) + except subprocess.TimeoutExpired as error: + stdout = ( + error.stdout.decode(errors="replace") + if isinstance(error.stdout, bytes) + else error.stdout + ) + stderr = ( + error.stderr.decode(errors="replace") + if isinstance(error.stderr, bytes) + else error.stderr + ) + raise AssertionError( + "Tiny Qwen Puzzletron setup timed out after 300 seconds.\n" + f"stdout:\n{(stdout or '')[-12000:]}\n" + f"stderr:\n{(stderr or '')[-12000:]}" + ) from error if setup.returncode != 0: raise AssertionError( "Tiny Qwen Puzzletron setup failed.\n" diff --git a/tests/unit/torch/puzzletron/test_setup_v2_quick.py b/tests/unit/torch/puzzletron/test_setup_v2_quick.py index 00331e4c227..c4561e6d7fa 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_quick.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_quick.py @@ -369,10 +369,11 @@ def test_non_interactive_backend_uses_semantic_defaults() -> None: choices = [PromptChoice("First", "first"), PromptChoice("Second", "second")] assert backend.text("Path:", "/resolved/path") == "/resolved/path" + assert backend.text("Optional commands:", "") == "" assert backend.select("Choice:", choices, "second") == "second" assert backend.checkbox("Choices:", choices, ["first"]) == ["first"] with pytest.raises(SetupError, match="requires a default"): - backend.text("Path:", "") + backend.text("Path:", None) def test_cli_forwards_full_to_the_wizard(monkeypatch, tmp_path):