# Copyright 2026 The Orbax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Types and constructs for PyTree metadata and serialization."""
from __future__ import annotations
import abc
import copy
import dataclasses
import enum
from typing import Any, Callable, Optional, Protocol, Sequence
from absl import logging
from etils import epath
import jax
import jax.numpy as jnp
import numpy as np
from orbax.checkpoint._src.arrays import types as arrays_types
from orbax.checkpoint._src.futures import future
from orbax.checkpoint._src.metadata import empty_values
from orbax.checkpoint._src.metadata import pytree_metadata_options as pytree_metadata_options_lib
from orbax.checkpoint._src.metadata import value as value_metadata
from orbax.checkpoint._src.path import types as path_types
from orbax.checkpoint._src.serialization import limits
from orbax.checkpoint._src.tree import types as tree_types
import tensorstore as ts
PyTreeMetadataOptions = pytree_metadata_options_lib.PyTreeMetadataOptions
class TransferPriority(enum.Enum):
"""Prioritization of arrays for saving.
NOTE: If the option is ASYNCHRONOUS_*, values may be altered if updated
concurrently before transfer.
SYNCHRONOUS: The device to host transfer is scheduled and executed
synchronously before returning to the caller. This ensures values are
never corrupted by concurrent updates or donations.
ASYNCHRONOUS_PRIORITIZED: The device to host transfer is scheduled
asynchronously in the background, but prioritized ahead of deprioritized
transfers.
ASYNCHRONOUS_DEPRIORITIZED: The device to host transfer is scheduled
asynchronously in the background after all prioritized transfers have
completed.
UNKNOWN: The transfer priority is unknown or unspecified.
"""
SYNCHRONOUS = 0
ASYNCHRONOUS_PRIORITIZED = 1
ASYNCHRONOUS_DEPRIORITIZED = 2
UNKNOWN = 3
class IoDirection(enum.Enum):
READ = 'read'
WRITE = 'write'
class SerializationStatusCallback(Protocol):
"""Callback for tracking serialization status of PyTree parameters.
NOTE: Callback methods can be called concurrently from parallel background
threads. Implementations must ensure thread safety if they access or modify
shared state.
Usage:
Users can implement this protocol to customize parameter serialization
behavior, such as transfer prioritization from device to host or to
receive notifications about transfer and write completion.
To use a custom callback, implement the protocol and register it with
`ArrayHandler` for `jax.Array` type before saving:
```
from orbax.checkpoint._src.serialization import types as serialization_types
from orbax.checkpoint._src.tree import types as tree_types
class MyCallback(serialization_types.SerializationStatusCallback):
def key_priority(self, keypath: tree_types.PyTreePath):
if 'large_param' in jax.tree_util.keystr(keypath):
return serialization_types.TransferPriority.ASYNCHRONOUS_DEPRIORITIZED
return serialization_types.TransferPriority.ASYNCHRONOUS_PRIORITIZED
def on_transfer_end(self, keypath: tree_types.PyTreePath):
print(f'Transfer ended for {jax.tree_util.keystr(keypath)}')
def on_write_end(self, keypath: tree_types.PyTreePath):
print(f'Write ended for {jax.tree_util.keystr(keypath)}')
ocp.type_handlers.register_type_handler(
jax.Array,
ocp.type_handlers.ArrayHandler(callback=MyCallback()),
override=True,
)
```
"""
def key_priority(self, keypath: tree_types.PyTreePath) -> TransferPriority:
"""Callback for when the parameter is checked for transfer priority.
Must return a TransferPriority indicating how the transfer will be
scheduled.
Args:
keypath: The keypath of the parameter being checked.
"""
...
def on_transfer_start(self, keypath: tree_types.PyTreePath) -> None:
"""Callback for when the parameter transfer to host starts.
Note: Currently not wired up. Reserved for future use.
Args:
keypath: The keypath of the parameter being transferred.
"""
...
def on_transfer_end(self, keypath: tree_types.PyTreePath) -> None:
"""Callback for when the parameter transfer to host ends.
Args:
keypath: The keypath of the parameter whose transfer has ended.
"""
...
def on_write_start(self, keypath: tree_types.PyTreePath) -> None:
"""Callback for when the parameter write starts.
Note: Currently not wired up. Reserved for future use.
Args:
keypath: The keypath of the parameter being written.
"""
...
def on_write_end(self, keypath: tree_types.PyTreePath) -> None:
"""Callback for when the parameter write ends.
Args:
keypath: The keypath of the parameter whose write has ended.
"""
...
class DefaultSerializationStatusCallback(SerializationStatusCallback):
"""Default no-op callback for serialization and deserialization."""
def key_priority(self, _: tree_types.PyTreePath) -> TransferPriority:
"""Callback for when the parameter is checked for transfer priority."""
return TransferPriority.SYNCHRONOUS
def on_transfer_start(self, _: tree_types.PyTreePath) -> None:
"""Callback for when the parameter transfer to host starts."""
pass
def on_transfer_end(self, _: tree_types.PyTreePath) -> None:
"""Callback for when the parameter transfer to host ends."""
pass
def on_write_start(self, _: tree_types.PyTreePath) -> None:
"""Callback for when the parameter write starts."""
pass
def on_write_end(self, _: tree_types.PyTreePath) -> None:
"""Callback for when the parameter write ends."""
pass
def is_supported_type(
value: Any,
pytree_metadata_options: PyTreeMetadataOptions = (
pytree_metadata_options_lib.PYTREE_METADATA_OPTIONS
),
) -> bool:
"""Determines if the `value` is supported without custom TypeHandler."""
return isinstance(
value,
(str, int, float, np.number, np.ndarray, bytes, jax.Array),
) or empty_values.is_supported_empty_value(value, pytree_metadata_options)
def check_input_arguments(*args):
l = None
for arg in args:
if l == 0:
raise ValueError('Cannot pass TypeHandler input of length 0.')
if l is None:
l = len(arg)
elif len(arg) != l:
raise ValueError('Found input args with mismatched lengths.')
class ParamInfo:
"""Information describing a parameter in a PyTree.
Note that ParamInfo is distinct from SaveArgs and RestoreArgs in that in
represents information not provided by a user, and should be computed
internally.
Attributes:
name: Name of the parameter.
parent_dir: A path providing location where all files under the same
checkpoint should be saved under. All ParamInfo provided to a given
TypeHandler should have the same parent_dir. The parent_dir is assumed to
be a directory. Accessing this property raises ValueError if the
underlying path is a PathAwaitingCreation that has not been resolved via
await_path_creation().
path: Automatically set to parent_dir / name. Accessing this property raises
ValueError if the underlying path is a PathAwaitingCreation that has not
been resolved via await_path_creation().
skip_deserialize: If specified, skips deserialization of the given parameter
using the TypeHandler.
byte_limiter: Object to limit the number of bytes that can be read or
written in parallel.
device_host_byte_limiter: Object to limit the number of bytes that can be
transferred from device to host memory in parallel.
is_ocdbt_checkpoint: Indicates whether the checkpoint path uses OCDBT format
or not. Only used for restoration.
use_compression: When True, turn on zstd compression. Default is True.
use_zarr3: If True, use Zarr ver3 otherwise ver2.
ocdbt_target_data_file_size: Specifies the target size (in bytes) of each
OCDBT data file.
ts_context: Tensorstore context to use for reading/writing.
value_typestr: Stores the original value's typestr (from TypeHandler).
enable_pinned_host_transfer: If False, disables transfer to pinned host.
raise_array_data_missing_error: Only used for restoring.
write_shape: Shape of the array shard. Used in the subchunking context.
is_prioritized_key_fn: See ``IsPrioritizedKeyFn`` definition.
keypath: Tuple of keys identifying the parameter's position in the PyTree.
"""
def __init__(
self,
*,
name: str,
parent_dir: epath.Path | path_types.PathAwaitingCreation,
path: epath.Path | path_types.PathAwaitingCreation | None = None,
keypath: tree_types.PyTreePath | None = None,
skip_deserialize: Optional[bool] = None,
byte_limiter: Optional[limits.ByteLimiter] = None,
device_host_byte_limiter: Optional[limits.ByteLimiter] = None,
is_ocdbt_checkpoint: Optional[bool] = None,
use_compression: bool | None = True,
use_zarr3: Optional[bool] = False,
ocdbt_target_data_file_size: Optional[int] = None,
ts_context: Optional[ts.Context] = None,
value_typestr: Optional[str] = None,
enable_pinned_host_transfer: bool = False,
raise_array_data_missing_error: bool = True,
write_shape: arrays_types.Shape | None = None,
is_prioritized_key_fn: Optional[IsPrioritizedKeyFn] = None,
):
self.name = name
self._parent_dir = parent_dir
self._path = path if path is not None else parent_dir / name
self.keypath = keypath
self.skip_deserialize = skip_deserialize
self.byte_limiter = byte_limiter
self.device_host_byte_limiter = device_host_byte_limiter
self.is_ocdbt_checkpoint = is_ocdbt_checkpoint
self.use_compression = use_compression
self.use_zarr3 = use_zarr3
self.ocdbt_target_data_file_size = ocdbt_target_data_file_size
if ts_context is None:
logging.warning(
'ParamInfo(name=%r) created with ts_context=None. Passing unshared or'
' default TensorStore contexts across multiple parameters can'
' duplicate OCDBT driver and B-tree node caches in memory.',
name,
)
self.ts_context = ts_context
self.value_typestr = value_typestr
self.enable_pinned_host_transfer = enable_pinned_host_transfer
self.raise_array_data_missing_error = raise_array_data_missing_error
self.write_shape = write_shape
self.is_prioritized_key_fn = is_prioritized_key_fn
@property
def parent_dir(self) -> epath.Path:
if isinstance(self._parent_dir, path_types.PathAwaitingCreation):
raise ValueError(
'parent_dir is a PathAwaitingCreation and has not been resolved yet. '
'Call `await info.await_path_creation()` first.'
)
return self._parent_dir
@parent_dir.setter
def parent_dir(self, value: epath.Path | path_types.PathAwaitingCreation):
self._parent_dir = value
@property
def path(self) -> epath.Path:
if isinstance(self._path, path_types.PathAwaitingCreation):
raise ValueError(
'path is a PathAwaitingCreation and has not been resolved yet. '
'Call `await info.await_path_creation()` first.'
)
if self._path is None:
raise ValueError('path is None.')
return self._path
@path.setter
def path(self, value: epath.Path | path_types.PathAwaitingCreation | None):
self._path = value # pyrefly: ignore[bad-assignment]
def replace(self, **kwargs) -> ParamInfo:
new = copy.copy(self)
for k, v in kwargs.items():
setattr(new, k, v)
return new
async def await_path_creation(self) -> None:
if isinstance(self._parent_dir, path_types.PathAwaitingCreation):
resolved_path = await self._parent_dir.await_creation()
self._parent_dir = resolved_path
self._path = resolved_path / self.name
[docs]
@dataclasses.dataclass
class SaveArgs:
"""Extra arguments that can be provided for saving.
aggregate:
Deprecated, please use custom TypeHandler
(https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.type_handlers.html#typehandler)
or contact Orbax team to migrate before August 1st, 2024. If true, saves the
given
parameter in an aggregated tree format rather than individually. See
AggregateHandler.
dtype:
If provided, casts the parameter to the given dtype before saving.
Note that the parameter must be compatible with the given type (e.g.
jnp.bfloat16 is not compatible with np.ndarray).
chunk_byte_size:
This is an experimental feature that automatically chooses the largest chunk
shape possible, while keeping the chunk byte size less than or equal to the
specified chunk_byte_size. Both the write_chunk_shape and read_chunk_shape
are automatically set to the chosen shape. This uses a greedy algorithm that
prioritizes splitting the largest dimensions first.
shard_axes: An optional list of axes that should be prioritized when
sharding array for storage. If empty, storage sharding implementation will
prioritize axes which are already sharded.
"""
aggregate: bool = False
dtype: Optional[jnp.dtype] = None
chunk_byte_size: Optional[int] = None
shard_axes: tuple[int, ...] = tuple()
def __post_init__(self):
if self.aggregate:
jax.monitoring.record_event('/jax/orbax/deprecation/aggregate')
logging.log_every_n_seconds(
logging.WARNING,
'The `aggregate` option is deprecated and will be ignored.',
n_seconds=12 * 60 * 60, # once every 12 hours
)
[docs]
@dataclasses.dataclass
class RestoreArgs:
"""Extra arguments that can be provided for restoration.
restore_type:
Specifies the object type of the restored parameter. The type
must have a corresponding TypeHandler for restoration. Ignored if the
parameter is restored from an aggregated checkpoint file.
dtype:
If provided, casts the parameter to the given dtype after restoring.
Note that the parameter must be compatible with the given type (e.g.
jnp.bfloat16 is not compatible with np.ndarray).
"""
restore_type: Optional[Any] = None
dtype: Optional[jnp.dtype] = None
[docs]
class TypeHandler(abc.ABC):
"""Interface for reading and writing a PyTree leaf."""
[docs]
@abc.abstractmethod
def typestr(self) -> str:
"""A string representation of the type.
Cannot conflict with other types.
Returns:
The type as a string.
"""
pass
[docs]
@abc.abstractmethod
async def serialize(
self,
values: Sequence[Any],
infos: Sequence[ParamInfo],
args: Optional[Sequence[SaveArgs]] = None,
) -> Sequence[future.Future]:
"""Writes the parameter to a storage location.
This method is responsible for copying the parameter from a remote device in
a synchronous fashion (if applicable). It should then return a list of
futures which can be later awaited to complete the final commit operation
to a storage location.
Note: Any operations writing to storage location should be done by using
`future.CommitFutureAwaitingContractedSignals` to wait for the directories
to be created.
The function can be used in a multihost setting, but should not implement
extra logic to ensure atomicity.
Args:
values: a sequence of parameters to save.
infos: a sequence of ParamInfo containing relevant information for
serialization of each value.
args: a sequence of additional arguments for serialization, provided by
the user.
Returns:
Sequence of commit futures which can be awaited to complete the save
operation.
"""
pass
[docs]
@abc.abstractmethod
async def deserialize(
self,
infos: Sequence[ParamInfo],
args: Optional[Sequence[RestoreArgs]] = None,
) -> Sequence[Any]:
"""Reads the parameter from a storage location.
Args:
infos: Sequence of ParamInfo for deserialization.
args: Sequence of user-provided restoration information.
Returns:
The deserialized parameters.
"""
pass
[docs]
def finalize(self, directory: epath.Path):
"""Performs any logic to finalize parameter files written by this class.
By default, does nothing.
Args:
directory: A path to the location of the checkpoint. This corresponds to
`param_info.parent_dir`.
"""
pass
[docs]
class TypeHandlerRegistry(Protocol):
"""A registry for TypeHandlers.
This internal base class is used for the global registry which serves as a
default for any type not found in a local registry. It is also accessed
through the module function get/set/has_type_handler.
"""
def add(
self,
ty: Any,
handler: TypeHandler,
func: Optional[Callable[[Any], bool]] = None,
override: bool = False,
ignore_warnings: bool = False,
):
"""Registers a type for serialization/deserialization with a given handler.
Note that it is possible for a type to match multiple different entries in
the registry, each with a different handler. In this case, only the first
match is used.
Args:
ty: A type to register.
handler: a TypeHandler capable of reading and writing parameters of type
`ty`.
func: A function that accepts a type and returns True if the type should
be handled by the provided TypeHandler. If this parameter is not
specified, defaults to `lambda t: issubclass(t, ty)`.
override: if True, will override an existing mapping of type to handler.
ignore_warnings: if True, will ignore warnings when replacing an existing
handler.
Raises:
ValueError if a type is already registered and override is False.
"""
...
def get(self, ty: Any) -> TypeHandler:
"""Returns the handler registered for a given type, if available.
Args:
ty: an object type (or string representation of the type.)
Returns:
The TypeHandler that is registered for the given type.
Raises:
ValueError if the given type has no registered handler.
"""
...
def has(self, ty: Any) -> bool:
"""Checks if a type is registered.
Args:
ty: an object type (or string representation of the type.)
Returns:
A boolean indicating if ty is registered.
"""
...
class IsPrioritizedKeyFn(Protocol):
"""Protocol for checking the transfer prioritization of a key.
The function accepts a PyTree keypath (obtained
using jax.tree.map_with_path) and returns a TransferPriority enum value
indicating how the D2H transfer should be scheduled.
For SYNCHRONOUS keys, the D2H transfer is scheduled before returning
to the caller, so the values will never be corrupted by a concurrent update
or donation.
For the rest of the keys, the D2H transfer is scheduled asynchronously and
the function returns immediately. ASYNCHRONOUS_PRIORITIZED keys will be
scheduled
before ASYNCHRONOUS_DEPRIORITIZED keys. Keys that are not prioritized will not
be scheduled for transfer until all prioritized keys have been fully
written to the checkpoint. This means that these values may be altered
if the values are updated concurrently.
Callers should take care to call
`wait_until_finished` before updating array values (e.g.
`apply_gradients`) if some keys are not prioritized. Note that any
synchronously transferred keys are assumed to be lightweight, and
`save_device_host_concurrent_gb` will be ignored for them.
"""
def __call__(self, keypath: tree_types.PyTreePath) -> TransferPriority:
"""Returns the prioritization of the key."""