# Copyright 2024 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.

"""BasePyTreeCheckpointHandler class.

Implementation of `CheckpointHandler` interface dealing with JAX PyTrees. Much
of the underlying reading/writing logic for individual leaf types can be
customized, and is delegated to the `TypeHandler` class.
"""

import asyncio
import collections
import dataclasses
import json
import os
import time
from typing import Any, Callable, List, Optional, Tuple, Union

from absl import logging
from etils import epath
import jax
import numpy as np
from orbax.checkpoint import aggregate_handlers
from orbax.checkpoint import async_checkpoint_handler
from orbax.checkpoint import checkpoint_args
from orbax.checkpoint import future
from orbax.checkpoint import transform_utils
from orbax.checkpoint import tree as tree_utils
from orbax.checkpoint import type_handlers
from orbax.checkpoint import utils
from orbax.checkpoint.metadata import tree as tree_metadata
from orbax.checkpoint.metadata import value as value_metadata
import tensorstore as ts

PyTree = Any
TupleKey = Tuple[str, ...]
RestoreArgs = type_handlers.RestoreArgs
ArrayRestoreArgs = type_handlers.ArrayRestoreArgs
SaveArgs = type_handlers.SaveArgs
ParamInfo = type_handlers.ParamInfo
TypeHandler = type_handlers.TypeHandler
TypeHandlerRegistry = type_handlers.TypeHandlerRegistry
AggregateHandler = aggregate_handlers.AggregateHandler
MsgpackHandler = aggregate_handlers.MsgpackHandler
LegacyTransformFn = Callable[[PyTree, PyTree, PyTree], Tuple[PyTree, PyTree]]
Transform = transform_utils.Transform
RestoreTransform = transform_utils.RestoreTransform
# TODO(b/298487158) Clean up protected access.
LimitInFlightBytes = type_handlers.LimitInFlightBytes
CheckpointArgs = checkpoint_args.CheckpointArgs
register_with_handler = checkpoint_args.register_with_handler


METADATA_FILE = '_METADATA'
_CHECKPOINT_FILE = 'checkpoint'


def get_byte_limiter(concurrent_gb: int):
  async def _create_byte_limiter():
    # Wrap creation in async function to avoid issues on python<=3.9.
    concurrent_bytes = concurrent_gb * 10**9
    # Construction must take place here so that it is within the same async
    # method, to prevent errors resulting from different event loops, and
    # cannot be created below this level because there must be a single object
    # for the entire restore call.
    return LimitInFlightBytes(concurrent_bytes)  # pylint: disable=protected-access

  return asyncio.run(_create_byte_limiter())


async def _create_param_save_dir(param_info: ParamInfo, args: SaveArgs):
  # Directory will be unused.
  path = param_info.path
  if path is None or args.aggregate:
    return
  # TODO(b/273803615): Note that keys with slashes ('/', generated by Haiku,
  # for example) will result in the creation of nested sub-directories, rather
  # than flat parameter directories like for a standard nested PyTree. Ensure
  # `use_ocdbt=True` in this scenario.
  await utils.async_makedirs(path, parents=True)


def _maybe_set_default_restore_args(args):
  if isinstance(args, RestoreArgs):
    return args
  return RestoreArgs(restore_type=None)


def _try_array_cast(arr, dtype):
  if dtype is not None:
    if utils.is_scalar(arr):
      arr = np.asarray(arr).astype(dtype).item()
    else:
      if hasattr(arr, 'astype'):
        arr = arr.astype(dtype)
  return arr


def _maybe_shard_array(value, args):
  if hasattr(value, 'reshape') and isinstance(args, ArrayRestoreArgs):
    value = value.reshape(args.global_shape)
    sharding = args.sharding or jax.sharding.NamedSharding(
        args.mesh, args.mesh_axes
    )
    value = jax.make_array_from_callback(
        value.shape, sharding, lambda idx: value[idx]
    )
  return value


def get_param_names(item: PyTree) -> PyTree:
  """Gets parameter names for PyTree elements."""

  def _param_name_from_keypath(keypath: Tuple[Any, ...]) -> str:
    return '.'.join([str(tree_utils.get_key_name(k)) for k in keypath])

  return jax.tree_util.tree_map_with_path(
      lambda kp, _: _param_name_from_keypath(kp),
      item,
      is_leaf=tree_utils.is_empty_or_leaf,
  )


def _keystr(key: Tuple[Any, ...]) -> str:
  return '/'.join(key)


@dataclasses.dataclass
class _InternalValueMetadata:
  restore_type: Optional[str]
  skip_deserialize: bool = False
  aggregate_value: Optional[Any] = None


def _get_restore_parameters(
    directory: epath.Path,
    structure: PyTree,
    restore_args: Optional[PyTree],
    byte_limiter: Optional[LimitInFlightBytes] = None,
    use_zarr3: bool = False,
) -> Tuple[PyTree, PyTree]:
  """Construct parameters needed for restoration.

  param_infos are
  constructed from the structure of the original checkpoint, and restore_args
  are serialized to a tree structure compatible with param_infos and structure.

  Args:
    directory: Checkpoint directory.
    structure: The structure of the original checkpoint.
    restore_args: User-provided restoration arguments. If None, they were not
      provided. Otherwise, the tree has the same structure as the desired output
      tree.
    byte_limiter: A _LimitInFlightBytes object.
    use_zarr3: If True, use Zarr ver3 otherwise Zarr ver2

  Returns:
    Tuple of param_infos, and restore_args.
  """
  flat_structure = tree_utils.to_flat_dict(structure, keep_empty_nodes=True)
  if restore_args is None:
    restore_args = jax.tree.map(lambda x: RestoreArgs(), structure)
  flat_param_infos = {}
  is_ocdbt_checkpoint = type_handlers.is_ocdbt_checkpoint(directory)

  def _get_param_info(
      nested_name: Tuple[str, ...],
      meta: _InternalValueMetadata,
  ) -> Union[ParamInfo, Any]:
    if utils.is_supported_empty_aggregation_type(meta):
      # Empty node, ParamInfo should not be returned.
      return meta
    name = '.'.join(nested_name)
    return ParamInfo(
        name=name,
        path=directory / name,
        parent_dir=directory,
        skip_deserialize=meta.skip_deserialize,
        is_ocdbt_checkpoint=is_ocdbt_checkpoint,
        byte_limiter=byte_limiter,
        use_zarr3=use_zarr3,
    )

  for key, meta in flat_structure.items():
    flat_param_infos[key] = _get_param_info(key, meta)
  restore_args = tree_utils.serialize_tree(restore_args, keep_empty_nodes=True)
  return (
      tree_utils.from_flat_dict(flat_param_infos, target=structure),
      restore_args,
  )


def _get_tree_for_aggregation(param_infos, save_args, item):
  """Get tree for aggregated checkpoint."""

  # TODO(b/283164080): These type checks result in logic from the lower layer
  # (TypeHandler/AggregateHandler) leaking into the upper layer
  # (CheckpointHandler). Ideally, AggregateHandler could define its own
  # supported values and error conditions.
  def _get_leaf_for_aggregation(param_info, arg, value):
    if arg.aggregate:  # Param was aggregated, return value after cast.
      if isinstance(value, jax.Array) and not value.is_fully_replicated:
        raise ValueError(
            'jax.Array must be fully replicated to be saved in aggregate file.'
        )
      if not utils.is_supported_aggregation_type(value):
        # Not an error because users' training states often have a bunch of
        # random unserializable objects in them (empty states, optimizer
        # objects, etc.).
        value = None
      return _try_array_cast(value, arg.dtype)
    else:  # Placeholder string for non-aggregated value.
      return utils.leaf_placeholder(param_info.name)

  return jax.tree.map(
      _get_leaf_for_aggregation, param_infos, save_args, item
  )


@dataclasses.dataclass
class _BatchRequest:
  """Represents a a request for batched serialization or deserialization.

  Attributes:
    handler: Used to serialize or deserialize the parameters.
    keys: Used to identify the original tree keys so that the PyTree can be
      reconstructed.
    values: Values to serialize.
    infos: ParamInfos.
    args: List of SaveArgs or RestoreArgs.
  """

  handler: TypeHandler
  keys: List[str]
  values: List[Any]
  infos: List[ParamInfo]
  args: List[Union[SaveArgs, RestoreArgs]]

  def __post_init__(self):
    length = len(self.values)
    if not all((
        length == len(self.infos),
        length == len(self.args),
        length == len(self.keys),
    )):
      raise AssertionError('Found `_BatchRequest` with mismatched parameters.')


def _batched_serialization_requests(
    tree: PyTree,
    param_infos: PyTree,
    args: PyTree,
    registry: TypeHandlerRegistry,
) -> List[_BatchRequest]:
  """Gets a list of batched serialization or deserialization requests."""
  grouped = {}

  def _group_value(
      keypath: Tuple[Any, ...],
      info: ParamInfo,
      value: Union[Any, _InternalValueMetadata],
      arg: RestoreArgs,
  ):
    nonlocal grouped
    tuple_key = tree_utils.tuple_path_from_keypath(keypath)
    # Exclude from serialize/deserialize with TypeHandler if aggregated.
    if info.skip_deserialize:
      return

    if isinstance(arg, RestoreArgs):
      assert isinstance(value, _InternalValueMetadata)
      restore_type = value.restore_type
      if arg.restore_type is not None:
        # Give user the chance to override restore_type if they want.
        restore_type = arg.restore_type
      type_for_registry_lookup = restore_type
    else:
      type_for_registry_lookup = type(value)
    try:
      handler = registry.get(type_for_registry_lookup)
    except ValueError as e:
      raise ValueError(
          f'TypeHandler lookup failed for: type={type_for_registry_lookup},'
          f' keypath={keypath}, ParamInfo={info}, RestoreArgs={arg},'
          f' value={value}'
      ) from e

    if handler not in grouped:
      grouped[handler] = _BatchRequest(handler, [], [], [], [])
    request = grouped[handler]
    grouped[handler] = dataclasses.replace(
        request,
        keys=request.keys + [tuple_key],
        values=request.values + [value],
        infos=request.infos + [info],
        args=request.args + [arg],
    )

  jax.tree_util.tree_map_with_path(
      _group_value,
      param_infos,
      tree,
      args,
  )
  return list(grouped.values())


class BasePyTreeCheckpointHandler(
    async_checkpoint_handler.AsyncCheckpointHandler
):
  """A CheckpointHandler implementation for any PyTree structure.

  See JAX documentation for more information on what consistutes a "PyTree".
  This handler is capable of saving and restoring any leaf object for which a
  `TypeHandler` (see documentation) is registered. By default, `TypeHandler`s
  for standard types like `np.ndarray`, `jax.Array`, Python scalars, and others
  are registered.

  As with all `CheckpointHandler` subclasses, `BasePyTreeCheckpointHandler`
  should only be used in conjunction with a `Checkpointer` (or subclass).
  By itself, the `CheckpointHandler` is non-atomic.

  Example::

    ckptr = Checkpointer(BasePyTreeCheckpointHandler())
  """

  def __init__(
      self,
      aggregate_filename: Optional[str] = None,
      concurrent_gb: int = 96,
      use_ocdbt: bool = True,
      use_zarr3: bool = False,
      primary_host: Optional[int] = 0,
      type_handler_registry: TypeHandlerRegistry = type_handlers.GLOBAL_TYPE_HANDLER_REGISTRY,
  ):
    """Creates BasePyTreeCheckpointHandler.

    Args:
      aggregate_filename: name that the aggregated checkpoint should be saved
        as.
      concurrent_gb: max concurrent GB that are allowed to be read. Can help to
        reduce the possibility of OOM's when large checkpoints are restored.
      use_ocdbt: enables Tensorstore OCDBT driver. This option allows using a
        different checkpoint format which is faster to read and write, as well
        as more space efficient.
      use_zarr3: If True, use Zarr ver3 otherwise Zarr ver2
      primary_host: the host id of the primary host.  Default to 0.  If it's set
        to None, then all hosts will be considered as primary.  It's useful in
        the case that all hosts are only working with local storage.
      type_handler_registry: a type_handlers.TypeHandlerRegistry. If not
        specified, the global type handler registry will be used.
    """
    self._aggregate_handler = MsgpackHandler(primary_host=primary_host)
    if aggregate_filename is None:
      aggregate_filename = _CHECKPOINT_FILE
    self._aggregate_filename = aggregate_filename
    self._concurrent_gb = concurrent_gb
    self._use_ocdbt = use_ocdbt
    self._use_zarr3 = use_zarr3
    self._primary_host = primary_host
    self._type_handler_registry = type_handler_registry


    if self._use_ocdbt:
      jax.monitoring.record_event(
          '/jax/orbax/pytree_checkpoint_handler/init/ocdbt'
      )

  def get_param_names(self, item: PyTree) -> PyTree:
    """Gets parameter names for PyTree elements."""
    return get_param_names(item)

  def _skip_deserialize(self, value: Any, args: SaveArgs) -> bool:
    """Returns True if _METADATA write is enabled and value is []/{}/None."""
    if utils.is_supported_empty_aggregation_type(value):
      # Skip deser if value is empty ([], {}, None) and _METADATA is enabled.
      # We don't want to write TypeHandlers for empty values, so will simply
      # identify them in metadata and skip deser.
      return True
    else:
      # Follow aggregate based flow: requires TypeHandler registry if
      # aggregate=False. Empty values will raise registry error if
      # aggregate=False. Users will be prompted to enable _METADATA to avoid
      # these errors for empty values.
      return args.aggregate

  def _get_param_infos(
      self,
      item: PyTree,
      directory: epath.Path,
      save_args: PyTree,
      ocdbt_target_data_file_size: Optional[int] = None,
  ) -> Tuple[PyTree, bool]:
    """Returns parameter information for elements in `item`.

    At minimum, this method should extract the names of each parameter for
    saving/restoring.

    Args:
      item: a PyTree to extract information from.
      directory: a directory where checkpoint files are located.
      save_args: PyTree matching item containing SaveArgs.
      ocdbt_target_data_file_size: Specifies the target size (in bytes) of each
        OCDBT data file.

    Returns:
      A PyTree matching `item` of ParamInfo, and a bool indicating whether all
      parameters were aggregated. The bool can enable us to skip some steps
      later, potentially saving time.
    """
    if not item:
      raise ValueError('Found empty item')
    names = self.get_param_names(item)
    all_params_aggregated = True

    def _param_info(value, name, args):
      nonlocal all_params_aggregated
      all_params_aggregated &= args.aggregate
      skip_deserialize = self._skip_deserialize(value, args)
      return ParamInfo(
          name=name,
          path=(directory / name),
          parent_dir=directory,
          skip_deserialize=skip_deserialize,
          is_ocdbt_checkpoint=self._use_ocdbt,
          use_zarr3=self._use_zarr3,
          ocdbt_target_data_file_size=ocdbt_target_data_file_size,
      )

    return (
        jax.tree.map(
            _param_info,
            item,
            names,
            save_args,
            is_leaf=tree_utils.is_empty_or_leaf,
        ),
        all_params_aggregated,
    )

  async def async_save(
      self,
      directory: epath.Path,
      args: Optional['BasePyTreeSaveArgs'] = None,
  ) -> Optional[List[future.Future]]:
    """Saves a PyTree to a given directory.

    This operation is compatible with a multi-host, multi-device setting. Tree
    leaf values must be supported by the type_handler_registry given in the
    constructor. Standard supported types include Python scalars, `np.ndarray`,
    `jax.Array`, and strings.

    After saving, all files will be located in "directory/". The exact files
    that are saved depend on the specific combination of options, including
    `use_ocdbt`. A JSON metadata file will be present to store the
    tree structure.
    In addition, a msgpack file may be present, allowing users to store
    aggregated values (see below).

    Example usage::

      ckptr = Checkpointer(BasePyTreeCheckpointHandler())
      item = {
          'layer0': {
              'w': np.ndarray(...),
              'b': np.ndarray(...),
          },
          'layer1': {
              'w': np.ndarray(...),
              'b': np.ndarray(...),
          },
      }
      # Note: save_args may be None if no customization is desired for saved
      # parameters.
      # In this case, we "aggregate" small parameters into a single file to
      # allow for greater file read/write efficiency (and potentially less)
      # wasted space). With OCDBT format active, this parameter is obsolete.
      save_args =
        jax.tree.map(
            lambda x: SaveArgs(aggregate=x.size < some_size), item)
      # Eventually calls through to `async_save`.
      ckptr.save(path, item, save_args)

    Args:
      directory: save location directory.
      args: `BasePyTreeSaveArgs` (see below).

    Returns:
      A Future that will commit the data to `directory` when awaited. Copying
      the data from its source will be awaited in this function.
    """
    args = args or BasePyTreeSaveArgs()
    item = args.item
    save_args = args.save_args
    ocdbt_target_data_file_size = args.ocdbt_target_data_file_size

    if ocdbt_target_data_file_size is not None and not self._use_zarr3:
      raise ValueError('`ocdbt_target_data_file_size` only works with Zarr3')

    # Because of empty states, the user-provided args may not contain
    # all necessary arguments. These should be filled in with default args.
    def _maybe_set_default_save_args(value, args_):
      # If already set, return.
      if isinstance(args_, SaveArgs):
        return args_
      if utils.is_supported_empty_aggregation_type(value):
        # If _METADATA is enabled and value is empty ([], {}, None) then stop
        # aggregating for a smooth SaveArgs.aggregate deprecation. Otherwise, we
        # will need to write TypeHandlers for empty values.
        return SaveArgs(aggregate=False)
      # Empty values will still raise TypeHandler registry error if _METADATA is
      # disabled. We will prompt users to enable _METADATA to avoid this error.
      aggregate = not self._type_handler_registry.has(type(value))
      return SaveArgs(aggregate=aggregate)

    save_args = jax.tree.map(
        _maybe_set_default_save_args,  # pylint: disable=protected-access
        item,
        item if save_args is None else save_args,
        is_leaf=tree_utils.is_empty_or_leaf,
    )
    param_infos, all_params_aggregated = self._get_param_infos(
        item, directory, save_args, ocdbt_target_data_file_size
    )
    assert all(
        leaf.parent_dir == directory
        for leaf in jax.tree.leaves(param_infos)
    )
    if not self._use_ocdbt and not all_params_aggregated:
      if utils.is_primary_host(self._primary_host):
        # Create directories in parallel.
        await asyncio.gather(
            *jax.tree.flatten(
                jax.tree.map(
                    _create_param_save_dir,
                    param_infos,
                    save_args,
                )
            )[0]
        )
      utils.sync_global_processes(
          'BasePyTreeCheckpointHandler:create_param_save_dirs'
      )

    if all_params_aggregated:
      commit_futures = []
    else:
      serialize_ops = []
      batch_requests = _batched_serialization_requests(
          item,
          param_infos,
          save_args,
          self._type_handler_registry,
      )
      for request in batch_requests:
        serialize_ops += [
            request.handler.serialize(
                request.values, request.infos, request.args
            )
        ]
      # Await copy futures. Returns list of lists.
      commit_futures = await asyncio.gather(*serialize_ops)
      commit_futures, _ = jax.tree.flatten(commit_futures)

    if logging.level_debug():
      logging.debug('param_info: %s', param_infos)
      logging.debug('save_args: %s', save_args)

    metadata_future = None
    if utils.is_primary_host(self._primary_host):
      metadata_write_start_time = time.time()
      metadata_future = await self._write_metadata_file(
          directory, item, save_args, self._use_zarr3
      )
      jax.monitoring.record_event_duration_secs(
          '/jax/checkpoint/write/async/metadata_write_duration_secs',
          time.time() - metadata_write_start_time,
      )

    aggregate_file_write_start_time = time.time()
    aggregate_commit_future = await self._write_aggregate_file(
        directory, item, param_infos, save_args
    )
    jax.monitoring.record_event_duration_secs(
        '/jax/checkpoint/write/async/aggregate_write_duration_secs',
        time.time() - aggregate_file_write_start_time,
    )
    return (
        commit_futures + [aggregate_commit_future] + [metadata_future]
        if metadata_future is not None
        else commit_futures + [aggregate_commit_future]
    )

  def save(self, directory: epath.Path, *args, **kwargs):
    """Saves the provided item.

    Blocks until both copy and commit complete.

    See async_save.

    Args:
      directory: the directory to save to.
      *args: additional arguments for save.
      **kwargs: additional arguments for save.
    """

    async def async_save(*args, **kwargs):
      commit_futures = await self.async_save(*args, **kwargs)  # pytype: disable=bad-return-type
      # Futures are already running, so sequential waiting is equivalent to
      # concurrent waiting.
      if commit_futures:  # May be None.
        for f in commit_futures:
          f.result()  # Block on result.

    asyncio.run(async_save(directory, *args, **kwargs))

  async def _maybe_deserialize(
      self, structure: PyTree, param_infos: PyTree, restore_args: PyTree
  ) -> PyTree:
    """Deserializes values or gets them from the aggregate file."""

    # Handle parameters from aggregate file.
    def _process_aggregated_value(info, meta, args):
      value = meta.aggregate_value
      if info.skip_deserialize:
        value = _try_array_cast(value, args.dtype)
        value = _maybe_shard_array(value, args)
      return value

    flat_aggregate = tree_utils.to_flat_dict(
        jax.tree.map(
            _process_aggregated_value, param_infos, structure, restore_args
        ),
    )

    batch_requests = _batched_serialization_requests(
        structure,
        param_infos,
        restore_args,
        self._type_handler_registry,
    )
    deserialized_batches = []
    deserialized_batches_ops = []
    for request in batch_requests:
      deserialized_batches_ops.append(
          request.handler.deserialize(request.infos, request.args)
      )
    deserialized_batches += await asyncio.gather(*deserialized_batches_ops)

    flat_restored = {}
    for request, deserialized in zip(batch_requests, deserialized_batches):
      for key, value in zip(request.keys, deserialized):
        flat_restored[key] = value
    # Add in any values which were not deserialized, coming from aggregate file.
    for key in flat_aggregate.keys():
      if key not in flat_restored:
        flat_restored[key] = flat_aggregate[key]
    return tree_utils.from_flat_dict(flat_restored, target=structure)

  def restore(
      self,
      directory: epath.Path,
      args: Optional['BasePyTreeRestoreArgs'] = None,
  ) -> PyTree:
    """Restores a PyTree from the checkpoint directory at the given path.

    In the most basic case, only `directory` is required. The tree will be
    restored exactly as saved, and all leaves will be restored as the correct
    types (assuming the tree metadata is present).

    However, `restore_args` is often required as well. This PyTree gives a
    `RestoreArgs` object (or subclass) for every leaf in the tree. Many types,
    such as string or `np.ndarray` do not require any special options for
    restoration. When restoring an individual leaf as `jax.Array`, however,
    some properties may be required.

    One example is `sharding`, which defines how a `jax.Array` in the restored
    tree should be partitioned. `mesh` and `mesh_axes` can also be used to
    specify `sharding`, but `sharding` is the preferred way of specifying this
    partition since `mesh` and `mesh_axes` only constructs
    `jax.sharding.NamedSharding`. For more information, see `ArrayTypeHandler`
    documentation and JAX sharding documentation.

    Example::

      ckptr = Checkpointer(BasePyTreeCheckpointHandler())
      restore_args = {
          'layer0': {
              'w': RestoreArgs(),
              'b': RestoreArgs(),
          },
          'layer1': {
              'w': ArrayRestoreArgs(
                  # Restores as jax.Array, regardless of how it was saved.
                  restore_type=jax.Array,
                  sharding=jax.sharding.Sharding(...),
                  # Warning: may truncate or pad!
                  global_shape=(x, y),
                ),
              'b': ArrayRestoreArgs(
                  restore_type=jax.Array,
                  sharding=jax.sharding.Sharding(...),
                  global_shape=(x, y),
                ),
          },
      }
      ckptr.restore(path, restore_args=restore_args)

    Providing `item` is typically only necessary when restoring a custom PyTree
    class (or when using transformations). In this case, the restored object
    will take on the same structure as `item`.

    Example::

      @flax.struct.dataclass
      class TrainState:
        layer0: dict[str, jax.Array]
        layer1: dict[str, jax.Array]

      ckptr = Checkpointer(BasePyTreeCheckpointHandler())
      train_state = TrainState(
          layer0={
              'w': jax.Array(...),  # zeros
              'b': jax.Array(...),  # zeros
          },
          layer1={
              'w': jax.Array(...),  # zeros
              'b': jax.Array(...),  # zeros
          },
      )
      restore_args = jax.tree.map(_make_restore_args, train_state)
      ckptr.restore(path, item=train_state, restore_args=restore_args)
      # restored tree is of type `TrainState`.

    Args:
      directory: saved checkpoint location directory.
      args: `BasePyTreeRestoreArgs` (see below).

    Returns:
      A PyTree matching the structure of `item`.

    Raises:
      FileNotFoundError: `directory` does not exist or is missing required files
      ValueError: `transforms` is provided without `item`.
      ValueError: `transforms` contains elements with `multi_value_fn`.
    """
    args = args or BasePyTreeRestoreArgs()
    item = args.item
    restore_args = args.restore_args
    logging.debug('directory=%s, restore_args=%s', directory, restore_args)
    if not directory.exists():
      raise FileNotFoundError(
          f'Requested directory for restore does not exist at {directory}'
      )
    byte_limiter = get_byte_limiter(self._concurrent_gb)
    structure, use_zarr3_metadata = self._get_internal_metadata(directory)
    # `checkpoint_restore_args` has a structure relative to the checkpoint,
    # while `restore_args` remains structured relative to the output.
    param_infos, checkpoint_restore_args = _get_restore_parameters(
        directory,
        structure,
        restore_args,
        byte_limiter=byte_limiter,
        use_zarr3=use_zarr3_metadata
        if use_zarr3_metadata is not None
        else self._use_zarr3,
    )

    def _maybe_set_default_restore_types(
        meta: _InternalValueMetadata, arg: RestoreArgs
    ):
      if not meta.skip_deserialize and meta.restore_type is None:
        return dataclasses.replace(
            meta, restore_type=type_handlers.default_restore_type(arg)
        )
      return meta

    # If metadata file was missing in the checkpoint, we need to decide
    # restore_type based on RestoreArgs.
    structure = jax.tree.map(
        _maybe_set_default_restore_types, structure, checkpoint_restore_args
    )

    restored_item = asyncio.run(
        self._maybe_deserialize(structure, param_infos, checkpoint_restore_args)
    )

    if logging.level_debug():
      logging.debug('param_infos: %s', param_infos)
      logging.debug('checkpoint_restore_args: %s', checkpoint_restore_args)
      logging.debug(
          'restored_item: %s', jax.tree.structure(restored_item)
      )
      logging.debug(
          'ts_metrics: %s',
          json.dumps(ts.experimental_collect_matching_metrics('/tensorstore/')),
      )

    if item is not None:
      return tree_utils.deserialize_tree(restored_item, item)
    return restored_item

  async def _write_aggregate_file(
      self,
      directory: epath.Path,
      item: PyTree,
      param_infos: PyTree,
      save_args: PyTree,
  ) -> future.Future:
    ser_item = _get_tree_for_aggregation(param_infos, save_args, item)
    return await self._aggregate_handler.serialize(
        directory / self._aggregate_filename, ser_item
    )

  def _read_aggregate_file(self, directory: epath.Path) -> PyTree:
    """Restores the aggregate file representing PyTree structure."""
    checkpoint_path = directory / self._aggregate_filename
    if checkpoint_path.exists():
      return self._aggregate_handler.deserialize(checkpoint_path)
    elif self._use_ocdbt:
      raise FileNotFoundError(
          f'Checkpoint structure file does not exist at {directory}.'
      )
    else:
      return utils.pytree_structure(directory)

  async def _write_metadata_file(
      self,
      directory: epath.Path,
      item: PyTree,
      save_args: PyTree,
      use_zarr3: bool = False,
  ) -> future.Future:
    tspec = type_handlers._get_tensorstore_spec(  # pylint: disable=protected-access
        os.fspath(directory), name=METADATA_FILE, use_ocdbt=False
    )['kvstore']
    txn = ts.Transaction()
    metadata_ts_context = type_handlers.get_ts_context(use_ocdbt=False)
    t = await ts.KvStore.open(
        tspec, context=metadata_ts_context
    )
    metadata_content = tree_metadata.TreeMetadata.build(
        item,
        save_args=save_args,
        type_handler_registry=self._type_handler_registry,
        use_zarr3=use_zarr3,
    )
    write_future = t.with_transaction(txn).write(
        '', json.dumps(metadata_content.to_json())
    )
    await write_future
    commit_future = txn.commit_async()
    return commit_future

  def _read_metadata_file(
      self, directory: epath.Path
  ) -> tree_metadata.TreeMetadata:
    """Reads metadata file and returns a tree of restore types.

    Args:
      directory: directory

    Returns:
      Tree with _InternalValueMetadata as values.

    Raises:
      FileNotFoundError: if the metadata file is not found.
    """
    path = directory / METADATA_FILE
    if not path.exists():
      raise FileNotFoundError(
          f'Metadata file (named {METADATA_FILE}) does not exist at'
          f' {directory}.'
      )
    return tree_metadata.TreeMetadata.from_json(json.loads(path.read_text()))

  def _get_internal_metadata(
      self, directory: epath.Path
  ) -> Tuple[PyTree, Optional[bool]]:
    """Gets limited information needed to fully restore the checkpoint.

    This information just consists of the restore type for each leaf, as well
    as the aggregated value (from the msgpack file) if present, and determines
    whether we need to deserialize the parameter using TypeHandler later.

    Args:
      directory: directory

    Returns:
      A PyTree of _InternalValueMetadata with the tree structure of the
      checkpoint.
    """
    aggregate_tree = self._read_aggregate_file(directory)
    flat_aggregate = tree_utils.to_flat_dict(
        aggregate_tree, keep_empty_nodes=True
    )
    try:
      metadata = self._read_metadata_file(directory)
      metadata_tree = metadata.as_nested_tree(keep_empty_nodes=True)
      flat_metadata = tree_utils.to_flat_dict(
          metadata_tree, keep_empty_nodes=True
      )
      use_zarr3 = metadata.use_zarr3
    except FileNotFoundError:
      metadata_tree = None
      flat_metadata = None
      use_zarr3 = None
    if flat_metadata is None:
      flat_metadata = jax.tree.map(
          lambda _: None, flat_aggregate, is_leaf=tree_utils.is_empty_or_leaf
      )

    def _get_internal_value_metadata(value_meta, value):
      if value_meta is None:
        if utils.is_supported_empty_aggregation_type(value):
          return value
        restore_type = None
        skip_deserialize = not utils.leaf_is_placeholder(value)
      else:
        if type_handlers.is_empty_typestr(value_meta.value_type):
          return type_handlers.get_empty_value_from_typestr(
              value_meta.value_type
          )
        restore_type, skip_deserialize = (
            value_meta.value_type,
            value_meta.skip_deserialize,
        )
      return _InternalValueMetadata(
          restore_type=restore_type,
          skip_deserialize=skip_deserialize,
          aggregate_value=value,
      )

    result = {}
    for tuple_key in flat_metadata.keys():
      result[tuple_key] = _get_internal_value_metadata(
          flat_metadata[tuple_key], flat_aggregate[tuple_key]
      )
    target = metadata_tree if metadata_tree is not None else aggregate_tree
    return tree_utils.from_flat_dict(result, target=target), use_zarr3

  def _get_user_metadata(self, directory: epath.Path) -> PyTree:
    """Reads metadata file and constructs user-friendly metadata.

    This will involve more file reads than are necessary for internal metadata.
    Typically, we will need to perform extra reads in order to get metadata
    about individual arrays.

    Args:
      directory: directory

    Returns:
      A PyTree of value_metadata.Metadata matching the checkpoint tree
      structure.
    """
    is_ocdbt_checkpoint = type_handlers.is_ocdbt_checkpoint(directory)

    flat_param_infos = {}
    flat_restore_types = {}
    metadata = self._read_metadata_file(directory)
    metadata_tree = metadata.as_nested_tree(keep_empty_nodes=False)
    for keypath, value_meta in tree_utils.to_flat_dict(metadata_tree).items():
      param_name = '.'.join(keypath)
      restore_type, skip_deserialize = (
          value_meta.value_type,
          value_meta.skip_deserialize,
      )
      flat_param_infos[keypath] = ParamInfo(
          name=param_name,
          path=directory / param_name,
          parent_dir=directory,
          skip_deserialize=skip_deserialize,
          is_ocdbt_checkpoint=is_ocdbt_checkpoint,
          use_zarr3=metadata.use_zarr3,
      )
      flat_restore_types[keypath] = restore_type

    flat_metadatas = {}
    batched_param_infos = collections.defaultdict(list)
    batched_keypaths = collections.defaultdict(list)
    for keypath in flat_param_infos:
      param_info = flat_param_infos[keypath]
      restore_type = flat_restore_types[keypath]
      if param_info.skip_deserialize:
        flat_metadatas[keypath] = value_metadata.Metadata(
            name=param_info.name, directory=directory
        )
      else:
        batched_keypaths[restore_type].append(keypath)
        batched_param_infos[restore_type].append(param_info)

    metadata_ops = []
    for restore_type, param_infos in batched_param_infos.items():
      handler = self._type_handler_registry.get(restore_type)
      metadata_ops.append(handler.metadata(param_infos))

    async def _get_metadata():
      return await asyncio.gather(*metadata_ops)

    batched_metadatas = asyncio.run(_get_metadata())
    for keypath_batch, metadata_batch in zip(
        batched_keypaths.values(), batched_metadatas
    ):
      for keypath, value in zip(keypath_batch, metadata_batch):
        flat_metadatas[keypath] = value
    return tree_utils.from_flat_dict(flat_metadatas, target=metadata_tree)

  def metadata(self, directory: epath.Path) -> Optional[PyTree]:
    """Returns tree metadata.

    The result will be a PyTree matching the structure of the saved checkpoint.
    Note that if the item saved was a custom class, the restored metadata will
    be returned as a nested dictionary representation.

    Example::

      {
        'layer0': {
            'w': ArrayMetadata(dtype=jnp.float32, shape=(8, 8), shards=(1, 2)),
            'b': ArrayMetadata(dtype=jnp.float32, shape=(8,), shards=(1,)),
        },
        'step': ScalarMetadata(dtype=jnp.int64),
      }

    If the required metadata file is not present, this method will raise an
    error.

    Args:
      directory: checkpoint location.

    Returns:
      tree containing metadata.
    """
    try:
      return self._get_user_metadata(directory)
    except FileNotFoundError as e:
      raise FileNotFoundError('Could not locate metadata file.') from e

  def finalize(self, directory: epath.Path) -> None:
    """Finalization step.

    Called automatically by the Checkpointer/AsyncCheckpointer just before the
    checkpoint is considered "finalized" in the sense of ensuring atomicity. See
    documentation for `type_handlers.merge_ocdbt_per_process_files`.

    Args:
      directory: Path where the checkpoint is located.
    """
    if not self._use_ocdbt:
      return
    merge_start_time = time.time()
    type_handlers.merge_ocdbt_per_process_files(directory)
    jax.monitoring.record_event_duration_secs(
        '/jax/checkpoint/write/async/ocdbt_merge_duration_secs',
        time.time() - merge_start_time,
    )

  def close(self):
    """Closes the handler. Called automatically by Checkpointer."""
    self._aggregate_handler.close()


@register_with_handler(BasePyTreeCheckpointHandler, for_save=True)
@dataclasses.dataclass
class BasePyTreeSaveArgs(CheckpointArgs):
  """Parameters for saving a PyTree.

  Attributes:
    item (required): a PyTree to be saved.
    save_args: a PyTree with the same structure of `item`, which consists of
      `ocp.SaveArgs` objects as values. `None` can be used for values where no
      `SaveArgs` are specified.
    ocdbt_target_data_file_size: Specifies the target size (in bytes) of each
      OCDBT data file.  It only applies when OCDBT is enabled and Zarr3 must be
      turned on.  If left unspecified, default size is 2GB.  A value of 0
      indicates no maximum file size limit.  For best results, ensure
      chunk_byte_size is smaller than this value.  For more details, refer to
      https://google.github.io/tensorstore/kvstore/ocdbt/index.html#json-kvstore/ocdbt.target_data_file_size
  """

  item: PyTree
  save_args: Optional[PyTree] = None
  ocdbt_target_data_file_size: Optional[int] = None


@register_with_handler(BasePyTreeCheckpointHandler, for_restore=True)
@dataclasses.dataclass
class BasePyTreeRestoreArgs(CheckpointArgs):
  """Parameters for restoring a PyTree.

  Attributes (all optional):
    item: provides the tree structure for the restored item. If not provided,
      will infer the structure from the saved checkpoint. Transformations will
      not be run in this case. Necessary particularly in the case where the
      caller needs to restore the tree as a custom object.
    restore_args: optional object containing additional arguments for
      restoration. It should be a PyTree matching the structure of `item`, or
      if `item` is not provided, then it should match the structure of the
      checkpoint. Each value in the tree should be a `RestoreArgs` object (OR
      a subclass of `RestoreArgs`). Importantly, note that when restoring a
      leaf as a certain type, a specific subclass of `RestoreArgs` may be
      required. `RestoreArgs` also provides the option to customize the
      restore type of an individual leaf.
  """

  item: Optional[PyTree] = None
  restore_args: Optional[PyTree] = None
