nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/backend/faults.py

@@ -0,0 +1,151 @@
+"""The device-fault channel: a kernel on a backend that does not honour `tl.device_assert` reports a contract violation
+through a device word, read at the sync points. Cut out of tensor.py (M2, 2026-10-29), unchanged."""
+
+from __future__ import annotations
+
+import ctypes
+import functools
+import threading
+
+from nbxtensor.backend import capabilities as _caps
+
+
+# Whether this backend's device memory is directly readable at its own
+# `data_ptr()` from the host. The fault channel below reads its status word
+# that way — after a flush that has already happened, so no copy and no
+# recursion through `memcpy`. Unified memory says yes; a discrete heap does
+# not, and must add its row (and a copy path) before the channel is armed
+# on it.
+_BACKEND_MEMORY_IS_HOST_READABLE = {"cuda": False, "hip": False, "metal": True}
+
+
+# ============================================================================
+# DEVICE-SIDE FAULT CHANNEL
+# ============================================================================
+#
+# The portable half of `tl.device_assert`, for backends that do not honour
+# it. A kernel that detects a broken contract stores its fault code into a
+# one-word status buffer; the code is a `tl.constexpr`, so on a backend that
+# DOES trap the wrapper passes 0 and not one instruction is emitted (the
+# CUDA path is unchanged down to the generated PTX bar one unused kernel
+# parameter).
+#
+# Why a device-side word rather than checking the indices on the host:
+# measured 2026-09-10 on this machine, a host readback of the index tensor
+# costs +0.24 ms (1 id) to +0.29 ms (512 ids) per call against an embedding
+# launch of 0.017–0.020 ms — 13x to 17x the op it guards, and flat in the
+# number of indices because the cost is the `flush()` a D2H copy performs,
+# not the data. That is the same trap the D2D measurement of 2026-09-07
+# found (1,086 host-ordered copies in an 8-token decode: 0.007 s -> 4.08 s).
+# The status word costs one conditional store on the failing path and a
+# 4-byte read at a sync that was already happening.
+#
+# It is READ where the host observes device results — after `sync_device()`
+# and after a copy to the host — which is the sticky-error contract
+# `sync_device` already documents: the fault is at or before that point.
+_FAULT_MESSAGES: list = []
+_FAULT_BUFFERS: dict = {}
+_FAULT_LOCK = threading.Lock()
+
+
+def register_device_fault(message: str) -> int:
+    """Reserve a fault code for `message`. Called at import by the kernel
+    that raises it, so codes are stable within a process and never zero
+    (zero is 'clean', and it is what a disarmed kernel is compiled with)."""
+    with _FAULT_LOCK:
+        if message not in _FAULT_MESSAGES:
+            _FAULT_MESSAGES.append(message)
+        return _FAULT_MESSAGES.index(message) + 1
+
+
+def device_fault_code(message: str, *, armed: bool = None) -> int:
+    """The code a kernel should be compiled with: its own where the backend
+    needs the channel, 0 where the backend traps by itself."""
+    if armed is None:
+        armed = not _caps.backend_traps_on_device_assert()
+    return register_device_fault(message) if armed else 0
+
+
+@functools.lru_cache(maxsize=None)
+def device_fault_code_cached(message: str) -> int:
+    """`device_fault_code` memoised per message — it runs at every launch of
+    an armed kernel, and the answer cannot change within a process."""
+    return device_fault_code(message)
+
+
+def fault_channel(message: str, spare: 'NBXTensor') -> tuple:
+    """`(pointer_argument, FAULT_CODE)` for a kernel that carries the channel.
+
+    Where the backend HONOURS `tl.device_assert`, the code is 0 and the
+    generated kernel contains no fault store at all — proven on the emitted
+    code, not assumed: with FAULT_CODE=0 the body carries neither the store nor
+    the reduction that feeds it. The pointer is then a bound-but-unread
+    argument, so `spare` — a tensor the call already owns — is handed over and
+    NOTHING IS ALLOCATED.
+
+    That matters because the fault buffer is never freed by contract (a frozen
+    replay plan records its raw pointer), so allocating it on a backend that
+    can never write to it is a permanent allocation for a dead path. Reported
+    from the other machine 2026-09-11, whose CUDA proof also says the guard
+    itself holds: three kernels refuse an out-of-range index by name, six
+    outputs identical, +0.13% wall.
+
+    `spare` is never dereferenced through this argument on that path; passing
+    the output tensor is therefore safe and costs nothing.
+    """
+    code = device_fault_code_cached(message)
+    if code == 0:
+        return spare, 0
+    return device_fault_buffer(spare._device_idx), code
+
+
+def device_fault_buffer(device_idx: int) -> 'NBXTensor':
+    """This device's zeroed one-word status buffer.
+
+    Handed to the kernel as a TENSOR, not as a raw address: the launcher
+    binds a tensor as a pointer argument, and a bare int arrives as an
+    int64 scalar that `tl.store` refuses.
+
+    Never freed: a kernel launch records this raw pointer, and a frozen
+    replay plan may hold it long after the call that produced it — the same
+    reason `_INT64_ARRAY_CACHE` is never evicted.
+    """
+    buf = _FAULT_BUFFERS.get(device_idx)          # hot path: no lock
+    if buf is None:
+        with _FAULT_LOCK:
+            buf = _FAULT_BUFFERS.get(device_idx)
+            if buf is None:
+                from nbxtensor.dtype.nbx_dtype import NBXDtype  # the tensor module imports this one
+                from nbxtensor.tensor import NBXTensor
+                buf = NBXTensor.zeros((1,), NBXDtype.int32, f"cuda:{device_idx}")
+                _FAULT_BUFFERS[device_idx] = buf
+    return buf
+
+
+def check_device_faults() -> None:
+    """Raise if any kernel reported a broken contract since the last check.
+
+    Reads the status word straight at its own address: this runs only after
+    a flush (a completed sync, or a copy that flushed), and only on a
+    backend whose device memory the host can read directly. The word is
+    cleared as it is read, so one violation raises once.
+    """
+    if not _FAULT_BUFFERS:
+        return
+    if not _caps._backend_capability(
+            _BACKEND_MEMORY_IS_HOST_READABLE, "_BACKEND_MEMORY_IS_HOST_READABLE",
+            "whether the host can read its device memory directly"):
+        return
+    for device_idx, buf in list(_FAULT_BUFFERS.items()):
+        word = ctypes.c_int32.from_address(buf.data_ptr())
+        code = word.value
+        if code:
+            word.value = 0
+            message = (_FAULT_MESSAGES[code - 1]
+                       if 1 <= code <= len(_FAULT_MESSAGES)
+                       else f"unregistered fault code {code}")
+            raise RuntimeError(
+                f"device-side contract violated on device {device_idx}: "
+                f"{message}. The kernel reported it through the fault "
+                f"channel because this backend does not honour "
+                f"`tl.device_assert`; the fault is at or before this point.")

nbxtensor/backend/host_memory.py

@@ -0,0 +1,185 @@
+"""What the machine actually has free, right now.
+
+A budget that describes what the hardware *recommends* is not a measure of
+what is *available*. On unified memory the two are unrelated: this M4 Pro
+reports `recommendedMaxWorkingSetSize` = 18 186 MB and does not lower it by
+one byte when another process holds a third of the machine. A user never has
+an empty machine — a browser, a mail client, sometimes a container or a VM —
+so a plan sized against the recommendation is accepted and then killed by the
+system under memory pressure, mid-render, without a word.
+
+This module measures the other number. It has two consumers and they want
+the same reading:
+
+  * the planner, which must refuse loudly and name BOTH figures rather than
+    start and be killed;
+  * the measurement protocol, because a green obtained on a half-occupied
+    machine is not the same green as one obtained on a free machine, and
+    without this line written at the time nobody can tell them apart later.
+
+Nothing is guessed. Where the platform is not one this module knows how to
+read, `available_mb` is None and `source` says so — the caller then announces
+that it could not be measured, which is honest, rather than substituting a
+number, which is not.
+"""
+from __future__ import annotations
+
+import re
+import subprocess
+import sys
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+_MB = 1024 * 1024
+
+
+@dataclass(frozen=True)
+class MemoryState:
+    """One reading of the machine, at one instant."""
+
+    total_mb: Optional[int] = None
+    available_mb: Optional[int] = None
+    swap_used_mb: Optional[int] = None
+    swap_total_mb: Optional[int] = None
+    #: (command, resident MB), largest first — this is what names the VM or
+    #: the browser holding the memory, without hardcoding either.
+    largest_residents: Tuple[Tuple[str, int], ...] = field(default=())
+    #: How it was read, or why it could not be.
+    source: str = "not measured"
+
+    @property
+    def measured(self) -> bool:
+        return self.available_mb is not None
+
+    def describe(self) -> str:
+        """The one line every measurement cell writes beside its result."""
+        if not self.measured:
+            return f"machine memory: NOT MEASURED ({self.source})"
+        parts = [f"{self.available_mb} MB free"]
+        if self.total_mb:
+            parts.append(f"of {self.total_mb} MB")
+        if self.swap_total_mb:
+            parts.append(f"swap {self.swap_used_mb}/{self.swap_total_mb} MB")
+        if self.largest_residents:
+            biggest = ", ".join(f"{name} {mb} MB"
+                                for name, mb in self.largest_residents)
+            parts.append(f"largest resident: {biggest}")
+        return "machine memory: " + ", ".join(parts)
+
+
+def _run(cmd) -> Optional[str]:
+    try:
+        out = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
+    except (OSError, subprocess.SubprocessError):
+        return None
+    return out.stdout if out.returncode == 0 else None
+
+
+def _largest_residents(count: int = 3) -> Tuple[Tuple[str, int], ...]:
+    """The biggest resident processes, by RSS. Named, never acted on."""
+    out = _run(["ps", "-Ao", "rss,comm", "-r"])
+    if not out:
+        return ()
+    found = []
+    for line in out.splitlines()[1:]:
+        line = line.strip()
+        if not line:
+            continue
+        rss, _, comm = line.partition(" ")
+        if not rss.isdigit():
+            continue
+        mb = int(rss) // 1024
+        if mb <= 0:
+            continue
+        found.append((comm.strip().rsplit("/", 1)[-1], mb))
+        if len(found) >= count:
+            break
+    return tuple(found)
+
+
+def _available_mb_from_vm_stat(vm: str) -> int:
+    """The available figure from a `vm_stat` text (the reader's one parsing, testable on a captured text)."""
+    page = re.search(r"page size of (\d+) bytes", vm)
+    page_size = int(page.group(1)) if page else 4096
+
+    def pages(label: str) -> int:
+        m = re.search(rf"{label}:\s+(\d+)", vm)
+        return int(m.group(1)) if m else 0
+
+    # The reclaimable set, what a new allocation gets back WITHOUT the compressor: free pages, purgeable
+    # pages, and the whole file cache — `File-backed pages`, active or inactive (Activity Monitor's
+    # "Cached Files"); the speculative queue is already inside it (measured on this M4 Pro:
+    # active + inactive + speculative == file-backed + anonymous). Anonymous pages are NOT counted,
+    # active or inactive: they come back only through the compressor and swap. The previous formula
+    # (free + inactive + purgeable + speculative) subtracted every ACTIVE file page — 4.4 GB of the
+    # container just copied when orpheus-3b was refused at placement (2026-09-29, inbox 68).
+    available_pages = (pages("Pages free") + pages("Pages purgeable") + pages("File-backed pages"))
+    return available_pages * page_size // _MB
+
+
+def _macos_state() -> MemoryState:
+    """`vm_stat` for pages, `sysctl` for the totals and for swap.
+
+    Available is free + purgeable + file-backed pages: the file cache is
+    reclaimable on demand whichever queue holds it, and is never subtracted;
+    anonymous, wired and compressed pages are not reclaimable without the
+    compressor and are deliberately excluded (`_available_mb_from_vm_stat`).
+    """
+    vm = _run(["vm_stat"])
+    if not vm:
+        return MemoryState(source="vm_stat unavailable")
+    available_mb = _available_mb_from_vm_stat(vm)
+
+    total_mb = None
+    memsize = _run(["sysctl", "-n", "hw.memsize"])
+    if memsize and memsize.strip().isdigit():
+        total_mb = int(memsize.strip()) // _MB
+
+    swap_used = swap_total = None
+    swap = _run(["sysctl", "-n", "vm.swapusage"])
+    if swap:
+        tot = re.search(r"total\s*=\s*([\d.]+)M", swap)
+        use = re.search(r"used\s*=\s*([\d.]+)M", swap)
+        if tot:
+            swap_total = int(float(tot.group(1)))
+        if use:
+            swap_used = int(float(use.group(1)))
+
+    return MemoryState(total_mb=total_mb, available_mb=available_mb,
+                       swap_used_mb=swap_used, swap_total_mb=swap_total,
+                       largest_residents=_largest_residents(),
+                       source="vm_stat + sysctl")
+
+
+def _linux_state() -> MemoryState:
+    """`MemAvailable` is the kernel's own estimate of what a new allocation
+    can have without swapping — exactly the question, already answered."""
+    try:
+        text = open("/proc/meminfo").read()
+    except OSError:
+        return MemoryState(source="/proc/meminfo unreadable")
+
+    def kb(label: str) -> Optional[int]:
+        m = re.search(rf"^{label}:\s+(\d+) kB", text, re.MULTILINE)
+        return int(m.group(1)) if m else None
+
+    total, avail = kb("MemTotal"), kb("MemAvailable")
+    swap_total, swap_free = kb("SwapTotal"), kb("SwapFree")
+    return MemoryState(
+        total_mb=total // 1024 if total else None,
+        available_mb=avail // 1024 if avail is not None else None,
+        swap_total_mb=swap_total // 1024 if swap_total else None,
+        swap_used_mb=((swap_total - swap_free) // 1024
+                      if swap_total is not None and swap_free is not None
+                      else None),
+        largest_residents=_largest_residents(),
+        source="/proc/meminfo")
+
+
+def memory_state() -> MemoryState:
+    """Read the machine now. Never cached: the whole point is that it moves."""
+    if sys.platform == "darwin":
+        return _macos_state()
+    if sys.platform.startswith("linux"):
+        return _linux_state()
+    return MemoryState(source=f"no reader for platform {sys.platform!r}")