diff --git a/hs_connectors/src/hs_connectors/__init__.py b/hs_connectors/src/hs_connectors/__init__.py index 35669ee83..ec7ec4015 100644 --- a/hs_connectors/src/hs_connectors/__init__.py +++ b/hs_connectors/src/hs_connectors/__init__.py @@ -3,6 +3,8 @@ FileTransfer, HiddenStatesBackend, HiddenStatesTransfer, + HttpBackend, + HttpTransfer, MooncakeBackend, MooncakeTransfer, ) @@ -12,6 +14,8 @@ "FileTransfer", "HiddenStatesBackend", "HiddenStatesTransfer", + "HttpBackend", + "HttpTransfer", "MooncakeBackend", "MooncakeTransfer", ] diff --git a/hs_connectors/src/hs_connectors/transfer.py b/hs_connectors/src/hs_connectors/transfer.py index cf3b69aea..d8237157a 100644 --- a/hs_connectors/src/hs_connectors/transfer.py +++ b/hs_connectors/src/hs_connectors/transfer.py @@ -8,11 +8,15 @@ import shutil import socket import time +import urllib.error +import urllib.request from abc import ABC, abstractmethod +from http import HTTPStatus from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar import torch +from safetensors.torch import load as load_safetensors_bytes from safetensors.torch import load_file from hs_connectors.mooncake_store import MooncakeHiddenStatesStore, MooncakeStoreConfig @@ -22,7 +26,7 @@ from collections.abc import Callable -def wait_for_lock(lock_path: str, timeout: float = 10.0, poll_interval: float = 0.1): +def wait_for_lock(lock_path: str, timeout: float = 10.0, poll_interval: float = 0.02): fd = os.open(lock_path, os.O_RDWR) try: deadline = time.monotonic() + timeout @@ -332,3 +336,124 @@ def build_kv_transfer_config(args: argparse.Namespace) -> dict[str, Any]: "mooncake": dataclasses.asdict(mooncake_cfg), }, } + + +# --------------------------------------------------------------------------- +# HTTP backend (vLLM writes to local fast disk; trainer fetches over HTTP) +# +# Motivation: pointing ``shared_storage_path`` at a shared network filesystem +# (e.g. cephfs) makes every connector write a cross-host op that contends on +# the MDS; under saturation, batches of writes stall together with 1-3 s tail +# latency, which punches through the trainer's prefetch buffer and produces +# tail training steps. This backend instead tells vLLM to write to a local +# disk on its own node (no MDS, no cross-host write bursts) and has the +# trainer pull the file back over plain HTTP from a tiny static file server +# (``scripts/serve_hs.py``) running next to vLLM. It also works without any +# shared filesystem between the two nodes. +# +# The connector side is unchanged: it still writes +# ``{shared_storage_path}/{req_id}.safetensors`` under a ``.lock`` flock and +# returns that absolute path as the handle. We only reinterpret the handle on +# the trainer side: strip the basename, prepend ``--hs-http-base``, GET it. +# The file server blocks on the same ``.lock`` flock until the write is done, +# preserving the existing synchronization semantics. +# --------------------------------------------------------------------------- + + +class HttpTransfer(HiddenStatesTransfer): + def __init__( + self, + hs_http_base: str, + hidden_states_path: Path, + timeout: float = 120.0, + ): + self.hs_http_base = hs_http_base.rstrip("/") + self.hidden_states_path = hidden_states_path + self.timeout = timeout + + def get_cached(self, file_idx: int) -> dict[str, torch.Tensor] | None: + path = self.hidden_states_path / f"hs_{file_idx}.safetensors" + return _load_hs_file(path) + + def _url_for(self, handle: str) -> str: + return f"{self.hs_http_base}/{os.path.basename(handle)}" + + def get_generated(self, handle: str) -> dict[str, torch.Tensor] | None: + url = self._url_for(handle) + try: + req = urllib.request.Request(url, method="GET") # noqa: S310 + with urllib.request.urlopen(req, timeout=self.timeout) as resp: # noqa: S310 + payload = resp.read() + except urllib.error.HTTPError as e: + if e.code == HTTPStatus.NOT_FOUND: + return None + raise + if not payload: + return None + return load_safetensors_bytes(payload) + + def cache(self, handle: str, file_idx: int) -> None: + raise NotImplementedError( + "HttpTransfer.cache() is unsupported: the hidden-states file lives" + " on the vLLM node. Use --hidden-states-backend file for" + " on_generate=cache, or keep on_generate=delete with the http" + " backend." + ) + + def delete(self, handle: str) -> None: + url = self._url_for(handle) + try: + req = urllib.request.Request(url, method="DELETE") # noqa: S310 + with urllib.request.urlopen(req, timeout=self.timeout) as resp: # noqa: S310 + resp.read() + except OSError: + pass # best-effort cleanup; the server TTL-sweeps leftovers + + +@HiddenStatesBackend.register("http") +class HttpBackend(HiddenStatesBackend): + @staticmethod + def add_train_args(parser): + parser.add_argument( + "--hs-http-base", + type=str, + default=None, + help="Base URL of the hidden-states file server, e.g. 'http://10.0.0.1:9010'.", + ) + parser.add_argument( + "--hs-http-timeout", + type=float, + default=120.0, + help="Per-request timeout (seconds) for HTTP hidden-states fetches.", + ) + + @staticmethod + def add_launch_args(parser): + pass + + @staticmethod + def from_train_args(args, data_path): + if not getattr(args, "hs_http_base", None): + raise ValueError( + "--hs-http-base is required when --hidden-states-backend=http" + ) + hs_path = ( + Path(args.hidden_states_path) + if args.hidden_states_path + else Path(data_path) / "hidden_states" + ) + return HttpTransfer( + hs_http_base=args.hs_http_base, + hidden_states_path=hs_path, + timeout=getattr(args, "hs_http_timeout", 120.0), + ) + + @staticmethod + def build_kv_transfer_config(args): + return { + "kv_connector": "ExampleHiddenStatesConnector", + "kv_role": "kv_producer", + "kv_connector_extra_config": { + "shared_storage_path": args.hidden_states_path, + }, + } diff --git a/scripts/serve_hs.py b/scripts/serve_hs.py new file mode 100755 index 000000000..12cdaf4ee --- /dev/null +++ b/scripts/serve_hs.py @@ -0,0 +1,345 @@ +#!/usr/bin/env python3 +"""Tiny static file server for HTTP-based hidden-states transfer. + +Run this on each vLLM node, pointing ``--root`` at the same local directory +the connector writes to (``shared_storage_path`` / ``--hidden-states-path``). +It pairs with the ``http`` hidden-states backend +(``hs_connectors.transfer.HttpBackend``): the trainer does + + GET http://:/.safetensors -> wait for write, stream + DELETE http://:/.safetensors -> drop file + lock + +Why it exists +------------ +The connector (``example_hidden_states_connector.py``) returns a request's +handle *before* its async DtoH copy + disk write finishes. To hide that race it +holds an advisory flock (``LOCK_EX``) on a companion ``.lock`` for the +duration of the write and releases it (closes the fd) once the write is done. +The file backend's reader re-acquires that lock to block until the write +completes. Over HTTP the trainer can no longer see the lock, so this server +re-acquires it on the trainer's behalf: a ``GET`` blocks until the lock is +released (= write complete), then streams the file. This preserves the exact +synchronization semantics of the file backend with zero connector changes. + +A background TTL sweeper deletes stale ``.safetensors`` files (and their +``.lock`` companions) older than ``--ttl`` seconds, so a trainer crash can never +fill the local disk with orphaned files. + +Usage +----- + python scripts/serve_hs.py --root /data/local_hs --port 9010 --host 0.0.0.0 + +Examples +-------- + # One server per vLLM instance, each serving that instance's own subdir: + python scripts/serve_hs.py --root /data/local_hs/8010 --port 9010 + python scripts/serve_hs.py --root /data/local_hs/8020 --port 9020 +""" + +from __future__ import annotations + +import argparse +import contextlib +import fcntl +import http.server +import math +import os +import socketserver +import threading +import time +from pathlib import Path + +DEFAULT_ROOT = "/tmp/hidden_states" # noqa: S108 +DEFAULT_LOCK_TIMEOUT = 120.0 # seconds to wait for the connector's write +DEFAULT_TTL = 600.0 # seconds; stale files older than this are swept +DEFAULT_TTL_INTERVAL = 60.0 # seconds between sweeper passes +CHUNK = 1024 * 1024 # 1 MiB streaming buffer + + +def _wait_for_file(path: Path, timeout: float, poll: float = 0.05) -> bool: + """Spin until ``path`` exists or ``timeout`` elapses. True if it appeared.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if path.exists(): + return True + time.sleep(poll) + return path.exists() + + +def _wait_for_write(data_path: Path, lock_timeout: float) -> str | None: + """Block until the connector finishes writing ``data_path``. + + Returns a short status string for logging ("locked"), or ``None`` if the + file is not available within the timeout. + + Readiness is driven entirely by the connector's ``.lock`` flock: the + connector creates ``.lock`` and holds LOCK_EX while writing, + then closes its fd (releasing the lock) once done. We wait for the lock + file to appear, then flock-wait (LOCK_EX) until we acquire it — acquiring + it means the writer released it = write done. This is race-free and never + serves a half-written file. + + There is deliberately no fixed cap on waiting for the lock file to appear: + it returns as soon as the lock shows up (normally within a scheduler step). + An earlier version capped this wait at 5s then fell back to serving the + data file by existence — which both burned a flat ~5s on every fetch + (observed: ``get_generated`` p50 == 5035ms) AND risked serving a + half-written file. We now just wait for the lock. + """ + lock_path = Path(str(data_path) + ".lock") + + # Wait for the lock file to appear (connector creates it when the write + # starts). No fixed cap — returns as soon as it appears. + if not _wait_for_file(lock_path, timeout=lock_timeout): + return None # lock never appeared; do NOT serve (would risk a partial read) + + try: + fd = os.open(str(lock_path), os.O_RDONLY) + except OSError: + # Lost a race with DELETE / TTL sweeper between the existence check + # and the open; treat as unavailable rather than crashing the handler. + return None + deadline = time.monotonic() + lock_timeout + try: + while time.monotonic() < deadline: + try: + # NB: LOCK_EX (not SH) so it blocks until the writer's exclusive + # lock is released, exactly like the file backend's reader. + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + return "locked" + except BlockingIOError: + time.sleep(0.02) + return None # timed out waiting for the writer to release + finally: + os.close(fd) # releases our (just-acquired) lock; leave the file in place + + +def _sweep(root: Path, ttl: float, interval: float) -> None: + """Daemon loop: delete ``.safetensors`` (and ``.lock``) older than ``ttl``.""" + while True: + time.sleep(interval) + # NB: st_mtime is Unix wall-clock time; compare against time.time(), + # NOT time.monotonic() (seconds since boot), or the age check never + # fires and stale files are never reclaimed. + now = time.time() + try: + entries = list(root.iterdir()) + except OSError: + continue + for entry in entries: + try: + mtime = entry.stat().st_mtime + except OSError: + continue + if now - mtime < ttl: + continue + if entry.suffix in {".safetensors", ".lock"}: + with contextlib.suppress(OSError): + entry.unlink() + + +class HiddenStatesHandler(http.server.BaseHTTPRequestHandler): + root: Path # injected via the server factory below + lock_timeout: float + + # One concise line per request (default http.server logs hit stderr with + # client address + date noise on every GET). + def log_message(self, fmt: str, *args) -> None: + print(f"[serve_hs] {self.address_string()} {fmt % args}", flush=True) + + # --- helpers ---------------------------------------------------------- + def _resolve(self, urlpath: str) -> Path | None: + """Map URL path to a safe absolute path under ``root`` (no traversal).""" + # Strip query/fragment, take the basename only. The connector writes a + # flat ``{req_id}.safetensors``, so we never serve nested paths. + name = os.path.basename(urlpath.partition("?")[0]) + if not name or name in (".", "..") or "/" in name or "\\" in name: + return None + candidate = (self.root / name).resolve() + try: + candidate.relative_to(self.root.resolve()) + except ValueError: + return None + return candidate + + # --- GET -------------------------------------------------------------- + def do_GET(self) -> None: + data_path = self._resolve(self.path) + if data_path is None or data_path.suffix != ".safetensors": + self.send_error(404, "Not found") + return + + status = _wait_for_write(data_path, self.lock_timeout) + if status is None or not data_path.exists(): + self.send_error(404, "Hidden states not available") + return + + # Open before sending the 200 so a DELETE / TTL-sweeper race after the + # lock-wait answers 404 instead of failing mid-transfer. The fd stays + # open across header-send and streaming, so no context manager. + try: + f = open(data_path, "rb") # noqa: SIM115 - closed in finally below + except OSError: + self.send_error(404, "Not found") + return + with contextlib.suppress(OSError): + size = os.fstat(f.fileno()).st_size + + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", str(size)) + self.end_headers() + try: + while True: + buf = f.read(CHUNK) + if not buf: + break + self.wfile.write(buf) + except (BrokenPipeError, ConnectionResetError): + # Client went away mid-transfer; nothing to do. + pass + finally: + with contextlib.suppress(OSError): + f.close() + + # --- DELETE ----------------------------------------------------------- + def do_DELETE(self) -> None: + data_path = self._resolve(self.path) + if data_path is None or data_path.suffix != ".safetensors": + self.send_error(404, "Not found") + return + lock_path = Path(str(data_path) + ".lock") + for p in (data_path, lock_path): + with contextlib.suppress(OSError): + p.unlink() + self.send_response(204) + self.end_headers() + + +class ThreadingHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer): + """One thread per request so concurrent GETs don't serialize.""" + + daemon_threads = True + allow_reuse_address = True + + # OS-level TCP listen backlog (SOMAXCONN is typically 128 on Linux); + # set explicitly so the server can accept bursts of concurrent GET/DELETE + # requests without ECONNREFUSED when multiple trainer workers hit the + # same vLLM node simultaneously. + request_queue_size = 128 + +def make_handler(root: Path, lock_timeout: float) -> type[HiddenStatesHandler]: + return type( + "BoundHiddenStatesHandler", + (HiddenStatesHandler,), + {"root": root, "lock_timeout": lock_timeout}, + ) + + +def _positive_float(value: str) -> float: + """argparse type: finite float strictly greater than zero.""" + try: + f = float(value) + except ValueError as e: + raise argparse.ArgumentTypeError(f"{value!r} is not a number") from e + if not math.isfinite(f) or f <= 0: + raise argparse.ArgumentTypeError(f"{value!r} must be a positive finite number") + return f + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Serve connector-written hidden-states files over HTTP for the " + "'http' hs_connectors backend. Blocks on the connector's .lock " + "flock until each write is complete, then streams the file." + ), + ) + parser.add_argument( + "--root", + type=str, + default=DEFAULT_ROOT, + help=( + "Directory the connector writes hidden states to (its " + "shared_storage_path / --hidden-states-path). Must be local fast " + f"disk. Default: {DEFAULT_ROOT}" + ), + ) + parser.add_argument( + "--port", + type=int, + required=True, + help="Port to listen on.", + ) + parser.add_argument( + "--host", + type=str, + default="127.0.0.1", + help=( + "Bind address. The trainer usually runs on another node; pass " + "--host 0.0.0.0 (or the node's IP) to accept cross-node " + "connections. Default: 127.0.0.1." + ), + ) + parser.add_argument( + "--lock-timeout", + type=_positive_float, + default=DEFAULT_LOCK_TIMEOUT, + help=( + "Seconds to wait for the connector to finish writing a file before " + f"giving up (GET). Default: {DEFAULT_LOCK_TIMEOUT}" + ), + ) + parser.add_argument( + "--ttl", + type=_positive_float, + default=DEFAULT_TTL, + help=( + "Stale .safetensors/.lock files older than this many seconds are " + f"deleted by the background sweeper. Default: {DEFAULT_TTL}" + ), + ) + parser.add_argument( + "--ttl-interval", + type=_positive_float, + default=DEFAULT_TTL_INTERVAL, + help=f"Seconds between sweeper passes. Default: {DEFAULT_TTL_INTERVAL}", + ) + parser.add_argument( + "--no-sweeper", + action="store_true", + help="Disable the background TTL sweeper (not recommended).", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + root = Path(args.root).resolve() + root.mkdir(parents=True, exist_ok=True) + print( + f"[serve_hs] root={root} host={args.host} port={args.port} " + f"lock_timeout={args.lock_timeout}s ttl={args.ttl}s", + flush=True, + ) + + if not args.no_sweeper: + t = threading.Thread( + target=_sweep, + args=(root, args.ttl, args.ttl_interval), + daemon=True, + ) + t.start() + + handler = make_handler(root, args.lock_timeout) + server = ThreadingHTTPServer((args.host, args.port), handler) + try: + server.serve_forever() + except KeyboardInterrupt: + pass + finally: + server.shutdown() + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/hs_connectors/test_http_roundtrip.py b/tests/e2e/hs_connectors/test_http_roundtrip.py new file mode 100644 index 000000000..551d2ad6f --- /dev/null +++ b/tests/e2e/hs_connectors/test_http_roundtrip.py @@ -0,0 +1,210 @@ +"""E2E smoke test for the HTTP hidden-states producer/consumer loop. + +A writer standing in for the vLLM connector writes a safetensors payload under +the ``.lock`` flock protocol to a directory served by ``scripts/serve_hs.py``; +the test then reads it back via ``HttpTransfer`` (standing in for the trainer +on another node) and validates shape and the lock-blocking semantics. + +Only depends on the standard library, so it always runs; the server is +launched automatically on a scratch port. +""" + +import fcntl +import os +import socket +import subprocess +import sys +import threading +import time +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + +hs_connectors = pytest.importorskip( + "hs_connectors.transfer", reason="hs_connectors not installed" +) +HttpTransfer = hs_connectors.HttpTransfer + +REPO_ROOT = Path(__file__).resolve().parents[3] + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +@pytest.fixture +def serve_hs(tmp_path: Path): + """Launch ``scripts/serve_hs.py`` on a scratch port; yield its base URL.""" + port = _free_port() + root = tmp_path / "hidden_states" + root.mkdir() + serve_script = str(REPO_ROOT / "scripts" / "serve_hs.py") + proc = subprocess.Popen( # noqa: S603 - fixed argv, repo-controlled script + [ + sys.executable, + serve_script, + "--root", + str(root), + "--port", + str(port), + "--no-sweeper", + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + base = f"http://127.0.0.1:{port}" + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline: + try: + with socket.create_connection(("127.0.0.1", port), timeout=0.2): + break + except OSError: + if proc.poll() is not None: + raise RuntimeError("serve_hs.py exited during startup") from None + time.sleep(0.05) + else: + proc.terminate() + raise RuntimeError("serve_hs.py did not start within 10s") + yield base, root + proc.terminate() + proc.wait(timeout=5) + + +def _write_under_lock(root: Path, name: str, tensors: dict[str, torch.Tensor]): + """Write ``tensors`` the way the vLLM connector does: flock held during write.""" + data_path = root / name + lock_path = Path(str(data_path) + ".lock") + fd = os.open(str(lock_path), os.O_CREAT | os.O_RDWR) + try: + fcntl.flock(fd, fcntl.LOCK_EX) + save_file(tensors, str(data_path)) + finally: + fcntl.flock(fd, fcntl.LOCK_UN) + os.close(fd) + return str(data_path) + + +@pytest.mark.e2e +def test_http_hidden_states_roundtrip(serve_hs): + """Producer writes under flock; HttpTransfer GETs, validates, and DELETEs.""" + base, root = serve_hs + hs = torch.randn(32, 4, 16, dtype=torch.bfloat16) + token_ids = torch.randint(0, 1000, (32,)) + handle = _write_under_lock( + root, "req_test.safetensors", {"hidden_states": hs, "token_ids": token_ids} + ) + + transfer = HttpTransfer(base, root, timeout=10.0) + out = transfer.get_generated(handle) + + assert out is not None, "get_generated returned None for an existing sample" + assert torch.equal(out["hidden_states"], hs) + assert torch.equal(out["token_ids"], token_ids) + + transfer.delete(handle) + assert not (root / "req_test.safetensors").exists() + assert not (root / "req_test.safetensors.lock").exists() + + +@pytest.mark.e2e +def test_http_get_blocks_until_lock_released(serve_hs): + """The server must hold the GET until the connector releases the flock.""" + base, root = serve_hs + hs = torch.randn(8, 2, 16) + data_path = root / "req_locked.safetensors" + lock_path = Path(str(data_path) + ".lock") + + fd = os.open(str(lock_path), os.O_CREAT | os.O_RDWR) + fcntl.flock(fd, fcntl.LOCK_EX) + save_file({"hidden_states": hs}, str(data_path)) + + hold_seconds = 1.0 + + def _release(): + fcntl.flock(fd, fcntl.LOCK_UN) + + timer = threading.Timer(hold_seconds, _release) + timer.start() + try: + transfer = HttpTransfer(base, root, timeout=10.0) + start = time.monotonic() + out = transfer.get_generated(str(data_path)) + elapsed = time.monotonic() - start + + assert out is not None + assert torch.equal(out["hidden_states"], hs) + assert elapsed >= hold_seconds, ( + f"GET returned in {elapsed:.2f}s, before the {hold_seconds}s lock release" + ) + finally: + # Cancel-and-join first so the timer can never fire on a closed fd. + timer.cancel() + timer.join() + os.close(fd) + + +@pytest.mark.e2e +def test_http_get_rejected_name_maps_404_to_none(serve_hs): + """A 404 must surface as ``None``, not an exception, so the trainer can retry. + + (The server answers 404 immediately for names it refuses to serve; for a + valid name whose write has not started yet it holds the GET until the + connector's lock file appears — see ``serve_hs.py``.) + """ + base, root = serve_hs + transfer = HttpTransfer(base, root, timeout=10.0) + assert transfer.get_generated(str(root / "not_a_safetensors_file.txt")) is None + + +@pytest.mark.e2e +def test_http_ttl_sweeper_reclaims_stale_files(tmp_path: Path): + """Files older than --ttl are deleted (trainer-crash cleanup).""" + port = _free_port() + root = tmp_path / "hidden_states" + root.mkdir() + proc = subprocess.Popen( # noqa: S603 - fixed argv, repo-controlled script + [ + sys.executable, + str(REPO_ROOT / "scripts" / "serve_hs.py"), + "--root", + str(root), + "--port", + str(port), + "--ttl", + "1", + "--ttl-interval", + "0.2", + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + stale = root / "stale.safetensors" + stale_lock = root / "stale.safetensors.lock" + fresh = root / "fresh.safetensors" + fresh_lock = root / "fresh.safetensors.lock" + for p in (stale, stale_lock, fresh, fresh_lock): + p.write_bytes(b"x") + + # Age the "stale" pair past the 1s TTL while "fresh" stays young. + stale_age = time.time() - 10 + os.utime(stale, (stale_age, stale_age)) + os.utime(stale_lock, (stale_age, stale_age)) + + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline: + if not stale.exists() and not stale_lock.exists(): + break + time.sleep(0.1) + + assert not stale.exists(), "stale .safetensors was not swept" + assert not stale_lock.exists(), "stale .lock was not swept" + assert fresh.exists(), "fresh file was swept prematurely" + assert fresh_lock.exists(), "fresh lock was swept prematurely" + finally: + proc.terminate() + proc.wait(timeout=5)