diff --git a/examples/puzzletron/configs/orchestration/execution.example.yaml b/examples/puzzletron/configs/orchestration/execution.example.yaml index 562d0cd5b84..2bfd3f675cb 100644 --- a/examples/puzzletron/configs/orchestration/execution.example.yaml +++ b/examples/puzzletron/configs/orchestration/execution.example.yaml @@ -7,6 +7,7 @@ execution: defaults: failure_policy: strict halt_policy: drain + artifact_settling_timeout_seconds: 300 gpus_per_node: 8 stages: convert: diff --git a/examples/puzzletron/distributed_eval/run_coordinator.sh b/examples/puzzletron/distributed_eval/run_coordinator.sh index ac22b6774fd..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}" @@ -88,6 +103,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 +124,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 +150,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/docs/v2_architecture.md b/examples/puzzletron/docs/v2_architecture.md index bfbb00de7c4..6464fab6ae5 100644 --- a/examples/puzzletron/docs/v2_architecture.md +++ b/examples/puzzletron/docs/v2_architecture.md @@ -366,6 +366,12 @@ durable manual decision gate. | Global KD sanity | Can the student overfit with the configured CE/KLD/MTP loss path? | Validates forward/backward and loss semantics before a long run | | Artifact completion | Are all expected identities, shards, candidates, and outputs present? | Partial work remains resumable progress, not a completed stage | +On filesystems with delayed visibility, the controller allows completed work a +bounded interval to publish valid stage artifacts. Advanced deployments can set +`execution.defaults.artifact_settling_timeout_seconds` in the execution config; +the default is 300 seconds. This changes only the post-completion settling +window, not scheduler or stage timeouts. + These gates answer different questions. Width ranking compares the quality of different reduced candidates at the same target geometry. Sort and slicing equivalence compare routes that are supposed to represent the same model 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..d0bca646dd5 100644 --- a/examples/puzzletron/finalize_replacement_scoring.py +++ b/examples/puzzletron/finalize_replacement_scoring.py @@ -1,19 +1,148 @@ #!/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.""" from __future__ import annotations import argparse +import json +import os +from collections.abc import Mapping from pathlib import Path -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 +__all__ = [ + "finalization_marker_is_current", + "finalize_replacement_scoring", + "finalize_replacement_scoring_diagnostics", + "main", + "write_finalization_marker", +] + + +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) + + +def _successful_manifest_identity_from_payload(manifest: object) -> str | 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") + return str(identity) if identity else None + + +def _successful_manifest_identity(manifest_path: str | Path) -> str | None: + try: + manifest = json.loads(Path(manifest_path).read_text()) + except (OSError, ValueError): + return None + return _successful_manifest_identity_from_payload(manifest) + + +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 + outputs = manifest.get("outputs") if isinstance(manifest, Mapping) else None + return bool( + marker_identity + and marker_identity == _successful_manifest_identity_from_payload(manifest) + and isinstance(outputs, Mapping) + and "report" in outputs + and summary == outputs["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 = stage_manifest_from_config("replacement_scoring", 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() @@ -21,33 +150,10 @@ def main() -> None: 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/examples/puzzletron/main.py b/examples/puzzletron/main.py index 2c6343f5601..43b451374fd 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, ) @@ -67,7 +67,12 @@ if __package__: from .acceptance_resume import build_payload, check_marker, marker_path, write_marker else: - from acceptance_resume import build_payload, check_marker, marker_path, write_marker + from acceptance_resume import ( # type: ignore[no-redef] + build_payload, + check_marker, + marker_path, + write_marker, + ) STAGES = stage_ids() PIPELINE_STAGE_ORDER = topological_stage_ids() @@ -415,6 +420,18 @@ 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 + + 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 +448,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, @@ -454,7 +466,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": @@ -470,7 +481,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..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 @@ -14,10 +29,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 +42,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 +133,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 +153,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 +203,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 +251,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 +269,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..8082c9b4aaf 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", @@ -136,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 ) @@ -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/orchestration/adapters/base.py b/modelopt/torch/puzzletron/orchestration/adapters/base.py index ebf71dd4de8..a274144e04b 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.""" @@ -22,7 +34,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 +87,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 binding a new attempt. + + Read-only currentness checks call only :meth:`execution_identity_projection`. + Adapters that require preparation must implement it here rather than mutate + state while projecting identity. + """ + def classify_failure( self, *, diff --git a/modelopt/torch/puzzletron/orchestration/adapters/pool.py b/modelopt/torch/puzzletron/orchestration/adapters/pool.py index 73a489fe9d3..1502a503a21 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/pool.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/pool.py @@ -1,12 +1,26 @@ # 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.""" from __future__ import annotations from pathlib import Path +from typing import cast +from ..identity import stable_hash from ..schema import ( AttemptSpec, CampaignPlan, @@ -19,6 +33,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 +145,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.""" @@ -149,7 +181,7 @@ def plan(self, plan: CampaignPlan, node: StagePlanNode) -> WorkPlan: work_id=( f"{node.stage_id}:gang" if len(widths) == 1 - else f"{node.stage_id}:width-{int(width):04d}" + else f"{node.stage_id}:width-{cast('int', width):04d}" ), stage_id=node.stage_id, shard_index=index, @@ -233,7 +265,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 +287,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,12 +298,10 @@ 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( - len(replacement_widths) - ), + "FINALIZE_EXPECTED_COMPLETIONS": str(len(replacement_widths)), } ) script = repo / "examples/puzzletron/distributed_eval/run_replacement_pool.sh" @@ -326,6 +358,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/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index 24fc77bc0a9..4d05c519f3e 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -17,11 +17,10 @@ from __future__ import annotations +import asyncio import json -import subprocess from pathlib import Path - -from puzzletron_orchestrator.post_mip.records import CandidateLedger +from typing import Any, Sequence from ..schema import ( AttemptSpec, @@ -41,6 +40,47 @@ __all__ = ["ManualInputRequired", "PostMIPAdapter"] +_DEFAULT_AGGREGATION_TIMEOUT_SECONDS = 300.0 + + +async def _communicate_with_timeout( + argv: Sequence[str], *, cwd: Path, timeout_seconds: float +) -> tuple[int, str, str]: + process = await asyncio.create_subprocess_exec( + *argv, + cwd=cwd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout_seconds) + except TimeoutError: + process.kill() + await process.communicate() + raise + return_code = process.returncode + if return_code is None: + raise RuntimeError("aggregation subprocess exited without a return code") + return return_code, stdout.decode(), stderr.decode() + + +def _run_aggregation_command( + argv: Sequence[str], *, cwd: Path, timeout_seconds: float +) -> tuple[int, str, str]: + """Run one argv-only aggregation process within a finite deadline.""" + + return asyncio.run(_communicate_with_timeout(argv, cwd=cwd, timeout_seconds=timeout_seconds)) + + +def _post_mip_identity_api() -> Any: + """Load the producer identity contract after orchestration initialization.""" + + 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 + return identity_api + class ManualInputRequired(RuntimeError): """A durable manual-filter review exists and needs a user decision.""" @@ -75,23 +115,19 @@ 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: + identity_api = _post_mip_identity_api() + try: + return identity_api.expected_post_mip_candidate_count(_identity_config(plan), stage_id) + except identity_api.PostMIPExecutionContractUnavailable: + registry = plan.puzzle_dir / "artifacts" / "post_mip" / "candidate_registry.json" + if registry.exists(): + raise + return None def _full_node_instance_count(node: StagePlanNode, count: int) -> int: @@ -110,12 +146,37 @@ 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 + _post_mip_identity_api().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 _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) 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( @@ -223,27 +284,38 @@ def aggregate( ) -> PublishedOutput | None: repo = Path(plan.runner.contract.repository) script = repo / "examples" / "puzzletron" / "run_post_mip_node.py" - result = subprocess.run( - ( - "python", - str(script), - "--config", - plan.experiment_config_path, - "--stage-id", - node.stage_id, - "--aggregate", - ), - cwd=repo, - capture_output=True, - text=True, - check=False, + argv = [ + "python", + str(script), + "--config", + plan.experiment_config_path, + "--stage-id", + node.stage_id, + "--aggregate", + ] + for override in plan.overrides: + argv.extend(["--override", override]) + timeout_seconds = float( + plan.execution_defaults.get( + "artifact_settling_timeout_seconds", + _DEFAULT_AGGREGATION_TIMEOUT_SECONDS, + ) ) - if result.returncode: + try: + return_code, stdout, stderr = _run_aggregation_command( + argv, + cwd=repo, + timeout_seconds=timeout_seconds, + ) + except TimeoutError as error: + raise RuntimeError( + f"{node.stage_id} aggregation timed out after {timeout_seconds:g}s" + ) from error + if return_code: raise RuntimeError( - f"{node.stage_id} aggregation failed: " - f"{result.stderr.strip() or result.stdout.strip()}" + f"{node.stage_id} aggregation failed: {stderr.strip() or stdout.strip()}" ) - output_lines = [line for line in result.stdout.splitlines() if line.strip()] + output_lines = [line for line in stdout.splitlines() if line.strip()] if not output_lines: raise RuntimeError(f"{node.stage_id} aggregation produced no summary") payload = json.loads(output_lines[-1]) diff --git a/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py b/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py index cc8ef517c42..8940bd91678 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py @@ -457,155 +457,25 @@ 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.""" + """Validate a node summary without importing torch.""" 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 (__package__ or "").startswith("puzzletron_orchestrator."): + from puzzletron_orchestrator.post_mip.identity import ( + expected_post_mip_execution_identity, ) - 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) + 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( + effective_config, + stage_id, + ) except (KeyError, OSError, RuntimeError, TypeError, ValueError): return False diff --git a/modelopt/torch/puzzletron/orchestration/compiler.py b/modelopt/torch/puzzletron/orchestration/compiler.py index 1eaf8992e49..b41d5e203a2 100644 --- a/modelopt/torch/puzzletron/orchestration/compiler.py +++ b/modelopt/torch/puzzletron/orchestration/compiler.py @@ -17,6 +17,7 @@ from __future__ import annotations +import math from dataclasses import asdict from pathlib import Path from typing import Any, Mapping @@ -62,6 +63,7 @@ ] _CONTROLLER_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_DEFAULT_ARTIFACT_SETTLING_TIMEOUT_SECONDS = 300.0 def _mapping(value: Any) -> dict[str, Any]: @@ -76,7 +78,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": { @@ -370,6 +372,29 @@ def load_execution_config(path: str | Path) -> dict[str, Any]: return _mapping(payload.get("execution")) +def _resolve_artifact_settling_timeout_seconds( + execution_defaults: Mapping[str, Any], +) -> float: + """Resolve the finite positive timeout for publishing completed-stage artifacts.""" + + value = execution_defaults.get( + "artifact_settling_timeout_seconds", + _DEFAULT_ARTIFACT_SETTLING_TIMEOUT_SECONDS, + ) + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError( + "execution.defaults.artifact_settling_timeout_seconds must be a positive " + f"finite number, got {value!r}" + ) + timeout_seconds = float(value) + if not math.isfinite(timeout_seconds) or timeout_seconds <= 0: + raise ValueError( + "execution.defaults.artifact_settling_timeout_seconds must be a positive " + f"finite number, got {value!r}" + ) + return timeout_seconds + + def _parse_mesh_override(payload: Mapping[str, Any] | None) -> ParallelMeshOverride | None: if not payload: return None @@ -408,6 +433,7 @@ def resolve_stage_execution_specs( """Resolve per-stage execution specs with defaults.""" defaults = _mapping(execution.get("defaults")) + _resolve_artifact_settling_timeout_seconds(defaults) default_gpus_per_node = int(defaults.get("gpus_per_node", 8)) default_policy = FailurePolicy(str(defaults.get("failure_policy", FailurePolicy.STRICT.value))) stage_payload = _mapping(execution.get("stages")) @@ -596,6 +622,7 @@ def compile_campaign_plan( execution_defaults=_mapping(execution.get("defaults")), stages=tuple(nodes), contract_hash=contract_hash, + overrides=tuple(overrides or ()), ) @@ -606,6 +633,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/config.py b/modelopt/torch/puzzletron/orchestration/config.py index 0e1c1594216..f5265eb8e36 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: @@ -144,17 +176,38 @@ def _resolve(value: Any, config: Mapping[str, Any]) -> Any: def _apply_override(config: dict[str, Any], override: str) -> None: + if override.startswith("~"): + raise ValueError(f"Deletion overrides are not supported: {override!r}") key, separator, raw_value = override.partition("=") if not separator: raise ValueError(f"Override must have KEY=VALUE form: {override!r}") - keys = key.lstrip("+").split(".") + addition_only = False + allow_missing = False + if key.startswith("++"): + key = key[2:] + allow_missing = True + elif key.startswith("+"): + key = key[1:] + addition_only = True + allow_missing = True + if not key or key.startswith(("+", "~")): + raise ValueError(f"Unsupported Hydra override form: {override!r}") + keys = key.split(".") target = config for part in keys[:-1]: - child = target.setdefault(part, {}) + if part not in target: + if not allow_missing: + raise ValueError(f"Override path does not exist: {override!r}") + target[part] = {} + child = target[part] 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) + if addition_only and keys[-1] in target: + raise ValueError(f"Addition override already exists: {override!r}") + if not allow_missing and keys[-1] not in target: + raise ValueError(f"Override key does not exist: {override!r}") + target[keys[-1]] = _load_yaml(raw_value) def load_experiment_config( diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index f36c1105c48..d98a18b7def 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.py @@ -13,32 +13,34 @@ # 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 import json +import math import signal 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 types import MappingProxyType +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 -from .compiler import plan_to_dict +from .compiler import _resolve_artifact_settling_timeout_seconds, 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 +49,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, @@ -74,6 +78,17 @@ def create_executor(plan: CampaignPlan, *, local: bool = False) -> Executor: raise ValueError(f"Unsupported runner kind: {plan.runner.kind}") +def _persisted_completion_timestamp(attempt: Mapping[str, Any]) -> float | None: + for field in ("completed_at", "submitted_at"): + value = attempt.get(field) + if isinstance(value, bool) or not isinstance(value, (int, float)): + continue + timestamp = float(value) + if math.isfinite(timestamp): + return timestamp + return None + + def _stage_dashboard_display_name( config: Mapping[str, Any], stage_id: str, @@ -116,11 +131,23 @@ 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, *, 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) @@ -133,7 +160,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( @@ -189,17 +216,28 @@ 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 + self._compiled_nodes: Mapping[str, Mapping[str, Any]] = MappingProxyType( + {str(stage["stage_id"]): stage for stage in plan_to_dict(self.plan)["stages"]} + ) + self._stage_execution_identity_cache: dict[str, str] | None = None defaults = dict(plan.execution_defaults or {}) + self.artifact_settling_timeout_seconds = _resolve_artifact_settling_timeout_seconds( + defaults + ) self._halt_policy = HaltPolicy(str(defaults.get("halt_policy", HaltPolicy.DRAIN.value))) def _recover_active_attempts(self) -> None: + active_states = { + JobState.RUNNING.value, + JobState.PENDING.value, + JobState.UNKNOWN.value, + } + nodes = {node.stage_id: node for node in self.plan.stages} + recoverable: list[tuple[dict[str, Any], JobHandle, bool]] = [] for attempt in self.store.list_attempts(): - if attempt.get("status") not in { - JobState.RUNNING.value, - JobState.PENDING.value, - JobState.UNKNOWN.value, - }: + if attempt.get("status") not in active_states: continue handle_payload = attempt.get("handle") if not isinstance(handle_payload, dict): @@ -210,12 +248,56 @@ def _recover_active_attempts(self) -> None: attempt_id=str(handle_payload["attempt_id"]), metadata=dict(handle_payload.get("metadata") or {}), ) + node = nodes.get(str(attempt.get("stage_id") or "")) + metadata = attempt.get("metadata") + current = False + if ( + node is not None + and attempt.get("contract_hash") == self.plan.contract_hash + and isinstance(metadata, Mapping) + ): + try: + current = metadata.get( + "stage_execution_identity" + ) == self._stage_execution_identity(node) + except ExecutionIdentityProjectionUnavailable: + pass + recoverable.append((attempt, handle, current)) + + stale_active_attempts: list[str] = [] + for attempt, handle, current in sorted(recoverable, key=lambda item: item[2]): + if current and stale_active_attempts: + continue status = self.executor.recover(handle) self.store.update_attempt_status( str(attempt["work_id"]), str(attempt["attempt_id"]), status, ) + if not current: + self.store.untrack_live_job(handle.handle_id) + if status.state.value in active_states: + self.executor.cancel([handle]) + stale_active_attempts.append(str(attempt["attempt_id"])) + self.store.append_event( + "stale_active_attempt_cancellation_requested", + { + "attempt_id": str(attempt["attempt_id"]), + "work_id": str(attempt["work_id"]), + "reported_state": status.state.value, + }, + ) + self.logger.warning( + f"requested cancellation of stale {attempt['work_id']} " + f"[{handle.handle_id}]; " + "waiting for scheduler confirmation before resubmission" + ) + else: + self.logger.warning( + f"ignored stale {attempt['work_id']} recovered as " + f"{status.state.value} [{handle.handle_id}]" + ) + continue if status.state in {JobState.RUNNING, JobState.PENDING, JobState.UNKNOWN}: tracked = JobHandle( backend=handle.backend, @@ -247,6 +329,11 @@ def _recover_active_attempts(self) -> None: self.logger.warning( f"recovered {attempt['work_id']} as {status.state.value} [{handle.handle_id}]" ) + if stale_active_attempts: + raise RuntimeError( + "stale active attempts were cancelled; rerun after the scheduler reports " + "them terminal" + ) def _parents_ready(self, node: StagePlanNode) -> bool: for parent in node.parents: @@ -311,22 +398,116 @@ 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: + 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 = [ + 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: + return _persisted_completion_timestamp(attempt) or 0.0 + + completed.append( + max( + matches, + key=_completion_time, + ) + ) + return completed + + def _stage_execution_identity( + self, + node: StagePlanNode, + work_plan: WorkPlan | None = None, + ) -> str: + cache = self._stage_execution_identity_cache + if cache is not None and node.stage_id in cache: + return cache[node.stage_id] + adapter = adapter_for_stage(node) + work_plan = work_plan or adapter.plan(self.plan, node) + payload = { + "execution_contract_hash": self.plan.contract_hash, + "semantic_config": semantic_stage_config(self.plan.experiment_config, node.stage_id), + "compiled_node": self._compiled_nodes[node.stage_id], + "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) + identity = stable_hash(payload, prefix=f"{node.stage_id}_execution") + if cache is not None: + cache[node.stage_id] = identity + return identity + + 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 _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: + timestamp = _persisted_completion_timestamp(attempt) + if timestamp is None: + return self.artifact_settling_timeout_seconds + completed_at.append(timestamp) + 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 +543,64 @@ 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 + # 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 False 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) + 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") + == stage_execution_identity + for attempt in record.attempts + ) + if not current_failure 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") + def _available_nodes(self) -> int | None: slurm = self.plan.runner.slurm if slurm is None or slurm.max_nodes is None: @@ -382,10 +613,13 @@ 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) + adapter.prepare_execution_identity_projection(plan=self.plan, node=node) + if self._stage_execution_identity_cache is not None: + self._stage_execution_identity_cache.pop(node.stage_id, None) work_plan = adapter.plan(self.plan, node) handles: list[tuple[JobHandle, str, str]] = [] available_nodes = self._available_nodes() @@ -394,7 +628,7 @@ def _submit_stage(self, node: StagePlanNode, *, overrides: list[str] | None = No f"{node.gpus_per_instance} GPU(s)/instance" ) for item in work_plan.items: - prior = [ + prior: list[Mapping[str, Any]] = [ attempt for attempt in self.store.list_attempts(node.stage_id) if attempt.get("work_id") == item.work_id @@ -409,13 +643,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=list(self.plan.overrides), + ), ) if available_nodes is not None and attempt.allocation_nodes > available_nodes: self.logger.wait( @@ -456,6 +694,18 @@ def _submit_stage(self, node: StagePlanNode, *, overrides: list[str] | None = No 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) @@ -464,12 +714,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 @@ -491,32 +736,142 @@ 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 ManualInputRequired as request: + return self._wait_for_manual_input(node, request) + 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 < self.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 = [ + 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( 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": self.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"{self.artifact_settling_timeout_seconds:g}s: {reason}" ) return True @@ -982,6 +1337,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 @@ -1029,12 +1389,14 @@ 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() self._interactive_ready = True while True: + self._stage_execution_identity_cache = {} try: if self._shutdown_requested: cancelled = True @@ -1052,18 +1414,33 @@ 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 ( + 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): - 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 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 - 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()) @@ -1122,6 +1499,7 @@ def _on_signal(signum: int, _frame: object | None) -> None: self.logger.shutdown(f"{signal_name} received; cancelling active jobs") self.shutdown(reason="keyboard interrupt") finally: + self._stage_execution_identity_cache = None self._interactive_ready = False self.terminal_controls.stop() if lease is not None and self._shutdown_requested and not self._shutting_down: @@ -1160,7 +1538,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/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/pipeline_config.py b/modelopt/torch/puzzletron/pipeline_config.py index e2d49695f54..e2c956d9c2a 100644 --- a/modelopt/torch/puzzletron/pipeline_config.py +++ b/modelopt/torch/puzzletron/pipeline_config.py @@ -17,7 +17,7 @@ from copy import deepcopy from pathlib import Path -from typing import Any +from typing import Any, cast import hydra from omegaconf import DictConfig, OmegaConf @@ -33,6 +33,7 @@ "load_runtime_hydra_config", "normalize_pipeline_config", "pipeline_config_from_path", + "rebase_authored_pipeline_config", ] @@ -115,7 +116,7 @@ def _config_root_and_name(path: Path) -> tuple[Path, str]: def _to_plain(config: DictConfig | dict[str, Any]) -> dict[str, Any]: if isinstance(config, DictConfig): - return OmegaConf.to_container(config, resolve=True) + return cast("dict[str, Any]", OmegaConf.to_container(config, resolve=True)) return deepcopy(dict(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/post_mip/identity.py b/modelopt/torch/puzzletron/post_mip/identity.py new file mode 100644 index 00000000000..3324a0947cb --- /dev/null +++ b/modelopt/torch/puzzletron/post_mip/identity.py @@ -0,0 +1,296 @@ +# 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.""" + +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__ or "").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 _dependency_owners(node: CompiledPostMIPNode) -> set[str]: + owners = { + reference.partition(".")[0] + for reference in node.metric_references + if not reference.startswith("mip.") + } + if node.model_source not in {"latest", "origin"}: + owners.add(node.model_source) + return owners + + +def _published_execution_identity(config: Mapping[str, Any], owner: str) -> str: + current_path = _puzzle_dir(config) / "artifacts" / "post_mip" / "nodes" / owner / "current.json" + try: + current = json.loads(current_path.read_text()) + execution_identity = current["execution_identity"] + except (FileNotFoundError, json.JSONDecodeError, KeyError, TypeError) as error: + raise PostMIPExecutionContractUnavailable( + f"post-MIP dependency {owner!r} has no published execution identity" + ) from error + if not isinstance(execution_identity, str) or not execution_identity: + raise ValueError(f"post-MIP dependency {owner!r} has an invalid execution identity") + return execution_identity + + +def _input_candidate_set( + ledger: CandidateLedger, + config: Mapping[str, Any], + node: CompiledPostMIPNode, + *, + input_execution_identity: str | None, +) -> 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, + execution_identity=input_execution_identity, + ) + 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") + if candidate_set.flow_id != node.flow_id or candidate_set.node_id != node.input_id: + raise RuntimeError( + f"post-MIP input node {node.input_id!r} candidate set has invalid context" + ) + if ( + input_execution_identity is not None + and candidate_set.producer_execution_identity != input_execution_identity + ): + raise RuntimeError( + f"post-MIP input node {node.input_id!r} candidate set does not match " + "its current execution" + ) + return candidate_set + + +def _active_mip_contract(config: Mapping[str, Any]) -> tuple[str, set[str]]: + active_path = _puzzle_dir(config) / "mip" / "active_profiles.json" + try: + active = json.loads(active_path.read_text()) + except (OSError, json.JSONDecodeError) as error: + raise PostMIPExecutionContractUnavailable( + "active MIP profile manifest is unavailable" + ) from error + 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) + dependency_executions = { + owner: _published_execution_identity(config, owner) + for owner in sorted(_dependency_owners(node)) + } + input_execution_identity = dependency_executions.get(node.input_id) + if node.input_id != "source" and input_execution_identity is None: + input_execution_identity = _published_execution_identity(config, node.input_id) + execution_identities = dict(dependency_executions) + if input_execution_identity is not None: + execution_identities[node.input_id] = input_execution_identity + ledger = CandidateLedger(_puzzle_dir(config) / "artifacts" / "post_mip") + ledger.load_execution_observations(execution_identities) + 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, + input_execution_identity=input_execution_identity, + ) + 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.""" + + try: + dependency_executions = { + owner: ledger.execution_identities[owner] for owner in sorted(_dependency_owners(node)) + } + except KeyError as error: + raise PostMIPExecutionContractUnavailable( + f"post-MIP dependency {error.args[0]!r} has no loaded execution identity" + ) from error + source_revisions = {} + for revision_id in candidate_set.revision_ids: + try: + source_revisions[revision_id] = ledger.source_revision( + revision_id, node.model_source + ).revision_id + except (KeyError, ValueError) as error: + raise PostMIPExecutionContractUnavailable( + f"post-MIP source revision for {revision_id!r} is unavailable" + ) from error + 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/records.py b/modelopt/torch/puzzletron/post_mip/records.py index 839e4c34fdf..0e73cf89703 100644 --- a/modelopt/torch/puzzletron/post_mip/records.py +++ b/modelopt/torch/puzzletron/post_mip/records.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. """Durable candidate lineage and observations for configurable post-MIP flows.""" @@ -106,6 +118,7 @@ class CandidateLedger: def __init__(self, root: str | Path): self.root = Path(root) self.registry_path = self.root / "candidate_registry.json" + self.execution_identities: dict[str, str] = {} self.architectures: dict[str, ArchitectureCandidate] = {} self.revisions: dict[str, CandidateRevision] = {} self.observations: dict[str, dict[str, NodeObservation]] = {} @@ -129,21 +142,27 @@ def _load(self) -> None: ) for key, value in dict(payload.get("revisions") or {}).items() } - self.active_mip_execution_identity = str( - payload.get("active_mip_execution_identity") or "" - ) - self.active_profile_ids = { - str(value) for value in payload.get("active_profile_ids") or () - } + self.active_mip_execution_identity = str(payload.get("active_mip_execution_identity") or "") + self.active_profile_ids = {str(value) for value in payload.get("active_profile_ids") or ()} + execution_identities = {} for current_path in sorted((self.root / "nodes").glob("*/current.json")): - node_id = current_path.parent.name current = json.loads(current_path.read_text()) - execution_identity = str(current["execution_identity"]) - path = current_path.parent / "executions" / execution_identity / "observations.json" - rows = json.loads(path.read_text()) + execution_identities[current_path.parent.name] = str(current["execution_identity"]) + self.load_execution_observations(execution_identities) + + def load_execution_observations( + self, + execution_identities: Mapping[str, str], + ) -> None: + """Load node observations from exact immutable executions.""" + + for node_id, execution_identity in execution_identities.items(): + path = self.root / "nodes" / node_id / "executions" / execution_identity + rows = json.loads((path / "observations.json").read_text()) self.observations[node_id] = { row["input_revision_id"]: NodeObservation(**row) for row in rows } + self.execution_identities[node_id] = execution_identity def publish(self) -> Path: payload = { @@ -184,19 +203,23 @@ def publish_node( index["current"] = execution_identity self._atomic_json(index_path, index) self._atomic_json(node_root / "current.json", {"execution_identity": execution_identity}) + self.execution_identities[node_id] = execution_identity self.publish() return observations_path, candidate_set_path - def load_candidate_set(self, node_id: str) -> CandidateSet: + def load_candidate_set( + self, + node_id: str, + *, + execution_identity: str | None = None, + ) -> CandidateSet: node_root = self.root / "nodes" / node_id - current = json.loads((node_root / "current.json").read_text()) + execution_identity = execution_identity or self.execution_identities.get(node_id) + if execution_identity is None: + current = json.loads((node_root / "current.json").read_text()) + execution_identity = str(current["execution_identity"]) payload = json.loads( - ( - node_root - / "executions" - / str(current["execution_identity"]) - / "candidate_set.json" - ).read_text() + (node_root / "executions" / execution_identity / "candidate_set.json").read_text() ) return CandidateSet( flow_id=str(payload["flow_id"]), @@ -212,8 +235,7 @@ def resolve_metric(self, revision_id: str, reference: str) -> float | None: owner, separator, metric = reference.partition(".") if not separator: raise ValueError( - "metric reference must be mip. or .: " - f"{reference}" + f"metric reference must be mip. or .: {reference}" ) if owner == "mip": root_revision = revision @@ -359,9 +381,7 @@ def candidate_metadata(self, revision_id: str) -> dict[str, Any]: observations[node_id] = canonicalize(asdict(observation)) return { "revision_id": revision_id, - "architecture": canonicalize( - asdict(self.architectures[revision.architecture_id]) - ), + "architecture": canonicalize(asdict(self.architectures[revision.architecture_id])), "lineage": lineage, "observations": observations, } @@ -433,15 +453,14 @@ def root_set(self, flow_id: str, source: Mapping[str, Any]) -> CandidateSet: origin for origin in architecture.origins if ( - origin.get("profile_id") in self.active_profile_ids - and origin.get("mip_execution_identity") - == self.active_mip_execution_identity - and origin.get("run_id") == run - and (variants == "all" or origin.get("variant_id") in variants) - and ( - objectives == "all" - or (origin.get("objective") or {}).get("metric") in objectives - ) + origin.get("profile_id") in self.active_profile_ids + and origin.get("mip_execution_identity") == self.active_mip_execution_identity + and origin.get("run_id") == run + and (variants == "all" or origin.get("variant_id") in variants) + and ( + objectives == "all" + or (origin.get("objective") or {}).get("metric") in objectives + ) ) ] if matching_origins: @@ -484,9 +503,7 @@ def _ingest_solution( prefix="architecture", ) costs = {} - for key, value in dict( - result.get("total_costs") or raw.get("total_costs") or {} - ).items(): + for key, value in dict(result.get("total_costs") or raw.get("total_costs") or {}).items(): if not isinstance(value, (int, float)) or isinstance(value, bool): continue key = str(key) diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 08aeff370a1..159ab200e9b 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_contract_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()) @@ -451,7 +400,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()) @@ -648,7 +599,8 @@ 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_contract = post_mip_execution_contract(config, node, candidate_set, ledger) + execution_identity = post_mip_execution_contract_identity(execution_contract) revision_ids = candidate_set.revision_ids[shard_index::shard_count] output_path = ( _execution_root(config, node, execution_identity) @@ -750,7 +702,8 @@ 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_contract = post_mip_execution_contract(config, node, input_set, ledger) + execution_identity = post_mip_execution_contract_identity(execution_contract) timed_out_candidates = [] if node.node_type == "filter": observations, output_set = _aggregate_filter(ledger, node, input_set, execution_identity) @@ -871,7 +824,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": execution_contract, "checkpoints": sorted( { str(ledger.revisions[revision_id].artifact["checkpoint"]) diff --git a/modelopt/torch/puzzletron/stage_runner.py b/modelopt/torch/puzzletron/stage_runner.py index be8e3cb9d79..8845b2bf04e 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, ) @@ -249,6 +248,7 @@ def run_stage( handler_map = DEFAULT_HANDLERS + assert handler_map is not None handler = handler_map.get(stage) if handler is None: raise RuntimeError(f"enabled stage {stage!r} has no registered handler") diff --git a/modelopt/torch/puzzletron/stages/graph.py b/modelopt/torch/puzzletron/stages/graph.py index d1e2785d084..754eafad2ff 100644 --- a/modelopt/torch/puzzletron/stages/graph.py +++ b/modelopt/torch/puzzletron/stages/graph.py @@ -514,18 +514,39 @@ 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. + + For identity projection, an absent section, a ``None`` section, and an + empty mapping are equivalent because none contains an authored option. + Other falsey values remain in the projection because they can be semantic. """ + 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] for key in sections if key in config} + return { + key: selected[key] + for key in sections + if key in selected + and selected[key] is not None + and not (isinstance(selected[key], Mapping) and not selected[key]) + } def stage_display_name(stage_id: str, *, granularity: str | None = None) -> str: 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..6ea9dead9e1 100644 --- a/tests/unit/torch/puzzletron/test_example_runner.py +++ b/tests/unit/torch/puzzletron/test_example_runner.py @@ -1,4 +1,44 @@ -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.pipeline_config import pipeline_config_from_path +from modelopt.torch.puzzletron.stage_runner import StageResult, run_stage def test_worker_command_propagates_gpu_count_to_composite_followups(): @@ -10,3 +50,162 @@ 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") + + +def test_loaded_stage_run_publishes_distinct_authored_and_effective_config(tmp_path: Path) -> None: + config_path = tmp_path / "experiment.yaml" + config_path.write_text( + "\n".join( + ( + f"puzzle_dir: {tmp_path}", + "experiment:", + f" dir: {tmp_path}", + "sort_sanity:", + " enabled: true", + "width_sanity:", + " enabled: true", + "slicing_sanity:", + " enabled: true", + ) + ) + + "\n" + ) + override = "+slicing_sanity.tolerance=0.25" + config = pipeline_config_from_path(config_path, overrides=[override]) + manifest_path = tmp_path / "manifests" / "slicing_sanity.json" + + def capture_handler(effective_config, manifest): + manifest.complete(outputs={}) + write_stage_manifest(manifest_path, manifest) + return StageResult( + stage="slicing_sanity", + status="success", + manifest_path=manifest_path, + message="captured", + ) + + run_stage(config, "slicing_sanity", handlers={"slicing_sanity": capture_handler}) + + pointer = json.loads(manifest_path.read_text()) + resolved = json.loads( + (tmp_path / pointer["execution_record"]["resolved_config_path"]).read_text() + ) + assert pointer["config"]["slicing_sanity"] == { + "enabled": True, + "tolerance": 0.25, + } + assert pointer["config"]["sort_sanity"] == {"enabled": True} + assert "search_space" not in pointer["config"] + assert resolved["resolved_stage_config"]["search_space"] == {"axes": {}} + assert resolved["provenance"]["overrides"] == [override] diff --git a/tests/unit/torch/puzzletron/test_orchestration_compiler.py b/tests/unit/torch/puzzletron/test_orchestration_compiler.py index c5c2af2eb40..a79335c3917 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_compiler.py +++ b/tests/unit/torch/puzzletron/test_orchestration_compiler.py @@ -168,6 +168,48 @@ def test_compile_campaign_plan_preserves_drain_halt_policy(tmp_configs): assert plan.execution_defaults["halt_policy"] == "drain" assert controller._halt_policy is HaltPolicy.DRAIN + assert controller.artifact_settling_timeout_seconds == 300.0 + + +def test_compile_campaign_plan_configures_artifact_settling_timeout(tmp_configs): + experiment_path, runner_path, execution_path = tmp_configs + execution = load_execution_config(execution_path) + execution["defaults"]["artifact_settling_timeout_seconds"] = 45 + + plan = compile_campaign_plan( + experiment_config_path=experiment_path, + runner=load_runner_config(runner_path), + execution=execution, + ) + + assert plan.execution_defaults["artifact_settling_timeout_seconds"] == 45 + assert CampaignController(plan).artifact_settling_timeout_seconds == 45.0 + + +@pytest.mark.parametrize( + ("value", "error_type"), + [ + (True, TypeError), + ("300", TypeError), + (0, ValueError), + (-1, ValueError), + (float("nan"), ValueError), + (float("inf"), ValueError), + ], +) +def test_compile_campaign_plan_rejects_invalid_artifact_settling_timeout( + tmp_configs, value, error_type +): + experiment_path, runner_path, execution_path = tmp_configs + execution = load_execution_config(execution_path) + execution["defaults"]["artifact_settling_timeout_seconds"] = value + + with pytest.raises(error_type, match="artifact_settling_timeout_seconds"): + compile_campaign_plan( + experiment_config_path=experiment_path, + runner=load_runner_config(runner_path), + execution=execution, + ) def test_compile_campaign_plan_uses_cpu_partition_without_gpus(tmp_configs): diff --git a/tests/unit/torch/puzzletron/test_orchestration_executors.py b/tests/unit/torch/puzzletron/test_orchestration_executors.py index 5e77a624586..617b53fb66f 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_executors.py +++ b/tests/unit/torch/puzzletron/test_orchestration_executors.py @@ -781,7 +781,10 @@ def test_replacement_pool_uses_one_four_node_gang_allocation(tmp_path: Path): assert attempt.command.argv[-1].endswith("run_replacement_pool.sh") -def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path): +def _replacement_width_attempts( + tmp_path: Path, + widths: list[int], +) -> tuple[CampaignPlan, WorkPlan, list[AttemptSpec]]: runner = RunnerEnvironment( kind="slurm", contract=ExecutionContract(repository=str(tmp_path), venv=str(tmp_path / ".venv")), @@ -806,7 +809,7 @@ def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path) experiment_config_path=str(tmp_path / "experiment.yaml"), puzzle_dir=tmp_path / "run", experiment_config={ - "embedding_pruning": {"enabled": True, "widths": [2048, 1792]}, + "embedding_pruning": {"enabled": True, "widths": widths}, "replacement_scoring": {"granularity": "subblock"}, }, runner=runner, @@ -816,13 +819,6 @@ def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path) ) adapter = adapter_for_stage(node) work_plan = adapter.plan(plan, node) - - assert [item.work_id for item in work_plan.items] == [ - "replacement_scoring:width-2048", - "replacement_scoring:width-1792", - ] - assert [item.metadata["worker_count"] for item in work_plan.items] == [4, 4] - attempts = [ adapter.command( plan=plan, @@ -830,14 +826,40 @@ 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) ] + return plan, work_plan, attempts + + +def test_replacement_pool_splits_workers_across_embedding_widths(tmp_path: Path): + plan, work_plan, attempts = _replacement_width_attempts(tmp_path, [2048, 1792]) + runner = plan.runner + node = plan.stages[0] + adapter = adapter_for_stage(node) + + assert [item.work_id for item in work_plan.items] == [ + "replacement_scoring:width-2048", + "replacement_scoring:width-1792", + ] + assert [item.metadata["worker_count"] for item in work_plan.items] == [4, 4] + assert [attempt.allocation_nodes for attempt in attempts] == [2, 2] assert [attempt.allocation_gpus for attempt in attempts] == [16, 16] 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,10 +872,62 @@ 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") +def test_replacement_pool_completion_identity_changes_with_embedding_widths(tmp_path: Path): + _, _, baseline_attempts = _replacement_width_attempts(tmp_path, [2048, 1792]) + _, _, changed_attempts = _replacement_width_attempts(tmp_path, [2048, 1792, 1536, 1280]) + + baseline_completion_dirs = { + attempt.command.env["FINALIZE_COMPLETION_DIR"] for attempt in baseline_attempts + } + changed_completion_dirs = { + attempt.command.env["FINALIZE_COMPLETION_DIR"] for attempt in changed_attempts + } + assert len(baseline_completion_dirs) == 1 + assert len(changed_completion_dirs) == 1 + assert changed_completion_dirs != baseline_completion_dirs + assert [attempt.command.env["FINALIZE_COMPLETION_MARKER"] for attempt in changed_attempts] == [ + "width-2048", + "width-1792", + "width-1536", + "width-1280", + ] + assert [ + attempt.command.env["FINALIZE_EXPECTED_COMPLETIONS"] for attempt in changed_attempts + ] == ["4", "4", "4", "4"] + + def test_stage_partition_override_forces_batch(tmp_path: Path): experiment = tmp_path / "experiment.yaml" experiment.write_text( diff --git a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py index 5ba0266162c..edf9bd78192 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_lightweight.py +++ b/tests/unit/torch/puzzletron/test_orchestration_lightweight.py @@ -22,8 +22,10 @@ import os import subprocess import sys +from copy import deepcopy from pathlib import Path +import pytest import yaml from puzzletron_orchestrator.adapters.stage_compat import stage_is_complete @@ -54,6 +56,28 @@ 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: + 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, + "-c", + "from modelopt.torch.puzzletron.pipeline_config import pipeline_config_from_path; " + "assert callable(pipeline_config_from_path)", + ], + cwd=REPOSITORY_ROOT, + env=environment, + 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 @@ -164,6 +188,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 +209,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 +218,77 @@ 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"], + ) + ) + + +@pytest.mark.parametrize("override", ["~value", "~value=1"]) +def test_load_experiment_config_rejects_deletion_overrides( + tmp_path: Path, + override: str, +) -> None: + experiment = tmp_path / "experiment.yaml" + experiment.write_text("value: 1\n") + + with pytest.raises(ValueError, match="^Deletion overrides are not supported"): + load_experiment_config(experiment, overrides=[override]) + + +def test_load_experiment_config_distinguishes_hydra_addition_modes( + tmp_path: Path, +) -> None: + experiment = tmp_path / "experiment.yaml" + experiment.write_text("value: 1\n") + + added = load_experiment_config(experiment, overrides=["+added.value=2"]) + with pytest.raises(ValueError, match="^Addition override already exists"): + load_experiment_config(experiment, overrides=["+value=2"]) + with pytest.raises(ValueError, match="^Override key does not exist"): + load_experiment_config(experiment, overrides=["missing=2"]) + with pytest.raises(ValueError, match="^Override path does not exist"): + load_experiment_config(experiment, overrides=["missing.value=2"]) + replaced = load_experiment_config(experiment, overrides=["++value=2"]) + created = load_experiment_config(experiment, overrides=["++created.value=3"]) + + assert added["added"] == {"value": 2} + assert replaced["value"] == 2 + assert created["created"] == {"value": 3} + + def test_convert_completeness_requires_runtime_subblock_library( tmp_path: Path, write_terminal_manifest ) -> None: @@ -313,8 +410,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("{}") @@ -326,6 +421,40 @@ 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: + 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..17704f2ad23 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py +++ b/tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.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 orchestrator shutdown and progress reporting.""" from __future__ import annotations @@ -41,8 +38,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 +264,7 @@ def submit(self, attempt: AttemptSpec) -> JobHandle: }, ) self._handles[handle.handle_id] = handle + self._attempts[handle.handle_id] = attempt return handle @staticmethod @@ -279,6 +284,76 @@ def _blocked_descendants(plan, failed_stages: set[str]) -> set[str]: return blocked +def _compile_test_plan( + tmp_path: Path, + *, + stage_filter: str | None = None, + overrides: list[str] | None = None, + execution_defaults: dict | None = None, +): + experiment, runner_path, execution_path = _write_configs(tmp_path) + execution = load_execution_config(execution_path) + execution["defaults"].update(execution_defaults or {}) + return compile_campaign_plan( + experiment_config_path=experiment, + runner=load_runner_config(runner_path), + execution=execution, + overrides=overrides, + stage_filter=stage_filter, + ) + + +def _compile_changed_convert_plans(tmp_path: Path): + experiment, runner_path, execution_path = _write_configs(tmp_path) + runner = load_runner_config(runner_path) + execution = load_execution_config(execution_path) + baseline = compile_campaign_plan( + experiment_config_path=experiment, + runner=runner, + execution=execution, + stage_filter="convert", + ) + config = yaml.safe_load(experiment.read_text()) + config["convert"]["model_path"] = "/models/replacement" + experiment.write_text(yaml.safe_dump(config)) + changed = compile_campaign_plan( + experiment_config_path=experiment, + runner=runner, + execution=execution, + stage_filter="convert", + ) + return baseline, changed + + +def _record_completed_attempt( + controller: CampaignController, + node, + *, + handle: JobHandle | None = None, +) -> AttemptSpec: + adapter = adapter_for_stage(node) + work_plan = adapter.plan(controller.plan, node) + item = work_plan.items[0] + attempt = controller._bind_attempt_to_stage_execution( + node, + work_plan, + adapter.command( + plan=controller.plan, + node=node, + item=item, + attempt_id="completed-attempt", + runner=controller.plan.runner, + ), + ) + controller.store.save_attempt(attempt, None, JobState.RUNNING.value) + controller.store.update_attempt_status( + item.work_id, + attempt.attempt_id, + JobStatus(handle=handle, state=JobState.COMPLETED), + ) + return attempt + + def test_depth_progress_reports_removal_and_candidate_counts(tmp_path: Path): iteration = tmp_path / "depth" / "iterative" / "iteration_00" iteration.mkdir(parents=True) @@ -424,13 +499,7 @@ def test_summarize_active_progress_prefers_stage_lines(tmp_path: Path): def test_controller_shutdown_cancels_active_jobs(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", - ) + plan = _compile_test_plan(tmp_path, stage_filter="convert") executor = _FakeExecutor() controller = CampaignController(plan, executor=executor, poll_interval_seconds=0.01) result = controller.run(once=True) @@ -448,13 +517,7 @@ def test_controller_shutdown_cancels_active_jobs(tmp_path: Path): def test_controller_waits_for_parent_job_after_artifact_appears(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="vllm_stats", - ) + plan = _compile_test_plan(tmp_path, stage_filter="vllm_stats") controller = CampaignController(plan, executor=_FakeExecutor()) child = next(node for node in plan.stages if node.stage_id == "vllm_stats") handle = JobHandle(backend="fake", handle_id="parent", attempt_id="a1") @@ -469,33 +532,430 @@ 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): - 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", +def test_controller_resubmits_legacy_completed_attempt_without_execution_identity( + tmp_path: Path, +): + plan = _compile_test_plan(tmp_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, ) - controller = CampaignController(plan, executor=_FakeExecutor()) + executor = _TrackingFakeExecutor() + controller = CampaignController(plan, executor=executor) + controller.store.save_attempt(attempt, None, JobState.COMPLETED.value) + controller.store.write_stage_record( + StageRunRecord( + stage_id="convert", + status=JobState.FAILED.value, + attempts=[ + PersistedAttempt( + attempt_id=attempt.attempt_id, + work_id=attempt.work_id, + stage_id=node.stage_id, + status=JobState.COMPLETED.value, + contract_hash=plan.contract_hash, + metadata={"stage_execution_identity_incompatible": True}, + ) + ], + ) + ) + + result = controller.run(once=True) + + assert result["halted"] is False + assert result["failed_stages"] == [] + assert executor.submitted_stage_ids == ["convert"] + attempts = controller.store.list_attempts("convert") + assert len(attempts) == 2 + attempts_by_id = {attempt["attempt_id"]: attempt for attempt in attempts} + legacy = attempts_by_id.pop("legacy-completed-attempt") + assert "stage_execution_identity" not in legacy["metadata"] + submitted = next(iter(attempts_by_id.values())) + assert submitted["metadata"]["stage_execution_identity"] == ( + controller._stage_execution_identity(node) + ) + assert not list( + controller.store.events_root.glob("*_stage_execution_identity_incompatible.json") + ) + + +def test_controller_resubmits_completed_work_when_stage_semantics_change( + tmp_path: Path, monkeypatch +): + plan_a, plan_b = _compile_changed_convert_plans(tmp_path) + 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.CANCELLED.value}, - {"work_id": "convert:0", "status": JobState.COMPLETED.value}, + { + "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_controller_cancels_stale_active_attempt_before_current_resubmission(tmp_path: Path): + old_plan, plan = _compile_changed_convert_plans(tmp_path) + old_node = old_plan.stages[0] + old_controller = CampaignController(old_plan, executor=_FakeExecutor()) + handle = JobHandle( + backend="fake", + handle_id="stale-active-handle", + attempt_id="completed-attempt", + metadata={"work_id": "convert:0"}, + ) + attempt = _record_completed_attempt(old_controller, old_node, handle=handle) + old_controller.store.save_attempt(attempt, handle, JobState.RUNNING.value) + old_controller.store.track_live_job(handle) - 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]) + class _FailsIfPolledExecutor(_TrackingFakeExecutor): + def poll(self, handles): + raise AssertionError("a stale active attempt must not enter the current poll set") + executor = _FailsIfPolledExecutor() + controller = CampaignController(plan, executor=executor) -def test_controller_completed_summary_excludes_store_only_completion(tmp_path: Path, monkeypatch): + with pytest.raises(RuntimeError, match="stale active attempts were cancelled"): + controller.run(once=True) + + assert executor.cancelled == [handle] + assert executor.submitted_stage_ids == [] + assert controller._active == {} + assert controller._failed_stages == set() + assert controller.store.load_attempt(attempt.work_id, attempt.attempt_id)["status"] == "running" + assert ( + len( + list( + controller.store.events_root.glob( + "*_stale_active_attempt_cancellation_requested.json" + ) + ) + ) + == 1 + ) + + +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 +): + plan = _compile_test_plan(tmp_path, stage_filter="convert") + node = plan.stages[0] + delegate = adapter_for_stage(node) + executor = _TrackingFakeExecutor() + controller = CampaignController(plan, executor=executor, poll_interval_seconds=45.0) + 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), + ) + _record_completed_attempt( + controller, + node, + handle=JobHandle( + backend="fake", + handle_id="completed-handle", + attempt_id="completed-attempt", + ), + ) + + 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_artifact_settling_deadline_survives_restart(tmp_path: Path, monkeypatch): + plan = _compile_test_plan( + tmp_path, + stage_filter="convert", + execution_defaults={"artifact_settling_timeout_seconds": 120}, + ) + node = plan.stages[0] + first_controller = CampaignController(plan, executor=_FakeExecutor()) + now = [1000.0] + monkeypatch.setattr("puzzletron_orchestrator.controller.time.time", lambda: now[0]) + _record_completed_attempt(first_controller, node) + + now[0] += 120.0 + restarted = CampaignController(plan, executor=_FakeExecutor()) + + assert ( + restarted._completed_work_artifact_settling_elapsed( + node, + restarted.store.list_attempts(node.stage_id), + ) + == 120.0 + ) + + +@pytest.mark.parametrize( + ("attempt", "expected_elapsed"), + [ + ({"completed_at": "invalid", "submitted_at": 990.0}, 10.0), + ({"completed_at": float("nan"), "submitted_at": False}, 120.0), + ], + ids=("submitted-at-fallback", "no-usable-timestamp"), +) +def test_controller_settling_elapsed_handles_legacy_attempt_timestamps( + tmp_path: Path, monkeypatch, attempt, expected_elapsed +): + plan = _compile_test_plan( + tmp_path, + stage_filter="convert", + execution_defaults={"artifact_settling_timeout_seconds": 120}, + ) + controller = CampaignController(plan, executor=_FakeExecutor()) + monkeypatch.setattr( + controller, + "_required_completed_attempts", + lambda _node, _attempts: [attempt], + ) + monkeypatch.setattr("puzzletron_orchestrator.controller.time.time", lambda: 1000.0) + + assert controller._completed_work_artifact_settling_elapsed(None, []) == expected_elapsed + + +@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 +): + plan = _compile_test_plan( + tmp_path, + stage_filter="convert", + execution_defaults={"artifact_settling_timeout_seconds": 120}, + ) + node = plan.stages[0] + executor = _TrackingFakeExecutor() + controller = CampaignController(plan, executor=executor, poll_interval_seconds=60.0) + now = [1000.0] + delegate = adapter_for_stage(node) + monkeypatch.setattr("puzzletron_orchestrator.controller.time.time", lambda: now[0]) + monkeypatch.setattr( + controller, + "_interruptible_sleep", + lambda seconds: now.__setitem__(0, now[0] + seconds), + ) + _record_completed_attempt( + controller, + node, + handle=JobHandle( + backend="fake", + handle_id="completed-handle", + attempt_id="completed-attempt", + ), + ) + + class _MissingArtifactsAdapter: + aggregation_ready = False + + def __getattr__(self, name): + return getattr(delegate, name) + + def aggregate(self, *, plan, node, work_plan): + 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: + assert self.aggregation_ready + return ValidatedResult(valid=True, reason="stage outputs present") + 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"] == 120.0 + assert event["payload"]["timeout_seconds"] == 120.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 + 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 (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 + + +def test_controller_ignores_failed_record_from_stale_stage_execution(tmp_path: Path): + old_plan, plan = _compile_changed_convert_plans(tmp_path) + old_identity = CampaignController(old_plan, executor=_FakeExecutor())._stage_execution_identity( + old_plan.stages[0] + ) + executor = _TrackingFakeExecutor() + controller = CampaignController(plan, executor=executor) + 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=plan.contract_hash, + metadata={"stage_execution_identity": old_identity}, + ) + ], + ) + ) + + result = controller.run(once=True) + + assert result["failed_stages"] == [] + assert executor.submitted_stage_ids == ["convert"] + + +def test_controller_completed_summary_excludes_store_only_completion(tmp_path: Path, monkeypatch): + plan = _compile_test_plan(tmp_path, stage_filter="convert") controller = CampaignController(plan, executor=_FakeExecutor()) controller.store.write_stage_record( StageRunRecord( @@ -517,13 +977,7 @@ def test_controller_completed_summary_excludes_store_only_completion(tmp_path: P def test_controller_aggregates_completed_work_before_resubmitting( tmp_path: Path, monkeypatch, write_terminal_manifest ): - 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", - ) + plan = _compile_test_plan(tmp_path, stage_filter="convert") node = plan.stages[0] delegate = adapter_for_stage(node) item = delegate.plan(plan, node).items[0] @@ -536,6 +990,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: @@ -559,13 +1014,7 @@ def aggregate(self, *, plan, node, work_plan): def test_controller_shutdown_cancels_store_tracked_jobs(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", - ) + plan = _compile_test_plan(tmp_path, stage_filter="convert") executor = _FakeExecutor() controller = CampaignController(plan, executor=executor, poll_interval_seconds=0.01) controller.run(once=True) @@ -579,13 +1028,7 @@ def test_controller_shutdown_cancels_store_tracked_jobs(tmp_path: Path): def test_controller_preserves_live_job_when_cancel_fails(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", - ) + plan = _compile_test_plan(tmp_path, stage_filter="convert") class _FailingCancelExecutor(_FakeExecutor): def cancel(self, handles: Sequence[JobHandle]) -> None: @@ -860,13 +1303,7 @@ def test_unknown_stage_progress_does_not_echo_arbitrary_log_tail(tmp_path: Path) def test_controller_keyboard_interrupt_cancels_jobs(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", - ) + plan = _compile_test_plan(tmp_path, stage_filter="convert") executor = _FakeExecutor() controller = CampaignController(plan, executor=executor, poll_interval_seconds=0.01) @@ -885,13 +1322,7 @@ def _raise_interrupt(_seconds: float) -> None: def test_controller_can_resume_quit_menu_then_detach_live_jobs(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", - ) + plan = _compile_test_plan(tmp_path, stage_filter="convert") executor = _FakeExecutor() class _Controls: @@ -1267,13 +1698,7 @@ def poll(self, handles: Sequence[JobHandle]) -> list[JobStatus]: def test_controller_sigint_during_recovery_cancels_jobs(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", - ) + plan = _compile_test_plan(tmp_path, stage_filter="convert") first_executor = _FakeExecutor() CampaignController(plan, executor=first_executor, poll_interval_seconds=0.01).run(once=True) @@ -1293,13 +1718,7 @@ def recover(self, handle: JobHandle) -> JobStatus: def test_controller_shutdown_flag_cancels_without_keyboard_interrupt(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", - ) + plan = _compile_test_plan(tmp_path, stage_filter="convert") executor = _FakeExecutor() controller = CampaignController(plan, executor=executor, poll_interval_seconds=0.05) diff --git a/tests/unit/torch/puzzletron/test_post_mip_adapter.py b/tests/unit/torch/puzzletron/test_post_mip_adapter.py index cffd6bf639e..cadbc951e94 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_adapter.py +++ b/tests/unit/torch/puzzletron/test_post_mip_adapter.py @@ -15,9 +15,17 @@ """Tests for post-MIP orchestration adapter launch policy.""" +import json +import os +import sys +from dataclasses import replace from pathlib import Path +from types import SimpleNamespace -from puzzletron_orchestrator.adapters.post_mip import PostMIPAdapter +import pytest + +import puzzletron_orchestrator.adapters.post_mip as post_mip_adapter +from puzzletron_orchestrator.adapters.post_mip import PostMIPAdapter, _run_aggregation_command from puzzletron_orchestrator.schema import ( CampaignPlan, ExecutionContract, @@ -27,6 +35,7 @@ StagePlanNode, TaskLauncher, WorkItem, + WorkPlan, ) @@ -66,6 +75,66 @@ def _plan(tmp_path: Path, *, stage_id: str, node_type: str) -> tuple[CampaignPla return plan, node +class _PostMIPExecutionContractUnavailableError(RuntimeError): + pass + + +def _candidate_count_api(candidate_count: int | Exception): + def expected_post_mip_candidate_count(_config, _stage_id): + if isinstance(candidate_count, Exception): + raise candidate_count + return candidate_count + + return SimpleNamespace( + PostMIPExecutionContractUnavailable=_PostMIPExecutionContractUnavailableError, + expected_post_mip_candidate_count=expected_post_mip_candidate_count, + ) + + +def test_post_mip_evaluation_preserves_pre_ledger_dry_run_fallback(tmp_path: Path, monkeypatch): + plan, node = _plan(tmp_path, stage_id="post.params.online_eval", node_type="evaluation") + identity_api = _candidate_count_api( + _PostMIPExecutionContractUnavailableError("post-MIP candidate registry is unavailable") + ) + monkeypatch.setattr(post_mip_adapter, "_post_mip_identity_api", lambda: identity_api) + + work_plan = PostMIPAdapter().plan(plan, node) + + assert work_plan.items[0].metadata["logical_shard_count"] == node.instances + + +@pytest.mark.parametrize( + "message", + [ + "post-MIP candidate registry does not reflect the active MIP execution", + "post-MIP inputs for 'post.params.online_eval' are unavailable", + ], + ids=("stale-registry", "missing-input"), +) +def test_post_mip_evaluation_with_existing_registry_fails_closed( + tmp_path: Path, monkeypatch, message: str +): + plan, node = _plan(tmp_path, stage_id="post.params.online_eval", node_type="evaluation") + registry = plan.puzzle_dir / "artifacts" / "post_mip" / "candidate_registry.json" + registry.parent.mkdir(parents=True) + registry.write_text("{}\n") + identity_api = _candidate_count_api(_PostMIPExecutionContractUnavailableError(message)) + monkeypatch.setattr(post_mip_adapter, "_post_mip_identity_api", lambda: identity_api) + + with pytest.raises(_PostMIPExecutionContractUnavailableError, match=message): + PostMIPAdapter().plan(plan, node) + + +def test_post_mip_evaluation_clamps_workers_to_available_candidates(tmp_path: Path, monkeypatch): + plan, node = _plan(tmp_path, stage_id="post.params.online_eval", node_type="evaluation") + identity_api = _candidate_count_api(1) + monkeypatch.setattr(post_mip_adapter, "_post_mip_identity_api", lambda: identity_api) + + work_plan = PostMIPAdapter().plan(plan, node) + + assert [item.work_id for item in work_plan.items] == [f"{node.stage_id}:0"] + + def test_post_mip_evaluation_uses_torchrun_for_single_gpu_workers(tmp_path: Path): plan, node = _plan(tmp_path, stage_id="post.params.online_eval", node_type="evaluation") attempt = PostMIPAdapter().command( @@ -139,3 +208,110 @@ 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 = [] + timeouts = [] + + def run(command, *, cwd, timeout_seconds): + assert cwd == tmp_path + commands.append(tuple(command)) + timeouts.append(timeout_seconds) + return 0, json.dumps({"status": "success"}), "" + + monkeypatch.setattr(post_mip_adapter, "_run_aggregation_command", 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 timeouts == [300.0] + assert publication is not None + assert publication.summary == {"status": "success"} + + +def test_post_mip_aggregation_timeout_is_bounded(tmp_path: Path, monkeypatch): + plan, node = _plan(tmp_path, stage_id="post.params.online_eval", node_type="evaluation") + plan = replace( + plan, + execution_defaults={ + **plan.execution_defaults, + "artifact_settling_timeout_seconds": 17, + }, + ) + + def time_out(_command, **kwargs): + assert kwargs == {"cwd": tmp_path, "timeout_seconds": 17.0} + raise TimeoutError + + monkeypatch.setattr(post_mip_adapter, "_run_aggregation_command", time_out) + + with pytest.raises(RuntimeError, match=r"aggregation timed out after 17s"): + PostMIPAdapter().aggregate( + plan=plan, + node=node, + work_plan=WorkPlan(stage_id=node.stage_id, strategy=node.strategy, items=()), + ) + + +def test_aggregation_runner_captures_output_and_return_code(tmp_path: Path): + return_code, stdout, stderr = _run_aggregation_command( + ( + sys.executable, + "-c", + "import sys; print('output'); print('error', file=sys.stderr); sys.exit(7)", + ), + cwd=tmp_path, + timeout_seconds=5, + ) + + assert (return_code, stdout, stderr) == (7, "output\n", "error\n") + + +def test_aggregation_runner_kills_a_timed_out_process(tmp_path: Path): + pid_path = tmp_path / "aggregation.pid" + + with pytest.raises(TimeoutError): + _run_aggregation_command( + ( + sys.executable, + "-c", + ( + "import os, pathlib, time; " + f"pathlib.Path({str(pid_path)!r}).write_text(str(os.getpid())); " + "time.sleep(60)" + ), + ), + cwd=tmp_path, + timeout_seconds=1, + ) + + process_id = int(pid_path.read_text()) + with pytest.raises(ProcessLookupError): + os.kill(process_id, 0) 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..98e2b502c6e --- /dev/null +++ b/tests/unit/torch/puzzletron/test_post_mip_execution_identity.py @@ -0,0 +1,643 @@ +# 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.""" + +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 import identity as post_mip_identity +from puzzletron_orchestrator.post_mip import runner as post_mip_runner +from puzzletron_orchestrator.post_mip.base import compile_post_mip_flows +from puzzletron_orchestrator.post_mip.identity import ( + PostMIPExecutionContractUnavailable, + expected_post_mip_execution_contract, +) +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 _publish_node( + ledger: CandidateLedger, + node_id: str, + input_revision_ids: tuple[str, ...], + *, + execution_identity: str, + output_revision_ids: tuple[str, ...] | None = None, + status: str = "success", + include_loss: bool = False, + flow_id: str = "params", + candidate_set_node_id: str | None = None, + producer_execution_identity: str | None = None, +) -> None: + output_revision_ids = output_revision_ids or input_revision_ids + observations = [] + for index, (input_revision_id, output_revision_id) in enumerate( + zip(input_revision_ids, output_revision_ids, strict=True) + ): + observations.append( + NodeObservation( + node_id=node_id, + input_revision_id=input_revision_id, + source_revision_id=input_revision_id, + output_revision_id=output_revision_id, + status=status, + metrics={"loss": float(index)} if include_loss else {}, + ) + ) + ledger.publish_node( + node_id, + observations, + CandidateSet.create( + flow_id, + candidate_set_node_id or node_id, + output_revision_ids, + producer_execution_identity=producer_execution_identity or execution_identity, + ), + execution_identity, + ) + + +def _materialize_revisions( + ledger: CandidateLedger, + roots: tuple[str, ...], + *, + execution_identity: str, + checkpoint_prefix: str = "/checkpoint", +) -> None: + materialized_ids = tuple( + ledger.add_revision( + architecture_id=f"architecture-{index}", + artifact_kind=ArtifactKind.CHECKPOINT, + artifact={"checkpoint": f"{checkpoint_prefix}/{index}"}, + parent_revision_id=revision_id, + producer_node="materialize", + ).revision_id + for index, revision_id in enumerate(roots) + ) + _publish_node( + ledger, + "materialize", + roots, + output_revision_ids=materialized_ids, + execution_identity=execution_identity, + ) + + +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() + _publish_node( + ledger, + "score", + roots[:2], + execution_identity="score-a", + include_loss=True, + ) + _publish_node( + ledger, + "select", + roots[:2], + execution_identity="select-a", + status="selected", + ) + _materialize_revisions(ledger, roots[:2], execution_identity="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_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_torn_active_mip_defers_execution_contract(tmp_path: Path): + config, _ledger, _roots = _identity_fixture(tmp_path) + active_path = Path(config["puzzle_dir"]) / "mip/active_profiles.json" + active_path.write_text("{") + + with pytest.raises(PostMIPExecutionContractUnavailable): + expected_post_mip_execution_contract(config, "post.params.select") + + +@pytest.mark.parametrize( + "payload", + ["{", "{}", "[]"], + ids=["torn-json", "missing-identity", "non-object"], +) +def test_unpublished_dependency_current_defers_execution_contract( + tmp_path: Path, + payload: str, +): + config, _ledger, _roots = _identity_fixture(tmp_path) + current_path = Path(config["puzzle_dir"]) / "artifacts/post_mip/nodes/materialize/current.json" + current_path.write_text(payload) + + with pytest.raises( + PostMIPExecutionContractUnavailable, + match="dependency 'materialize' has no published execution identity", + ): + expected_post_mip_execution_contract(config, "post.params.final") + + +def test_invalid_dependency_execution_identity_fails_closed(tmp_path: Path): + config, _ledger, _roots = _identity_fixture(tmp_path) + current_path = Path(config["puzzle_dir"]) / "artifacts/post_mip/nodes/materialize/current.json" + current_path.write_text('{"execution_identity": null}\n') + + with pytest.raises(ValueError, match="invalid execution identity"): + expected_post_mip_execution_contract(config, "post.params.final") + + +def test_input_candidate_set_must_match_current_producer_execution(tmp_path: Path): + config, ledger, roots = _identity_fixture(tmp_path) + _publish_node( + ledger, + "score", + roots[:2], + execution_identity="score-b", + include_loss=True, + producer_execution_identity="score-a", + ) + + with pytest.raises(RuntimeError, match="candidate set does not match its current execution"): + expected_post_mip_execution_contract(config, "post.params.select") + + +def test_dependency_identity_uses_the_validated_publication(tmp_path: Path, monkeypatch): + config, ledger, roots = _identity_fixture(tmp_path) + source_revisions = { + revision_id: ledger.source_revision(revision_id, "materialize").revision_id + for revision_id in roots[:2] + } + original = post_mip_identity._published_execution_identity + reads = [] + + def published_execution_identity(_config, owner): + reads.append(owner) + execution_identity = original(_config, owner) + if owner == "materialize": + _materialize_revisions( + ledger, + roots[:2], + execution_identity="materialize-b", + checkpoint_prefix="/checkpoint-b", + ) + return execution_identity + + monkeypatch.setattr( + post_mip_identity, + "_published_execution_identity", + published_execution_identity, + ) + + contract = expected_post_mip_execution_contract(config, "post.params.final") + + assert contract["dependency_executions"] == {"materialize": "materialize-a"} + assert contract["source_revisions"] == source_revisions + assert reads == ["materialize", "select"] + + +def test_post_mip_aggregation_reuses_one_dependency_snapshot(tmp_path: Path, monkeypatch): + config, ledger, roots = _identity_fixture(tmp_path) + original = post_mip_runner._input_set + + def input_set(snapshot, input_config, node): + candidate_set = original(snapshot, input_config, node) + _publish_node( + ledger, + "score", + roots[:2], + execution_identity="score-b", + include_loss=True, + ) + return candidate_set + + monkeypatch.setattr(post_mip_runner, "_input_set", input_set) + + summary = post_mip_runner.aggregate_post_mip_node(config, "post.params.select") + + assert summary["execution_contract"]["dependency_executions"] == {"score": "score-a"} + assert summary["execution_identity"] == post_mip_identity.post_mip_execution_contract_identity( + summary["execution_contract"] + ) + + +@pytest.mark.parametrize( + ("flow_id", "node_id"), + [("other-flow", "score"), ("params", "other-node")], + ids=["wrong-flow", "wrong-node"], +) +def test_input_candidate_set_must_match_compiled_context( + tmp_path: Path, + flow_id: str, + node_id: str, +): + config, ledger, roots = _identity_fixture(tmp_path) + execution_identity = "score-b" + _publish_node( + ledger, + "score", + roots[:2], + execution_identity=execution_identity, + include_loss=True, + flow_id=flow_id, + candidate_set_node_id=node_id, + ) + + with pytest.raises(RuntimeError, match="candidate set has invalid context"): + expected_post_mip_execution_contract(config, "post.params.select") + + +def test_incomplete_source_mapping_defers_execution_contract(tmp_path: Path): + config, _ledger, _roots = _identity_fixture(tmp_path) + observations_path = ( + Path(config["puzzle_dir"]) + / "artifacts/post_mip/nodes/materialize/executions/materialize-a/observations.json" + ) + observations = json.loads(observations_path.read_text()) + observations_path.write_text(json.dumps(observations[1:]) + "\n") + + with pytest.raises( + PostMIPExecutionContractUnavailable, + match="post-MIP source revision .* is unavailable", + ): + expected_post_mip_execution_contract(config, "post.params.final") + + +def _assert_changed_identity_resubmits( + config: dict, + changed_config: dict, + mutate=None, + *, + stage_id: str = "post.params.select", +) -> None: + plan_a, node_a = _plan(config, stage_id) + controller_a = CampaignController(plan_a, executor=object()) + adapter = adapter_for_stage(node_a) + work_plan = adapter.plan(plan_a, node_a) + for index, item in enumerate(work_plan.items): + attempt = controller_a._bind_attempt_to_stage_execution( + node_a, + work_plan, + adapter.command( + plan=plan_a, + node=node_a, + item=item, + attempt_id=f"attempt-{index}", + runner=plan_a.runner, + ), + ) + controller_a.store.save_attempt( + attempt, + JobHandle( + backend="fake", + handle_id=f"fake-attempt-{index}", + attempt_id=attempt.attempt_id, + ), + JobState.COMPLETED.value, + ) + if mutate is not None: + mutate() + + plan_b, node_b = _plan(changed_config, stage_id) + 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 + + +@pytest.mark.parametrize("change", ["config", "candidate_set"]) +def test_changed_post_mip_contract_resubmits_completed_work(tmp_path: Path, change: str): + config, ledger, roots = _identity_fixture(tmp_path) + changed_config = config + mutate = None + + if change == "config": + changed_config = copy.deepcopy(config) + changed_config["post_mip"]["flows"]["params"]["nodes"]["select"]["config"]["label"] = ( + "changed" + ) + + else: + + def replace_candidate_set() -> None: + _publish_node( + ledger, + "score", + roots[1:], + execution_identity="score-b", + include_loss=True, + ) + + mutate = replace_candidate_set + + _assert_changed_identity_resubmits(config, changed_config, mutate) + + +@pytest.mark.parametrize("change", ["dependency_execution", "source_revisions"]) +def test_changed_post_mip_lineage_resubmits_completed_work(tmp_path: Path, change: str): + config, ledger, roots = _identity_fixture(tmp_path) + + def republish_materialize() -> None: + observations = list(ledger.observations["materialize"].values()) + candidate_set = ledger.load_candidate_set("materialize") + execution_identity = "materialize-b" + output_revisions = candidate_set.revision_ids + if change == "source_revisions": + _materialize_revisions( + ledger, + roots[:2], + execution_identity=execution_identity, + checkpoint_prefix="/checkpoint/changed", + ) + return + candidate_set = CandidateSet.create( + "params", + "materialize", + output_revisions, + producer_execution_identity=execution_identity, + ) + ledger.publish_node( + "materialize", + observations, + candidate_set, + execution_identity, + ) + + before = expected_post_mip_execution_contract(config, "post.params.final") + _assert_changed_identity_resubmits( + config, + config, + republish_materialize, + stage_id="post.params.final", + ) + after = expected_post_mip_execution_contract(config, "post.params.final") + + assert before["dependency_executions"] != after["dependency_executions"] + if change == "source_revisions": + assert before["source_revisions"] != after["source_revisions"] + else: + assert before["source_revisions"] == after["source_revisions"] + + +def test_unresolved_future_post_mip_node_defers_failed_record_recovery( + tmp_path: Path, +): + 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=JobState.COMPLETED.value, + 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 47f6e17012e..4754d8dfe72 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -15,11 +15,13 @@ """Tests for post-MIP execution, including managed downstream evaluation.""" +import json from pathlib import Path from types import SimpleNamespace 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 +134,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 d9736ee7bd7..9cddc2a8bcf 100644 --- a/tests/unit/torch/puzzletron/test_stage_graph.py +++ b/tests/unit/torch/puzzletron/test_stage_graph.py @@ -143,6 +143,28 @@ def test_dynamic_stage_semantic_projection_keeps_stage_id_fallback() -> None: assert semantic_stage_config(config, "post.custom") == config +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_width_scenarios.py b/tests/unit/torch/puzzletron/test_width_scenarios.py index daa9be9a0dd..c65ca179478 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,176 @@ ) 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 _finalize_replacement_scoring(tmp_path, monkeypatch): + config_path = tmp_path / "experiment.yaml" + config_path.touch() + root_override = "+replacement_scoring.automodel.lm_head_backend=streaming" + loaded_overrides = [] + config = { + "puzzle_dir": str(tmp_path), + "model": {"path": "tiny-qwen"}, + "embedding_pruning": {"enabled": True, "widths": [256]}, + "replacement_scoring": { + "granularity": "subblock", + "automodel": {"lm_head_backend": "streaming"}, + }, + } + + def load_config(path, *, overrides=None): + assert path == config_path + loaded_overrides.extend(overrides or ()) + return config + + 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], + ) + return config, loaded_overrides, published_report + + +def test_replacement_scoring_finalizer_publishes_current_terminal_manifest(tmp_path, monkeypatch): + config, loaded_overrides, published_report = _finalize_replacement_scoring( + tmp_path, monkeypatch + ) + + manifest_path = tmp_path / "manifests" / "replacement_scoring.json" + manifest = json.loads(manifest_path.read_text()) + assert loaded_overrides == ["+replacement_scoring.automodel.lm_head_backend=streaming"] + assert published_report == {"scenario_count": 1, "widths": [256]} + assert manifest["stage"] == "replacement_scoring" + assert manifest["status"] == "success" + assert manifest["semantic_config"]["replacement_scoring"]["automodel"] == { + "lm_head_backend": "streaming" + } + assert manifest["outputs"]["report"] == published_report + assert stage_is_complete(config, "replacement_scoring") + + +@pytest.mark.parametrize( + "stale_input", + [ + "missing-summary", + "changed-identity", + "malformed-manifest", + "malformed-outputs", + "missing-report", + ], +) +def test_replacement_scoring_finalization_marker_rejects_stale_inputs( + tmp_path, monkeypatch, stale_input +): + _finalize_replacement_scoring(tmp_path, monkeypatch) + manifest_path = tmp_path / "manifests" / "replacement_scoring.json" + summary = tmp_path / "artifacts" / "replacement_scoring" / "summary.json" + + marker_a = tmp_path / "completion-a" / "finalized" + marker_a.parent.mkdir() + replacement_finalizer.write_finalization_marker(marker_a, manifest_path) + assert replacement_finalizer.finalization_marker_is_current(marker_a, manifest_path, summary) + + if stale_input == "missing-summary": + summary.unlink() + elif stale_input == "changed-identity": + manifest = json.loads(manifest_path.read_text()) + manifest["semantic_identity"] = "replacement_scoring_semantic_b" + manifest_path.write_text(json.dumps(manifest)) + elif stale_input == "malformed-manifest": + manifest_path.write_text("[]\n") + elif stale_input == "malformed-outputs": + manifest = json.loads(manifest_path.read_text()) + manifest["outputs"] = ["invalid"] + manifest_path.write_text(json.dumps(manifest)) + else: + manifest = json.loads(manifest_path.read_text()) + manifest["outputs"] = {} + manifest_path.write_text(json.dumps(manifest)) + summary.write_text("null\n") + + assert not replacement_finalizer.finalization_marker_is_current( + marker_a, manifest_path, summary + ) + + +def test_replacement_scoring_marker_uses_one_manifest_snapshot(tmp_path, monkeypatch): + _finalize_replacement_scoring(tmp_path, monkeypatch) + manifest_path = tmp_path / "manifests" / "replacement_scoring.json" + summary = tmp_path / "artifacts" / "replacement_scoring" / "summary.json" + marker = tmp_path / "completion" / "finalized" + marker.parent.mkdir() + replacement_finalizer.write_finalization_marker(marker, manifest_path) + original_read_text = Path.read_text + manifest_reads = 0 + + def republish_after_read(path, *args, **kwargs): + nonlocal manifest_reads + payload = original_read_text(path, *args, **kwargs) + if path == manifest_path: + manifest_reads += 1 + changed = json.loads(payload) + changed["semantic_identity"] = "replacement_scoring_semantic_b" + manifest_path.write_text(json.dumps(changed)) + return payload + + monkeypatch.setattr(Path, "read_text", republish_after_read) + + assert replacement_finalizer.finalization_marker_is_current(marker, manifest_path, summary) + assert manifest_reads == 1 + + +def test_replacement_scoring_finalizer_main_reads_root_overrides(monkeypatch): + captured = {} + overrides = [ + "+replacement_scoring.automodel.lm_head_backend=streaming", + "embedding_pruning.enabled=false", + ] + + def finalize(config_path, puzzle_dir, *, overrides=None): + captured.update( + config_path=config_path, + puzzle_dir=puzzle_dir, + overrides=overrides, + ) + + monkeypatch.setenv("FINALIZE_OVERRIDES", "\n".join(overrides)) + monkeypatch.setattr(replacement_finalizer, "finalize_replacement_scoring", finalize) + monkeypatch.setattr( + sys, + "argv", + [ + "finalize_replacement_scoring.py", + "--config", + "experiment.yaml", + "--puzzle-dir", + "run", + ], + ) + + replacement_finalizer.main() + + assert captured == { + "config_path": "experiment.yaml", + "puzzle_dir": "run", + "overrides": overrides, + } def _write_scenario_manifest( @@ -397,6 +568,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):