From fb374623041d177332bb70c24dd53f68cc2142c3 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 8 Aug 2026 21:25:15 -0700 Subject: [PATCH 1/5] Add exact resume-safe WWPGD projection dose telemetry Replace PR #129's ambiguous first-event flag with exact realized dose fields, zero-based scheduled-event labeling, and per-layer first actual application counts reconstructed from durable scientific projection history. Add old-schema rejection, checkpoint-reconciliation coverage, focused tests, and documentation. No optimizer or projection mathematics change. --- docs/WWPGD_PROJECTION_DOSE.md | 49 +++++++ src/wwgpt/checkpointing.py | 125 +++++++++++++++- tests/test_wwpgd_interval_config.py | 45 +++++- tests/test_wwpgd_projection_dose.py | 207 +++++++++++++++++++++++++++ tests/test_wwpgd_training_cadence.py | 38 ++++- 5 files changed, 457 insertions(+), 7 deletions(-) create mode 100644 docs/WWPGD_PROJECTION_DOSE.md create mode 100644 tests/test_wwpgd_projection_dose.py diff --git a/docs/WWPGD_PROJECTION_DOSE.md b/docs/WWPGD_PROJECTION_DOSE.md new file mode 100644 index 0000000..13ddfb5 --- /dev/null +++ b/docs/WWPGD_PROJECTION_DOSE.md @@ -0,0 +1,49 @@ +# WW-PGD projection dose telemetry + +The main scientific runner’s `wwpgd_projection.csv` rows report the realized intervention size and distinguish scheduled events from actual per-layer applications. + +## Scope + +This schema is applied by the shared durable artifact writer used by `run_scientific_single`. It covers the scientific Level 0–2 WW-PGD, delayed-onset, adaptive, and norm-matched-sham arms that write the root experiment artifact `wwpgd_projection.csv`. + +It does not silently redefine the separate invalid-for-science `run_single` smoke artifact or the isolated `level_0_wwpgd` package’s independent projection schema. Those paths have different lifecycle contracts and require separate schema changes if they are ever promoted to scientific evidence. + +## Realized dose + +For each projected matrix, let `W_before` be the live matrix immediately before the WW-PGD projection and `W_after` the matrix after the applied, trust-region-limited projection. The reported dose is + +```text +dose_value = ||W_after - W_before||_F / max(||W_before||_F, 1e-12) +``` + +The row fields are: + +- `dose_definition = applied_projection_delta_frobenius_over_preprojection_weight_frobenius` +- `dose_value`: the realized relative Frobenius change +- `dose_relative_frobenius`: compatibility alias of `dose_value` +- `dose_applied`: true only when the live matrix actually changed and the realized dose is positive + +This is an observed dose. It is not `blend_eta`, requested hardness, or a configured trust-region cap. + +## Scheduled event indexing + +Within the main scientific runner, `projection_event` is zero-based. Rows also include: + +- `projection_event_index_base = 0` +- `projection_event_number = projection_event + 1` +- `is_first_scheduled_projection_event = (projection_event == 0)` + +The scheduled-event flag is intentionally not called `is_first_projection_event`. Warm-up, delayed onset, adaptive gates, zero hardness, unchanged candidates, and trust-region behavior can make the first scheduled event different from the first actual application. + +## First actual application + +Actual applications are counted independently for each `layer_name` from the durable CSV prefix: + +- `layer_application_index`: one-based application count for an applied row; empty when no dose was applied +- `is_first_applied_projection`: true exactly when `layer_application_index == 1` + +The count is reconstructed from the durable projection-row history before every append. Checkpoint recovery first reconciles the CSV to its committed prefix, so the count remains correct after resume or process restart without adding process-local counters to the optimizer, controller, model, or checkpoint state. + +## Compatibility + +These fields extend the main scientific projection CSV schema but do not change model weights, optimizer state, projection mathematics, spectral targets, or defaults. An incomplete scientific run created with an older projection CSV header is rejected before append; start a fresh run under the new schema instead. diff --git a/src/wwgpt/checkpointing.py b/src/wwgpt/checkpointing.py index 9e85789..2c17cb6 100644 --- a/src/wwgpt/checkpointing.py +++ b/src/wwgpt/checkpointing.py @@ -1,5 +1,5 @@ from __future__ import annotations -import json, os, time, random, hashlib +import json, os, time, random, hashlib, math from pathlib import Path from typing import Any import numpy as np @@ -9,6 +9,11 @@ CODE_VERSION_COMPAT = ("git_commit", "git_dirty", "weightwatcher_version", "wwpgd_commit", "torch_version", "optimizer_implementation_version") REQUIRED_COMPAT=("configuration_hash","data_hash","tokenizer_hash","initialization_hash","model_configuration_hash","training_configuration_hash","wwpgd_configuration_hash","validation_probe_hash","training_probe_hash","scientific_schema_version","optimizer_fingerprint", *CODE_VERSION_COMPAT) +WWPGD_PROJECTION_DOSE_DEFINITION = ( + "applied_projection_delta_frobenius_over_preprojection_weight_frobenius" +) +WWPGD_PROJECTION_EVENT_INDEX_BASE = 0 + def stable_hash(obj: Any) -> str: return hashlib.sha256(json.dumps(obj, sort_keys=True, default=str, separators=(",", ":")).encode()).hexdigest() @@ -65,12 +70,130 @@ def _sha256_file(path: Path) -> str: h.update(chunk) return h.hexdigest() + +def _projection_bool(value: Any, *, field: str) -> bool: + if isinstance(value, bool): + return value + token = str(value).strip().lower() + if token in {"true", "1"}: + return True + if token in {"false", "0", ""}: + return False + raise ValueError(f"{field} must be boolean-compatible, got {value!r}") + + +def _projection_event_index(value: Any) -> int: + if isinstance(value, bool): + raise ValueError("projection_event must be an integer") + try: + numeric = float(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("projection_event must be an integer") from exc + if not math.isfinite(numeric) or not numeric.is_integer(): + raise ValueError("projection_event must be an integer") + event = int(numeric) + if event < 0: + raise ValueError("projection_event must be non-negative") + return event + + +def _existing_projection_application_counts(path: Path) -> dict[str, int]: + import csv + if not path.exists() or not path.stat().st_size: + return {} + with path.open(newline="") as handle: + reader = csv.DictReader(handle) + fields = set(reader.fieldnames or ()) + required = {"layer_name", "dose_applied", "layer_application_index"} + missing = sorted(required - fields) + if missing: + raise ValueError( + "existing WWPGD projection CSV uses an older telemetry schema; " + f"missing fields {missing}. Start a fresh run instead of appending." + ) + counts: dict[str, int] = {} + for row in reader: + layer_name = str(row.get("layer_name") or "").strip() + if not layer_name: + raise ValueError("existing WWPGD projection row is missing layer_name") + if _projection_bool(row.get("dose_applied", False), field="dose_applied"): + counts[layer_name] = counts.get(layer_name, 0) + 1 + return counts + + +def _enrich_wwpgd_projection_rows( + path: Path, + rows: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Add exact dose and resume-safe first-application identity to projection rows. + + ``projection_event`` is zero-based. First actual application is tracked per + layer from the durable CSV prefix rather than from process-local counters, so + the annotation remains correct across checkpoint resume and restart. + """ + application_counts = _existing_projection_application_counts(path) + enriched: list[dict[str, Any]] = [] + for source in rows: + row = dict(source) + if "is_first_projection_event" in row: + raise ValueError( + "ambiguous is_first_projection_event is unsupported; use " + "is_first_scheduled_projection_event and is_first_applied_projection" + ) + if "projection_event" not in row: + raise ValueError("WWPGD projection row is missing projection_event") + event = _projection_event_index(row["projection_event"]) + + raw_dose = row.get("relative_frobenius_change_applied") + if raw_dose is None or raw_dose == "" or isinstance(raw_dose, bool): + raise ValueError( + "WWPGD projection row is missing relative_frobenius_change_applied" + ) + try: + dose = float(raw_dose) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError( + "relative_frobenius_change_applied must be numeric" + ) from exc + if not math.isfinite(dose) or dose < 0.0: + raise ValueError( + "relative_frobenius_change_applied must be finite and non-negative" + ) + + layer_name = str(row.get("layer_name") or "").strip() + if not layer_name: + raise ValueError("WWPGD projection row is missing layer_name") + changed = _projection_bool(row.get("changed", False), field="changed") + dose_applied = changed and dose > 0.0 + application_index = None + if dose_applied: + application_index = application_counts.get(layer_name, 0) + 1 + application_counts[layer_name] = application_index + row.update( + { + "dose_definition": WWPGD_PROJECTION_DOSE_DEFINITION, + "dose_value": dose, + "dose_relative_frobenius": dose, + "dose_applied": dose_applied, + "layer_application_index": application_index, + "is_first_applied_projection": application_index == 1, + "projection_event_index_base": WWPGD_PROJECTION_EVENT_INDEX_BASE, + "projection_event_number": event + 1, + "is_first_scheduled_projection_event": event == 0, + } + ) + enriched.append(row) + return enriched + + def append_csv_records(path: Path, rows: list[dict[str, Any]]) -> None: """Durably append records. A flush boundary is a transaction boundary.""" import csv if not rows: return path = Path(path); path.parent.mkdir(parents=True, exist_ok=True) + if path.name == "wwpgd_projection.csv": + rows = _enrich_wwpgd_projection_rows(path, rows) # Cached-endpoint fast-relaxation rows already have the dedicated # wwpgd_endpoint_relaxation.csv stream. Do not duplicate them into # wwpgd_controller.csv, whose stable schema is the slow measurement record. diff --git a/tests/test_wwpgd_interval_config.py b/tests/test_wwpgd_interval_config.py index 8d6747c..b75f07a 100644 --- a/tests/test_wwpgd_interval_config.py +++ b/tests/test_wwpgd_interval_config.py @@ -51,12 +51,34 @@ def init_state(cfg): return {k: v.detach().clone() for k, v in m.state_dict().items()}, "init" +def _fake_projection_row(*, event_index, actual_step, actual_tokens_seen=None): + row = { + "projection_event": event_index, + "actual_step": actual_step, + "layer_name": "fake", + "relative_frobenius_change_applied": 0.01, + "changed": True, + } + if actual_tokens_seen is not None: + row["actual_tokens_seen"] = actual_tokens_seen + return row + + def test_extension_interval_cadences_and_skips(monkeypatch): calls = [] pre_calls = [] monkeypatch.setattr("wwgpt.train.weightwatcher_details", lambda model: pre_calls.append("pre") or object()) monkeypatch.setattr("wwgpt.train.build_stock_wwpgd_candidate", lambda *args, **kwargs: SimpleNamespace(internal_diagnostics=[])) - monkeypatch.setattr("wwgpt.train.apply_external_wwpgd", lambda *args, **kw: calls.append(kw) or [{"projection_event": kw["event_index"], "actual_step": kw["actual_step"], "actual_tokens_seen": kw["actual_tokens_seen"], "layer_name": "fake"}]) + monkeypatch.setattr( + "wwgpt.train.apply_external_wwpgd", + lambda *args, **kw: calls.append(kw) or [ + _fake_projection_row( + event_index=kw["event_index"], + actual_step=kw["actual_step"], + actual_tokens_seen=kw["actual_tokens_seen"], + ) + ], + ) ext = WWPGDExtension(WWPGDConfig(), interval=2) rows_by_step = [ext.after_optimizer_step(model=object(), optimizer_step=s, total_optimizer_steps=6, tokens_seen=s*10, collect_pre_details=True)[1] for s in range(1, 7)] @@ -82,7 +104,16 @@ def test_extension_rejects_non_positive_interval(): def test_training_interval_manifest_and_counts(monkeypatch, tmp_path): calls = [] monkeypatch.setattr("wwgpt.train.spectral_summary", lambda *a, **k: []) - monkeypatch.setattr("wwgpt.train.apply_external_wwpgd", lambda *a, **kw: calls.append(kw["actual_step"]) or [{"projection_event": kw["event_index"], "actual_step": kw["actual_step"], "actual_tokens_seen": kw["actual_tokens_seen"], "layer_name": "fake"}]) + monkeypatch.setattr( + "wwgpt.train.apply_external_wwpgd", + lambda *a, **kw: calls.append(kw["actual_step"]) or [ + _fake_projection_row( + event_index=kw["event_index"], + actual_step=kw["actual_step"], + actual_tokens_seen=kw["actual_tokens_seen"], + ) + ], + ) cfg = tiny_cfg(steps=6, interval=2) state, h = init_state(cfg) run = run_scientific_single(tmp_path, "adamw_wwpgd", 3, cfg, tiny_data(), "pair", state, h, 0, 1, device="cpu") @@ -100,7 +131,15 @@ def test_training_interval_manifest_and_counts(monkeypatch, tmp_path): def test_run_function_cli_override_precedes_config(monkeypatch, tmp_path): calls = [] monkeypatch.setattr("wwgpt.train.spectral_summary", lambda *a, **k: []) - monkeypatch.setattr("wwgpt.train.apply_external_wwpgd", lambda *a, **kw: calls.append(kw["actual_step"]) or [{"projection_event": kw["event_index"], "actual_step": kw["actual_step"], "layer_name": "fake"}]) + monkeypatch.setattr( + "wwgpt.train.apply_external_wwpgd", + lambda *a, **kw: calls.append(kw["actual_step"]) or [ + _fake_projection_row( + event_index=kw["event_index"], + actual_step=kw["actual_step"], + ) + ], + ) cfg = tiny_cfg(steps=8, interval=2) state, h = init_state(cfg) run_scientific_single(tmp_path, "adamw_wwpgd", 3, cfg, tiny_data(), "pair", state, h, 0, 1, device="cpu", ww_interval=4) diff --git a/tests/test_wwpgd_projection_dose.py b/tests/test_wwpgd_projection_dose.py new file mode 100644 index 0000000..8e04905 --- /dev/null +++ b/tests/test_wwpgd_projection_dose.py @@ -0,0 +1,207 @@ +"""Regression coverage for exact, resume-safe WW-PGD projection-dose telemetry.""" +from __future__ import annotations + +import csv +import math + +import pytest + +from wwgpt.checkpointing import append_csv_records, csv_commit, reconcile_csv_artifacts + + +DOSE_DEFINITION = ( + "applied_projection_delta_frobenius_over_preprojection_weight_frobenius" +) + + +def _read_rows(path): + with path.open(newline="", encoding="utf-8") as handle: + return list(csv.DictReader(handle)) + + +def _projection_row( + *, + event: int | float, + step: int, + dose: float, + changed: bool, + layer: str = "blocks.0.attn.query", +): + return { + "layer_name": layer, + "projection_event": event, + "actual_step": step, + "relative_frobenius_change_applied": dose, + "changed": changed, + } + + +def test_projection_csv_reports_exact_realized_dose_and_first_application(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + source = _projection_row(event=0, step=10, dose=0.125, changed=True) + + append_csv_records(path, [source]) + + row = _read_rows(path)[0] + assert row["dose_definition"] == DOSE_DEFINITION + assert float(row["dose_value"]) == pytest.approx(0.125) + assert float(row["dose_relative_frobenius"]) == pytest.approx(0.125) + assert row["dose_value"] == row["relative_frobenius_change_applied"] + assert row["dose_applied"] == "True" + assert int(row["layer_application_index"]) == 1 + assert row["is_first_applied_projection"] == "True" + assert int(row["projection_event_index_base"]) == 0 + assert int(row["projection_event_number"]) == 1 + assert row["is_first_scheduled_projection_event"] == "True" + assert "is_first_projection_event" not in row + assert "dose_value" not in source # durable enrichment must not mutate callers + + +def test_first_scheduled_event_can_precede_first_actual_application(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records( + path, + [ + _projection_row(event=0, step=10, dose=0.0, changed=False), + _projection_row(event=1, step=20, dose=0.05, changed=True), + ], + ) + + rows = _read_rows(path) + assert rows[0]["is_first_scheduled_projection_event"] == "True" + assert rows[0]["dose_applied"] == "False" + assert rows[0]["layer_application_index"] == "" + assert rows[0]["is_first_applied_projection"] == "False" + assert float(rows[0]["dose_value"]) == 0.0 + + assert rows[1]["is_first_scheduled_projection_event"] == "False" + assert rows[1]["dose_applied"] == "True" + assert int(rows[1]["layer_application_index"]) == 1 + assert rows[1]["is_first_applied_projection"] == "True" + assert int(rows[1]["projection_event_number"]) == 2 + + first_applied = next(row for row in rows if row["is_first_applied_projection"] == "True") + assert int(first_applied["projection_event"]) == 1 + assert int(first_applied["actual_step"]) == 20 + + +def test_layer_application_index_survives_sequential_appends(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records(path, [_projection_row(event=0, step=10, dose=0.0, changed=False)]) + append_csv_records(path, [_projection_row(event=1, step=20, dose=0.05, changed=True)]) + append_csv_records(path, [_projection_row(event=2, step=30, dose=0.04, changed=True)]) + + rows = _read_rows(path) + assert len(rows) == 3 + assert rows[0]["layer_application_index"] == "" + assert int(rows[1]["layer_application_index"]) == 1 + assert rows[1]["is_first_applied_projection"] == "True" + assert int(rows[2]["layer_application_index"]) == 2 + assert rows[2]["is_first_applied_projection"] == "False" + assert int(rows[2]["projection_event_number"]) == 3 + + +def test_reconciled_uncommitted_suffix_does_not_advance_application_count(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records(path, [_projection_row(event=0, step=10, dose=0.05, changed=True)]) + committed = csv_commit(path) + + # Simulate rows flushed after the last durable checkpoint, then a restart. + append_csv_records(path, [_projection_row(event=1, step=20, dose=0.04, changed=True)]) + assert int(_read_rows(path)[-1]["layer_application_index"]) == 2 + reconcile_csv_artifacts(tmp_path, {path.name: committed}) + + append_csv_records(path, [_projection_row(event=2, step=30, dose=0.03, changed=True)]) + rows = _read_rows(path) + assert len(rows) == 2 + assert int(rows[-1]["layer_application_index"]) == 2 + assert rows[-1]["is_first_applied_projection"] == "False" + + +def test_first_application_is_tracked_independently_per_layer(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records( + path, + [ + _projection_row( + event=0, + step=10, + dose=0.05, + changed=True, + layer="blocks.0.attn.query", + ), + _projection_row( + event=0, + step=10, + dose=0.03, + changed=True, + layer="blocks.0.attn.key", + ), + ], + ) + append_csv_records( + path, + [ + _projection_row( + event=1, + step=20, + dose=0.02, + changed=True, + layer="blocks.0.attn.query", + ) + ], + ) + + rows = _read_rows(path) + assert [int(rows[index]["layer_application_index"]) for index in (0, 1, 2)] == [1, 1, 2] + assert [rows[index]["is_first_applied_projection"] for index in (0, 1, 2)] == [ + "True", + "True", + "False", + ] + + +def test_non_projection_csv_rows_are_unchanged(tmp_path): + path = tmp_path / "metrics.csv" + append_csv_records(path, [{"step": 1, "loss": 2.5}]) + assert _read_rows(path) == [{"step": "1", "loss": "2.5"}] + + +def test_old_projection_schema_is_rejected_before_append(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter( + handle, + fieldnames=[ + "layer_name", + "projection_event", + "actual_step", + "relative_frobenius_change_applied", + "changed", + ], + ) + writer.writeheader() + writer.writerow(_projection_row(event=0, step=10, dose=0.05, changed=True)) + + before = path.read_bytes() + with pytest.raises(ValueError, match="older telemetry schema"): + append_csv_records(path, [_projection_row(event=1, step=20, dose=0.04, changed=True)]) + assert path.read_bytes() == before + + +@pytest.mark.parametrize( + "row, message", + [ + ({"relative_frobenius_change_applied": 0.1, "changed": True}, "projection_event"), + (_projection_row(event=-1, step=1, dose=0.1, changed=True), "non-negative"), + (_projection_row(event=0.5, step=1, dose=0.1, changed=True), "must be an integer"), + (_projection_row(event=0, step=1, dose=-0.1, changed=True), "finite and non-negative"), + (_projection_row(event=0, step=1, dose=math.nan, changed=True), "finite and non-negative"), + (_projection_row(event=0, step=1, dose=0.1, changed=True, layer=""), "layer_name"), + ], +) +def test_invalid_projection_telemetry_fails_before_writing(tmp_path, row, message): + path = tmp_path / "wwpgd_projection.csv" + with pytest.raises(ValueError, match=message): + append_csv_records(path, [row]) + assert not path.exists() diff --git a/tests/test_wwpgd_training_cadence.py b/tests/test_wwpgd_training_cadence.py index 30a5fab..2078da9 100644 --- a/tests/test_wwpgd_training_cadence.py +++ b/tests/test_wwpgd_training_cadence.py @@ -59,6 +59,19 @@ def _init_state(cfg: ExperimentConfig): return state, "tiny-init" +def _fake_projection_row(*, event_index, actual_step, actual_tokens_seen=None, layer_name="fake"): + row = { + "projection_event": event_index, + "actual_step": actual_step, + "layer_name": layer_name, + "relative_frobenius_change_applied": 0.01, + "changed": True, + } + if actual_tokens_seen is not None: + row["actual_tokens_seen"] = actual_tokens_seen + return row + + def test_wwpgd_runs_once_per_successful_optimizer_step(monkeypatch, tmp_path: Path): steps = 3 calls = [] @@ -67,8 +80,18 @@ def fake_apply(model, *, event_index, scheduled_token_fraction, actual_step, act assert stock_candidate is not None calls.append(actual_step) return [ - {"projection_event": event_index, "layer_name": "blocks.0.attn.key"}, - {"projection_event": event_index, "layer_name": "blocks.0.attn.value"}, + _fake_projection_row( + event_index=event_index, + actual_step=actual_step, + actual_tokens_seen=actual_tokens_seen, + layer_name="blocks.0.attn.key", + ), + _fake_projection_row( + event_index=event_index, + actual_step=actual_step, + actual_tokens_seen=actual_tokens_seen, + layer_name="blocks.0.attn.value", + ), ] monkeypatch.setattr("wwgpt.train.apply_external_wwpgd", fake_apply) @@ -136,7 +159,16 @@ def test_checkpoint_resume_is_deterministic_and_complete(monkeypatch, tmp_path: ) init_state, init_hash = _init_state(cfg) monkeypatch.setattr("wwgpt.train.spectral_summary", lambda *args, **kwargs: []) - monkeypatch.setattr("wwgpt.train.apply_external_wwpgd", lambda *args, **kwargs: [{"projection_event": kwargs["event_index"], "actual_step": kwargs["actual_step"], "actual_tokens_seen": kwargs["actual_tokens_seen"], "layer_name": "fake"}]) + monkeypatch.setattr( + "wwgpt.train.apply_external_wwpgd", + lambda *args, **kwargs: [ + _fake_projection_row( + event_index=kwargs["event_index"], + actual_step=kwargs["actual_step"], + actual_tokens_seen=kwargs["actual_tokens_seen"], + ) + ], + ) full = run_scientific_single(tmp_path / "full", "adamw_wwpgd", 7, cfg, _tiny_data(), "pair_tiny", init_state, init_hash, 0, 1, device="cpu") From dbb316332083b0c846138b8882d7cb0f9e30a412 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 8 Aug 2026 21:27:57 -0700 Subject: [PATCH 2/5] Avoid quadratic projection-history rescans --- src/wwgpt/checkpointing.py | 76 ++++++++++++++++++++++++++++++++++---- 1 file changed, 69 insertions(+), 7 deletions(-) diff --git a/src/wwgpt/checkpointing.py b/src/wwgpt/checkpointing.py index 2c17cb6..5f2aca8 100644 --- a/src/wwgpt/checkpointing.py +++ b/src/wwgpt/checkpointing.py @@ -13,6 +13,7 @@ "applied_projection_delta_frobenius_over_preprojection_weight_frobenius" ) WWPGD_PROJECTION_EVENT_INDEX_BASE = 0 +WWPGD_PROJECTION_APPLICATION_INDEX_SCHEMA_VERSION = 1 def stable_hash(obj: Any) -> str: @@ -97,7 +98,11 @@ def _projection_event_index(value: Any) -> int: return event -def _existing_projection_application_counts(path: Path) -> dict[str, int]: +def _projection_application_index_path(path: Path) -> Path: + return path.with_name(f".{path.name}.application_counts.json") + + +def _scan_projection_application_counts(path: Path) -> dict[str, int]: import csv if not path.exists() or not path.stat().st_size: return {} @@ -121,17 +126,71 @@ def _existing_projection_application_counts(path: Path) -> dict[str, int]: return counts +def _valid_projection_application_counts(value: Any) -> dict[str, int] | None: + if not isinstance(value, dict): + return None + counts: dict[str, int] = {} + for raw_name, raw_count in value.items(): + name = str(raw_name).strip() + if not name or isinstance(raw_count, bool): + return None + try: + count = int(raw_count) + except (TypeError, ValueError, OverflowError): + return None + if count < 0 or count != raw_count: + return None + counts[name] = count + return counts + + +def _load_projection_application_counts(path: Path) -> dict[str, int]: + if not path.exists() or not path.stat().st_size: + return {} + stat = path.stat() + index_path = _projection_application_index_path(path) + try: + index = json.loads(index_path.read_text()) + except (OSError, TypeError, json.JSONDecodeError): + index = None + if isinstance(index, dict): + counts = _valid_projection_application_counts(index.get("counts")) + if ( + int(index.get("schema_version", -1)) + == WWPGD_PROJECTION_APPLICATION_INDEX_SCHEMA_VERSION + and int(index.get("source_size_bytes", -1)) == stat.st_size + and int(index.get("source_mtime_ns", -1)) == stat.st_mtime_ns + and counts is not None + ): + return counts + return _scan_projection_application_counts(path) + + +def _store_projection_application_counts(path: Path, counts: dict[str, int]) -> None: + stat = path.stat() + _atomic_write_json( + _projection_application_index_path(path), + { + "schema_version": WWPGD_PROJECTION_APPLICATION_INDEX_SCHEMA_VERSION, + "source_size_bytes": stat.st_size, + "source_mtime_ns": stat.st_mtime_ns, + "counts": dict(sorted(counts.items())), + }, + ) + + def _enrich_wwpgd_projection_rows( path: Path, rows: list[dict[str, Any]], -) -> list[dict[str, Any]]: +) -> tuple[list[dict[str, Any]], dict[str, int]]: """Add exact dose and resume-safe first-application identity to projection rows. ``projection_event`` is zero-based. First actual application is tracked per - layer from the durable CSV prefix rather than from process-local counters, so - the annotation remains correct across checkpoint resume and restart. + layer from a self-validating durable-history index. The index is rebuilt + from the reconciled CSV only after a crash, truncation, or metadata mismatch, + avoiding an O(n^2) history scan during interval-one experiments. """ - application_counts = _existing_projection_application_counts(path) + application_counts = _load_projection_application_counts(path) enriched: list[dict[str, Any]] = [] for source in rows: row = dict(source) @@ -183,7 +242,7 @@ def _enrich_wwpgd_projection_rows( } ) enriched.append(row) - return enriched + return enriched, application_counts def append_csv_records(path: Path, rows: list[dict[str, Any]]) -> None: @@ -192,8 +251,9 @@ def append_csv_records(path: Path, rows: list[dict[str, Any]]) -> None: if not rows: return path = Path(path); path.parent.mkdir(parents=True, exist_ok=True) + projection_application_counts = None if path.name == "wwpgd_projection.csv": - rows = _enrich_wwpgd_projection_rows(path, rows) + rows, projection_application_counts = _enrich_wwpgd_projection_rows(path, rows) # Cached-endpoint fast-relaxation rows already have the dedicated # wwpgd_endpoint_relaxation.csv stream. Do not duplicate them into # wwpgd_controller.csv, whose stable schema is the slow measurement record. @@ -240,6 +300,8 @@ def append_csv_records(path: Path, rows: list[dict[str, Any]]) -> None: writer.writeheader() writer.writerows(rows) f.flush(); os.fsync(f.fileno()) + if projection_application_counts is not None: + _store_projection_application_counts(path, projection_application_counts) def csv_commit(path: Path) -> dict[str, Any]: """Describe the exact durable prefix referenced by a checkpoint.""" From 788bfd61d40aaadcc51ed630605711c3e34fc735 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 8 Aug 2026 21:28:49 -0700 Subject: [PATCH 3/5] Test projection application index cache and rebuild --- tests/test_wwpgd_projection_dose.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/test_wwpgd_projection_dose.py b/tests/test_wwpgd_projection_dose.py index 8e04905..bc641eb 100644 --- a/tests/test_wwpgd_projection_dose.py +++ b/tests/test_wwpgd_projection_dose.py @@ -6,6 +6,7 @@ import pytest +import wwgpt.checkpointing as checkpointing from wwgpt.checkpointing import append_csv_records, csv_commit, reconcile_csv_artifacts @@ -101,6 +102,30 @@ def test_layer_application_index_survives_sequential_appends(tmp_path): assert int(rows[2]["projection_event_number"]) == 3 +def test_application_count_index_avoids_full_history_rescan(monkeypatch, tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records(path, [_projection_row(event=0, step=10, dose=0.05, changed=True)]) + index_path = path.with_name(f".{path.name}.application_counts.json") + assert index_path.is_file() + + def fail_scan(_path): + raise AssertionError("valid sidecar index should avoid a full CSV rescan") + + monkeypatch.setattr(checkpointing, "_scan_projection_application_counts", fail_scan) + append_csv_records(path, [_projection_row(event=1, step=20, dose=0.04, changed=True)]) + assert int(_read_rows(path)[-1]["layer_application_index"]) == 2 + + +def test_invalid_sidecar_metadata_rebuilds_from_csv(tmp_path): + path = tmp_path / "wwpgd_projection.csv" + append_csv_records(path, [_projection_row(event=0, step=10, dose=0.05, changed=True)]) + index_path = path.with_name(f".{path.name}.application_counts.json") + index_path.write_text('{"schema_version": "broken", "counts": {}}') + + append_csv_records(path, [_projection_row(event=1, step=20, dose=0.04, changed=True)]) + assert int(_read_rows(path)[-1]["layer_application_index"]) == 2 + + def test_reconciled_uncommitted_suffix_does_not_advance_application_count(tmp_path): path = tmp_path / "wwpgd_projection.csv" append_csv_records(path, [_projection_row(event=0, step=10, dose=0.05, changed=True)]) From 94682a44a157ce5967858a5732e60594ea9a97d3 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 8 Aug 2026 21:32:14 -0700 Subject: [PATCH 4/5] Rebuild malformed projection application indexes safely --- src/wwgpt/checkpointing.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/src/wwgpt/checkpointing.py b/src/wwgpt/checkpointing.py index 5f2aca8..7386f97 100644 --- a/src/wwgpt/checkpointing.py +++ b/src/wwgpt/checkpointing.py @@ -151,18 +151,24 @@ def _load_projection_application_counts(path: Path) -> dict[str, int]: index_path = _projection_application_index_path(path) try: index = json.loads(index_path.read_text()) - except (OSError, TypeError, json.JSONDecodeError): + except (OSError, UnicodeError, TypeError, json.JSONDecodeError): index = None if isinstance(index, dict): counts = _valid_projection_application_counts(index.get("counts")) - if ( - int(index.get("schema_version", -1)) - == WWPGD_PROJECTION_APPLICATION_INDEX_SCHEMA_VERSION - and int(index.get("source_size_bytes", -1)) == stat.st_size - and int(index.get("source_mtime_ns", -1)) == stat.st_mtime_ns - and counts is not None - ): - return counts + try: + schema_version = int(index.get("schema_version", -1)) + source_size_bytes = int(index.get("source_size_bytes", -1)) + source_mtime_ns = int(index.get("source_mtime_ns", -1)) + except (TypeError, ValueError, OverflowError): + pass + else: + if ( + schema_version == WWPGD_PROJECTION_APPLICATION_INDEX_SCHEMA_VERSION + and source_size_bytes == stat.st_size + and source_mtime_ns == stat.st_mtime_ns + and counts is not None + ): + return counts return _scan_projection_application_counts(path) From 2bb8ceff94397548e0ad4c4120fa2011f81e8317 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Sat, 8 Aug 2026 21:32:43 -0700 Subject: [PATCH 5/5] Document efficient resume-safe application counting --- docs/WWPGD_PROJECTION_DOSE.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/WWPGD_PROJECTION_DOSE.md b/docs/WWPGD_PROJECTION_DOSE.md index 13ddfb5..3232669 100644 --- a/docs/WWPGD_PROJECTION_DOSE.md +++ b/docs/WWPGD_PROJECTION_DOSE.md @@ -37,12 +37,14 @@ The scheduled-event flag is intentionally not called `is_first_projection_event` ## First actual application -Actual applications are counted independently for each `layer_name` from the durable CSV prefix: +Actual applications are counted independently for each `layer_name`: - `layer_application_index`: one-based application count for an applied row; empty when no dose was applied - `is_first_applied_projection`: true exactly when `layer_application_index == 1` -The count is reconstructed from the durable projection-row history before every append. Checkpoint recovery first reconciles the CSV to its committed prefix, so the count remains correct after resume or process restart without adding process-local counters to the optimizer, controller, model, or checkpoint state. +The projection CSV remains the source of truth. A small hidden sidecar caches only the current per-layer counts together with the CSV byte size and nanosecond modification time. Normal interval-one runs therefore do not rescan the complete projection history after every optimizer step. If the sidecar is absent, malformed, stale, or disagrees with a reconciled/truncated CSV, it is discarded and rebuilt from the durable CSV before the next append. + +Checkpoint recovery first reconciles the CSV to its committed prefix. The first-application annotation therefore remains correct after interruption or process restart without adding counters to the optimizer, controller, model, or checkpoint state. ## Compatibility