fix(checkpoint): leave the calls a checkpoint abandoned out of the capture ledger

nemo_gym/_checkpoint/model.py

@@ -134,6 +134,10 @@ def mark_started(self, rollout_id: str) -> None: ...
 
     async def delete_unstarted(self, rollout_ids: Sequence[str]) -> dict: ...
 
+    def discard_calls(self, rollout_id: str, model_call_ids: Sequence[str]) -> int:
+        """Drop the rows of these calls from a live ledger; return how many were dropped."""
+        ...
+
 
 class EpisodeCut(BaseModel):
     """A durable cut of one undelivered call, keyed by the episode that will continue it."""
@@ -174,6 +178,10 @@ class GateSnapshot(BaseModel):
     # Undelivered model calls (model_call_id -> capture key): the commit leaves their ledger rows out.
     undelivered: dict[str, str]
     cuts: list[EpisodeCut]
+    # Calls whose client went away while the checkpoint held their response (model_call_id -> capture key).
+    # The agent that issued them was interrupted at the checkpoint and re-issues the call afterwards, so the
+    # commit leaves their rows out like an undelivered call's, and the resume discards them from the live ledger.
+    abandoned: dict[str, str] = {}
 
 
 class RestoredCuts(Protocol):
@@ -355,6 +363,8 @@ def __init__(
         self.served = False
         # Restored attempts no call has reached here yet, by capture key.
         self.restored_targets: set[str] = set()
+        # Calls abandoned by their client while a checkpoint held their response; see GateSnapshot.abandoned.
+        self.abandoned: dict[str, str] = {}
 
     async def admit(self, capture_key: Optional[str]) -> None:
         """Wait while a checkpoint is open, then admit the call unless its attempt was retired.
@@ -419,6 +429,11 @@ def enter(self, capture_key: str) -> _Ticket:
         return ticket
 
     async def exit(self, ticket: _Ticket) -> None:
+        if not self.accepting and not ticket.response_started and ticket.capture is not None:
+            # The client gave up on a held response: a harness interrupted at the checkpoint. Whatever the
+            # ledger recorded for the call (a row, or a failure row for a request that ended without
+            # coordinates) describes a call no trajectory contains.
+            self.abandoned[ticket.capture.model_call_id] = ticket.capture_key
         self.tickets.discard(ticket)
         try:
             await self.restored_cuts.settle(ticket)
@@ -525,8 +540,14 @@ def snapshot(self) -> GateSnapshot:
                 for ticket in undelivered
                 if ticket.cut_record is not None
             ],
+            abandoned=dict(self.abandoned),
         )
 
+    def take_abandoned(self) -> dict[str, str]:
+        """The calls abandoned during the checkpoint that is ending, cleared for the next one."""
+        abandoned, self.abandoned = self.abandoned, {}
+        return abandoned
+
     async def _cut_all(self, request: CheckpointRequest) -> None:
         by_backend: dict[str, list[_Ticket]] = {}
         for ticket in self.tickets:
@@ -625,7 +646,7 @@ def export_model_records(
     if episode_ids is None:
         raise ControlError("model commit needs the episode_ids the controller continues from this checkpoint")
     snapshots = list(snapshots)
-    undelivered = {call_id for snapshot in snapshots for call_id in snapshot.undelivered}
+    undelivered = {call_id for snapshot in snapshots for call_id in (*snapshot.undelivered, *snapshot.abandoned)}
     # One cut per episode: when a client retried a call the server was still running,
     # both calls are undelivered and cut, but only the latest is still awaited and will be re-issued after a restore.
     latest: dict[str, EpisodeCut] = {}
@@ -672,6 +693,29 @@ def staging_keys_by_episode(records: Iterable[CheckpointRecord]) -> dict[str, li
     return {record.episode_id.capture_key: retained_staging_keys([record]) for record in records}
 
 
+def discard_abandoned_calls(ledger: CheckpointableLedger, abandoned: Mapping[str, str]) -> None:
+    """Drop the live ledger rows of calls a checkpoint abandoned (``GateSnapshot.abandoned``).
+
+    The commit already left them out; this is for the episode that continues without a crash, whose live
+    ledger would otherwise carry a row, or a failure row, for a call its harness never consumed and re-issued
+    after the resume. Best effort: a ledger without ``discard_calls`` keeps the rows and the finalizer sees them.
+    """
+    discard = getattr(ledger, "discard_calls", None)
+    if discard is None:
+        LOGGER.warning("capture ledger %s cannot discard %d abandoned call(s)", type(ledger).__name__, len(abandoned))
+        return
+    by_key: dict[str, list[str]] = {}
+    for model_call_id, capture_key in abandoned.items():
+        by_key.setdefault(capture_key, []).append(model_call_id)
+    for capture_key, model_call_ids in by_key.items():
+        try:
+            dropped = discard(capture_key, model_call_ids)
+        except Exception:
+            LOGGER.warning("could not discard abandoned calls %s of %s", model_call_ids, capture_key, exc_info=True)
+            continue
+        LOGGER.info("discarded %d abandoned call row(s) of %s after the checkpoint", dropped, capture_key)
+
+
 async def retire_ledgers(ledger: Optional[CheckpointableLedger], episode_ids: Iterable[EpisodeId]) -> None:
     """Retire the ledgers of each episode and every earlier attempt of it, in one batch.
 
@@ -789,7 +833,10 @@ async def close_admission(self, request: CheckpointRequest) -> None:
         await self.gate.close(request)
 
     async def open_admission(self) -> None:
+        abandoned = self.gate.take_abandoned()
         self.gate.open()
+        if abandoned and self.ledger is not None:
+            await asyncio.to_thread(discard_abandoned_calls, self.ledger, abandoned)
 
     async def mark_retired(self, episode_ids: list[EpisodeId]) -> None:
         await super().mark_retired(episode_ids)

nemo_gym/_checkpoint/model_workers.py

@@ -37,6 +37,7 @@
     PolicyGate,
     _Ticket,
     covers,
+    discard_abandoned_calls,
     export_model_records,
     import_model_records,
     merge_reports,
@@ -284,6 +285,8 @@ def __init__(
         ledger: Optional[CheckpointableLedger] = None,
     ) -> None:
         super().__init__(socket_path=socket_path, on_coordinator_lost=on_coordinator_lost)
+        # This worker's view of the shared capture ledger, to drop the rows of calls it saw abandoned.
+        self.ledger = ledger
         self.restored_cuts = CoordinatedRestoredCuts(self)
         self.gate = PolicyGate(
             server_name=server_name,
@@ -308,7 +311,10 @@ async def close_local(self, request: CheckpointRequest) -> None:
 
     async def open_local(self, body: dict[str, Any]) -> None:
         self.adopt(body)
+        abandoned = self.gate.take_abandoned()
         self.gate.open()
+        if abandoned and self.ledger is not None:
+            await asyncio.to_thread(discard_abandoned_calls, self.ledger, abandoned)
 
     def mark_local(self, episode_ids: list[EpisodeId]) -> None:
         super().mark_local(episode_ids)

nemo_gym/token_id_capture/lineage.py

@@ -1237,6 +1237,31 @@ def export_rows(self, rollout_id: str) -> list[dict]:
         with self._locked(rollout_id):
             return list(self._read(rollout_id))
 
+    def discard_calls(self, rollout_id: str, model_call_ids: Sequence[str]) -> int:
+        """Drop the rows of these calls from a live ledger and return how many were dropped.
+
+        For calls a checkpoint abandoned: the harness that issued them was interrupted before it received the
+        response and re-issues the call, so their rows (or failure rows) describe nothing the trajectory
+        contains. The file is rewritten atomically under the ledger's lock; other rows keep their order.
+        """
+        drop = set(model_call_ids)
+        with self._locked(rollout_id):
+            rows = list(self._read(rollout_id))
+            kept = [row for row in rows if row.get("model_call_id") not in drop]
+            if len(kept) == len(rows):
+                return 0
+            path = self._ledger_path(rollout_id)
+            payload = b"".join(json.dumps(row, sort_keys=True, separators=(",", ":")).encode() + b"\n" for row in kept)
+            temporary = path.with_name(f".{path.name}.discard")
+            with temporary.open("wb") as handle:
+                handle.write(payload)
+                handle.flush()
+                os.fsync(handle.fileno())
+            os.replace(temporary, path)
+            with self._cache_guard:
+                self._ledger_cache.pop(rollout_id, None)
+            return len(rows) - len(kept)
+
     def import_rows(self, rollout_id: str, rows: list[dict]) -> None:
         """Install checkpointed rows as the complete ledger of ``rollout_id``; see ``import_rows_many``."""
         self.import_rows_many({rollout_id: rows})

tests/unit_tests/test_checkpoint_model.py

@@ -620,6 +620,41 @@ async def test_undelivered_rows_are_excluded_on_any_capture_path(tmp_path: Path)
     assert '"c1"' in rows and '"c2"' not in rows
 
 
+async def test_a_call_abandoned_while_held_is_left_out_of_the_commit_and_dropped_at_resume(tmp_path: Path) -> None:
+    """A harness interrupted at the checkpoint gives up on its held call. The ledger recorded the call (here a
+    row; in the external path often a failure row for a request that ended without coordinates). The commit
+    leaves it out, and the resume drops it from the live ledger, so neither the restored attempt nor the
+    continued one carries a call no trajectory contains."""
+    ledger = FileLineageStore(tmp_path / "store")
+    await ledger.record(_commit(_call_record("c1"), [USER_1], [ASSISTANT_1], rollout_id="r", staging_chain=("r/c1",)))
+    participant = PolicyModelParticipant(ledger)
+    controller = ParticipantControlPlane(participant, instance_name="policy", lease_grace_seconds=60)
+    release = asyncio.Event()
+    app = _captured_app(participant, ledger, release)
+    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://m") as client:
+        call = asyncio.create_task(client.post("/ng-rollout/r/v1/chat/completions"))
+        await asyncio.sleep(0.05)
+        await controller.prepare(CheckpointRequest(**control()))
+        # The client (the interrupted harness) goes away while the response is held.
+        call.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await call
+        await asyncio.sleep(0.05)
+        live_rows = [row["model_call_id"] for row in ledger.export_rows("r")]
+        snapshot = participant.gate.snapshot()
+        commit = await controller.commit(_commit_request("c1", tmp_path / "ckpt", [{"rollout_id": "r"}]))
+        await controller.resume(CheckpointRequest(**control()))
+        await asyncio.sleep(0.05)
+        after_resume = [row["model_call_id"] for row in ledger.export_rows("r")]
+
+    rows = (tmp_path / "ckpt" / "gym" / "model" / "policy" / commit["manifest"]["records_file"]).read_text()
+    assert live_rows == ["c1", "c2"], "the abandoned call's row was in the live ledger"
+    assert snapshot.abandoned == {"c2": "r"} and participant.gate.tickets == set()
+    assert '"c1"' in rows and '"c2"' not in rows, "the commit left the abandoned call out"
+    assert after_resume == ["c1"], "the resume dropped it from the live ledger"
+    assert participant.gate.abandoned == {}, "cleared for the next checkpoint"
+
+
 async def test_each_model_server_continues_only_its_own_restored_cut() -> None:
     first, second = PolicyModelParticipant(), PolicyModelParticipant()
     continuation = GenerationCutContinuation(