ocp.v1.path module

Public API for user-facing Path constructs.

Types

class orbax.checkpoint.experimental.v1.path.Path(*args: str | PathLike)

Abstract base class for pathlib.Path-like API.

See [pathlib.Path](https://docs.python.org/3/library/pathlib.html) documentation.

static __new__(cls, *args)

Create a new path.

`python path = abcpath.Path() `

We use __new__ instead of __init__ to allow subclassing, even though the usage of __init__ is possible from python>=3.12.

Parameters:

*args – Paths to create

Returns:

The registered path

Return type:

path

format(*args, **kwargs)

Apply str.format() to the path.

Return type:

~_T

exists()

Returns True if self exists.

Return type:

bool

is_dir()

Returns True if self is a dir.

Return type:

bool

is_file()

Returns True if self is a file.

Return type:

bool

iterdir()

Iterates over the directory.

Return type:

Iterator[~_T]

listdir()

Lists files/folders in a directory.

Colab usage only!

This is not in the official pathlib.Path API, but on Colab, using list(path.iterdir()) is such a common pattern that it’s convenient to have it.

Return type:

list[~_T]

glob(pattern)

Yields all matching files (of any kind).

Return type:

Iterator[~_T]

rglob(pattern)

Yields all matching files recursively (of any kind).

Return type:

Iterator[~_T]

expanduser()

Returns a new path with expanded ~ and ~user constructs.

Return type:

~_T

absolute()

Returns the absolute path.

Return type:

~_T

resolve(strict=False)

Returns the absolute path.

Return type:

~_T

open(mode='r', encoding=None, errors=None, **kwargs)

Opens the file.

Return type:

IO[AnyStr]

read_bytes()

Reads contents of self as bytes.

Return type:

bytes

read_text(encoding=None)

Reads contents of self as a string.

Return type:

str

stat()

Returns metadata for the file/directory.

Return type:

StatResult

mkdir(mode=None, parents=False, exist_ok=False)

Create a new directory at this given path.

Return type:

None

rmdir()

Remove the empty directory at this given path.

Return type:

None

rmtree(missing_ok=False)

Remove the directory, including all sub-files.

Return type:

None

unlink(missing_ok=False)

Remove this file or symbolic link.

Return type:

None

write_bytes(data)

Writes content as bytes.

Return type:

int

write_text(data, encoding=None, errors=None)

Writes content as str.

Return type:

int

touch(mode=438, exist_ok=True)

Create a file at this given path.

Return type:

None

rename(target)

Renames the path.

Return type:

~_T

replace(target)

Overwrites the destination path.

Return type:

~_T

copy(dst, overwrite=False)

Copy the current file to the given destination.

Return type:

~_T

classmethod __get_pydantic_core_schema__(source_type, handler)

Handles pydantic serialization similarly to a pathlib.PurePosixPath.

See https://docs.pydantic.dev/latest/concepts/types

Parameters:

• source_type – Any

• handler – pydantic.GetCoreSchemaHandler

Returns:

pydantic_core.CoreSchema

orbax.checkpoint.experimental.v1.path.PathLike

alias of Path | str

class orbax.checkpoint.experimental.v1.path.PathAwaitingCreation(*args, **kwargs)

A path that may not exist yet, but will exist after await_creation.

This construct is used to represent a path in the process of being created. The underlying path can be accessed logically, but the actual location in the filesystem should not be accessed until await_creation() is called.

Usage:

path: :py:class:`.PathAwaitingCreation` = ...
# Logical accesses are OK.
print(path.path)
# Block until the path is known to exist.
path = await path.await_creation()
path.exists()  # True.

async await_creation()

Waits for the directory to be created.

This is a blocking operation, though it should return immediately if the path has already been created. Be cautious about where this method is called, since implementations may trigger the creation if await_creation is called before the directory creation has been triggered.

It is recommended to only call this from background awaitables, or to delay as long as possible.

Return type:

Path

Returns:

The path that was created.

__init__(*args, **kwargs)

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

Step entities

class orbax.checkpoint.v1.path.step.NameFormat(*args, **kwargs)

Responsible for naming and querying steps.

build_name(step)

Returns step name.

Implementation hint: Implement it to build a name for the given step using the class’s custom formatting attributes. Since it is mainly meant for building names to save checkpoints, it can raise error if this NameFormat is just meant for finding already existing step paths.

Parameters:

step (int) – Step number.

Return type:

str

abstractmethod find_all(base_path)

Returns metadata of all steps.

NOTE: Ignores uncommitted checkpoints.

Implementation hint: Implement it to find all step folders under base_path performing IO operations if needed. Use build_step_metadatas(…) helper function to build all the MetadataT using the found step paths.

Parameters:

base_path (Union[str, PathLike]) – root Path under which Step folders are placed.

Return type:

Iterator[~MetadataT]

abstractmethod find_step(base_path, step)

Returns the metadata for step or raises ValueError.

NOTE: Ignores uncommitted checkpoints.

Implementation hint: Implement it to find the step folder under base_path performing IO operations if needed.

Parameters:

• base_path (Union[str, PathLike]) – root Path under which Step folders are placed.

• step (int) – Step number.

Raises:

ValueError if no committed paths for the requested step is found. –

Return type:

~MetadataT

__init__(*args, **kwargs)

classmethod __subclasshook__(other)

Abstract classes can override this to customize issubclass().

This is invoked early on by abc.ABCMeta.__subclasscheck__(). It should return True, False or NotImplemented. If it returns NotImplemented, the normal algorithm is used. Otherwise, it overrides the normal algorithm (and the outcome is cached).

orbax.checkpoint.v1.path.step.standard_name_format(*, step_prefix=None, step_format_fixed_length=None, single_host_load_and_broadcast=False)

Returns NameFormat for ‘standard’ steps for common Orbax use cases.

NOTE: Ignores uncommitted checkpoints.

Naming examples:

• step_prefix=None step_format_fixed_length=None -> 23

• step_prefix=None step_format_fixed_length=4 -> 0023

• step_prefix=step step_format_fixed_length=None -> step_23

• step_prefix=step step_format_fixed_length=4 -> step_0023

Parameters:

• step_prefix (UnionType[str, None]) – Optional fixed string prefixed to step. Note an underscore is appended before applying it.

• step_format_fixed_length (UnionType[int, None]) – Optional length of the zero padded step. e.g. 6 for 000123.

• single_host_load_and_broadcast (bool) – If True, the jax process=0 will list all steps and broadcast them to all other processes. NOTE: Ignored if jax backend is not multi controller.

Return type:

NameFormat[CheckpointMetadata[None]]

orbax.checkpoint.v1.path.step.composite_name_format(write_name_format, read_name_formats)

Returns composite NameFormat supporting multiple read/single write formats.

Parameters:

• write_name_format (NameFormat[CheckpointMetadata[None]]) – NameFormat used to build step names meant for writing checkpoints. Must be present in read_name_formats at a preferred priority position.

• read_name_formats (Sequence[NameFormat[CheckpointMetadata[None]]]) – Sequence (ordered) of NameFormats used to find steps for reading checkpoints. Please note that to resolve conflicts (and avoid raising errors) in case of multiple NameFormats matching a given step, the sequence should be provided in highest to lowest priority order: NameFormat appearing earlier in the sequence is preferred.

Return type:

NameFormat[CheckpointMetadata[None]]