[ananthsub]

• Branch: ananthsub/partial-ckpt-workplace-assistant
• Base branch: main
• Suggested labels: feature, area:env-infra

---

What changed and why

Workplace Assistant used the default restart_only mode, so every rollout with a live Workplace session was a restart and started over from scratch after a crash. This PR lets those sessions be checkpointed and continued.

• The server declares checkpoint_mode = "exported" and implements the three session hooks over its per-session state. That state is the tool environment created at /seed_session: a set of in-memory pandas tables for email, calendar, analytics, project management, the customer relationship manager, and the company directory.
• export_session_states encodes every table column by column. It keeps column dtypes, such as the boolean user_engaged column, and index labels, which matter because the analytics tool appends rows by label. It also keeps the difference between empty CSV cells, which load as NaN, and values a tool stored as None. JSON writes both as null, so NaN positions are listed separately. Encoding runs in a worker thread. Admission is closed during a commit, so no tool call changes the tables while they are encoded.
• restore_session_states rebuilds every session's tool environment first and installs them only after all of them are valid. A state with missing or extra containers or tables, or with mismatched column lengths, raises ValueError and installs nothing.
• Export leaves out a session the server no longer holds. /verify discards the session's tool environment in a finally block. If verify raises, the session is gone but the participant still tracks it, so the commit must not fail on it.
• retire_session_state drops the session's tool environment.
• checkpoint_verify stays at the default wait. Verify computes the reward only from the response and the ground truth, but it deletes the session's state, so it is not replay-safe.
• The toolkit list is now one module constant, _TOOLKITS, instead of being repeated in seed_session and mcp_tools.

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 tool environment is a set of tables. A checkpoint exports them whole, and a restore rebuilds them before any of them is installed.

sequenceDiagram
  participant A as Agent
  participant W as Workplace Assistant
  participant D as Checkpoint directory
  A->>W: /seed_session creates the session's tables
  A->>W: tool calls read and change the tables
  Note over W: commit - export_session_states writes each table column by column,<br/>with column types, index labels, and where cells are empty
  W->>D: one record per session
  Note over A,D: crash, then restore into a fresh server
  D-->>W: restore_session_states rebuilds every session, then installs them all
  A->>W: tool calls continue on the same cookie session
  A->>W: /verify scores and drops the session

Loading

Where this sits in the stack

This PR uses the resources server checkpoint participant from #3886, which has merged, and needs nothing else from the open part of the partial-rollout checkpointing stack. It is one of five PRs that port the environments the old stack checkpointed onto the new hooks. Workplace Assistant and the Gymnasium fix target main, and Blackjack builds on the Gymnasium fix. Indirect prompt injection and proof refinement are still based on the end-to-end suite (#3889):

• feat(checkpoint): checkpoint Workplace Assistant sessions #3894 (ananthsub/partial-ckpt-workplace-assistant): feat(checkpoint): checkpoint Workplace Assistant sessions (this PR)
• 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
• 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

• Ported from feat(checkpoint): adapt initial stateful environments #2946, which carried the Workplace Assistant adapter on the old hooks. It covers gap 7 of the re-cut, where this adapter was not carried.
• Dropped from feat(checkpoint): adapt initial stateful environments #2946:
  • the execution_to_session map and the identity lookups in the tool route and verify. The v2 hooks are keyed by the cookie session ID, so neither is needed.
  • route classification of tool calls as mutations. In v2, admission refuses every data route while a checkpoint is open.
  • synthetic checkpoint:<rollout>:a<attempt> session IDs on restore. The restored session keeps its original cookie session ID.
• The old encoding used pandas to_json(orient="split") and repaired types after reading the JSON back. That turned NaN into None and lost the difference between them. The new encoding keeps it.
• Not ported from feat(checkpoint): add durable turn-level rollout recovery #3349: the audit-trail test fixture (checkpoint_test_app.py) and its test. Both depend on old-stack rollout identities and on the test agent from that PR.

Issue

No tracking issue exists for this re-cut. The re-cut series and the old stack's PRs carry the discussion.

Validation

Run on this branch, on the current main:

• RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q resources_servers/workplace_assistant/tests: passed, including the new tests/test_checkpoint.py. Its tests drive the real control routes over ASGI:
  • A session is seeded, mutated by three tool calls, prepared, and committed. It is restored into a fresh server instance with the same cookies and continues with four more tool calls and /verify. Tool outputs, every table, and the reward all match an uncheckpointed run of the same steps.
  • When one of two sessions has an invalid state, restore raises and installs neither session.
  • A retired session's tool environment is gone and the next commit exports nothing. A retire stops and frees the session instead of fencing it, so a later tool call is refused and does not recreate the environment.
  • A session whose /verify raised is left out of the export. The commit succeeds with only the live session, which then continues after restore.
• A mutation check confirmed the equivalence test is not vacuous. Without restoring NaN cells, the test fails.
• RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q tests/unit_tests/test_checkpoint_*.py: passed.
• ruff check and ruff format --check on the changed files: clean.
• pre-commit run --files <changed files>: all hooks passed.
• The commit carries a Signed-off-by line.

Rollout evidence

Setup

The model was Qwen/Qwen3-4B-Instruct-2507, served by vLLM 0.24.0 with the hermes tool-call parser 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/workplace_assistant/data/example.jsonl, and ran Simple Agent with max_steps: 6 behind the flat-row wrapper.
Each tree ran in its own worktree, with server venvs created by gym env start.

Uninterrupted baseline

This change on main ran with checkpointing off, eight rollouts per task:

gym env start --config resources_servers/workplace_assistant/configs/workplace_assistant.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/workplace_assistant/configs/workplace_assistant.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 workplace_assistant_simple_agent --input resources_servers/workplace_assistant/data/example.jsonl \
    --output rollouts.jsonl --num-repeats 8 --concurrency 16 --temperature 1.0

The 40 rollouts had a mean reward of 0.25.
By task, the email reply succeeded 7 of 8 times, the CRM status move 3 of 8 times, and the calendar rename and the two other CRM tasks never succeeded.
The same command on the checkpoint stack with this change, still with checkpointing off, also gave 0.25 over 40 rollouts.
Neither run had failed or masked rollouts or tracebacks in the server logs.

Checkpoint, crash and restore

The tree was the checkpoint stack (ananthsub/partial-ckpt-e2e) with this change cherry-picked on top; it applied without conflicts.
A driver built on the process-level checkpoint suite (tests/e2e/checkpoint/checkpoint_deployment.py) started the policy model, the Workplace Assistant resources server, Simple Agent and the flat-row wrapper, each as its own process from its own venv.
It sent the rollouts to the wrapper's /run and drove the checkpoint through nemo_gym._checkpoint.coordination, as a training controller would:

1. Prepare and commit a checkpoint while rollouts are mid-turn.
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.

Four runs cut the checkpoint at different points, so that later cuts catch episodes after their state-changing tool calls:

Run	Rollouts	Finished before the checkpoint	Continued	Continued with changed Workplace state	Mean reward of the run
1	40	2	38	0	0.275
2	40	13	27	0	0.300
3	60	28	32	7	0.317
4	80	46	34	10	0.288

Every prepare returned in 8 ms or less and reported no restarts.
Simple Agent reported every in-flight episode at a boundary, and the Workplace server exported one session per continued episode.

The 17 episodes with changed state had already made one or two successful tool calls that changed their session's tables.
Twelve had sent the email reply and five had moved one or two CRM customers to "Lost" before the checkpoint.
Each export differs from a freshly seeded session exactly where those calls changed it.
An email episode's export has one added row, the reply "RE: Task Update on Develop prototype for report generation", and every other email row equal to the seeded data.
A CRM episode that had updated customers 00000037 and 00000195 exports both with status "Lost" where the seeded data has "Proposal", and no other cell differs.

After the restore, all 131 continued episodes resumed rather than restarted:

• All 131 replacements returned HTTP 200 with a reward.
• Each final trajectory begins with the tool calls made before the crash, with the same call IDs and arguments, and the model then continued from there.
• The second checkpoint, taken in the fresh processes before any replacement ran, exported the same tables as the first for all 131 sessions.
• The five CRM episodes kept updating customers on their restored sessions: three or four more updates each, every one answered "Customer updated successfully."
• None of the 276 tool results returned after the restore reported a missing session or a stale attempt.

At the end, the Workplace server and Simple Agent held no sessions.

The rewards are comparable to the uninterrupted run.
Over all 220 rollouts the mean was 0.295 (65 of 220), against 0.25 (10 of 40) uninterrupted.
By task, the email reply succeeded 44 of 44 times and the CRM status move 20 of 44 times, against 7 of 8 and 3 of 8.
The calendar rename succeeded once in 44 and the other two CRM tasks never, as in the baseline.
Of the 17 episodes caught after changing their state, 16 earned a reward of 1.

What is not covered

• /verify scores a trajectory by replaying its tool calls on a fresh environment, so the reward does not depend on the restored session.
  The restored state is shown instead by the second export and by the tool calls that ran on the restored sessions.
• Only the flat-row wrapper ran.
  The typed single_agent_turn request path did not run with Workplace Assistant.
• Only the five example tasks ran; the train and validation splits are downloaded from Hugging Face and were not used.
  The model rarely solves three of the five tasks.
• MCP tool exposure, token capture and generation cuts stayed off.

Compatibility

• Without the global checkpoint: block set to enabled: true, behavior is unchanged. The hooks are only called by the checkpoint control routes.
• With checkpointing on, Workplace Assistant sessions no longer block a checkpoint. They are exported, at about 350 KB of JSON per session.
• Tool behavior, verify, and rewards are unchanged.

[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.

[zyzhou5]

LGTM.

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/39343677971

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/39344263376

[ananthsub]

Rollout evidence rerun on vLLM 0.29.0, the version Gym's vllm extra pins. The evidence in the description ran on vLLM 0.24.0.

The tree was main at dc709bcfc3, which contains this change, #3895 and #3896. The model was Qwen/Qwen3-4B-Instruct-2507, served with the hermes tool-call parser on one NVIDIA RTX 6000 Ada Generation, with uses_reasoning_parser=false and temperature 1.0. The commands and the checkpoint driver are the ones in the description.

Uninterrupted baseline. 40 rollouts, eight per task, gave a mean reward of 0.275. The email reply succeeded 8 of 8 times, the CRM status move 3 of 8 times, and the other three tasks never succeeded. That is the same pattern as on vLLM 0.24.0, where the email reply succeeded 7 of 8 times. There were no failed or masked rollouts and no tracebacks.

Checkpoint, crash and restore. Two runs cut the checkpoint at different points. Each crashed every Gym process with SIGKILL, restored into fresh processes, took a second checkpoint of the restored state, and ran every unfinished rollout again as attempt 1.

Cut	Rollouts	Finished before the checkpoint	Continued	Continued with changed Workplace state	Mean reward of the run
Early	40	2	38	0	0.275
Late	80	45	35	12	0.288

• Prepare returned in 4 ms in both runs and reported no restarts. Simple Agent held every in-flight episode at a boundary, and the Workplace server exported one session for each.
• Six of the 12 episodes with changed state had sent the email reply. Each export adds exactly that reply and leaves every seeded email unchanged. The other six had updated one to three CRM customers, and each export differs from a freshly seeded session only in the cells those successful calls changed. For example, one episode had reassigned customers 00000076 and 00000198 to raj@company.com.
• In both runs, the second export of every restored session equals the first.
• All 73 continued episodes resumed as attempt 1 and returned 200. Each final trajectory starts with exactly the calls, arguments and outputs recorded at the checkpoint. Six of the episodes with changed CRM state went on to update more customers on the restored session. None of the 124 tool outputs returned after the restores reported a missing or stale session.

At the end of both runs, every participant was idle and the Workplace server held no sessions. No server log had a traceback.