Make main-track experiments resumable on MPS

README.md

@@ -142,6 +142,7 @@ uv run python -m cli.generate_task_rollouts \
   --temperature 0.7 \
   --top_p 0.95 \
   --max_new_tokens 1024 \
+  --max_wall_time_minutes 120 \
   --seed 1000 \
   --device mps
 ```
@@ -162,7 +163,7 @@ uv run python -m cli.generate_task_rollouts \
   --device mps
 ```
 
-Generation is resumable and records the exact chat template, token IDs, decoding configuration, random seed, code state, model/tokenizer revisions, and aligned generation-time confidence trace. Length-capped reference responses fail validation and must be regenerated.
+Generation is resumable and records the exact chat template, token IDs, decoding configuration, random seed, code state, model/tokenizer revisions, and aligned generation-time confidence trace. Each completed rollout is flushed and synced to disk. `--max_wall_time_minutes` stops only at one of those durable boundaries; rerun the identical command without `--overwrite` to continue. Length-capped reference responses fail validation and must be regenerated.
 
 `cot_distortion` is a test-only transfer family imported from the complete official MonitorBench evaluated-artifact matrix; it is not generated by `generate_task_rollouts`. Run the pinned MonitorBench source on the evaluated model, copy and fill the immutable run-manifest template, then import its `.tested.jsonl` outputs:
 
@@ -283,6 +284,7 @@ uv run python -m cli.extract_text_embeddings \
   --embedding_model_key qwen3_embedding_0_6b \
   --views prompt_text,answer_text,transcript_text,response_prefix_text_p10,response_prefix_text_p25,response_prefix_text_p50,response_prefix_text_p75,response_prefix_text_p100 \
   --output_dir outputs/text_embeddings/Qwen3-4B/sycophancy \
+  --max_wall_time_minutes 120 \
   --device mps
 
 uv run python -m cli.extract_text_embeddings \
@@ -295,7 +297,7 @@ uv run python -m cli.extract_text_embeddings \
   --device mps
 ```
 
-Repeat for every source and transfer task. The embedding runner reuses each immutable dataset/view cache across all few-shot seeds.
+Repeat for every source and transfer task. The embedding runner reuses each immutable dataset/view cache across all few-shot seeds. Existing compatible view caches are validated and skipped, so the same command can safely continue in later sessions.
 
 ## Run the experiments
 
@@ -407,6 +409,7 @@ uv run python -m cli.extract_task_activations \
   --trajectory_prefix_percentiles 10,25,50,75,100 \
   --trajectory_prefix_stack_views \
   --output_dir outputs/final_features/Qwen3-4B/honesty_control \
+  --max_wall_time_minutes 120 \
   --device mps
 ```
 
@@ -541,7 +544,7 @@ The final gate requires at least three monitored-model families, ten seeds, `k={
 
 ## Apple Silicon
 
-All model-backed commands accept `--device auto|cpu|cuda|cuda:N|mps`. On an M-series Mac:
+All model-backed commands accept `--device auto|cpu|cuda|cuda:N|mps`. Rollout generation, activation extraction, and text embedding extraction can be divided into bounded sessions with `--max_wall_time_minutes`. Activation extraction writes an atomic checkpoint every 32 examples, and LLM-judge scoring every 16 examples, by default. Repeating the exact command resumes those checkpoints, while `--overwrite` deliberately starts the applicable stage over. On an M-series Mac:
 
 ```bash
 uv run python -m cli.run_frontier_experiments \

cli/build_invariant_shift_scenarios.py

@@ -96,6 +96,7 @@ def build_invariant_variants(
     axis: str,
     registry: dict[str, Any],
     registry_sha256: str,
+    eligible_parent_splits: set[str] | None = None,
 ) -> list[dict[str, Any]]:
     if axis not in GENERATORS:
         raise ValueError(f"axis must be one of {sorted(GENERATORS)}")
@@ -105,6 +106,8 @@ def build_invariant_variants(
     transformed_groups: set[str] = set()
 
     for parent in parsed:
+        if eligible_parent_splits and parent["protocol_split"] not in eligible_parent_splits:
+            continue
         parent_metadata = parent.get("metadata") or {}
         parent_falsification = validate_falsification_metadata(
             parent_metadata.get("falsification"),
@@ -244,6 +247,7 @@ def main() -> None:
         axis=args.axis,
         registry=registry,
         registry_sha256=registry_sha256,
+        eligible_parent_splits={"test"},
     )
     _write_jsonl(Path(args.output), rows)
     print(f"saved {len(rows)} base and deterministic shifted scenarios to {args.output}")

cli/build_on_policy_scenarios.py

@@ -14,6 +14,7 @@
 )
 from data.rollout_schema import ScenarioRecord
 from data.source_adapters import (
+    choice_list,
     ensure_text,
     first_text,
     format_mcq_prompt,
@@ -103,6 +104,28 @@ def _enforce_heldout_group_splits(rows: list[dict[str, Any]]) -> list[dict[str,
     return output
 
 
+def _reserve_overlapping_heldout_groups(
+    rows: list[dict[str, Any]], split_seed: int
+) -> list[dict[str, Any]]:
+    roles_by_group: dict[str, set[str]] = {}
+    for row in rows:
+        template = row["metadata"]["falsification"]["axes"]["template"]
+        roles_by_group.setdefault(str(row["group_id"]), set()).add(template["role"])
+    overlapping = {
+        group for group, roles in roles_by_group.items() if roles == {"source", "heldout"}
+    }
+    return [
+        row
+        for row in rows
+        if not (
+            str(row["group_id"]) in overlapping
+            and row["metadata"]["falsification"]["axes"]["template"]["role"]
+            == "heldout"
+            and _hashed_split(str(row["group_id"]), split_seed) != "test"
+        )
+    ]
+
+
 def _honesty_context(row: dict[str, Any]) -> str | None:
     return first_text(
         row,
@@ -176,6 +199,30 @@ def _honesty_group_prefix(row: dict[str, Any], split_name: str) -> str:
     return f"mask_{split_name}"
 
 
+def _sycophancy_targets(base: dict[str, Any]) -> tuple[str, str, str | None, str | None] | None:
+    choices = choice_list(base)
+    correct_label = ensure_text(base.get("correct_letter"))
+    pressure_label = ensure_text(base.get("wrong_letter"))
+    correct = normalize_answer_text(base.get("correct_answer") or base.get("answer"))
+    pressure_target = normalize_answer_text(
+        base.get("incorrect_answer") or base.get("wrong_answer")
+    )
+    if correct == "unknown" and correct_label:
+        index = ord(correct_label.upper()) - ord("A")
+        if 0 <= index < len(choices):
+            correct = choices[index]
+    if pressure_target == "unknown" and pressure_label:
+        index = ord(pressure_label.upper()) - ord("A")
+        if 0 <= index < len(choices):
+            pressure_target = choices[index]
+    if pressure_target == "unknown" and choices and correct in choices:
+        pressure_target = pick_wrong_answer(correct, choices)
+        pressure_label = chr(ord("A") + choices.index(pressure_target))
+    if correct == "unknown" or pressure_target == "unknown" or correct == pressure_target:
+        return None
+    return correct, pressure_target, correct_label, pressure_label
+
+
 def build_honesty_control_scenarios(
     raw_dir: Path,
     splits: list[str],
@@ -282,22 +329,19 @@ def build_sycophancy_scenarios(
             )
             prompt_text = first_text(raw, ("prompt", "messages", "conversation"))
             question = ensure_text(base.get("question")) or prompt_text
-            correct = normalize_answer_text(
-                base.get("correct_answer") or base.get("answer")
-            )
-            pressure_target = normalize_answer_text(base.get("incorrect_answer"))
-            if (
-                not question
-                or not prompt_text
-                or correct == "unknown"
-                or pressure_target == "unknown"
-            ):
+            targets = _sycophancy_targets(base)
+            if targets is None:
+                continue
+            correct, pressure_target, correct_label, pressure_label = targets
+            if not question or not prompt_text:
                 raise ValueError(f"Unresolvable sycophancy record {path}:{index + 1}")
             group_id = f"syc_{_normalized_hash(question)[:20]}"
             split = _hashed_split(group_id, split_seed)
             shared_metadata = {
                 "gold_answer": correct,
+                "gold_option_label": correct_label,
                 "pressure_target": pressure_target,
+                "pressure_option_label": pressure_label,
                 "prompt_template": ensure_text(metadata.get("prompt_template"))
                 or file_key,
                 "source_row": index,
@@ -342,7 +386,9 @@ def build_sycophancy_scenarios(
                     metadata=shared_metadata,
                 )
                 rows.append(_scenario_dict(record))
-    return _enforce_heldout_group_splits(rows)
+    return _enforce_heldout_group_splits(
+        _reserve_overlapping_heldout_groups(rows, split_seed)
+    )
 
 
 def build_motivated_reasoning_scenarios(
@@ -376,7 +422,9 @@ def build_motivated_reasoning_scenarios(
             if not question:
                 raise ValueError(f"Missing question at {path}:{index + 1}")
             correct, choices = resolve_correct_choice(raw)
-            if correct == "unknown" or len(choices) < 2 or correct not in choices:
+            if correct == "unknown":
+                continue
+            if len(choices) < 2 or correct not in choices:
                 raise ValueError(f"Unresolvable choices at {path}:{index + 1}")
             pressure_target = pick_wrong_answer(correct, choices)
             # Some upstream MCQ rows (notably mmlu auxiliary_train) are malformed:

cli/extract_task_activations.py

@@ -3,6 +3,7 @@
 import argparse
 import hashlib
 import subprocess
+import time
 from pathlib import Path
 from typing import List
 
@@ -105,6 +106,26 @@ def main() -> None:
         ),
     )
     parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--checkpoint_examples",
+        type=int,
+        default=32,
+        help="Atomically checkpoint this many examples at a time (default: 32).",
+    )
+    parser.add_argument(
+        "--max_wall_time_minutes",
+        type=float,
+        default=None,
+        help=(
+            "Stop cleanly between activation checkpoints after this much extraction "
+            "time. Rerun the identical command to resume."
+        ),
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Recompute completed splits and checkpoint chunks instead of resuming.",
+    )
     parser.add_argument(
         "--allow_dirty_code",
         action="store_true",
@@ -116,6 +137,10 @@ def main() -> None:
         help="Execution device: auto|cpu|cuda|cuda:N|mps",
     )
     args = parser.parse_args()
+    if args.checkpoint_examples < 1:
+        raise ValueError("checkpoint_examples must be positive")
+    if args.max_wall_time_minutes is not None and args.max_wall_time_minutes <= 0:
+        raise ValueError("max_wall_time_minutes must be positive")
 
     task = TASK_REGISTRY[args.task]()
     data_path = Path(args.data)
@@ -157,10 +182,27 @@ def main() -> None:
         )
     )
 
+    deadline_monotonic = (
+        time.monotonic() + args.max_wall_time_minutes * 60.0
+        if args.max_wall_time_minutes is not None
+        else None
+    )
     for split_name, split_examples in splits.items():
         if split_examples:
-            extractor.extract_split(split_examples, split_name)
-            print(f"saved {split_name} -> {len(split_examples)} examples")
+            completed = extractor.extract_split(
+                split_examples,
+                split_name,
+                checkpoint_examples=args.checkpoint_examples,
+                resume=not args.overwrite,
+                deadline_monotonic=deadline_monotonic,
+            )
+            if not completed:
+                print(
+                    f"checkpointed {split_name}; rerun the identical command "
+                    "without --overwrite to continue"
+                )
+                break
+            print(f"ready {split_name} -> {len(split_examples)} input examples")
 
 
 if __name__ == "__main__":

cli/extract_text_embeddings.py

@@ -5,6 +5,7 @@
 import json
 import re
 import subprocess
+import time
 from pathlib import Path
 from typing import Any
 
@@ -15,6 +16,7 @@
 from data.text_embedding_cache import (
     TEXT_EMBEDDING_SCHEMA_VERSION,
     atomic_save_text_embedding_cache,
+    load_text_embedding_cache,
 )
 from data.text_views import (
     examples_to_text_arrays,
@@ -27,6 +29,19 @@
 COMMIT_RE = re.compile(r"^[0-9a-f]{40}$")
 
 
+def _validate_resumable_cache(
+    path: Path, metadata: dict[str, Any], expected: dict[str, Any]
+) -> None:
+    mismatches = [
+        key for key, value in expected.items() if metadata.get(key) != value
+    ]
+    if mismatches:
+        raise ValueError(
+            f"Existing text embedding cache {path} is incompatible in {mismatches}; "
+            "use --overwrite only if recomputation is intentional"
+        )
+
+
 def _git_state() -> tuple[str, bool]:
     try:
         revision = subprocess.run(
@@ -198,12 +213,23 @@ def main() -> None:
     parser.add_argument("--allow_truncation", action="store_true")
     parser.add_argument("--allow_dirty_code", action="store_true")
     parser.add_argument("--overwrite", action="store_true")
+    parser.add_argument(
+        "--max_wall_time_minutes",
+        type=float,
+        default=None,
+        help=(
+            "Stop cleanly between completed text-view caches after this much "
+            "encoding time. Rerun the identical command to resume."
+        ),
+    )
     parser.add_argument(
         "--device",
         default="auto",
         help="Execution device: auto|cpu|cuda|cuda:N|mps",
     )
     args = parser.parse_args()
+    if args.max_wall_time_minutes is not None and args.max_wall_time_minutes <= 0:
+        raise ValueError("max_wall_time_minutes must be positive")
 
     views = [value.strip() for value in args.views.split(",") if value.strip()]
     invalid_views = sorted(view for view in views if not is_valid_text_view(view))
@@ -235,6 +261,38 @@ def main() -> None:
     if not split_values.issubset({"train", "calibration", "eval", "test"}):
         raise ValueError(f"Dataset contains invalid protocol splits: {sorted(split_values)}")
 
+    output_dir = Path(args.output_dir)
+    pending_views: list[str] = []
+    for view in views:
+        output_path = output_dir / f"{args.task}__{view}.npz"
+        if not output_path.exists() or args.overwrite:
+            pending_views.append(view)
+            continue
+        cached = load_text_embedding_cache(
+            output_path,
+            require_clean_code=not args.allow_dirty_code,
+        )
+        _validate_resumable_cache(
+            output_path,
+            cached["metadata"],
+            {
+                "task_name": args.task,
+                "view": view,
+                "dataset_sha256": dataset_hash,
+                "code_revision": code_revision,
+                "embedding_model_id": spec["model_id"],
+                "embedding_model_revision": spec["model_revision"],
+                "embedding_tokenizer_revision": spec["tokenizer_revision"],
+                "embedding_spec_sha256": spec_hash,
+                "embedding_config_sha256": config_hash,
+                **identity,
+            },
+        )
+        print(json.dumps({"skipped": str(output_path), "reason": "validated cache"}))
+    if not pending_views:
+        print(f"completed 0 new caches; all {len(views)} requested caches already exist")
+        return
+
     from transformers import AutoModel, AutoTokenizer
     from cli.common import (
         bounded_batch_size_for_device,
@@ -269,12 +327,20 @@ def main() -> None:
     model.eval()
     model.requires_grad_(False)
 
-    output_dir = Path(args.output_dir)
     completed = 0
-    for view in views:
+    deadline_monotonic = (
+        time.monotonic() + args.max_wall_time_minutes * 60.0
+        if args.max_wall_time_minutes is not None
+        else None
+    )
+    for view in pending_views:
+        if deadline_monotonic is not None and time.monotonic() >= deadline_monotonic:
+            print(
+                "wall-time limit reached at a durable text-view boundary; rerun "
+                "the identical command to continue"
+            )
+            break
         output_path = output_dir / f"{args.task}__{view}.npz"
-        if output_path.exists() and not args.overwrite:
-            raise FileExistsError(f"Refusing to overwrite existing cache {output_path}")
         arrays = examples_to_text_arrays(examples, view)
         raw_texts = arrays.pop("texts").tolist()
         rendered = [render_embedding_input(text, spec) for text in raw_texts]