diff --git a/docs/measurements/checkpoint-continuation-costs-20261101.md b/docs/measurements/checkpoint-continuation-costs-20261101.md new file mode 100644 index 00000000..df5ba13a --- /dev/null +++ b/docs/measurements/checkpoint-continuation-costs-20261101.md @@ -0,0 +1,151 @@ +# Distributed checkpoint continuation: agent efficiency pilot + +**This case did not demonstrate a TrainCheck debugging-efficiency benefit.** All four agents localized the same checkpoint binding problem from source and submitted a correct repair before their first remote request. Both treatment agents chose ordinary Python diagnostics; neither imported or activated TrainCheck. This measures optional tool availability, and provides no estimate of the effect or overhead of active TrainCheck use. + +All four final selections and first submissions pass independent full-state and continuation checks: the original two-rank checkpoint, two fresh two-rank seeds, and a fresh single-rank compatibility case. There are 16 passing contract reports, independently replayed locally. This is a small synthetic development pilot, not an official RSI benchmark result or a general causal estimate. + +## Measured milestones + +| Trial | Supported localization (s) | First valid submitted code (s) | Tokens through that submission | GPU-s before that submission | +| --- | ---: | ---: | ---: | ---: | +| Control 1 | 26.79 | 102.55 | 380,339 | 0 | +| Optional TrainCheck 1 | 38.01 | 111.12 | 334,186 | 0 | +| Control 2 | 34.35 | 99.39 | 304,915 | 0 | +| Optional TrainCheck 2 | 29.49 | 104.53 | 263,008 | 0 | + +Localization is the first received visible statement with a sufficient causal argument, adjudicated after all agents froze. Static source evidence qualifies. All four statements explicitly derived shard order `0,2,1,3` being bound to model order `0,1,2,3`. No later runtime result was retroactively credited. + +Times start at trial launch and are observed upper bounds: message receipt for localization, and CLI completion for the earliest submitted snapshot later passing every independent contract. Snapshot-copy timing is retained separately. These are not measurements of private thought or the exact instant of an edit. Native token counters are exact at completed-turn boundaries; localization occurred within turn 1, so its token count is only bracketed between zero and that turn’s total. We do not interpolate it. + +### Tokens through the first valid submission + +| Trial | Input | Cached input, included in input | Uncached input | Output, including reasoning | +| --- | ---: | ---: | ---: | ---: | +| Control 1 | 374,612 | 335,360 | 39,252 | 5,727 | +| Optional TrainCheck 1 | 328,810 | 278,784 | 50,026 | 5,376 | +| Control 2 | 299,253 | 251,904 | 47,349 | 5,662 | +| Optional TrainCheck 2 | 257,047 | 217,600 | 39,447 | 5,961 | + +The milestone token total is input plus output across all internal model calls, not the length of the final answer. Cached input and reasoning output are subsets and are not added again. Native `turn.completed` counters accumulate across resumed sessions; each turn’s increment was independently recomputed from the raw JSONL. Fewer total tokens cannot be treated as a monetary saving without accounting for the cache distribution and actual pricing. No dollar-cost claim is made. + +## Complete agent trial costs + +Each agent used two decision turns and one remote diagnostic request. All requests completed successfully. Control 1, treatment 1, and control 2 each executed eight distributed optimizer updates. Treatment 2 executed reference eight plus candidate eight, matching its declared total of 16. Rank replicas are not counted as extra global updates. + +| Trial | Raw wall (s) | Shared-service wait (s) | Wall excluding that wait (s) | CLI elapsed (s) | Allocated GPU-s | +| --- | ---: | ---: | ---: | ---: | ---: | +| Control 1 | 206.76 | 0.00 | 206.76 | 149.86 | 24.748 | +| Optional TrainCheck 1 | 245.48 | 48.28 | 197.20 | 146.41 | 24.330 | +| Control 2 | 189.36 | 0.00 | 189.36 | 130.85 | 24.901 | +| Optional TrainCheck 2 | 247.88 | 53.32 | 194.55 | 137.47 | 24.798 | + +| Trial | Final input | Final cached input | Final uncached input | Final output | +| --- | ---: | ---: | ---: | ---: | +| Control 1 | 613,698 | 567,424 | 46,274 | 7,457 | +| Optional TrainCheck 1 | 566,261 | 508,288 | 57,973 | 6,868 | +| Control 2 | 461,659 | 410,240 | 51,419 | 6,894 | +| Optional TrainCheck 2 | 408,396 | 360,320 | 48,076 | 6,961 | + +The shared-service lock serializes transport as well as execution. This wait is not a cluster scheduler measurement. Only the measured wait is subtracted; each agent’s own transport, setup, diagnostic work, and tool use remain counted. Reversing pair 2’s launch order did not reverse request arrival: the control reached the service first in both pairs. Raw wall times are retained to expose that asymmetry. + +CLI elapsed includes local tools, startup, model/API waiting, and response generation. The service uploads frozen source and checkpoints for each request, so transfer costs are part of this harness. GPU-s is supervised worker lifetime multiplied by allocated devices, including initialization, state exports, and diagnostics; it is not active GPU utilization or pure training time. Standard audit snapshots execute equally in both arms. + +Localization timing favored different arms in the two pairs. The optional-tool arm’s first valid submission was later in both pairs, while it consumed fewer input-plus-output tokens; uncached-token differences were mixed. Queue-adjusted final completion also favored different arms. Two pairs and zero actual TrainCheck uses do not establish a stable efficiency effect. + +## Task and controls + +- Ring18, two NVIDIA A16 devices (0 and 1), driver 535.309.01, Python 3.10.14, PyTorch 2.2.2+cu121, FP32. Shared-machine runs use isolated directories and the existing environment. +- A 10,080-parameter, four-layer Transformer: width 16, two attention heads, FF width 32, dropout 0.1, learned positional embeddings, vocabulary 32, sequence length eight. Circular previous-token sequence tagging requires cross-position information. SGD uses LR 0.15 and momentum 0.9; eight examples per rank, accumulation four, 128 tokens per distributed update. +- The previous backend saves complete per-rank snapshots. The current backend saves shared non-layer state plus rank-owned full layer bundles. Rank-major reconstruction loses logical-name association for middle layers while shapes still match. The defect is in the synthetic initial task fixture, not a newly discovered production TrainCheck defect. +- Public training seed 20261101, independent evaluation seed 777000; 120 prefix updates and eight shown continuation updates. Agents receive real curves, both implementations, and five existing checkpoint files. The prompt describes the storage change and expected continuation contract, without saying to find a bug or naming an invariant. +- Four fresh native CLI sessions, requested model `gpt-6-astra`, high effort, priority service, CLI 0.160.0. Same normal workspace-write policy, no approval bypass, no delegation, no previous chat context. Limits: 900 CLI seconds, six requests, at most 16 resumed updates and 180 seconds per two-GPU request. Arbitrary diagnostic update counts are instructional and source-audited; every actual count meets its declaration and cap. +- Controls retain ordinary logging, assertions, state helpers, and custom Python. Treatments additionally receive a generic observation manual and 31 instrumenter/support files. Twenty common packet files match byte-for-byte apart from the declared optional-capability prompt sentence; treatment packets contain 52 files. No inference engine or target checker is supplied. +- Standard requests return logs, metrics, and a terminal-state export; automatic restore-time audit snapshots remain private. Either arm can explicitly capture that same state in a diagnostic, under the same request and time budgets. Every agent chose a diagnostic. Private grading began only after all four final selections froze. + +Packet read scope is an instruction, not an OS read-isolation claim. The independent audit examined all 47 completed commands and eight edit events; it found no observed external-context, delegation, or update-budget violation. No production typing/assertion contracts were weakened. Both treatments read the manual; treatment 1 also inspected instrumenter source after localization. Source and documentation reads alone do not count as TrainCheck activation: all trace-rank directories and archived traces were checked and contained zero trace files/records. + +## Selection and calibration limits + +All 12 original severity-screen configurations failed, and every source revision, phase result, and failed screen is retained. The selection was made after calibration and before any solver access: + +| Version | Change | Prefix / continuation configurations | Severity screens passed | +| --- | --- | --- | ---: | +| v1 | Optimizer-state sharding, cloned layer initialization | 12/8, 8/8, 16/8 | 0/3 | +| v2 | Independent layer initialization | 12/8, 8/8, 16/8 | 0/3 | +| v3 | Full model-plus-optimizer layer bundles | 12/8, 8/8, 16/8 | 0/3 | +| v4 | Contextual target and positional embeddings | 32/8, 64/8, 120/8 | 0/3 | + +V4 attempt 3 was selected as a learned checkpoint-contract continuation case. It was **not reclassified as passing the original severity screen**. Its original `qualified=false` report and two failed severity flags remain unchanged; a separate, hash-bound study-selection manifest records the changed study purpose. All 29 non-severity checks pass, including reference learning, exact reference restart versus uninterrupted training, exact known-correct restoration/continuation, and the original model/optimizer mismatch with other restored state intact. + +At the restart boundary, validation loss was 0.665012464 for the reference and 0.679926716 for the original current pipeline (+2.24%, reference denominator). At update 128 it was 0.413439389 versus 0.439974099 (+6.42%, reference denominator). Intermediate curves cross. These observations support a restoration-contract case, not a large or uniformly worsening quality-regression claim. + +## Independent validation and resource scope + +The frozen evaluator compares every named model/optimizer tensor, dtype/shape, group contract, RNG state, progress/data cursor, rank/topology, and continued trajectory. It does not import candidate code. It validates frozen source and artifact manifests; restore state must be exact, while subsequent trajectories use the declared atol 2e-6 / rtol 2e-5. No record, tensor content, or meaningful identity is dropped to obtain parity. + +All four distinct implementations pass the public checkpoint and new seeds 20261117/20261129 on two ranks, plus a fresh one-rank checkpoint at seed 20261117 (evaluation seed 777001). The one-rank case does not test changing the world size of an existing two-rank checkpoint. All 16 reports replay identically locally, with 25 successful GPU phases and unchanged frozen files. Both submitted turns of every agent map to passing snapshots. The original v4 implementation fails restoration and trajectory checks under the same evaluator; the independent known-correct control passes. + +- 239 focused repository tests pass in 12.16 seconds; pinned configured hooks pass. The hook’s mypy scope skips these tools, so all 17 committed tool/test source files were explicitly checked with mypy 1.9 using `--explicit-package-bases --check-untyped-defs --disallow-untyped-defs --disallow-any-generics`. +- Operator tests separately cover artifact delivery, failure/counter accounting, the selection manifest, and milestone accounting (20 synthetic milestone tests). Native CLI initial/resume preflights verify cumulative usage. A real two-rank service preflight verifies output delivery and one complete generic TrainCheck event pair per rank, preserving 104 tensor pairs per rank plus scalar/container and RNG state in the earlier v3 setup. That preflight is not treatment-agent usage. + +| Resource category | Allocated GPU-s | +| --- | ---: | +| Four agent trials | 98.777 | +| All calibration (72 phases) | 2010.000 | +| Two service preflight phases | 49.040 | +| Private grading (25 phases) | 724.607 | + +Qualification, preflight, and private grading costs are excluded from agent efficiency totals and reported separately. Compression, transfer, local audits, and benchmark implementation work are not GPU training costs. This study does not estimate total researcher/engineering cost or a monetary bill. + +## Reproduction and evidence + +Committed workload/interfaces and measurement tools: [checkpoint pilot](../../tools/checkpoint_pilot/README.md). The executed application source is signed revision `900bd8e71015a51d6b2d092b72a6624170fa37e9`; final neutral packet documentation is signed revision `3cb6b5b2107eb6bb231ae8c16e8390cb3999cc1f`. Every candidate is identified by its full submitted-file manifest in `retrospective-grading/plan.json`. + +Operator-only `protocol-final.json`, `study-selection.json`, `prelaunch-freeze.json`, `packet-manifests.json`, `environment.json`, and `usage-preflight.json` bind the setup. `trials/` retains prompts, raw CLI JSONL, receive timestamps, exact usage, every submission, request receipt, source/checkpoint hash, log, and diagnostic artifact. `milestone-annotations.json` binds independently adjudicated visible statements and passing-grade evidence; `milestone-summary.json`, `actual-accounting-audit.json`, `metrics.csv`, and `results-summary.json` contain the reported numbers. + +Original packets are reconstructed byte-exact from immutable common sources and hash-verified inputs, with provenance recorded; live candidate edits are not substituted for initial source. All failed calibration versions and independent negative/positive controls remain in the archive. Empty trace directories are preserved to distinguish unactivated recording from missing artifacts. + +For a new run, use fresh experiment paths in the archived machine-specific scripts. The following records the actual interpreter and host separation; it is not a command to overwrite this completed study. + +On the Mac, in the experiment directory: + +```sh +EXPERIMENT_ROOT=/Users/matt/gitrepos/traincheck/outputs/ring18-checkpoint-continuation-20261101 +CHECKPOINT_LOCAL_PY=/Users/matt/miniconda3/envs/traincheck/bin/python +cd "$EXPERIMENT_ROOT" +"$CHECKPOINT_LOCAL_PY" prepare_packets.py --selection study-selection.json +MPLCONFIGDIR=/tmp/traincheck-checkpoint-mpl XDG_CACHE_HOME=/tmp/traincheck-checkpoint-cache /Users/matt/miniconda3/envs/acamate/bin/python plot_curves.py +"$CHECKPOINT_LOCAL_PY" study_driver.py +"$CHECKPOINT_LOCAL_PY" retrospective_grading.py prepare --selection study-selection.json --source "$EXPERIMENT_ROOT/qualification-source-v4/source" --protocol protocol-final.json +``` + +After uploading the generated grading directory, run scoring **on Ring18**: + +```sh +REMOTE_GRADING_ROOT=/home/yuxuan/experiments/traincheck-checkpoint-20261101/retrospective-grading +CHECKPOINT_REMOTE_PY=/home/yuxuan/miniconda3/envs/DS-1801/bin/python +cd "$REMOTE_GRADING_ROOT" +"$CHECKPOINT_REMOTE_PY" retrospective_grading.py score --root "$REMOTE_GRADING_ROOT" +``` + +After retrieving the results, replay and summarize **on the Mac**: + +```sh +EXPERIMENT_ROOT=/Users/matt/gitrepos/traincheck/outputs/ring18-checkpoint-continuation-20261101 +LOCAL_GRADING_ROOT="$EXPERIMENT_ROOT/retrospective-grading" +CHECKPOINT_LOCAL_PY=/Users/matt/miniconda3/envs/traincheck/bin/python +cd "$EXPERIMENT_ROOT" +"$CHECKPOINT_LOCAL_PY" retrospective_grading.py audit --root "$LOCAL_GRADING_ROOT" +"$CHECKPOINT_LOCAL_PY" summarize_milestones.py --study-root "$EXPERIMENT_ROOT" --annotations milestone-annotations.json --output milestone-summary-replayed.json +``` + +The operator scripts are machine-specific and archived with exact invocation receipts. Preparation uses fresh directories and refuses to overwrite prior trials. Re-execution requires new experiment paths; restoring the archive and replaying its evaluator/accounting does not require another solver or training run. + +Local evidence root: `/Users/matt/gitrepos/traincheck/outputs/ring18-checkpoint-continuation-20261101`. +Ring18 evidence root: `/home/yuxuan/experiments/traincheck-checkpoint-20261101`. + +Sealed archive: `final-evidence-v2.tar.xz`, 10,011,932 bytes, SHA-256 `829aa1424f9b1c05ccbabcb4e275411adb8d54db48b22c4e302202eb33d76b1b`. All 3,113 evidence files plus the manifest and 1,480 directories were reopened and verified. The archive and its manifest were copied to the Ring18 evidence root and their remote hashes independently verified. + +## Interpretation + +The complete experiment works, and the agents preserve the distributed continuation contract. The case still has little need for dynamic observation: its decisive logical-name association is visible in a short source fragment, and checkpoints already expose the relevant state. A useful next validation should involve a realistic cross-library or cross-rank interaction whose cause requires runtime state to establish, while retaining ordinary diagnostics and a neutral prompt. These two pairs do not justify further tracing-performance investment on their own. diff --git a/tests/test_checkpoint_accounting.py b/tests/test_checkpoint_accounting.py new file mode 100644 index 00000000..bce15d11 --- /dev/null +++ b/tests/test_checkpoint_accounting.py @@ -0,0 +1,350 @@ +"""Regressions for unknown usage, cumulative counters, and cost boundaries.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping + +import pytest + +from tools.checkpoint_pilot.accounting import ( + CompletedTurn, + FailedExecutionError, + IncompleteExecutionError, + Interval, + TokenUsage, + UsageCheckpoint, + UsageMode, + account_turn, + blocked_seconds, + gpu_allocation_seconds, + parse_exec_jsonl, +) + + +def events(*rows: Mapping[str, object]) -> str: + return "\n".join(json.dumps(dict(row)) for row in rows) + "\n" + + +def stream(usage: Mapping[str, object], thread_id: str = "session-a") -> str: + return events( + {"type": "thread.started", "thread_id": thread_id}, + {"type": "turn.started"}, + {"type": "item.completed", "item": {"type": "agent_message", "text": "done"}}, + {"type": "turn.completed", "usage": dict(usage)}, + ) + + +def minimal_usage() -> dict[str, object]: + return {"input_tokens": 100, "cached_input_tokens": 40, "output_tokens": 20} + + +def test_actual_cli_fields_preserve_usage_and_subset_accounting() -> None: + usage = { + "input_tokens": 20766, + "cached_input_tokens": 12160, + "cache_write_input_tokens": 0, + "output_tokens": 46, + "reasoning_output_tokens": 39, + } + result = parse_exec_jsonl(stream(usage, "01a1978f-ef8b-77d0-9c02-bce593f7edef")) + assert result.thread_id == "01a1978f-ef8b-77d0-9c02-bce593f7edef" + assert result.usage == TokenUsage(20766, 12160, 46, 39, 0) + assert result.usage.total_tokens == 20812 + + +def test_absent_optional_usage_is_unknown_not_zero() -> None: + result = parse_exec_jsonl(stream(minimal_usage())) + assert result.usage.reasoning_output_tokens is None + assert result.usage.cache_write_input_tokens is None + + +def test_explicit_successful_zero_usage_remains_distinct_from_unknown() -> None: + usage = {"input_tokens": 0, "cached_input_tokens": 0, "output_tokens": 0} + assert parse_exec_jsonl(stream(usage)).usage == TokenUsage(0, 0, 0) + + +def test_unknown_event_and_usage_fields_are_extensible() -> None: + usage = minimal_usage() + usage["future_counter"] = {"arbitrary": "metadata"} + text = events( + {"type": "thread.started", "thread_id": "session-a"}, + {"type": "future.event", "value": [1, 2]}, + {"type": "turn.completed", "usage": usage}, + ) + assert parse_exec_jsonl(text).usage == TokenUsage(100, 40, 20) + + +@pytest.mark.parametrize( + "field", ["input_tokens", "cached_input_tokens", "output_tokens"] +) +def test_missing_required_usage_is_rejected(field: str) -> None: + usage = minimal_usage() + del usage[field] + with pytest.raises(ValueError, match="Missing required field"): + parse_exec_jsonl(stream(usage)) + + +@pytest.mark.parametrize( + "field", + [ + "input_tokens", + "cached_input_tokens", + "output_tokens", + "reasoning_output_tokens", + "cache_write_input_tokens", + ], +) +@pytest.mark.parametrize("value", [True, -1, 1.5, "1", None]) +def test_invalid_present_usage_is_rejected(field: str, value: object) -> None: + usage = minimal_usage() + usage[field] = value + with pytest.raises(ValueError, match="nonnegative integer"): + parse_exec_jsonl(stream(usage)) + + +def test_cached_tokens_cannot_exceed_input() -> None: + usage = minimal_usage() + usage["cached_input_tokens"] = 101 + with pytest.raises(ValueError, match="exceeds input"): + parse_exec_jsonl(stream(usage)) + + +@pytest.mark.parametrize( + "text", ["", "\n", events({"type": "thread.started", "thread_id": "x"})] +) +def test_incomplete_execution_never_has_zero_usage(text: str) -> None: + with pytest.raises(IncompleteExecutionError): + parse_exec_jsonl(text) + + +@pytest.mark.parametrize("kind", ["error", "turn.failed"]) +def test_failed_execution_rejected_even_if_followed_by_completion(kind: str) -> None: + text = events( + {"type": "thread.started", "thread_id": "x"}, + {"type": kind, "message": "transport failed"}, + {"type": "turn.completed", "usage": minimal_usage()}, + ) + with pytest.raises(FailedExecutionError): + parse_exec_jsonl(text) + + +@pytest.mark.parametrize( + "text", + [ + "not json", + "[]", + events({"type": 2}), + events({"type": "thread.started"}), + events({"type": "thread.started", "thread_id": " "}), + events({"type": "turn.completed", "usage": minimal_usage()}), + events({"type": "turn.started"}), + events( + {"type": "thread.started", "thread_id": "x"}, {"type": "turn.completed"} + ), + events( + {"type": "thread.started", "thread_id": "x"}, + {"type": "turn.completed", "usage": []}, + ), + '{"type":"thread.started","thread_id":"a","thread_id":"b"}', + ], +) +def test_malformed_json_or_identity_is_rejected(text: str) -> None: + with pytest.raises(ValueError): + parse_exec_jsonl(text) + + +def test_duplicate_starts_completions_and_trailing_events_are_rejected() -> None: + start = {"type": "thread.started", "thread_id": "x"} + with pytest.raises(ValueError, match="Duplicate"): + parse_exec_jsonl(events(start, start)) + with pytest.raises(ValueError, match="Ambiguous"): + parse_exec_jsonl( + events(start, {"type": "turn.started"}, {"type": "turn.started"}) + ) + tails: tuple[dict[str, object], ...] = ( + {"type": "turn.completed", "usage": minimal_usage()}, + {"type": "future.event"}, + ) + for tail in tails: + with pytest.raises(ValueError, match="after terminal"): + parse_exec_jsonl(stream(minimal_usage()) + events(tail)) + + +@pytest.mark.parametrize("mode", ["per_turn", "cumulative"]) +def test_first_turn_requires_explicit_initial_assertion(mode: UsageMode) -> None: + turn = CompletedTurn("s", TokenUsage(100, 40, 20)) + with pytest.raises(ValueError, match="Missing prior"): + account_turn(turn, mode=mode) + first = account_turn(turn, mode=mode, initial_turn=True) + assert first.delta == turn.usage + assert first.checkpoint == UsageCheckpoint("s", mode, turn.usage) + with pytest.raises(ValueError, match="Invalid prior"): + account_turn(turn, mode=mode, previous=first.checkpoint, initial_turn=True) + + +def test_per_turn_counters_sum_even_when_second_turn_is_smaller() -> None: + first = account_turn( + CompletedTurn("s", TokenUsage(100, 40, 20, 12, 4)), + mode="per_turn", + initial_turn=True, + ) + second = account_turn( + CompletedTurn("s", TokenUsage(30, 20, 10, 7, 2)), + mode="per_turn", + previous=first.checkpoint, + ) + assert second.delta == TokenUsage(30, 20, 10, 7, 2) + assert second.checkpoint.cumulative_usage == TokenUsage(130, 60, 30, 19, 6) + + +def test_cumulative_counters_compute_validated_delta() -> None: + first = account_turn( + CompletedTurn("s", TokenUsage(100, 40, 20, 12, 4)), + mode="cumulative", + initial_turn=True, + ) + second = account_turn( + CompletedTurn("s", TokenUsage(130, 60, 30, 19, 6)), + mode="cumulative", + previous=first.checkpoint, + ) + assert second.delta == TokenUsage(30, 20, 10, 7, 2) + assert second.checkpoint.cumulative_usage == TokenUsage(130, 60, 30, 19, 6) + + +def test_actual_cli_resume_reports_cumulative_counters() -> None: + first = account_turn( + CompletedTurn( + "01a1978f-ef8b-77d0-9c02-bce593f7edef", TokenUsage(20766, 12160, 46, 39, 0) + ), + mode="cumulative", + initial_turn=True, + ) + resumed = account_turn( + CompletedTurn( + "01a1978f-ef8b-77d0-9c02-bce593f7edef", TokenUsage(42606, 32768, 52, 39, 0) + ), + mode="cumulative", + previous=first.checkpoint, + ) + assert resumed.delta == TokenUsage(21840, 20608, 6, 0, 0) + assert resumed.checkpoint.cumulative_usage.total_tokens == 42658 + + +@pytest.mark.parametrize("mode", ["per_turn", "cumulative"]) +def test_session_identity_and_mode_must_match(mode: UsageMode) -> None: + turn = CompletedTurn("s", TokenUsage(100, 40, 20)) + wrong_thread = UsageCheckpoint("different", mode, turn.usage) + with pytest.raises(ValueError, match="different thread"): + account_turn(turn, mode=mode, previous=wrong_thread) + other_mode: UsageMode = "cumulative" if mode == "per_turn" else "per_turn" + wrong_mode = UsageCheckpoint("s", other_mode, turn.usage) + with pytest.raises(ValueError, match="mode changed"): + account_turn(turn, mode=mode, previous=wrong_mode) + + +@pytest.mark.parametrize( + "current", + [ + TokenUsage(99, 40, 20, 12, 4), + TokenUsage(100, 39, 20, 12, 4), + TokenUsage(100, 40, 19, 12, 4), + TokenUsage(100, 40, 20, 11, 4), + TokenUsage(100, 40, 20, 12, 3), + ], +) +def test_decreasing_cumulative_counters_are_rejected(current: TokenUsage) -> None: + previous = UsageCheckpoint("s", "cumulative", TokenUsage(100, 40, 20, 12, 4)) + with pytest.raises(ValueError, match="decreased"): + account_turn(CompletedTurn("s", current), mode="cumulative", previous=previous) + + +def test_cumulative_cache_delta_is_also_a_subset() -> None: + previous = UsageCheckpoint("s", "cumulative", TokenUsage(100, 40, 20)) + with pytest.raises(ValueError, match="exceeds input"): + account_turn( + CompletedTurn("s", TokenUsage(101, 45, 22)), + mode="cumulative", + previous=previous, + ) + + +def test_optional_unknown_counter_cannot_be_reconstructed_from_one_endpoint() -> None: + previous = UsageCheckpoint("s", "cumulative", TokenUsage(100, 40, 20)) + current = CompletedTurn("s", TokenUsage(130, 60, 30, 7, 2)) + result = account_turn(current, mode="cumulative", previous=previous) + assert result.delta.reasoning_output_tokens is None + assert result.delta.cache_write_input_tokens is None + assert result.checkpoint.cumulative_usage.reasoning_output_tokens == 7 + per_turn = account_turn( + current, + mode="per_turn", + previous=UsageCheckpoint("s", "per_turn", previous.cumulative_usage), + ) + assert per_turn.checkpoint.cumulative_usage.reasoning_output_tokens is None + assert per_turn.checkpoint.cumulative_usage.cache_write_input_tokens is None + + +def test_optional_counter_disappearing_makes_delta_and_new_total_unknown() -> None: + previous = UsageCheckpoint("s", "cumulative", TokenUsage(100, 40, 20, 12, 4)) + result = account_turn( + CompletedTurn("s", TokenUsage(130, 60, 30)), + mode="cumulative", + previous=previous, + ) + assert result.delta.reasoning_output_tokens is None + assert result.delta.cache_write_input_tokens is None + assert result.checkpoint.cumulative_usage.reasoning_output_tokens is None + assert result.checkpoint.cumulative_usage.cache_write_input_tokens is None + + +def test_blocked_time_uses_union_including_duplicates_and_touching_intervals() -> None: + trial = Interval(10, 40) + blocked = [ + Interval(20, 25), + Interval(11, 15), + Interval(14, 19), + Interval(19, 20), + Interval(11, 15), + Interval(28, 28), + Interval(30, 32), + ] + assert blocked_seconds(trial, blocked) == 16 + assert blocked_seconds(trial, []) == 0 + assert blocked_seconds(trial, [trial]) == trial.seconds + assert trial.seconds - blocked_seconds(trial, blocked) == 14 + + +@pytest.mark.parametrize("interval", [Interval(9, 11), Interval(39, 41)]) +def test_blocked_intervals_are_not_silently_clipped(interval: Interval) -> None: + with pytest.raises(ValueError, match="outside trial"): + blocked_seconds(Interval(10, 40), [interval]) + + +@pytest.mark.parametrize( + "start,end", [(-1, 2), (2, 1), (float("nan"), 2), (0, float("inf")), (True, 2)] +) +def test_invalid_clock_intervals_are_rejected(start: float, end: float) -> None: + with pytest.raises(ValueError): + Interval(start, end) + + +@pytest.mark.parametrize("count", [0, -1, True]) +def test_gpu_count_must_be_positive_integer(count: int) -> None: + with pytest.raises(ValueError): + gpu_allocation_seconds(count, Interval(0, 1)) + + +def test_gpu_cost_is_allocated_device_time_not_request_count() -> None: + assert gpu_allocation_seconds(2, Interval(10, 40)) == 60 + assert gpu_allocation_seconds(2, Interval(10, 10)) == 0 + + +def test_unrepresentable_clock_and_gpu_durations_are_rejected() -> None: + with pytest.raises(ValueError, match="clock range"): + Interval(0, 10**1000) + with pytest.raises(ValueError, match="supported range"): + gpu_allocation_seconds(2, Interval(0, 1e308)) + with pytest.raises(ValueError, match="supported range"): + gpu_allocation_seconds(10**1000, Interval(0, 1.0)) diff --git a/tests/test_checkpoint_cli_turn.py b/tests/test_checkpoint_cli_turn.py new file mode 100644 index 00000000..6b6cc81c --- /dev/null +++ b/tests/test_checkpoint_cli_turn.py @@ -0,0 +1,264 @@ +"""Exercise CLI transport/lifecycle with a tiny local executable, not a model.""" + +from __future__ import annotations + +import json +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +from tools.checkpoint_pilot.accounting import ( + FailedExecutionError, + IncompleteExecutionError, + TokenUsage, +) +from tools.checkpoint_pilot.cli_turn import CLISettings, command, run_turn + +FAKE_PRELUDE = """ +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +prompt = sys.stdin.read() +response = Path(sys.argv[sys.argv.index('-o') + 1]) +rows = [ + {'type': 'thread.started', 'thread_id': 'session-one'}, + {'type': 'turn.started'}, + {'type': 'future.event', 'detail': 'valid extension'}, + {'type': 'turn.completed', 'usage': { + 'input_tokens': 100, 'cached_input_tokens': 40, 'output_tokens': 20, + }}, +] +""" + + +def setup(root: Path, body: str) -> tuple[CLISettings, Path, Path]: + workspace = root / "workspace" + workspace.mkdir() + schema = root / "response-schema.json" + schema.write_text('{"type":"object"}\n') + executable = root / "fake-cli" + executable.write_text(f"#!{sys.executable}\n" + FAKE_PRELUDE + body) + executable.chmod(0o700) + return ( + CLISettings(executable, "test-model", "high", "priority", schema), + workspace, + root / "output", + ) + + +def read_object(path: Path) -> dict[str, object]: + value: object = json.loads(path.read_text()) + assert isinstance(value, dict) + result: dict[str, object] = {} + for key, item in value.items(): + assert isinstance(key, str) + result[key] = item + return result + + +def number(value: object) -> float: + assert isinstance(value, (int, float)) and not isinstance(value, bool) + return float(value) + + +def test_successful_stream_records_buffered_lines_without_waiting_for_exit( + tmp_path: Path, +) -> None: + settings, workspace, output = setup( + tmp_path, + """ +response.write_text(json.dumps({'decision': 'continue', 'prompt': prompt})) +sys.stderr.write('retained stderr\\n') +sys.stdout.write(''.join(json.dumps(row) + '\\n' for row in rows)) +sys.stdout.flush() +time.sleep(0.6) +""", + ) + receipt, parsed = run_turn(settings, workspace, output, "inspect the curve", 10) + assert receipt.returncode == 0 + assert receipt.error_type is None and receipt.error_message is None + assert not receipt.timed_out + assert parsed.thread_id == "session-one" + assert parsed.usage == TokenUsage(100, 40, 20) + assert read_object(output / "response.json")["prompt"] == "inspect the curve" + assert (output / "stderr.log").read_text() == "retained stderr\n" + received = [ + read_object_line(line) + for line in (output / "received.jsonl").read_text().splitlines() + ] + assert len(received) == 4 + first = number(received[0]["received_monotonic"]) + last = number(received[-1]["received_monotonic"]) + # All events arrived in one write, well before the producer's delayed exit. + assert last - first < 0.3 + assert receipt.ended_monotonic - last > 0.4 + assert receipt.started_monotonic <= first <= last <= receipt.ended_monotonic + assert read_object(output / "receipt.json")["returncode"] == 0 + + +def read_object_line(line: str) -> dict[str, object]: + value: object = json.loads(line) + assert isinstance(value, dict) + result: dict[str, object] = {} + for key, item in value.items(): + assert isinstance(key, str) + result[key] = item + return result + + +@pytest.mark.parametrize("failure", ["incomplete", "failed", "malformed"]) +def test_failure_usage_propagates_and_retains_receipt( + tmp_path: Path, failure: str +) -> None: + bodies = { + "incomplete": "print(json.dumps(rows[0]), flush=True)\n", + "failed": "print(json.dumps(rows[0]), flush=True)\nprint(json.dumps({'type':'turn.failed','error':{'message':'failed'}}), flush=True)\n", + "malformed": "rows[-1]['usage']['input_tokens'] = True\nprint('\\n'.join(json.dumps(row) for row in rows), flush=True)\n", + } + settings, workspace, output = setup( + tmp_path, "response.write_text('{}')\n" + bodies[failure] + ) + error_type: type[ValueError] = { + "incomplete": IncompleteExecutionError, + "failed": FailedExecutionError, + "malformed": ValueError, + }[failure] + with pytest.raises(error_type): + run_turn(settings, workspace, output, "continue", 3) + receipt = read_object(output / "receipt.json") + assert receipt["returncode"] == 0 + assert receipt["error_type"] == error_type.__name__ + assert (output / "events.jsonl").stat().st_size > 0 + assert "usage" not in receipt # Unknown usage was not fabricated as zero. + + +def test_nonzero_exit_is_failure_even_if_completion_event_exists( + tmp_path: Path, +) -> None: + settings, workspace, output = setup( + tmp_path, + "response.write_text('{}')\nprint('\\n'.join(json.dumps(row) for row in rows), flush=True)\nsys.exit(7)\n", + ) + with pytest.raises(RuntimeError, match="CLI decision failed"): + run_turn(settings, workspace, output, "continue", 3) + receipt = read_object(output / "receipt.json") + assert receipt["returncode"] == 7 + assert receipt["error_type"] == "RuntimeError" + assert receipt["timed_out"] is False + + +def test_partial_line_cannot_block_past_timeout(tmp_path: Path) -> None: + settings, workspace, output = setup( + tmp_path, + 'sys.stdout.write(\'{"type":"thread.started"\')\nsys.stdout.flush()\ntime.sleep(60)\n', + ) + started = time.monotonic() + with pytest.raises(RuntimeError, match="CLI decision failed"): + run_turn(settings, workspace, output, "continue", 2) + assert time.monotonic() - started < 5 + receipt = read_object(output / "receipt.json") + assert receipt["timed_out"] is True + assert isinstance(receipt["returncode"], int) and receipt["returncode"] < 0 + assert (output / "events.jsonl").read_text() == '{"type":"thread.started"' + + +def test_timeout_stops_child_in_owned_process_group(tmp_path: Path) -> None: + settings, workspace, output = setup( + tmp_path, + """ +child = subprocess.Popen([sys.executable, '-c', 'import signal,time; signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(60)']) +Path('child.pid').write_text(str(child.pid)) +print(json.dumps(rows[0]), flush=True) +time.sleep(60) +""", + ) + with pytest.raises(RuntimeError): + run_turn(settings, workspace, output, "continue", 2) + pid = int((workspace / "child.pid").read_text()) + # A briefly unreaped zombie is terminated; it is not a running leaked child. + result = subprocess.run( + ["ps", "-o", "stat=", "-p", str(pid)], + capture_output=True, + text=True, + check=False, + ) + state = result.stdout.strip() + assert not state or state.startswith("Z"), state + assert read_object(output / "receipt.json")["timed_out"] is True + + +def test_spawn_failure_retains_unknown_exit_and_exception(tmp_path: Path) -> None: + settings, workspace, output = setup(tmp_path, "") + settings.executable.unlink() + with pytest.raises(FileNotFoundError): + run_turn(settings, workspace, output, "continue", 3) + receipt = read_object(output / "receipt.json") + assert receipt["returncode"] is None + assert receipt["error_type"] == "FileNotFoundError" + assert receipt["timed_out"] is False + + +def test_resumed_session_identity_must_match(tmp_path: Path) -> None: + settings, workspace, output = setup( + tmp_path, + "response.write_text('{}')\nprint('\\n'.join(json.dumps(row) for row in rows), flush=True)\n", + ) + with pytest.raises(ValueError, match="different thread"): + run_turn(settings, workspace, output, "continue", 3, thread_id="different") + receipt = read_object(output / "receipt.json") + assert receipt["returncode"] == 0 + assert receipt["error_type"] == "ValueError" + + +@pytest.mark.parametrize("response", [None, "[]", "{not json}"]) +def test_missing_or_malformed_structured_response_fails( + tmp_path: Path, response: str | None +) -> None: + body = "print('\\n'.join(json.dumps(row) for row in rows), flush=True)\n" + if response is not None: + body += f"response.write_text({response!r})\n" + settings, workspace, output = setup(tmp_path, body) + with pytest.raises(ValueError): + run_turn(settings, workspace, output, "continue", 3) + receipt = read_object(output / "receipt.json") + assert receipt["returncode"] == 0 + assert receipt["error_type"] in ("ValueError", "JSONDecodeError") + + +def test_initial_and_resume_commands_preserve_policy_and_model_config( + tmp_path: Path, +) -> None: + settings, workspace, output = setup(tmp_path, "") + initial = command(settings, workspace, output, None) + resumed = command(settings, workspace, output, "session-one") + for arguments in (initial, resumed): + assert arguments[arguments.index("--sandbox") + 1] == "workspace-write" + assert 'approval_policy="never"' in arguments + assert arguments[arguments.index("-m") + 1] == "test-model" + assert 'model_reasoning_effort="high"' in arguments + assert 'service_tier="priority"' in arguments + assert "--ignore-user-config" in arguments + assert "--dangerously-bypass-approvals-and-sandbox" not in arguments + assert "--dangerously-bypass-hook-trust" not in arguments + assert arguments[-1] == "-" + resume_index = resumed.index("resume") + assert resumed[resume_index + 1] == "session-one" + assert resumed.index("-m") > resume_index + assert "resume" not in initial + + +@pytest.mark.parametrize("timeout", [0, -1, float("nan"), float("inf"), True]) +def test_invalid_timeout_is_rejected_before_execution( + tmp_path: Path, timeout: float +) -> None: + settings, workspace, output = setup(tmp_path, "") + with pytest.raises(ValueError, match="timeout_seconds"): + run_turn(settings, workspace, output, "continue", timeout) + assert not output.exists() diff --git a/tests/test_checkpoint_decision.py b/tests/test_checkpoint_decision.py new file mode 100644 index 00000000..5bfc773f --- /dev/null +++ b/tests/test_checkpoint_decision.py @@ -0,0 +1,223 @@ +"""Malformed or incomplete solver decisions must not become executable requests.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from tools.checkpoint_pilot.decision import parse_decision, parse_request, read_decision + + +def valid_request() -> dict[str, object]: + return { + "kind": "resume", + "backend": "current", + "steps": 8, + "script": None, + "purpose": "Measure the revised continuation.", + } + + +def valid_decision() -> dict[str, object]: + return { + "action": "run", + "summary": "Collect evidence before selecting the next revision.", + "observations": [], + "hypotheses": [], + "next_action": "Run the current continuation.", + "request": valid_request(), + } + + +def test_empty_observations_and_hypotheses_remain_empty_not_success() -> None: + decision = parse_decision(valid_decision()) + assert decision.observations == () + assert decision.hypotheses == () + assert decision.request is not None + assert decision.request.kind == "resume" + assert decision.request.script is None + + +def test_valid_observations_preserve_text_order_and_duplicate_evidence() -> None: + payload = valid_decision() + payload["observations"] = [ + {"claim": " Measured state is available. ", "evidence": ["a.json", "a.json"]}, + {"claim": "No supporting artifact is available yet.", "evidence": []}, + ] + payload["hypotheses"] = ["A possible explanation.", "An alternative explanation."] + decision = parse_decision(payload) + assert decision.observations[0].claim == " Measured state is available. " + assert decision.observations[0].evidence == ("a.json", "a.json") + assert decision.observations[1].evidence == () + assert decision.hypotheses == ( + "A possible explanation.", + "An alternative explanation.", + ) + + +def test_finish_can_retain_implementation_without_request(tmp_path: Path) -> None: + payload = valid_decision() + payload["action"] = "finish" + payload["request"] = None + payload["next_action"] = ( + "Retain the current implementation pending further evidence." + ) + path = tmp_path / "decision.json" + path.write_text(json.dumps(payload)) + decision = read_decision(path) + assert decision.action == "finish" + assert decision.request is None + + +def test_valid_diagnostic_can_select_reference_and_nested_script() -> None: + payload = valid_request() + payload.update( + kind="diagnostic", backend="reference", steps=16, script="checks/inspect.py" + ) + request = parse_request(payload) + assert request.kind == "diagnostic" + assert request.backend == "reference" + assert request.steps == 16 + assert request.script == "checks/inspect.py" + + +@pytest.mark.parametrize("field", list(valid_decision())) +def test_missing_decision_fields_are_rejected(field: str) -> None: + payload = valid_decision() + del payload[field] + with pytest.raises(ValueError, match="missing="): + parse_decision(payload) + + +@pytest.mark.parametrize("field", list(valid_request())) +def test_missing_request_fields_are_rejected(field: str) -> None: + payload = valid_request() + del payload[field] + with pytest.raises(ValueError, match="missing="): + parse_request(payload) + + +@pytest.mark.parametrize("where", ["decision", "request", "observation"]) +def test_unknown_fields_are_rejected(where: str) -> None: + payload = valid_decision() + if where == "decision": + payload["success"] = True + elif where == "request": + request = valid_request() + request["timeout"] = 9000 + payload["request"] = request + else: + payload["observations"] = [ + {"claim": "A claim.", "evidence": [], "verified": True} + ] + with pytest.raises(ValueError, match="extra="): + parse_decision(payload) + + +@pytest.mark.parametrize( + "action,requested", [("run", None), ("finish", valid_request())] +) +def test_action_request_combinations_are_required( + action: str, requested: object +) -> None: + payload = valid_decision() + payload.update(action=action, request=requested) + with pytest.raises(ValueError, match="requires"): + parse_decision(payload) + + +@pytest.mark.parametrize("steps", [True, False, 1.0, "8", None, 0, -1, 17]) +def test_invalid_update_counts_are_rejected(steps: object) -> None: + payload = valid_request() + payload["steps"] = steps + with pytest.raises((TypeError, ValueError), match="steps"): + parse_request(payload) + + +def test_update_count_subclasses_are_not_native_json_integers() -> None: + class Count(int): + pass + + payload = valid_request() + payload["steps"] = Count(8) + with pytest.raises(TypeError, match="steps"): + parse_request(payload) + + +@pytest.mark.parametrize( + "script", + [ + None, + "", + " ", + "/tmp/check.py", + "../check.py", + "checks/../check.py", + "check.txt", + "check.py\x00", + ], +) +def test_invalid_diagnostic_paths_are_rejected(script: object) -> None: + payload = valid_request() + payload.update(kind="diagnostic", script=script) + with pytest.raises((TypeError, ValueError)): + parse_request(payload) + + +def test_resume_cannot_execute_a_diagnostic_script() -> None: + payload = valid_request() + payload["script"] = "check.py" + with pytest.raises(ValueError, match="script=null"): + parse_request(payload) + + +@pytest.mark.parametrize( + "field,value", + [("kind", "train"), ("backend", "other"), ("purpose", " "), ("purpose", False)], +) +def test_invalid_request_text_is_rejected(field: str, value: object) -> None: + payload = valid_request() + payload[field] = value + with pytest.raises((TypeError, ValueError)): + parse_request(payload) + + +@pytest.mark.parametrize( + "field,value", + [ + ("action", "inspect"), + ("summary", ""), + ("summary", 4), + ("next_action", "\n"), + ("observations", None), + ("observations", ["a claim"]), + ("observations", [{"claim": "", "evidence": []}]), + ("observations", [{"claim": "A claim", "evidence": "file.py"}]), + ("hypotheses", [""]), + ("hypotheses", [False]), + ], +) +def test_malformed_decision_content_is_rejected(field: str, value: object) -> None: + payload = valid_decision() + payload[field] = value + with pytest.raises((TypeError, ValueError)): + parse_decision(payload) + + +def test_duplicate_json_fields_are_rejected(tmp_path: Path) -> None: + path = tmp_path / "decision.json" + payload = json.dumps(valid_decision()) + path.write_text(payload[:-1] + ', "action": "finish"}') + with pytest.raises(ValueError, match="Duplicate JSON object field"): + read_decision(path) + + +def test_invalid_json_and_missing_files_propagate(tmp_path: Path) -> None: + path = tmp_path / "decision.json" + with pytest.raises(FileNotFoundError): + read_decision(path) + path.write_text("{") + with pytest.raises(json.JSONDecodeError): + read_decision(path) diff --git a/tests/test_checkpoint_evaluator.py b/tests/test_checkpoint_evaluator.py new file mode 100644 index 00000000..6e2d4565 --- /dev/null +++ b/tests/test_checkpoint_evaluator.py @@ -0,0 +1,463 @@ +"""Private evaluator regressions; never export these into solver packets.""" + +from __future__ import annotations + +import copy +import json +from dataclasses import asdict +from pathlib import Path + +import pytest +import torch + +from tools.checkpoint_pilot import evaluate as evaluator + + +def contract(world_size: int = 2) -> evaluator.Contract: + return evaluator.Contract(world_size, 17, 23, 4, 4, 2, 0.15, False) + + +def state(rank: int, world_size: int, steps: int) -> dict[str, object]: + return { + "model": { + "layers.0.weight": torch.full((2, 2), 1.0 + steps / 10), + "layers.1.weight": torch.full((2, 2), 3.0 + steps / 10), + }, + "optimizer_state": { + "layers.0.weight": {"momentum_buffer": torch.full((2, 2), 10.0)}, + "layers.1.weight": {"momentum_buffer": torch.full((2, 2), 20.0)}, + }, + "optimizer_groups": [ + { + "params": ["layers.0.weight", "layers.1.weight"], + "lr": 0.15, + "momentum": 0.9, + "dampening": 0, + "weight_decay": 0, + "nesterov": False, + "maximize": False, + "foreach": None, + }, + ], + "progress": {"completed_steps": steps, "data_cursor": steps, "seed": 17}, + "rng": { + "cpu": torch.tensor([rank + 1, 3, 7, 9], dtype=torch.uint8), + "cuda": None, + }, + "rank": rank, + "world_size": world_size, + } + + +def metrics(rank: int, world_size: int, stage: str) -> dict[str, object]: + start = 0 if stage == "prefix" else 4 + return { + "stage": stage, + "backend": "reference", + "rank": rank, + "world_size": world_size, + "seed": 17, + "eval_seed": 23, + "accumulation": 2, + "lr": 0.15, + "parameter_count": 8, + "torch_version": "2.2.2", + "device": "cpu", + "examples_per_rank_update": 8, + "tokens_per_global_update": world_size * 64, + "requested_updates": 4, + "completed_steps": start + 4, + "training_seconds_including_validation": 0.1, + "steps": [ + {"step": step, "train_token_loss": 1 / (step + 1)} + for step in range(start + 1, start + 5) + ], + "validation_curve": [ + {"step": step, "validation_loss": 1 / (step + 1)} + for step in (start, start + 4) + ], + } + + +def make_run( + root: Path, name: str, stage: str, world_size: int = 2 +) -> evaluator.RunEvidence: + source = root / name / "source" + source.mkdir(parents=True) + (source / "worker.py").write_text("# trusted fixed workload\n") + source_hashes = {"worker.py": evaluator.sha256(source / "worker.py")} + phase = root / name / "phase" + phase.mkdir() + (phase / "phase-result.json").write_text( + json.dumps( + { + "exit_code": 0, + "stop_reason": None, + "source_unchanged": True, + } + ) + ) + (phase / "invocation.json").write_text(json.dumps({"source_sha256": source_hashes})) + for rank in range(world_size): + directory = phase / "payload" / f"rank{rank}" + directory.mkdir(parents=True) + (directory / "metrics.json").write_text( + json.dumps(metrics(rank, world_size, stage)) + ) + torch.save( + state(rank, world_size, 4 if stage == "prefix" else 8), + directory / "terminal_state.pt", + ) + if stage == "resume": + torch.save(state(rank, world_size, 4), directory / "restored_state.pt") + return evaluator.RunEvidence(phase, source, source_hashes, artifact_hashes(phase)) + + +def artifact_hashes(phase: Path) -> dict[str, str]: + return { + str(path.relative_to(phase)): evaluator.sha256(path) + for path in phase.rglob("*") + if path.is_file() + } + + +def refreshed(run: evaluator.RunEvidence) -> evaluator.RunEvidence: + return evaluator.RunEvidence( + run.phase_directory, + run.source_directory, + run.source_sha256, + artifact_hashes(run.phase_directory), + ) + + +def four_runs(root: Path, world_size: int = 2) -> tuple[ + evaluator.RunEvidence, + evaluator.RunEvidence, + evaluator.RunEvidence, + evaluator.RunEvidence, +]: + return ( + make_run(root, "reference-prefix", "prefix", world_size), + make_run(root, "reference-resume", "resume", world_size), + make_run(root, "candidate-prefix", "prefix", world_size), + make_run(root, "candidate-resume", "resume", world_size), + ) + + +@pytest.mark.parametrize("world_size", [1, 2]) +def test_complete_restore_and_trajectory_pass(tmp_path: Path, world_size: int) -> None: + result = evaluator.evaluate_pair( + *four_runs(tmp_path, world_size), contract(world_size) + ) + assert result["status"] == "passed" + assert result["contract_passed"] is True + assert result["source_and_artifact_manifests_verified"] is True + assert result["ranks"] is not None + assert len(result["ranks"]) == world_size + for row in result["ranks"]: + assert all(item["passed"] for item in row["restore"]) + tensors = [item for item in row["restore"] if item["kind"] == "tensor"] + assert len(tensors) == 5 + assert all( + item["reference_sha256"] == item["candidate_sha256"] for item in tensors + ) + + +def test_same_shaped_optimizer_permutation_fails_exact_restore(tmp_path: Path) -> None: + rp, rr, cp, cr = four_runs(tmp_path) + data = state(0, 2, 4) + optimizer = evaluator.record(data["optimizer_state"], "optimizer") + optimizer["layers.0.weight"], optimizer["layers.1.weight"] = ( + optimizer["layers.1.weight"], + optimizer["layers.0.weight"], + ) + data["optimizer_state"] = optimizer + torch.save(data, cr.phase_directory / "payload/rank0/restored_state.pt") + result = evaluator.evaluate_pair(rp, rr, cp, refreshed(cr), contract()) + assert result["status"] == "violated" + assert result["prefix_passed"] is True + assert result["restore_passed"] is False + assert ( + result["trajectory_passed"] is True + ) # Deliberately equal fixture final state. + assert result["ranks"] is not None + failures = [row for row in result["ranks"][0]["restore"] if not row["passed"]] + assert len(failures) == 2 + assert all("optimizer_state" in row["path"] for row in failures) + + +@pytest.mark.parametrize( + "kind", ["rng", "cursor", "group_order", "group_lr", "reset_momentum"] +) +def test_restore_contract_includes_other_state(tmp_path: Path, kind: str) -> None: + rp, rr, cp, cr = four_runs(tmp_path) + data = state(1, 2, 4) + if kind == "rng": + data["rng"] = { + "cpu": torch.tensor([99, 3, 7, 9], dtype=torch.uint8), + "cuda": None, + } + elif kind == "cursor": + data["progress"] = {"completed_steps": 4, "data_cursor": 0, "seed": 17} + elif kind in ("group_order", "group_lr"): + groups = evaluator.array(data["optimizer_groups"], "groups") + group = evaluator.record(groups[0], "group") + if kind == "group_order": + group["params"] = ["layers.1.weight", "layers.0.weight"] + else: + group["lr"] = 0.3 + data["optimizer_groups"] = [group] + else: + data["optimizer_state"] = { + name: {"momentum_buffer": torch.zeros((2, 2))} + for name in ("layers.0.weight", "layers.1.weight") + } + torch.save(data, cr.phase_directory / "payload/rank1/restored_state.pt") + result = evaluator.evaluate_pair(rp, rr, cp, refreshed(cr), contract()) + assert result["status"] == "violated" + assert result["restore_passed"] is False + + +def test_named_mapping_reordering_is_benign_but_recorded(tmp_path: Path) -> None: + rp, rr, cp, cr = four_runs(tmp_path) + data = state(0, 2, 4) + for key in ("model", "optimizer_state"): + mapping = evaluator.record(data[key], key) + data[key] = dict(reversed(list(mapping.items()))) + torch.save(data, cr.phase_directory / "payload/rank0/restored_state.pt") + result = evaluator.evaluate_pair(rp, rr, cp, refreshed(cr), contract()) + assert result["status"] == "passed" + assert result["ranks"] is not None + keys = [ + row + for row in result["ranks"][0]["restore"] + if row["path"] == "restore['optimizer_state'].keys" + ] + assert len(keys) == 1 + assert keys[0]["reference"] != keys[0]["candidate"] + assert keys[0]["passed"] + + +@pytest.mark.parametrize("difference,expected", [(1e-7, "passed"), (0.01, "violated")]) +def test_trajectory_tolerance_separate_from_restore( + tmp_path: Path, difference: float, expected: str +) -> None: + rp, rr, cp, cr = four_runs(tmp_path) + data = state(0, 2, 8) + model = evaluator.record(data["model"], "model") + value = model["layers.0.weight"] + assert isinstance(value, torch.Tensor) + model["layers.0.weight"] = value + difference + data["model"] = model + torch.save(data, cr.phase_directory / "payload/rank0/terminal_state.pt") + result = evaluator.evaluate_pair(rp, rr, cp, refreshed(cr), contract()) + assert result["status"] == expected + assert result["restore_passed"] is True + + +def test_tiny_restore_difference_does_not_use_trajectory_tolerance( + tmp_path: Path, +) -> None: + rp, rr, cp, cr = four_runs(tmp_path) + data = state(0, 2, 4) + model = evaluator.record(data["model"], "model") + value = model["layers.0.weight"] + assert isinstance(value, torch.Tensor) + model["layers.0.weight"] = value + 1e-7 + data["model"] = model + torch.save(data, cr.phase_directory / "payload/rank0/restored_state.pt") + result = evaluator.evaluate_pair(rp, rr, cp, refreshed(cr), contract()) + assert result["status"] == "violated" + assert result["restore_passed"] is False + + +@pytest.mark.parametrize( + "kind", + [ + "missing_field", + "bad_rank", + "bad_steps", + "bool_rank", + "missing_optimizer", + "bad_rng", + "nan", + ], +) +def test_malformed_state_is_unknown_not_pass(tmp_path: Path, kind: str) -> None: + rp, rr, cp, cr = four_runs(tmp_path) + data = state(0, 2, 4) + if kind == "missing_field": + del data["progress"] + elif kind == "bad_rank": + data["rank"] = 1 + elif kind == "bool_rank": + data["rank"] = False + elif kind == "bad_steps": + data["progress"] = {"completed_steps": 8, "data_cursor": 8, "seed": 17} + elif kind == "missing_optimizer": + data["optimizer_state"] = {} + elif kind == "bad_rng": + data["rng"] = {"cpu": None, "cuda": None} + else: + data["model"] = {"layers.0.weight": torch.full((2, 2), float("nan"))} + torch.save(data, cr.phase_directory / "payload/rank0/restored_state.pt") + result = evaluator.evaluate_pair(rp, rr, cp, refreshed(cr), contract()) + assert result["status"] == "error" + assert result["contract_passed"] is None + assert result["restore_passed"] is None + assert result["ranks"] is None + assert result["error"] is not None + + +@pytest.mark.parametrize( + "kind", + [ + "failed", + "timeout", + "source_changed", + "missing_rank", + "hash_changed", + "unlisted_source", + "invocation_source", + "missing_metrics", + "extra_rank", + ], +) +def test_incomplete_or_changed_evidence_is_unknown(tmp_path: Path, kind: str) -> None: + rp, rr, cp, cr = four_runs(tmp_path) + phase = cr.phase_directory + if kind in ("failed", "timeout", "source_changed"): + (phase / "phase-result.json").write_text( + json.dumps( + { + "exit_code": 1 if kind == "failed" else 0, + "stop_reason": "timeout" if kind == "timeout" else None, + "source_unchanged": kind != "source_changed", + } + ) + ) + cr = refreshed(cr) + elif kind == "missing_rank": + (phase / "payload/rank1/restored_state.pt").unlink() + cr = refreshed(cr) + elif kind == "hash_changed": + (phase / "payload/rank0/metrics.json").write_text("{}") + elif kind == "unlisted_source": + (cr.source_directory / "new_helper.py").write_text("pass\n") + elif kind == "invocation_source": + (phase / "invocation.json").write_text( + json.dumps({"source_sha256": {"worker.py": "a" * 64}}) + ) + cr = refreshed(cr) + elif kind == "extra_rank": + (phase / "payload/rank2").mkdir() + else: + (phase / "payload/rank0/metrics.json").write_text("{}") + cr = refreshed(cr) + result = evaluator.evaluate_pair(rp, rr, cp, cr, contract()) + assert result["status"] == "error" + assert result["contract_passed"] is None + + +def test_reference_restore_mismatch_invalidates_grading(tmp_path: Path) -> None: + rp, rr, cp, cr = four_runs(tmp_path) + data = state(0, 2, 4) + data["progress"] = {"completed_steps": 4, "data_cursor": 0, "seed": 17} + torch.save(data, rr.phase_directory / "payload/rank0/restored_state.pt") + result = evaluator.evaluate_pair(rp, refreshed(rr), cp, cr, contract()) + assert result["status"] == "error" + assert result["error"] is not None + assert "Reference restoration" in result["error"]["message"] + + +def test_curve_change_cannot_pass_with_matching_terminal_state(tmp_path: Path) -> None: + rp, rr, cp, cr = four_runs(tmp_path) + path = cr.phase_directory / "payload/rank0/metrics.json" + data = metrics(0, 2, "resume") + data["steps"] = [{"step": step, "train_token_loss": 10.0} for step in range(5, 9)] + path.write_text(json.dumps(data)) + result = evaluator.evaluate_pair(rp, rr, cp, refreshed(cr), contract()) + assert result["status"] == "violated" + assert result["restore_passed"] is True + assert result["trajectory_passed"] is False + + +def test_manifest_roundtrip(tmp_path: Path) -> None: + runs = four_runs(tmp_path) + data: dict[str, object] = {"schema_version": 1, "contract": asdict(contract())} + for name, run in zip( + ( + "reference_prefix", + "reference_resume", + "candidate_prefix", + "candidate_resume", + ), + runs, + ): + data[name] = { + "phase_directory": str(run.phase_directory), + "source_directory": str(run.source_directory), + "source_sha256": run.source_sha256, + "artifact_sha256": run.artifact_sha256, + } + manifest = tmp_path / "manifest.json" + manifest.write_text(json.dumps(data)) + result = evaluator.run_from_manifest(manifest) + assert result["status"] == "passed" + malformed = copy.deepcopy(data) + malformed["contract"] = {"world_size": 2} + manifest.write_text(json.dumps(malformed)) + with pytest.raises(ValueError, match="missing required field"): + evaluator.run_from_manifest(manifest) + + +def test_exact_tensor_comparison_preserves_signed_zero() -> None: + left = torch.tensor([0.0]) + right = torch.tensor([-0.0]) + rows = evaluator.compare_states(left, right, "state", exact=True, atol=1, rtol=1) + assert rows[0]["passed"] is False + assert rows[0]["reference_sha256"] != rows[0]["candidate_sha256"] + + +def test_real_single_rank_roundtrip_matches_reference(tmp_path: Path) -> None: + from tools.checkpoint_pilot.task.engine import Options, run + + runs: list[evaluator.RunEvidence] = [] + for backend in ("reference", "current"): + checkpoint = tmp_path / f"{backend}-checkpoint" + for stage in ("prefix", "resume"): + evidence = make_run(tmp_path, f"actual-{backend}-{stage}", stage, 1) + # The helper installs protocol receipts; use a fresh payload for + # actual engine outputs so no synthetic state can satisfy this test. + payload = evidence.phase_directory / "actual-payload" + options = Options( + stage="prefix" if stage == "prefix" else "resume", + backend="reference" if backend == "reference" else "current", + checkpoint=checkpoint, + output=payload, + steps=1, + seed=17, + eval_seed=23, + lr=0.15, + accumulation=2, + device="cpu", + ) + run(options) + for generated in (payload / "rank0").iterdir(): + generated.replace( + evidence.phase_directory / "payload/rank0" / generated.name + ) + runs.append(refreshed(evidence)) + actual_contract = evaluator.Contract(1, 17, 23, 1, 1, 2, 0.15, False) + result = evaluator.evaluate_pair( + runs[0], runs[1], runs[2], runs[3], actual_contract + ) + assert result["status"] == "passed", result["error"] + + +@pytest.mark.parametrize( + "manifest", [{"../outside.py": "a" * 64}, {"worker.py": "bad"}, {}] +) +def test_manifest_rejects_paths_and_hashes(manifest: dict[str, str]) -> None: + with pytest.raises(ValueError): + evaluator.validate_manifest(manifest, "manifest") diff --git a/tests/test_checkpoint_worker.py b/tests/test_checkpoint_worker.py new file mode 100644 index 00000000..54765194 --- /dev/null +++ b/tests/test_checkpoint_worker.py @@ -0,0 +1,137 @@ +"""Exercise the real worker bootstrap without importing a training framework.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from tools.checkpoint_pilot import run_worker + +PROBE = """import json +import os +from pathlib import Path +import probe_helper + +trace = Path(os.environ['TRAINCHECK_OUTPUT_DIR']) +assert trace.is_dir(), 'trace sink must exist before module execution' +result = {'trace': str(trace), 'cwd': str(Path.cwd()), + 'rank': os.environ['RANK'], 'helper': probe_helper.VALUE, + 'route': ROUTE} +(Path(os.environ['EXPERIMENT_OUTPUT']) / ('seen-' + os.environ['RANK'] + '.json')).write_text(json.dumps(result)) +""" + + +def packet(tmp_path: Path) -> tuple[Path, Path]: + workspace = tmp_path / "packet" + workspace.mkdir() + (workspace / "task").mkdir() + (workspace / "task/__init__.py").write_text("") + (workspace / "task/engine.py").write_text("ROUTE = 'engine'\n" + PROBE) + (workspace / "diagnostic.py").write_text("ROUTE = 'diagnostic'\n" + PROBE) + (workspace / "probe_helper.py").write_text("VALUE = 'workspace-imported'\n") + return workspace, tmp_path / "output" + + +def execute( + workspace: Path, + output: Path, + rank: str | None = "0", + script: str | None = None, +) -> subprocess.CompletedProcess[str]: + environment = dict(os.environ) + environment.update( + CHECKPOINT_WORKSPACE=str(workspace), + EXPERIMENT_OUTPUT=str(output), + TRAINCHECK_OUTPUT_DIR="/stale/shared/trace", + PYTHONDONTWRITEBYTECODE="1", + ) + environment.pop("RANK", None) + environment.pop("EXPERIMENT_SCRIPT", None) + if rank is not None: + environment["RANK"] = rank + if script is not None: + environment["EXPERIMENT_SCRIPT"] = script + return subprocess.run( + [sys.executable, str(Path(run_worker.__file__).resolve())], + cwd=workspace.parent, + env=environment, + capture_output=True, + text=True, + timeout=20, + check=False, + ) + + +@pytest.mark.parametrize( + "script,route", [(None, "engine"), ("diagnostic.py", "diagnostic")] +) +def test_routes_set_unique_rank_sink_and_workspace_before_import( + tmp_path: Path, script: str | None, route: str +) -> None: + workspace, output = packet(tmp_path) + for rank in ("0", "1"): + result = execute(workspace, output, rank=rank, script=script) + assert result.returncode == 0, result.stderr + seen: object = json.loads((output / f"seen-{rank}.json").read_text()) + assert seen == { + "trace": str(output / "traces" / f"rank{rank}"), + "cwd": str(workspace), + "rank": rank, + "helper": "workspace-imported", + "route": route, + } + + +@pytest.mark.parametrize("rank", [None, "-1", "not-a-rank"]) +def test_missing_or_invalid_rank_fails_before_import( + tmp_path: Path, rank: str | None +) -> None: + workspace, output = packet(tmp_path) + result = execute(workspace, output, rank=rank) + assert result.returncode != 0 + assert not list(output.glob("seen-*.json")) + assert not (output / "traces").exists() + + +@pytest.mark.parametrize("route", ["parent", "symlink", "extension", "directory"]) +def test_diagnostic_must_resolve_to_a_packet_python_file( + tmp_path: Path, route: str +) -> None: + workspace, output = packet(tmp_path) + outside = tmp_path / "outside.py" + outside.write_text("raise RuntimeError('escaped-code-executed')\n") + if route == "parent": + script = "../outside.py" + elif route == "symlink": + (workspace / "link.py").symlink_to(outside) + script = "link.py" + elif route == "extension": + (workspace / "diagnostic.txt").write_text( + "raise RuntimeError('escaped-code-executed')\n" + ) + script = "diagnostic.txt" + else: + (workspace / "directory.py").mkdir() + (workspace / "directory.py/__main__.py").write_text( + "raise RuntimeError('escaped-code-executed')\n" + ) + script = "directory.py" + result = execute(workspace, output, script=script) + assert result.returncode != 0 + assert "Diagnostic script must be a Python file in the packet" in result.stderr + assert "escaped-code-executed" not in result.stderr + assert not list(output.glob("seen-*.json")) + + +def test_existing_rank_trace_directory_cannot_be_reused(tmp_path: Path) -> None: + workspace, output = packet(tmp_path) + (output / "traces/rank0").mkdir(parents=True) + result = execute(workspace, output) + assert result.returncode != 0 + assert "FileExistsError" in result.stderr + assert not list(output.glob("seen-*.json")) diff --git a/tests/test_checkpoint_workload.py b/tests/test_checkpoint_workload.py new file mode 100644 index 00000000..1ae71b6b --- /dev/null +++ b/tests/test_checkpoint_workload.py @@ -0,0 +1,521 @@ +"""Private workload qualification; never included in solver packets.""" + +from __future__ import annotations + +import copy +import importlib.util +import sys +from pathlib import Path +from types import FunctionType + +import pytest +import torch + +from tools.checkpoint_pilot.task import current, reference, state +from tools.checkpoint_pilot.task.engine import Options, parse_options, run, run_update +from tools.checkpoint_pilot.task.model import ( + LAYERS, + SEQUENCE_LENGTH, + VOCABULARY, + Transformer, + evaluate, + make_batch, + microbatches, +) + +CPU = torch.device("cpu") + + +def test_layers_start_with_distinct_matrix_weights_and_identical_shapes() -> None: + with torch.random.fork_rng(devices=[]): + torch.manual_seed(73) + model = Transformer() + layers = [ + state.tensors(layer.state_dict(), "layer") for layer in model.encoder.layers + ] + assert len(layers) == LAYERS + matrix_names = [name for name, value in layers[0].items() if value.ndim > 1] + assert matrix_names + for index, layer in enumerate(layers): + assert list(layer) == list(layers[0]) + for name, value in layer.items(): + assert value.shape == layers[0][name].shape + assert value.dtype == layers[0][name].dtype + for other in layers[:index]: + assert all( + not torch.equal(layer[name], other[name]) for name in matrix_names + ) + torch.manual_seed(73) + repeated = Transformer() + assert equal(model.state_dict(), repeated.state_dict()) + + +def initialized() -> tuple[Transformer, torch.optim.SGD, state.Progress]: + torch.set_num_threads(1) + torch.manual_seed(43) + model = Transformer().train() + optimizer = torch.optim.SGD(model.parameters(), lr=0.15, momentum=0.9) + for step in range(1, 4): + run_update(model, optimizer, microbatches(make_batch(43, step, 0, CPU), 2)) + return model, optimizer, state.Progress(3, 3, 43) + + +def fresh() -> tuple[Transformer, torch.optim.SGD]: + torch.manual_seed(877) + model = Transformer().train() + return model, torch.optim.SGD(model.parameters(), lr=0.15, momentum=0.9) + + +def equal(left: object, right: object) -> bool: + if isinstance(left, torch.Tensor) and isinstance(right, torch.Tensor): + return left.dtype == right.dtype and torch.equal(left, right) + if isinstance(left, dict) and isinstance(right, dict): + return list(left) == list(right) and all( + equal(left[key], right[key]) for key in left + ) + if isinstance(left, (tuple, list)) and isinstance(right, (tuple, list)): + return ( + type(left) is type(right) + and len(left) == len(right) + and all(equal(a, b) for a, b in zip(left, right, strict=True)) + ) + return type(left) is type(right) and left == right + + +def sharded_checkpoint(directory: Path, world_size: int) -> state.Artifact: + model, optimizer, progress = initialized() + directory.mkdir() + result = state.artifact(model, optimizer, progress, 0, world_size, CPU) + for rank in range(world_size): + if rank: + torch.manual_seed(100 + rank) + current.save_checkpoint( + directory, model, optimizer, progress, rank, world_size, CPU + ) + return result + + +@pytest.mark.parametrize("backend", [reference, current]) +def test_single_rank_checkpoint_round_trip_preserves_complete_state( + tmp_path: Path, backend: object +) -> None: + model, optimizer, progress = initialized() + directory = tmp_path / "checkpoint" + directory.mkdir() + before = state.artifact(model, optimizer, progress, 0, 1, CPU) + if backend is reference: + reference.save_checkpoint(directory, model, optimizer, progress, 0, 1, CPU) + else: + assert backend is current + current.save_checkpoint(directory, model, optimizer, progress, 0, 1, CPU) + resumed, resumed_optimizer = fresh() + if backend is reference: + restored = reference.restore_checkpoint( + directory, resumed, resumed_optimizer, 0, 1, CPU + ) + else: + restored = current.restore_checkpoint( + directory, resumed, resumed_optimizer, 0, 1, CPU + ) + after = state.artifact(resumed, resumed_optimizer, restored, 0, 1, CPU) + assert equal(before, after) + + +def test_two_rank_shards_permute_model_and_momentum_together( + tmp_path: Path, +) -> None: + before = sharded_checkpoint(tmp_path / "checkpoint", 2) + model, optimizer = fresh() + progress = current.restore_checkpoint( + tmp_path / "checkpoint", model, optimizer, 0, 2, CPU + ) + after = state.artifact(model, optimizer, progress, 0, 2, CPU) + preserved: list[tuple[object, object]] = [ + (before["optimizer_groups"], after["optimizer_groups"]), + (before["progress"], after["progress"]), + (before["rng"], after["rng"]), + (before["rank"], after["rank"]), + (before["world_size"], after["world_size"]), + ] + assert all(equal(left, right) for left, right in preserved) + states_before = before["optimizer_state"] + states_after = after["optimizer_state"] + for name in states_before: + source_name = name.replace("encoder.layers.1.", "encoder.layers.2.") + if name.startswith("encoder.layers.2."): + source_name = name.replace("encoder.layers.2.", "encoder.layers.1.") + assert equal(states_before[source_name], states_after[name]) + assert torch.equal(before["model"][source_name], after["model"][name]) + assert not equal(states_before, states_after) + assert not equal(before["model"], after["model"]) + + +def test_sharded_files_store_complete_model_optimizer_bundles_once( + tmp_path: Path, +) -> None: + directory = tmp_path / "checkpoint" + before = sharded_checkpoint(directory, 2) + shared = state.read(directory / "shared.pt") + assert shared["schema_version"] == 2 + assert "model" not in shared + names = state.strings(shared["nonlayer_names"], "nonlayer_names") + raw_values = shared["nonlayer_parameters"] + assert isinstance(raw_values, list) + serialized = { + name: current.read_bundle(value, name) + for name, value in zip(names, raw_values, strict=True) + } + for rank in (0, 1): + shard = state.read(directory / f"rank{rank}.pt") + assert shard["schema_version"] == 2 + layers = shard["layers"] + assert isinstance(layers, list) + for raw_layer in layers: + layer = current.read_layer(raw_layer) + for name, value in zip( + layer["parameter_names"], layer["parameters"], strict=True + ): + assert name not in serialized + serialized[name] = value + assert set(serialized) == set(before["model"]) == set(before["optimizer_state"]) + for name, bundle in serialized.items(): + assert torch.equal(bundle["model"], before["model"][name]) + assert equal(bundle["optimizer"], before["optimizer_state"][name]) + + +def test_private_name_binding_revision_restores_complete_existing_checkpoint( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + before = sharded_checkpoint(tmp_path / "checkpoint", 2) + source = Path(current.__file__).read_text() + previous = "restored.update(zip(layer_names, layer_values, strict=True))" + corrected = "restored.update(zip(chunk_names, layer_values, strict=True))" + assert source.count(previous) == 1 + path = tmp_path / "operator_reference.py" + path.write_text(source.replace(previous, corrected)) + name = "tools.checkpoint_pilot.task._operator_reference" + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + restore: object = vars(module)["restore_checkpoint"] + assert isinstance(restore, FunctionType) + model, optimizer = fresh() + progress: object = restore(tmp_path / "checkpoint", model, optimizer, 0, 2, CPU) + assert isinstance(progress, state.Progress) + after = state.artifact(model, optimizer, progress, 0, 2, CPU) + assert equal(before, after) + + +def test_each_rank_restores_its_own_rng(tmp_path: Path) -> None: + sharded_checkpoint(tmp_path / "checkpoint", 2) + recorded: list[torch.Tensor] = [] + for rank in (0, 1): + model, optimizer = fresh() + current.restore_checkpoint( + tmp_path / "checkpoint", model, optimizer, rank, 2, CPU + ) + actual = torch.get_rng_state() + expected = state.read_rng( + state.read(tmp_path / "checkpoint" / f"rank{rank}.pt")["rng"], CPU + ) + assert torch.equal(actual, expected["cpu"]) + recorded.append(actual) + assert not torch.equal(recorded[0], recorded[1]) + + +def test_named_optimizer_export_uses_live_identity_not_parameter_ordinal() -> None: + model, optimizer, progress = initialized() + before = state.artifact(model, optimizer, progress, 0, 1, CPU) + optimizer.param_groups[0]["params"].reverse() + after = state.artifact(model, optimizer, progress, 0, 1, CPU) + assert set(before["optimizer_state"]) == set(after["optimizer_state"]) + for name in before["optimizer_state"]: + assert equal(before["optimizer_state"][name], after["optimizer_state"][name]) + assert after["optimizer_groups"][0]["params"] == list( + reversed(before["optimizer_state"]) + ) + + +def test_reference_split_run_equals_uninterrupted_with_dropout(tmp_path: Path) -> None: + run( + Options( + "prefix", + "reference", + tmp_path / "short-checkpoint", + tmp_path / "short", + steps=2, + seed=53, + device="cpu", + ) + ) + run( + Options( + "resume", + "reference", + tmp_path / "short-checkpoint", + tmp_path / "resumed", + steps=3, + seed=53, + device="cpu", + ) + ) + run( + Options( + "prefix", + "reference", + tmp_path / "full-checkpoint", + tmp_path / "full", + steps=5, + seed=53, + device="cpu", + ) + ) + prefix = state.read(tmp_path / "short/rank0/terminal_state.pt") + restored = state.read(tmp_path / "resumed/rank0/restored_state.pt") + resumed = state.read(tmp_path / "resumed/rank0/terminal_state.pt") + full = state.read(tmp_path / "full/rank0/terminal_state.pt") + assert equal(prefix, restored) + assert equal(resumed, full) + + +def test_evaluation_preserves_rng_gradients_state_and_module_modes() -> None: + model, optimizer, progress = initialized() + model.embedding.eval() + before = state.artifact(model, optimizer, progress, 0, 1, CPU) + gradients = [ + parameter.grad.clone() + for parameter in model.parameters() + if parameter.grad is not None + ] + modes = [module.training for module in model.modules()] + score = evaluate(model, 777, CPU) + assert score == evaluate(model, 777, CPU) + assert equal(before, state.artifact(model, optimizer, progress, 0, 1, CPU)) + assert modes == [module.training for module in model.modules()] + assert equal(gradients, [parameter.grad for parameter in model.parameters()]) + + +def test_complete_batch_is_independent_of_microbatch_partition_and_rng() -> None: + rng = torch.get_rng_state() + batch = make_batch(7, 3, 1, CPU) + assert torch.equal(rng, torch.get_rng_state()) + assert not torch.equal(batch.tokens, make_batch(7, 3, 0, CPU).tokens) + for accumulation in (1, 2, 4, 8): + parts = microbatches(batch, accumulation) + assert torch.equal(torch.cat([part.tokens for part in parts]), batch.tokens) + assert torch.equal(torch.cat([part.labels for part in parts]), batch.labels) + + +def test_targets_require_previous_position_and_wrap_at_the_sequence_boundary() -> None: + batch = make_batch(7, 3, 1, CPU) + assert batch.tokens.shape == batch.labels.shape == (8, SEQUENCE_LENGTH) + assert batch.tokens.dtype == batch.labels.dtype == torch.int64 + assert torch.equal(batch.labels[:, 1:], batch.tokens[:, :-1]) + assert torch.equal(batch.labels[:, 0], batch.tokens[:, -1]) + assert not torch.equal(batch.labels, batch.tokens) + + +def test_model_can_distinguish_positions_for_identical_input_tokens() -> None: + # Without positional information, a permutation-equivariant encoder cannot + # generally express the fixed previous-position target. This checks that + # symmetry obstruction is removed; learning remains a qualification result. + with torch.random.fork_rng(devices=[]), torch.no_grad(): + torch.manual_seed(73) + model = Transformer().eval() + tokens = torch.zeros((2, SEQUENCE_LENGTH), dtype=torch.int64) + output = model(tokens) + assert output.shape == (2, SEQUENCE_LENGTH, VOCABULARY) + assert bool(torch.isfinite(output).all()) + assert torch.equal(output[0], output[1]) + assert not torch.allclose(output[:, 0], output[:, 1], atol=1e-6, rtol=1e-6) + + +@pytest.mark.parametrize( + "shape", [(SEQUENCE_LENGTH,), (1, SEQUENCE_LENGTH - 1), (1, SEQUENCE_LENGTH + 1)] +) +def test_model_rejects_inputs_outside_the_fixed_sequence_contract( + shape: tuple[int, ...] +) -> None: + with pytest.raises(ValueError, match="fixed-length"): + Transformer()(torch.zeros(shape, dtype=torch.int64)) + + +@pytest.mark.parametrize( + "field,value", + [ + ("world_size", True), + ("world_size", 3), + ("schema_version", 1), + ("parameter_names", []), + ("nonlayer_names", ["unknown"]), + ], +) +def test_sharded_restore_rejects_malformed_shared_metadata( + tmp_path: Path, field: str, value: object +) -> None: + directory = tmp_path / "checkpoint" + sharded_checkpoint(directory, 2) + payload = state.read(directory / "shared.pt") + payload[field] = value + torch.save(payload, directory / "shared.pt") + model, optimizer = fresh() + with pytest.raises((ValueError, TypeError)): + current.restore_checkpoint(directory, model, optimizer, 0, 2, CPU) + + +def test_missing_required_field_is_rejected(tmp_path: Path) -> None: + directory = tmp_path / "checkpoint" + sharded_checkpoint(directory, 2) + payload = state.read(directory / "shared.pt") + del payload["optimizer_groups"] + torch.save(payload, directory / "shared.pt") + model, optimizer = fresh() + with pytest.raises(ValueError, match="fields differ"): + current.restore_checkpoint(directory, model, optimizer, 0, 2, CPU) + + +def test_missing_shard_is_not_silently_replaced(tmp_path: Path) -> None: + directory = tmp_path / "checkpoint" + sharded_checkpoint(directory, 2) + (directory / "rank1.pt").unlink() + model, optimizer = fresh() + with pytest.raises(FileNotFoundError): + current.restore_checkpoint(directory, model, optimizer, 0, 2, CPU) + + +@pytest.mark.parametrize( + "mutation", + ["missing_model", "missing_optimizer", "extra", "nonfinite", "shape", "dtype"], +) +def test_malformed_parameter_bundle_is_rejected_before_restoring_model( + tmp_path: Path, mutation: str +) -> None: + directory = tmp_path / "checkpoint" + sharded_checkpoint(directory, 2) + payload = state.read(directory / "rank0.pt") + layers = payload["layers"] + assert isinstance(layers, list) + layer = state.mapping(layers[0], "layer") + values = layer["parameters"] + assert isinstance(values, list) + bundle = state.mapping(values[0], "bundle") + model_tensor = bundle["model"] + assert isinstance(model_tensor, torch.Tensor) + if mutation == "missing_model": + del bundle["model"] + elif mutation == "missing_optimizer": + del bundle["optimizer"] + elif mutation == "extra": + bundle["unrecognized"] = True + elif mutation == "nonfinite": + bundle["model"] = torch.full_like(model_tensor, float("nan")) + elif mutation == "shape": + bundle["model"] = torch.zeros(1) + else: + assert mutation == "dtype" + bundle["model"] = model_tensor.double() + values[0] = bundle + layers[0] = layer + torch.save(payload, directory / "rank0.pt") + model, optimizer = fresh() + before = copy.deepcopy(model.state_dict()) + with pytest.raises((ValueError, TypeError)): + current.restore_checkpoint(directory, model, optimizer, 0, 2, CPU) + assert equal(before, model.state_dict()) + + +def test_sharded_writer_rejects_unhandled_model_buffers(tmp_path: Path) -> None: + model, optimizer, progress = initialized() + model.register_buffer("extra_state", torch.ones(1)) + with pytest.raises(ValueError, match="every model field to be a parameter"): + current.save_checkpoint(tmp_path, model, optimizer, progress, 0, 2, CPU) + assert not list(tmp_path.iterdir()) + + +def test_duplicate_layer_and_wrong_progress_are_rejected(tmp_path: Path) -> None: + directory = tmp_path / "checkpoint" + sharded_checkpoint(directory, 2) + original = state.read(directory / "rank0.pt") + duplicate = copy.deepcopy(original) + layers = duplicate["layers"] + assert isinstance(layers, list) + layers.append(copy.deepcopy(layers[0])) + torch.save(duplicate, directory / "rank0.pt") + model, optimizer = fresh() + with pytest.raises(ValueError, match="layer coverage"): + current.restore_checkpoint(directory, model, optimizer, 0, 2, CPU) + original["progress"] = {"completed_steps": 3, "data_cursor": 2, "seed": 43} + torch.save(original, directory / "rank0.pt") + with pytest.raises(ValueError, match="one complete batch"): + current.restore_checkpoint(directory, model, optimizer, 0, 2, CPU) + + +@pytest.mark.parametrize( + "arguments", + [ + ["--steps", "0"], + ["--steps", "129"], + ["--accumulation", "3"], + ["--seed", "-1"], + ["--lr", "nan"], + ], +) +def test_cli_rejects_invalid_configuration(arguments: list[str]) -> None: + with pytest.raises(SystemExit): + parse_options( + [ + "prefix", + "--backend", + "reference", + "--checkpoint", + "/unused", + "--output", + "/unused", + *arguments, + ] + ) + + +def test_optimizer_shape_nonfinite_and_foreign_parameters_are_rejected() -> None: + model, optimizer, _ = initialized() + _, states, groups = state.optimizer_snapshot(model, optimizer) + wrong_shape = copy.deepcopy(states) + wrong_shape[0]["momentum_buffer"] = torch.zeros(1) + with pytest.raises(ValueError, match="shape or dtype"): + state.load_optimizer(model, optimizer, wrong_shape, groups) + states[0]["momentum_buffer"].fill_(float("nan")) + with pytest.raises(ValueError, match="finite"): + state.load_optimizer(model, optimizer, states, groups) + foreign = torch.optim.SGD(Transformer().parameters(), lr=0.15, momentum=0.9) + with pytest.raises(ValueError, match="outside the model"): + state.optimizer_parameter_names(model, foreign) + + +@pytest.mark.parametrize( + "field,value", + [ + ("lr", -1.0), + ("momentum", True), + ("dampening", "zero"), + ("nesterov", 1), + ("foreach", "automatic"), + ("weight_decay", float("inf")), + ], +) +def test_optimizer_group_values_are_validated(field: str, value: object) -> None: + model, optimizer, _ = initialized() + _, _, groups = state.optimizer_snapshot(model, optimizer) + fields: dict[str, object] = dict(groups[0]) + fields[field] = value + with pytest.raises((ValueError, TypeError)): + state.groups([fields], len(list(model.parameters()))) + + +def test_missing_optimizer_group_field_is_rejected_early() -> None: + model, optimizer, _ = initialized() + _, _, groups = state.optimizer_snapshot(model, optimizer) + del groups[0]["dampening"] + with pytest.raises(ValueError, match="fields differ"): + state.groups(groups, len(list(model.parameters()))) diff --git a/tools/checkpoint_pilot/DECISION_SCHEMA.json b/tools/checkpoint_pilot/DECISION_SCHEMA.json new file mode 100644 index 00000000..2e64a98a --- /dev/null +++ b/tools/checkpoint_pilot/DECISION_SCHEMA.json @@ -0,0 +1,40 @@ +{ + "type": "object", + "additionalProperties": false, + "properties": { + "action": {"type": "string", "enum": ["run", "finish"]}, + "summary": {"type": "string"}, + "observations": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": false, + "properties": { + "claim": {"type": "string"}, + "evidence": {"type": "array", "items": {"type": "string"}} + }, + "required": ["claim", "evidence"] + } + }, + "hypotheses": {"type": "array", "items": {"type": "string"}}, + "next_action": {"type": "string"}, + "request": { + "anyOf": [ + {"type": "null"}, + { + "type": "object", + "additionalProperties": false, + "properties": { + "kind": {"type": "string", "enum": ["resume", "diagnostic"]}, + "backend": {"type": "string", "enum": ["reference", "current"]}, + "steps": {"type": "integer"}, + "script": {"anyOf": [{"type": "string"}, {"type": "null"}]}, + "purpose": {"type": "string"} + }, + "required": ["kind", "backend", "steps", "script", "purpose"] + } + ] + } + }, + "required": ["action", "summary", "observations", "hypotheses", "next_action", "request"] +} diff --git a/tools/checkpoint_pilot/README.md b/tools/checkpoint_pilot/README.md new file mode 100644 index 00000000..c78924c1 --- /dev/null +++ b/tools/checkpoint_pilot/README.md @@ -0,0 +1,87 @@ +# Distributed checkpoint continuation pilot + +This is a synthetic development experiment for measuring whether optional +TrainCheck observations help an agent continue a training-engineering task. +It is not an official RSI benchmark result. It does not use invariant inference. + +The [measured report](../../docs/measurements/checkpoint-continuation-costs-20261101.md) +records the completed two-pair study, including localization, repair, token, and +GPU costs. Neither treatment agent activated TrainCheck in that study. + +The solver receives the previous and current pipelines, existing checkpoints, +and measured training/validation curves. The prompt does not identify an error +or supply a target checker. The control can use ordinary source inspection, +logging, assertions, and custom diagnostics. The treatment additionally receives +the generic TrainCheck observation guide and instrumenter source. + +## Components + +- `task/`: a fixed small Transformer sequence-tagging workload, per-rank data and + RNG streams, full-snapshot reference backend, and layer-sharded current backend. + The current backend intentionally represents the initial task implementation; + it is not a change to TrainCheck's production checkpoint behavior. +- `TASK.md` and `DECISION_SCHEMA.json`: the neutral continuation instructions and + structured run/finish interface. Standard training results and explicitly + requested diagnostic outputs have a declared, equal delivery policy. +- `decision.py`: strict validation of decisions and bounded requests. +- `run_worker.py`: per-rank output setup before candidate or TrainCheck imports. +- `cli_turn.py`: fresh/resumed native CLI execution, raw event preservation, + monotonic receive timestamps, bounded process cleanup, and failure receipts. +- `accounting.py`: validated usage counters, per-turn differences for cumulative + sessions, timing intervals, and allocated GPU accounting. Cached input and + reasoning tokens are subsets, not additional tokens. +- `evaluate.py`: operator-only evidence validation. It never imports candidate + code. It verifies source/artifact manifests and compares logically named model + and optimizer state, parameter-group contracts, RNG/progress, and continuation + trajectories on every rank. Unavailable evidence is an error, not a pass. + +## Experiment boundaries + +Freeze the task, model configuration, prompts, source hashes, evaluator, scoring +rubric, and budgets before starting fresh solver sessions. Preserve unsuccessful +qualification configurations. Give both arms the same existing checkpoints; +testing only a candidate's newly written checkpoints does not establish backward +compatibility. + +Every run or remote diagnostic consumes the same request allowance. Record +failed requests, transport/setup time, measured queue waits, and allocated GPU +seconds. Retain raw wall time and report any queue-adjusted quantity separately. +Allocated GPU seconds include worker startup and the experimental state-export +harness; they are not measured GPU utilization or pure training time. Setup, +qualification, and private evaluation costs belong in separate totals. + +Record first causal localization and the first submitted implementation that +later passes independent evaluation. A sufficient source-based explanation can +qualify as localization; a dynamic trace is not required. Turn-boundary usage is +authoritative, whereas a statement within a turn has only a usage interval. +First correct submission and the agent's final selected submission are different +milestones. Evaluate only after every solver has frozen its final selection; +private grading must not become development feedback. + +Machine-specific transport/orchestration, raw sessions, traces, checkpoints, and +their manifests belong in the isolated experiment evidence directory. Solver +packets exclude this README, the private evaluator, operator protocol, and +regression tests. Packet read-scope rules are instructions, not an OS isolation +guarantee. The CLI uses its normal workspace write restrictions. + +## Validation + +Use the repository's configured hooks and explicitly type-check these reusable +tools and their tests: the current pre-commit mypy scope does not select `tools/`. +The focused suite is `tests/test_checkpoint_*.py`. It exercises malformed +decisions, cumulative-counter errors, subprocess timeouts/cleanup, rank output +separation, checkpoint corruption and compatibility, source/artifact integrity, +and full-state comparisons. GPU qualification is separate from local tests. + +For operator scoring, run: + +```sh +python -m tools.checkpoint_pilot.evaluate --manifest evaluation-input.json \ + --output evaluation-result.json +``` + +The manifest supplies the frozen contract and four `RunEvidence` records: +reference prefix/resume and candidate prefix/resume. For an existing-checkpoint +test, candidate prefix evidence describes the original supplied checkpoint, +while candidate resume evidence describes the submitted loader. Exit codes are +0 for passed, 1 for violated, and 2 for invalid or unavailable evidence. diff --git a/tools/checkpoint_pilot/TASK.md b/tools/checkpoint_pilot/TASK.md new file mode 100644 index 00000000..07499251 --- /dev/null +++ b/tools/checkpoint_pilot/TASK.md @@ -0,0 +1,71 @@ +We changed checkpoint storage from complete model/optimizer snapshots on each +rank to shared non-layer state plus full layer shards written by the training ranks. +The aim is to reduce duplicated checkpoint state while keeping the existing +two-GPU training behavior. We expected this storage change to preserve +continuation behavior. The supplied curves compare what we observed after +restarting the previous and current pipelines. + +Continue this engineering task. Decide what to investigate or try next, implement +the most justified revision, and explain the evidence supporting your decision. +You may also retain the implementation if the evidence does not support a change. +You have both implementations and checkpoints written by each version. Preserve +the ability to resume these existing checkpoint files, the model/data/optimizer +configuration, effective batch, rank count, and public save/restore interfaces. +The purpose of restarting is to continue the same training computation. + +Your editable implementation is `task/current.py`. You may add Python diagnostic +helpers or implementation helpers imported by it. The shared model, training +engine, state utilities, reference implementation, initial checkpoints, and run +profile are fixed. Keep annotations and explicit failure checks. Do not silently +turn unavailable or malformed state into a successful result. + +Work only in this packet. Do not inspect previous conversations, other solver +packets, parent repository/history, memory, private grading, or experiment +operator files. Do not delegate to another model, launch another agent, contact +external services, or access the training host directly. This packet scope is +an experimental instruction, not a claim of OS read isolation. The operator +handles transport, Git, and final independent evaluation. + +You can inspect source and write diagnostics with local tools. The PyTorch/CUDA +runtime is on the training host; do not train a model locally. Ordinary logging, +assertions, and custom Python analyses are allowed. All remote inspections and +training use the same request allowance: six requests, each at most 16 resumed +updates and 180 seconds of allocated two-GPU process time. Failed requests count. +Initial curves and checkpoints cost neither arm a request. There is also a +15-minute aggregate CLI execution budget, including your analysis, local tool +use, and startup; shared service waiting is recorded separately. + +Each decision turn ends with the supplied JSON response schema. Record your +observations and their evidence paths, current hypotheses, and next action. +Choose `action="run"` with one request or `action="finish"` with `request=null`. +The operator snapshots your source at each submission, executes the request, +and supplies the outputs described below in the next turn. Your final implementation +is the packet's `task/current.py` and its imported helpers when you finish. + +A `resume` request selects `backend="reference"` or `"current"`, the number of +resumed `steps` (1–16), and a purpose. Set `script=null`. It starts fresh workers +on the matching supplied checkpoint using the fixed public run profile. +Standard responses contain complete logs and errors, execution receipts, +metrics and curves, and an end-of-run state export (`terminal_state.pt`). +Automatic intermediate state exports belong to independent evaluation and +are not included in standard responses. Either arm may capture intermediate +state explicitly in a diagnostic request, using ordinary Python (including +the supplied state helpers), logging, assertions, or the available observation +tools. These diagnostics use the same request and time allowance. + +A `diagnostic` request names a relative `.py` script in this packet, a backend, +and a maximum resumed update count. Two fresh workers execute that script under +`torchrun`. It can inspect actual state, run checks, and write observations to +its output directory. The environment supplies `CHECKPOINT_INPUT`, +`EXPERIMENT_OUTPUT`, `EXPERIMENT_BACKEND`, `EXPERIMENT_STEPS`, `EXPERIMENT_SEED`, +`EXPERIMENT_EVAL_SEED`, `EXPERIMENT_LR`, `EXPERIMENT_ACCUMULATION`, plus standard +`RANK`, `LOCAL_RANK`, and `WORLD_SIZE`. The packet is the working directory and +on `PYTHONPATH`. Follow the fixed data/model/update budget. Output files are +returned after the workers exit. Each worker also gets its own +`TRAINCHECK_OUTPUT_DIR` before importing your code, if diagnostics use it. + +You are not given private grading feedback during the task. Ordinary training +metrics, checkpoints, exceptions, and your diagnostic outputs remain available. +Report what actually ran and distinguish observations from hypotheses. Final +submissions will be evaluated on fresh checkpoint/training seeds and on their +compatibility with the supplied checkpoint format. diff --git a/tools/checkpoint_pilot/accounting.py b/tools/checkpoint_pilot/accounting.py new file mode 100644 index 00000000..cef89eea --- /dev/null +++ b/tools/checkpoint_pilot/accounting.py @@ -0,0 +1,351 @@ +"""Validated CLI token accounting and nonoverlapping experiment costs. + +Usage mode is an operator-verified protocol choice, never inferred from counter +sizes. Cached input is a subset of input tokens; optional counters remain unknown +when the CLI does not report them. Durations use one monotonic clock per interval +set. These helpers do not estimate token prices or infer active GPU utilization. +""" + +from __future__ import annotations + +import json +import math +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Literal + +UsageMode = Literal["per_turn", "cumulative"] + + +class IncompleteExecutionError(ValueError): + """The stream does not establish a completed turn with known usage.""" + + +class FailedExecutionError(ValueError): + """The CLI explicitly reported an execution error or failed turn.""" + + +def _count(value: object, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{field} must be a nonnegative integer") + return value + + +def _identity(value: object) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError("thread_id must be a nonempty string") + return value + + +def _mode(value: object) -> UsageMode: + if not isinstance(value, str) or value not in ("per_turn", "cumulative"): + raise ValueError("usage mode must be per_turn or cumulative") + if value == "per_turn": + return "per_turn" + return "cumulative" + + +@dataclass(frozen=True) +class TokenUsage: + input_tokens: int + cached_input_tokens: int + output_tokens: int + reasoning_output_tokens: int | None = None + cache_write_input_tokens: int | None = None + + def __post_init__(self) -> None: + for field, value in ( + ("input_tokens", self.input_tokens), + ("cached_input_tokens", self.cached_input_tokens), + ("output_tokens", self.output_tokens), + ): + _count(value, field) + if self.cached_input_tokens > self.input_tokens: + raise ValueError("cached_input_tokens exceeds input_tokens") + for field, optional_value in ( + ("reasoning_output_tokens", self.reasoning_output_tokens), + ("cache_write_input_tokens", self.cache_write_input_tokens), + ): + if optional_value is not None: + _count(optional_value, field) + + @property + def total_tokens(self) -> int: + """Reported input plus output; do not add cached or reasoning subsets.""" + return self.input_tokens + self.output_tokens + + +@dataclass(frozen=True) +class CompletedTurn: + thread_id: str + usage: TokenUsage + + def __post_init__(self) -> None: + _identity(self.thread_id) + if not isinstance(self.usage, TokenUsage): + raise ValueError("usage must be validated TokenUsage") + + +@dataclass(frozen=True) +class UsageCheckpoint: + thread_id: str + mode: UsageMode + cumulative_usage: TokenUsage + + def __post_init__(self) -> None: + _identity(self.thread_id) + _mode(self.mode) + if not isinstance(self.cumulative_usage, TokenUsage): + raise ValueError("cumulative_usage must be validated TokenUsage") + + +@dataclass(frozen=True) +class AccountedTurn: + delta: TokenUsage + checkpoint: UsageCheckpoint + + +def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"Duplicate JSON field: {key}") + result[key] = value + return result + + +def _object(value: object, context: str) -> dict[str, object]: + if not isinstance(value, dict): + raise ValueError(f"{context} must be a JSON object") + result: dict[str, object] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise ValueError(f"{context} must have string keys") + result[key] = item + return result + + +def _required(record: dict[str, object], field: str) -> object: + if field not in record: + raise ValueError(f"Missing required field: {field}") + return record[field] + + +def _optional_count(record: dict[str, object], field: str) -> int | None: + # Absence means the provider did not expose this counter. A present null or + # malformed value is not a supported substitute for an integer counter. + return _count(record[field], field) if field in record else None + + +def parse_exec_jsonl(text: str) -> CompletedTurn: + """Parse one CLI exec/resume invocation, ending in exactly one completion. + + Unknown event kinds are accepted before completion for extensibility. Error + events, duplicate starts/completions, trailing events, and malformed usage + never become successful zero-cost turns. Raw JSONL should also be retained. + """ + thread_id: str | None = None + completed: CompletedTurn | None = None + started = False + for line_number, line in enumerate(text.splitlines(), start=1): + if not line.strip(): + continue + value: object = json.loads(line, object_pairs_hook=_unique_object) + event = _object(value, f"line {line_number}") + kind = _required(event, "type") + if not isinstance(kind, str) or not kind: + raise ValueError(f"line {line_number}: type must be a nonempty string") + if completed is not None: + raise ValueError("Event after terminal turn.completed") + if kind in ("error", "turn.failed"): + raise FailedExecutionError(f"CLI reported {kind} at line {line_number}") + if kind == "thread.started": + if thread_id is not None or started: + raise ValueError("Duplicate or late thread.started") + thread_id = _identity(_required(event, "thread_id")) + elif kind == "turn.started": + if thread_id is None or started: + raise ValueError("Ambiguous turn.started or missing thread.started") + started = True + elif kind == "turn.completed": + if thread_id is None: + raise ValueError("Completion has no thread.started identity") + usage = _object(_required(event, "usage"), "usage") + completed = CompletedTurn( + thread_id, + TokenUsage( + input_tokens=_count( + _required(usage, "input_tokens"), "input_tokens" + ), + cached_input_tokens=_count( + _required(usage, "cached_input_tokens"), "cached_input_tokens" + ), + output_tokens=_count( + _required(usage, "output_tokens"), "output_tokens" + ), + reasoning_output_tokens=_optional_count( + usage, "reasoning_output_tokens" + ), + cache_write_input_tokens=_optional_count( + usage, "cache_write_input_tokens" + ), + ), + ) + if completed is None: + raise IncompleteExecutionError("No terminal turn.completed with known usage") + return completed + + +def _optional_sum(left: int | None, right: int | None) -> int | None: + return None if left is None or right is None else left + right + + +def _difference(current: int, previous: int, field: str) -> int: + if current < previous: + raise ValueError(f"Cumulative {field} decreased") + return current - previous + + +def _optional_difference( + current: int | None, previous: int | None, field: str +) -> int | None: + if current is None or previous is None: + return None + return _difference(current, previous, field) + + +def account_turn( + turn: CompletedTurn, + *, + mode: UsageMode, + previous: UsageCheckpoint | None = None, + initial_turn: bool = False, +) -> AccountedTurn: + """Accumulate a completed turn using empirically verified CLI semantics. + + With no prior checkpoint, initial_turn=True explicitly asserts this is the + first observed turn of a newly created session. Never use it to invent a + zero baseline for an existing session with unobserved prior usage. + """ + mode = _mode(mode) + if not isinstance(turn, CompletedTurn): + raise ValueError("turn must be a validated CompletedTurn") + if not isinstance(initial_turn, bool): + raise ValueError("initial_turn must be boolean") + if previous is None: + if not initial_turn: + raise ValueError("Missing prior checkpoint; initial turn must be explicit") + return AccountedTurn( + turn.usage, UsageCheckpoint(turn.thread_id, mode, turn.usage) + ) + if initial_turn or not isinstance(previous, UsageCheckpoint): + raise ValueError("Invalid prior checkpoint for this turn") + if previous.thread_id != turn.thread_id: + raise ValueError("Prior checkpoint belongs to a different thread") + if previous.mode != mode: + raise ValueError("Usage mode changed within a session") + before, current = previous.cumulative_usage, turn.usage + if mode == "per_turn": + cumulative = TokenUsage( + before.input_tokens + current.input_tokens, + before.cached_input_tokens + current.cached_input_tokens, + before.output_tokens + current.output_tokens, + _optional_sum( + before.reasoning_output_tokens, current.reasoning_output_tokens + ), + _optional_sum( + before.cache_write_input_tokens, current.cache_write_input_tokens + ), + ) + return AccountedTurn(current, UsageCheckpoint(turn.thread_id, mode, cumulative)) + delta = TokenUsage( + _difference(current.input_tokens, before.input_tokens, "input_tokens"), + _difference( + current.cached_input_tokens, + before.cached_input_tokens, + "cached_input_tokens", + ), + _difference(current.output_tokens, before.output_tokens, "output_tokens"), + _optional_difference( + current.reasoning_output_tokens, + before.reasoning_output_tokens, + "reasoning_output_tokens", + ), + _optional_difference( + current.cache_write_input_tokens, + before.cache_write_input_tokens, + "cache_write_input_tokens", + ), + ) + return AccountedTurn(delta, UsageCheckpoint(turn.thread_id, mode, current)) + + +def _timestamp(value: object, field: str) -> None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{field} must be a finite nonnegative timestamp") + try: + finite = math.isfinite(value) + except OverflowError as error: + raise ValueError(f"{field} exceeds the supported clock range") from error + if value < 0 or not finite: + raise ValueError(f"{field} must be a finite nonnegative timestamp") + + +@dataclass(frozen=True) +class Interval: + """Start and end in seconds on the same nonnegative monotonic clock.""" + + start: float + end: float + + def __post_init__(self) -> None: + _timestamp(self.start, "start") + _timestamp(self.end, "end") + if self.end < self.start: + raise ValueError("Interval ends before it starts") + + @property + def seconds(self) -> float: + return self.end - self.start + + +def blocked_seconds(trial: Interval, intervals: Iterable[Interval]) -> float: + """Measure the union of genuinely blocked intervals, without clipping them.""" + if not isinstance(trial, Interval): + raise ValueError("trial must be a validated Interval") + ordered: list[Interval] = [] + for interval in intervals: + if not isinstance(interval, Interval): + raise ValueError("blocked intervals must be validated Intervals") + if interval.start < trial.start or interval.end > trial.end: + raise ValueError("Blocked interval lies outside trial") + ordered.append(interval) + ordered.sort(key=lambda interval: (interval.start, interval.end)) + if not ordered: + return 0.0 + total = 0.0 + start, end = ordered[0].start, ordered[0].end + for interval in ordered[1:]: + if interval.start <= end: + end = max(end, interval.end) + else: + total += end - start + start, end = interval.start, interval.end + return total + end - start + + +def gpu_allocation_seconds(gpu_count: int, interval: Interval) -> float: + """Allocated device-seconds, including idle time inside the allocation.""" + if _count(gpu_count, "gpu_count") == 0: + raise ValueError("gpu_count must be positive") + if not isinstance(interval, Interval): + raise ValueError("interval must be a validated Interval") + try: + seconds = gpu_count * interval.seconds + except OverflowError as error: + raise ValueError( + "GPU allocation duration exceeds the supported range" + ) from error + if not math.isfinite(seconds): + raise ValueError("GPU allocation duration exceeds the supported range") + return seconds diff --git a/tools/checkpoint_pilot/cli_turn.py b/tools/checkpoint_pilot/cli_turn.py new file mode 100644 index 00000000..86927503 --- /dev/null +++ b/tools/checkpoint_pilot/cli_turn.py @@ -0,0 +1,272 @@ +"""Run one sandboxed CLI decision turn and retain operator-clock evidence.""" + +from __future__ import annotations + +import json +import math +import os +import selectors +import signal +import subprocess +import time +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import BinaryIO, TextIO + +from tools.checkpoint_pilot.accounting import CompletedTurn, parse_exec_jsonl + + +@dataclass(frozen=True) +class CLISettings: + executable: Path + model: str + reasoning_effort: str + service_tier: str + response_schema: Path + + +@dataclass(frozen=True) +class TurnReceipt: + started_utc: str + ended_utc: str + started_monotonic: float + ended_monotonic: float + # None means the process never started or cleanup could not establish exit. + returncode: int | None + timed_out: bool + command: list[str] + workspace: str + error_type: str | None + error_message: str | None + + +def command( + settings: CLISettings, + workspace: Path, + output: Path, + thread_id: str | None, +) -> list[str]: + """Use the same policy/configuration on initial and resumed turns.""" + arguments = [ + str(settings.executable), + "exec", + "-C", + str(workspace.resolve(strict=True)), + "--sandbox", + "workspace-write", + "-c", + 'approval_policy="never"', + ] + if thread_id is not None: + if not thread_id.strip(): + raise ValueError("thread_id must not be empty") + arguments.extend(["resume", thread_id]) + arguments.extend( + [ + "--json", + "--ignore-user-config", + "--skip-git-repo-check", + "--disable", + "multi_agent", + "-m", + settings.model, + "-c", + f"model_reasoning_effort={json.dumps(settings.reasoning_effort)}", + "-c", + f"service_tier={json.dumps(settings.service_tier)}", + "--output-schema", + str(settings.response_schema.resolve(strict=True)), + "-o", + str(output / "response.json"), + "-", + ] + ) + return arguments + + +def stop_owned_process(process: subprocess.Popen[bytes]) -> None: + """Terminate the explicitly created process group and reap its leader.""" + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + pass + try: + process.wait(timeout=3) + except subprocess.TimeoutExpired: + pass + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=3) + + +def _record_lines(buffer: bytearray, received: TextIO, *, final: bool) -> None: + """Timestamp complete lines without waiting for another byte on the pipe.""" + while buffer: + boundary = buffer.find(b"\n") + if boundary < 0 and not final: + return + size = len(buffer) if boundary < 0 else boundary + 1 + line = bytes(buffer[:size]).decode("utf-8") + del buffer[:size] + received.write( + json.dumps( + { + "received_utc": datetime.now(timezone.utc).isoformat(), + "received_monotonic": time.monotonic(), + "raw_line": line.rstrip("\n"), + } + ) + + "\n" + ) + received.flush() + + +def _collect_events( + process: subprocess.Popen[bytes], + events: BinaryIO, + received: TextIO, + deadline: float, +) -> bool: + """Return timeout status; preserve exact bytes even for a partial last line.""" + assert process.stdout is not None + descriptor = process.stdout.fileno() + os.set_blocking(descriptor, False) + buffer = bytearray() + with selectors.DefaultSelector() as selector: + selector.register(descriptor, selectors.EVENT_READ) + while selector.get_map(): + remaining = deadline - time.monotonic() + if remaining <= 0: + # A partial line is already preserved in events.jsonl. It cannot + # be interpreted as a complete event or authoritative usage. + return True + for key, _ in selector.select(timeout=min(remaining, 0.5)): + try: + chunk = os.read(descriptor, 65536) + except BlockingIOError: + continue # Readiness can change before the nonblocking read. + if not chunk: + selector.unregister(key.fileobj) + _record_lines(buffer, received, final=True) + continue + events.write(chunk) + events.flush() + buffer.extend(chunk) + _record_lines(buffer, received, final=False) + return False + + +def _validate_response(path: Path) -> None: + if not path.is_file(): + raise ValueError("CLI completed without a structured response") + # The CLI applies response_schema; the study's response consumer must still + # validate its required fields. This boundary establishes JSON-object shape. + value: object = json.loads(path.read_text()) + if not isinstance(value, dict): + raise ValueError("CLI structured response must be a JSON object") + + +def run_turn( + settings: CLISettings, + workspace: Path, + output: Path, + prompt: str, + timeout_seconds: float, + thread_id: str | None = None, +) -> tuple[TurnReceipt, CompletedTurn]: + """Record raw events and receive times, then require authoritative usage. + + A timeout/failure preserves evidence and raises; callers must report it as + incomplete, never synthesize zero tokens or a successful decision. The + operator owns the output directory outside the agent's writable packet. + """ + if ( + isinstance(timeout_seconds, bool) + or not isinstance(timeout_seconds, (int, float)) + or not math.isfinite(timeout_seconds) + or not 0 < timeout_seconds <= 1800 + ): + raise ValueError("timeout_seconds must be in (0, 1800]") + if not isinstance(prompt, str) or not prompt.strip(): + raise ValueError("prompt must not be empty") + workspace = workspace.resolve(strict=True) + output = output.resolve() + arguments = command(settings, workspace, output, thread_id) + output.mkdir(parents=True, exist_ok=False) + (output / "prompt.txt").write_text(prompt) + started_utc = datetime.now(timezone.utc).isoformat() + started = time.monotonic() + timed_out = False + process: subprocess.Popen[bytes] | None = None + error_type: str | None = None + error_message: str | None = None + try: + with ( + (output / "prompt.txt").open("rb") as stdin, + (output / "stderr.log").open("wb") as stderr, + (output / "events.jsonl").open("wb") as events, + (output / "received.jsonl").open("w") as received, + ): + process = subprocess.Popen( + arguments, + cwd=workspace, + stdin=stdin, + stdout=subprocess.PIPE, + stderr=stderr, + start_new_session=True, + ) + timed_out = _collect_events( + process, events, received, started + timeout_seconds + ) + if not timed_out: + remaining = timeout_seconds - (time.monotonic() - started) + try: + process.wait(timeout=max(remaining, 0.001)) + except subprocess.TimeoutExpired: + timed_out = True + if timed_out or process.returncode != 0: + raise RuntimeError(f"CLI decision failed; evidence retained in {output}") + parsed = parse_exec_jsonl((output / "events.jsonl").read_text()) + if thread_id is not None and parsed.thread_id != thread_id: + raise ValueError("Resumed CLI returned a different thread") + _validate_response(output / "response.json") + except BaseException as error: + # Cleanup must also cover interruption; preserve and propagate failure. + error_type, error_message = type(error).__name__, str(error) + raise + finally: + try: + if process is not None: + # Reap the leader and stop any remaining members of the owned + # process group, including when the leader exited successfully. + stop_owned_process(process) + except BaseException as cleanup_error: + detail = f"Cleanup failed: {type(cleanup_error).__name__}: {cleanup_error}" + if error_type is None: + error_type = type(cleanup_error).__name__ + error_message = detail + else: + error_message = f"{error_message}; {detail}" + raise + finally: + if process is not None and process.stdout is not None: + process.stdout.close() + receipt = TurnReceipt( + started_utc, + datetime.now(timezone.utc).isoformat(), + started, + time.monotonic(), + None if process is None else process.returncode, + timed_out, + arguments, + str(workspace), + error_type, + error_message, + ) + (output / "receipt.json").write_text( + json.dumps(asdict(receipt), indent=2) + "\n" + ) + return receipt, parsed diff --git a/tools/checkpoint_pilot/decision.py b/tools/checkpoint_pilot/decision.py new file mode 100644 index 00000000..c9e8c1d0 --- /dev/null +++ b/tools/checkpoint_pilot/decision.py @@ -0,0 +1,189 @@ +"""Validate structured continuation decisions at the operator boundary.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import Literal, TypeAlias + +Action: TypeAlias = Literal["run", "finish"] +RequestKind: TypeAlias = Literal["resume", "diagnostic"] +Backend: TypeAlias = Literal["reference", "current"] + + +@dataclass(frozen=True) +class Observation: + claim: str + evidence: tuple[str, ...] + + +@dataclass(frozen=True) +class Request: + kind: RequestKind + backend: Backend + steps: int + script: str | None + purpose: str + + +@dataclass(frozen=True) +class Decision: + action: Action + summary: str + observations: tuple[Observation, ...] + hypotheses: tuple[str, ...] + next_action: str + request: Request | None + + +def _object(value: object, expected: set[str], label: str) -> dict[str, object]: + if not isinstance(value, dict): + raise TypeError(f"{label} must be an object") + fields: dict[str, object] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise TypeError(f"{label} fields must have string names") + fields[key] = item + if set(fields) != expected: + missing = sorted(expected - set(fields)) + extra = sorted(set(fields) - expected) + raise ValueError(f"{label} fields differ: missing={missing}, extra={extra}") + return fields + + +def _text(value: object, label: str) -> str: + if not isinstance(value, str): + raise TypeError(f"{label} must be a string") + if not value.strip(): + raise ValueError(f"{label} must be nonempty") + return value + + +def _texts(value: object, label: str) -> tuple[str, ...]: + if not isinstance(value, list): + raise TypeError(f"{label} must be an array") + return tuple(_text(item, f"{label}[{index}]") for index, item in enumerate(value)) + + +def _backend(value: object) -> Backend: + text = _text(value, "request.backend") + if text == "reference": + return "reference" + if text == "current": + return "current" + raise ValueError("request.backend must be reference or current") + + +def parse_request(value: object) -> Request: + fields = _object( + value, {"kind", "backend", "steps", "script", "purpose"}, "request" + ) + raw_kind = _text(fields["kind"], "request.kind") + kind: RequestKind + if raw_kind == "resume": + kind = "resume" + elif raw_kind == "diagnostic": + kind = "diagnostic" + else: + raise ValueError("request.kind must be resume or diagnostic") + steps = fields["steps"] + # JSON counts require a native int; bool and application-defined subclasses + # are outside this boundary contract, so an exact type comparison is needed. + if type(steps) is not int: # noqa: E721 + raise TypeError("request.steps must be an integer, excluding booleans") + assert isinstance(steps, int) + if not 1 <= steps <= 16: + raise ValueError("request.steps must be in [1,16]") + script: str | None + if kind == "resume": + if fields["script"] is not None: + raise ValueError("A resume request must have script=null") + script = None + else: + script = _text(fields["script"], "request.script") + path = PurePosixPath(script) + if ( + "\x00" in script + or path.is_absolute() + or ".." in path.parts + or path.suffix != ".py" + ): + raise ValueError( + "A diagnostic script must be a relative .py path without '..'" + ) + return Request( + kind=kind, + backend=_backend(fields["backend"]), + steps=steps, + script=script, + purpose=_text(fields["purpose"], "request.purpose"), + ) + + +def parse_decision(value: object) -> Decision: + """Reject malformed decisions while retaining empty evidence/hypothesis sets. + + Empty arrays mean the solver has not supplied observations, evidence, or + hypotheses yet. They are not converted into verified findings or success. + Diagnostic paths are lexical POSIX packet paths; execution must additionally + resolve them and verify containment/existence in the snapshotted workspace. + """ + fields = _object( + value, + {"action", "summary", "observations", "hypotheses", "next_action", "request"}, + "decision", + ) + raw_action = _text(fields["action"], "action") + action: Action + if raw_action == "run": + action = "run" + elif raw_action == "finish": + action = "finish" + else: + raise ValueError("action must be run or finish") + raw_observations = fields["observations"] + if not isinstance(raw_observations, list): + raise TypeError("observations must be an array") + observations: list[Observation] = [] + for index, raw in enumerate(raw_observations): + observation = _object(raw, {"claim", "evidence"}, f"observations[{index}]") + observations.append( + Observation( + claim=_text(observation["claim"], f"observations[{index}].claim"), + evidence=_texts( + observation["evidence"], f"observations[{index}].evidence" + ), + ) + ) + request: Request | None + if action == "run": + if fields["request"] is None: + raise ValueError("action=run requires a request") + request = parse_request(fields["request"]) + else: + if fields["request"] is not None: + raise ValueError("action=finish requires request=null") + request = None + return Decision( + action=action, + summary=_text(fields["summary"], "summary"), + observations=tuple(observations), + hypotheses=_texts(fields["hypotheses"], "hypotheses"), + next_action=_text(fields["next_action"], "next_action"), + request=request, + ) + + +def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"Duplicate JSON object field: {key}") + result[key] = value + return result + + +def read_decision(path: Path) -> Decision: + value: object = json.loads(path.read_text(), object_pairs_hook=_unique_object) + return parse_decision(value) diff --git a/tools/checkpoint_pilot/evaluate.py b/tools/checkpoint_pilot/evaluate.py new file mode 100644 index 00000000..ec83744c --- /dev/null +++ b/tools/checkpoint_pilot/evaluate.py @@ -0,0 +1,790 @@ +"""Operator-only comparison of independently produced checkpoint run artifacts. + +This module never imports a candidate. The operator must produce the reference +before running candidate code, freeze manifests, and run each prefix/resume in a +fresh process. The evaluator checks the recorded evidence; it is not a sandbox +against a candidate that can rewrite the operator's reference or manifests. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import re +import struct +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, TypeAlias, TypedDict + +import torch + +State: TypeAlias = ( + torch.Tensor + | None + | bool + | int + | float + | str + | list["State"] + | tuple["State", ...] + | dict[str, "State"] +) + + +@dataclass(frozen=True) +class Contract: + world_size: int + seed: int + eval_seed: int + prefix_steps: int + suffix_steps: int + accumulation: int + learning_rate: float + cuda_expected: bool + trajectory_atol: float = 2e-6 + trajectory_rtol: float = 2e-5 + + def __post_init__(self) -> None: + for name in ("world_size", "prefix_steps", "suffix_steps", "accumulation"): + if integer(getattr(self, name), name) <= 0: + raise ValueError(f"{name} must be positive") + for name in ("seed", "eval_seed"): + if integer(getattr(self, name), name) < 0: + raise ValueError(f"{name} must be nonnegative") + if number(self.learning_rate, "learning_rate") <= 0: + raise ValueError("learning_rate must be positive") + for name in ("trajectory_atol", "trajectory_rtol"): + if number(getattr(self, name), name) < 0: + raise ValueError(f"{name} must be nonnegative") + if not isinstance(self.cuda_expected, bool): + raise TypeError("cuda_expected must be a boolean") + + +@dataclass(frozen=True) +class RunEvidence: + phase_directory: Path + source_directory: Path + source_sha256: dict[str, str] + artifact_sha256: dict[str, str] + + +class Comparison(TypedDict): + path: str + kind: Literal["tensor", "scalar", "mapping_keys", "sequence_type", "length"] + passed: bool + reference: str + candidate: str + reference_sha256: str | None + candidate_sha256: str | None + maximum_absolute_error: float | None + + +class RankResult(TypedDict): + rank: int + prefix_passed: bool + restore_passed: bool + trajectory_passed: bool + prefix: list[Comparison] + restore: list[Comparison] + terminal: list[Comparison] + curves: list[Comparison] + + +class Error(TypedDict): + type: str + message: str + + +class Evaluation(TypedDict): + status: Literal["passed", "violated", "error"] + contract_passed: bool | None + prefix_passed: bool | None + restore_passed: bool | None + trajectory_passed: bool | None + ranks: list[RankResult] | None + source_and_artifact_manifests_verified: bool | None + error: Error | None + trajectory_atol: float + trajectory_rtol: float + limitations: list[str] + + +def record(value: object, context: str) -> dict[str, object]: + if not isinstance(value, dict): + raise TypeError(f"{context} must be a mapping") + result: dict[str, object] = {} + for key, child in value.items(): + if not isinstance(key, str): + raise TypeError(f"{context} must have string keys") + result[key] = child + return result + + +def integer(value: object, context: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{context} must be an integer") + return value + + +def number(value: object, context: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"{context} must be numeric") + if not math.isfinite(value): + raise ValueError(f"{context} must be finite") + return float(value) + + +def string(value: object, context: str) -> str: + if not isinstance(value, str) or not value: + raise TypeError(f"{context} must be a nonempty string") + return value + + +def array(value: object, context: str) -> list[object]: + if not isinstance(value, list): + raise TypeError(f"{context} must be an array") + return list(value) + + +def require(data: dict[str, object], key: str, context: str) -> object: + if key not in data: + raise ValueError(f"{context} missing required field {key}") + return data[key] + + +def read_json(path: Path) -> dict[str, object]: + value: object = json.loads(path.read_text()) + return record(value, str(path)) + + +def sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def validate_manifest(value: object, context: str) -> dict[str, str]: + mapping = record(value, context) + if not mapping: + raise ValueError(f"{context} cannot be empty") + result: dict[str, str] = {} + for relative, raw in mapping.items(): + path = Path(relative) + if path.is_absolute() or ".." in path.parts or str(path) != relative: + raise ValueError(f"{context} contains an invalid relative path") + digest = string(raw, context) + if re.fullmatch("[0-9a-f]{64}", digest) is None: + raise ValueError(f"{context} contains an invalid SHA-256") + result[relative] = digest + return result + + +def verify_files(root: Path, manifest: dict[str, str]) -> None: + resolved_root = root.resolve(strict=True) + for relative, expected in manifest.items(): + path = root / relative + resolved = path.resolve(strict=True) + if path.is_symlink() or not resolved.is_relative_to(resolved_root): + raise ValueError(f"Evidence path escapes root: {path}") + if not path.is_file() or sha256(path) != expected: + raise ValueError(f"Evidence hash mismatch: {path}") + + +def validate_run(run: RunEvidence, contract: Contract, stage: str) -> None: + source = validate_manifest(run.source_sha256, "source_sha256") + artifacts = validate_manifest(run.artifact_sha256, "artifact_sha256") + verify_files(run.source_directory, source) + actual_sources = { + str(path.relative_to(run.source_directory)) + for path in run.source_directory.rglob("*.py") + if path.is_file() + } + manifest_sources = {name for name in source if name.endswith(".py")} + if actual_sources != manifest_sources: + raise ValueError("Source manifest does not cover the exact Python source set") + required = {"phase-result.json", "invocation.json"} + for rank in range(contract.world_size): + required.update( + { + f"payload/rank{rank}/metrics.json", + f"payload/rank{rank}/terminal_state.pt", + } + ) + if stage == "resume": + required.add(f"payload/rank{rank}/restored_state.pt") + if not required.issubset(artifacts): + raise ValueError("Artifact manifest lacks a required phase/rank artifact") + verify_files(run.phase_directory, artifacts) + phase = read_json(run.phase_directory / "phase-result.json") + if ( + integer(require(phase, "exit_code", "phase"), "exit_code") != 0 + or require(phase, "stop_reason", "phase") is not None + or require(phase, "source_unchanged", "phase") is not True + ): + raise ValueError("Training phase did not complete successfully") + invocation = read_json(run.phase_directory / "invocation.json") + invocation_sources = validate_manifest( + require(invocation, "source_sha256", "invocation"), "invocation source_sha256" + ) + expected_sources = {name: source[name] for name in manifest_sources} + if invocation_sources != expected_sources: + raise ValueError("Run invocation does not match frozen Python source manifest") + rank_directories = { + path.name + for path in (run.phase_directory / "payload").iterdir() + if path.is_dir() and path.name.startswith("rank") + } + if rank_directories != {f"rank{rank}" for rank in range(contract.world_size)}: + raise ValueError("Incomplete or unexpected rank directories") + + +def snapshot(value: object, path: str) -> State: + if isinstance(value, torch.Tensor): + if value.layout != torch.strided or value.is_quantized: + raise TypeError(f"{path}: only dense nonquantized tensors are supported") + if (value.is_floating_point() or value.is_complex()) and not bool( + torch.isfinite(value).all() + ): + raise ValueError(f"{path}: nonfinite tensor") + return value.detach().cpu().clone() + if value is None or isinstance(value, (bool, int, str)): + return value + if isinstance(value, float): + if not math.isfinite(value): + raise ValueError(f"{path}: nonfinite scalar") + return value + if isinstance(value, list): + return [snapshot(item, f"{path}[{i}]") for i, item in enumerate(value)] + if isinstance(value, tuple): + return tuple(snapshot(item, f"{path}[{i}]") for i, item in enumerate(value)) + if isinstance(value, dict): + return { + key: snapshot(item, f"{path}[{key!r}]") + for key, item in record(value, path).items() + } + raise TypeError(f"{path}: unsupported value {type(value).__name__}") + + +def tensor_digest(value: torch.Tensor) -> str: + raw = value.detach().cpu().contiguous().reshape(-1).view(torch.uint8) + return hashlib.sha256(raw.numpy().tobytes()).hexdigest() + + +def comparison( + path: str, + kind: Literal["scalar", "mapping_keys", "sequence_type", "length"], + left: object, + right: object, + passed: bool, +) -> Comparison: + return { + "path": path, + "kind": kind, + "passed": passed, + "reference": repr(left), + "candidate": repr(right), + "reference_sha256": None, + "candidate_sha256": None, + "maximum_absolute_error": None, + } + + +def compare_states( + left: State, + right: State, + path: str, + *, + exact: bool, + atol: float, + rtol: float, +) -> list[Comparison]: + if isinstance(left, torch.Tensor) and isinstance(right, torch.Tensor): + metadata_equal = left.shape == right.shape and left.dtype == right.dtype + delta: float | None = None + if metadata_equal and (left.is_floating_point() or left.is_complex()): + delta = float((left - right).abs().max().item()) if left.numel() else 0.0 + passed = ( + tensor_digest(left) == tensor_digest(right) + if exact + else bool(torch.allclose(left, right, atol=atol, rtol=rtol)) + ) + else: + passed = metadata_equal and bool(torch.equal(left, right)) + return [ + { + "path": path, + "kind": "tensor", + "passed": passed, + "reference": f"shape={tuple(left.shape)},dtype={left.dtype}", + "candidate": f"shape={tuple(right.shape)},dtype={right.dtype}", + "reference_sha256": tensor_digest(left), + "candidate_sha256": tensor_digest(right), + "maximum_absolute_error": delta, + } + ] + if isinstance(left, dict) and isinstance(right, dict): + # Names bind state to logical parameters. Dictionary insertion order is + # retained as evidence but is not itself parameter identity; ordered + # parameter-group lists are compared separately, without reordering. + reports = [ + comparison( + path + ".keys", + "mapping_keys", + list(left), + list(right), + set(left) == set(right), + ) + ] + for key in left: + if key in right: + reports.extend( + compare_states( + left[key], + right[key], + f"{path}[{key!r}]", + exact=exact, + atol=atol, + rtol=rtol, + ) + ) + return reports + if isinstance(left, (list, tuple)) and isinstance(right, (list, tuple)): + reports = [ + comparison( + path + ".type", + "sequence_type", + type(left).__name__, + type(right).__name__, + type(left) is type(right), + ), + comparison( + path + ".length", + "length", + len(left), + len(right), + len(left) == len(right), + ), + ] + for index, (a, b) in enumerate(zip(left, right)): + reports.extend( + compare_states( + a, b, f"{path}[{index}]", exact=exact, atol=atol, rtol=rtol + ) + ) + return reports + if isinstance(left, torch.Tensor) or isinstance(right, torch.Tensor): + return [ + comparison( + path + ".type", + "scalar", + type(left).__name__, + type(right).__name__, + False, + ) + ] + passed = type(left) is type(right) and left == right + if isinstance(left, float) and isinstance(right, float): + passed = struct.pack("!d", left) == struct.pack("!d", right) + return [comparison(path, "scalar", left, right, passed is True)] + + +def load_state( + path: Path, contract: Contract, rank: int, completed_steps: int +) -> State: + # These are operator-owned workload outputs, not arbitrary downloaded pickle + # files. weights_only also rejects custom Python objects before validation. + raw: object = torch.load(path, map_location="cpu", weights_only=True) + data = record(raw, str(path)) + expected_fields = { + "model", + "optimizer_state", + "optimizer_groups", + "progress", + "rng", + "rank", + "world_size", + } + if set(data) != expected_fields: + raise ValueError("Unexpected or missing state fields") + for key, expected in (("rank", rank), ("world_size", contract.world_size)): + if integer(data[key], key) != expected: + raise ValueError(f"Unexpected state {key}") + progress = record(data["progress"], "progress") + if set(progress) != {"completed_steps", "data_cursor", "seed"}: + raise ValueError("Unexpected progress schema") + for key, expected in ( + ("completed_steps", completed_steps), + ("seed", contract.seed), + ): + if integer(progress[key], key) != expected: + raise ValueError(f"Unexpected progress {key}") + if integer(progress["data_cursor"], "data_cursor") < 0: + raise ValueError("Negative data cursor") + model = record(data["model"], "model") + if not model or not all( + isinstance(value, torch.Tensor) for value in model.values() + ): + raise TypeError("Model state must contain named tensors") + optimizer = record(data["optimizer_state"], "optimizer_state") + if not optimizer: + raise ValueError("Missing optimizer state") + for name, value in optimizer.items(): + if name not in model or not record(value, "parameter optimizer state"): + raise ValueError("Optimizer state has unknown or empty parameter state") + names: list[str] = [] + for raw_group in array(data["optimizer_groups"], "optimizer_groups"): + group = record(raw_group, "optimizer group") + names.extend( + string(name, "parameter name") + for name in array( + require(group, "params", "optimizer group"), "group params" + ) + ) + if not names or len(set(names)) != len(names) or set(names) != set(optimizer): + raise ValueError( + "Optimizer groups do not bind all states to unique parameter names" + ) + rng = record(data["rng"], "rng") + if set(rng) != {"cpu", "cuda"}: + raise ValueError("Unexpected RNG schema") + for name in ("cpu", "cuda"): + value = rng[name] + if name == "cuda" and not contract.cuda_expected: + if value is not None: + raise ValueError("Unexpected CUDA RNG in CPU scenario") + continue + if ( + not isinstance(value, torch.Tensor) + or value.dtype != torch.uint8 + or value.ndim != 1 + or value.numel() == 0 + ): + raise TypeError( + f"{name} RNG must be a nonempty one-dimensional byte tensor" + ) + return snapshot(data, str(path)) + + +def curves(path: Path, contract: Contract, rank: int, stage: str) -> State: + data = read_json(path) + start = 0 if stage == "prefix" else contract.prefix_steps + count = contract.prefix_steps if stage == "prefix" else contract.suffix_steps + expected: dict[str, object] = { + "stage": stage, + "rank": rank, + "world_size": contract.world_size, + "seed": contract.seed, + "eval_seed": contract.eval_seed, + "accumulation": contract.accumulation, + "examples_per_rank_update": 8, + "tokens_per_global_update": contract.world_size * 64, + "requested_updates": count, + "completed_steps": start + count, + "device": f"cuda:{rank}" if contract.cuda_expected else "cpu", + } + for key, value in expected.items(): + actual = require(data, key, "metrics") + if type(actual) is not type(value) or actual != value: + raise ValueError(f"Unexpected metrics {key}") + string(require(data, "backend", "metrics"), "backend") + if number(require(data, "lr", "metrics"), "lr") != contract.learning_rate: + raise ValueError("Unexpected learning rate") + parameter_count = integer( + require(data, "parameter_count", "metrics"), "parameter_count" + ) + if parameter_count <= 0: + raise ValueError("Invalid parameter count") + torch_version = string(require(data, "torch_version", "metrics"), "torch_version") + if ( + number( + require(data, "training_seconds_including_validation", "metrics"), + "training seconds", + ) + < 0 + ): + raise ValueError("Negative training time") + training = array(require(data, "steps", "metrics"), "steps") + if len(training) != count: + raise ValueError("Incomplete training curve") + for offset, item in enumerate(training, 1): + row = record(item, "training row") + if ( + set(row) != {"step", "train_token_loss"} + or integer(row["step"], "step") != start + offset + ): + raise ValueError("Unexpected training curve step/schema") + number(row["train_token_loss"], "training loss") + validation = array(require(data, "validation_curve", "metrics"), "validation_curve") + expected_steps = [start] + [ + step + for step in range(start + 1, start + count + 1) + if step % 4 == 0 or step == start + count + ] + if len(validation) != len(expected_steps): + raise ValueError("Incomplete validation curve") + for item, step in zip(validation, expected_steps): + row = record(item, "validation row") + if ( + set(row) != {"step", "validation_loss"} + or integer(row["step"], "step") != step + ): + raise ValueError("Unexpected validation curve step/schema") + number(row["validation_loss"], "validation loss") + # Convert only measured numerical curve values to tensors so declared + # trajectory tolerance applies. Keep step positions as exact integer lists. + return { + "parameter_count": parameter_count, + "torch_version": torch_version, + "training_steps": [ + integer(record(row, "row")["step"], "step") for row in training + ], + "training_losses": torch.tensor( + [ + number(record(row, "row")["train_token_loss"], "loss") + for row in training + ], + dtype=torch.float64, + ), + "validation_steps": [ + integer(record(row, "row")["step"], "step") for row in validation + ], + "validation_losses": torch.tensor( + [ + number(record(row, "row")["validation_loss"], "loss") + for row in validation + ], + dtype=torch.float64, + ), + } + + +def _evaluate_pair( + reference_prefix: RunEvidence, + reference_resume: RunEvidence, + candidate_prefix: RunEvidence, + candidate_resume: RunEvidence, + contract: Contract, +) -> list[RankResult]: + runs = ( + (reference_prefix, "prefix"), + (reference_resume, "resume"), + (candidate_prefix, "prefix"), + (candidate_resume, "resume"), + ) + for run, stage in runs: + validate_run(run, contract, stage) + results: list[RankResult] = [] + for rank in range(contract.world_size): + directories = [ + run.phase_directory / "payload" / f"rank{rank}" for run, _ in runs + ] + reference_start = load_state( + directories[0] / "terminal_state.pt", contract, rank, contract.prefix_steps + ) + reference_restored = load_state( + directories[1] / "restored_state.pt", contract, rank, contract.prefix_steps + ) + reference_end = load_state( + directories[1] / "terminal_state.pt", + contract, + rank, + contract.prefix_steps + contract.suffix_steps, + ) + candidate_start = load_state( + directories[2] / "terminal_state.pt", contract, rank, contract.prefix_steps + ) + candidate_restored = load_state( + directories[3] / "restored_state.pt", contract, rank, contract.prefix_steps + ) + candidate_end = load_state( + directories[3] / "terminal_state.pt", + contract, + rank, + contract.prefix_steps + contract.suffix_steps, + ) + options = {"atol": contract.trajectory_atol, "rtol": contract.trajectory_rtol} + baseline = compare_states( + reference_start, + reference_restored, + "reference_restore", + exact=True, + **options, + ) + if not all(row["passed"] for row in baseline): + raise ValueError( + "Reference restoration does not match its own checkpoint state" + ) + prefix = compare_states( + reference_start, candidate_start, "prefix", exact=True, **options + ) + restored = compare_states( + reference_start, candidate_restored, "restore", exact=True, **options + ) + terminal = compare_states( + reference_end, candidate_end, "terminal", exact=False, **options + ) + curve_values = [ + curves(directory / "metrics.json", contract, rank, stage) + for directory, (_, stage) in zip(directories, runs) + ] + prefix.extend( + compare_states( + curve_values[0], curve_values[2], "prefix_curves", exact=True, **options + ) + ) + trajectory = compare_states( + curve_values[1], curve_values[3], "resume_curves", exact=False, **options + ) + results.append( + { + "rank": rank, + "prefix_passed": all(row["passed"] for row in prefix), + "restore_passed": all(row["passed"] for row in restored), + "trajectory_passed": all( + row["passed"] for row in terminal + trajectory + ), + "prefix": prefix, + "restore": restored, + "terminal": terminal, + "curves": trajectory, + } + ) + # Recheck after all reads so changes during evaluation cannot be accepted. + for run, stage in runs: + validate_run(run, contract, stage) + return results + + +def evaluate_pair( + reference_prefix: RunEvidence, + reference_resume: RunEvidence, + candidate_prefix: RunEvidence, + candidate_resume: RunEvidence, + contract: Contract, +) -> Evaluation: + limits = [ + "Trusted operator manifests and independently generated reference are required; this is not an adversarial execution sandbox", + "Exact restore compares all saved values and logical parameter identity, not storage aliasing or original device placement", + "Named mapping insertion order is recorded but nonsemantic; parameter-group and all sequence orders are semantic", + "Finite post-resume tensors and curves use declared tolerances; RNG, scalar state and progress remain exact", + "Passing finite scenarios does not certify arbitrary collective behavior, long-run training, or untested checkpoint formats", + ] + try: + ranks = _evaluate_pair( + reference_prefix, + reference_resume, + candidate_prefix, + candidate_resume, + contract, + ) + except Exception as error: + # IO, deserialization, and validation failures are explicit unknown + # outcomes. The caller must inspect status/contract_passed and CLI exits2. + return { + "status": "error", + "contract_passed": None, + "prefix_passed": None, + "restore_passed": None, + "trajectory_passed": None, + "ranks": None, + "source_and_artifact_manifests_verified": None, + "error": {"type": type(error).__name__, "message": str(error)}, + "trajectory_atol": contract.trajectory_atol, + "trajectory_rtol": contract.trajectory_rtol, + "limitations": limits, + } + prefix = all(row["prefix_passed"] for row in ranks) + restored = all(row["restore_passed"] for row in ranks) + trajectory = all(row["trajectory_passed"] for row in ranks) + passed = prefix and restored and trajectory + return { + "status": "passed" if passed else "violated", + "contract_passed": passed, + "prefix_passed": prefix, + "restore_passed": restored, + "trajectory_passed": trajectory, + "ranks": ranks, + "source_and_artifact_manifests_verified": True, + "error": None, + "trajectory_atol": contract.trajectory_atol, + "trajectory_rtol": contract.trajectory_rtol, + "limitations": limits, + } + + +def run_from_manifest(path: Path) -> Evaluation: + manifest = read_json(path) + if integer(require(manifest, "schema_version", "manifest"), "schema_version") != 1: + raise ValueError("Unsupported manifest schema") + raw = record(require(manifest, "contract", "manifest"), "contract") + cuda = require(raw, "cuda_expected", "contract") + if not isinstance(cuda, bool): + raise TypeError("cuda_expected must be boolean") + contract = Contract( + world_size=integer(require(raw, "world_size", "contract"), "world_size"), + seed=integer(require(raw, "seed", "contract"), "seed"), + eval_seed=integer(require(raw, "eval_seed", "contract"), "eval_seed"), + prefix_steps=integer(require(raw, "prefix_steps", "contract"), "prefix_steps"), + suffix_steps=integer(require(raw, "suffix_steps", "contract"), "suffix_steps"), + accumulation=integer(require(raw, "accumulation", "contract"), "accumulation"), + learning_rate=number( + require(raw, "learning_rate", "contract"), "learning_rate" + ), + cuda_expected=cuda, + trajectory_atol=number( + require(raw, "trajectory_atol", "contract"), "trajectory_atol" + ), + trajectory_rtol=number( + require(raw, "trajectory_rtol", "contract"), "trajectory_rtol" + ), + ) + runs: list[RunEvidence] = [] + for name in ( + "reference_prefix", + "reference_resume", + "candidate_prefix", + "candidate_resume", + ): + run = record(require(manifest, name, "manifest"), name) + runs.append( + RunEvidence( + Path(string(require(run, "phase_directory", name), "phase_directory")), + Path( + string(require(run, "source_directory", name), "source_directory") + ), + validate_manifest(require(run, "source_sha256", name), "source_sha256"), + validate_manifest( + require(run, "artifact_sha256", name), "artifact_sha256" + ), + ) + ) + return evaluate_pair(runs[0], runs[1], runs[2], runs[3], contract) + + +class Options(argparse.Namespace): + manifest: Path + output: Path + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + options = parser.parse_args(namespace=Options()) + try: + result: object = run_from_manifest(options.manifest) + except Exception as error: + # A malformed operator manifest has no evaluable contract. Preserve an + # explicit error record rather than inventing defaults for its fields. + result = { + "status": "error", + "contract_passed": None, + "error": {"type": type(error).__name__, "message": str(error)}, + } + options.output.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n") + status = record(result, "evaluation")["status"] + return 0 if status == "passed" else 1 if status == "violated" else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/checkpoint_pilot/run_worker.py b/tools/checkpoint_pilot/run_worker.py new file mode 100644 index 00000000..d970df7f --- /dev/null +++ b/tools/checkpoint_pilot/run_worker.py @@ -0,0 +1,39 @@ +"""Set per-rank observation paths before importing training or diagnostic code.""" + +from __future__ import annotations + +import os +import runpy +import sys +from pathlib import Path + + +def main() -> None: + workspace = Path(os.environ["CHECKPOINT_WORKSPACE"]).resolve(strict=True) + output = Path(os.environ["EXPERIMENT_OUTPUT"]).resolve() + rank = int(os.environ["RANK"]) + if rank < 0: + raise ValueError("RANK must be nonnegative") + trace = output / "traces" / f"rank{rank}" + trace.mkdir(parents=True, exist_ok=False) + os.environ["TRAINCHECK_OUTPUT_DIR"] = str(trace) + sys.path.insert(0, str(workspace)) + os.chdir(workspace) + # Absence selects the fixed training engine; presence must name a packet + # Python file and may not escape the explicitly supplied workspace. + script = os.environ.get("EXPERIMENT_SCRIPT") + if script is None: + runpy.run_module("task.engine", run_name="__main__") + else: + path = (workspace / script).resolve(strict=True) + if ( + not path.is_relative_to(workspace) + or path.suffix != ".py" + or not path.is_file() + ): + raise ValueError("Diagnostic script must be a Python file in the packet") + runpy.run_path(str(path), run_name="__main__") + + +if __name__ == "__main__": + main() diff --git a/tools/checkpoint_pilot/task/__init__.py b/tools/checkpoint_pilot/task/__init__.py new file mode 100644 index 00000000..bad1a4b6 --- /dev/null +++ b/tools/checkpoint_pilot/task/__init__.py @@ -0,0 +1 @@ +"""Small distributed training continuation task.""" diff --git a/tools/checkpoint_pilot/task/current.py b/tools/checkpoint_pilot/task/current.py new file mode 100644 index 00000000..b88fad53 --- /dev/null +++ b/tools/checkpoint_pilot/task/current.py @@ -0,0 +1,228 @@ +"""Layer-sharded checkpoint backend.""" + +from __future__ import annotations + +from pathlib import Path +from typing import TypedDict + +import torch +import torch.nn as nn + +from . import state +from .model import LAYERS + + +class ParameterBundle(TypedDict): + model: torch.Tensor + optimizer: state.TensorMap + + +class LayerChunk(TypedDict): + layer_name: str + parameter_names: list[str] + parameters: list[ParameterBundle] + + +class SharedPayload(TypedDict): + schema_version: int + parameter_names: list[str] + optimizer_groups: list[state.Group] + nonlayer_names: list[str] + nonlayer_parameters: list[ParameterBundle] + progress: state.ProgressFields + world_size: int + + +class RankPayload(TypedDict): + schema_version: int + rank: int + world_size: int + layers: list[LayerChunk] + progress: state.ProgressFields + rng: state.RngFields + + +def layer_name(index: int) -> str: + return f"encoder.layers.{index}" + + +def save_checkpoint( + directory: Path, + model: nn.Module, + optimizer: torch.optim.Optimizer, + progress: state.Progress, + rank: int, + world_size: int, + device: torch.device, +) -> None: + state.validate_rank(rank, world_size) + names, values, groups = state.optimizer_snapshot(model, optimizer) + raw_model: object = model.state_dict() + model_values = state.tensors(raw_model, "model") + if set(model_values) != set(names): + raise ValueError( + "This checkpoint schema requires every model field to be a parameter" + ) + by_name: dict[str, ParameterBundle] = { + name: {"model": model_values[name].detach().cpu().clone(), "optimizer": value} + for name, value in zip(names, values, strict=True) + } + nonlayer_names = [name for name in names if not name.startswith("encoder.layers.")] + if rank == 0: + shared: SharedPayload = { + "schema_version": 2, + "parameter_names": names, + "optimizer_groups": groups, + "nonlayer_names": nonlayer_names, + "nonlayer_parameters": [by_name[name] for name in nonlayer_names], + "progress": progress.fields(), + "world_size": world_size, + } + torch.save(shared, directory / "shared.pt") + chunks: list[LayerChunk] = [] + for index in range(rank, LAYERS, world_size): + prefix = layer_name(index) + selected = [name for name in names if name.startswith(prefix + ".")] + if not selected: + raise ValueError("Expected a nonempty Transformer layer") + chunks.append( + { + "layer_name": prefix, + "parameter_names": selected, + "parameters": [by_name[name] for name in selected], + } + ) + shard: RankPayload = { + "schema_version": 2, + "rank": rank, + "world_size": world_size, + "layers": chunks, + "progress": progress.fields(), + "rng": state.capture_rng(device), + } + torch.save(shard, directory / f"rank{rank}.pt") + + +def read_bundle(value: object, label: str) -> ParameterBundle: + fields = state.mapping(value, label) + state.exact_keys(fields, {"model", "optimizer"}, label) + model = state.tensors({"model": fields["model"]}, label)["model"] + optimizer = state.momentum(fields["optimizer"], f"{label}.optimizer") + momentum = optimizer["momentum_buffer"] + if model.shape != momentum.shape or model.dtype != momentum.dtype: + raise ValueError(f"{label} model and momentum shape or dtype differ") + return {"model": model, "optimizer": optimizer} + + +def read_layer(value: object) -> LayerChunk: + chunk = state.mapping(value, "layer") + state.exact_keys(chunk, {"layer_name", "parameter_names", "parameters"}, "layer") + prefix = chunk["layer_name"] + if not isinstance(prefix, str) or prefix not in { + layer_name(i) for i in range(LAYERS) + }: + raise ValueError("Unknown checkpoint layer") + names = state.strings(chunk["parameter_names"], "layer parameter_names") + values = chunk["parameters"] + if not names or not all(name.startswith(prefix + ".") for name in names): + raise ValueError("Layer parameter names differ") + if not isinstance(values, list) or len(values) != len(names): + raise ValueError("Incomplete layer state") + return { + "layer_name": prefix, + "parameter_names": names, + "parameters": [read_bundle(v, n) for n, v in zip(names, values, strict=True)], + } + + +def restore_checkpoint( + directory: Path, + model: nn.Module, + optimizer: torch.optim.Optimizer, + rank: int, + world_size: int, + device: torch.device, +) -> state.Progress: + state.validate_rank(rank, world_size) + shared = state.read(directory / "shared.pt") + state.exact_keys( + shared, + { + "schema_version", + "parameter_names", + "optimizer_groups", + "nonlayer_names", + "nonlayer_parameters", + "progress", + "world_size", + }, + "shared checkpoint", + ) + if state.integer(shared["schema_version"], "schema_version") != 2: + raise ValueError("Unsupported sharded checkpoint version") + if state.integer(shared["world_size"], "world_size", 1) != world_size: + raise ValueError("Checkpoint topology differs") + names = state.strings(shared["parameter_names"], "parameter_names") + if names != state.optimizer_parameter_names(model, optimizer): + raise ValueError("Checkpoint optimizer parameter names differ") + groups = state.groups(shared["optimizer_groups"], len(names)) + progress = state.read_progress(shared["progress"]) + chunks: list[LayerChunk] = [] + rank_rng: state.RngFields | None = None + for source_rank in range(world_size): + shard = state.read(directory / f"rank{source_rank}.pt") + state.exact_keys( + shard, + {"schema_version", "rank", "world_size", "layers", "progress", "rng"}, + "rank shard", + ) + if state.integer(shard["schema_version"], "schema_version") != 2: + raise ValueError("Unsupported rank shard version") + if ( + state.integer(shard["rank"], "rank") != source_rank + or state.integer(shard["world_size"], "world_size", 1) != world_size + ): + raise ValueError("Shard topology differs") + if state.read_progress(shard["progress"]) != progress: + raise ValueError("Shard progress differs") + layers = shard["layers"] + if not isinstance(layers, list): + raise TypeError("Shard layers must be a list") + parsed = [read_layer(value) for value in layers] + expected = [ + layer_name(index) for index in range(source_rank, LAYERS, world_size) + ] + if [chunk["layer_name"] for chunk in parsed] != expected: + raise ValueError("Shard layer coverage differs") + chunks.extend(parsed) + rng = state.read_rng(shard["rng"], device) + if source_rank == rank: + rank_rng = rng + assert rank_rng is not None + nonlayer_names = state.strings(shared["nonlayer_names"], "nonlayer_names") + nonlayer_values = shared["nonlayer_parameters"] + expected_nonlayer = [ + name for name in names if not name.startswith("encoder.layers.") + ] + if ( + nonlayer_names != expected_nonlayer + or not isinstance(nonlayer_values, list) + or len(nonlayer_values) != len(nonlayer_names) + ): + raise ValueError("Nonlayer parameter coverage differs") + restored = { + name: read_bundle(value, name) + for name, value in zip(nonlayer_names, nonlayer_values, strict=True) + } + layer_names = [name for name in names if name.startswith("encoder.layers.")] + chunk_names = [name for chunk in chunks for name in chunk["parameter_names"]] + if len(chunk_names) != len(layer_names) or set(chunk_names) != set(layer_names): + raise ValueError("Layer parameter coverage differs") + layer_values = [value for chunk in chunks for value in chunk["parameters"]] + restored.update(zip(layer_names, layer_values, strict=True)) + state.load_model(model, {name: restored[name]["model"] for name in names}) + state.load_optimizer( + model, optimizer, [restored[name]["optimizer"] for name in names], groups + ) + state.restore_rng(rank_rng, device) + return progress diff --git a/tools/checkpoint_pilot/task/engine.py b/tools/checkpoint_pilot/task/engine.py new file mode 100644 index 00000000..fc605c2f --- /dev/null +++ b/tools/checkpoint_pilot/task/engine.py @@ -0,0 +1,280 @@ +"""Prefix training and fresh-process checkpoint continuation.""" + +from __future__ import annotations + +import argparse +import datetime +import json +import math +import os +import time +from contextlib import nullcontext +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, TypedDict + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.nn.parallel import DistributedDataParallel + +from . import current, reference, state +from .model import ( + Batch, + Transformer, + evaluate, + loss_for_batch, + make_batch, + microbatches, +) + + +class TrainPoint(TypedDict): + step: int + train_token_loss: float + + +class ValidationPoint(TypedDict): + step: int + validation_loss: float + + +@dataclass(frozen=True) +class Options: + stage: Literal["prefix", "resume"] + backend: Literal["reference", "current"] + checkpoint: Path + output: Path + steps: int = 12 + seed: int = 20261101 + eval_seed: int = 777000 + lr: float = 0.15 + accumulation: int = 4 + device: Literal["cpu", "cuda"] = "cuda" + + def validate(self) -> None: + if self.stage not in ("prefix", "resume") or self.backend not in ( + "reference", + "current", + ): + raise ValueError("Invalid stage or backend") + if self.device not in ("cpu", "cuda"): + raise ValueError("Invalid device") + for name, value in ( + ("steps", self.steps), + ("seed", self.seed), + ("eval_seed", self.eval_seed), + ("accumulation", self.accumulation), + ): + state.integer(value, name) + if not 1 <= self.steps <= 128: + raise ValueError("steps must be in [1,128]") + if self.seed >= 2**31 or self.eval_seed >= 2**31: + raise ValueError("Seeds must be less than 2**31") + if self.accumulation not in (1, 2, 4, 8): + raise ValueError("accumulation must be a positive divisor of eight") + if ( + isinstance(self.lr, bool) + or not isinstance(self.lr, (int, float)) + or not math.isfinite(self.lr) + or self.lr <= 0 + ): + raise ValueError("lr must be finite and positive") + + +def run_update( + model: nn.Module, optimizer: torch.optim.Optimizer, batches: list[Batch] +) -> float: + if not batches: + raise ValueError("An update requires microbatches") + optimizer.zero_grad(set_to_none=True) + total = torch.zeros((), device=batches[0].tokens.device) + for index, batch in enumerate(batches): + synchronization = ( + model.no_sync() + if isinstance(model, DistributedDataParallel) and index + 1 < len(batches) + else nullcontext() + ) + with synchronization: + loss = loss_for_batch(model, batch) + total += loss.detach() + (loss / len(batches)).backward() + for parameter in model.parameters(): + if parameter.grad is None or not bool(torch.isfinite(parameter.grad).all()): + raise ValueError("Every parameter must have a finite gradient") + optimizer.step() + return float((total / len(batches)).item()) + + +def run(options: Options) -> None: + options.validate() + # Absence means ordinary single-process execution, matching torchrun's API. + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + state.validate_rank(rank, world_size) + if local_rank < 0: + raise ValueError("LOCAL_RANK must be nonnegative") + torch.set_num_threads(1) + device = ( + torch.device("cuda", local_rank) + if options.device == "cuda" + else torch.device("cpu") + ) + if device.type == "cuda": + torch.cuda.set_device(device) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.use_deterministic_algorithms(True) + torch.manual_seed(options.seed) + if world_size > 1: + dist.init_process_group( + "nccl" if device.type == "cuda" else "gloo", + timeout=datetime.timedelta(seconds=60), + ) + try: + output = options.output / f"rank{rank}" + output.mkdir(parents=True, exist_ok=False) + if options.stage == "prefix": + if rank == 0: + options.checkpoint.mkdir(parents=True, exist_ok=False) + if world_size > 1: + dist.barrier() + elif not options.checkpoint.is_dir(): + raise FileNotFoundError(options.checkpoint) + model = Transformer().to(device).train() + wrapped: nn.Module = model + if world_size > 1: + wrapped = DistributedDataParallel( + model, device_ids=[local_rank] if device.type == "cuda" else None + ) + # Initialization is shared; stochastic training has a stream per rank. + torch.manual_seed(options.seed + 104729 * rank + 1) + optimizer = torch.optim.SGD(model.parameters(), lr=options.lr, momentum=0.9) + backend = reference if options.backend == "reference" else current + progress = state.Progress(0, 0, options.seed) + if options.stage == "resume": + progress = backend.restore_checkpoint( + options.checkpoint, model, optimizer, rank, world_size, device + ) + if progress.seed != options.seed: + raise ValueError("Checkpoint training seed differs from requested seed") + _, _, restored_groups = state.optimizer_snapshot(model, optimizer) + if restored_groups[0]["lr"] != options.lr: + raise ValueError("Requested learning rate differs from checkpoint") + restored = state.artifact( + model, optimizer, progress, rank, world_size, device + ) + torch.save(restored, output / "restored_state.pt") + curve: list[ValidationPoint] = [ + { + "step": progress.completed_steps, + "validation_loss": evaluate(model, options.eval_seed, device), + } + ] + points: list[TrainPoint] = [] + if device.type == "cuda": + torch.cuda.synchronize(device) + started = time.perf_counter() + for _ in range(options.steps): + batch = make_batch(progress.seed, progress.data_cursor + 1, rank, device) + local_loss = run_update( + wrapped, optimizer, microbatches(batch, options.accumulation) + ) + global_loss = torch.tensor(local_loss, dtype=torch.float64, device=device) + if world_size > 1: + dist.all_reduce(global_loss) + global_loss /= world_size + loss = float(global_loss.item()) + if not math.isfinite(loss): + raise ValueError("Training loss is not finite") + progress = state.Progress( + progress.completed_steps + 1, progress.data_cursor + 1, progress.seed + ) + points.append({"step": progress.completed_steps, "train_token_loss": loss}) + if progress.completed_steps % 4 == 0 or len(points) == options.steps: + curve.append( + { + "step": progress.completed_steps, + "validation_loss": evaluate(model, options.eval_seed, device), + } + ) + if device.type == "cuda": + torch.cuda.synchronize(device) + seconds = time.perf_counter() - started + terminal = state.artifact(model, optimizer, progress, rank, world_size, device) + torch.save(terminal, output / "terminal_state.pt") + if options.stage == "prefix": + backend.save_checkpoint( + options.checkpoint, model, optimizer, progress, rank, world_size, device + ) + if world_size > 1: + dist.barrier() + metrics: dict[str, object] = { + "stage": options.stage, + "backend": options.backend, + "rank": rank, + "world_size": world_size, + "seed": options.seed, + "eval_seed": options.eval_seed, + "lr": options.lr, + "accumulation": options.accumulation, + "device": str(device), + "torch_version": str(torch.__version__), + "parameter_count": sum( + parameter.numel() for parameter in model.parameters() + ), + "examples_per_rank_update": 8, + "tokens_per_global_update": world_size * 64, + "requested_updates": options.steps, + "completed_steps": progress.completed_steps, + "training_seconds_including_validation": seconds, + "steps": points, + "validation_curve": curve, + } + (output / "metrics.json").write_text( + json.dumps(metrics, indent=2, allow_nan=False) + "\n" + ) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def parse_options(argv: list[str] | None = None) -> Options: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("stage", choices=("prefix", "resume")) + parser.add_argument("--backend", choices=("reference", "current"), required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--steps", type=int, default=12) + parser.add_argument("--seed", type=int, default=20261101) + parser.add_argument("--eval-seed", type=int, default=777000) + parser.add_argument("--lr", type=float, default=0.15) + parser.add_argument("--accumulation", type=int, default=4) + parser.add_argument("--device", choices=("cpu", "cuda"), default="cuda") + args = parser.parse_args(argv) + options = Options( + args.stage, + args.backend, + args.checkpoint, + args.output, + args.steps, + args.seed, + args.eval_seed, + args.lr, + args.accumulation, + args.device, + ) + try: + options.validate() + except (ValueError, TypeError) as error: + parser.error(str(error)) + return options + + +def main() -> None: + run(parse_options()) + + +if __name__ == "__main__": + main() diff --git a/tools/checkpoint_pilot/task/model.py b/tools/checkpoint_pilot/task/model.py new file mode 100644 index 00000000..72c7690d --- /dev/null +++ b/tools/checkpoint_pilot/task/model.py @@ -0,0 +1,112 @@ +"""Fixed model, objective, and reproducible data for the continuation task.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F + +VOCABULARY = 32 +SEQUENCE_LENGTH = 8 +EXAMPLES_PER_RANK = 8 +LAYERS = 4 + + +@dataclass(frozen=True) +class Batch: + tokens: torch.Tensor + labels: torch.Tensor + + +def make_layer() -> nn.TransformerEncoderLayer: + return nn.TransformerEncoderLayer( + d_model=16, nhead=2, dim_feedforward=32, dropout=0.1, batch_first=True + ) + + +class Transformer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.embedding = nn.Embedding(VOCABULARY, 16) + self.position_embedding = nn.Embedding(SEQUENCE_LENGTH, 16) + self.encoder = nn.TransformerEncoder(make_layer(), num_layers=LAYERS) + # TransformerEncoder clones its prototype. Construct the remaining + # blocks independently while retaining each module's default initializer. + for index in range(1, LAYERS): + self.encoder.layers[index] = make_layer() + self.readout = nn.Linear(16, VOCABULARY) + + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + if tokens.ndim != 2 or tokens.shape[1] != SEQUENCE_LENGTH: + raise ValueError( + "The task requires a batch of fixed-length token sequences" + ) + positions = torch.arange(SEQUENCE_LENGTH, device=tokens.device) + embedded = self.embedding(tokens) + self.position_embedding(positions) + value: object = self.readout(self.encoder(embedded)) + if not isinstance(value, torch.Tensor): + raise TypeError("Transformer output must be a tensor") + return value + + +def make_batch(seed: int, cursor: int, rank: int, device: torch.device) -> Batch: + """Predict each position's previous token, wrapping position zero to the end. + + The labels require cross-position information. This is bidirectional sequence + tagging with a known circular boundary, not causal next-token generation. + """ + if seed < 0 or cursor < 0 or rank < 0: + raise ValueError("Seed, cursor, and rank must be nonnegative") + generator = torch.Generator(device="cpu") + generator.manual_seed(seed + 10000 * cursor + 100 * rank) + tokens = torch.randint( + VOCABULARY, + (EXAMPLES_PER_RANK, SEQUENCE_LENGTH), + generator=generator, + ) + return Batch(tokens.to(device), tokens.roll(1, dims=1).to(device)) + + +def microbatches(batch: Batch, accumulation: int) -> list[Batch]: + if accumulation < 1 or EXAMPLES_PER_RANK % accumulation: + raise ValueError("Accumulation must be a positive divisor of eight") + expected = (EXAMPLES_PER_RANK, SEQUENCE_LENGTH) + if tuple(batch.tokens.shape) != expected or tuple(batch.labels.shape) != expected: + raise ValueError("Unexpected complete batch shape") + size = EXAMPLES_PER_RANK // accumulation + return [ + Batch(batch.tokens[start : start + size], batch.labels[start : start + size]) + for start in range(0, EXAMPLES_PER_RANK, size) + ] + + +def loss_for_batch(model: nn.Module, batch: Batch) -> torch.Tensor: + output: object = model(batch.tokens) + if not isinstance(output, torch.Tensor): + raise TypeError("Model output must be a tensor") + result: object = F.cross_entropy( + output.reshape(-1, VOCABULARY), batch.labels.reshape(-1), reduction="mean" + ) + if not isinstance(result, torch.Tensor): + raise TypeError("Cross entropy must be a tensor") + return result + + +def evaluate(model: nn.Module, seed: int, device: torch.device) -> float: + modes = [(module, module.training) for module in model.modules()] + model.eval() + try: + with torch.no_grad(): + losses = [ + loss_for_batch(model, make_batch(seed, cursor, 0, device)).double() + for cursor in range(8) + ] + result = float(torch.stack(losses).mean().item()) + if not torch.isfinite(torch.tensor(result)): + raise ValueError("Validation loss is not finite") + return result + finally: + for module, training in modes: + module.training = training diff --git a/tools/checkpoint_pilot/task/reference.py b/tools/checkpoint_pilot/task/reference.py new file mode 100644 index 00000000..a92482ad --- /dev/null +++ b/tools/checkpoint_pilot/task/reference.py @@ -0,0 +1,99 @@ +"""Replicated checkpoint backend.""" + +from __future__ import annotations + +from pathlib import Path +from typing import TypedDict + +import torch +import torch.nn as nn + +from . import state + + +class Payload(TypedDict): + schema_version: int + model: state.TensorMap + parameter_names: list[str] + optimizer_states: list[state.TensorMap] + optimizer_groups: list[state.Group] + progress: state.ProgressFields + rng: state.RngFields + rank: int + world_size: int + + +def save_checkpoint( + directory: Path, + model: nn.Module, + optimizer: torch.optim.Optimizer, + progress: state.Progress, + rank: int, + world_size: int, + device: torch.device, +) -> None: + state.validate_rank(rank, world_size) + names, values, groups = state.optimizer_snapshot(model, optimizer) + raw_model: object = model.state_dict() + payload: Payload = { + "schema_version": 1, + "model": state.tensor_snapshot(state.tensors(raw_model, "model")), + "parameter_names": names, + "optimizer_states": values, + "optimizer_groups": groups, + "progress": progress.fields(), + "rng": state.capture_rng(device), + "rank": rank, + "world_size": world_size, + } + torch.save(payload, directory / f"rank{rank}.pt") + + +def restore_checkpoint( + directory: Path, + model: nn.Module, + optimizer: torch.optim.Optimizer, + rank: int, + world_size: int, + device: torch.device, +) -> state.Progress: + state.validate_rank(rank, world_size) + payload = state.read(directory / f"rank{rank}.pt") + state.exact_keys( + payload, + { + "schema_version", + "model", + "parameter_names", + "optimizer_states", + "optimizer_groups", + "progress", + "rng", + "rank", + "world_size", + }, + "replicated checkpoint", + ) + if state.integer(payload["schema_version"], "schema_version") != 1: + raise ValueError("Unsupported replicated checkpoint version") + if ( + state.integer(payload["rank"], "rank") != rank + or state.integer(payload["world_size"], "world_size", 1) != world_size + ): + raise ValueError("Checkpoint topology differs") + names = state.strings(payload["parameter_names"], "parameter_names") + if names != state.optimizer_parameter_names(model, optimizer): + raise ValueError("Checkpoint optimizer parameter names differ") + values = payload["optimizer_states"] + if not isinstance(values, list) or len(values) != len(names): + raise ValueError("Incomplete optimizer states") + states = [ + state.momentum(value, name) for name, value in zip(names, values, strict=True) + ] + groups = state.groups(payload["optimizer_groups"], len(names)) + progress = state.read_progress(payload["progress"]) + rng = state.read_rng(payload["rng"], device) + state.load_model(model, payload["model"]) + state.load_optimizer(model, optimizer, states, groups) + state.restore_rng(rng, device) + return progress diff --git a/tools/checkpoint_pilot/task/state.py b/tools/checkpoint_pilot/task/state.py new file mode 100644 index 00000000..30d14c4d --- /dev/null +++ b/tools/checkpoint_pilot/task/state.py @@ -0,0 +1,357 @@ +"""Validated state boundaries shared by checkpoint backends.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path +from typing import TypeAlias, TypedDict + +import torch +import torch.nn as nn + +TensorMap: TypeAlias = dict[str, torch.Tensor] +GroupScalar: TypeAlias = bool | int | float | str | None +Group: TypeAlias = dict[str, GroupScalar | list[int]] + + +class ProgressFields(TypedDict): + completed_steps: int + data_cursor: int + seed: int + + +class RngFields(TypedDict): + cpu: torch.Tensor + cuda: torch.Tensor | None + + +class OptimizerFields(TypedDict): + state: dict[int, TensorMap] + param_groups: list[Group] + + +NamedGroup: TypeAlias = dict[str, GroupScalar | list[str]] + + +class Artifact(TypedDict): + model: TensorMap + optimizer_state: dict[str, TensorMap] + optimizer_groups: list[NamedGroup] + progress: ProgressFields + rng: RngFields + rank: int + world_size: int + + +@dataclass(frozen=True) +class Progress: + completed_steps: int + data_cursor: int + seed: int + + def fields(self) -> ProgressFields: + if min(self.completed_steps, self.data_cursor, self.seed) < 0: + raise ValueError("Progress fields must be nonnegative") + if self.completed_steps != self.data_cursor: + raise ValueError("This task consumes one complete batch per update") + return { + "completed_steps": self.completed_steps, + "data_cursor": self.data_cursor, + "seed": self.seed, + } + + +def mapping(value: object, label: str) -> dict[str, object]: + if not isinstance(value, dict): + raise TypeError(f"{label} must be a mapping") + result: dict[str, object] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise TypeError(f"{label} keys must be strings") + result[key] = item + return result + + +def exact_keys(value: dict[str, object], keys: set[str], label: str) -> None: + if set(value) != keys: + raise ValueError(f"{label} fields differ: expected {sorted(keys)}") + + +def integer(value: object, label: str, minimum: int = 0) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{label} must be an integer") + if value < minimum: + raise ValueError(f"{label} must be at least {minimum}") + return value + + +def strings(value: object, label: str) -> list[str]: + if not isinstance(value, list): + raise TypeError(f"{label} must be a list") + result: list[str] = [] + for item in value: + if not isinstance(item, str) or not item: + raise TypeError(f"{label} entries must be nonempty strings") + result.append(item) + if len(set(result)) != len(result): + raise ValueError(f"{label} entries must be unique") + return result + + +def tensors(value: object, label: str) -> TensorMap: + result: TensorMap = {} + for name, item in mapping(value, label).items(): + if not isinstance(item, torch.Tensor): + raise TypeError(f"{label}.{name} must be a tensor") + if item.layout != torch.strided or item.is_quantized: + raise ValueError(f"{label}.{name} must be a dense ordinary tensor") + if not bool(torch.isfinite(item).all()): + raise ValueError(f"{label}.{name} must be finite") + result[name] = item + return result + + +def tensor_snapshot(values: TensorMap) -> TensorMap: + return {name: value.detach().cpu().clone() for name, value in values.items()} + + +def read(path: Path) -> dict[str, object]: + raw: object = torch.load(path, map_location="cpu", weights_only=True) + return mapping(raw, str(path)) + + +def read_progress(value: object) -> Progress: + fields = mapping(value, "progress") + exact_keys(fields, {"completed_steps", "data_cursor", "seed"}, "progress") + progress = Progress( + integer(fields["completed_steps"], "completed_steps"), + integer(fields["data_cursor"], "data_cursor"), + integer(fields["seed"], "seed"), + ) + progress.fields() + return progress + + +def capture_rng(device: torch.device) -> RngFields: + return { + "cpu": torch.get_rng_state().clone(), + "cuda": ( + torch.cuda.get_rng_state(device).clone() if device.type == "cuda" else None + ), + } + + +def read_rng(value: object, device: torch.device) -> RngFields: + fields = mapping(value, "rng") + exact_keys(fields, {"cpu", "cuda"}, "rng") + cpu = fields["cpu"] + cuda = fields["cuda"] + for name, item in (("cpu", cpu), ("cuda", cuda)): + if item is None and name == "cuda" and device.type == "cpu": + continue + if ( + not isinstance(item, torch.Tensor) + or item.dtype != torch.uint8 + or item.ndim != 1 + or item.device.type != "cpu" + ): + raise TypeError(f"Invalid {name} RNG state") + if not isinstance(cpu, torch.Tensor): + raise TypeError("Missing CPU RNG state") + if device.type == "cpu" and cuda is not None: + raise ValueError("CPU checkpoint must not contain CUDA RNG state") + if cuda is not None and not isinstance(cuda, torch.Tensor): + raise TypeError("Invalid CUDA RNG state") + return {"cpu": cpu, "cuda": cuda} + + +def restore_rng(value: RngFields, device: torch.device) -> None: + torch.set_rng_state(value["cpu"]) + if device.type == "cuda": + cuda = value["cuda"] + if cuda is None: + raise ValueError("CUDA checkpoint requires CUDA RNG state") + torch.cuda.set_rng_state(cuda, device) + + +def groups(value: object, expected_count: int) -> list[Group]: + if not isinstance(value, list) or len(value) != 1: + raise ValueError("This task uses exactly one optimizer group") + raw = mapping(value[0], "optimizer group") + # This task uses the explicit SGD checkpoint schema from pinned PyTorch2.2. + # Unknown settings require a deliberate compatibility review, not a fallback. + exact_keys( + raw, + { + "params", + "lr", + "momentum", + "dampening", + "weight_decay", + "nesterov", + "maximize", + "foreach", + "differentiable", + }, + "optimizer group", + ) + parameter_ids = raw["params"] + if not isinstance(parameter_ids, list): + raise TypeError("Optimizer params must be a list") + ids = [integer(item, "parameter ID") for item in parameter_ids] + if len(ids) != expected_count or len(set(ids)) != len(ids): + raise ValueError("Optimizer parameter IDs are incomplete or duplicate") + result: Group = {"params": ids} + for name, item in raw.items(): + if name == "params": + continue + if item is not None and not isinstance(item, (bool, int, float, str)): + raise TypeError(f"Unsupported optimizer group field {name}") + if isinstance(item, float) and not math.isfinite(item): + raise ValueError(f"Nonfinite optimizer group field {name}") + result[name] = item + for name in ("lr", "momentum", "dampening", "weight_decay"): + item = result[name] + if isinstance(item, bool) or not isinstance(item, (int, float)): + raise TypeError(f"Optimizer {name} must be numeric") + if item < 0 or (name == "lr" and item == 0): + raise ValueError(f"Optimizer {name} is outside the supported range") + for name in ("nesterov", "maximize", "differentiable"): + if not isinstance(result[name], bool): + raise TypeError(f"Optimizer {name} must be boolean") + foreach = result["foreach"] + if foreach is not None and not isinstance(foreach, bool): + raise TypeError("Optimizer foreach must be boolean or unavailable") + if result["nesterov"] and (result["momentum"] == 0 or result["dampening"] != 0): + raise ValueError("Nesterov requires positive momentum and zero dampening") + return [result] + + +def parameter_ids(group_values: list[Group]) -> list[int]: + assert len(group_values) == 1 + ids = group_values[0]["params"] + assert isinstance(ids, list) + return ids + + +def momentum(value: object, label: str) -> TensorMap: + result = tensors(value, label) + if set(result) != {"momentum_buffer"}: + raise ValueError(f"{label} must contain exactly one momentum buffer") + return result + + +def optimizer_snapshot( + model: nn.Module, optimizer: torch.optim.Optimizer +) -> tuple[list[str], list[TensorMap], list[Group]]: + names = optimizer_parameter_names(model, optimizer) + raw: object = optimizer.state_dict() + fields = mapping(raw, "optimizer") + exact_keys(fields, {"state", "param_groups"}, "optimizer") + group_values = groups(fields["param_groups"], len(names)) + raw_states = fields["state"] + if not isinstance(raw_states, dict): + raise TypeError("Optimizer state must be a mapping") + states: dict[int, TensorMap] = {} + for key, value in raw_states.items(): + index = integer(key, "optimizer state ID") + states[index] = momentum(value, f"optimizer state {index}") + ids = parameter_ids(group_values) + if set(states) != set(ids): + raise ValueError("Checkpoint requires initialized state for every parameter") + return names, [tensor_snapshot(states[index]) for index in ids], group_values + + +def optimizer_parameter_names( + model: nn.Module, optimizer: torch.optim.Optimizer +) -> list[str]: + by_identity = {id(parameter): name for name, parameter in model.named_parameters()} + live_groups: object = optimizer.param_groups + if not isinstance(live_groups, list) or len(live_groups) != 1: + raise ValueError("This task uses exactly one live optimizer group") + group = mapping(live_groups[0], "live optimizer group") + parameters = group["params"] + if not isinstance(parameters, list): + raise TypeError("Live optimizer params must be a list") + names: list[str] = [] + for parameter in parameters: + if not isinstance(parameter, torch.Tensor) or id(parameter) not in by_identity: + raise ValueError("Optimizer references a parameter outside the model") + names.append(by_identity[id(parameter)]) + if len(names) != len(by_identity) or len(set(names)) != len(names): + raise ValueError("Optimizer must cover each model parameter exactly once") + return names + + +def load_model(model: nn.Module, value: object) -> None: + source = tensors(value, "model") + expected: object = model.state_dict() + target = tensors(expected, "live model") + if source.keys() != target.keys(): + raise ValueError("Checkpoint model names differ") + for name, item in source.items(): + if item.shape != target[name].shape or item.dtype != target[name].dtype: + raise ValueError(f"Checkpoint model shape or dtype differs: {name}") + model.load_state_dict(source, strict=True) + + +def load_optimizer( + model: nn.Module, + optimizer: torch.optim.Optimizer, + state_values: list[TensorMap], + group_values: list[Group], +) -> None: + by_name = dict(model.named_parameters()) + parameters = [by_name[name] for name in optimizer_parameter_names(model, optimizer)] + ids = parameter_ids(group_values) + if len(state_values) != len(parameters) or len(ids) != len(parameters): + raise ValueError("Optimizer states do not cover every parameter") + for parameter, state in zip(parameters, state_values, strict=True): + buffer = momentum(state, "optimizer state")["momentum_buffer"] + if buffer.shape != parameter.shape or buffer.dtype != parameter.dtype: + raise ValueError("Optimizer state shape or dtype differs from parameter") + payload: OptimizerFields = { + "state": dict(zip(ids, state_values, strict=True)), + "param_groups": group_values, + } + optimizer.load_state_dict(dict(payload)) + + +def validate_rank(rank: int, world_size: int) -> None: + integer(rank, "rank") + integer(world_size, "world_size", 1) + if world_size < 1 or rank < 0 or rank >= world_size: + raise ValueError("Invalid distributed rank metadata") + + +def artifact( + model: nn.Module, + optimizer: torch.optim.Optimizer, + progress: Progress, + rank: int, + world_size: int, + device: torch.device, +) -> Artifact: + validate_rank(rank, world_size) + names, states, group_values = optimizer_snapshot(model, optimizer) + named_groups: list[NamedGroup] = [] + for group in group_values: + named: NamedGroup = {} + for key, value in group.items(): + if key == "params": + named[key] = list(names) + else: + assert not isinstance(value, list) + named[key] = value + named_groups.append(named) + raw_model: object = model.state_dict() + return { + "model": tensor_snapshot(tensors(raw_model, "model")), + "optimizer_state": dict(zip(names, states, strict=True)), + "optimizer_groups": named_groups, + "progress": progress.fields(), + "rng": capture_rng(device), + "rank": rank, + "world_size": world_size, + }