trainmeter 0.0.2__py3-none-any.whl

trainmeter/agent/sender.py

@@ -0,0 +1,68 @@
+"""One-way, non-blocking, lossy delivery of records to the supervisor's datagram socket.
+
+Standard library only. It never blocks and never raises: a full buffer drops the record, and
+after repeated failures (no supervisor any more) the sender goes quiet for good.
+"""
+
+from __future__ import annotations
+
+import os
+import socket
+import time
+from typing import Any
+
+from ..records import encode, make_record
+
+SOCK_ENV = "TRAINMETER_SOCK"
+MAX_FAILURES = 20
+
+
+class Sender:
+    def __init__(self, path: str) -> None:
+        self._path = path
+        self._sock: socket.socket | None = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
+        self._sock.setblocking(False)
+        self._failures = 0
+
+    def send(self, kind: str, src: str, d: dict[str, Any], rank: int | None = None) -> None:
+        sock = self._sock
+        if sock is None:
+            return
+        try:
+            data = encode(make_record(kind, src, time.monotonic(), d, pid=os.getpid(), rank=rank))
+            sock.sendto(data.encode("utf-8"), self._path)
+            self._failures = 0
+        except BlockingIOError:
+            return  # buffer full: drop
+        except (ConnectionRefusedError, FileNotFoundError):
+            self._failures += 1
+            if self._failures >= MAX_FAILURES:
+                self.close()
+        except (OSError, TypeError, ValueError):
+            return
+
+    def close(self) -> None:
+        sock, self._sock = self._sock, None
+        if sock is not None:
+            try:
+                sock.close()
+            except OSError:
+                pass
+
+
+_sender: Sender | None = None
+_sender_pid: int | None = None
+
+
+def get_sender() -> Sender | None:
+    """The process's sender, or None when not running under `tm`. Re-created after a fork."""
+    global _sender, _sender_pid
+    path = os.environ.get(SOCK_ENV)
+    if not path:
+        return None
+    if _sender is None or _sender_pid != os.getpid():
+        try:
+            _sender, _sender_pid = Sender(path), os.getpid()
+        except OSError:
+            return None
+    return _sender

trainmeter/agent/taps.py

@@ -0,0 +1,144 @@
+"""Logger taps: the series a loop already sends to wandb, TensorBoard or MLflow.
+
+Each tap wraps the tracker's public logging call, reads the numbers, reports them and forwards the
+call unchanged. It reads only Python and NumPy numbers and 0-dimensional CPU tensors, so a loss
+that still lives on an accelerator is skipped and never converted. The raw key is forwarded and
+the engine maps it to a canonical name.
+
+The tap never changes what the tracker receives, and a failure in the tap disables that one tap
+and warns once. An exception raised by the tracker itself is the user's and propagates untouched.
+
+Verified against wandb 0.30.0 (offline mode), torch 2.14.0 with tensorboard 2.21.0, and
+mlflow-skinny 3.16.1, on CPU. Trackers that bypass these entry points (tensorboardX, a raw event
+writer) are not tapped.
+"""
+
+from __future__ import annotations
+
+import importlib
+import threading
+from collections.abc import Callable, Iterator
+from typing import Any
+
+from ..emit import as_number
+from .wrap import patch
+
+Report = Callable[[str, float, "int | None", str], None]
+Fail = Callable[[str, BaseException], None]
+
+_busy = threading.local()  # a tracker that logs through itself must not be counted twice
+
+
+def _flatten(prefix: str, value: Any, sep: str) -> Iterator[tuple[str, Any]]:
+    if isinstance(value, dict):
+        for k, v in value.items():
+            yield from _flatten(f"{prefix}{sep}{k}" if prefix else str(k), v, sep)
+    else:
+        yield prefix, value
+
+
+def _step(args: tuple, kwargs: dict, position: int, name: str) -> int | None:
+    value = args[position] if len(args) > position else kwargs.get(name)
+    return value if isinstance(value, int) and not isinstance(value, bool) else None
+
+
+def _arg(args: tuple, kwargs: dict, position: int, name: str) -> Any:
+    return args[position] if len(args) > position else kwargs.get(name)
+
+
+def _scalars(pairs: Iterator[tuple[str, Any]]) -> Iterator[tuple[str, float]]:
+    for key, value in pairs:
+        number = as_number(value)
+        if number is not None:
+            yield key, number
+
+
+def _tap(
+    owner: Any,
+    name: str,
+    via: str,
+    extract: Callable[[tuple, dict], tuple[Iterator[tuple[str, float]], int | None]],
+    report: Report,
+    fail: Fail,
+    disabled: Callable[[str], bool],
+) -> bool:
+    def factory(orig: Callable) -> Callable:
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if getattr(_busy, "on", False):
+                return orig(*args, **kwargs)  # called by the tracker itself: already counted
+            _busy.on = True
+            try:
+                if not disabled(via):
+                    try:
+                        pairs, step = extract(args, kwargs)
+                        for key, value in pairs:
+                            report(key, value, step, via)
+                    except Exception as exc:  # noqa: BLE001 - the tap must never cost a run
+                        fail(via, exc)
+                return orig(*args, **kwargs)
+            finally:
+                _busy.on = False
+
+        return wrapper
+
+    return patch(owner, name, factory) is not None
+
+
+# ---- one adapter per tracker --------------------------------------------------------------
+
+
+def attach_wandb(report: Report, fail: Fail, disabled: Callable[[str], bool]) -> None:
+    run_module = importlib.import_module("wandb.sdk.wandb_run")
+
+    def extract(args: tuple, kwargs: dict):
+        data = _arg(args, kwargs, 1, "data")
+        return _scalars(_flatten("", data, ".")), _step(args, kwargs, 2, "step")
+
+    _tap(run_module.Run, "log", "wandb.log", extract, report, fail, disabled)
+
+
+def attach_tensorboard(report: Report, fail: Fail, disabled: Callable[[str], bool]) -> None:
+    writer = importlib.import_module("torch.utils.tensorboard.writer").SummaryWriter
+
+    def add_scalar(args: tuple, kwargs: dict):
+        pair = [(_arg(args, kwargs, 1, "tag"), _arg(args, kwargs, 2, "scalar_value"))]
+        return _scalars(iter(pair)), _step(args, kwargs, 3, "global_step")
+
+    def add_scalars(args: tuple, kwargs: dict):
+        main = _arg(args, kwargs, 1, "main_tag")
+        values = _arg(args, kwargs, 2, "tag_scalar_dict") or {}
+        pairs = ((f"{main}/{k}", v) for k, v in values.items())
+        return _scalars(pairs), _step(args, kwargs, 3, "global_step")
+
+    _tap(writer, "add_scalar", "SummaryWriter.add_scalar", add_scalar, report, fail, disabled)
+    _tap(writer, "add_scalars", "SummaryWriter.add_scalars", add_scalars, report, fail, disabled)
+
+
+def attach_mlflow(report: Report, fail: Fail, disabled: Callable[[str], bool]) -> None:
+    fluent = importlib.import_module("mlflow.tracking.fluent")
+    top = importlib.import_module("mlflow")
+
+    def log_metric(args: tuple, kwargs: dict):
+        pair = [(_arg(args, kwargs, 0, "key"), _arg(args, kwargs, 1, "value"))]
+        return _scalars(iter(pair)), _step(args, kwargs, 2, "step")
+
+    def log_metrics(args: tuple, kwargs: dict):
+        values = _arg(args, kwargs, 0, "metrics") or {}
+        return _scalars(iter(values.items())), _step(args, kwargs, 1, "step")
+
+    for name, extract, via in (
+        ("log_metric", log_metric, "mlflow.log_metric"),
+        ("log_metrics", log_metrics, "mlflow.log_metrics"),
+    ):
+        original = getattr(fluent, name)
+        if _tap(fluent, name, via, extract, report, fail, disabled):
+            if getattr(top, name, None) is original:  # `mlflow.log_metric` is the same function
+                setattr(top, name, getattr(fluent, name))
+
+
+# Module whose import triggers the adapter, and the adapter.
+ADAPTERS: dict[str, Callable[[Report, Fail, Callable[[str], bool]], None]] = {
+    "wandb": attach_wandb,
+    "torch.utils.tensorboard.writer": attach_tensorboard,
+    "mlflow": attach_mlflow,
+}

trainmeter/agent/wrap.py

@@ -0,0 +1,23 @@
+"""Wrapping a callable in place, once, keeping its identity for introspection."""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Callable
+from typing import Any
+
+
+def patch(owner: Any, name: str, factory: Callable[[Callable], Callable]) -> Callable | None:
+    """Replace `owner.name` with `factory(original)`. Returns the original, or None if already
+    wrapped or missing, so a second attach is a no-op."""
+    original = getattr(owner, name, None)
+    if original is None or getattr(original, "_trainmeter_wrapped", False):
+        return None
+    wrapped = factory(original)
+    try:
+        wrapped = functools.wraps(original)(wrapped)
+    except (AttributeError, TypeError):
+        pass
+    wrapped._trainmeter_wrapped = True  # type: ignore[attr-defined]
+    setattr(owner, name, wrapped)
+    return original

trainmeter/catalog.py

@@ -0,0 +1,100 @@
+"""Canonical names for the series a training loop reports, and the aliases that map to them.
+
+The agent forwards the raw key the loop used. The engine maps it here, so a wrong guess never
+changes what was recorded.
+"""
+
+from __future__ import annotations
+
+ALIASES: dict[str, tuple[str, ...]] = {
+    "train_loss": ("train_loss", "train/loss", "loss/train", "training_loss"),
+    "val_loss": (
+        "val_loss",
+        "val/loss",
+        "loss/val",
+        "eval_loss",
+        "eval/loss",
+        "valid_loss",
+        "validation_loss",
+    ),
+    "lr": ("lr", "learning_rate", "train/lr", "train/learning_rate"),
+    "grad_norm": ("grad_norm", "train/grad_norm", "gradient_norm"),
+}
+
+# A bare `loss` means train_loss, but only while no explicit train key has been seen.
+BARE_LOSS = "loss"
+
+_BY_KEY = {alias.lower(): name for name, aliases in ALIASES.items() for alias in aliases}
+
+
+def canonical(key: str) -> str | None:
+    """The canonical series name for a raw key, or None if no alias matches."""
+    return _BY_KEY.get(key.lower())
+
+
+# name, label, unit, group, format. The dashboard renders any series from this list, so a new
+# metric needs no page change. Formats: percent (a fraction shown as %), si (12.3 G), float,
+# sci (1.2e-4), seconds, bytes.
+METRICS: tuple[tuple[str, str, str, str, str], ...] = (
+    ("train_loss", "Train loss", "", "training", "float"),
+    ("val_loss", "Validation loss", "", "training", "float"),
+    ("lr", "Learning rate", "", "training", "sci"),
+    ("grad_norm", "Grad norm", "", "training", "float"),
+    ("tokens_per_s", "Tokens per second", "tok/s", "training", "si"),
+    ("step_time_s", "Step time", "s", "training", "seconds"),
+    ("progress", "Progress", "", "training", "percent"),
+    ("flops_spent", "FLOPs spent (modeled)", "FLOPs", "compute", "si"),
+    ("flops_observed", "FLOPs observed (counters)", "FLOPs", "compute", "si"),
+    ("mfu", "MFU", "", "efficiency", "percent"),
+    ("ofu", "OFU", "", "efficiency", "percent"),
+    ("mfu_ofu_gap", "OFU minus MFU", "", "efficiency", "percent"),
+    ("arithmetic_intensity", "Arithmetic intensity", "FLOP/byte", "efficiency", "si"),
+    ("goodput", "Goodput", "", "time", "percent"),
+    ("host_cpu_cores", "Host CPU, process tree", "cores", "host", "float"),
+    ("host_rss_b", "Host memory (RSS), process tree", "B", "host", "bytes"),
+    ("gpu_util", "GPU-Util", "", "ladder", "percent"),
+    ("sm_active", "SM active", "", "ladder", "percent"),
+    ("tensor_active", "Tensor active", "", "ladder", "percent"),
+)
+
+# Per-GPU series are named `gpu.<index>.<field>`.
+GPU_FIELDS: tuple[tuple[str, str, str, str], ...] = (
+    ("gpu_util", "GPU-Util", "", "percent"),
+    ("sm_active", "SM active", "", "percent"),
+    ("tensor_active", "Tensor active", "", "percent"),
+    ("dram_active", "DRAM active", "", "percent"),
+    ("power_w", "Power", "W", "float"),
+    ("sm_clock_mhz", "SM clock", "MHz", "float"),
+    ("mem_clock_mhz", "Memory clock", "MHz", "float"),
+    ("mem_used_b", "Memory used", "B", "bytes"),
+    ("temp_c", "Temperature", "C", "float"),
+    ("nvlink_rx_bps", "NVLink receive", "B/s", "si"),
+    ("nvlink_tx_bps", "NVLink transmit", "B/s", "si"),
+)
+
+
+def describe() -> dict:
+    """The body of `GET /api/catalog`."""
+    return {
+        "metrics": [
+            {
+                "name": name,
+                "label": label,
+                "unit": unit,
+                "group": group,
+                "format": fmt,
+                "aliases": list(ALIASES.get(name, ())),
+            }
+            for name, label, unit, group, fmt in METRICS
+        ],
+        "gpu_fields": [
+            {"name": name, "label": label, "unit": unit, "format": fmt}
+            for name, label, unit, fmt in GPU_FIELDS
+        ],
+        "x_axes": [
+            {"name": "step", "label": "Step"},
+            {"name": "tokens", "label": "Tokens"},
+            {"name": "flops", "label": "FLOPs spent (modeled)"},
+            {"name": "time", "label": "Wall time"},
+        ],
+    }