From a921f108194e3cabfd2d10adb9acc9f33f439ae3 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Fri, 2 Oct 2026 06:21:07 -0700 Subject: [PATCH 1/6] feat(token-capture): retire and delete external-staging capture ledgers With external staging the model server keeps a capture ledger per rollout and never removes it, because only the training framework knows when a rollout's manifest is no longer needed. Give the ledger the two operations the complete-record token store already has. CaptureLedger.retire removes a rollout's ledger and leaves an empty .lineage.retired fence, like the fence TokenSource.drop leaves: the framework retires every rollout it is done with, finished or abandoned, and later calls for it are discarded instead of starting a new ledger. That covers an abandoned agent that keeps running and a client retry that ran a finished rollout twice. CaptureLedger.delete removes the ledger and the fence, like TokenCaptureStore.delete: the framework deletes retired rollouts once nothing of them can still write, and deletes a rollout ID before reusing it. Both are batched, validate the whole batch first, reject a bare string, and are safe to retry, behind bearer-protected POST .../rollouts/retire and .../rollouts/delete routes with matching RolloutControlClient methods. Writes look for the fence only when they would create the ledger, so a captured call costs no extra metadata lookup. Staged token data is untouched; the framework that staged it removes it. Signed-off-by: Ananth Subramaniam --- .../pages/model-server/model-call-capture.mdx | 44 ++++++ nemo_gym/base_responses_api_model.py | 2 +- nemo_gym/token_id_capture/control_routes.py | 67 ++++++++- nemo_gym/token_id_capture/lineage.py | 116 +++++++++++++- nemo_gym/token_id_capture/protocols.py | 31 ++++ nemo_gym/token_id_capture/staging/records.py | 9 ++ nemo_gym/token_id_capture/store.py | 11 ++ .../test_token_capture_control_routes.py | 141 ++++++++++++++++++ tests/unit_tests/test_token_capture_ledger.py | 136 ++++++++++++++++- 9 files changed, 546 insertions(+), 11 deletions(-) create mode 100644 tests/unit_tests/test_token_capture_control_routes.py diff --git a/fern/versions/latest/pages/model-server/model-call-capture.mdx b/fern/versions/latest/pages/model-server/model-call-capture.mdx index 399a052578f..0017672678d 100644 --- a/fern/versions/latest/pages/model-server/model-call-capture.mdx +++ b/fern/versions/latest/pages/model-server/model-call-capture.mdx @@ -293,6 +293,50 @@ digest-covered `media` summary. Deploy the paired RL changes with this adapter; checkpoints written in either older format need a separate migration. Supported processors and modalities are enforced by the framework's inference worker. +## Retire capture ledgers + +With external staging, the model server keeps a capture ledger for each rollout: one `.lineage.jsonl` file in the lineage store's directory, recording the model calls the rollout made without their token arrays. +The training framework reads it once, through the manifest route, when the rollout ends. +Gym can't tell when the framework is done with a rollout, so the framework removes each ledger itself, using the same two operations the complete-record token store has: retire and delete. + +### Retire a rollout the framework is done with + +Retiring deletes the rollout's ledger and leaves an empty `.lineage.retired` file in its place. +That file acts as a fence: if a model call for the rollout finishes later, Gym discards its row instead of starting a new ledger. +Late calls are more common than they sound, because an agent the framework gave up on can keep running, and a client retry can run a finished rollout a second time. +`TokenSource.drop` leaves the same kind of fence in the complete-record token store. + +Retire every rollout the framework is done with, whether it finished or was abandoned. +A finished rollout can be retired once its manifest has been read and its receipt built, unless a crash before the receipt is saved would make the framework rerun it under the same ID; in that case, wait until the receipt is saved. +An abandoned rollout should be retired as soon as the framework gives up on it. + +### Delete fences once nothing can still write + +Deleting removes a rollout's ledger together with its fence, which is only safe once no call from that rollout can still arrive. +In practice that means the end of the run, or a point where every Gym server sharing the lineage directory has restarted. +Until then, the fences accumulate at one empty file per retired rollout. +Also delete a rollout ID before reusing it, so that the new run starts from an empty ledger, as `TokenCaptureStore.delete` does for the complete-record store. + +### Send the requests + +Send each request to a model server whose lineage directory holds the rollout's ledger, either the server that served the rollout or any other server configured with the same directory. +A model server only sees the ledgers in its own lineage directory and reports any other rollout as `absent`. + +```http +POST /training-token-capture/control/rollouts/retire +POST /training-token-capture/control/rollouts/delete +Authorization: Bearer + +{"rollout_ids": ["", ...]} +``` + +Each request takes at most 4096 rollout IDs, and `RolloutControlClient.retire` and `RolloutControlClient.delete` split longer lists into batches for you. +The reply lists the rollouts whose ledger was removed under `removed` and the rest under `absent`. +Because repeating either operation changes nothing, a failed batch can simply be sent again. + +Neither operation touches staged token data. +The framework that staged the data removes it, and it must keep any staged data that a checkpoint it may restore still refers to. + ## Limitations The model-server boundary observes model HTTP requests and responses. It can record tool calls, diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 3b891a5b1cf..ad6b596af77 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -1883,7 +1883,7 @@ def install_model_call_capture( if not isinstance(lineage_store, CaptureLedger): raise ValueError( "token_id_capture.external_staging requires the lineage store to implement " - "the CaptureLedger protocol (record, record_failure, manifest, has_rows)" + "the CaptureLedger protocol (record, record_failure, manifest, has_rows, retire, delete)" ) capture_ledger: CaptureLedger = lineage_store install_rollout_control_routes( diff --git a/nemo_gym/token_id_capture/control_routes.py b/nemo_gym/token_id_capture/control_routes.py index eb9cd8fc05c..356240be28f 100644 --- a/nemo_gym/token_id_capture/control_routes.py +++ b/nemo_gym/token_id_capture/control_routes.py @@ -1,26 +1,40 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Expose rollout capture manifests over HTTP. +"""Expose rollout capture manifests and ledger removal over HTTP. The training framework reads one manifest for each rollout. The endpoint returns call metadata and does not read staged token data. The framework builds receipts and removes staged data after use. +It then retires the rollout's ledger, which Gym cannot otherwise know is no longer needed. """ from __future__ import annotations import asyncio import hmac +from collections.abc import Sequence from typing import Any from fastapi import APIRouter, Header, HTTPException +from pydantic import BaseModel, ConfigDict, Field from nemo_gym.token_id_capture.protocols import CaptureLedger -from nemo_gym.token_id_capture.staging.records import RolloutManifest +from nemo_gym.token_id_capture.staging.records import RolloutManifest, RolloutRemoval CONTROL_ROUTE_PREFIX = "/training-token-capture/control" +# Bounds one retire or delete request so it finishes well within a control timeout. +# The client splits larger batches. +MAX_LEDGER_BATCH = 4096 + + +class LedgerBatchRequest(BaseModel): + """Body of the retire and delete routes.""" + + model_config = ConfigDict(extra="forbid") + + rollout_ids: list[str] = Field(min_length=1, max_length=MAX_LEDGER_BATCH) def install_rollout_control_routes( @@ -29,7 +43,7 @@ def install_rollout_control_routes( *, auth_token: str, ) -> None: - """Install the bearer-protected manifest route.""" + """Install the bearer-protected manifest, retire, and delete routes.""" if not auth_token: raise ValueError("rollout control routes require a non-empty auth token") expected = f"Bearer {auth_token}" @@ -53,11 +67,33 @@ async def rollout_manifest( except ValueError as error: raise HTTPException(status_code=400, detail=str(error)) from error + @router.post("/rollouts/retire") + async def retire_rollouts( + body: LedgerBatchRequest, + authorization: str | None = Header(default=None), + ) -> dict: + check_auth(authorization) + try: + return await lineage_store.retire(body.rollout_ids) + except ValueError as error: + raise HTTPException(status_code=400, detail=str(error)) from error + + @router.post("/rollouts/delete") + async def delete_rollouts( + body: LedgerBatchRequest, + authorization: str | None = Header(default=None), + ) -> dict: + check_auth(authorization) + try: + return await lineage_store.delete(body.rollout_ids) + except ValueError as error: + raise HTTPException(status_code=400, detail=str(error)) from error + app.include_router(router) class RolloutControlClient: - """Framework-side client for the read-only manifest route.""" + """Framework-side client for the manifest, retire, and delete routes.""" def __init__( self, @@ -80,6 +116,29 @@ async def manifest(self, rollout_id: str) -> RolloutManifest: ) return RolloutManifest.model_validate(await response.json()) + async def retire(self, rollout_ids: Sequence[str]) -> RolloutRemoval: + """Retire ledgers; see ``CaptureLedger.retire``. Retrying after a failure is safe.""" + return await self._remove("retire", rollout_ids) + + async def delete(self, rollout_ids: Sequence[str]) -> RolloutRemoval: + """Delete ledgers and their fences; see ``CaptureLedger.delete``. Retrying after a failure is safe.""" + return await self._remove("delete", rollout_ids) + + async def _remove(self, action: str, rollout_ids: Sequence[str]) -> RolloutRemoval: + rollout_ids = list(rollout_ids) + result = RolloutRemoval() + for start in range(0, len(rollout_ids), MAX_LEDGER_BATCH): + batch = rollout_ids[start : start + MAX_LEDGER_BATCH] + response = await self._request("POST", f"/rollouts/{action}", json={"rollout_ids": batch}) + if response.status != 200: + raise RuntimeError( + f"{action} of {len(batch)} rollout ledgers failed: HTTP {response.status} {await response.text()}" + ) + removal = RolloutRemoval.model_validate(await response.json()) + result.removed.extend(removal.removed) + result.absent.extend(removal.absent) + return result + async def _request(self, method: str, path: str, **kwargs: Any): # Deferred because server_utils loads Gym's aiohttp/server stack. from nemo_gym.server_utils import request diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py index d1e12283f4d..255206c5d90 100644 --- a/nemo_gym/token_id_capture/lineage.py +++ b/nemo_gym/token_id_capture/lineage.py @@ -43,7 +43,9 @@ import asyncio import json +import logging import os +from collections.abc import Sequence from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any @@ -66,6 +68,19 @@ from nemo_gym.token_id_capture.staging.records import CallRecord, CaptureLedgerCommit +logger = logging.getLogger(__name__) + + +def _log_late_write_dropped(rollout_id: str, model_call_id: str) -> None: + # A call that finished after its rollout was retired, such as one from an abandoned attempt or a request + # the client had already retried. Nobody will read its row, so this is not an error. + logger.warning( + "Discarding the ledger row of model call %s: rollout %s is retired.", + model_call_id, + rollout_id, + ) + + # Names of the token-free CallRecord custody columns a ledger row carries. # ``staging_digest`` is the worker's staged-record digest; the row's ``digest`` # key remains the lineage digest over the cumulative token IDs. @@ -444,6 +459,7 @@ class InMemoryLineageStore: def __init__(self, max_rollouts: int = 512, max_tokens: int = 8_000_000) -> None: self.index = LineageIndex(max_rollouts=max_rollouts, max_tokens=max_tokens) self._ledgers: dict[str, list[dict]] = {} + self._retired: set[str] = set() async def resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: return self.index.for_rollout(rollout_id).resolve(request_items) @@ -460,6 +476,9 @@ async def record(self, commit: CaptureLedgerCommit) -> None: # hash covers continuity, so the index keeps tokens only for # lineage-only local-capture rows that inject prompt prefixes. record = commit.record + if commit.rollout_id in self._retired: + _log_late_write_dropped(commit.rollout_id, record.model_call_id) + return rows = self._ledgers.setdefault(commit.rollout_id, []) row = {"model_call_id": record.model_call_id, **_custody_columns(record, commit.staging_chain)} # Write-once per model call (CaptureLedger contract): an identical @@ -484,6 +503,9 @@ async def record(self, commit: CaptureLedgerCommit) -> None: rows.append(row) async def record_failure(self, rollout_id: str, model_call_id: str, reason: str) -> None: + if rollout_id in self._retired: + _log_late_write_dropped(rollout_id, model_call_id) + return rows = self._ledgers.setdefault(rollout_id, []) row = {"model_call_id": model_call_id, "failure_reason": reason} if not any(existing == row for existing in rows): @@ -497,9 +519,31 @@ async def has_rows(self, rollout_id: str) -> bool: return True return bool(self.index.for_rollout(rollout_id).by_call_id) + async def retire(self, rollout_ids: Sequence[str]) -> dict: + from nemo_gym.token_id_capture.store import validate_rollout_ids + + rollout_ids = validate_rollout_ids(rollout_ids) + self._retired.update(rollout_ids) + return self._remove(rollout_ids) + + async def delete(self, rollout_ids: Sequence[str]) -> dict: + from nemo_gym.token_id_capture.store import validate_rollout_ids + + rollout_ids = validate_rollout_ids(rollout_ids) + self._retired.difference_update(rollout_ids) + return self._remove(rollout_ids) + + def _remove(self, rollout_ids: list[str]) -> dict: + removed, absent = [], [] + for rollout_id in rollout_ids: + (removed if self._ledgers.pop(rollout_id, None) is not None else absent).append(rollout_id) + self.index.drop(rollout_id) + return {"removed": removed, "absent": absent} + async def close(self) -> None: self.index.clear() self._ledgers.clear() + self._retired.clear() class IncrementalLineageStore: @@ -793,6 +837,19 @@ def _ledger_path(self, rollout_id: str) -> Path: return self._ledger_root / f"{validate_rollout_id(rollout_id)}.lineage.jsonl" + def _retired_path(self, rollout_id: str) -> Path: + # Present while the rollout is retired: its writes are discarded. ``delete`` removes it. + from nemo_gym.token_id_capture.store import validate_rollout_id + + return self._ledger_root / f"{validate_rollout_id(rollout_id)}.lineage.retired" + + def _is_retired(self, rollout_id: str) -> bool: + # Call with the rollout's lock held, only when its ledger is empty or missing. ``retire`` writes every + # marker before it removes any ledger, so a rollout whose ledger still has rows was not retired when + # the ledger was read, or its retire failed partway and will be retried. Skipping the check + # otherwise saves a metadata lookup on every captured call. + return self._retired_path(rollout_id).exists() + def _locked(self, rollout_id: str): # Ledger rows share the token store's per-rollout lock file. The two # record families are never both written for one rollout, so one lock @@ -826,6 +883,13 @@ def _read(self, rollout_id: str) -> list[dict]: self._ledger_cache[rollout_id] = (inode, offset, records) return records + def _fsync_ledger_root(self) -> None: + directory_fd = os.open(self._ledger_root, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + def _append(self, rollout_id: str, record: dict, records: list[dict]) -> None: path = self._ledger_path(rollout_id) created = not path.exists() @@ -837,11 +901,7 @@ def _append(self, rollout_id: str, record: dict, records: list[dict]) -> None: offset = handle.tell() inode = os.fstat(handle.fileno()).st_ino if created: - directory_fd = os.open(self._ledger_root, os.O_RDONLY) - try: - os.fsync(directory_fd) - finally: - os.close(directory_fd) + self._fsync_ledger_root() records.append(record) self._ledger_cache[rollout_id] = (inode, offset, records) @@ -903,6 +963,9 @@ def _record(self, commit: CaptureLedgerCommit) -> None: } with self._locked(commit.rollout_id): records = self._read(commit.rollout_id) + if not records and self._is_retired(commit.rollout_id): + _log_late_write_dropped(commit.rollout_id, model_call_id) + return matches = [existing for existing in records if existing["model_call_id"] == model_call_id] if matches: if matches[0] != record: @@ -919,6 +982,9 @@ def _record_failure(self, rollout_id: str, model_call_id: str, reason: str) -> N record = {"model_call_id": model_call_id, "failure_reason": reason} with self._locked(rollout_id): records = self._read(rollout_id) + if not records and self._is_retired(rollout_id): + _log_late_write_dropped(rollout_id, model_call_id) + return if any(existing == record for existing in records): return self._append(rollout_id, record, records) @@ -937,3 +1003,43 @@ async def has_rows(self, rollout_id: str) -> bool: def _has_rows(self, rollout_id: str) -> bool: with self._locked(rollout_id): return bool(self._read(rollout_id)) + + async def retire(self, rollout_ids: Sequence[str]) -> dict: + return await asyncio.to_thread(self._retire, rollout_ids) + + def _retire(self, rollout_ids: Sequence[str]) -> dict: + from nemo_gym.token_id_capture.store import validate_rollout_ids + + rollout_ids = validate_rollout_ids(rollout_ids) + # Make every retired marker durable before deleting any ledger, so a crash in between leaves a fence. + for rollout_id in rollout_ids: + with self._locked(rollout_id): + self._retired_path(rollout_id).touch() + if rollout_ids: + self._fsync_ledger_root() + return self._remove(rollout_ids) + + async def delete(self, rollout_ids: Sequence[str]) -> dict: + return await asyncio.to_thread(self._delete, rollout_ids) + + def _delete(self, rollout_ids: Sequence[str]) -> dict: + from nemo_gym.token_id_capture.store import validate_rollout_ids + + return self._remove(validate_rollout_ids(rollout_ids), unretire=True) + + def _remove(self, rollout_ids: list[str], *, unretire: bool = False) -> dict: + removed, absent = [], [] + for rollout_id in rollout_ids: + with self._locked(rollout_id): + try: + self._ledger_path(rollout_id).unlink() + removed.append(rollout_id) + except FileNotFoundError: + absent.append(rollout_id) + if unretire: + self._retired_path(rollout_id).unlink(missing_ok=True) + with self._cache_guard: + self._ledger_cache.pop(rollout_id, None) + if rollout_ids: + self._fsync_ledger_root() + return {"removed": removed, "absent": absent} diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py index 7252c22f4a8..ed223862544 100644 --- a/nemo_gym/token_id_capture/protocols.py +++ b/nemo_gym/token_id_capture/protocols.py @@ -41,6 +41,7 @@ from __future__ import annotations +from collections.abc import Sequence from dataclasses import dataclass from typing import Protocol, runtime_checkable @@ -146,6 +147,7 @@ class CaptureLedger(LineageResolver, Protocol): ``record`` publishes a successfully staged call for parent resolution. ``record_failure`` records a call that did not commit. ``manifest`` returns both kinds of rows without including token arrays. + ``retire`` and ``delete`` remove the rows of rollouts the framework no longer needs. Records must be visible to every serving worker before ``record`` returns. """ @@ -180,6 +182,35 @@ async def has_rows(self, rollout_id: str) -> bool: """Return whether any ledger row (committed or failed) exists.""" ... + async def retire(self, rollout_ids: Sequence[str]) -> dict: + """Remove the ledgers of rollouts the framework is done with, and fence them. + + Retire a finished rollout once its manifest has been read and the receipt built from it is + durable, and an abandoned rollout as soon as the framework gives up on it. Like + ``TokenSource.drop`` on the complete-record store, retiring keeps a fence: later ``record`` and + ``record_failure`` calls for the rollout are discarded instead of starting a new ledger. Even a + finished rollout can still write: a client retry can run it twice, and the second execution + writes after the first one returned. + Retiring again is a no-op, so retrying a batch is safe. + Staged token data is untouched: the framework that staged it also removes it. + + The result validates as ``staging.records.RolloutRemoval``. + """ + ... + + async def delete(self, rollout_ids: Sequence[str]) -> dict: + """Remove the ledgers and fences of rollouts, as ``TokenCaptureStore.delete`` does. + + Delete retired rollouts once nothing of those attempts can still write, for example after every + Gym server has restarted or at the end of the run, which removes their fences. Delete a rollout + ID before reusing it, so the new execution starts from an empty ledger. Deleting again is a + no-op. + Staged token data is untouched: the framework that staged it also removes it. + + The result validates as ``staging.records.RolloutRemoval``. + """ + ... + @runtime_checkable class TokenSink(Protocol): diff --git a/nemo_gym/token_id_capture/staging/records.py b/nemo_gym/token_id_capture/staging/records.py index a4b27018ebf..18cf579286a 100644 --- a/nemo_gym/token_id_capture/staging/records.py +++ b/nemo_gym/token_id_capture/staging/records.py @@ -356,6 +356,15 @@ class RolloutManifest(_WireModel): failures: list[ManifestFailure] = Field(default_factory=list) +class RolloutRemoval(_WireModel): + """The result of retiring or deleting rollouts in a capture store or ledger.""" + + # Rollouts whose captured data existed and was removed. + removed: list[Identifier] = Field(default_factory=list) + # Rollouts with no captured data: already removed, or never recorded a call. + absent: list[Identifier] = Field(default_factory=list) + + class RolloutReceipt(_DigestWireModel): """Token-free immutable manifest the framework assembles at rollout end.""" diff --git a/nemo_gym/token_id_capture/store.py b/nemo_gym/token_id_capture/store.py index f80b5bdaa86..06275fdf6a8 100644 --- a/nemo_gym/token_id_capture/store.py +++ b/nemo_gym/token_id_capture/store.py @@ -33,6 +33,7 @@ import os import tempfile import time +from collections.abc import Sequence from contextlib import contextmanager from pathlib import Path from typing import Any @@ -54,6 +55,16 @@ def validate_rollout_id(rollout_id: str) -> str: return rollout_id +def validate_rollout_ids(rollout_ids: Sequence[str]) -> list[str]: + """Validate a whole batch before anything changes, and drop duplicates. + + A bare string is rejected: it is a sequence of one-character rollout IDs. + """ + if isinstance(rollout_ids, str): + raise TypeError(f"expected a sequence of rollout ids, got the string {rollout_ids!r}") + return list(dict.fromkeys(validate_rollout_id(rollout_id) for rollout_id in rollout_ids)) + + class TokenCaptureStore: """Durable, rollout-keyed JSONL sink for ``TokenEntry`` records.""" diff --git a/tests/unit_tests/test_token_capture_control_routes.py b/tests/unit_tests/test_token_capture_control_routes.py new file mode 100644 index 00000000000..8edd6ab0712 --- /dev/null +++ b/tests/unit_tests/test_token_capture_control_routes.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The bearer-protected control routes a training framework calls on a capture ledger.""" + +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from nemo_gym.token_id_capture import control_routes +from nemo_gym.token_id_capture.control_routes import ( + CONTROL_ROUTE_PREFIX, + RolloutControlClient, + install_rollout_control_routes, +) +from nemo_gym.token_id_capture.lineage import FileLineageStore + + +TOKEN = "secret" +AUTH = {"Authorization": f"Bearer {TOKEN}"} +RETIRE = f"{CONTROL_ROUTE_PREFIX}/rollouts/retire" +DELETE = f"{CONTROL_ROUTE_PREFIX}/rollouts/delete" + + +@pytest.fixture +def ledger(tmp_path): + return FileLineageStore(tmp_path) + + +@pytest.fixture +def client(ledger): + app = FastAPI() + install_rollout_control_routes(app, ledger, auth_token=TOKEN) + return TestClient(app) + + +async def _record_failure(ledger, rollout_id: str) -> None: + await ledger.record_failure(rollout_id, "c1", "worker_capture_failed") + + +@pytest.mark.asyncio +async def test_retire_route_removes_ledgers_and_fences_them(client, ledger, tmp_path): + await _record_failure(ledger, "r1") + + response = client.post(RETIRE, json={"rollout_ids": ["r1", "r2"]}, headers=AUTH) + + assert response.status_code == 200 + assert response.json() == {"removed": ["r1"], "absent": ["r2"]} + assert not (tmp_path / "r1.lineage.jsonl").exists() + await _record_failure(ledger, "r1") + assert not await ledger.has_rows("r1") + + +@pytest.mark.asyncio +async def test_delete_route_removes_the_fence(client, ledger): + await _record_failure(ledger, "r1") + client.post(RETIRE, json={"rollout_ids": ["r1"]}, headers=AUTH) + + response = client.post(DELETE, json={"rollout_ids": ["r1"]}, headers=AUTH) + + assert response.status_code == 200 + await _record_failure(ledger, "r1") + assert await ledger.has_rows("r1") + + +@pytest.mark.parametrize("route", [RETIRE, DELETE]) +def test_ledger_routes_require_the_bearer_token(client, route): + for headers in ({}, {"Authorization": "Bearer wrong"}): + assert client.post(route, json={"rollout_ids": ["r1"]}, headers=headers).status_code == 401 + + +@pytest.mark.parametrize("route", [RETIRE, DELETE]) +@pytest.mark.parametrize( + "body", + [ + {"rollout_ids": []}, + {"rollout_ids": [f"r{index}" for index in range(control_routes.MAX_LEDGER_BATCH + 1)]}, + {"rollout_ids": ["r1"], "unknown": True}, + ], +) +def test_ledger_routes_reject_malformed_batches(client, route, body): + assert client.post(route, json=body, headers=AUTH).status_code == 422 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", [RETIRE, DELETE]) +async def test_ledger_routes_reject_an_invalid_rollout_id_without_changing_anything(client, ledger, route): + await _record_failure(ledger, "r1") + response = client.post(route, json={"rollout_ids": ["r1", "a/b"]}, headers=AUTH) + assert response.status_code == 400 + assert await ledger.has_rows("r1") + + +class _Response: + def __init__(self, status: int, payload: dict) -> None: + self.status = status + self._payload = payload + + async def json(self) -> dict: + return self._payload + + async def text(self) -> str: + return str(self._payload) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["retire", "delete"]) +async def test_client_splits_large_batches(monkeypatch, action): + monkeypatch.setattr(control_routes, "MAX_LEDGER_BATCH", 2) + calls = [] + + async def fake_request(method, path, **kwargs): + body = kwargs["json"] + calls.append((method, path, body)) + return _Response(200, {"removed": body["rollout_ids"][:1], "absent": body["rollout_ids"][1:]}) + + client = RolloutControlClient("http://model", auth_token=TOKEN, request_timeout_s=1.0) + monkeypatch.setattr(client, "_request", fake_request) + + result = await getattr(client, action)(["r1", "r2", "r3"]) + + assert calls == [ + ("POST", f"/rollouts/{action}", {"rollout_ids": ["r1", "r2"]}), + ("POST", f"/rollouts/{action}", {"rollout_ids": ["r3"]}), + ] + assert result.removed == ["r1", "r3"] + assert result.absent == ["r2"] + + +@pytest.mark.asyncio +async def test_client_raises_on_a_failed_request(monkeypatch): + async def fake_request(method, path, **kwargs): + return _Response(400, {"detail": "Invalid rollout id"}) + + client = RolloutControlClient("http://model", auth_token=TOKEN, request_timeout_s=1.0) + monkeypatch.setattr(client, "_request", fake_request) + + with pytest.raises(RuntimeError, match="HTTP 400"): + await client.retire(["r1"]) diff --git a/tests/unit_tests/test_token_capture_ledger.py b/tests/unit_tests/test_token_capture_ledger.py index eb51b91a84f..34a994638db 100644 --- a/tests/unit_tests/test_token_capture_ledger.py +++ b/tests/unit_tests/test_token_capture_ledger.py @@ -10,6 +10,8 @@ from __future__ import annotations +from unittest.mock import patch + import pytest from nemo_gym.token_id_capture.fingerprint import FINGERPRINT_VERSION @@ -28,7 +30,12 @@ compute_chain_hash, hash_token_ids, ) -from nemo_gym.token_id_capture.staging.records import CallRecord, CaptureLedgerCommit, RolloutManifest +from nemo_gym.token_id_capture.staging.records import ( + CallRecord, + CaptureLedgerCommit, + RolloutManifest, + RolloutRemoval, +) USER_1 = {"role": "user", "content": "solve the task"} @@ -471,3 +478,130 @@ async def test_unversioned_legacy_token_carrying_row_cannot_resolve_or_anchor_a_ assert sorted(failure.reason for failure in manifest.failures) == sorted( ["ledger_row_missing_response_id", UNRESOLVED_PARENT_REASON] ) + + +# --- retire and delete ------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_retire_removes_ledgers_and_reports_absent_rollouts(store): + await _record_call_1(store) + await store.record_failure("r2", "c1", "worker_capture_failed") + + result = RolloutRemoval.model_validate(await store.retire(["r1", "r2", "r-none", "r1"])) + + # Duplicates in one batch count once. + assert result.removed == ["r1", "r2"] + assert result.absent == ["r-none"] + for rollout_id in ("r1", "r2"): + assert not await store.has_rows(rollout_id) + manifest = RolloutManifest.model_validate(await store.manifest(rollout_id)) + assert manifest.records == [] and manifest.failures == [] + # The retired call can no longer anchor a continuation. + assert (await store.resolve("r1", [USER_1, ASSISTANT_1, USER_2])).match is None + + +@pytest.mark.asyncio +async def test_retiring_again_is_a_no_op(store): + await _record_call_1(store) + await store.retire(["r1"]) + result = RolloutRemoval.model_validate(await store.retire(["r1"])) + assert result.removed == [] and result.absent == ["r1"] + + +@pytest.mark.asyncio +async def test_retired_rollouts_discard_later_records_and_failures(store): + await _record_call_1(store) + # Retiring a rollout that never recorded a call fences it as well. + await store.retire(["r1", "r-unstarted"]) + + await _record_call_1(store) + await store.record_failure("r1", "c2", "worker_capture_failed") + await _record_call_1(store, rollout_id="r-unstarted") + + for rollout_id in ("r1", "r-unstarted"): + assert not await store.has_rows(rollout_id) + # Other rollouts are unaffected. + await _record_call_1(store, rollout_id="r2") + assert await store.has_rows("r2") + + +@pytest.mark.asyncio +async def test_delete_removes_the_fence_so_the_rollout_id_can_be_reused(store): + await _record_call_1(store) + await store.retire(["r1"]) + + result = RolloutRemoval.model_validate(await store.delete(["r1"])) + assert result.removed == [] and result.absent == ["r1"] + + await _record_call_1(store) + manifest = RolloutManifest.model_validate(await store.manifest("r1")) + assert [record.model_call_id for record in manifest.records] == ["c1"] + + +@pytest.mark.asyncio +async def test_delete_removes_a_live_ledger(store): + await _record_call_1(store) + result = RolloutRemoval.model_validate(await store.delete(["r1"])) + assert result.removed == ["r1"] + assert not await store.has_rows("r1") + + +@pytest.mark.asyncio +async def test_file_store_retirement_is_visible_to_every_worker(tmp_path): + retirer = FileLineageStore(tmp_path) + late_writer = FileLineageStore(tmp_path) + await _record_call_1(late_writer) + + await retirer.retire(["r1"]) + await late_writer.record(_commit(_call_record("c2"), [USER_1], [ASSISTANT_1])) + # A fresh process sees retirements made before it started. + await _record_call_1(FileLineageStore(tmp_path)) + await late_writer.record_failure("r1", "c3", "worker_capture_failed") + + # No ledger was recreated; only the fence remains. + assert sorted(path.name for path in tmp_path.glob("*.lineage.*")) == ["r1.lineage.retired"] + + +@pytest.mark.asyncio +async def test_file_store_checks_the_fence_only_before_creating_a_ledger(tmp_path): + store = FileLineageStore(tmp_path) + with patch.object(store, "_is_retired", wraps=store._is_retired) as is_retired: + await _record_call_1(store) + await store.record(_commit(_call_record("c2"), [USER_1], [ASSISTANT_1])) + await store.record_failure("r1", "c3", "worker_capture_failed") + # Only the first write, which creates the ledger, looks for a fence: a ledger with rows is not retired. + assert is_retired.call_count == 1 + + +@pytest.mark.asyncio +async def test_file_store_delete_removes_every_ledger_and_fence_file(tmp_path): + store = FileLineageStore(tmp_path) + for rollout_id in ("r1", "r2"): + await _record_call_1(store, rollout_id=rollout_id) + await store.retire(["r1"]) + + await store.delete(["r1", "r2"]) + + assert list(tmp_path.glob("*.lineage.*")) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["retire", "delete"]) +async def test_a_bare_rollout_id_string_is_rejected(store, method): + """A string is a sequence of one-character IDs; treating it as a batch would remove the wrong rollouts.""" + await _record_call_1(store) + with pytest.raises(TypeError, match="sequence of rollout ids"): + await getattr(store, method)("r1") + assert await store.has_rows("r1") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["retire", "delete"]) +async def test_file_store_rejects_the_whole_batch_on_an_invalid_rollout_id(tmp_path, method): + store = FileLineageStore(tmp_path) + await _record_call_1(store) + with pytest.raises(ValueError, match="Invalid rollout id"): + await getattr(store, method)(["r1", "../escape"]) + assert await store.has_rows("r1") + assert not (tmp_path / "r1.lineage.retired").exists() From 465ee92871522a3bee660ba75ffe497ebce75ebb Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Fri, 9 Oct 2026 09:02:10 -0700 Subject: [PATCH 2/6] fix(token-capture): validate batches up front and skip no-op ledger syncs RolloutControlClient now validates and deduplicates the whole list of rollout IDs before sending the first batch. Before, a bare string was split into one-character IDs, an invalid ID in a later batch failed only after earlier batches had been retired, and an ID repeated across batches was reported both removed and absent. Retire no longer touches fences that already exist, and retire and delete sync the directory only when they changed something. A retire with a ledger left to delete still syncs first, so a fence left by a crashed earlier retire is durable before its ledger goes. The manifest of a retired rollout now raises RolloutRetiredError, which the manifest route returns as HTTP 410, instead of an empty manifest that reads like a rollout with no calls. The warning for a discarded late row names the call's staging key, since its staged tokens are in no manifest. retire and delete declare their result as the RolloutRemovalPayload TypedDict instead of dict. Signed-off-by: Ananth Subramaniam --- .../pages/model-server/model-call-capture.mdx | 4 + nemo_gym/token_id_capture/control_routes.py | 16 +++- nemo_gym/token_id_capture/lineage.py | 65 +++++++++++---- nemo_gym/token_id_capture/protocols.py | 36 ++++++-- .../test_token_capture_control_routes.py | 59 +++++++++++++ tests/unit_tests/test_token_capture_ledger.py | 83 ++++++++++++++++++- 6 files changed, 229 insertions(+), 34 deletions(-) diff --git a/fern/versions/latest/pages/model-server/model-call-capture.mdx b/fern/versions/latest/pages/model-server/model-call-capture.mdx index 0017672678d..165a57b43ae 100644 --- a/fern/versions/latest/pages/model-server/model-call-capture.mdx +++ b/fern/versions/latest/pages/model-server/model-call-capture.mdx @@ -305,6 +305,10 @@ Retiring deletes the rollout's ledger and leaves an empty `.lineage. That file acts as a fence: if a model call for the rollout finishes later, Gym discards its row instead of starting a new ledger. Late calls are more common than they sound, because an agent the framework gave up on can keep running, and a client retry can run a finished rollout a second time. `TokenSource.drop` leaves the same kind of fence in the complete-record token store. +After a rollout is retired, its manifest route returns HTTP 410 instead of an empty manifest, so a framework that reads it too late finds out rather than building an empty receipt. + +A late call may already have staged its tokens before Gym discarded its row, and no manifest names those tokens. +Gym logs a warning with the call's staging key, and the framework that staged the tokens has to remove them, for example by sweeping staged rows that no receipt references. Retire every rollout the framework is done with, whether it finished or was abandoned. A finished rollout can be retired once its manifest has been read and its receipt built, unless a crash before the receipt is saved would make the framework rerun it under the same ID; in that case, wait until the receipt is saved. diff --git a/nemo_gym/token_id_capture/control_routes.py b/nemo_gym/token_id_capture/control_routes.py index 356240be28f..04c684a9b29 100644 --- a/nemo_gym/token_id_capture/control_routes.py +++ b/nemo_gym/token_id_capture/control_routes.py @@ -19,7 +19,7 @@ from fastapi import APIRouter, Header, HTTPException from pydantic import BaseModel, ConfigDict, Field -from nemo_gym.token_id_capture.protocols import CaptureLedger +from nemo_gym.token_id_capture.protocols import CaptureLedger, RolloutRemovalPayload, RolloutRetiredError from nemo_gym.token_id_capture.staging.records import RolloutManifest, RolloutRemoval @@ -64,6 +64,9 @@ async def rollout_manifest( check_auth(authorization) try: return await lineage_store.manifest(rollout_id) + except RolloutRetiredError as error: + # Gone, not empty: an empty manifest would read as a rollout that made no calls. + raise HTTPException(status_code=410, detail=str(error)) from error except ValueError as error: raise HTTPException(status_code=400, detail=str(error)) from error @@ -71,7 +74,7 @@ async def rollout_manifest( async def retire_rollouts( body: LedgerBatchRequest, authorization: str | None = Header(default=None), - ) -> dict: + ) -> RolloutRemovalPayload: check_auth(authorization) try: return await lineage_store.retire(body.rollout_ids) @@ -82,7 +85,7 @@ async def retire_rollouts( async def delete_rollouts( body: LedgerBatchRequest, authorization: str | None = Header(default=None), - ) -> dict: + ) -> RolloutRemovalPayload: check_auth(authorization) try: return await lineage_store.delete(body.rollout_ids) @@ -125,7 +128,12 @@ async def delete(self, rollout_ids: Sequence[str]) -> RolloutRemoval: return await self._remove("delete", rollout_ids) async def _remove(self, action: str, rollout_ids: Sequence[str]) -> RolloutRemoval: - rollout_ids = list(rollout_ids) + from nemo_gym.token_id_capture.store import validate_rollout_ids + + # Validate and deduplicate the whole list before the first request. Batch by batch, a bare string + # would become one-character IDs, an invalid ID in a later batch would fail after earlier batches + # had changed state, and an ID repeated across batches would be reported both removed and absent. + rollout_ids = validate_rollout_ids(rollout_ids) result = RolloutRemoval() for start in range(0, len(rollout_ids), MAX_LEDGER_BATCH): batch = rollout_ids[start : start + MAX_LEDGER_BATCH] diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py index 255206c5d90..b54416afd48 100644 --- a/nemo_gym/token_id_capture/lineage.py +++ b/nemo_gym/token_id_capture/lineage.py @@ -60,7 +60,12 @@ from nemo_gym.token_id_capture.fingerprint import ( canonicalize_tool_arguments as canonicalize_tool_arguments, ) -from nemo_gym.token_id_capture.protocols import LineageMatch, LineageResolution +from nemo_gym.token_id_capture.protocols import ( + LineageMatch, + LineageResolution, + RolloutRemovalPayload, + RolloutRetiredError, +) from nemo_gym.token_id_capture.records import ParentResolutionStatus, TokenEntry, cumulative_tokens @@ -71,13 +76,20 @@ logger = logging.getLogger(__name__) -def _log_late_write_dropped(rollout_id: str, model_call_id: str) -> None: +def _log_late_write_dropped(rollout_id: str, model_call_id: str, staging_key: str | None = None) -> None: # A call that finished after its rollout was retired, such as one from an abandoned attempt or a request # the client had already retried. Nobody will read its row, so this is not an error. + # The worker staged the call's tokens before the ledger refused the row, so no manifest names them; + # the warning carries the staging key for the framework that has to remove them. + if staging_key is None: + logger.warning("Discarding the ledger row of model call %s: rollout %s is retired.", model_call_id, rollout_id) + return logger.warning( - "Discarding the ledger row of model call %s: rollout %s is retired.", + "Discarding the ledger row of model call %s: rollout %s is retired. " + "Its staged tokens (staging key %s) are in no manifest, so the framework must remove them.", model_call_id, rollout_id, + staging_key, ) @@ -477,7 +489,7 @@ async def record(self, commit: CaptureLedgerCommit) -> None: # lineage-only local-capture rows that inject prompt prefixes. record = commit.record if commit.rollout_id in self._retired: - _log_late_write_dropped(commit.rollout_id, record.model_call_id) + _log_late_write_dropped(commit.rollout_id, record.model_call_id, record.staging_key) return rows = self._ledgers.setdefault(commit.rollout_id, []) row = {"model_call_id": record.model_call_id, **_custody_columns(record, commit.staging_chain)} @@ -512,6 +524,8 @@ async def record_failure(self, rollout_id: str, model_call_id: str, reason: str) rows.append(row) async def manifest(self, rollout_id: str) -> dict: + if rollout_id in self._retired: + raise RolloutRetiredError(f"rollout {rollout_id} is retired") return _manifest_from_rows(rollout_id, list(self._ledgers.get(rollout_id) or [])) async def has_rows(self, rollout_id: str) -> bool: @@ -519,21 +533,21 @@ async def has_rows(self, rollout_id: str) -> bool: return True return bool(self.index.for_rollout(rollout_id).by_call_id) - async def retire(self, rollout_ids: Sequence[str]) -> dict: + async def retire(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: from nemo_gym.token_id_capture.store import validate_rollout_ids rollout_ids = validate_rollout_ids(rollout_ids) self._retired.update(rollout_ids) return self._remove(rollout_ids) - async def delete(self, rollout_ids: Sequence[str]) -> dict: + async def delete(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: from nemo_gym.token_id_capture.store import validate_rollout_ids rollout_ids = validate_rollout_ids(rollout_ids) self._retired.difference_update(rollout_ids) return self._remove(rollout_ids) - def _remove(self, rollout_ids: list[str]) -> dict: + def _remove(self, rollout_ids: list[str]) -> RolloutRemovalPayload: removed, absent = [], [] for rollout_id in rollout_ids: (removed if self._ledgers.pop(rollout_id, None) is not None else absent).append(rollout_id) @@ -964,7 +978,7 @@ def _record(self, commit: CaptureLedgerCommit) -> None: with self._locked(commit.rollout_id): records = self._read(commit.rollout_id) if not records and self._is_retired(commit.rollout_id): - _log_late_write_dropped(commit.rollout_id, model_call_id) + _log_late_write_dropped(commit.rollout_id, model_call_id, commit.record.staging_key) return matches = [existing for existing in records if existing["model_call_id"] == model_call_id] if matches: @@ -995,6 +1009,8 @@ async def manifest(self, rollout_id: str) -> dict: def _manifest(self, rollout_id: str) -> dict: with self._locked(rollout_id): rows = list(self._read(rollout_id)) + if not rows and self._is_retired(rollout_id): + raise RolloutRetiredError(f"rollout {rollout_id} is retired") return _manifest_from_rows(rollout_id, rows) async def has_rows(self, rollout_id: str) -> bool: @@ -1004,42 +1020,55 @@ def _has_rows(self, rollout_id: str) -> bool: with self._locked(rollout_id): return bool(self._read(rollout_id)) - async def retire(self, rollout_ids: Sequence[str]) -> dict: + async def retire(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: return await asyncio.to_thread(self._retire, rollout_ids) - def _retire(self, rollout_ids: Sequence[str]) -> dict: + def _retire(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: from nemo_gym.token_id_capture.store import validate_rollout_ids rollout_ids = validate_rollout_ids(rollout_ids) - # Make every retired marker durable before deleting any ledger, so a crash in between leaves a fence. + created_fence = has_ledger = False for rollout_id in rollout_ids: with self._locked(rollout_id): - self._retired_path(rollout_id).touch() - if rollout_ids: + fence = self._retired_path(rollout_id) + if not fence.exists(): + fence.touch() + created_fence = True + has_ledger = has_ledger or self._ledger_path(rollout_id).exists() + # Make every fence durable before deleting any ledger, so a crash in between leaves a fence. That + # includes fences that already existed: a retire that crashed earlier may have created them without + # syncing. A retried batch with nothing left to delete changes nothing, so it syncs nothing. + if created_fence or has_ledger: self._fsync_ledger_root() return self._remove(rollout_ids) - async def delete(self, rollout_ids: Sequence[str]) -> dict: + async def delete(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: return await asyncio.to_thread(self._delete, rollout_ids) - def _delete(self, rollout_ids: Sequence[str]) -> dict: + def _delete(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: from nemo_gym.token_id_capture.store import validate_rollout_ids return self._remove(validate_rollout_ids(rollout_ids), unretire=True) - def _remove(self, rollout_ids: list[str], *, unretire: bool = False) -> dict: + def _remove(self, rollout_ids: list[str], *, unretire: bool = False) -> RolloutRemovalPayload: removed, absent = [], [] + changed = False for rollout_id in rollout_ids: with self._locked(rollout_id): try: self._ledger_path(rollout_id).unlink() removed.append(rollout_id) + changed = True except FileNotFoundError: absent.append(rollout_id) if unretire: - self._retired_path(rollout_id).unlink(missing_ok=True) + try: + self._retired_path(rollout_id).unlink() + changed = True + except FileNotFoundError: + pass with self._cache_guard: self._ledger_cache.pop(rollout_id, None) - if rollout_ids: + if changed: self._fsync_ledger_root() return {"removed": removed, "absent": absent} diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py index ed223862544..32cebc57236 100644 --- a/nemo_gym/token_id_capture/protocols.py +++ b/nemo_gym/token_id_capture/protocols.py @@ -43,12 +43,33 @@ from collections.abc import Sequence from dataclasses import dataclass -from typing import Protocol, runtime_checkable +from typing import Protocol, TypedDict, runtime_checkable from nemo_gym.token_id_capture.records import ParentResolutionStatus, TokenEntry from nemo_gym.token_id_capture.staging.records import CaptureLedgerCommit +class RolloutRemovalPayload(TypedDict): + """The result of retiring or deleting rollouts, as plain wire data. + + It validates as ``staging.records.RolloutRemoval``. + """ + + # Rollouts whose captured data existed and was removed. + removed: list[str] + # Rollouts with no captured data: already removed, or never recorded a call. + absent: list[str] + + +class RolloutRetiredError(LookupError): + """Reject reading the ledger of a rollout that was retired. + + A retired rollout's ledger is gone, so an empty manifest would look like a rollout that made no calls. + Reading it is a framework error, such as fetching a manifest after retiring the rollout. + Deleting the rollout removes its fence, after which the ID reads as unused again. + """ + + class TokenCaptureFrozenError(RuntimeError): """Reject a write that arrived after its rollout's capture was frozen. @@ -175,6 +196,7 @@ async def manifest(self, rollout_id: str) -> dict: committed rows under ``records`` (each a ``CallRecord`` payload) and poison rows under ``failures``. Cumulative token IDs never appear in the manifest. + A retired rollout raises ``RolloutRetiredError`` instead of returning an empty manifest. """ ... @@ -182,23 +204,21 @@ async def has_rows(self, rollout_id: str) -> bool: """Return whether any ledger row (committed or failed) exists.""" ... - async def retire(self, rollout_ids: Sequence[str]) -> dict: + async def retire(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: """Remove the ledgers of rollouts the framework is done with, and fence them. - Retire a finished rollout once its manifest has been read and the receipt built from it is - durable, and an abandoned rollout as soon as the framework gives up on it. Like + Retire a finished rollout once its manifest has been read and the receipt built from it, + and an abandoned rollout as soon as the framework gives up on it. Like ``TokenSource.drop`` on the complete-record store, retiring keeps a fence: later ``record`` and ``record_failure`` calls for the rollout are discarded instead of starting a new ledger. Even a finished rollout can still write: a client retry can run it twice, and the second execution writes after the first one returned. Retiring again is a no-op, so retrying a batch is safe. Staged token data is untouched: the framework that staged it also removes it. - - The result validates as ``staging.records.RolloutRemoval``. """ ... - async def delete(self, rollout_ids: Sequence[str]) -> dict: + async def delete(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: """Remove the ledgers and fences of rollouts, as ``TokenCaptureStore.delete`` does. Delete retired rollouts once nothing of those attempts can still write, for example after every @@ -206,8 +226,6 @@ async def delete(self, rollout_ids: Sequence[str]) -> dict: ID before reusing it, so the new execution starts from an empty ledger. Deleting again is a no-op. Staged token data is untouched: the framework that staged it also removes it. - - The result validates as ``staging.records.RolloutRemoval``. """ ... diff --git a/tests/unit_tests/test_token_capture_control_routes.py b/tests/unit_tests/test_token_capture_control_routes.py index 8edd6ab0712..c72eb09c828 100644 --- a/tests/unit_tests/test_token_capture_control_routes.py +++ b/tests/unit_tests/test_token_capture_control_routes.py @@ -139,3 +139,62 @@ async def fake_request(method, path, **kwargs): with pytest.raises(RuntimeError, match="HTTP 400"): await client.retire(["r1"]) + + +def _client_over(client_app, monkeypatch) -> RolloutControlClient: + """A control client whose requests reach the real routes and ledger.""" + + async def request(method, path, **kwargs): + response = client_app.request(method, f"{CONTROL_ROUTE_PREFIX}{path}", json=kwargs.get("json"), headers=AUTH) + return _Response(response.status_code, response.json()) + + control = RolloutControlClient("http://model", auth_token=TOKEN, request_timeout_s=1.0) + monkeypatch.setattr(control, "_request", request) + return control + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["retire", "delete"]) +async def test_client_rejects_a_bare_rollout_id_before_sending_anything(client, ledger, monkeypatch, action): + await _record_failure(ledger, "r") + control = _client_over(client, monkeypatch) + + # A string is a sequence of one-character IDs: "rollout-abc" would remove rollouts "r", "o", "l", ... + with pytest.raises(TypeError, match="sequence of rollout ids"): + await getattr(control, action)("rollout-abc") + assert await ledger.has_rows("r") + + +@pytest.mark.asyncio +async def test_client_validates_every_batch_before_changing_anything(client, ledger, monkeypatch): + monkeypatch.setattr(control_routes, "MAX_LEDGER_BATCH", 2) + await _record_failure(ledger, "r1") + control = _client_over(client, monkeypatch) + + with pytest.raises(ValueError, match="Invalid rollout id"): + await control.retire(["r1", "r2", "a/b"]) + assert await ledger.has_rows("r1") + + +@pytest.mark.asyncio +async def test_client_reports_an_id_repeated_across_batches_once(client, ledger, monkeypatch): + monkeypatch.setattr(control_routes, "MAX_LEDGER_BATCH", 2) + await _record_failure(ledger, "r1") + control = _client_over(client, monkeypatch) + + result = await control.retire(["r1", "r2", "r1"]) + + assert result.removed == ["r1"] + assert result.absent == ["r2"] + + +@pytest.mark.asyncio +async def test_manifest_route_reports_a_retired_rollout_as_gone(client, ledger, monkeypatch): + await _record_failure(ledger, "r1") + client.post(RETIRE, json={"rollout_ids": ["r1"]}, headers=AUTH) + + response = client.get(f"{CONTROL_ROUTE_PREFIX}/rollouts/r1/manifest", headers=AUTH) + + assert response.status_code == 410 + with pytest.raises(RuntimeError, match="HTTP 410"): + await _client_over(client, monkeypatch).manifest("r1") diff --git a/tests/unit_tests/test_token_capture_ledger.py b/tests/unit_tests/test_token_capture_ledger.py index 34a994638db..a187313b927 100644 --- a/tests/unit_tests/test_token_capture_ledger.py +++ b/tests/unit_tests/test_token_capture_ledger.py @@ -16,7 +16,7 @@ from nemo_gym.token_id_capture.fingerprint import FINGERPRINT_VERSION from nemo_gym.token_id_capture.lineage import FileLineageStore, InMemoryLineageStore, _custody_columns -from nemo_gym.token_id_capture.protocols import CaptureLedger +from nemo_gym.token_id_capture.protocols import CaptureLedger, RolloutRemovalPayload, RolloutRetiredError from nemo_gym.token_id_capture.records import ParentResolutionStatus, compute_digest from nemo_gym.token_id_capture.sink import ( UNRESOLVED_PARENT_REASON, @@ -495,8 +495,8 @@ async def test_retire_removes_ledgers_and_reports_absent_rollouts(store): assert result.absent == ["r-none"] for rollout_id in ("r1", "r2"): assert not await store.has_rows(rollout_id) - manifest = RolloutManifest.model_validate(await store.manifest(rollout_id)) - assert manifest.records == [] and manifest.failures == [] + with pytest.raises(RolloutRetiredError): + await store.manifest(rollout_id) # The retired call can no longer anchor a continuation. assert (await store.resolve("r1", [USER_1, ASSISTANT_1, USER_2])).match is None @@ -605,3 +605,80 @@ async def test_file_store_rejects_the_whole_batch_on_an_invalid_rollout_id(tmp_p await getattr(store, method)(["r1", "../escape"]) assert await store.has_rows("r1") assert not (tmp_path / "r1.lineage.retired").exists() + + +@pytest.mark.asyncio +async def test_retiring_again_does_not_touch_fences_or_sync_the_directory(tmp_path): + store = FileLineageStore(tmp_path) + await _record_call_1(store) + await store.retire(["r1"]) + fence = tmp_path / "r1.lineage.retired" + fence_mtime = fence.stat().st_mtime_ns + + with patch.object(store, "_fsync_ledger_root", wraps=store._fsync_ledger_root) as fsync_root: + result = await store.retire(["r1"]) + await store.delete(["never-recorded"]) + + # A retried batch changes nothing on disk, so it writes no metadata and syncs nothing. + assert result == {"removed": [], "absent": ["r1"]} + assert fence.stat().st_mtime_ns == fence_mtime + assert fsync_root.call_count == 0 + + +@pytest.mark.asyncio +async def test_retire_syncs_an_existing_fence_before_deleting_the_ledger(tmp_path): + """A crashed retire can leave a fence that was never synced next to a ledger that was never deleted.""" + store = FileLineageStore(tmp_path) + await _record_call_1(store) + (tmp_path / "r1.lineage.retired").touch() + events = [] + real_fsync, real_unlink = store._fsync_ledger_root, type(tmp_path).unlink + + def fsync_root(): + events.append("fsync") + real_fsync() + + def unlink(path, *args, **kwargs): + if path.name == "r1.lineage.jsonl": + events.append("unlink ledger") + return real_unlink(path, *args, **kwargs) + + with patch.object(store, "_fsync_ledger_root", fsync_root), patch.object(type(tmp_path), "unlink", unlink): + result = await store.retire(["r1"]) + + assert result == {"removed": ["r1"], "absent": []} + assert events[:2] == ["fsync", "unlink ledger"] + + +@pytest.mark.asyncio +async def test_a_discarded_late_row_names_its_staging_key(tmp_path, caplog): + """The worker staged the late call's tokens; the warning is the only remaining pointer to them.""" + store = FileLineageStore(tmp_path) + await _record_call_1(store) + await store.retire(["r1"]) + + with caplog.at_level("WARNING"): + await store.record(_commit(_call_record("c2"), [USER_1], [ASSISTANT_1], staging_chain=("r1/c1", "r1/c2"))) + + assert any("r1/c2" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_reading_a_retired_rollouts_manifest_fails(store): + """An empty manifest would look like a rollout that made no calls.""" + await _record_call_1(store) + await store.retire(["r1", "never-recorded"]) + + for rollout_id in ("r1", "never-recorded"): + with pytest.raises(RolloutRetiredError): + await store.manifest(rollout_id) + # Delete removes the fence, so the ID reads as unused again. + await store.delete(["r1"]) + assert RolloutManifest.model_validate(await store.manifest("r1")).records == [] + + +def test_retire_and_delete_declare_their_result_shape(): + import typing + + for method in (CaptureLedger.retire, CaptureLedger.delete, FileLineageStore.retire, InMemoryLineageStore.delete): + assert typing.get_type_hints(method)["return"] is RolloutRemovalPayload From 97c02589565f07fe002dd78039f67f4080f010e4 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Fri, 9 Oct 2026 10:29:30 -0700 Subject: [PATCH 3/6] fix(token-capture): refuse capture for retired rollouts and type the client's manifest errors A call that started after its rollout was retired was admitted as a new text root: the retired rollout had no ledger rows, so admission saw a fresh rollout. The worker then staged the call's tokens and the ledger discarded its row, leaving staged tokens that no manifest names. The same happened to the first call of a duplicate run, which resolves as a root. has_rows now raises RolloutRetiredError for a retired rollout, and admission skips capture for it. RolloutControlClient.manifest now validates the rollout ID before building the URL (dot segments are normalized, so "../rollouts/r1" fetched r1's manifest), raises RolloutRetiredError for HTTP 410 instead of a plain RuntimeError, and rejects a manifest for a different rollout. Signed-off-by: Ananth Subramaniam --- .../pages/model-server/model-call-capture.mdx | 3 +- nemo_gym/token_id_capture/control_routes.py | 12 +++++- nemo_gym/token_id_capture/lineage.py | 8 +++- nemo_gym/token_id_capture/protocols.py | 5 ++- nemo_gym/token_id_capture/sink.py | 14 +++++- .../test_token_capture_control_routes.py | 43 ++++++++++++++++++- tests/unit_tests/test_token_capture_ledger.py | 20 ++++++++- 7 files changed, 96 insertions(+), 9 deletions(-) diff --git a/fern/versions/latest/pages/model-server/model-call-capture.mdx b/fern/versions/latest/pages/model-server/model-call-capture.mdx index 165a57b43ae..e3bdca9b9bb 100644 --- a/fern/versions/latest/pages/model-server/model-call-capture.mdx +++ b/fern/versions/latest/pages/model-server/model-call-capture.mdx @@ -307,7 +307,8 @@ Late calls are more common than they sound, because an agent the framework gave `TokenSource.drop` leaves the same kind of fence in the complete-record token store. After a rollout is retired, its manifest route returns HTTP 410 instead of an empty manifest, so a framework that reads it too late finds out rather than building an empty receipt. -A late call may already have staged its tokens before Gym discarded its row, and no manifest names those tokens. +A call that starts after its rollout was retired is not admitted for capture, so it stages nothing. +A call that was already in flight, however, may have staged its tokens before Gym discarded its row, and no manifest names those tokens. Gym logs a warning with the call's staging key, and the framework that staged the tokens has to remove them, for example by sweeping staged rows that no receipt references. Retire every rollout the framework is done with, whether it finished or was abandoned. diff --git a/nemo_gym/token_id_capture/control_routes.py b/nemo_gym/token_id_capture/control_routes.py index 04c684a9b29..d4e192aa6fa 100644 --- a/nemo_gym/token_id_capture/control_routes.py +++ b/nemo_gym/token_id_capture/control_routes.py @@ -112,12 +112,22 @@ def __init__( self._request_timeout_s = request_timeout_s async def manifest(self, rollout_id: str) -> RolloutManifest: + """Fetch a rollout's manifest. A retired rollout raises ``RolloutRetiredError``.""" + from nemo_gym.token_id_capture.store import validate_rollout_id + + # Validate before building the URL: dot segments are normalized, so "../rollouts/r1" would fetch r1. + validate_rollout_id(rollout_id) response = await self._request("GET", f"/rollouts/{rollout_id}/manifest") + if response.status == 410: + raise RolloutRetiredError(f"rollout {rollout_id} is retired: {await response.text()}") if response.status != 200: raise RuntimeError( f"rollout {rollout_id} manifest fetch failed: HTTP {response.status} {await response.text()}" ) - return RolloutManifest.model_validate(await response.json()) + manifest = RolloutManifest.model_validate(await response.json()) + if manifest.rollout_id != rollout_id: + raise ValueError(f"asked for the manifest of rollout {rollout_id} but got rollout {manifest.rollout_id}") + return manifest async def retire(self, rollout_ids: Sequence[str]) -> RolloutRemoval: """Retire ledgers; see ``CaptureLedger.retire``. Retrying after a failure is safe.""" diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py index b54416afd48..ba5b0a6680b 100644 --- a/nemo_gym/token_id_capture/lineage.py +++ b/nemo_gym/token_id_capture/lineage.py @@ -531,6 +531,8 @@ async def manifest(self, rollout_id: str) -> dict: async def has_rows(self, rollout_id: str) -> bool: if self._ledgers.get(rollout_id): return True + if rollout_id in self._retired: + raise RolloutRetiredError(f"rollout {rollout_id} is retired") return bool(self.index.for_rollout(rollout_id).by_call_id) async def retire(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: @@ -1018,7 +1020,11 @@ async def has_rows(self, rollout_id: str) -> bool: def _has_rows(self, rollout_id: str) -> bool: with self._locked(rollout_id): - return bool(self._read(rollout_id)) + if self._read(rollout_id): + return True + if self._is_retired(rollout_id): + raise RolloutRetiredError(f"rollout {rollout_id} is retired") + return False async def retire(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: return await asyncio.to_thread(self._retire, rollout_ids) diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py index 32cebc57236..0ee7ab6c6c4 100644 --- a/nemo_gym/token_id_capture/protocols.py +++ b/nemo_gym/token_id_capture/protocols.py @@ -201,7 +201,10 @@ async def manifest(self, rollout_id: str) -> dict: ... async def has_rows(self, rollout_id: str) -> bool: - """Return whether any ledger row (committed or failed) exists.""" + """Return whether any ledger row (committed or failed) exists. + + A retired rollout raises ``RolloutRetiredError``, so admission can refuse to capture its late calls. + """ ... async def retire(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index af78638e20e..98491fa6f05 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -39,6 +39,7 @@ CaptureLedger, LineageResolution, LineageResolver, + RolloutRetiredError, TokenCaptureFrozenError, TokenSink, ) @@ -286,7 +287,18 @@ async def resolve_parent(request_messages: list | None) -> None: ) return is_root = context.parent_resolution is not None and context.parent_resolution.status == ParentResolutionStatus.ROOT - if is_root or not await ledger.has_rows(context.rollout_id): + try: + has_rows = await ledger.has_rows(context.rollout_id) + except RolloutRetiredError: + # A late call of an abandoned attempt, or the first call of a duplicate run. Admitting it would stage + # tokens that the ledger then refuses to record, so no manifest would ever name them. + logger.warning( + "Model call %s of rollout %s arrived after the rollout was retired; its tokens are not captured.", + context.model_call_id, + context.rollout_id, + ) + return + if is_root or not has_rows: context.capture_admission = CaptureAdmission( rollout_id=context.rollout_id, model_call_id=context.model_call_id, diff --git a/tests/unit_tests/test_token_capture_control_routes.py b/tests/unit_tests/test_token_capture_control_routes.py index c72eb09c828..7ad8534d9a6 100644 --- a/tests/unit_tests/test_token_capture_control_routes.py +++ b/tests/unit_tests/test_token_capture_control_routes.py @@ -16,6 +16,7 @@ install_rollout_control_routes, ) from nemo_gym.token_id_capture.lineage import FileLineageStore +from nemo_gym.token_id_capture.protocols import RolloutRetiredError TOKEN = "secret" @@ -50,7 +51,8 @@ async def test_retire_route_removes_ledgers_and_fences_them(client, ledger, tmp_ assert response.json() == {"removed": ["r1"], "absent": ["r2"]} assert not (tmp_path / "r1.lineage.jsonl").exists() await _record_failure(ledger, "r1") - assert not await ledger.has_rows("r1") + with pytest.raises(RolloutRetiredError): + await ledger.has_rows("r1") @pytest.mark.asyncio @@ -196,5 +198,42 @@ async def test_manifest_route_reports_a_retired_rollout_as_gone(client, ledger, response = client.get(f"{CONTROL_ROUTE_PREFIX}/rollouts/r1/manifest", headers=AUTH) assert response.status_code == 410 - with pytest.raises(RuntimeError, match="HTTP 410"): + + +@pytest.mark.asyncio +async def test_client_reports_a_retired_manifest_with_the_typed_error(client, ledger, monkeypatch): + await _record_failure(ledger, "r1") + client.post(RETIRE, json={"rollout_ids": ["r1"]}, headers=AUTH) + + with pytest.raises(RolloutRetiredError): await _client_over(client, monkeypatch).manifest("r1") + + +@pytest.mark.asyncio +async def test_client_validates_a_manifest_rollout_id_before_sending(monkeypatch): + control = RolloutControlClient("http://model", auth_token=TOKEN, request_timeout_s=1.0) + sent = [] + + async def request(method, path, **kwargs): + sent.append(path) + return _Response(200, {"rollout_id": "r1", "records": [], "failures": []}) + + monkeypatch.setattr(control, "_request", request) + + # Dot segments are normalized in the URL, so "../rollouts/r1" would fetch rollout r1's manifest. + with pytest.raises(ValueError, match="Invalid rollout id"): + await control.manifest("../rollouts/r1") + assert sent == [] + + +@pytest.mark.asyncio +async def test_client_rejects_a_manifest_for_a_different_rollout(monkeypatch): + control = RolloutControlClient("http://model", auth_token=TOKEN, request_timeout_s=1.0) + + async def request(method, path, **kwargs): + return _Response(200, {"rollout_id": "r2", "records": [], "failures": []}) + + monkeypatch.setattr(control, "_request", request) + + with pytest.raises(ValueError, match="r2"): + await control.manifest("r1") diff --git a/tests/unit_tests/test_token_capture_ledger.py b/tests/unit_tests/test_token_capture_ledger.py index a187313b927..5d9b29008d0 100644 --- a/tests/unit_tests/test_token_capture_ledger.py +++ b/tests/unit_tests/test_token_capture_ledger.py @@ -494,7 +494,8 @@ async def test_retire_removes_ledgers_and_reports_absent_rollouts(store): assert result.removed == ["r1", "r2"] assert result.absent == ["r-none"] for rollout_id in ("r1", "r2"): - assert not await store.has_rows(rollout_id) + with pytest.raises(RolloutRetiredError): + await store.has_rows(rollout_id) with pytest.raises(RolloutRetiredError): await store.manifest(rollout_id) # The retired call can no longer anchor a continuation. @@ -520,7 +521,8 @@ async def test_retired_rollouts_discard_later_records_and_failures(store): await _record_call_1(store, rollout_id="r-unstarted") for rollout_id in ("r1", "r-unstarted"): - assert not await store.has_rows(rollout_id) + with pytest.raises(RolloutRetiredError): + await store.has_rows(rollout_id) # Other rollouts are unaffected. await _record_call_1(store, rollout_id="r2") assert await store.has_rows("r2") @@ -682,3 +684,17 @@ def test_retire_and_delete_declare_their_result_shape(): for method in (CaptureLedger.retire, CaptureLedger.delete, FileLineageStore.retire, InMemoryLineageStore.delete): assert typing.get_type_hints(method)["return"] is RolloutRemovalPayload + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_items", [[USER_1, ASSISTANT_1, USER_2], [USER_1]], ids=["continuation", "root"]) +async def test_a_call_on_a_retired_rollout_is_not_admitted_for_staging(store, request_items): + """A late continuation, or the first call of a duplicate run, would stage tokens no manifest names.""" + await _record_call_1(store) + await store.retire(["r1"]) + + context = await _admit(store, request_items) + + assert context.capture_admission is None + with pytest.raises(RolloutRetiredError): + await store.has_rows("r1") From 5ba8f437c5ab6b47ff99abf6183958768db61325 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Fri, 9 Oct 2026 10:48:12 -0700 Subject: [PATCH 4/6] fix(token-capture): remove cached ledger rows through one helper Retire and delete removed cached ledger rows with a direct dictionary pop. A cache that also tracks its size, as the bounded ledger cache in NVIDIA-NeMo/Gym#4268 does, would then keep counting the removed rows and gradually evict everything else. Removals now go through _ledger_cache_pop, the helper that cache defines, and a test checks that removing ledgers leaves no tracked size behind. Signed-off-by: Ananth Subramaniam --- nemo_gym/token_id_capture/lineage.py | 10 ++++++++-- tests/unit_tests/test_token_capture_ledger.py | 14 ++++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py index ba5b0a6680b..178901cf294 100644 --- a/nemo_gym/token_id_capture/lineage.py +++ b/nemo_gym/token_id_capture/lineage.py @@ -872,6 +872,13 @@ def _locked(self, rollout_id: str): # discipline covers both and no second lock file is minted. return self._store._locked(rollout_id) + def _ledger_cache_pop(self, rollout_id: str) -> None: + # Forget a rollout's cached rows. Every removal goes through here, so that a cache that also tracks + # its size (NVIDIA-NeMo/Gym#4268) stays consistent. Different rollouts are removed concurrently from + # worker threads, so the dictionary needs the guard. + with self._cache_guard: + self._ledger_cache.pop(rollout_id, None) + def _read(self, rollout_id: str) -> list[dict]: path = self._ledger_path(rollout_id) if not path.exists(): @@ -1073,8 +1080,7 @@ def _remove(self, rollout_ids: list[str], *, unretire: bool = False) -> RolloutR changed = True except FileNotFoundError: pass - with self._cache_guard: - self._ledger_cache.pop(rollout_id, None) + self._ledger_cache_pop(rollout_id) if changed: self._fsync_ledger_root() return {"removed": removed, "absent": absent} diff --git a/tests/unit_tests/test_token_capture_ledger.py b/tests/unit_tests/test_token_capture_ledger.py index 5d9b29008d0..0944de62b5d 100644 --- a/tests/unit_tests/test_token_capture_ledger.py +++ b/tests/unit_tests/test_token_capture_ledger.py @@ -698,3 +698,17 @@ async def test_a_call_on_a_retired_rollout_is_not_admitted_for_staging(store, re assert context.capture_admission is None with pytest.raises(RolloutRetiredError): await store.has_rows("r1") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["retire", "delete"]) +async def test_removing_ledgers_forgets_their_cached_rows(tmp_path, method): + store = FileLineageStore(tmp_path) + for index in range(5): + await _record_call_1(store, rollout_id=f"r{index}") + + await getattr(store, method)([f"r{index}" for index in range(5)]) + + assert store._ledger_cache == {} + # A cache that also tracks its size must account for every removal, or it slowly fills with nothing. + assert getattr(store, "_ledger_cache_weight", 0) == 0 From 8dcbe4bcf1040fe5ae9514cb42b7899c159b3ea6 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Fri, 9 Oct 2026 14:16:47 -0700 Subject: [PATCH 5/6] fix(token-capture): sync every non-empty ledger delete A delete that found nothing to remove skipped the directory sync. When it overlapped another delete of the same rollout that had removed the files but not yet synced, it returned success while the removal was not durable, so a crash could bring back the old fence or ledger after the caller had started reusing the rollout ID. Every non-empty delete batch now syncs; a retire with nothing left to remove still syncs nothing, because its fences are already durable. Signed-off-by: Ananth Subramaniam --- nemo_gym/token_id_capture/lineage.py | 5 ++++- tests/unit_tests/test_token_capture_ledger.py | 16 ++++++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py index 178901cf294..11af24595db 100644 --- a/nemo_gym/token_id_capture/lineage.py +++ b/nemo_gym/token_id_capture/lineage.py @@ -1081,6 +1081,9 @@ def _remove(self, rollout_ids: list[str], *, unretire: bool = False) -> RolloutR except FileNotFoundError: pass self._ledger_cache_pop(rollout_id) - if changed: + # A retire with nothing left to remove syncs nothing, since its fences are already durable. A delete + # always syncs a non-empty batch: an overlapping delete may have removed the files without syncing yet, + # and a caller that reuses the rollout ID must not see the old fence or ledger come back after a crash. + if changed or (unretire and rollout_ids): self._fsync_ledger_root() return {"removed": removed, "absent": absent} diff --git a/tests/unit_tests/test_token_capture_ledger.py b/tests/unit_tests/test_token_capture_ledger.py index 0944de62b5d..6b85ef68973 100644 --- a/tests/unit_tests/test_token_capture_ledger.py +++ b/tests/unit_tests/test_token_capture_ledger.py @@ -619,9 +619,8 @@ async def test_retiring_again_does_not_touch_fences_or_sync_the_directory(tmp_pa with patch.object(store, "_fsync_ledger_root", wraps=store._fsync_ledger_root) as fsync_root: result = await store.retire(["r1"]) - await store.delete(["never-recorded"]) - # A retried batch changes nothing on disk, so it writes no metadata and syncs nothing. + # A retried retire changes nothing on disk, so it writes no metadata and syncs nothing. assert result == {"removed": [], "absent": ["r1"]} assert fence.stat().st_mtime_ns == fence_mtime assert fsync_root.call_count == 0 @@ -712,3 +711,16 @@ async def test_removing_ledgers_forgets_their_cached_rows(tmp_path, method): assert store._ledger_cache == {} # A cache that also tracks its size must account for every removal, or it slowly fills with nothing. assert getattr(store, "_ledger_cache_weight", 0) == 0 + + +@pytest.mark.asyncio +async def test_file_store_delete_syncs_every_non_empty_batch(tmp_path): + """An overlapping delete may have removed the files without syncing yet, so a delete that finds nothing still syncs.""" + store = FileLineageStore(tmp_path) + await _record_call_1(store) + + with patch.object(store, "_fsync_ledger_root", wraps=store._fsync_ledger_root) as fsync_root: + await store.delete(["r1"]) + await store.delete(["r1", "never-recorded"]) + + assert fsync_root.call_count == 2 From c5489510944c818591ab9d93904cc029f9e96a03 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Fri, 9 Oct 2026 18:37:55 -0700 Subject: [PATCH 6/6] fix(token-capture): sync every ledger retire and keep a reused ID's ledger A retire that found every fence already in place and no ledger to delete returned without syncing. If an overlapping retire had created those fences without syncing yet, a crash could lose them after this call had reported success, and a late write could start a new ledger. Every non-empty retire batch now syncs its fences before removing anything and before returning; a retried batch still writes nothing. A retire removes ledgers after releasing and retaking each rollout's lock. In between, a delete could clear the fence and a reused rollout ID record a new ledger, which the retire then deleted. Removal now re-checks the fence under the lock and leaves an unfenced rollout's ledger alone. Signed-off-by: Ananth Subramaniam --- nemo_gym/token_id_capture/lineage.py | 22 +++++++------ tests/unit_tests/test_token_capture_ledger.py | 31 +++++++++++++++++-- 2 files changed, 40 insertions(+), 13 deletions(-) diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py index 11af24595db..c54f5b57e63 100644 --- a/nemo_gym/token_id_capture/lineage.py +++ b/nemo_gym/token_id_capture/lineage.py @@ -1040,18 +1040,15 @@ def _retire(self, rollout_ids: Sequence[str]) -> RolloutRemovalPayload: from nemo_gym.token_id_capture.store import validate_rollout_ids rollout_ids = validate_rollout_ids(rollout_ids) - created_fence = has_ledger = False for rollout_id in rollout_ids: with self._locked(rollout_id): fence = self._retired_path(rollout_id) if not fence.exists(): fence.touch() - created_fence = True - has_ledger = has_ledger or self._ledger_path(rollout_id).exists() - # Make every fence durable before deleting any ledger, so a crash in between leaves a fence. That - # includes fences that already existed: a retire that crashed earlier may have created them without - # syncing. A retried batch with nothing left to delete changes nothing, so it syncs nothing. - if created_fence or has_ledger: + # Make every fence durable before deleting any ledger, and before returning, so a crash leaves a fence. + # That includes fences that already existed: an overlapping or crashed retire may have created them + # without syncing. A retried batch writes nothing, but still syncs once. + if rollout_ids: self._fsync_ledger_root() return self._remove(rollout_ids) @@ -1068,6 +1065,11 @@ def _remove(self, rollout_ids: list[str], *, unretire: bool = False) -> RolloutR changed = False for rollout_id in rollout_ids: with self._locked(rollout_id): + if not unretire and not self._retired_path(rollout_id).exists(): + # A delete cleared the fence after this retire wrote it, and the rollout ID may already be + # in use again: its ledger now belongs to the new attempt, so leave it. + absent.append(rollout_id) + continue try: self._ledger_path(rollout_id).unlink() removed.append(rollout_id) @@ -1081,9 +1083,9 @@ def _remove(self, rollout_ids: list[str], *, unretire: bool = False) -> RolloutR except FileNotFoundError: pass self._ledger_cache_pop(rollout_id) - # A retire with nothing left to remove syncs nothing, since its fences are already durable. A delete - # always syncs a non-empty batch: an overlapping delete may have removed the files without syncing yet, - # and a caller that reuses the rollout ID must not see the old fence or ledger come back after a crash. + # A retire with nothing left to remove has already synced its fences. A delete always syncs a non-empty + # batch: an overlapping delete may have removed the files without syncing yet, and a caller that reuses + # the rollout ID must not see the old fence or ledger come back after a crash. if changed or (unretire and rollout_ids): self._fsync_ledger_root() return {"removed": removed, "absent": absent} diff --git a/tests/unit_tests/test_token_capture_ledger.py b/tests/unit_tests/test_token_capture_ledger.py index 6b85ef68973..ffbb766f936 100644 --- a/tests/unit_tests/test_token_capture_ledger.py +++ b/tests/unit_tests/test_token_capture_ledger.py @@ -610,7 +610,7 @@ async def test_file_store_rejects_the_whole_batch_on_an_invalid_rollout_id(tmp_p @pytest.mark.asyncio -async def test_retiring_again_does_not_touch_fences_or_sync_the_directory(tmp_path): +async def test_retiring_again_writes_nothing_but_syncs_once(tmp_path): store = FileLineageStore(tmp_path) await _record_call_1(store) await store.retire(["r1"]) @@ -620,10 +620,11 @@ async def test_retiring_again_does_not_touch_fences_or_sync_the_directory(tmp_pa with patch.object(store, "_fsync_ledger_root", wraps=store._fsync_ledger_root) as fsync_root: result = await store.retire(["r1"]) - # A retried retire changes nothing on disk, so it writes no metadata and syncs nothing. + # A retried retire writes no metadata. It still syncs once: an overlapping retire may have written the + # fence without syncing it yet, and this call must not report success before the fence is durable. assert result == {"removed": [], "absent": ["r1"]} assert fence.stat().st_mtime_ns == fence_mtime - assert fsync_root.call_count == 0 + assert fsync_root.call_count == 1 @pytest.mark.asyncio @@ -724,3 +725,27 @@ async def test_file_store_delete_syncs_every_non_empty_batch(tmp_path): await store.delete(["r1", "never-recorded"]) assert fsync_root.call_count == 2 + + +@pytest.mark.asyncio +async def test_retire_does_not_remove_a_ledger_recreated_after_a_delete_between_its_phases(tmp_path): + """Between fencing and removing, a delete can clear the fence and a reused rollout ID can record again.""" + store = FileLineageStore(tmp_path) + await _record_call_1(store) + real_fsync = store._fsync_ledger_root + reused = [] + + def fsync_root(): + real_fsync() + if not reused: + reused.append(True) + store._delete(["r1"]) + store._record(_commit(_call_record("c1"), [USER_1], [ASSISTANT_1])) + + with patch.object(store, "_fsync_ledger_root", fsync_root): + await store.retire(["r1"]) + + # The new attempt's ledger survives, unfenced. + assert not (tmp_path / "r1.lineage.retired").exists() + manifest = RolloutManifest.model_validate(await store.manifest("r1")) + assert [record.model_call_id for record in manifest.records] == ["c1"]