Checkpointing and Resuming#

Aspire provides a few simple patterns to resume long runs.

Saving checkpoints while sampling#

Pass checkpoint_path (an HDF5 file) to aspire.Aspire.sample_posterior() to write checkpoints as the sampler runs. Use checkpoint_every to control frequency and checkpoint_save_config/flow to control what metadata is saved.
For a convenience wrapper, wrap your sampling in with aspire.auto_checkpoint("run.h5", every=1): .... Inside the context, sample_posterior will default to checkpointing to that file, and the config/flow will be updated as needed.

What gets saved#

The sampler stores checkpoints under /checkpoint/state in the HDF5 file.
Aspire writes /aspire_config for Aspire-level configuration, /sampler_config for sampler metadata, and /flow for the saved flow. If these already exist, they are overwritten when saving.
resume_from_file and auto_checkpoint(..., resume=True) also accept legacy files where sampler metadata was embedded inside /aspire_config instead of stored under /sampler_config.

Resuming from a file#

Use aspire.Aspire.resume_from_file() to rebuild an Aspire instance and flow from a checkpoint file:

aspire = Aspire.resume_from_file(
    "run.h5",
    log_likelihood=log_likelihood,
    log_prior=log_prior,
)
# Optionally continue checkpointing to the same file
with aspire.auto_checkpoint("run.h5", every=1):
    samples = aspire.sample_posterior(...)

resume_from_file loads config, flow, and the last checkpoint (if present), and primes the instance to resume sampling; you can still override sampler kwargs when calling sample_posterior.

Resuming with `auto_checkpoint`#

If you use auto_checkpoint(..., resume=True), it will attempt to load the flow and config from the checkpoint file when entering the context. If a checkpoint is found, it will also load that and prime the instance to resume sampling. If no checkpoint is found, it will simply load the flow and config (if present) and start fresh sampling.

aspire = Aspire(...)

with aspire.auto_checkpoint("run.h5", every=1, resume=True):
    # Flow with be loaded from the file if present and training will be skipped
    aspire.fit(...)
    # If a checkpoint was found, sampling will resume from that; otherwise it will start from scratch
    samples = aspire.sample_posterior(...)

Manual resume via `sample_posterior` args#

If you have a checkpoint blob (bytes or dict) already, you can pass it directly:

samples = aspire.sample_posterior(
    n_samples=...,
    sampler="smc",
    resume_from=checkpoint_bytes_or_dict,
    checkpoint_path="run.h5",  # optional: keep writing checkpoints
)

To resume from a file without using resume_from_file, load the checkpoint bytes and flow yourself, then call sample_posterior:

from aspire.utils import AspireFile

aspire = Aspire(..., flow_backend="zuko")
with AspireFile("run.h5", "r") as f:
    aspire.load_flow(f, path="flow")
    # Standard layout is /checkpoint/state; adjust if you used a different path
    checkpoint_bytes = f["checkpoint"]["state"][...].tobytes()
samples = aspire.sample_posterior(
    n_samples=..., sampler="smc", resume_from="run.h5"
    # or resume_from=checkpoint_bytes if you prefer to pass bytes directly
)

Notes and tips#

Checkpoint files must be HDF5 (.h5/.hdf5).
If a checkpoint is missing in the file (e.g., sampling never wrote one), the flow and config are still loaded; you can simply start sampling again and checkpointing will continue to the same file.
For manual control, you can always call save_config / save_flow yourself on an aspire.utils.AspireFile.
SMC samplers also accept a custom checkpoint_callback and checkpoint_every if you want full control over how checkpoints are persisted or inspected. Provide a callable that accepts the checkpoint state dict; from there you can, for example, serialize to another format or push to remote storage.
When resuming, you must still specify the keyword arguments for sample_posterior.

Checkpointing and Resuming

Contents