feat(checkpoint): export Blackjack game state for partial-rollout checkpoints

[ananthsub]

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

---

What changed and why

Blackjack is the first Gymnasium environment to opt in to checkpointing. Without it, a live Blackjack game is restart_only: it blocks a checkpoint until its rollout is retired and replayed from the start.

The rest of a Blackjack game depends only on the player's hand, the dealer's hand, and the session's random number generator. Blackjack sets checkpoint_mode = "exported" and overrides the two Gymnasium hooks:

• export_env_state writes both hands and the generator state (random.Random.getstate()) as JSON.
• restore_env_state checks that the hands are lists of card ranks, rebuilds the generator with setstate, and raises ValueError for any malformed state. Nothing is installed if any session is invalid.

The Gymnasium base handles everything else: tracking sessions from /reset to the terminal /step, exporting every live session at commit, and closing a retired session.

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, fencing, lease, storage"]
  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")]
  C --> CO --> K
  K --> E & M & A & R
  M --> W
  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
  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
  C->>G: resume
  C->>G: /run as attempt + 1 continues each episode from its boundary
  G-->>C: a late call from attempt 0 gets 409 stale_attempt

Loading

This PR

The game state that decides the rest of a hand, saved at a checkpoint and restored into a fresh server.

sequenceDiagram
  participant A as Agent
  participant B as Blackjack
  A->>B: /reset deals the hands
  A->>B: /step hit
  Note over B: commit - export_env_state saves the player's hand,<br/>the dealer's hand, and the random number generator state
  Note over A,B: crash, then restore into a fresh server
  Note over B: restore_env_state checks the hands and rebuilds the generator
  A->>B: /step stand
  B-->>A: the same outcome as an uncheckpointed game, including the dealer's draws

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) (this 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 #2946 onto the v2 checkpoint hooks. The old version serialized the same three fields through the Gymnasium base's serialize_session_state and deserialize_session_state, keyed by rollout attempt. Here the state is keyed by the cookie session ID, and restore also validates the hands and turns a bad generator state into a clear error.

Issue

No separate issue. This is the Blackjack part of the partial-rollout checkpointing port.

Validation

• RAY_TMPDIR=/tmp python -m pytest -q resources_servers/blackjack/tests: 23 passed (16 existing, 7 new). The new tests drive a checkpoint-enabled server through its control routes:
  • With a fixed seed, reset and hit, commit, restore into a fresh server instance, then stand. The restored game's final step equals an uncheckpointed game with the same seed and actions, including the dealer's cards drawn after the restore. Breaking the generator restore makes this test fail.
  • Invalid state written at commit fails the restore, and nothing is installed.
  • Malformed states (missing generator, bad generator state, a hand that is not a list) are rejected.
  • A retired session is closed, and its next step is refused as stale.
  • A session the server already dropped is left out of the export.
• RAY_TMPDIR=/tmp python -m pytest -q tests/unit_tests/test_checkpoint_*.py resources_servers/gymnasium/tests responses_api_agents/gymnasium_agent/tests: 145 passed.
• ruff check, ruff format, and pre-commit run --files on the changed files.

Rollout evidence

Pending. Before marking ready: run Blackjack rollouts through gymnasium_agent with a real model and checkpointing on, with a checkpoint taken mid-game and a restore into a fresh process, and confirm the continued games finish normally. The process-level checkpoint e2e suite does not cover Gymnasium environments yet.

Compatibility

• With checkpointing off, nothing changes.
• With checkpointing on, live Blackjack games are exported instead of blocking prepare until their rollouts are retired.
• Game behavior, data, 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.