From 9d093889242c4759b575bbe0b8152fb9f18e3302 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 15:34:47 +0000 Subject: [PATCH 1/9] Preserve planner refusal attribution and compact report context --- dev/trainer_rank_planner_design.md | 13 ++- src/art/trainer_rank/_impl.py | 24 ++++ src/art/trainer_rank/_planner_misses.py | 72 ++++++++++-- .../test_trainer_rank_planner_evidence.py | 107 +++++++++++++++++- 4 files changed, 203 insertions(+), 13 deletions(-) diff --git a/dev/trainer_rank_planner_design.md b/dev/trainer_rank_planner_design.md index d63532cd0..c747cdddf 100644 --- a/dev/trainer_rank_planner_design.md +++ b/dev/trainer_rank_planner_design.md @@ -889,8 +889,17 @@ wave/participant identities, execution breadcrumbs, reserved summary storage, exact delivery ACKs, and aggregate spool budgets remain follow-up work. Current spool-full/process-loss limits still apply. The selected-plan replay keeps its existing completeness limits; this is not a full GPU failure reproduction claim. -New planning-event reports cap replay size at 256KiB, preserving scalar evidence -with an explicit incomplete reason when replay exceeds that cap. Ordinary +Planning-event reports cap size at 256KiB. Oversized replay drops bulk request +and layout arrays first, retaining whole compact source, rank/device, model and +estimator fields that fit; `omitted_fields` identifies every removed field. +Such reports remain explicitly incomplete, including when an individual compact +field itself exceeds the cap. This does not add a second capture or upload. +After an exhausted split ladder, the ordinary refusal may retain an unsplit +context plan but the final split-rung check. Its report labels +`candidate_matches_check=false`, leaves the denied candidate prediction unknown, +and refuses complete replay; it never attributes that check to the context plan. +No extra rejected plan is materialized and admission behavior is unchanged. +Ordinary refusals/planning errors can use only the first 64 entries/16MiB of the rank spool, counting all existing entries. Thus they cannot alone consume the original 1024-entry/256MiB allowance used by misses and OOMs (including a planning event diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index d5d2a3c4c..f83987dee 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1077,6 +1077,7 @@ class _ForwardRefusal: message: str overridable: bool = True candidate: Any = None + check_matches_plan: bool = True def error(self, context: str) -> TrainerRankMemoryError: return _memory_error( @@ -2736,6 +2737,9 @@ def _find_admissible_forward( "the bounded ladder (2, 4, ..., one request per subforward) is " "predicted to exceed available memory once all returned graphs are " "live together", + # Without the override the rejected rung is not retained. Its + # check must not be attributed to the unsplit context plan. + check_matches_plan=getattr(self, "_allow_oversized_batches", False), ) def _expert_parallel_active(self) -> bool: @@ -5797,6 +5801,26 @@ def _report_planning_failure( }, ) self._fill_planner_snapshot(refusal.plan, refusal.check, observation) + snapshot = observation["replay"] + + def replay() -> dict[str, Any]: + payload = snapshot() + return { + **payload, + "candidate_matches_check": refusal.check_matches_plan, + "incomplete_reasons": [ + *payload.get("incomplete_reasons", []), + *( + [] + if refusal.check_matches_plan + else ["candidate does not describe denying check"] + ), + ], + } + + observation["replay"] = replay + if not refusal.check_matches_plan: + observation["predicted"] = None self._planner_reporter.report( predicted_peak_bytes=observation["predicted"], observed_peak_bytes=None, diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index 679e15fe9..55a65b0b4 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -83,17 +83,65 @@ def _warn(reason: str) -> None: pass -def _encode(record: dict[str, Any]) -> bytes: +class _ReportTooLarge(ValueError): + pass + + +def _encode(record: dict[str, Any], *, limit: int = MAX_REPORT_BYTES) -> bytes: chunks = bytearray() encoder = json.JSONEncoder(sort_keys=True, separators=(",", ":"), allow_nan=False) for chunk in encoder.iterencode(record): encoded = chunk.encode("utf-8") - if len(chunks) + len(encoded) + 1 > MAX_REPORT_BYTES: - raise ValueError("report exceeds byte limit") + if len(chunks) + len(encoded) + 1 > limit: + raise _ReportTooLarge("report exceeds byte limit") chunks.extend(encoded) return bytes(chunks) + b"\n" +def _compact_planning_record(record: dict[str, Any]) -> dict[str, Any]: + """Keep whole compact facts before bulk inputs; never claim partial replay.""" + payload = record["replay"] + reasons = [ + *record["incomplete_reasons"], + "planning replay exceeds report limit", + ] + omitted: list[str] = [] + compact: dict[str, Any] = { + "incomplete_reasons": reasons, + "omitted_fields": omitted, + } + result = { + **record, + "replay": compact, + "replay_complete": False, + "incomplete_reasons": reasons, + } + priority = ( + "source_files", + "source_scope", + "rank", + "device", + "model", + "model_identity", + "candidate_matches_check", + "memory_replay", + ) + for key in dict.fromkeys((*priority, *payload)): + if key not in payload or key == "incomplete_reasons": + continue + if key in {"requests", "layouts", "omitted_fields"}: + omitted.append(key) + continue + compact[key] = payload[key] + try: + # Reserve space for explicit omission names as fields are rejected. + _encode(result, limit=MAX_PLANNING_REPORT_BYTES - 4096) + except _ReportTooLarge: + del compact[key] + omitted.append(key) + return result + + def _source_files() -> dict[str, dict[str, str | int]]: """Fingerprint current module files, not an attestation of loaded bytecode.""" result = {} @@ -361,7 +409,18 @@ def report( ) if not record["replay_complete"] and not record["incomplete_reasons"]: record["incomplete_reasons"] = ["memory replay inputs unavailable"] - raw = _encode(record) + try: + raw = _encode( + record, + limit=MAX_PLANNING_REPORT_BYTES + if planning + else MAX_REPORT_BYTES, + ) + except _ReportTooLarge: + if not planning: + raise + record = _compact_planning_record(record) + raw = _encode(record, limit=MAX_PLANNING_REPORT_BYTES) except Exception as exc: record["replay"] = None record["replay_complete"] = False @@ -369,11 +428,6 @@ def report( f"replay unavailable: {type(exc).__name__}" ] raw = _encode(record) - if planning and len(raw) > MAX_PLANNING_REPORT_BYTES: - record["replay"] = None - record["replay_complete"] = False - record["incomplete_reasons"] = ["planning replay exceeds report limit"] - raw = _encode(record) path = persist_report( raw, self.spool_dir, diff --git a/tests/unit/test_trainer_rank_planner_evidence.py b/tests/unit/test_trainer_rank_planner_evidence.py index 263fa8150..31f45448d 100644 --- a/tests/unit/test_trainer_rank_planner_evidence.py +++ b/tests/unit/test_trainer_rank_planner_evidence.py @@ -156,6 +156,53 @@ def test_final_refusal_emits_without_forward_or_memory_window(monkeypatch, tmp_p assert not getattr(rank, "_planner_active_observations", {}) +@pytest.mark.parametrize("override", [False, True]) +def test_exhausted_ladder_identifies_unmatched_check(monkeypatch, override): + rank = _oversized(monkeypatch) + rank._allow_oversized_batches = override + found = rank._find_admissible_forward( + [_request(i) for i in range(4)], + checkpoint=tr.Unset, + refusal_prefix="test refusal", + ) + assert isinstance(found, tr._ForwardRefusal) + assert found.check_matches_plan is override + + +@pytest.mark.parametrize("microbatch", [False, True]) +def test_unmatched_refusal_never_claims_denial_replay( + monkeypatch, tmp_path, microbatch +): + rank = _oversized(monkeypatch) + rank._allow_oversized_batches = False + rank._planner_reporter = reports.Reporter(5, spool_dir=tmp_path / "reports") + + def snapshot(plan, check, observation): + # Even if the context candidate's estimator is fully reproducible, + # its plan is not the final rejected rung that produced this check. + observation.update( + predicted=42, + replay=lambda: {"memory_replay": {"estimates": [42]}}, + ) + + monkeypatch.setattr(rank, "_fill_planner_snapshot", snapshot) + with pytest.raises(tr.TrainerRankMemoryError): + requests = [_request(i) for i in range(4)] + if microbatch: + rank._select_next_micro_batch([requests], 0) + else: + rank.dp_rank_forward(requests) + record = reports.validate_report( + next(rank._planner_reporter.spool_dir.glob("*.json")).read_bytes() + ) + assert record["replay"]["candidate_matches_check"] is False + assert record["predicted_peak_bytes"] is None + assert record["replay_complete"] is False + assert "candidate does not describe denying check" in record["incomplete_reasons"] + with pytest.raises(ValueError, match="incomplete replay"): + reports.replay(record) + + def test_planning_error_original_identity_and_no_prediction(scalar, tmp_path): rank, _, _, _ = scalar rank._planner_reporter = reports.Reporter(5, spool_dir=tmp_path / "reports") @@ -332,8 +379,64 @@ def test_large_planning_replay_keeps_bounded_scalar_event(tmp_path): ) assert path is not None and path.stat().st_size <= 256 * 1024 record = reports.validate_report(path.read_bytes()) - assert record["event"] == "planning_error" and record["replay"] is None - assert record["incomplete_reasons"] == ["planning replay exceeds report limit"] + assert record["event"] == "planning_error" + assert record["replay"]["omitted_fields"] == ["payload"] + assert "planning replay exceeds report limit" in record["incomplete_reasons"] + + +@pytest.mark.parametrize("tokens", [88_576, 5_000_000]) +def test_planning_cap_preserves_compact_context(tmp_path, tokens): + inputs = [1] * tokens + compact = { + "rank": 3, + "device": {"device": "cuda:3"}, + "model": "test-model", + "model_identity": {"revision": "original"}, + "memory_replay": {"rank": {"num_layers": 3}, "estimates": [{"coefficient": 2}]}, + "candidate_matches_check": False, + } + payload = { + **compact, + "requests": [{"input_tokens": inputs}], + "layouts": [{"input_tokens": [inputs]}], + "incomplete_reasons": ["candidate does not describe denying check"], + } + path = reports.Reporter(5, spool_dir=tmp_path).report( + predicted_peak_bytes=None, + observed_peak_bytes=None, + phase="planning", + event="admission_refused", + replay_factory=lambda: payload, + ) + assert path is not None and path.stat().st_size <= reports.MAX_PLANNING_REPORT_BYTES + record = reports.validate_report(path.read_bytes()) + for key, value in compact.items(): + assert record["replay"][key] == value + assert record["replay"]["source_files"] == reports._source_files() + assert set(record["replay"]["omitted_fields"]) == {"requests", "layouts"} + assert "candidate does not describe denying check" in record["incomplete_reasons"] + assert not record["replay_complete"] + assert payload["requests"][0]["input_tokens"] is inputs + + +def test_oversized_compact_field_is_omitted_whole(tmp_path): + path = reports.Reporter(5, spool_dir=tmp_path).report( + predicted_peak_bytes=None, + observed_peak_bytes=None, + phase="planning", + event="planning_error", + replay_factory=lambda: { + "rank": 7, + "memory_replay": {"estimates": ["x" * 300_000]}, + }, + ) + assert path is not None and path.stat().st_size <= reports.MAX_PLANNING_REPORT_BYTES + record = reports.validate_report(path.read_bytes()) + assert record["replay"]["rank"] == 7 + assert record["replay"]["source_files"] == reports._source_files() + assert record["replay"]["omitted_fields"] == ["memory_replay"] + assert "memory_replay" not in record["replay"] + assert record["replay_complete"] is False def test_summary_trimming_preserves_first_selected(monkeypatch): From edf0e199d84a002cd29851913a1567cccec775e4 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 15:56:31 +0000 Subject: [PATCH 2/9] Bound compact report field processing and preserve error reasons --- dev/trainer_rank_planner_design.md | 4 ++- src/art/trainer_rank/_planner_misses.py | 27 ++++++++++---- .../test_trainer_rank_planner_evidence.py | 36 +++++++++++++++++++ .../unit/test_trainer_rank_planner_reports.py | 2 ++ 4 files changed, 62 insertions(+), 7 deletions(-) diff --git a/dev/trainer_rank_planner_design.md b/dev/trainer_rank_planner_design.md index c747cdddf..9be81bb35 100644 --- a/dev/trainer_rank_planner_design.md +++ b/dev/trainer_rank_planner_design.md @@ -891,7 +891,9 @@ spool-full/process-loss limits still apply. The selected-plan replay keeps its existing completeness limits; this is not a full GPU failure reproduction claim. Planning-event reports cap size at 256KiB. Oversized replay drops bulk request and layout arrays first, retaining whole compact source, rank/device, model and -estimator fields that fit; `omitted_fields` identifies every removed field. +estimator fields that fit. Omission names and optional field inspection are +bounded too; `omitted_fields`, `unlisted_fields` and +`omitted_field_names_truncated` disclose that coverage. Such reports remain explicitly incomplete, including when an individual compact field itself exceeds the cap. This does not add a second capture or upload. After an exhausted split ladder, the ordinary refusal may retain an unsplit diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index 55a65b0b4..123d42a9e 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -11,6 +11,7 @@ from dataclasses import dataclass from datetime import datetime, timezone import hashlib +from itertools import islice import json import logging import math @@ -109,6 +110,8 @@ def _compact_planning_record(record: dict[str, Any]) -> dict[str, Any]: compact: dict[str, Any] = { "incomplete_reasons": reasons, "omitted_fields": omitted, + "unlisted_fields": 0, + "omitted_field_names_truncated": 0, } result = { **record, @@ -126,19 +129,31 @@ def _compact_planning_record(record: dict[str, Any]) -> dict[str, Any]: "candidate_matches_check", "memory_replay", ) - for key in dict.fromkeys((*priority, *payload)): + # The maintained snapshot has fewer than 32 fields. Bound optional factory + # fields too: neither omission names nor repeated encoding may grow without + # limit just because the report itself exceeds its cap. + selected = dict.fromkeys((*priority, *islice(payload, 64))) + compact["unlisted_fields"] = len(payload) - sum(key in payload for key in selected) + + def omit(key: str) -> None: + if len(key) > 32: + compact["omitted_field_names_truncated"] += 1 + key = "" + omitted.append(key) + + for key in selected: if key not in payload or key == "incomplete_reasons": continue if key in {"requests", "layouts", "omitted_fields"}: - omitted.append(key) + omit(key) continue compact[key] = payload[key] try: - # Reserve space for explicit omission names as fields are rejected. - _encode(result, limit=MAX_PLANNING_REPORT_BYTES - 4096) + # At most 72 names, each <=32 characters (<=12 JSON bytes/character). + _encode(result, limit=MAX_PLANNING_REPORT_BYTES - 32 * 1024) except _ReportTooLarge: del compact[key] - omitted.append(key) + omit(key) return result @@ -425,7 +440,7 @@ def report( record["replay"] = None record["replay_complete"] = False record["incomplete_reasons"] = [ - f"replay unavailable: {type(exc).__name__}" + f"replay unavailable: {'ValueError' if isinstance(exc, _ReportTooLarge) else type(exc).__name__}" ] raw = _encode(record) path = persist_report( diff --git a/tests/unit/test_trainer_rank_planner_evidence.py b/tests/unit/test_trainer_rank_planner_evidence.py index 31f45448d..c1c2e582a 100644 --- a/tests/unit/test_trainer_rank_planner_evidence.py +++ b/tests/unit/test_trainer_rank_planner_evidence.py @@ -439,6 +439,42 @@ def test_oversized_compact_field_is_omitted_whole(tmp_path): assert record["replay_complete"] is False +@pytest.mark.parametrize( + "payload", + [ + {"k" * 300_000: 1}, + {str(i): "x" * 10_000 for i in range(1000)}, + ], +) +def test_compaction_bounds_field_names_and_total_encoding_work( + tmp_path, monkeypatch, payload +): + original = reports._encode + calls = [] + + def encode(*args, **kwargs): + calls.append(True) + return original(*args, **kwargs) + + monkeypatch.setattr(reports, "_encode", encode) + path = reports.Reporter(5, spool_dir=tmp_path).report( + predicted_peak_bytes=None, + observed_peak_bytes=None, + phase="planning", + event="planning_error", + replay_factory=lambda: payload, + ) + assert path is not None and path.stat().st_size <= reports.MAX_PLANNING_REPORT_BYTES + assert len(calls) <= 80 + record = reports.validate_report(path.read_bytes()) + assert record["replay"]["source_files"] == reports._source_files() + assert record["replay_complete"] is False + assert ( + record["replay"]["unlisted_fields"] + or record["replay"]["omitted_field_names_truncated"] + ) + + def test_summary_trimming_preserves_first_selected(monkeypatch): decision = evidence.Decision("dp_rank_forward", sync_across_dp=False) decision.first = decision.selected = sample(decision) diff --git a/tests/unit/test_trainer_rank_planner_reports.py b/tests/unit/test_trainer_rank_planner_reports.py index 43989c0aa..28cf073f4 100644 --- a/tests/unit/test_trainer_rank_planner_reports.py +++ b/tests/unit/test_trainer_rank_planner_reports.py @@ -131,6 +131,8 @@ def factory(): assert record["replay"] is None assert record["replay_complete"] is False assert record["incomplete_reasons"] + if failure == "overflow": + assert record["incomplete_reasons"] == ["replay unavailable: ValueError"] assert b"private payload" not in path.read_bytes() From 024626e394bdfa5f0bfe86d20405b12d653db63e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 16:01:32 +0000 Subject: [PATCH 3/9] Enforce cumulative assigned planner report retention budgets --- dev/trainer_rank_planner_design.md | 11 ++ src/art/trainer_rank/_planner_misses.py | 31 +++- src/art/trainer_rank/_planner_retention.py | 193 ++++++++++++++++++++ tests/unit/test_planner_retention_budget.py | 170 +++++++++++++++++ 4 files changed, 402 insertions(+), 3 deletions(-) create mode 100644 src/art/trainer_rank/_planner_retention.py create mode 100644 tests/unit/test_planner_retention_budget.py diff --git a/dev/trainer_rank_planner_design.md b/dev/trainer_rank_planner_design.md index 9be81bb35..32cb6067b 100644 --- a/dev/trainer_rank_planner_design.md +++ b/dev/trainer_rank_planner_design.md @@ -909,5 +909,16 @@ whose recorded failure type is `OutOfMemoryError`). Beyond the smaller cap, planning reports are dropped with the existing warning/failure counter. Driver delivery retains its original limits; this is not a cross-process storage quota. +An execution owner can explicitly assign `RetentionLimits` using +`report_retention_scope`. This chooses a private spool and cumulative report/byte +allowance before emission; no scope keeps standalone behavior. An atomic, fsynced +charge ledger precedes each payload write. Exact duplicate bytes cost nothing new, +but deletion, failed payload writes and process restart never replenish the grant. +Changed execution, producer or allowance identity refuses. Exhaustion keeps bounded +omitted-attempt/byte counters, including a zero allowance, without another spool. +Ledger/lock metadata needs a separate owner reservation. This hook neither reserves +capacity across producers nor acknowledges transport: those are the execution +owner's responsibilities. It adds no GPU operation or full-input capture. + Caladan's format-2 reader must land before this producer is enabled, and both rank and driver need the updated ART validator. Old format-1 JSON remains readable. diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index 123d42a9e..12541d2f5 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -24,7 +24,9 @@ from typing import Any import uuid -from . import _planner_evidence +from . import _planner_evidence, _planner_retention +from ._planner_retention import RetentionLimits as RetentionLimits +from ._planner_retention import report_retention_scope as report_retention_scope ALLOW_OVERSIZED_ENV = "ART_TRAINER_RANK_ALLOW_OVERSIZED_BATCHES" MISS_THRESHOLD_ENV = "ART_TRAINER_RANK_PLANNER_MISS_THRESHOLD_PCT" @@ -41,6 +43,7 @@ "_prefix_tree_performance_search.py", "_planner_misses.py", "_planner_evidence.py", + "_planner_retention.py", ) _SOURCE_BYTE_LIMIT = 1024 * 1024 _REPORT_KEYS = frozenset( @@ -273,10 +276,24 @@ def pairs(items: list[tuple[str, Any]]) -> dict[str, Any]: def persist_report( - raw: bytes, spool_dir: Path, *, planning_budget: bool = False + raw: bytes, + spool_dir: Path, + *, + planning_budget: bool = False, + retention: RetentionLimits | None = None, ) -> Path: """Durably retain exact bytes; duplicate delivery is safe, conflicts refuse.""" record = validate_report(raw) + if retention is not None: + if spool_dir != retention.spool_dir: + raise ValueError("planner report spool differs from assigned allowance") + _planner_retention.charge( + retention, + record["id"], + raw, + count_limit=MAX_PLANNING_REPORTS if planning_budget else MAX_SPOOL_REPORTS, + byte_limit=MAX_PLANNING_SPOOL_BYTES if planning_budget else MAX_SPOOL_BYTES, + ) with _spool_lock: spool_dir.mkdir(mode=0o700, parents=True, exist_ok=True) info = spool_dir.lstat() @@ -306,6 +323,12 @@ def persist_report( ) size = count = 0 for entry in spool_dir.iterdir(): + if retention is not None and entry.name in { + ".retention.lock", + ".retention.json", + }: + # Assigned metadata has its separate reserved allowance. + continue try: item = entry.lstat() except FileNotFoundError: @@ -443,11 +466,13 @@ def report( f"replay unavailable: {'ValueError' if isinstance(exc, _ReportTooLarge) else type(exc).__name__}" ] raw = _encode(record) + retention = _planner_retention.current_limits() path = persist_report( raw, - self.spool_dir, + self.spool_dir if retention is None else retention.spool_dir, planning_budget=planning and not (failure is not None and failure["type"] == "OutOfMemoryError"), + retention=retention, ) except Exception as exc: self.failures += 1 diff --git a/src/art/trainer_rank/_planner_retention.py b/src/art/trainer_rank/_planner_retention.py new file mode 100644 index 000000000..ca8eb8953 --- /dev/null +++ b/src/art/trainer_rank/_planner_retention.py @@ -0,0 +1,193 @@ +"""Execution-assigned cumulative report allowances; no transport or quota refund.""" + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +import fcntl +import hashlib +import json +import os +from pathlib import Path +import re +import stat +import tempfile +from typing import Any, Iterator + +_UUID = re.compile(r"[0-9a-f]{32}\Z") +_SHA256 = re.compile(r"[0-9a-f]{64}\Z") +_LEDGER_LIMIT = 512 * 1024 + + +@dataclass(frozen=True) +class RetentionLimits: + spool_dir: Path + max_bytes: int + max_reports: int + allowance_id: str + execution_id: str + producer_id: str + + def __post_init__(self) -> None: + if not isinstance(self.spool_dir, Path) or not self.spool_dir.is_absolute(): + raise ValueError("assigned planner spool must be an absolute path") + if ( + type(self.max_bytes) is not int + or not 0 <= self.max_bytes <= 256 * 1024 * 1024 + or type(self.max_reports) is not int + or not 0 <= self.max_reports <= 1024 + or any( + not isinstance(value, str) or not _UUID.fullmatch(value) + for value in (self.allowance_id, self.execution_id, self.producer_id) + ) + ): + raise ValueError("invalid assigned planner allowance") + + def identity(self) -> dict[str, Any]: + return { + "format": 1, + "allowance_id": self.allowance_id, + "execution_id": self.execution_id, + "producer_id": self.producer_id, + "max_bytes": self.max_bytes, + "max_reports": self.max_reports, + } + + +_current: ContextVar[RetentionLimits | None] = ContextVar( + "planner_retention", default=None +) + + +@contextmanager +def report_retention_scope(limits: RetentionLimits | None) -> Iterator[None]: + if limits is not None and not isinstance(limits, RetentionLimits): + raise ValueError("invalid assigned planner retention scope") + token = _current.set(limits) + try: + yield + finally: + _current.reset(token) + + +def current_limits() -> RetentionLimits | None: + return _current.get() + + +def _encode(value: dict[str, Any]) -> bytes: + raw = ( + json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + "\n" + ).encode() + if len(raw) > _LEDGER_LIMIT: + raise ValueError("planner charge ledger exceeds limit") + return raw + + +def _write(path: Path, value: dict[str, Any]) -> None: + raw = _encode(value) + descriptor, name = tempfile.mkstemp(prefix=".retention-", dir=path.parent) + temporary = Path(name) + try: + with os.fdopen(descriptor, "wb") as target: + target.write(raw) + target.flush() + os.fsync(target.fileno()) + os.replace(temporary, path) + directory = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory) + finally: + os.close(directory) + finally: + temporary.unlink(missing_ok=True) + + +def charge( + limits: RetentionLimits, + event_id: str, + raw: bytes, + *, + count_limit: int, + byte_limit: int, +) -> None: + """Commit a charge before the payload write; ambiguous writes never refund it. + + Payload reclamation does not change this ledger. The caller reserves bounded + ledger/lock metadata separately from cumulative captured payload bytes. + """ + if not _UUID.fullmatch(event_id): + raise ValueError("invalid planner charge identity") + root = limits.spool_dir + root.mkdir(mode=0o700, parents=True, exist_ok=True) + info = root.lstat() + if not stat.S_ISDIR(info.st_mode) or info.st_mode & 0o077: + raise ValueError("assigned planner spool must be a private directory") + lock = os.open( + root / ".retention.lock", os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600 + ) + try: + if not stat.S_ISREG(os.fstat(lock).st_mode): + raise ValueError("planner retention lock is not regular") + fcntl.flock(lock, fcntl.LOCK_EX) + path = root / ".retention.json" + try: + descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK) + except FileNotFoundError: + if any(entry.name != ".retention.lock" for entry in root.iterdir()): + raise ValueError("assigned spool contains unaccounted evidence") + ledger = { + "allowance": limits.identity(), + "charges": {}, + "omitted": 0, + "omitted_bytes": 0, + } + else: + with os.fdopen(descriptor, "rb") as source: + if not stat.S_ISREG(os.fstat(source.fileno()).st_mode): + raise ValueError("planner charge ledger is not regular") + saved = source.read(_LEDGER_LIMIT + 1) + ledger = json.loads(saved) + if ( + _encode(ledger) != saved + or set(ledger) != {"allowance", "charges", "omitted", "omitted_bytes"} + or _encode(ledger["allowance"]) != _encode(limits.identity()) + or not isinstance(ledger["charges"], dict) + or type(ledger["omitted"]) is not int + or not 0 <= ledger["omitted"] <= 2**63 - 1 + or type(ledger["omitted_bytes"]) is not int + or not 0 <= ledger["omitted_bytes"] <= 2**63 - 1 + ): + raise ValueError("planner charge ledger differs from enrollment") + charges = ledger["charges"] + total = 0 + for identifier, item in charges.items(): + if ( + not _UUID.fullmatch(identifier) + or not isinstance(item, list) + or len(item) != 2 + or not isinstance(item[0], str) + or not _SHA256.fullmatch(item[0]) + or type(item[1]) is not int + or not 0 < item[1] <= 16 * 1024 * 1024 + ): + raise ValueError("invalid retained planner charge") + total += item[1] + if len(charges) > limits.max_reports or total > limits.max_bytes: + raise ValueError("planner charge ledger exceeds enrollment") + identity = [hashlib.sha256(raw).hexdigest(), len(raw)] + if event_id in charges: + if charges[event_id] != identity: + raise ValueError("planner report ID has conflicting charged bytes") + return + if len(charges) >= min(limits.max_reports, count_limit) or total + len( + raw + ) > min(limits.max_bytes, byte_limit): + ledger["omitted"] = min(ledger["omitted"] + 1, 2**63 - 1) + ledger["omitted_bytes"] = min(ledger["omitted_bytes"] + len(raw), 2**63 - 1) + _write(path, ledger) + raise ValueError("assigned planner retention exhausted") + charges[event_id] = identity + _write(path, ledger) + finally: + os.close(lock) diff --git a/tests/unit/test_planner_retention_budget.py b/tests/unit/test_planner_retention_budget.py new file mode 100644 index 000000000..3939045d3 --- /dev/null +++ b/tests/unit/test_planner_retention_budget.py @@ -0,0 +1,170 @@ +"""Assigned quotas count captured bytes once, including after source reclamation.""" + +from dataclasses import replace +import json + +import pytest + +from art.trainer_rank import _planner_misses as reports + + +def limits(tmp_path, **values): + return reports.RetentionLimits( + spool_dir=tmp_path / "assigned", + max_bytes=values.pop("max_bytes", 1024 * 1024), + max_reports=values.pop("max_reports", 4), + allowance_id="a" * 32, + execution_id="b" * 32, + producer_id="c" * 32, + **values, + ) + + +def emit(reporter, **values): + return reporter.report( + **{ + "predicted_peak_bytes": 100, + "observed_peak_bytes": 200, + "phase": "forward", + "replay_factory": lambda: {}, + **values, + } + ) + + +def ledger(bound): + return json.loads((bound.spool_dir / ".retention.json").read_bytes()) + + +def test_reclamation_never_refunds_cumulative_budget(tmp_path): + bound = limits(tmp_path, max_reports=1) + reporter = reports.Reporter(5, spool_dir=tmp_path / "standalone") + with reports.report_retention_scope(bound): + path = emit(reporter) + assert path is not None and path.parent == bound.spool_dir + raw = path.read_bytes() + original = ledger(bound) + path.unlink() + assert emit(reporter) is None + assert not reporter.spool_dir.exists() + assert ledger(bound)["charges"] == original["charges"] + assert ledger(bound)["omitted"] == 1 + assert ledger(bound)["omitted_bytes"] > 0 + # Retransmission of the same captured bytes spends nothing new, even if + # recreating a previously reclaimed source copy for the original report. + assert ( + reports.persist_report(raw, bound.spool_dir, retention=bound).read_bytes() + == raw + ) + assert ledger(bound)["charges"] == original["charges"] + + +def test_origin_scope_is_nested_and_exception_safe(tmp_path): + a = limits(tmp_path) + b = replace(a, spool_dir=tmp_path / "other", allowance_id="d" * 32) + reporter = reports.Reporter(5, spool_dir=tmp_path / "standalone") + with reports.report_retention_scope(a): + with pytest.raises(RuntimeError): + with reports.report_retention_scope(b): + assert emit(reporter).parent == b.spool_dir + raise RuntimeError("science") + assert emit(reporter).parent == a.spool_dir + assert emit(reporter).parent == reporter.spool_dir + + +@pytest.mark.parametrize( + "field,value", + [ + ("producer_id", "d" * 32), + ("execution_id", "e" * 32), + ("allowance_id", "f" * 32), + ("max_bytes", 2 * 1024 * 1024), + ("max_reports", 5), + ], +) +def test_enrollment_cannot_be_replaced_or_replenished(tmp_path, field, value): + bound = limits(tmp_path) + reporter = reports.Reporter(5) + with reports.report_retention_scope(bound): + assert emit(reporter) is not None + original = (bound.spool_dir / ".retention.json").read_bytes() + with reports.report_retention_scope(replace(bound, **{field: value})): + assert emit(reporter) is None + assert (bound.spool_dir / ".retention.json").read_bytes() == original + + +def test_ambiguous_payload_write_keeps_charge(tmp_path, monkeypatch): + bound = limits(tmp_path, max_reports=1) + reporter = reports.Reporter(5) + + def fail(*args): + raise OSError("synthetic failure after charge") + + with reports.report_retention_scope(bound): + with monkeypatch.context() as patch: + patch.setattr(reports.os, "link", fail) + assert emit(reporter) is None + assert len(ledger(bound)["charges"]) == 1 + assert emit(reporter) is None + assert not list(bound.spool_dir.glob("[0-9a-f]*.json")) + + assert ledger(bound)["omitted"] == 1 + assert ledger(bound)["omitted_bytes"] > 0 + assert reporter.failures == 2 + + +def test_planning_flood_keeps_oom_headroom_after_reclaim(tmp_path, monkeypatch): + monkeypatch.setattr(reports, "MAX_PLANNING_REPORTS", 1) + bound = limits(tmp_path) + reporter = reports.Reporter(5) + planning = { + "event": "planning_error", + "phase": "planning", + "observed_peak_bytes": None, + } + with reports.report_retention_scope(bound): + first = emit(reporter, **planning) + assert first is not None + first.unlink() + assert emit(reporter, **planning) is None + assert ( + emit(reporter, oom=True, observed_peak_bytes=None, partial_peak_bytes=10) + is not None + ) + assert len(ledger(bound)["charges"]) == 2 + + +@pytest.mark.parametrize( + "bound_values", [{"max_bytes": 0}, {"max_reports": 0}, {"max_bytes": 1}] +) +def test_exhausted_allowance_has_no_fallback(tmp_path, bound_values): + bound = limits(tmp_path, **bound_values) + reporter = reports.Reporter(5, spool_dir=tmp_path / "standalone") + with reports.report_retention_scope(bound): + assert emit(reporter) is None + assert not reporter.spool_dir.exists() + assert not list(bound.spool_dir.glob("[0-9a-f]*.json")) + assert ledger(bound)["omitted"] == 1 + assert ledger(bound)["omitted_bytes"] > 0 + + +def test_unaccounted_spool_and_corrupt_ledger_refuse(tmp_path): + bound = limits(tmp_path) + bound.spool_dir.mkdir(mode=0o700) + foreign = bound.spool_dir / "foreign.json" + foreign.write_text("retained") + reporter = reports.Reporter(5) + with reports.report_retention_scope(bound): + assert emit(reporter) is None + assert foreign.read_text() == "retained" + foreign.unlink() + assert emit(reporter) is not None + (bound.spool_dir / ".retention.json").write_text("partial") + assert emit(reporter) is None + + +def test_disabled_reporting_never_creates_assigned_spool(tmp_path): + bound = limits(tmp_path) + with reports.report_retention_scope(bound): + assert emit(reports.Reporter()) is None + assert not bound.spool_dir.exists() From 87ea8526fe418a3c248b724e0536a7835f2f771d Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 16:12:56 +0000 Subject: [PATCH 4/9] Protect compact planner evidence bookkeeping from payload fields --- src/art/trainer_rank/_planner_misses.py | 8 +++++++- tests/unit/test_trainer_rank_planner_evidence.py | 2 ++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index 123d42a9e..ecfa9cd93 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -144,7 +144,13 @@ def omit(key: str) -> None: for key in selected: if key not in payload or key == "incomplete_reasons": continue - if key in {"requests", "layouts", "omitted_fields"}: + if key in { + "requests", + "layouts", + "omitted_fields", + "unlisted_fields", + "omitted_field_names_truncated", + }: omit(key) continue compact[key] = payload[key] diff --git a/tests/unit/test_trainer_rank_planner_evidence.py b/tests/unit/test_trainer_rank_planner_evidence.py index c1c2e582a..b3f2cdd5f 100644 --- a/tests/unit/test_trainer_rank_planner_evidence.py +++ b/tests/unit/test_trainer_rank_planner_evidence.py @@ -444,6 +444,8 @@ def test_oversized_compact_field_is_omitted_whole(tmp_path): [ {"k" * 300_000: 1}, {str(i): "x" * 10_000 for i in range(1000)}, + {"unlisted_fields": 0, **{str(i): "x" * 10_000 for i in range(1000)}}, + {"omitted_field_names_truncated": -1, "k" * 300_000: 1}, ], ) def test_compaction_bounds_field_names_and_total_encoding_work( From d3f1a6f6577cd7687339ae0b8b325d36bde17dc6 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 16:13:20 +0000 Subject: [PATCH 5/9] Keep assigned spool occupancy within quota after interrupted writes --- src/art/trainer_rank/_planner_misses.py | 15 +++++++++++---- src/art/trainer_rank/_planner_retention.py | 5 ++++- tests/unit/test_planner_retention_budget.py | 21 +++++++++++++++++++++ 3 files changed, 36 insertions(+), 5 deletions(-) diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index 12541d2f5..3e85a8496 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -8,6 +8,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping +from contextlib import nullcontext from dataclasses import dataclass from datetime import datetime, timezone import hashlib @@ -284,9 +285,9 @@ def persist_report( ) -> Path: """Durably retain exact bytes; duplicate delivery is safe, conflicts refuse.""" record = validate_report(raw) - if retention is not None: - if spool_dir != retention.spool_dir: - raise ValueError("planner report spool differs from assigned allowance") + if retention is not None and spool_dir != retention.spool_dir: + raise ValueError("planner report spool differs from assigned allowance") + charged = ( _planner_retention.charge( retention, record["id"], @@ -294,7 +295,10 @@ def persist_report( count_limit=MAX_PLANNING_REPORTS if planning_budget else MAX_SPOOL_REPORTS, byte_limit=MAX_PLANNING_SPOOL_BYTES if planning_budget else MAX_SPOOL_BYTES, ) - with _spool_lock: + if retention is not None + else nullcontext() + ) + with _spool_lock, charged: spool_dir.mkdir(mode=0o700, parents=True, exist_ok=True) info = spool_dir.lstat() if not stat.S_ISDIR(info.st_mode) or info.st_mode & 0o077: @@ -321,6 +325,9 @@ def persist_report( if planning_budget else MAX_SPOOL_BYTES ) + if retention is not None: + count_limit = min(count_limit, retention.max_reports) + byte_limit = min(byte_limit, retention.max_bytes) size = count = 0 for entry in spool_dir.iterdir(): if retention is not None and entry.name in { diff --git a/src/art/trainer_rank/_planner_retention.py b/src/art/trainer_rank/_planner_retention.py index ca8eb8953..775807160 100644 --- a/src/art/trainer_rank/_planner_retention.py +++ b/src/art/trainer_rank/_planner_retention.py @@ -103,6 +103,7 @@ def _write(path: Path, value: dict[str, Any]) -> None: temporary.unlink(missing_ok=True) +@contextmanager def charge( limits: RetentionLimits, event_id: str, @@ -110,7 +111,7 @@ def charge( *, count_limit: int, byte_limit: int, -) -> None: +) -> Iterator[None]: """Commit a charge before the payload write; ambiguous writes never refund it. Payload reclamation does not change this ledger. The caller reserves bounded @@ -179,6 +180,7 @@ def charge( if event_id in charges: if charges[event_id] != identity: raise ValueError("planner report ID has conflicting charged bytes") + yield return if len(charges) >= min(limits.max_reports, count_limit) or total + len( raw @@ -189,5 +191,6 @@ def charge( raise ValueError("assigned planner retention exhausted") charges[event_id] = identity _write(path, ledger) + yield finally: os.close(lock) diff --git a/tests/unit/test_planner_retention_budget.py b/tests/unit/test_planner_retention_budget.py index 3939045d3..f1bd5e84a 100644 --- a/tests/unit/test_planner_retention_budget.py +++ b/tests/unit/test_planner_retention_budget.py @@ -168,3 +168,24 @@ def test_disabled_reporting_never_creates_assigned_spool(tmp_path): with reports.report_retention_scope(bound): assert emit(reports.Reporter()) is None assert not bound.spool_dir.exists() + + +@pytest.mark.parametrize("bounded_by", ["bytes", "reports"]) +def test_interrupted_payload_is_not_duplicated_by_retry(tmp_path, bounded_by): + seed = emit(reports.Reporter(5, spool_dir=tmp_path / "seed")) + assert seed is not None + raw = seed.read_bytes() + bound = limits( + tmp_path, + **({"max_bytes": len(raw)} if bounded_by == "bytes" else {"max_reports": 1}), + ) + committed = reports.persist_report(raw, bound.spool_dir, retention=bound) + charge = ledger(bound) + # Same durable state as death after payload fsync but before final link. + orphan = bound.spool_dir / ".pending-interrupted" + committed.rename(orphan) + with pytest.raises(ValueError, match="spool is full"): + reports.persist_report(raw, bound.spool_dir, retention=bound) + assert orphan.read_bytes() == raw + assert not committed.exists() + assert ledger(bound) == charge From 1149f800379c9c686a5e20502eb73db6840881ff Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 16:57:25 +0000 Subject: [PATCH 6/9] Count failed retention attempts after durable quota charge --- dev/trainer_rank_planner_design.md | 10 ++++- src/art/trainer_rank/_planner_retention.py | 30 +++++++------ tests/unit/test_planner_retention_budget.py | 47 ++++++++++++++++++++- 3 files changed, 71 insertions(+), 16 deletions(-) diff --git a/dev/trainer_rank_planner_design.md b/dev/trainer_rank_planner_design.md index 32cb6067b..e09558fb3 100644 --- a/dev/trainer_rank_planner_design.md +++ b/dev/trainer_rank_planner_design.md @@ -915,7 +915,15 @@ allowance before emission; no scope keeps standalone behavior. An atomic, fsynce charge ledger precedes each payload write. Exact duplicate bytes cost nothing new, but deletion, failed payload writes and process restart never replenish the grant. Changed execution, producer or allowance identity refuses. Exhaustion keeps bounded -omitted-attempt/byte counters, including a zero allowance, without another spool. +omitted-attempt/byte counters, including a zero allowance and ordinary failures +following a charge, without another spool. These count failed persistence attempts, +not unique lost reports or certain payload absence. Counter writes are best effort: +an unreadable ledger, disk failure or process death can leave coverage incomplete. +Unknown crash leftovers are retained and count against occupancy; a first ledger +write interrupted before publication requires owner reconciliation, not automatic +cleanup or a new grant. Enter the context in the actual reporting thread: plain +threads and `run_in_executor` do not inherit it. Spools belong to one producer; +the file lock can block if that contract is violated or storage stalls. Ledger/lock metadata needs a separate owner reservation. This hook neither reserves capacity across producers nor acknowledges transport: those are the execution owner's responsibilities. It adds no GPU operation or full-input capture. diff --git a/src/art/trainer_rank/_planner_retention.py b/src/art/trainer_rank/_planner_retention.py index 775807160..fe80e14f8 100644 --- a/src/art/trainer_rank/_planner_retention.py +++ b/src/art/trainer_rank/_planner_retention.py @@ -2,7 +2,7 @@ from __future__ import annotations -from contextlib import contextmanager +from contextlib import contextmanager, suppress from contextvars import ContextVar from dataclasses import dataclass import fcntl @@ -177,20 +177,24 @@ def charge( if len(charges) > limits.max_reports or total > limits.max_bytes: raise ValueError("planner charge ledger exceeds enrollment") identity = [hashlib.sha256(raw).hexdigest(), len(raw)] - if event_id in charges: - if charges[event_id] != identity: - raise ValueError("planner report ID has conflicting charged bytes") + if event_id in charges and charges[event_id] != identity: + raise ValueError("planner report ID has conflicting charged bytes") + try: + if event_id not in charges: + if len(charges) >= min(limits.max_reports, count_limit) or total + len( + raw + ) > min(limits.max_bytes, byte_limit): + raise ValueError("assigned planner retention exhausted") + charges[event_id] = identity + _write(path, ledger) yield - return - if len(charges) >= min(limits.max_reports, count_limit) or total + len( - raw - ) > min(limits.max_bytes, byte_limit): + except Exception: + # Count failed attempts even after charging (including duplicate + # retries blocked by crash leftovers); never refund uncertain writes. ledger["omitted"] = min(ledger["omitted"] + 1, 2**63 - 1) ledger["omitted_bytes"] = min(ledger["omitted_bytes"] + len(raw), 2**63 - 1) - _write(path, ledger) - raise ValueError("assigned planner retention exhausted") - charges[event_id] = identity - _write(path, ledger) - yield + with suppress(Exception): + _write(path, ledger) + raise finally: os.close(lock) diff --git a/tests/unit/test_planner_retention_budget.py b/tests/unit/test_planner_retention_budget.py index f1bd5e84a..50c396c88 100644 --- a/tests/unit/test_planner_retention_budget.py +++ b/tests/unit/test_planner_retention_budget.py @@ -108,7 +108,7 @@ def fail(*args): assert emit(reporter) is None assert not list(bound.spool_dir.glob("[0-9a-f]*.json")) - assert ledger(bound)["omitted"] == 1 + assert ledger(bound)["omitted"] == 2 assert ledger(bound)["omitted_bytes"] > 0 assert reporter.failures == 2 @@ -188,4 +188,47 @@ def test_interrupted_payload_is_not_duplicated_by_retry(tmp_path, bounded_by): reports.persist_report(raw, bound.spool_dir, retention=bound) assert orphan.read_bytes() == raw assert not committed.exists() - assert ledger(bound) == charge + assert ledger(bound)["charges"] == charge["charges"] + assert ledger(bound)["omitted"] == charge["omitted"] + 1 + assert ledger(bound)["omitted_bytes"] == charge["omitted_bytes"] + len(raw) + + +def test_post_charge_refusal_counts_omission_without_refund(tmp_path): + bound = limits(tmp_path) + reporter = reports.Reporter(5) + with reports.report_retention_scope(bound): + assert emit(reporter) is not None + original = ledger(bound) + orphan = bound.spool_dir / ".pending-orphan" + orphan.write_bytes(b"x" * bound.max_bytes) + assert emit(reporter) is None + checked = ledger(bound) + assert len(checked["charges"]) == len(original["charges"]) + 1 + assert checked["omitted"] == 1 + assert checked["omitted_bytes"] > 0 + assert orphan.stat().st_size == bound.max_bytes + + +def test_omission_write_failure_preserves_original_failure(tmp_path, monkeypatch): + bound = limits(tmp_path) + seed = emit(reports.Reporter(5, spool_dir=tmp_path / "seed")) + assert seed is not None + raw = seed.read_bytes() + original_write = reports._planner_retention._write + calls = [] + + def write(path, value): + calls.append(None) + if len(calls) == 2: + raise OSError("omission ledger unavailable") + original_write(path, value) + + def fail(*args): + raise ValueError("payload failed") + + monkeypatch.setattr(reports._planner_retention, "_write", write) + monkeypatch.setattr(reports.os, "link", fail) + with pytest.raises(ValueError, match="payload failed"): + reports.persist_report(raw, bound.spool_dir, retention=bound) + assert len(calls) == 2 + assert len(ledger(bound)["charges"]) == 1 From 903ed28c64cc3d3029dee5713eb368fb4c50b78d Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 17:46:43 +0000 Subject: [PATCH 7/9] Suppress planner capture explicitly before execution enrollment --- dev/trainer_rank_planner_design.md | 9 ++- src/art/trainer_rank/_planner_misses.py | 2 +- src/art/trainer_rank/_planner_retention.py | 18 +++-- tests/unit/test_planner_retention_budget.py | 75 +++++++++++++++++++++ 4 files changed, 97 insertions(+), 7 deletions(-) diff --git a/dev/trainer_rank_planner_design.md b/dev/trainer_rank_planner_design.md index e09558fb3..533d12c26 100644 --- a/dev/trainer_rank_planner_design.md +++ b/dev/trainer_rank_planner_design.md @@ -911,7 +911,14 @@ delivery retains its original limits; this is not a cross-process storage quota. An execution owner can explicitly assign `RetentionLimits` using `report_retention_scope`. This chooses a private spool and cumulative report/byte -allowance before emission; no scope keeps standalone behavior. An atomic, fsynced +allowance before emission; no scope keeps standalone behavior. An explicit +`capture=False` suppresses report creation before replay/source hashing or I/O, +including during construction before a rank can be bound. It differs from `None` +(standalone retention) and a zero grant (which still accounts for omitted attempts). +Suppression survives nested retention scopes and restores on exit; it follows +ordinary ContextVar propagation, including `asyncio.to_thread`, but not arbitrary +new threads. It does not disable the planner or change admission/training. +An atomic, fsynced charge ledger precedes each payload write. Exact duplicate bytes cost nothing new, but deletion, failed payload writes and process restart never replenish the grant. Changed execution, producer or allowance identity refuses. Exhaustion keeps bounded diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index e93bf6f44..607fbe313 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -401,7 +401,7 @@ def report( ) -> Path | None: """Serialize only a miss; ordinary observation failures never escape.""" threshold = self.threshold_pct - if threshold is None: + if threshold is None or not _planner_retention.capture_enabled(): return None event = event or ("oom" if oom else "estimate_miss") planning = event in {"admission_refused", "planning_error"} diff --git a/src/art/trainer_rank/_planner_retention.py b/src/art/trainer_rank/_planner_retention.py index fe80e14f8..128cdfac8 100644 --- a/src/art/trainer_rank/_planner_retention.py +++ b/src/art/trainer_rank/_planner_retention.py @@ -55,16 +55,20 @@ def identity(self) -> dict[str, Any]: } -_current: ContextVar[RetentionLimits | None] = ContextVar( - "planner_retention", default=None +_current: ContextVar[tuple[RetentionLimits | None, bool]] = ContextVar( + "planner_retention", default=(None, True) ) @contextmanager -def report_retention_scope(limits: RetentionLimits | None) -> Iterator[None]: +def report_retention_scope( + limits: RetentionLimits | None, *, capture: bool = True +) -> Iterator[None]: + if type(capture) is not bool: + raise ValueError("planner capture must be a boolean") if limits is not None and not isinstance(limits, RetentionLimits): raise ValueError("invalid assigned planner retention scope") - token = _current.set(limits) + token = _current.set((limits, capture and capture_enabled())) try: yield finally: @@ -72,7 +76,11 @@ def report_retention_scope(limits: RetentionLimits | None) -> Iterator[None]: def current_limits() -> RetentionLimits | None: - return _current.get() + return _current.get()[0] + + +def capture_enabled() -> bool: + return _current.get()[1] def _encode(value: dict[str, Any]) -> bytes: diff --git a/tests/unit/test_planner_retention_budget.py b/tests/unit/test_planner_retention_budget.py index 50c396c88..032310ba7 100644 --- a/tests/unit/test_planner_retention_budget.py +++ b/tests/unit/test_planner_retention_budget.py @@ -1,5 +1,6 @@ """Assigned quotas count captured bytes once, including after source reclamation.""" +import asyncio from dataclasses import replace import json @@ -232,3 +233,77 @@ def fail(*args): reports.persist_report(raw, bound.spool_dir, retention=bound) assert len(calls) == 2 assert len(ledger(bound)["charges"]) == 1 + + +@pytest.mark.parametrize("assigned", [False, True]) +def test_capture_off_suppresses_constructor_before_replay_or_io( + tmp_path, monkeypatch, assigned +): + bound = limits(tmp_path, max_bytes=0, max_reports=0) if assigned else None + replayed = [] + sourced = [] + monkeypatch.setattr(reports, "_source_files", lambda: sourced.append(True)) + + class Constructed: + def __init__(self): + self.reporter = reports.Reporter(5, spool_dir=tmp_path / "standalone") + self.path = emit( + self.reporter, replay_factory=lambda: replayed.append(True) + ) + + with reports.report_retention_scope(bound, capture=False): + constructed = Constructed() + assert constructed.path is None + assert constructed.reporter.failures == 0 + assert replayed == sourced == [] + assert list(tmp_path.iterdir()) == [] + + +def test_capture_suppression_nests_and_restores_after_exception(tmp_path): + reporter = reports.Reporter(5, spool_dir=tmp_path / "standalone") + bound = limits(tmp_path) + with pytest.raises(RuntimeError, match="science"): + with reports.report_retention_scope(None, capture=False): + with reports.report_retention_scope(bound): + assert emit(reporter) is None + raise RuntimeError("science") + assert not bound.spool_dir.exists() + with reports.report_retention_scope(bound): + with reports.report_retention_scope(None, capture=False): + assert emit(reporter) is None + assert emit(reporter).parent == bound.spool_dir + assert emit(reporter).parent == reporter.spool_dir + + +def test_capture_off_differs_from_zero_allowance(tmp_path): + bound = limits(tmp_path, max_bytes=0, max_reports=0) + reporter = reports.Reporter(5) + with reports.report_retention_scope(bound, capture=False): + assert emit(reporter) is None + assert not bound.spool_dir.exists() + with reports.report_retention_scope(bound): + assert emit(reporter) is None + assert ledger(bound)["omitted"] == 1 + + +def test_capture_scope_inherits_into_to_thread(tmp_path): + bound = limits(tmp_path) + reporter = reports.Reporter(5, spool_dir=tmp_path / "standalone") + + async def run(): + with reports.report_retention_scope(bound, capture=False): + assert await asyncio.to_thread(emit, reporter) is None + assert not bound.spool_dir.exists() + with reports.report_retention_scope(bound): + assert (await asyncio.to_thread(emit, reporter)).parent == bound.spool_dir + + asyncio.run(run()) + assert not reporter.spool_dir.exists() + + +@pytest.mark.parametrize("capture", [None, 0, "false"]) +def test_capture_flag_refuses_non_boolean_without_changing_scope(tmp_path, capture): + with pytest.raises(ValueError, match="capture must be a boolean"): + with reports.report_retention_scope(None, capture=capture): + pytest.fail("invalid capture flag entered") + assert emit(reports.Reporter(5, spool_dir=tmp_path / "standalone")) is not None From f8f551cbe658ed2717706adf2e49c09377080895 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 17:47:21 +0000 Subject: [PATCH 8/9] Keep capture flag refusal control independent of pytest overload typing --- tests/unit/test_planner_retention_budget.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_planner_retention_budget.py b/tests/unit/test_planner_retention_budget.py index 032310ba7..c5a3ae6e3 100644 --- a/tests/unit/test_planner_retention_budget.py +++ b/tests/unit/test_planner_retention_budget.py @@ -305,5 +305,5 @@ async def run(): def test_capture_flag_refuses_non_boolean_without_changing_scope(tmp_path, capture): with pytest.raises(ValueError, match="capture must be a boolean"): with reports.report_retention_scope(None, capture=capture): - pytest.fail("invalid capture flag entered") + raise AssertionError("invalid capture flag entered") assert emit(reports.Reporter(5, spool_dir=tmp_path / "standalone")) is not None From a2e4b7aed4637e2b32b4b50915656e393c5ef5b2 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 23 Sep 2026 18:11:44 +0000 Subject: [PATCH 9/9] Reestablish planner report durability on duplicate retries --- src/art/trainer_rank/_planner_misses.py | 24 +++++++--- tests/unit/test_planner_retention_budget.py | 52 +++++++++++++++++++++ 2 files changed, 69 insertions(+), 7 deletions(-) diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index 607fbe313..9ec5feefb 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -311,13 +311,23 @@ def persist_report( raise ValueError("report spool must be a private directory") path = spool_dir / f"{record['id']}.json" if path.exists() or path.is_symlink(): - info = path.lstat() - if ( - not stat.S_ISREG(info.st_mode) - or info.st_size != len(raw) - or path.read_bytes() != raw - ): - raise ValueError("existing report identity has different bytes") + descriptor = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK) + with os.fdopen(descriptor, "rb") as existing: + info = os.fstat(existing.fileno()) + if ( + not stat.S_ISREG(info.st_mode) + or info.st_size != len(raw) + or existing.read(len(raw) + 1) != raw + ): + raise ValueError("existing report identity has different bytes") + # A prior attempt may have linked the file but failed its + # durability barrier. Visibility alone cannot acknowledge it. + os.fsync(existing.fileno()) + directory = os.open(spool_dir, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory) + finally: + os.close(directory) return path # Rank-side ordinary planning failures may use only the first small # part of the spool. OOMs/misses and delivery keep their original limit. diff --git a/tests/unit/test_planner_retention_budget.py b/tests/unit/test_planner_retention_budget.py index c5a3ae6e3..fdfe69c1d 100644 --- a/tests/unit/test_planner_retention_budget.py +++ b/tests/unit/test_planner_retention_budget.py @@ -3,6 +3,7 @@ import asyncio from dataclasses import replace import json +import stat import pytest @@ -235,6 +236,57 @@ def fail(*args): assert len(ledger(bound)["charges"]) == 1 +@pytest.mark.parametrize("assigned", [False, True]) +@pytest.mark.parametrize("failed_barrier", ["file", "directory"]) +def test_duplicate_retry_reestablishes_durability( + tmp_path, monkeypatch, assigned, failed_barrier +): + seed = emit(reports.Reporter(5, spool_dir=tmp_path / "seed")) + assert seed is not None + raw = seed.read_bytes() + bound = limits(tmp_path) + path = bound.spool_dir / seed.name + retention = bound if assigned else None + fsync = reports.os.fsync + + def fail_after_link(fd): + if path.exists() and stat.S_ISDIR(reports.os.fstat(fd).st_mode): + raise OSError("directory barrier unavailable") + fsync(fd) + + monkeypatch.setattr(reports.os, "fsync", fail_after_link) + with pytest.raises(OSError, match="directory barrier unavailable"): + reports.persist_report(raw, bound.spool_dir, retention=retention) + assert path.read_bytes() == raw # Visible is not proof of durable custody. + charges = ledger(bound)["charges"] if assigned else None + barriers = [] + + def retry_barrier(fd): + info = reports.os.fstat(fd) + if stat.S_ISDIR(info.st_mode): + barrier = "directory" + elif info.st_ino == path.stat().st_ino: + barrier = "file" + else: + return fsync(fd) + barriers.append(barrier) + if barrier == failed_barrier: + raise OSError("retry barrier unavailable") + fsync(fd) + + monkeypatch.setattr(reports.os, "fsync", retry_barrier) + with pytest.raises(OSError, match="retry barrier unavailable"): + reports.persist_report(raw, bound.spool_dir, retention=retention) + assert failed_barrier in barriers + failed_barrier = None + barriers.clear() + assert reports.persist_report(raw, bound.spool_dir, retention=retention) == path + assert barriers == ["file", "directory"] + assert path.read_bytes() == raw + if assigned: + assert ledger(bound)["charges"] == charges # Never charge a retry twice. + + @pytest.mark.parametrize("assigned", [False, True]) def test_capture_off_suppresses_constructor_before_replay_or_io( tmp_path, monkeypatch, assigned