Harden Puzzletron orchestration state integrity

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:

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" \

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

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",

examples/puzzletron/finalize_replacement_scoring.py

@@ -1,53 +1,159 @@
 #!/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()
     parser.add_argument("--config", required=True)
     parser.add_argument("--puzzle-dir", required=True)
     args = parser.parse_args()
 
-    config = pipeline_config_from_path(args.config)
-    config["puzzle_dir"] = args.puzzle_dir
-    embedding = config.get("embedding_pruning") or {}
-    if bool(embedding.get("enabled", False)):
-        finalize_replacement_scoring_diagnostics(config)
-        return
-
-    puzzle_dir = Path(args.puzzle_dir)
-    scoring = config.get("replacement_scoring") or {}
-    granularity = str(scoring.get("granularity", "block"))
-    stem = (
-        "single_subblock_replacement_solutions"
-        if granularity == "subblock"
-        else "single_sequence_replacement_solutions"
-    )
-    generate_replace_block_report(
-        puzzle_dir,
-        scores_dir=puzzle_dir / f"{stem}--validation",
-        output_dir=puzzle_dir / "artifacts" / "replacement_scoring",
-        granularity=granularity,
-        default_metric=str(
-            scoring.get("default_metric", "normalized_mse_loss_hidden_states")
-        ),
-        default_layer_count=int(scoring.get("default_layer_count", 5)),
-        anchor_count=int(scoring.get("anchor_count", 3)),
-        trend_relative_tolerance=float(scoring.get("trend_relative_tolerance", 0.02)),
-    )
+    overrides = [
+        override for override in os.environ.get("FINALIZE_OVERRIDES", "").splitlines() if override
+    ]
+    finalize_replacement_scoring(args.config, args.puzzle_dir, overrides=overrides)
 
 
 if __name__ == "__main__":

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(