# 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.
"""Registry for checkpoint layouts."""
from __future__ import annotations
import asyncio
from typing import Any
from absl import logging
from orbax.checkpoint._src import asyncio_utils
from orbax.checkpoint._src.path import fs_probe
from orbax.checkpoint.experimental.v1._src.context import context as context_lib
from orbax.checkpoint.experimental.v1._src.context import options as options_lib
from orbax.checkpoint.experimental.v1._src.layout import checkpoint_layout
from orbax.checkpoint.experimental.v1._src.layout import orbax_layout
from orbax.checkpoint.experimental.v1._src.layout import orbax_v0_layout
from orbax.checkpoint.experimental.v1._src.layout import safetensors_layout
from orbax.checkpoint.experimental.v1._src.path import types as path_types
InvalidLayoutError = checkpoint_layout.InvalidLayoutError
CheckpointLayout = checkpoint_layout.CheckpointLayout
CheckpointLayoutEnum = options_lib.CheckpointLayout
ORBAX_LAYOUT_CLASSES = [
orbax_layout.OrbaxLayout,
orbax_v0_layout.OrbaxV0Layout,
]
async def _is_orbax_checkpoint_async(path: path_types.PathLike) -> bool:
"""Checks asynchronously whether the path is an Orbax checkpoint.
Args:
path: Path to the checkpoint to check.
Returns:
True if the checkpoint matches any registered Orbax layout class.
"""
ctx = context_lib.get_context()
path = ctx.file_options.path_class(path)
tasks = []
for layout_cls in ORBAX_LAYOUT_CLASSES:
tasks.append(layout_cls().validate_checkpointables(path))
results = await asyncio.gather(*tasks, return_exceptions=True)
return any(not isinstance(r, Exception) for r in results)
[docs]
def is_orbax_checkpoint(path: path_types.PathLike) -> bool:
"""Returns True if the path is an Orbax checkpoint.
Args:
path: Path to the checkpoint to check.
Returns:
True if the path is recognized as an Orbax checkpoint.
"""
return asyncio_utils.run_sync(_is_orbax_checkpoint_async(path))
async def detect_layout(
path: path_types.PathLike,
*,
root_index: fs_probe.DirectoryIndex | None = None,
) -> CheckpointLayoutEnum:
"""Detects the checkpoint layout from filesystem markers.
Layouts are evaluated concurrently for primary markers. Orbax is checked first
or prioritized according to specific layout markers, falling back to
Safetensors if primary layouts do not match.
Args:
path: The path to the checkpoint directory or file.
root_index: Optional pre-computed directory index of `path`.
Returns:
The detected CheckpointLayout enum.
Raises:
InvalidLayoutError: If the path does not match any registered layout.
"""
ctx = context_lib.get_context()
resolved_path = ctx.file_options.path_class(path)
if root_index is None:
root_index = await fs_probe.index_directory(resolved_path)
if await orbax_layout.matches_markers(root_index):
return CheckpointLayoutEnum.ORBAX
if await safetensors_layout.matches_markers(root_index):
return CheckpointLayoutEnum.SAFETENSORS
tried = [
CheckpointLayoutEnum.ORBAX.value,
CheckpointLayoutEnum.SAFETENSORS.value,
]
raise InvalidLayoutError(
f"Could not auto-detect checkpoint layout at {path}. "
f"Tried layouts: {tried}."
)
async def get_layout_class(
layout_enum: CheckpointLayoutEnum, path: path_types.PathLike | None = None
) -> type[CheckpointLayout]:
"""Returns the layout class for the given layout enum.
Args:
layout_enum: Layout enum identifying the checkpoint format.
path: Optional checkpoint path used for version or format detection.
Returns:
The concrete CheckpointLayout class.
Raises:
ValueError: If layout_enum is not recognized.
ImportError: If the requested layout engine (e.g. Roc) is not linked.
"""
match layout_enum:
case CheckpointLayoutEnum.AUTO:
if path is None:
# When saving, there is no existing checkpoint to detect; default to
# ORBAX ("detect on read, always write Orbax").
return orbax_layout.OrbaxLayout
detected_enum = await detect_layout(path)
return await get_layout_class(detected_enum, path)
case CheckpointLayoutEnum.ORBAX:
if path is None or (
await orbax_layout.checkpoint_version(path)
== orbax_layout.CheckpointVersion.V1
):
return orbax_layout.OrbaxLayout
else:
return orbax_v0_layout.OrbaxV0Layout
case CheckpointLayoutEnum.SAFETENSORS:
return safetensors_layout.SafetensorsLayout
case _:
raise ValueError(f"Unsupported checkpoint layout: {layout_enum}")
async def get_checkpoint_layout(
path: path_types.PathLike, layout_enum: CheckpointLayoutEnum
) -> CheckpointLayout:
"""Returns the checkpoint layout class for the given path and validates it.
Args:
path: The path to the checkpoint.
layout_enum: The checkpoint layout to use.
Returns:
The class of the matching :py:class:`.CheckpointLayout`.
Raises:
InvalidLayoutError: If the path is not a valid checkpoint for any registered
layout, with details from each layout's validation attempt.
"""
ctx = context_lib.get_context()
path = ctx.file_options.path_class(path)
layout_class = await get_layout_class(layout_enum, path)
try:
layout = layout_class()
await layout.validate_checkpointables(path)
return layout
except InvalidLayoutError as e:
raise InvalidLayoutError(
f"Could not recognize the checkpoint at {path} as a valid"
f" {layout_enum.value.capitalize()} checkpoint. If you are trying to"
" load a checkpoint that does not conform to the standard Orbax"
" format, use `ctx.checkpoint_layout = ...` to specify the expected"
" checkpoint layout."
) from e
async def _resolve_auto_pytree_name(
layout: CheckpointLayout, path: path_types.Path
) -> str | None:
"""Discovers and validates a pytree checkpointable name at path.
Args:
layout: Validated CheckpointLayout instance.
path: Path to the checkpoint directory.
Returns:
The resolved checkpointable name (str or None).
Raises:
InvalidLayoutError: If no valid PyTree checkpointable can be found.
"""
names = await layout.get_checkpointable_names(path)
for name in names:
try:
await layout.validate(path, name)
logging.info(
"AUTO resolution mode successfully identified a pytree with"
" checkpointable name '%s' at path '%s'. Attempting to load with"
" this name. If this is not the desired checkpointable, please"
" specify the name explicitly.",
name,
path,
)
return name
except InvalidLayoutError:
continue
if not isinstance(layout, orbax_layout.OrbaxLayout):
try:
await layout.validate(path, None)
logging.info(
"AUTO resolution mode successfully identified a pytree at path"
" '%s'. Attempting to load as a flat layout checkpoint with"
" checkpointable_name=None.",
path,
)
return None
except InvalidLayoutError:
pass
raise InvalidLayoutError(
"Failed to load checkpoint using AUTO resolution mode on"
f" path='{path}'. No valid PyTree checkpointable found."
)
class CheckpointLayoutResolver:
"""Resolves the layout and pytree name for a checkpoint."""
def __init__(
self,
path: path_types.PathLike,
layout_enum: CheckpointLayoutEnum,
layout: CheckpointLayout,
resolved_pytree_name: str | None,
):
self._path = path
self._layout_enum = layout_enum
self._layout = layout
self._resolved_pytree_name = resolved_pytree_name
@classmethod
async def resolve(
cls,
path: path_types.PathLike,
layout_enum: CheckpointLayoutEnum,
*,
pytree_name: str | None = None,
) -> CheckpointLayoutResolver:
"""Resolves the layout and pytree name for a checkpoint.
Args:
path: The path to the checkpoint.
layout_enum: The checkpoint layout to use.
pytree_name: The name of the pytree to load. If
`checkpoint_layout.AUTO_CHECKPOINTABLE_KEY`, the name will be
auto-resolved.
Returns:
A CheckpointLayoutResolver instance.
Raises:
InvalidLayoutError: If the checkpoint layout is invalid or if a valid
PyTree checkpointable cannot be found.
"""
ctx = context_lib.get_context()
path = ctx.file_options.path_class(path)
layout = await get_checkpoint_layout(path, layout_enum)
if pytree_name == checkpoint_layout.AUTO_CHECKPOINTABLE_KEY:
resolved_name = await _resolve_auto_pytree_name(layout, path)
return cls(path, layout_enum, layout, resolved_name)
await layout.validate(path, pytree_name)
return cls(path, layout_enum, layout, pytree_name)
@property
def layout(self) -> CheckpointLayout:
"""Returns the checkpoint layout."""
return self._layout
@property
def pytree_name(self) -> str | None:
"""Returns the resolved pytree name."""
return self._resolved_pytree_name