trainmeter 0.0.2__py3-none-any.whl

trainmeter/sources/gpu.py

@@ -0,0 +1,458 @@
+"""GPU sampler: NVML and GPM counters of the selected GPUs, read from outside the training process.
+
+The hardware sits behind the `Backend` seam. `NvmlBackend` talks to `nvidia-ml-py` (the optional
+`gpu` extra, imported lazily); `FakeBackend` serves scripted numbers so everything else runs
+without a GPU. Nothing here initializes CUDA, allocates on a GPU or synchronizes a device: NVML
+is a management library and reading it does not create a context.
+
+Units in a payload: activity and utilization are fractions in [0, 1], memory is in bytes, power
+in watts, clocks in MHz, NVLink and PCIe throughput in bytes per second.
+
+`NvmlBackend` is verified only against fakes. It has NOT run on real hardware: the scale of the
+GPM values (percent, divided by 100 here), the metric ids and the per-generation availability
+are taken from the `nvidia-ml-py` 13.615.71 sources and the NVML headers, and need the
+checklist in `docs/hardware-checklist.md` on a real node.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import threading
+import time
+import warnings
+from collections.abc import Callable, Sequence
+from typing import Any, Protocol
+
+# GPM metric ids (NVML_GPM_METRIC_*), from nvidia-ml-py.
+GPM_SM_UTIL = 2
+GPM_ANY_TENSOR_UTIL = 5
+GPM_DRAM_BW_UTIL = 10
+GPM_PCIE_TX = 20
+GPM_PCIE_RX = 21
+GPM_NVLINK_RX = 60
+GPM_NVLINK_TX = 61
+
+_GPM_FIELDS = {
+    GPM_SM_UTIL: "sm_active",
+    GPM_ANY_TENSOR_UTIL: "tensor_active",
+    GPM_DRAM_BW_UTIL: "dram_active",
+    GPM_PCIE_TX: "pcie_tx_bps",
+    GPM_PCIE_RX: "pcie_rx_bps",
+    GPM_NVLINK_RX: "nvlink_rx_bps",
+    GPM_NVLINK_TX: "nvlink_tx_bps",
+}
+_FRACTIONS = {"sm_active", "tensor_active", "dram_active"}  # NVML reports these in percent
+
+
+class Backend(Protocol):
+    def open(self) -> None:
+        """Raises `GpuUnavailable` with a reason when there is nothing to read."""
+
+    def close(self) -> None: ...
+
+    def versions(self) -> dict[str, str]:
+        """`driver` and `nvml` version strings, whichever are known."""
+
+    def identities(self) -> list[dict[str, Any]]:
+        """Per index: `name`, `uuid`, `memory_total_b`, `max_sm_clock_mhz`."""
+
+    def basic(self, index: int) -> dict[str, Any]:
+        """`gpu_util`, `mem_used_b`, `power_w`, `temp_c`, `sm_clock_mhz`, `mem_clock_mhz`."""
+
+    def gpm_support(self, index: int) -> str | None:
+        """None when GPM works on this device, else the reason it does not."""
+
+    def gpm(self, index: int) -> dict[str, float] | None:
+        """Counters over the window since the previous call, or None on the first call."""
+
+    def processes(self, index: int) -> set[int] | None:
+        """Pids of the compute processes on this GPU, or None when NVML cannot say.
+
+        Inside a container these are host pids, which do not match the container's own.
+        """
+
+
+class GpuUnavailable(Exception):
+    """NVML cannot be used on this node. The message is the reason shown to the owner."""
+
+
+# ---- selection ----------------------------------------------------------------------------
+
+_INDEX_LIST = re.compile(r"\d+(,\d+)*")
+_UUID_LIST = re.compile(r"(GPU|MIG)-[0-9A-Za-z-]+(,(GPU|MIG)-[0-9A-Za-z-]+)*")
+
+
+def valid_gpus_spec(spec: str) -> bool:
+    """`0,1` or a list of GPU UUIDs, the two forms `--gpus` and `CUDA_VISIBLE_DEVICES` share."""
+    return bool(_INDEX_LIST.fullmatch(spec) or _UUID_LIST.fullmatch(spec))
+
+
+def select(
+    identities: Sequence[dict[str, Any]], flag: str | None, cuda_visible: str | None
+) -> tuple[list[int], str, str | None]:
+    """Choose the GPUs to sample: `--gpus`, else `CUDA_VISIBLE_DEVICES`, else all.
+
+    Returns (indices, how, warning). `how` is `flag` (declared), `env` or `all` (both assumed).
+    An unusable `CUDA_VISIBLE_DEVICES` falls back to all GPUs with a warning; an unusable
+    `--gpus` selects nothing, because the owner asked for something specific.
+
+    `CUDA_VISIBLE_DEVICES` numbers devices in CUDA order, `--gpus` and NVML in PCI bus order.
+    They agree when `CUDA_DEVICE_ORDER=PCI_BUS_ID`; the warning says so when it is not set.
+    """
+    everything = list(range(len(identities)))
+    uuids = {ident.get("uuid"): i for i, ident in enumerate(identities)}
+
+    def resolve(spec: str) -> tuple[list[int], str | None]:
+        picked: list[int] = []
+        for part in spec.split(","):
+            part = part.strip()
+            index = int(part) if part.isdigit() else uuids.get(part, uuids.get(part.upper()))
+            if index is None or not 0 <= index < len(identities):
+                return [], f"{part!r} is not a GPU on this node ({len(identities)} found)"
+            picked.append(index)
+        return sorted(set(picked)), None
+
+    if flag is not None:
+        picked, problem = resolve(flag)
+        return picked, "flag", (f"--gpus: {problem}" if problem else None)
+    if cuda_visible is not None and cuda_visible.strip():
+        picked, problem = resolve(cuda_visible)
+        if problem:
+            return everything, "all", f"CUDA_VISIBLE_DEVICES: {problem}; sampling every GPU"
+        note = None
+        if os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID" and cuda_visible[0].isdigit():
+            note = (
+                "CUDA_VISIBLE_DEVICES indices are matched to NVML order; "
+                "set CUDA_DEVICE_ORDER=PCI_BUS_ID to be sure"
+            )
+        return picked, "env", note
+    return everything, "all", None
+
+
+# ---- NVML ---------------------------------------------------------------------------------
+
+
+class NvmlBackend:
+    """NVML and GPM through `nvidia-ml-py`. Not verified on real hardware, see the module doc."""
+
+    def __init__(self) -> None:
+        self._n: Any = None
+        self._handles: list[Any] = []
+        self._prev: dict[int, Any] = {}
+        self._gpm_ok: dict[int, str | None] = {}
+
+    def open(self) -> None:
+        try:
+            import pynvml  # nvidia-ml-py installs the module under this name
+        except ImportError as exc:
+            raise GpuUnavailable(
+                "nvidia-ml-py is not installed (uv add 'trainmeter[gpu]')"
+            ) from exc
+        self._n = pynvml
+        try:
+            pynvml.nvmlInit()
+            count = pynvml.nvmlDeviceGetCount()
+            self._handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(count)]
+        except pynvml.NVMLError as exc:
+            raise GpuUnavailable(f"NVML failed to initialize: {exc}") from exc
+        if not self._handles:
+            raise GpuUnavailable("NVML sees no GPUs on this node")
+
+    def close(self) -> None:
+        n = self._n
+        if n is None:
+            return
+        for sample in self._prev.values():
+            try:
+                n.nvmlGpmSampleFree(sample)
+            except Exception:  # noqa: BLE001
+                pass
+        self._prev.clear()
+        try:
+            n.nvmlShutdown()
+        except Exception:  # noqa: BLE001
+            pass
+        self._n = None
+
+    def versions(self) -> dict[str, str]:
+        n, out = self._n, {}
+        for key, read in (
+            ("driver", lambda: n.nvmlSystemGetDriverVersion()),
+            ("nvml", lambda: n.nvmlSystemGetNVMLVersion()),
+        ):
+            try:
+                out[key] = self._text(read())
+            except n.NVMLError:
+                pass
+        return out
+
+    @staticmethod
+    def _text(value: Any) -> str:
+        return value.decode() if isinstance(value, bytes) else str(value)
+
+    def identities(self) -> list[dict[str, Any]]:
+        n, out = self._n, []
+        for h in self._handles:
+            ident: dict[str, Any] = {
+                "name": self._text(n.nvmlDeviceGetName(h)),
+                "uuid": self._text(n.nvmlDeviceGetUUID(h)),
+            }
+            for key, read in (
+                ("memory_total_b", lambda h=h: n.nvmlDeviceGetMemoryInfo(h).total),
+                (
+                    "max_sm_clock_mhz",
+                    lambda h=h: n.nvmlDeviceGetMaxClockInfo(h, n.NVML_CLOCK_SM),
+                ),
+            ):
+                try:
+                    ident[key] = read()
+                except n.NVMLError:
+                    pass
+            out.append(ident)
+        return out
+
+    def basic(self, index: int) -> dict[str, Any]:
+        n, h = self._n, self._handles[index]
+        out: dict[str, Any] = {}
+
+        def put(key: str, read: Callable[[], Any]) -> None:
+            try:
+                out[key] = read()
+            except n.NVMLError:
+                pass  # this counter is not offered by the device, the key is simply absent
+
+        put("gpu_util", lambda: n.nvmlDeviceGetUtilizationRates(h).gpu / 100)
+        put("mem_used_b", lambda: n.nvmlDeviceGetMemoryInfo(h).used)
+        put("power_w", lambda: n.nvmlDeviceGetPowerUsage(h) / 1000)  # milliwatts
+        put("temp_c", lambda: n.nvmlDeviceGetTemperature(h, n.NVML_TEMPERATURE_GPU))
+        put("sm_clock_mhz", lambda: n.nvmlDeviceGetClockInfo(h, n.NVML_CLOCK_SM))
+        put("mem_clock_mhz", lambda: n.nvmlDeviceGetClockInfo(h, n.NVML_CLOCK_MEM))
+        return out
+
+    def gpm_support(self, index: int) -> str | None:
+        if index in self._gpm_ok:
+            return self._gpm_ok[index]
+        n, reason = self._n, None
+        try:
+            if not n.nvmlGpmQueryDeviceSupport(self._handles[index]).isSupportedDevice:
+                reason = "GPM is not supported by this GPU (it needs Hopper or newer)"
+        except (n.NVMLError, AttributeError) as exc:
+            reason = f"GPM is not available: {exc}"
+        self._gpm_ok[index] = reason
+        return reason
+
+    def processes(self, index: int) -> set[int] | None:
+        n = self._n
+        try:
+            return {
+                int(p.pid) for p in n.nvmlDeviceGetComputeRunningProcesses(self._handles[index])
+            }
+        except (n.NVMLError, AttributeError):
+            return None
+
+    def gpm(self, index: int) -> dict[str, float] | None:
+        n, h = self._n, self._handles[index]
+        new = n.nvmlGpmSampleAlloc()
+        try:
+            n.nvmlGpmSampleGet(h, new)
+        except Exception:
+            n.nvmlGpmSampleFree(new)
+            raise
+        old = self._prev.get(index)
+        self._prev[index] = new
+        if old is None:
+            return None
+        try:
+            request = n.c_nvmlGpmMetricsGet_t()
+            request.version = n.NVML_GPM_METRICS_GET_VERSION
+            request.sample1, request.sample2 = old, new
+            ids = list(_GPM_FIELDS)
+            request.numMetrics = len(ids)
+            for slot, metric_id in enumerate(ids):
+                request.metrics[slot].metricId = metric_id
+            n.nvmlGpmMetricsGet(request)
+            out: dict[str, float] = {}
+            for slot, metric_id in enumerate(ids):
+                metric = request.metrics[slot]
+                if metric.nvmlReturn != n.NVML_SUCCESS:
+                    continue
+                field = _GPM_FIELDS[metric_id]
+                out[field] = metric.value / 100 if field in _FRACTIONS else metric.value
+            return out
+        finally:
+            n.nvmlGpmSampleFree(old)
+
+
+# ---- fake ---------------------------------------------------------------------------------
+
+
+class FakeBackend:
+    """Scripted backend for tests and for developing without a GPU.
+
+    `gpm_rows[index]` is a list of counter dicts served in order, one per `gpm` call after the
+    first; the last row repeats. `open_error` makes `open` raise `GpuUnavailable`.
+    """
+
+    def __init__(
+        self,
+        identities: Sequence[dict[str, Any]] = (),
+        basic: dict[int, dict[str, Any]] | None = None,
+        gpm_rows: dict[int, list[dict[str, float]]] | None = None,
+        gpm_unsupported: str | None = None,
+        open_error: str | None = None,
+        versions: dict[str, str] | None = None,
+        processes: dict[int, set[int]] | None = None,
+    ) -> None:
+        self._versions = versions or {}
+        self._processes = processes
+        self._identities = [dict(i) for i in identities]
+        self._basic = basic or {}
+        self._rows = gpm_rows or {}
+        self._gpm_unsupported = gpm_unsupported
+        self._open_error = open_error
+        self._calls: dict[int, int] = {}
+        self.opened = self.closed = False
+
+    def open(self) -> None:
+        if self._open_error:
+            raise GpuUnavailable(self._open_error)
+        self.opened = True
+
+    def close(self) -> None:
+        self.closed = True
+
+    def versions(self) -> dict[str, str]:
+        return dict(self._versions)
+
+    def identities(self) -> list[dict[str, Any]]:
+        return [dict(i) for i in self._identities]
+
+    def basic(self, index: int) -> dict[str, Any]:
+        return dict(self._basic.get(index, {}))
+
+    def gpm_support(self, index: int) -> str | None:
+        return self._gpm_unsupported
+
+    def processes(self, index: int) -> set[int] | None:
+        return None if self._processes is None else set(self._processes.get(index, set()))
+
+    def gpm(self, index: int) -> dict[str, float] | None:
+        n = self._calls.get(index, 0)
+        self._calls[index] = n + 1
+        rows = self._rows.get(index)
+        if n == 0 or not rows:
+            return None
+        return dict(rows[min(n - 1, len(rows) - 1)])
+
+
+def default_backend() -> Backend:
+    return NvmlBackend()
+
+
+# ---- sampler ------------------------------------------------------------------------------
+
+
+class GpuSampler:
+    """Calls `emit(payload)` about every `interval` seconds with one `sample.gpu` payload.
+
+    A payload is `{"gpus": [...], "selected_by": ..., "window_s": ...}` plus `gpm_unavailable`
+    (a reason) when the counters are missing. When NVML cannot be used at all the one payload
+    is `{"unavailable": reason}`. Every failure is reported once and the thread goes quiet.
+    """
+
+    def __init__(
+        self,
+        emit: Callable[[dict[str, Any]], None],
+        *,
+        gpus: str | None = None,
+        cuda_visible: str | None = None,
+        backend: Backend | None = None,
+        interval: float = 1.0,
+        clock: Callable[[], float] | None = None,
+        pids: Callable[[], set[int]] | None = None,
+    ) -> None:
+        self._pids = pids  # `tm attach`: the GPUs the watched process tree is using
+        self._emit = emit
+        self._flag, self._env = gpus, cuda_visible
+        self._backend = backend if backend is not None else default_backend()
+        self._interval = interval
+        self._clock = clock or time.monotonic
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._run, name="trainmeter-gpu", daemon=True)
+
+    def start(self) -> None:
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._stop.set()
+        if self._thread.is_alive():
+            self._thread.join(timeout=2.0)
+
+    def _run(self) -> None:
+        try:
+            try:
+                self._backend.open()
+            except GpuUnavailable as exc:
+                self._emit({"unavailable": str(exc)})
+                return
+            idents = self._backend.identities()
+            indices, how, warning = self._choose(idents)
+            if warning:
+                warnings.warn(f"trainmeter: {warning}", stacklevel=2)
+            if not indices:
+                self._emit({"unavailable": warning or "no GPU selected"})
+                return
+            self._loop(idents, indices, how)
+        except Exception as exc:  # noqa: BLE001 - monitoring must not raise into anything
+            warnings.warn(f"trainmeter GPU sampler stopped: {exc!r}", stacklevel=2)
+        finally:
+            try:
+                self._backend.close()
+            except Exception:  # noqa: BLE001
+                pass
+
+    def _choose(self, idents: list[dict[str, Any]]) -> tuple[list[int], str, str | None]:
+        """`--gpus`, else `CUDA_VISIBLE_DEVICES`, else (attach) the GPUs in use, else all."""
+        indices, how, warning = select(idents, self._flag, self._env)
+        if self._pids is None or how != "all":
+            return indices, how, warning
+        mine = self._pids()
+        owned = [i for i in range(len(idents)) if mine & (self._backend.processes(i) or set())]
+        if owned:
+            return owned, "process", None
+        return (
+            indices,
+            "all",
+            "no GPU lists a process of the watched program (inside a container NVML shows host "
+            "pids); sampling every GPU, pass --gpus to choose",
+        )
+
+    def _loop(self, idents: list[dict[str, Any]], indices: list[int], how: str) -> None:
+        last = {i: self._clock() for i in indices}
+        sent_versions = False
+        while True:
+            gpus: list[dict[str, Any]] = []
+            reasons: set[str] = set()
+            for i in indices:
+                row: dict[str, Any] = {"index": i, **idents[i], **self._backend.basic(i)}
+                reason = self._backend.gpm_support(i)
+                if reason is None:
+                    now = self._clock()
+                    counters = self._backend.gpm(i)
+                    if counters is not None:
+                        row.update(counters)
+                        row["window_s"] = now - last[i]
+                    last[i] = now
+                else:
+                    reasons.add(reason)
+                gpus.append(row)
+            payload: dict[str, Any] = {"gpus": gpus, "selected_by": how}
+            if not sent_versions:
+                sent_versions = True
+                payload["versions"] = self._backend.versions()
+            if reasons:
+                payload["gpm_unavailable"] = "; ".join(sorted(reasons))
+            self._emit(payload)
+            if self._stop.wait(self._interval):
+                return

trainmeter/sources/host.py

@@ -0,0 +1,135 @@
+"""Host sampler: uptime and the resource use of the child process tree, read from /proc.
+
+Linux only. Where /proc is missing every reader returns None and the sampler stays quiet.
+Nothing here attaches to or traces the child: it only reads files.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+import warnings
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+PROC = Path("/proc")
+
+
+def read_uptime(proc_root: Path = PROC) -> float | None:
+    """Seconds since the host booted."""
+    try:
+        return float((Path(proc_root) / "uptime").read_text().split()[0])
+    except (OSError, ValueError, IndexError):
+        return None
+
+
+def _stat(pid: int, proc_root: Path) -> tuple[int, int, int] | None:
+    """(ppid, utime ticks, stime ticks) of a process, or None if it is gone."""
+    try:
+        text = (proc_root / str(pid) / "stat").read_text()
+        # The command name is in parentheses and may contain spaces, so split after the last one.
+        fields = text[text.rindex(")") + 2 :].split()
+        return int(fields[1]), int(fields[11]), int(fields[12])
+    except (OSError, ValueError, IndexError):
+        return None
+
+
+def _rss_pages(pid: int, proc_root: Path) -> int:
+    try:
+        return int((proc_root / str(pid) / "statm").read_text().split()[1])
+    except (OSError, ValueError, IndexError):
+        return 0
+
+
+def descendants(root_pid: int, proc_root: Path = PROC) -> list[int]:
+    """`root_pid` and every process below it, or an empty list if the root is gone."""
+    proc_root = Path(proc_root)
+    parent_of: dict[int, int] = {}
+    try:
+        names = [p.name for p in proc_root.iterdir() if p.name.isdigit()]
+    except OSError:
+        return []
+    for name in names:
+        stat = _stat(int(name), proc_root)
+        if stat is not None:
+            parent_of[int(name)] = stat[0]
+    if root_pid not in parent_of:
+        return []
+    children: dict[int, list[int]] = {}
+    for pid, ppid in parent_of.items():
+        children.setdefault(ppid, []).append(pid)
+    tree, todo = [], [root_pid]
+    while todo:
+        pid = todo.pop()
+        tree.append(pid)
+        todo.extend(children.get(pid, []))
+    return tree
+
+
+def sample(
+    root_pid: int,
+    proc_root: Path = PROC,
+    clk_tck: int | None = None,
+    page_size: int | None = None,
+) -> dict[str, Any] | None:
+    """One `sample.host` payload for the tree under `root_pid`, or None if it has exited.
+
+    `cpu_s` and `rss_kb` cover the processes alive now. CPU time of processes that already
+    exited is not included.
+    """
+    proc_root = Path(proc_root)
+    tree = descendants(root_pid, proc_root)
+    if not tree:
+        return None
+    tck = clk_tck or os.sysconf("SC_CLK_TCK")
+    page = page_size or os.sysconf("SC_PAGE_SIZE")
+    ticks = pages = 0
+    for pid in tree:
+        stat = _stat(pid, proc_root)
+        if stat is not None:
+            ticks += stat[1] + stat[2]
+            pages += _rss_pages(pid, proc_root)
+    return {
+        "procs": len(tree),
+        "cpu_s": ticks / tck,
+        "rss_kb": pages * page // 1024,
+        "uptime_s": read_uptime(proc_root),
+    }
+
+
+class HostSampler:
+    """Calls `emit(payload)` for the tree under `root_pid` now and then every `interval` seconds."""
+
+    def __init__(
+        self,
+        root_pid: int,
+        emit: Callable[[dict[str, Any]], None],
+        interval: float = 1.0,
+        proc_root: Path = PROC,
+    ) -> None:
+        self._root = root_pid
+        self._emit = emit
+        self._interval = interval
+        self._proc_root = Path(proc_root)
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._run, name="trainmeter-host", daemon=True)
+
+    def start(self) -> None:
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._stop.set()
+        if self._thread.is_alive():
+            self._thread.join(timeout=2.0)
+
+    def _run(self) -> None:
+        try:
+            while True:
+                payload = sample(self._root, self._proc_root)
+                if payload is not None:
+                    self._emit(payload)
+                if self._stop.wait(self._interval):
+                    return
+        except Exception as exc:  # noqa: BLE001 - monitoring must not raise into anything
+            warnings.warn(f"trainmeter host sampler stopped: {exc!r}", stacklevel=2)

trainmeter/supervisor/__init__.py

@@ -0,0 +1 @@
+"""The `tm` supervisor: launcher and run log."""

trainmeter/supervisor/ingest.py

@@ -0,0 +1,86 @@
+"""Receives the agents' datagrams and feeds one consumer: the single writer of state.
+
+The socket lives in a private directory (mode 0700). Malformed datagrams are dropped and
+counted. Nothing here ever raises into the supervisor's main thread.
+"""
+
+from __future__ import annotations
+
+import queue
+import shutil
+import socket
+import tempfile
+import threading
+import warnings
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+from ..records import decode
+
+_STOP = object()
+
+
+class Pipeline:
+    """Ingest thread -> queue -> one consumer thread calling `handle(record)`."""
+
+    def __init__(self, handle: Callable[[dict[str, Any]], None]) -> None:
+        self._handle = handle
+        self._queue: queue.Queue[Any] = queue.Queue()
+        self._dir = Path(tempfile.mkdtemp(prefix="tm-"))  # mkdtemp is 0700
+        self.path = str(self._dir / "events.sock")
+        self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
+        self._sock.bind(self.path)
+        self._sock.settimeout(0.1)
+        self._stop = threading.Event()
+        self.dropped = 0
+        self._ingest = threading.Thread(target=self._receive, name="trainmeter-ingest", daemon=True)
+        self._consumer = threading.Thread(
+            target=self._consume, name="trainmeter-engine", daemon=True
+        )
+
+    def start(self) -> None:
+        self._ingest.start()
+        self._consumer.start()
+
+    def put(self, record: dict[str, Any]) -> None:
+        self._queue.put(record)
+
+    def _receive(self) -> None:
+        while True:
+            try:
+                data = self._sock.recv(65536)
+            except TimeoutError:
+                if self._stop.is_set():
+                    return  # stopped and idle: everything already sent has been read
+                continue
+            except OSError:
+                return
+            try:
+                self._queue.put(decode(data.decode("utf-8")))
+            except (ValueError, UnicodeDecodeError):
+                self.dropped += 1
+
+    def _consume(self) -> None:
+        warned = False
+        while True:
+            item = self._queue.get()
+            if item is _STOP:
+                return
+            try:
+                self._handle(item)
+            except Exception as exc:  # noqa: BLE001 - one bad record must not stop the run
+                if not warned:
+                    warned = True
+                    warnings.warn(f"trainmeter engine skipped a record: {exc!r}", stacklevel=1)
+
+    def close(self) -> None:
+        """Drain what is already queued, then stop and remove the socket directory."""
+        self._stop.set()
+        self._ingest.join(timeout=2.0)
+        self._queue.put(_STOP)
+        self._consumer.join(timeout=5.0)
+        try:
+            self._sock.close()
+        finally:
+            shutil.rmtree(self._dir, ignore_errors=True)