Step entities#
Orbax step storage entities.
Please reserve this file for entities related to NameFormat. Other existing functions should be moved out over time.
Metadata#
- class orbax.checkpoint.path.step.Metadata(step, path)[source]#
Metadata of a step.
- step#
step number of the checkpoint.
- Type:
int
- path#
path to the checkpoint.
- 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,None]
- 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,None]
- 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)[source]#
Protocol responsible for naming and querying steps.
- build_name(step)[source]#
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)[source]#
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)[source]#
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).
Factories for various NameFormats#
- orbax.checkpoint.path.step.standard_name_format(*, step_prefix=None, step_format_fixed_length=None, single_host_load_and_broadcast=False)[source]#
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,None]) – Optional fixed string prefixed to step. Note an underscore is appended before applying it.step_format_fixed_length (
Optional[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:
- orbax.checkpoint.path.step.composite_name_format(write_name_format, read_name_formats)[source]#
Returns composite NameFormat supporting multiple read/single write formats.
- Parameters:
write_name_format (
NameFormat[Metadata]) – 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[Metadata]]) – 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:
Helper functions#
- orbax.checkpoint.path.step.build_step_path(base_path, name_format, step)[source]#
Returns step path under base_path for step name_format.
- Return type:
- orbax.checkpoint.path.step.checkpoint_steps_paths(checkpoint_dir)[source]#
Returns a list of finalized checkpoint paths in the directory.
- Return type:
List[Path]
- orbax.checkpoint.path.step.checkpoint_steps(checkpoint_dir)[source]#
Returns a list of finalized checkpoint steps in the directory.
- Return type:
List[int]
- orbax.checkpoint.path.step.any_checkpoint_step(checkpoint_dir)[source]#
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,None]- Returns:
Any finalized checkpoint step in the directory or None.
Helper functions (experimental)#
- orbax.checkpoint.path.step.find_step_path(base_path, name_format, *, step, include_uncommitted=False)[source]#
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[Metadata]) – 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:
- 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)[source]#
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,None]) – Item name (‘params’, ‘state’, ‘dataset’, etc.).step_prefix (
Optional[str,None]) – 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,None]) – 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[Metadata],None]) – 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:
- Returns:
A directory.
- orbax.checkpoint.path.step.is_tmp_checkpoint(*a, **k)#
- orbax.checkpoint.path.step.is_checkpoint_finalized(*a, **k)#
- orbax.checkpoint.path.step.tmp_checkpoints(*a, **k)#