fix(engram): serialize UVA registration across DP ranks and register in 1 GiB chunks

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,24 @@
 # 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 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.
+_held_lock_fds: list[int] = []
+
+
+def _lock_heartbeat() -> None:
+    """Touch the held registration lock file, so waiters see progress."""
+    for fd in _held_lock_fds:
+        with contextlib.suppress(OSError):
+            os.utime(fd)
 
 
 @functools.cache
@@ -513,7 +539,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 +592,145 @@ 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 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
+            _held_lock_fds.append(fd)
+            _lock_heartbeat()
+    try:
+        yield
+    finally:
+        if held:
+            _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 +756,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,11 +864,15 @@ 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."""
+    libc = _libc()
     pieces = _chunks(lo, hi - lo, _CHUNK_BYTES)
 
     def fault(piece) -> bool:
@@ -703,21 +910,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 +950,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)