Checkpointing Context

Defines exported symbols for orbax.checkpoint.experimental.v1.

Prefer to use the style:

import orbax.checkpoint.experimental.v1 as ocp

Context

final class orbax.checkpoint.experimental.v1.Context(context=None, *, pytree_options=None, array_options=None, async_options=None, multiprocessing_options=None, file_options=None, checkpointables_options=None, pathways_options=None, checkpoint_layout=None, deletion_options=None, memory_options=None, safetensors_options=None)

Context for customized checkpointing.

This class manages the configuration options (e.g., async, multiprocessing, array handling) used during Orbax checkpoint operations using a mutable namespace pattern.

Creating a new Context within an existing Context sets all parameters from scratch by default. To inherit properties from a parent Context, pass the parent context as the first positional or explicit context keyword argument. The new context will inherit the parent’s properties but can be mutated independently.

WARNING: Context variables are thread-local by default. If you dispatch a checkpointing operation to a raw worker thread (e.g. threading.Thread), that thread will not inherit the context and will fall back to default settings. Furthermore, when background tasks or coroutines (e.g. asyncio.create_task, save_async) inherit the active context, they inherit a reference to the exact same underlying Context instance in memory. Consequently, if the main thread exits the with ctx: block and mutates the configuration of ctx, those changes will take effect immediately across any ongoing background asynchronous operations. To avoid unintended side effects, prefer creating a new Context instance (ctx = ocp.Context(parent_ctx)) for separate asynchronous operations rather than mutating a shared context instance mid-flight.

Note: When testing or mixing checkpointer instances and free functions, explicitly wrap free functions inside their own with ocp.Context(…) block, or pass explicit contexts to Checkpointer constructors, to ensure each actor receives its correct active configuration independent of the surrounding context.

Example

Basic usage and explicit inheritance:

from orbax.checkpoint import v1 as ocp

# Basic usage
ctx = ocp.Context()
ctx.pytree.loading.partial_load = True
with ctx:
  ocp.save(directory, tree)

# Inheriting properties from an existing context
ctx1 = ocp.Context()
ctx1.pytree.loading.partial_load = True
with ctx1 as outer_ctx:
  # inner_ctx inherits partial_load, but mutates array saving
  ctx2 = ocp.Context(outer_ctx)
  ctx2.array.saving.use_zarr3 = False
  with ctx2 as inner_ctx:
    ocp.save(directory, tree)

Context is not shared across threads:

from concurrent.futures import ThreadPoolExecutor
from orbax.checkpoint import v1 as ocp

executor = ThreadPoolExecutor(max_workers=1)
ctx = ocp.Context()
ctx.pytree.loading.partial_load = True
with ctx:  # Thread #1 creates Context.
  # The following save call is executed in Thread #2, which sees
  # a "default" Context, NOT the one created above.
  executor.submit(ocp.save, directory, tree)

pytree

Options for PyTree checkpointing. See PyTreeOptions.

array

Options for saving and loading array (and array-like objects). See ArrayOptions.

asynchronous

Options for controlling asynchronous behavior. See AsyncOptions.

multiprocessing

Options for multiprocessing behavior. See MultiprocessingOptions.

file

Options for working with the file system. See FileOptions.

checkpointables

Options for controlling checkpointables behavior. See CheckpointablesOptions.

pathways

Options for Pathways checkpointing. See PathwaysOptions.

checkpoint_layout

The layout of the checkpoint. Defaults to ORBAX. See CheckpointLayout.

deletion

Options for controlling deletion behavior. See DeletionOptions.

memory

Options for controlling memory limits during save / load. See MemoryOptions.

__init__(context=None, *, pytree_options=None, array_options=None, async_options=None, multiprocessing_options=None, file_options=None, checkpointables_options=None, pathways_options=None, checkpoint_layout=None, deletion_options=None, memory_options=None, safetensors_options=None)