Enforce cumulative execution allowances for planner reports

dev/trainer_rank_planner_design.md

@@ -909,5 +909,31 @@ 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 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
+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.
+
 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.

src/art/trainer_rank/_planner_misses.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Mapping
+from contextlib import nullcontext
 from dataclasses import asdict, dataclass
 from datetime import datetime, timezone
 import hashlib
@@ -30,7 +31,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"
@@ -48,6 +51,7 @@
     "_planner_misses.py",
     "_gdn_memory.py",
     "_planner_evidence.py",
+    "_planner_retention.py",
 )
 _SOURCE_BYTE_LIMIT = 1024 * 1024
 _REPORT_KEYS = frozenset(
@@ -286,24 +290,51 @@ 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)
-    with _spool_lock:
+    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"],
+            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,
+        )
+        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:
             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.
@@ -317,8 +348,17 @@ 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 {
+                ".retention.lock",
+                ".retention.json",
+            }:
+                # Assigned metadata has its separate reserved allowance.
+                continue
             try:
                 item = entry.lstat()
             except FileNotFoundError:
@@ -378,7 +418,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"}
@@ -456,11 +496,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

src/art/trainer_rank/_planner_retention.py

@@ -0,0 +1,208 @@
+"""Execution-assigned cumulative report allowances; no transport or quota refund."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager, suppress
+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[tuple[RetentionLimits | None, bool]] = ContextVar(
+    "planner_retention", default=(None, True)
+)
+
+
+@contextmanager
+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, capture and capture_enabled()))
+    try:
+        yield
+    finally:
+        _current.reset(token)
+
+
+def current_limits() -> RetentionLimits | None:
+    return _current.get()[0]
+
+
+def capture_enabled() -> bool:
+    return _current.get()[1]
+
+
+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)
+
+
+@contextmanager
+def charge(
+    limits: RetentionLimits,
+    event_id: str,
+    raw: bytes,
+    *,
+    count_limit: int,
+    byte_limit: int,
+) -> 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
+    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 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
+        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)
+            with suppress(Exception):
+                _write(path, ledger)
+            raise
+    finally:
+        os.close(lock)