From 856cbc2a088b96e16778b81079d5fb34eb8fe50a Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 1 Sep 2026 11:05:07 -0700 Subject: [PATCH 1/5] feat(checkpoint): add policy model participant Pause policy generation at a quiescent cut and checkpoint token-free capture custody while leaving framework-owned staged token arrays to TransferQueue. Signed-off-by: Ananth Subramaniam --- nemo_gym/_checkpoint/__init__.py | 26 ++ nemo_gym/_checkpoint/ledger.py | 397 ++++++++++++++++ nemo_gym/_checkpoint/model_admission.py | 198 ++++++++ nemo_gym/base_responses_api_model.py | 83 +++- nemo_gym/token_id_capture/lineage.py | 5 + tests/unit_tests/test_checkpoint_admission.py | 439 ++++++++++++++++++ tests/unit_tests/test_checkpoint_ledger.py | 249 ++++++++++ 7 files changed, 1392 insertions(+), 5 deletions(-) create mode 100644 nemo_gym/_checkpoint/ledger.py create mode 100644 nemo_gym/_checkpoint/model_admission.py create mode 100644 tests/unit_tests/test_checkpoint_admission.py create mode 100644 tests/unit_tests/test_checkpoint_ledger.py diff --git a/nemo_gym/_checkpoint/__init__.py b/nemo_gym/_checkpoint/__init__.py index 2195f18c38e..42d05196023 100644 --- a/nemo_gym/_checkpoint/__init__.py +++ b/nemo_gym/_checkpoint/__init__.py @@ -63,6 +63,21 @@ WorkerRegistrationError, build_coordinator_control_app, ) +from nemo_gym._checkpoint.ledger import ( + LEDGER_MANIFEST_NAME, + MODEL_CHECKPOINT_URL_PREFIX, + MODEL_LEDGER_SUBDIR, + CaptureLedgerCheckpointer, + CheckpointableCaptureLedger, + LedgerMismatchError, + LedgerNotCheckpointableError, + LedgerNotQuiescentError, + install_model_checkpoint, +) +from nemo_gym._checkpoint.model_admission import ( + NotPolicyInstanceError, + install_model_admission, +) from nemo_gym._checkpoint.model_control_contracts import ( MODEL_ADMISSION_URL_PREFIX, GenerationCutAck, @@ -84,7 +99,10 @@ "CONTROL_SCHEMA_VERSION", "CONTROL_URL_PREFIX", "GATED_MODEL_ROUTE_SUFFIXES", + "LEDGER_MANIFEST_NAME", "MODEL_ADMISSION_URL_PREFIX", + "MODEL_CHECKPOINT_URL_PREFIX", + "MODEL_LEDGER_SUBDIR", "PLANE_HEADER", "AdmissionLimiter", "AdmissionMiddleware", @@ -117,11 +135,19 @@ "ControlError", "ControlFence", "Deadline", + "CaptureLedgerCheckpointer", + "CheckpointableCaptureLedger", "InvalidPhaseError", + "LedgerMismatchError", + "LedgerNotCheckpointableError", + "LedgerNotQuiescentError", "MultiProcessCapability", + "NotPolicyInstanceError", "StaleAttemptError", "StaleCheckpointError", "install_control_plane", + "install_model_admission", + "install_model_checkpoint", "mark_current_generation_safe", "mark_current_generation_started", "multi_process_capability_from_num_workers", diff --git a/nemo_gym/_checkpoint/ledger.py b/nemo_gym/_checkpoint/ledger.py new file mode 100644 index 00000000000..6f06f992ca5 --- /dev/null +++ b/nemo_gym/_checkpoint/ledger.py @@ -0,0 +1,397 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Token-free capture-ledger checkpoint commit and restore. + +PR #2872 separates token custody from token storage. The generation worker +stages token arrays in the training framework's TransferQueue. The Gym model +server stores token-free lineage rows that identify those staged entries. +This participant checkpoints only those lineage rows. TransferQueue owns its +own checkpoint and restore. + +Three properties make the copy a checkpoint rather than a backup: + +- **Tombstone exclusion.** A rollout attempt force-closed at the prepare + deadline must not restore: its rows describe an execution the restored run + replaces with a fresh dispatch. Commit skips tombstoned attempts and + records the tombstones in the manifest so the restored server re-installs + the fence before serving anything. +- **Manifest-last ordering.** Every ledger file is written and fsynced + before the manifest appears (temporary name, fsync, rename). A commit that + died partway leaves no manifest, and restore refuses the directory instead + of installing a torn ledger. +- **Digest verification.** The manifest records each rollout file's SHA-256. + Restore verifies every installed file against it, so silent corruption in + transit fails loudly at restore instead of surfacing as wrong training + data later. +""" + +import asyncio +import hashlib +import json +import os +import shutil +import tempfile +from pathlib import Path +from typing import Any, Callable, Literal, Optional, Protocol, runtime_checkable + +from fastapi import FastAPI, Header + +from nemo_gym._checkpoint.admission import AdmissionLimiter +from nemo_gym._checkpoint.control import ( + CONTROL_URL_PREFIX, + CheckpointControlRequest, + CheckpointPhase, + ControlError, + ControlFence, +) +from nemo_gym._checkpoint.model_admission import NotPolicyInstanceError +from nemo_gym.rollout_correlation import capture_key_for +from nemo_gym.token_id_capture.control_routes import require_control_auth +from nemo_gym.token_id_capture.protocols import CaptureLedger + + +MODEL_CHECKPOINT_URL_PREFIX = f"{CONTROL_URL_PREFIX}/model-checkpoint" +MODEL_LEDGER_SUBDIR = "model-ledger" +LEDGER_MANIFEST_NAME = "manifest.json" +LEDGER_SCHEMA_VERSION = 1 + +# FileLineageStore writes one token-free custody file per rollout. +# Lock files and token-store files are not part of this participant. +_LEDGER_SUFFIX = ".lineage.jsonl" + + +class LedgerMismatchError(ControlError): + """The checkpoint directory does not match its manifest. + + A missing manifest means the commit tore partway; a digest mismatch + means a file changed after commit. Either way the ledger must not be + installed: restored custody would refer to rows that do not exist as + committed. + """ + + code = "ledger_mismatch" + + +class LedgerNotCheckpointableError(ControlError): + """The configured capture backend has no checkpoint lifecycle.""" + + code = "ledger_not_checkpointable" + + +class LedgerNotQuiescentError(ControlError): + """The model participant still has accepted generation requests.""" + + code = "ledger_not_quiescent" + + +@runtime_checkable +class CheckpointableCaptureLedger(CaptureLedger, Protocol): + """Optional lifecycle implemented by framework-owned capture backends. + + The backend snapshots token-free custody and its private parent-resolution + state. The framework checkpoints staged token arrays separately. + """ + + async def checkpoint_capture_ledger( + self, + checkpoint_dir: Path, + *, + checkpoint_id: str, + tombstones: tuple[tuple[str, int], ...], + ) -> dict[str, Any]: ... + + async def restore_capture_ledger(self, checkpoint_dir: Path) -> dict[str, Any]: ... + + +def _file_digest(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _copy_fsynced(source: Path, target: Path) -> None: + with tempfile.NamedTemporaryFile(dir=target.parent, prefix=".ledger-", delete=False) as handle: + temporary = Path(handle.name) + try: + with source.open("rb") as src: + shutil.copyfileobj(src, handle) + handle.flush() + os.fsync(handle.fileno()) + except BaseException: + temporary.unlink(missing_ok=True) + raise + os.replace(temporary, target) + + +def _fsync_dir(path: Path) -> None: + descriptor = os.open(path, os.O_RDONLY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +class CaptureLedgerCheckpointer: + """Commit and restore one token-capture store directory.""" + + def __init__(self, store_root: Path) -> None: + self.store_root = Path(store_root) + + def _rollout_ids(self) -> list[str]: + return sorted(path.name[: -len(_LEDGER_SUFFIX)] for path in self.store_root.glob(f"*{_LEDGER_SUFFIX}")) + + def commit(self, checkpoint_dir: Path, *, checkpoint_id: str, tombstones: list[tuple[str, int]]) -> dict[str, Any]: + """Copy the ledger into ``checkpoint_dir``; the caller has already drained. + + The store must be quiescent (admission paused) when this runs: the + copy takes no locks because nothing may be writing. + """ + ledger_dir = Path(checkpoint_dir) / MODEL_LEDGER_SUBDIR + if (ledger_dir / LEDGER_MANIFEST_NAME).exists(): + raise LedgerMismatchError(f"checkpoint ledger already committed at {ledger_dir}") + ledger_dir.mkdir(parents=True, exist_ok=True) + fenced = {capture_key_for(rollout_id, attempt_index) for rollout_id, attempt_index in tombstones} + + rollouts: dict[str, dict[str, Any]] = {} + excluded = 0 + total_rows = 0 + for rollout_id in self._rollout_ids(): + if rollout_id in fenced: + excluded += 1 + continue + files: dict[str, str] = {} + rows = 0 + source = self.store_root / f"{rollout_id}{_LEDGER_SUFFIX}" + target = ledger_dir / source.name + _copy_fsynced(source, target) + files[source.name] = _file_digest(target) + rows = sum(1 for line in target.read_bytes().splitlines() if line.strip()) + rollouts[rollout_id] = {"files": files, "rows": rows} + total_rows += rows + _fsync_dir(ledger_dir) + + manifest = { + "schema_version": LEDGER_SCHEMA_VERSION, + "checkpoint_id": checkpoint_id, + "rollouts": rollouts, + "tombstones": [ + {"rollout_id": rollout_id, "attempt_index": attempt} for rollout_id, attempt in sorted(tombstones) + ], + } + payload = json.dumps(manifest, sort_keys=True, indent=1).encode() + with tempfile.NamedTemporaryFile(dir=ledger_dir, prefix=".manifest-", delete=False) as handle: + temporary = Path(handle.name) + try: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + except BaseException: + temporary.unlink(missing_ok=True) + raise + os.replace(temporary, ledger_dir / LEDGER_MANIFEST_NAME) + _fsync_dir(ledger_dir) + + return { + "rollouts": len(rollouts), + "rows": total_rows, + "excluded_tombstoned": excluded, + "manifest_digest": hashlib.sha256(payload).hexdigest(), + } + + def restore(self, checkpoint_dir: Path) -> dict[str, Any]: + """Install a committed ledger into this store root and verify it.""" + ledger_dir = Path(checkpoint_dir) / MODEL_LEDGER_SUBDIR + manifest_path = ledger_dir / LEDGER_MANIFEST_NAME + if not manifest_path.exists(): + raise LedgerMismatchError( + f"no ledger manifest at {manifest_path}; the commit tore partway or never ran, " + f"so this directory must not be installed" + ) + manifest = json.loads(manifest_path.read_text()) + if manifest.get("schema_version", 0) > LEDGER_SCHEMA_VERSION: + raise LedgerMismatchError( + f"ledger manifest schema_version {manifest.get('schema_version')} is newer than this " + f"reader ({LEDGER_SCHEMA_VERSION})" + ) + + expected_names = {name for metadata in manifest["rollouts"].values() for name in metadata["files"]} + existing_names = {path.name for path in self.store_root.glob(f"*{_LEDGER_SUFFIX}")} + unexpected = existing_names - expected_names + if unexpected: + raise LedgerMismatchError( + "restore requires a fresh capture-ledger namespace; " + f"found files absent from the checkpoint: {sorted(unexpected)}" + ) + + # Validate the complete source before changing the live namespace. + validated: list[tuple[Path, str]] = [] + total_rows = 0 + for rollout_id, meta in manifest["rollouts"].items(): + for name, digest in meta["files"].items(): + source = ledger_dir / name + if not source.exists() or _file_digest(source) != digest: + raise LedgerMismatchError( + f"ledger file {name} for rollout {rollout_id!r} is missing or does not match " + f"its committed digest; refusing to install a corrupted ledger" + ) + validated.append((source, name)) + total_rows += int(meta.get("rows", 0)) + + self.store_root.mkdir(parents=True, exist_ok=True) + for source, name in validated: + _copy_fsynced(source, self.store_root / name) + _fsync_dir(self.store_root) + + return { + "rollouts": len(manifest["rollouts"]), + "rows": total_rows, + "checkpoint_id": manifest.get("checkpoint_id"), + "tombstones": list(manifest.get("tombstones", ())), + } + + +class ModelCheckpointCommitRequest(CheckpointControlRequest): + checkpoint_dir: str + + +class ModelCheckpointRestoreRequest(CheckpointControlRequest): + checkpoint_dir: str + + +def install_model_checkpoint( + app: FastAPI, + *, + fence: ControlFence, + limiter: AdmissionLimiter, + ledger_provider: Callable[[], Optional[CaptureLedger]], + file_ledger_root_provider: Callable[[], Optional[Path]], + instance_role: Literal["policy", "auxiliary"], + auth_token: str, +) -> None: + """Register ``/ng-control/v1/model-checkpoint`` on a model-server app. + + Commit requires the prepared (drained) phase; restore runs on a freshly + started server and leaves it paused, so nothing serves until the + coordinator has restored every component and explicitly resumes. + """ + + def _require_policy() -> None: + if instance_role != "policy": + raise NotPolicyInstanceError( + "this model-server instance is auxiliary (judge or simulator traffic); " + "it produces no training tokens and has no capture ledger to checkpoint" + ) + + def _require_quiescent() -> None: + counts = limiter.counts() + if counts["state"] != "paused" or counts["inflight_total"] != 0: + raise LedgerNotQuiescentError( + "capture-ledger commit requires paused admission and zero in-flight generation requests" + ) + + async def _with_ledger( + checkpoint_dir: Path, + *, + checkpoint_id: str, + operation: Literal["commit", "restore"], + ) -> dict[str, Any]: + ledger = ledger_provider() + if isinstance(ledger, CheckpointableCaptureLedger): + if operation == "commit": + return await ledger.checkpoint_capture_ledger( + checkpoint_dir, + checkpoint_id=checkpoint_id, + tombstones=tuple(limiter.tombstones()), + ) + return await ledger.restore_capture_ledger(checkpoint_dir) + + file_root = file_ledger_root_provider() + if file_root is None: + raise LedgerNotCheckpointableError( + "the configured CaptureLedger must implement CheckpointableCaptureLedger; " + "Gym cannot infer how to snapshot a framework-owned backend" + ) + checkpointer = CaptureLedgerCheckpointer(file_root) + if operation == "commit": + return await _run_sync( + lambda: checkpointer.commit( + checkpoint_dir, + checkpoint_id=checkpoint_id, + tombstones=limiter.tombstones(), + ) + ) + return await _run_sync(lambda: checkpointer.restore(checkpoint_dir)) + + @app.post(f"{MODEL_CHECKPOINT_URL_PREFIX}/commit") + async def model_checkpoint_commit( + body: ModelCheckpointCommitRequest, + authorization: str | None = Header(default=None), + ) -> dict[str, Any]: + require_control_auth(authorization, auth_token) + _require_policy() + + async def run() -> dict[str, Any]: + _require_quiescent() + return await _with_ledger( + Path(body.checkpoint_dir), + checkpoint_id=body.checkpoint_id, + operation="commit", + ) + + return await fence.run_operation( + body.checkpoint_id, + "model-checkpoint/commit", + allowed_phases=frozenset({CheckpointPhase.PREPARED}), + phase_during=CheckpointPhase.COMMITTING, + phase_after=CheckpointPhase.COMMITTED_PAUSED, + run=run, + ) + + @app.post(f"{MODEL_CHECKPOINT_URL_PREFIX}/restore") + async def model_checkpoint_restore( + body: ModelCheckpointRestoreRequest, + authorization: str | None = Header(default=None), + ) -> dict[str, Any]: + require_control_auth(authorization, auth_token) + _require_policy() + + async def run() -> dict[str, Any]: + # The restored server boots into the paused state: nothing may be + # admitted until every component is restored and the coordinator + # explicitly resumes. + limiter.close() + result = await _with_ledger( + Path(body.checkpoint_dir), + checkpoint_id=body.checkpoint_id, + operation="restore", + ) + for tombstone in result["tombstones"]: + limiter.install_tombstone(tombstone["rollout_id"], tombstone["attempt_index"]) + return result + + return await fence.run_operation( + body.checkpoint_id, + "model-checkpoint/restore", + allowed_phases=frozenset({CheckpointPhase.IDLE}), + phase_during=CheckpointPhase.RESTORING, + phase_after=CheckpointPhase.RESTORED_PAUSED, + run=run, + ) + + +async def _run_sync(operation: Callable[[], dict[str, Any]]) -> dict[str, Any]: + return await asyncio.to_thread(operation) diff --git a/nemo_gym/_checkpoint/model_admission.py b/nemo_gym/_checkpoint/model_admission.py new file mode 100644 index 00000000000..b9a1ed62fbb --- /dev/null +++ b/nemo_gym/_checkpoint/model_admission.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Model-server admission control routes. + +The checkpoint coordinator pauses only policy model-server instances: the +instances whose generations produce training tokens. Judge and simulator +instances keep serving environment traffic through the whole checkpoint so +accepted operations can finish draining; pausing them would deadlock the +drain. An instance declares its role in configuration and a pause sent to an +auxiliary instance is rejected instead of silently honored. + +Pause closes admission and returns promptly with the drain still in +progress; ``status`` long-polls until the server reports ``paused`` (nothing +in flight). ``abort_inflight`` is the deadline escape hatch: it tombstones a +rollout attempt that cannot finish draining in time so the checkpoint can +proceed and the restored run dispatches a replacement attempt. +""" + +from typing import Any, Literal, Optional + +from fastapi import FastAPI, Header, Query + +from nemo_gym._checkpoint.admission import AdmissionLimiter +from nemo_gym._checkpoint.control import ( + AdmissionState, + CheckpointPhase, + ControlError, + ControlFence, + Deadline, +) +from nemo_gym._checkpoint.model_control_contracts import ( + MODEL_ADMISSION_URL_PREFIX, + ModelAbortInflightRequest, + ModelAdmissionPauseRequest, + ModelAdmissionResumeRequest, +) +from nemo_gym.token_id_capture.control_routes import require_control_auth + + +class NotPolicyInstanceError(ControlError): + code = "not_a_policy_instance" + + +def install_model_admission( + app: FastAPI, + *, + limiter: AdmissionLimiter, + fence: ControlFence, + instance_role: Literal["policy", "auxiliary"], + auth_token: str, +) -> None: + """Register ``/ng-control/v1/model-admission`` on a model-server app. + + The routes are fenced by checkpoint id (idempotent replay, stale-id + rejection) and remain reachable while the data plane is paused. The + single-worker deployment reports itself as its only worker; the + multi-worker coordinator replaces these numbers with aggregated ones. + """ + + def _require_policy() -> None: + if instance_role != "policy": + raise NotPolicyInstanceError( + "this model-server instance is auxiliary (judge or simulator traffic); " + "only policy instances participate in checkpoint admission control" + ) + + def _workers() -> dict[str, int]: + return {"acknowledged": 1, "expected": 1} + + @app.post(f"{MODEL_ADMISSION_URL_PREFIX}/pause") + async def model_admission_pause( + body: ModelAdmissionPauseRequest, + authorization: str | None = Header(default=None), + ) -> dict[str, Any]: + require_control_auth(authorization, auth_token) + _require_policy() + + async def run() -> dict[str, Any]: + limiter.close() + counts = limiter.counts() + return { + "state": counts["state"], + "workers": _workers(), + "inflight_total": counts["inflight_total"], + "waiters_total": counts["waiters_total"], + } + + result = await fence.run_operation( + body.checkpoint_id, + "model-admission/pause", + allowed_phases=frozenset({CheckpointPhase.IDLE}), + phase_during=CheckpointPhase.PREPARING, + phase_after=CheckpointPhase.PREPARING, + run=run, + deadline=body, + ) + if result["state"] == AdmissionState.PAUSED.value: + fence.mark_prepared(body.checkpoint_id) + return result + + @app.get(f"{MODEL_ADMISSION_URL_PREFIX}/status") + async def model_admission_status( + checkpoint_id: str = Query(min_length=1, max_length=128, pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$"), + deadline_ts: float = Query(), + wait_state: Optional[Literal["paused"]] = None, + timeout_s: float = Query(default=0.0, ge=0.0), + authorization: str | None = Header(default=None), + ) -> dict[str, Any]: + require_control_auth(authorization, auth_token) + fence.require_phase( + checkpoint_id, + frozenset({CheckpointPhase.PREPARING, CheckpointPhase.PREPARED}), + ) + deadline = Deadline(deadline_ts=deadline_ts) + if wait_state == "paused" and timeout_s > 0: + await limiter.wait_for_drained(min(timeout_s, deadline.remaining())) + counts = limiter.counts() + if counts["state"] == AdmissionState.PAUSED.value: + fence.mark_prepared(checkpoint_id) + return { + "checkpoint_id": checkpoint_id, + "state": counts["state"], + "per_worker": {"0": {"state": counts["state"], "inflight": counts["inflight_total"]}}, + "inflight_total": counts["inflight_total"], + "waiters_total": counts["waiters_total"], + "inflight": counts["inflight"], + "tombstones": [ + {"rollout_id": rollout_id, "attempt_index": attempt} for rollout_id, attempt in limiter.tombstones() + ], + } + + @app.post(f"{MODEL_ADMISSION_URL_PREFIX}/resume") + async def model_admission_resume( + body: ModelAdmissionResumeRequest, + authorization: str | None = Header(default=None), + ) -> dict[str, Any]: + require_control_auth(authorization, auth_token) + _require_policy() + + async def run() -> dict[str, Any]: + limiter.resume() + return {"state": limiter.state.value, "workers": _workers(), "released_waiters": 0} + + return await fence.run_operation( + body.checkpoint_id, + "model-admission/resume", + allowed_phases=frozenset( + { + CheckpointPhase.PREPARING, + CheckpointPhase.PREPARED, + CheckpointPhase.COMMITTED_PAUSED, + CheckpointPhase.RESTORED_PAUSED, + } + ), + phase_during=CheckpointPhase.PREPARED, + phase_after=CheckpointPhase.IDLE, + run=run, + retire_outcome="resumed", + ) + + @app.post(f"{MODEL_ADMISSION_URL_PREFIX}/abort_inflight") + async def model_admission_abort_inflight( + body: ModelAbortInflightRequest, + authorization: str | None = Header(default=None), + ) -> dict[str, Any]: + require_control_auth(authorization, auth_token) + _require_policy() + entry_phase = fence.phase + + async def run() -> dict[str, Any]: + aborted = limiter.abort_inflight(body.rollout_id, body.attempt_index) + counts = limiter.counts() + return { + "state": counts["state"], + "aborted_inflight": len(aborted), + "inflight_total": counts["inflight_total"], + } + + return await fence.run_operation( + body.checkpoint_id, + f"model-admission/abort_inflight:{body.rollout_id}:{body.attempt_index}", + allowed_phases=frozenset({CheckpointPhase.PREPARING, CheckpointPhase.PREPARED}), + phase_during=entry_phase, + phase_after=entry_phase, + run=run, + ) diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index cf7b5b05133..98e1fceb9de 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -37,15 +37,19 @@ from abc import abstractmethod from contextlib import asynccontextmanager from pathlib import Path -from typing import Any, Iterable, Mapping, Optional +from typing import Any, Iterable, Literal, Mapping, Optional from uuid import uuid4 import orjson from fastapi import Body, FastAPI, Request, Response from fastapi.exceptions import RequestValidationError from fastapi.responses import StreamingResponse -from pydantic import BaseModel, Field, ValidationError, model_validator +from pydantic import BaseModel, Field, PrivateAttr, ValidationError, model_validator +from nemo_gym._checkpoint.admission import GATED_MODEL_ROUTE_SUFFIXES, AdmissionLimiter, AdmissionMiddleware +from nemo_gym._checkpoint.control import AdmissionState, ControlCapabilities +from nemo_gym._checkpoint.ledger import install_model_checkpoint +from nemo_gym._checkpoint.model_admission import install_model_admission from nemo_gym.anthropic_converter import AnthropicConverter from nemo_gym.chat_streaming import sanitize_streaming_chat_body, synthesize_chat_completion_sse from nemo_gym.config_types import ROLLOUT_PATH_PREFIX, TOKEN_CAPTURE_PATH_SEGMENT, ModelServerRef @@ -178,7 +182,12 @@ def _orjson_dispatch_response(content: Any) -> Any: class BaseResponsesAPIModelConfig(BaseRunServerInstanceConfig): - pass + # Checkpoint role of this instance. 'policy' instances serve the + # generations that produce training tokens and gate admission during a + # checkpoint. 'auxiliary' instances serve environment traffic (judges, + # user or tool simulators) and never pause: accepted operations must be + # able to finish their nested calls while the policy instance drains. + instance_role: Literal["policy", "auxiliary"] = "policy" class BaseResponsesAPIModel(BaseServer): @@ -187,6 +196,7 @@ class BaseResponsesAPIModel(BaseServer): class SimpleResponsesAPIModel(BaseResponsesAPIModel, SimpleServer): _CONTROL_COMPONENT = "responses_api_models" + _admission_limiter: Optional[AdmissionLimiter] = PrivateAttr(default=None) async def _finalize_served_response(self, response: Any) -> None: """Finalize capture after conversion to the response returned to the client.""" @@ -206,13 +216,14 @@ def setup_webserver(self) -> FastAPI: self.setup_session_middleware(app) self.setup_control_plane(app) capture_config = ModelCallCaptureConfig.model_validate(self.server_client.global_config_dict) - install_model_call_capture( + capture_ledger = install_model_call_capture( app, capture_config, model_server_name=self.config.name, global_config_dict=self.server_client.global_config_dict, num_workers=self.config.num_workers, ) + app.state.nemo_gym_capture_ledger = capture_ledger model_attributes = {"nemo.gym.server.name": self.config.name} app.post("/v1/chat/completions")( @@ -233,8 +244,69 @@ def setup_webserver(self) -> FastAPI: traced_endpoint(GymSpanGroup.MODEL_CALL, "gym.model.messages", self.messages, model_attributes) ) + self.setup_model_admission(app) + return app + def admission_limiter(self) -> AdmissionLimiter: + if self._admission_limiter is None: + self._admission_limiter = AdmissionLimiter() + return self._admission_limiter + + def setup_model_admission(self, app: FastAPI) -> None: + """Register admission control on this model server. + + External token staging provides the authenticated control plane used by + the first checkpoint milestone. + """ + auth_token = self.checkpoint_control_auth_token() + if auth_token is None: + return + install_model_admission( + app, + limiter=self.admission_limiter(), + fence=self.checkpoint_fence(), + instance_role=self.config.instance_role, + auth_token=auth_token, + ) + install_model_checkpoint( + app, + fence=self.checkpoint_fence(), + limiter=self.admission_limiter(), + ledger_provider=lambda: getattr(app.state, "nemo_gym_capture_ledger", None), + file_ledger_root_provider=lambda: ( + ledger.checkpoint_root + if isinstance((ledger := getattr(app.state, "nemo_gym_capture_ledger", None)), FileLineageStore) + else None + ), + instance_role=self.config.instance_role, + auth_token=auth_token, + ) + if self.config.instance_role == "policy": + app.add_middleware( + AdmissionMiddleware, + limiter=self.admission_limiter(), + gated_suffixes=GATED_MODEL_ROUTE_SUFFIXES, + ) + + def checkpoint_control_auth_token(self) -> Optional[str]: + settings = token_id_capture_config(self.server_client.global_config_dict) + if settings is None or not settings.token_id_capture.external_staging: + return None + return settings.token_id_capture.resolve_control_auth_token() + + def control_capabilities(self) -> ControlCapabilities: + capabilities = super().control_capabilities() + capabilities.instance_role = self.config.instance_role + if self.config.instance_role == "policy" and self.checkpoint_control_auth_token() is not None: + capabilities.admission_states = [ + AdmissionState.ACCEPTING, + AdmissionState.DRAINING, + AdmissionState.PAUSED, + ] + capabilities.checkpoint_mode = "export_restore" + return capabilities + @abstractmethod async def chat_completions( self, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body() @@ -1563,7 +1635,7 @@ def install_model_call_capture( model_server_name: str | None = None, global_config_dict: Any = None, num_workers: int | None = None, -) -> None: +) -> Optional[CaptureLedger]: """Install model-call capture middleware. Always strip ``/ng-rollout//...`` before routing. @@ -1652,6 +1724,7 @@ async def _capture_lifespan(application): external_staging=external_staging, token_capture_enabled=capture_settings.enabled if capture_settings is not None else False, ) + return lineage_store if isinstance(lineage_store, CaptureLedger) else None # --- Run-level capture helpers (rollout-collection side) --- diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py index d1e12283f4d..db5744fafb3 100644 --- a/nemo_gym/token_id_capture/lineage.py +++ b/nemo_gym/token_id_capture/lineage.py @@ -741,6 +741,11 @@ def __init__( self._ledger_root.mkdir(parents=True, exist_ok=True) self._ledger_cache: dict[str, tuple[int, int, list[dict]]] = {} + @property + def checkpoint_root(self) -> Path: + """Directory containing token-free custody rows.""" + return self._ledger_root + def _read_locked(self, rollout_id: str): return self._store._locked(rollout_id, shared=True) diff --git a/tests/unit_tests/test_checkpoint_admission.py b/tests/unit_tests/test_checkpoint_admission.py new file mode 100644 index 00000000000..f1be5e85ef0 --- /dev/null +++ b/tests/unit_tests/test_checkpoint_admission.py @@ -0,0 +1,439 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Admission control drains one policy model-server participant. + +The gate closes atomically with admission. +Accepted requests remain counted until their ASGI task exits. +Cancelled attempts are fenced from later admission. +""" + +import asyncio +from unittest.mock import MagicMock + +import pytest +from fastapi import FastAPI +from starlette.testclient import TestClient + +from nemo_gym._checkpoint import ( + CONTROL_URL_PREFIX, + GATED_MODEL_ROUTE_SUFFIXES, + MODEL_ADMISSION_URL_PREFIX, + AdmissionLimiter, + AdmissionMiddleware, + AdmissionParkedError, + AdmissionState, + StaleAttemptError, +) +from nemo_gym.base_responses_api_model import BaseResponsesAPIModelConfig, SimpleResponsesAPIModel +from nemo_gym.openai_utils import NeMoGymResponse, NeMoGymResponseCreateParamsNonStreaming +from nemo_gym.rollout_correlation import ATTEMPT_INDEX_HEADER, ROLLOUT_ID_HEADER +from nemo_gym.server_utils import ServerClient + + +AUTH_HEADERS = {"authorization": "Bearer test-control-token"} + + +def _identity_headers(rollout_id: str, attempt_index: int = 0) -> dict[str, str]: + return { + ROLLOUT_ID_HEADER: rollout_id, + ATTEMPT_INDEX_HEADER: str(attempt_index), + } + + +# --- limiter --- + + +def test_close_with_nothing_inflight_pauses_immediately() -> None: + limiter = AdmissionLimiter() + limiter.close() + assert limiter.state == AdmissionState.PAUSED + + +def test_drain_completes_when_last_inflight_releases() -> None: + limiter = AdmissionLimiter() + ticket = limiter.admit(rollout_id="4-2", attempt_index=0) + limiter.close() + assert limiter.state == AdmissionState.DRAINING + limiter.release(ticket) + assert limiter.state == AdmissionState.PAUSED + + +def test_new_root_operation_parks_while_draining() -> None: + limiter = AdmissionLimiter() + held = limiter.admit(rollout_id="4-2", attempt_index=0) + limiter.close() + with pytest.raises(AdmissionParkedError): + limiter.admit(rollout_id="9-9", attempt_index=0) + limiter.release(held) + + +def test_resume_reopens_admission() -> None: + limiter = AdmissionLimiter() + limiter.close() + limiter.resume() + assert limiter.state == AdmissionState.ACCEPTING + limiter.release(limiter.admit(rollout_id="4-2", attempt_index=0)) + + +def test_abort_inflight_tombstones_but_waits_for_request_exit() -> None: + limiter = AdmissionLimiter() + stuck = limiter.admit(rollout_id="7-1", attempt_index=2) + limiter.close() + aborted = limiter.abort_inflight("7-1", 2) + assert aborted == [stuck.ticket_id] + # Cancellation does not establish quiescence. + # The request remains counted until its ASGI task exits. + assert limiter.state == AdmissionState.DRAINING + limiter.release(stuck) + assert limiter.counts()["inflight_total"] == 0 + + limiter.resume() + with pytest.raises(StaleAttemptError): + limiter.admit(rollout_id="7-1", attempt_index=2) + limiter.release(limiter.admit(rollout_id="7-1", attempt_index=3)) + assert limiter.tombstones() == [("7-1", 2)] + + +@pytest.mark.asyncio +async def test_abort_cancels_request_task_before_drain_completes() -> None: + limiter = AdmissionLimiter() + admitted = asyncio.Event() + + async def operation() -> None: + ticket = limiter.admit( + rollout_id="7-1", + attempt_index=2, + task=asyncio.current_task(), + ) + admitted.set() + try: + await asyncio.Event().wait() + finally: + limiter.release(ticket) + + task = asyncio.create_task(operation()) + await admitted.wait() + limiter.close() + limiter.abort_inflight("7-1", 2) + with pytest.raises(asyncio.CancelledError): + await task + assert limiter.state == AdmissionState.PAUSED + assert limiter.counts()["inflight_total"] == 0 + + +@pytest.mark.asyncio +async def test_wait_for_drained_long_poll() -> None: + limiter = AdmissionLimiter() + ticket = limiter.admit(rollout_id="4-2", attempt_index=0) + limiter.close() + + async def release_soon() -> None: + await asyncio.sleep(0.01) + limiter.release(ticket) + + releaser = asyncio.create_task(release_soon()) + assert await limiter.wait_for_drained(timeout_s=1.0) + await releaser + assert limiter.state == AdmissionState.PAUSED + + +@pytest.mark.asyncio +async def test_wait_for_drained_times_out_with_stragglers() -> None: + limiter = AdmissionLimiter() + held = limiter.admit(rollout_id="4-2", attempt_index=0) + limiter.close() + assert not await limiter.wait_for_drained(timeout_s=0.01) + limiter.release(held) + + +@pytest.mark.asyncio +async def test_pause_admission_linearization_race() -> None: + for index in range(100): + limiter = AdmissionLimiter() + start = asyncio.Event() + admitted = asyncio.Event() + release = asyncio.Event() + side_effects = 0 + + async def request() -> str: + nonlocal side_effects + await start.wait() + try: + ticket = limiter.admit(rollout_id=f"rollout-{index}", attempt_index=0) + except AdmissionParkedError: + return "parked" + side_effects += 1 + admitted.set() + await release.wait() + limiter.release(ticket) + return "admitted" + + async def pause() -> None: + await start.wait() + limiter.close() + + request_task = asyncio.create_task(request()) + pause_task = asyncio.create_task(pause()) + start.set() + await pause_task + + if admitted.is_set(): + assert limiter.counts()["inflight_total"] == 1 + assert side_effects == 1 + else: + assert limiter.counts()["inflight_total"] == 0 + assert side_effects == 0 + release.set() + outcome = await request_task + assert (outcome, side_effects) in {("admitted", 1), ("parked", 0)} + + +# --- middleware --- + + +def _gated_app(limiter: AdmissionLimiter) -> TestClient: + app = FastAPI() + + @app.post("/v1/responses") + async def responses() -> dict: + return {"inflight": limiter.counts()["inflight_total"]} + + @app.get("/other") + async def other() -> dict: + return {"ok": True} + + app.add_middleware(AdmissionMiddleware, limiter=limiter, gated_suffixes=GATED_MODEL_ROUTE_SUFFIXES) + return TestClient(app) + + +def test_middleware_counts_request_until_response_completes() -> None: + limiter = AdmissionLimiter() + body = _gated_app(limiter).post("/v1/responses").json() + # The request was in flight while handled and released afterwards. + assert body["inflight"] == 1 + assert limiter.counts()["inflight_total"] == 0 + + +def test_middleware_parks_new_calls_when_closed() -> None: + limiter = AdmissionLimiter() + client = _gated_app(limiter) + limiter.close() + response = client.post("/v1/responses", headers=_identity_headers("4-2")) + assert response.status_code == 409 + assert response.json()["error"]["code"] == "checkpoint_parked" + assert response.headers["retry-after"] == "1" + + +def test_middleware_rejects_stale_attempt() -> None: + limiter = AdmissionLimiter() + client = _gated_app(limiter) + limiter.close() + limiter.abort_inflight("7-1", 0) + limiter.resume() + response = client.post("/v1/responses", headers=_identity_headers("7-1")) + assert response.status_code == 409 + assert response.json()["error"]["code"] == "stale_attempt" + + +def test_middleware_resolves_blackbox_identity_from_capture_path() -> None: + limiter = AdmissionLimiter() + client = _gated_app(limiter) + limiter.install_tombstone("7-1", 2) + response = client.post("/ng-rollout/7-1-a2/v1/responses") + assert response.status_code == 409 + assert response.json()["error"]["code"] == "stale_attempt" + + +def test_middleware_rejects_capture_path_identity_disagreement() -> None: + limiter = AdmissionLimiter() + response = _gated_app(limiter).post( + "/ng-rollout/7-1-a2/v1/responses", + headers=_identity_headers("7-1", 3), + ) + assert response.status_code == 409 + assert response.json()["error"]["code"] == "execution_identity_mismatch" + + +def test_middleware_leaves_ungated_paths_open_while_paused() -> None: + limiter = AdmissionLimiter() + client = _gated_app(limiter) + limiter.close() + assert client.get("/other").json() == {"ok": True} + + +# --- model server routes --- + + +def _model_server(instance_role: str) -> SimpleResponsesAPIModel: + class _Model(SimpleResponsesAPIModel): + def checkpoint_control_auth_token(self) -> str: + return "test-control-token" + + async def chat_completions(self, request): + raise NotImplementedError + + async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming) -> NeMoGymResponse: + return NeMoGymResponse.model_validate( + { + "id": "resp-1", + "created_at": 0.0, + "model": "m", + "object": "response", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + } + ) + + server_client = MagicMock(spec=ServerClient) + server_client.global_config_dict = {} + return _Model( + config=BaseResponsesAPIModelConfig(host="", port=0, entrypoint="", name="policy", instance_role=instance_role), + server_client=server_client, + ) + + +def test_policy_model_server_pause_drain_resume_cycle() -> None: + client = TestClient(_model_server("policy").setup_webserver()) + + capabilities = client.get(f"{CONTROL_URL_PREFIX}/capabilities").json() + assert capabilities["instance_role"] == "policy" + assert capabilities["admission_states"] == ["accepting", "draining", "paused"] + + # Generation works while accepting. + assert client.post("/v1/responses", json={"input": "hi"}).status_code == 200 + + pause_body = {"checkpoint_id": "ckpt-1", "deadline_ts": 4e9} + pause = client.post(f"{MODEL_ADMISSION_URL_PREFIX}/pause", json=pause_body, headers=AUTH_HEADERS) + assert pause.status_code == 200 + assert pause.json() == { + "state": "paused", + "workers": {"acknowledged": 1, "expected": 1}, + "inflight_total": 0, + "waiters_total": 0, + } + + # New generation parks; control routes stay reachable. + parked = client.post("/v1/responses", json={"input": "hi"}) + assert parked.status_code == 409 + assert parked.json()["error"]["code"] == "checkpoint_parked" + status = client.get( + f"{MODEL_ADMISSION_URL_PREFIX}/status", + params={"checkpoint_id": "ckpt-1", "deadline_ts": 4e9}, + headers=AUTH_HEADERS, + ).json() + assert status["state"] == "paused" + assert status["per_worker"] == {"0": {"state": "paused", "inflight": 0}} + + # A duplicate pause replays; a competing checkpoint conflicts. + replay = client.post(f"{MODEL_ADMISSION_URL_PREFIX}/pause", json=pause_body, headers=AUTH_HEADERS) + assert replay.json() == pause.json() + conflict = client.post( + f"{MODEL_ADMISSION_URL_PREFIX}/pause", + json={"checkpoint_id": "ckpt-2", "deadline_ts": 4e9}, + headers=AUTH_HEADERS, + ) + assert conflict.status_code == 409 + assert conflict.json()["error"]["code"] == "checkpoint_conflict" + + resume = client.post( + f"{MODEL_ADMISSION_URL_PREFIX}/resume", + json={"checkpoint_id": "ckpt-1", "deadline_ts": 4e9}, + headers=AUTH_HEADERS, + ) + assert resume.json() == { + "state": "accepting", + "workers": {"acknowledged": 1, "expected": 1}, + "released_waiters": 0, + } + assert client.post("/v1/responses", json={"input": "hi"}).status_code == 200 + + # The finished checkpoint id is stale for new operations; a new checkpoint proceeds. + stale = client.post(f"{MODEL_ADMISSION_URL_PREFIX}/pause", json=pause_body, headers=AUTH_HEADERS) + assert stale.status_code == 409 + assert stale.json()["error"]["code"] == "stale_checkpoint" + assert client.post( + f"{MODEL_ADMISSION_URL_PREFIX}/pause", + json={"checkpoint_id": "ckpt-3", "deadline_ts": 4e9}, + headers=AUTH_HEADERS, + ).json()["state"] in {"paused", "draining"} + + +def test_model_admission_requires_existing_control_bearer() -> None: + client = TestClient(_model_server("policy").setup_webserver()) + response = client.post( + f"{MODEL_ADMISSION_URL_PREFIX}/pause", + json={"checkpoint_id": "ckpt-1", "deadline_ts": 4e9}, + ) + assert response.status_code == 401 + + +def test_abort_inflight_tombstones_via_route() -> None: + server = _model_server("policy") + client = TestClient(server.setup_webserver()) + client.post( + f"{MODEL_ADMISSION_URL_PREFIX}/pause", + json={"checkpoint_id": "ckpt-1", "deadline_ts": 4e9}, + headers=AUTH_HEADERS, + ) + + abort = client.post( + f"{MODEL_ADMISSION_URL_PREFIX}/abort_inflight", + json={ + "checkpoint_id": "ckpt-1", + "deadline_ts": 4e9, + "rollout_id": "7-1", + "attempt_index": 2, + }, + headers=AUTH_HEADERS, + ) + assert abort.status_code == 200 + status = client.get( + f"{MODEL_ADMISSION_URL_PREFIX}/status", + params={"checkpoint_id": "ckpt-1", "deadline_ts": 4e9}, + headers=AUTH_HEADERS, + ).json() + assert status["tombstones"] == [{"rollout_id": "7-1", "attempt_index": 2}] + + client.post( + f"{MODEL_ADMISSION_URL_PREFIX}/resume", + json={"checkpoint_id": "ckpt-1", "deadline_ts": 4e9}, + headers=AUTH_HEADERS, + ) + stale = client.post("/v1/responses", json={"input": "hi"}, headers=_identity_headers("7-1", 2)) + assert stale.status_code == 409 + assert stale.json()["error"]["code"] == "stale_attempt" + fresh = client.post("/v1/responses", json={"input": "hi"}, headers=_identity_headers("7-1", 3)) + assert fresh.status_code == 200 + + +def test_auxiliary_model_server_never_gates_and_rejects_pause() -> None: + client = TestClient(_model_server("auxiliary").setup_webserver()) + + capabilities = client.get(f"{CONTROL_URL_PREFIX}/capabilities").json() + assert capabilities["instance_role"] == "auxiliary" + assert capabilities["admission_states"] == ["accepting"] + + pause = client.post( + f"{MODEL_ADMISSION_URL_PREFIX}/pause", + json={"checkpoint_id": "ckpt-1", "deadline_ts": 4e9}, + headers=AUTH_HEADERS, + ) + assert pause.status_code == 409 + assert pause.json()["error"]["code"] == "not_a_policy_instance" + + # Judge traffic keeps flowing: there is no admission middleware to close. + assert client.post("/v1/responses", json={"input": "grade"}).status_code == 200 diff --git a/tests/unit_tests/test_checkpoint_ledger.py b/tests/unit_tests/test_checkpoint_ledger.py new file mode 100644 index 00000000000..68788a04dd5 --- /dev/null +++ b/tests/unit_tests/test_checkpoint_ledger.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Checkpoint token-free model custody without copying staged token arrays.""" + +import json + +import pytest +from fastapi import FastAPI +from starlette.testclient import TestClient + +from nemo_gym._checkpoint import ( + LEDGER_MANIFEST_NAME, + MODEL_ADMISSION_URL_PREFIX, + MODEL_CHECKPOINT_URL_PREFIX, + MODEL_LEDGER_SUBDIR, + AdmissionLimiter, + CaptureLedgerCheckpointer, + ControlCapabilities, + ControlFence, + LedgerMismatchError, + MultiProcessCapability, + StaleAttemptError, + install_control_plane, + install_model_admission, + install_model_checkpoint, +) +from nemo_gym.token_id_capture.lineage import FileLineageStore + + +AUTH_TOKEN = "checkpoint-token" +AUTH_HEADERS = {"authorization": f"Bearer {AUTH_TOKEN}"} + + +def _write_custody(root, rollout_id: str, call_count: int = 2) -> bytes: + root.mkdir(parents=True, exist_ok=True) + rows = [ + { + "model_call_id": f"{rollout_id}-call-{index}", + "staging_key": f"opaque-{rollout_id}-{index}", + "staging_digest": f"digest-{index}", + "parent_call_id": None if index == 0 else f"{rollout_id}-call-{index - 1}", + } + for index in range(call_count) + ] + payload = b"".join(json.dumps(row, sort_keys=True).encode() + b"\n" for row in rows) + (root / f"{rollout_id}.lineage.jsonl").write_bytes(payload) + return payload + + +def test_commit_restore_preserves_only_token_free_custody(tmp_path) -> None: + root = tmp_path / "ledger-a" + expected = _write_custody(root, "rollout-a") + _write_custody(root, "rollout-b-a2") + + summary = CaptureLedgerCheckpointer(root).commit( + tmp_path / "checkpoint", + checkpoint_id="checkpoint-1", + tombstones=[("rollout-b", 2)], + ) + assert summary == { + "rollouts": 1, + "rows": 2, + "excluded_tombstoned": 1, + "manifest_digest": summary["manifest_digest"], + } + + ledger_dir = tmp_path / "checkpoint" / MODEL_LEDGER_SUBDIR + assert (ledger_dir / "rollout-a.lineage.jsonl").read_bytes() == expected + assert not (ledger_dir / "rollout-b-a2.lineage.jsonl").exists() + assert not list(ledger_dir.glob("*.tokens.*")) + + restored_root = tmp_path / "ledger-b" + result = CaptureLedgerCheckpointer(restored_root).restore(tmp_path / "checkpoint") + assert result["tombstones"] == [{"rollout_id": "rollout-b", "attempt_index": 2}] + assert (restored_root / "rollout-a.lineage.jsonl").read_bytes() == expected + + +def test_restore_validates_all_files_before_installing_any(tmp_path) -> None: + source = tmp_path / "source" + _write_custody(source, "rollout-a") + _write_custody(source, "rollout-b") + checkpoint = tmp_path / "checkpoint" + CaptureLedgerCheckpointer(source).commit(checkpoint, checkpoint_id="checkpoint-1", tombstones=[]) + (checkpoint / MODEL_LEDGER_SUBDIR / "rollout-b.lineage.jsonl").write_text("corrupt") + + restored = tmp_path / "restored" + with pytest.raises(LedgerMismatchError): + CaptureLedgerCheckpointer(restored).restore(checkpoint) + assert not restored.exists() + + +def test_restore_rejects_uncommitted_and_nonfresh_namespaces(tmp_path) -> None: + checkpoint = tmp_path / "checkpoint" + ledger_dir = checkpoint / MODEL_LEDGER_SUBDIR + ledger_dir.mkdir(parents=True) + _write_custody(ledger_dir, "rollout-a") + with pytest.raises(LedgerMismatchError): + CaptureLedgerCheckpointer(tmp_path / "restored").restore(checkpoint) + + source = tmp_path / "source" + _write_custody(source, "rollout-a") + CaptureLedgerCheckpointer(source).commit(checkpoint, checkpoint_id="checkpoint-1", tombstones=[]) + restored = tmp_path / "nonfresh" + _write_custody(restored, "old-rollout") + with pytest.raises(LedgerMismatchError): + CaptureLedgerCheckpointer(restored).restore(checkpoint) + + +def _participant(root) -> tuple[TestClient, AdmissionLimiter]: + app = FastAPI() + limiter = AdmissionLimiter() + fence = ControlFence() + ledger = FileLineageStore(root) + install_control_plane( + app, + capabilities=ControlCapabilities( + component="responses_api_models", + name="policy", + multi_process=MultiProcessCapability(mode="single_worker", num_workers=1), + instance_role="policy", + ), + fence=fence, + ) + install_model_admission( + app, + limiter=limiter, + fence=fence, + instance_role="policy", + auth_token=AUTH_TOKEN, + ) + install_model_checkpoint( + app, + fence=fence, + limiter=limiter, + ledger_provider=lambda: ledger, + file_ledger_root_provider=lambda: ledger.checkpoint_root, + instance_role="policy", + auth_token=AUTH_TOKEN, + ) + return TestClient(app), limiter + + +def test_commit_requires_completed_drain_and_restore_stays_paused(tmp_path) -> None: + source_client, source_limiter = _participant(tmp_path / "ledger-a") + _write_custody(tmp_path / "ledger-a", "rollout-a") + held = source_limiter.admit(rollout_id="rollout-a", attempt_index=0) + + pause_body = {"checkpoint_id": "checkpoint-1", "deadline_ts": 4e9} + pause = source_client.post( + f"{MODEL_ADMISSION_URL_PREFIX}/pause", + json=pause_body, + headers=AUTH_HEADERS, + ) + assert pause.json()["state"] == "draining" + commit_body = {**pause_body, "checkpoint_dir": str(tmp_path / "checkpoint")} + early = source_client.post( + f"{MODEL_CHECKPOINT_URL_PREFIX}/commit", + json=commit_body, + headers=AUTH_HEADERS, + ) + assert early.status_code == 409 + assert early.json()["error"]["code"] in {"invalid_phase", "ledger_not_quiescent"} + + source_limiter.release(held) + status = source_client.get( + f"{MODEL_ADMISSION_URL_PREFIX}/status", + params=pause_body, + headers=AUTH_HEADERS, + ) + assert status.json()["state"] == "paused" + commit = source_client.post( + f"{MODEL_CHECKPOINT_URL_PREFIX}/commit", + json=commit_body, + headers=AUTH_HEADERS, + ) + assert commit.status_code == 200 + + restored_client, restored_limiter = _participant(tmp_path / "ledger-b") + restore = restored_client.post( + f"{MODEL_CHECKPOINT_URL_PREFIX}/restore", + json={ + "checkpoint_id": "restore-1", + "deadline_ts": 4e9, + "checkpoint_dir": str(tmp_path / "checkpoint"), + }, + headers=AUTH_HEADERS, + ) + assert restore.status_code == 200 + assert restored_limiter.counts()["state"] == "paused" + + +def test_restored_tombstone_fences_exact_attempt(tmp_path) -> None: + source = tmp_path / "source" + _write_custody(source, "run-a1") + CaptureLedgerCheckpointer(source).commit( + tmp_path / "checkpoint", + checkpoint_id="checkpoint-1", + tombstones=[("run-a1", 0)], + ) + restored = tmp_path / "restored" + result = CaptureLedgerCheckpointer(restored).restore(tmp_path / "checkpoint") + limiter = AdmissionLimiter() + for tombstone in result["tombstones"]: + limiter.install_tombstone(tombstone["rollout_id"], tombstone["attempt_index"]) + + with pytest.raises(StaleAttemptError): + limiter.admit(rollout_id="run-a1", attempt_index=0) + limiter.release(limiter.admit(rollout_id="run", attempt_index=1)) + + +def test_checkpoint_routes_require_control_bearer(tmp_path) -> None: + client, _ = _participant(tmp_path / "ledger") + response = client.post( + f"{MODEL_CHECKPOINT_URL_PREFIX}/restore", + json={"checkpoint_id": "restore-1", "deadline_ts": 4e9, "checkpoint_dir": str(tmp_path)}, + ) + assert response.status_code == 401 + + +def test_manifest_is_published_last(tmp_path, monkeypatch) -> None: + source = tmp_path / "source" + _write_custody(source, "rollout-a") + checkpoint = tmp_path / "checkpoint" + + import nemo_gym._checkpoint.ledger as ledger_module + + real_replace = ledger_module.os.replace + + def fail_manifest_replace(source_path, target_path) -> None: + if target_path.name == LEDGER_MANIFEST_NAME: + raise RuntimeError("injected manifest publication failure") + real_replace(source_path, target_path) + + monkeypatch.setattr(ledger_module.os, "replace", fail_manifest_replace) + with pytest.raises(RuntimeError, match="injected"): + CaptureLedgerCheckpointer(source).commit(checkpoint, checkpoint_id="checkpoint-1", tombstones=[]) + assert not (checkpoint / MODEL_LEDGER_SUBDIR / LEDGER_MANIFEST_NAME).exists() From 36ad4261963925c1ef94ff0d6b3838233776c1e2 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Thu, 3 Sep 2026 01:52:52 -0700 Subject: [PATCH 2/5] fix(checkpoint): preserve model lineage across restore Namespace model artifacts, reject changing commit exclusions, and link the first restored call to its verified parent exactly once. Signed-off-by: Ananth Subramaniam --- nemo_gym/_checkpoint/ledger.py | 155 +++++++++++++-- nemo_gym/_checkpoint/model_admission.py | 17 +- nemo_gym/base_responses_api_model.py | 38 +++- nemo_gym/rollout_correlation.py | 121 +++++++++--- nemo_gym/server_utils.py | 33 ++++ nemo_gym/token_id_capture/lineage.py | 126 ++++++++++++ nemo_gym/token_id_capture/protocols.py | 17 ++ nemo_gym/token_id_capture/sink.py | 33 +++- tests/unit_tests/test_checkpoint_admission.py | 68 +++++++ tests/unit_tests/test_checkpoint_ledger.py | 182 ++++++++++++++++++ tests/unit_tests/test_server_utils.py | 61 +++++- tests/unit_tests/test_token_capture_ledger.py | 109 ++++++++++- 12 files changed, 904 insertions(+), 56 deletions(-) diff --git a/nemo_gym/_checkpoint/ledger.py b/nemo_gym/_checkpoint/ledger.py index 6f06f992ca5..6bf2af34c3c 100644 --- a/nemo_gym/_checkpoint/ledger.py +++ b/nemo_gym/_checkpoint/ledger.py @@ -47,6 +47,7 @@ from typing import Any, Callable, Literal, Optional, Protocol, runtime_checkable from fastapi import FastAPI, Header +from pydantic import BaseModel, ConfigDict, Field from nemo_gym._checkpoint.admission import AdmissionLimiter from nemo_gym._checkpoint.control import ( @@ -57,7 +58,7 @@ ControlFence, ) from nemo_gym._checkpoint.model_admission import NotPolicyInstanceError -from nemo_gym.rollout_correlation import capture_key_for +from nemo_gym.rollout_correlation import ROLLOUT_ID_PATTERN, capture_key_for from nemo_gym.token_id_capture.control_routes import require_control_auth from nemo_gym.token_id_capture.protocols import CaptureLedger @@ -96,12 +97,40 @@ class LedgerNotQuiescentError(ControlError): code = "ledger_not_quiescent" +class AttemptIdentity(BaseModel): + model_config = ConfigDict(extra="forbid") + + rollout_id: str + attempt_index: int = Field(ge=0) + + +class CaptureLedgerCommitResult(BaseModel): + model_config = ConfigDict(extra="forbid") + + rollouts: int = Field(ge=0) + rows: int = Field(ge=0) + excluded_tombstoned: int = Field(ge=0) + manifest_digest: str + + +class CaptureLedgerRestoreResult(BaseModel): + model_config = ConfigDict(extra="forbid") + + rollouts: int = Field(ge=0) + rows: int = Field(ge=0) + checkpoint_id: Optional[str] = None + tombstones: list[AttemptIdentity] = Field(default_factory=list) + source_attempts: list[AttemptIdentity] = Field(default_factory=list) + + @runtime_checkable class CheckpointableCaptureLedger(CaptureLedger, Protocol): """Optional lifecycle implemented by framework-owned capture backends. The backend snapshots token-free custody and its private parent-resolution state. The framework checkpoints staged token arrays separately. + Gym supplies a server-specific directory and ``server_name``. Because the + backend owns its manifest, it must record and validate that identity. """ async def checkpoint_capture_ledger( @@ -109,10 +138,17 @@ async def checkpoint_capture_ledger( checkpoint_dir: Path, *, checkpoint_id: str, + server_name: str, tombstones: tuple[tuple[str, int], ...], - ) -> dict[str, Any]: ... + source_attempts: tuple[tuple[str, int], ...], + ) -> CaptureLedgerCommitResult: ... - async def restore_capture_ledger(self, checkpoint_dir: Path) -> dict[str, Any]: ... + async def restore_capture_ledger( + self, + checkpoint_dir: Path, + *, + server_name: str, + ) -> CaptureLedgerRestoreResult: ... def _file_digest(path: Path) -> str: @@ -148,21 +184,43 @@ def _fsync_dir(path: Path) -> None: class CaptureLedgerCheckpointer: """Commit and restore one token-capture store directory.""" - def __init__(self, store_root: Path) -> None: + def __init__(self, store_root: Path, *, server_name: Optional[str] = None) -> None: self.store_root = Path(store_root) + self.server_name = _validate_server_name(server_name) if server_name is not None else None + + def _ledger_dir(self, checkpoint_dir: Path) -> Path: + directory = Path(checkpoint_dir) / MODEL_LEDGER_SUBDIR + return directory / self.server_name if self.server_name is not None else directory def _rollout_ids(self) -> list[str]: return sorted(path.name[: -len(_LEDGER_SUFFIX)] for path in self.store_root.glob(f"*{_LEDGER_SUFFIX}")) - def commit(self, checkpoint_dir: Path, *, checkpoint_id: str, tombstones: list[tuple[str, int]]) -> dict[str, Any]: + def commit( + self, + checkpoint_dir: Path, + *, + checkpoint_id: str, + tombstones: list[tuple[str, int]], + source_attempts: Optional[list[tuple[str, int]]] = None, + ) -> dict[str, Any]: """Copy the ledger into ``checkpoint_dir``; the caller has already drained. The store must be quiescent (admission paused) when this runs: the copy takes no locks because nothing may be writing. """ - ledger_dir = Path(checkpoint_dir) / MODEL_LEDGER_SUBDIR + ledger_dir = self._ledger_dir(checkpoint_dir) if (ledger_dir / LEDGER_MANIFEST_NAME).exists(): - raise LedgerMismatchError(f"checkpoint ledger already committed at {ledger_dir}") + result = self._validate_committed( + ledger_dir, + checkpoint_id=checkpoint_id, + server_name=self.server_name, + tombstones=tombstones, + source_attempts=source_attempts or [], + ) + # A previous attempt may have renamed the manifest and then + # failed its final directory fsync. Retry that durability barrier. + _fsync_dir(ledger_dir) + return result ledger_dir.mkdir(parents=True, exist_ok=True) fenced = {capture_key_for(rollout_id, attempt_index) for rollout_id, attempt_index in tombstones} @@ -187,10 +245,15 @@ def commit(self, checkpoint_dir: Path, *, checkpoint_id: str, tombstones: list[t manifest = { "schema_version": LEDGER_SCHEMA_VERSION, "checkpoint_id": checkpoint_id, + "server_name": self.server_name, "rollouts": rollouts, "tombstones": [ {"rollout_id": rollout_id, "attempt_index": attempt} for rollout_id, attempt in sorted(tombstones) ], + "source_attempts": [ + {"rollout_id": rollout_id, "attempt_index": attempt} + for rollout_id, attempt in sorted(source_attempts or []) + ], } payload = json.dumps(manifest, sort_keys=True, indent=1).encode() with tempfile.NamedTemporaryFile(dir=ledger_dir, prefix=".manifest-", delete=False) as handle: @@ -212,9 +275,47 @@ def commit(self, checkpoint_dir: Path, *, checkpoint_id: str, tombstones: list[t "manifest_digest": hashlib.sha256(payload).hexdigest(), } + @staticmethod + def _validate_committed( + ledger_dir: Path, + *, + checkpoint_id: str, + server_name: Optional[str], + tombstones: list[tuple[str, int]], + source_attempts: list[tuple[str, int]], + ) -> dict[str, Any]: + manifest_path = ledger_dir / LEDGER_MANIFEST_NAME + payload = manifest_path.read_bytes() + manifest = json.loads(payload) + if manifest.get("checkpoint_id") != checkpoint_id or manifest.get("server_name") != server_name: + raise LedgerMismatchError("ledger directory belongs to a different checkpoint transaction or model server") + expected_tombstones = [ + {"rollout_id": rollout_id, "attempt_index": attempt} for rollout_id, attempt in sorted(tombstones) + ] + expected_source_attempts = [ + {"rollout_id": rollout_id, "attempt_index": attempt} for rollout_id, attempt in sorted(source_attempts) + ] + if manifest.get("tombstones", []) != expected_tombstones: + raise LedgerMismatchError("committed ledger abort exclusions changed before commit retry") + if manifest.get("source_attempts", []) != expected_source_attempts: + raise LedgerMismatchError("committed ledger source attempts changed before commit retry") + total_rows = 0 + for rollout_id, metadata in manifest.get("rollouts", {}).items(): + for name, digest in metadata.get("files", {}).items(): + path = ledger_dir / name + if not path.exists() or _file_digest(path) != digest: + raise LedgerMismatchError(f"committed ledger file {name!r} for {rollout_id!r} is corrupted") + total_rows += int(metadata.get("rows", 0)) + return { + "rollouts": len(manifest.get("rollouts", {})), + "rows": total_rows, + "excluded_tombstoned": len(manifest.get("tombstones", [])), + "manifest_digest": hashlib.sha256(payload).hexdigest(), + } + def restore(self, checkpoint_dir: Path) -> dict[str, Any]: """Install a committed ledger into this store root and verify it.""" - ledger_dir = Path(checkpoint_dir) / MODEL_LEDGER_SUBDIR + ledger_dir = self._ledger_dir(checkpoint_dir) manifest_path = ledger_dir / LEDGER_MANIFEST_NAME if not manifest_path.exists(): raise LedgerMismatchError( @@ -227,6 +328,8 @@ def restore(self, checkpoint_dir: Path) -> dict[str, Any]: f"ledger manifest schema_version {manifest.get('schema_version')} is newer than this " f"reader ({LEDGER_SCHEMA_VERSION})" ) + if manifest.get("server_name") != self.server_name: + raise LedgerMismatchError("ledger checkpoint belongs to a different model server") expected_names = {name for metadata in manifest["rollouts"].values() for name in metadata["files"]} existing_names = {path.name for path in self.store_root.glob(f"*{_LEDGER_SUFFIX}")} @@ -261,6 +364,7 @@ def restore(self, checkpoint_dir: Path) -> dict[str, Any]: "rows": total_rows, "checkpoint_id": manifest.get("checkpoint_id"), "tombstones": list(manifest.get("tombstones", ())), + "source_attempts": list(manifest.get("source_attempts", ())), } @@ -272,6 +376,15 @@ class ModelCheckpointRestoreRequest(CheckpointControlRequest): checkpoint_dir: str +def _validate_server_name(server_name: str) -> str: + if ROLLOUT_ID_PATTERN.fullmatch(server_name) is None: + raise ValueError( + "model server name must contain only letters, digits, dots, dashes, or underscores " + "and start with a letter or digit" + ) + return server_name + + def install_model_checkpoint( app: FastAPI, *, @@ -280,6 +393,7 @@ def install_model_checkpoint( ledger_provider: Callable[[], Optional[CaptureLedger]], file_ledger_root_provider: Callable[[], Optional[Path]], instance_role: Literal["policy", "auxiliary"], + server_name: str, auth_token: str, ) -> None: """Register ``/ng-control/v1/model-checkpoint`` on a model-server app. @@ -289,6 +403,8 @@ def install_model_checkpoint( coordinator has restored every component and explicitly resumes. """ + server_name = _validate_server_name(server_name) + def _require_policy() -> None: if instance_role != "policy": raise NotPolicyInstanceError( @@ -311,13 +427,18 @@ async def _with_ledger( ) -> dict[str, Any]: ledger = ledger_provider() if isinstance(ledger, CheckpointableCaptureLedger): + participant_dir = checkpoint_dir / MODEL_LEDGER_SUBDIR / server_name if operation == "commit": - return await ledger.checkpoint_capture_ledger( - checkpoint_dir, + result = await ledger.checkpoint_capture_ledger( + participant_dir, checkpoint_id=checkpoint_id, - tombstones=tuple(limiter.tombstones()), + server_name=server_name, + tombstones=tuple(limiter.checkpoint_exclusions()), + source_attempts=tuple(limiter.seen_attempts()), ) - return await ledger.restore_capture_ledger(checkpoint_dir) + return CaptureLedgerCommitResult.model_validate(result).model_dump(mode="json") + result = await ledger.restore_capture_ledger(participant_dir, server_name=server_name) + return CaptureLedgerRestoreResult.model_validate(result).model_dump(mode="json") file_root = file_ledger_root_provider() if file_root is None: @@ -325,13 +446,14 @@ async def _with_ledger( "the configured CaptureLedger must implement CheckpointableCaptureLedger; " "Gym cannot infer how to snapshot a framework-owned backend" ) - checkpointer = CaptureLedgerCheckpointer(file_root) + checkpointer = CaptureLedgerCheckpointer(file_root, server_name=server_name) if operation == "commit": return await _run_sync( lambda: checkpointer.commit( checkpoint_dir, checkpoint_id=checkpoint_id, - tombstones=limiter.tombstones(), + tombstones=limiter.checkpoint_exclusions(), + source_attempts=limiter.seen_attempts(), ) ) return await _run_sync(lambda: checkpointer.restore(checkpoint_dir)) @@ -381,15 +503,18 @@ async def run() -> dict[str, Any]: ) for tombstone in result["tombstones"]: limiter.install_tombstone(tombstone["rollout_id"], tombstone["attempt_index"]) + for source_attempt in result.get("source_attempts", []): + limiter.install_tombstone(source_attempt["rollout_id"], source_attempt["attempt_index"]) return result return await fence.run_operation( body.checkpoint_id, "model-checkpoint/restore", - allowed_phases=frozenset({CheckpointPhase.IDLE}), + allowed_phases=frozenset({CheckpointPhase.IDLE, CheckpointPhase.RESTORE_FAILED_PAUSED}), phase_during=CheckpointPhase.RESTORING, phase_after=CheckpointPhase.RESTORED_PAUSED, run=run, + phase_on_failure=CheckpointPhase.RESTORE_FAILED_PAUSED, ) diff --git a/nemo_gym/_checkpoint/model_admission.py b/nemo_gym/_checkpoint/model_admission.py index b9a1ed62fbb..74834fc6857 100644 --- a/nemo_gym/_checkpoint/model_admission.py +++ b/nemo_gym/_checkpoint/model_admission.py @@ -121,13 +121,23 @@ async def model_admission_status( require_control_auth(authorization, auth_token) fence.require_phase( checkpoint_id, - frozenset({CheckpointPhase.PREPARING, CheckpointPhase.PREPARED}), + frozenset( + { + CheckpointPhase.PREPARING, + CheckpointPhase.PREPARED, + CheckpointPhase.COMMITTING, + CheckpointPhase.COMMITTED_PAUSED, + CheckpointPhase.RESTORING, + CheckpointPhase.RESTORE_FAILED_PAUSED, + CheckpointPhase.RESTORED_PAUSED, + } + ), ) deadline = Deadline(deadline_ts=deadline_ts) if wait_state == "paused" and timeout_s > 0: await limiter.wait_for_drained(min(timeout_s, deadline.remaining())) counts = limiter.counts() - if counts["state"] == AdmissionState.PAUSED.value: + if counts["state"] == AdmissionState.PAUSED.value and fence.phase == CheckpointPhase.PREPARING: fence.mark_prepared(checkpoint_id) return { "checkpoint_id": checkpoint_id, @@ -161,10 +171,11 @@ async def run() -> dict[str, Any]: CheckpointPhase.PREPARING, CheckpointPhase.PREPARED, CheckpointPhase.COMMITTED_PAUSED, + CheckpointPhase.RESTORE_FAILED_PAUSED, CheckpointPhase.RESTORED_PAUSED, } ), - phase_during=CheckpointPhase.PREPARED, + phase_during=fence.phase, phase_after=CheckpointPhase.IDLE, run=run, retire_outcome="resumed", diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index 98e1fceb9de..df790203920 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -47,7 +47,13 @@ from pydantic import BaseModel, Field, PrivateAttr, ValidationError, model_validator from nemo_gym._checkpoint.admission import GATED_MODEL_ROUTE_SUFFIXES, AdmissionLimiter, AdmissionMiddleware -from nemo_gym._checkpoint.control import AdmissionState, ControlCapabilities +from nemo_gym._checkpoint.control import ( + AdmissionState, + ControlCapabilities, +) +from nemo_gym._checkpoint.control import ( + checkpoint_control_auth_token as resolve_checkpoint_control_auth_token, +) from nemo_gym._checkpoint.ledger import install_model_checkpoint from nemo_gym._checkpoint.model_admission import install_model_admission from nemo_gym.anthropic_converter import AnthropicConverter @@ -67,7 +73,13 @@ synthesize_responses_sse, validate_streaming_responses_params, ) -from nemo_gym.rollout_correlation import MODEL_CALL_ID_HEADER, maybe_rollout_id_from_run_body +from nemo_gym.rollout_correlation import ( + LOGICAL_REQUEST_HEADER, + MODEL_CALL_ID_HEADER, + PARENT_MODEL_CALL_ID_HEADER, + SOURCE_CAPTURE_KEY_HEADER, + maybe_rollout_id_from_run_body, +) from nemo_gym.rollout_observability import AgentObservationBundle, ObservationGap, join_model_call_observations from nemo_gym.server_utils import ( BaseRunServerInstanceConfig, @@ -280,6 +292,7 @@ def setup_model_admission(self, app: FastAPI) -> None: else None ), instance_role=self.config.instance_role, + server_name=self.config.name, auth_token=auth_token, ) if self.config.instance_role == "policy": @@ -290,10 +303,7 @@ def setup_model_admission(self, app: FastAPI) -> None: ) def checkpoint_control_auth_token(self) -> Optional[str]: - settings = token_id_capture_config(self.server_client.global_config_dict) - if settings is None or not settings.token_id_capture.external_staging: - return None - return settings.token_id_capture.resolve_control_auth_token() + return resolve_checkpoint_control_auth_token(getattr(self.server_client, "global_config_dict", None)) def control_capabilities(self) -> ControlCapabilities: capabilities = super().control_capabilities() @@ -969,6 +979,19 @@ def _unique_request_header(headers: list, name: bytes) -> Optional[str]: return values.pop() if len(values) == 1 else None +def _scope_header(scope: dict[str, Any], name: str) -> str | None: + """Read one unambiguous ASGI header value.""" + encoded_name = name.lower().encode("ascii") + values = {value.decode("latin-1") for key, value in scope.get("headers", []) if key.lower() == encoded_name} + if not values: + return None + if len(values) == 1: + return values.pop() + # A joined value cannot authenticate. + # Logical ID validation rejects the comma. + return ",".join(sorted(values)) + + def _consume_terminal_sse_event(buffer: bytearray, dialect: str) -> Optional[str]: blocks = re.split(rb"\r?\n\r?\n", bytes(buffer)) buffer[:] = blocks.pop() @@ -1439,6 +1462,9 @@ async def _send_with_model_call_id(message: dict[str, Any]) -> None: lineage_store=self._capture_ledger if self._external_staging else self._lineage_store, delta_records=self._delta_records, external_staging=self._external_staging, + logical_request_id=_scope_header(scope, LOGICAL_REQUEST_HEADER), + source_capture_key=_scope_header(scope, SOURCE_CAPTURE_KEY_HEADER), + explicit_parent_call_id=_scope_header(scope, PARENT_MODEL_CALL_ID_HEADER), admitted_at=time.time(), ) sink_token = set_token_sink(capture_context) diff --git a/nemo_gym/rollout_correlation.py b/nemo_gym/rollout_correlation.py index b7476560417..cadb23fad8b 100644 --- a/nemo_gym/rollout_correlation.py +++ b/nemo_gym/rollout_correlation.py @@ -14,8 +14,9 @@ # limitations under the License. import re from collections.abc import Iterator, Mapping -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from contextvars import ContextVar +from dataclasses import dataclass from typing import Any, Optional from pydantic import BaseModel @@ -34,12 +35,27 @@ _LOGICAL_ROLLOUT_ID: ContextVar[Optional[str]] = ContextVar("nemo_gym_logical_rollout_id", default=None) _ATTEMPT_INDEX: ContextVar[Optional[int]] = ContextVar("nemo_gym_attempt_index", default=None) + +@dataclass +class _CheckpointParentState: + source_capture_key: str + parent_model_call_id: str + consumed: bool = False + + +_CHECKPOINT_PARENT: ContextVar[Optional[_CheckpointParentState]] = ContextVar( + "nemo_gym_checkpoint_parent", + default=None, +) + # These headers propagate execution identity independently of capture. # The rollout header is the stable logical ID. # The attempt header identifies one dispatch of that rollout. ROLLOUT_ID_HEADER = "x-nemo-gym-rollout-id" ATTEMPT_INDEX_HEADER = "x-nemo-gym-attempt-index" MODEL_CALL_ID_HEADER = "x-nemo-gym-model-call-id" +SOURCE_CAPTURE_KEY_HEADER = "x-nemo-gym-source-capture-key" +PARENT_MODEL_CALL_ID_HEADER = "x-nemo-gym-parent-model-call-id" # The transport id appends ``-a{n}`` for re-dispatch attempts. The suffix is a # capture and routing key, never the logical identity. This pattern recovers @@ -184,6 +200,34 @@ def current_execution_identity() -> tuple[Optional[str], Optional[int]]: return current_logical_rollout_id(), current_attempt_index() +def take_checkpoint_parent() -> tuple[Optional[str], Optional[str]]: + """Consume the restored parent for the first policy model request.""" + state = _CHECKPOINT_PARENT.get() + if state is None or state.consumed: + return None, None + state.consumed = True + return state.source_capture_key, state.parent_model_call_id + + +@contextmanager +def checkpoint_parent_context(source_capture_key: str, parent_model_call_id: str) -> Iterator[None]: + """Propagate an explicit cross-attempt parent to downstream model calls.""" + if ROLLOUT_ID_PATTERN.fullmatch(source_capture_key) is None: + raise ValueError("invalid source capture key") + if not parent_model_call_id or "\r" in parent_model_call_id or "\n" in parent_model_call_id: + raise ValueError("invalid parent model-call ID") + token = _CHECKPOINT_PARENT.set( + _CheckpointParentState( + source_capture_key=source_capture_key, + parent_model_call_id=parent_model_call_id, + ) + ) + try: + yield + finally: + _CHECKPOINT_PARENT.reset(token) + + @contextmanager def rollout_context( rollout_id: Optional[str], @@ -216,31 +260,50 @@ def __init__(self, app: Any) -> None: self._app = app @staticmethod - def _identity_from_headers(scope: dict[str, Any]) -> tuple[Optional[str], Optional[int], Optional[str]]: - values: dict[str, set[str]] = {ROLLOUT_ID_HEADER: set(), ATTEMPT_INDEX_HEADER: set()} + def _identity_from_headers( + scope: dict[str, Any], + ) -> tuple[Optional[str], Optional[int], Optional[str], Optional[str], Optional[str]]: + values: dict[str, set[str]] = { + ROLLOUT_ID_HEADER: set(), + ATTEMPT_INDEX_HEADER: set(), + SOURCE_CAPTURE_KEY_HEADER: set(), + PARENT_MODEL_CALL_ID_HEADER: set(), + } for name, value in scope.get("headers") or (): key = name.decode("latin-1").lower() if key in values: values[key].add(value.decode("latin-1")) if any(len(entries) > 1 for entries in values.values()): - return None, None, "conflicting duplicate execution identity headers" + return None, None, None, None, "conflicting duplicate correlation headers" + source_values = values[SOURCE_CAPTURE_KEY_HEADER] + parent_values = values[PARENT_MODEL_CALL_ID_HEADER] + if bool(source_values) != bool(parent_values): + return None, None, None, None, "source capture key and parent model-call ID must be sent together" + source_capture_key = next(iter(source_values), None) + parent_model_call_id = next(iter(parent_values), None) + if source_capture_key is not None and ROLLOUT_ID_PATTERN.fullmatch(source_capture_key) is None: + return None, None, None, None, "invalid source capture key header" + if parent_model_call_id is not None and ( + not parent_model_call_id or "\r" in parent_model_call_id or "\n" in parent_model_call_id + ): + return None, None, None, None, "invalid parent model-call ID header" rollout_values = values[ROLLOUT_ID_HEADER] attempt_values = values[ATTEMPT_INDEX_HEADER] if bool(rollout_values) != bool(attempt_values): - return None, None, "rollout ID and attempt index headers must be sent together" + return None, None, None, None, "rollout ID and attempt index headers must be sent together" if not rollout_values: - return None, None, None + return None, None, source_capture_key, parent_model_call_id, None rollout_id = next(iter(rollout_values)) attempt_raw = next(iter(attempt_values)) if ROLLOUT_ID_PATTERN.fullmatch(rollout_id) is None: - return None, None, "invalid logical rollout ID header" + return None, None, None, None, "invalid logical rollout ID header" try: attempt_index = int(attempt_raw) except ValueError: - return None, None, "invalid attempt index header" + return None, None, None, None, "invalid attempt index header" if attempt_index < 0: - return None, None, "attempt index must be non-negative" - return rollout_id, attempt_index, None + return None, None, None, None, "attempt index must be non-negative" + return rollout_id, attempt_index, source_capture_key, parent_model_call_id, None @staticmethod async def _reject(scope: dict[str, Any], receive: Any, send: Any, detail: str) -> None: @@ -256,22 +319,31 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None return match = self._PREFIX.match(scope.get("path", "")) - logical_rollout_id, attempt_index, identity_error = self._identity_from_headers(scope) + logical_rollout_id, attempt_index, source_capture_key, parent_model_call_id, identity_error = ( + self._identity_from_headers(scope) + ) if identity_error is not None: await self._reject(scope, receive, send, identity_error) return + parent_context = ( + checkpoint_parent_context(source_capture_key, parent_model_call_id) + if source_capture_key is not None and parent_model_call_id is not None + else nullcontext() + ) if match is None: if logical_rollout_id is None or attempt_index is None: - await self._app(scope, receive, send) + with parent_context: + await self._app(scope, receive, send) return capture_key = capture_key_for(logical_rollout_id, attempt_index) - with rollout_context( - capture_key, - attempt_index=attempt_index, - logical_rollout_id=logical_rollout_id, - ): - await self._app(scope, receive, send) + with parent_context: + with rollout_context( + capture_key, + attempt_index=attempt_index, + logical_rollout_id=logical_rollout_id, + ): + await self._app(scope, receive, send) return capture_key = match.group("rollout_id") @@ -288,9 +360,10 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None path = match.group("rest") scope = {**scope, "path": path, "raw_path": path.encode()} - with rollout_context( - capture_key, - attempt_index=attempt_index, - logical_rollout_id=logical_rollout_id, - ): - await self._app(scope, receive, send) + with parent_context: + with rollout_context( + capture_key, + attempt_index=attempt_index, + logical_rollout_id=logical_rollout_id, + ): + await self._app(scope, receive, send) diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py index d8f5623a2aa..8a1d8b43577 100644 --- a/nemo_gym/server_utils.py +++ b/nemo_gym/server_utils.py @@ -60,6 +60,7 @@ from nemo_gym import WORKING_DIR from nemo_gym._checkpoint.control import ( + CONTROL_URL_PREFIX, ControlCapabilities, ControlFence, install_control_plane, @@ -87,12 +88,15 @@ from nemo_gym.profiling import Profiler from nemo_gym.rollout_correlation import ( ATTEMPT_INDEX_HEADER, + PARENT_MODEL_CALL_ID_HEADER, ROLLOUT_ID_HEADER, + SOURCE_CAPTURE_KEY_HEADER, current_attempt_index, current_logical_rollout_id, current_rollout_id, execution_identity_from_run_body, maybe_rollout_id_from_run_body, + take_checkpoint_parent, ) from nemo_gym.telemetry._fallbacks import is_span_group_enabled, safe_set_span_attributes from nemo_gym.telemetry.span_groups import GymSpanGroup @@ -647,6 +651,7 @@ def load_from_global_config(cls, head_server_config: Optional[BaseServerConfig] async def request( self, server_name: str, url_path: str, method: str, **kwargs: Unpack[_RequestOptions] ) -> ClientResponse: + request_path = url_path.partition("?")[0] model_server_name = getenv(NEMO_GYM_MODEL_SERVER_NAME_ENV_VAR_NAME) model_server_base_url = getenv(NEMO_GYM_MODEL_SERVER_BASE_URL_ENV_VAR_NAME) if model_server_base_url and server_name == model_server_name: @@ -698,6 +703,31 @@ async def request( headers[ATTEMPT_INDEX_HEADER] = str(attempt_index) kwargs["headers"] = headers + is_policy_generation = ( + server_entry is not None + and "responses_api_models" in server_entry + and request_path.endswith(("/v1/responses", "/v1/chat/completions", "/v1/messages")) + and get_first_server_config_dict(self.global_config_dict, server_name).get("instance_role", "policy") + == "policy" + ) + is_agent_parent_relay = ( + server_entry is not None + and "responses_api_agents" in server_entry + and request_path.endswith("/v1/responses") + ) + source_capture_key, parent_model_call_id = ( + take_checkpoint_parent() if is_policy_generation or is_agent_parent_relay else (None, None) + ) + if source_capture_key is not None and parent_model_call_id is not None: + headers = dict(kwargs.get("headers") or {}) + if headers.get(SOURCE_CAPTURE_KEY_HEADER, source_capture_key) != source_capture_key: + raise ValueError("caller-supplied source capture key disagrees with the restored continuation") + if headers.get(PARENT_MODEL_CALL_ID_HEADER, parent_model_call_id) != parent_model_call_id: + raise ValueError("caller-supplied parent model-call ID disagrees with the restored continuation") + headers[SOURCE_CAPTURE_KEY_HEADER] = source_capture_key + headers[PARENT_MODEL_CALL_ID_HEADER] = parent_model_call_id + kwargs["headers"] = headers + return await request(method=method, url=f"{base_url}{url_path}", _internal=True, **kwargs) async def get( @@ -1279,6 +1309,9 @@ def run_webserver(cls) -> Optional[FastAPI]: # pragma: no cover server.setup_telemetry() app = server.setup_webserver() + capabilities_path = f"{CONTROL_URL_PREFIX}/capabilities" + if not any(getattr(route, "path", None) == capabilities_path for route in app.routes): + server.setup_control_plane(app) # After the app is fully built so subclass routes are present. Only resources servers expose tools over MCP, # so gating the lazy import on their config keeps the MCP SDK out of agent/model processes that never need it. if getattr(getattr(server, "config", None), "expose_tools_over_mcp", False): diff --git a/nemo_gym/token_id_capture/lineage.py b/nemo_gym/token_id_capture/lineage.py index db5744fafb3..65d8f6ebc2a 100644 --- a/nemo_gym/token_id_capture/lineage.py +++ b/nemo_gym/token_id_capture/lineage.py @@ -195,6 +195,7 @@ def _manifest_from_rows(rollout_id: str, rows: list[dict]) -> dict: @dataclass class LineageNode: call_id: str + fingerprint: str # ``None`` means the index is metadata-only. # A resolved match loads tokens from ``entry_offset``. cum_tokens: list[int] | None @@ -312,6 +313,7 @@ def add_entry(self, entry: TokenEntry, *, store_tokens: bool = True, entry_offse raise ValueError("delta records require a durable-log-backed lineage store") node = LineageNode( call_id=entry.model_call_id, + fingerprint=entry.continuation_fingerprint, cum_tokens=cumulative_tokens(entry) if store_tokens else None, cum_len=entry.cum_len if entry.cum_len is not None else len(cumulative_tokens(entry)), digest=entry.digest or "", @@ -320,6 +322,9 @@ def add_entry(self, entry: TokenEntry, *, store_tokens: bool = True, entry_offse context_digest=entry.continuation_context_digest, parent_call_id=entry.parent_call_id, prompt_is_delta=entry.prompt_is_delta, + staging_key=str(getattr(entry, "staging_key", "") or ""), + staging_chain=list(getattr(entry, "staging_chain", None) or []), + chain_hash=str(getattr(entry, "chain_hash", "") or ""), ) previous = self.by_call_id.get(entry.model_call_id) if previous is not None: @@ -352,6 +357,7 @@ def record( """ node = LineageNode( call_id=call_id, + fingerprint=assistant_fingerprint(messages), cum_tokens=list(cum_tokens), cum_len=cum_len if cum_len is not None else len(cum_tokens), digest=digest, @@ -448,6 +454,32 @@ def __init__(self, max_rollouts: int = 512, max_tokens: int = 8_000_000) -> None async def resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: return self.index.for_rollout(rollout_id).resolve(request_items) + async def resolve_explicit( + self, + source_capture_key: str, + parent_call_id: str, + request_items: list[dict], + ) -> LineageResolution: + lineage = self.index.for_rollout(source_capture_key) + node = lineage.by_call_id.get(parent_call_id) + if node is None: + return LineageResolution(ParentResolutionStatus.UNRESOLVED, reason="explicit_parent_missing") + if node.fingerprint != assistant_fingerprint(request_items) or not lineage._continues(node, request_items): + return LineageResolution(ParentResolutionStatus.UNRESOLVED, reason="explicit_parent_fingerprint_mismatch") + if node.cum_tokens is None: + return LineageResolution(ParentResolutionStatus.UNRESOLVED, reason="explicit_parent_tokens_missing") + return LineageResolution( + ParentResolutionStatus.RESOLVED, + match=LineageMatch( + model_call_id=node.call_id, + cumulative_token_ids=tuple(node.cum_tokens), + digest=node.digest, + staging_chain=tuple(node.staging_chain), + prev_len=node.cum_len, + chain_hash=node.chain_hash, + ), + ) + async def put(self, entry: TokenEntry) -> None: """Publish one committed entry to the worker-local index.""" self.index.for_rollout(entry.rollout_id).add_entry(entry) @@ -687,6 +719,52 @@ def _materialize( async def resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: return await asyncio.to_thread(self._resolve, rollout_id, request_items) + async def resolve_explicit( + self, + source_capture_key: str, + parent_call_id: str, + request_items: list[dict], + ) -> LineageResolution: + return await asyncio.to_thread( + self._resolve_explicit, + source_capture_key, + parent_call_id, + request_items, + ) + + def _resolve_explicit( + self, + source_capture_key: str, + parent_call_id: str, + request_items: list[dict], + ) -> LineageResolution: + with self._rollout_lock(source_capture_key), self._read_locked(source_capture_key): + refs, lineage = self._refresh(source_capture_key) + node = lineage.by_call_id.get(parent_call_id) + if node is None: + return LineageResolution(ParentResolutionStatus.UNRESOLVED, reason="explicit_parent_missing") + if node.fingerprint != assistant_fingerprint(request_items) or not lineage._continues(node, request_items): + return LineageResolution( + ParentResolutionStatus.UNRESOLVED, + reason="explicit_parent_fingerprint_mismatch", + ) + tokens = ( + node.cum_tokens + if node.cum_tokens is not None + else self._materialize(source_capture_key, node, refs, lineage) + ) + return LineageResolution( + ParentResolutionStatus.RESOLVED, + match=LineageMatch( + model_call_id=node.call_id, + cumulative_token_ids=tuple(tokens), + digest=node.digest, + staging_chain=tuple(node.staging_chain), + prev_len=node.cum_len, + chain_hash=node.chain_hash, + ), + ) + def _resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolution: with self._rollout_lock(rollout_id), self._read_locked(rollout_id): refs, lineage = self._refresh(rollout_id) @@ -702,6 +780,9 @@ def _resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolut model_call_id=node.call_id, cumulative_token_ids=tuple(tokens), digest=node.digest, + staging_chain=tuple(node.staging_chain), + prev_len=node.cum_len, + chain_hash=node.chain_hash, ), ) @@ -859,6 +940,51 @@ def _resolve(self, rollout_id: str, request_items: list[dict]) -> LineageResolut return LineageResolution(ParentResolutionStatus.RESOLVED, match=match) return resolution + def _resolve_explicit( + self, + source_capture_key: str, + parent_call_id: str, + request_items: list[dict], + ) -> LineageResolution: + resolution = super()._resolve_explicit(source_capture_key, parent_call_id, request_items) + if resolution.status == ParentResolutionStatus.RESOLVED: + return resolution + with self._locked(source_capture_key): + records = [ + record + for record in self._read(source_capture_key) + if record.get("model_call_id") == parent_call_id and record.get("failure_reason") is None + ] + if len(records) != 1: + return LineageResolution(ParentResolutionStatus.UNRESOLVED, reason="explicit_parent_missing") + record = records[0] + fingerprint = assistant_fingerprint(request_items) + if record.get("fingerprint") != fingerprint: + return LineageResolution( + ParentResolutionStatus.UNRESOLVED, + reason="explicit_parent_fingerprint_mismatch", + ) + context_len = int(record["context_len"]) + if ( + len(request_items) < context_len + or conversation_digest(request_items[:context_len]) != record["context_digest"] + ): + return LineageResolution( + ParentResolutionStatus.UNRESOLVED, + reason="explicit_parent_fingerprint_mismatch", + ) + return LineageResolution( + ParentResolutionStatus.RESOLVED, + match=LineageMatch( + model_call_id=str(record["model_call_id"]), + cumulative_token_ids=tuple(int(token) for token in record.get("cumulative_token_ids") or ()), + digest=str(record["digest"]), + staging_chain=tuple(record.get("staging_chain") or []), + prev_len=int(record.get("cum_len") or 0), + chain_hash=str(record.get("chain_hash") or ""), + ), + ) + def _resolve_row(self, rollout_id: str, request_items: list[dict]) -> LineageMatch | None: fingerprint = assistant_fingerprint(request_items) if not fingerprint: diff --git a/nemo_gym/token_id_capture/protocols.py b/nemo_gym/token_id_capture/protocols.py index 4c630ec1ffd..0027ad1c1b1 100644 --- a/nemo_gym/token_id_capture/protocols.py +++ b/nemo_gym/token_id_capture/protocols.py @@ -126,6 +126,23 @@ async def close(self) -> None: ... +@runtime_checkable +class ExplicitParentLineageStore(Protocol): + """Optional exact-parent lookup for restored cross-attempt continuations. + + Implementations must verify that ``request_items`` continue the named + parent. A resolved match must preserve staged-prefix coordinates: + ``staging_chain``, ``prev_len``, and ``chain_hash``. + """ + + async def resolve_explicit( + self, + source_capture_key: str, + parent_call_id: str, + request_items: list[dict], + ) -> LineageResolution: ... + + @runtime_checkable class CaptureLedger(LineageResolver, Protocol): """Store metadata for calls whose token data is staged externally. diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index 1cd7fa3c721..48fcd97d819 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -35,7 +35,13 @@ from nemo_gym.token_id_capture.fingerprint import assistant_fingerprint from nemo_gym.token_id_capture.lineage import stamp_continuation -from nemo_gym.token_id_capture.protocols import CaptureLedger, LineageResolution, LineageResolver, TokenSink +from nemo_gym.token_id_capture.protocols import ( + CaptureLedger, + ExplicitParentLineageStore, + LineageResolution, + LineageResolver, + TokenSink, +) from nemo_gym.token_id_capture.records import ( UNRESOLVED_PARENT_REASON, ParentResolutionStatus, @@ -94,6 +100,9 @@ class CaptureContext: # store doubles as the rollout's capture ledger and admission is the # strict tri-state of the lineage result. external_staging: bool = False + logical_request_id: str | None = None + source_capture_key: str | None = None + explicit_parent_call_id: str | None = None # Stamped once when the middleware admits the call. The ledger row reuses # this value on every commit retry so idempotent re-records stay # byte-identical. @@ -201,11 +210,11 @@ async def resolve_parent(request_messages: list | None) -> None: context = _CAPTURE_CONTEXT.get() if context is None or request_messages is None: return + if (context.source_capture_key is None) != (context.explicit_parent_call_id is None): + raise RuntimeError("source capture key and parent model-call ID must be provided together") context.request_items = list(request_messages) try: - if not assistant_fingerprint(request_messages): - context.parent_resolution = LineageResolution(ParentResolutionStatus.ROOT) - elif context.lineage_store is None: + if context.lineage_store is None: context.parent_resolution = LineageResolution( ParentResolutionStatus.UNRESOLVED, reason="resolver_unavailable", @@ -220,6 +229,20 @@ async def resolve_parent(request_messages: list | None) -> None: "No lineage resolver is available: every continuation resolves UNRESOLVED " "and multi-call rollouts will be masked (allow_unresolved_continuations is set)." ) + elif context.source_capture_key is not None and context.explicit_parent_call_id is not None: + if not isinstance(context.lineage_store, ExplicitParentLineageStore): + context.parent_resolution = LineageResolution( + ParentResolutionStatus.UNRESOLVED, + reason="explicit_parent_unsupported", + ) + else: + context.parent_resolution = await context.lineage_store.resolve_explicit( + context.source_capture_key, + context.explicit_parent_call_id, + request_messages, + ) + elif not assistant_fingerprint(request_messages): + context.parent_resolution = LineageResolution(ParentResolutionStatus.ROOT) else: context.parent_resolution = await context.lineage_store.resolve(context.rollout_id, request_messages) _count_resolution(context.parent_resolution.status.value) @@ -277,7 +300,7 @@ 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): + if is_root or (context.source_capture_key is None and not await ledger.has_rows(context.rollout_id)): context.capture_admission = CaptureAdmission( rollout_id=context.rollout_id, model_call_id=context.model_call_id, diff --git a/tests/unit_tests/test_checkpoint_admission.py b/tests/unit_tests/test_checkpoint_admission.py index f1be5e85ef0..93c29f026bb 100644 --- a/tests/unit_tests/test_checkpoint_admission.py +++ b/tests/unit_tests/test_checkpoint_admission.py @@ -76,6 +76,7 @@ def test_new_root_operation_parks_while_draining() -> None: limiter.close() with pytest.raises(AdmissionParkedError): limiter.admit(rollout_id="9-9", attempt_index=0) + assert ("9-9", 0) not in limiter.seen_attempts() limiter.release(held) @@ -133,6 +134,73 @@ async def operation() -> None: assert limiter.counts()["inflight_total"] == 0 +@pytest.mark.asyncio +async def test_abort_before_response_start_returns_stale_attempt_response() -> None: + limiter = AdmissionLimiter() + entered = asyncio.Event() + + async def app(scope, receive, send) -> None: + entered.set() + await asyncio.Event().wait() + + middleware = AdmissionMiddleware(app, limiter, GATED_MODEL_ROUTE_SUFFIXES) + messages = [] + + async def send(message) -> None: + messages.append(message) + + scope = { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [ + (ROLLOUT_ID_HEADER.encode(), b"rollout-a"), + (ATTEMPT_INDEX_HEADER.encode(), b"0"), + ], + } + task = asyncio.create_task(middleware(scope, None, send)) + await entered.wait() + limiter.close() + limiter.abort_inflight("rollout-a", 0) + await task + assert messages[0]["status"] == 409 + assert b"stale_attempt" in messages[1]["body"] + + +@pytest.mark.asyncio +async def test_abort_after_response_start_closes_transport_without_success_body() -> None: + limiter = AdmissionLimiter() + started = asyncio.Event() + + async def app(scope, receive, send) -> None: + await send({"type": "http.response.start", "status": 200, "headers": []}) + started.set() + await asyncio.Event().wait() + + middleware = AdmissionMiddleware(app, limiter, GATED_MODEL_ROUTE_SUFFIXES) + messages = [] + + async def send(message) -> None: + messages.append(message) + + scope = { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [ + (ROLLOUT_ID_HEADER.encode(), b"rollout-a"), + (ATTEMPT_INDEX_HEADER.encode(), b"0"), + ], + } + task = asyncio.create_task(middleware(scope, None, send)) + await started.wait() + limiter.close() + limiter.abort_inflight("rollout-a", 0) + with pytest.raises(asyncio.CancelledError): + await task + assert messages == [{"type": "http.response.start", "status": 200, "headers": []}] + + @pytest.mark.asyncio async def test_wait_for_drained_long_poll() -> None: limiter = AdmissionLimiter() diff --git a/tests/unit_tests/test_checkpoint_ledger.py b/tests/unit_tests/test_checkpoint_ledger.py index 68788a04dd5..8f781402b0f 100644 --- a/tests/unit_tests/test_checkpoint_ledger.py +++ b/tests/unit_tests/test_checkpoint_ledger.py @@ -15,6 +15,7 @@ """Checkpoint token-free model custody without copying staged token arrays.""" import json +import shutil import pytest from fastapi import FastAPI @@ -147,6 +148,7 @@ def _participant(root) -> tuple[TestClient, AdmissionLimiter]: ledger_provider=lambda: ledger, file_ledger_root_provider=lambda: ledger.checkpoint_root, instance_role="policy", + server_name="policy", auth_token=AUTH_TOKEN, ) return TestClient(app), limiter @@ -186,6 +188,19 @@ def test_commit_requires_completed_drain_and_restore_stays_paused(tmp_path) -> N headers=AUTH_HEADERS, ) assert commit.status_code == 200 + retry = source_client.post( + f"{MODEL_CHECKPOINT_URL_PREFIX}/commit", + json=commit_body, + headers=AUTH_HEADERS, + ) + assert retry.json() == commit.json() + committed_status = source_client.get( + f"{MODEL_ADMISSION_URL_PREFIX}/status", + params=pause_body, + headers=AUTH_HEADERS, + ) + assert committed_status.status_code == 200 + assert committed_status.json()["state"] == "paused" restored_client, restored_limiter = _participant(tmp_path / "ledger-b") restore = restored_client.post( @@ -199,6 +214,15 @@ def test_commit_requires_completed_drain_and_restore_stays_paused(tmp_path) -> N ) assert restore.status_code == 200 assert restored_limiter.counts()["state"] == "paused" + restored_status = restored_client.get( + f"{MODEL_ADMISSION_URL_PREFIX}/status", + params={"checkpoint_id": "restore-1", "deadline_ts": 4e9}, + headers=AUTH_HEADERS, + ) + assert restored_status.status_code == 200 + assert restored_status.json()["state"] == "paused" + with pytest.raises(StaleAttemptError): + restored_limiter.admit(rollout_id="rollout-a", attempt_index=0) def test_restored_tombstone_fences_exact_attempt(tmp_path) -> None: @@ -247,3 +271,161 @@ def fail_manifest_replace(source_path, target_path) -> None: with pytest.raises(RuntimeError, match="injected"): CaptureLedgerCheckpointer(source).commit(checkpoint, checkpoint_id="checkpoint-1", tombstones=[]) assert not (checkpoint / MODEL_LEDGER_SUBDIR / LEDGER_MANIFEST_NAME).exists() + + +def test_model_checkpoint_artifacts_are_namespaced_by_server(tmp_path) -> None: + first = tmp_path / "first" + second = tmp_path / "second" + _write_custody(first, "rollout-a") + _write_custody(second, "rollout-b") + checkpoint = tmp_path / "checkpoint" + + CaptureLedgerCheckpointer(first, server_name="policy-a").commit( + checkpoint, + checkpoint_id="checkpoint-1", + tombstones=[], + ) + CaptureLedgerCheckpointer(second, server_name="policy-b").commit( + checkpoint, + checkpoint_id="checkpoint-1", + tombstones=[], + ) + + root = checkpoint / MODEL_LEDGER_SUBDIR + assert (root / "policy-a" / "rollout-a.lineage.jsonl").exists() + assert (root / "policy-b" / "rollout-b.lineage.jsonl").exists() + with pytest.raises(ValueError, match="model server name"): + CaptureLedgerCheckpointer(first, server_name="../policy") + + shutil.copytree(root / "policy-a", root / "policy-copy") + with pytest.raises(LedgerMismatchError, match="different model server"): + CaptureLedgerCheckpointer(tmp_path / "restore-copy", server_name="policy-copy").restore(checkpoint) + + +@pytest.mark.parametrize( + ("retry_tombstones", "retry_source_attempts", "message"), + [ + ([("rollout-b", 1)], [("rollout-a", 0)], "abort exclusions"), + ([], [("rollout-a", 0), ("rollout-c", 2)], "source attempts"), + ], +) +def test_commit_retry_rejects_changed_semantic_sets_after_final_fsync_failure( + tmp_path, + monkeypatch, + retry_tombstones, + retry_source_attempts, + message, +) -> None: + source = tmp_path / "source" + _write_custody(source, "rollout-a") + checkpoint = tmp_path / "checkpoint" + checkpointer = CaptureLedgerCheckpointer(source, server_name="policy") + + import nemo_gym._checkpoint.ledger as ledger_module + + real_fsync_dir = ledger_module._fsync_dir + calls = 0 + + def fail_final_fsync(path) -> None: + nonlocal calls + calls += 1 + if calls == 2: + raise OSError("injected final directory fsync failure") + real_fsync_dir(path) + + monkeypatch.setattr(ledger_module, "_fsync_dir", fail_final_fsync) + with pytest.raises(OSError, match="injected"): + checkpointer.commit( + checkpoint, + checkpoint_id="checkpoint-1", + tombstones=[], + source_attempts=[("rollout-a", 0)], + ) + assert (checkpoint / MODEL_LEDGER_SUBDIR / "policy" / LEDGER_MANIFEST_NAME).exists() + + with pytest.raises(LedgerMismatchError, match=message): + checkpointer.commit( + checkpoint, + checkpoint_id="checkpoint-1", + tombstones=retry_tombstones, + source_attempts=retry_source_attempts, + ) + + +def test_restored_source_ledger_survives_the_next_commit(tmp_path) -> None: + source_client, source_limiter = _participant(tmp_path / "source") + expected = _write_custody(tmp_path / "source", "rollout-a") + source_limiter.release(source_limiter.admit(rollout_id="rollout-a", attempt_index=0)) + pause = {"checkpoint_id": "checkpoint-1", "deadline_ts": 4e9} + source_client.post(f"{MODEL_ADMISSION_URL_PREFIX}/pause", json=pause, headers=AUTH_HEADERS) + first_checkpoint = tmp_path / "checkpoint-1" + assert ( + source_client.post( + f"{MODEL_CHECKPOINT_URL_PREFIX}/commit", + json={**pause, "checkpoint_dir": str(first_checkpoint)}, + headers=AUTH_HEADERS, + ).status_code + == 200 + ) + + restored_client, restored_limiter = _participant(tmp_path / "restored") + restore = {"checkpoint_id": "restore-1", "deadline_ts": 4e9} + assert ( + restored_client.post( + f"{MODEL_CHECKPOINT_URL_PREFIX}/restore", + json={**restore, "checkpoint_dir": str(first_checkpoint)}, + headers=AUTH_HEADERS, + ).status_code + == 200 + ) + assert ( + restored_client.post( + f"{MODEL_ADMISSION_URL_PREFIX}/resume", + json=restore, + headers=AUTH_HEADERS, + ).status_code + == 200 + ) + with pytest.raises(StaleAttemptError): + restored_limiter.admit(rollout_id="rollout-a", attempt_index=0) + + second = {"checkpoint_id": "checkpoint-2", "deadline_ts": 4e9} + restored_client.post(f"{MODEL_ADMISSION_URL_PREFIX}/pause", json=second, headers=AUTH_HEADERS) + second_checkpoint = tmp_path / "checkpoint-2" + assert ( + restored_client.post( + f"{MODEL_CHECKPOINT_URL_PREFIX}/commit", + json={**second, "checkpoint_dir": str(second_checkpoint)}, + headers=AUTH_HEADERS, + ).status_code + == 200 + ) + assert (second_checkpoint / MODEL_LEDGER_SUBDIR / "policy" / "rollout-a.lineage.jsonl").read_bytes() == expected + + +def test_failed_restore_is_paused_observable_and_recoverable(tmp_path) -> None: + client, limiter = _participant(tmp_path / "restored") + body = { + "checkpoint_id": "restore-1", + "deadline_ts": 4e9, + "checkpoint_dir": str(tmp_path / "missing"), + } + failed = client.post(f"{MODEL_CHECKPOINT_URL_PREFIX}/restore", json=body, headers=AUTH_HEADERS) + assert failed.status_code == 409 + assert limiter.state.value == "paused" + capabilities = client.get("/ng-control/v1/capabilities").json() + assert capabilities["phase"] == "restore_failed_paused" + status = client.get( + f"{MODEL_ADMISSION_URL_PREFIX}/status", + params={"checkpoint_id": "restore-1", "deadline_ts": 4e9}, + headers=AUTH_HEADERS, + ) + assert status.status_code == 200 + assert status.json()["state"] == "paused" + resumed = client.post( + f"{MODEL_ADMISSION_URL_PREFIX}/resume", + json={"checkpoint_id": "restore-1", "deadline_ts": 4e9}, + headers=AUTH_HEADERS, + ) + assert resumed.status_code == 200 + assert limiter.state.value == "accepting" diff --git a/tests/unit_tests/test_server_utils.py b/tests/unit_tests/test_server_utils.py index 05f31be7c80..31bc4d831ad 100644 --- a/tests/unit_tests/test_server_utils.py +++ b/tests/unit_tests/test_server_utils.py @@ -38,6 +38,12 @@ NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME, NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME, ) +from nemo_gym.rollout_correlation import ( + PARENT_MODEL_CALL_ID_HEADER, + SOURCE_CAPTURE_KEY_HEADER, + RolloutContextMiddleware, + checkpoint_parent_context, +) from nemo_gym.server_utils import ( NEMO_GYM_MODEL_SERVER_BASE_URL_ENV_VAR_NAME, NEMO_GYM_MODEL_SERVER_NAME_ENV_VAR_NAME, @@ -302,14 +308,65 @@ async def test_ServerClient_get_post_sanity(self, monkeypatch: MonkeyPatch) -> N url_path="blah blah", ) assert "my mock response" == actual_response - actual_response = await server_client.post( server_name="my_server", url_path="blah blah", ) assert "my mock response" == actual_response - async def test_ServerClient_preserves_external_capture_url(self, monkeypatch: MonkeyPatch) -> None: + async def test_ServerClient_propagates_restored_parent_headers(self, monkeypatch: MonkeyPatch) -> None: + server_client = ServerClient( + head_server_config=BaseServerConfig(host="head", port=80), + global_config_dict=DictConfig( + { + "agent": {"responses_api_agents": {"agent": {"host": "agent", "port": 80}}}, + "judge": { + "responses_api_models": {"model": {"host": "judge", "port": 80, "instance_role": "auxiliary"}} + }, + "policy": {"responses_api_models": {"model": {"host": "policy", "port": 80}}}, + } + ), + ) + request_mock = AsyncMock(return_value="response") + monkeypatch.setattr(nemo_gym.server_utils, "request", request_mock) + + with checkpoint_parent_context("rollout-a2", "call-1"): + await server_client.post( + server_name="agent", + url_path="/ng-rollout/rollout-a3/v1/responses", + json={}, + ) + + relay_headers = request_mock.await_args.kwargs["headers"] + + async def relayed_agent_request(scope, receive, send) -> None: + await server_client.post(server_name="judge", url_path="/v1/responses", json={}) + await server_client.post(server_name="policy", url_path="/v1/responses", json={}) + await server_client.post(server_name="policy", url_path="/v1/responses", json={}) + + await RolloutContextMiddleware(relayed_agent_request)( + { + "type": "http", + "path": "/ng-rollout/rollout-a3/v1/responses", + "headers": [(name.encode(), value.encode()) for name, value in relay_headers.items()], + }, + None, + None, + ) + + calls = request_mock.await_args_list + relay_headers = calls[0].kwargs["headers"] + assert relay_headers[SOURCE_CAPTURE_KEY_HEADER] == "rollout-a2" + assert relay_headers[PARENT_MODEL_CALL_ID_HEADER] == "call-1" + assert SOURCE_CAPTURE_KEY_HEADER not in calls[1].kwargs["headers"] + assert PARENT_MODEL_CALL_ID_HEADER not in calls[1].kwargs["headers"] + first_policy_headers = calls[2].kwargs["headers"] + assert first_policy_headers[SOURCE_CAPTURE_KEY_HEADER] == "rollout-a2" + assert first_policy_headers[PARENT_MODEL_CALL_ID_HEADER] == "call-1" + assert SOURCE_CAPTURE_KEY_HEADER not in calls[3].kwargs["headers"] + assert PARENT_MODEL_CALL_ID_HEADER not in calls[3].kwargs["headers"] + + async def test_ServerClient_preserves_external_capture_url_and_capability(self, monkeypatch: MonkeyPatch) -> None: server_client = ServerClient( head_server_config=BaseServerConfig(host="head", port=12345), global_config_dict=DictConfig( diff --git a/tests/unit_tests/test_token_capture_ledger.py b/tests/unit_tests/test_token_capture_ledger.py index eb51b91a84f..2f548afcea8 100644 --- a/tests/unit_tests/test_token_capture_ledger.py +++ b/tests/unit_tests/test_token_capture_ledger.py @@ -210,13 +210,23 @@ async def test_has_rows_is_false_for_untouched_rollout(store): assert RolloutManifest.model_validate(await store.manifest("r-none")).records == [] -async def _admit(store, request_items, rollout_id="r1", model_call_id="c2"): +async def _admit( + store, + request_items, + rollout_id="r1", + model_call_id="c2", + *, + source_capture_key=None, + explicit_parent_call_id=None, +): context = CaptureContext( rollout_id=rollout_id, model_call_id=model_call_id, token_sink=None, lineage_store=store, external_staging=True, + source_capture_key=source_capture_key, + explicit_parent_call_id=explicit_parent_call_id, ) token = set_token_sink(context) try: @@ -242,6 +252,103 @@ async def test_admission_match_uses_staging_chain_without_wire_prefix(store): assert context.request_items == [USER_1, ASSISTANT_1, USER_2] +@pytest.mark.asyncio +async def test_cross_attempt_admission_uses_explicit_verified_parent(store): + await _record_call_1(store, rollout_id="r1") + context = await _admit( + store, + [USER_1, ASSISTANT_1, USER_2], + rollout_id="r1-a1", + source_capture_key="r1", + explicit_parent_call_id="c1", + ) + admission = context.capture_admission + assert admission is not None + assert admission.rollout_id == "r1-a1" + assert admission.parent_call_id == "c1" + assert admission.staging_chain == ["r1/c1"] + assert admission.prev_len == len(TOKENS_1) + assert admission.parent_chain_hash == CHAIN_HASH_1 + + tokens_2 = TOKENS_1 + [901, 902] + chain_hash_2 = compute_chain_hash(CHAIN_HASH_1, [901, 902]) + await store.record( + "r1-a1", + "c2", + [USER_1, ASSISTANT_1, USER_2], + [ASSISTANT_2], + [], + hash_token_ids(tokens_2), + parent_call_id="c1", + staging_key="r1-a1/c2", + weight_version=17, + prev_len=len(TOKENS_1), + delta_len=2, + cum_len=len(tokens_2), + staging_digest=STAGING_DIGEST, + extras_digest=EMPTY_EXTRAS_DIGEST, + mode="token_in", + staging_chain=["r1/c1", "r1-a1/c2"], + chain_hash=chain_hash_2, + cumulative_hash=hash_token_ids(tokens_2), + ) + next_context = await _admit( + store, + [USER_1, ASSISTANT_1, USER_2, ASSISTANT_2, USER_3], + rollout_id="r1-a1", + model_call_id="c3", + ) + next_admission = next_context.capture_admission + assert next_admission is not None + assert next_admission.parent_call_id == "c2" + assert next_admission.staging_chain == ["r1/c1", "r1-a1/c2"] + assert next_admission.prev_len == len(tokens_2) + assert next_admission.parent_chain_hash == chain_hash_2 + + +@pytest.mark.asyncio +async def test_cross_attempt_parent_requires_matching_fingerprint(store): + await _record_call_1(store, rollout_id="r1") + context = await _admit( + store, + [USER_1, ASSISTANT_SEEDED, USER_2], + rollout_id="r1-a1", + source_capture_key="r1", + explicit_parent_call_id="c1", + ) + assert context.capture_admission is None + manifest = RolloutManifest.model_validate(await store.manifest("r1-a1")) + assert [failure.reason for failure in manifest.failures] == [UNRESOLVED_PARENT_REASON] + + +@pytest.mark.asyncio +async def test_custom_store_without_explicit_lookup_fails_closed(): + class LegacyCustomStore: + def __init__(self) -> None: + self.inner = InMemoryLineageStore() + + async def resolve(self, rollout_id, request_items): + return await self.inner.resolve(rollout_id, request_items) + + async def has_rows(self, rollout_id): + return await self.inner.has_rows(rollout_id) + + async def record_failure(self, rollout_id, model_call_id, reason): + await self.inner.record_failure(rollout_id, model_call_id, reason) + + store = LegacyCustomStore() + context = await _admit( + store, + [USER_1, ASSISTANT_1, USER_2], + rollout_id="r1-a1", + source_capture_key="r1", + explicit_parent_call_id="c1", + ) + assert context.capture_admission is None + manifest = RolloutManifest.model_validate(await store.inner.manifest("r1-a1")) + assert [failure.reason for failure in manifest.failures] == [UNRESOLVED_PARENT_REASON] + + @pytest.mark.asyncio async def test_staging_chain_grows_across_external_calls(store): await _record_call_1(store) From 9c6e3e297af7ca229ba9c04341911eb2396163b4 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Thu, 10 Sep 2026 13:29:23 -0700 Subject: [PATCH 3/5] feat(checkpoint): bind model cuts to capture ledger Signed-off-by: Ananth Subramaniam --- nemo_gym/_checkpoint/ledger.py | 224 +++++++++++- nemo_gym/_checkpoint/model_admission.py | 14 +- nemo_gym/base_responses_api_model.py | 27 +- nemo_gym/token_id_capture/__init__.py | 2 + nemo_gym/token_id_capture/external_capture.py | 13 + nemo_gym/token_id_capture/sink.py | 31 ++ tests/unit_tests/test_checkpoint_ledger.py | 324 +++++++++++++++++- 7 files changed, 612 insertions(+), 23 deletions(-) diff --git a/nemo_gym/_checkpoint/ledger.py b/nemo_gym/_checkpoint/ledger.py index 6bf2af34c3c..4a16bc963bc 100644 --- a/nemo_gym/_checkpoint/ledger.py +++ b/nemo_gym/_checkpoint/ledger.py @@ -58,6 +58,10 @@ ControlFence, ) from nemo_gym._checkpoint.model_admission import NotPolicyInstanceError +from nemo_gym._checkpoint.model_control_contracts import ( + GenerationCutCoordinatorProof, + GenerationCutReceipt, +) from nemo_gym.rollout_correlation import ROLLOUT_ID_PATTERN, capture_key_for from nemo_gym.token_id_capture.control_routes import require_control_auth from nemo_gym.token_id_capture.protocols import CaptureLedger @@ -66,7 +70,9 @@ MODEL_CHECKPOINT_URL_PREFIX = f"{CONTROL_URL_PREFIX}/model-checkpoint" MODEL_LEDGER_SUBDIR = "model-ledger" LEDGER_MANIFEST_NAME = "manifest.json" -LEDGER_SCHEMA_VERSION = 1 +GENERATION_CUT_COORDINATOR_PROOF_NAME = "generation-cut-workers.json" +LEGACY_GENERATION_CUT_ACK_NAME = "generation-cut.json" +LEDGER_SCHEMA_VERSION = 2 # FileLineageStore writes one token-free custody file per rollout. # Lock files and token-store files are not part of this participant. @@ -111,6 +117,8 @@ class CaptureLedgerCommitResult(BaseModel): rows: int = Field(ge=0) excluded_tombstoned: int = Field(ge=0) manifest_digest: str + generation_cut_receipt: GenerationCutReceipt | None = None + generation_cut_proof: GenerationCutCoordinatorProof | None = None class CaptureLedgerRestoreResult(BaseModel): @@ -121,6 +129,8 @@ class CaptureLedgerRestoreResult(BaseModel): checkpoint_id: Optional[str] = None tombstones: list[AttemptIdentity] = Field(default_factory=list) source_attempts: list[AttemptIdentity] = Field(default_factory=list) + generation_cut_receipt: GenerationCutReceipt | None = None + generation_cut_proof: GenerationCutCoordinatorProof | None = None @runtime_checkable @@ -181,6 +191,32 @@ def _fsync_dir(path: Path) -> None: os.close(descriptor) +def _validate_ledger_schema( + manifest: dict[str, Any], + ledger_dir: Path, + *, + commit_retry: bool = False, +) -> None: + version = manifest.get("schema_version") + if not isinstance(version, int) or version < 1: + raise LedgerMismatchError(f"ledger manifest has invalid schema_version {version!r}") + if version > LEDGER_SCHEMA_VERSION: + raise LedgerMismatchError( + f"ledger manifest schema_version {version} is newer than this reader ({LEDGER_SCHEMA_VERSION})" + ) + legacy_cut = "generation_cut_ack" in manifest or (ledger_dir / LEGACY_GENERATION_CUT_ACK_NAME).exists() + if legacy_cut: + raise LedgerMismatchError( + "generation-cut checkpoint schema v1 cannot be migrated safely because its ack and sidecar " + "did not prove a final durable backend snapshot; recreate the checkpoint with schema_version 2" + ) + if commit_retry and version != LEDGER_SCHEMA_VERSION: + raise LedgerMismatchError( + f"cannot retry a schema_version {version} ledger commit with schema_version {LEDGER_SCHEMA_VERSION}; " + "restore the old checkpoint first and create a new checkpoint" + ) + + class CaptureLedgerCheckpointer: """Commit and restore one token-capture store directory.""" @@ -202,12 +238,22 @@ def commit( checkpoint_id: str, tombstones: list[tuple[str, int]], source_attempts: Optional[list[tuple[str, int]]] = None, + generation_cut_receipt: GenerationCutReceipt | None = None, + generation_cut_proof: GenerationCutCoordinatorProof | None = None, ) -> dict[str, Any]: """Copy the ledger into ``checkpoint_dir``; the caller has already drained. The store must be quiescent (admission paused) when this runs: the copy takes no locks because nothing may be writing. """ + if generation_cut_receipt is not None and generation_cut_proof is not None: + raise ValueError("single-worker cut receipt and coordinator cut proof are mutually exclusive") + if generation_cut_proof is not None: + generation_cut_proof = GenerationCutCoordinatorProof.model_validate( + generation_cut_proof.model_dump(mode="json") + ) + if generation_cut_proof.checkpoint_id != checkpoint_id: + raise ValueError("coordinator generation-cut proof belongs to a different checkpoint") ledger_dir = self._ledger_dir(checkpoint_dir) if (ledger_dir / LEDGER_MANIFEST_NAME).exists(): result = self._validate_committed( @@ -216,6 +262,8 @@ def commit( server_name=self.server_name, tombstones=tombstones, source_attempts=source_attempts or [], + generation_cut_receipt=generation_cut_receipt, + generation_cut_proof=generation_cut_proof, ) # A previous attempt may have renamed the manifest and then # failed its final directory fsync. Retry that durability barrier. @@ -255,6 +303,10 @@ def commit( for rollout_id, attempt in sorted(source_attempts or []) ], } + if generation_cut_receipt is not None: + manifest["generation_cut_receipt"] = generation_cut_receipt.model_dump(mode="json") + if generation_cut_proof is not None: + manifest["generation_cut_proof"] = generation_cut_proof.model_dump(mode="json") payload = json.dumps(manifest, sort_keys=True, indent=1).encode() with tempfile.NamedTemporaryFile(dir=ledger_dir, prefix=".manifest-", delete=False) as handle: temporary = Path(handle.name) @@ -268,12 +320,17 @@ def commit( os.replace(temporary, ledger_dir / LEDGER_MANIFEST_NAME) _fsync_dir(ledger_dir) - return { + result = { "rollouts": len(rollouts), "rows": total_rows, "excluded_tombstoned": excluded, "manifest_digest": hashlib.sha256(payload).hexdigest(), } + if generation_cut_receipt is not None: + result["generation_cut_receipt"] = generation_cut_receipt.model_dump(mode="json") + if generation_cut_proof is not None: + result["generation_cut_proof"] = generation_cut_proof.model_dump(mode="json") + return result @staticmethod def _validate_committed( @@ -283,10 +340,13 @@ def _validate_committed( server_name: Optional[str], tombstones: list[tuple[str, int]], source_attempts: list[tuple[str, int]], + generation_cut_receipt: GenerationCutReceipt | None, + generation_cut_proof: GenerationCutCoordinatorProof | None, ) -> dict[str, Any]: manifest_path = ledger_dir / LEDGER_MANIFEST_NAME payload = manifest_path.read_bytes() manifest = json.loads(payload) + _validate_ledger_schema(manifest, ledger_dir, commit_retry=True) if manifest.get("checkpoint_id") != checkpoint_id or manifest.get("server_name") != server_name: raise LedgerMismatchError("ledger directory belongs to a different checkpoint transaction or model server") expected_tombstones = [ @@ -299,6 +359,12 @@ def _validate_committed( raise LedgerMismatchError("committed ledger abort exclusions changed before commit retry") if manifest.get("source_attempts", []) != expected_source_attempts: raise LedgerMismatchError("committed ledger source attempts changed before commit retry") + expected_cut = generation_cut_receipt.model_dump(mode="json") if generation_cut_receipt is not None else None + if manifest.get("generation_cut_receipt") != expected_cut: + raise LedgerMismatchError("committed ledger generation cut changed before commit retry") + expected_proof = generation_cut_proof.model_dump(mode="json") if generation_cut_proof is not None else None + if manifest.get("generation_cut_proof") != expected_proof: + raise LedgerMismatchError("committed ledger worker generation-cut proof changed before commit retry") total_rows = 0 for rollout_id, metadata in manifest.get("rollouts", {}).items(): for name, digest in metadata.get("files", {}).items(): @@ -306,12 +372,17 @@ def _validate_committed( if not path.exists() or _file_digest(path) != digest: raise LedgerMismatchError(f"committed ledger file {name!r} for {rollout_id!r} is corrupted") total_rows += int(metadata.get("rows", 0)) - return { + result = { "rollouts": len(manifest.get("rollouts", {})), "rows": total_rows, "excluded_tombstoned": len(manifest.get("tombstones", [])), "manifest_digest": hashlib.sha256(payload).hexdigest(), } + if manifest.get("generation_cut_receipt") is not None: + result["generation_cut_receipt"] = manifest["generation_cut_receipt"] + if manifest.get("generation_cut_proof") is not None: + result["generation_cut_proof"] = manifest["generation_cut_proof"] + return result def restore(self, checkpoint_dir: Path) -> dict[str, Any]: """Install a committed ledger into this store root and verify it.""" @@ -323,11 +394,7 @@ def restore(self, checkpoint_dir: Path) -> dict[str, Any]: f"so this directory must not be installed" ) manifest = json.loads(manifest_path.read_text()) - if manifest.get("schema_version", 0) > LEDGER_SCHEMA_VERSION: - raise LedgerMismatchError( - f"ledger manifest schema_version {manifest.get('schema_version')} is newer than this " - f"reader ({LEDGER_SCHEMA_VERSION})" - ) + _validate_ledger_schema(manifest, ledger_dir) if manifest.get("server_name") != self.server_name: raise LedgerMismatchError("ledger checkpoint belongs to a different model server") @@ -359,13 +426,20 @@ def restore(self, checkpoint_dir: Path) -> dict[str, Any]: _copy_fsynced(source, self.store_root / name) _fsync_dir(self.store_root) - return { + result = { "rollouts": len(manifest["rollouts"]), "rows": total_rows, "checkpoint_id": manifest.get("checkpoint_id"), "tombstones": list(manifest.get("tombstones", ())), "source_attempts": list(manifest.get("source_attempts", ())), } + if manifest.get("generation_cut_receipt") is not None: + receipt = GenerationCutReceipt.model_validate(manifest["generation_cut_receipt"]) + result["generation_cut_receipt"] = receipt.model_dump(mode="json") + if manifest.get("generation_cut_proof") is not None: + proof = GenerationCutCoordinatorProof.model_validate(manifest["generation_cut_proof"]) + result["generation_cut_proof"] = proof.model_dump(mode="json") + return result class ModelCheckpointCommitRequest(CheckpointControlRequest): @@ -385,6 +459,47 @@ def _validate_server_name(server_name: str) -> str: return server_name +async def _restore_generation_cut( + backend: Any, + expected: GenerationCutReceipt | None, +) -> None: + if expected is None: + return + if backend is None: + raise LedgerMismatchError("checkpoint contains a generation cut but no cut backend is configured") + restored = GenerationCutReceipt.model_validate(await backend.restore_generation_cut(expected)) + if restored != expected: + raise LedgerMismatchError("restored generation cut does not match the model-ledger manifest") + + +def _store_generation_cut_proof(directory: Path, proof: GenerationCutCoordinatorProof) -> None: + directory.mkdir(parents=True, exist_ok=True) + path = directory / GENERATION_CUT_COORDINATOR_PROOF_NAME + payload = json.dumps(proof.model_dump(mode="json"), sort_keys=True, indent=1).encode() + if path.exists(): + if path.read_bytes() != payload: + raise LedgerMismatchError("coordinator generation-cut proof changed before checkpoint retry") + return + with tempfile.NamedTemporaryFile(dir=directory, prefix=".generation-cut-workers-", delete=False) as handle: + temporary = Path(handle.name) + try: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + except BaseException: + temporary.unlink(missing_ok=True) + raise + os.replace(temporary, path) + _fsync_dir(directory) + + +def _load_generation_cut_proof(directory: Path) -> GenerationCutCoordinatorProof | None: + path = directory / GENERATION_CUT_COORDINATOR_PROOF_NAME + if not path.exists(): + return None + return GenerationCutCoordinatorProof.model_validate_json(path.read_bytes()) + + def install_model_checkpoint( app: FastAPI, *, @@ -395,6 +510,8 @@ def install_model_checkpoint( instance_role: Literal["policy", "auxiliary"], server_name: str, auth_token: str, + generation_cut_proof_provider: Callable[[], GenerationCutCoordinatorProof] | None = None, + expected_workers: int = 1, ) -> None: """Register ``/ng-control/v1/model-checkpoint`` on a model-server app. @@ -404,6 +521,8 @@ def install_model_checkpoint( """ server_name = _validate_server_name(server_name) + if expected_workers < 1: + raise ValueError("expected_workers must be positive") def _require_policy() -> None: if instance_role != "policy": @@ -412,22 +531,45 @@ def _require_policy() -> None: "it produces no training tokens and has no capture ledger to checkpoint" ) - def _require_quiescent() -> None: + def _require_quiescent(checkpoint_id: str) -> GenerationCutCoordinatorProof | None: + if expected_workers > 1 and generation_cut_proof_provider is None: + raise LedgerNotQuiescentError( + "multi-worker ledger commit requires the coordinator's complete generation-cut proof" + ) + if generation_cut_proof_provider is not None: + try: + supplied_proof = generation_cut_proof_provider() + proof = GenerationCutCoordinatorProof.model_validate(supplied_proof.model_dump(mode="json")) + except (AttributeError, TypeError, ValueError) as error: + raise LedgerNotQuiescentError(f"coordinator generation-cut proof is incomplete: {error}") from error + if proof.checkpoint_id != checkpoint_id: + raise LedgerNotQuiescentError("coordinator generation-cut proof belongs to a different checkpoint") + if proof.expected_workers != expected_workers: + raise LedgerNotQuiescentError( + "coordinator generation-cut proof does not match configured worker membership" + ) + return proof counts = limiter.counts() - if counts["state"] != "paused" or counts["inflight_total"] != 0: + if counts["state"] != "paused" or not limiter.is_prepare_safe(): raise LedgerNotQuiescentError( - "capture-ledger commit requires paused admission and zero in-flight generation requests" + "capture-ledger commit requires paused admission and a generation-safe frozen cut" ) + return None async def _with_ledger( checkpoint_dir: Path, *, checkpoint_id: str, operation: Literal["commit", "restore"], + generation_cut_proof: GenerationCutCoordinatorProof | None = None, ) -> dict[str, Any]: ledger = ledger_provider() + cut_backend = limiter.generation_cut_backend + cut_receipt = None if generation_cut_proof is not None else limiter.generation_cut_receipt + participant_dir = checkpoint_dir / MODEL_LEDGER_SUBDIR / server_name + if operation == "commit" and generation_cut_proof is not None: + await _run_sync(lambda: _store_generation_cut_proof(participant_dir, generation_cut_proof)) if isinstance(ledger, CheckpointableCaptureLedger): - participant_dir = checkpoint_dir / MODEL_LEDGER_SUBDIR / server_name if operation == "commit": result = await ledger.checkpoint_capture_ledger( participant_dir, @@ -436,9 +578,31 @@ async def _with_ledger( tombstones=tuple(limiter.checkpoint_exclusions()), source_attempts=tuple(limiter.seen_attempts()), ) - return CaptureLedgerCommitResult.model_validate(result).model_dump(mode="json") + validated = CaptureLedgerCommitResult.model_validate(result) + if ( + validated.generation_cut_proof is not None + and validated.generation_cut_proof != generation_cut_proof + ): + raise LedgerMismatchError("capture ledger and coordinator generation-cut proof disagree") + validated.generation_cut_proof = generation_cut_proof + if validated.generation_cut_receipt is not None and validated.generation_cut_receipt != cut_receipt: + raise LedgerMismatchError("capture ledger and generation-cut receipt disagree") + validated.generation_cut_receipt = cut_receipt + return validated.model_dump(mode="json") result = await ledger.restore_capture_ledger(participant_dir, server_name=server_name) - return CaptureLedgerRestoreResult.model_validate(result).model_dump(mode="json") + validated = CaptureLedgerRestoreResult.model_validate(result) + sidecar_proof = await _run_sync(lambda: _load_generation_cut_proof(participant_dir)) + if validated.generation_cut_proof is not None and sidecar_proof not in ( + None, + validated.generation_cut_proof, + ): + raise LedgerMismatchError("capture ledger and coordinator generation-cut proof disagree") + validated.generation_cut_proof = validated.generation_cut_proof or sidecar_proof + await _restore_generation_cut( + cut_backend, + validated.generation_cut_receipt, + ) + return validated.model_dump(mode="json") file_root = file_ledger_root_provider() if file_root is None: @@ -448,15 +612,38 @@ async def _with_ledger( ) checkpointer = CaptureLedgerCheckpointer(file_root, server_name=server_name) if operation == "commit": - return await _run_sync( + result = await _run_sync( lambda: checkpointer.commit( checkpoint_dir, checkpoint_id=checkpoint_id, tombstones=limiter.checkpoint_exclusions(), source_attempts=limiter.seen_attempts(), + generation_cut_receipt=cut_receipt, + generation_cut_proof=generation_cut_proof, ) ) - return await _run_sync(lambda: checkpointer.restore(checkpoint_dir)) + return result + result = await _run_sync(lambda: checkpointer.restore(checkpoint_dir)) + sidecar_proof = await _run_sync(lambda: _load_generation_cut_proof(participant_dir)) + manifest_proof = ( + GenerationCutCoordinatorProof.model_validate(result["generation_cut_proof"]) + if result.get("generation_cut_proof") is not None + else None + ) + if sidecar_proof is not None and manifest_proof not in (None, sidecar_proof): + raise LedgerMismatchError("model ledger and coordinator generation-cut proof disagree") + if sidecar_proof is not None: + result["generation_cut_proof"] = sidecar_proof.model_dump(mode="json") + expected_receipt = ( + GenerationCutReceipt.model_validate(result["generation_cut_receipt"]) + if result.get("generation_cut_receipt") is not None + else None + ) + await _restore_generation_cut( + cut_backend, + expected_receipt, + ) + return result @app.post(f"{MODEL_CHECKPOINT_URL_PREFIX}/commit") async def model_checkpoint_commit( @@ -467,11 +654,12 @@ async def model_checkpoint_commit( _require_policy() async def run() -> dict[str, Any]: - _require_quiescent() + generation_cut_proof = _require_quiescent(body.checkpoint_id) return await _with_ledger( Path(body.checkpoint_dir), checkpoint_id=body.checkpoint_id, operation="commit", + generation_cut_proof=generation_cut_proof, ) return await fence.run_operation( diff --git a/nemo_gym/_checkpoint/model_admission.py b/nemo_gym/_checkpoint/model_admission.py index 74834fc6857..0ae634ce8e3 100644 --- a/nemo_gym/_checkpoint/model_admission.py +++ b/nemo_gym/_checkpoint/model_admission.py @@ -59,6 +59,7 @@ def install_model_admission( limiter: AdmissionLimiter, fence: ControlFence, instance_role: Literal["policy", "auxiliary"], + server_name: str = "policy", auth_token: str, ) -> None: """Register ``/ng-control/v1/model-admission`` on a model-server app. @@ -88,7 +89,16 @@ async def model_admission_pause( _require_policy() async def run() -> dict[str, Any]: - limiter.close() + limiter.close(body.checkpoint_id) + try: + await limiter.prepare_generation_cut( + body.checkpoint_id, + server_name=server_name, + timeout_s=body.remaining(), + ) + except BaseException: + limiter.resume() + raise counts = limiter.counts() return { "state": counts["state"], @@ -144,6 +154,8 @@ async def model_admission_status( "state": counts["state"], "per_worker": {"0": {"state": counts["state"], "inflight": counts["inflight_total"]}}, "inflight_total": counts["inflight_total"], + "response_inflight_total": counts["response_inflight_total"], + "generation_pending_total": counts["generation_pending_total"], "waiters_total": counts["waiters_total"], "inflight": counts["inflight"], "tombstones": [ diff --git a/nemo_gym/base_responses_api_model.py b/nemo_gym/base_responses_api_model.py index df790203920..5f16069be33 100644 --- a/nemo_gym/base_responses_api_model.py +++ b/nemo_gym/base_responses_api_model.py @@ -46,7 +46,13 @@ from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field, PrivateAttr, ValidationError, model_validator -from nemo_gym._checkpoint.admission import GATED_MODEL_ROUTE_SUFFIXES, AdmissionLimiter, AdmissionMiddleware +from nemo_gym._checkpoint.admission import ( + GATED_MODEL_ROUTE_SUFFIXES, + AdmissionLimiter, + AdmissionMiddleware, + bind_current_model_call, + mark_current_generation_started, +) from nemo_gym._checkpoint.control import ( AdmissionState, ControlCapabilities, @@ -56,6 +62,7 @@ ) from nemo_gym._checkpoint.ledger import install_model_checkpoint from nemo_gym._checkpoint.model_admission import install_model_admission +from nemo_gym._checkpoint.model_control_contracts import GenerationCutBackend from nemo_gym.anthropic_converter import AnthropicConverter from nemo_gym.chat_streaming import sanitize_streaming_chat_body, synthesize_chat_completion_sse from nemo_gym.config_types import ROLLOUT_PATH_PREFIX, TOKEN_CAPTURE_PATH_SEGMENT, ModelServerRef @@ -94,6 +101,7 @@ current_capture_context, installed_lineage_store, installed_token_sink, + mark_external_staging_failed, register_call_intent, reset_token_sink, resolve_parent, @@ -262,9 +270,13 @@ def setup_webserver(self) -> FastAPI: def admission_limiter(self) -> AdmissionLimiter: if self._admission_limiter is None: - self._admission_limiter = AdmissionLimiter() + self._admission_limiter = AdmissionLimiter(self.generation_cut_backend()) return self._admission_limiter + def generation_cut_backend(self) -> GenerationCutBackend | None: + """Return a framework-owned cut backend, if this server has one.""" + return None + def setup_model_admission(self, app: FastAPI) -> None: """Register admission control on this model server. @@ -279,6 +291,7 @@ def setup_model_admission(self, app: FastAPI) -> None: limiter=self.admission_limiter(), fence=self.checkpoint_fence(), instance_role=self.config.instance_role, + server_name=self.config.name, auth_token=auth_token, ) install_model_checkpoint( @@ -294,6 +307,7 @@ def setup_model_admission(self, app: FastAPI) -> None: instance_role=self.config.instance_role, server_name=self.config.name, auth_token=auth_token, + expected_workers=self.config.num_workers or 1, ) if self.config.instance_role == "policy": app.add_middleware( @@ -423,8 +437,10 @@ async def _invoke_chat_completions( else: request_messages = None if "request" in inspect.signature(self.chat_completions).parameters: + mark_current_generation_started() completion = await self.chat_completions(request=request, body=params) else: + mark_current_generation_started() completion = await self.chat_completions(body=params) await capture_tokens( completion, @@ -469,8 +485,10 @@ async def _invoke_responses( else: request_messages = None if "request" in inspect.signature(self.responses).parameters: + mark_current_generation_started() response = await self.responses(request=request, body=params) else: + mark_current_generation_started() response = await self.responses(body=params) # Capture before streaming dispatch wraps the response. # Anthropic mapping drops the token fields. @@ -1336,6 +1354,10 @@ async def _fail_uncommitted_external_call(context: CaptureContext | None) -> Non context.model_call_id, UNCOMMITTED_CALL_REASON, ) + mark_external_staging_failed( + rollout_id=context.rollout_id, + model_call_id=context.model_call_id, + ) except Exception: logger.exception( "Could not poison uncommitted call %s for rollout %s.", @@ -1439,6 +1461,7 @@ async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None rollout_id = rollout_from_path model_call_id = uuid4().hex client_session_id = _unique_request_header(scope.get("headers") or [], _CLIENT_SESSION_HEADER) + bind_current_model_call(model_call_id) async def _send_with_model_call_id(message: dict[str, Any]) -> None: if message.get("type") == "http.response.start": diff --git a/nemo_gym/token_id_capture/__init__.py b/nemo_gym/token_id_capture/__init__.py index 1b446c3a53d..af45db99c97 100644 --- a/nemo_gym/token_id_capture/__init__.py +++ b/nemo_gym/token_id_capture/__init__.py @@ -96,6 +96,7 @@ commit_entry, current_capture_context, mark_external_staging_committed, + mark_external_staging_failed, register_call_intent, reset_token_sink, resolve_parent, @@ -152,6 +153,7 @@ "commit_entry", "current_capture_context", "mark_external_staging_committed", + "mark_external_staging_failed", "FileLineageStore", "IncrementalLineageStore", "InMemoryLineageStore", diff --git a/nemo_gym/token_id_capture/external_capture.py b/nemo_gym/token_id_capture/external_capture.py index ee8fed83f2d..50f6b5960d2 100644 --- a/nemo_gym/token_id_capture/external_capture.py +++ b/nemo_gym/token_id_capture/external_capture.py @@ -23,6 +23,7 @@ CaptureContext, current_capture_context, mark_external_staging_committed, + mark_external_staging_failed, ) from nemo_gym.token_id_capture.staging.records import ( INVALID_COMMIT_COORDS_REASON, @@ -161,6 +162,10 @@ async def finalize_response(self, served_payload: dict[str, Any]) -> None: context.model_call_id, self._INVALID_CAPTURE_REASON, ) + mark_external_staging_failed( + rollout_id=context.rollout_id, + model_call_id=context.model_call_id, + ) except Exception: LOGGER.exception( f"Could not poison rollout %s call %s after a failed {self._BACKEND_LABEL} worker acknowledgement", @@ -196,6 +201,10 @@ async def _finalize_admitted_response( context.model_call_id, WORKER_MISSING_COMMIT_COORDS_REASON, ) + mark_external_staging_failed( + rollout_id=context.rollout_id, + model_call_id=context.model_call_id, + ) return coords = CommitCoords.model_validate(coords_payload) if coords.rollout_id != context.rollout_id or coords.model_call_id != context.model_call_id: @@ -209,6 +218,10 @@ async def _finalize_admitted_response( context.model_call_id, WORKER_CAPTURE_FAILED_REASON, ) + mark_external_staging_failed( + rollout_id=context.rollout_id, + model_call_id=context.model_call_id, + ) return if coords.parent_call_id != admission.parent_call_id or coords.prev_len != admission.prev_len: raise ValueError(f"coordinates for {coords.model_call_id} diverge from admission") diff --git a/nemo_gym/token_id_capture/sink.py b/nemo_gym/token_id_capture/sink.py index 48fcd97d819..68e02fcbbb9 100644 --- a/nemo_gym/token_id_capture/sink.py +++ b/nemo_gym/token_id_capture/sink.py @@ -181,6 +181,26 @@ def mark_external_staging_committed(*, rollout_id: str, model_call_id: str) -> N f"{context.rollout_id}/{context.model_call_id})" ) context.committed = True + from nemo_gym._checkpoint.admission import mark_current_generation_safe + + mark_current_generation_safe("durable_completed") + + +def mark_external_staging_failed(*, rollout_id: str, model_call_id: str) -> None: + """Mark a matching call safe after its poison row becomes durable.""" + context = _CAPTURE_CONTEXT.get() + if context is None: + raise RuntimeError("no training-token capture context is active") + if context.rollout_id != rollout_id or context.model_call_id != model_call_id: + raise ValueError( + "external staging failure does not match the active capture " + f"context ({rollout_id}/{model_call_id} != " + f"{context.rollout_id}/{context.model_call_id})" + ) + context.committed = True + from nemo_gym._checkpoint.admission import mark_current_generation_safe + + mark_current_generation_safe("durable_failure") def reset_token_sink(token: Token) -> None: @@ -298,6 +318,10 @@ async def resolve_parent(request_messages: list | None) -> None: context.model_call_id, UNRESOLVED_PARENT_REASON, ) + mark_external_staging_failed( + rollout_id=context.rollout_id, + model_call_id=context.model_call_id, + ) return is_root = context.parent_resolution is not None and context.parent_resolution.status == ParentResolutionStatus.ROOT if is_root or (context.source_capture_key is None and not await ledger.has_rows(context.rollout_id)): @@ -317,6 +341,10 @@ async def resolve_parent(request_messages: list | None) -> None: context.model_call_id, UNRESOLVED_PARENT_REASON, ) + mark_external_staging_failed( + rollout_id=context.rollout_id, + model_call_id=context.model_call_id, + ) async def register_call_intent() -> None: @@ -475,6 +503,9 @@ async def commit_entry( entry.parent_resolution_reason = resolution.reason or "" await context.token_sink.put(entry) context.committed = True + from nemo_gym._checkpoint.admission import mark_current_generation_safe + + mark_current_generation_safe("durable_completed") except Exception: await _capture_failed(context, "write") diff --git a/tests/unit_tests/test_checkpoint_ledger.py b/tests/unit_tests/test_checkpoint_ledger.py index 8f781402b0f..5bece4b8c2b 100644 --- a/tests/unit_tests/test_checkpoint_ledger.py +++ b/tests/unit_tests/test_checkpoint_ledger.py @@ -30,6 +30,13 @@ CaptureLedgerCheckpointer, ControlCapabilities, ControlFence, + GenerationCutCoordinatorProof, + GenerationCutFrozenTicket, + GenerationCutInventory, + GenerationCutPrefix, + GenerationCutPrefixAck, + GenerationCutReceipt, + GenerationCutWorkerProof, LedgerMismatchError, MultiProcessCapability, StaleAttemptError, @@ -102,6 +109,242 @@ def test_restore_validates_all_files_before_installing_any(tmp_path) -> None: assert not restored.exists() +def test_generation_cut_receipt_is_bound_to_ledger_commit_and_restore(tmp_path) -> None: + source = tmp_path / "source" + _write_custody(source, "rollout-a") + checkpoint = tmp_path / "checkpoint" + inventory = GenerationCutInventory.build( + checkpoint_id="checkpoint-1", + server_name="policy", + active_prefixes=[], + ) + receipt = GenerationCutReceipt( + checkpoint_id="checkpoint-1", + cut_id="cut-1", + inventory_digest=inventory.inventory_digest, + inventory=inventory, + backend_snapshot_id="snapshot-1", + ) + checkpointer = CaptureLedgerCheckpointer(source) + + committed = checkpointer.commit( + checkpoint, + checkpoint_id="checkpoint-1", + tombstones=[], + generation_cut_receipt=receipt, + ) + manifest = json.loads((checkpoint / MODEL_LEDGER_SUBDIR / LEDGER_MANIFEST_NAME).read_text()) + assert manifest["schema_version"] == 2 + assert "generation_cut_receipt" in manifest + assert "generation_cut_ack" not in manifest + assert committed["generation_cut_receipt"] == receipt.model_dump(mode="json") + restored = CaptureLedgerCheckpointer(tmp_path / "restored").restore(checkpoint) + assert restored["generation_cut_receipt"] == receipt.model_dump(mode="json") + + changed = receipt.model_copy(update={"cut_id": "cut-2"}) + with pytest.raises(LedgerMismatchError, match="generation cut changed"): + checkpointer.commit( + checkpoint, + checkpoint_id="checkpoint-1", + tombstones=[], + generation_cut_receipt=changed, + ) + + +def test_restore_rejects_legacy_generation_cut_ack_manifest_with_migration_guidance(tmp_path) -> None: + checkpoint = tmp_path / "checkpoint" + CaptureLedgerCheckpointer(tmp_path / "source").commit( + checkpoint, + checkpoint_id="checkpoint-1", + tombstones=[], + ) + manifest_path = checkpoint / MODEL_LEDGER_SUBDIR / LEDGER_MANIFEST_NAME + manifest = json.loads(manifest_path.read_text()) + manifest["schema_version"] = 1 + manifest["generation_cut_ack"] = {"checkpoint_id": "checkpoint-1", "cut_id": "legacy-cut"} + manifest_path.write_text(json.dumps(manifest)) + + with pytest.raises( + LedgerMismatchError, + match=r"cannot be migrated safely.*recreate the checkpoint with schema_version 2", + ): + CaptureLedgerCheckpointer(tmp_path / "restored").restore(checkpoint) + + +def test_restore_rejects_legacy_generation_cut_sidecar_with_migration_guidance(tmp_path) -> None: + checkpoint = tmp_path / "checkpoint" + CaptureLedgerCheckpointer(tmp_path / "source").commit( + checkpoint, + checkpoint_id="checkpoint-1", + tombstones=[], + ) + ledger_dir = checkpoint / MODEL_LEDGER_SUBDIR + manifest_path = ledger_dir / LEDGER_MANIFEST_NAME + manifest = json.loads(manifest_path.read_text()) + manifest["schema_version"] = 1 + manifest_path.write_text(json.dumps(manifest)) + (ledger_dir / "generation-cut.json").write_text( + json.dumps({"ack": {"cut_id": "legacy-cut"}, "receipt": {"backend_snapshot_id": "legacy-snapshot"}}) + ) + + with pytest.raises( + LedgerMismatchError, + match=r"cannot be migrated safely.*recreate the checkpoint with schema_version 2", + ): + CaptureLedgerCheckpointer(tmp_path / "restored").restore(checkpoint) + + +def test_restore_keeps_legacy_cut_free_ledger_compatibility(tmp_path) -> None: + source = tmp_path / "source" + _write_custody(source, "rollout-a") + checkpoint = tmp_path / "checkpoint" + CaptureLedgerCheckpointer(source).commit(checkpoint, checkpoint_id="checkpoint-1", tombstones=[]) + manifest_path = checkpoint / MODEL_LEDGER_SUBDIR / LEDGER_MANIFEST_NAME + manifest = json.loads(manifest_path.read_text()) + manifest["schema_version"] = 1 + manifest_path.write_text(json.dumps(manifest)) + + restored = CaptureLedgerCheckpointer(tmp_path / "restored").restore(checkpoint) + assert restored["checkpoint_id"] == "checkpoint-1" + + +def test_multi_worker_generation_cut_proof_is_persisted_and_validated(tmp_path) -> None: + workers = [] + for index in range(2): + frozen_ticket = GenerationCutFrozenTicket( + ticket_id=f"ticket-{index}", + rollout_id=f"rollout-{index}", + attempt_index=0, + model_call_id=f"call-{index}", + generation_started=True, + response_started=False, + ) + inventory = GenerationCutInventory.build( + checkpoint_id="checkpoint-1", + server_name="policy", + active_prefixes=[ + GenerationCutPrefix( + ticket_id=frozen_ticket.ticket_id, + rollout_id=frozen_ticket.rollout_id, + attempt_index=frozen_ticket.attempt_index, + model_call_id=frozen_ticket.model_call_id, + admitted_at=float(index), + ) + ], + ) + receipt = GenerationCutReceipt( + checkpoint_id="checkpoint-1", + cut_id=f"cut-{index}", + inventory_digest=inventory.inventory_digest, + inventory=inventory, + backend_snapshot_id=f"snapshot-{index}", + prefixes=( + GenerationCutPrefixAck( + ticket_id=frozen_ticket.ticket_id, + rollout_id=frozen_ticket.rollout_id, + attempt_index=frozen_ticket.attempt_index, + model_call_id=frozen_ticket.model_call_id, + admitted_at=float(index), + disposition="durable_prefix", + frozen_buffer_id=f"buffer-{index}", + staging_key=f"staging-{index}", + prefix_token_count=1, + prefix_digest=f"{index}" * 64, + ), + ), + ) + workers.append( + GenerationCutWorkerProof.build( + checkpoint_id="checkpoint-1", + coordinator_sequence=4, + worker_id=f"w{index}", + frozen_tickets=[frozen_ticket], + ready_ticket_ids=[frozen_ticket.ticket_id], + generation_cut_receipt=receipt, + ) + ) + proof = GenerationCutCoordinatorProof.build( + checkpoint_id="checkpoint-1", + coordinator_sequence=4, + expected_workers=2, + frozen_worker_ids=("w0", "w1"), + workers=workers, + ) + source = tmp_path / "source" + checkpointer = CaptureLedgerCheckpointer(source) + checkpoint = tmp_path / "checkpoint" + + committed = checkpointer.commit( + checkpoint, + checkpoint_id="checkpoint-1", + tombstones=[], + generation_cut_proof=proof, + ) + assert committed["generation_cut_proof"] == proof.model_dump(mode="json") + restored = CaptureLedgerCheckpointer(tmp_path / "restored").restore(checkpoint) + assert restored["generation_cut_proof"] == proof.model_dump(mode="json") + + omitted = proof.model_copy(update={"workers": (workers[0],)}) + with pytest.raises(ValueError, match="does not match frozen worker IDs"): + CaptureLedgerCheckpointer(source).commit( + tmp_path / "omitted", + checkpoint_id="checkpoint-1", + tombstones=[], + generation_cut_proof=omitted, + ) + + mismatched_worker = GenerationCutWorkerProof.build( + checkpoint_id=workers[1].checkpoint_id, + coordinator_sequence=5, + worker_id=workers[1].worker_id, + frozen_tickets=list(workers[1].frozen_tickets), + ready_ticket_ids=list(workers[1].ready_ticket_ids), + generation_cut_receipt=workers[1].generation_cut_receipt, + ) + mismatched = proof.model_copy(update={"workers": (workers[0], mismatched_worker)}) + with pytest.raises(ValueError, match="mismatched checkpoint or coordinator sequence"): + CaptureLedgerCheckpointer(source).commit( + tmp_path / "mismatched", + checkpoint_id="checkpoint-1", + tombstones=[], + generation_cut_proof=mismatched, + ) + + replacement_worker = GenerationCutWorkerProof.build( + checkpoint_id=workers[1].checkpoint_id, + coordinator_sequence=workers[1].coordinator_sequence, + worker_id="w2", + frozen_tickets=list(workers[1].frozen_tickets), + ready_ticket_ids=list(workers[1].ready_ticket_ids), + generation_cut_receipt=workers[1].generation_cut_receipt, + ) + replaced = proof.model_copy(update={"workers": (workers[0], replacement_worker)}) + with pytest.raises(ValueError, match="does not match frozen worker IDs"): + CaptureLedgerCheckpointer(source).commit( + tmp_path / "replaced", + checkpoint_id="checkpoint-1", + tombstones=[], + generation_cut_proof=replaced, + ) + + client, _ = _participant( + tmp_path / "route-ledger", + generation_cut_proof_provider=lambda: proof, + expected_workers=2, + ) + control = {"checkpoint_id": "checkpoint-1", "deadline_ts": 4e9} + assert client.post(f"{MODEL_ADMISSION_URL_PREFIX}/pause", json=control, headers=AUTH_HEADERS).status_code == 200 + route_checkpoint = tmp_path / "route-checkpoint" + committed = client.post( + f"{MODEL_CHECKPOINT_URL_PREFIX}/commit", + json={**control, "checkpoint_dir": str(route_checkpoint)}, + headers=AUTH_HEADERS, + ) + assert committed.status_code == 200 + assert committed.json()["generation_cut_proof"] == proof.model_dump(mode="json") + assert (route_checkpoint / MODEL_LEDGER_SUBDIR / "policy" / "generation-cut-workers.json").exists() + + def test_restore_rejects_uncommitted_and_nonfresh_namespaces(tmp_path) -> None: checkpoint = tmp_path / "checkpoint" ledger_dir = checkpoint / MODEL_LEDGER_SUBDIR @@ -119,9 +362,15 @@ def test_restore_rejects_uncommitted_and_nonfresh_namespaces(tmp_path) -> None: CaptureLedgerCheckpointer(restored).restore(checkpoint) -def _participant(root) -> tuple[TestClient, AdmissionLimiter]: +def _participant( + root, + *, + generation_cut_backend=None, + generation_cut_proof_provider=None, + expected_workers: int = 1, +) -> tuple[TestClient, AdmissionLimiter]: app = FastAPI() - limiter = AdmissionLimiter() + limiter = AdmissionLimiter(generation_cut_backend) fence = ControlFence() ledger = FileLineageStore(root) install_control_plane( @@ -150,10 +399,81 @@ def _participant(root) -> tuple[TestClient, AdmissionLimiter]: instance_role="policy", server_name="policy", auth_token=AUTH_TOKEN, + generation_cut_proof_provider=generation_cut_proof_provider, + expected_workers=expected_workers, ) return TestClient(app), limiter +def test_multi_worker_commit_without_coordinator_proof_fails_closed(tmp_path) -> None: + client, _ = _participant(tmp_path / "ledger", expected_workers=2) + control = {"checkpoint_id": "checkpoint-1", "deadline_ts": 4e9} + pause = client.post(f"{MODEL_ADMISSION_URL_PREFIX}/pause", json=control, headers=AUTH_HEADERS) + assert pause.json()["state"] == "paused" + + commit = client.post( + f"{MODEL_CHECKPOINT_URL_PREFIX}/commit", + json={**control, "checkpoint_dir": str(tmp_path / "checkpoint")}, + headers=AUTH_HEADERS, + ) + assert commit.status_code == 409 + assert commit.json()["error"]["code"] == "ledger_not_quiescent" + + +def test_generation_cut_receipt_is_final_before_ledger_commit(tmp_path) -> None: + class DurableReceiptBackend: + def __init__(self) -> None: + self.checkpoint_calls = 0 + + async def checkpoint_generation_cut(self, inventory): + self.checkpoint_calls += 1 + return GenerationCutReceipt( + checkpoint_id=inventory.checkpoint_id, + cut_id="cut-1", + inventory_digest=inventory.inventory_digest, + inventory=inventory, + backend_snapshot_id="snapshot-1", + prefixes=tuple( + GenerationCutPrefixAck( + ticket_id=prefix.ticket_id, + rollout_id=prefix.rollout_id, + attempt_index=prefix.attempt_index, + model_call_id=prefix.model_call_id, + admitted_at=prefix.admitted_at, + disposition="durable_prefix", + frozen_buffer_id=f"buffer/{prefix.ticket_id}", + staging_key=f"staging/{prefix.ticket_id}", + prefix_token_count=2, + prefix_digest="d" * 64, + ) + for prefix in inventory.active_prefixes + ), + ) + + async def restore_generation_cut(self, receipt): + return receipt + + backend = DurableReceiptBackend() + client, limiter = _participant(tmp_path / "ledger", generation_cut_backend=backend) + ticket = limiter.admit(rollout_id="rollout-a", attempt_index=0) + ticket.generation_started = True + ticket.model_call_id = "call-1" + pause_body = {"checkpoint_id": "checkpoint-1", "deadline_ts": 4e9} + pause = client.post(f"{MODEL_ADMISSION_URL_PREFIX}/pause", json=pause_body, headers=AUTH_HEADERS) + assert pause.json()["state"] == "paused" + assert backend.checkpoint_calls == 1 + + committed = client.post( + f"{MODEL_CHECKPOINT_URL_PREFIX}/commit", + json={**pause_body, "checkpoint_dir": str(tmp_path / "checkpoint")}, + headers=AUTH_HEADERS, + ) + assert committed.status_code == 200 + assert committed.json()["generation_cut_receipt"]["backend_snapshot_id"] == "snapshot-1" + assert backend.checkpoint_calls == 1 + limiter.release(ticket) + + def test_commit_requires_completed_drain_and_restore_stays_paused(tmp_path) -> None: source_client, source_limiter = _participant(tmp_path / "ledger-a") _write_custody(tmp_path / "ledger-a", "rollout-a") From 2f7e74362c19c70973ebb45d25600fb5257df925 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 22 Sep 2026 07:12:02 -0700 Subject: [PATCH 4/5] fix(checkpoint): restore lineage request header Signed-off-by: Ananth Subramaniam --- nemo_gym/rollout_correlation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nemo_gym/rollout_correlation.py b/nemo_gym/rollout_correlation.py index cadb23fad8b..469a81b4257 100644 --- a/nemo_gym/rollout_correlation.py +++ b/nemo_gym/rollout_correlation.py @@ -56,6 +56,7 @@ class _CheckpointParentState: MODEL_CALL_ID_HEADER = "x-nemo-gym-model-call-id" SOURCE_CAPTURE_KEY_HEADER = "x-nemo-gym-source-capture-key" PARENT_MODEL_CALL_ID_HEADER = "x-nemo-gym-parent-model-call-id" +LOGICAL_REQUEST_HEADER = "x-nemo-gym-logical-request-id" # The transport id appends ``-a{n}`` for re-dispatch attempts. The suffix is a # capture and routing key, never the logical identity. This pattern recovers From e3590a2a8685e7e51355f7bb03b85867b96b0cf2 Mon Sep 17 00:00:00 2001 From: Ananth Subramaniam Date: Tue, 22 Sep 2026 08:57:52 -0700 Subject: [PATCH 5/5] test(checkpoint): use typed lineage commits Signed-off-by: Ananth Subramaniam --- tests/unit_tests/test_token_capture_ledger.py | 44 +++++++++++-------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/tests/unit_tests/test_token_capture_ledger.py b/tests/unit_tests/test_token_capture_ledger.py index 2f548afcea8..dac1b96abee 100644 --- a/tests/unit_tests/test_token_capture_ledger.py +++ b/tests/unit_tests/test_token_capture_ledger.py @@ -273,24 +273,20 @@ async def test_cross_attempt_admission_uses_explicit_verified_parent(store): tokens_2 = TOKENS_1 + [901, 902] chain_hash_2 = compute_chain_hash(CHAIN_HASH_1, [901, 902]) await store.record( - "r1-a1", - "c2", - [USER_1, ASSISTANT_1, USER_2], - [ASSISTANT_2], - [], - hash_token_ids(tokens_2), - parent_call_id="c1", - staging_key="r1-a1/c2", - weight_version=17, - prev_len=len(TOKENS_1), - delta_len=2, - cum_len=len(tokens_2), - staging_digest=STAGING_DIGEST, - extras_digest=EMPTY_EXTRAS_DIGEST, - mode="token_in", - staging_chain=["r1/c1", "r1-a1/c2"], - chain_hash=chain_hash_2, - cumulative_hash=hash_token_ids(tokens_2), + _commit( + _call_record( + "c2", + parent_call_id="c1", + prev_len=len(TOKENS_1), + delta_len=2, + chain_hash=chain_hash_2, + cumulative_hash=hash_token_ids(tokens_2), + ).model_copy(update={"staging_key": "r1-a1/c2"}), + [USER_1, ASSISTANT_1, USER_2], + [ASSISTANT_2], + rollout_id="r1-a1", + staging_chain=("r1/c1", "r1-a1/c2"), + ) ) next_context = await _admit( store, @@ -330,9 +326,21 @@ def __init__(self) -> None: async def resolve(self, rollout_id, request_items): return await self.inner.resolve(rollout_id, request_items) + def is_process_shared(self): + return self.inner.is_process_shared() + + async def close(self): + await self.inner.close() + + async def record(self, commit): + await self.inner.record(commit) + async def has_rows(self, rollout_id): return await self.inner.has_rows(rollout_id) + async def manifest(self, rollout_id): + return await self.inner.manifest(rollout_id) + async def record_failure(self, rollout_id, model_call_id, reason): await self.inner.record_failure(rollout_id, model_call_id, reason)