diff --git a/bench/bench_live_metrics.py b/bench/bench_live_metrics.py new file mode 100644 index 00000000..bbda39b4 --- /dev/null +++ b/bench/bench_live_metrics.py @@ -0,0 +1,353 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark what live metrics cost a training loop, healthy and when its outputs fail. + +Every scenario runs the same synthetic training loop and logs through +``TrainingLogger``, the interface trainers use, with TensorBoard on in every +scenario: the difference to ``off`` (live metrics disabled) is the cost of live +metrics alone. Faulty scenarios break the jsonl disk (slow, stuck for good, +failing) or the ingest service (slow, never answering, refusing connections), +and one breaks both at once. + +The report gives the wall time of each ``scalars`` call on the training +thread, the loop's throughput against ``off`` (``off, control`` repeats ``off`` +and shows the noise floor), and what reached the disk and the service. Only the +training loop is measured: closing is not timed, and a faulty scenario closes +without waiting so it never sits out the close timeout (at most 5 s, once, at +the end of a real run; see wiki/design/live-metrics.md). The ingest service runs +in a separate process, like the real one, so its work does not compete for this +process's GIL. ``--workload python`` holds the GIL in the training loop itself, +the worst case for the sink threads. ``--max-overhead`` turns the run into a +check that exits non-zero when a scenario slows the loop down by more. +""" + +from __future__ import annotations + +import argparse +import gzip +import json +import logging +import multiprocessing +import socket +import statistics +import sys +import tempfile +import threading +import time +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import numpy as np +from report import render_table + +from motrix_rl.metrics import JSONL_FILENAME, HttpSink, JsonlSink, MetricsLogger +from motrix_rl.training_logger import TrainingLogger + +_HEADERS = [ + "scenario", + "loop it/s", + "overhead", + "emit p50 us", + "emit p99 us", + "emit max us", + "jsonl lines", + "delivered", +] + +# Disk faults, standing in for a slow, hung or broken run_dir. +_SLOW, _STUCK, _ERROR = "slow", "stuck", "error" +# Service faults: answers after a delay, accepts but never answers, refuses connections. +_OK, _SLOW_SERVICE, _SILENT, _DOWN = "ok", "slow", "silent", "down" + + +@dataclass(frozen=True, slots=True) +class _Scenario: + name: str + jsonl: str | None = None # None: off; "ok" or a disk fault + http: str | None = None # None: off; a service mode + + @property + def faulty(self) -> bool: + return self.jsonl not in (None, _OK) or self.http not in (None, _OK) + + +@dataclass(slots=True) +class _Result: + its_per_s: float + emit_ns: list[int] + jsonl_lines: int | None + delivered: int | None + + +SCENARIOS = [ + _Scenario("off"), + _Scenario("off, control"), # the same as "off": its overhead is the noise floor + _Scenario("jsonl", jsonl=_OK), + _Scenario("jsonl, slow disk", jsonl=_SLOW), + _Scenario("jsonl, stuck disk", jsonl=_STUCK), + _Scenario("jsonl, disk error", jsonl=_ERROR), + _Scenario("http", http=_OK), + _Scenario("http, slow service", http=_SLOW_SERVICE), + _Scenario("http, no response", http=_SILENT), + _Scenario("http, service down", http=_DOWN), + _Scenario("jsonl + http", jsonl=_OK, http=_OK), + _Scenario("stuck disk + no response", jsonl=_STUCK, http=_SILENT), +] + + +class _FaultyDisk(JsonlSink): + """A jsonl sink whose batch writes stall for *delay* seconds, or forever.""" + + def __init__(self, path: Path, delay: float | None) -> None: + self._delay = delay + self._never = threading.Event() # never set: a write that hangs for good + super().__init__(path) + + def _write_batch(self, batch): + if self._delay is None: + self._never.wait() + else: + time.sleep(self._delay) + super()._write_batch(batch) + + +def _serve(port_queue, delivered, delay: float | None) -> None: + """Ingest service for the child process: counts records, answers 200 after *delay* (never if None).""" + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + body = gzip.decompress(self.rfile.read(int(self.headers["Content-Length"]))) + if delay is None: + threading.Event().wait() # hold the connection open without ever answering + if delay: + time.sleep(delay) + with delivered.get_lock(): + delivered.value += len(json.loads(body)["records"]) + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + server.daemon_threads = True + port_queue.put(server.server_address[1]) + server.serve_forever() + + +class _Service: + def __init__(self, ctx, delay: float | None) -> None: + port_queue = ctx.Queue() + self.delivered = ctx.Value("q", 0) + self._process = ctx.Process(target=_serve, args=(port_queue, self.delivered, delay), daemon=True) + self._process.start() + self.url = f"http://127.0.0.1:{port_queue.get(timeout=30)}" + + def take_count(self) -> int: + with self.delivered.get_lock(): + count, self.delivered.value = self.delivered.value, 0 + return count + + def stop(self) -> None: + self._process.terminate() + self._process.join() + + +def _unreachable_url() -> str: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + return f"http://127.0.0.1:{port}" # nothing listens once the socket is closed + + +def _numpy_workload(step_ms: float, size: int): + """A numpy step of about *step_ms* milliseconds; BLAS releases the GIL, like most of a torch step.""" + rng = np.random.default_rng(0) + a = rng.standard_normal((size, size), dtype=np.float32) + b = rng.standard_normal((size, size), dtype=np.float32) + for _ in range(50): # warm up BLAS threads before timing + a @ b + + def once(): + a @ b + + return _repeat_to(once, step_ms) + + +def _python_workload(step_ms: float): + """A pure-Python step of about *step_ms* milliseconds that holds the GIL throughout.""" + + def once(): + total = 0 + for i in range(1000): + total += i * i + return total + + return _repeat_to(once, step_ms) + + +def _repeat_to(once, step_ms: float): + count = 0 + started = time.perf_counter() + while time.perf_counter() - started < 0.5: + once() + count += 1 + repeats = max(1, round(step_ms / 1000 / ((time.perf_counter() - started) / count))) + + def step(): + for _ in range(repeats): + once() + + return step + + +def _live(scenario: _Scenario, run_dir: Path, args, services: dict[str, _Service], down_url: str) -> MetricsLogger: + jsonl = None + path = run_dir / JSONL_FILENAME + if scenario.jsonl == _OK: + jsonl = JsonlSink(path) + elif scenario.jsonl == _SLOW: + jsonl = _FaultyDisk(path, args.disk_delay) + elif scenario.jsonl == _STUCK: + jsonl = _FaultyDisk(path, None) + elif scenario.jsonl == _ERROR: + (run_dir / "not-a-dir").write_text("") + jsonl = JsonlSink(run_dir / "not-a-dir" / JSONL_FILENAME) # every write fails with NotADirectoryError + http = None + if scenario.http is not None: + http = HttpSink(down_url if scenario.http == _DOWN else services[scenario.http].url) + return MetricsLogger("bench", jsonl=jsonl, http=http) if jsonl or http else MetricsLogger.disabled() + + +def _run(scenario: _Scenario, args, workload, services: dict[str, _Service], down_url: str) -> _Result: + with tempfile.TemporaryDirectory() as tmp: + run_dir = Path(tmp) + live = _live(scenario, run_dir, args, services, down_url) + log = TrainingLogger(run_dir, live) + values = {f"train/key_{index}": float(index) for index in range(args.keys)} + sys_samples = {"sys/cpu_percent": 50.0, "sys/mem_used_gb": 8.0} + log.scalars({}, 0, env_steps=0) # opens TensorBoard (a one-off startup cost) before timing; sends nothing + + emit_ns = [] + started = time.perf_counter() + for iteration in range(args.iters): + workload() + if iteration % args.log_every == 0: + values["train/key_0"] = float(iteration) + emit_started = time.perf_counter_ns() + log.scalars(values, iteration, env_steps=iteration * 4096, live_extra=lambda: sys_samples) + emit_ns.append(time.perf_counter_ns() - emit_started) + elapsed = time.perf_counter() - started + + # Not timed. Healthy sinks flush in milliseconds; a faulty one would only + # sit out the close timeout, so it is left behind instead. + log.close() + live.close() + if not scenario.faulty: + live.join(5) + path = run_dir / JSONL_FILENAME + lines = len(path.read_text().splitlines()) if scenario.jsonl and path.exists() else None + if scenario.jsonl is not None and lines is None: + lines = 0 + delivered = None + if scenario.http in services: + delivered = services[scenario.http].take_count() + elif scenario.http == _DOWN: + delivered = 0 + return _Result(args.iters / elapsed, emit_ns, lines, delivered) + + +def _percentile(values: list[int], q: float) -> float: + ordered = sorted(values) + return ordered[min(len(ordered) - 1, int(q * len(ordered)))] / 1000 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--iters", type=int, default=2000, help="Training-loop iterations per run.") + parser.add_argument("--step-ms", type=float, default=2.0, help="Target duration of one workload step.") + parser.add_argument("--workload", choices=["numpy", "python"], default="numpy", help="What a step spends time on.") + parser.add_argument("--matrix", type=int, default=256, help="Side of the square matmul in the numpy workload.") + parser.add_argument("--log-every", type=int, default=1, help="Iterations between scalars() calls (1 = worst case).") + parser.add_argument("--keys", type=int, default=40, help="Scalars per record.") + parser.add_argument("--repeats", type=int, default=5, help="Runs per scenario; throughput is the median.") + parser.add_argument("--disk-delay", type=float, default=0.5, help="Stall per batch for the slow-disk scenario.") + parser.add_argument("--service-delay", type=float, default=0.5, help="Reply delay for the slow-service scenario.") + parser.add_argument( + "--only", nargs="+", metavar="NAME", help="Run only scenarios whose name contains one of these." + ) + parser.add_argument( + "--max-overhead", + type=float, + default=None, + help="Fail (exit 1) when a scenario's median throughput is this many percent below 'off'.", + ) + args = parser.parse_args() + # The faulty scenarios warn once per sink by design; keep the report readable. + logging.getLogger("motrix_rl").setLevel(logging.ERROR) + + scenarios = SCENARIOS + if args.only: + scenarios = [s for s in SCENARIOS if s.name.startswith("off") or any(o in s.name for o in args.only)] + ctx = multiprocessing.get_context("spawn") + services = {_OK: _Service(ctx, 0.0), _SLOW_SERVICE: _Service(ctx, args.service_delay), _SILENT: _Service(ctx, None)} + down_url = _unreachable_url() + workload = ( + _numpy_workload(args.step_ms, args.matrix) if args.workload == "numpy" else _python_workload(args.step_ms) + ) + results: dict[str, list[_Result]] = {scenario.name: [] for scenario in scenarios} + try: + for repeat in range(args.repeats): + # Rotate the order so drift (thermal, other load) does not always hit the same scenario. + shift = repeat % len(scenarios) + for scenario in scenarios[shift:] + scenarios[:shift]: + results[scenario.name].append(_run(scenario, args, workload, services, down_url)) + finally: + for service in services.values(): + service.stop() + + baseline = statistics.median(result.its_per_s for result in results["off"]) + rows = [] + worst = 0.0 + for scenario in scenarios: + runs = results[scenario.name] + its = statistics.median(run.its_per_s for run in runs) + overhead = (1 - its / baseline) * 100 + if not scenario.name.startswith("off"): + worst = max(worst, overhead) + emits = [ns for run in runs for ns in run.emit_ns] + last = runs[-1] + rows.append( + [ + scenario.name, + f"{its:,.1f}", + "—" if scenario.name == "off" else f"{overhead:+.2f}%", + f"{_percentile(emits, 0.50):.1f}", + f"{_percentile(emits, 0.99):.1f}", + f"{max(emits) / 1000:.1f}", + "—" if last.jsonl_lines is None else str(last.jsonl_lines), + "—" if last.delivered is None else str(last.delivered), + ] + ) + + print(render_table(_HEADERS, rows)) + emitted = len(range(0, args.iters, args.log_every)) + print( + f"loop: workload={args.workload} iters={args.iters} step~{args.step_ms}ms log_every={args.log_every} " + f"keys={args.keys} repeats={args.repeats}; {emitted} records per run; TensorBoard on in every scenario" + ) + print(f"slow disk: {args.disk_delay}s stall per batch; slow service: {args.service_delay}s per reply") + print( + "overhead: median loop throughput against 'off', read against the 'off, control' noise floor; " + "jsonl lines / delivered: from the last run" + ) + if args.max_overhead is not None and worst > args.max_overhead: + print(f"FAIL: worst overhead {worst:.2f}% exceeds --max-overhead {args.max_overhead}%") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/motrix_rl/src/motrix_rl/contracts.py b/motrix_rl/src/motrix_rl/contracts.py index 90633324..004120a3 100644 --- a/motrix_rl/src/motrix_rl/contracts.py +++ b/motrix_rl/src/motrix_rl/contracts.py @@ -6,13 +6,17 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Iterable -from dataclasses import dataclass, is_dataclass +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field, is_dataclass from pathlib import Path from typing import Any, Generic, TypeVar +from motrix_rl.metrics import MetricsLogger + CfgT = TypeVar("CfgT") +_FAILURE_REASON_LIMIT = 2048 + class TrainerBase(ABC): """Runnable trainer contract implemented by RL integrations.""" @@ -44,6 +48,8 @@ class TrainerContext(Generic[CfgT]): checkpoint: Any render: Any = None resume_from: str | None = None + # Live metrics for this run; disabled unless the control plane enabled it for training. + metrics: MetricsLogger = field(default_factory=MetricsLogger.disabled) @dataclass(frozen=True) @@ -53,10 +59,33 @@ class TrainerHandle: run: Any trainer: TrainerBase result_factory: Any = None + metrics: MetricsLogger = field(default_factory=MetricsLogger.disabled) + # Extra fields of the ``running`` record (see wiki/design/live-metrics.md). + run_info: Mapping[str, Any] = field(default_factory=dict) def train(self): - """Run training and return the control plane's result view.""" - self.trainer.train() + """Run training and return the control plane's result view. + + Reports ``running``, then ``finished``, ``killed`` (Ctrl+C) or ``failed`` + to live metrics; exceptions propagate unchanged. Closing live metrics + does not wait: the sinks finish in the background. + """ + self.metrics.run("running", **self.run_info) + try: + self.trainer.train() + except KeyboardInterrupt: + self.metrics.run("killed") + raise + except BaseException as exc: + if isinstance(exc, SystemExit) and exc.code in (None, 0): + self.metrics.run("finished") + elif self.metrics.enabled: + self.metrics.run("failed", reason=_failure_reason(exc, Path(self.run.run_dir))) + raise + else: + self.metrics.run("finished") + finally: + self.metrics.close() if self.result_factory is None: return self.run return self.result_factory(self.run) @@ -66,6 +95,14 @@ def play(self, policy: str) -> None: self.trainer.play(policy) +def _failure_reason(exc: BaseException, run_dir: Path) -> str: + """``Type: message`` within 2 KB, pointing at any worker error logs the run left behind.""" + logs = sorted(path.relative_to(run_dir).as_posix() for path in (run_dir / "async_errors").glob("*.log")) + suffix = f" (see {', '.join(logs)})" if logs else "" + head = f"{type(exc).__name__}: {exc}"[: max(0, _FAILURE_REASON_LIMIT - len(suffix))] + return (head + suffix)[:_FAILURE_REASON_LIMIT] + + class AgentProvider(Generic[CfgT], ABC): """Provider for one algorithm and one train backend.""" @@ -102,6 +139,10 @@ def create_policy_exporter(self) -> Any: """Create an optional policy exporter owned by the control plane.""" return None + def total_env_steps(self, config: CfgT, num_envs: int) -> int | None: + """Environment steps a training run with *config* covers, for progress reporting.""" + return None + @dataclass(frozen=True) class AgentRegistration: diff --git a/motrix_rl/src/motrix_rl/frameworks.py b/motrix_rl/src/motrix_rl/frameworks.py index f0779c41..53675cde 100644 --- a/motrix_rl/src/motrix_rl/frameworks.py +++ b/motrix_rl/src/motrix_rl/frameworks.py @@ -24,6 +24,7 @@ TrainerContext, TrainerHandle, ) +from motrix_rl.metrics import MetricsLogger __all__ = [ "AgentProvider", @@ -156,6 +157,7 @@ def create_trainer(context: TrainerContext[Any]) -> TrainerHandle: run=context.run, trainer=provider.create_trainer(typed_context), result_factory=TrainResult, + metrics=typed_context.metrics, ) @@ -170,6 +172,7 @@ def create_trainer_context( checkpoint: CheckpointConfig, render: Any = None, resume_from: str | None = None, + metrics: MetricsLogger | None = None, ) -> TrainerContext[Any]: _validate_runtime_config(logging, checkpoint) metadata = run.metadata @@ -188,6 +191,7 @@ def create_trainer_context( checkpoint=checkpoint, render=render, resume_from=resume_from, + metrics=metrics or MetricsLogger.disabled(), ) diff --git a/motrix_rl/src/motrix_rl/metrics.py b/motrix_rl/src/motrix_rl/metrics.py new file mode 100644 index 00000000..e633835b --- /dev/null +++ b/motrix_rl/src/motrix_rl/metrics.py @@ -0,0 +1,515 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Live training metrics in the ``motrix.metrics/v1`` protocol. + +Setting ``MOTRIX_METRICS_URL`` makes a training run push its records to an +ingest service compatible with ``motrix.metrics/v1``; setting +``MOTRIX_METRICS_JSONL=1`` also (or instead) appends them to +``run_dir/metrics.jsonl``. With neither the logger is disabled: no file, no +thread. Both sinks write from daemon threads, so neither a slow service nor a +slow disk blocks or fails training. See ``wiki/design/live-metrics.md``. +""" + +from __future__ import annotations + +import atexit +import collections +import gzip +import json +import logging +import math +import os +import re +import socket +import threading +import time +import urllib.error +import urllib.request +import weakref +from collections.abc import Callable, Mapping +from http.client import HTTPException +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + +SCHEMA = "motrix.metrics/v1" +JSONL_FILENAME = "metrics.jsonl" + +_TRUTHY = {"1", "true", "yes", "on"} + +_RUN_ID_MAX = 200 +_RUN_ID_INVALID = re.compile(r"[^A-Za-z0-9._-]") + +# How long the interpreter waits at exit for sinks still writing or sending (see _join_at_exit). +EXIT_TIMEOUT = 5.0 +_OPEN: weakref.WeakSet[MetricsLogger] = weakref.WeakSet() + + +class MetricsLogger: + """Emits ``motrix.metrics/v1`` records to a jsonl file and an ingest service. + + Every record gets one ``seq`` and goes to every sink, so a jsonl written + next to the service can later backfill it record for record. Holds a thread and locks: it cannot + be pickled into worker processes. + """ + + def __init__( + self, + run_id: str, + *, + role: str = "trainer", + jsonl: JsonlSink | None = None, + http: HttpSink | None = None, + project: str | None = None, + group: str | None = None, + meta: Mapping[str, Any] | None = None, + ) -> None: + self.run_id = run_id + self.writer = f"{socket.gethostname()}:{os.getpid()}:{role}" + self._jsonl = jsonl + self._http = http + self._identity = {"project": project, "group": group, "meta": dict(meta or {})} + # Starts at the wall clock rather than 0: a restarted container can reuse + # hostname and pid, and the service drops records by (writer, seq). + self._seq = time.time_ns() // 1000 + self._seq_lock = threading.Lock() + self._emit_lock = threading.Lock() + self._nonfinite_keys: set[str] = set() + self._invalid_keys: set[str] = set() + self._closed = False + if http is not None: + http.start(lambda: self._record("heartbeat", {})) + if self.enabled: + _OPEN.add(self) + + @classmethod + def disabled(cls) -> MetricsLogger: + """A logger whose methods do nothing.""" + return cls("disabled") + + @classmethod + def from_env(cls, run_dir: str | Path, *, role: str = "trainer") -> MetricsLogger: + """Build the logger the ``MOTRIX_METRICS_*`` environment asks for.""" + url = os.environ.get("MOTRIX_METRICS_URL", "") + jsonl = os.environ.get("MOTRIX_METRICS_JSONL", "").strip().lower() in _TRUTHY + if not url and not jsonl: + return cls.disabled() + run_dir = Path(run_dir) + meta: dict[str, Any] = {} + raw_meta = os.environ.get("MOTRIX_METRICS_META", "") + if raw_meta: + try: + meta = dict(json.loads(raw_meta)) + except (ValueError, TypeError): + logger.warning("metrics: ignoring MOTRIX_METRICS_META, not a JSON object") + return cls( + sanitize_run_id(os.environ.get("MOTRIX_METRICS_RUN_ID") or "__".join(run_dir.parts[-5:])), + role=role, + jsonl=JsonlSink(run_dir / JSONL_FILENAME) if jsonl else None, + http=HttpSink(url, os.environ.get("MOTRIX_METRICS_TOKEN", "")) if url else None, + project=os.environ.get("MOTRIX_METRICS_PROJECT") or None, + group=os.environ.get("MOTRIX_METRICS_GROUP") or None, + meta=meta, + ) + + @property + def enabled(self) -> bool: + return self._jsonl is not None or self._http is not None + + def run(self, state: str, **data: Any) -> None: + """Report the run's state; ``running`` also carries project, group and meta.""" + if not self.enabled: + return + if state == "running": + identity = {key: value for key, value in self._identity.items() if value is not None} + data = {**identity, **data} + self._emit("run", {"state": state, **data}) + + def scalars(self, values: Mapping[str, Any], step: int, *, env_steps: int | None = None) -> None: + """Report scalar values at *step*; NaN and Inf become ``null``. + + A value that is not a single number (a multi-element tensor, a string) + is left out with a one-time warning per key. + """ + if not self.enabled: + return + data: dict[str, float | None] = {} + new_nonfinite = [] + for key, value in values.items(): + try: + number = None if value is None else float(value) + except Exception: # noqa: BLE001 - a bad value must not fail training + if key not in self._invalid_keys: + self._invalid_keys.add(key) + logger.warning("metrics: dropping %r, %s is not a single number", key, type(value).__name__) + continue + if number is not None and not math.isfinite(number): + number = None + if key not in self._nonfinite_keys: + self._nonfinite_keys.add(key) + new_nonfinite.append(key) + data[str(key)] = number + if not data: + return + axes = None if env_steps is None else {"env_steps": int(env_steps)} + self._emit("scalars", data, step=step, axes=axes, clean=False) # already floats or None + for key in new_nonfinite: + self.event("nonfinite", step=step, key=key) + + def event(self, name: str, step: int | None = None, **data: Any) -> None: + """Report a named event such as ``checkpoint_saved``.""" + if not self.enabled: + return + self._emit("event", {"name": name, **data}, step=step) + + def close(self) -> None: + """No more records: each sink sends or writes what it holds at its own pace, then exits. + + Returns at once; nothing waits for the sinks unless :meth:`join` is + called. Idempotent, never raises. + """ + with self._emit_lock: + if self._closed: + return + self._closed = True + for sink in self._sinks(): + try: + sink.close() + except Exception as exc: # noqa: BLE001 - shutting down must not mask the run's outcome + logger.warning("metrics: closing %s failed: %s", type(sink).__name__, exc) + + def join(self, timeout: float) -> None: + """Wait up to *timeout* seconds in total for the sinks to finish after :meth:`close`.""" + deadline = time.monotonic() + timeout + for sink in self._sinks(): + sink.join(max(0.0, deadline - time.monotonic())) + + def _sinks(self) -> list[JsonlSink | HttpSink]: + return [sink for sink in (self._jsonl, self._http) if sink is not None] + + def _record(self, kind: str, data: dict[str, Any], step: int | None = None, axes: dict | None = None) -> dict: + with self._seq_lock: + self._seq += 1 + seq = self._seq + record: dict[str, Any] = { + "v": 1, + "run_id": self.run_id, + "writer": self.writer, + "seq": seq, + "t": time.time(), + "kind": kind, + } + if step is not None: + record["step"] = int(step) + if axes is not None: + record["axes"] = axes + record["data"] = data + return record + + def _emit( + self, kind: str, data: dict[str, Any], step: int | None = None, axes: dict | None = None, *, clean: bool = True + ) -> None: + if self._closed: + return + if clean: + data = _clean(data) + # Both sinks under one lock: the jsonl stays in seq order and holds + # exactly the records the service was sent. + with self._emit_lock: + if self._closed: # close() ran while *data* was being cleaned + return + record = self._record(kind, data, step, axes) + if self._jsonl is not None: + self._jsonl.write(record) + if self._http is not None: + self._http.put(record) + + +def sanitize_run_id(raw: str) -> str: + """Fit *raw* to ``^[A-Za-z0-9._-]{1,200}$``, keeping its distinctive tail.""" + cleaned = _RUN_ID_INVALID.sub("_", raw)[-_RUN_ID_MAX:] + return cleaned or "run" + + +def _clean(value: Any) -> Any: + """Make *value* JSON-safe: numpy/torch to Python, NaN/Inf to None, recursively.""" + if isinstance(value, float): + return value if math.isfinite(value) else None + if value is None or isinstance(value, (bool, int, str)): + return value + if isinstance(value, Mapping): + return {str(key): _clean(item) for key, item in value.items()} + if isinstance(value, (list, tuple, set)): + return [_clean(item) for item in value] + if hasattr(value, "tolist"): # numpy scalars/arrays and torch tensors + try: + return _clean(value.tolist()) + except Exception: # noqa: BLE001 - fall back to the text form + pass + return str(value) + + +class JsonlSink: + """Appends records to a local file, one line each, from a daemon thread. + + Like TensorBoard's event writer, ``write`` only queues: the thread + serializes, writes and flushes whole batches, so a slow or networked disk + never stalls the training thread. Unlike it, a full queue drops the oldest + scalars instead of blocking. After :meth:`close` the writer writes what is + queued and exits. Errors are logged once, never raised. + """ + + def __init__(self, path: Path, *, flush_interval: float = 2.0, max_queue: int = 100_000) -> None: + self.path = path + self._flush_interval = flush_interval + self._max_queue = max_queue + self._queue: collections.deque[dict[str, Any]] = collections.deque() + self._cond = threading.Condition() + self._closing = False + self._file = None + self._warned: set[str] = set() + self.dropped = 0 + self._thread = threading.Thread(target=self._loop, name="motrix-metrics-jsonl", daemon=True) + self._thread.start() + + def write(self, record: dict[str, Any]) -> None: + """Queue *record* without blocking.""" + with self._cond: + if self._closing: + return + if not _bounded_append(self._queue, record, self._max_queue): + self.dropped += 1 + self._warn_once("full", f"metrics: {self.path} is not keeping up, dropping oldest scalars") + if len(self._queue) >= self._max_queue // 2: + self._cond.notify() + + def close(self) -> None: + """No more records: write what is queued, then exit; does not wait.""" + with self._cond: + self._closing = True + self._cond.notify() + + def join(self, timeout: float) -> None: + """Wait up to *timeout* seconds for the writer to finish after :meth:`close`.""" + if self._thread.is_alive(): + self._thread.join(timeout) + + def _loop(self) -> None: + while True: + with self._cond: + # One write and flush per interval, not per record; close() or a + # half-full queue cut the wait short. + deadline = time.monotonic() + self._flush_interval + while not (self._closing or len(self._queue) >= self._max_queue // 2): + remaining = deadline - time.monotonic() + if remaining <= 0: + break + self._cond.wait(remaining) + batch = list(self._queue) + self._queue.clear() + closing = self._closing + if batch: + self._write_batch(batch) + with self._cond: + if closing and not self._queue: + break + if self._file is not None: + try: + self._file.close() + except Exception as exc: # noqa: BLE001 - a local write must never stop training + self._warn_once("io", f"metrics: cannot write {self.path}: {exc}") + + def _write_batch(self, batch: list[dict[str, Any]]) -> None: + lines = [] + for record in batch: + try: + lines.append(json.dumps(record, allow_nan=False, separators=(",", ":"))) + except Exception as exc: # noqa: BLE001 - one bad record must not lose the batch + self._warn_once("encode", f"metrics: cannot encode a record for {self.path}: {exc}") + if not lines: + return + try: + if self._file is None: + self.path.parent.mkdir(parents=True, exist_ok=True) + self._file = self.path.open("a", encoding="utf-8") + self._file.write("\n".join(lines) + "\n") + self._file.flush() + except Exception as exc: # noqa: BLE001 - a local write must never stop training + self._warn_once("io", f"metrics: cannot write {self.path}: {exc}") + + def _warn_once(self, key: str, message: str) -> None: + if key not in self._warned: + self._warned.add(key) + logger.warning(message) + + +class HttpSink: + """Pushes records to ``{url}/api/v1/ingest`` in gzip batches from a daemon thread. + + The queue is bounded; when full the oldest scalars are dropped, never run + or event records. Network errors, 408, 429 and 5xx are retried with + backoff (503 honours ``Retry-After``); other 4xx drop the batch. After + :meth:`close` the sender sends what is queued and exits, retrying for at + most ``flush_budget`` seconds before it drops the rest. + """ + + def __init__( + self, + url: str, + token: str = "", + *, + batch_size: int = 500, + flush_interval: float = 2.0, + request_timeout: float = 5.0, + heartbeat_interval: float = 30.0, + max_queue: int = 10_000, + backoff: float = 1.0, + max_backoff: float = 30.0, + flush_budget: float = 60.0, + ) -> None: + self.endpoint = url.rstrip("/") + "/api/v1/ingest" + self._headers = {"Content-Type": "application/json", "Content-Encoding": "gzip"} + if token: + self._headers["Authorization"] = f"Bearer {token}" + self._batch_size = batch_size + self._flush_interval = flush_interval + self._request_timeout = request_timeout + self._heartbeat_interval = heartbeat_interval + self._max_queue = max_queue + self._backoff = backoff + self._max_backoff = max_backoff + self._flush_budget = flush_budget + self._queue: collections.deque[dict[str, Any]] = collections.deque() + self._cond = threading.Condition() + self._deadline: float | None = None # set by close() + self._heartbeat: Callable[[], dict[str, Any]] | None = None + self._thread = threading.Thread(target=self._loop, name="motrix-metrics-http", daemon=True) + self._warned: set[str] = set() + self.dropped = 0 + + def start(self, heartbeat: Callable[[], dict[str, Any]]) -> None: + self._heartbeat = heartbeat + self._thread.start() + + def put(self, record: dict[str, Any]) -> None: + """Queue *record* without blocking.""" + with self._cond: + if not _bounded_append(self._queue, record, self._max_queue): + self.dropped += 1 + self._warn_once("full", "metrics: send queue full, dropping oldest scalars") + if len(self._queue) >= self._batch_size: + self._cond.notify() + + def close(self) -> None: + """No more records: send what is queued within ``flush_budget``, then exit; does not wait.""" + with self._cond: + self._deadline = time.monotonic() + self._flush_budget + self._cond.notify() + + def join(self, timeout: float) -> None: + """Wait up to *timeout* seconds for the sender to finish after :meth:`close`.""" + if self._thread.is_alive(): + self._thread.join(timeout) + + def _loop(self) -> None: + last_sent = time.monotonic() + while True: + with self._cond: + if self._deadline is None and len(self._queue) < self._batch_size: + self._cond.wait(self._flush_interval) + if self._deadline is not None and not self._queue: + return + if not self._queue and time.monotonic() - last_sent < self._heartbeat_interval: + continue + batch = [self._queue.popleft() for _ in range(min(self._batch_size, len(self._queue)))] + try: + if not batch: + batch = [self._heartbeat()] # idle: tell the service the run is alive + if not self._send_with_retry(batch): + return + except Exception as exc: # noqa: BLE001 - the thread must outlive any one batch + self._warn_once("crash", f"metrics: dropping a batch after an unexpected error: {exc!r}") + last_sent = time.monotonic() + + def _send_with_retry(self, batch: list[dict[str, Any]]) -> bool: + """Deliver or drop *batch*; False when the flush budget ran out first.""" + body = gzip.compress(json.dumps({"schema": SCHEMA, "records": batch}, allow_nan=False).encode("utf-8")) + backoff = self._backoff + while True: + done, retry_after = self._post(body) + if done: + return True + delay = backoff if retry_after is None else retry_after + backoff = min(backoff * 2, self._max_backoff) + if not self._sleep(delay): + return False + + def _sleep(self, delay: float) -> bool: + """Wait *delay* seconds, cut short only by the flush budget; False once it has passed.""" + until = time.monotonic() + delay + with self._cond: + while True: + # put() and close() notify too; only the flush budget may shorten the wait. + end = until if self._deadline is None else min(until, self._deadline) + remaining = end - time.monotonic() + if remaining <= 0: + return self._deadline is None or time.monotonic() < self._deadline + self._cond.wait(remaining) + + def _post(self, body: bytes) -> tuple[bool, float | None]: + """Send once. Returns (done, Retry-After); done covers delivered and dropped batches.""" + request = urllib.request.Request(self.endpoint, data=body, headers=self._headers, method="POST") + try: + with urllib.request.urlopen(request, timeout=self._request_timeout): + return True, None + except urllib.error.HTTPError as exc: + if exc.code in (408, 429) or exc.code >= 500: + self._warn_once("retry", f"metrics: ingest answered {exc.code}, retrying") + return False, _seconds(exc.headers.get("Retry-After")) if exc.code == 503 else None + self._warn_once(f"drop{exc.code}", f"metrics: ingest rejected a batch with {exc.code}; dropping it") + return True, None + except (OSError, ValueError, HTTPException) as exc: # URLError, timeouts, resets, bad replies + self._warn_once("network", f"metrics: cannot reach {self.endpoint} ({exc}), retrying") + return False, None + + def _warn_once(self, key: str, message: str) -> None: + if key not in self._warned: + self._warned.add(key) + logger.warning(message) + + +def _bounded_append(queue: collections.deque[dict[str, Any]], record: dict[str, Any], limit: int) -> bool: + """Append *record*; when *queue* is full, drop its oldest scalars (or *record*) first. False if one was dropped.""" + if len(queue) < limit: + queue.append(record) + return True + victim = next((item for item in queue if item["kind"] == "scalars"), None) + if victim is not None: + queue.remove(victim) + elif record["kind"] == "scalars": + return False + queue.append(record) + return victim is None + + +@atexit.register +def _join_at_exit() -> None: + """Give sinks still writing or sending up to :data:`EXIT_TIMEOUT` seconds in total. + + Sinks run in daemon threads, which the interpreter kills at exit; without + this a script that exits right after training would drop its last records, + the terminal run state among them. Healthy sinks finish in milliseconds. + """ + deadline = time.monotonic() + EXIT_TIMEOUT + for metrics in list(_OPEN): + metrics.close() + metrics.join(max(0.0, deadline - time.monotonic())) + + +def _seconds(value: str | None) -> float | None: + try: + return max(float(value), 0.0) if value is not None else None + except ValueError: + return None diff --git a/motrix_rl/src/motrix_rl/runner.py b/motrix_rl/src/motrix_rl/runner.py index 63c5cd11..65254072 100644 --- a/motrix_rl/src/motrix_rl/runner.py +++ b/motrix_rl/src/motrix_rl/runner.py @@ -1,19 +1,25 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 +import importlib.metadata import logging import time -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path +from typing import Any from motrix_env_core.renderer import RenderConfig from motrix_rl import backend_runtime, checkpoints, frameworks, runs, utils from motrix_rl.config import TaskConfig, TrainConfig from motrix_rl.method import RlMethod +from motrix_rl.metrics import MetricsLogger from motrix_rl.result import TrainResult logger = logging.getLogger(__name__) +# Learning-curve family the live dashboard picks its charts by, keyed by agent name. +_METRICS_FAMILIES = {"fastsac": "rl.offpolicy", "ppo": "rl.onpolicy"} + @dataclass(frozen=True) class TrainRequest: @@ -79,13 +85,22 @@ def create_training_handle(request: TrainRequest) -> frameworks.TrainerHandle: resume_from = None if config.resume is not None: resume_from = str(checkpoints.resolve_resume_checkpoint_path(config.resume)) - handle = _create_handle( - run, - provider, - config, - render=request.render, - resume_from=resume_from, - ) + metrics = MetricsLogger.from_env(run.run_dir) + try: + handle = _create_handle( + run, + provider, + config, + render=request.render, + resume_from=resume_from, + metrics=metrics, + ) + if metrics.enabled: + runs.annotate_run(run.run_dir, live={"run_id": metrics.run_id}) + handle = replace(handle, run_info=_running_fields(run, provider, config)) + except BaseException: + metrics.close() + raise logger.info( "Training startup: trainer handle ready in %.3fs", time.perf_counter() - startup_started, @@ -138,6 +153,7 @@ def _create_handle( randomize_seed: bool = False, render: RenderConfig | None = None, resume_from: str | None = None, + metrics: MetricsLogger | None = None, ) -> frameworks.TrainerHandle: metadata = run.metadata if ( @@ -160,5 +176,32 @@ def _create_handle( checkpoint=task_cfg.checkpoint, render=render, resume_from=resume_from, + metrics=metrics, ) ) + + +def _running_fields(run: runs.RunContext, provider: frameworks.AgentProvider, config: TrainConfig) -> dict[str, Any]: + """Fields of the live ``running`` record besides state, project, group and meta.""" + metadata = run.metadata + try: + version = importlib.metadata.version("motrix-rl") + except importlib.metadata.PackageNotFoundError: + version = None + return { + "env_name": metadata.env_name, + "rllib": metadata.rllib, + "train_backend": metadata.train_backend, + "algo": metadata.algo, + "seed": metadata.seed, + "sim": metadata.sim, + "checkpoint_format": metadata.checkpoint_format, + "family": _METRICS_FAMILIES.get(metadata.algo), + "job_type": "train", + "primary_axis": "env_steps", + "total_steps": provider.total_env_steps(config.algo, config.num_envs), + "num_envs": config.num_envs, + "run_dir": str(run.run_dir), + "config": runs.task_config_snapshot(config), + "motrixlab_version": version, + } diff --git a/motrix_rl/src/motrix_rl/runs.py b/motrix_rl/src/motrix_rl/runs.py index a221f8b7..003dc72c 100644 --- a/motrix_rl/src/motrix_rl/runs.py +++ b/motrix_rl/src/motrix_rl/runs.py @@ -157,15 +157,19 @@ def task_config_path(run_dir: str | Path) -> Path: return Path(run_dir) / TASK_CONFIG_FILENAME -def write_task_config(run_dir: str | Path, config: TaskConfig) -> Path: - """Persist the resolved task recipe used to create a training run.""" +def task_config_snapshot(config: TaskConfig) -> dict[str, Any]: + """Return the resolved task recipe exactly as ``write_task_config`` persists it.""" if not isinstance(config, TaskConfig): raise TypeError(f"Expected TaskConfig, got {type(config).__name__}") - structured = OmegaConf.structured(config) - resolved = _resolved_mapping(structured) - snapshot = { + resolved = _resolved_mapping(OmegaConf.structured(config)) + return { key: resolved[key] for key in ("task", "num_envs", "play_num_envs", "seed", "logging", "checkpoint", "algo") } + + +def write_task_config(run_dir: str | Path, config: TaskConfig) -> Path: + """Persist the resolved task recipe used to create a training run.""" + snapshot = task_config_snapshot(config) path = task_config_path(run_dir) path.parent.mkdir(parents=True, exist_ok=True) OmegaConf.save(OmegaConf.create(snapshot), path) diff --git a/motrix_rl/src/motrix_rl/training_logger.py b/motrix_rl/src/motrix_rl/training_logger.py new file mode 100644 index 00000000..3f5c3147 --- /dev/null +++ b/motrix_rl/src/motrix_rl/training_logger.py @@ -0,0 +1,91 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""The one interface a trainer logs its training metrics through. + +``TrainingLogger`` puts TensorBoard and live metrics (``motrix.metrics/v1``, +see :mod:`motrix_rl.metrics`) behind one call per logging window, so a trainer +builds its scalars once and never talks to either output directly. See +``wiki/design/live-metrics.md``. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable, Collection, Mapping +from pathlib import Path +from typing import Any + +from motrix_rl.metrics import MetricsLogger + +logger = logging.getLogger(__name__) + + +class TrainingLogger: + """Writes a trainer's scalars to TensorBoard and sends them, with its events, to live metrics. + + The TensorBoard writer under ``log_dir`` is opened on the first + :meth:`scalars` call, so a run that fails before logging anything leaves no + event file. TensorBoard is optional: when it cannot be opened, scalars still + go live. The live :class:`MetricsLogger` belongs to the control plane, which + reports the run state and closes it; :meth:`close` closes only TensorBoard. + """ + + def __init__(self, log_dir: str | Path, live: MetricsLogger) -> None: + self._log_dir = Path(log_dir) + self.live = live + self._tensorboard: Any = None + self._tensorboard_opened = False + + def scalars( + self, + values: Mapping[str, float], + step: int, + *, + env_steps: int, + live_extra: Callable[[], Mapping[str, float]] | None = None, + placeholders: Collection[str] = (), + ) -> None: + """Log one window: *values* to TensorBoard under their keys as tags, and live. + + *live_extra* returns scalars sent live only (host samples TensorBoard + never had); it is called only when live metrics are on. + *placeholders* names keys whose value is not defined yet, such as a NaN + mean before the first episode ends: TensorBoard still gets them, live + leaves them out so its once-per-key ``nonfinite`` event stays free for a + real divergence. + """ + tensorboard = self._open_tensorboard() + if tensorboard is not None: + for tag, value in values.items(): + tensorboard.add_scalar(tag, value, step) + if self.live.enabled: + live = {key: value for key, value in values.items() if key not in placeholders} + if live_extra is not None: + live.update(live_extra()) + self.live.scalars(live, step, env_steps=env_steps) + + def event(self, name: str, step: int | None = None, **data: Any) -> None: + """Report a named event such as ``checkpoint_saved``; events go live only.""" + self.live.event(name, step, **data) + + def close(self) -> None: + """Flush and close TensorBoard; never raises, so it is safe in a ``finally``.""" + if self._tensorboard is None: + return + try: + self._tensorboard.close() + except Exception as exc: # noqa: BLE001 - closing must not mask the run's outcome + logger.warning("closing TensorBoard failed: %s", exc) + self._tensorboard = None + + def _open_tensorboard(self) -> Any: + if not self._tensorboard_opened: + self._tensorboard_opened = True + try: + from torch.utils.tensorboard import SummaryWriter + + self._tensorboard = SummaryWriter(log_dir=str(self._log_dir)) + except Exception as exc: # noqa: BLE001 - TensorBoard is optional + logger.warning("TensorBoard is unavailable, logging live metrics only: %s", exc) + return self._tensorboard diff --git a/motrix_rl/tests/fixtures/metrics_v1/motrixlab_event.json b/motrix_rl/tests/fixtures/metrics_v1/motrixlab_event.json new file mode 100644 index 00000000..8c6f9cb1 --- /dev/null +++ b/motrix_rl/tests/fixtures/metrics_v1/motrixlab_event.json @@ -0,0 +1,61 @@ +{ + "body": { + "records": [ + { + "data": { + "name": "run_started", + "workers": 3 + }, + "kind": "event", + "run_id": "master__51f2b01__1712", + "seq": 1791798330000004, + "step": 0, + "t": 1791798331.0, + "v": 1, + "writer": "trainer-host:4242:trainer" + }, + { + "data": { + "name": "warmup_done" + }, + "kind": "event", + "run_id": "master__51f2b01__1712", + "seq": 1791798330000005, + "step": 50, + "t": 1791798332.0, + "v": 1, + "writer": "trainer-host:4242:trainer" + }, + { + "data": { + "name": "checkpoint_saved", + "path": "runs/cartpole/skrl/torch/ppo/26-10-12_10-00-00-000000/checkpoints/model_0000500.pt" + }, + "kind": "event", + "run_id": "master__51f2b01__1712", + "seq": 1791798330000006, + "step": 500, + "t": 1791798333.0, + "v": 1, + "writer": "trainer-host:4242:trainer" + }, + { + "data": { + "state": "finished" + }, + "kind": "run", + "run_id": "master__51f2b01__1712", + "seq": 1791798330000007, + "t": 1791798334.0, + "v": 1, + "writer": "trainer-host:4242:trainer" + } + ], + "schema": "motrix.metrics/v1" + }, + "headers": { + "Authorization": "Bearer fixture-token", + "Content-Encoding": "gzip", + "Content-Type": "application/json" + } +} diff --git a/motrix_rl/tests/fixtures/metrics_v1/motrixlab_running.json b/motrix_rl/tests/fixtures/metrics_v1/motrixlab_running.json new file mode 100644 index 00000000..c60ddc71 --- /dev/null +++ b/motrix_rl/tests/fixtures/metrics_v1/motrixlab_running.json @@ -0,0 +1,70 @@ +{ + "body": { + "records": [ + { + "data": { + "algo": "ppo", + "checkpoint_format": "pt", + "config": { + "algo": { + "trainer": { + "timesteps": 1000 + } + }, + "checkpoint": { + "interval": 500 + }, + "logging": { + "backend": "tensorboard", + "interval": 100 + }, + "num_envs": 2048, + "play_num_envs": 16, + "seed": 42, + "task": { + "algo": "ppo", + "env": "cartpole", + "rllib": "skrl", + "train_backend": "torch" + } + }, + "env_name": "cartpole", + "family": "rl.onpolicy", + "group": "locomotion__1712", + "job_type": "train", + "meta": { + "branch": "master", + "commit_sha": "51f2b010c0ffee", + "exp_name": "locomotion", + "flow": "MotrixLabFlow", + "flow_run_id": "1712", + "job_type": "train" + }, + "motrixlab_version": "0.4.0b1", + "num_envs": 2048, + "primary_axis": "env_steps", + "project": "locomotion", + "rllib": "skrl", + "run_dir": "runs/cartpole/skrl/torch/ppo/26-10-12_10-00-00-000000", + "seed": 42, + "sim": null, + "state": "running", + "total_steps": 2048000, + "train_backend": "torch" + }, + "kind": "run", + "run_id": "master__51f2b01__1712", + "seq": 1791798330000001, + "t": 1791798331.0, + "v": 1, + "writer": "trainer-host:4242:trainer" + } + ], + "schema": "motrix.metrics/v1" + }, + "headers": { + "Authorization": "Bearer fixture-token", + "Content-Encoding": "gzip", + "Content-Type": "application/json" + } +} diff --git a/motrix_rl/tests/fixtures/metrics_v1/motrixlab_scalars.json b/motrix_rl/tests/fixtures/metrics_v1/motrixlab_scalars.json new file mode 100644 index 00000000..c41d0c93 --- /dev/null +++ b/motrix_rl/tests/fixtures/metrics_v1/motrixlab_scalars.json @@ -0,0 +1,44 @@ +{ + "body": { + "records": [ + { + "axes": { + "env_steps": 204800 + }, + "data": { + "reward/alive": 1.0, + "rollout/mean_ep_len": 12.5, + "rollout/mean_return": null, + "sys/cpu_percent": 37.5, + "train/policy_loss": 0.25 + }, + "kind": "scalars", + "run_id": "master__51f2b01__1712", + "seq": 1791798330000002, + "step": 100, + "t": 1791798331.0, + "v": 1, + "writer": "trainer-host:4242:trainer" + }, + { + "data": { + "key": "rollout/mean_return", + "name": "nonfinite" + }, + "kind": "event", + "run_id": "master__51f2b01__1712", + "seq": 1791798330000003, + "step": 100, + "t": 1791798332.0, + "v": 1, + "writer": "trainer-host:4242:trainer" + } + ], + "schema": "motrix.metrics/v1" + }, + "headers": { + "Authorization": "Bearer fixture-token", + "Content-Encoding": "gzip", + "Content-Type": "application/json" + } +} diff --git a/motrix_rl/tests/fixtures/metrics_v1/pipeline_failed.json b/motrix_rl/tests/fixtures/metrics_v1/pipeline_failed.json new file mode 100644 index 00000000..a24d721d --- /dev/null +++ b/motrix_rl/tests/fixtures/metrics_v1/pipeline_failed.json @@ -0,0 +1,24 @@ +{ + "body": { + "records": [ + { + "data": { + "reason": "VolcanoJob finished with phase: Failed", + "state": "failed" + }, + "kind": "run", + "run_id": "master__51f2b01__1712", + "seq": 1791798324430449, + "t": 1791798324.4304497, + "v": 1, + "writer": "orchestrator:4242:pipeline" + } + ], + "schema": "motrix.metrics/v1" + }, + "headers": { + "Authorization": "Bearer fixture-token", + "Content-Encoding": "gzip", + "Content-Type": "application/json" + } +} diff --git a/motrix_rl/tests/fixtures/metrics_v1/pipeline_pending.json b/motrix_rl/tests/fixtures/metrics_v1/pipeline_pending.json new file mode 100644 index 00000000..b8fc03dc --- /dev/null +++ b/motrix_rl/tests/fixtures/metrics_v1/pipeline_pending.json @@ -0,0 +1,33 @@ +{ + "body": { + "records": [ + { + "data": { + "group": "locomotion__1712", + "meta": { + "branch": "master", + "commit_sha": "51f2b010c0ffee", + "exp_name": "locomotion", + "flow": "MotrixLabFlow", + "flow_run_id": "1712", + "job_type": "train" + }, + "project": "locomotion", + "state": "pending" + }, + "kind": "run", + "run_id": "master__51f2b01__1712", + "seq": 1791798324428939, + "t": 1791798324.4289403, + "v": 1, + "writer": "orchestrator:4242:pipeline" + } + ], + "schema": "motrix.metrics/v1" + }, + "headers": { + "Authorization": "Bearer fixture-token", + "Content-Encoding": "gzip", + "Content-Type": "application/json" + } +} diff --git a/motrix_rl/tests/test_metrics.py b/motrix_rl/tests/test_metrics.py new file mode 100644 index 00000000..62748f03 --- /dev/null +++ b/motrix_rl/tests/test_metrics.py @@ -0,0 +1,600 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +import gzip +import json +import re +import socket +import threading +import time +import weakref +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import numpy as np +import pytest +import torch + +from motrix_rl import metrics as metrics_module +from motrix_rl.metrics import JSONL_FILENAME, SCHEMA, HttpSink, JsonlSink, MetricsLogger + +FIXTURES = Path(__file__).parent / "fixtures" / "metrics_v1" +RUN_ID_PATTERN = re.compile(r"^[A-Za-z0-9._-]{1,200}$") +WRITER_PATTERN = re.compile(r"^[^:]+:\d+:[a-z]+$") +KINDS = {"run", "scalars", "event", "eval", "media", "heartbeat"} +ENV_VARS = ( + "MOTRIX_METRICS_URL", + "MOTRIX_METRICS_JSONL", + "MOTRIX_METRICS_TOKEN", + "MOTRIX_METRICS_RUN_ID", + "MOTRIX_METRICS_PROJECT", + "MOTRIX_METRICS_GROUP", + "MOTRIX_METRICS_META", +) + + +class _Ingest: + """A local ingest service that records every request and answers from a script.""" + + def __init__(self, statuses=(), headers=None): + self.requests = [] + self._statuses = list(statuses) + self._headers = headers or {} + ingest = self + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + raw = self.rfile.read(int(self.headers["Content-Length"])) + ingest.requests.append( + {"path": self.path, "headers": dict(self.headers), "body": json.loads(gzip.decompress(raw))} + ) + status = ingest._statuses.pop(0) if ingest._statuses else 200 + self.send_response(status) + for key, value in ingest._headers.items(): + self.send_header(key, value) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, *args): + pass + + self._server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self.url = f"http://127.0.0.1:{self._server.server_address[1]}" + threading.Thread(target=self._server.serve_forever, args=(0.05,), daemon=True).start() + + def records(self): + return [record for request in self.requests for record in request["body"]["records"]] + + def stop(self): + self._server.shutdown() + self._server.server_close() + + +@pytest.fixture +def ingest(): + server = _Ingest() + yield server + server.stop() + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch): + for name in ENV_VARS: + monkeypatch.delenv(name, raising=False) + + +def _logger(tmp_path, url=None, token="", **http_kwargs): + http = HttpSink(url, token, **http_kwargs) if url else None + return MetricsLogger("run-1", jsonl=JsonlSink(tmp_path / JSONL_FILENAME), http=http) + + +def _wait_until(predicate, timeout=5.0): + """Return as soon as *predicate* holds; the timeout is only reached when the test fails.""" + deadline = time.monotonic() + timeout + while not predicate(): + assert time.monotonic() < deadline, "condition not reached" + time.sleep(0.005) + + +def _jsonl(tmp_path): + return [json.loads(line) for line in (tmp_path / JSONL_FILENAME).read_text().splitlines()] + + +def _assert_v1_record(record): + assert record["v"] == 1 + assert RUN_ID_PATTERN.match(record["run_id"]) + assert WRITER_PATTERN.match(record["writer"]) + assert isinstance(record["seq"], int) + assert isinstance(record["t"], float) + assert record["kind"] in KINDS + assert isinstance(record["data"], dict) + assert set(record) <= {"v", "run_id", "writer", "seq", "t", "kind", "step", "axes", "data"} + if record["kind"] == "run": + assert record["data"]["state"] in {"pending", "running", "finished", "failed", "killed"} + + +def _metrics_threads(): + return [thread for thread in threading.enumerate() if thread.name.startswith("motrix-metrics-")] + + +def test_without_url_or_jsonl_logger_is_disabled_and_inert(tmp_path): + threads_before = _metrics_threads() + logger = MetricsLogger.from_env(tmp_path) + + logger.run("running", total_steps=1) + logger.scalars({"a": 1.0}, step=0) + logger.event("run_started") + logger.close() + logger.join(5) + + assert not logger.enabled + assert list(tmp_path.iterdir()) == [] + assert set(_metrics_threads()) <= set(threads_before) # started none; others' may finish meanwhile + + +def test_records_carry_protocol_fields_and_monotonic_seq(tmp_path): + before = time.time_ns() // 1000 + logger = _logger(tmp_path) + + def emit(worker): + for step in range(50): + logger.scalars({f"w{worker}": step}, step=step, env_steps=step * 4) + + threads = [threading.Thread(target=emit, args=(i,)) for i in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + logger.event("checkpoint_saved", step=49, path="ckpt.pt") + logger.close() + logger.join(5) + + records = _jsonl(tmp_path) + for record in records: + _assert_v1_record(record) + seqs = [record["seq"] for record in records] + assert len(records) == 201 + assert seqs == sorted(set(seqs)) + assert seqs[0] > before + assert records[0]["writer"].endswith(":trainer") + assert records[0]["axes"] == {"env_steps": records[0]["step"] * 4} + assert records[-1]["data"] == {"name": "checkpoint_saved", "path": "ckpt.pt"} + assert "axes" not in records[-1] + + +def test_values_become_floats_or_null_and_nonfinite_is_reported_once(tmp_path): + logger = _logger(tmp_path) + + logger.scalars( + {"nan": float("nan"), "inf": np.float32("inf"), "np": np.float64(2.5), "torch": torch.tensor(1.5), "int": 3}, + step=0, + ) + logger.scalars({"nan": float("nan"), "inf": 1.0}, step=1) + logger.event("resumed", step=1, stats={"x": np.array([1.0, float("nan")]), "y": torch.tensor(2)}) + logger.close() + logger.join(5) + + records = _jsonl(tmp_path) + assert records[0]["data"] == {"nan": None, "inf": None, "np": 2.5, "torch": 1.5, "int": 3.0} + nonfinite = [record for record in records if record["kind"] == "event" and record["data"]["name"] == "nonfinite"] + assert sorted(record["data"]["key"] for record in nonfinite) == ["inf", "nan"] + assert all(record["step"] == 0 for record in nonfinite) + assert records[-2]["data"] == {"nan": None, "inf": 1.0} + assert records[-1]["data"] == {"name": "resumed", "stats": {"x": [1.0, None], "y": 2}} + + +def test_values_that_are_not_one_number_are_dropped_without_raising(tmp_path): + logger = _logger(tmp_path) + + for step in range(2): + logger.scalars({"vec": torch.ones(4), "arr": np.zeros(3), "text": "fast", "ok": torch.tensor([2.0])}, step=step) + logger.scalars({"vec": torch.ones(4)}, step=2) + logger.close() + logger.join(5) + + assert [record["data"] for record in _jsonl(tmp_path)] == [{"ok": 2.0}, {"ok": 2.0}] + + +def test_records_emitted_while_closing_are_discarded(tmp_path): + logger = _logger(tmp_path) + + cleaning, release = threading.Event(), threading.Event() + + class SlowValue: + def tolist(self): + cleaning.set() + release.wait(5) + return 1.0 + + logger.run("finished") + late = threading.Thread(target=logger.event, args=("late",), kwargs={"value": SlowValue()}) + late.start() + cleaning.wait(5) + logger.close() + logger.join(5) + release.set() + late.join() + + assert [record["data"]["state"] for record in _jsonl(tmp_path)] == ["finished"] + + +def test_close_returns_at_once_and_join_waits_for_both_sinks_against_one_deadline(): + calls = [] + + class Stuck: + """A sink that never finishes: join() uses up all the time it is given.""" + + def __init__(self, name): + self.name = name + + def start(self, heartbeat): + pass + + def close(self): + calls.append(("close", self.name)) + + def join(self, timeout): + calls.append(("join", self.name, timeout)) + time.sleep(timeout) + + logger = MetricsLogger("run-1", jsonl=Stuck("jsonl"), http=Stuck("http")) + logger.close() + assert calls == [("close", "jsonl"), ("close", "http")] # told, not waited for + + logger.join(0.1) + + assert [call[:2] for call in calls[2:]] == [("join", "jsonl"), ("join", "http")] + assert calls[3][2] < 0.05 # the second sink gets what is left, not another full timeout + + +def test_exit_writes_what_a_logger_still_holds(monkeypatch, tmp_path): + monkeypatch.setattr(metrics_module, "_OPEN", weakref.WeakSet()) + logger = MetricsLogger("run-1", jsonl=JsonlSink(tmp_path / JSONL_FILENAME, flush_interval=60)) + logger.run("finished") + + metrics_module._join_at_exit() # never closed: the exit hook closes it and waits + + assert [record["data"]["state"] for record in _jsonl(tmp_path)] == ["finished"] + + +def test_run_id_prefers_environment_and_is_sanitized(monkeypatch, tmp_path): + monkeypatch.setenv("MOTRIX_METRICS_URL", "http://127.0.0.1:9") + run_dir = tmp_path / "runs" / "cartpole" / "skrl" / "torch" / "ppo" / "26-10-12_10-00-00-000000" + + derived = MetricsLogger.from_env(run_dir) + derived.close() + assert derived.run_id == "cartpole__skrl__torch__ppo__26-10-12_10-00-00-000000" + + monkeypatch.setenv("MOTRIX_METRICS_RUN_ID", "a b/c:" + "x" * 300 + "-tail") + explicit = MetricsLogger.from_env(run_dir) + explicit.close() + assert RUN_ID_PATTERN.match(explicit.run_id) + assert len(explicit.run_id) == 200 + assert explicit.run_id.endswith("x-tail") + + monkeypatch.setenv("MOTRIX_METRICS_RUN_ID", "ok.run_1-2") + kept = MetricsLogger.from_env(run_dir) + kept.close() + assert kept.run_id == "ok.run_1-2" + + +def test_url_alone_pushes_without_writing_a_local_file(monkeypatch, tmp_path, ingest): + monkeypatch.setenv("MOTRIX_METRICS_URL", ingest.url) + logger = MetricsLogger.from_env(tmp_path) + + logger.run("running") + logger.close() + logger.join(5) + + assert [record["data"]["state"] for record in ingest.records()] == ["running"] + assert not (tmp_path / JSONL_FILENAME).exists() + + +@pytest.mark.parametrize("flag", ["1", "true", "YES", "on"]) +def test_jsonl_flag_alone_writes_locally_without_http(monkeypatch, tmp_path, flag): + monkeypatch.setenv("MOTRIX_METRICS_JSONL", flag) + before = set(threading.enumerate()) # other tests' senders may still be finishing in the background + logger = MetricsLogger.from_env(tmp_path) + + assert logger.enabled + assert [thread.name for thread in set(threading.enumerate()) - before] == ["motrix-metrics-jsonl"] + logger.run("running") + logger.close() + logger.join(5) + + assert [record["data"]["state"] for record in _jsonl(tmp_path)] == ["running"] + + +def test_jsonl_flag_off_values_keep_the_logger_disabled(monkeypatch, tmp_path): + monkeypatch.setenv("MOTRIX_METRICS_JSONL", "0") + + assert not MetricsLogger.from_env(tmp_path).enabled + + +def test_stuck_disk_does_not_block_emit_and_loses_nothing(tmp_path): + writing, unstick = threading.Event(), threading.Event() + + class StuckSink(JsonlSink): + def __init__(self, path): + self.batches = [] + super().__init__(path, flush_interval=0.01) + + def _write_batch(self, batch): + writing.set() + unstick.wait(5) # a hung or networked run_dir + self.batches.append(len(batch)) + super()._write_batch(batch) + + sink = StuckSink(tmp_path / JSONL_FILENAME) + logger = MetricsLogger("run-1", jsonl=sink) + logger.scalars({"a": 0.0}, step=0) + writing.wait(5) # the writer took the first record and is stuck on it + + for step in range(1, 500): + logger.scalars({"a": float(step)}, step=step) + finished_while_stuck = sink.batches == [] + unstick.set() + logger.close() + logger.join(5) + + assert finished_while_stuck # every emit returned while the disk was still stuck + assert [record["step"] for record in _jsonl(tmp_path)] == list(range(500)) + assert sink.batches == [1, 499] # the backlog goes out as one write and flush, not one per record + + +def test_from_env_pushes_gzip_batches_with_bearer_token(monkeypatch, tmp_path, ingest): + monkeypatch.setenv("MOTRIX_METRICS_URL", ingest.url + "/") + monkeypatch.setenv("MOTRIX_METRICS_JSONL", "1") + monkeypatch.setenv("MOTRIX_METRICS_TOKEN", "secret-token") + monkeypatch.setenv("MOTRIX_METRICS_PROJECT", "proj") + monkeypatch.setenv("MOTRIX_METRICS_GROUP", "grp") + monkeypatch.setenv("MOTRIX_METRICS_META", '{"flow_run_id": "7"}') + logger = MetricsLogger.from_env(tmp_path) + + logger.run("running", total_steps=10) + logger.scalars({"rollout/mean_return": 1.0}, step=1, env_steps=8) + logger.run("finished") + logger.close() + logger.join(5) + + assert len(ingest.requests) == 1 + request = ingest.requests[0] + assert request["path"] == "/api/v1/ingest" + assert request["headers"]["Authorization"] == "Bearer secret-token" + assert request["headers"]["Content-Encoding"] == "gzip" + assert request["headers"]["Content-Type"] == "application/json" + assert request["body"]["schema"] == SCHEMA + assert request["body"]["records"] == _jsonl(tmp_path) + running = request["body"]["records"][0]["data"] + assert running == { + "state": "running", + "project": "proj", + "group": "grp", + "meta": {"flow_run_id": "7"}, + "total_steps": 10, + } + assert "secret-token" not in (tmp_path / JSONL_FILENAME).read_text() + + +def test_server_errors_are_retried_without_duplicating_records(tmp_path): + ingest = _Ingest(statuses=[500]) + try: + # Nothing is sent before close(), so both records share the first batch. + logger = _logger(tmp_path, ingest.url, flush_interval=60, backoff=0.01) + logger.scalars({"a": 1.0}, step=0) + logger.event("warmup_done", step=0) + logger.close() + logger.join(10) + finally: + ingest.stop() + + assert len(ingest.requests) == 2 + assert ingest.requests[0]["body"] == ingest.requests[1]["body"] + assert [record["seq"] for record in ingest.records()[2:]] == [record["seq"] for record in _jsonl(tmp_path)] + + +def test_service_unavailable_honours_retry_after(tmp_path): + ingest = _Ingest(statuses=[503], headers={"Retry-After": "0"}) + try: + # Only Retry-After: 0 can explain a prompt retry under a 60 s backoff. + logger = _logger(tmp_path, ingest.url, flush_interval=0.01, backoff=60) + logger.event("run_started", step=0) + _wait_until(lambda: len(ingest.requests) == 2) + logger.close() + logger.join(5) + finally: + ingest.stop() + + +def test_records_queued_during_backoff_do_not_trigger_early_retries(tmp_path): + ingest = _Ingest(statuses=[503] * 10, headers={"Retry-After": "30"}) + try: + logger = _logger(tmp_path, ingest.url, flush_interval=60, batch_size=2) + logger.scalars({"a": 0.0}, step=0) + logger.scalars({"a": 1.0}, step=1) + _wait_until(lambda: len(ingest.requests) == 1) # now backing off for 30 s + for step in range(2, 50): # every put sees a full batch waiting + logger.scalars({"a": float(step)}, step=step) + logger.close() + logger.join(0.1) + finally: + ingest.stop() + + assert len(ingest.requests) == 1 + + +class _GarbageThenOk: + """A TCP server whose first reply is not HTTP; later requests get 200.""" + + def __init__(self): + self.bodies = [] + self._sock = socket.create_server(("127.0.0.1", 0)) + self.url = f"http://127.0.0.1:{self._sock.getsockname()[1]}" + threading.Thread(target=self._serve, daemon=True).start() + + def _serve(self): + replies = 0 + while True: + try: + conn, _ = self._sock.accept() + except OSError: + return + with conn: + head, _, body = conn.recv(65536).partition(b"\r\n\r\n") + length = int(re.search(rb"Content-Length: (\d+)", head).group(1)) + while len(body) < length: + body += conn.recv(65536) + replies += 1 + if replies == 1: + conn.sendall(b"garbage\r\n\r\n") + else: + self.bodies.append(json.loads(gzip.decompress(body))) + conn.sendall(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") + + def stop(self): + self._sock.close() + + +def test_malformed_replies_are_retried_and_the_sink_keeps_running(tmp_path): + server = _GarbageThenOk() + try: + logger = _logger(tmp_path, server.url, flush_interval=0.01, backoff=0.01) + logger.event("run_started", step=0) + _wait_until(lambda: server.bodies) + logger.event("warmup_done", step=1) + logger.close() + logger.join(5) + finally: + server.stop() + + sent = [record["data"]["name"] for body in server.bodies for record in body["records"]] + assert sent == ["run_started", "warmup_done"] + + +def test_unexpected_sender_error_does_not_stop_delivery(tmp_path, ingest): + calls = [] + + def heartbeat(): + calls.append(None) + raise RuntimeError("boom") + + http = HttpSink(ingest.url, flush_interval=0.02, heartbeat_interval=0.0) + http.start(heartbeat) + _wait_until(lambda: calls) + http.put({"kind": "event", "data": {"name": "warmup_done"}}) + http.close() + http.join(5) + + assert calls + assert [record["data"]["name"] for record in ingest.records()] == ["warmup_done"] + + +def test_client_errors_drop_the_batch(tmp_path): + ingest = _Ingest(statuses=[401, 401]) + try: + logger = _logger(tmp_path, ingest.url, flush_interval=0.01) + logger.event("run_started", step=0) + _wait_until(lambda: len(ingest.requests) == 1) + logger.event("warmup_done", step=1) + logger.close() + logger.join(5) + finally: + ingest.stop() + + assert [[r["data"]["name"] for r in req["body"]["records"]] for req in ingest.requests] == [ + ["run_started"], + ["warmup_done"], + ] + assert len(_jsonl(tmp_path)) == 2 + + +def test_unreachable_service_keeps_the_jsonl_and_gives_up_within_its_flush_budget(tmp_path): + # Emit cost under load is measured by bench/bench_live_metrics.py, not asserted here. + http = HttpSink("http://127.0.0.1:9", flush_interval=0.01, backoff=0.01, flush_budget=0.2) + logger = MetricsLogger("run-1", jsonl=JsonlSink(tmp_path / JSONL_FILENAME), http=http) + + for step in range(1000): + logger.scalars({"a": float(step)}, step=step) + logger.close() + logger.join(5) + + assert not http._thread.is_alive() # gave up after its budget, well before the join timeout + assert len(_jsonl(tmp_path)) == 1000 + + +def test_full_queue_drops_oldest_scalars_and_keeps_events(tmp_path, ingest): + logger = _logger(tmp_path, ingest.url, flush_interval=60, max_queue=3) + + logger.event("run_started", step=0) + for step in range(5): + logger.scalars({"a": float(step)}, step=step) + logger.close() + logger.join(5) + + sent = ingest.records() + assert [record["kind"] for record in sent] == ["event", "scalars", "scalars"] + assert [record["step"] for record in sent[1:]] == [3, 4] + assert len(_jsonl(tmp_path)) == 6 + + +def test_idle_sink_sends_heartbeats_only_over_http(tmp_path, ingest): + logger = _logger(tmp_path, ingest.url, flush_interval=0.01, heartbeat_interval=0.01) + + logger.run("running") + _wait_until(lambda: any(record["kind"] == "heartbeat" for record in ingest.records())) + logger.close() + logger.join(5) + + heartbeats = [record for record in ingest.records() if record["kind"] == "heartbeat"] + assert heartbeats + for record in heartbeats: + _assert_v1_record(record) + assert record["data"] == {} + assert [record["kind"] for record in _jsonl(tmp_path)] == ["run"] + + +@pytest.mark.parametrize("name", ["pipeline_pending.json", "pipeline_failed.json"]) +def test_orchestrator_fixtures_share_the_record_format(name, tmp_path, ingest): + fixture = json.loads((FIXTURES / name).read_text()) + logger = MetricsLogger("run-1", http=HttpSink(ingest.url, "fixture-token")) + logger.run("running") + logger.close() + logger.join(5) + + assert fixture["body"]["schema"] == SCHEMA + assert fixture["headers"] == {key: ingest.requests[0]["headers"][key] for key in fixture["headers"]} + for record in fixture["body"]["records"]: + _assert_v1_record(record) + assert record["writer"].endswith(":pipeline") + + +@pytest.mark.parametrize("name", ["motrixlab_running.json", "motrixlab_scalars.json", "motrixlab_event.json"]) +def test_logger_reproduces_motrixlab_fixtures(name, tmp_path, ingest): + fixture = json.loads((FIXTURES / name).read_text()) + expected = fixture["body"]["records"] + logger = MetricsLogger(expected[0]["run_id"], http=HttpSink(ingest.url, "fixture-token", flush_interval=60)) + + for record in expected: + data = dict(record["data"]) + if record["kind"] == "run": + logger.run(data.pop("state"), **data) + elif record["kind"] == "scalars": + values = {key: float("nan") if value is None else value for key, value in data.items()} + logger.scalars(values, record["step"], env_steps=record["axes"]["env_steps"]) + elif data["name"] != "nonfinite": # emitted by the logger itself + logger.event(data.pop("name"), step=record.get("step"), **data) + logger.close() + logger.join(5) + + assert len(ingest.requests) == 1 + request = ingest.requests[0] + assert fixture["headers"] == {key: request["headers"][key] for key in fixture["headers"]} + sent = request["body"]["records"] + for record in sent: + _assert_v1_record(record) + assert [record["seq"] for record in sent] == sorted(record["seq"] for record in sent) + normalized = [ + {**record, "writer": want["writer"], "seq": want["seq"], "t": want["t"]} + for record, want in zip(sent, expected, strict=True) + ] + assert {**request["body"], "records": normalized} == fixture["body"] diff --git a/motrix_rl/tests/test_rl_contracts.py b/motrix_rl/tests/test_rl_contracts.py index c5687186..67f705a2 100644 --- a/motrix_rl/tests/test_rl_contracts.py +++ b/motrix_rl/tests/test_rl_contracts.py @@ -1,9 +1,12 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 +import json +import threading from dataclasses import dataclass -from motrix_rl.contracts import AgentProvider, RlFramework, TrainerContext +from motrix_rl.contracts import AgentProvider, RlFramework, TrainerBase, TrainerContext, TrainerHandle +from motrix_rl.metrics import JSONL_FILENAME, JsonlSink, MetricsLogger @dataclass @@ -43,3 +46,92 @@ def test_provider_and_framework_contracts(): assert framework.supported_train_backends("unit") == ("test",) assert framework.get_agent_provider("unit", "test").checkpoint_format == "bin" assert Provider().validate_config(Config(3)).value == 3 + + +class _Run: + def __init__(self, run_dir): + self.run_dir = run_dir + + +class _Trainer(TrainerBase): + def __init__(self, error=None): + self.error = error + + def train(self): + if self.error is not None: + raise self.error + + def play(self, policy): + raise NotImplementedError + + +def _train(tmp_path, error=None): + metrics = MetricsLogger("run-1", jsonl=JsonlSink(tmp_path / JSONL_FILENAME)) + handle = TrainerHandle(run=_Run(tmp_path), trainer=_Trainer(error), metrics=metrics, run_info={"total_steps": 8}) + raised = None + try: + handle.train() + except BaseException as exc: # noqa: BLE001 - the test inspects what propagated + raised = exc + metrics.event("after_close") # dropped: the handle closed the logger + metrics.join(5) # the handle closed the logger without waiting for the sink; the test waits + records = [json.loads(line) for line in (tmp_path / JSONL_FILENAME).read_text().splitlines()] + return raised, [record["data"] for record in records] + + +def test_trainer_handle_reports_finished_after_running(tmp_path): + raised, states = _train(tmp_path) + + assert raised is None + assert states == [{"state": "running", "meta": {}, "total_steps": 8}, {"state": "finished"}] + + +def test_trainer_handle_does_not_wait_for_a_stuck_sink(tmp_path): + unstick = threading.Event() + + class StuckDisk(JsonlSink): + def _write_batch(self, batch): + unstick.wait(5) + super()._write_batch(batch) + + sink = StuckDisk(tmp_path / JSONL_FILENAME) + handle = TrainerHandle(run=_Run(tmp_path), trainer=_Trainer(), metrics=MetricsLogger("run-1", jsonl=sink)) + + handle.train() + + still_writing = sink._thread.is_alive() # train() returned while the disk still held the records + unstick.set() + handle.metrics.join(5) + assert still_writing + states = [json.loads(line)["data"]["state"] for line in (tmp_path / JSONL_FILENAME).read_text().splitlines()] + assert states == ["running", "finished"] + + +def test_trainer_handle_reports_killed_on_keyboard_interrupt(tmp_path): + error = KeyboardInterrupt() + raised, states = _train(tmp_path, error) + + assert raised is error + assert [state["state"] for state in states] == ["running", "killed"] + + +def test_trainer_handle_reports_failed_with_reason_and_reraises(tmp_path): + (tmp_path / "async_errors").mkdir() + (tmp_path / "async_errors" / "learner_error.log").write_text("Traceback") + error = RuntimeError("x" * 5000) + raised, states = _train(tmp_path, error) + + assert raised is error + assert [state["state"] for state in states] == ["running", "failed"] + reason = states[-1]["reason"] + assert reason.startswith("RuntimeError: xxx") + assert reason.endswith("(see async_errors/learner_error.log)") + assert len(reason) == 2048 + + +def test_trainer_handle_without_metrics_writes_nothing(tmp_path): + handle = TrainerHandle(run=_Run(tmp_path), trainer=_Trainer()) + + assert handle.train() is handle.run + assert not handle.metrics.enabled + assert list(tmp_path.iterdir()) == [] diff --git a/motrix_rl/tests/test_runner.py b/motrix_rl/tests/test_runner.py index 1ce8ba2e..eeac7fa9 100644 --- a/motrix_rl/tests/test_runner.py +++ b/motrix_rl/tests/test_runner.py @@ -1,12 +1,22 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 +import gzip +import json +import threading from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import pytest from motrix_env_core import registry as env_registry from motrix_rl import frameworks, runner, runs, utils from motrix_rl.config import CheckpointConfig, LoggingConfig, TaskMeta, TrainConfig from motrix_rl.frameworks import RlFramework +from motrix_rl.metrics import JSONL_FILENAME + +RUNNING_FIXTURE = Path(__file__).parent / "fixtures" / "metrics_v1" / "motrixlab_running.json" @dataclass @@ -14,30 +24,39 @@ class DemoCfg: value: int = 1 -def test_runner_train_creates_run_and_returns_result(monkeypatch, tmp_path): +@pytest.fixture +def demo_framework(monkeypatch): monkeypatch.setattr(frameworks, "_frameworks", {}) monkeypatch.setattr("motrix_rl.plugins.load_plugins", lambda: None) monkeypatch.setattr(env_registry, "contains", lambda env_name: env_name == "demo-env") monkeypatch.setattr(utils, "get_device_supports", lambda: utils.DeviceSupports(torch=True)) - + monkeypatch.delenv("MOTRIX_METRICS_URL", raising=False) + monkeypatch.delenv("MOTRIX_METRICS_JSONL", raising=False) + monkeypatch.delenv("MOTRIX_METRICS_RUN_ID", raising=False) provider = _RunnerProvider() frameworks.register_framework(_RunnerFramework("external", (provider,))) - - result = runner.train( - runner.TrainRequest( - config=TrainConfig( - task=TaskMeta(env="demo-env", rllib="external", algo="demo", train_backend="custom"), - num_envs=2, - play_num_envs=3, - seed=7, - algo=DemoCfg(), - logging=LoggingConfig(backend="tensorboard", interval=20), - checkpoint=CheckpointConfig(interval=50), - ), - runs_root=tmp_path, - ) + return provider + + +def _train_request(runs_root): + return runner.TrainRequest( + config=TrainConfig( + task=TaskMeta(env="demo-env", rllib="external", algo="demo", train_backend="custom"), + num_envs=2, + play_num_envs=3, + seed=7, + algo=DemoCfg(), + logging=LoggingConfig(backend="tensorboard", interval=20), + checkpoint=CheckpointConfig(interval=50), + ), + runs_root=runs_root, ) + +def test_runner_train_creates_run_and_returns_result(demo_framework, tmp_path): + provider = demo_framework + result = runner.train(_train_request(tmp_path)) + assert provider.trainer is not None assert provider.trainer.train_called assert result.run.metadata.env_name == "demo-env" @@ -65,6 +84,56 @@ def test_runner_train_creates_run_and_returns_result(monkeypatch, tmp_path): assert handle.trainer.context.rl_cfg == DemoCfg(value=2) assert handle.trainer.context.logging == LoggingConfig(backend="tensorboard", interval=20) assert handle.trainer.context.checkpoint == CheckpointConfig(interval=50) + assert not provider.trainer.context.metrics.enabled + assert not (result.run.run_dir / JSONL_FILENAME).exists() + + +def test_training_reports_live_metrics_but_play_does_not(demo_framework, monkeypatch, tmp_path): + received = [] + + class Ingest(BaseHTTPRequestHandler): + def do_POST(self): + body = gzip.decompress(self.rfile.read(int(self.headers["Content-Length"]))) + received.extend(json.loads(body)["records"]) + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Ingest) + threading.Thread(target=server.serve_forever, args=(0.05,), daemon=True).start() + monkeypatch.setenv("MOTRIX_METRICS_URL", f"http://127.0.0.1:{server.server_address[1]}") + monkeypatch.setenv("MOTRIX_METRICS_JSONL", "1") + monkeypatch.setenv("MOTRIX_METRICS_RUN_ID", "demo-run") + monkeypatch.setenv("MOTRIX_METRICS_PROJECT", "demo") + monkeypatch.setenv("MOTRIX_METRICS_GROUP", "demo__1") + monkeypatch.setenv("MOTRIX_METRICS_META", '{"flow_run_id": "1"}') + try: + result = runner.train(_train_request(tmp_path)) + metadata = (result.run.run_dir / runs.METADATA_FILENAME).read_text() + train_metrics = demo_framework.trainer.context.metrics + train_metrics.join(5) # training closed the logger without waiting for the sinks + play = runner.create_run_handle(result.run) + finally: + server.shutdown() + server.server_close() + + records = [json.loads(line) for line in (result.run.run_dir / JSONL_FILENAME).read_text().splitlines()] + assert records == received + assert [record["data"]["state"] for record in records] == ["running", "finished"] + running = records[0]["data"] + expected = json.loads(RUNNING_FIXTURE.read_text())["body"]["records"][0]["data"] + assert set(running) == set(expected) + assert running["config"] == runs.task_config_snapshot(_train_request(tmp_path).config) + assert running["run_dir"] == str(result.run.run_dir) + assert train_metrics.run_id == "demo-run" + assert json.loads(metadata)["live"] == {"run_id": "demo-run"} + + assert not play.metrics.enabled + assert not play.trainer.context.metrics.enabled + assert (result.run.run_dir / runs.METADATA_FILENAME).read_text() == metadata class _RunnerTrainer(frameworks.TrainerBase): diff --git a/motrix_rl/tests/test_training_logger.py b/motrix_rl/tests/test_training_logger.py new file mode 100644 index 00000000..16fe9e50 --- /dev/null +++ b/motrix_rl/tests/test_training_logger.py @@ -0,0 +1,102 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +import math +from unittest.mock import Mock + +import pytest +import torch.utils.tensorboard + +from motrix_rl.metrics import MetricsLogger +from motrix_rl.training_logger import TrainingLogger + + +class _Live(MetricsLogger): + """An enabled live logger that keeps what it is sent.""" + + def __init__(self): + super().__init__("run-1") + self.sent = [] + + @property + def enabled(self): + return True + + def scalars(self, values, step, *, env_steps=None): + self.sent.append(("scalars", dict(values), step, env_steps)) + + def event(self, name, step=None, **data): + self.sent.append(("event", name, step, data)) + + +@pytest.fixture +def tensorboard(monkeypatch): + writer = Mock() + factory = Mock(return_value=writer) + monkeypatch.setattr(torch.utils.tensorboard, "SummaryWriter", factory) + return factory, writer + + +def test_scalars_reach_tensorboard_and_live_from_one_call(tmp_path, tensorboard): + factory, writer = tensorboard + live = _Live() + log = TrainingLogger(tmp_path, live) + + log.scalars( + {"rollout/mean_return": math.nan, "train/qf_loss": 0.5}, + 7, + env_steps=56, + live_extra=lambda: {"sys/cpu_percent": 12.0}, + placeholders=("rollout/mean_return",), + ) + + factory.assert_called_once_with(log_dir=str(tmp_path)) + tags = [(call.args[0], call.args[2]) for call in writer.add_scalar.call_args_list] + assert tags == [("rollout/mean_return", 7), ("train/qf_loss", 7)] # placeholders still reach TensorBoard + assert live.sent == [("scalars", {"train/qf_loss": 0.5, "sys/cpu_percent": 12.0}, 7, 56)] + + +def test_tensorboard_opens_on_the_first_scalars_only(tmp_path, tensorboard): + factory, writer = tensorboard + live = _Live() + log = TrainingLogger(tmp_path, live) + + log.event("run_started", step=0) + log.close() + + factory.assert_not_called() # a run that never logged leaves no event file + assert live.sent == [("event", "run_started", 0, {})] + + +def test_close_closes_tensorboard_once_and_never_raises(tmp_path, tensorboard): + _, writer = tensorboard + writer.close.side_effect = OSError("disk gone") + log = TrainingLogger(tmp_path, _Live()) + log.scalars({"a": 1.0}, 0, env_steps=1) + + log.close() + log.close() + + writer.close.assert_called_once_with() + + +def test_live_still_gets_scalars_without_tensorboard(tmp_path, monkeypatch): + monkeypatch.setattr(torch.utils.tensorboard, "SummaryWriter", Mock(side_effect=ImportError("no tensorboard"))) + live = _Live() + log = TrainingLogger(tmp_path, live) + + log.scalars({"a": 1.0}, 0, env_steps=1) + log.scalars({"a": 2.0}, 1, env_steps=2) + + assert [sent[1] for sent in live.sent] == [{"a": 1.0}, {"a": 2.0}] + + +def test_disabled_live_writes_tensorboard_only(tmp_path, tensorboard): + _, writer = tensorboard + log = TrainingLogger(tmp_path, MetricsLogger.disabled()) + + log.scalars({"a": 1.0}, 0, env_steps=1, live_extra=Mock(side_effect=AssertionError("sampled for live"))) + log.event("run_started", step=0) + + writer.add_scalar.assert_called_once_with("a", 1.0, 0) + assert list(tmp_path.iterdir()) == [] diff --git a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/train.py b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/train.py index 29ca11df..cde4afe3 100644 --- a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/train.py +++ b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/train.py @@ -28,6 +28,7 @@ from motrix_env_core import registry as env_registry from motrix_env_core.renderer import RenderConfig +from motrix_rl import checkpoints from motrix_rl.console import TrainingPanelStats, emit_training_panel, open_training_live from motrix_rl.contracts import TrainerBase, TrainerContext from motrix_rl.device import resolve_default_device @@ -40,6 +41,7 @@ capture_system_info, sample_gpu_devices, ) +from motrix_rl.training_logger import TrainingLogger from motrix_rl_builtin.fastsac.agent import FastSacAgent from motrix_rl_builtin.fastsac.async_impl.numa import spawn_placement from motrix_rl_builtin.fastsac.async_impl.panels import BootPanel @@ -61,6 +63,7 @@ run_learner_process, ) from motrix_rl_builtin.fastsac.config import FastSacCfg +from motrix_rl_builtin.fastsac.live_metrics import ROLLOUT_MEANS, system_scalars from motrix_rl_builtin.fastsac.wrap import FastSacEnvWrap torch.set_float32_matmul_precision("high") @@ -140,6 +143,8 @@ def train(self) -> None: async_options = cfg.trainer.async_options logging_interval = self._context.logging.interval save_interval = self._context.checkpoint.interval + # Parent-only: the live logger holds a thread and locks and never reaches the workers. + training_log = TrainingLogger(self._context.run_dir, self._context.metrics) train_started_at = time.perf_counter() if self._context.logging.backend != "tensorboard": @@ -171,6 +176,12 @@ def train(self) -> None: # Falls back to the OS default placement on single-node hosts. cpus_per_collector = async_options.cpus_per_collector num_envs = self._context.num_envs + # Env-step accounting follows the collector side: each collector + # step ingests one batch of its own env shard, so the aggregate + # counter times the shard size is the true transition count (the + # learner's global_step is an iteration counter that can drift + # from collection under non-strict UTD). + per_collector_envs = num_envs // num_collectors # Shared-memory primitives allocated in the parent, inherited by children. # One SPSC ring + one weight channel per collector: every shared quantity # keeps exactly one producer and one consumer, so the lock-free @@ -244,6 +255,7 @@ def train(self) -> None: is_resume = resume_step > 0 control.resume_collector_steps(resume_step) control.global_step = resume_step + training_log.event("resumed", step=resume_step, path=str(self._resume_from), global_step=resume_step) ctx = mp.get_context("spawn") stats_queues = [ctx.Queue(maxsize=8) for _ in range(num_collectors)] @@ -340,6 +352,7 @@ def _drain_child_errors() -> list[tuple[str, str]]: "num_learners": num_learners, "rendezvous_file": rendezvous_file, "learner_device": str(topology.learners[rank].device), + "notify_checkpoints": training_log.live.enabled, "all_rings": rings, "weight_ipc": [ c.weight_ipc for c in topology.collectors[rank * per_learner : (rank + 1) * per_learner] @@ -387,6 +400,11 @@ def _drain_child_errors() -> list[tuple[str, str]]: ) for i in range(num_collectors) ] + # Process name -> the role a worker reports its traceback under (see worker.py). + error_roles = { + p.name: ("learner" if num_learners == 1 else f"learner[{rank}]") for rank, p in enumerate(p_learners) + } + error_roles.update({p.name: f"collector[{i}]" for i, p in enumerate(p_collectors)}) # Start each child while the parent is pre-placed on the child's NUMA # node: spawn children re-import torch/simulator before their entry @@ -446,7 +464,12 @@ def _drain_worker_stats() -> None: try: while True: payload = panel_queue.get_nowait() - learner_payloads[payload["rank"]] = payload + notice = payload.get("checkpoint_saved") + if notice is not None: + # Reported at the learner's own step as soon as it arrives. + training_log.event("checkpoint_saved", step=notice["step"], path=notice["path"]) + else: + learner_payloads[payload["rank"]] = payload except Empty: pass @@ -540,6 +563,7 @@ def _handoff_gate_line() -> str: print(f"[motrix.fastsac async] {abort_reason}; aborting.") raise RuntimeError(abort_reason) print(f"[motrix.fastsac async] all {len(expected)} workers ready — training starts") + training_log.event("run_started", step=resume_step, workers=len(expected)) finally: handshake.release() # idempotent; frees workers on an aborted boot # -------------------------------------------------- parent-owned panel @@ -549,12 +573,6 @@ def _handoff_gate_line() -> str: # workers stay peers with no "primary renderer" role. console, live = None, None # opened lazily once real data flows (see below) panel_open_attempted = False - try: - from torch.utils.tensorboard import SummaryWriter - - tb_writer = SummaryWriter(log_dir=str(self._context.run_dir)) - except Exception: - tb_writer = None panel_cpu_sampler = CpuLoadSampler() panel_gpu_sampler = GpuUtilizationSampler() panel_memory_sampler = MemoryUsageSampler() @@ -573,12 +591,14 @@ def _handoff_gate_line() -> str: panel_next_log = ( ((resume_step // max(logging_interval, 1)) + 1) * logging_interval if logging_interval > 0 else 0 ) + # warmup_done edge, detected on the panel's own state. + panel_was_warming = resume_step < cfg.agent.learning_starts def _render_panel(force: bool = False) -> None: nonlocal panel_anchored, panel_start_time, panel_last_render, console, live, panel_open_attempted nonlocal panel_rate_samples nonlocal panel_last_log_time, panel_last_log_step, panel_last_updates, panel_next_log - nonlocal panel_last_weight_version + nonlocal panel_last_weight_version, panel_was_warming # ALWAYS drain first: the quiescence gate below waits for learner # payloads that can only arrive through this drain — draining # after the gate would deadlock the workers on a full queue. @@ -600,7 +620,8 @@ def _render_panel(force: bool = False) -> None: if not force and time.time() - panel_last_render < 1.0: return panel_last_render = time.time() - step = control.collector_steps // num_collectors + collector_steps = control.collector_steps + step = collector_steps // num_collectors if step > resume_step and not panel_anchored: panel_start_time = time.time() panel_last_log_time = panel_start_time @@ -680,49 +701,57 @@ def _render_panel(force: bool = False) -> None: checkpoint_path=primary["checkpoint_path"] if primary else None, ) emit_training_panel(live, stats, title=f"{self._env_name}/motrix.fastsac") - if tb_writer is not None and step >= panel_next_log: - tb_writer.add_scalar("rollout/mean_return", last_stats["return"], step) - tb_writer.add_scalar("rollout/mean_ep_len", last_stats["ep_len"], step) + if panel_was_warming and not warming: + training_log.event("warmup_done", step=step) + panel_was_warming = warming + if step >= panel_next_log: + scalars = { + "rollout/mean_return": last_stats["return"], + "rollout/mean_ep_len": last_stats["ep_len"], + } if sps is not None: - tb_writer.add_scalar("perf/env_steps_per_s", sps, step) - tb_writer.add_scalar( - "perf/updates_per_s", (updates - panel_last_updates) / max(now - panel_last_log_time, 1e-6), step - ) - tb_writer.add_scalar("async/policy_lag", last_stats["policy_lag"], step) + scalars["perf/env_steps_per_s"] = sps + scalars["perf/updates_per_s"] = (updates - panel_last_updates) / max(now - panel_last_log_time, 1e-6) + scalars["async/policy_lag"] = last_stats["policy_lag"] if primary and primary.get("weight_version") is not None: - tb_writer.add_scalar( - "async/weight_publishes_per_iter", - (primary["weight_version"] - panel_last_weight_version) / max(step - panel_last_log_step, 1), - step, - ) + scalars["async/weight_publishes_per_iter"] = ( + primary["weight_version"] - panel_last_weight_version + ) / max(step - panel_last_log_step, 1) panel_last_weight_version = primary["weight_version"] if primary: fills = primary.get("ring_fill") or [] - tb_writer.add_scalar("async/ring_fill", sum(f for f in fills if f is not None), step) + scalars["async/ring_fill"] = sum(f for f in fills if f is not None) if primary.get("weight_version") is not None: - tb_writer.add_scalar("async/weight_version", primary["weight_version"], step) - tb_writer.add_scalar("async/utd", utd, step) + scalars["async/weight_version"] = primary["weight_version"] + scalars["async/utd"] = utd if num_collectors > 1: for i in range(num_collectors): stats_i = per_collector_stats.get(i) if stats_i: - tb_writer.add_scalar(f"async/policy_lag_collector{i}", stats_i["policy_lag"], step) + scalars[f"async/policy_lag_collector{i}"] = stats_i["policy_lag"] if primary: fill = (primary.get("ring_fill") or [None] * num_collectors)[i] if fill is not None: - tb_writer.add_scalar(f"async/ring_fill_collector{i}", fill, step) - tb_writer.add_scalar("perf/collect_ms_per_batch", collector_timing_ms.get("collect", 0.0), step) + scalars[f"async/ring_fill_collector{i}"] = fill + scalars["perf/collect_ms_per_batch"] = collector_timing_ms.get("collect", 0.0) for k, v in collector_timing_detail_ms.items(): - tb_writer.add_scalar(f"perf/collector_{k}_ms", v, step) - tb_writer.add_scalar("perf/learn_ms_total", learn_ms, step) - tb_writer.add_scalar("perf/learn_pct", learn_pct, step) + scalars[f"perf/collector_{k}_ms"] = v + scalars["perf/learn_ms_total"] = learn_ms + scalars["perf/learn_pct"] = learn_pct for k, v in last_stats["env_metrics"].items(): - tb_writer.add_scalar(f"metrics/{k}", v, step) + scalars[f"metrics/{k}"] = v for k, v in last_stats["reward_terms"].items(): - tb_writer.add_scalar(f"reward/{k}", v, step) + scalars[f"reward/{k}"] = v if primary and primary["metrics"] is not None: for k, v in primary["metrics"].items(): - tb_writer.add_scalar(f"train/{k}", v, step) + scalars[f"train/{k}"] = v + training_log.scalars( + scalars, + step, + env_steps=collector_steps * per_collector_envs, + live_extra=lambda: system_scalars(stats), + placeholders=() if last_stats["episodes"] else ROLLOUT_MEANS, + ) panel_last_log_time, panel_last_log_step, panel_last_updates = now, step, updates panel_next_log += logging_interval @@ -735,6 +764,16 @@ def _render_panel(force: bool = False) -> None: for p in crashed: print(f"[motrix.fastsac async] {p.name} crashed (exit {p.exitcode}); stopping.") _drain_child_errors() + crash_step = control.collector_steps // num_collectors + for p in crashed: + error_log = Path("async_errors") / f"{error_roles[p.name]}_error.log" + training_log.event( + "worker_crashed", + step=crash_step, + worker=p.name, + exitcode=p.exitcode, + error_log=error_log.as_posix() if (self._context.run_dir / error_log).exists() else None, + ) break if learners_done: _drain_child_errors() @@ -782,12 +821,6 @@ def _shutdown(processes, grace_s): except Exception: pass wall_time_s = time.perf_counter() - train_started_at - # Env-step accounting follows the collector side: each collector - # step ingests one batch of its own env shard, so the aggregate - # counter times the shard size is the true transition count (the - # learner's global_step is an iteration counter that can drift - # from collection under non-strict UTD). - per_collector_envs = num_envs // num_collectors total_env_steps = control.collector_steps * per_collector_envs annotate_run( self._context.run_dir, @@ -799,11 +832,7 @@ def _shutdown(processes, grace_s): "mean_env_steps_per_s": int(total_env_steps / max(wall_time_s, 1e-9)), }, ) - if tb_writer is not None: - try: - tb_writer.close() - except Exception: - pass + training_log.close() try: panel_queue.close() except Exception: @@ -829,6 +858,10 @@ def _shutdown(processes, grace_s): for i, p in enumerate(p_collectors): if p.exitcode not in (0, None): raise RuntimeError(f"motrix.fastsac async collector {i} process failed with exit code {p.exitcode}") + # The primary learner wrote the final checkpoint before exiting cleanly. + final_path = checkpoints.final_checkpoint_path(self._context.checkpoint_format, self._context.run_dir) + if final_path.exists(): + training_log.event("checkpoint_saved", step=control.collector_steps // num_collectors, path=str(final_path)) # ------------------------------------------------------------------ play def play(self, policy: str) -> None: diff --git a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/worker.py b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/worker.py index 6394d2b9..1774fabb 100644 --- a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/worker.py +++ b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/worker.py @@ -481,6 +481,7 @@ def run_learner_process( rendezvous_file: str | None = None, learner_device: str | None = None, final_stores: list[FinalObsStorage] | None = None, + notify_checkpoints: bool = False, ) -> None: """One learner process; ``rank 0`` additionally owns stats aggregation, logging and checkpointing. @@ -789,6 +790,15 @@ def run_learner_process( ) last_checkpoint_path = str(path) logger.info("saved checkpoint %s", path) + # A notice per save for live metrics, separate from the log-window + # payload: this step's payload already went out above, and a save + # after the last log window would otherwise never reach the parent. + # Never blocks the learner, like the log payloads. + if notify_checkpoints: + try: + panel_queue.put_nowait({"rank": rank, "checkpoint_saved": {"step": step, "path": str(path)}}) + except Full: + logger.warning("panel queue full; checkpoint notice for %s dropped", path) next_save += save_interval if num_learners > 1: diff --git a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/framework.py b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/framework.py index e6d5a080..7b3b2a01 100644 --- a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/framework.py +++ b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/framework.py @@ -30,6 +30,9 @@ def create_trainer(self, context: TrainerContext[FastSacCfg]) -> TrainerBase: return Trainer(context=context) + def total_env_steps(self, config: FastSacCfg, num_envs: int) -> int: + return config.trainer.num_learning_iterations * num_envs + def create_policy_exporter(self) -> OnnxPolicyExporter: from motrix_rl_builtin.fastsac.export import FastSacOnnxExporter diff --git a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/live_metrics.py b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/live_metrics.py new file mode 100644 index 00000000..4409e426 --- /dev/null +++ b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/live_metrics.py @@ -0,0 +1,32 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Host scalars FastSAC adds to its live metrics (see ``wiki/design/live-metrics.md``).""" + +from __future__ import annotations + +from motrix_rl.console import TrainingPanelStats + +_GIB = 1024**3 + +# Rollout means that are a NaN placeholder until the first episode ends. +ROLLOUT_MEANS = ("rollout/mean_return", "rollout/mean_ep_len") + + +def system_scalars(stats: TrainingPanelStats) -> dict[str, float]: + """The fixed ``sys/*`` keys from the samples a panel frame already holds. + + Memory is reported in GiB, like the panel. Per-core CPU load is left out on + purpose: it would spend the service's per-run key budget on large hosts. + """ + scalars: dict[str, float] = {} + if stats.cpu_load is not None: + scalars["sys/cpu_percent"] = stats.cpu_load.utilization_percent + if stats.memory_usage is not None: + scalars["sys/mem_used_gb"] = stats.memory_usage.used_bytes / _GIB + for device in stats.gpu_devices or (): + if device.utilization_percent is not None: + scalars[f"sys/gpu{device.index}/util_percent"] = device.utilization_percent + if device.memory is not None: + scalars[f"sys/gpu{device.index}/mem_used_gb"] = device.memory.used_bytes / _GIB + return scalars diff --git a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/sync/train.py b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/sync/train.py index a0c49c31..995a494b 100644 --- a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/sync/train.py +++ b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/sync/train.py @@ -23,8 +23,10 @@ MemoryUsageSampler, sample_gpu_devices, ) +from motrix_rl.training_logger import TrainingLogger from motrix_rl_builtin.fastsac.agent import FastSacAgent from motrix_rl_builtin.fastsac.config import FastSacCfg +from motrix_rl_builtin.fastsac.live_metrics import ROLLOUT_MEANS, system_scalars from motrix_rl_builtin.fastsac.wrap import FastSacEnvWrap from motrix_rl_builtin.fastsac.wrap_np import FastSacNpEnvWrap from motrix_rl_builtin.fastsac.wrap_torch import FastSacTorchEnvWrap @@ -48,7 +50,6 @@ def __init__( self._render = context.render self._resume_from = context.resume_from self._context = context - self._writer = None # ------------------------------------------------------------------ setup def _device(self) -> torch.device: @@ -88,7 +89,6 @@ def _make_agent(self, env: FastSacEnvWrap) -> FastSacAgent: device=device, action_scale=action_scale, action_bias=action_bias, - writer=self._writer, ) def _set_seed(self, seed: int | None) -> None: @@ -108,80 +108,84 @@ def train(self) -> None: raise ValueError("FastSAC supports only the 'tensorboard' logging backend.") run_dir = self._context.run_dir + training_log = TrainingLogger(run_dir, self._context.metrics) try: - from torch.utils.tensorboard import SummaryWriter - - self._writer = SummaryWriter(log_dir=str(run_dir)) - except Exception: # tensorboard optional - self._writer = None + env = self._make_env(self._context.num_envs, render=self._render) + agent = self._make_agent(env) + num_iterations = self._rlcfg.trainer.num_learning_iterations + + if self._resume_from: + ckpt = torch.load(self._resume_from, map_location=agent.device, weights_only=False) + agent.load_state_dict(ckpt, load_optimizers=True) + print(f"[motrix.fastsac sync] resumed from {self._resume_from} at global_step={agent.global_step}") + training_log.event( + "resumed", step=agent.global_step, path=str(self._resume_from), global_step=agent.global_step + ) + if agent.global_step >= num_iterations: + print( + f"[motrix.fastsac sync] WARNING: checkpoint global_step ({agent.global_step}) >= target " + f"iters ({num_iterations}); nothing to do. Increase trainer iterations in the RL config to " + "train further." + ) - env = self._make_env(self._context.num_envs, render=self._render) - agent = self._make_agent(env) - num_iterations = self._rlcfg.trainer.num_learning_iterations - - if self._resume_from: - ckpt = torch.load(self._resume_from, map_location=agent.device, weights_only=False) - agent.load_state_dict(ckpt, load_optimizers=True) - print(f"[motrix.fastsac sync] resumed from {self._resume_from} at global_step={agent.global_step}") - if agent.global_step >= num_iterations: - print( - f"[motrix.fastsac sync] WARNING: checkpoint global_step ({agent.global_step}) >= target iters " - f"({num_iterations}); nothing to do. Increase trainer iterations in the RL config to train further." + print( + f"[motrix.fastsac sync] training '{self._env_name}' on {agent.device} " + f"num_envs={env.num_envs} iters={num_iterations} (from {agent.global_step})" + ) + training_log.event("run_started", step=agent.global_step) + + def _record_periodic_checkpoint(path) -> None: + # Record each periodic checkpoint as the latest resumable training + # state so an interrupted run can still be resumed / played from the + # manifest, not just after train() completes. + checkpoints.record_checkpoint_artifact( + run_dir, + checkpoints.LATEST_TRAINING_STATE, + path, + checkpoints.TRAINING_STATE, + checkpoint_format=self._context.checkpoint_format, ) - print( - f"[motrix.fastsac sync] training '{self._env_name}' on {agent.device} " - f"num_envs={env.num_envs} iters={num_iterations} (from {agent.global_step})" - ) + self._run_loop( + agent, + env, + training_log, + num_iterations=num_iterations, + logging_interval=self._context.logging.interval, + save_interval=self._context.checkpoint.interval, + record_checkpoint=_record_periodic_checkpoint, + ) - def _record_periodic_checkpoint(path) -> None: - # Record each periodic checkpoint as the latest resumable training - # state so an interrupted run can still be resumed / played from the - # manifest, not just after train() completes. + ckpt_path = checkpoints.final_checkpoint_path(self._context.checkpoint_format, run_dir) + ckpt_path.parent.mkdir(parents=True, exist_ok=True) + torch.save(agent.state_dict(), ckpt_path) checkpoints.record_checkpoint_artifact( run_dir, checkpoints.LATEST_TRAINING_STATE, - path, + ckpt_path, checkpoints.TRAINING_STATE, checkpoint_format=self._context.checkpoint_format, ) - - self._run_loop( - agent, - env, - num_iterations=num_iterations, - logging_interval=self._context.logging.interval, - save_interval=self._context.checkpoint.interval, - record_checkpoint=_record_periodic_checkpoint, - ) - - ckpt_path = checkpoints.final_checkpoint_path(self._context.checkpoint_format, run_dir) - ckpt_path.parent.mkdir(parents=True, exist_ok=True) - torch.save(agent.state_dict(), ckpt_path) - checkpoints.record_checkpoint_artifact( - run_dir, - checkpoints.LATEST_TRAINING_STATE, - ckpt_path, - checkpoints.TRAINING_STATE, - checkpoint_format=self._context.checkpoint_format, - ) - checkpoints.record_checkpoint_artifact( - run_dir, - checkpoints.BEST_POLICY, - ckpt_path, - checkpoints.POLICY, - checkpoint_format=self._context.checkpoint_format, - ) - print(f"[motrix.fastsac sync] saved checkpoint to {ckpt_path}") - if self._writer is not None: - self._writer.close() + checkpoints.record_checkpoint_artifact( + run_dir, + checkpoints.BEST_POLICY, + ckpt_path, + checkpoints.POLICY, + checkpoint_format=self._context.checkpoint_format, + ) + print(f"[motrix.fastsac sync] saved checkpoint to {ckpt_path}") + training_log.event("checkpoint_saved", step=agent.global_step, path=str(ckpt_path)) + finally: + # Close on failure too, so buffered TensorBoard events are flushed. + training_log.close() # -------------------------------------------------------- sync train loop def _run_loop( self, agent: FastSacAgent, env: FastSacEnvWrap, + training_log: TrainingLogger, *, num_iterations: int, logging_interval: int, @@ -250,6 +254,8 @@ def emit_msg(msg: str) -> None: # ---------------- collector phase: decision + env.step + buffer write + bookkeeping t0 = time.perf_counter() warming = local < cfg.learning_starts + if cfg.learning_starts > 0 and local == cfg.learning_starts: + training_log.event("warmup_done", step=agent.global_step) if warming and not is_resume: actions = torch.empty(agent.num_envs, agent.act_dim, device=device).uniform_(-1.0, 1.0) actions = actions * agent.actor.action_scale + agent.actor.action_bias @@ -341,21 +347,28 @@ def emit_msg(msg: str) -> None: checkpoint_path=last_checkpoint_path, ) 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) - self._writer.add_scalar("rollout/env_reward", float(rewards.mean()), agent.global_step) - self._writer.add_scalar("perf/env_steps_per_s", sps, agent.global_step) - self._writer.add_scalar("perf/collect_ms_per_iter", collect_ms, agent.global_step) - self._writer.add_scalar("perf/learn_ms_per_iter", learn_ms, agent.global_step) - self._writer.add_scalar("perf/learn_pct", learn_pct, agent.global_step) - for k, v in env_metrics.items(): - self._writer.add_scalar(f"metrics/{k}", v, agent.global_step) - for k, v in term_means.items(): - self._writer.add_scalar(f"reward/{k}", v, agent.global_step) - if metrics_log is not None: - for k, v in metrics_log.items(): - self._writer.add_scalar(f"train/{k}", v, agent.global_step) + scalars = { + "rollout/mean_return": mean_ret, + "rollout/mean_ep_len": mean_len, + "rollout/env_reward": float(rewards.mean()), + "perf/env_steps_per_s": sps, + "perf/collect_ms_per_iter": collect_ms, + "perf/learn_ms_per_iter": learn_ms, + "perf/learn_pct": learn_pct, + **{f"metrics/{k}": v for k, v in env_metrics.items()}, + **{f"reward/{k}": v for k, v in term_means.items()}, + **{f"train/{k}": v for k, v in (metrics_log or {}).items()}, + } + # Iteration global_step has already collected (the counter + # advances at the end of the loop), so it is the + # (global_step + 1)-th batch of num_envs transitions. + training_log.scalars( + scalars, + agent.global_step, + env_steps=(agent.global_step + 1) * agent.num_envs, + live_extra=lambda: system_scalars(stats), + placeholders=() if recent_returns else ROLLOUT_MEANS, + ) last_log_time, last_log_step = now, agent.global_step term_accum, term_count = {}, 0 t_collect, t_learn = 0.0, 0.0 @@ -365,6 +378,7 @@ 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) + training_log.event("checkpoint_saved", step=agent.global_step, path=str(path)) if console is not None: last_checkpoint_path = str(path) else: @@ -379,7 +393,6 @@ def emit_msg(msg: str) -> None: # ------------------------------------------------------------------ play def play(self, policy: str) -> None: self._set_seed(self._context.seed) - self._writer = None env = self._make_env( self._context.play_num_envs, render=self._render, diff --git a/motrix_rl_builtin/tests/fastsac_async_mocks.py b/motrix_rl_builtin/tests/fastsac_async_mocks.py index 7ecf26e4..c8fe04df 100644 --- a/motrix_rl_builtin/tests/fastsac_async_mocks.py +++ b/motrix_rl_builtin/tests/fastsac_async_mocks.py @@ -1,16 +1,25 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""Shared white-box mocks for the async FastSAC learner's process-free tests. +"""Shared white-box mocks for the async FastSAC trainer's process-free tests. ``make_mock_learner`` builds a ``Learner`` without running ``__init__`` (no agent build, no CUDA): the drain path only needs the rings, a capturing replay-buffer stub and the generation assembler. Used by test_fastsac_async_multi.py and test_fastsac_pipeline_equivalence.py. + +``ThreadSpawnContext`` stands in for the ``spawn`` multiprocessing context of +the parent trainer: workers become threads running a test-provided script, +so the parent's boot barrier, monitor loop and panel run unchanged. Used by +test_fastsac_live_metrics.py. """ from __future__ import annotations +import inspect +import queue +import threading +from collections.abc import Callable from types import SimpleNamespace import torch @@ -34,3 +43,69 @@ def make_mock_learner(rings: list, extends: list | None) -> Learner: learner._staging = None learner._pending_copy = False return learner + + +class ThreadQueue(queue.Queue): + """``queue.Queue`` with the ``close()`` of a multiprocessing queue.""" + + def close(self) -> None: + pass + + +class ThreadProcess: + """A worker "process" that runs ``script(self)`` on a thread; its return value is the exit code.""" + + def __init__(self, script: Callable[[ThreadProcess], int], target, args=(), kwargs=None, name=None) -> None: + self.target = target + self.args = tuple(args) + self.kwargs = dict(kwargs or {}) + self.name = name + self.exitcode: int | None = None + self._script = script + self._thread = threading.Thread(target=self._run, name=name, daemon=True) + + @property + def arguments(self) -> dict: + """The worker entry point's arguments by parameter name.""" + return inspect.signature(self.target).bind(*self.args, **self.kwargs).arguments + + def _run(self) -> None: + self.exitcode = self._script(self) + + def start(self) -> None: + self._thread.start() + + def is_alive(self) -> bool: + return self._thread.is_alive() + + def join(self, timeout: float | None = None) -> None: + self._thread.join(timeout) + + def terminate(self) -> None: + pass + + def kill(self) -> None: + pass + + +class ThreadSpawnContext: + """Thread-backed stand-in for ``mp.get_context("spawn")``. + + ``scripts`` maps the worker entry point (``run_learner_process`` / + ``run_collector_process``) to the script its threads run. + """ + + def __init__(self, scripts: dict[Callable, Callable[[ThreadProcess], int]]) -> None: + self.scripts = scripts + self.processes: list[ThreadProcess] = [] + + def Queue(self, maxsize: int = 0) -> ThreadQueue: # noqa: N802 - mirrors the multiprocessing API + return ThreadQueue(maxsize) + + def Event(self) -> threading.Event: # noqa: N802 - mirrors the multiprocessing API + return threading.Event() + + def Process(self, target, args=(), kwargs=None, name=None) -> ThreadProcess: # noqa: N802 + process = ThreadProcess(self.scripts[target], target, args, kwargs, name) + self.processes.append(process) + return process diff --git a/motrix_rl_builtin/tests/test_builtin_plugin.py b/motrix_rl_builtin/tests/test_builtin_plugin.py index 753f2afa..58690a29 100644 --- a/motrix_rl_builtin/tests/test_builtin_plugin.py +++ b/motrix_rl_builtin/tests/test_builtin_plugin.py @@ -1,6 +1,7 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 +import pytest from motrix_rl_builtin.fastsac.config import FastSacCfg from motrix_rl_builtin.fastsac.framework import FastSacProvider @@ -12,3 +13,11 @@ def test_builtin_plugin_registers_fastsac(): assert frameworks.get_config_type("motrix", "fastsac") is FastSacCfg assert frameworks.supported_train_backends("motrix", "fastsac") == ("torch",) assert isinstance(frameworks.get_agent_provider("motrix", "fastsac", "torch"), FastSacProvider) + + +@pytest.mark.parametrize(("iterations", "num_envs", "expected"), [(1, 1, 1), (100, 64, 6400), (0, 8, 0)]) +def test_fastsac_total_env_steps_counts_one_batch_per_iteration(iterations, num_envs, expected): + cfg = FastSacCfg() + cfg.trainer.num_learning_iterations = iterations + + assert FastSacProvider().total_env_steps(cfg, num_envs) == expected diff --git a/motrix_rl_builtin/tests/test_fastsac_live_metrics.py b/motrix_rl_builtin/tests/test_fastsac_live_metrics.py new file mode 100644 index 00000000..27212ad0 --- /dev/null +++ b/motrix_rl_builtin/tests/test_fastsac_live_metrics.py @@ -0,0 +1,689 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Live metrics emitted by the FastSAC sync and async trainers. + +The TensorBoard tag sets below are a compatibility contract: live scalars reuse +the tags as canonical keys (see wiki/design/live-metrics.md), so renaming a tag +breaks dashboards as well as existing TensorBoard comparisons. +""" + +from __future__ import annotations + +import contextlib +import math +import re +import threading +import time +from pathlib import Path +from types import SimpleNamespace +from typing import NamedTuple + +import numpy as np +import pytest +import torch +import torch.utils.tensorboard +from fastsac_async_mocks import ThreadSpawnContext +from motrix_rl_builtin.fastsac.wrap import FastSacStepResult + +from motrix_rl import checkpoints +from motrix_rl.config import CheckpointConfig, LoggingConfig +from motrix_rl.contracts import TrainerContext +from motrix_rl.metrics import JsonlSink, MetricsLogger +from motrix_rl.system_metrics import CpuLoad, GpuDeviceUsage, MemoryUsage + +SYNC_TB_TAGS = { + "rollout/mean_return", + "rollout/mean_ep_len", + "rollout/env_reward", + "perf/env_steps_per_s", + "perf/collect_ms_per_iter", + "perf/learn_ms_per_iter", + "perf/learn_pct", + "metrics/height", + "reward/alive", + "train/qf_loss", + "train/actor_loss", +} + +ASYNC_TB_TAGS = { + "rollout/mean_return", + "rollout/mean_ep_len", + "perf/env_steps_per_s", + "perf/updates_per_s", + "async/policy_lag", + "async/weight_publishes_per_iter", + "async/ring_fill", + "async/weight_version", + "async/utd", + "async/policy_lag_collector0", + "async/policy_lag_collector1", + "async/ring_fill_collector0", + "async/ring_fill_collector1", + "perf/collect_ms_per_batch", + "perf/collector_env_step_ms", + "perf/collector_sync_ms", + "perf/collector_sync.wait_writer_ms", + "perf/learn_ms_total", + "perf/learn_pct", + "metrics/height", + "reward/alive", + "train/qf_loss", + "train/actor_loss", +} + +SYS_KEY = re.compile(r"^sys/(cpu_percent|mem_used_gb|gpu\d+/(util_percent|mem_used_gb))$") + + +# ---------------------------------------------------------------- shared helpers +class _RecordingSink(JsonlSink): + """The real jsonl sink, also keeping each record so a test can wait on them.""" + + def __init__(self, path) -> None: + super().__init__(path) + self.records: list[dict] = [] + self._lock = threading.Lock() + + def write(self, record) -> None: + super().write(record) + with self._lock: + self.records.append(record) + + def count(self, kind: str) -> int: + with self._lock: + return sum(1 for record in self.records if record["kind"] == kind) + + +class _RecordingWriter: + """SummaryWriter stand-in recording ``(tag, value, step)``.""" + + def __init__(self) -> None: + self.scalars: list[tuple[str, float, int]] = [] + + def add_scalar(self, tag, value, step) -> None: + self.scalars.append((tag, float(value), int(step))) + + def close(self) -> None: + pass + + +def _context(tmp_path, rl_cfg, metrics, *, num_envs, logging_interval, save_interval, resume_from=None): + import motrix_envs # noqa: F401 registers environments used by the trainer + + (tmp_path / "checkpoints").mkdir(exist_ok=True) + return TrainerContext( + run=None, + env_name="cartpole", + run_dir=tmp_path, + checkpoint_dir=tmp_path / "checkpoints", + sim=None, + checkpoint_format="pt", + num_envs=num_envs, + play_num_envs=1, + seed=None, + rl_cfg=rl_cfg, + logging=LoggingConfig(backend="tensorboard", interval=logging_interval), + checkpoint=CheckpointConfig(interval=save_interval), + resume_from=resume_from, + metrics=metrics, + ) + + +def _logger(tmp_path) -> tuple[MetricsLogger, _RecordingSink]: + sink = _RecordingSink(tmp_path / "metrics.jsonl") + return MetricsLogger("unit-run", jsonl=sink), sink + + +def _events(sink: _RecordingSink) -> list[tuple[str, int]]: + return [(record["data"]["name"], record["step"]) for record in sink.records if record["kind"] == "event"] + + +def _paths(sink: _RecordingSink) -> list[tuple[str, int]]: + return [ + (record["data"]["path"], record["step"]) + for record in sink.records + if record["kind"] == "event" and record["data"]["name"] == "checkpoint_saved" + ] + + +ROLLOUT_MEANS = {"rollout/mean_return", "rollout/mean_ep_len"} + + +def _event(sink: _RecordingSink, name: str) -> dict: + (record,) = [r for r in sink.records if r["kind"] == "event" and r["data"]["name"] == name] + return record["data"] + + +def _assert_live_matches_tensorboard( + sink: _RecordingSink, writer: _RecordingWriter, *, no_episode_steps: tuple[int, ...] = () +) -> None: + """Each TB log point has one live point: the same tags and values plus fixed ``sys/*`` keys. + + The one difference: at ``no_episode_steps`` (no episode has ended yet) the + rollout means are a NaN placeholder that TensorBoard keeps and live leaves out. + """ + log_points: list[tuple[int, dict[str, float]]] = [] + for tag, value, step in writer.scalars: + if not log_points or log_points[-1][0] != step or tag in log_points[-1][1]: + log_points.append((step, {})) + log_points[-1][1][tag] = value + records = [r for r in sink.records if r["kind"] == "scalars"] + assert [r["step"] for r in records] == [step for step, _ in log_points] + for record, (step, tensorboard) in zip(records, log_points): + live = record["data"] + placeholders = ROLLOUT_MEANS if step in no_episode_steps else set() + assert {key for key in live if not key.startswith("sys/")} == set(tensorboard) - placeholders + assert all(SYS_KEY.match(key) for key in live if key.startswith("sys/")) + for tag, value in tensorboard.items(): + if tag in placeholders: + continue + assert live[tag] == (None if math.isnan(value) else pytest.approx(value)), tag + + +def test_system_scalars_use_the_fixed_key_list(): + from motrix_rl_builtin.fastsac.live_metrics import system_scalars + + gib = 1024**3 + stats = SimpleNamespace( + cpu_load=CpuLoad( + utilization_percent=37.5, + used_logical_cpus=3.0, + logical_cpu_count=8, + physical_core_count=4, + iowait_percent=0.0, + steal_percent=0.0, + per_core_percent=(37.5,) * 8, + ), + memory_usage=MemoryUsage(used_bytes=3 * gib, total_bytes=16 * gib), + gpu_devices=[ + GpuDeviceUsage(index=0, utilization_percent=55.0, memory=MemoryUsage(used_bytes=2 * gib, total_bytes=gib)), + GpuDeviceUsage(index=1, utilization_percent=None, memory=None), + ], + ) + + assert system_scalars(stats) == { + "sys/cpu_percent": 37.5, + "sys/mem_used_gb": 3.0, + "sys/gpu0/util_percent": 55.0, + "sys/gpu0/mem_used_gb": 2.0, + } + assert system_scalars(SimpleNamespace(cpu_load=None, memory_usage=None, gpu_devices=None)) == {} + + +# ---------------------------------------------------------------- sync trainer +class _SyncEnv: + """Two-step episodes, one reward term and one env metric; rewards turn NaN from step ``diverge_at``.""" + + def __init__(self, num_envs: int, diverge_at: int | None = None) -> None: + self.num_envs = num_envs + self._t = 0 + self._diverge_at = diverge_at + + def reset(self): + return torch.zeros(self.num_envs, 2), torch.zeros(self.num_envs, 3) + + def step(self, actions): + self._t += 1 + n = self.num_envs + diverged = self._diverge_at is not None and self._t >= self._diverge_at + return FastSacStepResult( + obs=torch.zeros(n, 2), + critic_obs=torch.zeros(n, 3), + reward=torch.full((n,), math.nan if diverged else 1.0), + terminated=torch.full((n,), self._t % 2 == 0), + truncated=torch.zeros(n, dtype=torch.bool), + final_obs=None, + info={"Reward": {"alive": np.full(n, 0.5)}, "metrics": {"height": np.full(n, 0.9)}}, + ) + + +class _SyncAgent: + """The agent primitives the sync loop orchestrates.""" + + def __init__(self, num_envs: int, learning_starts: int) -> None: + self.cfg = SimpleNamespace(learning_starts=learning_starts, num_updates=1) + self.device = torch.device("cpu") + self.num_envs = num_envs + self.act_dim = 1 + self.actor = SimpleNamespace(action_scale=torch.ones(1), action_bias=torch.zeros(1)) + self.rb = SimpleNamespace(extend=lambda *columns: None, num_stored=0, buffer_size=4) + self.global_step = 0 + self.window_metrics = SimpleNamespace(stats_and_reset=lambda: {"qf_loss": 0.25, "actor_loss": -1.5}) + + def set_train_mode(self) -> None: + pass + + def act(self, obs, deterministic): + return torch.zeros(self.num_envs, 1) + + def update(self, num_updates) -> None: + pass + + def state_dict(self): + return {"global_step": self.global_step} + + def load_state_dict(self, state, load_optimizers=False) -> None: + self.global_step = state["global_step"] + + +def _train_sync(tmp_path, monkeypatch, *, resume_from=None, diverge_at=None): + from motrix_rl_builtin.fastsac.sync.train import Trainer + + writer = _RecordingWriter() + monkeypatch.setattr(torch.utils.tensorboard, "SummaryWriter", lambda **kwargs: writer) + num_envs = 3 + monkeypatch.setattr(Trainer, "_make_env", lambda self, *args, **kwargs: _SyncEnv(num_envs, diverge_at)) + monkeypatch.setattr(Trainer, "_make_agent", lambda self, env: _SyncAgent(num_envs, learning_starts=3)) + rl_cfg = SimpleNamespace(trainer=SimpleNamespace(num_learning_iterations=8)) + metrics, sink = _logger(tmp_path) + context = _context( + tmp_path, rl_cfg, metrics, num_envs=num_envs, logging_interval=2, save_interval=4, resume_from=resume_from + ) + + Trainer(context=context).train() + return sink, writer + + +def test_sync_trainer_emits_tensorboard_tags_as_live_scalars(tmp_path, monkeypatch): + sink, writer = _train_sync(tmp_path, monkeypatch) + + scalars = [r for r in sink.records if r["kind"] == "scalars"] + assert [r["step"] for r in scalars] == [0, 2, 4, 6] + # Iteration global_step has already collected at its log point: env_steps = (global_step + 1) * num_envs, + # so the last point of 8 iterations reaches total_steps = 8 * num_envs. + assert [r["axes"]["env_steps"] for r in scalars] == [3, 9, 15, 21] + assert {tag for tag, _, _ in writer.scalars} == SYNC_TB_TAGS + _assert_live_matches_tensorboard(sink, writer, no_episode_steps=(0,)) + + +def test_sync_trainer_emits_lifecycle_events(tmp_path, monkeypatch): + sink, _ = _train_sync(tmp_path, monkeypatch) + + assert _events(sink) == [ + ("run_started", 0), + ("warmup_done", 3), # learning_starts iterations in + ("checkpoint_saved", 4), + ("checkpoint_saved", 8), + ] + assert _paths(sink) == [ + (str(tmp_path / "checkpoints" / "model_0000004.pt"), 4), + (str(checkpoints.final_checkpoint_path("pt", tmp_path)), 8), + ] + + +def test_sync_trainer_leaves_placeholder_rollout_means_out_of_live(tmp_path, monkeypatch): + sink, writer = _train_sync(tmp_path, monkeypatch) + + # No episode has ended at step 0: TensorBoard keeps its NaN, live leaves the keys out + # instead of spending their one nonfinite event on the placeholder. + assert all(math.isnan(v) for tag, v, step in writer.scalars if tag in ROLLOUT_MEANS and step == 0) + first = next(r for r in sink.records if r["kind"] == "scalars") + assert first["step"] == 0 and not ROLLOUT_MEANS & set(first["data"]) + assert "nonfinite" not in {name for name, _ in _events(sink)} + + +def test_sync_trainer_flags_a_real_divergence(tmp_path, monkeypatch): + sink, _ = _train_sync(tmp_path, monkeypatch, diverge_at=4) + + nonfinite = { + r["data"]["key"]: r["step"] for r in sink.records if r["kind"] == "event" and r["data"]["name"] == "nonfinite" + } + # Episodes ending at env step 4 and later carry NaN returns (logged at global_step 4). + assert nonfinite["rollout/mean_return"] == 4 + + +def test_sync_trainer_reports_resume(tmp_path, monkeypatch): + resume_from = tmp_path / "resume.pt" + torch.save({"global_step": 4}, resume_from) + + sink, _ = _train_sync(tmp_path, monkeypatch, resume_from=str(resume_from)) + + assert _event(sink, "resumed") == {"name": "resumed", "path": str(resume_from), "global_step": 4} + assert _events(sink) == [ + ("resumed", 4), + ("run_started", 4), + ("checkpoint_saved", 4), + ("warmup_done", 7), # warmup restarts on resume: the replay buffer is empty + ("checkpoint_saved", 8), + ] + assert [(r["step"], r["axes"]["env_steps"]) for r in sink.records if r["kind"] == "scalars"] == [(4, 15), (6, 21)] + + +# ---------------------------------------------------------------- async trainer (parent process) +NUM_ENVS = 6 +NUM_COLLECTORS = 2 +LEARNING_STARTS = 20 + + +class _FastTime: + """``time`` for the async trainer module, running faster so its throttles and sleeps stay short.""" + + def __init__(self, speedup: float) -> None: + self._origin = time.time() + self._speedup = speedup + + def time(self) -> float: + return self._origin + (time.time() - self._origin) * self._speedup + + def sleep(self, seconds: float) -> None: + time.sleep(seconds / self._speedup) + + @staticmethod + def perf_counter() -> float: + return time.perf_counter() + + +class _ProbeEnv: + env = SimpleNamespace( + policy_observation_space=SimpleNamespace(shape=(2,)), + value_observation_space=SimpleNamespace(shape=(3,)), + action_space=SimpleNamespace(shape=(1,)), + ) + action_low = torch.tensor([-1.0]) + action_high = torch.tensor([1.0]) + max_episode_steps = None + + def close(self) -> None: + pass + + +class _Stage(NamedTuple): + """One scripted learner stage. + + The learner moves the per-collector step counters to ``steps``, publishes a + window of collector snapshots (``episodes`` finished per collector, rollout + means ``mean_return``) and, if ``payload``, a learner payload, then saves + one periodic checkpoint per entry of ``saves``. That is the real learner + order: the payload of a step goes out before that step's checkpoint. + """ + + steps: tuple[int, int] + saves: tuple[int, ...] = () + payload: bool = True + episodes: int = 2 + mean_return: float = 1.5 + + +def _snapshot(collector_id: int, stage: _Stage) -> dict: + no_episode = stage.episodes == 0 + return { + "collector_id": collector_id, + "return": math.nan if no_episode else stage.mean_return, + "ep_len": math.nan if no_episode else 12.0, + "episodes": stage.episodes, + "reward_terms": {"alive": 0.5}, + "env_metrics": {"height": 0.9}, + "policy_lag": collector_id + 1, + "timing_ms": {"collect": 2.0, "env_step": 1.0, "sync": 0.5, "sync.wait_writer": 0.2}, + } + + +def _payload(step: int, checkpoint_path: str | None) -> dict: + warming = step < LEARNING_STARTS + return { + "rank": 0, + "step": step, + "sps": 100.0, + "updates": 2 * step, + "metrics": None if warming else {"qf_loss": 0.25, "actor_loss": -1.5}, + "warming": warming, + "learn_ms": 1.0, + "learn_pct": 40.0, + "learner_timing": {}, + "buffer_size": 6, + "buffer_capacity": 12, + "ring_fill": [1, 2], + "weight_version": step, + "checkpoint_path": checkpoint_path, + } + + +class _Workers: + """Scripted learner and collectors. + + The learner drives each ``_Stage`` and, after each payload, waits until + the parent has logged it live (or, with ``live=False``, taken it off the + queue). After the last stage it either + writes the final checkpoint and exits cleanly, or makes ``crash_collector`` + report a traceback and exit with code 1. + """ + + def __init__(self, sink: _RecordingSink, stages, *, crash_collector: int | None = None, live: bool = True) -> None: + from motrix_rl_builtin.fastsac.async_impl.worker import run_collector_process, run_learner_process + + self.sink = sink + self.stages = stages + self.crash_collector = crash_collector + self.live = live + self.exitcodes: dict[int, int] = {} + self.ctx = ThreadSpawnContext({run_learner_process: self.learner, run_collector_process: self.collector}) + + def learner(self, process) -> int: + args = process.arguments + control, handshake = args["control"], args["handshake"] + handshake.report_ready("learner", 0) + handshake.wait_for_start(control) + collectors = [p.arguments for p in self.ctx.processes if p is not process] + logged = self.sink.count("scalars") + checkpoint_path = None + for stage in self.stages: + for collector_id, target in enumerate(stage.steps): + while control.collector_steps_at(collector_id) < target: + control.inc_collector_steps(collector_id) + for collector in collectors: + collector["stats_queue"].put(_snapshot(collector["collector_id"], stage)) + if stage.payload: + args["panel_queue"].put(_payload(control.collector_steps // NUM_COLLECTORS, checkpoint_path)) + logged += 1 + if self.live: + self._wait(lambda: self.sink.count("scalars") >= logged or control.stop) + else: # nothing is logged live; the parent taking the payload is all there is to wait for + self._wait(lambda: args["panel_queue"].empty() or control.stop) + for step in stage.saves: + checkpoint_path = str(Path(args["checkpoint_dir"]) / f"model_{step:07d}.pt") + args["panel_queue"].put({"rank": 0, "checkpoint_saved": {"step": step, "path": checkpoint_path}}) + if self.crash_collector is not None: + args["error_queue"].put((f"collector[{self.crash_collector}]", "Traceback: boom\n")) + self.exitcodes[self.crash_collector] = 1 + self._wait(lambda: control.stop) + return 0 + final_path = checkpoints.final_checkpoint_path("pt", args["run_dir"]) + final_path.parent.mkdir(parents=True, exist_ok=True) + final_path.write_bytes(b"") + return 0 + + def collector(self, process) -> int: + args = process.arguments + control, handshake, collector_id = args["control"], args["handshake"], args["collector_id"] + handshake.report_ready("collector", collector_id) + handshake.wait_for_start(control) + self._wait(lambda: control.stop or collector_id in self.exitcodes) + return self.exitcodes.get(collector_id, 0) + + @staticmethod + def _wait(predicate, timeout: float = 30.0) -> None: + deadline = time.monotonic() + timeout + while not predicate() and time.monotonic() < deadline: + time.sleep(0.005) + + +def _train_async(tmp_path, monkeypatch, workers: _Workers, *, num_iterations: int, resume_from=None, live=True): + from motrix_rl_builtin.fastsac.async_impl import train as async_train + from motrix_rl_builtin.fastsac.async_impl.topology import CollectorInfo, LearnerInfo, TrainerTopology + + cpu = torch.device("cpu") + topology = TrainerTopology( + learners=[LearnerInfo(rank=0, device=cpu, numa_node=None, cpus=[])], + collectors=[ + CollectorInfo(i, 0, NUM_ENVS // NUM_COLLECTORS, cpu, None, [], ring_ipc=False, weight_ipc=False) + for i in range(NUM_COLLECTORS) + ], + ) + writer = _RecordingWriter() + monkeypatch.setattr(torch.utils.tensorboard, "SummaryWriter", lambda **kwargs: writer) + monkeypatch.setattr(async_train, "time", _FastTime(10.0)) + monkeypatch.setattr(async_train, "mp", SimpleNamespace(get_context=lambda method: workers.ctx)) + monkeypatch.setattr(async_train, "resolve_trainer_topology", lambda *args, **kwargs: topology) + monkeypatch.setattr(async_train, "actor_param_numel", lambda *args: 0) + monkeypatch.setattr(async_train, "capture_system_info", lambda **kwargs: {}) + monkeypatch.setattr(async_train, "spawn_placement", lambda *args: contextlib.nullcontext()) + monkeypatch.setattr(async_train, "inherit_log_stdio", lambda *args: contextlib.nullcontext()) + monkeypatch.setattr(async_train.Trainer, "_make_env", lambda self, *args, **kwargs: _ProbeEnv()) + (tmp_path / "metadata.json").write_text("{}") + rl_cfg = SimpleNamespace( + device="cpu", + agent=SimpleNamespace(batch_size=8, learning_starts=LEARNING_STARTS), + trainer=SimpleNamespace( + num_learning_iterations=num_iterations, + async_options=SimpleNamespace( + num_collectors=NUM_COLLECTORS, + num_learners=1, + cpus_per_collector=None, + learner_devices=None, + collector_inference_device="cpu", + ring_capacity=2, + utd_mode="strict", + ), + ), + ) + metrics = MetricsLogger("unit-run", jsonl=workers.sink) if live else MetricsLogger.disabled() + context = _context( + tmp_path, rl_cfg, metrics, num_envs=NUM_ENVS, logging_interval=10, save_interval=0, resume_from=resume_from + ) + return async_train.Trainer(context=context), writer + + +def test_async_trainer_emits_tensorboard_tags_events_and_collector_env_steps(tmp_path, monkeypatch): + sink = _RecordingSink(tmp_path / "metrics.jsonl") + # Uneven collectors in the first stage: 21 batches -> step 10, but 21 * 3 env steps. + workers = _Workers(sink, [_Stage((11, 10)), _Stage((20, 20), saves=(20,)), _Stage((40, 40))]) + trainer, writer = _train_async(tmp_path, monkeypatch, workers, num_iterations=40) + + trainer.train() + + scalars = [r for r in sink.records if r["kind"] == "scalars"] + assert sorted({r["step"] for r in scalars}) == [10, 20, 40] + # env_steps = collector_steps * (num_envs // num_collectors), as performance.total_env_steps + collector_steps = {10: 21, 20: 40, 40: 80} + assert all(r["axes"]["env_steps"] == collector_steps[r["step"]] * 3 for r in scalars) + assert {tag for tag, _, _ in writer.scalars} == ASYNC_TB_TAGS + _assert_live_matches_tensorboard(sink, writer) + + # The checkpoint notice and the step-20 payload can land in one drain: their order is not fixed. + assert sorted(_events(sink), key=lambda event: event[1]) in ( + [("run_started", 0), ("warmup_done", 20), ("checkpoint_saved", 20), ("checkpoint_saved", 40)], + [("run_started", 0), ("checkpoint_saved", 20), ("warmup_done", 20), ("checkpoint_saved", 40)], + ) + assert _event(sink, "run_started")["workers"] == NUM_COLLECTORS + 1 + assert _paths(sink) == [ + (str(tmp_path / "checkpoints" / "model_0000020.pt"), 20), + (str(checkpoints.final_checkpoint_path("pt", tmp_path)), 40), + ] + + # The logger holds a thread and locks: it must never be pickled into a worker. + assert len(workers.ctx.processes) == NUM_COLLECTORS + 1 + for process in workers.ctx.processes: + values = (*process.args, *process.kwargs.values()) + assert not any(isinstance(value, (MetricsLogger, TrainerContext)) for value in values), process.name + + +@pytest.mark.parametrize("live", [True, False]) +def test_learner_sends_checkpoint_notices_only_for_live_metrics(tmp_path, monkeypatch, live): + # A notice is a put on the learner's hot loop; without live metrics nothing reads it. + sink = _RecordingSink(tmp_path / "metrics.jsonl") + workers = _Workers(sink, [_Stage((20, 20))], live=live) + trainer, _ = _train_async(tmp_path, monkeypatch, workers, num_iterations=20, live=live) + + trainer.train() + + (learner,) = [process for process in workers.ctx.processes if "learner" in process.name] + assert learner.kwargs["notify_checkpoints"] is live + + +def test_async_trainer_reports_worker_crash(tmp_path, monkeypatch): + sink = _RecordingSink(tmp_path / "metrics.jsonl") + workers = _Workers(sink, [_Stage((11, 10))], crash_collector=1) + trainer, _ = _train_async(tmp_path, monkeypatch, workers, num_iterations=40) + + with pytest.raises(RuntimeError, match="collector 1"): + trainer.train() + + crash = _event(sink, "worker_crashed") + assert crash == { + "name": "worker_crashed", + "worker": "fastsac-async-collector-1", + "exitcode": 1, + "error_log": "async_errors/collector[1]_error.log", + } + assert (tmp_path / crash["error_log"]).read_text() == "Traceback: boom\n" + assert "checkpoint_saved" not in {name for name, _ in _events(sink)} + + +def test_async_trainer_reports_resume(tmp_path, monkeypatch): + resume_from = tmp_path / "resume.pt" + torch.save({"global_step": 10}, resume_from) + sink = _RecordingSink(tmp_path / "metrics.jsonl") + workers = _Workers(sink, [_Stage((20, 20))]) + trainer, _ = _train_async(tmp_path, monkeypatch, workers, num_iterations=20, resume_from=str(resume_from)) + + trainer.train() + + assert _event(sink, "resumed") == {"name": "resumed", "path": str(resume_from), "global_step": 10} + assert _events(sink)[:3] == [("resumed", 10), ("run_started", 10), ("warmup_done", 20)] + assert {(r["step"], r["axes"]["env_steps"]) for r in sink.records if r["kind"] == "scalars"} == {(20, 120)} + + +def test_async_trainer_emits_live_scalars_without_tensorboard(tmp_path, monkeypatch): + sink = _RecordingSink(tmp_path / "metrics.jsonl") + workers = _Workers(sink, [_Stage((11, 10)), _Stage((20, 20))]) + trainer, _ = _train_async(tmp_path, monkeypatch, workers, num_iterations=20) + + def _unavailable(**kwargs): + raise ImportError("tensorboard is not installed") + + monkeypatch.setattr(torch.utils.tensorboard, "SummaryWriter", _unavailable) + trainer.train() + + scalars = [r for r in sink.records if r["kind"] == "scalars"] + assert sorted({r["step"] for r in scalars}) == [10, 20] + assert {"rollout/mean_return", "train/qf_loss"} <= set(scalars[-1]["data"]) + + +def test_async_trainer_reports_checkpoints_saved_after_the_last_log_window(tmp_path, monkeypatch): + sink = _RecordingSink(tmp_path / "metrics.jsonl") + # Real learner order: the step-10 payload goes out before the step-10 save, then two more + # saves land after it and the loop ends before the next log window (step 20) is reached. + workers = _Workers(sink, [_Stage((10, 10), saves=(10,)), _Stage((14, 14), saves=(12, 14), payload=False)]) + trainer, _ = _train_async(tmp_path, monkeypatch, workers, num_iterations=15) + + trainer.train() + + checkpoint_dir = tmp_path / "checkpoints" + assert _paths(sink) == [ + (str(checkpoint_dir / "model_0000010.pt"), 10), + (str(checkpoint_dir / "model_0000012.pt"), 12), + (str(checkpoint_dir / "model_0000014.pt"), 14), + (str(checkpoints.final_checkpoint_path("pt", tmp_path)), 14), + ] + + +def test_async_trainer_leaves_placeholder_rollout_means_out_of_live(tmp_path, monkeypatch): + sink = _RecordingSink(tmp_path / "metrics.jsonl") + stages = [ + _Stage((10, 10), episodes=0), # nothing finished yet: NaN placeholders + _Stage((20, 20)), + _Stage((30, 30), mean_return=math.nan), # a real divergence + ] + workers = _Workers(sink, stages) + trainer, writer = _train_async(tmp_path, monkeypatch, workers, num_iterations=30) + + trainer.train() + + _assert_live_matches_tensorboard(sink, writer, no_episode_steps=(10,)) + scalars = {r["step"]: r["data"] for r in sink.records if r["kind"] == "scalars"} + assert not ROLLOUT_MEANS & set(scalars[10]) + assert scalars[20]["rollout/mean_return"] == pytest.approx(1.5) + nonfinite = [ + (r["data"]["key"], r["step"]) for r in sink.records if r["kind"] == "event" and r["data"]["name"] == "nonfinite" + ] + assert nonfinite == [("rollout/mean_return", 30)] diff --git a/motrix_rl_builtin/tests/test_fastsac_sync_trainer.py b/motrix_rl_builtin/tests/test_fastsac_sync_trainer.py new file mode 100644 index 00000000..626901eb --- /dev/null +++ b/motrix_rl_builtin/tests/test_fastsac_sync_trainer.py @@ -0,0 +1,47 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch.utils.tensorboard + +from motrix_rl.config import CheckpointConfig, LoggingConfig +from motrix_rl.contracts import TrainerContext + + +def test_sync_train_closes_tensorboard_writer_when_training_fails(tmp_path, monkeypatch): + from motrix_rl_builtin.fastsac.sync.train import Trainer + + import motrix_envs # noqa: F401 registers environments used by the trainer + + writer = Mock() + monkeypatch.setattr(torch.utils.tensorboard, "SummaryWriter", Mock(return_value=writer)) + monkeypatch.setattr(Trainer, "_make_env", Mock()) + monkeypatch.setattr(Trainer, "_make_agent", Mock(return_value=SimpleNamespace(device="cpu", global_step=0))) + + def fail_after_logging(self, agent, env, training_log, **kwargs): + training_log.scalars({"train/qf_loss": 1.0}, 0, env_steps=1) # opens TensorBoard + raise RuntimeError("training failed") + + monkeypatch.setattr(Trainer, "_run_loop", fail_after_logging) + context = TrainerContext( + run=None, + env_name="cartpole", + run_dir=tmp_path, + checkpoint_dir=tmp_path / "checkpoints", + sim=None, + checkpoint_format="pt", + num_envs=1, + play_num_envs=1, + seed=None, + rl_cfg=SimpleNamespace(trainer=SimpleNamespace(num_learning_iterations=1)), + logging=LoggingConfig(backend="tensorboard", interval=1), + checkpoint=CheckpointConfig(interval=0), + ) + + with pytest.raises(RuntimeError, match="training failed"): + Trainer(context=context).train() + + writer.close.assert_called_once_with() diff --git a/wiki/design/index.md b/wiki/design/index.md index 87ba5c81..ed743f04 100644 --- a/wiki/design/index.md +++ b/wiki/design/index.md @@ -21,6 +21,8 @@ 在 RL 集成层引入 `rllib/train_backend/algo` 正交维度的架构：配置注册表、`RlFramework`/`AgentProvider`、`TrainerContext` 与通用 runner、算法配置与模型/memory 工厂、run metadata、checkpoint manifest 与 play 自动发现，以及内置 framework 矩阵（skrl / rslrl / motrix）；明确 `motrix_rl` 控制平面与 `motrix_rl.contracts` 契约、三个 provider 插件包的单向依赖和 entry-point discovery 边界。 - [FastSAC 异构（Collector/Learner 分进程）训练器设计](./fastsac-async-heterogeneous-trainer.md) 把仿真采样与网络训练拆到两进程、经共享内存交换数据的异构 FastSAC 训练器：统一注册为 `motrix.fastsac`，由 `algo.asynchronous` 选择执行拓扑，并包含 SPSC 有界背压 replay 环、update-to-data 比例治理、seqlock 双缓冲权重/normalizer 快照，以及进程生命周期与 checkpoint 兼容。 +- [训练实时指标（Live Metrics）设计](./live-metrics.md) + 设置 `MOTRIX_METRICS_URL` 后训练进程把 `motrix.metrics/v1` 记录同时写入 `run_dir/metrics.jsonl` 并批量推送给兼容的 ingest 服务：公开协议（环境变量、记录格式、状态机、幂等、重试规则）、规范 key 与 FastSAC/SKRL/RSLRL 映射、控制平面 `MetricsLogger` 与生命周期接入、失败语义与测试策略。 - [Framework / Task 配置分离设计](./framework-task-split.md) 训练入口通过 Hydra task group 直接组合环境、算法与运行配置；外部应用使用自己的 Hydra config root，不维护 Python task registry 或额外 config-root 状态。 - [ConfigClass 配置装饰器设计](./configclass.md) diff --git a/wiki/design/live-metrics.md b/wiki/design/live-metrics.md new file mode 100644 index 00000000..418b6df0 --- /dev/null +++ b/wiki/design/live-metrics.md @@ -0,0 +1,311 @@ +# 训练实时指标（Live Metrics）设计 + +## 摘要 + +本文描述 MotrixLab 训练进程把训练指标原生推送到外部 ingest 服务的设计。设置环境变量 `MOTRIX_METRICS_URL` +后，训练进程把 `motrix.metrics/v1` 记录通过 HTTP 批量推送给兼容 `motrix.metrics/v1` 的 ingest 服务；设置 +`MOTRIX_METRICS_JSONL=1` 后，同一份记录还会写入 `run_dir/metrics.jsonl`。两者都未设置时功能完全不生效。两路输出都在 +后台线程写出，训练线程只入队。训练器只通过控制平面的 `TrainingLogger` 一个接口记录指标，由它同时写 TensorBoard（tag +不变）与 live metrics。终端面板保持原样。 + +范围：公开协议 `motrix.metrics/v1`（环境变量、记录格式、请求体、状态机、幂等、重试规则）、规范 key 与各 +provider 的映射、控制平面中的 `MetricsLogger` 与 `TrainingLogger`、FastSAC sync/async 的接入点、失败语义与测试 +策略。非目标：SKRL / RSLRL 接入（见 §5.4）、wandb sink、media 上传、pod sidecar 采集、新增 `logging.backend` 取值、 +服务端实现。 + +RL 控制平面与 provider 插件的分层见 [RL 多算法架构设计](./rl-multi-algorithm-architecture.md)；FastSAC async +的进程结构见 [FastSAC 异构训练器设计](./fastsac-async-heterogeneous-trainer.md)。 + +## 1. 目标与非目标 + +**目标** + +- 训练中以秒级粒度在外部看板上看到进度、吞吐、loss、reward 曲线与关键事件（warmup 结束、checkpoint、崩溃）。 +- 一份协议覆盖所有 RL 算法，跨算法对比统一按 `env_steps` 对齐。本期接入 FastSAC（off-policy）；SKRL / RSLRL 后续 + 接入（§5.4）。 +- 训练器只面对一个日志接口 `TrainingLogger`，不分别调用 TensorBoard 与 live metrics。 +- 只用标准库，不新增依赖；服务不可达、慢或返回错误时，训练行为与耗时不受影响。 +- 开关只认环境变量，`task_config` 结构不变，旧 run 的 `task_config.yaml` 与 `metadata.json` 读取不受影响。 + +**非目标** + +- 不替换 TensorBoard：它是 `TrainingLogger` 的一路输出，tag 不变。终端面板消费结构化的 `TrainingPanelStats`，不经过 + `TrainingLogger`。 +- 本期不接入 SKRL / RSLRL：它们的日志由框架内部的 writer 产生，接入只能拦截框架 writer。后续改为只使用它们的纯算法 + 部分，日志统一走 `TrainingLogger`，届时再接入（§5.4）。 +- 不新增 `logging.backend` 取值。FastSAC sync/async 只接受 `tensorboard`，live metrics 是附加输出。 +- 不做 media 上传、eval 记录生产、wandb sink；协议为它们预留了 kind，但本仓库当前不产生。 +- 不在库内安装信号处理器（见 §6）。 + +## 2. 协议 `motrix.metrics/v1` + +本节是 MotrixLab 侧对协议的唯一定义。服务端是任意兼容 `motrix.metrics/v1` 的 ingest 服务；协议变化需要与服务端和 +外部编排器同步，测试用 golden fixtures 锁定（§7）。 + +### 2.1 环境变量 + +| 变量 | 说明 | +| ------------------------ | ------------------------------------------------------------------------------------------------------------ | +| `MOTRIX_METRICS_URL` | ingest 服务基础地址。设置后启用 http 输出 | +| `MOTRIX_METRICS_JSONL` | `1` / `true` / `yes` / `on` 时启用本地 `run_dir/metrics.jsonl` 输出，默认关闭。与 URL 都未设置时 logger 为 disabled，不建文件、不起线程 | +| `MOTRIX_METRICS_TOKEN` | 写 token，只放在 `Authorization` 请求头，不写入 `task_config.yaml`、`metadata.json` 或 jsonl | +| `MOTRIX_METRICS_RUN_ID` | 满足 `^[A-Za-z0-9._-]{1,200}$`。未设置时由 `run_dir` 相对路径派生：非法字符替换为 `_`，截断到 200 字符 | +| `MOTRIX_METRICS_PROJECT` | 看板上的项目名 | +| `MOTRIX_METRICS_GROUP` | 对比分组名 | +| `MOTRIX_METRICS_META` | JSON 对象，原样放进 running 记录的 `meta`；解析失败时告警并忽略 | + +HTTP 请求使用 `urllib` 默认的 proxy handler，代理与 `no_proxy` 由调用方环境决定。 + +### 2.2 记录格式 + +jsonl 每行一条记录，HTTP 每次发一批记录，两路都启用时是同一份记录： + +```json +{ + "v": 1, + "run_id": "...", + "writer": "host:pid:role", + "seq": 1791798324428939, + "t": 1791798324.43, + "kind": "run", + "step": 12000, + "axes": { "env_steps": 49152000 }, + "data": {} +} +``` + +- `writer` 为 `{hostname}:{pid}:{role}`，训练进程的 role 为 `trainer`。 +- `seq` 在同一 writer 内单调递增，**初值取构造时的 `time.time_ns() // 1000`**，之后每条加 1（加锁）。同一容器重启后 + hostname 与 pid 可能完全相同，从 0 开始会让新记录被服务端按 `(writer, seq)` 当成重复丢弃。 +- `t` 为 Unix 秒（float）。`step` 与 `axes` 只在 scalars / event 上出现。 +- 数值统一转成 float；NaN 与 ±Inf 编码为 `null`。清洗对 `data` 递归生效（包括 running 的 `config`、event 的附加字段）； + numpy 与 torch 标量先转 Python float。jsonl 序列化（在写出线程中）使用 `allow_nan=False` 作兜底，序列化异常吞掉并告警。 + +| kind | data | +| ----------- | -------------------------------------------------------------------------------------------------------------- | +| `run` | `state` 与元信息，字段见 §2.4 | +| `scalars` | `{key: value}`，key 见 §3 | +| `event` | `name` 加附加字段，`name` 取值见下表 | +| `eval` | 预留：`suite`、`task`、`ckpt_step`、`episodes`、`success_rate`、`metrics{}`，本仓库当前不产生 | +| `media` | 预留：run_dir 下的相对路径，本仓库当前不产生 | +| `heartbeat` | 空。只走 HTTP，不写 jsonl；空闲时补发，避免 `torch.compile` 等长时间无输出期间被服务端判为 lost | + +**事件名** + +| name | 时机 | 附加字段 | +| ------------------ | ------------------------------------------------------------------------------------- | ------------------------------ | +| `run_started` | provider 完成启动、进入训练主循环（FastSAC async 为所有 worker ready 之后） | provider 相关，例如 `workers` | +| `warmup_done` | off-policy warmup 结束（`warming` 由 True 变 False） | — | +| `checkpoint_saved` | 周期 checkpoint 与最终 checkpoint 写完 | `path` | +| `resumed` | 从 checkpoint 恢复 | `path`、`global_step` | +| `worker_crashed` | FastSAC async 的 learner / collector 子进程异常退出 | `worker`、`exitcode`、`error_log` | +| `nonfinite` | 某个 scalar key **第一次**出现 NaN/Inf；同一 logger 内每个 key 只发一次 | `key` | + +`run_finished` / `run_failed` 由服务端在状态迁移时合成，客户端不发。 + +### 2.3 端点与请求 + +- `POST {URL}/api/v1/ingest`：请求体 `{"schema": "motrix.metrics/v1", "records": [...]}`，`Content-Type: application/json`， + `Content-Encoding: gzip`；设置了 token 时带 `Authorization: Bearer `。 +- `GET {URL}/healthz`：服务存活探测，客户端不依赖。 +- 每批最多 500 条。 +- 2xx 表示已接收，响应体里单独列出的被拒记录不重试。 +- **重试规则**：只在网络错误、408、429、5xx 时指数退避重试，503 遵守 `Retry-After`；其他 4xx（例如缺 token 的 401） + 丢弃该批并只告警一次，避免无限重试堵塞队列。 + +### 2.4 状态机与 run 记录字段 + +状态单调前进：`pending → running → finished | failed | killed`。训练进程只发 `running` 与三种终态；`pending` 与编排侧 +`failed` 由外部编排器（`writer` 的 role 为 `pipeline`）发出。`lost` 由服务端推导：只有 `running` 超过 120 s 未收到任何 +记录才判为 lost，`pending` 永不变 lost。训练进程的终态可以覆盖编排器写的终态，被覆盖的原因由服务端保存为 +`infra_reason`。 + +| 时机 | `data` 字段 | +| -------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| 开始训练 | `state: "running"`、`project`、`group`、`meta`；`metadata.json` 的字段（env_name、rllib、train_backend、algo、seed、sim、checkpoint_format）；`family`；`job_type: "train"`；`primary_axis: "env_steps"`；`total_steps`；`num_envs`；`run_dir`；`config`；`motrixlab_version` | +| 正常结束 | `state: "finished"` | +| 异常 | `state: "failed"`、`reason`（`异常类型: 消息`，截断到 2 KB；FastSAC async 存在 `async_errors/*.log` 时附上路径） | +| Ctrl+C | `state: "killed"` | + +- `family`：FastSAC 为 `rl.offpolicy`，PPO 为 `rl.onpolicy`（PPO 目前只有 run 生命周期记录，见 §5.4）。 +- `config`：`task_config` 的 resolved dict（与 `task_config.yaml` 同源）。 +- `motrixlab_version`：`importlib.metadata.version("motrix-rl")`。 + +### 2.5 幂等 + +- scalars 由服务端按 `(run_id, key, step)` upsert，同一 step 后到的为准；event 按 `(run_id, writer, seq)` 去重。 +- 两路都启用时，同一条记录（同一 seq）必须同时进入 jsonl 与 HTTP，保证事后用 jsonl 回填时可按 `(writer, seq)` 去重，两边 + 记录逐条一致。jsonl 包含 run / scalars / event 全部 kind（heartbeat 除外）。需要回填保底的部署应同时设置 + `MOTRIX_METRICS_JSONL=1`。 + +## 3. 规范 key 与 provider 映射 + +规范 key 直接沿用 FastSAC 现有 TensorBoard tag，FastSAC 不需要映射。 + +| 规范 key | FastSAC | +| --------------------------------------- | ----------------------------------- | +| `rollout/mean_return` | 同名 | +| `rollout/mean_ep_len` | 同名 | +| `perf/env_steps_per_s` | 同名 | +| `train/` | `qf_loss`、`actor_loss`、`alpha` 等 | +| `reward/`、`metrics/` | 同名 | +| `train/learning_rate`、`train/action_std` | 预留给后续接入的 on-policy 算法 | + +- `reward/*` 的口径是"每步瞬时值在日志窗口内的均值"。 +- `sys/*` 是固定清单，只发到 live，不写 TensorBoard：`sys/cpu_percent`、`sys/mem_used_gb`、`sys/gpu{i}/util_percent`、 + `sys/gpu{i}/mem_used_gb`（内存单位与终端面板一致，为 GiB）。不发每核 CPU 数据，避免大核数机器占用服务端 key 配额。 + +**step 与轴** + +| provider | `step` | `axes.env_steps` | 频率 | +| ------------- | ------------------------------------- | -------------------------------------------------------- | -------------------------- | +| FastSAC sync | `global_step` | `(global_step + 1) * num_envs`（日志点上第 `global_step` 次迭代已采集完） | `logging.interval` | +| FastSAC async | `collector_steps // num_collectors` | `collector_steps * (num_envs // num_collectors)`，与 `performance.total_env_steps` 同口径 | 与 TensorBoard 写出点一致 | + +`running` 记录里的 `total_steps` 由 `AgentProvider.total_env_steps(config, num_envs)` 给出（默认 `None`）：FastSAC 为 +`num_learning_iterations * num_envs`，FastSAC sync 与 `env_steps` 同口径。 +FastSAC async 只是近似：训练按 learner 的迭代数停止，而 `env_steps` 按 collector 计数，非严格 UTD 下两者会漂移， +`num_envs` 不能被 `num_collectors` 整除时也会少计，所以进度可能略超或略低于 100%。看板按进度展示时应截断到 +`[0, 1]`，不能把它当作结束判据（结束以 `run` 状态为准）。 + +## 4. 组件：`motrix_rl/src/motrix_rl/metrics.py` + +控制平面只新增这一个模块，只用标准库（json、gzip、urllib.request、threading、queue、socket）。 + +```python +class MetricsLogger: + @classmethod + def from_env(cls, run_dir: Path, *, role: str = "trainer") -> MetricsLogger: ... + @classmethod + def disabled(cls) -> MetricsLogger: ... + enabled: bool + run_id: str + def run(self, state: str, **data) -> None: ... + def scalars(self, values: Mapping[str, float], step: int, *, env_steps: int | None = None) -> None: ... + def event(self, name: str, step: int | None = None, **data) -> None: ... + def close(self) -> None: ... # 没有后续数据了；立即返回，幂等，异常吞掉 + def join(self, timeout: float) -> None: ... # 可选：等 sink 收尾，最多 timeout +``` + +- **disabled**：所有方法是 no-op，不建文件、不起线程。 +- **`JsonlSink`**：与 TensorBoard 的 `EventFileWriter` 相同，采用异步写出：`write` 只入队，daemon 线程负责 JSON 序列化、 + 追加写 `run_dir/metrics.jsonl`（每条记录一行），每 2 s 批量写一次并 flush 一次，慢盘或网络文件系统不阻塞训练线程。 + 与 TensorBoard 队列满时阻塞调用方不同，队列（10 万条）满时丢最旧的 scalars 并告警，run / event 不丢。`close` 之后写完队列即 + 退出。序列化或 IO 异常吞掉并告警，不打断训练。 +- **`HttpSink`**：daemon 线程 + 单一有界队列；入队一律 `put_nowait`，满了丢最旧的 scalars 并计数告警（run / event + 记录不丢）；每 2 s 或攒满 500 条发一批，请求超时 5 s，重试规则见 §2.3；空闲时补发 heartbeat。 +- 两个 sink 都是具体类，不建基类或注册表。`MetricsLogger._emit` 在训练线程只做数据清洗、分配 seq 与入队；scalars + 在 `scalars()` 中已转成 float，不再二次清洗。 +- **关闭是异步的**：`close()` 只表示没有后续数据，立即返回。每个 sink 按自己的节奏收尾：jsonl 写完队列后退出；http 发完 + 队列后退出，失败按 §2.3 重试，最多重试 `flush_budget`（60 s）后丢弃剩余。训练结束不等待任何 sink；下一个训练任务用 + 新的 logger，与这一个无关。`join(timeout)` 只供需要等待的调用方（测试、进程退出）使用，两个 sink 共用一个截止时间。 +- **进程退出**：sink 是 daemon 线程，解释器退出时会被直接杀掉。`atexit` 钩子 `_join_at_exit` 在此之前关闭仍打开的 + logger 并 `join`，合计最多 `EXIT_TIMEOUT`（5 s），保证训练完立即退出的脚本也能送出最后几条记录（含终态）。健康时 + 只需毫秒。 +- logger 带线程与锁，不可 pickle；FastSAC async 只有父进程持有 logger，worker 子进程不接触。 + +### 4.1 统一日志接口：`motrix_rl/src/motrix_rl/training_logger.py` + +```python +class TrainingLogger: + def __init__(self, log_dir: Path, live: MetricsLogger) -> None: ... + def scalars(self, values: Mapping[str, float], step: int, *, env_steps: int, + live_extra: Callable[[], Mapping[str, float]] | None = None, + placeholders: Collection[str] = ()) -> None: ... + def event(self, name: str, step: int | None = None, **data) -> None: ... + def close(self) -> None: ... +``` + +- 训练器每个日志窗口只调用一次 `scalars`：`values` 以 key 为 tag 写 TensorBoard，同一份值发 live；指标只组装、归约一次。 +- `live_extra` 返回只发 live 的值（例如 `sys/*`，TensorBoard 原本没有），只在 live 启用时才调用，关闭 live 时不做任何 + 额外计算；`placeholders` 列出尚未定义的值（例如第一个 episode + 结束前 NaN 的 rollout 均值）：TensorBoard 照写，live 不发，`nonfinite` 名额留给真正的发散。 +- `event` 只发 live。`close` 只关 TensorBoard（异常吞掉，可放在 `finally`）；live logger 的终态与关闭归控制平面 + （`TrainerHandle`）。 +- TensorBoard writer 在第一次 `scalars` 时才打开（`torch.utils.tensorboard.SummaryWriter`，打不开时告警，只发 live）， + 训练在记录任何指标之前失败不会留下空 event 文件。 + +## 5. 接入点 + +### 5.1 控制平面 + +- **构造**：只在 `runner.create_training_handle` 中用 `MetricsLogger.from_env(run.run_dir)` 构造，经 `metrics=` 参数传给 + `create_trainer_context`。play 路径（`create_run_handle`）一律使用 disabled，避免设置了环境变量时 play 也起线程或 + 改写原训练 run 的 `metadata.json`。 +- **契约字段**：`TrainerContext.metrics` 与 `TrainerHandle.metrics`，均有 disabled 默认值，第三方 provider 不接入也能 + 得到状态记录。`create_trainer` 从 typed context 填入 `TrainerHandle.metrics`。 +- **生命周期**（`TrainerHandle.train()`）：发 `running` → `trainer.train()` → 发 `finished`；`KeyboardInterrupt` 发 + `killed`；`SystemExit` 非 0 与其他异常发 `failed`；最后在 `finally` 中 `close()`，不等待 sink。异常照常 + 向上抛。 +- **run 标注**：`annotate_run(run_dir, live={"run_id": ...})`，只写 run_id，不写 URL 与 token。旧版本读取 + `metadata.json` 时忽略未知顶层 key。 + +### 5.2 FastSAC sync（`fastsac/sync/train.py`） + +- 日志点组装一份 `scalars` 字典（tag 集合逐字不变）交给 `TrainingLogger.scalars`，`sys/*`（复用面板已采好的样本） + 经 `live_extra` 只发 live。 +- 还没有 episode 结束时，`rollout/mean_return`、`rollout/mean_ep_len` 是 NaN 占位，经 `placeholders` 传入：TensorBoard + 照写，live 不发。否则每次训练开局都会把这两个 key 的 `nonfinite` 名额用掉，之后真正的发散反而不报。 +- 事件：`run_started`、`warmup_done`、`checkpoint_saved`（周期与最终）、`resumed`。 +- `TrainingLogger.close()` 放在 `finally` 中，训练异常时也关闭 TensorBoard。 + +### 5.3 FastSAC async（父进程 `fastsac/async_impl/train.py`，learner 只多发 checkpoint 通知） + +- 与 sync 相同，日志窗口经 `TrainingLogger.scalars` 写出；汇总的 episode 数为 0 时 rollout 均值作为 `placeholders`。 + TensorBoard 不可用时 live 照常输出。 +- 事件：所有 worker ready 后发 `run_started`；resume 时发 `resumed`；`warming` 由 True 变 False 时发 `warmup_done`； + primary learner 每写完一个周期 checkpoint，就经 panel 队列单独发一条通知（仅当 live 启用，由 learner 参数 + `notify_checkpoints` 控制；`put_nowait`，队列满时丢弃并告警，不阻塞 learner）（`{"rank", "checkpoint_saved": {"step", + "path"}}`），父进程 drain 到即按其中的 step 发 `checkpoint_saved`。不能借用日志窗口 payload 的 `checkpoint_path`： + learner 在同一轮里先发 payload 后存 checkpoint，路径要等下一个日志窗口才送达，最后一个窗口之后的 checkpoint 会丢。 + 所有 worker 正常退出且最终 checkpoint 存在时，父进程再发一次；监控循环检测到子进程崩溃时发 `worker_crashed`。 + 终态 `failed` 由父进程抛出的异常经 `TrainerHandle` 统一发出。 + +### 5.4 SKRL 与 RSLRL（后续） + +本期不接入。它们的 TensorBoard 由框架内部的 writer 写出，接入只能替换或包裹框架 writer，日志逻辑仍分散在框架里。 +后续计划只使用两者的纯算法部分（更新、网络与损失），rollout、日志窗口与指标由 MotrixLab 自己组织，统一经 +`TrainingLogger` 写出，与 FastSAC 共用一个接口与一套规范 key（`train/learning_rate`、`train/action_std` 已预留）。 + +在此之前，SKRL / RSLRL 训练仍会得到控制平面发出的 run 生命周期记录（`running` / `finished` / `failed` / `killed`， +`TrainerHandle` 与 provider 无关），但没有 scalars 与事件；`total_steps` 为 `None`。 + +## 6. 失败语义 + +- **服务不可达或慢**：入队非阻塞，发送在 daemon 线程中进行，单次请求超时 5 s；训练结束不等待发送，进程退出最多多等 + `EXIT_TIMEOUT`。未送出的记录丢弃；启用了 jsonl 时本地记录仍保留，可事后回填。 +- **服务拒绝**：非重试类 4xx 丢弃该批并告警一次，训练继续。 +- **队列满**：丢最旧的 scalars，计数并告警；run / event 记录保留。 +- **本地写失败或慢盘**：jsonl 在后台线程写出，序列化或 IO 异常吞掉并告警；慢盘只会让队列积压，不阻塞训练，训练结束 + 也不等待；磁盘卡死时进程退出最多多等 `EXIT_TIMEOUT`。 +- **SIGTERM（例如 pod 被删除）**：Python 默认处理直接退出，`finally` 不执行，训练进程不会发出终态，jsonl 也可能来不及 + 被外部归档。库内不安装信号处理器（侵入性过强），由外部编排器写 `failed` 与服务端的 lost 推导兜底。这是接受的缺口。 +- **分布式**：只有 rank0 持有 enabled logger。FastSAC async 只在父进程。 + +## 7. 测试策略 + +测试面向公开契约，不测私有实现，不新增 conftest 与 marker。 + +- `motrix_rl/tests/test_metrics.py`：未设置 URL 时不生效（不建文件、不起线程）；记录字段、writer 格式与 seq + 初值和单调性；NaN / Inf / numpy / torch 值与嵌套 `data` 的清洗；`nonfinite` 每 key 一次；run_id 清洗与截断； + 用标准库 `http.server` 校验 gzip、Bearer 与请求体；500 后补发成功且 seq 不重复；4xx 丢弃；服务不可达时 `close` + 按时返回；heartbeat；与 golden fixtures（`motrix_rl/tests/fixtures/metrics_v1/`，外部编排器的 pending / failed 请求 + 体与 MotrixLab 产出的 running / scalars / event 样例）结构一致。 +- `motrix_rl/tests/test_rl_contracts.py`：成功、异常、Ctrl+C 分别发出 finished、failed（带 reason）、killed，异常照常 + 抛出，`close` 在终态之后执行；磁盘卡住时 `train()` 照常立即返回，记录随后写完。`test_metrics.py`：`close` 立即返回、 + `join` 共用一个截止时间；退出钩子写完从未关闭的 logger 的记录；服务不可达时 http 在 `flush_budget` 内放弃。`test_runner.py`：play 不构造 enabled logger；未设置 `MOTRIX_METRICS_JSONL` 时不产生 jsonl。 + `test_metrics.py` 另覆盖 URL 与 jsonl 开关的组合，以及磁盘卡住时 emit 不阻塞、记录不丢且按批写出。 +- 单元测试只断言行为，不断言耗时：等待一律用 Event 或条件轮询（条件满足立即返回，超时只在失败时出现），重试测试把 + `HttpSink(backoff=...)` 调小，不依赖真实的 1 s 退避。 +- 性能由 `bench/bench_live_metrics.py` 单独度量，不进 CI。同一个合成训练循环（`--workload numpy`：BLAS 释放 GIL， + 接近 torch 训练步；`--workload python`：训练线程全程持有 GIL，后台线程的最坏情况）每 `--log-every` 次迭代经 + `TrainingLogger.scalars` 记录一次，所有场景都开着 TensorBoard，与 off 的差值即 live 本身的代价。场景：off（与重复一次的 + `off, control` 噪声基线）、jsonl、jsonl 慢盘（每批卡 0.5 s）、jsonl 磁盘卡死（永不返回）、jsonl 写入报错、http、 + http 慢服务（0.5 s 才应答）、http 无响应（接受连接但永不应答）、http 拒绝连接、jsonl + http、磁盘卡死 + 无响应。 + 只度量训练循环：报告训练线程上每次 emit 的耗时分位数、循环吞吐相对 off 的开销、落盘与送达的记录数；关闭不计时， + 故障场景只 `close()` 不 `join`，不必每轮等满超时。ingest 服务在独立进程中运行，不与训练进程争 GIL。`--max-overhead <百分比>` 超限时以非零退出码结束。 +- FastSAC：锁定 sync / async 的 TensorBoard tag 集合；scalars key 为 TB tag 加 `sys/*`；`env_steps` 口径；事件； + async worker 进程参数中不出现 logger；learner 的 `notify_checkpoints` 只在 live 启用时为真。 +- `motrix_rl/tests/test_training_logger.py`：一次 `scalars` 同时写 TensorBoard 与 live；`placeholders` 只写 + TensorBoard、`live_extra` 只发 live；只发事件时不打开 TensorBoard；TensorBoard 不可用时 live 照常；live 关闭时只写 + TensorBoard、不产生文件；`close` 只关一次且不抛异常。