feat(checkpoint): adapt initial stateful environments

[ananthsub]

Note

This checkpoint implementation is experimental. Its Python modules live under nemo_gym._checkpoint; they are not a supported public import surface yet.

Summary

• add checkpoint adapters for Blackjack, session-state management, and Workplace Assistant
• require each environment subclass to opt into checkpoint continuation explicitly
• preserve RNG, counters, tabular state, and execution-to-session mappings
• add a cross-layer prepare, commit, restore, resume, lineage, and straggler-retirement test

What this layer adds

The shared participant in #2945 controls admission, locking, revisions, files, and restore activation. Each environment still needs a codec that translates its own live session object into versioned serializable state. This PR supplies the first concrete adapters and demonstrates several state shapes.

flowchart LR
    C[Resources checkpoint participant] -->|export frozen session| G[Gymnasium base adapter]
    G --> B[Blackjack: environment state plus RNG position]
    C --> S[Session-state example: counter value]
    C --> W[Workplace Assistant: toolkit DataFrames]

    B --> M[Versioned ResourceSnapshot]
    S --> M
    W --> M
    M --> D[Controller-provided checkpoint directory]
    D --> V[Validate complete replacement set]
    V --> R[Restore sessions under attempt N+1]
    R --> X[Rebuild execution-to-session mappings]

Loading

The Gymnasium base deep-copies supported session state and tracks the mapping from (rollout_id, attempt_index) to the environment session. Checkpoint continuation defaults to disabled and each supported subclass opts in after defining its state contract. Blackjack adds explicit serialization for random.Random. The counter example persists its logical value. Workplace Assistant serializes and restores its pandas-backed tool state.

Cross-layer recovery test

sequenceDiagram
    participant C as Checkpoint controller
    participant M as Policy model participant
    participant A as Agent participant
    participant R as Resources participant
    participant D as Controller-provided directory

    C->>M: pause and drain accepted generations
    C->>A: prepare active executions
    A-->>C: committed boundary or explicit blocker
    C->>R: prepare bound resource sessions
    C->>M: retire excluded straggler and commit
    C->>A: retire excluded straggler and commit
    C->>R: retire excluded straggler and commit
    M->>D: model-server lineage manifest
    A->>D: agent-instance boundary manifest
    R->>D: resources-server state manifest
    C->>M: restore attempt N+1 paused
    C->>A: restore attempt N+1 paused
    C->>R: restore attempt N+1 paused
    C->>M: explicit resume
    C->>A: explicit resume
    C->>R: explicit resume
    A->>M: first resumed call with attempt N parent
    A->>M: later call linked within attempt N+1

Loading

The test uses production participant and control-route implementations. It proves that the three manifests occupy separate namespaces, source attempts remain fenced, resources revisions survive restore, and the restored-parent relay applies only to the first policy call. TransferQueue durability and the training framework's controller remain simulated because they are outside this Gym stack.

Stack

• depends on: #2945
• stack ends here

Test plan

• environment adapter tests
• complete cross-participant checkpoint-cycle test
• real two-worker Uvicorn admission simulation
• scoped pre-commit checks

Remaining integration work

• a training framework must provide durable rollout IDs and attempts, maintain its live-execution registry, and orchestrate prepare, commit, restore, and resume
• Gym production startup must launch and route the service-level coordinator before multi-worker checkpointing can be advertised
• the current model participant drains accepted generations; live prefix-cut recovery requires additional generation-worker and TransferQueue contracts
• durable acknowledgement and release of completed /run results belongs to the training-framework integration layer

Generation-safe checkpoint update

• adds deterministic active-generation, abort, multi-worker, receipt-replay, and restore races
• validates complete agent/resources/model recovery across the initial adapters

Layer validation: 20 tests passed. Full stack validation: 624 tests passed; Ruff and formatting clean.

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

[ananthsub]

Related to the performance roadmap in #2998. The stateful-environment adapters exercise cleanup and restored-session behavior needed by the lifecycle work in #3037.