Checkpointing

Writing checkpoint files

Long-running simulations commonly write checkpoint files so that runs can be restarted from a recent saved state. In JustPIC, checkpoints are written in JLD2 format.

At the lowest level, you can serialize arrays manually:

jldsave(
    "my_file.jld2";
    particles     = Array(particles),
    phases        = to_cpu(phases),
    phase_ratios  = Array(phase_ratios),
    particle_args = to_cpu.(particle_args),
)

This saves particle information to my_file.jld2, ready to be reloaded later.

If file size matters more than exact restart reproducibility, you can downcast to Float32 before writing:

jldsave(
    "my_file.jld2";
    particles     = Array(Float32, particles),
    phases        = to_cpu(Float32, phases),
    phase_ratios  = Array(Float32, phase_ratios),
    particle_args = to_cpu.(Float32, particle_args),
)

For routine use, prefer the built-in helper:

checkpointing_particles(
    dst,
    particles;
    phases = nothing,
    phase_ratios = nothing,
    chain = nothing,
    t = nothing,
    dt = nothing,
    particle_args = nothing,
    kwargs...,
)

On multiple MPI ranks, pass the rank id to get rank-local filenames:

checkpointing_particles(
    dst,
    particles,
    me;
    phases = nothing,
    phase_ratios = nothing,
    chain = nothing,
    t = nothing,
    dt = nothing,
    particle_args = nothing,
    kwargs...,
)

Any additional keyword arguments are written into the checkpoint as extra fields. Arrays, CellArrays and JustPIC containers are moved to host memory, tuples and named tuples are converted element by element, and numbers, strings, symbols and nothing are stored as they are. Any other value is rejected with an ArgumentError before the checkpoint file is touched.

The checkpoint is first written to a temporary file in the destination directory and then renamed over the previous checkpoint. On POSIX filesystems this replacement is atomic: a write that fails or is interrupted leaves the previous checkpoint intact.

Loading a checkpoint file

To restart a simulation, load the file with load_checkpoint and cast the stored arrays back to the active backend:

data          = load_checkpoint("my_file.jld2")
particles     = TA(backend)(Float64, data["particles"])
phase_ratios  = TA(backend)(Float64, data["phase_ratios"])
phases        = data["phases"]
particle_args = data["particle_args"]

TA(backend)(T, x) moves a JustPIC container such as Particles or PhaseRatios to backend, converting its values to T, so the same checkpoint can be restored onto CPU or accelerator arrays. Plain CellArrays such as phases and particle_args load as CPU CellArrays. On a GPU backend, move them with the same call, TA(backend)(Float64, data["phases"]) and TA(backend).(Float64, data["particle_args"]); on the CPU, use them as loaded or change their precision with to_cpu(Float64, data["phases"]).

On multiple MPI ranks, each rank reads its own file with load_checkpoint(dst, me).

Compatibility

Every checkpoint records a schema version ("checkpoint_schema"), the JustPIC version that wrote it ("justpic_version") and, for rank-local files, the MPI partition ("partition": rank, number of ranks, process grid, Cartesian coordinates and global grid size). load_checkpoint enforces the following:

• A checkpoint with a newer schema version than the running JustPIC supports is rejected with an error.

• A checkpoint without a schema version, written by JustPIC 0.7.9 or earlier, loads with a warning and without further checks.

• A rank-local checkpoint must have been written by the same rank. When ImplicitGlobalGrid is initialized, the number of ranks, process grid, Cartesian coordinates and global grid size must match the ones recorded in the file: restarting on a different decomposition is not supported.

Checkpoints remain plain JLD2 files, so JLD2.load still reads them without these checks.

API

JustPIC.checkpointing_particles Function

checkpointing_particles(dst, particles; phases=nothing, phase_ratios=nothing, chain=nothing, t=nothing, dt=nothing, particle_args=nothing, kwargs...)
checkpointing_particles(dst, particles, me::Integer; kwargs...)
checkpointing_particles(dst, particles, fname::AbstractString; kwargs...)

Write particle state and optional companion data to a JLD2 checkpoint.

By default the file is saved as particles_checkpoint.jld2 in dst. Passing me writes rank-local files named after the zero-based MPI rank: particles_checkpoint0000.jld2, particles_checkpoint0001.jld2, and so on. Passing fname writes to that path instead.

Common keywords

• phases: per-particle phase labels.

• phase_ratios: PhaseRatios container to checkpoint.

• chain: marker-chain state.

• t: simulation time, stored as "time".

• dt: timestep size, stored as "timestep".

• particle_args: tuple of extra particle-carried fields.

Additional keyword arguments are stored under their own names.

Payload

Every value is converted to host memory before writing: arrays and CellArrays become CPU arrays, JustPIC containers are converted with Array, and Tuples and NamedTuples are converted element by element. Numbers, strings, symbols and nothing are stored as they are. Any other value throws an ArgumentError before the file is touched. The keys checkpoint_schema, justpic_version and partition hold metadata that load_checkpoint checks, and cannot be used as keyword names; neither can time and timestep.

Replacement

The checkpoint is written to a temporary file in the destination directory and then renamed over fname. On filesystems where a rename within one directory is atomic (POSIX filesystems), a failed or interrupted write leaves the previous checkpoint intact.

source

JustPIC.load_checkpoint Function

load_checkpoint(fname::AbstractString)
load_checkpoint(dst, me::Integer)

Read a checkpoint written by checkpointing_particles and return a Dict{String, Any} with the same keys that JLD2.load returns. Stored arrays and containers are on the CPU; convert them to the active backend with TA(backend).

The second form reads the rank-local file of rank me in dst.

Compatibility

• Checkpoints whose schema version is newer than this JustPIC supports are rejected with an ArgumentError.

• Checkpoints without schema metadata (written before schema versioning) load with a warning, without compatibility checks.

• The rank-local form checks that the file was written by rank me. When ImplicitGlobalGrid is initialized and the file records a topology, the number of ranks, the process grid, the rank's Cartesian coordinates and the global grid size must match the current ones.

source