nemo_rl.utils.checkpoint_storage

Opt-in MSC checkpoint storage; ordinary filesystem paths remain unchanged.

Module Contents

Classes

CheckpointStorage	How a checkpoint root creates, publishes and lists step_N directories.
LocalStorage
MscStorage

Functions

_parse_step	Return N for a published step_N directory name, else None.
make_storage	Pick the storage that matches root (msc:// URL or local directory).

Data

_STEP_DIR

API

nemo_rl.utils.checkpoint_storage._STEP_DIR

‘compile(…)’

nemo_rl.utils.checkpoint_storage._parse_step(dir_name: str) → int | None

Return N for a published step_N directory name, else None.

class nemo_rl.utils.checkpoint_storage.CheckpointStorage

Bases: typing.Protocol

How a checkpoint root creates, publishes and lists step_N directories.

Local filesystems publish atomically by renaming tmp_step_N to step_N. Object stores cannot rename a prefix, so MSC writes step_N in place and publishes it by writing a _SUCCESS marker last.

begin(step: int) → str | pathlib.Path

Create and return the directory the checkpoint for step is written into.

publish(path: str | os.PathLike[str]) → None

Make the directory returned by begin visible to published_steps.

published_steps() → list[tuple[int, str]]

Return (step, path) for every published checkpoint under the root.

class nemo_rl.utils.checkpoint_storage.LocalStorage

root: pathlib.Path

None

begin(step: int) → pathlib.Path

publish(path: str | os.PathLike[str]) → None

Rename tmp_step_N to step_N.

If step_N already exists (defensive guard for edge cases, e.g. resuming training), performs a pseudo-atomic swap via an intermediate old_step_N directory.

published_steps() → list[tuple[int, str]]

class nemo_rl.utils.checkpoint_storage.MscStorage

root: pathlib.Path

None

begin(step: int) → str

publish(path: str | os.PathLike[str]) → None

published_steps() → list[tuple[int, str]]

nemo_rl.utils.checkpoint_storage.make_storage(

root: str | os.PathLike[str],

) → nemo_rl.utils.checkpoint_storage.CheckpointStorage

Pick the storage that matches root (msc:// URL or local directory).