CheckpointableHandler Registration

CheckpointableHandler Registration#

Public API for CheckpointableHandlers.

CheckpointableHandlerRegistry#

Utilities#

orbax.checkpoint.experimental.v1.handlers.global_registry()[source][source]#

Returns the global registry.

The global registry serves as the default, singleton storage for all handlers registered throughout the application’s lifecycle via register_handler.

Example

Retrieve the global registry to inspect available handlers:

from orbax.checkpoint.v1 import handlers

# Fetch the singleton global registry
registry = handlers.global_registry()

# Check if a specific handler name is registered globally
is_registered = registry.has("my_custom_model_handler")
Returns:

The global singleton registry instance.

Return type:

CheckpointableHandlerRegistry

orbax.checkpoint.experimental.v1.handlers.local_registry(other_registry=None, *, include_global_registry=True)[source][source]#

Creates a local registry.

This function builds a new registry by optionally combining the existing global registry with a provided custom registry. It is highly useful for overriding handlers for a specific checkpointer operation without mutating the global state.

Example

Create a registry with custom handlers, potentially including global ones:

from orbax.checkpoint.v1 import handlers

class MyHandler(handlers.CheckpointableHandler):
  pass

# Create a registry and add a handler. By default, it includes
# globally-registered handlers.
my_registry = handlers.local_registry()
my_registry.add(MyHandler)

# To start with an empty registry, use:
# my_registry = handlers.local_registry(include_global_registry=False)
Parameters:

  • other_registry (UnionType[CheckpointableHandlerRegistry, None]) – An optional registry of handlers to include in the returned registry.

  • include_global_registry (bool) – If True, includes globally-registered handlers in the returned registry by default.

Return type:

CheckpointableHandlerRegistry

Returns:

A local registry.

orbax.checkpoint.experimental.v1.handlers.register_handler(cls, *, checkpointable_name=None, secondary_typestrs=None)[source][source]#

Registers a CheckpointableHandler globally.

The order in which handlers are registered strictly matters. If multiple handlers could potentially be used to save or load an object (i.e., are capable of handling the checkpointable according to is_handleable/ is_abstract_handleable for save/load, respectively), the framework resolves them in Last-In, First-Out (LIFO) order. This means the handler added most recently will be selected.

Example

Registering a custom handler using a direct function call. Note the import path from the v1 namespace:

from orbax.checkpoint.v1 import handlers

class BarHandler(handlers.CheckpointableHandler):
  pass

handlers.register_handler(BarHandler)
Parameters:

  • cls (~CheckpointableHandlerType) – The handler class to register globally.

  • checkpointable_name (UnionType[str, None]) – The checkpointable name. If not-None, the registered handler will be scoped to that specific name. Otherwise, the handler will be available for any checkpointable name.

  • secondary_typestrs (Optional[Sequence[str], None]) – A sequence of alternate handler typestrs that serve as secondary identifiers for the handler.

Return type:

~CheckpointableHandlerType

Returns:

The handler class.

Contents