Step entities

Orbax step storage entities.

Metadata

class orbax.checkpoint.path.step.Metadata(step, path)

Metadata of a step.

step

step number of the checkpoint.

Type:

int

path

path to the checkpoint.

Type:

etils.epath.abstract_path.Path

property init_timestamp_nsecs: int | None

Returns init timestamp of uncommitted checkpoint of this step.

It is specified as nano seconds since epoch.

Return type:

Optional[int]

property commit_timestamp_nsecs: int | None

Returns commit timestamp of the checkpoint of this step.

It is specified as nano seconds since epoch.

Return type:

Optional[int]

property commit_timestamp: datetime

Returns commit timestamp of the checkpoint of this step.

It is specified as datetime in UTC timezone.

Return type:

datetime

__delattr__(name)

Implement delattr(self, name).

__eq__(other)

Return self==value.

__hash__()

Return hash(self).

__init__(step, path)

__setattr__(name, value)

Implement setattr(self, name, value).

NameFormat

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

Protocol responsible for naming and querying steps.

abstract 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

abstract build_metadata(step_path, *args, **kwargs)

Returns metadata for given step_path if it is valid or None.

NOTE: Ignores uncommitted checkpoints.

Implementation hint: Implement it to build the MetadataT instance using the given step_path. It should be called from find_metadata(…), find_step(…) and find_all(…). This method may not perform IO operations.

Parameters:

• step_path (Path) – Path to folder containing step data.

• *args – Overridable args meant to provide custom params.

• **kwargs – Overridable kwargs meant to provide custom params.

Return type:

Optional[~MetadataT]

abstract find_metadata(base_path, step, *args, **kwargs)

Returns metadata for given base_path and step or None.

NOTE: Ignores uncommitted checkpoints.

Implementation hint: Implement it to find the step folder under base_path performing IO operations if needed. Use build_metadata(…) to build the MetadataT using the found step path.

Parameters:

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

• step (int) – Step number.

• *args – Overridable args meant to provide custom params.

• **kwargs – Overridable kwargs meant to provide custom params.

Return type:

Optional[~MetadataT]

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_metadata(…) and build_step_metadatas(…) 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]

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. Use build_metadata(…) to build the MetadataT using the found step path.

Parameters:

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

• step (int) – Step number.

Return type:

~MetadataT

__init__(*args, **kwargs)

__subclasshook__()

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).

Factories for various NameFormats

orbax.checkpoint.path.step.standard_name_format(step_prefix=None, step_format_fixed_length=None)

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 (Optional[str]) – Optional fixed string prefixed to step. Note an underscore is appended before applying it.

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

Return type:

NameFormat

orbax.checkpoint.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) – 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]) – 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

Helper functions

orbax.checkpoint.path.step.build_step_path(base_path, name_format, step)

Returns step path under base_path for step name_format.

Return type:

Path

orbax.checkpoint.path.step.checkpoint_steps_paths(checkpoint_dir)

Returns a list of finalized checkpoint paths in the directory.

Return type:

List[Path]

orbax.checkpoint.path.step.checkpoint_steps(checkpoint_dir, single_host_load_and_broadcast=False)

Returns a list of finalized checkpoint steps in the directory.

Return type:

List[int]

orbax.checkpoint.path.step.any_checkpoint_step(checkpoint_dir)

Returns any finalized checkpoint step in the directory or None.

This avoids iterating over the entire directory.

Parameters:

checkpoint_dir (Union[str, PathLike]) – Checkpoint directory.

Return type:

Optional[int]

Returns:

Any finalized checkpoint step in the directory or None.

orbax.checkpoint.path.step.latest_step_metadata(root_path, name_format)

Returns step.MetadataT of the latest step in root_path.

Return type:

Optional[~MetadataT]

orbax.checkpoint.path.step.step_metadata_of_checkpoint_path(checkpoint_path, name_format)

Returns step.MetadataT of given checkpoint_path.

Return type:

~MetadataT

Helper functions (experimental)

orbax.checkpoint.path.step.find_step_path(base_path, name_format, *, step, include_uncommitted=False)

Returns step path under base_path for step name_format.

NOTE: Experimental function, subject to change.

Parameters:

• base_path (Union[str, PathLike]) – directory path containing step subdirs.

• name_format (NameFormat) – NameFormat of the target step.

• step (int) – target step number.

• include_uncommitted (bool) – if True then uncommitted steps are considered in search too, otherwise only committed steps are looked up.

Raises:

ValueError if the target step path does not exist. –

Return type:

Path

orbax.checkpoint.path.step.is_gcs_path(path)

Return type:

bool

orbax.checkpoint.path.step.get_save_directory(step, directory, name=None, step_prefix=None, override_directory=None, step_format_fixed_length=None, step_name_format=None)

Returns the standardized path to a save directory for a single item.

Parameters:

• step (int) – Step number.

• directory (Union[str, PathLike]) – Top level checkpoint directory.

• name (Optional[str]) – Item name (‘params’, ‘state’, ‘dataset’, etc.).

• step_prefix (Optional[str]) – Prefix applied to step (e.g. ‘checkpoint’).

• override_directory (Union[str, PathLike, None]) – If provided, step, directory, and step_prefix are ignored.

• step_format_fixed_length (Optional[int]) – Uses a fixed number of digits with leading zeros to represent the step number. If None, there are no leading zeros.

• step_name_format (Optional[NameFormat]) – NameFormat used to define step name for step and under given root directory. If provided, step_prefix and step_format_fixed_length are ignored.

Return type:

Path

Returns:

A directory.

orbax.checkpoint.path.step.is_tmp_checkpoint(path)

Determines whether a directory is a tmp checkpoint path.

Return type:

bool

orbax.checkpoint.path.step.is_checkpoint_finalized(path)

Determines if the given path represents a finalized checkpoint.

Path takes the form:

path/to/my/dir/<name>.orbax-checkpoint-tmp-<timestamp>/  # not finalized
path/to/my/dir/<name>/  # finalized

Alternatively:

gs://path/to/my/dir/<name>/  # finalized
  commit_success.txt
  ...
gs://<path/to/my/dir/<name>/  # not finalized
  ...

Parameters:

path (Union[str, PathLike]) – Directory.

Return type:

bool

Returns:

True if the checkpoint is finalized.

Raises:

• ValueError if the provided path is not a directory. Valid checkpoint paths –

• must be a directory. –

orbax.checkpoint.path.step.create_tmp_directory(final_dir, *, primary_host=0, active_processes=None)

Creates a non-deterministic tmp directory for saving for given final_dir.

Also writes checkpoint metadata in the tmp directory.

Parameters:

• final_dir (Union[str, PathLike]) – The eventual directory path where checkpoint will be committed.

• primary_host (Optional[int]) – primary host id, default=0.

• active_processes (Optional[Set[int]]) – Ids of active processes. default=None

Return type:

Path

Returns:

The tmp directory.

Raises:

FileExistsError – if tmp directory already exists.

orbax.checkpoint.path.step.tmp_checkpoints(checkpoint_dir)

Returns a list of tmp checkpoint dir names in checkpoint_dir.

Return type:

List[str]

orbax.checkpoint.path.step.cleanup_tmp_directories(directory, primary_host=0, active_processes=None)

Cleanup steps in directory with tmp files, as these are not finalized.

orbax.checkpoint.path.step.get_tmp_directory(path)

Returns a non-deterministic tmp directory for path without creating it.

Return type:

Path

orbax.checkpoint.path.step.step_from_checkpoint_name(name)

Returns the step from a checkpoint name. Also works for tmp checkpoints.

Return type:

int

Helper functions to implement NameFormat classes

orbax.checkpoint.path.step.build_step_metadatas(step_paths, build_metadata)

Yields filtered metadata mapped with step_paths.

Parameters:

• step_paths (Iterator[Path]) – Iterator of step paths.

• build_metadata (Callable[[Path], Optional[~MetadataT]]) – Callable to match and build step metadata from step_paths elements. If a step_paths element doesn’t match then it returns None.

Yields:

Step metadata.

Return type:

Iterator[~MetadataT]

orbax.checkpoint.path.step.step_prefix_with_underscore(step_prefix)

Returns step_prefix appended with underscore or <empty> if None.

Return type:

str