"""Run directory and checkpoint resolution helpers."""
from __future__ import annotations
import datetime
import os
import traceback
from collections.abc import Iterator
from contextlib import contextmanager
from os import PathLike
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
from omegaconf import DictConfig, OmegaConf
from unisim.backend.base import normalize_play_render_mode
from unilab.rl.runtime.distributed.context import UNILAB_DP_LOG_DIR
from unilab.utils.checkpoint import (
TEST_LOG_ROOT_ENV,
get_latest_checkpoint,
get_latest_run,
normalize_load_run,
resolve_checkpoint_path,
)
if TYPE_CHECKING:
from unilab.training.tensor_diagnostics import NanGuardCfg
[docs]
def build_run_dir_name(timestamp: str, sim_backend: str, *, world_size: int = 1) -> str:
"""Return the canonical run directory name shared by all training entries."""
gpu_suffix = f"_gpux{world_size}" if world_size > 1 else ""
return f"{timestamp}_{sim_backend}{gpu_suffix}"
[docs]
def algo_config_dict(cfg: DictConfig) -> dict[str, Any]:
"""Resolve the composed ``cfg.algo`` subtree into a plain mutable dict."""
raw = OmegaConf.to_container(cfg.algo, resolve=True)
if not isinstance(raw, dict):
raise TypeError("cfg.algo must resolve to a dict")
return cast(dict[str, Any], raw)
[docs]
def resolve_nan_guard_cfg(training_cfg: Any) -> NanGuardCfg | None:
"""Build the shared tensor diagnostics config, or ``None``."""
nan_guard_cfg = getattr(training_cfg, "nan_guard", None)
if nan_guard_cfg is None or not getattr(nan_guard_cfg, "enabled", False):
return None
from unilab.training.tensor_diagnostics import NanGuardCfg
return NanGuardCfg(
enabled=True,
buffer_size=int(getattr(nan_guard_cfg, "buffer_size", 100)),
max_envs_to_dump=int(getattr(nan_guard_cfg, "max_envs_to_dump", 5)),
output_dir=getattr(nan_guard_cfg, "output_dir", None),
)
[docs]
def apply_env_nan_guard(env: Any, training_cfg: Any) -> None:
"""Attach a ``NanGuard`` to ``env`` when diagnostics are enabled."""
nan_guard_cfg = resolve_nan_guard_cfg(training_cfg)
if nan_guard_cfg is None:
return
from unilab.training.tensor_diagnostics import NanGuard
env.set_nan_guard(
NanGuard(
nan_guard_cfg,
num_envs=env.num_envs,
supports_state_playback=env.play_capabilities.supports_physics_state_playback,
)
)
[docs]
def is_viser_play_render_mode(play_render_mode: str | None) -> bool:
"""Return whether the configured mode selects the browser-based viser viewer.
``viser`` is a UniLab-level playback mode layered on top of the unisim
``PLAY_RENDER_MODES`` contract; it renders through the shared MuJoCo
playback shell instead of a backend-native renderer, so it never reaches
``SimBackend.resolve_play_render_plan``.
"""
return play_render_mode is not None and str(play_render_mode).strip().lower() == "viser"
[docs]
def should_run_playback(*, play_only: bool, no_play: bool, play_render_mode: str | None) -> bool:
"""Return whether train/eval should enter playback for the configured mode."""
if is_viser_play_render_mode(play_render_mode):
return bool(play_only) or not bool(no_play)
if normalize_play_render_mode(play_render_mode) == "none":
return False
return bool(play_only) or not bool(no_play)
[docs]
@contextmanager
def nonfatal_play_step(step: str) -> Iterator[None]:
"""Keep independent post-training play artifacts (ONNX export, video render)
from blocking each other: log the failure and continue with the rest."""
try:
yield
except Exception:
print(f"WARNING: {step} failed; continuing with the remaining play steps.")
traceback.print_exc()
[docs]
def get_log_root(
root_dir: str | Path,
cfg: DictConfig | None = None,
*,
algo_log_name: str | None = None,
log_root: str | Path | None = None,
) -> Path:
"""Resolve the algorithm log root.
The configured ``training.log_root`` (or the explicit ``log_root``
argument) wins; otherwise the ``UNILAB_TEST_LOG_ROOT`` override applies;
otherwise the default is ``<root_dir>/logs/<algo_log_name>``. Pass either
a composed Hydra ``cfg`` or an explicit ``algo_log_name``.
"""
if cfg is not None:
if log_root is None:
log_root = OmegaConf.select(cfg, "training.log_root")
if algo_log_name is None:
selected = OmegaConf.select(cfg, "algo.algo_log_name")
algo_log_name = str(selected) if selected is not None else None
if algo_log_name is None:
raise ValueError("get_log_root requires a Hydra cfg or an explicit algo_log_name")
if log_root:
configured_root = Path(log_root)
return (
configured_root if configured_root.is_absolute() else Path(root_dir) / configured_root
)
test_log_root = os.environ.get(TEST_LOG_ROOT_ENV)
if test_log_root:
return Path(test_log_root) / algo_log_name
return Path(root_dir) / "logs" / algo_log_name
[docs]
def resolve_run_log_dir(
cfg: DictConfig,
*,
world_size: int = 1,
timestamp: str | None = None,
root_dir: str | Path | None = None,
) -> str:
"""Resolve one canonical run directory shared by all distributed ranks."""
distributed_log_dir = os.environ.get(UNILAB_DP_LOG_DIR)
if distributed_log_dir:
return distributed_log_dir
configured_log_dir = OmegaConf.select(cfg, "training.log_dir", default=None)
if configured_log_dir:
return str(configured_log_dir)
timestamp = timestamp or datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
return str(
get_log_root(Path.cwd() if root_dir is None else root_dir, cfg)
/ str(cfg.training.task_name)
/ build_run_dir_name(
timestamp,
str(cfg.training.sim_backend),
world_size=world_size,
)
)