| |
| """Host side of the async hand-off (torch-free): libdsv41_hostflag.so (hostflag.c), the page with the gathered / |
| released / claimed words and the reply area, the proxy's wait and publish, fail-all and the native watchdog. |
| |
| Split out of streamsync.py unchanged so that the out-of-process proxy (proxy_proc.py) can use it without importing |
| torch; streamsync.py re-exports every name. The page is either allocated here (one process: the thread proxy) or lies |
| in a shared region another object owns (shmregion.py: the process proxy), see HostFlagCore(page=...). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import atexit |
| import ctypes |
| import hashlib |
| import os |
| import shutil |
| import stat |
| import subprocess |
| import tempfile |
|
|
| import numpy as np |
|
|
| CU_STREAM_WAIT_VALUE_GEQ = 0x0 |
| CU_STREAM_WRITE_VALUE_DEFAULT = 0x0 |
| CU_MEMHOSTREGISTER_PORTABLE = 0x01 |
| CU_MEMHOSTREGISTER_DEVICEMAP = 0x02 |
|
|
| HOSTFLAG_LIB = "libdsv41_hostflag.so" |
| HOSTFLAG_VERSION = 2 |
| PAGE_BYTES = 65536 |
| OFF_GATHERED = 0 |
| OFF_RELEASED = 128 |
| OFF_CLAIMED = 256 |
| OFF_CONSUMED = 384 |
| OFF_REPLY = 4096 |
| REPLY_CAPACITY = PAGE_BYTES - OFF_REPLY |
| FAR_AHEAD = 1 << 30 |
| _M32 = 0xFFFFFFFF |
| _SRC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "hostflag.c") |
|
|
|
|
| class StreamSyncError(RuntimeError): |
| pass |
|
|
|
|
| class StreamStuckError(RuntimeError): |
| """The engine stream stayed parked after the self-test's rescue: this device does not observe host stores. Not |
| recoverable (the stream cannot be unparked), so the caller must fail loudly, not fall back (S9 F4).""" |
|
|
|
|
| |
|
|
| _V = ctypes.c_void_p |
|
|
|
|
| def _bind(lib: ctypes.CDLL) -> ctypes.CDLL: |
| try: |
| sig = { |
| "dsv41_hf_version": (ctypes.c_int, []), |
| "dsv41_hf_alloc": (_V, [ctypes.c_size_t]), |
| "dsv41_hf_free": (None, [_V]), |
| "dsv41_hf_lock": (ctypes.c_int, [_V, ctypes.c_size_t]), |
| "dsv41_hf_load": (ctypes.c_uint32, [_V]), |
| "dsv41_hf_release_max": (None, [_V, ctypes.c_uint32]), |
| "dsv41_hf_store": (None, [_V, ctypes.c_uint32]), |
| "dsv41_hf_set": (None, [_V, ctypes.c_uint32]), |
| "dsv41_hf_claim": (ctypes.c_int, [_V, ctypes.c_uint32]), |
| "dsv41_hf_publish": (ctypes.c_int, [_V, _V, ctypes.c_uint32, _V, _V, ctypes.c_size_t]), |
| "dsv41_hf_fail_release": (ctypes.c_int, [_V, _V, ctypes.c_uint32, _V, ctypes.c_size_t, ctypes.c_int32]), |
| "dsv41_hf_wait_geq": (ctypes.c_int, [_V, ctypes.c_uint32, ctypes.c_int64, ctypes.c_int64, ctypes.c_int64]), |
| "dsv41_hf_watchdog_start": (_V, [_V, _V, _V, _V, ctypes.c_size_t, ctypes.c_int32, ctypes.c_int64, |
| ctypes.c_int64]), |
| "dsv41_hf_watchdog_stats": (None, [_V, ctypes.POINTER(ctypes.c_uint64)]), |
| "dsv41_hf_watchdog_stop": (ctypes.c_int, [_V]), |
| } |
| for name, (res, args) in sig.items(): |
| f = getattr(lib, name) |
| f.restype = res |
| f.argtypes = args |
| except AttributeError as e: |
| raise StreamSyncError(f"{lib._name}: {e}") from e |
| v = lib.dsv41_hf_version() |
| if v != HOSTFLAG_VERSION: |
| raise StreamSyncError(f"{lib._name}: hostflag version {v}, expected {HOSTFLAG_VERSION}") |
| return lib |
|
|
|
|
| def _private_dir(path: str) -> None: |
| """Create (or accept) a build directory only this user can write.""" |
| os.makedirs(path, mode=0o700, exist_ok=True) |
| st = os.lstat(path) |
| if not stat.S_ISDIR(st.st_mode) or st.st_uid != os.getuid() or st.st_mode & 0o022: |
| raise StreamSyncError(f"hostflag build directory {path} is not a private directory of uid {os.getuid()}") |
|
|
|
|
| def _compile(cc: str, out_dir: str) -> str: |
| with open(_SRC, "rb") as f: |
| digest = hashlib.sha256(f.read()).hexdigest()[:16] |
| d = os.path.join(out_dir, f"dsv41-hostflag-{os.getuid()}-{digest}") |
| _private_dir(d) |
| so = os.path.join(d, HOSTFLAG_LIB) |
| if not os.path.exists(so): |
| tmp = so + f".tmp{os.getpid()}" |
| try: |
| r = subprocess.run([cc, "-O2", "-Wall", "-shared", "-fPIC", "-o", tmp, _SRC, "-lpthread"], |
| capture_output=True, text=True, timeout=120) |
| except (subprocess.TimeoutExpired, OSError) as e: |
| raise StreamSyncError(f"compiling {_SRC} with {cc}: {e}") from e |
| if r.returncode != 0: |
| raise StreamSyncError(f"compiling {_SRC} with {cc} failed ({r.returncode}): {r.stderr.strip()[:400]}") |
| os.replace(tmp, so) |
| return so |
|
|
|
|
| def load_hostflag(lib_path: str = "") -> tuple[ctypes.CDLL, str]: |
| """libdsv41_hostflag.so: an explicit path, else VLLM_DSPARK_REMOTE_OFL_DIR's copy, else hostflag.c compiled with |
| the system C compiler into a private directory under the temp dir. Returns (library, where it came from).""" |
| if lib_path: |
| if not os.path.exists(lib_path): |
| raise StreamSyncError(f"VLLM_DSPARK_REMOTE_HOSTFLAG_LIB={lib_path} does not exist") |
| return _bind(ctypes.CDLL(lib_path)), lib_path |
| ofl_dir = (os.environ.get("VLLM_DSPARK_REMOTE_OFL_DIR") or "").strip() |
| if ofl_dir and os.path.exists(os.path.join(ofl_dir, HOSTFLAG_LIB)): |
| p = os.path.join(ofl_dir, HOSTFLAG_LIB) |
| return _bind(ctypes.CDLL(p)), p |
| cc = shutil.which("cc") or shutil.which("gcc") |
| if cc is None: |
| raise StreamSyncError(f"no {HOSTFLAG_LIB} (set VLLM_DSPARK_REMOTE_HOSTFLAG_LIB or put it in " |
| "VLLM_DSPARK_REMOTE_OFL_DIR) and no C compiler to build hostflag.c") |
| p = _compile(cc, tempfile.gettempdir()) |
| return _bind(ctypes.CDLL(p)), f"{p} (compiled at startup)" |
|
|
|
|
| PR_SET_TIMERSLACK = 29 |
|
|
|
|
| def set_timer_slack(ns: int = 1000) -> bool: |
| """Linux timer slack of the calling thread: with VLLM_DSPARK_REMOTE_ASYNC_NAP_US > 0 the proxy's waits nap in |
| short nanosleeps, which the default 50 us slack would stretch. False where prctl is unavailable.""" |
| try: |
| libc = ctypes.CDLL(None, use_errno=True) |
| return libc.prctl(PR_SET_TIMERSLACK, ctypes.c_ulong(int(ns)), ctypes.c_ulong(0), ctypes.c_ulong(0), |
| ctypes.c_ulong(0)) == 0 |
| except (OSError, AttributeError): |
| return False |
|
|
|
|
| |
|
|
| class HostFlagCore: |
| """The host side of the hand-off (no CUDA): the C-owned page with the three words and the reply area, the proxy's |
| wait and publish, fail-all, and the native watchdog. The page is never freed while the process may still copy from |
| it (close() only after the round stopped and the stream drained).""" |
|
|
| cuda_free_proxy = True |
|
|
| def __init__(self, lib_path: str = "", spin_us: int = 1000, nap_us: int = 0, page: int = 0): |
| """page: 0 = allocate the page here (freed by close()); else the address of a PAGE_BYTES page, 4096-aligned, |
| in memory the caller owns and keeps mapped (a shared region: shmregion.py). An external page is never freed |
| here, and its words are not reset: the region's creator zeroed them.""" |
| self.lib, self.lib_origin = load_hostflag(lib_path) |
| self.spin_ns = int(spin_us) * 1000 |
| self.nap_ns = int(nap_us) * 1000 |
| self.owns_page = not page |
| if page: |
| if page % 4096: |
| raise StreamSyncError(f"external hostflag page {page:#x} is not 4096-aligned") |
| else: |
| page = self.lib.dsv41_hf_alloc(PAGE_BYTES) |
| if not page: |
| raise StreamSyncError("hostflag page allocation failed") |
| self.page = int(page) |
| self.mlock_errno = int(self.lib.dsv41_hf_lock(_V(self.page), PAGE_BYTES)) |
| self.gathered_host = self.page + OFF_GATHERED |
| self.released_host = self.page + OFF_RELEASED |
| self.claimed_host = self.page + OFF_CLAIMED |
| self.consumed_host = self.page + OFF_CONSUMED |
| self.reply_host = self.page + OFF_REPLY |
| self.reply_bytes = 0 |
| self.fail_status = 0 |
| self.deadline_ms = 0 |
| self._wd: int | None = None |
| self._atexit_registered = False |
|
|
| |
|
|
| def wait_gathered(self, value: int, timeout_s: float) -> bool: |
| """True once gathered >= value; False after timeout_s. Releases the GIL while it waits.""" |
| r = self.lib.dsv41_hf_wait_geq(_V(self.gathered_host), value & _M32, int(timeout_s * 1e9), |
| self.spin_ns, self.nap_ns) |
| return r == 0 |
|
|
| def publish(self, value: int, src: np.ndarray) -> bool: |
| """Claim round value, copy src into the reply area, release value: one GIL-free C call. False when the round |
| had already been claimed (the watchdog failed it first); nothing is written then.""" |
| if not src.flags.c_contiguous or src.nbytes > REPLY_CAPACITY: |
| raise ValueError(f"reply of {src.nbytes} bytes (contiguous {src.flags.c_contiguous}) does not fit") |
| return bool(self.lib.dsv41_hf_publish(_V(self.claimed_host), _V(self.released_host), value & _M32, |
| _V(self.reply_host), _V(src.ctypes.data), src.nbytes)) |
|
|
| def taken(self, value: int) -> bool: |
| """True once round value (or a later one) has been claimed: by this proxy's publish, the watchdog or a |
| fail-all. A proxy that finds its round taken before it starts must not read that round's staging.""" |
| c = int(self.lib.dsv41_hf_load(_V(self.claimed_host))) |
| return ((c - value) & _M32) < (1 << 31) |
|
|
| def reply_consumed(self, value: int) -> bool: |
| """The consumer acknowledgment: True once the engine stream has copied round value's reply out of the page |
| (it writes the word right after the copy). Round t's reply may be written only when round t-1's is consumed; |
| gathered >= t implies it by stream order, and the proxy checks the word before every publish.""" |
| c = int(self.lib.dsv41_hf_load(_V(self.consumed_host))) |
| return ((c - value) & _M32) < (1 << 31) |
|
|
| def consumed(self, value: int, stream=None) -> None: |
| """Engine side, after copy_reply(): acknowledge round value's reply. CPU stand-ins copy synchronously, so the |
| store follows the copy; HostFlagSync enqueues it on the stream instead.""" |
| self.lib.dsv41_hf_release_max(_V(self.consumed_host), value & _M32) |
|
|
| def release(self, value: int) -> None: |
| """Monotonic release (cyclic max) without a reply: the self-test's rescue.""" |
| self.lib.dsv41_hf_release_max(_V(self.released_host), value & _M32) |
|
|
| def drain_marker(self): |
| """Engine side: a marker after everything enqueued on the engine stream so far (process proxy respawn, S11 P-4). |
| CPU stand-ins run their "stream" synchronously, so nothing is ever pending: None.""" |
| return None |
|
|
| def drained(self, marker) -> bool: |
| """True once everything before the marker ran (never waits).""" |
| return True |
|
|
| |
|
|
| def fail_all(self, value: int) -> bool: |
| """Claim value, write the failed reply, release value. With value FAR_AHEAD of every enqueued round this fails |
| every pending and later round at once (proxy death, close, exit). False if value was already claimed.""" |
| n = self.reply_bytes or REPLY_CAPACITY |
| return bool(self.lib.dsv41_hf_fail_release(_V(self.claimed_host), _V(self.released_host), value & _M32, |
| _V(self.reply_host), n, self.fail_status)) |
|
|
| def fail_ahead(self) -> bool: |
| """fail_all at FAR_AHEAD past the furthest of the three words.""" |
| g, r, c = self.values() |
| base = g |
| for w in (r, c): |
| if ((w - base) & _M32) < (1 << 31): |
| base = w |
| return self.fail_all(base + FAR_AHEAD) |
|
|
| def values(self) -> tuple[int, int, int]: |
| ld = self.lib.dsv41_hf_load |
| return (int(ld(_V(self.gathered_host))), int(ld(_V(self.released_host))), int(ld(_V(self.claimed_host)))) |
|
|
| |
|
|
| def attach(self, reply_bytes: int, fail_status: int, deadline_ms: int, poll_us: int = 1000) -> None: |
| """Size the reply, write a failed reply into it (what the stream copies before any round was published) and |
| start the watchdog.""" |
| if self._wd is not None: |
| raise StreamSyncError("hostflag watchdog already attached") |
| if not 4 <= reply_bytes <= REPLY_CAPACITY: |
| raise StreamSyncError(f"reply of {reply_bytes} bytes does not fit the {REPLY_CAPACITY}-byte area") |
| if deadline_ms <= 0: |
| raise StreamSyncError(f"watchdog deadline {deadline_ms} ms") |
| self.configure_reply(reply_bytes, fail_status) |
| self.deadline_ms = int(deadline_ms) |
| failed = np.zeros(self.reply_bytes // 4, dtype=np.int32) |
| failed[0] = self.fail_status |
| ctypes.memmove(self.reply_host, failed.ctypes.data, self.reply_bytes) |
| poll_ns = max(50_000, min(int(poll_us) * 1000, self.deadline_ms * 100_000)) |
| wd = self.lib.dsv41_hf_watchdog_start(_V(self.gathered_host), _V(self.released_host), |
| _V(self.claimed_host), _V(self.reply_host), self.reply_bytes, |
| self.fail_status, self.deadline_ms * 1_000_000, poll_ns) |
| if not wd: |
| raise StreamSyncError(f"hostflag watchdog did not start (errno {ctypes.get_errno()})") |
| self._wd = int(wd) |
| if not self._atexit_registered: |
| atexit.register(self._at_exit) |
| self._atexit_registered = True |
|
|
| def configure_reply(self, reply_bytes: int, fail_status: int) -> None: |
| """Reply size and failed status, without touching the page (attach() and a process proxy's view of a page |
| the engine process attached).""" |
| if not 4 <= reply_bytes <= REPLY_CAPACITY: |
| raise StreamSyncError(f"reply of {reply_bytes} bytes does not fit the {REPLY_CAPACITY}-byte area") |
| self.reply_bytes = int(reply_bytes) |
| self.fail_status = int(fail_status) |
|
|
| def watchdog_stats(self) -> dict: |
| if self._wd is None: |
| return {"running": False} |
| out = (ctypes.c_uint64 * 3)() |
| self.lib.dsv41_hf_watchdog_stats(_V(self._wd), out) |
| return {"running": True, "deadline_ms": self.deadline_ms, "fires": int(out[0]), |
| "last_failed_round": int(out[1]), "max_late_ms": round(out[2] / 1e6, 3)} |
|
|
| def watchdog_mirror(self, dst: int) -> None: |
| """Copy the watchdog's raw stats (fires, last failed round, max lateness in ns: 3 x u64) to address dst (the |
| process proxy's shared control block). No-op without a watchdog.""" |
| if self._wd is not None: |
| self.lib.dsv41_hf_watchdog_stats(_V(self._wd), ctypes.cast(dst, ctypes.POINTER(ctypes.c_uint64))) |
|
|
| def stop_watchdog(self) -> None: |
| if self._wd is not None: |
| wd, self._wd = self._wd, None |
| self.lib.dsv41_hf_watchdog_stop(_V(wd)) |
|
|
| def reset(self) -> None: |
| """All words to 0. Only with no round in flight and no watchdog (setup, self-test).""" |
| if self._wd is not None: |
| raise StreamSyncError("reset with the watchdog running") |
| for a in (self.gathered_host, self.released_host, self.claimed_host, self.consumed_host): |
| self.lib.dsv41_hf_set(_V(a), 0) |
|
|
| def describe(self) -> dict: |
| return {"mechanism": "host_flags", "lib": self.lib_origin, "spin_us": self.spin_ns // 1000, |
| "nap_us": self.nap_ns // 1000, "mlock": "ok" if self.mlock_errno == 0 else f"errno {self.mlock_errno}", |
| "watchdog_deadline_ms": self.deadline_ms or None} |
|
|
| def _at_exit(self) -> None: |
| |
| |
| if self.page: |
| self.fail_ahead() |
| self.stop_watchdog() |
|
|
| def close(self) -> None: |
| """Only after the round stopped and the stream drained (nothing may copy from the page afterwards).""" |
| self.stop_watchdog() |
| if self.page: |
| if self.owns_page: |
| self.lib.dsv41_hf_free(_V(self.page)) |
| self.page = 0 |
|
|
|
|
