# 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.
"""Provides utils for PytreeCheckpointHandler."""
from __future__ import annotations
import asyncio
import copy
from typing import Any, Dict, Optional, Sequence, TypeAlias, Union
from absl import logging
import jax
import numpy as np
from orbax.checkpoint._src import asyncio_utils
from orbax.checkpoint._src.futures import future
from orbax.checkpoint._src.metadata import value as value_metadata
from orbax.checkpoint._src.multihost import multihost
from orbax.checkpoint._src.path import format_utils
from orbax.checkpoint._src.serialization import jax_array_handlers
from orbax.checkpoint._src.serialization import ocdbt_utils
from orbax.checkpoint._src.serialization import tensorstore_utils as ts_utils
from orbax.checkpoint._src.serialization import types
import tensorstore as ts
ParamInfo: TypeAlias = types.ParamInfo
SaveArgs: TypeAlias = types.SaveArgs
RestoreArgs: TypeAlias = types.RestoreArgs
TypeHandler: TypeAlias = types.TypeHandler
TypeHandlerRegistry: TypeAlias = types.TypeHandlerRegistry
PLACEHOLDER = ...
PLACEHOLDER_TYPESTR = 'placeholder'
JAX_ARRAY_TYPE_STR = jax_array_handlers.JAX_ARRAY_TYPE_STR
Scalar = Union[int, float, np.number]
ScalarMetadata = value_metadata.ScalarMetadata
StringMetadata = value_metadata.StringMetadata
is_ocdbt_checkpoint = format_utils.is_ocdbt_checkpoint
check_array_values = jax_array_handlers.check_array_values
represents_jax_array = jax_array_handlers.represents_jax_array
any_jax_array_param_info = jax_array_handlers.any_jax_array_param_info
ArrayRestoreArgs = jax_array_handlers.ArrayRestoreArgs
ArrayHandler = jax_array_handlers.ArrayHandler
SingleReplicaArrayRestoreArgs = jax_array_handlers.SingleReplicaArrayRestoreArgs
SingleReplicaArrayHandler = jax_array_handlers.SingleReplicaArrayHandler
InvalidShardingError = jax_array_handlers.InvalidShardingError
[docs]
class NumpyHandler(types.TypeHandler):
"""Provides an implementation of TypeHandler for replicated numpy arrays."""
[docs]
def __init__(
self,
metadata_key: Optional[str] = None,
ocdbt_process_id: str | None = None,
deepcopy_host_arrays: bool = True,
):
"""Constructor.
Args:
metadata_key: name to give to Tensorstore metadata files.
ocdbt_process_id: name of the process id to be used by single controller
systems to write in OCDBT format. The checkpoints are written in a
subdir with this name to avoid collisions with the subdir names used by
other host processes managed by this controller.
deepcopy_host_arrays: Whether to deepcopy host arrays before
serialization. Defaults to True as a safety measure to guard against
concurrent modification, but can be disabled to avoid excess memory
consumption.
"""
self._metadata_key = metadata_key
self._override_ocdbt_process_id = ocdbt_process_id
self._deepcopy_host_arrays = deepcopy_host_arrays
[docs]
def typestr(self) -> str:
return 'np.ndarray'
async def _open_and_write(
self, value: np.ndarray, tspec: Dict[str, Any], ts_context: ts.Context
):
"""Opens and writes using Tensorstore."""
t = await ts.open(
ts.Spec(tspec), create=True, open=True, context=ts_context
)
await t.write(value, can_reference_source_data_indefinitely=True)
async def _background_serialize(
self,
values: Sequence[np.ndarray],
infos: Sequence[types.ParamInfo],
args: Optional[Sequence[types.SaveArgs]] = None,
):
"""Serializes numpy arrays in a background thread."""
write_coros = []
for value, info, arg in zip(values, infos, args): # pyrefly: ignore[bad-argument-type]
await info.await_path_creation()
array_write_spec = ts_utils.build_array_write_spec(
info=info,
arg=arg,
global_shape=value.shape,
local_shape=value.shape,
dtype=value.dtype,
use_ocdbt=info.is_ocdbt_checkpoint, # pyrefly: ignore[bad-argument-type]
process_index=ocdbt_utils.get_process_index_for_subdir(
use_ocdbt=info.is_ocdbt_checkpoint, # pyrefly: ignore[bad-argument-type]
override_ocdbt_process_id=self._override_ocdbt_process_id,
),
metadata_key=self._metadata_key,
)
tspec = array_write_spec.json
if logging.vlog_is_on(1):
logging.vlog(1, 'tspec = %s', tspec)
logging.vlog(1, 'infos = %s', info)
logging.vlog(1, 'args = %s', arg)
if multihost.process_index() == 0:
ts_context = info.ts_context
write_coros.append(self._open_and_write(value, tspec, ts_context)) # pyrefly: ignore[bad-argument-type]
gather_future = asyncio.gather(*write_coros)
await asyncio_utils.cancellable(
gather_future,
message='[process=%s] Numpy handler gather was cancelled.',
process_index=multihost.process_index(),
)
[docs]
async def serialize(
self,
values: Sequence[np.ndarray],
infos: Sequence[types.ParamInfo],
args: Optional[Sequence[types.SaveArgs]] = None,
) -> Sequence[future.Future]:
"""Uses Tensorstore to serialize a numpy array."""
args = args or [types.SaveArgs()] * len(values)
types.check_input_arguments(values, infos, args)
check_array_values(values, infos)
if logging.vlog_is_on(1):
ts_utils.print_ts_debug_data(self._metadata_key, infos)
if self._deepcopy_host_arrays:
values = [copy.deepcopy(v) for v in values]
return [
future.CommitFutureAwaitingContractedSignals(
self._background_serialize(values, infos, args),
name='np_type_handler',
)
]
[docs]
async def deserialize(
self,
infos: Sequence[types.ParamInfo],
args: Optional[Sequence[RestoreArgs]] = None,
) -> Sequence[np.ndarray]:
"""Deserializes the array using Tensorstore."""
args = args or [RestoreArgs()] * len(infos)
types.check_input_arguments(infos, args)
open_futures = []
for info, arg in zip(infos, args):
await info.await_path_creation()
if not info.is_ocdbt_checkpoint:
await ts_utils.assert_parameter_files_exist(
info.parent_dir / info.name, self._metadata_key, info.use_zarr3 # pyrefly: ignore[bad-argument-type]
)
# Use OCDBT flag from the existing checkpoint.
use_ocdbt = info.is_ocdbt_checkpoint
array_read_spec = ts_utils.build_array_read_spec(
info,
use_ocdbt=use_ocdbt, # pyrefly: ignore[bad-argument-type]
metadata_key=self._metadata_key,
raise_array_data_missing_error=info.raise_array_data_missing_error,
target_dtype=arg.dtype,
)
tspec = array_read_spec.json
if logging.vlog_is_on(1):
logging.vlog(1, 'tspec = %s', tspec)
logging.vlog(1, 'infos = %s', infos)
logging.vlog(1, 'args = %s', args)
open_futures += [
ts.open(ts.Spec(tspec), open=True, context=info.ts_context)
]
tensorstores = await asyncio.gather(*open_futures)
read_ops = [t.read() for t in tensorstores]
ret = await asyncio.gather(*read_ops)
if logging.vlog_is_on(1):
for a in ret:
logging.vlog(
1, 'restored ndarray.shape = %s, array.dtype = %s', a.shape, a.dtype
)
ts_utils.print_ts_debug_data(self._metadata_key, infos)
return ret
[docs]
class ScalarHandler(NumpyHandler):
"""A wrapper around NumpyHandler to deal with scalar types (int, float, etc.)."""
[docs]
def __init__(
self,
metadata_key: Optional[str] = None,
ocdbt_process_id: str | None = None,
):
super().__init__(
metadata_key=metadata_key,
ocdbt_process_id=ocdbt_process_id,
deepcopy_host_arrays=False,
)
[docs]
def typestr(self) -> str:
return 'scalar'
[docs]
async def serialize( # pyrefly: ignore[bad-override]
self,
values: Sequence[Scalar],
infos: Sequence[types.ParamInfo],
args: Optional[Sequence[types.SaveArgs]] = None,
) -> Sequence[future.Future]:
"""See superclass documentation."""
values = [np.asarray(v) for v in values] # pyrefly: ignore[bad-assignment]
return await super().serialize(values, infos, args) # pyrefly: ignore[bad-argument-type]
[docs]
async def deserialize( # pyrefly: ignore[bad-override]
self,
infos: Sequence[types.ParamInfo],
args: Optional[Sequence[RestoreArgs]] = None,
) -> Sequence[Scalar]:
"""See superclass documentation."""
results = await super().deserialize(infos, args)
for r in results:
if r.ndim != 0:
raise ValueError('Restored result is not a scalar.')
results = [r.item() for r in results]
if args:
# Cast to the intended `restore_type` if it is provided.
return [
a.restore_type(r) if a.restore_type else r
for a, r in zip(args, results)
]
return results
[docs]
class StringHandler(types.TypeHandler):
"""TypeHandler for strings."""
[docs]
def __init__(
self,
filename: Optional[str] = None,
):
self._filename = filename or '_strings.json'
self._ts_context = ts_utils.get_ts_context(use_ocdbt=False)
def _get_json_tspec(
self,
info: types.ParamInfo,
) -> Dict[str, Any]:
"""Gets Tensorstore spec in JSON format."""
if info.parent_dir is None:
raise ValueError('Must provide info.parent_dir.')
directory = (info.parent_dir / self._filename).as_posix()
kvstore_tspec = ts_utils.build_kvstore_tspec(directory, use_ocdbt=False)
tspec = {
'driver': 'json',
'kvstore': kvstore_tspec,
'json_pointer': '/' + info.name,
}
return tspec
[docs]
def typestr(self) -> str:
return 'string'
async def _convert_to_string(self, tensorstore):
result = await tensorstore.read()
return str(result)
async def _background_serialize(
self,
values: Sequence[str],
infos: Sequence[types.ParamInfo],
):
"""Writes strings using Tensorstore in the background thread."""
types.check_input_arguments(values, infos)
write_coros = []
txn = ts.Transaction()
for (
info,
value,
) in zip(infos, values):
await info.await_path_creation()
tspec = self._get_json_tspec(info)
if multihost.process_index() == 0:
t = await ts.open(
tspec,
open=True,
context=self._ts_context,
)
write_coros.append(t.with_transaction(txn).write(value))
gather_future = asyncio.gather(*write_coros)
await asyncio_utils.cancellable(
gather_future,
message='[process=%s] String handler gather was cancelled.',
process_index=multihost.process_index(),
)
commit_future = txn.commit_async()
await asyncio_utils.cancellable(
commit_future,
message='[process=%s] String handler commit was cancelled.',
process_index=multihost.process_index(),
)
[docs]
async def serialize(
self,
values: Sequence[str],
infos: Sequence[types.ParamInfo],
args: Optional[Sequence[types.SaveArgs]] = None,
) -> Sequence[future.Future]:
"""See superclass documentation."""
del args
# Copy is not needed since strings are passed by value.
return [
future.CommitFutureAwaitingContractedSignals(
self._background_serialize(values, infos),
name='string_type_handler',
)
]
[docs]
async def deserialize(
self,
infos: Sequence[types.ParamInfo],
args: Optional[Sequence[RestoreArgs]] = None,
) -> Sequence[Optional[str]]:
"""See superclass documentation."""
del args
types.check_input_arguments(infos)
open_futures = []
for info in infos:
await info.await_path_creation()
tspec = self._get_json_tspec(info)
open_future = ts.open(
tspec, open=True, read=True, context=self._ts_context
)
open_futures += [open_future]
tensorstores = await asyncio.gather(*open_futures)
read_ops = [self._convert_to_string(t) for t in tensorstores]
return await asyncio.gather(*read_ops)
def is_placeholder(value: Any) -> bool:
return value is PLACEHOLDER
class PlaceholderHandler(types.TypeHandler):
"""TypeHandler for placeholders."""
def typestr(self) -> str:
return PLACEHOLDER_TYPESTR
async def metadata(self, infos: Sequence[types.ParamInfo]) -> Sequence[Any]:
raise NotImplementedError('Placeholders do not have metadata.')
async def serialize(
self,
values: Sequence[str],
infos: Sequence[types.ParamInfo],
args: Sequence[types.SaveArgs] | None = None,
) -> Sequence[future.Future]:
"""See superclass documentation."""
raise NotImplementedError('Placeholders cannot be serialized.')
async def deserialize(
self,
infos: Sequence[ParamInfo],
args: Sequence[RestoreArgs] | None = None,
) -> Sequence[Any]:
"""See superclass documentation."""
del args
return [PLACEHOLDER] * len(infos)
# TODO(b/253238305) Deprecate when all checkpoints have saved types.
def default_restore_type(args: RestoreArgs) -> Any:
if isinstance(args, ArrayRestoreArgs):
return jax.Array
elif isinstance(args, RestoreArgs):
return np.ndarray
else:
raise ValueError(f'Unsupported restore_args type: {type(args)}')