Add a versioned, topology-aware CheckpointStore

[joglekara]

Motivation

Checkpointing is currently entangled with save/postprocess paths. Downstream runs have OOMed while materializing full-state saves, failed after the solve because a checkpoint dependency was absent, produced non-restorable checkpoints, or lost outputs after late artifact failures.

A checkpoint is durable execution state, not a diagnostic observation or an MLflow artifact convention.

Scope

Define a CheckpointStore protocol with save, restore, validate, list, and latest operations. Add versioned metadata including:

• checkpoint/schema version;
• program identity and compatibility version;
• config/provenance and structural fingerprints;
• simulation time, step, and chunk identity;
• PyTree paths, shapes, dtypes, and logical sharding;
• package/code versions.

Require atomic commit semantics. Save the solver carry directly rather than routing it through ObservationPlan/SaveAt buffers. Support restoration into a compatible target topology. Implement initial null and local/Orbax-backed stores plus checkpoint policy in a run plan.

Acceptance criteria

• An interrupted serial run resumes to the uninterrupted result within solver tolerance.
• A sharded integration test round-trips state and restores onto a different compatible layout.
• Injected failure during save leaves the prior checkpoint valid and excludes partial output from latest.
• Missing dependencies, unwritable paths, incompatible schemas, and incompatible trees fail during preflight/restore with actionable errors.
• Asynchronous save completion is explicitly awaited and restore-validated in tests.
• Checkpoint failure does not erase already-produced metrics or artifacts.
• Version migrations are registered/tested; unsupported versions fail clearly.
• MLflow is not the primary checkpoint protocol.

Dependencies

• Stable state/result PyTrees from Define an explicit pure-JAX JaxProgram boundary for continuous and discrete systems #350 and Make Params, State, Inputs, RNG, and objectives explicit #352.
• Multi-host commit coordination from the RunPlan/Executor issue.

Out of scope

• Transparent restoration across arbitrary physics/schema changes.
• A general distributed object store beyond the initial filesystem/Orbax backend.