Checkpointing a Deterministic Forecast¶
Basic inference workflow checkpointing.
This example shows how to use earth2studio.utils.checkpoint.Checkpoint
to restart a deterministic forecast after it stops partway through a run. See the
checkpointing user guide for more
information on checkpoint catalogs and restart policies.
The example uses earth2studio.data.Random and
earth2studio.models.px.FCN, the FourCastNet AFNO prognostic model.
In this example you will learn:
- Creating a persistent checkpoint
- Running a forecast that stops before the requested final horizon
- Re-opening the IO backend and checkpoint
- Resuming the deterministic workflow from the latest completed lead time
Set Up¶
A restartable forecast needs two persistent locations: one for forecast fields and one for the checkpoint. The IO backend owns the forecast arrays. The checkpoint owns restart metadata plus any model state required to continue the rollout. Model weights and forecast fields are not copied into the checkpoint.
Warning
Model checkpoint state is opt-in. Before relying on restartable inference, verify that the model you plan to use documents checkpoint support. If a model does not support checkpointing yet, open a feature request on the Earth2Studio GitHub.
import os
import shutil
from collections import OrderedDict
from pathlib import Path
import numpy as np
import torch
import earth2studio.run as run
from earth2studio.data import Random
from earth2studio.io import ZarrBackend
from earth2studio.models.px import FCN
from earth2studio.utils.checkpoint import Checkpoint
from earth2studio.utils.time import to_time_array
os.makedirs("outputs", exist_ok=True)
forecast_store = Path("outputs/04_checkpoint_restart.zarr")
checkpoint_store = Path("outputs/04_checkpoint_restart_checkpoint")
# Clean up any left over checkpoints
for path in (forecast_store, checkpoint_store):
if path.exists():
shutil.rmtree(path)
Console output9 lines
/__w/earth2studio/earth2studio/.venv/lib/python3.13/site-packages/torch/cuda/__init__.py:64: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
import pynvml # type: ignore[import]
WARNING[XFORMERS]: xFormers can't load C++/CUDA extensions. xFormers was built for:
PyTorch 2.10.0+cu128 with CUDA 1208 (you have 2.13.0+cu130)
Python 3.10.19 (you have 3.13.13)
Please reinstall xformers (see https://github.com/facebookresearch/xformers#installing-xformers)
Memory-efficient attention, SwiGLU, sparse and more won't be available.
Set XFORMERS_MORE_DETAILS=1 for more details
CuPy distance computation test failed with error: cuVS >= 24.12 or pylibraft < 24.12 should be installed to use this featureLoad the packaged FCN/AFNO forecast model package. The default package points at the FourCastNet model artifacts. Instantiate FCN inside checkpoint contexts so its restart state binds to the active checkpoint session.
Rollout checkpoint state can be staged on the same device used for inference.
Setting device to the current CUDA device can reduce CPU/GPU transfers for
restart tensors during a run. Set it to torch.device("cpu") for CPU-only
development. Since FCN level 2 checkpoints store the complete autoregressive
state, keep history_size small when using mode="append".
compute_device = torch.device(
f"cuda:{torch.cuda.current_device()}" if torch.cuda.is_available() else "cpu"
)
model_package = FCN.load_default_package()
output_variables = np.array(["t2m", "u10m"])
time = ["2024-01-01T00:00:00"]
final_nsteps = 3
first_attempt_nsteps = 1
Preallocate the full output store. A real full-length deterministic run does
this before the first model step. We do it explicitly here because this example
simulates a mid-run stop by intentionally running only the first forecast step.
The IO store writes only two variables, while the checkpoint keeps FCN's complete
restart state internally when level=2 is used.
def deterministic_output_coords(model, time, nsteps, variables):
"""Full output coords for entire roll out"""
input_coords = model.input_coords()
output_coords = model.output_coords(input_coords).copy()
for key, value in model.output_coords(input_coords).items():
if value.shape == (0,):
del output_coords[key]
output_coords["time"] = to_time_array(time)
output_coords["lead_time"] = np.asarray(
[model.output_coords(input_coords)["lead_time"] * i for i in range(nsteps + 1)]
).flatten()
output_coords["variable"] = variables
output_coords.move_to_end("lead_time", last=False)
output_coords.move_to_end("time", last=False)
return output_coords
def model_domain_coords(model):
"""Small helper"""
coords = model.input_coords().copy()
for key in ("batch", "lead_time", "variable"):
coords.pop(key)
return coords
First Attempt¶
Every restartable run should be performed inside a checkpoint context. On an
empty checkpoint, with checkpoint opens a new session for future writes.
Construct restart-aware components inside that context so their state binds to
the active checkpoint session. The workflow records a checkpoint row
after each successful IO write because flush_interval=1 and
mode="append" keeps each row in the printed checkpoint table.
checkpoint = Checkpoint(
"restart-demo",
path=checkpoint_store,
mode="append",
flush_interval=1,
history_size=4,
level=2,
device=compute_device,
)
with checkpoint as ckpt:
model = FCN.load_model(model_package)
domain_coords = model_domain_coords(model)
io = ZarrBackend(str(forecast_store), backend_kwargs={"overwrite": True})
coords = deterministic_output_coords(model, time, final_nsteps, output_variables)
var_names = coords.pop("variable")
io.add_array(coords, var_names)
data = Random(domain_coords=domain_coords)
run.deterministic(
time=time,
nsteps=first_attempt_nsteps,
prognostic=model,
data=data,
io=io,
output_coords=OrderedDict({"variable": output_variables}),
device=compute_device,
verbose=False,
checkpoint=ckpt,
)
print("Checkpoint after the stopped run:")
print(checkpoint)
Console output18 lines
2026-08-15 04:38:18.672 | INFO | earth2studio.run:deterministic:85 - Running simple workflow!
2026-08-15 04:38:18.672 | INFO | earth2studio.run:deterministic:92 - Inference device: cuda:0
2026-08-15 04:38:18.732 | WARNING | earth2studio.io.zarr:add_array:206 - t2m is already in Zarr Store. Skipping add_array.
2026-08-15 04:38:18.733 | WARNING | earth2studio.io.zarr:add_array:206 - u10m is already in Zarr Store. Skipping add_array.
2026-08-15 04:38:19.325 | SUCCESS | earth2studio.run:deterministic:154 - Fetched data from Random
2026-08-15 04:38:19.326 | INFO | earth2studio.run:deterministic:162 - Inference starting!
2026-08-15 04:38:20.179 | SUCCESS | earth2studio.run:deterministic:189 -
Inference complete
Checkpoint after the stopped run:
Checkpoint("restart-demo")
path: outputs/04_checkpoint_restart_checkpoint
mode: append
level: 2
rank: 0/1
id lead_time write_count saved_at
0 0 hours 1 2026-08-15T04:38:19.502654+00:00
1 6 hours 2 2026-08-15T04:38:20.029458+00:00Resume¶
In a new process, re-open the same IO store and checkpoint. The printout above
shows the available row ids. Select -1 to resume from the latest row.
The selected checkpoint session is used as a context manager so the chosen row
is the active restart state while components are constructed and while the
workflow runs. FCN restores its restart state during construction. Its
iterator consumes the selected checkpoint boundary internally and yields the
next forecast state, while the workflow still fetches the normal initial
condition and feeds it to the iterator.
io = ZarrBackend(str(forecast_store))
checkpoint = Checkpoint(
"restart-demo",
path=checkpoint_store,
mode="append",
history_size=4,
level=2,
device=compute_device,
)
with checkpoint.select(-1) as ckpt:
model = FCN.load_model(model_package)
domain_coords = model_domain_coords(model)
data = Random(domain_coords=domain_coords)
run.deterministic(
time=time,
nsteps=final_nsteps,
prognostic=model,
data=data,
io=io,
output_coords=OrderedDict({"variable": output_variables}),
device=compute_device,
verbose=False,
checkpoint=ckpt,
)
print("Checkpoint after resume:")
print(checkpoint)
print(io.root.tree())
Console output27 lines
2026-08-15 04:38:22.312 | INFO | earth2studio.run:deterministic:85 - Running simple workflow!
2026-08-15 04:38:22.312 | INFO | earth2studio.run:deterministic:92 - Inference device: cuda:0
2026-08-15 04:38:22.362 | WARNING | earth2studio.io.zarr:add_array:206 - t2m is already in Zarr Store. Skipping add_array.
2026-08-15 04:38:22.364 | WARNING | earth2studio.io.zarr:add_array:206 - u10m is already in Zarr Store. Skipping add_array.
2026-08-15 04:38:22.950 | SUCCESS | earth2studio.run:deterministic:154 - Fetched data from Random
2026-08-15 04:38:22.951 | INFO | earth2studio.run:deterministic:162 - Inference starting!
2026-08-15 04:38:23.512 | SUCCESS | earth2studio.run:deterministic:189 -
Inference complete
Checkpoint after resume:
Checkpoint("restart-demo")
path: outputs/04_checkpoint_restart_checkpoint
mode: append
level: 2
rank: 0/1
id lead_time write_count saved_at
0 0 hours 1 2026-08-15T04:38:19.502654+00:00
1 6 hours 2 2026-08-15T04:38:20.029458+00:00
2 12 hours 3 2026-08-15T04:38:23.091748+00:00
3 18 hours 4 2026-08-15T04:38:23.363178+00:00
/
├── lat (720,) float64
├── lead_time (4,) timedelta64[h]
├── lon (1440,) float64
├── t2m (1, 4, 720, 1440) float32
├── time (1,) datetime64[ns]
└── u10m (1, 4, 720, 1440) float32