Integrate current Puzzletron v2 base

Apply the reviewed PR 2125 base integration, corrections, and review feedback as one fast-forward commit.

Signed-off-by: Johannes Rausch <jrausch@nvidia.com>

Author: j-rausch

examples/puzzletron/README.md

@@ -461,8 +461,7 @@ python examples/puzzletron/main.py \
   --gpus-per-node 8
 ```
 
-The authoritative dependencies and enabled-stage rules live in
-[`stages/graph.py`](../../modelopt/torch/puzzletron/stages/graph.py).
+The dependency-free [`StageSpec` registry](../../modelopt/torch/puzzletron/stages/graph.py) is the authoritative contract for every public stage's identity, dependencies, enablement, semantic config sections, and static completion artifacts. Add or change those properties there. Default execution strategies remain scheduler-specific and live in the [orchestration compiler](../../modelopt/torch/puzzletron/orchestration/compiler.py); handlers, scheduler adapters, mesh resolution, and heavyweight artifact validators remain separate runtime concerns.
 
 | Stage | Purpose |
 |---|---|

examples/puzzletron/embedding_pipeline.py

@@ -272,6 +272,9 @@ def run_embedding_stage(
             ),
             check=True,
         )
+        # Workers validate their scenario-local runtime statistics before
+        # building a replacement library, so seed them from the root aggregate.
+        _project_vllm_stats_to_scenarios(config)
         _run_commands(
             scenario_worker_commands(
                 config_path=config_path,

examples/puzzletron/main.py

@@ -45,13 +45,21 @@
     semantic_stage_config,
     write_stage_manifest,
 )
+from modelopt.torch.puzzletron.orchestration.adapters.stage_compat import (
+    stage_is_complete as artifacts_are_complete,
+)
+from modelopt.torch.puzzletron.orchestration.adapters.stage_compat import (
+    stage_output_patterns as canonical_stage_output_patterns,
+)
 from modelopt.torch.puzzletron.stages.graph import (
+    StageStatus,
     configured_parent_stage_ids,
     distributed_stage_ids,
     enabled_stage_ids,
     required_stage_ids,
     stage_ids,
     stage_is_enabled,
+    stage_terminal_state,
     topological_stage_ids,
 )
 
@@ -64,30 +72,6 @@
 PIPELINE_STAGE_ORDER = topological_stage_ids()
 REQUIRED_STAGES = frozenset(required_stage_ids())
 DISTRIBUTED_STAGES = frozenset(distributed_stage_ids())
-REQUIRED_OUTPUT_PATTERNS = {
-    "convert": ("ckpts/teacher/config.json",),
-    "tokenize_data": ("dataset_cache/*.tokens", "dataset_cache/*.tokens.json"),
-    "sort": ("ckpts/sorted_teacher/config.json",),
-    "slicing_sanity": (
-        "artifacts/width_slice_equivalence/manifest.json",
-        "artifacts/width_slice_equivalence/summary.json",
-        "artifacts/width_slice_equivalence/cases/**/*.json",
-        "artifacts/width_slice_equivalence/comparisons/*.safetensors",
-    ),
-    "depth_importance": ("depth/iterative/trajectory.json",),
-    "build_library": (
-        "replacement_library.json",
-        "candidate_library.json",
-        "subblock_stats.json",
-    ),
-    "vllm_stats": ("artifacts/vllm_stats/summary.json",),
-    "replacement_scoring": ("artifacts/replacement_scoring/summary.json",),
-    "mip": ("mip/**/*.json",),
-    "zero_shot_evaluation": ("artifacts/**/evaluation_summary.json",),
-    "aiperf": ("artifacts/aiperf/**/*.json",),
-    "global_distillation_sanity": ("artifacts/global_distillation_sanity/**/*.json",),
-    "global_distillation": ("artifacts/global_distillation/**/*.json",),
-}
 
 
 def _register_faulthandler() -> None:
@@ -216,50 +200,17 @@ def _report_model_name(config: dict) -> str:
     )
 
 
-def _manifest_is_complete(config: dict, stage: str) -> bool:
+def _manifest_terminal_state(config: dict, stage: str):
     puzzle_dir = Path(config.get("puzzle_dir") or (config.get("experiment") or {})["dir"])
     path = puzzle_dir / "manifests" / f"{stage}.json"
     try:
         payload = json.loads(path.read_text())
     except (OSError, ValueError):
-        return False
-    return payload.get("status") in {"success", "imported"}
-
-
-def _runtime_stats_filename(config: dict) -> str:
-    stats = config.get("vllm_stats") or {}
-    return str(stats.get("subblock_stats_filename", "subblock_stats.json"))
-
-
-def _stage_output_patterns(config: dict, stage: str) -> tuple[str, ...]:
-    if stage == "vllm_stats":
-        return (_runtime_stats_filename(config),)
-    if stage == "slicing_sanity":
-        slicing_cfg = config.get("slicing_sanity") or {}
-        if slicing_cfg.get("backend") == "distributed_parent_sweep":
-            return ("artifacts/slicing_sanity/summary.json",)
-    patterns = REQUIRED_OUTPUT_PATTERNS.get(stage, ())
-    if stage == "build_library":
-        resolved = [
-            _runtime_stats_filename(config) if pattern == "subblock_stats.json" else pattern
-            for pattern in patterns
-        ]
-        embedding = config.get("embedding_pruning") or {}
-        if bool(embedding.get("enabled", False)):
-            resolved.append("scenarios/width_scenarios.json")
-            for configured_width in embedding.get("widths", ()):
-                scenario = f"scenarios/width-{int(configured_width):04d}/depth-00"
-                resolved.extend(
-                    (
-                        f"{scenario}/scenario_manifest.json",
-                        f"{scenario}/replacement_library.json",
-                        f"{scenario}/candidate_library.json",
-                        f"{scenario}/{_runtime_stats_filename(config)}",
-                        f"{scenario}/manifests/build_library.json",
-                    )
-                )
-        return tuple(resolved)
-    return patterns
+        return None
+    state = stage_terminal_state(payload, expected_stage=stage)
+    if state is None or not state.allows_completion(stage, config):
+        return None
+    return state
 
 
 def _resume_kwargs(config: dict, config_path: str | Path, stage: str) -> dict:
@@ -280,7 +231,7 @@ def _resume_kwargs(config: dict, config_path: str | Path, stage: str) -> dict:
         "depth": None,
         "required_patterns": (
             f"manifests/{stage}.json",
-            *_stage_output_patterns(config, stage),
+            *canonical_stage_output_patterns(config, stage),
         ),
         "upstream_markers": upstream,
         "stage_config": semantic_stage_config(config, stage),
@@ -289,19 +240,76 @@ def _resume_kwargs(config: dict, config_path: str | Path, stage: str) -> dict:
 
 
 def _completion_is_valid(config: dict, config_path: str | Path, stage: str) -> bool:
-    if not _manifest_is_complete(config, stage):
+    state = _manifest_terminal_state(config, stage)
+    if state is None:
+        return False
+    if state.status is StageStatus.SKIPPED:
+        return True
+    if not artifacts_are_complete(config, stage):
         return False
     kwargs = _resume_kwargs(config, config_path, stage)
     return check_marker(marker_path(kwargs["root"], stage, None, None), **kwargs)
 
 
 def _mark_completion(config: dict, config_path: str | Path, stage: str) -> None:
-    if not _manifest_is_complete(config, stage):
-        raise RuntimeError(f"enabled stage {stage!r} did not write a successful manifest")
+    state = _manifest_terminal_state(config, stage)
+    if state is None:
+        raise RuntimeError(f"stage {stage!r} did not write an accepted terminal manifest")
+    if state.status is StageStatus.SKIPPED:
+        return
+    if not artifacts_are_complete(config, stage):
+        raise RuntimeError(f"stage {stage!r} failed canonical artifact validation")
     kwargs = _resume_kwargs(config, config_path, stage)
     write_marker(kwargs["root"], stage, build_payload(**kwargs))
 
 
+def _validate_worker_result(config: dict, result, *, expected_stage: str | None = None) -> None:
+    """Fail the worker unless its result, manifest, and required artifacts agree."""
+
+    expected_stage = expected_stage or result.stage
+    if result.stage != expected_stage:
+        raise RuntimeError(
+            f"worker stage {expected_stage!r} returned result for stage {result.stage!r}"
+        )
+    puzzle_dir = Path(config.get("puzzle_dir") or (config.get("experiment") or {})["dir"])
+    expected_manifest_path = puzzle_dir / "manifests" / f"{expected_stage}.json"
+    if Path(result.manifest_path).resolve() != expected_manifest_path.resolve():
+        raise RuntimeError(
+            f"stage {expected_stage!r} returned manifest path {result.manifest_path!s}; "
+            f"expected {expected_manifest_path!s}"
+        )
+    try:
+        payload = json.loads(expected_manifest_path.read_text())
+    except (OSError, ValueError) as exc:
+        raise RuntimeError(f"stage {expected_stage!r} wrote an unreadable manifest") from exc
+    if payload.get("stage") != expected_stage:
+        raise RuntimeError(
+            f"stage {expected_stage!r} manifest identifies stage {payload.get('stage')!r}"
+        )
+    state = stage_terminal_state(payload, expected_stage=expected_stage)
+    if state is None or not state.allows_completion(expected_stage, config):
+        raise RuntimeError(f"stage {expected_stage!r} wrote an invalid terminal manifest")
+    if result.status != state.status.value:
+        raise RuntimeError(
+            f"stage {expected_stage!r} result status {result.status!r} disagrees with "
+            f"manifest status {state.status.value!r}"
+        )
+    expected_reason = state.skip_reason.value if state.skip_reason is not None else None
+    if result.skip_reason != expected_reason:
+        raise RuntimeError(
+            f"stage {expected_stage!r} result skip reason {result.skip_reason!r} disagrees with "
+            f"manifest skip reason {expected_reason!r}"
+        )
+    if not state.produced_artifacts:
+        return
+    if not artifacts_are_complete(config, expected_stage):
+        expected = canonical_stage_output_patterns(config, expected_stage)
+        raise RuntimeError(
+            f"stage {expected_stage!r} failed canonical artifact validation; expected: "
+            + (", ".join(expected) or "stage-specific outputs")
+        )
+
+
 def run_pipeline(
     *,
     config_path: str | Path,
@@ -369,7 +377,9 @@ def _run_worker(args: argparse.Namespace) -> None:
     )
     gpus_per_node = int(args.gpus_per_node or (cfg.get("execution") or {}).get("gpus_per_node", 8))
     composite_only = {"replacement_scoring", "mip"}
-    if args.worker_stage == "tokenize_data":
+    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:
@@ -405,13 +415,15 @@ def _run_worker(args: argparse.Namespace) -> None:
             )
             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":
+            refresh_campaign_report(cfg)
+        _validate_worker_result(cfg, result, expected_stage=args.worker_stage)
     refresh_campaign_report(cfg)
     mtpz.tools.mprint(
         f"Puzzletron stage {result.stage!r} finished with status {result.status}: "
         f"{result.manifest_path}"
     )
-    if result.status not in {"success", "skipped"}:
-        raise RuntimeError(f"stage {result.stage!r} finished with status {result.status!r}")
 
 
 def _complete_composite_stage(config: dict, stage: str, outputs: dict):

examples/puzzletron/tokenize_data.py

@@ -13,79 +13,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tokenize the fixed data caches consumed by Puzzletron stages."""
+"""Optionally materialize fixed-token caches ahead of Puzzletron stages."""
 
 from __future__ import annotations
 
 import subprocess
 import sys
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING:
-    from collections.abc import Mapping
 
 from modelopt.torch.puzzletron.manifest import StageManifest, 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
 
 __all__ = ["resolve_tokenize_caches", "tokenize_data_stage"]
 
 
-def resolve_tokenize_caches(config: Mapping[str, Any]) -> list[dict[str, Any]]:
-    """Return explicit tokenize caches, or defaults from campaign token paths."""
-
-    stage_config = config.get("tokenize_data") or {}
-    caches = [dict(cache) for cache in stage_config.get("caches") or ()]
-    if caches:
-        return caches
-
-    data_cfg = config.get("data") or {}
-    calibration = data_cfg.get("calibration") or {}
-    train_samples = int(calibration.get("num_samples") or 32768)
-    train_seq = int(calibration.get("seq_len") or data_cfg.get("max_sample_length") or 4096)
-    scoring = data_cfg.get("replacement_scoring") or {}
-    val_samples = int(scoring.get("num_samples") or 128)
-    configured_seed = (config.get("pruning") or {}).get("shuffle_seed")
-    train_seed = 444 if configured_seed is None else int(configured_seed)
-
-    defaults: list[dict[str, Any]] = []
-    train_path = config.get("train_token_cache_path")
-    if train_path:
-        defaults.append(
-            {
-                "output": str(train_path),
-                "split": "train",
-                "num_samples": train_samples,
-                "seq_length": train_seq,
-                "shuffle_seed": train_seed,
-            }
-        )
-    validation_path = config.get("validation_token_cache_path")
-    if validation_path:
-        defaults.append(
-            {
-                "output": str(validation_path),
-                "split": "validation",
-                "num_samples": val_samples,
-                "seq_length": train_seq,
-                "shuffle_seed": train_seed + 1,
-            }
-        )
-    return defaults
-
-
 def tokenize_data_stage(config: dict) -> StageResult:
-    """Build every configured cache and record a normal stage manifest."""
+    """Materialize every configured fixed-token cache ahead of its consumers."""
 
     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)
-    if not bool(stage_config.get("enabled", False)):
-        manifest.complete(outputs={"enabled": False}, status="skipped")
+    if not stage_is_enabled("tokenize_data", config):
+        skip_reason = StageSkipReason.DISABLED
+        manifest.complete(
+            outputs={"enabled": False},
+            status="skipped",
+            skip_reason=skip_reason,
+        )
         write_stage_manifest(manifest_path, manifest)
         return StageResult(
-            "tokenize_data", "skipped", manifest_path, "Data tokenization is disabled."
+            "tokenize_data",
+            "skipped",
+            manifest_path,
+            "Ahead-of-time fixed-token cache materialization is disabled.",
+            skip_reason.value,
         )
 
     caches = resolve_tokenize_caches(config)