diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index d6acdcd..b54558d 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -28,7 +28,8 @@ separately downloaded optional asset. ## Runtime dependencies The following are installed from external distributions and remain under their -own licenses: MotrixSim, MuJoCo, Gymnasium, NumPy, Numba, Hydra/OmegaConf, +own licenses: MotrixSim, MuJoCo, Gymnasium, NumPy, Numba, nvidia-ml-py (NVML +bindings used for training-panel GPU metrics), Hydra/OmegaConf, SKRL, RSL-RL, PyTorch/JAX, ONNX Runtime, TensorBoard, and the Unitree SDK2 Python package. Their licenses are not replaced by the MotrixLab license. The release process should generate a dependency license report from the final diff --git a/motrix_rl/pyproject.toml b/motrix_rl/pyproject.toml index 7b28539..8480a3c 100644 --- a/motrix_rl/pyproject.toml +++ b/motrix_rl/pyproject.toml @@ -20,6 +20,7 @@ dependencies = [ "motrix-deploy", "numpy>=1.26", "omegaconf>=2.3,<2.4", + "nvidia-ml-py>=13.610.43", ] [project.optional-dependencies] diff --git a/motrix_rl/src/motrix_rl/console.py b/motrix_rl/src/motrix_rl/console.py index 037d06e..2355feb 100644 --- a/motrix_rl/src/motrix_rl/console.py +++ b/motrix_rl/src/motrix_rl/console.py @@ -6,19 +6,29 @@ from __future__ import annotations import math +import os import re +import shutil import sys from collections.abc import Mapping from dataclasses import dataclass, field from typing import Any -from motrix_rl.system_metrics import CpuLoad +from motrix_rl.system_metrics import CpuLoad, MemoryUsage + +try: # cbreak keyboard input needs a POSIX terminal; other platforms keep a plain Live + import select + import termios + import tty + + _POSIX_TTY = True +except ImportError: + _POSIX_TTY = False try: # optional pretty console; callers fall back to plain text if unavailable from rich.console import Console, Group from rich.live import Live from rich.panel import Panel - from rich.rule import Rule from rich.table import Table _RICH = True @@ -56,6 +66,58 @@ class TrainingPanelStats: # as an indented sub-tree (e.g. per-process timing with stage breakdowns). timing_groups: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) cpu_load: CpuLoad | None = None + gpu_utilization_percent: float | None = None + memory_usage: MemoryUsage | None = None + gpu_memory_usage: MemoryUsage | None = None + checkpoint_path: str | None = None + + +class _InputLive(Live): + """Rich Live display with direct single-key input from the controlling TTY.""" + + def start(self): + # One try/except spans the whole setup: a termios failure mid-way must + # still restore the saved TTY state and close an opened /dev/tty FD. + self._stdin_state = None + self._input_fd = None + self._owns_input_fd = False + try: + if sys.stdin.isatty(): + self._input_fd = sys.stdin.fileno() + else: + try: + self._input_fd = os.open("/dev/tty", os.O_RDONLY | os.O_NONBLOCK) + self._owns_input_fd = True + except OSError: + self._input_fd = None + if self._input_fd is not None: + fd = self._input_fd + self._stdin_state = (fd, termios.tcgetattr(fd)) + tty.setcbreak(fd) + attrs = termios.tcgetattr(fd) + attrs[3] &= ~termios.ECHO + termios.tcsetattr(fd, termios.TCSANOW, attrs) + return super().start() + except BaseException: + self._restore_stdin() + raise + + def stop(self): + try: + return super().stop() + finally: + self._restore_stdin() + + def _restore_stdin(self): + state = self._stdin_state + if state is not None: + fd, attrs = state + termios.tcsetattr(fd, termios.TCSANOW, attrs) + self._stdin_state = None + if self._owns_input_fd and self._input_fd is not None: + os.close(self._input_fd) + self._input_fd = None + self._owns_input_fd = False def open_training_live(): @@ -63,15 +125,30 @@ def open_training_live(): if not _RICH or not sys.stdout.isatty(): return None, None console = Console() - live = Live(console=console, auto_refresh=False, vertical_overflow="visible") + live = (_InputLive if _POSIX_TTY else Live)(console=console, auto_refresh=False, vertical_overflow="visible") live.start() return console, live def emit_training_panel(live, stats: TrainingPanelStats, *, title: str = "rl") -> None: - """Render one training stats panel using rich when ``live`` is available.""" + """Render one training panel; 1/2 switch overview and timing views.""" if live is not None: - live.update(render_training_panel(stats, title=title), refresh=True) + detail = bool(getattr(live, "_motrix_detail", False)) + input_fd = getattr(live, "_input_fd", None) + if input_fd is not None: + ready, _, _ = select.select([input_fd], [], [], 0) + if ready: + try: + command = os.read(input_fd, 1).decode("utf-8", errors="ignore").lower() + except OSError: + command = "" + if command == "2": + detail = True + elif command == "1": + detail = False + live._motrix_detail = detail + panel = render_training_panel(stats, title=title, detail=detail) + live.update(panel, refresh=True) else: print(format_training_panel(stats, title=title)) @@ -240,142 +317,341 @@ def render_block( render_block(pending, lead, row_prefix, row_prefix) -def render_training_panel(stats: TrainingPanelStats, *, title: str = "rl"): - """Build a rich panel for RL training stats.""" - if not _RICH: - raise RuntimeError("rich is not available") +def _timing_totals(stats: TrainingPanelStats) -> tuple[float, float]: + """Return the canonical parent totals stored on TrainingPanelStats.""" + return stats.collect_ms, stats.learn_ms + + +def _timing_group_total(group: str, stats: TrainingPanelStats, index: int) -> float: + name = _strip_markup(group).lower() + if name.startswith("collector"): + return stats.collect_ms + if name.startswith("learner") and "idle" not in name: + return stats.learn_ms + return stats.collect_ms if index == 0 else stats.learn_ms + + +# --------------------------------------------------------------------------- +# Prototype layout renderer +# +# Keep the plain-text helpers above for callers that imported them during the +# initial console experiment. The definitions below are the public renderer +# used by ``emit_training_panel`` and deliberately mirror the compact/timing +# prototypes in docs/prototypes. +def _prototype_metric_cell( + key: str, + value: Any, + *, + precision: int = 3, + signed: bool = False, + label_width: int = 22, +): + from rich.text import Text + + name = _strip_markup(str(key)) + # Reward terms commonly have descriptive names, so keep a wider label + # while still bounding unusual keys to preserve the grid layout. + if len(name) > label_width: + name = name[: label_width - 3] + "..." + text = Text(f"{name:<{label_width}} ", style="white") + color = "white" + if signed and isinstance(value, (float, int)): + color = "green" if value >= 0 else "red" + text.append(f"{_compact_metric_value(value, precision=precision, signed=signed):>9}", style=color) + return text - def hms(t: float) -> str: - t = int(t) - h, m, sec = t // 3600, (t % 3600) // 60, t % 60 - return f"{h}h{m:02d}m" if h else f"{m}m{sec:02d}s" - def si(n: float) -> str: - for unit in ("", "k", "M"): - if abs(n) < 1000: - return f"{n:.0f}{unit}" if unit == "" else f"{n:.1f}{unit}" - n /= 1000.0 - return f"{n:.1f}G" +def _compact_metric_value(value: Any, *, precision: int = 3, signed: bool = False) -> str: + """Fit a metric value into the fixed 9-char cell so metric rows never wrap. - pct = 100.0 * stats.iteration / max(stats.total_iterations, 1) - head = Table.grid(expand=True, padding=(0, 1)) - head.add_column(ratio=1) - head.add_column(justify="right") - head.add_row( - f"[bold]iter[/] {stats.iteration:,}/{stats.total_iterations:,} [dim]({pct:.1f}%)[/]", - f"[cyan]{stats.steps_per_second:.0f}[/] env-steps/s [dim]{hms(stats.elapsed_seconds)}[/]", - ) - if stats.cpu_load is not None: - load = stats.cpu_load - color = "green" if load.utilization_percent < 70.0 else "yellow" if load.utilization_percent < 90.0 else "red" - cores = f" · {load.physical_core_count}C" if load.physical_core_count is not None else "" - head.add_row( - f"[bold]cpu[/] [{color}]{load.utilization_percent:.1f}%[/] " - f"[dim]{load.used_logical_cpus:.1f}/{load.logical_cpu_count}T{cores}[/]", - f"[dim]iowait {load.iowait_percent:.1f}% steal {load.steal_percent:.1f}%[/]", - ) + Values that outgrow the cell (e.g. ``-3.6200e-05`` at precision 4) fall + back to two-digit scientific notation, keeping the grid geometry stable + whatever magnitudes are currently on screen. + """ + formatted = _format_value(value, precision=precision, signed=signed) + if len(formatted) <= 9 or not isinstance(value, (int, float)) or not math.isfinite(value): + return formatted + return f"{value:+.2e}" if signed else f"{value:.2e}" + + +def _rich_compact_metric_grid(items: Mapping[str, Any], *, limit: int | None, precision: int = 3, signed: bool = False): + grid = Table.grid(expand=True, padding=(0, 1)) + # A metric cell is roughly 25–30 terminal columns wide. Adapt the number + # of columns to the real TTY instead of forcing every group into two. + terminal_width = shutil.get_terminal_size((120, 24)).columns + column_count = max(2, min(4, terminal_width // 30)) + for index in range(column_count): + grid.add_column(ratio=1) + if index < column_count - 1: + grid.add_column(width=1, justify="center") + pairs = list(items.items()) + visible = pairs if limit is None else pairs[:limit] + # Compact grids render one Panel deep (Training/Rewards/Environment), so + # budget for that panel's border and padding. Each metric cell needs label + # + separator + 9-char value; deriving the label width from the exact rich + # column arithmetic keeps every cell on one line down to narrow widths. + column_width = (terminal_width - 8 - 3 * column_count) // column_count + label_width = max(12, min(32, column_width - 11)) + cells = [ + _prototype_metric_cell(k, v, precision=precision, signed=signed, label_width=label_width) for k, v in visible + ] + omitted = 0 if limit is None else len(pairs) - len(visible) + if omitted: + from rich.text import Text + + cells.append(Text(f"+{omitted} more", style="cyan")) + for index in range(0, len(cells), column_count): + row: list[Any] = [] + for column, cell in enumerate(cells[index : index + column_count]): + if column: + from rich.text import Text + + row.append(Text("│", style="grey50")) + row.append(cell) + grid.add_row(*row) + return grid - roll = Table.grid(expand=True, padding=(0, 2)) - for _ in range(3): - roll.add_column(ratio=1) - roll.add_row( - f"return [bold green]{stats.mean_return:.2f}[/]", - f"ep_len [yellow]{stats.mean_episode_length:.1f}[/]", - f"episodes {stats.episodes:,}", - ) - buf = f"buffer [magenta]{si(stats.buffer_size)}[/]/{si(stats.buffer_capacity)}" - blocks = [head, Rule(style="grey37"), roll] - metrics = stats.training_metrics - if metrics is None: - message = "warmup - filling replay buffer" if stats.warming else "training metrics unavailable" - blocks += [Rule("training", style="grey37", align="left"), f"[dim]{message}[/] {buf}"] - else: - train_grid = _rich_key_value_grid(metrics) - blocks += [Rule("training", style="grey37", align="left"), train_grid, buf] +def _rich_timing_grid(items: Mapping[str, Any], *, total: float | None = None): + from rich.text import Text - terms = stats.reward_terms - if terms: - sorted_terms = dict(sorted(terms.items(), key=lambda kv: -abs(float(kv[1])))) - grid = _rich_key_value_grid(sorted_terms, signed=True, precision=4) - blocks += [Rule("rewards x dt", style="grey37", align="left"), grid] + grid = Table.grid(expand=True, padding=(0, 1)) + grid.add_column(ratio=1) + grid.add_column(justify="right", width=11) + grid.add_column(justify="right", width=8) + grid.add_row(Text("STAGE", style="dim"), Text("MEAN", style="dim"), Text("SHARE", style="dim")) + + def append(values: Mapping[str, Any], level: int = 0) -> None: + entries = list(values.items()) + for index, (key, value) in enumerate(entries): + branch = "'- " if index == len(entries) - 1 else "|- " + prefix = " " * level + branch + if isinstance(value, Mapping): + children = dict(value) + node_total = children.pop("total", None) + if node_total is not None: + share = f"{100.0 * float(node_total) / total:.0f}%" if total else "-" + grid.add_row(Text(prefix + str(key), style="white"), f"{_format_value(node_total)} ms", share) + append(children, level + 1) + else: + share = f"{100.0 * float(value) / total:.0f}%" if total else "-" + grid.add_row(Text(prefix + str(key), style="white"), f"{_format_value(value)} ms", share) - env_metrics = stats.env_metrics - if env_metrics: - metric_grid = _rich_key_value_grid(dict(sorted(env_metrics.items())), signed=True) - blocks += [Rule("metrics", style="grey37", align="left"), metric_grid] - - if stats.timing_groups: - # Tree panels: one umbrella section, per-process groups as bold - # sub-titles (titles may carry markup for value highlighting). - blocks += [Rule("timing", style="grey37", align="left")] - for group, items in stats.timing_groups.items(): - blocks += [f"[bold]{group}[/]", _rich_timing_grid(items)] - else: - timing = Table.grid(expand=True, padding=(0, 2)) - timing_cells = [ - f"collect [yellow]{stats.collect_ms:.1f}[/]ms", - f"learn [magenta]{stats.learn_ms:.1f}[/]ms", - ] - for _ in timing_cells: - timing.add_column(ratio=1) - timing.add_row(*timing_cells) - blocks += [Rule("timing", style="grey37", align="left"), timing] - if stats.timing_metrics: - blocks.append(_rich_key_value_grid(stats.timing_metrics)) - if stats.diagnostics: - blocks += [Rule("diagnostics", style="grey37", align="left"), _rich_key_value_grid(stats.diagnostics)] + append(items) + return grid - return Panel(Group(*blocks), title=f"[bold]{title}[/]", border_style="cyan", padding=(0, 1)) +def _prototype_bar(fraction: float, *, width: int = 22, style: str = "cyan"): + from rich.text import Text -def _rich_timing_grid(items: Mapping[str, Any]): - """Two-column key/value grid; nested mappings render as indented sub-blocks.""" - grid = Table.grid(expand=True, padding=(0, 3)) - grid.add_column(ratio=1) - grid.add_column(ratio=1) - rows: list[list[str]] = [] - pending: list[str] = [] + filled = max(0, min(width, round(width * fraction))) + result = Text("━" * filled, style=style) + result.append("━" * (width - filled), style="grey23") + return result - def flush() -> None: - while pending: - row = pending[:2] - del pending[:2] - rows.append(row + [""] * (2 - len(row))) - entries = list(items.items()) - for idx, (key, value) in enumerate(entries): - if isinstance(value, Mapping): - flush() - branch = "└─" if idx == len(entries) - 1 else "├─" - children = dict(value) - node_total = children.pop("total", None) - tail = f" {_format_value(node_total)}" if node_total is not None else "" - rows.append([f"[dim]{branch} {key}[/]{tail}", ""]) - pending.extend(_rich_key_value_cell(f" {child}", child_value) for child, child_value in children.items()) - else: - pending.append(_rich_key_value_cell(key, value)) - flush() - for row in rows: - grid.add_row(*row) - return grid +def _prototype_tabs(stats: TrainingPanelStats, detail: bool): + from rich.text import Text + tabs = Table.grid(expand=True, padding=(0, 2)) + tabs.add_column() + tabs.add_column() + tabs.add_column(ratio=1, justify="right") + active = "Timing" if detail else "Overview" + labels = (("Overview", ""), ("Timing", "")) + row = [] + for label, count in labels: + item = Text(label, style="bold cyan" if label == active else "dim") + if count: + item.append(f" {count}", style="grey50") + row.append(item) + if _POSIX_TTY: # key handling is POSIX-only; don't advertise it elsewhere + row.append(Text("keyboard: 1/2 switch tabs", style="dim")) + tabs.add_row(*row) + return tabs -def _rich_key_value_grid(items: Mapping[str, Any], *, precision: int = 3, signed: bool = False): - grid = Table.grid(expand=True, padding=(0, 3)) - grid.add_column(ratio=1) - grid.add_column(ratio=1) - pairs = list(items.items()) - for i in range(0, len(pairs), 2): - left = _rich_key_value_cell(*pairs[i], precision=precision, signed=signed) - right = _rich_key_value_cell(*pairs[i + 1], precision=precision, signed=signed) if i + 1 < len(pairs) else "" - grid.add_row(left, right) - return grid +def _format_memory(memory: MemoryUsage | None) -> str: + if memory is None: + return "n/a" + gib = 1024**3 + return f"{memory.used_bytes / gib:.1f}/{memory.total_bytes / gib:.1f} GiB" -def _rich_key_value_cell(key: str, value: Any, *, precision: int = 3, signed: bool = False) -> str: - color = "" - if signed and isinstance(value, (float, int)): - color = "green" if value >= 0 else "red" - formatted = _format_value(value, precision=precision, signed=signed) - return f"{key:<22}[{color}]{formatted}[/]" if color else f"{key:<22}{formatted}" + +def render_training_panel(stats: TrainingPanelStats, *, title: str = "rl", detail: bool = False): + if not _RICH: + raise RuntimeError("rich is not available") + from rich.text import Text + + progress = max(0.0, min(1.0, stats.iteration / max(stats.total_iterations, 1))) + collect, learn = _timing_totals(stats) + + def card(label: str, body: Any, footer: Any = "") -> Panel: + parts: list[Any] = [body] + if footer: + parts.append(footer) + return Panel(Group(*parts), title=label, border_style="grey37", padding=(0, 1)) + + # Overview keeps the five operator-facing cards on one row. Replay-buffer + # occupancy remains available in the training/metrics data, but is not a + # headline card. + summary = Table.grid(expand=True, padding=(0, 1)) + # System health contains two side-by-side stacks (CPU/GPU and RAM/VRAM), + # so give it 1.5x the width of the other cards instead of forcing each + # value into a narrow half-column. + for ratio in (2, 2, 2, 2, 3): + summary.add_column(ratio=ratio) + load = stats.cpu_load + + def load_style(value: float | None) -> str: + if value is None: + return "dim" + return "green" if value < 70.0 else "yellow" if value < 90.0 else "red" + + def memory_style(memory: MemoryUsage | None) -> str: + if memory is None or memory.total_bytes <= 0: + return "dim" + ratio = memory.used_bytes / memory.total_bytes + return "green" if ratio < 0.70 else "yellow" if ratio < 0.90 else "red" + + cpu_text = f"CPU {load.utilization_percent:.0f}%" if load is not None else "CPU n/a" + progress_row = Table.grid(expand=True, padding=(0, 1)) + progress_row.add_column(ratio=1) + progress_row.add_row(_prototype_bar(progress, width=18)) + system_health = Table.grid(expand=True, padding=(0, 1)) + system_health.add_column(ratio=1) + system_health.add_column(ratio=1) + system_health.add_row( + Group( + Text(cpu_text, style=f"bold {load_style(load.utilization_percent if load else None)}"), + Text( + f"GPU {stats.gpu_utilization_percent:.0f}%" if stats.gpu_utilization_percent is not None else "GPU n/a", + style=f"{load_style(stats.gpu_utilization_percent)}", + ), + ), + Group( + Text(f"RAM {_format_memory(stats.memory_usage)}", style=memory_style(stats.memory_usage)), + Text(f"VRAM {_format_memory(stats.gpu_memory_usage)}", style=memory_style(stats.gpu_memory_usage)), + ), + ) + summary.add_row( + card( + f"Run progress ({progress * 100:.1f}%)", + Group(Text(f"{stats.iteration:,}/{stats.total_iterations:,} iters", style="white"), progress_row), + ), + card( + "Episode stats", + Group( + Text(f"return {stats.mean_return:+.2f}", style="bold green"), + Text(f"length {stats.mean_episode_length:.1f}", style="white"), + ), + ), + card( + "Throughput", + Group( + Text(f"{stats.steps_per_second:,.0f} env-steps/s", style="bold cyan"), + Text(f"{stats.iteration / max(stats.elapsed_seconds, 1e-9):,.0f} iter/s", style="white"), + ), + ), + card( + "Timing", + Group( + Text(f"Collect {collect:.1f} ms", style="bold yellow"), + Text(f"Learn {learn:.1f} ms", style="magenta"), + ), + ), + card("System health", system_health), + ) + + # UTD (update-to-data ratio) describes learner/training efficiency, so + # keep it with algorithm metrics rather than environment observations. + utd_items = {k: v for k, v in stats.diagnostics.items() if str(k).strip().lower() == "utd"} + train_items = dict(stats.training_metrics or {}) + train_items.update(utd_items) + train_parts: list[Any] = [] + if train_items: + train_parts.append(_rich_compact_metric_grid(train_items, limit=None)) + else: + train_parts.append( + Text("warmup - filling replay buffer" if stats.warming else "training metrics unavailable", style="grey70") + ) + training = Panel(Group(*train_parts), title=f"Training ({len(train_items)})", border_style="grey37", padding=(0, 1)) + left_blocks: list[Any] = [training] + if stats.reward_terms: + rewards = dict(sorted(stats.reward_terms.items(), key=lambda kv: -abs(float(kv[1])))) + left_blocks.append( + Panel( + _rich_compact_metric_grid(rewards, limit=None, precision=4, signed=True), + title=f"Rewards ({len(rewards)})", + border_style="grey37", + padding=(0, 1), + ) + ) + + env_items = dict(sorted(stats.env_metrics.items())) + env_parts: list[Any] = [] + if env_items: + env_parts.append(_rich_compact_metric_grid(env_items, limit=None, signed=True)) + else: + env_parts.append(Text("no environment metrics reported", style="grey70")) + other_diagnostics = {k: v for k, v in stats.diagnostics.items() if str(k).strip().lower() != "utd"} + if other_diagnostics: + env_parts.append( + Text(" · ".join(f"{k} {_format_value(v)}" for k, v in other_diagnostics.items()), style="cyan") + ) + environment = Panel( + Group(*env_parts), title=f"Environment metrics ({len(env_items)})", border_style="grey37", padding=(0, 1) + ) + + if detail: + timing_blocks: list[Any] = [] + columns = Table.grid(expand=True, padding=(0, 1)) + columns.add_column(ratio=1) + columns.add_column(ratio=1) + groups = list(stats.timing_groups.items()) + for index in range(0, len(groups), 2): + row: list[Any] = [] + for offset in (0, 1): + if index + offset >= len(groups): + row.append("") + continue + group, items = groups[index + offset] + total = _timing_group_total(group, stats, index + offset) + row.append( + Panel( + _rich_timing_grid(items, total=total), + title=f"{_strip_markup(group)} {total:.1f} ms", + border_style="grey37", + padding=(0, 1), + ) + ) + columns.add_row(*row) + if groups: + timing_blocks.append(columns) + if stats.timing_metrics: + timing_blocks.append( + Panel( + _rich_compact_metric_grid(stats.timing_metrics, limit=None), + title="Timing detail", + border_style="grey37", + padding=(0, 1), + ) + ) + if stats.diagnostics: + timing_blocks.append( + Panel( + _rich_compact_metric_grid(stats.diagnostics, limit=None, signed=True), + title="Diagnostics", + border_style="grey37", + padding=(0, 1), + ) + ) + lower = Group(*timing_blocks) if timing_blocks else Text("timing details unavailable", style="grey70") + else: + lower = Group(*left_blocks, environment) + blocks = [summary, lower, _prototype_tabs(stats, detail)] + if stats.checkpoint_path: + blocks.insert(-1, Text(f"✓ saved checkpoint {stats.checkpoint_path}", style="green")) + # Let Rich use the actual terminal width. The card bodies are intentionally + # single-line so wide terminals gain space without introducing vertical gaps. + return Panel(Group(*blocks), title=_strip_markup(title), border_style="cyan", padding=(0, 1)) diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py index cd6b5c3..c88c407 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/train.py @@ -198,6 +198,7 @@ def _drain_child_errors() -> list[tuple[str, str]]: logging_interval, save_interval, str(self._context.run_dir), + self._env_name, str(self._context.checkpoint_dir), self._context.checkpoint_format, self._resume_from, diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py index c330e89..deee8c7 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py @@ -39,7 +39,7 @@ from motrix_rl.fastsac.wrap import FastSacEnvWrap from motrix_rl.fastsac.wrap_np import FastSacNpEnvWrap from motrix_rl.fastsac.wrap_torch import FastSacTorchEnvWrap -from motrix_rl.system_metrics import CpuLoadSampler +from motrix_rl.system_metrics import CpuLoadSampler, GpuMemoryUsageSampler, GpuUtilizationSampler, MemoryUsageSampler def _timing_mean(values: list[float]) -> float: @@ -190,6 +190,7 @@ def run_learner_process( logging_interval: int, save_interval: int, run_dir: str, + env_name: str, checkpoint_dir: str, checkpoint_format: str, resume_from: str | None, @@ -233,12 +234,17 @@ def run_learner_process( last_metrics = None next_log = ((resume_step // logging_interval) + 1) * logging_interval if logging_interval > 0 else 0 next_save = ((resume_step // save_interval) + 1) * save_interval if save_interval > 0 else 0 - t_learn_win = 0.0 # wall-clock spent in gradient updates this log window + t_learn_win = 0.0 # wall-clock spent in learner train calls this log window + learner_train_samples_ms: list[float] = [] learner_drain_samples_ms: list[float] = [] learner_breakdown_samples_ms: dict[str, list[float]] = {} learner_ring_wait_samples_ms: list[float] = [] learner_gate_wait_samples_ms: list[float] = [] cpu_sampler = CpuLoadSampler() + gpu_sampler = GpuUtilizationSampler() + memory_sampler = MemoryUsageSampler() + gpu_memory_sampler = GpuMemoryUsageSampler() + last_checkpoint_path: str | None = None def _drain_stats(): nonlocal last_stats @@ -259,6 +265,10 @@ def _drain_stats(): last_metrics = metrics elapsed_learn_s = time.perf_counter() - t_l t_learn_win += elapsed_learn_s + learner_train_samples_ms.append(elapsed_learn_s * 1000.0) + # Keep the raw total cost of one agent.update(n) call. The + # timing tree uses total-cost semantics, not per-gradient-step + # normalization. for key, value in learner.agent._last_update_timing_ms.items(): learner_breakdown_samples_ms.setdefault(key, []).append(value) learner_breakdown_samples_ms.setdefault("publish", []).append(learner._last_publish_ms) @@ -288,23 +298,22 @@ def _drain_stats(): # these do NOT sum to 100% like the sync panel): # timing_ms[collect] — collector's avg ms per env-step batch (from queue) # timing_ms[wait] — avg ring-backpressure wait per batch - # learn_ms — learner's avg ms per gradient update + # learn_ms — learner's avg ms per train call (one UTD + # execution, which may run several gradient + # updates); idle waits are not included # learn_pct — fraction of learner wall-clock spent updating vs # idle/starved (≈100% when GPU-bound, lower if the # collector can't keep the buffer fed) # Window means only: live-panel percentiles are noise at these # sample counts (benchmarks own the tail statistics). - updates_delta = updates - last_update_idx - learn_ms = t_learn_win * 1000.0 / max(updates_delta, 1) learn_pct = 100.0 * t_learn_win / max(now - last_log_time, 1e-9) collector_timing_ms = last_stats.get("timing_ms", {}) collector_timing_detail_ms = { key: value for key, value in collector_timing_ms.items() if key != "collect" } - # Panel tree is per-process; the headline per-batch/per-update - # means are folded into the group titles, sub-stages nest under - # "sync" / "update" branches. - collect_ms = collector_timing_ms.get("collect", 0.0) + # Panel tree is per-process; the headline collect/learn means + # live on TrainingPanelStats, sub-stages nest under "sync" / + # "update" branches. collector_items: dict[str, Any] = {} sync_items: dict[str, float] = {} for key, value in collector_timing_detail_ms.items(): @@ -314,20 +323,27 @@ def _drain_stats(): collector_items[key] = value if sync_items: collector_items["sync"] = {"total": collector_items.pop("sync", 0.0), **sync_items} - timing_groups = {f"collector [yellow]{collect_ms:.1f}[/]ms": collector_items} + timing_groups = {"collector": collector_items} learner_items: dict[str, Any] = {} + drain_ms = _timing_mean(learner_drain_samples_ms) if learner_drain_samples_ms else 0.0 + ring_wait_ms = _timing_mean(learner_ring_wait_samples_ms) if learner_ring_wait_samples_ms else 0.0 + gate_wait_ms = _timing_mean(learner_gate_wait_samples_ms) if learner_gate_wait_samples_ms else 0.0 if learner_drain_samples_ms: - learner_items["drain"] = _timing_mean(learner_drain_samples_ms) + learner_items["drain"] = drain_ms if learner_ring_wait_samples_ms: - learner_items["ring wait"] = _timing_mean(learner_ring_wait_samples_ms) + learner_items["ring wait"] = ring_wait_ms if learner_gate_wait_samples_ms: - learner_items["gate wait"] = _timing_mean(learner_gate_wait_samples_ms) - if learner_breakdown_samples_ms: - learner_items["update"] = { - key: _timing_mean(values) for key, values in learner_breakdown_samples_ms.items() - } + learner_items["gate wait"] = gate_wait_ms + update_items = {key: _timing_mean(values) for key, values in learner_breakdown_samples_ms.items()} + if update_items: + # publish is a child stage of the learner update in the + # panel, so include it in the displayed update total too. + if "publish" in update_items and "total" in update_items: + update_items["total"] += update_items["publish"] + learner_items["update"] = update_items if learner_items: - timing_groups[f"learner [magenta]{learn_ms:.1f}[/]ms"] = learner_items + timing_groups["learner"] = learner_items + learn_ms = _timing_mean(learner_train_samples_ms) if learner_train_samples_ms else 0.0 stats = TrainingPanelStats( iteration=step, total_iterations=num_iterations, @@ -348,8 +364,12 @@ def _drain_stats(): timing_groups=timing_groups, diagnostics={"UTD": utd}, cpu_load=cpu_sampler.sample(), + gpu_utilization_percent=gpu_sampler.sample(), + memory_usage=memory_sampler.sample(), + gpu_memory_usage=gpu_memory_sampler.sample(), + checkpoint_path=last_checkpoint_path, ) - emit_training_panel(live, stats, title="motrix.fastsac (async)") + emit_training_panel(live, stats, title=f"{env_name}/motrix.fastsac") if writer is not None: writer.add_scalar("rollout/mean_return", last_stats["return"], step) writer.add_scalar("rollout/mean_ep_len", last_stats["ep_len"], step) @@ -364,7 +384,7 @@ def _drain_stats(): writer.add_scalar("perf/collect_ms_per_batch", collector_timing_ms.get("collect", 0.0), step) for k, v in collector_timing_detail_ms.items(): writer.add_scalar(f"perf/collector_{k}_ms", v, step) - writer.add_scalar("perf/learn_ms_per_update", learn_ms, step) + writer.add_scalar("perf/learn_ms_total", learn_ms, step) writer.add_scalar("perf/learn_pct", learn_pct, step) for k, v in last_stats["env_metrics"].items(): writer.add_scalar(f"metrics/{k}", v, step) @@ -375,6 +395,7 @@ def _drain_stats(): writer.add_scalar(f"train/{k}", v, step) last_log_time, last_log_step, last_update_idx = now, step, updates t_learn_win = 0.0 + learner_train_samples_ms = [] learner_drain_samples_ms = [] learner_breakdown_samples_ms = {} learner_ring_wait_samples_ms = [] @@ -392,7 +413,10 @@ def _drain_stats(): checkpoints.TRAINING_STATE, checkpoint_format=checkpoint_format, ) - (console.print if console else print)(f"[motrix.fastsac async] saved checkpoint {path}") + if console is not None: + last_checkpoint_path = str(path) + else: + print(f"[motrix.fastsac async] saved checkpoint {path}") next_save += save_interval # final checkpoint (identical structure to sync fastsac) diff --git a/motrix_rl/src/motrix_rl/fastsac/sync/train.py b/motrix_rl/src/motrix_rl/fastsac/sync/train.py index 14a2e40..5189c95 100644 --- a/motrix_rl/src/motrix_rl/fastsac/sync/train.py +++ b/motrix_rl/src/motrix_rl/fastsac/sync/train.py @@ -20,7 +20,7 @@ from motrix_rl.fastsac.wrap_np import FastSacNpEnvWrap from motrix_rl.fastsac.wrap_torch import FastSacTorchEnvWrap from motrix_rl.frameworks import TrainerBase, TrainerContext -from motrix_rl.system_metrics import CpuLoadSampler +from motrix_rl.system_metrics import CpuLoadSampler, GpuMemoryUsageSampler, GpuUtilizationSampler, MemoryUsageSampler # Enable TF32 matmul on Ampere+ GPUs. SAC training has no precision concern with # TF32 (10 mantissa bits), and the speedup is meaningful when AMP is off. @@ -215,6 +215,10 @@ def _run_loop( t_collect = 0.0 t_learn = 0.0 cpu_sampler = CpuLoadSampler() + gpu_sampler = GpuUtilizationSampler() + memory_sampler = MemoryUsageSampler() + gpu_memory_sampler = GpuMemoryUsageSampler() + last_checkpoint_path: str | None = None # optional live console that refreshes one panel in place console, live = open_training_live() @@ -321,8 +325,12 @@ def emit_msg(msg: str) -> None: reward_terms=term_means, env_metrics=env_metrics, cpu_load=cpu_sampler.sample(), + gpu_utilization_percent=gpu_sampler.sample(), + memory_usage=memory_sampler.sample(), + gpu_memory_usage=gpu_memory_sampler.sample(), + checkpoint_path=last_checkpoint_path, ) - emit_training_panel(live, stats, title="motrix.fastsac (sync)") + emit_training_panel(live, stats, title=f"{self._env_name}/motrix.fastsac") if self._writer is not None: self._writer.add_scalar("rollout/mean_return", mean_ret, agent.global_step) self._writer.add_scalar("rollout/mean_ep_len", mean_len, agent.global_step) @@ -347,9 +355,10 @@ def emit_msg(msg: str) -> None: path = Path(self._context.checkpoint_dir) / f"model_{agent.global_step:07d}.pt" torch.save(agent.state_dict(), path) record_checkpoint(path) - emit_msg( - f"[green]✓[/] saved checkpoint [dim]{path}[/]" if console else f"saved checkpoint {path}" - ) + if console is not None: + last_checkpoint_path = str(path) + else: + emit_msg(f"saved checkpoint {path}") local += 1 agent.global_step += 1 diff --git a/motrix_rl/src/motrix_rl/system_metrics.py b/motrix_rl/src/motrix_rl/system_metrics.py index 28d3fe1..280643c 100644 --- a/motrix_rl/src/motrix_rl/system_metrics.py +++ b/motrix_rl/src/motrix_rl/system_metrics.py @@ -1,13 +1,22 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""Low-overhead host metrics sampled at training-panel refresh boundaries.""" +"""Low-overhead host metrics sampled at training-panel refresh boundaries. + +CPU samplers read Linux ``/proc`` interfaces and return ``None`` where they +are unavailable, so panels degrade to ``n/a`` fields; memory sampling also +supports Windows via ``GlobalMemoryStatusEx``. The GPU samplers use NVML, +which works on any platform with an NVIDIA driver. +""" from __future__ import annotations +import ctypes import os +import sys from dataclasses import dataclass from pathlib import Path +from typing import Any @dataclass(frozen=True) @@ -121,3 +130,116 @@ def _read_physical_core_count(self) -> int | None: except (OSError, ValueError): return None return len(cores) + + +@dataclass(frozen=True) +class MemoryUsage: + """Memory usage in bytes for a host or accelerator device.""" + + used_bytes: int + total_bytes: int + + +class _MemoryStatusEx(ctypes.Structure): + """``MEMORYSTATUSEX`` layout for the Windows ``GlobalMemoryStatusEx`` call.""" + + _fields_ = [ + ("dwLength", ctypes.c_ulong), + ("dwMemoryLoad", ctypes.c_ulong), + ("ullTotalPhys", ctypes.c_ulonglong), + ("ullAvailPhys", ctypes.c_ulonglong), + ("ullTotalPageFile", ctypes.c_ulonglong), + ("ullAvailPageFile", ctypes.c_ulonglong), + ("ullTotalVirtual", ctypes.c_ulonglong), + ("ullAvailVirtual", ctypes.c_ulonglong), + ("ullAvailExtendedVirtual", ctypes.c_ulonglong), + ] + + +def _windows_memory_status() -> MemoryUsage | None: + """Physical memory usage via ``GlobalMemoryStatusEx``, mirroring the /proc semantics.""" + status = _MemoryStatusEx() + status.dwLength = ctypes.sizeof(_MemoryStatusEx) + if not ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)): + return None + return MemoryUsage(used_bytes=status.ullTotalPhys - status.ullAvailPhys, total_bytes=status.ullTotalPhys) + + +class MemoryUsageSampler: + """Read host memory usage from Linux ``/proc/meminfo`` or the Windows memory API.""" + + def __init__(self, *, meminfo_path: str | Path = "/proc/meminfo") -> None: + self._meminfo_path = Path(meminfo_path) + + def sample(self) -> MemoryUsage | None: + if sys.platform == "win32": + return _windows_memory_status() + try: + values: dict[str, int] = {} + for line in self._meminfo_path.read_text().splitlines(): + key, value, *_ = line.split() + if key in {"MemTotal:", "MemAvailable:"}: + values[key] = int(value) * 1024 + except (OSError, ValueError): + return None + total = values.get("MemTotal:") + available = values.get("MemAvailable:") + if total is None or available is None: + return None + return MemoryUsage(used_bytes=max(0, total - available), total_bytes=total) + + +# NVML reads the same counters nvidia-smi reports, but in-process at +# microsecond cost instead of a subprocess spawn per query. +_nvml_state: tuple[Any, list[Any]] | tuple[()] | None = None # None: untried; (): unavailable + + +def _nvml() -> tuple[Any, list[Any]] | None: + """Lazily initialize NVML and return ``(module, device_handles)``, or ``None``.""" + global _nvml_state + if _nvml_state is None: + try: + import pynvml + + pynvml.nvmlInit() + devices = [pynvml.nvmlDeviceGetHandleByIndex(index) for index in range(pynvml.nvmlDeviceGetCount())] + _nvml_state = (pynvml, devices) + except Exception: # ImportError (pynvml missing) or NVML init failure (no driver/GPU) + _nvml_state = () + return _nvml_state or None + + +class GpuMemoryUsageSampler: + """Read aggregate NVIDIA memory usage across all visible GPUs via NVML.""" + + def sample(self) -> MemoryUsage | None: + session = _nvml() + if session is None: + return None + pynvml, devices = session + used = total = 0 + try: + for device in devices: + info = pynvml.nvmlDeviceGetMemoryInfo(device) + used += info.used + total += info.total + except pynvml.NVMLError: + return None + return MemoryUsage(used_bytes=used, total_bytes=total) if total > 0 else None + + +class GpuUtilizationSampler: + """Read aggregate NVIDIA GPU utilization across all visible GPUs via NVML.""" + + def sample(self) -> float | None: + session = _nvml() + if session is None: + return None + pynvml, devices = session + values: list[int] = [] + try: + for device in devices: + values.append(pynvml.nvmlDeviceGetUtilizationRates(device).gpu) + except pynvml.NVMLError: + return None + return sum(values) / len(values) if values else None diff --git a/motrix_rl/tests/test_console.py b/motrix_rl/tests/test_console.py index 09542b9..aef5a6a 100644 --- a/motrix_rl/tests/test_console.py +++ b/motrix_rl/tests/test_console.py @@ -1,8 +1,24 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -from motrix_rl.console import TrainingPanelStats, _format_value, format_training_panel -from motrix_rl.system_metrics import CpuLoad +import importlib +import sys +from collections.abc import Iterator +from typing import Any + +import pytest +from rich.console import Console + +import motrix_rl.console as console_module +from motrix_rl.console import ( + TrainingPanelStats, + _compact_metric_value, + _format_memory, + _format_value, + format_training_panel, + render_training_panel, +) +from motrix_rl.system_metrics import CpuLoad, MemoryUsage def test_format_value_uses_fixed_notation_for_regular_floats() -> None: @@ -131,3 +147,157 @@ def test_format_training_panel_shows_cpu_load_with_explicit_topology() -> None: assert "cpu 50.0% (16.0/32T, 16C)" in panel assert "iowait 1.5%" in panel assert "steal 0.2%" in panel + + +def _panel_stats(**overrides: Any) -> TrainingPanelStats: + values: dict[str, Any] = dict( + iteration=2, + total_iterations=10, + steps_per_second=100.0, + elapsed_seconds=1.0, + mean_return=1.5, + mean_episode_length=8.0, + episodes=4, + buffer_size=512, + buffer_capacity=100_000, + collect_ms=2.0, + learn_ms=3.0, + learn_percent=50.0, + ) + values.update(overrides) + return TrainingPanelStats(**values) + + +def _render_panel(stats: TrainingPanelStats, *, detail: bool = False, width: int = 200) -> str: + console = Console(width=width) + with console.capture() as capture: + console.print(render_training_panel(stats, detail=detail)) + return capture.get() + + +def test_render_training_panel_overview_keeps_timing_tree_hidden() -> None: + stats = _panel_stats( + timing_groups={ + "collector": {"env_step": 1.0, "sync": {"total": 0.5, "weights": 0.25}}, + "learner": {"update": {"total": 3.0, "critic": 2.0}}, + }, + ) + + panel = _render_panel(stats) + + assert "Run progress" in panel + assert "Episode stats" in panel + assert "Throughput" in panel + assert "System health" in panel + assert "Training (" in panel + assert "Environment metrics (" in panel + # per-stage timing belongs to the detail view only + assert "STAGE" not in panel + assert "env_step" not in panel + + +def test_render_training_panel_detail_view_shows_timing_tree_with_shares() -> None: + stats = _panel_stats( + timing_groups={ + "collector": {"env_step": 1.0, "sync": {"total": 0.5, "weights": 0.25}}, + "learner": {"update": {"total": 3.0, "critic": 2.0}}, + }, + timing_metrics={"queue_depth": 1.0}, + diagnostics={"UTD": 0.5}, + ) + + overview = _render_panel(stats) + detail = _render_panel(stats, detail=True) + + assert "STAGE" not in overview + assert "STAGE" in detail + assert "collector" in detail + assert "learner" in detail + assert "env_step" in detail + assert "Timing detail" in detail + assert "Diagnostics" in detail + # known group totals render a share column + assert "%" in detail + + +def test_render_training_panel_shows_saved_checkpoint_path() -> None: + panel = _render_panel(_panel_stats(checkpoint_path="/runs/cartpole/model_0000002.pt")) + + assert "saved checkpoint" in panel + assert "/runs/cartpole/model_0000002.pt" in panel + + +def test_render_training_panel_reports_system_health() -> None: + stats = _panel_stats( + gpu_utilization_percent=85.0, + memory_usage=MemoryUsage(used_bytes=1024**3, total_bytes=2 * 1024**3), + ) + + panel = _render_panel(stats) + + assert "CPU" in panel + assert "GPU 85%" in panel + assert "RAM 1.0/2.0 GiB" in panel + assert "VRAM n/a" in panel + + +def test_format_memory_renders_gib_and_missing_values() -> None: + assert _format_memory(None) == "n/a" + assert _format_memory(MemoryUsage(used_bytes=1024**3, total_bytes=4 * 1024**3)) == "1.0/4.0 GiB" + + +def test_compact_metric_value_fits_the_nine_char_value_cell() -> None: + assert _compact_metric_value(0.4, precision=4, signed=True) == "+0.4000" + # precision-4 magnitudes below 1e-4 would render as "-3.6200e-05" + assert _compact_metric_value(-3.62e-05, precision=4, signed=True) == "-3.62e-05" + # precision-3 values already fit until signed ("+1.235e-04" is 10 chars) + assert _compact_metric_value(1.235e-04, precision=3) == "1.235e-04" + assert _compact_metric_value(4.56e-04, precision=3, signed=True) == "+4.56e-04" + for value in (-3.62e-05, -0.0123, 0.4, 1.2345, 999999.5, 12_345_678.9, 123456789): + assert len(_compact_metric_value(value, precision=4, signed=True)) <= 9 + + +def test_render_training_panel_keeps_metric_values_on_their_label_line(monkeypatch) -> None: + stats = _panel_stats( + reward_terms={"torque": -3.62e-05, "action_rate": -0.0123}, + training_metrics={"alpha_loss": 1.235e-04}, + ) + + for width in (120, 150, 190): + monkeypatch.setenv("COLUMNS", str(width)) + lines = _render_panel(stats, width=width).splitlines() + assert any("torque" in line and "e-05" in line for line in lines), f"wrapped at width {width}" + assert any("action_rate" in line and "-0.0123" in line for line in lines), f"wrapped at width {width}" + assert any("alpha_loss" in line and "e-04" in line for line in lines), f"wrapped at width {width}" + + +def test_render_training_panel_only_advertises_keyboard_on_posix_tty(monkeypatch) -> None: + # 1/2 key handling needs a POSIX TTY; other platforms get a plain Live + stats = _panel_stats() + monkeypatch.setattr(console_module, "_POSIX_TTY", True) + assert "keyboard: 1/2 switch tabs" in _render_panel(stats) + monkeypatch.setattr(console_module, "_POSIX_TTY", False) + assert "keyboard" not in _render_panel(stats) + + +@pytest.fixture +def _reload_console_module() -> Iterator[None]: + """Reload the console module after the test, once monkeypatch undid import blocks. + + Listed before ``monkeypatch`` in the test signature so this teardown runs last. + """ + yield + importlib.reload(console_module) + + +def test_console_module_degrades_without_posix_tty_support(_reload_console_module, monkeypatch) -> None: + # Windows has no termios/tty; hidden modules raise ImportError on import. + for name in ("termios", "tty", "select"): + monkeypatch.setitem(sys.modules, name, None) + reloaded = importlib.reload(console_module) + + assert reloaded._POSIX_TTY is False + # rendering and the plain-text fallback stay fully functional + stats = _panel_stats() + assert "Run progress" in _render_panel(stats) + assert "iter" in reloaded.format_training_panel(stats) diff --git a/motrix_rl/tests/test_system_metrics.py b/motrix_rl/tests/test_system_metrics.py index 19429f2..b2387de 100644 --- a/motrix_rl/tests/test_system_metrics.py +++ b/motrix_rl/tests/test_system_metrics.py @@ -1,7 +1,18 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -from motrix_rl.system_metrics import CpuLoadSampler +import ctypes +import sys +import types + +import motrix_rl.system_metrics as system_metrics +from motrix_rl.system_metrics import ( + CpuLoadSampler, + GpuMemoryUsageSampler, + GpuUtilizationSampler, + MemoryUsage, + MemoryUsageSampler, +) def test_cpu_load_sampler_uses_counter_deltas_for_available_cpus(tmp_path) -> None: @@ -34,3 +45,109 @@ def test_cpu_load_sampler_returns_none_without_elapsed_cpu_time(tmp_path) -> Non stat_path.write_text(contents) assert sampler.sample() is None + + +def _fake_nvml(monkeypatch, handles, utilization, memory, error=None) -> None: + """Install a fake ``(pynvml, handles)`` session with per-handle metric tables.""" + + def utilization_rates(handle): + if error is not None: + raise error + return types.SimpleNamespace(gpu=utilization[handle]) + + def memory_info(handle): + if error is not None: + raise error + return types.SimpleNamespace(used=memory[handle][0], total=memory[handle][1]) + + fake_pynvml = types.SimpleNamespace( + nvmlDeviceGetUtilizationRates=utilization_rates, + nvmlDeviceGetMemoryInfo=memory_info, + NVMLError=RuntimeError, + ) + monkeypatch.setattr(system_metrics, "_nvml_state", (fake_pynvml, handles)) + + +def test_gpu_samplers_aggregate_utilization_mean_and_memory_sum(monkeypatch) -> None: + handles = ["gpu0", "gpu1"] + _fake_nvml( + monkeypatch, + handles, + utilization={"gpu0": 10, "gpu1": 30}, + memory={"gpu0": (100 * 1024**2, 200 * 1024**2), "gpu1": (300 * 1024**2, 400 * 1024**2)}, + ) + + assert GpuUtilizationSampler().sample() == 20.0 + assert GpuMemoryUsageSampler().sample() == MemoryUsage(used_bytes=400 * 1024**2, total_bytes=600 * 1024**2) + + +def test_gpu_samplers_return_none_on_nvml_error(monkeypatch) -> None: + handles = ["gpu0"] + _fake_nvml( + monkeypatch, + handles, + utilization={"gpu0": 10}, + memory={"gpu0": (1, 2)}, + error=RuntimeError("driver failure"), + ) + + assert GpuUtilizationSampler().sample() is None + assert GpuMemoryUsageSampler().sample() is None + + +def test_gpu_samplers_return_none_without_nvml(monkeypatch) -> None: + monkeypatch.setattr(system_metrics, "_nvml_state", ()) + + assert GpuUtilizationSampler().sample() is None + assert GpuMemoryUsageSampler().sample() is None + + +def test_memory_usage_sampler_reads_proc_meminfo(tmp_path) -> None: + meminfo = tmp_path / "meminfo" + meminfo.write_text( + "MemTotal: 32768000 kB\nMemFree: 1024000 kB\nCached: 8192000 kB\n" + "MemAvailable: 16384000 kB\nSwapTotal: 0 kB\n" + ) + sampler = MemoryUsageSampler(meminfo_path=meminfo) + + assert sampler.sample() == MemoryUsage(used_bytes=(32768000 - 16384000) * 1024, total_bytes=32768000 * 1024) + + +def test_memory_usage_sampler_returns_none_when_meminfo_missing(tmp_path) -> None: + sampler = MemoryUsageSampler(meminfo_path=tmp_path / "missing") + + assert sampler.sample() is None + + +def test_memory_usage_sampler_dispatches_to_windows_api(monkeypatch) -> None: + windows_usage = MemoryUsage(used_bytes=7, total_bytes=9) + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(system_metrics, "_windows_memory_status", lambda: windows_usage) + + assert MemoryUsageSampler().sample() == windows_usage + + +def _patch_windows_memory_api(monkeypatch, *, succeed: bool) -> None: + def global_memory_status_ex(pointer) -> int: + status = ctypes.cast(pointer, ctypes.POINTER(system_metrics._MemoryStatusEx)).contents + assert status.dwLength == ctypes.sizeof(system_metrics._MemoryStatusEx) + if not succeed: + return 0 + status.ullTotalPhys = 100 + status.ullAvailPhys = 25 + return 1 + + kernel32 = types.SimpleNamespace(GlobalMemoryStatusEx=global_memory_status_ex) + monkeypatch.setattr(ctypes, "windll", types.SimpleNamespace(kernel32=kernel32), raising=False) + + +def test_windows_memory_status_maps_total_and_available_physical_memory(monkeypatch) -> None: + _patch_windows_memory_api(monkeypatch, succeed=True) + + assert system_metrics._windows_memory_status() == MemoryUsage(used_bytes=75, total_bytes=100) + + +def test_windows_memory_status_returns_none_when_api_reports_failure(monkeypatch) -> None: + _patch_windows_memory_api(monkeypatch, succeed=False) + + assert system_metrics._windows_memory_status() is None diff --git a/uv.lock b/uv.lock index 70eccc1..8809e2b 100644 --- a/uv.lock +++ b/uv.lock @@ -1090,6 +1090,7 @@ dependencies = [ { name = "motrix-env-core" }, { name = "motrix-env-motrixsim" }, { name = "numpy" }, + { name = "nvidia-ml-py" }, { name = "omegaconf" }, { name = "python-abc" }, { name = "rich" }, @@ -1142,6 +1143,7 @@ requires-dist = [ { name = "motrix-env-core", editable = "motrix_env_core" }, { name = "motrix-env-motrixsim", editable = "motrix_env_motrixsim" }, { name = "numpy", specifier = ">=1.26" }, + { name = "nvidia-ml-py", specifier = ">=13.610.43" }, { name = "omegaconf", specifier = ">=2.3,<2.4" }, { name = "onnx", marker = "extra == 'onnx'", specifier = "==1.20.1" }, { name = "onnx", marker = "extra == 'rslrl'", specifier = "==1.20.1" }, @@ -1466,6 +1468,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3b/9a/72ef35b399b0e183bc2e8f6f558036922d453c4d8237dab26c666a04244b/nvidia_cusparselt_cu12-0.6.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:e5c8a26c36445dd2e6812f1177978a24e2d37cacce7e090f297a688d1ec44f46", size = 156785796, upload-time = "2024-10-15T21:29:17.709Z" }, ] +[[package]] +name = "nvidia-ml-py" +version = "13.610.43" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f0/b5/a8fbc356f768fa5c9cfd646668fd7d34bf55bdd1c6e20754642a64d930d4/nvidia_ml_py-13.610.43.tar.gz", hash = "sha256:65437eb73d68d0c62c931ca4d45038472faff03bd0b8729abba4b899f70d60f2", size = 52109, upload-time = "2026-06-01T18:54:08.829Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/45/caa600acfab94560807a20a64b5830d2cd3c3202b7f1328644d70b7d6bd8/nvidia_ml_py-13.610.43-py3-none-any.whl", hash = "sha256:f13c72698edef492f985cc225f14faafe68ae065a2e407f45bdf6f4b9b43fde8", size = 53163, upload-time = "2026-06-01T18:54:07.704Z" }, +] + [[package]] name = "nvidia-nccl-cu12" version = "2.26.2"