ocp.v1.handlers module

`ocp.v1.handlers` module#

Public API for CheckpointableHandlers.

Types#

class orbax.checkpoint.experimental.v1.handlers.CheckpointableHandler(*args, **kwargs)[source][source]#

An interface that defines save/load logic for a checkpointable object.

NOTE: Prefer to use StatefulCheckpointable interface when possible.

A PyTree of arrays, representing model parameters, is the most basic “checkpointable”. A singular array is also a checkpointable.

In most contexts, when dealing with just a PyTree, the API of choice is:

ocp.save(directory, pytree)

The concept of “checkpointable” is not so obvious in this case. When dealing with multiple objects, we can use:

ocp.save_checkpointables(
    directory,
    dict(
        pytree=model_params,
        dataset=dataset_iterator,
        # other checkpointables, e.g. extra metadata, etc.
    ),
)

Now, it is easy to simply skip loading the dataset, as is commonly desired when running evals or inference:

ocp.load_checkpointables(
    directory,
    dict(
        pytree=abstract_model_params,
    ),
)
# Equivalently,
ocp.load(directory, abstract_model_params)

With the methods defined in this Protocol (save, load), logic within the method itself is executed in the main thread, in a blocking fashion. Additional logic can be executed in the background by returning an Awaitable function (which itself may return a result).

Let’s look at some suggestions on how to implement a CheckpointableHandler.

To create a custom handler, you must define a class that implements the methods defined in this Protocol. The class should be generic over the concrete type Checkpointable (the object being saved/loaded) and the abstract type AbstractCheckpointable (the lightweight metadata representation).

Crucially, once implemented, the handler must be registered with the global registry or a context-local registry so that save_checkpointables and load_checkpointables can automatically detect and use it for the corresponding types. Use register_handler() for global registration, or provide handlers via CheckpointablesOptions for context-local registration.

First, take a look at handler_utils.py for some toy implementations used for unit testing.

Here are some details on how to implement is_handleable and is_abstract_handleable.

For example, if a handler may be defined as follows:

class FooHandler(CheckpointableHandler[Foo, AbstractFoo]):

  def is_handleable(self, checkpointable: Foo) -> bool:
    return isinstance(foo, Foo)

  def is_abstract_handleable(
      self, abstract_checkpointable: AbstractFoo) -> bool:
    return isinstance(abstract_foo, AbstractFoo)

This is simple because the handler only works with Foo and AbstractFoo. But the handler may work on more generic types. In a toy example, let’s say we’ve developed an improved way of storing very large arrays, which is still suboptimal for more normal-sized arrays. We can implement the handler as:

class FooHandler(CheckpointableHandler[jax.Array, jax.ShapeDtypeStruct]):

  def is_handleable(self, checkpointable: jax.Array) -> bool:
    return (
        isinstance(checkpointable, jax.Array)
        and checkpointable.size > LARGE_ARRAY_THRESHOLD
    )

  def is_abstract_handleable(
      self, abstract_checkpointable: jax.ShapeDtypeStruct) -> bool:
    return (
        isinstance(abstract_checkpointable, jax.ShapeDtypeStruct)
        and abstract_checkpointable.size > LARGE_ARRAY_THRESHOLD
    )

In many cases, no information is needed for loading. In this case, AbstractCheckpointable may be defined as None. For example:

class FooHandler(CheckpointableHandler[Foo, None]):

  def is_handleable(self, checkpointable: Foo) -> bool:
    return isinstance(checkpointable, Foo)

  def is_abstract_handleable(self, abstract_checkpointable: None) -> bool:
    return abstract_checkpointable is None

class orbax.checkpoint.experimental.v1.handlers.StatefulCheckpointable(*args, **kwargs)[source][source]#: An interface that defines save/load logic for a checkpointable object.

Handlers#

final class orbax.checkpoint.experimental.v1.handlers.PyTreeHandler(*, context=None, array_metadata_validator=<orbax.checkpoint._src.metadata.array_metadata_store.Validator object>, leaf_handler_registry=None, partial_save_mode=False)[source][source]#

An implementation of CheckpointableHandler for PyTrees.

PyTreeHandler manages the decomposition of JAX PyTree structures into leaf- level parameters for persistence. It utilizes an asynchronous two-tier execution model to allow for background I/O, ensuring that heavy array serialization does not block the main training process.

Note: Users are encouraged NEVER to instantiate or use this handler directly. Always use the top-level APIs like ocp.save_checkpointables and ocp.load_checkpointables. Orbax uses this handler by default for standard JAX PyTrees (like nested dictionaries of arrays).

To configure a specific serialization context for a PyTree and aggressively force Orbax to use the customized PyTreeHandler, the recommended approach is to use ocp.Context with CheckpointablesOptions. This allows you to bind the handler to a specific dictionary key within the Context scope.

See CheckpointablesOptions for more details on handler registration.

Usage Example:

Save a state dictionary configuration:

import orbax.checkpoint as ocp

state_pytree = {'weights': [1.0, 2.0], 'bias': 0.0}

registry = ocp.handlers.local_registry()
registry.add(
    ocp.handlers.PyTreeHandler, checkpointable_name='model_state'
)
ctx = ocp.Context()
ctx.checkpointables.registry = registry
with ctx:
    ocp.save_checkpointables(path, dict(model_state=state_pytree))

context#

Optional V1 Context providing configuration for serialization, array options, and multiprocessing coordination.

Type:: Optional[Context]

array_metadata_validator#

A validator object used to verify consistency of array metadata during restoration.

Type:: Validator

final class orbax.checkpoint.experimental.v1.handlers.ProtoHandler(filename='proto.pbtxt')[source][source]#

Implementation of CheckpointableHandler for protocol buffers.

ProtoHandler manages the serialization and deserialization of Protocol Buffer messages in text format. It utilizes an asynchronous two-tier execution model to offload I/O operations, ensuring background writing does not block the main process. In distributed environments, it provides multihost coordination to ensure that only the primary host performs the write operation.

Note: Users are encouraged NEVER to instantiate or use this handler directly. Always use the top-level APIs like ocp.save_checkpointables and ocp.load_checkpointables. Orbax uses this handler by default for standard protocol buffer messages.

To save a custom Protocol Buffer message and aggressively force Orbax to use the ProtoHandler (e.g., to specify a custom filename), the recommended approach is to use ocp.Context with CheckpointablesOptions. This allows you to bind the handler to a specific dictionary key within the Context scope.

See CheckpointablesOptions for more details on handler registration.

Example Usage:

Save a protobuf message configuration:

import orbax.checkpoint as ocp

# Assuming MyProtoMessage is your compiled protobuf class
my_proto_msg = MyProtoMessage(config_field="value")

registry = ocp.handlers.local_registry()
registry.add(
    ocp.handlers.ProtoHandler, checkpointable_name="proto_config"
)
ctx = ocp.Context()
ctx.checkpointables.registry = registry
with ctx:
    ocp.save_checkpointables(path, dict(proto_config=my_proto_msg))

filename#

An optional filename used for saving and loading the protobuf data. If not provided, it defaults to a standard internal default filename.

Type:: str

final class orbax.checkpoint.experimental.v1.handlers.JsonHandler(filename=None)[source][source]#

An implementation of CheckpointableHandler for Json.

JsonHandler enables the persistence of standard Python structures (dicts, lists, and primitives) that are JSON-serializable. It utilizes an asynchronous two-tier execution model to offload I/O operations, ensuring background writing does not block the main process. It also provides multihost coordination to ensure that only the primary host performs the write operation.

Note: Users are encouraged NEVER to instantiate or use this handler directly. Always use the top-level APIs like ocp.save_checkpointables and ocp.load_checkpointables. Orbax uses this handler by default for standard JSON-serializable objects.

To save a custom JSON-serializable object (like a specific dictionary containing metadata) and aggressively force Orbax to use the JsonHandler, the recommended approach is to use ocp.Context with CheckpointablesOptions, which only applies to save/load operations strictly within the Context scope.

See CheckpointablesOptions for more details on handler registration.

Example Usage:

Save a dictionary configuration:

import orbax.checkpoint as ocp

config = {'learning_rate': 0.01, 'batch_size': 32}

registry = ocp.handlers.local_registry()
registry.add(
    ocp.handlers.JsonHandler,
    checkpointable_name='experiment_config',
)
ctx = ocp.Context()
ctx.checkpointables.registry = registry
with ctx:
    ocp.save_checkpointables(path, dict(experiment_config=config))

filename#: An optional specific filename to use for saving and loading the JSON data. If not provided, the handler will fall back to a default set of supported JSON filenames.

ocp.v1.handlers module

Contents

ocp.v1.handlers module#

Types#

Handlers#

`ocp.v1.handlers` module#