diff --git a/atom/model_ops/engram/host_memory.py b/atom/model_ops/engram/host_memory.py index 0dae35d5b6..a49a87206c 100644 --- a/atom/model_ops/engram/host_memory.py +++ b/atom/model_ops/engram/host_memory.py @@ -12,6 +12,9 @@ page runs (`_PageRuns`): ranges are widened to page boundaries, so ranges stored back to back (a layer's scales right before its values) share a page, and HIP refuses to register a range overlapping one already registered. +Gaps are registered in 1 GiB pieces, and a table every DP-attention rank +shares registers and unregisters under a per-checkpoint file lock +(`_register_lock_file`), one process at a time. Every refused or failed HIP call clears HIP's last error (`_clear_hip_error`): HIP also leaves it as the runtime's sticky last error, which torch raises at @@ -33,14 +36,19 @@ during a capture is skipped, with a warning, rather than breaking it. """ +import contextlib import ctypes +import dataclasses import errno +import fcntl import functools +import hashlib import logging import os import resource import sys import threading +import time from concurrent.futures import ThreadPoolExecutor, wait import numpy as np @@ -77,6 +85,29 @@ # source is still mapped before it is dropped (`_source_dropper`). _CHUNK_BYTES = 256 << 20 _MAX_COPY_THREADS = 16 +# Largest single cudaHostRegister. One ioctl covering a whole Engram table +# does not return when several ranks issue it together. +_REG_CHUNK_BYTES = 1 << 30 +# How long a waiter tolerates a registration lock whose holder shows no +# progress (the holder touches the file after every faulted-in piece and every +# registered chunk) before giving up. Once a path has timed out it is never +# waited on again in this process. +_REG_LOCK_TIMEOUT_S = 600 +_logged_lock_paths: set[str] = set() +_timed_out_lock_paths: set[str] = set() +# Descriptors of the registration lock files this process holds. Fault-in +# worker threads touch them, so the list and each touch are under the lock: a +# descriptor leaves the list before it is closed and is never touched after. +_held_lock_fds: list[int] = [] +_held_lock_fds_lock = threading.Lock() + + +def _lock_heartbeat() -> None: + """Touch the held registration lock files, so waiters see progress.""" + with _held_lock_fds_lock: + for fd in _held_lock_fds: + with contextlib.suppress(OSError): + os.utime(fd) @functools.cache @@ -513,7 +544,7 @@ def pin_rows( if how == "copy": out = _copy(allocator, sources, leases) else: - out = _register(sources, leases) + out = _register(sources, leases, shared=shared) except Exception: # any failure is a fallback logger.warning("engram: page-locking the UVA shard failed", exc_info=True) out = None @@ -566,6 +597,148 @@ def _read_only(lo: int, hi: int) -> bool: return False +def _reg_chunk_bytes() -> int: + """Largest single `cudaHostRegister`, in bytes.""" + return _REG_CHUNK_BYTES + + +def _mapping_path(addr: int) -> str | None: + """Regular file mapped at `addr` (`/proc/self/maps`), or None.""" + try: + with open("/proc/self/maps") as maps: + for line in maps: + fields = line.split(maxsplit=5) + lo, hi = (int(x, 16) for x in fields[0].split("-")) + if not lo <= addr < hi: + continue + name = fields[5].strip() if len(fields) > 5 else "" + if name.startswith("/") and not name.startswith("/dev/"): + return name + return None + except (OSError, ValueError, IndexError): + return None + return None + + +def _registration_lock_path(sources) -> str: + """Lock file for `sources`. + + `ATOM_ENGRAM_REG_LOCK`, when set. Otherwise one file per checkpoint under + `/dev/shm`, so two servers loading different checkpoints do not share it. + """ + explicit = envs.ATOM_ENGRAM_REG_LOCK + if explicit: + return explicit + mapped = _mapping_path(sources[0].data_ptr()) if sources else None + key = os.path.realpath(mapped) if mapped else "anonymous" + digest = hashlib.sha256(key.encode()).hexdigest()[:16] + return f"/dev/shm/engram-uva-{digest}.lock" + + +def _flock_logging_wait(fd: int, path: str) -> None: + """Exclusive `flock` on `fd`. + + A wait is logged. It is bounded by the holder's progress, not by its + total hold time: the holder touches the file after every faulted-in piece + and every registered chunk, and only `_REG_LOCK_TIMEOUT_S` without a + touch raises `TimeoutError` (an `OSError` subclass; callers must not + swallow it with `except OSError`). + """ + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + return + except BlockingIOError: + pass + logger.info("engram: waiting for %s; another process is page-locking", path) + start = time.monotonic() + wait_from = time.time() + while True: + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + last = max(wait_from, os.fstat(fd).st_mtime) + if time.time() - last >= _REG_LOCK_TIMEOUT_S: + logger.error( + "engram: gave up on %s: its holder made no progress for " + "%d s; not waiting on it again in this process", + path, + _REG_LOCK_TIMEOUT_S, + ) + raise TimeoutError( + f"engram registration lock {path}: holder made no progress " + f"for {_REG_LOCK_TIMEOUT_S}s" + ) + time.sleep(0.05) + else: + logger.info("engram: took %s after %.1f s", path, time.monotonic() - start) + return + + +@contextlib.contextmanager +def _register_lock_file(path: str): + """Hold the registration lock file `path`. + + A file that cannot be opened or locked is logged and skipped. The `yield` + is outside the `except`, so an error during registration is not chained + onto that failure. A holder that makes no progress for + `_REG_LOCK_TIMEOUT_S` raises `TimeoutError`, and after that the path is + never waited on again in this process: every later call raises at once. + """ + if path in _timed_out_lock_paths: + raise TimeoutError(f"engram registration lock {path} timed out earlier") + fd = None + held = False + open_error = None + try: + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o666) + except OSError as exc: + open_error = exc + if open_error is not None: + logger.warning( + "engram: cannot open registration lock %s (%s); page-locking without it", + path, + open_error, + ) + else: + if path not in _logged_lock_paths: + _logged_lock_paths.add(path) + logger.info("engram: page-locking serialized on %s", path) + try: + _flock_logging_wait(fd, path) + except TimeoutError: + _timed_out_lock_paths.add(path) + os.close(fd) + fd = None + raise + except OSError as exc: + logger.warning( + "engram: cannot lock %s (%s); page-locking without it", path, exc + ) + else: + held = True + with _held_lock_fds_lock: + _held_lock_fds.append(fd) + _lock_heartbeat() + try: + yield + finally: + if held: + with _held_lock_fds_lock: + _held_lock_fds.remove(fd) + fcntl.flock(fd, fcntl.LOCK_UN) + if fd is not None: + os.close(fd) + + +@dataclasses.dataclass +class _Run: + """One in-place registration `[base, end)` and its references.""" + + end: int + refs: int + + class _PageRuns: """This process's in-place page-locked runs, reference-counted, each `mlock`ed for as long as it is registered. @@ -591,50 +764,88 @@ class _PageRuns: """ def __init__(self): - self._runs: dict[int, list[int]] = {} # base -> [end, references] + self._runs: dict[int, _Run] = {} self._lock = threading.RLock() def acquire(self, rt, lo: int, hi: int, flags: int = 0) -> list[int] | None: """References on the runs covering `[lo, hi)` (page-aligned), the - uncovered gaps locked and registered first; None, holding nothing, if - HIP refuses one. A read-only range (`flags`) is faulted in even where - its lock is refused.""" + uncovered gaps locked and registered first, in `_REG_CHUNK_BYTES` + pieces; None, holding nothing, if HIP refuses one. A read-only range + (`flags`) is faulted in even where its lock is refused. Rollback goes + through `_release_locked` and does not take the registration lock.""" with self._lock: held: list[int] = [] + created = created_bytes = 0 cursor = lo overlapping = sorted( - b for b, (e, _) in self._runs.items() if b < hi and e > lo + b for b, run in self._runs.items() if b < hi and run.end > lo ) for base in [*overlapping, None]: gap_end = hi if base is None else base if cursor < gap_end: read_only = bool(flags & _HOST_REGISTER_READ_ONLY) _fault_and_lock(cursor, gap_end, read_only) - if int(rt.cudaHostRegister(cursor, gap_end - cursor, flags)): - _clear_hip_error() - _libc().munlock(cursor, gap_end - cursor) - self.release(rt, held) - return None - self._runs[cursor] = [gap_end, 1] - held.append(cursor) + c = cursor + while c < gap_end: + e = min(gap_end, c + _reg_chunk_bytes()) + logger.debug( + "engram: cudaHostRegister %.2f GiB", (e - c) / 2**30 + ) + if int(rt.cudaHostRegister(c, e - c, flags)): + _clear_hip_error() + _libc().munlock(c, gap_end - c) + self._release_locked(rt, held) + return None + self._runs[c] = _Run(e, 1) + held.append(c) + created += 1 + created_bytes += e - c + _lock_heartbeat() + c = e if base is None: break - self._runs[base][1] += 1 + self._runs[base].refs += 1 held.append(base) - cursor = max(cursor, self._runs[base][0]) + cursor = max(cursor, self._runs[base].end) + if created_bytes >= _REG_CHUNK_BYTES: + logger.info( + "engram: registered %.2f GiB in-place as %d userptr BO(s)", + created_bytes / 2**30, + created, + ) return held - def release(self, rt, bases) -> None: + def release(self, rt, bases, lock_path: str | None = None) -> None: + """Drop references taken by `acquire`. Lock order: the registration + lock file `lock_path` (a shared registration), then `self._lock`. + + Best effort: if the lock file times out the runs are released without + it, so they never stay registered and locked for the process's life. + """ + if lock_path: + try: + with _register_lock_file(lock_path), self._lock: + self._release_locked(rt, bases) + return + except TimeoutError: + logger.warning( + "engram: unregistering without %s, which timed out", lock_path + ) with self._lock: - for base in bases: - run = self._runs[base] - run[1] -= 1 - if run[1] == 0: - del self._runs[base] - _libc().munlock(base, run[0] - base) - if int(rt.cudaHostUnregister(base) or 0): - _clear_hip_error() - logger.warning("engram: cudaHostUnregister failed") + self._release_locked(rt, bases) + + def _release_locked(self, rt, bases) -> None: + """`release` with `self._lock` already held by the caller.""" + for base in bases: + run = self._runs[base] + run.refs -= 1 + if run.refs == 0: + del self._runs[base] + _libc().munlock(base, run.end - base) + if int(rt.cudaHostUnregister(base) or 0): + _clear_hip_error() + logger.warning("engram: cudaHostUnregister failed") + _lock_heartbeat() def __len__(self) -> int: return len(self._runs) @@ -661,16 +872,27 @@ def _fault_and_lock(lo: int, hi: int, populate: bool) -> None: and registration faults the pages as before. A refused lock warns once per process and the range is still faulted in. """ - libc = _libc() - if libc.mlock2(lo, hi - lo, _MLOCK_ONFAULT) != 0: + if _libc().mlock2(lo, hi - lo, _MLOCK_ONFAULT) != 0: _warn_unlocked(ctypes.get_errno(), hi - lo) - if not populate: - return + if populate: + _populate(lo, hi) + + +def _populate(lo: int, hi: int) -> None: + """`MADV_POPULATE_READ` over `[lo, hi)`, one chunk per thread. + + Each finished chunk touches any registration lock this process holds: + `acquire` faults in under the lock what reclaim took back, and on a cold + cache that can outlast `_REG_LOCK_TIMEOUT_S` with no chunk registered. + """ + libc = _libc() pieces = _chunks(lo, hi - lo, _CHUNK_BYTES) def fault(piece) -> bool: o, n = piece - return libc.madvise(lo + o, n, _MADV_POPULATE_READ) == 0 + ok = libc.madvise(lo + o, n, _MADV_POPULATE_READ) == 0 + _lock_heartbeat() + return ok threads = min(copy_threads(hi - lo), len(pieces)) if threads <= 1: @@ -703,21 +925,38 @@ def _warn_unlocked(err: int, nbytes: int) -> None: _UNLOCKED_WARNED: list = [] -def _register(sources, leases): +def _register(sources, leases, *, shared: bool = False): """Register the mapping in place; each range is widened to page boundaries, which registration requires (the extra bytes are neighbouring rows the kernel never addresses), and shares any page another range of this process already holds (`_PageRuns`). A read-only mapping (the shared checkpoint mapping) is faulted in on several threads first and registered - read-only.""" + read-only. + + A `shared` table takes the registration lock file around `acquire` and + the matching `release`. Its fault-in runs before the lock, so ranks read + the page cache concurrently; `acquire` then faults in only what reclaim + took back before its `mlock`.""" rt = torch.cuda.cudart() + spans = [] + for src in sources: + lo, hi = _page_span(src) + flags = _HOST_REGISTER_READ_ONLY if _read_only(lo, hi) else 0 + if shared and flags: + _populate(lo, hi) + spans.append((lo, hi, flags)) + if not shared: + return _register_spans(rt, sources, spans, leases, None) + lock_path = _registration_lock_path(sources) + with _register_lock_file(lock_path): + return _register_spans(rt, sources, spans, leases, lock_path) + +def _register_spans(rt, sources, spans, leases, lock_path): out = [] - for src in sources: + for src, (lo, hi, flags) in zip(sources, spans, strict=True): base = src.data_ptr() nbytes = src.numel() * src.element_size() - lo, hi = _page_span(src) - flags = _HOST_REGISTER_READ_ONLY if _read_only(lo, hi) else 0 runs = _PAGE_RUNS.acquire(rt, lo, hi, flags) if runs is None: logger.info("engram: page-lock registration of %d bytes refused", hi - lo) @@ -726,7 +965,7 @@ def _register(sources, leases): def unregister(_ptr, runs=runs): if torch.cuda.is_initialized(): torch.cuda.synchronize() - _PAGE_RUNS.release(rt, runs) + _PAGE_RUNS.release(rt, runs, lock_path) lease = _Lease(lo, hi - lo, unregister, keep=src) leases.append(lease) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 380b713827..c988f7edca 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -1263,6 +1263,9 @@ def _finite_float_env(name: str, default: float, *, allow_zero: bool) -> float: # rather than copied once per rank -- and copies only if that is refused. # `tools/engram/bench_uva_shard.py` measures both. "ATOM_ENGRAM_UVA_COPY": lambda: os.getenv("ATOM_ENGRAM_UVA_COPY", "1") == "1", + # Optional override of the file that serializes shared in-place + # registration. Empty: host_memory names one from the checkpoint path. + "ATOM_ENGRAM_REG_LOCK": lambda: os.getenv("ATOM_ENGRAM_REG_LOCK", ""), # Where the compressed-vocab table is cached between runs. The table is # reproducible from the tokenizer, so this only trades startup time for # disk; point it at shared storage to let several servers build it once. diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 6baf69ef15..60fdb824f3 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -318,7 +318,8 @@ some layers and the device path for others, which is the confusing state. | Variable | Type | Default | Description | |----------|------|---------|-------------| | **ATOM_ENGRAM_UVA** | bool | 1 (true) | Page-lock this rank's shard of the hash tables (copied into page-locked memory by default, see `ATOM_ENGRAM_UVA_COPY`) and let a device kernel read the rows it needs across the bus, dequantizing there. No HBM for the table. `0` falls back to gathering the rows on the host, which returns the same rows but costs ~50 ms of CPU per decode step with the GPU idle behind it. Anything that would make the device path unsafe — no CUDA, more TP ranks than hash heads, a registration that will not fit — falls back on its own, so the switch is for taking the host path deliberately. The fallback is the whole TP group's: the lookup ends in an all-gather, so one rank that cannot register turns every rank around rather than leaving the others in a collective it never enters. | -| **ATOM_ENGRAM_UVA_COPY** | bool | unset: 1, except under DP attention | How the UVA path page-locks this rank's shard. Unset, a TP or single-GPU rank copies (`1`), and a DP-attention rank -- whose TP group is one GPU, so it holds the whole table, the same one as every other DP rank of the node -- registers in place (`0`), copying only if that is refused: a copy per rank is ~183 GiB each for V4.1-Flash (~730 GiB on DEP4), the in-place registration pins the shared checkpoint page cache once per node. Set explicitly, it decides for every rank. `1` copies the rows into page-locked memory of exactly their size (`hipHostMalloc`, on base pages: with transparent huge pages it compacts memory once the page cache has fragmented it, ~197 s per TP2 rank instead of ~12 s): one `copy_` below 512 MiB, otherwise one thread per 128 MiB but at least 8, at most 16 and the CPUs this process may run on (two or three threads lose to one `copy_`); each copied chunk of the checkpoint is dropped from the process (`madvise(DONTNEED)`; it stays in the page cache, reclaimable), so the shard is resident once and the peak stays near one shard. `0` registers the checkpoint mapping in place (`hipHostRegister`). The Engram tables are mapped shared and read-only, so this pins the page cache itself -- the same pages for every process that maps the file, no copy -- faulted in on several threads first and registered read-only, and `mlock`ed while registered (faulted in and locked in one parallel pass, unlocked when the last tensor over it goes). The lock covers the shared checkpoint page cache node-wide, so host memory pressure ends in the OOM killer rather than in reclaim. A refused `mlock` (`RLIMIT_MEMLOCK` without `CAP_IPC_LOCK`) warns once and the registration stands. **Required host settings for in-place registration:** `vm.compact_unevictable_allowed=0` and `vm.compaction_proactiveness=0`; with the defaults kcompactd migrates the registered pages, locked or not, and each migration evicts the process's GPU queues (decode stalls of seconds, or a hang; watch `/sys/class/kfd/kfd/proc//stats_*/evicted_ms`). Ranges that share a page (a layer's scales right before its values) share one reference-counted registration; a refused registration (or any failed HIP call) clears HIP's last error so the host-path fallback still runs. The page-locked memory is owned by the tensors the kernel reads and released (and unlocked) with the last of them. Needs the HIP runtime -- the one torch loaded, found in `/proc/self/maps` -- and without it the shard is registered in place. `tools/engram/bench_uva_shard.py` reproduces per-rank time, memory and gather latency for either mode. | +| **ATOM_ENGRAM_UVA_COPY** | bool | unset: 1, except under DP attention | How the UVA path page-locks this rank's shard. Unset, a TP or single-GPU rank copies (`1`), and a DP-attention rank -- whose TP group is one GPU, so it holds the whole table, the same one as every other DP rank of the node -- registers in place (`0`), copying only if that is refused: a copy per rank is ~183 GiB each for V4.1-Flash (~730 GiB on DEP4), the in-place registration pins the shared checkpoint page cache once per node. Set explicitly, it decides for every rank. `1` copies the rows into page-locked memory of exactly their size (`hipHostMalloc`, on base pages: with transparent huge pages it compacts memory once the page cache has fragmented it, ~197 s per TP2 rank instead of ~12 s): one `copy_` below 512 MiB, otherwise one thread per 128 MiB but at least 8, at most 16 and the CPUs this process may run on (two or three threads lose to one `copy_`); each copied chunk of the checkpoint is dropped from the process (`madvise(DONTNEED)`; it stays in the page cache, reclaimable), so the shard is resident once and the peak stays near one shard. `0` registers the checkpoint mapping in place (`hipHostRegister`). The Engram tables are mapped shared and read-only, so this pins the page cache itself -- the same pages for every process that maps the file, no copy -- faulted in on several threads first and registered read-only in 1 GiB pieces (one `cudaHostRegister` covering a whole ~92 GiB table, issued by several DP ranks at once, does not return), and `mlock`ed while registered (faulted in and locked in one parallel pass, unlocked when the last tensor over it goes). The lock covers the shared checkpoint page cache node-wide, so host memory pressure ends in the OOM killer rather than in reclaim. A refused `mlock` (`RLIMIT_MEMLOCK` without `CAP_IPC_LOCK`) warns once and the registration stands. **Required host settings for in-place registration:** `vm.compact_unevictable_allowed=0` and `vm.compaction_proactiveness=0`; with the defaults kcompactd migrates the registered pages, locked or not, and each migration evicts the process's GPU queues (decode stalls of seconds, or a hang; watch `/sys/class/kfd/kfd/proc//stats_*/evicted_ms`). Ranges that share a page (a layer's scales right before its values) share one reference-counted registration; a refused registration (or any failed HIP call) clears HIP's last error so the host-path fallback still runs. The page-locked memory is owned by the tensors the kernel reads and released (and unlocked) with the last of them. Needs the HIP runtime -- the one torch loaded, found in `/proc/self/maps` -- and without it the shard is registered in place. `tools/engram/bench_uva_shard.py` reproduces per-rank time, memory and gather latency for either mode. | +| **ATOM_ENGRAM_REG_LOCK** | path | per checkpoint, under `/dev/shm` | Lock file for in-place registration of a shared Engram table (DP attention: every rank page-locks the same checkpoint). Unset, the path is `/dev/shm/engram-uva-.lock`, hashed from that checkpoint, so two servers loading different checkpoints do not queue on each other; set it to force one path. Taken only around `cudaHostRegister` / `cudaHostUnregister` (fault-in runs before it), and only for that shared registration — a TP rank registering its own shard does not take it. A wait is logged. The holder touches the file after every 1 GiB chunk, so a slow but progressing holder (a cold page cache, minutes per rank) is waited on for as long as it takes; only a holder with no progress for 600 s makes the waiter give up and fall back to copying. After that the path is never waited on again in that process: its later tables copy at once, and its releases unregister without the lock. A lock file that cannot be opened or locked is logged and skipped, and the rank registers without it. `/dev/shm` is private to a container unless shared (`docker run --ipc=host`), so ranks that do not share it must set this to a path they all see. | | **ATOM_ENGRAM_CACHE_DIR** | path | `~/.cache/atom/engram` | Where the compressed-vocab table is cached between runs. The table is reproducible from the tokenizer, so this only trades startup time for disk; point it at shared storage to let several servers build it once. A truncated or stale cache is rebuilt rather than raised. | ## Attention side streams (DeepSeek-V4.1) diff --git a/tests/model_ops/engram/test_host_path.py b/tests/model_ops/engram/test_host_path.py index b90e8698b4..1075707f09 100644 --- a/tests/model_ops/engram/test_host_path.py +++ b/tests/model_ops/engram/test_host_path.py @@ -5,7 +5,14 @@ which only happens if the prime search runs in exactly the reference order. """ +import contextlib +import os +import subprocess +import sys +import threading +import time from concurrent.futures import ThreadPoolExecutor +from itertools import pairwise import numpy as np import pytest @@ -604,7 +611,7 @@ def test_staging_rejects_padding_that_truncates_or_exceeds_capacity(): @pytest.fixture(autouse=True) -def _no_real_host_allocator(monkeypatch): +def _no_real_host_allocator(monkeypatch, tmp_path): """Keep every test off the real `hipHostMalloc`: the registration tests script `cudaHostRegister`, and the copy tests install `FakeHostAllocator` themselves.""" @@ -615,6 +622,9 @@ def _no_real_host_allocator(monkeypatch): # copy policy as a single-GPU or TP rank sees it by default. monkeypatch.setattr(host_memory, "_PAGE_RUNS", host_memory._PageRuns()) monkeypatch.delenv("ATOM_ENGRAM_UVA_COPY", raising=False) + # A per-test registration lock file, so tests never share /dev/shm. + monkeypatch.setenv("ATOM_ENGRAM_REG_LOCK", str(tmp_path / "register.lock")) + monkeypatch.setattr(host_memory, "_timed_out_lock_paths", set()) class FakeCudart: @@ -691,6 +701,410 @@ def test_enabling_the_same_range_twice_registers_once(monkeypatch): assert cudart.calls == 2 +WIDE_VALUES = torch.ones(4, 4096) + + +def wide_table() -> HostEmbeddingTable: + """One values range of 64 KiB, so a chunk of a page or two splits it.""" + return HostEmbeddingTable(WIDE_VALUES, num_rows=4, head_dim=4096) + + +def registered_pages(values: torch.Tensor) -> int: + from atom.model_ops.engram import host_memory + + base = values.data_ptr() + nbytes = values.numel() * values.element_size() + lo = base - base % host_memory._PAGE + hi = -(-(base + nbytes) // host_memory._PAGE) * host_memory._PAGE + return (hi - lo) // host_memory._PAGE + + +def test_registration_is_split_into_chunks(monkeypatch): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + monkeypatch.setattr(host_memory, "_reg_chunk_bytes", lambda: 2 * host_memory._PAGE) + table = wide_table() + assert table.enable_uva(0, 4) is True + assert cudart.calls == -(-registered_pages(WIDE_VALUES) // 2) + table.disable_uva() + assert cudart.registered == [] + + +def test_refused_chunk_gives_back_the_chunks_before_it(monkeypatch): + """A refusal part way through a gap releases the pieces already taken. + + The rollback runs while this call holds the registration lock, and must + not take it again (that flock is not reentrant). + """ + from atom.model_ops.engram import host_memory + + cudart = FakeCudart(refuse_call=3) + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + monkeypatch.setattr(host_memory, "_reg_chunk_bytes", lambda: host_memory._PAGE) + monkeypatch.setattr(host_memory, "_REG_LOCK_TIMEOUT_S", 2) + table = wide_table() + assert table.enable_uva(0, 4, shared=True) is False + assert cudart.registered == [] + assert table._uva is None + + +def test_unopenable_lock_file_still_registers_in_place(monkeypatch, tmp_path): + """A lock file that cannot be opened is skipped, not treated as a refused + registration (which would send the rank to the copy path).""" + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + blocker = tmp_path / "not-a-dir" + blocker.write_text("") + monkeypatch.setenv("ATOM_ENGRAM_REG_LOCK", str(blocker / "register.lock")) + table = wide_table() + assert table.enable_uva(0, 4, shared=True) is True + assert len(cudart.registered) >= 1 + + +def test_registration_exception_is_not_chained_to_a_skipped_lock(tmp_path): + from atom.model_ops.engram import host_memory + + blocker = tmp_path / "not-a-dir" + blocker.write_text("x") + path = str(blocker / "register.lock") + with ( + pytest.raises(RuntimeError) as caught, + host_memory._register_lock_file(path), + ): + raise RuntimeError("during registration") + assert caught.value.__cause__ is None + assert caught.value.__context__ is None + + +def test_registration_and_unregistration_take_the_lock(monkeypatch): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + taken = [] + + @contextlib.contextmanager + def counting_lock_file(path): + taken.append((path, len(cudart.registered))) + yield + + monkeypatch.setattr(host_memory, "_register_lock_file", counting_lock_file) + table = wide_table() + assert table.enable_uva(0, 4, shared=True) is True + table.disable_uva() + # Once around registering (nothing held yet), once around releasing. + assert taken[0][1] == 0 + assert len(taken) >= 2 + assert cudart.registered == [] + + +def test_unshared_registration_skips_the_lock(monkeypatch): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + taken = [] + + @contextlib.contextmanager + def counting_lock_file(path): + taken.append(path) + yield + + monkeypatch.setattr(host_memory, "_register_lock_file", counting_lock_file) + table = wide_table() + assert table.enable_uva(0, 4) is True + table.disable_uva() + assert taken == [] + assert cudart.registered == [] + + +def test_lock_wait_is_bounded(tmp_path, monkeypatch): + from atom.model_ops.engram import host_memory + + monkeypatch.setattr(host_memory, "_REG_LOCK_TIMEOUT_S", 0.3) + path = str(tmp_path / "register.lock") + held, release = threading.Event(), threading.Event() + + def holder(): + with host_memory._register_lock_file(path): + held.set() + release.wait(10) + + thread = threading.Thread(target=holder) + thread.start() + assert held.wait(5) + try: + started = time.monotonic() + with pytest.raises(TimeoutError), host_memory._register_lock_file(path): + pass + assert time.monotonic() - started < 2 + # Latched: the next table does not wait out the timeout again. + monkeypatch.setattr(host_memory, "_REG_LOCK_TIMEOUT_S", 30) + started = time.monotonic() + with pytest.raises(TimeoutError), host_memory._register_lock_file(path): + pass + assert time.monotonic() - started < 1 + finally: + release.set() + thread.join(5) + assert not thread.is_alive() + + +def test_a_progressing_holder_is_waited_on_past_the_timeout(tmp_path, monkeypatch): + """The timeout bounds a holder's silence, not its hold time: a cold-cache + rank holding the lock for minutes, chunk by chunk, must not send the + ranks behind it to the copy path.""" + from atom.model_ops.engram import host_memory + + monkeypatch.setattr(host_memory, "_REG_LOCK_TIMEOUT_S", 0.3) + path = str(tmp_path / "register.lock") + held = threading.Event() + + def holder(): + with host_memory._register_lock_file(path): + held.set() + for _ in range(10): # ~1 s of chunks, three timeouts long + time.sleep(0.1) + host_memory._lock_heartbeat() + + thread = threading.Thread(target=holder) + thread.start() + assert held.wait(5) + started = time.monotonic() + with host_memory._register_lock_file(path): + waited = time.monotonic() - started + thread.join(5) + assert waited > 0.6 + assert path not in host_memory._timed_out_lock_paths + + +def test_fault_in_touches_a_held_lock(tmp_path, monkeypatch): + """`acquire` faults in under the lock what reclaim took back; on a cold + cache that can outlast the timeout before a chunk is registered, so every + faulted-in piece touches the held lock file.""" + from atom.model_ops.engram import host_memory + + advised = [] + + class Libc: + def madvise(self, addr, nbytes, advice): + advised.append(nbytes) + return 0 + + monkeypatch.setattr(host_memory, "_libc", lambda: Libc()) + beats = [] + heartbeat = host_memory._lock_heartbeat + monkeypatch.setattr( + host_memory, "_lock_heartbeat", lambda: (beats.append(1), heartbeat()) + ) + path = str(tmp_path / "register.lock") + pieces = 4 + lo = 1 << 30 + with host_memory._register_lock_file(path): + os.utime(path, (0, 0)) + beats.clear() + host_memory._populate(lo, lo + pieces * host_memory._CHUNK_BYTES) + assert os.stat(path).st_mtime > 0 + assert len(advised) == pieces + assert len(beats) == pieces + + +def test_heartbeat_touches_under_the_fd_list_lock(): + """Fault-in workers touch the held lock files while the main thread adds + and removes them: a touch waits for the list's lock, so a descriptor that + is being removed and closed is never touched.""" + from atom.model_ops.engram import host_memory + + done = threading.Event() + + def beat(): + host_memory._lock_heartbeat() + done.set() + + with host_memory._held_lock_fds_lock: + thread = threading.Thread(target=beat) + thread.start() + assert not done.wait(0.2) + assert done.wait(5) + thread.join(5) + + +def test_release_after_a_lock_timeout_still_unregisters(monkeypatch, tmp_path): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + path = str(tmp_path / "register.lock") + monkeypatch.setenv("ATOM_ENGRAM_REG_LOCK", path) + table = wide_table() + assert table.enable_uva(0, 4, shared=True) is True + assert cudart.registered + host_memory._timed_out_lock_paths.add(path) + table.disable_uva() + assert cudart.registered == [] + assert len(host_memory._PAGE_RUNS) == 0 + + +def test_lock_path_uses_the_override_or_the_mapped_file(monkeypatch, tmp_path): + import mmap + + from atom.model_ops.engram import host_memory + + monkeypatch.delenv("ATOM_ENGRAM_REG_LOCK", raising=False) + anon = torch.ones(4) + anonymous = host_memory._registration_lock_path([anon]) + assert anonymous.startswith("/dev/shm/engram-uva-") + assert host_memory._registration_lock_path([torch.ones(4)]) == anonymous + monkeypatch.setenv("ATOM_ENGRAM_REG_LOCK", str(tmp_path / "forced.lock")) + assert host_memory._registration_lock_path([anon]) == str(tmp_path / "forced.lock") + + monkeypatch.delenv("ATOM_ENGRAM_REG_LOCK", raising=False) + with contextlib.ExitStack() as stack: + paths = [] + for name in ("a.bin", "b.bin"): + raw = tmp_path / name + raw.write_bytes(b"\0" * 4096) + handle = stack.enter_context(open(raw, "rb")) + mapped = stack.enter_context( + mmap.mmap(handle.fileno(), 4096, prot=mmap.PROT_READ) + ) + paths.append( + host_memory._registration_lock_path( + [torch.frombuffer(mapped, dtype=torch.uint8)] + ) + ) + assert paths[0] != paths[1] + assert paths[0] != anonymous + + +_LOCK_WORKER = r""" +import os, sys, time, traceback +os.environ["ATOM_ENGRAM_REG_LOCK"] = os.environ["LOCK_PATH"] +os.environ["ATOM_ENGRAM_UVA_COPY"] = "0" +try: + import torch + from atom.model_ops.engram import host_memory + from atom.model_ops.engram.tables import HostEmbeddingTable + hold = float(os.environ.get("HOLD_S", "0.4")) + + class _Rt: + def cudaHostRegister(self, ptr, size, flags): + start = time.monotonic() + time.sleep(hold) + with open(os.environ["RECORD"], "a") as fh: + fh.write(f"{start:.6f} {time.monotonic():.6f}\n") + return 0 + + def cudaHostUnregister(self, ptr): + return 0 + + torch.cuda.cudart = lambda: _Rt() + values = torch.ones(4, 4096) + table = HostEmbeddingTable(values, num_rows=4, head_dim=4096) + if not table.enable_uva(0, 4, shared=True): + sys.exit("enable_uva failed") + table.disable_uva() +except Exception: + traceback.print_exc() + sys.exit(1) +""" + + +def test_two_processes_do_not_register_at_once(tmp_path): + env = os.environ.copy() + env["LOCK_PATH"] = str(tmp_path / "register.lock") + env["HOLD_S"] = "0.4" + env["HIP_VISIBLE_DEVICES"] = "" + procs = [] + records = [] + for i in range(2): + record = tmp_path / f"rec{i}" + records.append(record) + child = env.copy() + child["RECORD"] = str(record) + procs.append( + subprocess.Popen( + [sys.executable, "-c", _LOCK_WORKER], + env=child, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + ) + for proc in procs: + _out, err = proc.communicate(timeout=60) + assert proc.returncode == 0, err + spans = [] + for record in records: + text = record.read_text().strip().splitlines() + assert text, record + for line in text: + start, end = (float(part) for part in line.split()) + spans.append((start, end)) + spans.sort() + assert len(spans) >= 2 + for (_start, end), (nxt, _end) in pairwise(spans): + assert end <= nxt + + +class RecordingLibc: + """`mlock2` / `munlock` / `madvise` spans, with no effect.""" + + def __init__(self): + self.locked, self.unlocked = [], [] + + def mlock2(self, addr, n, flags): + self.locked.append((addr, addr + n)) + return 0 + + def munlock(self, addr, n): + self.unlocked.append((addr, addr + n)) + return 0 + + def madvise(self, addr, n, advice): + return 0 + + +def covered_bytes(ranges) -> int: + return sum(hi - lo for lo, hi in ranges) + + +def test_chunked_runs_are_unlocked_with_their_last_reference(monkeypatch): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart() + libc = RecordingLibc() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + monkeypatch.setattr(host_memory, "_libc", lambda: libc) + monkeypatch.setattr(host_memory, "_reg_chunk_bytes", lambda: 2 * host_memory._PAGE) + table = wide_table() + assert table.enable_uva(0, 4) is True + assert ( + covered_bytes(libc.locked) == registered_pages(WIDE_VALUES) * host_memory._PAGE + ) + assert libc.unlocked == [] + table.disable_uva() + assert cudart.registered == [] + assert len(libc.unlocked) == cudart.calls + assert covered_bytes(libc.unlocked) == covered_bytes(libc.locked) + + +def test_refused_chunk_leaves_nothing_locked(monkeypatch): + from atom.model_ops.engram import host_memory + + cudart = FakeCudart(refuse_call=3) + libc = RecordingLibc() + monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) + monkeypatch.setattr(host_memory, "_libc", lambda: libc) + monkeypatch.setattr(host_memory, "_reg_chunk_bytes", lambda: host_memory._PAGE) + table = wide_table() + assert table.enable_uva(0, 4) is False + assert cudart.registered == [] + assert covered_bytes(libc.unlocked) == covered_bytes(libc.locked) + + class FakeHostAllocator: """Page-locked allocations backed by ctypes buffers kept alive here."""