[ananthsub]

What changed and why

The indirect prompt injection resources server keeps a mutable environment for each session: tool calls update patient charts, send messages, and so on. Before this change it was restart-only by default, so every rollout with a live session started over from its input after a crash. Now a rollout interrupted by a crash can continue from its last checkpoint.

• The server declares checkpoint_mode = "exported" and implements the three session hooks, keyed by the cookie session ID:
  • export_session_states returns each live session's environment as an IPICheckpointState;
  • restore_session_states validates every state first, then installs them all, so an invalid state installs nothing;
  • retire_session_state drops the session.
• checkpoint_verify stays at the default wait, because /verify deletes the session's environment and cannot be replayed.
• Export leaves out a session the server no longer holds. /verify drops the session in its finally block even when verification raises, and the participant only stops tracking a session on a successful /verify. Without this, the next commit after a failed verification would fail.
• The environment is copied through JSON, which keeps its key order and rejects values JSON cannot carry. Tool outputs serialize the environment as stored, so a restored session answers byte for byte as the original would; a test compares every tool output and the verify reply against an uncheckpointed run.

How it works

Where this PR sits in the overall flow

The highlighted part is what this PR adds.

flowchart LR
  C["Controller<br/>NeMo RL, or rollout collection"]
  CO["Coordination<br/>prepare, commit, restore, resume,<br/>retire, forget"]
  K["Control plane on every server<br/>phases, lease, storage,<br/>retire: refuse, stop, free;<br/>forget"]
  subgraph G["One participant per Gym server"]
    E["Environment server<br/>episode steps"]
    M["Policy model<br/>held responses, generation cuts"]
    A["Agent<br/>sessions parked at boundaries"]
    R["Resources server<br/>session state"]
  end
  W["Inference worker<br/>stages cut prefixes"]
  D[("Checkpoint directory<br/>records, then manifest")]
  L[("Capture ledger<br/>retire and delete from #3938, #3939")]
  C --> CO --> K
  K --> E & M & A & R
  M --> W
  M --> L
  G --> D
  classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937
  class R this

Loading

One checkpoint, a crash, and the restore, end to end:

sequenceDiagram
  participant C as Controller
  participant G as Gym participants
  participant D as Checkpoint directory
  C->>G: prepare, in order environment, model, agent, resources
  Note over G: admission closes, in-flight work parks at a boundary,<br/>undelivered model responses are held
  G-->>C: prepared, or blockers at the deadline
  C->>G: commit with the episodes the controller continues
  G->>D: each participant writes its records, then its manifest
  Note over G: restored state the commit's scope leaves out is deleted
  C->>C: publish the checkpoint with the controller's own state
  C->>G: resume, in order resources, agent, model, environment
  Note over C,G: crash - every Gym process dies
  C->>G: restore the checkpoint in fresh processes, all or nothing
  D-->>G: records installed under attempt N+1, attempt N's capture ledger retired
  C->>G: resume
  C->>G: /run as attempt N+1 continues each episode from its boundary
  Note over C,G: retiring an episode, only while no checkpoint is open
  C->>G: retire - environment, then agent, then model and resources
  Note over G: each server marks the attempt, stops its work, waits, frees its state, then replies
  C->>G: forget, once the rollout is done for good - same order
  Note over G: the mark refusing the retired attempt goes away

Loading

This PR

Each session's environment, saved at a checkpoint. It is copied through JSON, which keeps its key order, so a restored session answers byte for byte as the original would.

sequenceDiagram
  participant A as Agent
  participant I as Indirect prompt injection server
  A->>I: /seed_session with the environment
  A->>I: tool calls read and change the environment
  Note over I: commit - export_session_states saves each environment
  Note over A,I: crash, then restore into a fresh server
  Note over I: restore_session_states validates every state, then installs them all
  A->>I: tool calls continue on the same cookie session
  A->>I: /verify scores and drops the session, so verify mode stays wait

Loading

Where this sits in the stack

This PR builds on the partial-rollout checkpointing stack: #3882 (core) through #3889 (end-to-end suite), plus #3893 (rollout collection). It is one of five PRs that port the environments the old stack checkpointed onto the new hooks. Each one is based on #3889 and can be reviewed on its own, except Blackjack, which builds on the Gymnasium fix:

• feat(checkpoint): checkpoint Workplace Assistant sessions #3894 (ananthsub/partial-ckpt-workplace-assistant): feat(checkpoint): checkpoint Workplace Assistant sessions
• fix(gymnasium): build Gymnasium servers on the shared resources server setup #3895 (ananthsub/partial-ckpt-gymnasium): fix(gymnasium): build Gymnasium servers on the shared resources server setup
• feat(checkpoint): export Blackjack game state for partial-rollout checkpoints #3896 (ananthsub/partial-ckpt-blackjack): feat(checkpoint): export Blackjack game state for partial-rollout checkpoints (on the Gymnasium PR)
• feat(checkpoint): export indirect prompt injection session state #3897 (ananthsub/partial-ckpt-indirect-prompt-injection): feat(checkpoint): export indirect prompt injection session state (this PR)
• feat(checkpoint): continue proof refinement agent turns from their boundaries #3898 (ananthsub/partial-ckpt-proof-refinement): feat(checkpoint): continue proof refinement agent turns from their boundaries

Relationship to the old stack

This supersedes #3548, which added the same recovery on the old (v1) checkpoint stack. Ported from #3548 onto the v2 checkpoint hooks.

Carried over:

• the strict, versioned checkpoint state model;
• validate-then-install restore;
• the restore-and-continue test parametrized over the checkpoint point (after seed, after a chart update, after a message), plus the aliasing, invalid-state, and non-JSON export tests.

Not carried, because the v2 stack removes the mechanisms they served:

• the execution_to_session map and the identity lookups from rollout headers, since v2 keys resources state by the cookie session ID;
• mutation receipts, revisions, and the receipt-replay test, since v2 relies on ordering (a tool call runs between agent boundaries, so it blocks prepare);
• the verify identity check and the explicit terminal marking, since the participant learns a session ended from a successful /verify;
• the duplicate-identity and stale-attempt-identity tests. A restore batch is keyed by session ID, so it cannot hold duplicates, and a retire stops and frees a session itself.

Issue

No separate issue. This is part of the partial-rollout checkpointing work, as the port of draft #3548.

Validation

Run on this branch, on the current main:

• pytest -q resources_servers/indirect_prompt_injection/tests: passed, including the new test_checkpoint.py. They drive the server through its real /ng-control/v1/checkpoint/* routes over ASGI:
  • a seeded session is checkpointed after 0, 1, or 2 tool calls, restored into a fresh server instance, and continued through the remaining tool calls and /verify. Tool outputs and the verify reply match an uncheckpointed run exactly. After verification, the restored server and its participant hold no session;
  • an invalid checkpoint fails restore through the control route and installs nothing;
  • restore with any invalid state in the batch installs nothing. The batch puts a valid environment first, passed as a dictionary as the state model requires, so a restore that installed each state as it validated it would leave that environment behind; the invalid cases are a newer schema version, an unexpected field, the environment as JSON text, a value JSON cannot carry, a list, and a non-object. With an intentionally non-atomic restore, this test fails;
  • a retired session is gone from the server, and a later call with its cookie is refused and does not recreate it, while another session keeps working;
  • a session dropped by a failed verification is left out of the next commit, and that commit succeeds;
  • export and restore do not alias live state, and export rejects values JSON cannot carry.
• pytest -q tests/unit_tests/test_checkpoint_*.py (with RAY_TMPDIR=/tmp): passed.
• ruff check and ruff format on the server: clean.
• pre-commit run --files on the changed files: passed.

Rollout evidence

Setup

The model was Qwen/Qwen3-4B-Instruct-2507, served by vLLM 0.29.0 (the version Gym's vllm extra pins, with torch 2.13.0+cu130) using the hermes tool-call parser, --max-model-len 32768 and --gpu-memory-utilization 0.45 on one NVIDIA RTX 6000 Ada Generation (48 GB).
Gym's vllm_model server ran with uses_reasoning_parser=false, because the served model is an instruct model without a reasoning parser.
Every rollout sampled at temperature 1.0, used the five tasks in resources_servers/indirect_prompt_injection/data/example.jsonl, and ran Simple Agent with max_steps: 5, behind the environment server that relays flat rows to the agent's own /run episode.
Each tree ran in its own worktree, with server venvs created by gym env start.

Uninterrupted baseline

This change (156851958c) and its base (80ace9f0f9) each ran with checkpointing off, eight rollouts per task:

gym env start --config resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml --model-type vllm_model \
    --model Qwen/Qwen3-4B-Instruct-2507 --model-url http://127.0.0.1:8100/v1 --model-api-key dummy \
    ++policy_model.responses_api_models.vllm_model.uses_reasoning_parser=false
gym eval run --no-serve --config resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml --model-type vllm_model \
    --model Qwen/Qwen3-4B-Instruct-2507 --model-url http://127.0.0.1:8100/v1 --model-api-key dummy \
    ++policy_model.responses_api_models.vllm_model.uses_reasoning_parser=false \
    --agent indirect_prompt_injection_simple_agent --input resources_servers/indirect_prompt_injection/data/example.jsonl \
    --output rollouts.jsonl --num-repeats 8 --concurrency 16 --temperature 1.0

Both trees gave the same result over 40 rollouts: a mean reward of 0.80, no injection followed, and every rollout safe.
Four tasks succeeded 8 of 8 times on both trees. The first task's user message in example.jsonl is the placeholder ..., so the model never calls the tools the task requires and its utility reward is 0 on both trees.
Neither run had failed or masked rollouts or tracebacks in the server logs.

Checkpoint, crash and restore

The tree was this change, which already contains the process-level checkpoint suite.
A driver built on that suite (tests/e2e/checkpoint/checkpoint_deployment.py) started the policy model, the indirect prompt injection resources server, Simple Agent and the flat-row environment server, each as its own process from its own venv.
It sent 80 rollouts (16 per task) to the environment server's /run and drove the checkpoint through nemo_gym._checkpoint.coordination, as a training controller would:

1. Prepare and commit a checkpoint once 21 rollouts had finished and 59 were mid-episode.
2. Kill every Gym server process group with SIGKILL, start fresh processes from the same configs, restore the committed episodes and resume.
3. Take a second checkpoint of the restored state before the replacements start, then resume again.
4. Run every unfinished rollout again as attempt 1, which continues from its checkpointed boundary.

Prepare returned in 36 ms and reported no restarts.
Simple Agent reported all 59 in-flight episodes at a boundary. The resources server exported 58 sessions: the remaining episode was parked after its last model call, ready to return, and the server no longer held its session.

28 episodes had already made a successful state-changing tool call before the checkpoint.
Seven had prescribed medication, 11 had moved candidate C-4552 to "Interview Scheduled", and 10 had moved candidate C-1002 to "offered".
Each export differs from the environment its task row seeds exactly where that call changed it. For example, a prescription episode's export differs only under prescriptions.P-7742, and a candidate episode's export differs only in candidates.C-4552.status.
Every exported environment keeps the seeded key order, and the second checkpoint exported every restored session unchanged.

After the restore, all 59 continued episodes resumed as attempt 1 and returned 200.
Each final trajectory starts with exactly the tool calls, arguments and outputs recorded at the checkpoint.
16 episodes made their state-changing call on the restored session after the restore: nine prescriptions and seven candidate status updates. Each call succeeded.
13 tool calls made after the restore have an uninterrupted baseline call with the same arguments, and each returned the same output, byte for byte.
The run's mean reward over all 80 rollouts was 0.80, the same as the uninterrupted baseline, with the same per-task results.
No server log had a traceback.

The runs only exercised commit, restore and resume. No episode was reported as a restart, so no retire ran.

Compatibility

• Nothing changes when checkpointing is off. The hooks only run when the checkpoint participant is installed.
• With checkpointing on, live sessions of this server are exported and restored instead of restarting their rollouts.
• The stored state format (schema_version: 1, the environment as a JSON object) is new and does not read checkpoints written by feat(checkpoint): add indirect prompt injection state recovery #3548.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.