feat(checkpoint): checkpoint Workplace Assistant sessions

[ananthsub]

• Branch: ananthsub/partial-ckpt-workplace-assistant
• Base branch: ananthsub/partial-ckpt-e2e
• Suggested labels: feature, area:env-infra (draft, so no state label yet)

---

What changed and why

Workplace Assistant sessions used to block a checkpoint, because the server used the default restart_only mode. Every rollout with a live Workplace session had to be retired and restarted from scratch. 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, retire"]
  K["Control plane on every server<br/>phases, lease, storage,<br/>retire: stop, then free"]
  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 released
  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 + 1, attempt N's capture ledger retired
  C->>G: resume
  C->>G: /run as attempt + 1 continues each episode from its boundary
  Note over C,G: dropping an episode, only while no checkpoint is open
  C->>G: retire - environment, then agent, then model and resources
  Note over G: each server stops the attempt's work, waits, frees its state, then replies

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 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 (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:

• RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q resources_servers/workplace_assistant/tests: 36 passed, 4 of them new in tests/test_checkpoint.py. The new 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 gets the server's own error 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: 136 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

• No process-level run with Workplace Assistant yet. Checkpointing a live episode end to end also needs the agent lane (Simple Agent continuation), which the base branch includes.
• Real-model rollout evidence: pending. The next step is a checkpoint, kill, and restore during a Workplace Assistant rollout with a real policy model, comparing the reward against an uninterrupted run.

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.