opentau.utils.train_utils

Utilities for training checkpoint management and state persistence.

This module provides functions for saving and loading training checkpoints, managing checkpoint directories, and pruning old checkpoints.

Functions

find_missing_rng_state_ranks(checkpoint_dir, ...)	Return the process indices whose random_states_<i>.pkl is absent from a checkpoint.
get_step_checkpoint_dir(output_dir, ...)	Returns the checkpoint sub-directory corresponding to the step number.
get_step_identifier(step, total_steps)	Generate a zero-padded step identifier string.
load_running_best_state(output_dir, ...)	Load a running-best tracker, self-healing against pruned directories.
load_training_state(checkpoint_dir)	Load training state including step, optimizer, scheduler, and RNG state.
load_training_step(save_dir)	Load the training step number from a file.
log_output_dir(out_dir)	Log the output directory path with colored formatting.
prune_old_checkpoints(latest_checkpoint_path)	Delete all checkpoint directories except the specified one (and any protected ones).
reseed_new_ranks_on_resume(checkpoint_dir, ...)	Re-seed ranks that have no per-rank RNG file in the checkpoint.
running_best_is_improvement(score, best, ...)	Return whether score strictly improves on best.
save_checkpoint(checkpoint_dir, step, cfg)	Save training checkpoint including config and RNG state.
save_running_best_state(tracker, metric)	Persist the running-best state next to the checkpoint directories.
save_training_step(step, save_dir)	Save the current training step number to a file.
update_last_checkpoint(checkpoint_dir)	Update the symlink pointing to the last checkpoint.

Classes

RunningBestTracker(output_dir, total_steps, ...)

Rank-0 bookkeeping for running-best checkpoints.

class opentau.utils.train_utils.RunningBestTracker(output_dir: Path, total_steps: int, higher_is_better: bool, max_count: int, best: float | None = None, steps: list[dict] | None = None)

Bases: object

Rank-0 bookkeeping for running-best checkpoints.

Tracks the best metric value seen so far and the pool of step numbers whose checkpoints are currently retained as running bests (most-recent last). Every method is pure Python / filesystem-only – no distributed collectives – so it runs on the main process and is unit testable on CPU. The actual (collective) accelerator.save_state is performed by the caller.

Each pool entry stores the step and whether a regular save_freq checkpoint was written at that step (captured at registration time). On eviction, a regular checkpoint’s directory is left to the normal retention logic; only running-best-only directories are deleted.

__init__(output_dir: Path, total_steps: int, higher_is_better: bool, max_count: int, best: float | None = None, steps: list[dict] | None = None)

is_improvement(score: float | None) → bool

Return whether score strictly improves on the current best (no mutation).

protected_dirs() → set[Path]

Return the set of checkpoint directories currently retained as running bests.

register(step: int, score: float, is_regular: bool) → list[Path]

Record a new best at step and return the checkpoint directories to delete.

Updates the high-water mark, adds step to the pool (de-duplicating, moving it to most-recent), then evicts the oldest entries beyond max_count. An evicted step’s directory is scheduled for deletion only if it was not a regular save_freq checkpoint and the directory still exists.

Parameters:

• step – The training step that achieved the new best.

• score – The (finite) metric value at this step; becomes the new high-water mark.

• is_regular – Whether a regular save_freq checkpoint was also written at this step.

Returns:

List of running-best-only checkpoint directories the caller should rmtree.

step_dir(step: int) → Path

Return the checkpoint directory for a given step.

to_dict() → dict

Serialize to a JSON-friendly dict (high-water mark + pool).

opentau.utils.train_utils.find_missing_rng_state_ranks(checkpoint_dir: Path, world_size: int) → list[int]

Return the process indices whose random_states_<i>.pkl is absent from a checkpoint.

Accelerate writes exactly one random_states_<process_index>.pkl per process on every backend, so a complete checkpoint contains world_size of them. A non-empty result means some rank saved into a different directory or not at all (e.g. a divergent output_dir across ranks), leaving the checkpoint unresumable.

Parameters:

• checkpoint_dir – Directory accelerate saved the per-rank RNG files into.

• world_size – Expected number of processes (accelerator.num_processes).

Returns:

Sorted list of missing process indices; empty when the checkpoint is complete.

opentau.utils.train_utils.get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) → Path

Returns the checkpoint sub-directory corresponding to the step number.

opentau.utils.train_utils.get_step_identifier(step: int, total_steps: int) → str

Generate a zero-padded step identifier string.

Parameters:

• step – Current step number.

• total_steps – Total number of steps (used to determine padding width).

Returns:

Zero-padded string representation of the step number.

opentau.utils.train_utils.load_running_best_state(output_dir: Path, total_steps: int, higher_is_better: bool, max_count: int) → RunningBestTracker

Load a running-best tracker, self-healing against pruned directories.

If the state file is absent (a fresh run, or a resume of a run that predates the feature), a fresh tracker is returned. Pool entries whose checkpoint directory no longer exists are dropped (a regular-coincident best may have been pruned later by the normal retention logic), but the high-water mark is preserved regardless.

If the persisted optimization direction disagrees with higher_is_better (a resume that flipped running_best_metric between maximize and minimize), the persisted high-water mark is meaningless under the new direction, so it is discarded (reset to -inf/+inf) while the existing pool is kept; a warning is emitted.

Parameters:

• output_dir – The run output directory (its checkpoints/ subdir holds the state file).

• total_steps – Total training steps (for step-dir zero-padding).

• higher_is_better – True to maximize the metric, False to minimize.

• max_count – Maximum number of running bests to keep.

Returns:

A RunningBestTracker restored from disk, or a fresh one if no state exists.

opentau.utils.train_utils.load_training_state(checkpoint_dir: Path) → tuple[int, Optimizer, LRScheduler | None]

Load training state including step, optimizer, scheduler, and RNG state.

This is used to resume a training run. Note: optimizer and scheduler states are loaded by accelerate, not by this function.

Parameters:

checkpoint_dir – The checkpoint directory. Should contain a ‘training_state’ dir.

Returns:

optimizer and scheduler are loaded separately by accelerate.

Return type:

Tuple containing the training step number. Note

Raises:

NotADirectoryError – If checkpoint_dir doesn’t exist or is not a directory.

opentau.utils.train_utils.load_training_step(save_dir: Path) → int

Load the training step number from a file.

Parameters:

save_dir – Directory containing the step file.

Returns:

Training step number.

opentau.utils.train_utils.log_output_dir(out_dir)

Log the output directory path with colored formatting.

Parameters:

out_dir – Output directory path to log.

opentau.utils.train_utils.prune_old_checkpoints(latest_checkpoint_path: str, protected_paths: set[Path] | None = None) → None

Delete all checkpoint directories except the specified one (and any protected ones).

Recursively deletes all checkpoint directories in a parent folder except for the specified one. This function is designed to clean up old model checkpoints, preserving only the most recent one. It includes safety checks to ensure it only deletes directories and handles potential filesystem errors.

Parameters:

• latest_checkpoint_path – The full path to the checkpoint directory that should be kept.

• protected_paths – Additional checkpoint directories that must NOT be deleted (e.g. running-best checkpoints). Compared by resolved path. Defaults to None.

opentau.utils.train_utils.reseed_new_ranks_on_resume(checkpoint_dir: Path, accelerator: accelerate.Accelerator, seed: int | None) → None

Re-seed ranks that have no per-rank RNG file in the checkpoint.

Accelerate’s load_state reads random_states_<process_index>.pkl per rank and silently logs "Could not load random states" at INFO level when the file is missing. When resuming on more GPUs than the checkpoint was saved with, the new ranks are left with whatever default RNG state the process happened to start with – correlated across ranks and irreproducible. This helper detects that case and assigns each new rank a deterministic, decorrelated seed via set_seed.

Surviving ranks (whose files loaded successfully) are not touched.

Parameters:

• checkpoint_dir – Directory holding the saved random_states_*.pkl files.

• accelerator – The current Accelerator (provides num_processes and process_index).

• seed – Base training seed (cfg.seed). If None, no re-seeding is performed and a warning is emitted instead.

opentau.utils.train_utils.running_best_is_improvement(score: float | None, best: float, higher_is_better: bool) → bool

Return whether score strictly improves on best.

Strict comparison (> / <): a value equal to the current best does not count as an improvement, so a plateaued metric does not thrash the running-best pool. None and NaN scores never count as an improvement (e.g. an eval that produced no episodes reports NaN).

Parameters:

• score – The new metric value (may be None or NaN).

• best – The current best metric value.

• higher_is_better – True to maximize (e.g. success rate), False to minimize (e.g. loss).

Returns:

True iff score is finite and strictly better than best.

opentau.utils.train_utils.save_checkpoint(checkpoint_dir: Path, step: int, cfg: TrainPipelineConfig) → None

Save training checkpoint including config and RNG state.

Note: accelerate saves the model and training run. This method saves all other auxiliary objects such as configs and RNG state.

Parameters:

• checkpoint_dir – Directory where the checkpoint will be saved.

• step – Current training step number.

• cfg – The training config used for this run.

opentau.utils.train_utils.save_running_best_state(tracker: RunningBestTracker, metric: str) → None

Persist the running-best state next to the checkpoint directories.

Written as a plain file in <output_dir>/checkpoints/ (not inside a step dir), so it survives prune_old_checkpoints (which only deletes directories).

Parameters:

• tracker – The running-best tracker to persist (provides output_dir).

• metric – The resolved driving metric name (stored for provenance/debugging).

opentau.utils.train_utils.save_training_step(step: int, save_dir: Path) → None

Save the current training step number to a file.

Parameters:

• step – Current training step number.

• save_dir – Directory where the step file will be saved.

opentau.utils.train_utils.update_last_checkpoint(checkpoint_dir: Path) → Path

Update the symlink pointing to the last checkpoint.

Parameters:

checkpoint_dir – Path to the checkpoint directory to link.

Returns:

Path to the symlink that was created or updated.