From ea57e3c0c0e25c4e6fe378569dd949252f01b0b0 Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 26 Oct 2026 20:53:14 +0800 Subject: [PATCH 01/29] Protect sealed weight snapshots and daemon runtime ownership --- src/infiniloader/cache.py | 112 +++++++++++++++-------- src/infiniloader/cli.py | 149 ++++++++++++++++++------------- src/infiniloader/daemon.py | 132 +++++++++++++++++++++------ src/infiniloader/paths.py | 17 ++-- src/infiniloader/runtime.py | 106 ++++++++++++++++++++++ tests/test_cli.py | 6 ++ tests/test_daemon.py | 3 + tests/test_integrity.py | 68 ++++++++++++++ tests/test_lifecycle.py | 98 ++++++++++++++++++++ tests/test_review_regressions.py | 40 +++++++++ tests/test_review_runtime.py | 69 ++++++++++++++ 11 files changed, 668 insertions(+), 132 deletions(-) create mode 100644 src/infiniloader/runtime.py create mode 100644 tests/test_integrity.py create mode 100644 tests/test_lifecycle.py create mode 100644 tests/test_review_regressions.py create mode 100644 tests/test_review_runtime.py diff --git a/src/infiniloader/cache.py b/src/infiniloader/cache.py index 3dd6206..9371474 100644 --- a/src/infiniloader/cache.py +++ b/src/infiniloader/cache.py @@ -7,10 +7,12 @@ import shutil import stat import time +import tempfile from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from pathlib import Path from typing import Callable, Iterable +from .runtime import read_json DEFAULT_EXCLUDES = ( @@ -39,6 +41,9 @@ class SourceFile: size: int mode: int mtime_ns: int + ctime_ns: int = 0 + device: int = 0 + inode: int = 0 @property def rel_posix(self) -> str: @@ -56,8 +61,9 @@ def proc_path(self) -> str: return f"/proc/{os.getpid()}/fd/{self.fd}" def close(self) -> None: + fd, self.fd = self.fd, -1 try: - os.close(self.fd) + os.close(fd) except OSError: pass @@ -148,16 +154,15 @@ def collect_source_files( files: list[SourceFile] = [] for path in sorted(source_root.rglob("*")): - try: - st = path.stat() - except OSError: - continue - if not stat.S_ISREG(st.st_mode): - continue rel = path.relative_to(source_root) if not should_include(rel, include_patterns, exclude_patterns): continue - files.append(SourceFile(path=path, rel=rel, size=st.st_size, mode=st.st_mode, mtime_ns=st.st_mtime_ns)) + st = path.stat() + if not stat.S_ISREG(st.st_mode): + continue + files.append(SourceFile(path=path, rel=rel, size=st.st_size, mode=st.st_mode, + mtime_ns=st.st_mtime_ns, ctime_ns=st.st_ctime_ns, + device=st.st_dev, inode=st.st_ino)) return files @@ -184,6 +189,19 @@ def inspect_source( } +def source_identity(source_root: Path, include_patterns: tuple[str, ...] = (), + exclude_patterns: tuple[str, ...] = DEFAULT_EXCLUDES) -> dict: + root = source_root.expanduser().resolve() + files = collect_source_files(root, include_patterns, exclude_patterns) + description = { + "source_root": str(root), "include_patterns": list(include_patterns), + "exclude_patterns": list(exclude_patterns), + "files": [(f.rel_posix, f.size, f.mtime_ns, f.ctime_ns, f.device, f.inode) for f in files], + } + return {"source_root": str(root), "fingerprint": hashlib.sha256( + json.dumps(description, sort_keys=True).encode()).hexdigest()} + + def _memfd_name(rel: Path) -> str: rel_text = rel.as_posix() digest = hashlib.sha1(rel_text.encode("utf-8")).hexdigest()[:12] @@ -207,7 +225,7 @@ def _copy_with_copy_file_range(src_fd: int, dst_fd: int, size: int, chunk_size: def _copy_with_read_write(src_fd: int, dst_fd: int, size: int, chunk_size: int) -> int: copied = 0 - while True: + while copied < size: data = os.read(src_fd, min(chunk_size, max(size - copied, 1))) if not data: break @@ -216,6 +234,7 @@ def _copy_with_read_write(src_fd: int, dst_fd: int, size: int, chunk_size: int) written = os.write(dst_fd, view) view = view[written:] copied += len(data) + del view, data return copied @@ -223,11 +242,20 @@ def copy_file_to_memfd(source: SourceFile, chunk_size: int = 64 * 1024 * 1024) - if not hasattr(os, "memfd_create"): raise RuntimeError("memfd_create is not available; InfiniLoader requires Linux") - flags = getattr(os, "MFD_CLOEXEC", 0) - dst_fd = os.memfd_create(_memfd_name(source.rel), flags) - digest = hashlib.sha256() + import fcntl + + if chunk_size <= 0: + raise ValueError("chunk_size must be positive") src_fd = os.open(source.path, os.O_RDONLY) + dst_fd = None try: + before = os.fstat(src_fd) + if (before.st_size, before.st_mtime_ns) != (source.size, source.mtime_ns): + raise RuntimeError(f"source changed before preload: {source.path}") + if source.inode and (before.st_dev, before.st_ino, before.st_ctime_ns) != (source.device, source.inode, source.ctime_ns): + raise RuntimeError(f"source replaced before preload: {source.path}") + flags = os.MFD_CLOEXEC | os.MFD_ALLOW_SEALING + dst_fd = os.memfd_create(_memfd_name(source.rel), flags) copied = 0 try: copied = _copy_with_copy_file_range(src_fd, dst_fd, source.size, chunk_size) @@ -244,17 +272,31 @@ def copy_file_to_memfd(source: SourceFile, chunk_size: int = 64 * 1024 * 1024) - if copied != source.size: raise IOError(f"short copy for {source.path}: {copied} != {source.size}") + after = os.fstat(src_fd) + if (before.st_size, before.st_mtime_ns, before.st_ctime_ns) != (after.st_size, after.st_mtime_ns, after.st_ctime_ns): + raise RuntimeError(f"source changed during preload: {source.path}") os.lseek(dst_fd, 0, os.SEEK_SET) # Hash only a small prefix for diagnostics; full hashing doubles load time. - prefix = os.read(dst_fd, min(source.size, 1024 * 1024)) - digest.update(prefix) + digest = hashlib.sha256() + remaining = min(source.size, 1024 * 1024) + while remaining: + data = os.read(dst_fd, min(chunk_size, remaining)) + if not data: + raise IOError("short read while hashing cached prefix") + digest.update(data) + remaining -= len(data) + del data os.lseek(dst_fd, 0, os.SEEK_SET) os.ftruncate(dst_fd, source.size) + fcntl.fcntl(dst_fd, fcntl.F_ADD_SEALS, + fcntl.F_SEAL_WRITE | fcntl.F_SEAL_GROW | fcntl.F_SEAL_SHRINK | fcntl.F_SEAL_SEAL) + os.fchmod(dst_fd, 0o400) return CachedFile(source=source, fd=dst_fd, digest_prefix=digest.hexdigest()[:16]) except Exception: try: - os.close(dst_fd) + if dst_fd is not None: + os.close(dst_fd) finally: raise finally: @@ -271,6 +313,8 @@ def load_model_into_memory( chunk_size: int = 64 * 1024 * 1024, workers: int = 1, ) -> CachedModel: + if chunk_size <= 0 or workers < 1 or not 0 < max_memory_fraction <= 1: + raise ValueError("chunk_size and workers must be positive; memory fraction must be in (0, 1]") source_root = source_root.expanduser().resolve() source_files = collect_source_files(source_root, include_patterns, exclude_patterns) total = sum(item.size for item in source_files) @@ -311,6 +355,8 @@ def load_model_into_memory( copied += item.size if progress: progress(copied, total, item.rel) + if collect_source_files(source_root, include_patterns, exclude_patterns) != source_files: + raise RuntimeError("source model changed during preload; retry using an immutable model revision") cached = [item for item in cached_by_index if item is not None] return CachedModel( source_root=source_root, @@ -340,38 +386,30 @@ def load_model_into_memory( raise -def _safe_remove_view(path: Path) -> None: +def _safe_remove_view(path: Path, expected_pid: int | None = None) -> None: if not path.exists() and not path.is_symlink(): return resolved = path.resolve(strict=False) if str(resolved) in {"/", str(Path.home())}: raise RuntimeError(f"refusing to remove unsafe view path {path}") - is_tmp = ".tmp." in path.name - is_iloader_view = (path / ".iloader" / "manifest.json").exists() if path.is_dir() else False - if path.is_dir() and not is_tmp and not is_iloader_view: - try: - path.rmdir() - return - except OSError as exc: - raise RuntimeError(f"refusing to remove non-InfiniLoader directory {path}") from exc - if path.is_symlink() or path.is_file(): - if not is_tmp: - raise RuntimeError(f"refusing to remove non-InfiniLoader file {path}") - path.unlink() - else: - shutil.rmtree(path) + manifest = read_json(path / ".iloader" / "manifest.json") if path.is_dir() else None + if path.is_symlink() or not manifest or manifest.get("format_version") != 1: + raise RuntimeError(f"refusing to remove non-InfiniLoader directory {path}") + if expected_pid is not None and manifest.get("daemon_pid") != expected_pid: + raise RuntimeError(f"view belongs to another daemon: {path}") + shutil.rmtree(path) def build_view(cached: CachedModel, view_path: Path, replace: bool = False) -> Path: - view_path = view_path.expanduser().resolve(strict=False) - if view_path == cached.source_root: - raise RuntimeError("view path must be different from source model path") + view_path = view_path.expanduser().absolute() + if view_path.is_symlink(): + raise RuntimeError("view path must not be a symlink") + if view_path.resolve().is_relative_to(cached.source_root) or cached.source_root.is_relative_to(view_path.resolve()): + raise RuntimeError("view and source model directories must not overlap") parent = view_path.parent parent.mkdir(parents=True, exist_ok=True) - tmp_path = parent / f".{view_path.name}.tmp.{os.getpid()}" - _safe_remove_view(tmp_path) - tmp_path.mkdir(parents=True, exist_ok=False) + tmp_path = Path(tempfile.mkdtemp(prefix=f".{view_path.name}.tmp.", dir=parent)) try: for item in cached.files: @@ -391,5 +429,5 @@ def build_view(cached: CachedModel, view_path: Path, replace: bool = False) -> P os.replace(tmp_path, view_path) return view_path except Exception: - _safe_remove_view(tmp_path) + shutil.rmtree(tmp_path) raise diff --git a/src/infiniloader/cli.py b/src/infiniloader/cli.py index 615b7ba..9d183b5 100644 --- a/src/infiniloader/cli.py +++ b/src/infiniloader/cli.py @@ -13,7 +13,8 @@ from . import __version__ from .benchmark import benchmark_path -from .cache import DEFAULT_EXCLUDES, build_view, human_bytes, inspect_source, load_model_into_memory +from .cache import DEFAULT_EXCLUDES, build_view, human_bytes, inspect_source, load_model_into_memory, source_identity, _safe_remove_view +from .runtime import FileLock, process_matches, read_json, signal_owned_process, validate_name from .daemon import WeightLoaderDaemon from .distributed import build_tp_plan, create_rank_views, simulate_distributed_read from .ipc import LoaderConnectionError, request @@ -106,6 +107,8 @@ def _starting_status(args: argparse.Namespace, error: Exception) -> dict | None: result.update(state) result["state"] = state.get("state") or "starting" result["pid"] = pid + if not state or not process_matches(pid, state.get("process_token")): + return None return result @@ -234,6 +237,7 @@ def cmd_serve(args: argparse.Namespace) -> int: chunk_size=args.chunk_size, workers=args.workers, quiet=args.quiet, + identity=json.loads(args.identity) if getattr(args, "identity", None) else None, ) return daemon.serve_forever() @@ -274,80 +278,82 @@ def _daemon_command(args: argparse.Namespace, model_path: Path | None = None, up command.extend(["--exclude", pattern]) if args.force: command.append("--force") + command.extend(["--identity", json.dumps(args.identity, sort_keys=True)]) return command def cmd_start(args: argparse.Namespace) -> int: + args.identity = source_identity(Path(args.model), _patterns(args.pattern), tuple(args.exclude or DEFAULT_EXCLUDES)) + with FileLock(_socket(args).with_suffix(".start.lock"), timeout=args.wait_timeout): + return _start_locked(args) + + +def _require_identity(args: argparse.Namespace, status: dict) -> None: + if status.get("identity") != args.identity: + raise RuntimeError("name already belongs to a different model, revision or file selection; use another name or stop it first") + if args.view and Path(status["view"]).absolute() != Path(args.view).expanduser().absolute(): + raise RuntimeError("name already belongs to a different view path") + if args.refresh_cache: + raise RuntimeError("stop the running cache before requesting --refresh-cache") + + +def _start_locked(args: argparse.Namespace) -> int: sock = _socket(args) try: status = request(sock, {"command": "status"}, timeout=1.0)["status"] + except LoaderConnectionError: + status = None + if status is not None: + _require_identity(args, status) print(status["view"]) print(f"already running: name={args.name} pid={status['pid']}", file=sys.stderr) return 0 - except Exception: - pass - - source, upstream_source, local_cache_path = _prepare_model_source(args) - log_path = Path(args.log_file).expanduser() if args.log_file else default_log_path(args.name, getattr(args, "base_dir", None)) - log_path.parent.mkdir(parents=True, exist_ok=True) - log = log_path.open("ab", buffering=0) - proc = subprocess.Popen( - _daemon_command(args, source, upstream_source, local_cache_path), - stdout=log, - stderr=subprocess.STDOUT, - stdin=subprocess.DEVNULL, - start_new_session=True, - close_fds=True, - ) - pid_path = _pid(args) - pid_path.parent.mkdir(parents=True, exist_ok=True) - pid_path.write_text(str(proc.pid), encoding="utf-8") - state_path = _state(args) - state_path.parent.mkdir(parents=True, exist_ok=True) - state_path.write_text( - json.dumps( - { - "name": args.name, - "pid": proc.pid, - "state": "starting", - "view": str(_view(args)), - "socket": str(_socket(args)), - "log": str(log_path), - "started_at": time.time(), - }, - indent=2, - sort_keys=True, - ), - encoding="utf-8", - ) - print(f"starting name={args.name} pid={proc.pid} log={log_path}", file=sys.stderr) - if args.no_wait: - print(_view(args)) - return 0 + starting = _starting_status(args, LoaderConnectionError("not ready")) + proc = None + if starting: + _require_identity(args, starting) + if args.no_wait: + print(starting["view"]) + return 0 + else: + source, upstream_source, local_cache_path = _prepare_model_source(args) + if source_identity(Path(args.model), _patterns(args.pattern), tuple(args.exclude or DEFAULT_EXCLUDES)) != args.identity: + raise RuntimeError("source model changed while preparing local cache; retry") + log_path = Path(args.log_file).expanduser() if args.log_file else default_log_path(args.name, getattr(args, "base_dir", None)) + log_path.parent.mkdir(parents=True, exist_ok=True) + with log_path.open("ab", buffering=0) as log: + proc = subprocess.Popen( + _daemon_command(args, source, upstream_source, local_cache_path), + stdout=log, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, + start_new_session=True, close_fds=True, + ) + print(f"starting name={args.name} pid={proc.pid} log={log_path}", file=sys.stderr) - deadline = time.time() + args.wait_timeout + deadline = time.monotonic() + args.wait_timeout last_progress = None - while time.time() < deadline: - if proc.poll() is not None: - if sys.stderr.isatty() and last_progress: - print(file=sys.stderr) - print(f"daemon exited with code {proc.returncode}; see {log_path}", file=sys.stderr) - return proc.returncode or 1 + while time.monotonic() < deadline: + if proc is not None and proc.poll() is not None: + raise RuntimeError(f"daemon exited with code {proc.returncode}; see {log_path}") + if args.no_wait: + registered = _starting_status(args, LoaderConnectionError("not ready")) + if registered and registered["pid"] == proc.pid: + _require_identity(args, registered) + print(registered["view"]) + return 0 try: status = request(sock, {"command": "status"}, timeout=1.0)["status"] + except LoaderConnectionError: last_progress = _print_progress_update(args, last_progress) - if sys.stderr.isatty() and last_progress: - print(file=sys.stderr) - print(status["view"]) - return 0 - except Exception: - last_progress = _print_progress_update(args, last_progress) - time.sleep(1.0) - if sys.stderr.isatty() and last_progress: - print(file=sys.stderr) - print(f"timeout waiting for daemon readiness; pid={proc.pid} log={log_path}", file=sys.stderr) - print(_view(args)) - return 2 + if proc is None and not _starting_status(args, LoaderConnectionError("not ready")): + raise RuntimeError("daemon exited before becoming ready") + time.sleep(0.1) + continue + _require_identity(args, status) + if sys.stderr.isatty() and last_progress: + print(file=sys.stderr) + print(status["view"]) + return 0 + raise TimeoutError("timeout waiting for daemon readiness; inspect status and log before retrying") def cmd_status(args: argparse.Namespace) -> int: @@ -476,9 +482,18 @@ def _stop_one(args: argparse.Namespace) -> int: if not starting: print(f"{args.name}: {exc}", file=sys.stderr) return 1 - os.kill(int(starting["pid"]), signal.SIGTERM) + if starting.get("state") not in ("starting", "loading") and not getattr(args, "force", False): + raise RuntimeError("ready daemon is unresponsive; no signal was sent, retry or use stop --force") + signal_owned_process(int(starting["pid"]), starting.get("process_token")) if args.remove_view: - shutil.rmtree(_view(args), ignore_errors=True) + deadline = time.monotonic() + args.timeout + while process_matches(int(starting["pid"]), starting.get("process_token")): + if time.monotonic() >= deadline: + raise TimeoutError("daemon has not exited; view was preserved") + time.sleep(.05) + view = Path(starting["view"]) + if view.exists(): + _safe_remove_view(view, expected_pid=int(starting["pid"])) print(f"{args.name}: stopping pid={starting['pid']}") return 0 @@ -505,7 +520,11 @@ def cmd_stop(args: argparse.Namespace) -> int: timeout=args.timeout, remove_view=args.remove_view, ) - rc = max(rc, _stop_one(ns)) + try: + rc = max(rc, _stop_one(ns)) + except (RuntimeError, OSError, ValueError) as exc: + print(f"{row['name']}: {exc}", file=sys.stderr) + rc = max(rc, 1) return rc return _stop_one(args) @@ -813,6 +832,7 @@ def build_parser() -> argparse.ArgumentParser: serve_p.add_argument("--base-dir") serve_p.add_argument("--upstream-source", help=argparse.SUPPRESS) serve_p.add_argument("--local-cache-path", help=argparse.SUPPRESS) + serve_p.add_argument("--identity", help=argparse.SUPPRESS) serve_p.add_argument("--pattern", action="append") serve_p.add_argument("--exclude", action="append") serve_p.add_argument("--replace", action="store_true") @@ -964,6 +984,9 @@ def main(argv: list[str] | None = None) -> int: parser = build_parser() args = parser.parse_args(argv) try: + for attribute in ("name", "name_arg"): + if getattr(args, attribute, None): + validate_name(getattr(args, attribute)) return int(args.func(args)) except KeyboardInterrupt: return 130 diff --git a/src/infiniloader/daemon.py b/src/infiniloader/daemon.py index 13c9e6b..8712cbd 100644 --- a/src/infiniloader/daemon.py +++ b/src/infiniloader/daemon.py @@ -3,13 +3,16 @@ import json import os import signal +import stat import socket import time import traceback +from contextlib import ExitStack from pathlib import Path from typing import Any -from .cache import CachedModel, build_view, human_bytes, load_model_into_memory +from .cache import CachedModel, build_view, human_bytes, load_model_into_memory, source_identity, _safe_remove_view +from .runtime import FileLock, atomic_json, atomic_text, process_matches, read_json, process_token class WeightLoaderDaemon: @@ -32,6 +35,7 @@ def __init__( chunk_size: int, workers: int, quiet: bool = False, + identity: dict | None = None, ) -> None: self.name = name self.source = source @@ -49,12 +53,17 @@ def __init__( self.chunk_size = chunk_size self.workers = workers self.quiet = quiet + self.identity = identity or {} + self.process_token = process_token(os.getpid()) + self._owns_runtime = False self.cached: CachedModel | None = None self.started_at = time.time() self.load_started_at: float | None = None self.load_finished_at: float | None = None self._stop = False self._remove_view_on_exit = False + self._claimed_runtime: dict[Path, tuple[int, int]] = {} + self._socket_identity: dict | None = None def log(self, message: str) -> None: if not self.quiet: @@ -62,15 +71,17 @@ def log(self, message: str) -> None: def _write_state(self, state: dict[str, Any]) -> None: state.setdefault("name", self.name) + state.setdefault("runtime_format", 1) state.setdefault("pid", os.getpid()) state.setdefault("source", str(self.source)) state.setdefault("view", str(self.view)) state.setdefault("socket", str(self.socket_path)) state.setdefault("started_at", self.started_at) - self.state_path.parent.mkdir(parents=True, exist_ok=True) - tmp = self.state_path.with_name(f".{self.state_path.name}.tmp.{os.getpid()}") - tmp.write_text(json.dumps(state, indent=2, sort_keys=True), encoding="utf-8") - os.replace(tmp, self.state_path) + state.setdefault("identity", self.identity) + state.setdefault("process_token", self.process_token) + state.setdefault("runtime_socket", self._socket_identity) + atomic_json(self.state_path, state) + self._claim_runtime_file(self.state_path) def _progress(self, copied: int, total: int, rel: Path) -> None: pct = copied / total * 100 if total else 100.0 @@ -135,6 +146,8 @@ def _status(self) -> dict[str, Any]: assert self.cached is not None return { "name": self.name, + "identity": self.identity, + "process_token": self.process_token, "pid": os.getpid(), "source": str(self.source), "upstream_source": str(self.upstream_source) if self.upstream_source else None, @@ -168,21 +181,67 @@ def _handle(self, payload: dict[str, Any]) -> dict[str, Any]: return {"ok": False, "error": f"unknown command: {command!r}"} def _cleanup_socket(self) -> None: - try: - if self.socket_path.exists() or self.socket_path.is_socket(): - self.socket_path.unlink() - except OSError: - pass - try: - if self.pid_path.exists(): - self.pid_path.unlink() - except OSError: - pass - try: - if self.state_path.exists(): - self.state_path.unlink() - except OSError: - pass + if not self._owns_runtime: + return + for path, identity in self._claimed_runtime.items(): + try: + st = path.lstat() + if (st.st_dev, st.st_ino) == identity: + path.unlink() + except OSError: + pass + + def _claim_runtime_file(self, path: Path) -> None: + st = path.lstat() + self._claimed_runtime[path] = (st.st_dev, st.st_ino) + + def _check_runtime_paths(self) -> None: + paths = (self.socket_path, self.pid_path, self.state_path) + if len({path.absolute() for path in paths}) != len(paths): + raise RuntimeError("socket, PID and state paths must be distinct") + for path in paths: + for source in (self.source, self.upstream_source): + if source and path.resolve().is_relative_to(source.resolve()): + raise RuntimeError(f"runtime files must be outside the model directory: {path}") + if path.is_symlink(): + raise RuntimeError(f"refusing to replace a runtime symlink: {path}") + if self.socket_path.exists() and not self.socket_path.is_socket(): + raise RuntimeError(f"refusing to replace a non-socket file: {self.socket_path}") + for path in (self.pid_path, self.state_path): + if path.exists() and not stat.S_ISREG(path.lstat().st_mode): + raise RuntimeError(f"runtime metadata must be a regular file: {path}") + for path in paths: + if path.exists() and path.lstat().st_uid != os.getuid(): + raise RuntimeError(f"runtime path belongs to another user: {path}") + state = read_json(self.state_path) + if self.state_path.exists(): + if (not state or state.get("runtime_format") != 1 or state.get("name") != self.name + or state.get("socket") != str(self.socket_path) or not isinstance(state.get("pid"), int)): + raise RuntimeError(f"refusing to overwrite unverified runtime state: {self.state_path}") + if process_matches(state["pid"], state.get("process_token")): + raise RuntimeError("runtime metadata belongs to a live daemon") + if self.pid_path.exists(): + try: + pid = int(self.pid_path.read_text().strip()) + except (OSError, ValueError) as exc: + raise RuntimeError(f"refusing to overwrite an unverified PID file: {self.pid_path}") from exc + if not state or pid != state.get("pid") or state.get("runtime_format") != 1: + raise RuntimeError(f"PID file has no matching owned state: {self.pid_path}") + if self.socket_path.exists(): + st = self.socket_path.lstat() + if not state or state.get("runtime_socket") != {"device": st.st_dev, "inode": st.st_ino}: + raise RuntimeError(f"socket has no matching owned state: {self.socket_path}") + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: + probe.settimeout(.2) + try: + probe.connect(str(self.socket_path)) + except ConnectionRefusedError: + pass # Only a verified, non-listening endpoint may be reclaimed. + except OSError as exc: + raise RuntimeError(f"cannot establish that runtime socket is stale: {self.socket_path}") from exc + else: + raise RuntimeError(f"runtime socket still has a listener: {self.socket_path}") + def _install_signal_handlers(self) -> None: signal.signal(signal.SIGTERM, self._handle_signal) @@ -196,24 +255,47 @@ def _handle_signal(self, signum: int, _frame: object) -> None: self._stop = True def serve_forever(self) -> int: + with ExitStack() as locks: + locks.enter_context(FileLock(self.socket_path.with_suffix(".owner.lock"), timeout=0)) + locks.enter_context(FileLock(self.view.parent / f".{self.view.name}.owner.lock", timeout=0)) + for path in (self.pid_path, self.state_path): + locks.enter_context(FileLock(path.with_name(f".{path.name}.owner.lock"), timeout=0)) + self._check_runtime_paths() + self._owns_runtime = True + try: + return self._serve_owned() + finally: + self._owns_runtime = False + + def _serve_owned(self) -> int: + if not self.identity: + self.identity = source_identity(self.upstream_source or self.source, self.include_patterns, self.exclude_patterns) self._install_signal_handlers() self.socket_path.parent.mkdir(parents=True, exist_ok=True) self.pid_path.parent.mkdir(parents=True, exist_ok=True) self.state_path.parent.mkdir(parents=True, exist_ok=True) if self.socket_path.exists(): self.socket_path.unlink() - self.pid_path.write_text(str(os.getpid()), encoding="utf-8") self._write_state({"state": "starting"}) + atomic_text(self.pid_path, str(os.getpid())) + self._claim_runtime_file(self.pid_path) server: socket.socket | None = None try: self.load() + if source_identity(self.upstream_source or self.source, self.include_patterns, self.exclude_patterns) != self.identity: + raise RuntimeError("source model changed during startup; stop and retry") server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) server.settimeout(1.0) server.bind(str(self.socket_path)) + self._claim_runtime_file(self.socket_path) + st = self.socket_path.lstat() + self._socket_identity = {"device": st.st_dev, "inode": st.st_ino} server.listen(16) + self._write_state({"state": "ready", "progress": 1.0, + "copied_bytes": self.cached.total_bytes, **self._status()}) self.log(f"[listen] socket={self.socket_path}") while not self._stop: @@ -243,11 +325,9 @@ def serve_forever(self) -> int: if self.cached is not None: self.cached.close() if self._remove_view_on_exit: - import shutil - try: - shutil.rmtree(self.view) - except OSError: - pass + _safe_remove_view(self.view, expected_pid=os.getpid()) + except (OSError, RuntimeError) as exc: + self.log(f"[cleanup] {exc}") self.log("[stop] daemon exited") return 0 diff --git a/src/infiniloader/paths.py b/src/infiniloader/paths.py index dda2340..d5a811c 100644 --- a/src/infiniloader/paths.py +++ b/src/infiniloader/paths.py @@ -1,7 +1,9 @@ from __future__ import annotations import os +import tempfile from pathlib import Path +from .runtime import validate_name APP_DIR_NAME = "infiniloader" @@ -10,10 +12,8 @@ def _writable_dir(path: Path) -> bool: try: path.mkdir(parents=True, exist_ok=True) - probe = path / ".write-test" - with probe.open("w", encoding="utf-8") as f: - f.write("ok") - probe.unlink(missing_ok=True) + with tempfile.TemporaryFile(dir=path) as f: + f.write(b"ok") return True except OSError: return False @@ -35,8 +35,8 @@ def default_base_dir() -> Path: for candidate in candidates: if _writable_dir(candidate): return candidate - fallback = Path(os.environ.get("TMPDIR", "/tmp")) / APP_DIR_NAME - fallback.mkdir(parents=True, exist_ok=True) + fallback = Path(os.environ.get("TMPDIR", "/tmp")) / f"{APP_DIR_NAME}-{os.getuid()}" + fallback.mkdir(parents=True, exist_ok=True, mode=0o700) return fallback @@ -76,20 +76,25 @@ def model_cache_root(base_dir: str | os.PathLike[str] | None = None) -> Path: def default_socket_path(name: str, base_dir: str | os.PathLike[str] | None = None) -> Path: + validate_name(name) return runtime_dir(base_dir) / f"{name}.sock" def default_pid_path(name: str, base_dir: str | os.PathLike[str] | None = None) -> Path: + validate_name(name) return runtime_dir(base_dir) / f"{name}.pid" def default_state_path(name: str, base_dir: str | os.PathLike[str] | None = None) -> Path: + validate_name(name) return runtime_dir(base_dir) / f"{name}.state.json" def default_log_path(name: str, base_dir: str | os.PathLike[str] | None = None) -> Path: + validate_name(name) return logs_dir(base_dir) / f"{name}.log" def default_view_path(name: str, base_dir: str | os.PathLike[str] | None = None) -> Path: + validate_name(name) return view_root(base_dir) / name diff --git a/src/infiniloader/runtime.py b/src/infiniloader/runtime.py new file mode 100644 index 0000000..e788040 --- /dev/null +++ b/src/infiniloader/runtime.py @@ -0,0 +1,106 @@ +"""Local process ownership, atomic metadata, and advisory locks.""" +from __future__ import annotations + +import json +import os +from pathlib import Path +import re +import signal +import tempfile +import time +from typing import Any + + +def validate_name(name: str) -> str: + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,95}", name) or ".." in name: + raise ValueError("name must be 1-96 ASCII letters, digits, '.', '_' or '-', without '..'") + return name + + +def read_json(path: Path) -> dict[str, Any] | None: + try: + value = json.loads(path.read_text(encoding="utf-8")) + return value if isinstance(value, dict) else None + except (OSError, ValueError): + return None + + +def atomic_json(path: Path, value: dict[str, Any]) -> None: + atomic_text(path, json.dumps(value, indent=2, sort_keys=True)) + + +def atomic_text(path: Path, value: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as stream: + stream.write(value) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + finally: + Path(temporary).unlink(missing_ok=True) + + +class FileLock: + """Never unlink lock files: waiters must keep locking the same inode.""" + + def __init__(self, path: Path, timeout: float = 30.0, shared: bool = False): + self.path, self.timeout, self.shared = path, timeout, shared + self.fd: int | None = None + + def __enter__(self): + import fcntl + + self.path.parent.mkdir(parents=True, exist_ok=True) + self.fd = os.open(self.path, os.O_CREAT | os.O_RDWR | os.O_CLOEXEC | os.O_NOFOLLOW, 0o600) + deadline = time.monotonic() + self.timeout + operation = fcntl.LOCK_SH if self.shared else fcntl.LOCK_EX + try: + while True: + try: + fcntl.flock(self.fd, operation | fcntl.LOCK_NB) + return self + except BlockingIOError: + if time.monotonic() >= deadline: + raise TimeoutError(f"another operation owns {self.path}") + time.sleep(0.05) + except BaseException: + os.close(self.fd) + self.fd = None + raise + + def __exit__(self, *_): + if self.fd is not None: + os.close(self.fd) + self.fd = None + + +def process_token(pid: int) -> str | None: + try: + fields = Path(f"/proc/{pid}/stat").read_text().rsplit(")", 1)[1].split() + if fields[0] == "Z": + return None + boot = Path("/proc/sys/kernel/random/boot_id").read_text().strip() + return f"{boot}:{fields[19]}" + except (OSError, IndexError, ValueError): + return None + + +def process_matches(pid: int, token: str | None) -> bool: + return token is not None and process_token(pid) == token + + +def signal_owned_process(pid: int, token: str | None) -> None: + if hasattr(os, "pidfd_open") and hasattr(signal, "pidfd_send_signal"): + fd = os.pidfd_open(pid) + try: + if not process_matches(pid, token): + raise RuntimeError("refusing to signal a stale or unverified PID") + signal.pidfd_send_signal(fd, signal.SIGTERM) + finally: + os.close(fd) + else: + if not process_matches(pid, token): + raise RuntimeError("refusing to signal a stale or unverified PID") + os.kill(pid, signal.SIGTERM) diff --git a/tests/test_cli.py b/tests/test_cli.py index 78e6015..c7a796c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -5,6 +5,7 @@ from conftest import make_model_dir from infiniloader.cli import build_parser, main +from infiniloader.runtime import process_token def test_main_inspect_json_outputs_model_summary(tmp_path, capsys): @@ -36,6 +37,7 @@ def test_status_reports_starting_when_pid_exists_without_socket(tmp_path, monkey json.dumps( { "state": "loading", + "process_token": process_token(os.getpid()), "copied_bytes": 5, "total_bytes": 10, "rate_bytes_per_second": 2, @@ -60,6 +62,7 @@ def test_list_reports_active_starting_cache_name(tmp_path, monkeypatch, capsys): run_dir = tmp_path / "state" / "run" run_dir.mkdir(parents=True) (run_dir / "unit.pid").write_text(str(os.getpid()), encoding="utf-8") + (run_dir / "unit.state.json").write_text(json.dumps({"process_token": process_token(os.getpid())})) code = main(["list", "--base-dir", str(tmp_path / "state"), "--names-only"]) captured = capsys.readouterr() @@ -72,9 +75,12 @@ def test_stop_all_can_stop_starting_daemon(tmp_path, monkeypatch, capsys): killed = [] monkeypatch.setattr("infiniloader.cli._pid_is_running", lambda _pid: True) monkeypatch.setattr("infiniloader.cli.os.kill", lambda pid, sig: killed.append((pid, sig))) + monkeypatch.setattr("infiniloader.cli.process_matches", lambda pid, token: True) + monkeypatch.setattr("infiniloader.cli.signal_owned_process", lambda pid, token: killed.append((pid, token))) run_dir = tmp_path / "state" / "run" run_dir.mkdir(parents=True) (run_dir / "unit.pid").write_text("12345", encoding="utf-8") + (run_dir / "unit.state.json").write_text(json.dumps({"process_token": "test"})) code = main(["stop", "--all", "--base-dir", str(tmp_path / "state")]) captured = capsys.readouterr() diff --git a/tests/test_daemon.py b/tests/test_daemon.py index b6b42f6..741ad49 100644 --- a/tests/test_daemon.py +++ b/tests/test_daemon.py @@ -39,6 +39,9 @@ def fake_exit(code: int) -> None: chunk_size=1024, workers=1, ) + daemon._owns_runtime = True + daemon._claim_runtime_file(pid_path) + daemon._claim_runtime_file(state_path) with pytest.raises(SystemExit): daemon._handle_signal(signal.SIGTERM, None) diff --git a/tests/test_integrity.py b/tests/test_integrity.py new file mode 100644 index 0000000..df74995 --- /dev/null +++ b/tests/test_integrity.py @@ -0,0 +1,68 @@ +import errno +import os + +import pytest + +from conftest import make_model_dir +from infiniloader.cache import build_view, collect_source_files, copy_file_to_memfd, load_model_into_memory + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux memfd tests") + + +def test_published_weights_cannot_be_modified_or_truncated(tmp_path): + import fcntl + + model = make_model_dir(tmp_path / "model") + cached = load_model_into_memory(model) + try: + view = build_view(cached, tmp_path / "view") + target = view / "model.safetensors" + original = target.read_bytes() + item = next(f for f in cached.files if f.source.rel.name == target.name) + assert fcntl.fcntl(item.fd, fcntl.F_GET_SEALS) & fcntl.F_SEAL_WRITE + with pytest.raises(OSError): + with target.open("r+b") as stream: + stream.write(b"bad") + stream.flush() + with pytest.raises(OSError): + os.ftruncate(item.fd, 0) + assert target.read_bytes() == original + finally: + cached.close() + + +def test_missing_source_does_not_leak_memfd(tmp_path): + model = make_model_dir(tmp_path / "model") + source = collect_source_files(model)[0] + source.path.unlink() + before = len(os.listdir("/proc/self/fd")) + with pytest.raises(FileNotFoundError): + copy_file_to_memfd(source) + assert len(os.listdir("/proc/self/fd")) == before + + +def test_source_mutation_aborts_preload(tmp_path, monkeypatch): + from infiniloader import cache + + model = make_model_dir(tmp_path / "model") + original = cache._copy_with_read_write + def copy_and_mutate(src, dst, size, chunk): + result = original(src, dst, size, chunk) + (model / "model.safetensors").write_bytes(b"different-source") + return result + monkeypatch.setattr(cache, "_copy_with_copy_file_range", lambda *a: (_ for _ in ()).throw(OSError(errno.EXDEV, "fallback"))) + monkeypatch.setattr(cache, "_copy_with_read_write", copy_and_mutate) + before = len(os.listdir("/proc/self/fd")) + with pytest.raises(RuntimeError, match="changed"): + load_model_into_memory(model) + assert len(os.listdir("/proc/self/fd")) == before + + +def test_source_replacement_with_same_size_and_mtime_is_detected(tmp_path): + model = make_model_dir(tmp_path / "model") + source = collect_source_files(model)[0] + source.path.unlink() + source.path.write_bytes(b"x" * source.size) + os.utime(source.path, ns=(source.mtime_ns, source.mtime_ns)) + with pytest.raises(RuntimeError, match="replaced"): + copy_file_to_memfd(source) diff --git a/tests/test_lifecycle.py b/tests/test_lifecycle.py new file mode 100644 index 0000000..7e58383 --- /dev/null +++ b/tests/test_lifecycle.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from contextlib import contextmanager +import json +import os +from pathlib import Path +import subprocess +import sys +import time + +import pytest + +from conftest import make_model_dir +from infiniloader.cli import main +from infiniloader.ipc import request +from infiniloader.runtime import process_token, signal_owned_process + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux integration tests") + + +def wait_for(check, timeout=8): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if check(): + return + time.sleep(.02) + raise AssertionError("daemon did not reach expected state") + + +@contextmanager +def running_daemon(tmp_path, model, *, slow=False): + state = tmp_path / "state" + env = dict(os.environ, PYTHONPATH=str(Path(__file__).resolve().parents[1] / "src")) + code = ( + "import time; from infiniloader.daemon import WeightLoaderDaemon; " + "from infiniloader.cli import main; " + "original=WeightLoaderDaemon.load; " + "WeightLoaderDaemon.load=lambda self: (time.sleep(1), original(self))[1]; " + "raise SystemExit(main())" + ) + command = [sys.executable, "-c", code] if slow else [sys.executable, "-m", "infiniloader"] + command += ["_serve", str(model), "--name", "unit", "--base-dir", str(state), "--replace", "--quiet"] + with (tmp_path / "daemon.log").open("w+") as log: + proc = subprocess.Popen(command, env=env, stdout=log, stderr=log) + try: + wait_for(lambda: (state / "run/unit.state.json").exists()) + if not slow: + wait_for(lambda: (state / "run/unit.sock").exists()) + yield proc, state + finally: + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait(timeout=3) + + +def test_different_model_and_revision_cannot_reuse_name(tmp_path, capsys): + a = make_model_dir(tmp_path / "a") + b = make_model_dir(tmp_path / "b") + with running_daemon(tmp_path, a) as (_, state): + assert main(["start", str(b), "--name", "unit", "--base-dir", str(state)]) == 1 + (a / "config.json").write_text('{"changed":true}') + assert main(["start", str(a), "--name", "unit", "--base-dir", str(state)]) == 1 + assert "different model" in capsys.readouterr().err + + +def test_start_joins_in_progress_daemon(tmp_path, capsys): + model = make_model_dir(tmp_path / "model") + with running_daemon(tmp_path, model, slow=True) as (proc, state): + assert main(["start", str(model), "--name", "unit", "--base-dir", str(state)]) == 0 + status = request(state / "run/unit.sock", {"command": "status"})["status"] + assert status["pid"] == proc.pid + assert (Path(status["view"]) / "model.safetensors").read_bytes() == b"0123456789abcdef" + + +def test_stop_loading_preserves_unmanaged_directory(tmp_path): + model = make_model_dir(tmp_path / "model") + ordinary = tmp_path / "state/views/unit" + ordinary.mkdir(parents=True) + (ordinary / "important").write_text("keep") + with running_daemon(tmp_path, model, slow=True) as (proc, state): + assert main(["stop", "unit", "--base-dir", str(state), "--remove-view"]) == 1 + proc.wait(timeout=3) + assert (ordinary / "important").read_text() == "keep" + + +def test_reused_pid_is_never_signalled(): + with pytest.raises(RuntimeError, match="unverified PID"): + signal_owned_process(os.getpid(), "wrong-start-time") + assert process_token(os.getpid()) is not None + + +@pytest.mark.parametrize("name", ["../model", "/tmp/model", "a/b", "a..b", "", "x" * 97]) +def test_invalid_names_are_rejected(name, tmp_path): + assert main(["start", str(tmp_path), "--name", name, "--base-dir", str(tmp_path / "state")]) == 1 diff --git a/tests/test_review_regressions.py b/tests/test_review_regressions.py new file mode 100644 index 0000000..d672d8b --- /dev/null +++ b/tests/test_review_regressions.py @@ -0,0 +1,40 @@ +"""Cache integrity and allocation regressions from independent expert review.""" +import json +import os +import tracemalloc + +import pytest + +from conftest import make_model_dir +from infiniloader import cache +from infiniloader.runtime import atomic_json + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux review regressions") + + + + + + + + +@pytest.mark.parametrize("backend", ["memfd"]) +def test_copy_workers_hold_only_one_chunk_at_a_time(tmp_path, monkeypatch, backend): + chunk = 1024 * 1024 + model = make_model_dir(tmp_path / "model", payload=b"x" * (3 * chunk)) + source = next(item for item in cache.collect_source_files(model) if item.rel.name == "model.safetensors") + def fallback(*_): + raise OSError("force read/write fallback") + monkeypatch.setattr(cache, "_copy_with_copy_file_range", fallback) + tracemalloc.start() + try: + if backend == "memfd": + cached = cache.copy_file_to_memfd(source, chunk_size=chunk) + cached.close() + else: + from infiniloader.local_cache import _copy_one + _copy_one(source, tmp_path / "copy", chunk) + _, peak = tracemalloc.get_traced_memory() + assert peak < chunk * 1.5, peak + finally: + tracemalloc.stop() diff --git a/tests/test_review_runtime.py b/tests/test_review_runtime.py new file mode 100644 index 0000000..4b692b3 --- /dev/null +++ b/tests/test_review_runtime.py @@ -0,0 +1,69 @@ +"""Runtime and CLI regressions reproduced during independent expert review.""" +import os +from pathlib import Path +import socket +import subprocess +import sys +import time + +import pytest + +from conftest import make_model_dir +from infiniloader.cli import main +from infiniloader.ipc import request +from infiniloader.runtime import process_token +from test_lifecycle import running_daemon + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux review regressions") + + + + +@pytest.mark.parametrize("option", ["--socket", "--pid-file", "--state-file"]) +@pytest.mark.parametrize("symlink", [False, True]) +def test_runtime_collisions_never_overwrite_ordinary_files(tmp_path, option, symlink): + model = make_model_dir(tmp_path / "model") + valuable = tmp_path / "important.txt" + valuable.write_text("preserve this data") + destination = tmp_path / "alias" if symlink else valuable + if symlink: + destination.symlink_to(valuable) + env = dict(os.environ, PYTHONPATH=str(Path(__file__).resolve().parents[1] / "src")) + result = subprocess.run([sys.executable, "-m", "infiniloader", "_serve", str(model), "--name", "unit", + "--base-dir", str(tmp_path / "state"), option, str(destination)], + env=env, capture_output=True, timeout=15) + assert result.returncode == 1, result.stderr.decode() + assert valuable.read_text() == "preserve this data" + if symlink: + assert destination.is_symlink() + + + + +def test_existing_unrelated_unix_listener_is_preserved(tmp_path): + model = make_model_dir(tmp_path / "model") + path = tmp_path / "service.sock" + env = dict(os.environ, PYTHONPATH=str(Path(__file__).resolve().parents[1] / "src")) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as server: + server.bind(str(path)) + server.listen(1) + inode = path.stat().st_ino + result = subprocess.run([sys.executable, "-m", "infiniloader", "_serve", str(model), + "--name", "unit", "--base-dir", str(tmp_path / "state"), "--socket", str(path)], + capture_output=True, env=env, timeout=15) + assert result.returncode == 1 + assert path.stat().st_ino == inode + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(str(path)) + connection, _ = server.accept() + connection.close() + + +def test_shutdown_preserves_replaced_pid_file(tmp_path): + with running_daemon(tmp_path, make_model_dir(tmp_path / "model")) as (proc, state): + replacement = tmp_path / "replacement" + replacement.write_text("unrelated replacement data") + replacement.replace(state / "run/unit.pid") + request(state / "run/unit.sock", {"command": "stop"}) + proc.wait(timeout=4) + assert (state / "run/unit.pid").read_text() == "unrelated replacement data" From 63e7ea42ac5dcef3f8384a455f859510015f704b Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 26 Oct 2026 20:53:36 +0800 Subject: [PATCH 02/29] Bound IPC requests and cancel failed distributed reads --- src/infiniloader/cli.py | 6 +++ src/infiniloader/daemon.py | 62 ++++++++++++++++++------- src/infiniloader/distributed.py | 80 ++++++++++++++++++++++++++------- src/infiniloader/ipc.py | 9 +++- tests/test_boundaries.py | 32 +++++++++++++ tests/test_ipc_failures.py | 51 +++++++++++++++++++++ tests/test_review_runtime.py | 9 ++++ 7 files changed, 215 insertions(+), 34 deletions(-) create mode 100644 tests/test_boundaries.py create mode 100644 tests/test_ipc_failures.py diff --git a/src/infiniloader/cli.py b/src/infiniloader/cli.py index 9d183b5..112652d 100644 --- a/src/infiniloader/cli.py +++ b/src/infiniloader/cli.py @@ -709,6 +709,7 @@ def cmd_dist_simulate(args: argparse.Namespace) -> int: include_patterns=_patterns(args.pattern) or ("*.safetensors",), exclude_patterns=tuple(args.exclude or DEFAULT_EXCLUDES), chunk_size=args.chunk_size, + timeout=args.timeout, ) if args.json: _print_json(result) @@ -927,6 +928,7 @@ def build_parser() -> argparse.ArgumentParser: dist_sim_p.add_argument("--pattern", action="append") dist_sim_p.add_argument("--exclude", action="append") dist_sim_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) + dist_sim_p.add_argument("--timeout", type=float, default=60.0) dist_sim_p.add_argument("--json", action="store_true") dist_sim_p.set_defaults(func=cmd_dist_simulate) @@ -984,6 +986,10 @@ def main(argv: list[str] | None = None) -> int: parser = build_parser() args = parser.parse_args(argv) try: + for attribute in ("workers", "cache_workers", "chunk_size", "ranks", "tp", "repeat", "timeout", "wait_timeout"): + value = getattr(args, attribute, None) + if value is not None and value <= 0: + raise ValueError(f"{attribute} must be positive") for attribute in ("name", "name_arg"): if getattr(args, attribute, None): validate_name(getattr(args, attribute)) diff --git a/src/infiniloader/daemon.py b/src/infiniloader/daemon.py index 8712cbd..3046a7b 100644 --- a/src/infiniloader/daemon.py +++ b/src/infiniloader/daemon.py @@ -6,8 +6,10 @@ import stat import socket import time -import traceback from contextlib import ExitStack +from concurrent.futures import ThreadPoolExecutor +import struct +import threading from pathlib import Path from typing import Any @@ -64,6 +66,7 @@ def __init__( self._remove_view_on_exit = False self._claimed_runtime: dict[Path, tuple[int, int]] = {} self._socket_identity: dict | None = None + self._connections = threading.BoundedSemaphore(16) def log(self, message: str) -> None: if not self.quiet: @@ -180,6 +183,41 @@ def _handle(self, payload: dict[str, Any]) -> dict[str, Any]: return {"ok": True, "stopping": True} return {"ok": False, "error": f"unknown command: {command!r}"} + def _serve_connection(self, conn: socket.socket) -> None: + try: + with conn: + conn.settimeout(1.0) + if hasattr(socket, "SO_PEERCRED"): + _, uid, _ = struct.unpack("3i", conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12)) + if uid != os.getuid(): + return + data = bytearray() + deadline = time.monotonic() + 1.0 + while b"\n" not in data: + remaining = deadline - time.monotonic() + if remaining <= 0: + return + conn.settimeout(remaining) + chunk = conn.recv(min(65536 - len(data), 4096)) + if not chunk: + return + data.extend(chunk) + if len(data) >= 65536: + return + try: + payload = json.loads(data.split(b"\n", 1)[0]) + if not isinstance(payload, dict): + raise ValueError("request must be a JSON object") + response = self._handle(payload) + except (ValueError, TypeError, RuntimeError) as exc: + response = {"ok": False, "error": str(exc)} + conn.settimeout(1.0) + conn.sendall(json.dumps(response).encode("utf-8") + b"\n") + except OSError: + pass # A stalled or disconnected client must not terminate the daemon. + finally: + self._connections.release() + def _cleanup_socket(self) -> None: if not self._owns_runtime: return @@ -281,6 +319,7 @@ def _serve_owned(self) -> int: self._claim_runtime_file(self.pid_path) server: socket.socket | None = None + clients = ThreadPoolExecutor(max_workers=4, thread_name_prefix="iloader-ipc") try: self.load() @@ -293,6 +332,7 @@ def _serve_owned(self) -> int: self._claim_runtime_file(self.socket_path) st = self.socket_path.lstat() self._socket_identity = {"device": st.st_dev, "inode": st.st_ino} + os.chmod(self.socket_path, 0o600) server.listen(16) self._write_state({"state": "ready", "progress": 1.0, "copied_bytes": self.cached.total_bytes, **self._status()}) @@ -303,22 +343,12 @@ def _serve_owned(self) -> int: conn, _ = server.accept() except socket.timeout: continue - with conn: - data = bytearray() - while True: - chunk = conn.recv(65536) - if not chunk: - break - data.extend(chunk) - if b"\n" in chunk: - break - try: - payload = json.loads(data.split(b"\n", 1)[0].decode("utf-8")) - response = self._handle(payload) - except Exception as exc: # pragma: no cover - defensive server path - response = {"ok": False, "error": str(exc), "traceback": traceback.format_exc()} - conn.sendall(json.dumps(response).encode("utf-8") + b"\n") + if self._connections.acquire(blocking=False): + clients.submit(self._serve_connection, conn) + else: + conn.close() finally: + clients.shutdown(wait=True) if server is not None: server.close() self._cleanup_socket() diff --git a/src/infiniloader/distributed.py b/src/infiniloader/distributed.py index 494997a..f1b8dd3 100644 --- a/src/infiniloader/distributed.py +++ b/src/infiniloader/distributed.py @@ -6,12 +6,15 @@ import shutil import struct import time +import queue as queue_module +import tempfile from dataclasses import dataclass from pathlib import Path from typing import Any from .cache import DEFAULT_EXCLUDES, collect_source_files, human_bytes from .paths import dist_view_root +from .runtime import validate_name, read_json DTYPE_BYTES = { @@ -57,6 +60,8 @@ def parse_safetensors_header(path: Path) -> dict[str, Any]: if len(raw_len) != 8: raise ValueError(f"{path} is too small to be safetensors") (header_len,) = struct.unpack(" 100 * 1024 * 1024 or header_len > path.stat().st_size - 8: + raise ValueError(f"invalid safetensors header length: {path}") header = f.read(header_len) return json.loads(header.decode("utf-8")) @@ -149,6 +154,8 @@ def infer_tp_split(tensor: TensorInfo, tp: int) -> dict[str, Any]: def build_tp_plan(model_dir: Path, tp: int) -> dict[str, Any]: + if tp < 1: + raise ValueError("tp must be positive") tensors = collect_safetensors_metadata(model_dir) planned = [infer_tp_split(tensor, tp) for tensor in tensors] split = [item for item in planned if item["action"] == "split"] @@ -179,8 +186,8 @@ def build_tp_plan(model_dir: Path, tp: int) -> dict[str, Any]: def _safe_remove_dist_view(path: Path) -> None: if not path.exists() and not path.is_symlink(): return - marker = path / ".iloader-dist" / "manifest.json" - if not marker.exists() and ".tmp." not in path.name: + marker = read_json(path / ".iloader-dist" / "manifest.json") + if path.is_symlink() or not marker or marker.get("mode") != "full-hf-compatible": raise RuntimeError(f"refusing to remove non-InfiniLoader distributed view {path}") shutil.rmtree(path) @@ -201,19 +208,17 @@ def create_rank_views( base_dir: str | os.PathLike[str] | None = None, replace: bool = False, ) -> dict[str, Any]: + validate_name(name) + if ranks < 1: + raise ValueError("ranks must be positive") source = source.expanduser().resolve() if not source.is_dir(): raise NotADirectoryError(source) root = dist_view_root(base_dir) / name - if root.exists(): - if not replace: - raise FileExistsError(f"{root} already exists; pass --replace") - _safe_remove_dist_view(root) - tmp = root.parent / f".{root.name}.tmp.{os.getpid()}" - _safe_remove_dist_view(tmp) - tmp.mkdir(parents=True) - + if root.resolve().is_relative_to(source) or source.is_relative_to(root.resolve()): + raise RuntimeError("distributed view and source directories must not overlap") files = collect_source_files(source, (), DEFAULT_EXCLUDES) + tmp = Path(tempfile.mkdtemp(prefix=f".{root.name}.tmp.", dir=root.parent)) try: for rank in range(ranks): rank_dir = tmp / f"rank{rank}" @@ -247,11 +252,15 @@ def create_rank_views( "total_human": human_bytes(sum(item.size for item in files)), } (root_marker / "manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8") + if root.exists() or root.is_symlink(): + if not replace: + raise FileExistsError(f"{root} already exists; pass --replace") + _safe_remove_dist_view(root) os.replace(tmp, root) manifest["rank_paths"] = [str((root / f"rank{rank}").resolve()) for rank in range(ranks)] return manifest except Exception: - _safe_remove_dist_view(tmp) + shutil.rmtree(tmp) raise @@ -260,7 +269,7 @@ def _read_files_worker(rank: int, root: str, include_patterns: tuple[str, ...], path = Path(root) files = collect_source_files(path, include_patterns, exclude_patterns) total = sum(item.size for item in files) - barrier.wait() + barrier.wait(timeout=30) start = time.perf_counter() read_bytes = 0 for item in files: @@ -271,8 +280,11 @@ def _read_files_worker(rank: int, root: str, include_patterns: tuple[str, ...], break read_bytes += len(data) elapsed = time.perf_counter() - start + if read_bytes != total: + raise IOError(f"read {read_bytes} bytes, expected {total}") queue.put({"rank": rank, "path": root, "bytes": read_bytes, "expected_bytes": total, "seconds": elapsed, "ok": True}) except Exception as exc: + barrier.abort() queue.put({"rank": rank, "path": root, "ok": False, "error": str(exc)}) @@ -284,7 +296,10 @@ def simulate_distributed_read( include_patterns: tuple[str, ...] = ("*.safetensors",), exclude_patterns: tuple[str, ...] = DEFAULT_EXCLUDES, chunk_size: int = 64 * 1024 * 1024, + timeout: float = 60.0, ) -> dict[str, Any]: + if ranks < 1 or chunk_size <= 0 or timeout <= 0: + raise ValueError("ranks, chunk_size and timeout must be positive") roots = [] if rank_view_root: for rank in range(ranks): @@ -301,11 +316,42 @@ def simulate_distributed_read( for rank in range(ranks) ] start = time.perf_counter() - for proc in processes: - proc.start() - rows = [queue.get() for _ in processes] - for proc in processes: - proc.join() + rows = [] + try: + for proc in processes: + proc.start() + deadline = time.monotonic() + timeout + while len(rows) < ranks: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + row = queue.get(timeout=min(remaining, .1)) + rows.append(row) + if not row.get("ok"): + break + except queue_module.Empty: + if any(proc.exitcode not in (None, 0) for proc in processes): + break + received = {row["rank"] for row in rows} + for rank in range(ranks): + if rank not in received: + rows.append({"rank": rank, "path": roots[rank], "ok": False, + "error": "rank failed, timed out or was cancelled"}) + finally: + barrier.abort() + for proc in processes: + if proc.pid is None: + continue + proc.join(timeout=.2) + if proc.is_alive(): + proc.terminate() + proc.join(timeout=1) + if proc.is_alive(): + proc.kill() + proc.join(timeout=1) + queue.close() + queue.cancel_join_thread() elapsed = time.perf_counter() - start ok = all(row.get("ok") for row in rows) aggregate_bytes = sum(row.get("bytes", 0) for row in rows) diff --git a/src/infiniloader/ipc.py b/src/infiniloader/ipc.py index f1684be..f7d33a0 100644 --- a/src/infiniloader/ipc.py +++ b/src/infiniloader/ipc.py @@ -23,6 +23,8 @@ def request(socket_path: str | Path, payload: dict[str, Any], timeout: float = 5 if not chunk: break data.extend(chunk) + if len(data) > 64 * 1024 * 1024: + raise LoaderConnectionError("daemon response exceeds 64 MiB") if b"\n" in chunk: break except OSError as exc: @@ -30,7 +32,12 @@ def request(socket_path: str | Path, payload: dict[str, Any], timeout: float = 5 if not data: raise LoaderConnectionError(f"empty response from {sock_path}") - response = json.loads(data.split(b"\n", 1)[0].decode("utf-8")) + try: + response = json.loads(data.split(b"\n", 1)[0]) + if not isinstance(response, dict): + raise ValueError("response must be a JSON object") + except ValueError as exc: + raise LoaderConnectionError(f"invalid response from {sock_path}: {exc}") from exc if not response.get("ok", False): raise RuntimeError(response.get("error", "InfiniLoader daemon request failed")) return response diff --git a/tests/test_boundaries.py b/tests/test_boundaries.py new file mode 100644 index 0000000..dfd4c65 --- /dev/null +++ b/tests/test_boundaries.py @@ -0,0 +1,32 @@ +import os + +import pytest + +from conftest import make_model_dir +from infiniloader.cache import build_view, load_model_into_memory +from infiniloader.distributed import create_rank_views + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux view tests") + + +def test_view_must_not_overlap_source(tmp_path): + source = make_model_dir(tmp_path / "model") + cached = load_model_into_memory(source) + try: + for view in (source, source / "view", source.parent): + with pytest.raises(RuntimeError, match="overlap"): + build_view(cached, view, replace=True) + assert (source / "model.safetensors").exists() + finally: + cached.close() + + +def test_tmp_name_does_not_authorize_deleting_ordinary_dist_directory(tmp_path): + source = make_model_dir(tmp_path / "model") + base = tmp_path / "state" + ordinary = base / "dist-views" / "normal.tmp.name" + ordinary.mkdir(parents=True) + (ordinary / "important").write_text("keep") + with pytest.raises(RuntimeError, match="non-InfiniLoader"): + create_rank_views(source, name="normal.tmp.name", ranks=2, base_dir=base, replace=True) + assert (ordinary / "important").read_text() == "keep" diff --git a/tests/test_ipc_failures.py b/tests/test_ipc_failures.py new file mode 100644 index 0000000..b984724 --- /dev/null +++ b/tests/test_ipc_failures.py @@ -0,0 +1,51 @@ +import os +import socket +import time + +import pytest + +from conftest import make_model_dir +from infiniloader.distributed import simulate_distributed_read +from infiniloader.ipc import request +from test_lifecycle import running_daemon + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux IPC tests") + + +def test_disconnected_malformed_and_oversized_clients_do_not_kill_daemon(tmp_path): + model = make_model_dir(tmp_path / "model") + with running_daemon(tmp_path, model) as (proc, state): + sock = state / "run/unit.sock" + for payload in (b"", b"[]\n", b"not-json\n", b"x" * 65536): + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(str(sock)) + client.sendall(payload) + client.shutdown(socket.SHUT_RDWR) + assert request(sock, {"command": "ping"})["pong"] + assert proc.poll() is None + assert sock.stat().st_mode & 0o777 == 0o600 + + +def test_idle_client_does_not_block_status_or_stop(tmp_path): + model = make_model_dir(tmp_path / "model") + with running_daemon(tmp_path, model) as (proc, state): + sock = state / "run/unit.sock" + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(str(sock)) + assert request(sock, {"command": "status"}, timeout=.5)["ok"] + assert request(sock, {"command": "stop"}, timeout=.5)["stopping"] + proc.wait(timeout=4) + assert proc.returncode == 0 + + +def test_rank_failure_cancels_barrier_and_reaps_children(tmp_path): + import multiprocessing as mp + + make_model_dir(tmp_path / "rank0") + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + result = simulate_distributed_read(tmp_path, ranks=2, rank_view_root=tmp_path, timeout=2) + assert result["ok"] is False + assert len(result["rank_results"]) == 2 + assert time.monotonic() - start < 5 + assert {p.pid for p in mp.active_children()} == before diff --git a/tests/test_review_runtime.py b/tests/test_review_runtime.py index 4b692b3..0d8c8e6 100644 --- a/tests/test_review_runtime.py +++ b/tests/test_review_runtime.py @@ -67,3 +67,12 @@ def test_shutdown_preserves_replaced_pid_file(tmp_path): request(state / "run/unit.sock", {"command": "stop"}) proc.wait(timeout=4) assert (state / "run/unit.pid").read_text() == "unrelated replacement data" + + +def test_dist_simulate_failure_text_has_all_paths(tmp_path, capsys): + make_model_dir(tmp_path / "rank0") + assert main(["dist-simulate", str(tmp_path), "--ranks", "2", "--rank-view-root", str(tmp_path), "--timeout", "2"]) == 1 + output = capsys.readouterr() + assert "rank 0: failed" in output.out + assert "rank 1: failed" in output.out + assert "error: 'path'" not in output.err From 97335f6f718dc8eea9724bb90908f874ab3cd828 Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 26 Oct 2026 20:54:07 +0800 Subject: [PATCH 03/29] Publish verified disk cache generations with lease-aware collection --- src/infiniloader/cli.py | 47 ++- src/infiniloader/daemon.py | 5 + src/infiniloader/local_cache.py | 502 ++++++++++++++++--------------- tests/test_lifecycle.py | 3 +- tests/test_local_versions.py | 94 ++++++ tests/test_review_regressions.py | 47 ++- 6 files changed, 453 insertions(+), 245 deletions(-) create mode 100644 tests/test_local_versions.py diff --git a/src/infiniloader/cli.py b/src/infiniloader/cli.py index 112652d..d5e17c7 100644 --- a/src/infiniloader/cli.py +++ b/src/infiniloader/cli.py @@ -18,7 +18,7 @@ from .daemon import WeightLoaderDaemon from .distributed import build_tp_plan, create_rank_views, simulate_distributed_read from .ipc import LoaderConnectionError, request -from .local_cache import cache_status, list_model_caches, materialize_model_cache +from .local_cache import cache_status, list_model_caches, materialize_model_cache, verify_cache, garbage_collect, lease_path, touch_cache from .paths import ( default_log_path, default_pid_path, @@ -599,12 +599,40 @@ def cmd_cache_list(args: argparse.Namespace) -> int: return 0 +def cmd_verify(args: argparse.Namespace) -> int: + path = Path(args.path).expanduser().absolute() + with FileLock(lease_path(path), shared=True): + data = verify_cache(path, full=not args.quick) + if args.json: + _print_json(data) + else: + print(f"valid: {data['valid']} ({data['checked_files']} files)") + for error in data["errors"]: + print(error, file=sys.stderr) + return 0 if data["valid"] else 1 + + +def cmd_gc(args: argparse.Namespace) -> int: + data = garbage_collect(cache_root=args.cache_root, base_dir=args.base_dir, + max_bytes=args.max_bytes, ttl=args.ttl, apply=args.apply) + if args.json: + _print_json(data) + else: + print("removed:" if args.apply else "would remove:") + for row in data["candidates"]: + print(f" {row['cache_path']} ({row['total_human']})") + for row in data["skipped"]: + print(f" skipped: {row['cache_path']} ({row['reason']})") + print(f"remaining: {human_bytes(data['remaining_bytes'])}; budget met: {data['budget_met']}") + return 0 if data["budget_met"] else 1 + + BASH_COMPLETION = r'''_iloader_complete() { local cur prev sub commands opts COMPREPLY=() cur="${COMP_WORDS[COMP_CWORD]}" prev="${COMP_WORDS[COMP_CWORD-1]}" - commands="inspect start list status view manifest stop bench selftest cache cache-status cache-list dist-views dist-simulate dist-plan completion" + commands="inspect start list status view manifest stop bench selftest cache cache-status cache-list verify gc dist-views dist-simulate dist-plan completion" if [[ ${COMP_CWORD} -eq 1 ]]; then COMPREPLY=( $(compgen -W "${commands}" -- "${cur}") ) @@ -905,6 +933,21 @@ def build_parser() -> argparse.ArgumentParser: cache_list_p.add_argument("--json", action="store_true") cache_list_p.set_defaults(func=cmd_cache_list) + verify_p = sub.add_parser("verify", help="verify a versioned local cache using SHA-256") + verify_p.add_argument("path") + verify_p.add_argument("--quick", action="store_true", help="check file metadata without reading all bytes") + verify_p.add_argument("--json", action="store_true") + verify_p.set_defaults(func=cmd_verify) + + gc_p = sub.add_parser("gc", help="preview or apply LRU/TTL collection of unleased disk caches") + gc_p.add_argument("--base-dir") + gc_p.add_argument("--cache-root") + gc_p.add_argument("--max-bytes", type=int, help="target total disk cache bytes") + gc_p.add_argument("--ttl", type=float, help="collect caches unused for this many seconds") + gc_p.add_argument("--apply", action="store_true", help="delete the selected caches; default is dry run") + gc_p.add_argument("--json", action="store_true") + gc_p.set_defaults(func=cmd_gc) + list_p = sub.add_parser("list", help="list active resident caches") list_p.add_argument("--base-dir") list_p.add_argument("--timeout", type=float, default=1.0) diff --git a/src/infiniloader/daemon.py b/src/infiniloader/daemon.py index 3046a7b..736543b 100644 --- a/src/infiniloader/daemon.py +++ b/src/infiniloader/daemon.py @@ -15,6 +15,7 @@ from .cache import CachedModel, build_view, human_bytes, load_model_into_memory, source_identity, _safe_remove_view from .runtime import FileLock, atomic_json, atomic_text, process_matches, read_json, process_token +from .local_cache import lease_path, verify_cache class WeightLoaderDaemon: @@ -299,6 +300,10 @@ def serve_forever(self) -> int: for path in (self.pid_path, self.state_path): locks.enter_context(FileLock(path.with_name(f".{path.name}.owner.lock"), timeout=0)) self._check_runtime_paths() + if self.local_cache_path: + locks.enter_context(FileLock(lease_path(self.local_cache_path), shared=True)) + if not verify_cache(self.local_cache_path, full=False)["valid"]: + raise RuntimeError("local cache changed or was removed before daemon startup") self._owns_runtime = True try: return self._serve_owned() diff --git a/src/infiniloader/local_cache.py b/src/infiniloader/local_cache.py index 71fcaf3..c2858e5 100644 --- a/src/infiniloader/local_cache.py +++ b/src/infiniloader/local_cache.py @@ -1,17 +1,22 @@ from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass import hashlib import json import os +from pathlib import Path import shutil +import stat +import tempfile import time -from concurrent.futures import ThreadPoolExecutor, as_completed -from dataclasses import dataclass -from pathlib import Path from typing import Callable +import uuid +import math from .cache import DEFAULT_EXCLUDES, SourceFile, collect_source_files, human_bytes from .paths import model_cache_root +from .runtime import FileLock, atomic_json, read_json @dataclass @@ -24,40 +29,24 @@ class LocalCacheResult: elapsed_seconds: float def to_dict(self) -> dict: - return { - "source_root": str(self.source_root), - "cache_path": str(self.cache_path), - "total_bytes": self.total_bytes, - "total_human": human_bytes(self.total_bytes), - "file_count": self.file_count, - "reused": self.reused, - "elapsed_seconds": self.elapsed_seconds, - } + return dict(source_root=str(self.source_root), cache_path=str(self.cache_path), + total_bytes=self.total_bytes, total_human=human_bytes(self.total_bytes), + file_count=self.file_count, reused=self.reused, elapsed_seconds=self.elapsed_seconds) def _safe_name(text: str) -> str: - cleaned = "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in text) - cleaned = cleaned.strip("._-") - return cleaned or "model" + return "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in text).strip("._-")[:64] or "model" def cache_key(source_root: Path, name: str | None = None) -> str: source_root = source_root.expanduser().resolve() - digest = hashlib.sha1(str(source_root).encode("utf-8")).hexdigest()[:12] - prefix = _safe_name(name or source_root.name) - return f"{prefix}-{digest}" - - -def cache_path_for( - source_root: Path, - *, - name: str | None = None, - cache_root: str | os.PathLike[str] | None = None, - base_dir: str | os.PathLike[str] | None = None, -) -> Path: - root = Path(cache_root).expanduser() if cache_root else model_cache_root(base_dir) + return f"{_safe_name(name or source_root.name)}-{hashlib.sha256(str(source_root).encode()).hexdigest()[:12]}" + + +def _root(cache_root=None, base_dir=None) -> Path: + root = Path(cache_root).expanduser().resolve() if cache_root else model_cache_root(base_dir).resolve() root.mkdir(parents=True, exist_ok=True) - return root / cache_key(source_root, name) + return root def _manifest_path(path: Path) -> Path: @@ -65,236 +54,267 @@ def _manifest_path(path: Path) -> Path: def _files_manifest(files: list[SourceFile]) -> list[dict]: - return [ - { - "relative_path": item.rel_posix, - "size": item.size, - "mtime_ns": item.mtime_ns, - } - for item in files - ] - - -def _build_manifest( - source_root: Path, - files: list[SourceFile], - include_patterns: tuple[str, ...], - exclude_patterns: tuple[str, ...], -) -> dict: - return { - "format_version": 1, - "created_at": time.time(), - "source_root": str(source_root), - "file_count": len(files), - "total_bytes": sum(item.size for item in files), - "include_patterns": list(include_patterns), - "exclude_patterns": list(exclude_patterns), - "files": _files_manifest(files), - } - - -def _load_manifest(path: Path) -> dict | None: + return [dict(relative_path=f.rel_posix, size=f.size, mtime_ns=f.mtime_ns, + ctime_ns=f.ctime_ns, device=f.device, inode=f.inode) for f in files] + + +def _revision(files, include_patterns, exclude_patterns) -> str: + data = dict(files=_files_manifest(files), include=list(include_patterns), exclude=list(exclude_patterns)) + return hashlib.sha256(json.dumps(data, sort_keys=True).encode()).hexdigest()[:24] + + +def lease_path(cache_path: Path) -> Path: + return cache_path.parent / ".leases" / f"{cache_path.name}.lock" + + +def touch_cache(cache_path: Path) -> None: + atomic_json(cache_path.parent / ".usage" / f"{cache_path.name}.json", {"last_used_at": time.time()}) + + +def _safe_file(root: Path, relative: str) -> Path: + rel = Path(relative) + if rel.is_absolute() or not rel.parts or ".." in rel.parts: + raise ValueError("invalid relative path in cache manifest") + path = root / rel + if path.is_symlink() or not path.resolve().is_relative_to(root.resolve()): + raise ValueError("cache file escapes its generation directory") + return path + + +def _hash_stream(stream) -> str: + digest = hashlib.sha256() + for data in iter(lambda: stream.read(8 * 1024 * 1024), b""): + digest.update(data) + return digest.hexdigest() + + +def verify_cache(cache_path: Path, *, full: bool = True) -> dict: + cache_path = cache_path.expanduser().absolute() + manifest = read_json(_manifest_path(cache_path)) + errors, checked = [], 0 + if cache_path.is_symlink() or not manifest or manifest.get("format_version") != 2: + errors.append("not a version-2 InfiniLoader disk cache; rebuild legacy caches to verify") + elif (not isinstance(manifest.get("source_root"), str) + or type(manifest.get("file_count")) is not int or manifest["file_count"] < 1 + or type(manifest.get("total_bytes")) is not int or manifest["total_bytes"] < 0): + errors.append("invalid cache manifest metadata") + else: + records = manifest.get("cached_files") + if (not isinstance(records, list) or not records or len(records) != manifest.get("file_count") + or not all(isinstance(item, dict) for item in records)): + errors.append("incomplete file manifest") + else: + names = [item.get("relative_path") for item in records] + if (any(type(item.get("size")) is not int or item["size"] < 0 for item in records) + or sum(item["size"] for item in records) != manifest["total_bytes"]): + errors.append("cached file sizes do not match total_bytes") + source_records = manifest.get("files") + source_names = [item.get("relative_path") for item in source_records if isinstance(item, dict)] if isinstance(source_records, list) else [] + if (not all(isinstance(name, str) for name in names) or len(set(names)) != len(names) + or not isinstance(source_records, list) + or not all(isinstance(item, dict) for item in source_records) + or not all(isinstance(name, str) for name in source_names) + or sorted(names) != sorted(source_names)): + errors.append("file manifest names are missing, duplicated or inconsistent") + return dict(cache_path=str(cache_path), valid=False, full=full, checked_files=0, errors=errors) + try: + actual = {path.relative_to(cache_path).as_posix() for path in cache_path.rglob("*") + if path.is_symlink() or not path.is_dir()} + expected = set(names) | {".iloader-cache/manifest.json"} + if actual != expected: + errors.append(f"generation file set changed: extra={sorted(actual - expected)}, missing={sorted(expected - actual)}") + except OSError as exc: + errors.append(f"cannot inventory generation: {exc}") + for item in records: + try: + file = _safe_file(cache_path, item["relative_path"]) + st = file.stat() + if not stat.S_ISREG(st.st_mode) or (st.st_size, st.st_mtime_ns, st.st_ctime_ns) != ( + item["size"], item["mtime_ns"], item["ctime_ns"]): + raise ValueError("file metadata changed") + if full: + with file.open("rb") as stream: + digest = _hash_stream(stream) + if digest != item["sha256"]: + raise ValueError("SHA-256 mismatch") + checked += 1 + except (OSError, ValueError, KeyError, TypeError) as exc: + errors.append(f"{item.get('relative_path', '?')}: {exc}") + return dict(cache_path=str(cache_path), valid=not errors, full=full, checked_files=checked, errors=errors) + + +def _created_at(path: Path) -> float: + value = (read_json(_manifest_path(path)) or {}).get("created_at", 0) try: - return json.loads(_manifest_path(path).read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError): - return None - - -def cache_is_valid( - source_root: Path, - cache_path: Path, - files: list[SourceFile], - include_patterns: tuple[str, ...], - exclude_patterns: tuple[str, ...], -) -> bool: - manifest = _load_manifest(cache_path) - if not manifest: - return False - if manifest.get("source_root") != str(source_root): - return False - if tuple(manifest.get("include_patterns", ())) != include_patterns: - return False - if tuple(manifest.get("exclude_patterns", ())) != exclude_patterns: + parsed = float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else 0.0 + return parsed if math.isfinite(parsed) else 0.0 + except (ValueError, OverflowError): + return 0.0 + + +def cache_is_valid(source_root, cache_path, files, include_patterns, exclude_patterns) -> bool: + manifest = read_json(_manifest_path(cache_path)) + if not manifest or manifest.get("source_root") != str(source_root): return False - if manifest.get("files") != _files_manifest(files): + if manifest.get("include_patterns") != list(include_patterns) or manifest.get("exclude_patterns") != list(exclude_patterns): return False - for item in files: - dest = cache_path / item.rel - try: - st = dest.stat() - except OSError: - return False - if st.st_size != item.size: - return False - return True + return manifest.get("files") == _files_manifest(files) and verify_cache(cache_path, full=False)["valid"] -def _safe_remove_cache(path: Path) -> None: - if not path.exists() and not path.is_symlink(): - return - if not _manifest_path(path).exists() and ".tmp." not in path.name: - raise RuntimeError(f"refusing to remove non-InfiniLoader cache directory {path}") - shutil.rmtree(path) +def cache_path_for(source_root: Path, *, name=None, cache_root=None, base_dir=None, + include_patterns=(), exclude_patterns=DEFAULT_EXCLUDES) -> Path: + source_root = source_root.expanduser().resolve() + files = collect_source_files(source_root, include_patterns, exclude_patterns) + root = _root(cache_root, base_dir) + prefix = f"{cache_key(source_root, name)}-{_revision(files, include_patterns, exclude_patterns)}" + for candidate in sorted(root.glob(f"{prefix}*"), + key=_created_at, reverse=True): + if cache_is_valid(source_root, candidate, files, include_patterns, exclude_patterns): + return candidate + return root / prefix -def _copy_one(item: SourceFile, source_root: Path, cache_path: Path, chunk_size: int) -> int: - src = source_root / item.rel +def _copy_one(item: SourceFile, cache_path: Path, chunk_size: int) -> dict: dst = cache_path / item.rel dst.parent.mkdir(parents=True, exist_ok=True) - copy_file_range = getattr(os, "copy_file_range", None) - - src_fd = os.open(src, os.O_RDONLY) - try: - dst_fd = os.open(dst, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, item.mode & 0o777) - try: - copied = 0 - if copy_file_range is not None: - try: - while copied < item.size: - sent = copy_file_range(src_fd, dst_fd, min(chunk_size, item.size - copied)) - if sent == 0: - break - copied += sent - except OSError: - os.lseek(src_fd, 0, os.SEEK_SET) - os.ftruncate(dst_fd, 0) - os.lseek(dst_fd, 0, os.SEEK_SET) - copied = 0 - - if copied != item.size: - os.ftruncate(dst_fd, 0) - os.lseek(src_fd, 0, os.SEEK_SET) - os.lseek(dst_fd, 0, os.SEEK_SET) - copied = 0 - while copied < item.size: - data = os.read(src_fd, min(chunk_size, max(item.size - copied, 1))) - if not data: - break - view = memoryview(data) - while view: - written = os.write(dst_fd, view) - view = view[written:] - copied += len(data) - - if copied != item.size: - raise IOError(f"short copy for {src}: {copied} != {item.size}") - os.fsync(dst_fd) - finally: - os.close(dst_fd) - finally: - os.close(src_fd) - + digest, copied = hashlib.sha256(), 0 + with item.path.open("rb", buffering=0) as src, dst.open("xb", buffering=0) as target: + before = os.fstat(src.fileno()) + if (before.st_size, before.st_mtime_ns, before.st_ctime_ns, before.st_ino) != (item.size, item.mtime_ns, item.ctime_ns, item.inode): + raise RuntimeError(f"source changed before cache copy: {item.path}") + while copied < item.size: + data = src.read(min(chunk_size, item.size - copied)) + if not data: + raise IOError(f"short copy: {item.path}") + digest.update(data) + remaining = memoryview(data) + while remaining: + written = target.write(remaining) + if not written: + raise IOError(f"short write: {dst}") + remaining = remaining[written:] + copied += len(data) + del remaining, data + after = os.fstat(src.fileno()) + if (before.st_size, before.st_mtime_ns, before.st_ctime_ns) != (after.st_size, after.st_mtime_ns, after.st_ctime_ns): + raise RuntimeError(f"source changed during cache copy: {item.path}") + os.fsync(target.fileno()) os.utime(dst, ns=(item.mtime_ns, item.mtime_ns)) - return item.size - - -def materialize_model_cache( - source_root: Path, - *, - name: str | None = None, - cache_root: str | os.PathLike[str] | None = None, - base_dir: str | os.PathLike[str] | None = None, - include_patterns: tuple[str, ...] = (), - exclude_patterns: tuple[str, ...] = DEFAULT_EXCLUDES, - refresh: bool = False, - workers: int = 4, - chunk_size: int = 64 * 1024 * 1024, - progress: Callable[[int, int, Path, bool], None] | None = None, -) -> LocalCacheResult: + dst.chmod(0o444) + st = dst.stat() + return dict(relative_path=item.rel_posix, size=st.st_size, mtime_ns=st.st_mtime_ns, + ctime_ns=st.st_ctime_ns, sha256=digest.hexdigest()) + + +def materialize_model_cache(source_root: Path, *, name=None, cache_root=None, base_dir=None, + include_patterns=(), exclude_patterns=DEFAULT_EXCLUDES, refresh=False, + workers=4, chunk_size=64 * 1024 * 1024, + progress: Callable | None = None) -> LocalCacheResult: + if workers < 1 or chunk_size < 1: + raise ValueError("workers and chunk_size must be positive") start = time.perf_counter() source_root = source_root.expanduser().resolve() - files = collect_source_files(source_root, include_patterns, exclude_patterns) - if not files: - raise RuntimeError(f"no regular files selected under {source_root}") - total = sum(item.size for item in files) - cache_path = cache_path_for(source_root, name=name, cache_root=cache_root, base_dir=base_dir).resolve(strict=False) - - if not refresh and cache_is_valid(source_root, cache_path, files, include_patterns, exclude_patterns): - return LocalCacheResult(source_root, cache_path, total, len(files), True, time.perf_counter() - start) - - cache_path.parent.mkdir(parents=True, exist_ok=True) - tmp_path = cache_path.parent / f".{cache_path.name}.tmp.{os.getpid()}" - _safe_remove_cache(tmp_path) - tmp_path.mkdir(parents=True, exist_ok=False) - - copied = 0 - try: - if workers <= 1: - for item in files: - copied += _copy_one(item, source_root, tmp_path, chunk_size) - if progress: - progress(copied, total, item.rel, False) - else: + root = _root(cache_root, base_dir) + with FileLock(root / ".cache.lock", timeout=600): + files = collect_source_files(source_root, include_patterns, exclude_patterns) + if not files: + raise RuntimeError(f"no regular files selected under {source_root}") + total = sum(item.size for item in files) + path = cache_path_for(source_root, name=name, cache_root=root, + include_patterns=include_patterns, exclude_patterns=exclude_patterns) + if not refresh and cache_is_valid(source_root, path, files, include_patterns, exclude_patterns): + touch_cache(path) + return LocalCacheResult(source_root, path, total, len(files), True, time.perf_counter() - start) + if shutil.disk_usage(root).free < total: + raise OSError("insufficient disk space for a new cache generation; inspect iloader gc") + if path.exists(): + prefix = f"{cache_key(source_root, name)}-{_revision(files, include_patterns, exclude_patterns)}" + path = root / f"{prefix}-r{uuid.uuid4().hex[:12]}" + temporary = Path(tempfile.mkdtemp(prefix=".build-", dir=root)) + try: + copied, records = 0, [] with ThreadPoolExecutor(max_workers=workers) as pool: - futures = {pool.submit(_copy_one, item, source_root, tmp_path, chunk_size): item for item in files} - for future in as_completed(futures): - item = futures[future] - copied += future.result() + for item, record in zip(files, pool.map(lambda f: _copy_one(f, temporary, chunk_size), files)): + records.append(record) + copied += item.size if progress: progress(copied, total, item.rel, False) - - manifest = _build_manifest(source_root, files, include_patterns, exclude_patterns) - manifest_dir = tmp_path / ".iloader-cache" - manifest_dir.mkdir(exist_ok=True) - (manifest_dir / "manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8") - - if cache_path.exists(): - _safe_remove_cache(cache_path) - os.replace(tmp_path, cache_path) - except Exception: - _safe_remove_cache(tmp_path) - raise - - return LocalCacheResult(source_root, cache_path, total, len(files), False, time.perf_counter() - start) - - -def cache_status( - source_root: Path, - *, - name: str | None = None, - cache_root: str | os.PathLike[str] | None = None, - base_dir: str | os.PathLike[str] | None = None, - include_patterns: tuple[str, ...] = (), - exclude_patterns: tuple[str, ...] = DEFAULT_EXCLUDES, -) -> dict: + if collect_source_files(source_root, include_patterns, exclude_patterns) != files: + raise RuntimeError("source model changed during cache copy") + manifest = dict(format_version=2, created_at=time.time(), source_root=str(source_root), + file_count=len(files), total_bytes=total, + include_patterns=list(include_patterns), exclude_patterns=list(exclude_patterns), + files=_files_manifest(files), cached_files=records) + atomic_json(_manifest_path(temporary), manifest) + os.rename(temporary, path) + touch_cache(path) + except BaseException: + if temporary.exists(): + shutil.rmtree(temporary) + raise + return LocalCacheResult(source_root, path, total, len(files), False, time.perf_counter() - start) + + +def cache_status(source_root: Path, *, name=None, cache_root=None, base_dir=None, + include_patterns=(), exclude_patterns=DEFAULT_EXCLUDES) -> dict: source_root = source_root.expanduser().resolve() files = collect_source_files(source_root, include_patterns, exclude_patterns) - cache_path = cache_path_for(source_root, name=name, cache_root=cache_root, base_dir=base_dir).resolve(strict=False) - valid = cache_is_valid(source_root, cache_path, files, include_patterns, exclude_patterns) - return { - "source_root": str(source_root), - "cache_path": str(cache_path), - "exists": cache_path.exists(), - "valid": valid, - "file_count": len(files), - "total_bytes": sum(item.size for item in files), - "total_human": human_bytes(sum(item.size for item in files)), - } - - -def list_model_caches( - *, - cache_root: str | os.PathLike[str] | None = None, - base_dir: str | os.PathLike[str] | None = None, -) -> list[dict]: - root = Path(cache_root).expanduser() if cache_root else model_cache_root(base_dir) - rows: list[dict] = [] - if not root.exists(): - return rows + path = cache_path_for(source_root, name=name, cache_root=cache_root, base_dir=base_dir, + include_patterns=include_patterns, exclude_patterns=exclude_patterns) + total = sum(f.size for f in files) + return dict(source_root=str(source_root), cache_path=str(path), exists=path.exists(), + valid=cache_is_valid(source_root, path, files, include_patterns, exclude_patterns), + file_count=len(files), total_bytes=total, total_human=human_bytes(total)) + + +def list_model_caches(*, cache_root=None, base_dir=None) -> list[dict]: + root = _root(cache_root, base_dir) + rows = [] for path in sorted(root.iterdir()): - if not path.is_dir(): + if path.name.startswith(".") or path.is_symlink() or not path.is_dir(): + continue + manifest = read_json(_manifest_path(path)) + if not manifest or manifest.get("format_version") not in (1, 2): continue - manifest = _load_manifest(path) - if not manifest: + usage = read_json(root / ".usage" / f"{path.name}.json") or {} + # Collection accounts actual model files, even when a damaged manifest lies + # about total_bytes or extra files have appeared after publication. + try: + total = sum(file.lstat().st_size for file in path.rglob("*") + if file.is_file() and not file.is_symlink() and file != _manifest_path(path)) + except OSError: continue - rows.append( - { - "cache_path": str(path), - "source_root": manifest.get("source_root"), - "file_count": manifest.get("file_count"), - "total_bytes": manifest.get("total_bytes"), - "total_human": human_bytes(int(manifest.get("total_bytes") or 0)), - "created_at": manifest.get("created_at"), - } - ) + rows.append(dict(cache_path=str(path), source_root=manifest.get("source_root"), + file_count=manifest.get("file_count"), total_bytes=total, total_human=human_bytes(total), + format_version=manifest["format_version"], created_at=_created_at(path), + last_used_at=usage.get("last_used_at", _created_at(path)))) return rows + + +def garbage_collect(*, cache_root=None, base_dir=None, max_bytes=None, ttl=None, apply=False) -> dict: + if max_bytes is not None and max_bytes < 0 or ttl is not None and ttl < 0: + raise ValueError("max_bytes and ttl must be non-negative") + root = _root(cache_root, base_dir) + removed, skipped = [], [] + with FileLock(root / ".cache.lock", timeout=30): + rows = sorted(list_model_caches(cache_root=root), key=lambda row: row["last_used_at"]) + total, now = sum(row["total_bytes"] for row in rows), time.time() + for row in rows: + expired = ttl is not None and now - row["last_used_at"] >= ttl + over_budget = max_bytes is not None and total > max_bytes + if not (expired or over_budget): + continue + path = Path(row["cache_path"]) + try: + with FileLock(lease_path(path), timeout=0): + if path.is_symlink() or path.parent != root or not read_json(_manifest_path(path)): + raise RuntimeError("cache ownership changed during collection") + if apply: + shutil.rmtree(path) + (root / ".usage" / f"{path.name}.json").unlink(missing_ok=True) + removed.append(row) + total -= row["total_bytes"] + except TimeoutError: + skipped.append(dict(cache_path=str(path), reason="active lease")) + return dict(apply=apply, candidates=removed, skipped=skipped, remaining_bytes=total, + budget_met=max_bytes is None or total <= max_bytes) diff --git a/tests/test_lifecycle.py b/tests/test_lifecycle.py index 7e58383..64fce53 100644 --- a/tests/test_lifecycle.py +++ b/tests/test_lifecycle.py @@ -28,7 +28,7 @@ def wait_for(check, timeout=8): @contextmanager -def running_daemon(tmp_path, model, *, slow=False): +def running_daemon(tmp_path, model, *, slow=False, extra=()): state = tmp_path / "state" env = dict(os.environ, PYTHONPATH=str(Path(__file__).resolve().parents[1] / "src")) code = ( @@ -40,6 +40,7 @@ def running_daemon(tmp_path, model, *, slow=False): ) command = [sys.executable, "-c", code] if slow else [sys.executable, "-m", "infiniloader"] command += ["_serve", str(model), "--name", "unit", "--base-dir", str(state), "--replace", "--quiet"] + command += list(extra) with (tmp_path / "daemon.log").open("w+") as log: proc = subprocess.Popen(command, env=env, stdout=log, stderr=log) try: diff --git a/tests/test_local_versions.py b/tests/test_local_versions.py new file mode 100644 index 0000000..bdea26e --- /dev/null +++ b/tests/test_local_versions.py @@ -0,0 +1,94 @@ +import json +import os +from pathlib import Path +import time + +import pytest + +from conftest import make_model_dir +from infiniloader.cli import main +from infiniloader.local_cache import garbage_collect, lease_path, materialize_model_cache, verify_cache +from infiniloader.runtime import FileLock, atomic_json +from test_lifecycle import running_daemon + +pytestmark = pytest.mark.skipif(os.name != "posix", reason="POSIX file locks") + + +def test_corruption_is_rejected_and_rebuilt_as_new_generation(tmp_path): + model = make_model_dir(tmp_path / "model") + first = materialize_model_cache(model, base_dir=tmp_path / "state") + file = first.cache_path / "model.safetensors" + before = file.stat() + file.chmod(0o644) + file.write_bytes(b"X" * before.st_size) + os.utime(file, ns=(before.st_atime_ns, before.st_mtime_ns)) + assert not verify_cache(first.cache_path, full=False)["valid"] + assert not verify_cache(first.cache_path)["valid"] + second = materialize_model_cache(model, base_dir=tmp_path / "state") + assert not second.reused + assert second.cache_path != first.cache_path + assert verify_cache(second.cache_path)["valid"] + assert first.cache_path.exists() + + +def test_refresh_and_source_revision_preserve_old_reader_paths(tmp_path): + model = make_model_dir(tmp_path / "model") + first = materialize_model_cache(model, base_dir=tmp_path / "state") + refreshed = materialize_model_cache(model, base_dir=tmp_path / "state", refresh=True) + assert refreshed.cache_path != first.cache_path + assert materialize_model_cache(model, base_dir=tmp_path / "state").cache_path == refreshed.cache_path + (model / "model.safetensors").write_bytes(b"new model") + changed = materialize_model_cache(model, base_dir=tmp_path / "state") + assert changed.cache_path != refreshed.cache_path + assert (first.cache_path / "model.safetensors").read_bytes() == b"0123456789abcdef" + assert (changed.cache_path / "model.safetensors").read_bytes() == b"new model" + + +def test_gc_dry_run_lru_and_active_leases(tmp_path): + root = tmp_path / "cache" + first = materialize_model_cache(make_model_dir(tmp_path / "a"), cache_root=root) + second = materialize_model_cache(make_model_dir(tmp_path / "b"), cache_root=root) + atomic_json(root / ".usage" / f"{first.cache_path.name}.json", {"last_used_at": 1}) + preview = garbage_collect(cache_root=root, max_bytes=second.total_bytes) + assert [x["cache_path"] for x in preview["candidates"]] == [str(first.cache_path)] + assert first.cache_path.exists() + with FileLock(lease_path(first.cache_path), shared=True): + result = garbage_collect(cache_root=root, max_bytes=0, apply=True) + assert not result["budget_met"] + assert result["skipped"][0]["reason"] == "active lease" + assert first.cache_path.exists() + assert not second.cache_path.exists() + assert garbage_collect(cache_root=root, ttl=0, apply=True)["remaining_bytes"] == 0 + + +def test_daemon_pins_local_disk_generation(tmp_path): + model = make_model_dir(tmp_path / "model") + local = materialize_model_cache(model, base_dir=tmp_path / "state") + with running_daemon(tmp_path, local.cache_path, extra=["--local-cache-path", str(local.cache_path), + "--upstream-source", str(model)]): + assert not garbage_collect(base_dir=tmp_path / "state", max_bytes=0, apply=True)["budget_met"] + assert main(["verify", str(local.cache_path), "--json"]) == 0 + + +def test_verify_rejects_manifest_traversal(tmp_path): + local = materialize_model_cache(make_model_dir(tmp_path / "model"), base_dir=tmp_path / "state") + manifest_file = local.cache_path / ".iloader-cache/manifest.json" + manifest = json.loads(manifest_file.read_text()) + manifest["cached_files"][0]["relative_path"] = "../outside" + atomic_json(manifest_file, manifest) + assert not verify_cache(local.cache_path)["valid"] + + +def test_full_verify_checks_content_even_if_metadata_matches_manifest(tmp_path): + local = materialize_model_cache(make_model_dir(tmp_path / "model"), base_dir=tmp_path / "state") + manifest_file = local.cache_path / ".iloader-cache/manifest.json" + manifest = json.loads(manifest_file.read_text()) + item = manifest["cached_files"][0] + file = local.cache_path / item["relative_path"] + file.chmod(0o644) + file.write_bytes(b"X" * item["size"]) + current = file.stat() + item.update(mtime_ns=current.st_mtime_ns, ctime_ns=current.st_ctime_ns) + atomic_json(manifest_file, manifest) + assert verify_cache(local.cache_path, full=False)["valid"] + assert any("SHA-256 mismatch" in error for error in verify_cache(local.cache_path)["errors"]) diff --git a/tests/test_review_regressions.py b/tests/test_review_regressions.py index d672d8b..7059b27 100644 --- a/tests/test_review_regressions.py +++ b/tests/test_review_regressions.py @@ -7,18 +7,63 @@ from conftest import make_model_dir from infiniloader import cache +from infiniloader.local_cache import garbage_collect, materialize_model_cache, verify_cache from infiniloader.runtime import atomic_json pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux review regressions") +def test_extra_model_file_invalidates_and_rebuilds_generation(tmp_path): + model = make_model_dir(tmp_path / "model") + first = materialize_model_cache(model, base_dir=tmp_path / "state") + (first.cache_path / "unexpected.safetensors").write_bytes(b"extra weights") + assert not verify_cache(first.cache_path, full=False)["valid"] + assert not verify_cache(first.cache_path)["valid"] + second = materialize_model_cache(model, base_dir=tmp_path / "state") + assert not second.reused + assert second.cache_path != first.cache_path + assert not (second.cache_path / "unexpected.safetensors").exists() +@pytest.mark.parametrize("malformation", ["null", "duplicate", "invalid_source_name", "bad_created_at", "overflow_time", "wrong_total", "bool_total"]) +def test_corrupt_manifest_is_reported_and_refresh_can_recover(tmp_path, malformation): + model = make_model_dir(tmp_path / "model") + first = materialize_model_cache(model, base_dir=tmp_path / "state") + file = first.cache_path / ".iloader-cache/manifest.json" + data = json.loads(file.read_text()) + if malformation == "null": + data["cached_files"][0] = None + elif malformation == "duplicate": + data["cached_files"][0] = data["cached_files"][1] + elif malformation == "invalid_source_name": + data["files"][0]["relative_path"] = None + elif malformation == "bad_created_at": + data["created_at"] = {} + data["cached_files"][0] = None + elif malformation == "overflow_time": + data["created_at"] = 10**1000 + data["cached_files"][0] = None + else: + data["total_bytes"] = False if malformation == "bool_total" else 0 + atomic_json(file, data) + assert not verify_cache(first.cache_path)["valid"] + second = materialize_model_cache(model, base_dir=tmp_path / "state", refresh=True) + assert second.cache_path != first.cache_path + assert verify_cache(second.cache_path)["valid"] +def test_gc_counts_files_even_if_manifest_total_is_wrong(tmp_path): + local = materialize_model_cache(make_model_dir(tmp_path / "model"), base_dir=tmp_path / "state") + file = local.cache_path / ".iloader-cache/manifest.json" + data = json.loads(file.read_text()) + data["total_bytes"] = 0 + atomic_json(file, data) + result = garbage_collect(base_dir=tmp_path / "state", max_bytes=0) + assert len(result["candidates"]) == 1 + assert result["candidates"][0]["total_bytes"] == local.total_bytes -@pytest.mark.parametrize("backend", ["memfd"]) +@pytest.mark.parametrize("backend", ["memfd", "disk"]) def test_copy_workers_hold_only_one_chunk_at_a_time(tmp_path, monkeypatch, backend): chunk = 1024 * 1024 model = make_model_dir(tmp_path / "model", payload=b"x" * (3 * chunk)) From 8cca2e6fdc7f5c5d15075ecb20519073502f01eb Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 26 Oct 2026 20:54:35 +0800 Subject: [PATCH 04/29] Enforce container memory limits and shared resident budgets --- src/infiniloader/cache.py | 22 ++---- src/infiniloader/cli.py | 17 +++++ src/infiniloader/daemon.py | 25 ++++++- src/infiniloader/memory.py | 134 +++++++++++++++++++++++++++++++++++++ tests/test_memory.py | 68 +++++++++++++++++++ 5 files changed, 250 insertions(+), 16 deletions(-) create mode 100644 src/infiniloader/memory.py create mode 100644 tests/test_memory.py diff --git a/src/infiniloader/cache.py b/src/infiniloader/cache.py index 9371474..7c66309 100644 --- a/src/infiniloader/cache.py +++ b/src/infiniloader/cache.py @@ -13,6 +13,7 @@ from pathlib import Path from typing import Callable, Iterable from .runtime import read_json +from .memory import memory_info, check_memory DEFAULT_EXCLUDES = ( @@ -118,14 +119,7 @@ def human_bytes(value: int) -> str: def available_memory_bytes() -> int | None: - meminfo = Path("/proc/meminfo") - try: - for line in meminfo.read_text(encoding="utf-8").splitlines(): - if line.startswith("MemAvailable:"): - return int(line.split()[1]) * 1024 - except OSError: - return None - return None + return memory_info()["available_bytes"] def _matches_any(rel: str, patterns: Iterable[str]) -> bool: @@ -312,6 +306,8 @@ def load_model_into_memory( progress: ProgressFn | None = None, chunk_size: int = 64 * 1024 * 1024, workers: int = 1, + max_memory_bytes: int | None = None, + reserve_bytes: int = 256 * 1024 * 1024, ) -> CachedModel: if chunk_size <= 0 or workers < 1 or not 0 < max_memory_fraction <= 1: raise ValueError("chunk_size and workers must be positive; memory fraction must be in (0, 1]") @@ -323,13 +319,9 @@ def load_model_into_memory( if progress: progress(0, total, Path(".")) - available = available_memory_bytes() - if available is not None and total > int(available * max_memory_fraction) and not force: - raise MemoryError( - "selected model is larger than the memory budget: " - f"{human_bytes(total)} selected, {human_bytes(available)} available, " - f"budget fraction {max_memory_fraction:.2f}; rerun with --force to bypass" - ) + buffers = min(workers, len(source_files)) * min(chunk_size, max(f.size for f in source_files)) + check_memory(total + buffers, fraction=max_memory_fraction, max_bytes=max_memory_bytes, + reserve_bytes=reserve_bytes, force=force) copied = 0 cached_by_index: list[CachedFile | None] = [None] * len(source_files) diff --git a/src/infiniloader/cli.py b/src/infiniloader/cli.py index d5e17c7..483e018 100644 --- a/src/infiniloader/cli.py +++ b/src/infiniloader/cli.py @@ -15,6 +15,7 @@ from .benchmark import benchmark_path from .cache import DEFAULT_EXCLUDES, build_view, human_bytes, inspect_source, load_model_into_memory, source_identity, _safe_remove_view from .runtime import FileLock, process_matches, read_json, signal_owned_process, validate_name +from .memory import parse_bytes from .daemon import WeightLoaderDaemon from .distributed import build_tp_plan, create_rank_views, simulate_distributed_read from .ipc import LoaderConnectionError, request @@ -238,6 +239,9 @@ def cmd_serve(args: argparse.Namespace) -> int: workers=args.workers, quiet=args.quiet, identity=json.loads(args.identity) if getattr(args, "identity", None) else None, + max_memory_bytes=args.max_memory_bytes, + reserve_bytes=args.reserve_bytes, + max_resident_bytes=args.max_resident_bytes, ) return daemon.serve_forever() @@ -279,6 +283,10 @@ def _daemon_command(args: argparse.Namespace, model_path: Path | None = None, up if args.force: command.append("--force") command.extend(["--identity", json.dumps(args.identity, sort_keys=True)]) + command.extend(["--reserve-bytes", str(args.reserve_bytes)]) + for option in ("max_memory_bytes", "max_resident_bytes"): + if getattr(args, option, None) is not None: + command.extend(["--" + option.replace("_", "-"), str(getattr(args, option))]) return command @@ -295,6 +303,10 @@ def _require_identity(args: argparse.Namespace, status: dict) -> None: raise RuntimeError("name already belongs to a different view path") if args.refresh_cache: raise RuntimeError("stop the running cache before requesting --refresh-cache") + for option in ("max_memory_bytes", "max_resident_bytes"): + requested = getattr(args, option, None) + if requested is not None and (status.get("memory_budget") or {}).get(option) != requested: + raise RuntimeError(f"stop the running cache before changing --{option.replace('_', '-')}") def _start_locked(args: argparse.Namespace) -> int: @@ -901,6 +913,11 @@ def build_parser() -> argparse.ArgumentParser: start_p.add_argument("--cache-workers", type=int, default=4) start_p.set_defaults(func=cmd_start) + for memory_parser in (serve_p, start_p): + memory_parser.add_argument("--max-memory-bytes", type=parse_bytes, help="per-cache preload cap including worker buffers") + memory_parser.add_argument("--reserve-bytes", type=parse_bytes, default=256 * 1024 * 1024, help="headroom left for the engine; default 256 MiB") + memory_parser.add_argument("--max-resident-bytes", type=parse_bytes, help="set persistent total resident quota for this base directory") + cache_p = sub.add_parser("cache", help="materialize a local model cache") cache_p.add_argument("model") cache_p.add_argument("--name") diff --git a/src/infiniloader/daemon.py b/src/infiniloader/daemon.py index 736543b..3e268d2 100644 --- a/src/infiniloader/daemon.py +++ b/src/infiniloader/daemon.py @@ -13,9 +13,10 @@ from pathlib import Path from typing import Any -from .cache import CachedModel, build_view, human_bytes, load_model_into_memory, source_identity, _safe_remove_view +from .cache import CachedModel, build_view, human_bytes, load_model_into_memory, source_identity, _safe_remove_view, collect_source_files from .runtime import FileLock, atomic_json, atomic_text, process_matches, read_json, process_token from .local_cache import lease_path, verify_cache +from .memory import ResidentReservation class WeightLoaderDaemon: @@ -39,6 +40,9 @@ def __init__( workers: int, quiet: bool = False, identity: dict | None = None, + max_memory_bytes: int | None = None, + reserve_bytes: int = 256 * 1024 * 1024, + max_resident_bytes: int | None = None, ) -> None: self.name = name self.source = source @@ -59,6 +63,9 @@ def __init__( self.identity = identity or {} self.process_token = process_token(os.getpid()) self._owns_runtime = False + self.max_memory_bytes, self.reserve_bytes = max_memory_bytes, reserve_bytes + self.max_resident_bytes = max_resident_bytes + self._reservation = None self.cached: CachedModel | None = None self.started_at = time.time() self.load_started_at: float | None = None @@ -83,6 +90,7 @@ def _write_state(self, state: dict[str, Any]) -> None: state.setdefault("started_at", self.started_at) state.setdefault("identity", self.identity) state.setdefault("process_token", self.process_token) + state.setdefault("memory_budget", self._reservation.info if self._reservation else None) state.setdefault("runtime_socket", self._socket_identity) atomic_json(self.state_path, state) self._claim_runtime_file(self.state_path) @@ -124,6 +132,8 @@ def load(self) -> None: progress=self._progress, chunk_size=self.chunk_size, workers=self.workers, + max_memory_bytes=self.max_memory_bytes, + reserve_bytes=self.reserve_bytes, ) self.load_finished_at = time.time() build_view(self.cached, self.view, replace=self.replace) @@ -167,6 +177,9 @@ def _status(self) -> dict[str, Any]: "include_patterns": list(self.include_patterns), "exclude_patterns": list(self.exclude_patterns), "workers": self.workers, + "memory_budget": self._reservation.info if self._reservation else None, + "sealed": True, + "physical_residency_guaranteed": False, } def _handle(self, payload: dict[str, Any]) -> dict[str, Any]: @@ -304,6 +317,15 @@ def serve_forever(self) -> int: locks.enter_context(FileLock(lease_path(self.local_cache_path), shared=True)) if not verify_cache(self.local_cache_path, full=False)["valid"]: raise RuntimeError("local cache changed or was removed before daemon startup") + files = collect_source_files(self.source, self.include_patterns, self.exclude_patterns) + if not files: + raise RuntimeError("no files selected") + total = sum(f.size for f in files) + buffers = min(self.workers, len(files)) * min(self.chunk_size, max(f.size for f in files)) + self._reservation = locks.enter_context(ResidentReservation( + self.socket_path.parent, self.name, total, buffers, + max_resident_bytes=self.max_resident_bytes, fraction=self.max_memory_fraction, + max_bytes=self.max_memory_bytes, reserve_bytes=self.reserve_bytes, force=self.force)) self._owns_runtime = True try: return self._serve_owned() @@ -328,6 +350,7 @@ def _serve_owned(self) -> int: try: self.load() + self._reservation.ready() if source_identity(self.upstream_source or self.source, self.include_patterns, self.exclude_patterns) != self.identity: raise RuntimeError("source model changed during startup; stop and retry") diff --git a/src/infiniloader/memory.py b/src/infiniloader/memory.py new file mode 100644 index 0000000..84a7e85 --- /dev/null +++ b/src/infiniloader/memory.py @@ -0,0 +1,134 @@ +"""Container-aware headroom and per-base-directory resident reservations.""" +from __future__ import annotations + +import os +from pathlib import Path +import re + +from .runtime import FileLock, atomic_json, process_matches, process_token, read_json + + +def parse_bytes(value: str) -> int: + match = re.fullmatch(r"(\d+)([KMGT](?:I?B)?|B)?", value.upper()) + if not match: + raise ValueError("size must be an integer optionally followed by B, KiB, MiB, GiB or TiB") + unit = (match[2] or "B")[0] + return int(match[1]) * (1024 ** ("BKMGT".index(unit))) + + +def _number(path: Path) -> int | None: + try: + value = int(path.read_text().strip()) + return value if 0 <= value < 1 << 60 else None + except (OSError, ValueError): + return None + + +def memory_info(proc_root: Path = Path("/proc")) -> dict: + host = None + try: + for line in (proc_root / "meminfo").read_text().splitlines(): + if line.startswith("MemAvailable:"): + host = int(line.split()[1]) * 1024 + except (OSError, ValueError): + pass + groups = [] + try: + memberships = [line.split(":", 2) for line in (proc_root / "self/cgroup").read_text().splitlines()] + for line in (proc_root / "self/mountinfo").read_text().splitlines(): + before, after = line.split(" - ", 1) + fields, fs = before.split(), after.split() + if fs[0] not in ("cgroup", "cgroup2"): + continue + v2 = fs[0] == "cgroup2" + if not v2 and "memory" not in fs[2].split(","): + continue + membership = next((item[2] for item in memberships if (item[0] == "0" if v2 else "memory" in item[1].split(","))), None) + if membership is None: + continue + mount = Path(fields[4].replace("\\040", " ")) + root = Path(fields[3]) + member = Path(membership) + relative = member.relative_to(root) if member.is_relative_to(root) else Path(".") + if ".." in relative.parts: + relative = Path(".") + leaf = mount / relative + if not leaf.is_dir(): + leaf = mount # Namespaced cgroup mounts may already represent the leaf. + while leaf.is_relative_to(mount): + limit = _number(leaf / ("memory.max" if v2 else "memory.limit_in_bytes")) + used = _number(leaf / ("memory.current" if v2 else "memory.usage_in_bytes")) + if limit is not None and used is not None: + groups.append(dict(path=str(leaf), limit_bytes=limit, used_bytes=used, + available_bytes=max(0, limit - used))) + if leaf == mount: + break + leaf = leaf.parent + except (OSError, ValueError, IndexError): + pass + candidates = [group["available_bytes"] for group in groups] + if host is not None: + candidates.append(host) + return dict(host_available_bytes=host, cgroups=groups, + available_bytes=min(candidates) if candidates else None) + + +def check_memory(required: int, *, fraction=.9, max_bytes=None, reserve_bytes=256 * 1024 * 1024, + force=False, pending_bytes=0) -> dict: + if not 0 < fraction <= 1 or reserve_bytes < 0 or max_bytes is not None and max_bytes <= 0: + raise ValueError("invalid memory budget") + info = memory_info() + available = info["available_bytes"] + allowed = max(0, int(available * fraction) - reserve_bytes - pending_bytes) if available is not None else None + if max_bytes is not None and required > max_bytes: + raise MemoryError(f"preload requires {required} bytes, exceeding --max-memory-bytes={max_bytes}") + if not force and allowed is not None and required > allowed: + raise MemoryError(f"preload requires {required} bytes including worker buffers; effective budget is {allowed} bytes") + return dict(**info, budget_bytes=allowed, required_bytes=required, reserve_bytes=reserve_bytes, + pending_bytes=pending_bytes, max_memory_bytes=max_bytes, available_check_bypassed=force) + + +class ResidentReservation: + def __init__(self, run_dir: Path, name: str, total_bytes: int, buffer_bytes: int, *, + max_resident_bytes=None, fraction=.9, max_bytes=None, reserve_bytes=256 * 1024 * 1024, force=False): + self.run_dir, self.name, self.total, self.buffers = run_dir, name, total_bytes, buffer_bytes + self.cap = max_resident_bytes + self.options = dict(fraction=fraction, max_bytes=max_bytes, reserve_bytes=reserve_bytes, force=force) + self.token, self.pid = process_token(os.getpid()), os.getpid() + self.key = f"{self.name}:{self.pid}:{self.token}" + self.info = {} + + def _read(self): + data = read_json(self.run_dir / "resident.json") or {"entries": {}} + data["entries"] = {key: row for key, row in data.get("entries", {}).items() + if process_matches(row["pid"], row.get("process_token"))} + return data + + def __enter__(self): + with FileLock(self.run_dir / "resident.lock"): + data = self._read() + cap = self.cap if self.cap is not None else data.get("max_resident_bytes") + reserved = sum(row["reserved_bytes"] for row in data["entries"].values()) + required = self.total + self.buffers + if cap is not None and reserved + required > cap: + raise MemoryError(f"resident quota exceeded: {reserved} reserved + {required} requested > {cap}") + pending = sum(row["reserved_bytes"] for row in data["entries"].values() if row["pending"]) + self.info = check_memory(required, pending_bytes=pending, **self.options) + self.info["max_resident_bytes"] = cap + data["max_resident_bytes"] = cap + data["entries"][self.key] = dict(name=self.name, pid=self.pid, process_token=self.token, + model_bytes=self.total, reserved_bytes=required, pending=True) + atomic_json(self.run_dir / "resident.json", data) + return self + + def ready(self): + with FileLock(self.run_dir / "resident.lock"): + data = self._read() + data["entries"][self.key].update(reserved_bytes=self.total, pending=False) + atomic_json(self.run_dir / "resident.json", data) + + def __exit__(self, *_): + with FileLock(self.run_dir / "resident.lock"): + data = self._read() + data["entries"].pop(self.key, None) + atomic_json(self.run_dir / "resident.json", data) diff --git a/tests/test_memory.py b/tests/test_memory.py new file mode 100644 index 0000000..d70ebd7 --- /dev/null +++ b/tests/test_memory.py @@ -0,0 +1,68 @@ +import json +import os +from pathlib import Path + +import pytest + +from infiniloader import memory +from infiniloader.memory import ResidentReservation, check_memory, memory_info, parse_bytes +from infiniloader.runtime import atomic_json + + +def fake_cgroup(tmp_path, *, v1=False): + proc = tmp_path / "proc" + (proc / "self").mkdir(parents=True) + (proc / "meminfo").write_text("MemAvailable: 1000000 kB\n") + mount = tmp_path / "cgroup" + leaf = mount / "job" + leaf.mkdir(parents=True) + if v1: + (proc / "self/cgroup").write_text("3:memory:/job\n") + (proc / "self/mountinfo").write_text(f"1 0 0:0 / {mount} rw - cgroup cgroup rw,memory\n") + limit, used = "memory.limit_in_bytes", "memory.usage_in_bytes" + else: + (proc / "self/cgroup").write_text("0::/job\n") + (proc / "self/mountinfo").write_text(f"1 0 0:0 / {mount} rw - cgroup2 cgroup rw\n") + limit, used = "memory.max", "memory.current" + (mount / limit).write_text("3000") + (mount / used).write_text("2800") + (leaf / limit).write_text("2000") + (leaf / used).write_text("500") + return proc + + +@pytest.mark.parametrize("v1", [False, True]) +def test_effective_headroom_honors_all_cgroup_ancestors(tmp_path, v1): + info = memory_info(fake_cgroup(tmp_path, v1=v1)) + assert info["host_available_bytes"] == 1024000000 + assert info["available_bytes"] == 200 + assert len(info["cgroups"]) == 2 + + +def test_force_does_not_bypass_absolute_cap(monkeypatch): + monkeypatch.setattr(memory, "memory_info", lambda: dict(available_bytes=100, cgroups=[], host_available_bytes=100)) + with pytest.raises(MemoryError): + check_memory(90, reserve_bytes=20) + assert check_memory(90, reserve_bytes=20, force=True)["available_check_bypassed"] + with pytest.raises(MemoryError, match="max-memory-bytes"): + check_memory(90, max_bytes=80, force=True) + + +@pytest.mark.skipif(os.name != "posix", reason="POSIX file locks") +def test_reservations_enforce_shared_quota_and_remove_stale_pids(tmp_path, monkeypatch): + monkeypatch.setattr(memory, "memory_info", lambda: dict(available_bytes=10000, cgroups=[], host_available_bytes=10000)) + atomic_json(tmp_path / "resident.json", {"entries": {"dead": {"pid": 99999999, "process_token": "old", "reserved_bytes": 99999}}}) + with ResidentReservation(tmp_path, "a", 60, 10, max_resident_bytes=100, reserve_bytes=0) as a: + with pytest.raises(MemoryError, match="quota"): + with ResidentReservation(tmp_path, "b", 40, 0, reserve_bytes=0): + pass + a.ready() + with ResidentReservation(tmp_path, "b", 40, 0, reserve_bytes=0): + data = json.loads((tmp_path / "resident.json").read_text()) + assert len(data["entries"]) == 2 + assert not json.loads((tmp_path / "resident.json").read_text())["entries"] + + +@pytest.mark.parametrize("value, expected", [("0", 0), ("2GiB", 2 * 1024**3), ("256M", 256 * 1024**2)]) +def test_human_readable_sizes(value, expected): + assert parse_bytes(value) == expected From e2c20659880966071aded2a41905adbe8f5832d8 Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 26 Oct 2026 20:55:31 +0800 Subject: [PATCH 05/29] Add leased commands and cache lifecycle diagnostics --- src/infiniloader/benchmark.py | 8 ++ src/infiniloader/cli.py | 156 ++++++++++++++++++++++++++++++-- src/infiniloader/daemon.py | 66 +++++++++++++- src/infiniloader/diagnostics.py | 62 +++++++++++++ tests/test_lifecycle.py | 33 ++++++- tests/test_management.py | 61 +++++++++++++ tests/test_review_runtime.py | 30 ++++++ 7 files changed, 401 insertions(+), 15 deletions(-) create mode 100644 src/infiniloader/diagnostics.py create mode 100644 tests/test_management.py diff --git a/src/infiniloader/benchmark.py b/src/infiniloader/benchmark.py index 604c655..37071f6 100644 --- a/src/infiniloader/benchmark.py +++ b/src/infiniloader/benchmark.py @@ -2,6 +2,7 @@ import os import time +import statistics from pathlib import Path from .cache import DEFAULT_EXCLUDES, collect_source_files, human_bytes @@ -25,6 +26,8 @@ def benchmark_path( repeat: int = 3, chunk_size: int = 64 * 1024 * 1024, ) -> dict: + if repeat < 1 or chunk_size < 1: + raise ValueError("repeat and chunk_size must be positive") root = root.expanduser().resolve() files = collect_source_files(root, include_patterns, exclude_patterns) total_size = sum(item.size for item in files) @@ -35,6 +38,8 @@ def benchmark_path( for item in files: read_size += read_all_bytes(root / item.rel, chunk_size=chunk_size) elapsed = time.perf_counter() - start + if read_size != total_size: + raise IOError(f"read {read_size} bytes, expected {total_size}") gib = read_size / (1024**3) runs.append( { @@ -54,5 +59,8 @@ def benchmark_path( "repeat": repeat, "runs": runs, "best_gib_per_second": best, + "median_seconds": statistics.median(run["seconds"] for run in runs), + "first_run_seconds": runs[0]["seconds"], + "cache_state": "uncontrolled; this benchmark does not clear OS page cache", "pid": os.getpid(), } diff --git a/src/infiniloader/cli.py b/src/infiniloader/cli.py index 483e018..c61bb3b 100644 --- a/src/infiniloader/cli.py +++ b/src/infiniloader/cli.py @@ -14,7 +14,7 @@ from . import __version__ from .benchmark import benchmark_path from .cache import DEFAULT_EXCLUDES, build_view, human_bytes, inspect_source, load_model_into_memory, source_identity, _safe_remove_view -from .runtime import FileLock, process_matches, read_json, signal_owned_process, validate_name +from .runtime import FileLock, process_matches, process_token, read_json, signal_owned_process, validate_name from .memory import parse_bytes from .daemon import WeightLoaderDaemon from .distributed import build_tp_plan, create_rank_views, simulate_distributed_read @@ -242,6 +242,7 @@ def cmd_serve(args: argparse.Namespace) -> int: max_memory_bytes=args.max_memory_bytes, reserve_bytes=args.reserve_bytes, max_resident_bytes=args.max_resident_bytes, + idle_ttl=args.idle_ttl, ) return daemon.serve_forever() @@ -287,6 +288,8 @@ def _daemon_command(args: argparse.Namespace, model_path: Path | None = None, up for option in ("max_memory_bytes", "max_resident_bytes"): if getattr(args, option, None) is not None: command.extend(["--" + option.replace("_", "-"), str(getattr(args, option))]) + if args.idle_ttl is not None: + command.extend(["--idle-ttl", str(args.idle_ttl)]) return command @@ -303,6 +306,8 @@ def _require_identity(args: argparse.Namespace, status: dict) -> None: raise RuntimeError("name already belongs to a different view path") if args.refresh_cache: raise RuntimeError("stop the running cache before requesting --refresh-cache") + if args.idle_ttl is not None and status.get("idle_ttl") != args.idle_ttl: + raise RuntimeError("stop the running cache before changing --idle-ttl") for option in ("max_memory_bytes", "max_resident_bytes"): requested = getattr(args, option, None) if requested is not None and (status.get("memory_budget") or {}).get(option) != requested: @@ -317,6 +322,7 @@ def _start_locked(args: argparse.Namespace) -> int: status = None if status is not None: _require_identity(args, status) + request(sock, {"command": "touch"}) print(status["view"]) print(f"already running: name={args.name} pid={status['pid']}", file=sys.stderr) return 0 @@ -361,6 +367,7 @@ def _start_locked(args: argparse.Namespace) -> int: time.sleep(0.1) continue _require_identity(args, status) + request(sock, {"command": "touch"}) if sys.stderr.isatty() and last_progress: print(file=sys.stderr) print(status["view"]) @@ -410,6 +417,8 @@ def cmd_status(args: argparse.Namespace) -> int: if status.get("load_seconds") is not None: print(f"load: {status['load_seconds']:.3f}s") print(f"uptime: {status['uptime_seconds']:.1f}s") + print(f"leases: {status.get('active_leases', 0)}; reuse requests: {status.get('reuse_requests', 0)}") + print(f"sealed: {status.get('sealed', False)}; physical residency guaranteed: {status.get('physical_residency_guaranteed', False)}") return 0 @@ -486,7 +495,8 @@ def cmd_list(args: argparse.Namespace) -> int: def _stop_one(args: argparse.Namespace) -> int: _resolve_name(args) try: - response = request(_socket(args), {"command": "stop", "remove_view": args.remove_view}, timeout=args.timeout) + response = request(_socket(args), {"command": "stop", "remove_view": args.remove_view, + "force": getattr(args, "force", False)}, timeout=args.timeout) print(f"{args.name}: {'stopping' if response.get('stopping') else 'ok'}") return 0 except LoaderConnectionError as exc: @@ -531,6 +541,7 @@ def cmd_stop(args: argparse.Namespace) -> int: view=None, timeout=args.timeout, remove_view=args.remove_view, + force=getattr(args, "force", False), ) try: rc = max(rc, _stop_one(ns)) @@ -625,8 +636,13 @@ def cmd_verify(args: argparse.Namespace) -> int: def cmd_gc(args: argparse.Namespace) -> int: - data = garbage_collect(cache_root=args.cache_root, base_dir=args.base_dir, - max_bytes=args.max_bytes, ttl=args.ttl, apply=args.apply) + if args.max_bytes is not None and args.max_bytes < 0 or args.ttl is not None and args.ttl < 0: + raise ValueError("max-bytes and ttl must be non-negative") + if args.resident and args.cache_root: + raise ValueError("--cache-root applies to disk caches; use --base-dir for resident caches") + data = _gc_resident(args) if args.resident else garbage_collect( + cache_root=args.cache_root, base_dir=args.base_dir, + max_bytes=args.max_bytes, ttl=args.ttl, apply=args.apply) if args.json: _print_json(data) else: @@ -639,12 +655,100 @@ def cmd_gc(args: argparse.Namespace) -> int: return 0 if data["budget_met"] else 1 +def _gc_resident(args: argparse.Namespace) -> dict: + rows = _active_daemon_rows(args) + total = sum(row.get("total_bytes", 0) for row in rows) + candidates, skipped = [], [] + for row in sorted(rows, key=lambda r: r.get("last_used_at", float("inf"))): + expired = args.ttl is not None and time.time() - row.get("last_used_at", time.time()) >= args.ttl + if not expired and not (args.max_bytes is not None and total > args.max_bytes): + continue + if row.get("state") != "ready" or row.get("active_leases"): + skipped.append(dict(cache_path=row.get("view", row["name"]), reason="loading, unavailable or leased")) + continue + if args.apply: + try: + request(row["socket"], {"command": "evict", "last_used_at": row["last_used_at"], "remove_view": True}) + except (LoaderConnectionError, RuntimeError) as exc: + skipped.append(dict(cache_path=row["view"], reason=str(exc))) + continue + candidates.append(dict(cache_path=row["view"], total_human=row["total_human"])) + total -= row["total_bytes"] + return dict(apply=args.apply, candidates=candidates, skipped=skipped, remaining_bytes=total, + budget_met=args.max_bytes is None or total <= args.max_bytes) + + +def cmd_doctor(args: argparse.Namespace) -> int: + from .diagnostics import diagnose + + data = diagnose(Path(args.base_dir) if args.base_dir else None) + if args.json: + _print_json(data) + else: + for check in data["checks"]: + print(f"{'ok' if check['ok'] else 'FAIL'}: {check['name']}: {check['detail']}") + print(f"effective memory headroom: {data['memory']['available_bytes']} bytes") + for warning in data["warnings"]: + print(f"note: {warning}") + return 0 if data["ok"] else 1 + + +def cmd_run(args: argparse.Namespace) -> int: + program = args.program[1:] if args.program[:1] == ["--"] else args.program + if args.cache_path and args.name_arg: + program = [args.name_arg] + program + if not program or not (args.name_arg or args.cache_path): + raise ValueError("use run NAME -- COMMAND ... or run --cache-path PATH -- COMMAND ...") + if args.cache_path: + path = Path(args.cache_path).expanduser().resolve() + with FileLock(lease_path(path), shared=True) as lease: + if not verify_cache(path, full=False)["valid"]: + raise RuntimeError("cache is missing or invalid; run verify or rebuild it") + touch_cache(path) + return _run_program(program, path, pass_fds=(lease.fd,)) + args.name = _resolve_name(args) + sock = _socket(args) + lease = request(sock, {"command": "lease-acquire", "pid": os.getpid(), "process_token": process_token(os.getpid())}) + try: + return _run_program(program, Path(lease["view"]), socket_path=sock, lease=lease["lease"]) + finally: + try: + request(sock, {"command": "lease-release", "lease": lease["lease"]}) + except (LoaderConnectionError, RuntimeError): + pass + + +def _run_program(program, path, *, pass_fds=(), socket_path=None, lease=None) -> int: + env = dict(os.environ, ILOADER_MODEL_PATH=str(path)) + proc = subprocess.Popen([part.replace("{model}", str(path)) for part in program], env=env, pass_fds=pass_fds) + try: + if lease: + token = process_token(proc.pid) + if proc.poll() is None: + try: + request(socket_path, {"command": "lease-update", "lease": lease, "pid": proc.pid, + "process_token": token}) + except RuntimeError: + if proc.poll() is None: + raise + return proc.wait() + except BaseException: + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + raise + + BASH_COMPLETION = r'''_iloader_complete() { local cur prev sub commands opts COMPREPLY=() cur="${COMP_WORDS[COMP_CWORD]}" prev="${COMP_WORDS[COMP_CWORD-1]}" - commands="inspect start list status view manifest stop bench selftest cache cache-status cache-list verify gc dist-views dist-simulate dist-plan completion" + commands="inspect start list status view manifest stop bench selftest cache cache-status cache-list verify gc run doctor dist-views dist-simulate dist-plan completion" if [[ ${COMP_CWORD} -eq 1 ]]; then COMPREPLY=( $(compgen -W "${commands}" -- "${cur}") ) @@ -670,7 +774,7 @@ def cmd_gc(args: argparse.Namespace) -> int: sub="${COMP_WORDS[1]}" case "${sub}" in start) - opts="--name --view --socket --pid-file --base-dir --log-file --pattern --exclude --force --no-wait --wait-timeout --chunk-size --workers --max-memory-fraction --local-cache --cache-root --cache-name --refresh-cache --cache-workers --help" + opts="--name --view --socket --pid-file --base-dir --log-file --pattern --exclude --force --no-wait --wait-timeout --chunk-size --workers --max-memory-fraction --max-memory-bytes --max-resident-bytes --reserve-bytes --idle-ttl --local-cache --cache-root --cache-name --refresh-cache --cache-workers --help" ;; list) opts="--base-dir --timeout --json --names-only --help" @@ -679,7 +783,7 @@ def cmd_gc(args: argparse.Namespace) -> int: opts="--name --socket --base-dir --timeout --json --help" ;; stop) - opts="--name --all --socket --base-dir --timeout --remove-view --help" + opts="--name --all --socket --base-dir --timeout --remove-view --force --help" ;; cache|cache-status) opts="--name --cache-name --cache-root --base-dir --pattern --exclude --refresh --cache-workers --chunk-size --quiet --json --help" @@ -687,11 +791,23 @@ def cmd_gc(args: argparse.Namespace) -> int: cache-list) opts="--cache-root --base-dir --json --help" ;; + verify) + opts="--quick --json --help" + ;; + gc) + opts="--base-dir --cache-root --max-bytes --ttl --resident --apply --json --help" + ;; + run) + opts="--base-dir --cache-path --help" + ;; + doctor) + opts="--base-dir --json --help" + ;; dist-views) opts="--name --ranks --base-dir --replace --json --help" ;; dist-simulate) - opts="--ranks --rank-view-root --pattern --exclude --chunk-size --json --help" + opts="--ranks --rank-view-root --pattern --exclude --chunk-size --timeout --json --help" ;; dist-plan) opts="--tp --limit --json --help" @@ -813,6 +929,7 @@ def cmd_bench(args: argparse.Namespace) -> int: for run in result["runs"]: print(f"run {run['run']}: {run['seconds']:.3f}s, {run['gib_per_second']:.3f} GiB/s") print(f"best: {result['best_gib_per_second']:.3f} GiB/s") + print(f"median: {result['median_seconds']:.3f}s; cache state: {result['cache_state']}") return 0 @@ -917,6 +1034,7 @@ def build_parser() -> argparse.ArgumentParser: memory_parser.add_argument("--max-memory-bytes", type=parse_bytes, help="per-cache preload cap including worker buffers") memory_parser.add_argument("--reserve-bytes", type=parse_bytes, default=256 * 1024 * 1024, help="headroom left for the engine; default 256 MiB") memory_parser.add_argument("--max-resident-bytes", type=parse_bytes, help="set persistent total resident quota for this base directory") + memory_parser.add_argument("--idle-ttl", type=float, help="evict after this many idle seconds; use run to protect readers") cache_p = sub.add_parser("cache", help="materialize a local model cache") cache_p.add_argument("model") @@ -959,12 +1077,25 @@ def build_parser() -> argparse.ArgumentParser: gc_p = sub.add_parser("gc", help="preview or apply LRU/TTL collection of unleased disk caches") gc_p.add_argument("--base-dir") gc_p.add_argument("--cache-root") - gc_p.add_argument("--max-bytes", type=int, help="target total disk cache bytes") + gc_p.add_argument("--max-bytes", type=parse_bytes, help="target total cache bytes; accepts GiB/MiB") + gc_p.add_argument("--resident", action="store_true", help="collect resident daemons instead of disk caches") gc_p.add_argument("--ttl", type=float, help="collect caches unused for this many seconds") gc_p.add_argument("--apply", action="store_true", help="delete the selected caches; default is dry run") gc_p.add_argument("--json", action="store_true") gc_p.set_defaults(func=cmd_gc) + run_p = sub.add_parser("run", help="run a foreground command while protecting its cache from eviction") + run_p.add_argument("--base-dir") + run_p.add_argument("--cache-path") + run_p.add_argument("name_arg", nargs="?") + run_p.add_argument("program", nargs=argparse.REMAINDER) + run_p.set_defaults(func=cmd_run) + + doctor_p = sub.add_parser("doctor", help="check Linux, memfd, proc access, locks and resource limits") + doctor_p.add_argument("--base-dir") + doctor_p.add_argument("--json", action="store_true") + doctor_p.set_defaults(func=cmd_doctor) + list_p = sub.add_parser("list", help="list active resident caches") list_p.add_argument("--base-dir") list_p.add_argument("--timeout", type=float, default=1.0) @@ -1022,6 +1153,7 @@ def build_parser() -> argparse.ArgumentParser: stop_p.add_argument("--base-dir") stop_p.add_argument("--timeout", type=float, default=5.0) stop_p.add_argument("--remove-view", action="store_true") + stop_p.add_argument("--force", action="store_true", help="stop even when commands hold active leases") stop_p.set_defaults(func=cmd_stop) bench_p = sub.add_parser("bench", help="sequential read benchmark") @@ -1051,8 +1183,14 @@ def main(argv: list[str] | None = None) -> int: if value is not None and value <= 0: raise ValueError(f"{attribute} must be positive") for attribute in ("name", "name_arg"): + if args.command == "run" and getattr(args, "cache_path", None) and attribute == "name_arg": + continue if getattr(args, attribute, None): validate_name(getattr(args, attribute)) + for attribute in ("idle_ttl", "max_memory_bytes", "max_resident_bytes"): + value = getattr(args, attribute, None) + if value is not None and value <= 0: + raise ValueError(f"{attribute} must be positive") return int(args.func(args)) except KeyboardInterrupt: return 130 diff --git a/src/infiniloader/daemon.py b/src/infiniloader/daemon.py index 3e268d2..2f3a755 100644 --- a/src/infiniloader/daemon.py +++ b/src/infiniloader/daemon.py @@ -3,18 +3,19 @@ import json import os import signal -import stat import socket import time from contextlib import ExitStack from concurrent.futures import ThreadPoolExecutor import struct import threading +import uuid +import stat from pathlib import Path from typing import Any from .cache import CachedModel, build_view, human_bytes, load_model_into_memory, source_identity, _safe_remove_view, collect_source_files -from .runtime import FileLock, atomic_json, atomic_text, process_matches, read_json, process_token +from .runtime import FileLock, atomic_json, atomic_text, process_token, process_matches, read_json from .local_cache import lease_path, verify_cache from .memory import ResidentReservation @@ -43,6 +44,7 @@ def __init__( max_memory_bytes: int | None = None, reserve_bytes: int = 256 * 1024 * 1024, max_resident_bytes: int | None = None, + idle_ttl: float | None = None, ) -> None: self.name = name self.source = source @@ -66,15 +68,21 @@ def __init__( self.max_memory_bytes, self.reserve_bytes = max_memory_bytes, reserve_bytes self.max_resident_bytes = max_resident_bytes self._reservation = None + self.idle_ttl = idle_ttl + self._management_lock = threading.RLock() + self._leases: dict[str, dict] = {} + self.last_used_at = time.time() + self._last_use = time.monotonic() + self.reuse_requests = 0 self.cached: CachedModel | None = None self.started_at = time.time() self.load_started_at: float | None = None self.load_finished_at: float | None = None self._stop = False self._remove_view_on_exit = False + self._connections = threading.BoundedSemaphore(16) self._claimed_runtime: dict[Path, tuple[int, int]] = {} self._socket_identity: dict | None = None - self._connections = threading.BoundedSemaphore(16) def log(self, message: str) -> None: if not self.quiet: @@ -90,6 +98,7 @@ def _write_state(self, state: dict[str, Any]) -> None: state.setdefault("started_at", self.started_at) state.setdefault("identity", self.identity) state.setdefault("process_token", self.process_token) + state.setdefault("idle_ttl", self.idle_ttl) state.setdefault("memory_budget", self._reservation.info if self._reservation else None) state.setdefault("runtime_socket", self._socket_identity) atomic_json(self.state_path, state) @@ -180,10 +189,28 @@ def _status(self) -> dict[str, Any]: "memory_budget": self._reservation.info if self._reservation else None, "sealed": True, "physical_residency_guaranteed": False, + "idle_ttl": self.idle_ttl, + "last_used_at": self.last_used_at, + "active_leases": len(self._active_leases()), + "reuse_requests": self.reuse_requests, } def _handle(self, payload: dict[str, Any]) -> dict[str, Any]: + with self._management_lock: + return self._handle_locked(payload) + + def _active_leases(self) -> dict: + self._leases = {key: lease for key, lease in self._leases.items() + if process_matches(lease["pid"], lease["process_token"])} + return self._leases + + def _touch(self) -> None: + self.last_used_at, self._last_use = time.time(), time.monotonic() + + def _handle_locked(self, payload: dict[str, Any]) -> dict[str, Any]: command = payload.get("command") + if self._stop and command not in ("ping", "status", "lease-release"): + raise RuntimeError("daemon is stopping; retry after it exits") if command == "ping": return {"ok": True, "pong": True} if command == "status": @@ -191,7 +218,32 @@ def _handle(self, payload: dict[str, Any]) -> dict[str, Any]: if command == "manifest": assert self.cached is not None return {"ok": True, "manifest": self.cached.to_manifest(self.view)} - if command == "stop": + if command == "touch": + self._touch() + self.reuse_requests += 1 + return {"ok": True} + if command in ("lease-acquire", "lease-update"): + pid, token = int(payload["pid"]), payload.get("process_token") + if not process_matches(pid, token): + raise RuntimeError("lease owner is not a live process") + if len(self._active_leases()) >= 4096: + raise RuntimeError("too many active leases") + key = payload.get("lease") if command == "lease-update" else uuid.uuid4().hex + if command == "lease-update" and key not in self._leases: + raise RuntimeError("unknown lease") + self._leases[key] = dict(pid=pid, process_token=token) + self._touch() + self.reuse_requests += command == "lease-acquire" + return {"ok": True, "lease": key, "view": str(self.view)} + if command == "lease-release": + self._leases.pop(payload.get("lease"), None) + self._touch() + return {"ok": True} + if command in ("stop", "evict"): + if self._active_leases() and not payload.get("force", False): + raise RuntimeError("cache has active leases; wait or use stop --force") + if command == "evict" and self.last_used_at != payload.get("last_used_at"): + raise RuntimeError("cache was used after collection was planned") self._remove_view_on_exit = bool(payload.get("remove_view", False)) self._stop = True return {"ok": True, "stopping": True} @@ -294,7 +346,6 @@ def _check_runtime_paths(self) -> None: else: raise RuntimeError(f"runtime socket still has a listener: {self.socket_path}") - def _install_signal_handlers(self) -> None: signal.signal(signal.SIGTERM, self._handle_signal) signal.signal(signal.SIGINT, self._handle_signal) @@ -351,6 +402,7 @@ def _serve_owned(self) -> int: try: self.load() self._reservation.ready() + self._touch() if source_identity(self.upstream_source or self.source, self.include_patterns, self.exclude_patterns) != self.identity: raise RuntimeError("source model changed during startup; stop and retry") @@ -367,6 +419,10 @@ def _serve_owned(self) -> int: self.log(f"[listen] socket={self.socket_path}") while not self._stop: + with self._management_lock: + if self.idle_ttl is not None and not self._active_leases() and time.monotonic() - self._last_use >= self.idle_ttl: + self._stop, self._remove_view_on_exit = True, True + break try: conn, _ = server.accept() except socket.timeout: diff --git a/src/infiniloader/diagnostics.py b/src/infiniloader/diagnostics.py new file mode 100644 index 0000000..6a1b3be --- /dev/null +++ b/src/infiniloader/diagnostics.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import os +from pathlib import Path +import platform +import resource +import shutil +import socket +import subprocess +import sys +import tempfile + +from .memory import memory_info +from .paths import default_base_dir +from .runtime import FileLock + + +def diagnose(base_dir: Path | None = None) -> dict: + checks = [] + def record(name, ok, detail): + checks.append(dict(name=name, ok=ok, detail=str(detail))) + + record("linux", sys.platform == "linux", platform.platform()) + base = (base_dir or default_base_dir()).expanduser().resolve() + try: + base.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix="doctor-", dir=base) as tmp: + root = Path(tmp) + with FileLock(root / "probe.lock"): + try: + with FileLock(root / "probe.lock", timeout=0): + record("file_lock", False, "exclusive locks did not exclude another opener") + except TimeoutError: + record("file_lock", True, "exclusive locks work") + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: + sock.bind(str(root / "probe.sock")) + record("unix_socket", True, "socket paths under base directory are usable") + fd = os.memfd_create("iloader-doctor", os.MFD_CLOEXEC | os.MFD_ALLOW_SEALING) + try: + import fcntl + + os.write(fd, b"iloader") + fcntl.fcntl(fd, fcntl.F_ADD_SEALS, + fcntl.F_SEAL_WRITE | fcntl.F_SEAL_GROW | fcntl.F_SEAL_SHRINK | fcntl.F_SEAL_SEAL) + result = subprocess.run([sys.executable, "-c", "import sys; assert open(sys.argv[1],'rb').read()==b'iloader'", + f"/proc/{os.getpid()}/fd/{fd}"], capture_output=True, timeout=5) + record("memfd_seals_and_proc_reopen", result.returncode == 0, + "cross-process reopening works" if result.returncode == 0 else result.stderr.decode(errors="replace")) + finally: + os.close(fd) + except (OSError, RuntimeError, AttributeError, subprocess.TimeoutExpired) as exc: + record("runtime_probe", False, exc) + limits = dict(file_descriptors=list(resource.getrlimit(resource.RLIMIT_NOFILE)), + locked_memory=list(resource.getrlimit(resource.RLIMIT_MEMLOCK))) + usage = shutil.disk_usage(base) if base.exists() else None + warnings = ["memfd pages are not mlocked; swap and memory pressure can affect latency", + "use a node-local base directory; containers must access the daemon PID and /proc fd paths"] + soft = limits["file_descriptors"][0] + if soft != resource.RLIM_INFINITY and soft < 4096: + warnings.append(f"RLIMIT_NOFILE={soft}; checkpoints with many files may require a higher limit") + return dict(ok=all(row["ok"] for row in checks), base_dir=str(base), checks=checks, + memory=memory_info(), limits=limits, disk_free_bytes=usage.free if usage else None, warnings=warnings) diff --git a/tests/test_lifecycle.py b/tests/test_lifecycle.py index 64fce53..c9230a3 100644 --- a/tests/test_lifecycle.py +++ b/tests/test_lifecycle.py @@ -44,7 +44,11 @@ def running_daemon(tmp_path, model, *, slow=False, extra=()): with (tmp_path / "daemon.log").open("w+") as log: proc = subprocess.Popen(command, env=env, stdout=log, stderr=log) try: - wait_for(lambda: (state / "run/unit.state.json").exists()) + def registered(): + if proc.poll() is not None: + raise AssertionError(f"daemon exited {proc.returncode}: {(tmp_path / 'daemon.log').read_text()}") + return (state / "run/unit.state.json").exists() + wait_for(registered, timeout=15) if not slow: wait_for(lambda: (state / "run/unit.sock").exists()) yield proc, state @@ -94,6 +98,33 @@ def test_reused_pid_is_never_signalled(): assert process_token(os.getpid()) is not None +def test_two_simultaneous_cli_starts_share_one_daemon(tmp_path): + model = make_model_dir(tmp_path / "model", payload=b"x" * 1024 * 1024) + state = tmp_path / "state" + env = dict(os.environ, PYTHONPATH=str(Path(__file__).resolve().parents[1] / "src")) + command = [sys.executable, "-m", "infiniloader", "start", str(model), "--name", "unit", + "--base-dir", str(state), "--no-wait"] + clients = [subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env) for _ in range(2)] + daemon = None + try: + outputs = [client.communicate(timeout=20) for client in clients] + assert all(client.returncode == 0 for client in clients), outputs + assert outputs[0][0] == outputs[1][0] + wait_for(lambda: (state / "run/unit.sock").exists()) + daemon = request(state / "run/unit.sock", {"command": "status"})["status"] + entries = json.loads((state / "run/resident.json").read_text())["entries"] + assert len(entries) == 1 + assert next(iter(entries.values()))["pid"] == daemon["pid"] + finally: + for client in clients: + if client.poll() is None: + client.kill() + client.wait(timeout=3) + if daemon: + signal_owned_process(daemon["pid"], daemon["process_token"]) + wait_for(lambda: process_token(daemon["pid"]) is None) + + @pytest.mark.parametrize("name", ["../model", "/tmp/model", "a/b", "a..b", "", "x" * 97]) def test_invalid_names_are_rejected(name, tmp_path): assert main(["start", str(tmp_path), "--name", name, "--base-dir", str(tmp_path / "state")]) == 1 diff --git a/tests/test_management.py b/tests/test_management.py new file mode 100644 index 0000000..61da33c --- /dev/null +++ b/tests/test_management.py @@ -0,0 +1,61 @@ +import os +import sys +import time + +import pytest + +from conftest import make_model_dir +from infiniloader.cli import main +from infiniloader.diagnostics import diagnose +from infiniloader.ipc import request +from infiniloader.local_cache import materialize_model_cache +from infiniloader.runtime import process_token +from test_lifecycle import running_daemon + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux management tests") + + +def test_leases_protect_stop_gc_and_idle_ttl(tmp_path): + model = make_model_dir(tmp_path / "model") + with running_daemon(tmp_path, model, extra=["--idle-ttl", "2"]) as (proc, state): + sock = state / "run/unit.sock" + lease = request(sock, {"command": "lease-acquire", "pid": os.getpid(), "process_token": process_token(os.getpid())}) + time.sleep(2.1) + assert request(sock, {"command": "status"})["status"]["active_leases"] == 1 + assert main(["stop", "unit", "--base-dir", str(state)]) == 1 + assert main(["gc", "--resident", "--base-dir", str(state), "--max-bytes", "0", "--apply"]) == 1 + assert proc.poll() is None + request(sock, {"command": "lease-release", "lease": lease["lease"]}) + proc.wait(timeout=4) + assert not (state / "views/unit").exists() + + +def test_run_substitutes_model_path_and_releases_lease(tmp_path): + model = make_model_dir(tmp_path / "model") + out = tmp_path / "result" + program = [sys.executable, "-c", "import pathlib,sys; pathlib.Path(sys.argv[2]).write_bytes((pathlib.Path(sys.argv[1])/'model.safetensors').read_bytes())", "{model}", str(out)] + with running_daemon(tmp_path, model) as (_, state): + assert main(["run", "--base-dir", str(state), "unit", "--", *program]) == 0 + assert out.read_bytes() == b"0123456789abcdef" + assert request(state / "run/unit.sock", {"command": "status"})["status"]["active_leases"] == 0 + + +def test_run_disk_cache_preserves_command_exit_code(tmp_path): + local = materialize_model_cache(make_model_dir(tmp_path / "model"), base_dir=tmp_path / "state") + assert main(["run", "--cache-path", str(local.cache_path), "--", sys.executable, "-c", "raise SystemExit(7)"]) == 7 + + +def test_doctor_checks_actual_linux_capabilities(tmp_path): + result = diagnose(tmp_path) + assert result["ok"], result + assert result["memory"]["available_bytes"] is not None + assert "locked_memory" in result["limits"] + + +def test_resident_gc_removes_unleased_cache(tmp_path): + with running_daemon(tmp_path, make_model_dir(tmp_path / "model")) as (proc, state): + assert main(["gc", "--resident", "--base-dir", str(state), "--ttl", "0"]) == 0 + assert proc.poll() is None + assert main(["gc", "--resident", "--base-dir", str(state), "--ttl", "0", "--apply"]) == 0 + proc.wait(timeout=4) + assert not (state / "views/unit").exists() diff --git a/tests/test_review_runtime.py b/tests/test_review_runtime.py index 0d8c8e6..9164c7f 100644 --- a/tests/test_review_runtime.py +++ b/tests/test_review_runtime.py @@ -17,6 +17,23 @@ pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux review regressions") +def test_ipc_timeout_cannot_signal_a_ready_leased_daemon(tmp_path): + with running_daemon(tmp_path, make_model_dir(tmp_path / "model")) as (proc, state): + sock = state / "run/unit.sock" + request(sock, {"command": "lease-acquire", "pid": os.getpid(), "process_token": process_token(os.getpid())}) + clients = [] + try: + for _ in range(4): + client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + client.connect(str(sock)) + clients.append(client) + time.sleep(.05) + assert main(["stop", "unit", "--base-dir", str(state), "--timeout", ".05"]) == 1 + assert proc.poll() is None + finally: + for client in clients: + client.close() + assert request(sock, {"command": "status"})["status"]["active_leases"] == 1 @pytest.mark.parametrize("option", ["--socket", "--pid-file", "--state-file"]) @@ -38,6 +55,19 @@ def test_runtime_collisions_never_overwrite_ordinary_files(tmp_path, option, sym assert destination.is_symlink() +def test_stop_all_continues_after_a_leased_instance(tmp_path, monkeypatch): + first_root, second_root = tmp_path / "first", tmp_path / "second" + first_root.mkdir() + second_root.mkdir() + with running_daemon(first_root, make_model_dir(first_root / "model")) as (a, first): + with running_daemon(second_root, make_model_dir(second_root / "model")) as (b, second): + request(first / "run/unit.sock", {"command": "lease-acquire", "pid": os.getpid(), "process_token": process_token(os.getpid())}) + rows = [dict(name="a", state="ready"), dict(name="b", state="ready")] + monkeypatch.setattr("infiniloader.cli._active_daemon_rows", lambda _: rows) + monkeypatch.setattr("infiniloader.cli._socket", lambda args: (first if args.name == "a" else second) / "run/unit.sock") + assert main(["stop", "--all", "--base-dir", str(tmp_path)]) == 1 + b.wait(timeout=4) + assert a.poll() is None def test_existing_unrelated_unix_listener_is_preserved(tmp_path): From 7454e9cb73f500ec37dd06ad2df88849b1d2c12f Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 26 Oct 2026 20:57:57 +0800 Subject: [PATCH 06/29] Document cache operations and add Linux regression CI --- .github/workflows/tests.yml | 23 ++++++ README.md | 18 +++++ README.zh-CN.md | 17 ++++ deploy/Dockerfile | 11 +++ docs/expert-review.md | 71 +++++++++++++++++ docs/operations.md | 149 ++++++++++++++++++++++++++++++++++++ docs/validation.md | 49 ++++++++++++ 7 files changed, 338 insertions(+) create mode 100644 .github/workflows/tests.yml create mode 100644 deploy/Dockerfile create mode 100644 docs/expert-review.md create mode 100644 docs/operations.md create mode 100644 docs/validation.md diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..5a17a8a --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,23 @@ +name: Linux regression tests +on: [push, pull_request] +permissions: + contents: read +jobs: + tests: + runs-on: ubuntu-latest + timeout-minutes: 10 + strategy: + fail-fast: false + matrix: + python: ['3.10', '3.11', '3.12', '3.13', '3.14'] + steps: + - uses: actions/checkout@v7 + with: + persist-credentials: false + - uses: actions/setup-python@v7 + with: + python-version: ${{ matrix.python }} + - run: python -m pip install -e '.[test]' + - run: python -m pytest -q + - run: iloader selftest + - run: iloader doctor --json diff --git a/README.md b/README.md index c8ad4f9..909ea48 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,24 @@ [中文](README.zh-CN.md) +The development branch adds model identity checks, single-owner startup, sealed +memfd files, versioned disk caches, container-aware budgets, and lease-aware cache +management. See [operations and migration](docs/operations.md) for `doctor`, +`verify`, `run`, `gc`, and container guidance. + +```bash +iloader doctor +iloader start "$MODEL" --name dev --local-cache \ + --max-memory-bytes 32GiB --reserve-bytes 4GiB --max-resident-bytes 64GiB +iloader run dev -- vllm serve '{model}' +iloader gc --resident --max-bytes 32GiB # preview +iloader gc --resident --max-bytes 32GiB --apply +``` + +Use `iloader cache` plus `iloader run --cache-path PATH -- COMMAND` for disk-only +caching. Existing direct-path readers remain supported; eviction protection requires +`run`, so leave automatic idle eviction disabled for unregistered readers. + `InfiniLoader` is a userspace weight loading manager for LLM inference framework development. It is designed for the loop where you repeatedly restart vLLM, SGLang, InfiniLM, or a custom engine while debugging loader, kernel, diff --git a/README.zh-CN.md b/README.zh-CN.md index b24eebb..22aa82b 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -2,6 +2,23 @@ [English README](README.md) +开发分支新增了模型身份校验、daemon 启动互斥、memfd 写封印、不可变磁盘缓存版本、 +容器内存预算及租约保护。新增 `doctor`、`verify`、`run`、`gc` 命令，详见 +[运行管理与迁移说明](docs/operations.md)。 + +```bash +iloader doctor +iloader start "$MODEL" --name dev --local-cache \ + --max-memory-bytes 32GiB --reserve-bytes 4GiB --max-resident-bytes 64GiB +iloader run dev -- vllm serve '{model}' +iloader gc --resident --max-bytes 32GiB # 只预览 +iloader gc --resident --max-bytes 32GiB --apply +``` + +内存紧张时，可只用 `iloader cache` 构建本地磁盘缓存，再用 +`iloader run --cache-path PATH -- COMMAND` 启动框架。原有直接传 view 路径的方式仍可使用； +自动清理无法识别未注册的读者，启用闲置过期或执行清理前应通过 `run` 保护任务。 + `InfiniLoader` 是一个面向大模型推理框架开发调试的用户态权重加载管理器。它解决的核心问题是：调试 vLLM、SGLang、InfiniLM 或自研推理框架时，经常需要反复停止、修改代码、重新启动框架，而每次启动都要重新从存储读取模型权重。 默认模式下，InfiniLoader 会把模型目录常驻在内存里，并暴露出一个看起来像普通 Hugging Face 模型目录的路径。框架仍然使用原来的文件读取逻辑，通常只需要把模型路径换成 `iloader start/view` 输出的路径。 diff --git a/deploy/Dockerfile b/deploy/Dockerfile new file mode 100644 index 0000000..a9efb57 --- /dev/null +++ b/deploy/Dockerfile @@ -0,0 +1,11 @@ +FROM python:3.12-slim +WORKDIR /opt/infiniloader +COPY pyproject.toml README.md LICENSE ./ +COPY src ./src +RUN pip install --no-cache-dir . \ + && mkdir -p /var/cache/infiniloader \ + && chown 10001:10001 /var/cache/infiniloader +ENV ILOADER_BASE_DIR=/var/cache/infiniloader +USER 10001:10001 +ENTRYPOINT ["iloader"] +CMD ["doctor", "--json"] diff --git a/docs/expert-review.md b/docs/expert-review.md new file mode 100644 index 0000000..a163eb9 --- /dev/null +++ b/docs/expert-review.md @@ -0,0 +1,71 @@ +# Independent expert review + +> Historical evidence for the original 2026-10-18 commits. Commit IDs and +> branch names below identify the archived source history, not this reorganized +> branch. See `history-reorganization.md` in the final checkpoint for the new +> commit mapping and independently repeated validation. + +Date: 2026-10-18 (Asia/Shanghai). Base: upstream master `5014a6a`. +Branch: `codex/reliability-usability`. No merge into master was performed. + +Three independent reviewer agents inspected the complete original nine-commit range, +with complementary specializations. They used read-only source review and isolated +Linux probes. Implementation and fixes were performed by the primary agent, then +returned to the reviewers for independent verification. + +## Findings and resolution + +| Priority | Finding | Resolution and regression evidence | +| --- | --- | --- | +| P1 | IPC timeout could signal a ready daemon despite active leases | Signal fallback now requires genuine pre-ready state or explicit force. Real saturated-socket regression preserves the leased daemon. | +| P1 | Custom runtime paths could replace ordinary files, symlinks or other Unix services | Preflight validates resource types, marked stale state, owner and endpoint identity. Existing listeners are refused. Atomic metadata writes and inode claims prevent cleanup of replaced resources. | +| P2 | Verification accepted extra unverified model files | Full inventory must match the manifest, including duplicate/source-name checks. Extra safetensors cause invalidation and a new generation. | +| P2 | Malformed manifest records prevented verification and refresh | Schema failures return invalid status. Timestamp overflow is normalized; totals must be non-boolean integers matching recorded sizes. GC measures actual file bytes. | +| P2 | Copy buffers could exceed the advertised reservation | Byte buffers and memoryviews are released before the next read. Prefix hashing is chunked. Tracemalloc regressions check both memfd and disk copy paths. | +| P2 | Bulk stop aborted after the first leased cache | Per-instance errors are reported while subsequent instances are processed. A two-daemon regression verifies continuation. | +| P2 | Distributed failure text omitted synthetic rank paths and raised KeyError | Every cancellation row includes its rank path. Text-mode failure diagnostics are covered. | + +Follow-up fixes are deliberately separate commits: + +- `7e1b248`: runtime resource ownership, leased stop fallback, bulk stop and rank diagnostics. +- `cc67790`: generation integrity, malformed manifests, GC accounting and bounded copy buffers. + +## Per-commit independence + +The CLI/validation expert exported each original commit to a fresh temporary archive +and ran its complete test suite, independently of the working checkout. All passed +on Linux / Python 3.14.4: + +| Commit | Passing tests | +| --- | ---: | +| `2f7204b` | 27 | +| `c67ba9e` | 31 | +| `deaa940` | 34 | +| `d726100` | 39 | +| `2f54480` | 46 | +| `fb2f057` | 51 | +| `8fab6cd` | 55 | +| `a17a6e9` | 55 | +| `6a33d83` | 55 | +| `7e1b248` | 66 | +| `cc67790` | 77 | + +## Final independent verdicts + +- Lifecycle expert approved `origin/master..cc67790` within daemon lifecycle, IPC, + path/resource ownership and distributed-cancellation scope. The updated targeted + set passed 44 tests. Additional probes verified ordinary-file/listener preservation, + lease timeout protection, and recovery after SIGKILL of an owned daemon. +- Cache/memory expert approved `origin/master..cc67790` within cache integrity, + generation publication, leases/GC and memory accounting scope. The final independent + targeted set passed 28 relevant tests; all original and follow-up findings resolved. +- CLI/validation expert confirmed the three CLI findings resolved, with 28 targeted + tests passing, and independently verified each commit's full suite as above. + Final approval is tied to `cc67790f7b4d65b20cbeef051f5ee9fc6189088c`; its isolated + complete suite passed 77 tests in 19.88 seconds (preceding runtime fix: 66 tests). + +The implementation is considered merge-ready within this branch's agreed reliability +and local-usability scope once its final remote CI is green. This is a code-review +judgement, not an approval to merge. Native shard conversion and multi-node features +remain deferred. Real inference-engine/GPU compatibility, cross-container sharing +and performance are outside these CPU/file-service validations. diff --git a/docs/operations.md b/docs/operations.md new file mode 100644 index 0000000..f9f36a8 --- /dev/null +++ b/docs/operations.md @@ -0,0 +1,149 @@ +# Operating the development version + +This branch improves reliability and local usability. TP checkpoint conversion and +multi-node transport remain future work. Runtime dependencies remain limited to +Python's standard library. Linux and Python 3.10+ are required. + +## Choose a backend + +```bash +MODEL=/models/checkpoint +iloader doctor +VIEW=$(iloader start "$MODEL" --name dev --local-cache) +iloader run dev -- vllm serve '{model}' --tensor-parallel-size 4 +``` + +`run` replaces literal `{model}` placeholders in arguments and sets +`ILOADER_MODEL_PATH`. It executes arguments without a shell and returns the command's +exit code. Put InfiniLoader options before the cache name. Use a foreground engine +command that remains alive while its workers run. + +For disk-only caching without another model-sized RAM allocation: + +```bash +LOCAL=$(iloader cache "$MODEL" --name dev) +iloader verify "$LOCAL" +iloader run --cache-path "$LOCAL" -- vllm serve '{model}' +``` + +The OS controls disk cache residency. Memfd files are sealed against writes, growth +and truncation, but their pages are not mlocked and can be swapped. Status reports +logical file capacity, not measured physical residency. Parsing, CPU tensor allocations, +weight transformation and CPU-to-device transfers remain the engine's responsibility. + +## Identity, verification and migration + +A resident name identifies one source path, file selection and source fingerprint. +The fingerprint includes sizes, nanosecond timestamps, device and inode IDs. Reusing +a name for another revision fails; stop the cache or choose another name. Changes +while copying also fail. Use immutable checkpoints: metadata checks cannot make an +external writer's multi-file update transactional. + +Disk generations have version-2 manifests with per-file SHA-256 digests. Normal reuse +checks source and destination metadata; `verify PATH` reads all files and checks hashes. +`verify PATH --quick` only checks metadata. The digests describe copied bytes; they do +not authenticate a publisher or validate model semantics. + +```bash +NEW_LOCAL=$(iloader cache "$MODEL" --name dev --refresh) +iloader verify "$NEW_LOCAL" --json +``` + +Refresh publishes a new generation, preserves old paths, and needs additional disk +space. SHA-256 is computed during the copy and its CPU cost belongs in preparation +benchmarks. Version-1 caches remain listed but are not reused; rebuild them, verify +the new path, and use `gc` to remove unleased old versions. + +## Memory budgets + +```bash +iloader start "$MODEL" --name dev --workers 4 \ + --max-memory-bytes 32GiB --reserve-bytes 4GiB --max-resident-bytes 64GiB +``` + +- Effective available memory is the minimum of host MemAvailable and remaining + cgroup v1/v2 limits, including visible ancestors. +- The default budget is 90% of effective availability minus a 256 MiB reserve. + Increase the reserve for engine CPU weights and other allocations. +- `--max-memory-bytes` caps one preload including estimated worker buffers. +- `--max-resident-bytes` sets a persistent quota for all daemons in the same runtime + directory. A subsequent start can explicitly update the cap. In-flight loads reserve + model and buffer bytes; ready caches reserve model bytes. Dead reservations are + pruned using boot ID and PID start time. +- `--force` bypasses the available-memory check, while explicit per-cache and shared + caps remain enforced. Quotas account for InfiniLoader, not engine allocations. +- Byte options accept integers and binary suffixes such as `256MiB` and `32GiB`. + +Use a node-local `--base-dir`. A runtime registry shared between nodes is unsupported. + +## Leases, LRU and TTL + +`run` protects a cache against ordinary stop and eviction. Local-cache daemons also +hold a disk-generation lease. Foreground disk commands inherit its OS lock descriptor; +resident commands register the child process's PID and start time. + +```bash +iloader gc --max-bytes 200GiB # disk LRU preview +iloader gc --ttl 604800 # disk generations idle for 7 days +iloader gc --max-bytes 200GiB --apply +iloader gc --resident --max-bytes 32GiB # resident LRU preview +iloader gc --resident --ttl 3600 --apply +``` + +Without `--apply`, collection only previews. An unmet capacity target returns nonzero, +for example if leased caches exceed the target. Resident collection requests shutdown; +reported remaining capacity is projected until the daemons exit. Unmarked directories +and staging directories left by a killed writer are never automatically deleted. + +```bash +iloader start "$MODEL" --name dev --idle-ttl 3600 +iloader run dev -- vllm serve '{model}' +iloader stop dev --remove-view +``` + +Idle eviction is opt-in. Direct-path readers are not registered: keep idle eviction +disabled and avoid collecting their caches, or launch them through `run`. `stop --force` +overrides active resident leases. External signals also stop a leased daemon. Stale +PID metadata alone is never enough to signal a process. + +## Diagnostics and containers + +```bash +iloader doctor --json +iloader status dev --json +iloader list --json +iloader dist-simulate "$VIEW" --ranks 4 --timeout 60 +``` + +Doctor exercises locks, Unix sockets, memfd sealing and reopening in another process. +It reports cgroup headroom, disk space, FD and memlock limits. Status includes identity, +file bytes, the startup budget snapshot, active leases, reuse requests and last use. +Distributed simulation is a bounded single-host read test, not a GPU benchmark. + +Prefer installing InfiniLoader into the engine container and using the same UID. +Separate containers require the reader to see the daemon PID in its proc namespace +and pass proc and user-namespace access checks. A shared volume alone is insufficient. +Test an actual engine open in the deployment environment. Resolving individual view +file symlinks to their final target can break reopening; use the published file paths. +After a daemon exits, existing open descriptors/mappings may survive, but new opens +through its view fail. Automatic failover is not provided. + +The non-root diagnostic image contains InfiniLoader, without an engine or model: + +```bash +docker build -f deploy/Dockerfile -t infiniloader:dev . +docker run --rm --memory=2g infiniloader:dev +``` + +## Validation + +Run `python -m pytest`, `iloader selftest`, and `iloader doctor`. CI tests Python +3.10-3.14 on Ubuntu. Tests exercise real daemon subprocesses, source mutations, +disconnected clients, immutable generations, quotas, leases, and failed ranks. + +`bench` reports each run, first-run and median time, and best throughput. It does not +clear OS page cache. Compare cold/warm shared storage, cold/warm local disk, memfd and +engine-native loaders separately. Include preparation, end-to-end engine readiness, +peak CPU memory, per-rank bytes and CPU-to-device transfers. Historical README numbers +are not new measurements for this branch; CPU tests do not establish GPU or multi-node +performance. diff --git a/docs/validation.md b/docs/validation.md new file mode 100644 index 0000000..55fa79f --- /dev/null +++ b/docs/validation.md @@ -0,0 +1,49 @@ +# Reliability and usability validation + +> Historical evidence for the original 2026-10-18 commits. Commit IDs and +> branch names below identify the archived source history, not this reorganized +> branch. See `history-reorganization.md` in the final checkpoint for the new +> commit mapping and independently repeated validation. + +Date: 2026-10-18 (Asia/Shanghai). Baseline: `5014a6a` on upstream master. + +## Local results + +- Ubuntu on WSL2, Linux 6.18.33.2, Python 3.14.4: **77 tests passed** after independent + review fixes (25.92 seconds). The baseline contained 17 tests; the original nine + implementation commits ended with 55 passing tests. +- Built `infiniloader-0.1.0-py3-none-any.whl` with Python 3.12.14/setuptools 84.0.0. + The package version is unchanged; this branch has not been published as a release. +- Executed `python -m infiniloader selftest` from the built wheel on Linux: passed + the disk-cache reuse, sealed memfd view, child-process open and mmap checks. +- Executed `doctor` from the built wheel: Linux, advisory locks, Unix sockets, + memfd seals and cross-process proc reopening all passed. +- `git diff --check`: passed. + +## Regression coverage + +Real child processes cover concurrent CLI starts, joining a daemon during preload, +model/revision mismatch rejection, lease-protected commands, idle expiry, safe +resident eviction, and stalled/disconnected IPC clients. Other tests check sealed +weights, FD cleanup, mutation during preload, unsafe path rejection, disk corruption, +SHA-256 verification, immutable refresh, LRU/TTL and leases, cgroup v1/v2 ancestor +limits, pending resident reservations, and failed-rank cancellation and reaping. + +The idle-expiry integration test uses a two-second TTL so process startup scheduling +does not consume an unrealistically short test deadline. The test waits beyond the +TTL while a lease is active, then verifies eviction after lease release. + +## Scope and remaining validation + +The original branch HEAD `6a33d83` passed the complete remote Ubuntu Python 3.10-3.14 +matrix in [Actions run 39101870656](https://github.com/InfiniTensor/InfiniLoader/actions/runs/39101870656). +Check the newer branch workflow for the subsequent review fixes. Independent review +and per-commit test details are recorded in [expert-review.md](expert-review.md). +The Dockerfile is a non-root diagnostic image; Docker was unavailable locally, so +the image was not built or exercised in this environment. + +These are CPU/file-service tests using synthetic model files. They do not establish +vLLM/SGLang/InfiniLM version compatibility, GPU transfer throughput, real TP/PP loading, +cross-container namespace sharing, or multi-node performance. Existing README timing +numbers are historical. Native shard conversion and multi-node coordination are +deferred to subsequent work as agreed for this branch. From 48615de1981afb1ff17475423227813932112e43 Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 26 Oct 2026 20:58:54 +0800 Subject: [PATCH 07/29] Add explicit tensor sharding with atomic streaming export --- src/infiniloader/cli.py | 2 + src/infiniloader/shards.py | 347 +++++++++++++++++++++++++++++++++++++ tests/test_shards.py | 275 +++++++++++++++++++++++++++++ 3 files changed, 624 insertions(+) create mode 100644 src/infiniloader/shards.py create mode 100644 tests/test_shards.py diff --git a/src/infiniloader/cli.py b/src/infiniloader/cli.py index c61bb3b..ed2bc6a 100644 --- a/src/infiniloader/cli.py +++ b/src/infiniloader/cli.py @@ -979,6 +979,8 @@ def build_parser() -> argparse.ArgumentParser: inspect_p.add_argument("--exclude", action="append", help="exclude glob; replaces defaults when provided") inspect_p.add_argument("--json", action="store_true") inspect_p.set_defaults(func=cmd_inspect) + from .shards import configure_parser as configure_shards + configure_shards(sub) serve_p = sub.add_parser("_serve", help=argparse.SUPPRESS) serve_p.add_argument("model") diff --git a/src/infiniloader/shards.py b/src/infiniloader/shards.py new file mode 100644 index 0000000..8a8269c --- /dev/null +++ b/src/infiniloader/shards.py @@ -0,0 +1,347 @@ +"""Explicit, portable tensor layouts. No model-name based partition inference. + +The output is an InfiniLoader checkpoint, not a drop-in sharded HF checkpoint. +An engine must implement the supplied TP/PP contract (see the torch example). +""" +from __future__ import annotations + +import hashlib +import ctypes +import errno +import json +import math +import os +import shutil +import stat +import struct +import tempfile +from pathlib import Path + +from .distributed import DTYPE_BYTES + +CHUNK = 1024 * 1024 +MAX_HEADER = 100_000_000 + + +def _unique(pairs): + result = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate JSON key: {key}") + result[key] = value + return result + + +def read_json(path: Path): + return json.loads(path.read_text(encoding="utf-8"), object_pairs_hook=_unique) + + +def canonical(value) -> bytes: + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True, + allow_nan=False).encode() + + +def digest(value) -> str: + return hashlib.sha256(canonical(value)).hexdigest() + + +def file_hash(path: Path) -> str: + h = hashlib.sha256() + with path.open("rb") as stream: + while block := stream.read(CHUNK): + h.update(block) + return h.hexdigest() + + +def _signature(path: Path): + s = path.stat() + return s.st_dev, s.st_ino, s.st_size, s.st_mtime_ns, s.st_ctime_ns + + +def _int(value, *, minimum=0): + return type(value) is int and value >= minimum + + +def tensor_header(path: Path) -> tuple[dict, int]: + """Reject holes, overlaps, truncated payloads and unsupported packed dtypes.""" + with path.open("rb") as stream: + size = os.fstat(stream.fileno()).st_size + raw = stream.read(8) + if len(raw) != 8: + raise ValueError(f"truncated safetensors: {path}") + length, = struct.unpack(" dict: + """Content identity is independent of hostname, inode and source directory.""" + source = source.expanduser().resolve() + files, tensors = {}, {} + paths = sorted(source.rglob("*.safetensors")) + signatures = {p: _signature(p) for p in paths} + for path in paths: + before = signatures[path] + header, base = tensor_header(path) + rel = path.relative_to(source).as_posix() + files[rel] = {"sha256": file_hash(path), "size": before[2], "data_start": base} + if before != _signature(path): + raise RuntimeError(f"source changed while cataloging: {path}") + for name, spec in header.items(): + if name in tensors: + raise ValueError(f"duplicate tensor across files: {name}") + tensors[name] = dict(spec, file=rel) + if not files or not tensors: + raise ValueError("no safetensors tensors found") + index_path = source / "model.safetensors.index.json" + if index_path.exists(): + index = read_json(index_path) + if not isinstance(index, dict) or index.get("weight_map") != {n: t["file"] for n, t in tensors.items()}: + raise ValueError("safetensors index does not match tensor inventory") + if paths != sorted(source.rglob("*.safetensors")) or any(_signature(p) != s for p, s in signatures.items()): + raise RuntimeError("source inventory changed while cataloging") + result = {"format": 1, "files": files, "tensors": tensors} + result["content_id"] = digest(result) + result["tensor_bytes"] = sum(t["data_offsets"][1] - t["data_offsets"][0] for t in tensors.values()) + return result + + +def build_plan(source: Path, layout: dict, *, mode="sharded", capacities: list[int] | None = None) -> dict: + cat = catalog(source) + if not isinstance(layout, dict) or layout.get("format") != 1: + raise ValueError("layout requires format=1") + tp, pp = layout.get("tp"), layout.get("pp") + rules = layout.get("tensors") + if not _int(tp, minimum=1) or not _int(pp, minimum=1) or not isinstance(rules, dict): + raise ValueError("layout requires positive tp, pp and explicit tensor rules") + if mode not in ("sharded", "replicated") or set(rules) != set(cat["tensors"]): + raise ValueError("invalid mode or layout must explicitly cover exactly all tensors") + for name, spec in cat["tensors"].items(): + rule = rules[name] + if not isinstance(rule, dict) or set(rule) != {"stage", "split_dim"}: + raise ValueError(f"rule requires stage and split_dim: {name}") + stage, dim = rule["stage"], rule["split_dim"] + if not _int(stage) or stage >= pp: + raise ValueError(f"invalid pipeline stage: {name}") + if dim is not None and (not _int(dim) or dim >= len(spec["shape"]) or spec["shape"][dim] % tp): + raise ValueError(f"invalid or nondivisible split dimension: {name}") + ranks = [] + for rank in range(tp * pp): + stage, lane = divmod(rank, tp) + parts = {} + for name, spec in cat["tensors"].items(): + rule = rules[name] + if mode == "sharded" and rule["stage"] != stage: + continue + dim = rule["split_dim"] if mode == "sharded" else None + shape = list(spec["shape"]) + start, stop = 0, 0 + if dim is not None: + width = shape[dim] // tp + start, stop, shape[dim] = lane * width, (lane + 1) * width, width + parts[name] = {"dim": dim, "start": start, "stop": stop, "shape": shape, "dtype": spec["dtype"], + "bytes": math.prod(shape) * DTYPE_BYTES[spec["dtype"]]} + ranks.append({"rank": rank, "stage": stage, "lane": lane, "tensors": parts, + "tensor_bytes": sum(p["bytes"] for p in parts.values())}) + if any(not r["tensors"] for r in ranks): + raise ValueError("every pipeline stage must own at least one tensor") + if capacities is not None: + if len(capacities) != len(ranks) or not all(_int(v) for v in capacities): + raise ValueError("capacities requires one nonnegative byte limit per rank") + if any(r["tensor_bytes"] > c for r, c in zip(ranks, capacities)): + raise MemoryError("rank tensor payload exceeds node capacity (headers/buffers are additional)") + plan = {"format": 1, "content_id": cat["content_id"], "mode": mode, "layout": layout, + "unique_tensor_bytes": cat["tensor_bytes"], "ranks": ranks, + "aggregate_tensor_bytes": sum(r["tensor_bytes"] for r in ranks)} + plan["plan_id"] = digest(plan) + return plan + + +def iter_part(source: Path, cat: dict, name: str, part: dict): + if not part["bytes"]: + return + spec = cat["tensors"][name] + begin = cat["files"][spec["file"]]["data_start"] + spec["data_offsets"][0] + dim = part["dim"] + if dim is None: + rows, stride, offset, width = 1, 0, 0, part["bytes"] + else: + shape = spec["shape"] + inner = math.prod(shape[dim + 1:]) * DTYPE_BYTES[spec["dtype"]] + rows, stride = math.prod(shape[:dim]), shape[dim] * inner + offset, width = part["start"] * inner, (part["stop"] - part["start"]) * inner + with (source / spec["file"]).open("rb") as stream: + for row in range(rows): + stream.seek(begin + row * stride + offset) + left = width + while left: + block = stream.read(min(left, CHUNK)) + if not block: + raise IOError("source truncated during shard export") + left -= len(block) + yield block + + +def _publish_directory(temp: Path, destination: Path): + if os.name == "nt": + os.rename(temp, destination) # Windows rename never replaces a destination. + return + # POSIX rename would replace a concurrently created empty directory. Some NFS + # servers lack RENAME_NOREPLACE: publish a relative symlink to the complete, + # private staging directory instead. symlink creation is atomic and exclusive. + libc = ctypes.CDLL(None, use_errno=True) + rename = getattr(libc, "renameat2", None) + if rename is None: + os.symlink(temp.name, destination, target_is_directory=True) + return + rename.argtypes = [ctypes.c_int, ctypes.c_char_p, ctypes.c_int, ctypes.c_char_p, ctypes.c_uint] + rename.restype = ctypes.c_int + if rename(-100, os.fsencode(temp), -100, os.fsencode(destination), 1): + code = ctypes.get_errno() + if code in (errno.ENOSYS, errno.EINVAL, errno.EOPNOTSUPP): + os.symlink(temp.name, destination, target_is_directory=True) + return + raise OSError(code, os.strerror(code), destination) + + +def export_rank(source: Path, plan: dict, rank: int, destination: Path) -> dict: + """Stream slices with bounded buffers; publish a new directory atomically.""" + source, destination = source.expanduser().resolve(), destination.expanduser().absolute() + if destination.is_symlink() or destination.exists(): + raise FileExistsError(destination) + resolved = destination.resolve() + if resolved.is_relative_to(source) or source.is_relative_to(resolved): + raise ValueError("export and source directories must not overlap") + expected = build_plan(source, plan["layout"], mode=plan["mode"]) + if plan != expected: + raise ValueError("stale or modified shard plan") + if not _int(rank) or rank >= len(plan["ranks"]): + raise ValueError("invalid rank") + cat = catalog(source) + if cat["content_id"] != plan["content_id"]: + raise ValueError("source changed since plan validation") + signatures = {p: _signature(source / p) for p in cat["files"]} + selected = plan["ranks"][rank] + header, offset = {}, 0 + for name, part in selected["tensors"].items(): + header[name] = {"dtype": cat["tensors"][name]["dtype"], "shape": part["shape"], + "data_offsets": [offset, offset + part["bytes"]]} + offset += part["bytes"] + raw = canonical(header) + raw += b" " * (-len(raw) % 8) + if len(raw) > MAX_HEADER: + raise ValueError("exported safetensors header exceeds format limit") + destination.parent.mkdir(parents=True, exist_ok=True) + temp = Path(tempfile.mkdtemp(prefix=".iloader-export-", dir=destination.parent)) + try: + weights = temp / "weights.safetensors" + with weights.open("wb") as stream: + stream.write(struct.pack(" dict: + receipt = read_json(path / "shard.json") + if (not isinstance(receipt, dict) or receipt.get("format") != 1 + or receipt.get("kind") != "infiniloader-explicit-shard"): + raise ValueError("invalid InfiniLoader shard receipt") + plan = receipt["plan"] + if plan.get("plan_id") != digest({k: v for k, v in plan.items() if k != "plan_id"}): + raise ValueError("invalid plan identity") + rank = receipt["rank"] + if not _int(rank) or rank >= len(plan["ranks"]): + raise ValueError("invalid shard rank") + if file_hash(path / "weights.safetensors") != receipt["weights_sha256"]: + raise ValueError("shard checksum mismatch") + header, _ = tensor_header(path / "weights.safetensors") + parts = plan["ranks"][rank]["tensors"] + if set(header) != set(parts) or any(t["shape"] != parts[n]["shape"] or t["dtype"] != parts[n]["dtype"] + for n, t in header.items()): + raise ValueError("shard tensor inventory/layout mismatch") + return receipt + + +def configure_parser(sub): + parser = sub.add_parser("shard", help="explicit tensor layout planning, export and verification") + actions = parser.add_subparsers(dest="shard_action", required=True) + plan = actions.add_parser("plan") + plan.add_argument("source", type=Path) + plan.add_argument("--layout", type=Path, required=True) + plan.add_argument("--mode", choices=("sharded", "replicated"), default="sharded") + plan.add_argument("--output", type=Path, required=True) + export = actions.add_parser("export") + export.add_argument("source", type=Path) + export.add_argument("--plan", type=Path, required=True) + export.add_argument("--rank", type=int, required=True) + export.add_argument("--output", type=Path, required=True) + verify = actions.add_parser("verify") + verify.add_argument("path", type=Path) + + def run(args): + if args.shard_action == "plan": + result = build_plan(args.source, read_json(args.layout), mode=args.mode) + # Never silently replace an existing deployment plan. + with args.output.open("xb") as stream: + stream.write(canonical(result)) + elif args.shard_action == "export": + result = export_rank(args.source, read_json(args.plan), args.rank, args.output) + else: + result = verify_rank(args.path) + print(json.dumps(result, indent=2, sort_keys=True)) + return 0 + + for command in (plan, export, verify): + command.set_defaults(func=run) diff --git a/tests/test_shards.py b/tests/test_shards.py new file mode 100644 index 0000000..aa040f9 --- /dev/null +++ b/tests/test_shards.py @@ -0,0 +1,275 @@ +from __future__ import annotations + +import json +import shutil +import struct + +import pytest + +from infiniloader.shards import (build_plan, catalog, canonical, export_rank, + tensor_header, verify_rank) + + +def write_tensor(path, shape=(4, 6, 2), name="weight"): + path.mkdir() + values = list(range(__import__("math").prod(shape))) + payload = struct.pack("<" + "f" * len(values), *values) + header = canonical({name: {"dtype": "F32", "shape": list(shape), "data_offsets": [0, len(payload)]}}) + (path / "weights.safetensors").write_bytes(struct.pack(" plan["aggregate_tensor_bytes"] + + +@pytest.mark.parametrize("mutate", [ + lambda x: x["tensors"].clear(), + lambda x: x["tensors"]["weight"].update(split_dim=3), + lambda x: x.update(tp=5), + lambda x: x.update(tp=True), + lambda x: x["tensors"]["weight"].update(stage=1), +]) +def test_reject_ambiguous_or_invalid_layout(tmp_path, mutate): + source = tmp_path / "source" + write_tensor(source) + rules = layout(0) + mutate(rules) + with pytest.raises(ValueError): + build_plan(source, rules) + + +def test_no_overwrite_stale_plan_and_corruption(tmp_path): + source = tmp_path / "source" + write_tensor(source) + plan = build_plan(source, layout(0)) + dest = tmp_path / "out" + export_rank(source, plan, 0, dest) + with pytest.raises(FileExistsError): + export_rank(source, plan, 0, dest) + weights = dest / "weights.safetensors" + weights.write_bytes(weights.read_bytes()[:-1] + b"x") + with pytest.raises(ValueError, match="checksum"): + verify_rank(dest) + weights = source / "weights.safetensors" + weights.write_bytes(weights.read_bytes()[:-1] + b"x") + with pytest.raises(ValueError, match="stale"): + export_rank(source, plan, 0, tmp_path / "other") + + +@pytest.mark.parametrize("header,payload", [ + (b'{"a":{},"a":{}}', b""), + (canonical({"a": {"dtype": "F32", "shape": [True], "data_offsets": [0, 4]}}), b"1234"), + (canonical({"a": {"dtype": "F32", "shape": [1], "data_offsets": [1, 5]}}), b"12345"), + (canonical({"a": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}}), b"12345"), + (canonical({"__metadata__": {"x": 42}}), b""), + (canonical({"a": {"dtype": "F4", "shape": [2], "data_offsets": [0, 1]}}), b"x"), +]) +def test_reject_invalid_safetensors(tmp_path, header, payload): + path = tmp_path / "bad.safetensors" + path.write_bytes(struct.pack(" Date: Mon, 26 Oct 2026 21:00:10 +0800 Subject: [PATCH 08/29] Share sealed CPU tensors with live topology slicing --- pyproject.toml | 1 + src/infiniloader/torch_store.py | 119 ++++++++++++++++++++++++++++++++ tests/test_torch_store.py | 115 ++++++++++++++++++++++++++++++ 3 files changed, 235 insertions(+) create mode 100644 src/infiniloader/torch_store.py create mode 100644 tests/test_torch_store.py diff --git a/pyproject.toml b/pyproject.toml index 680cb3a..6813f8d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,7 @@ dependencies = [] [project.optional-dependencies] test = ["pytest>=8"] +torch = ["torch>=2.2"] [project.scripts] iloader = "infiniloader.cli:main" diff --git a/src/infiniloader/torch_store.py b/src/infiniloader/torch_store.py new file mode 100644 index 0000000..bfc584e --- /dev/null +++ b/src/infiniloader/torch_store.py @@ -0,0 +1,119 @@ +"""Optional PyTorch access to canonical CPU weights and explicit TP/PP slices. + +Private COW mappings protect the shared checkpoint against in-place tensor writes. +No pinning, dtype conversion, GPU transfer or framework weight packing is implicit. +""" +from __future__ import annotations + +import math +import mmap +import hashlib +from pathlib import Path + +from .shards import catalog, verify_rank, read_json + + +class TensorStore: + def __init__(self, path: str | Path): + import torch + self.torch = torch + self.path = Path(path).expanduser().resolve() + marker = self.path / "shard.json" + has_receipt = marker.exists() + self.receipt = verify_rank(self.path) if has_receipt else None + self.catalog = catalog(self.path) + if self.receipt and (set(self.catalog["files"]) != {"weights.safetensors"} + or self.catalog["files"]["weights.safetensors"]["sha256"] != self.receipt["weights_sha256"]): + raise ValueError("mapped catalog does not match shard receipt") + self.content_id = self.receipt["plan"]["content_id"] if self.receipt else self.catalog["content_id"] + self._maps = {} + self._closed = False + import fcntl + seals = fcntl.F_SEAL_WRITE | fcntl.F_SEAL_GROW | fcntl.F_SEAL_SHRINK + try: + for file, record in self.catalog["files"].items(): + with (self.path / file).open("rb") as stream: + try: + sealed = fcntl.fcntl(stream.fileno(), fcntl.F_GET_SEALS) & seals == seals + except OSError: + sealed = False + if not sealed: + raise ValueError("TensorStore requires sealed InfiniLoader weights; preload and use the resident view") + mapping = mmap.mmap(stream.fileno(), 0, access=mmap.ACCESS_COPY) + self._maps[file] = mapping + if hashlib.sha256(mapping).hexdigest() != record["sha256"]: + raise ValueError("weights changed between catalog and mapping") + if marker.exists() != has_receipt or has_receipt and read_json(marker) != self.receipt: + raise ValueError("shard receipt changed during store initialization") + except BaseException: + self.close() + raise + + def tensor(self, name: str): + if self._closed: + raise RuntimeError("tensor store is closed") + spec = self.catalog["tensors"][name] + types = {"F64": "float64", "F32": "float32", "F16": "float16", "BF16": "bfloat16", + "I64": "int64", "I32": "int32", "I16": "int16", "I8": "int8", "U8": "uint8", + "BOOL": "bool", "U16": "uint16", "U32": "uint32", "U64": "uint64", + "F8_E4M3": "float8_e4m3fn", "F8_E5M2": "float8_e5m2"} + dtype = getattr(self.torch, types[spec["dtype"]], None) + if dtype is None: + raise ValueError(f"installed torch does not support {spec['dtype']}") + count = math.prod(spec["shape"]) + if not count: + return self.torch.empty(spec["shape"], dtype=dtype) + file = spec["file"] + offset = self.catalog["files"][file]["data_start"] + spec["data_offsets"][0] + return self.torch.frombuffer(self._maps[file], dtype=dtype, count=count, offset=offset).reshape(spec["shape"]) + + def rank_tensors(self, layout: dict, rank: int) -> dict: + """Canonical replicas can change layout without rereading/rebuilding weights. + + Pre-sharded stores are bound to their exported layout and global rank. + The caller implements the actual model's collective/bias semantics. + """ + if self._closed: + raise RuntimeError("tensor store is closed") + tp, pp, rules = layout.get("tp"), layout.get("pp"), layout.get("tensors") + if (layout.get("format") != 1 or type(tp) is not int or tp < 1 + or type(pp) is not int or pp < 1 or type(rank) is not int or not 0 <= rank < tp * pp + or not isinstance(rules, dict)): + raise ValueError("invalid tensor layout/rank") + if self.receipt and self.receipt["plan"]["mode"] == "sharded": + if layout != self.receipt["plan"]["layout"] or rank != self.receipt["rank"]: + raise ValueError("sharded store requires its exported layout and rank") + return {n: self.tensor(n) for n in self.catalog["tensors"]} + if set(rules) != set(self.catalog["tensors"]): + raise ValueError("layout must cover exactly all canonical tensors") + result = {} + for name, spec in self.catalog["tensors"].items(): + rule = rules[name] + if not isinstance(rule, dict) or set(rule) != {"stage", "split_dim"}: + raise ValueError(f"invalid tensor rule: {name}") + stage, dim = rule["stage"], rule["split_dim"] + if type(stage) is not int or not 0 <= stage < pp: + raise ValueError(f"invalid stage: {name}") + if dim is not None and (type(dim) is not int or not 0 <= dim < len(spec["shape"]) + or spec["shape"][dim] % tp): + raise ValueError(f"invalid split: {name}") + if stage != rank // tp: + continue + tensor = self.tensor(name) + if dim is not None: + width = tensor.shape[dim] // tp + tensor = tensor.narrow(dim, rank % tp * width, width) + result[name] = tensor + return result + + def close(self): + self._closed = True + # torch.frombuffer keeps its buffer alive through storage lifetime, including + # tensor views. Never explicitly close mappings still referenced by tensors. + self._maps.clear() + + def __enter__(self): + return self + + def __exit__(self, *_): + self.close() diff --git a/tests/test_torch_store.py b/tests/test_torch_store.py new file mode 100644 index 0000000..1283f07 --- /dev/null +++ b/tests/test_torch_store.py @@ -0,0 +1,115 @@ +import gc +from contextlib import contextmanager + +import pytest + +torch = pytest.importorskip("torch") + +from test_shards import layout, write_tensor +from infiniloader.shards import build_plan, export_rank +from infiniloader.torch_store import TensorStore +from infiniloader.cache import load_model_into_memory, build_view + + +@contextmanager +def resident(path, view): + model = load_model_into_memory(path, chunk_size=4096) + try: + build_view(model, view) + yield view + finally: + model.close() + + +def test_cow_protection_and_outliving_store(tmp_path): + source = tmp_path / "source" + write_tensor(source) + with resident(source, tmp_path / "view") as view: + store = TensorStore(view) + first = store.tensor("weight") + other = TensorStore(view) + untouched = other.tensor("weight") + first.add_(10) + assert untouched.flatten()[0].item() == 0 + tail = first[1:] + del first + store.close() + gc.collect() + assert tail.flatten()[0].item() == 22 + with pytest.raises(RuntimeError, match="closed"): + store.tensor("weight") + + +def test_live_layout_switch_does_not_reopen_source(tmp_path, monkeypatch): + source = tmp_path / "source" + write_tensor(source) + with resident(source, tmp_path / "view") as view: + store = TensorStore(view) + tensor = store.tensor("weight") + # Once mapped, the original path is no longer needed for topology changes. + (source / "weights.safetensors").unlink() + ptr = tensor.untyped_storage().data_ptr() + for tp in (4, 2, 1, 4): + outputs = [store.rank_tensors(layout(0, tp), r)["weight"] for r in range(tp)] + assert all(t.untyped_storage().data_ptr() == ptr for t in outputs) + torch.testing.assert_close(torch.cat(outputs, dim=0), tensor) + # Also test a strided inner-axis view without a hidden CPU contiguous copy. + outputs = [store.rank_tensors(layout(1), r)["weight"] for r in range(2)] + assert all(t.untyped_storage().data_ptr() == ptr for t in outputs) + torch.testing.assert_close(torch.cat(outputs, dim=1), tensor) + + +def test_sharded_store_enforces_exported_contract(tmp_path): + source = tmp_path / "source" + write_tensor(source) + rules = layout(1) + plan = build_plan(source, rules) + dest = tmp_path / "shard" + export_rank(source, plan, 1, dest) + with resident(dest, tmp_path / "shard-view") as view, TensorStore(view) as store: + actual = store.rank_tensors(rules, 1)["weight"] + with resident(source, tmp_path / "full-view") as full_view, TensorStore(full_view) as full: + torch.testing.assert_close(actual, full.tensor("weight")[:, 3:, :]) + with pytest.raises(ValueError, match="exported"): + store.rank_tensors(rules, 0) + + +def test_scalar_and_empty_cpu_view(tmp_path): + scalar, empty = tmp_path / "scalar", tmp_path / "empty" + write_tensor(scalar, shape=()) + write_tensor(empty, shape=(0, 2)) + with resident(scalar, tmp_path / "scalar-view") as view: + assert TensorStore(view).tensor("weight").shape == () + with resident(empty, tmp_path / "empty-view") as view: + assert TensorStore(view).tensor("weight").numel() == 0 + + +def test_reject_mutable_source(tmp_path): + source = tmp_path / "source" + write_tensor(source) + with pytest.raises(ValueError, match="sealed"): + TensorStore(source) + + +def test_receipt_is_bound_to_mapped_weights(tmp_path, monkeypatch): + import infiniloader.torch_store as module + import struct + source = tmp_path / "source" + write_tensor(source) + plan = build_plan(source, layout(0)) + first, second = tmp_path / "first", tmp_path / "second" + export_rank(source, plan, 0, first) + weights = source / "weights.safetensors" + weights.write_bytes(weights.read_bytes()[:-4] + struct.pack(" Date: Mon, 26 Oct 2026 21:07:53 +0800 Subject: [PATCH 09/29] Validate distributed TP and PP execution after coordinated preload --- .github/workflows/tests.yml | 13 ++ examples/distributed_mlp.py | 336 ++++++++++++++++++++++++++++++++ tests/test_torch_distributed.py | 79 ++++++++ 3 files changed, 428 insertions(+) create mode 100644 examples/distributed_mlp.py create mode 100644 tests/test_torch_distributed.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 5a17a8a..bdea10f 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -21,3 +21,16 @@ jobs: - run: python -m pytest -q - run: iloader selftest - run: iloader doctor --json + torch-cpu: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v7 + with: + persist-credentials: false + - uses: actions/setup-python@v7 + with: + python-version: '3.12' + - run: python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + - run: python -m pip install -e '.[test]' safetensors numpy + - run: python -m pytest -q diff --git a/examples/distributed_mlp.py b/examples/distributed_mlp.py new file mode 100644 index 0000000..cb90600 --- /dev/null +++ b/examples/distributed_mlp.py @@ -0,0 +1,336 @@ +"""Numerical TP/PP validation: local processes or one worker per SSH host. + +Uses real torch.distributed collectives and activation transfers, not IO simulation. +Run `python examples/distributed_mlp.py --help`. Only generated fixtures are used. +""" +from __future__ import annotations + +import argparse +import datetime +import json +import os +import shutil +import socket +import subprocess +import sys +import tempfile +import time +import uuid +from pathlib import Path + +from infiniloader.cache import build_view, load_model_into_memory +from infiniloader.shards import build_plan, canonical, digest, export_rank, read_json +from infiniloader.torch_store import TensorStore +from infiniloader.runtime import atomic_json, process_token + +DIM, HIDDEN, LAYERS = 32, 64, 4 + + +def fixture(): + import torch + g = torch.Generator().manual_seed(20261018) + weights = {} + for layer in range(LAYERS): + for role, shape in {"up": (HIDDEN, DIM), "down": (DIM, HIDDEN), + "up_bias": (HIDDEN,), "down_bias": (DIM,)}.items(): + weights[f"layers.{layer}.{role}"] = torch.randn(shape, generator=g) * .05 + return weights + + +def layout(tp, pp): + if LAYERS % pp: + raise ValueError("fixture layers must be divisible by pp") + return {"format": 1, "tp": tp, "pp": pp, "tensors": { + f"layers.{layer}.{role}": {"stage": layer // (LAYERS // pp), "split_dim": dim} + for layer in range(LAYERS) for role, dim in {"up": 0, "down": 1, "up_bias": 0, "down_bias": None}.items()}} + + +def input_tensor(): + import torch + return torch.arange(8 * DIM, dtype=torch.float32).reshape(8, DIM) / 256 + + +def reference(): + import torch + weights, x = fixture(), input_tensor() + for layer in range(LAYERS): + p = {r: weights[f"layers.{layer}.{r}"] for r in ("up", "down", "up_bias", "down_bias")} + x = torch.nn.functional.linear(torch.nn.functional.linear(x, p["up"], p["up_bias"]).relu(), p["down"], p["down_bias"]) + return x + + +def switch_cases(world): + cases = [(world // pp, pp) for pp in (1, 2, 4) if world % pp == 0 and pp <= LAYERS] + return cases + cases[:1] if len(cases) > 1 else cases + + +def validate_outputs(rows): + import torch + if not rows or not rows[0]["cases"] or any(len(r["cases"]) != len(rows[0]["cases"]) for r in rows): + raise RuntimeError("missing topology results") + errors, expected = [], reference() + for index in range(len(rows[0]["cases"])): + cases = [r["cases"][index] for r in rows] + if len({(c["tp"], c["pp"]) for c in cases}) != 1: + raise RuntimeError("ranks disagree on topology") + outputs = [c["output"] for c in cases if c["output"] is not None] + if len(outputs) != 1: + raise RuntimeError("each topology must produce exactly one terminal output") + actual = torch.tensor(outputs[0]) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5) + errors.append((actual - expected).abs().max().item()) + return errors + + +def prepare(root: Path, mode, tp, pp): + from safetensors.torch import save_file + root.mkdir(parents=True, exist_ok=False) + full = root / "canonical" + full.mkdir() + save_file(fixture(), str(full / "weights.safetensors")) + plan = build_plan(full, layout(tp, pp), mode=mode) + for rank in range(tp * pp): + export_rank(full, plan, rank, root / f"rank{rank}") + (root / "fixture.json").write_bytes(canonical({"format": 1, "kind": "generated-mlp", "plan": plan})) + return plan + + +def memory(): + result = {} + for line in Path("/proc/self/smaps_rollup").read_text().splitlines(): + key, _, value = line.partition(":") + if key in ("Rss", "Pss", "Private_Dirty", "Private_Clean", "Shared_Dirty", "Shared_Clean"): + result[key + "_bytes"] = int(value.split()[0]) * 1024 + return result + + +def publish_start(directory: Path, world: int, plan_id: str, timeout: float): + """Controller-side release after every independently launched rank is ready.""" + deadline = time.monotonic() + timeout + paths = [directory / f"ready{r}.json" for r in range(world)] + while not all(p.exists() for p in paths): + if time.monotonic() >= deadline: + raise TimeoutError("not every rank reached the startup barrier") + time.sleep(.05) + rows = [read_json(p) for p in paths] + if any(r.get("rank") != i or r.get("plan_id") != plan_id or not isinstance(r.get("nonce"), str) + for i, r in enumerate(rows)) or len({r["nonce"] for r in rows}) != world: + raise ValueError("mixed versions/ranks at startup barrier") + marker = directory / "start.json" + if marker.exists(): + raise FileExistsError("startup barrier was already released") + atomic_json(marker, {"format": 1, "plan_id": plan_id, + "participants": {str(r["rank"]): r["nonce"] for r in rows}}) + return rows + + +def wait_start(directory: Path, rank: int, world: int, plan_id: str, timeout: float): + if not directory.is_dir() or directory.is_symlink(): + raise ValueError("startup barrier requires a fresh shared directory") + ready = {"rank": rank, "plan_id": plan_id, "nonce": uuid.uuid4().hex, + "hostname": socket.gethostname(), "pid": os.getpid(), "process_token": process_token(os.getpid())} + path = directory / f"ready{rank}.json" + # Each operation uses a fresh barrier; never join stale readiness records. + if path.exists(): + raise FileExistsError(path) + atomic_json(path, ready) + deadline = time.monotonic() + timeout + while True: + marker = directory / "start.json" + if marker.exists(): + start = read_json(marker) + if (start.get("format") != 1 or start.get("plan_id") != plan_id + or set(start.get("participants", {})) != {str(r) for r in range(world)} + or start["participants"].get(str(rank)) != ready["nonce"]): + raise ValueError("stale or mismatched startup release") + return ready + if time.monotonic() >= deadline: + raise TimeoutError("controller did not release startup barrier") + time.sleep(.05) + + +def execute(store, rule, rank, device): + import torch + import torch.distributed as dist + tp, pp = rule["tp"], rule["pp"] + stage, lane = divmod(rank, tp) + groups = [dist.new_group(list(range(s * tp, (s + 1) * tp))) for s in range(pp)] + try: + # .to('cpu') preserves the CPU COW views; CUDA makes only this rank's copies. + tensors = {n: t.to(device) for n, t in store.rank_tensors(rule, rank).items()} + x = input_tensor().to(device) if stage == 0 else torch.empty((8, DIM), device=device) + if stage: + dist.recv(x, src=rank - tp) + for layer in range(stage * (LAYERS // pp), (stage + 1) * (LAYERS // pp)): + p = {r: tensors[f"layers.{layer}.{r}"] for r in ("up", "down", "up_bias", "down_bias")} + hidden = torch.nn.functional.linear(x, p["up"], p["up_bias"]).relu() + x = torch.nn.functional.linear(hidden, p["down"]) + if tp > 1: + dist.all_reduce(x, group=groups[stage]) + x += p["down_bias"] # Add replicated row-parallel bias exactly once. + if stage < pp - 1: + dist.send(x, dst=rank + tp) + return x.cpu().tolist() if stage == pp - 1 and lane == 0 else None + finally: + for group in groups: + if group != dist.GroupMember.NON_GROUP_MEMBER: + dist.destroy_process_group(group) + + +def worker(args): + import torch + import torch.distributed as dist + torch.set_num_threads(1) + if args.device != "cpu": + torch.cuda.set_device(torch.device(args.device)) + # Finish device runtime startup before reporting readiness to the controller. + warmup = torch.zeros(1, device=args.device) + torch.cuda.synchronize() + del warmup + meta = read_json(args.root / "fixture.json") + if meta.get("kind") != "generated-mlp": + raise ValueError("only generated fixture directories may be used") + plan = meta["plan"] + source = args.root / f"rank{args.rank}" + if (source / "weights.safetensors.offline").exists(): + raise ValueError("fixture already consumed; prepare a fresh one") + with tempfile.TemporaryDirectory(prefix="iloader-mlp-") as temp: + model, store = None, None + try: + start = time.perf_counter() + model = load_model_into_memory(source, chunk_size=1024 * 1024, max_memory_bytes=256 * 1024**2) + view = Path(temp) / "view" + build_view(model, view) + store = TensorStore(view) + if store.content_id != plan["content_id"]: + raise ValueError("fixture content mismatch") + # Remove original fixture access. All later TP/PP switches must use RAM. + (source / "weights.safetensors").rename(source / "weights.safetensors.offline") + info = {"rank": args.rank, "hostname": socket.gethostname(), + "boot_id": Path("/proc/sys/kernel/random/boot_id").read_text().strip(), + "content_id": store.content_id, "plan_id": plan["plan_id"], + "backend": args.backend, "device": args.device, "torch": torch.__version__, + "canonical_tensor_bytes": store.catalog["tensor_bytes"], + "resident_file_bytes": model.total_bytes, "preload_seconds": time.perf_counter() - start, + "memory_after_map": memory(), "source_available": False} + if args.ready_dir: + info["startup"] = wait_start(args.ready_dir, args.rank, args.world_size, plan["plan_id"], args.timeout) + dist.init_process_group(args.backend, init_method=args.init, rank=args.rank, + world_size=args.world_size, timeout=datetime.timedelta(seconds=args.timeout)) + identities = [None] * args.world_size + dist.all_gather_object(identities, (store.content_id, plan["plan_id"])) + if len(set(identities)) != 1 or args.world_size != len(plan["ranks"]): + raise ValueError("mixed checkpoint/layout versions or world size") + cases = [(plan["layout"]["tp"], plan["layout"]["pp"])] + if plan["mode"] == "replicated" and args.switch: + cases = switch_cases(args.world_size) + results = [] + for tp, pp in cases: + before = time.perf_counter() + output = execute(store, layout(tp, pp), args.rank, args.device) + dist.barrier() + if args.device != "cpu": + torch.cuda.synchronize() + results.append({"tp": tp, "pp": pp, "seconds": time.perf_counter() - before, + "output": output, "memory": memory()}) + info["cases"] = results + info["ok"] = True + args.output.write_text(json.dumps(info, indent=2)) + finally: + if store is not None: + store.close() + if model is not None: + model.close() + if dist.is_initialized(): + dist.destroy_process_group() + + +def run_local(args): + import torch + root = args.root + plan = prepare(root, args.mode, args.tp, args.pp) + world = args.tp * args.pp + barrier = root / "startup" + barrier.mkdir() + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + port = probe.getsockname()[1] + processes, logs = [], [] + start = time.perf_counter() + try: + for rank in range(world): + output = root / f"result{rank}.json" + log = (root / f"worker{rank}.log").open("w") + logs.append(log) + cmd = [sys.executable, str(Path(__file__).resolve()), "worker", "--root", str(root), + "--rank", str(rank), "--world-size", str(world), "--init", f"tcp://127.0.0.1:{port}", + "--backend", args.backend, "--device", args.device, "--output", str(output), + "--timeout", str(args.timeout)] + cmd.extend(["--ready-dir", str(barrier)]) + if args.mode == "replicated": + cmd.append("--switch") + processes.append(subprocess.Popen(cmd, stdout=log, stderr=log)) + try: + publish_start(barrier, world, plan["plan_id"], args.timeout) + except TimeoutError: + if any(p.poll() not in (None, 0) for p in processes): + raise RuntimeError(f"worker failed before startup; inspect {root}/worker*.log") from None + raise + deadline = time.monotonic() + args.timeout + 20 + for p in processes: + p.wait(timeout=max(.1, deadline - time.monotonic())) + if any(p.returncode for p in processes): + raise RuntimeError(f"worker failed; inspect {root}/worker*.log") + rows = [read_json(root / f"result{r}.json") for r in range(world)] + errors = validate_outputs(rows) + result = {"ok": True, "scope": "single-host torch.distributed", "mode": args.mode, + "world_size": world, "plan": plan, "max_abs_error": max(errors), "ranks": rows, + "wall_seconds": time.perf_counter() - start} + args.output.write_text(json.dumps(result, indent=2)) + print(json.dumps({k: v for k, v in result.items() if k not in ("ranks", "plan")}, indent=2)) + finally: + for process in processes: + if process.poll() is None: + process.terminate() + for process in processes: + try: + process.wait(timeout=3) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + for log in logs: + log.close() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="action", required=True) + for action in ("prepare", "local", "worker"): + p = sub.add_parser(action) + p.add_argument("--root", type=Path, required=True) + if action != "worker": + p.add_argument("--mode", choices=("sharded", "replicated"), default="sharded") + p.add_argument("--tp", type=int, default=4) + p.add_argument("--pp", type=int, default=1) + if action != "prepare": + p.add_argument("--output", type=Path, required=True) + p.add_argument("--backend", choices=("gloo", "nccl"), default="gloo") + p.add_argument("--device", default="cpu") + p.add_argument("--timeout", type=int, default=40) + if action == "worker": + p.add_argument("--rank", type=int, required=True) + p.add_argument("--world-size", type=int, required=True) + p.add_argument("--init", required=True) + p.add_argument("--switch", action="store_true") + p.add_argument("--ready-dir", type=Path, help="wait for controller release after all ranks preload") + args = parser.parse_args() + if args.action == "prepare": + print(json.dumps(prepare(args.root, args.mode, args.tp, args.pp), indent=2)) + elif args.action == "worker": + worker(args) + else: + run_local(args) + + +if __name__ == "__main__": + main() diff --git a/tests/test_torch_distributed.py b/tests/test_torch_distributed.py new file mode 100644 index 0000000..cb1d7e0 --- /dev/null +++ b/tests/test_torch_distributed.py @@ -0,0 +1,79 @@ +import argparse +import importlib.util +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") +safe = pytest.importorskip("safetensors.torch") +pytest.importorskip("numpy") + +script = Path(__file__).resolve().parents[1] / "examples" / "distributed_mlp.py" +spec = importlib.util.spec_from_file_location("distributed_mlp_example", script) +mlp = importlib.util.module_from_spec(spec) +spec.loader.exec_module(mlp) + + +@pytest.mark.parametrize("mode,tp,pp", [("sharded", 4, 1), ("sharded", 1, 4), + ("sharded", 2, 2), ("replicated", 4, 1), ("replicated", 1, 1)]) +def test_real_collectives_and_layout_switches(tmp_path, mode, tp, pp): + args = argparse.Namespace(root=tmp_path / "fixture", mode=mode, tp=tp, pp=pp, + output=tmp_path / "result.json", backend="gloo", device="cpu", timeout=25) + mlp.run_local(args) + result = mlp.read_json(args.output) + assert result["ok"] and result["max_abs_error"] < 1e-6 + assert all(not r["source_available"] for r in result["ranks"]) + assert sum(r["canonical_tensor_bytes"] for r in result["ranks"]) == result["plan"]["aggregate_tensor_bytes"] + if mode == "replicated": + assert all(len(r["cases"]) == len(mlp.switch_cases(tp * pp)) for r in result["ranks"]) + # Independent official parser accepts each exported safetensors checkpoint. + for rank in range(tp * pp): + official = safe.load_file(str(args.root / f"rank{rank}" / "weights.safetensors.offline")) + parts = result["plan"]["ranks"][rank]["tensors"] + assert set(official) == set(parts) + for name, tensor in official.items(): + assert list(tensor.shape) == parts[name]["shape"] + + +def test_corrupt_rank_fails_and_all_workers_are_reaped(tmp_path, monkeypatch): + original_prepare = mlp.prepare + def corrupt(root, *args): + plan = original_prepare(root, *args) + file = root / "rank1" / "weights.safetensors" + file.write_bytes(file.read_bytes()[:-1]) + return plan + original_popen, children = mlp.subprocess.Popen, [] + def remember(*args, **kwargs): + p = original_popen(*args, **kwargs) + children.append(p) + return p + monkeypatch.setattr(mlp, "prepare", corrupt) + monkeypatch.setattr(mlp.subprocess, "Popen", remember) + args = argparse.Namespace(root=tmp_path / "fixture", mode="sharded", tp=2, pp=1, + output=tmp_path / "result.json", backend="gloo", device="cpu", timeout=20) + with pytest.raises(RuntimeError, match="worker failed"): + mlp.run_local(args) + assert len(children) == 2 and all(p.poll() is not None for p in children) + assert not args.output.exists() + + +def test_every_switch_needs_one_output_and_legal_small_world(): + assert mlp.switch_cases(1) == [(1, 1)] + assert mlp.switch_cases(2) == [(2, 1), (1, 2), (2, 1)] + assert all(pp <= mlp.LAYERS for tp, pp in mlp.switch_cases(8)) + output = mlp.reference().tolist() + rows = [{"cases": [{"tp": 1, "pp": 1, "output": output}, + {"tp": 1, "pp": 1, "output": None}]}] + with pytest.raises(RuntimeError, match="exactly one"): + mlp.validate_outputs(rows) + + +def test_startup_barrier_rejects_missing_and_mixed_versions(tmp_path): + with pytest.raises(TimeoutError): + mlp.publish_start(tmp_path, 2, "plan", .01) + for rank in range(2): + (tmp_path / f"ready{rank}.json").write_text(__import__("json").dumps({ + "rank": rank, "plan_id": "wrong" if rank else "plan", "nonce": str(rank)})) + with pytest.raises(ValueError, match="mixed"): + mlp.publish_start(tmp_path, 2, "plan", .1) + assert not (tmp_path / "start.json").exists() From 895c5efce64030dacc635eb26331d1aa606417b2 Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 26 Oct 2026 21:10:07 +0800 Subject: [PATCH 10/29] Manage GPU expert residency with CPU fallback and safe eviction --- .github/workflows/tests.yml | 1 + examples/hybrid_moe.py | 85 ++++++++++++++ src/infiniloader/experts.py | 224 ++++++++++++++++++++++++++++++++++++ tests/test_experts.py | 120 +++++++++++++++++++ 4 files changed, 430 insertions(+) create mode 100644 examples/hybrid_moe.py create mode 100644 src/infiniloader/experts.py create mode 100644 tests/test_experts.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index bdea10f..5e19ae6 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -34,3 +34,4 @@ jobs: - run: python -m pip install torch --index-url https://download.pytorch.org/whl/cpu - run: python -m pip install -e '.[test]' safetensors numpy - run: python -m pytest -q + - run: python examples/hybrid_moe.py --device cpu --output hybrid-cpu.json diff --git a/examples/hybrid_moe.py b/examples/hybrid_moe.py new file mode 100644 index 0000000..4dbe6cf --- /dev/null +++ b/examples/hybrid_moe.py @@ -0,0 +1,85 @@ +"""Small real CPU/GPU MoE validation; no model download or global GPU changes.""" +from __future__ import annotations + +import argparse +import json +import tempfile +import time +from pathlib import Path + +from infiniloader.cache import build_view, load_model_into_memory +from infiniloader.experts import ExpertPool, hybrid_moe +from infiniloader.torch_store import TensorStore + + +def run(device): + import torch + from safetensors.torch import save_file + torch.set_num_threads(1) + generator = torch.Generator().manual_seed(23) + groups, weights = {}, {} + for expert in range(8): + groups[str(expert)] = {} + for role, shape in {"up": (128, 64), "down": (64, 128)}.items(): + name = f"experts.{expert}.{role}" + groups[str(expert)][role] = name + weights[name] = torch.randn(shape, generator=generator) * .05 + with tempfile.TemporaryDirectory(prefix="iloader-moe-") as directory: + root = Path(directory) + source = root / "source" + source.mkdir() + save_file(weights, str(source / "weights.safetensors")) + del weights + model = load_model_into_memory(source, chunk_size=1024**2) + try: + build_view(model, root / "view") + with TensorStore(root / "view") as store: + # Ensure every cold lookup also comes from the immutable RAM store. + (source / "weights.safetensors").unlink() + one = 2 * 128 * 64 * 4 + x = torch.randn(32, 64, generator=generator) + ids = torch.arange(64).reshape(32, 2) % 8 + gates = torch.softmax(torch.randn(32, 2, generator=generator), dim=1) + expected = torch.zeros_like(x) + for token in range(len(x)): + for slot in range(2): + names = groups[str(ids[token, slot].item())] + h = (x[token] @ store.tensor(names["up"]).T).relu() + expected[token] += gates[token, slot] * (h @ store.tensor(names["down"]).T) + if device is not None: + torch.cuda.set_device(torch.device(device)) + torch.cuda.reset_peak_memory_stats() + baseline_allocated = torch.cuda.memory_allocated() + with ExpertPool(store, groups, max_gpu_bytes=2 * one, device=device) as pool: + phases = [] + for hot in (("0", "1"), ("2", "3"), ("6", "7"), ("0", "1")): + start = time.perf_counter() + admitted = [pool.promote(e) for e in hot] + output = hybrid_moe(pool, x, ids, gates) + torch.testing.assert_close(output, expected, rtol=1e-5, atol=1e-6) + phases.append(dict(hot=list(hot), admitted=admitted, + max_abs_error=(output - expected).abs().max().item(), + seconds=time.perf_counter() - start, pool=pool.status())) + result = {"ok": True, "gpu_validated": device is not None, + "scope": "reference dense ReLU MoE; no KTransformers kernel integration", + "cpu_tensor_bytes": store.catalog["tensor_bytes"], "cpu_source_available": False, + "gpu_weight_budget_bytes": 2 * one, "phases": phases, "torch": torch.__version__} + if device is not None: + result["gpu"] = {"name": torch.cuda.get_device_name(), + "runtime": "rocm" if torch.version.hip else "cuda", + "allocated_baseline": baseline_allocated, + "peak_allocated": torch.cuda.max_memory_allocated(), + "peak_reserved": torch.cuda.max_memory_reserved()} + return result + finally: + model.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cpu", help="cpu, cuda:0, ... (PyTorch ROCm also uses cuda)") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + result = run(None if args.device == "cpu" else args.device) + args.output.write_text(json.dumps(result, indent=2)) + print(json.dumps(result, indent=2)) diff --git a/src/infiniloader/experts.py b/src/infiniloader/experts.py new file mode 100644 index 0000000..3c0af96 --- /dev/null +++ b/src/infiniloader/experts.py @@ -0,0 +1,224 @@ +"""Bounded, lease-protected GPU expert copies over a canonical CPU TensorStore. + +This is an engine integration API and reference dense MoE, not a KTransformers +kernel/quantization adapter. CPU cold experts execute without GPU admission. +""" +from __future__ import annotations + +import math +import threading +from collections import OrderedDict +from contextlib import contextmanager +from dataclasses import dataclass, field + +from .distributed import DTYPE_BYTES + + +@dataclass +class _Entry: + weights: dict + nbytes: int + users: int = 0 + events: list = field(default_factory=list) + + +@dataclass +class ExpertLease: + weights: dict + device: object + placement: str + _streams: list = field(default_factory=list) + + def record_stream(self, stream): + """Register additional CUDA consumer streams before leaving the lease.""" + if self.placement != "gpu": + raise ValueError("CPU leases do not use CUDA streams") + if stream.device != self.device: + raise ValueError("consumer stream belongs to another device") + if stream not in self._streams: + self._streams.append(stream) + + +class ExpertPool: + def __init__(self, store, groups: dict[str, dict[str, str]], *, max_gpu_bytes: int, device=None): + self.torch, self.store = store.torch, store + if type(max_gpu_bytes) is not int or max_gpu_bytes < 0: + raise ValueError("max_gpu_bytes must be a nonnegative integer") + if not groups or any(not isinstance(k, str) or not v for k, v in groups.items()): + raise ValueError("expert groups require string IDs and nonempty parameter maps") + self.groups = {k: dict(v) for k, v in groups.items()} + self.sizes = {k: sum(math.prod(store.catalog["tensors"][n]["shape"]) * + DTYPE_BYTES[store.catalog["tensors"][n]["dtype"]] for n in v.values()) + for k, v in self.groups.items()} + self.limit = max_gpu_bytes + self.device = self.torch.device(device) if device is not None else None + if self.device is not None: + if self.device.type != "cuda" or not self.torch.cuda.is_available(): + raise ValueError("GPU expert pool requires an available CUDA/ROCm device") + index = self.device.index if self.device.index is not None else self.torch.cuda.current_device() + self.device = self.torch.device("cuda", index) + self.entries = OrderedDict() + self._lock = threading.RLock() + self._users = 0 + self._closed = False + self._stats = dict(gpu_hits=0, cpu_hits=0, promotions=0, evictions=0, uploaded_bytes=0) + + def _check(self): + if self._closed: + raise RuntimeError("expert pool is closed") + + def _cpu(self, expert): + return {role: self.store.tensor(name) for role, name in self.groups[expert].items()} + + def _evict(self, expert): + entry = self.entries[expert] + if entry.users: + raise RuntimeError("expert is leased") + for event in entry.events: + event.synchronize() + del self.entries[expert] + self._stats["evictions"] += 1 + + def promote(self, expert: str) -> bool: + """Admit the complete expert or return False (CPU fallback is available). + + Transfers are synchronous. The lock reserves capacity across the entire + allocation, preventing concurrent duplicate uploads and partial publication. + """ + with self._lock: + self._check() + size = self.sizes[expert] + if expert in self.entries: + self.entries.move_to_end(expert) + return True + if self.device is None or size > self.limit: + return False + while sum(e.nbytes for e in self.entries.values()) + size > self.limit: + victim = next((k for k, e in self.entries.items() if not e.users), None) + if victim is None: + return False + self._evict(victim) + weights = {} + try: + with self.torch.cuda.device(self.device): + for role, tensor in self._cpu(expert).items(): + weights[role] = tensor.to(self.device, non_blocking=False) + self.torch.cuda.current_stream(self.device).synchronize() + except BaseException: + if weights: + self.torch.cuda.synchronize(self.device) + weights.clear() + raise + self.entries[expert] = _Entry(weights, size) + self._stats["promotions"] += 1 + self._stats["uploaded_bytes"] += size + return True + + @contextmanager + def acquire(self, expert: str, *, stream=None): + """Execute inside this scope; GPU work defaults to the selected stream. + + Additional streams must be registered using lease.record_stream(). Tensor + references retained outside the lease must not be used for new operations. + """ + with self._lock: + self._check() + entry = self.entries.get(expert) + if entry: + consumer = stream or self.torch.cuda.current_stream(self.device) + if getattr(consumer, "device", None) != self.device: + raise ValueError("consumer stream belongs to another device") + lease = ExpertLease(dict(entry.weights), self.device, "gpu", [consumer]) + entry.users += 1 + self.entries.move_to_end(expert) + else: + if stream is not None: + raise ValueError("CPU fallback does not accept a CUDA stream") + lease = ExpertLease(self._cpu(expert), self.torch.device("cpu"), "cpu") + self._users += 1 + self._stats["gpu_hits" if entry else "cpu_hits"] += 1 + try: + if entry: + with self.torch.cuda.device(self.device), self.torch.cuda.stream(consumer): + yield lease + else: + yield lease + finally: + with self._lock: + if entry: + try: + for consumer_stream in lease._streams: + event = self.torch.cuda.Event() + event.record(consumer_stream) + entry.events.append(event) + for tensor in entry.weights.values(): + tensor.record_stream(consumer_stream) + # Completed events do not need to accumulate over hot hits. + entry.events[:] = [e for e in entry.events if not e.query()] + except BaseException: + self.torch.cuda.synchronize(self.device) + raise + finally: + lease.weights.clear() + entry.users -= 1 + self._users -= 1 + else: + lease.weights.clear() + self._users -= 1 + + def evict(self, expert): + with self._lock: + self._check() + if expert in self.entries: + self._evict(expert) + + def status(self): + with self._lock: + return dict(self._stats, gpu_weight_bytes=sum(e.nbytes for e in self.entries.values()), + max_gpu_bytes=self.limit, active_leases=self._users, + resident_experts=list(self.entries), device=str(self.device)) + + def close(self): + with self._lock: + if self._users: + raise RuntimeError("expert pool has active leases") + for expert in list(self.entries): + self._evict(expert) + self._closed = True + + def __enter__(self): + return self + + def __exit__(self, *_): + self.close() + + +def hybrid_moe(pool: ExpertPool, inputs, expert_ids, gates): + """Reference ReLU MoE with externally provided top-k routing, including repeats. + + Each group maps up/down and optionally up_bias/down_bias to tensor names. + Returns output on the input device. Engine adapters can instead use acquire() + around their own CPU/GPU kernels, packed formats and routing implementation. + """ + torch = pool.torch + if (inputs.ndim != 2 or expert_ids.ndim != 2 or gates.shape != expert_ids.shape + or expert_ids.shape[0] != inputs.shape[0] or expert_ids.dtype not in (torch.int32, torch.int64)): + raise ValueError("expected inputs[tokens,dim], integer IDs and gates[tokens,top_k]") + ids = expert_ids.to("cpu") + if any(str(i) not in pool.groups for i in ids.unique().tolist()): + raise ValueError("unknown routed expert") + first = next(iter(pool.groups.values())) + out_dim = pool.store.catalog["tensors"][first["down"]]["shape"][0] + output = torch.zeros((inputs.shape[0], out_dim), dtype=inputs.dtype, device=inputs.device) + with torch.inference_mode(): + for expert in ids.unique().tolist(): + tokens, slots = (ids == expert).nonzero(as_tuple=True) + with pool.acquire(str(expert)) as lease: + weights = lease.weights + activation = inputs.index_select(0, tokens.to(inputs.device)).to(lease.device) + hidden = torch.nn.functional.linear(activation, weights["up"], weights.get("up_bias")).relu() + result = torch.nn.functional.linear(hidden, weights["down"], weights.get("down_bias")) + factors = gates[tokens.to(gates.device), slots.to(gates.device)].to(result.device) + result = (result * factors[:, None]).to(device=output.device, dtype=output.dtype) + output.index_add_(0, tokens.to(output.device), result) + return output diff --git a/tests/test_experts.py b/tests/test_experts.py new file mode 100644 index 0000000..3089ed1 --- /dev/null +++ b/tests/test_experts.py @@ -0,0 +1,120 @@ +from contextlib import contextmanager + +import pytest + +torch = pytest.importorskip("torch") +safe = pytest.importorskip("safetensors.torch") + +from test_torch_store import resident +from infiniloader.torch_store import TensorStore +from infiniloader.experts import ExpertPool, hybrid_moe + + +@contextmanager +def moe(tmp_path, device=None, dtype=torch.float32): + source = tmp_path / "source" + source.mkdir() + generator = torch.Generator().manual_seed(17) + weights, groups = {}, {} + for i in range(4): + groups[str(i)] = {} + for role, shape in {"up": (16, 8), "down": (8, 16), "up_bias": (16,), "down_bias": (8,)}.items(): + name = f"experts.{i}.{role}" + groups[str(i)][role] = name + weights[name] = (torch.randn(shape, generator=generator) * 0.1).to(dtype) + safe.save_file(weights, str(source / "model.safetensors")) + with resident(source, tmp_path / "view") as view, TensorStore(view) as store: + one = sum(weights[n].numel() * weights[n].element_size() for n in groups["0"].values()) + with ExpertPool(store, groups, max_gpu_bytes=2 * one, device=device) as pool: + yield pool, weights, groups, one + + +def reference(inputs, ids, gates, weights, groups): + # Independent full CPU implementation, deliberately token/route based. + output = torch.zeros_like(inputs) + for token in range(len(inputs)): + for slot in range(ids.shape[1]): + p = {role: weights[name] for role, name in groups[str(ids[token, slot].item())].items()} + hidden = (inputs[token] @ p["up"].T + p["up_bias"]).relu() + output[token] += gates[token, slot] * (hidden @ p["down"].T + p["down_bias"]) + return output + + +def test_cpu_cold_experts_topk_repeats_and_empty_routes(tmp_path): + with moe(tmp_path) as (pool, weights, groups, one): + x = torch.arange(32, dtype=torch.float32).reshape(4, 8) / 32 + ids = torch.tensor([[0, 1], [2, 3], [1, 1], [3, 0]]) + gates = torch.tensor([[.2, .8], [.7, .3], [.5, .5], [.9, .1]]) + torch.testing.assert_close(hybrid_moe(pool, x, ids, gates), reference(x, ids, gates, weights, groups)) + assert pool.status()["cpu_hits"] == 4 and pool.status()["gpu_hits"] == 0 + assert not pool.promote("0") + assert hybrid_moe(pool, x[:0], ids[:0], gates[:0]).shape == (0, 8) + with pool.acquire("0") as lease: + with pytest.raises(RuntimeError, match="active leases"): + pool.close() + assert not lease.weights + + +@pytest.mark.parametrize("dtype,gate_dtype", [(torch.float32, torch.float64), (torch.bfloat16, torch.float32)]) +def test_mixed_router_precision(tmp_path, dtype, gate_dtype): + with moe(tmp_path, dtype=dtype) as (pool, weights, groups, one): + inputs = torch.ones((1, 8), dtype=dtype) + ids = torch.tensor([[0, 1]]) + gates = torch.tensor([[.25, .75]], dtype=gate_dtype) + actual = hybrid_moe(pool, inputs, ids, gates) + assert actual.dtype == inputs.dtype + torch.testing.assert_close(actual, reference(inputs, ids, gates, weights, groups), atol=2e-3, rtol=1e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires real CUDA/ROCm device") +def test_gpu_hot_cpu_cold_lru_budget_and_lease(tmp_path): + with moe(tmp_path, "cuda:0") as (pool, weights, groups, one): + assert pool.promote("0") and pool.promote("1") + with pytest.raises(ValueError, match="stream"): + with pool.acquire("0", stream=object()): + pass + assert pool.status()["active_leases"] == 0 + assert pool.entries["0"].users == 0 + assert pool.status()["gpu_weight_bytes"] == 2 * one + with pool.acquire("0") as lease: + assert lease.device.type == "cuda" + with pytest.raises(RuntimeError, match="leased"): + pool.evict("0") + assert pool.promote("2") # evicts 1, keeps leased 0 + with pool.acquire("2"): + assert not pool.promote("3") # all capacity is currently leased + x = torch.randn(5, 8) + ids = torch.tensor([[0, 1], [2, 3], [1, 1], [3, 0], [2, 0]]) + gates = torch.softmax(torch.randn(5, 2), dim=1) + torch.testing.assert_close(hybrid_moe(pool, x, ids, gates), reference(x, ids, gates, weights, groups), rtol=1e-5, atol=1e-6) + assert pool.status()["cpu_hits"] >= 2 and pool.status()["gpu_hits"] >= 2 + assert pool.status()["gpu_weight_bytes"] <= pool.limit + pool.limit = one - 1 + pool.evict("0") + pool.evict("2") + assert not pool.promote("3") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires real CUDA/ROCm device") +def test_gpu_cross_stream_completion_and_failed_group_is_not_published(tmp_path, monkeypatch): + with moe(tmp_path, "cuda:0") as (pool, weights, groups, one): + pool.promote("0") + stream = torch.cuda.Stream() + done = torch.cuda.Event() + with pool.acquire("0") as lease: + lease.record_stream(stream) + with torch.cuda.stream(stream): + result = lease.weights["up"] @ torch.ones(8, device="cuda") + done.record(stream) + pool.evict("0") + assert done.query() + torch.testing.assert_close(result.cpu(), weights[groups["0"]["up"]] @ torch.ones(8)) + real = pool._cpu + class Failure: + def to(self, *args, **kwargs): + raise RuntimeError("injected transfer failure") + monkeypatch.setattr(pool, "_cpu", lambda expert: {**real(expert), "down": Failure()}) + with pytest.raises(RuntimeError, match="injected"): + pool.promote("1") + assert pool.status()["gpu_weight_bytes"] == 0 + assert "1" not in pool.status()["resident_experts"] From 33b176e937ce6035ef39d4aab7414fd74a247424 Mon Sep 17 00:00:00 2001 From: whjthu Date: Mon, 26 Oct 2026 21:13:47 +0800 Subject: [PATCH 11/29] Document validated history and development priorities --- README.md | 13 +- README.zh-CN.md | 7 +- docs/distributed-results.json | 306 ++++++++++++++++++++++++++ docs/distributed-validation.md | 133 +++++++++++ docs/distributed.md | 202 +++++++++++++++++ docs/history-reorganization.json | 363 +++++++++++++++++++++++++++++++ docs/history-reorganization.md | 90 ++++++++ docs/roadmap.zh-CN.md | 47 ++++ 8 files changed, 1158 insertions(+), 3 deletions(-) create mode 100644 docs/distributed-results.json create mode 100644 docs/distributed-validation.md create mode 100644 docs/distributed.md create mode 100644 docs/history-reorganization.json create mode 100644 docs/history-reorganization.md create mode 100644 docs/roadmap.zh-CN.md diff --git a/README.md b/README.md index 909ea48..d211f07 100644 --- a/README.md +++ b/README.md @@ -240,6 +240,9 @@ iloader dist-plan "$MODEL" --tp 4 --limit 20 ## Command Summary +Explicit tensor sharding, live CPU layout selection and the CPU/GPU expert API +are documented in [Distributed and hybrid residency](docs/distributed.md). + ```bash iloader inspect MODEL_DIR iloader start MODEL_DIR --name NAME @@ -296,5 +299,11 @@ local cache plus memfd views. - The default view duplicates selected model bytes in RAM. - Fully transparent same-path interception is not included in 0.1.0; pass the printed view path to the engine. -- Physical rank-sharded checkpoint conversion is not included in this cleaned - release; `dist-plan` only analyzes candidate tensor splits. +- `shard plan/export/verify` supports explicit InfiniLoader TP/PP checkpoints. + `dist-plan` still only analyzes candidates. Loading these checkpoints into a + production model engine requires an adapter; see the reference PyTorch MLP. + +## Development notes + +See [commit history and validation](docs/history-reorganization.md) and the +[development roadmap](docs/roadmap.zh-CN.md). diff --git a/README.zh-CN.md b/README.zh-CN.md index 22aa82b..8357d52 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -258,4 +258,9 @@ iloader selftest - 框架读取 view 时，InfiniLoader daemon 必须保持运行。 - 默认 view 会在内存中保留一份选中的模型文件内容。 - 0.1.0 不包含完全透明的同路径拦截；需要把输出的 view 路径传给推理引擎。 -- 当前清理版不包含物理 rank-sharded checkpoint 转换；`dist-plan` 只做候选 tensor 切分分析。 +- `shard plan/export/verify` 已支持显式 TP/PP checkpoint；`dist-plan` 仍仅做候选分析。生产引擎加载这些分片需要适配器。 + +## 开发说明 + +显式分片、CPU TensorStore 与专家池接口见 [分布式与混合驻留](docs/distributed.md)。 +提交线及验证见 [提交重整记录](docs/history-reorganization.md)，后续能力见 [开发路线](docs/roadmap.zh-CN.md)。 diff --git a/docs/distributed-results.json b/docs/distributed-results.json new file mode 100644 index 0000000..64a9fbc --- /dev/null +++ b/docs/distributed-results.json @@ -0,0 +1,306 @@ +{ + "date": "2026-10-18", + "scope": "Small deterministic integration fixtures; not production LLM benchmarks", + "final_code_commit": "758f8c866a4de4dc4a8a9057ce14680e78e44b8d", + "environment": { + "nodes": 4, + "python": "3.11.9", + "torch": "2.9.0", + "hip": "6.3.26093", + "device": "Hygon BW DCU, approximately 64 GiB", + "gpu_devices_used_per_node": 1, + "transport": "Gloo/RCCL sockets; no RDMA qualification" + }, + "local_regression": { + "passed": 119, + "skipped_real_device": 2 + }, + "ci": "https://github.com/InfiniTensor/InfiniLoader/actions/runs/39107252516", + "distributed_cases": [ + { + "case": "tp4", + "commit": "123c6e2dd3f9e69a1df6cd9422d27ba0bd2628f5", + "backend": "gloo", + "device": "cpu", + "unique_tensor_bytes": 67072, + "aggregate_tensor_bytes": 68608, + "per_rank_tensor_bytes": [ + 17152, + 17152, + 17152, + 17152 + ], + "per_rank_resident_file_bytes": [ + 25505, + 25505, + 25505, + 25505 + ], + "distinct_host_count": 4, + "distinct_boot_count": 4, + "model_content_id": "9412c39d7796ccf3d177504ce019a335e9ebfcb0e67e2b991700553dec485f4b", + "layout_id": "8827169242156e4ed206ef26947e99cc1ae82fadb59534fda9d8f4a0a8e43ba8", + "original_rank_path_available": false, + "topologies": [ + { + "tp": 4, + "pp": 1, + "max_abs_error": 7.450580596923828e-09 + } + ], + "pss_after_map_bytes": [ + 329339904, + 331081728, + 329254912, + 328738816 + ] + }, + { + "case": "pp4", + "commit": "123c6e2dd3f9e69a1df6cd9422d27ba0bd2628f5", + "backend": "gloo", + "device": "cpu", + "unique_tensor_bytes": 67072, + "aggregate_tensor_bytes": 67072, + "per_rank_tensor_bytes": [ + 16768, + 16768, + 16768, + 16768 + ], + "per_rank_resident_file_bytes": [ + 19917, + 19917, + 19917, + 19917 + ], + "distinct_host_count": 4, + "distinct_boot_count": 4, + "model_content_id": "9412c39d7796ccf3d177504ce019a335e9ebfcb0e67e2b991700553dec485f4b", + "layout_id": "54d4ce8c585635f0024d84c8511a3ea86b3ef900485f6104b7d2e86b98369598", + "original_rank_path_available": false, + "topologies": [ + { + "tp": 1, + "pp": 4, + "max_abs_error": 7.450580596923828e-09 + } + ], + "pss_after_map_bytes": [ + 448529408, + 448201728, + 448129024, + 448111616 + ] + }, + { + "case": "tp2pp2-startup180", + "commit": "758f8c866a4de4dc4a8a9057ce14680e78e44b8d", + "backend": "gloo", + "device": "cpu", + "unique_tensor_bytes": 67072, + "aggregate_tensor_bytes": 67584, + "per_rank_tensor_bytes": [ + 16896, + 16896, + 16896, + 16896 + ], + "per_rank_resident_file_bytes": [ + 21785, + 21785, + 21785, + 21785 + ], + "distinct_host_count": 4, + "distinct_boot_count": 4, + "model_content_id": "9412c39d7796ccf3d177504ce019a335e9ebfcb0e67e2b991700553dec485f4b", + "layout_id": "09340703d164c34492330c70b70276f3948de8f1cdf7a82f0e026d941413c98c", + "original_rank_path_available": false, + "topologies": [ + { + "tp": 2, + "pp": 2, + "max_abs_error": 7.450580596923828e-09 + } + ], + "pss_after_map_bytes": [ + 448126976, + 447916032, + 447982592, + 447778816 + ] + }, + { + "case": "switch-owncollection", + "commit": "758f8c866a4de4dc4a8a9057ce14680e78e44b8d", + "backend": "gloo", + "device": "cpu", + "unique_tensor_bytes": 67072, + "aggregate_tensor_bytes": 268288, + "per_rank_tensor_bytes": [ + 67072, + 67072, + 67072, + 67072 + ], + "per_rank_resident_file_bytes": [ + 75521, + 75521, + 75521, + 75521 + ], + "distinct_host_count": 4, + "distinct_boot_count": 4, + "model_content_id": "9412c39d7796ccf3d177504ce019a335e9ebfcb0e67e2b991700553dec485f4b", + "layout_id": "9302c0d3822e6b80e0c75008086fcfd6c79d87ef86f6a04c4dabc585820028fc", + "original_rank_path_available": false, + "topologies": [ + { + "tp": 4, + "pp": 1, + "max_abs_error": 7.450580596923828e-09 + }, + { + "tp": 2, + "pp": 2, + "max_abs_error": 7.450580596923828e-09 + }, + { + "tp": 1, + "pp": 4, + "max_abs_error": 7.450580596923828e-09 + }, + { + "tp": 4, + "pp": 1, + "max_abs_error": 7.450580596923828e-09 + } + ], + "pss_after_map_bytes": [ + 448324608, + 448145408, + 448058368, + 447993856 + ] + }, + { + "case": "gpu-tp4", + "commit": "758f8c866a4de4dc4a8a9057ce14680e78e44b8d", + "backend": "nccl", + "device": "cuda:0", + "unique_tensor_bytes": 67072, + "aggregate_tensor_bytes": 68608, + "per_rank_tensor_bytes": [ + 17152, + 17152, + 17152, + 17152 + ], + "per_rank_resident_file_bytes": [ + 25505, + 25505, + 25505, + 25505 + ], + "distinct_host_count": 4, + "distinct_boot_count": 4, + "model_content_id": "9412c39d7796ccf3d177504ce019a335e9ebfcb0e67e2b991700553dec485f4b", + "layout_id": "8827169242156e4ed206ef26947e99cc1ae82fadb59534fda9d8f4a0a8e43ba8", + "original_rank_path_available": false, + "topologies": [ + { + "tp": 4, + "pp": 1, + "max_abs_error": 7.450580596923828e-09 + } + ], + "pss_after_map_bytes": [ + 2524764160, + 2409260032, + 2405049344, + 2528457728 + ] + }, + { + "case": "gpu-pp4", + "commit": "758f8c866a4de4dc4a8a9057ce14680e78e44b8d", + "backend": "nccl", + "device": "cuda:0", + "unique_tensor_bytes": 67072, + "aggregate_tensor_bytes": 67072, + "per_rank_tensor_bytes": [ + 16768, + 16768, + 16768, + 16768 + ], + "per_rank_resident_file_bytes": [ + 19917, + 19917, + 19917, + 19917 + ], + "distinct_host_count": 4, + "distinct_boot_count": 4, + "model_content_id": "9412c39d7796ccf3d177504ce019a335e9ebfcb0e67e2b991700553dec485f4b", + "layout_id": "54d4ce8c585635f0024d84c8511a3ea86b3ef900485f6104b7d2e86b98369598", + "original_rank_path_available": false, + "topologies": [ + { + "tp": 1, + "pp": 4, + "max_abs_error": 7.450580596923828e-09 + } + ], + "pss_after_map_bytes": [ + 2524720128, + 2528584704, + 2524154880, + 2528459776 + ] + } + ], + "expert_pool": { + "code_commit": "123c6e2dd3f9e69a1df6cd9422d27ba0bd2628f5", + "nodes": 4, + "tests_per_node": 11, + "cpu_tensor_bytes": 524288, + "gpu_weight_budget_bytes": 131072, + "max_abs_error": 1.7881393432617188e-07, + "peak_allocated_bytes": 33700352, + "peak_reserved_bytes": 35651584, + "final_gpu_calls_per_node": [ + 8, + 8, + 8, + 8 + ], + "final_cpu_calls_per_node": [ + 24, + 24, + 24, + 24 + ], + "final_evictions_per_node": [ + 6, + 6, + 6, + 6 + ] + }, + "closure": { + "reservation_rounds": 2, + "all_cleanup_confirmed": true, + "active_worker_count": 0 + }, + "private_raw_evidence_sha256": { + "infiniloader-four-node-verified-cpu.json": "482e663e45fbddaece5d78be997ebea14eb7186203a2ac29d33445ee0697891c", + "infiniloader-four-node-verified-rccl.json": "2ce13dfef202d9d5d009c4ba3c4758c784e09be861ea538943c4267616799783", + "infiniloader-four-node-gpu.json": "def4222f684338159d102c617ee321845abf6a90a6d47117f841a3bc5fdab8e5", + "infiniloader-own-node-results.json": "4e2b0ead848f81350fde98568bccbde88c4254e1873dc95ff509cba25bd3eeea", + "infiniloader-reservation-initial-closed-receipt.json": "c2ae7c7e08a08db158dc5c6b844a0bcedda77aea29de20e779e6d46a06666b1e", + "infiniloader-cluster-cleanup.json": "a6e6297e608a832be52e3f4adc54a73c54806eae7097f60beca5be3635f30a12", + "infiniloader-reservation-round2-closed-receipt.json": "7b368fba5e463a14eebd98015ade04f7653f02f0bcd78c471cfae11cbb8f3af3" + } +} diff --git a/docs/distributed-validation.md b/docs/distributed-validation.md new file mode 100644 index 0000000..b165353 --- /dev/null +++ b/docs/distributed-validation.md @@ -0,0 +1,133 @@ +# Distributed and hybrid residency validation + +> Historical 2026-10-18 measurements retain their original tested commit IDs. +> The reorganized history preserves production code and examples byte-for-byte; +> see [history-reorganization.md](history-reorganization.md) for provenance +> and independently repeated CPU tests. No new four-node/device run is claimed. + +Date: 2026-10-18, Asia/Shanghai. Feature branch: `codex/distributed-residency`, based +on reviewed reliability commit `7442271`. No merge into master was performed. + +## What was implemented and tested + +| User scenario | Implementation and evidence | Remaining boundary | +| --- | --- | --- | +| Each server holds part of the weights | Explicit tensor TP axes and PP stages; streaming safetensors export; four-node TP4, PP4 and TP2×PP2 numerical execution | Minimum **logical tensor payload** for the selected layout, not a proof of global physical RAM optimality. Production engines need semantic adapters. | +| Every server holds all weights for topology changes | Per-node canonical sealed CPU storage; COW tensor views; four-node TP4 → TP2/PP2 → PP4 → TP4 with the original rank filename unavailable | GPU transfers, communications and framework initialization remain; no large-model speedup claim or automatic old-shard redistribution. | +| CPU weights plus GPU hot/CPU cold experts | Grouped, bounded GPU copies; stream-aware leases and LRU; real CPU/GPU ReLU MoE, repeated hot-set changes and failure injection on four DCUs | Sequential reference expert execution; no CPU/GPU overlap claim, distributed EP/MoE, quantized kernel or native KTransformers adapter. | + +The detailed protocol, commands, contracts and source references are in +[distributed.md](distributed.md). Sanitized numerical evidence is provided in +[distributed-results.json](distributed-results.json). + +## Environments and methodology + +- Local: Ubuntu WSL2, Linux 6.18.33.2, Python 3.14.4, PyTorch 2.14.1+cpu, + safetensors 0.8.0 and NumPy 2.5.4. +- Remote: four distinct authorized SCNet workers, different host and boot + identities, Linux containers, Python 3.11.9, PyTorch 2.9.0, HIP 6.3.26093. + Each exposes eight Hygon BW DCUs of approximately 64 GiB; experiments used only + the one device reserved on each worker. Socket-based RCCL was tested; RDMA was not. +- Existing strict SSH host keys, pod UID and image identity checks were retained. + Private keys, addresses, hardware serials and lease tokens are not committed. +- Small deterministic four-layer MLP and eight-expert ReLU MoE fixtures were used, + with independent full CPU reference execution. These are integration and + correctness tests, not performance results for a production LLM. +- Workers report per-rank model/layout identities, original-path availability, + tensor and file bytes, PSS/RSS, topology, runtime and timings. Unique host/boot + identities and complete rank sets were independently checked. + +## Results + +| Four-node case | Cluster CPU tensor payload | Maximum absolute output error | +| --- | ---: | ---: | +| Gloo TP4 | 68,608 B | 7.45e-9 | +| Gloo PP4 | 67,072 B | 7.45e-9 | +| Gloo TP2×PP2 | 67,584 B | 7.45e-9 | +| Full replicas, four successive TP/PP layouts | 268,288 B | 7.45e-9 in every layout | +| RCCL TP4, one DCU per node | 68,608 B | 7.45e-9 | +| RCCL PP4, one DCU per node | 67,072 B | 7.45e-9 | + +Unique MLP payload is 67,072 B. PP4 has no replicated tensor data; TP4 adds 1,536 B +of replicated output biases; TP2×PP2 adds 512 B. The full-replica case is exactly +four times the unique payload. Files, receipts, page cache and runtime memory add +cost beyond this table. Same-node readers share canonical pages; COW writes and +CUDA transfers create additional storage. + +CPU TP4/PP4 used code commit `123c6e2`; CPU TP2×PP2 and replica switching, plus +the static RCCL TP4/PP4 runs, used `758f8c8`. GPU TP2×PP2 and GPU dynamic topology +switching were not validated. The exact per-case commits appear in the JSON. + +The replica run kept each worker and TensorStore alive for all four layouts and +renamed its original rank weight filename to `.offline` after mapping. Each layout +had exactly one terminal output and passed numerical comparison. This proves the +observed switch reused its RAM mappings; it does not mean every disk copy was +deleted or physically inaccessible. + +Four remote devices each passed **11 targeted tests without skips**, including +real GPU expert admission, active-lease protection, capacity rejection, cross-stream +completion and injected transfer failure. Each eight-expert demo retained 512 KiB +on CPU and used a 128 KiB GPU weight pool. Four hot-set phases produced eight GPU +and 24 CPU expert calls, eight promotions and six evictions per node. Maximum +absolute error was **1.79e-7**. PyTorch peak allocated memory was **33,700,352 B** +and peak reserved memory **35,651,584 B**; the 128 KiB weight budget is not a total +application VRAM cap. The demo's CPU source file was unlinked after preload. + +Final local regression: **119 passed, 2 real-device tests skipped** in 41.16 seconds. +Both device tests were exercised in the remote 11-test runs. The built +wheel passed Linux `selftest` and `doctor` independently of the source directory. +All six jobs for code commit `758f8c8` passed in +[GitHub Actions run 39107252516](https://github.com/InfiniTensor/InfiniLoader/actions/runs/39107252516): +Python 3.10–3.14 plus the optional PyTorch CPU integration job. + +## Failures found and addressed + +Real NFS did not support `renameat2(RENAME_NOREPLACE)`. The export path now uses +an atomic exclusive relative symlink to a completed private backing directory on +such filesystems. Independent probes also found publication-after-interrupt and +uncertain `lstat` cases; successful/possibly successful publications are preserved. + +Initial multi-node runs had staggered SSH/device startup exceeding the selected +window. A plan/rank/nonce startup barrier now releases workers after preload; +the final slow-environment runs used a 180-second window. The barrier does not +replace a remote liveness service or make arbitrarily delayed startup succeed. + +Master-only result collection also missed complete files visible on the owning +nodes. Direct per-node retrieval recovered the original outputs, with matching +startup nonces/PID tokens, and each output was independently recomputed. This was +result reconciliation, not replay of an execution with uncertain completion. +Shared metadata absence must not trigger an automatic duplicate start. + +## Independent review and operational closure + +Every code commit below was independently reviewed; identified issues were fixed +in unpublished commits before approval and SSH push. + +| Commit | Change | Independent review evidence | +| --- | --- | --- | +| `d8fe312` | Strict catalog, layout and streaming export | Two experts; source-race/format probes; 100-test isolated suite | +| `4ab6d4a` | Sealed CPU tensor store and live slicing | Two experts; six tensor tests; daemon-death and final-storage-release probe | +| `674f168` | GPU expert pool and hybrid reference | Two experts; mixed precision/lease/stream checks; later real-device tests | +| `717997e` | Actual TP/PP validation and CPU CI | Validation expert; seven isolated integration tests | +| `123c6e2` | NFS atomic publication and uncertain cleanup | Two experts; 25 tests and post-publication interrupt/EIO probes | +| `758f8c8` | Startup coordination after preload | Two experts; eight isolated integration tests and stale-marker probes | + +Two separately identified device reservation rounds were used. All experiment GPU +workers were confirmed exited and every reservation reached closed state with +`cleanup_confirmed=true`. CPU worker cleanup was also checked at completion. No +platform service, host configuration, unrelated process or master branch was changed. +The reservation interface is infrastructure-specific and not part of InfiniLoader. + +## Useful next increments + +1. Add a concrete production engine adapter, including GQA/packed projections, + quantization scales and tied-weight semantics; exercise an actual model. +2. Integrate KTransformers/kt-kernel formats and CPU kernels behind its explicit + ownership contract; profile NUMA placement and conversion/pinning costs. +3. Add asynchronous prefetch and measured overlap, then routing-based hot expert + selection; compare end-to-end latency against native engine offload. +4. Add authenticated shard transfer, compact per-rank manifests, batched catalog + preparation and old-topology-to-new-topology redistribution with failure recovery. +5. Extend larger-model, multi-GPU, long-running stress and controlled physical + memory measurements. Native NVIDIA CUDA, RDMA and distributed expert parallelism + remain unvalidated by the DCU/socket tests here. diff --git a/docs/distributed.md b/docs/distributed.md new file mode 100644 index 0000000..fad02f6 --- /dev/null +++ b/docs/distributed.md @@ -0,0 +1,202 @@ +# Distributed CPU residency and hybrid expert integration + +InfiniLoader now provides explicit safetensors partitioning, a PyTorch CPU tensor +store, and a bounded GPU expert pool. The reference examples execute real TP/PP +and a dense ReLU MoE. Framework-specific Llama/GQA, quantized kernels and +KTransformers integration remain separate adapter work. + +## When to use each mode + +All three modes address real workloads: + +| Mode | Suitable workload | Capacity and tradeoff | +| --- | --- | --- | +| Sharded CPU weights | Fixed TP/PP deployment; models larger than one node's RAM | Cluster payload `U + R`: unique tensor bytes plus execution-required replicas. PP can achieve `R=0`; TP often replicates some parameters. A new topology may require preparing/transferring new shards. | +| Full CPU replica per node | Repeated debugging, worker restart, frequent TP/PP changes | `N × U` for N nodes, with readers on a node sharing its canonical store. Slice existing CPU mappings for new layouts; GPU copies, collectives and engine initialization still cost time. | +| CPU canonical store plus GPU hot experts | Sparse MoE inference with limited VRAM | CPU retains the full model, GPU retains selected complete expert groups. Cold experts can compute on CPU. Eviction drops the GPU copy; immutable weights need no D2H writeback. | + +`U` counts actual validated checkpoint tensor payload, including separately stored +quantization tensors and duplicate stored parameters. It does not deduplicate +semantic tied weights. `R` is dictated by the explicit execution layout. These are +logical capacities, not guarantees of minimum physical system RAM: file headers, +page cache, Python, COW pages, activations, staging and allocator state are additional. +Use PSS plus cgroup/system measurements rather than summing shared-page RSS. + +The demand is independently supported by [vLLM's rank-specific sharded loader](https://docs.vllm.ai/en/latest/api/vllm/model_executor/model_loader/sharded_state_loader/), +[PyTorch distributed checkpoint resharding](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html), +and [KTransformers' CPU/GPU heterogeneous expert placement](https://github.com/kvcache-ai/ktransformers). +The policy choices above are our interpretation of those workloads, not a claim +that these engines consume the new InfiniLoader shard format directly. + +## Explicit checkpoint preparation + +The legacy `dist-plan` remains a heuristic report. The new executable path requires +an exact tensor inventory and explicit TP dimension/PP stage for **every** tensor: + +```json +{ + "format": 1, "tp": 2, "pp": 1, + "tensors": { + "up.weight": {"stage": 0, "split_dim": 0}, + "down.weight": {"stage": 0, "split_dim": 1}, + "down.bias": {"stage": 0, "split_dim": null} + } +} +``` + +```bash +iloader shard plan /models/source --layout layout.json --mode sharded --output plan.json +iloader shard export /models/source --plan plan.json --rank 0 --output /cache/rank0 +iloader shard verify /cache/rank0 +``` + +`--mode replicated` exports a complete canonical checkpoint for each node. A plain +full safetensors model can also be preloaded directly without exporting replicas. +Use one daemon per node. Each reader opens its own TensorStore over that daemon's +canonical memfds, sharing the underlying CPU pages across processes. A daemon per +local rank unnecessarily duplicates the canonical CPU allocation. + +Plans carry path-independent SHA-256 content identity and a separate layout ID. +Catalogs reject duplicate tensors/JSON keys, invalid offsets, shapes, dtypes and +index mismatches. Nonzero-axis slicing streams the appropriate row intervals. +Preparation performs multiple full source hashes to reject concurrent changes; +this is an offline cost, **not** a claim that preparation reads only rank bytes. +Export uses bounded chunks; the destination contains only that rank's tensors. + +On filesystems without Linux `RENAME_NOREPLACE`, publication atomically creates a +relative rank symlink to a completed private `.iloader-export-*` backing directory. +Keep the link and its backing together. To copy a rank independently, dereference +the link (for example `cp -aL rank0 copy-of-rank0`). If publication is interrupted, +verify an existing destination before retrying; existing output is never replaced. +An uncertain NFS lookup preserves possibly published data for reconciliation. + +The output is an **InfiniLoader explicit-layout checkpoint**. Being valid +safetensors does not make it a native HF/vLLM/InfiniLM checkpoint. A model adapter +must implement the specified collectives, packed projections, GQA, biases and +quantization semantics. No name heuristic silently decides these semantics. + +## Use the canonical CPU store + +Install a PyTorch build appropriate for the hardware, then install this project. +Core commands still have no PyTorch dependency. + +```bash +iloader start /cache/rank0 --name rank0 --workers 1 --chunk-size 1048576 +iloader run rank0 -- python adapter.py '{model}' +``` + +```python +from pathlib import Path +import sys +from infiniloader.shards import read_json +from infiniloader.torch_store import TensorStore + +resident_view = sys.argv[1] +layout = read_json(Path("plan.json"))["layout"] +with TensorStore(resident_view) as store: + tensors = store.rank_tensors(layout, rank=0) + # Pass CPU tensors to an adapter; .to(device) transfers only required slices. +``` + +The store accepts sealed InfiniLoader memfds, validates their contents and maps +them privately with COW. Reads share canonical pages; in-place writes consume +private pages without modifying the daemon. `torch.frombuffer` retains the mapping +through tensor storage lifetime. Closing a store prevents new reads but does not +invalidate already returned tensors/views. Completed mappings also survive daemon +exit; new opens through a dead view do not. Such client-held mappings can outlive +the daemon's logical reservation and continue using physical RAM. + +Use `iloader run` to protect initialization from daemon idle expiry/stop. The +TensorStore path API itself does not automatically acquire a daemon lease. + +A replicated store can call `rank_tensors` repeatedly with different valid +layouts without rereading the original source or rebuilding CPU weights. A +pre-sharded store is bound to its original layout and rank. Redistribution from +an already sharded old topology to another sharded topology is not implemented. + +## CPU/GPU expert integration + +```python +from infiniloader.experts import ExpertPool + +groups = {"0": {"up": "experts.0.up", "down": "experts.0.down"}} +with ExpertPool(store, groups, max_gpu_bytes=256 * 1024**2, device="cuda:0") as pool: + pool.promote("0") # Whole-group synchronous upload; False means CPU fallback. + with pool.acquire("0") as lease: + output = engine_expert_kernel(inputs.to(lease.device), lease.weights) +``` + +This API is for **inference**. Do not retain weights or submit work after a lease +ends. Autograd that saves weights beyond the lease is outside the contract. GPU +weights default to the current consumer stream; use `acquire(..., stream=stream)` +and `lease.record_stream(other_stream)` for every additional consumer stream. +Before eviction, the pool waits for recorded work and protects allocator reuse. +Active leases cannot be evicted. Failed uploads never publish half an expert. + +The budget counts pool-owned weight payload, including admission in progress; it +does not cap activations, external tensor references, CUDA context or allocator +reserved memory. Record both weight bytes and PyTorch allocated/reserved peaks. +Promotions are currently synchronous; asynchronous prefetch and routing prediction +are future work. The caller selects hot experts, while the pool supplies admission, +CPU fallback and lease-safe LRU replacement. + +`hybrid_moe` is a reference dense ReLU execution path, with top-k routing and +mixed CPU/GPU output aggregation. Native KTransformers/kt-kernel uses specialized +formats and kernels; its conversion/NUMA/kernel adapter has not been implemented. +Pinning, quantization conversion and CPU→GPU transfer are not zero-copy operations. + +## Reproduce numerical validation + +Install `pytest`, `safetensors` and `numpy` beside the correct PyTorch build. + +```bash +python -m pytest -q +python examples/distributed_mlp.py local --root /tmp/tp4-new --mode sharded --tp 4 --pp 1 --output tp4.json +python examples/distributed_mlp.py local --root /tmp/switch-new --mode replicated --tp 4 --pp 1 --output switch.json +python examples/hybrid_moe.py --device cuda:0 --output hybrid.json +``` + +The MLP fixture uses true all-reduce and point-to-point activation transfers, +including correct post-reduction bias addition. Replica switching removes the +source filename after preload and requires every topology to produce exactly one +numerically checked terminal output. The fixture is intentionally small and does +not establish large-model speedup or optimal physical RAM. + +For multiple hosts, prepare on shared storage or transfer each rank checkpoint, +then invoke one `worker` per host with the same rendezvous and different rank: + +```bash +python examples/distributed_mlp.py prepare --root /shared/new-run --mode sharded --tp 2 --pp 2 +mkdir -m 700 /shared/new-run/startup +# Run separately on hosts 0..3, substituting RANK and a reachable rendezvous address. +python examples/distributed_mlp.py worker --root /shared/new-run --rank RANK --world-size 4 \ + --init tcp://HOST0:PORT --backend gloo --device cpu --output /shared/new-run/resultRANK.json \ + --ready-dir /shared/new-run/startup --timeout 180 +# After launching all workers, release them from one trusted controller: +PYTHONPATH=src:examples python - <<'PY' +from pathlib import Path +from distributed_mlp import publish_start +from infiniloader.shards import read_json +root = Path('/shared/new-run') +plan = read_json(root / 'fixture.json')['plan'] +publish_start(root / 'startup', 4, plan['plan_id'], 180) +PY +``` + +Use a fresh private barrier directory and a single controller. The participant +nonces prevent joining an old release; the recorded PID/token is diagnostic and +does not authenticate remote liveness. The startup and collective waits are both +bounded by `--timeout`; choose a window that covers slow SSH and device warmup. +Collect results from each owning node. In the tested NFS environment, the master +temporarily missed records that were complete on the generating node; absence +from a shared-directory lookup did not establish that execution had failed. + +For GPU communication use the device/framework's `nccl` backend (RCCL under ROCm) +and an assigned device such as `cuda:0`. Memfd views must remain node-local; shared +NFS paths do not make another host's `/proc/PID/fd` valid. External launchers must +verify host/checkpoint identity, allocate devices, bound execution and collect all +rank results. The example is a validation harness, not a fault-tolerant cluster +controller, authenticated tensor-transfer service or inference serving framework. + +See [distributed validation](distributed-validation.md) for the tested hardware, +results, independent review and remaining limitations. diff --git a/docs/history-reorganization.json b/docs/history-reorganization.json new file mode 100644 index 0000000..e8f2661 --- /dev/null +++ b/docs/history-reorganization.json @@ -0,0 +1,363 @@ +{ + "date": "2026-10-26", + "branch": "improve-weight-residency", + "base": "5014a6a3459f0af4629097b92bddf9d8c14e97fd", + "archived_source_head": "88d6bac6a80a424201da1688989d1d87e3217a93", + "implementation_checkpoint": "895c5efce64030dacc635eb26331d1aa606417b2", + "author": "whjthu ", + "committer": "whjthu ", + "implementation_git_objects": { + "src": { + "original_git_object": "75d6fb1fa850c4869ab2cbd1d9926d5db909e7c2", + "reorganized_git_object": "75d6fb1fa850c4869ab2cbd1d9926d5db909e7c2", + "equal": true + }, + "tests": { + "original_git_object": "e7766d1d8ecdd158659aa6b5e36df6b52b149476", + "reorganized_git_object": "f1e98e2fe02aa9ff4af3b1c34d95fd0fa445b0e7", + "equal": false, + "only_difference": "Corrupt-rank test startup timeout increased from 3 to 20 seconds; all failure/reaping assertions preserved." + }, + "examples": { + "original_git_object": "f5317c389c7db9019e9054a9772cbab070d2aaaa", + "reorganized_git_object": "f5317c389c7db9019e9054a9772cbab070d2aaaa", + "equal": true + }, + ".github": { + "original_git_object": "1055429eb9eaecee7c918f3ec13b346992b5c002", + "reorganized_git_object": "1055429eb9eaecee7c918f3ec13b346992b5c002", + "equal": true + }, + "deploy": { + "original_git_object": "4108c1d2d31a0311432a6bbcbaffc67b92cc941b", + "reorganized_git_object": "4108c1d2d31a0311432a6bbcbaffc67b92cc941b", + "equal": true + }, + "pyproject.toml": { + "original_git_object": "6813f8d38b81f6076e28518d5850c3f6c43bfe11", + "reorganized_git_object": "6813f8d38b81f6076e28518d5850c3f6c43bfe11", + "equal": true + } + }, + "independent_per_commit_tests": [ + { + "step": 1, + "commit": "ea57e3c0c0e25c4e6fe378569dd949252f01b0b0", + "subject": "Protect sealed weight snapshots and daemon runtime ownership", + "author_name": "whjthu", + "author_email": "haojie0429@gmail.com", + "committer_name": "whjthu", + "committer_email": "haojie0429@gmail.com", + "identity_ok": true, + "single_english_subject_ok": true, + "exit_code": 0, + "elapsed_seconds": 9.4, + "log": "01-ea57e3c-pytest.txt", + "summary": "40 passed in 7.09s", + "raw_pytest_sha256": "b874aa1a71328c3ff5e072077c10885b30852792437568bdf2ff73a5d9fea7ad" + }, + { + "step": 2, + "commit": "63e7ea42ac5dcef3f8384a455f859510015f704b", + "subject": "Bound IPC requests and cancel failed distributed reads", + "author_name": "whjthu", + "author_email": "haojie0429@gmail.com", + "committer_name": "whjthu", + "committer_email": "haojie0429@gmail.com", + "identity_ok": true, + "single_english_subject_ok": true, + "exit_code": 0, + "elapsed_seconds": 14.57, + "log": "02-63e7ea4-pytest.txt", + "summary": "46 passed in 11.71s", + "raw_pytest_sha256": "6759f5949b4abcff97d8c4fc5e9b7f400a559c66e8298a91dd67f2d09743cf6d" + }, + { + "step": 3, + "commit": "97335f6f718dc8eea9724bb90908f874ab3cd828", + "subject": "Publish verified disk cache generations with lease-aware collection", + "author_name": "whjthu", + "author_email": "haojie0429@gmail.com", + "committer_name": "whjthu", + "committer_email": "haojie0429@gmail.com", + "identity_ok": true, + "single_english_subject_ok": true, + "exit_code": 0, + "elapsed_seconds": 16.45, + "log": "03-97335f6-pytest.txt", + "summary": "62 passed in 12.15s", + "raw_pytest_sha256": "cb594ab2b711f786748d039fcfbcaf3ce47fd77f019e3d15648893f69e9c8299" + }, + { + "step": 4, + "commit": "8cca2e6fdc7f5c5d15075ecb20519073502f01eb", + "subject": "Enforce container memory limits and shared resident budgets", + "author_name": "whjthu", + "author_email": "haojie0429@gmail.com", + "committer_name": "whjthu", + "committer_email": "haojie0429@gmail.com", + "identity_ok": true, + "single_english_subject_ok": true, + "exit_code": 0, + "elapsed_seconds": 15.97, + "log": "04-8cca2e6-pytest.txt", + "summary": "69 passed in 12.73s", + "raw_pytest_sha256": "e1d9f33f1dad63eb375860e97c22505c697d400d885fc1345dec0fd8f7ec3570" + }, + { + "step": 5, + "commit": "e2c20659880966071aded2a41905adbe8f5832d8", + "subject": "Add leased commands and cache lifecycle diagnostics", + "author_name": "whjthu", + "author_email": "haojie0429@gmail.com", + "committer_name": "whjthu", + "committer_email": "haojie0429@gmail.com", + "identity_ok": true, + "single_english_subject_ok": true, + "exit_code": 0, + "elapsed_seconds": 33.62, + "log": "05-e2c2065-pytest.txt", + "summary": "77 passed in 25.08s", + "raw_pytest_sha256": "be4e1d73ed316bcae8d74340ba31710b809114a221b1e2a28bd938c9b38bf551" + }, + { + "step": 6, + "commit": "7454e9cb73f500ec37dd06ad2df88849b1d2c12f", + "subject": "Document cache operations and add Linux regression CI", + "author_name": "whjthu", + "author_email": "haojie0429@gmail.com", + "committer_name": "whjthu", + "committer_email": "haojie0429@gmail.com", + "identity_ok": true, + "single_english_subject_ok": true, + "exit_code": 0, + "elapsed_seconds": 28.95, + "log": "06-7454e9c-pytest.txt", + "summary": "77 passed in 23.76s", + "raw_pytest_sha256": "fa9f44d93921cba94a2aecc120ef11fb5813ac2ed3b5a4b4fda4371738063f15" + }, + { + "step": 7, + "commit": "48615de1981afb1ff17475423227813932112e43", + "subject": "Add explicit tensor sharding with atomic streaming export", + "author_name": "whjthu", + "author_email": "haojie0429@gmail.com", + "committer_name": "whjthu", + "committer_email": "haojie0429@gmail.com", + "identity_ok": true, + "single_english_subject_ok": true, + "exit_code": 0, + "elapsed_seconds": 30.45, + "log": "07-48615de-pytest.txt", + "summary": "102 passed in 27.24s", + "raw_pytest_sha256": "d73753398c5d1dbc093e837b2bf0cf7d6a9597e272876838c2f1fe9a3fe34ecd" + }, + { + "step": 8, + "commit": "43d7a75dead6bcb4f7a82cc8734ef4622ce13333", + "subject": "Share sealed CPU tensors with live topology slicing", + "author_name": "whjthu", + "author_email": "haojie0429@gmail.com", + "committer_name": "whjthu", + "committer_email": "haojie0429@gmail.com", + "identity_ok": true, + "single_english_subject_ok": true, + "exit_code": 0, + "elapsed_seconds": 40.33, + "log": "08-43d7a75-pytest.txt", + "summary": "108 passed in 31.83s", + "raw_pytest_sha256": "b437bfedec5274814745face31df619df1e401776d8bc29e31f6e6cc814b92b2" + }, + { + "step": 9, + "commit": "70c0979fd37c0a9a085a63312b3c2f3e47515bc9", + "subject": "Validate distributed TP and PP execution after coordinated preload", + "author_name": "whjthu", + "author_email": "haojie0429@gmail.com", + "committer_name": "whjthu", + "committer_email": "haojie0429@gmail.com", + "identity_ok": true, + "single_english_subject_ok": true, + "exit_code": 0, + "elapsed_seconds": 84.46, + "log": "09-70c0979-pytest.txt", + "summary": "116 passed in 78.24s (0:01:18)", + "raw_pytest_sha256": "a725e5a9dc829b1c817fc60240fc335f4074d53b48c5bf1d8f3a257f659d524a" + }, + { + "step": 10, + "commit": "895c5efce64030dacc635eb26331d1aa606417b2", + "subject": "Manage GPU expert residency with CPU fallback and safe eviction", + "author_name": "whjthu", + "author_email": "haojie0429@gmail.com", + "committer_name": "whjthu", + "committer_email": "haojie0429@gmail.com", + "identity_ok": true, + "single_english_subject_ok": true, + "exit_code": 0, + "elapsed_seconds": 85.2, + "log": "10-895c5ef-pytest.txt", + "summary": "119 passed, 2 skipped in 80.49s (0:01:20)", + "raw_pytest_sha256": "13b46cb1a09aa63c5df2558ff6e45b3378869e3ced4f97a7ee5d8eb5fd556d7d" + } + ], + "original_to_new_steps": [ + { + "original_commit": "2f7204be58a2c53d43117a8c1e5d90bba14c0125", + "new_steps": [ + 1 + ], + "reason": "Startup serialization and runtime ownership" + }, + { + "original_commit": "c67ba9eb9048f44d00f98b86cd8abdacf1c0446d", + "new_steps": [ + 1 + ], + "reason": "Sealed immutable snapshots" + }, + { + "original_commit": "deaa940d2178f0cb7d2211061a156f88068b9b71", + "new_steps": [ + 2 + ], + "reason": "Bounded IPC and failed-rank cancellation" + }, + { + "original_commit": "d7261009fe6da8bf1c09c9a813ceb43ce7c230dd", + "new_steps": [ + 1, + 3 + ], + "reason": "Identity rechecks moved to startup; verified disk generations" + }, + { + "original_commit": "2f544804c91667c9af2d26e4d10b1260fdcd4de2", + "new_steps": [ + 4 + ], + "reason": "Cgroup and shared resident budgets" + }, + { + "original_commit": "fb2f0570b8ac62a83667534dbfc1030b3a54b96f", + "new_steps": [ + 5 + ], + "reason": "Leased execution and diagnostics" + }, + { + "original_commit": "8fab6cd5ff4d0b0fe640403c8d583c4c3e012f40", + "new_steps": [ + 1, + 2, + 3, + 4, + 5 + ], + "reason": "Boundary, stop, GC and budget corrections folded into owners" + }, + { + "original_commit": "a17a6e91f42cee30e6418a31a212a6143413501c", + "new_steps": [ + 5 + ], + "reason": "Benchmark medians and cache uncertainty" + }, + { + "original_commit": "6a33d83e2b6834302c63661e9c44e7099e2b0b56", + "new_steps": [ + 6 + ], + "reason": "Operations and core CI" + }, + { + "original_commit": "7e1b2485b7608ef4df8100aa2eb1384134757a81", + "new_steps": [ + 1, + 2, + 5 + ], + "reason": "Runtime ownership, cancellation diagnostics and lease-safe stop" + }, + { + "original_commit": "cc67790f7b4d65b20cbeef051f5ee9fc6189088c", + "new_steps": [ + 1, + 3 + ], + "reason": "Bounded copy buffers and complete disk generation verification" + }, + { + "original_commit": "7442271c0698a9e50ebd07c00feb1dad95f7e198", + "new_steps": [ + 5, + 6 + ], + "reason": "Regression alignment and historical expert evidence" + }, + { + "original_commit": "d8fe312a58ae612d6bfa25a3bf5ce76ddf0cc898", + "new_steps": [ + 7 + ], + "reason": "Strict explicit sharding" + }, + { + "original_commit": "4ab6d4a067ae41e764264162928498006f926d2b", + "new_steps": [ + 8 + ], + "reason": "CPU TensorStore" + }, + { + "original_commit": "674f16871907f489d3fa0ad4e12650874d330ed0", + "new_steps": [ + 10 + ], + "reason": "GPU expert pool and hybrid reference" + }, + { + "original_commit": "717997e5fa108f57f9efc646ad79da444fc52ae3", + "new_steps": [ + 9, + 10 + ], + "reason": "TP/PP tests; hybrid CI smoke kept with its example" + }, + { + "original_commit": "123c6e2dd3f9e69a1df6cd9422d27ba0bd2628f5", + "new_steps": [ + 7 + ], + "reason": "NFS publication fix included in initial shard feature" + }, + { + "original_commit": "758f8c866a4de4dc4a8a9057ce14680e78e44b8d", + "new_steps": [ + 9 + ], + "reason": "Startup barrier included in initial TP/PP harness" + }, + { + "original_commit": "88d6bac6a80a424201da1688989d1d87e3217a93", + "new_steps": [ + 11 + ], + "reason": "Distributed contracts, historical numerical evidence and README" + } + ], + "documentation_step": 11, + "documentation_commit": "The commit containing this file; inspect git log to avoid a self-referential hash.", + "historical_remote_tests": "Original measured commit IDs are preserved. No new GPU or four-node execution is claimed for this history rewrite.", + "independent_test_command": [ + "python3", + "-B", + "-m", + "pytest", + "-q" + ], + "independent_test_environment": { + "PYTEST_DISABLE_PLUGIN_AUTOLOAD": "1", + "PYTHONDONTWRITEBYTECODE": "1" + } +} diff --git a/docs/history-reorganization.md b/docs/history-reorganization.md new file mode 100644 index 0000000..0e3acdb --- /dev/null +++ b/docs/history-reorganization.md @@ -0,0 +1,90 @@ +# 提交线重整与验证 + +日期：2026-10-26（Asia/Shanghai）。 + +将 master `5014a6a` 之后的 19 个工作提交重整为 11 个主题提交。原远程备份 +`codex/all-changes-backup-20261020` 及原功能分支保留；不改写其历史，不合并 master。 +新分支为 `improve-weight-residency`。新提交均使用一句英文 subject，Author 和 +Committer 均为 `whjthu `。 + +## 新提交顺序与独立验证 + +测试专家对每个精确 hash 单独执行 git archive，在 Linux 原生临时目录运行完整 +`python3 -B -m pytest -q`，未依赖变化中的工作树。设置 +`PYTEST_DISABLE_PLUGIN_AUTOLOAD=1`、`PYTHONDONTWRITEBYTECODE=1`。环境为 WSL Ubuntu、 +Python 3.14.4、PyTorch 2.14.1+cpu、NumPy 2.5.4、safetensors 0.8.0、pytest 9.1.1。 + +| 顺序 | 提交 | 一句话 subject | 独立完整测试原始摘要 | +| --- | --- | --- | --- | +| 1 | `ea57e3c` | Protect sealed weight snapshots and daemon runtime ownership | 40 passed in 7.09s | +| 2 | `63e7ea4` | Bound IPC requests and cancel failed distributed reads | 46 passed in 11.71s | +| 3 | `97335f6` | Publish verified disk cache generations with lease-aware collection | 62 passed in 12.15s | +| 4 | `8cca2e6` | Enforce container memory limits and shared resident budgets | 69 passed in 12.73s | +| 5 | `e2c2065` | Add leased commands and cache lifecycle diagnostics | 77 passed in 25.08s | +| 6 | `7454e9c` | Document cache operations and add Linux regression CI | 77 passed in 23.76s | +| 7 | `48615de` | Add explicit tensor sharding with atomic streaming export | 102 passed in 27.24s | +| 8 | `43d7a75` | Share sealed CPU tensors with live topology slicing | 108 passed in 31.83s | +| 9 | `70c0979` | Validate distributed TP and PP execution after coordinated preload | 116 passed in 78.24s (0:01:18) | +| 10 | `895c5ef` | Manage GPU expert residency with CPU fallback and safe eviction | 119 passed, 2 skipped in 80.49s (0:01:20) | +| 11 | 本文件所在提交 | Document validated history and development priorities | 文档与来源映射；最终 hash 的审核结果随交付报告给出 | + +逐提交元数据、原始 pytest 日志 SHA-256 及完整映射见 +[history-reorganization.json](history-reorganization.json)。原始日志保存在本次任务的 +`InfiniLoader-commit-review-20261026` 证据目录；该本地目录不是公开下载链接。 + +生命周期专家审核了 1–5；缓存/内存专家审核了 1、3、4、7、8、10；CLI/验证专家审核 +提交依赖、元数据、逐提交完整测试及 6、9 的 CI/验证流程。每个提交都须完成对应独立 +审核后才推送。最终文档提交还需独立终审，此处不预先记为已批准。 + +## 为什么这样重整 + +- 把运行文件所有权、PID/state 原子更新和 inode 清理放在启动主题中；把 prepare 后与 + daemon load 后的身份复核一并提前，避免中间提交加载新内容却发布旧身份。 +- 把单 chunk 内存约束纳入快照拷贝，把完整 inventory/schema 和真实容量回收纳入磁盘代际。 +- 分片功能首次出现即具备 NFS 发布与中断/EIO 保护；TP/PP 示例首次出现即带启动屏障。 +- TensorStore 先于专家池；CI 的 hybrid 示例步骤随实际示例一起引入。 + +## 实现等价与历史证据 + +第 10 提交 `895c5ef` 的 `src/`、`examples/`、`.github/`、 +`deploy/` 和 `pyproject.toml` Git 对象与原终版 `88d6bac` 完全相同。最终提交仅调整文档。 +测试唯一差异是故障注入用例的启动等待窗由 3 秒改为 20 秒；源位于 Windows 挂载目录时 +首次完整回归在 worker 报告损坏前已耗尽 3 秒，因此出现与预期错误类型不同的 barrier +超时。错误类型、损坏检测和全部子进程回收断言均保留；修正随首次引入 TP/PP 测试提交。 +可独立复核： + +```bash +git diff --exit-code 88d6bac improve-weight-residency -- src examples .github deploy pyproject.toml +git diff 88d6bac improve-weight-residency -- tests +``` + +原四机 CPU/GPU 测量保留其真实代码版本（例如 `123c6e2`、`758f8c8`），不改写为新 hash。 +源码相同说明实现被保留，不代表本次重新执行了 GPU、RCCL 或四机实验。两个真实设备 +测试在本地 CPU 环境跳过，原设备验证见 [distributed-validation.md](distributed-validation.md)。 +当前分支的远端 CI 应查看该分支自己的 Actions；原报告中的旧 Actions 链接仍只证明旧提交。 + +## 原提交归属 + +| 原提交 | 新主题顺序 | 归属说明 | +| --- | --- | --- | +| `2f7204b` | 1 | Startup serialization and runtime ownership | +| `c67ba9e` | 1 | Sealed immutable snapshots | +| `deaa940` | 2 | Bounded IPC and failed-rank cancellation | +| `d726100` | 1, 3 | Identity rechecks moved to startup; verified disk generations | +| `2f54480` | 4 | Cgroup and shared resident budgets | +| `fb2f057` | 5 | Leased execution and diagnostics | +| `8fab6cd` | 1, 2, 3, 4, 5 | Boundary, stop, GC and budget corrections folded into owners | +| `a17a6e9` | 5 | Benchmark medians and cache uncertainty | +| `6a33d83` | 6 | Operations and core CI | +| `7e1b248` | 1, 2, 5 | Runtime ownership, cancellation diagnostics and lease-safe stop | +| `cc67790` | 1, 3 | Bounded copy buffers and complete disk generation verification | +| `7442271` | 5, 6 | Regression alignment and historical expert evidence | +| `d8fe312` | 7 | Strict explicit sharding | +| `4ab6d4a` | 8 | CPU TensorStore | +| `674f168` | 10 | GPU expert pool and hybrid reference | +| `717997e` | 9, 10 | TP/PP tests; hybrid CI smoke kept with its example | +| `123c6e2` | 7 | NFS publication fix included in initial shard feature | +| `758f8c8` | 9 | Startup barrier included in initial TP/PP harness | +| `88d6bac` | 11 | Distributed contracts, historical numerical evidence and README | + +后续能力优先级与验收条件见 [roadmap.zh-CN.md](roadmap.zh-CN.md)。 diff --git a/docs/roadmap.zh-CN.md b/docs/roadmap.zh-CN.md new file mode 100644 index 0000000..973787e --- /dev/null +++ b/docs/roadmap.zh-CN.md @@ -0,0 +1,47 @@ +# InfiniLoader 后续能力路线 + +评估日期：2026-10-26。以下是基于当前实现和官方上游接口的开发建议，尚未实现，也不是性能承诺。 + +## 建议定位 + +InfiniLoader 适合继续发展成供推理引擎复用的权重存储与驻留层：管理内容身份、不可变代际、进程间共享、布局和设备副本生命周期。计算图、算子、KV cache 和请求调度仍由推理引擎负责。下一轮最有价值的交付是一个真实模型适配器和可信的端到端性能证据。 + +当前已经具备 sealed memfd、租约、磁盘代际、cgroup 内存预算、显式 TP/PP 切片、共享 CPU TensorStore 和同步 GPU ExpertPool。四机测试证明了小型 MLP 的 TP/PP 数值路径及参考 MoE 换入换出；它们没有验证生产 LLM、量化引擎兼容性、GPU 动态拓扑切换或分布式专家并行。 + +## 优先级与验收条件 + +| 优先级 | 能力 | 当前缺口与具体交付 | 验收依据 | +| --- | --- | --- | --- | +| P0 | 真实模型与引擎适配 | 固定一个模型族和引擎版本，从 PyTorch/Transformers 接入开始，再扩展 vLLM/InfiniLM；显式处理 GQA、packed QKV、bias、tied weights、dtype 和非连续 slice | 与原生加载比较 logits/生成结果；验证 TP/PP；重启和重复加载；实际模型的冷启动、热启动与内存结果 | +| P0 | 端到端性能与物理内存观测 | 将读盘、hash、映射、参数安装、转换、H2D、引擎初始化和首 token 分开计时；增加共享 PSS、cgroup、GPU allocated/reserved 与 NUMA 观测 | 同模型、同设备、同软件版本下比较原生 mmap、热 page cache、InfiniLoader；报告重复样本及 p50/p95，不用读文件带宽代替启动收益 | +| P1 | 批量 catalog、导出和复用 | 当前 export_rank 在导出前后分别 catalog 全量源；TensorStore 也会校验全量内容。设计一次 snapshot/catalog、多 rank 有界导出和代际级校验复用 | 记录总读字节与 hash 字节；rank 数增加时避免重复全量扫描；源变化、进程退出、损坏和重试测试仍通过 | +| P1 | 引擎原生格式与 KTransformers 适配 | 按内容 ID、格式、转换参数和内核版本缓存派生产物；提供原生 tensor/格式适配，明确 C++ 内核是否复制、持有还是借用存储 | 一个真实 MoE 端到端运行；精度检查、格式失效和重启复用；同时报告 CPU canonical、转换副本、GPU 权重与总显存 | +| P1 | 有界异步预取与冷热策略 | 当前 promote 同步，hot set 由调用方指定，参考执行按专家串行。先增加有界 staging、异步句柄、取消/失败状态，再对接路由统计和原生 CPU 内核 | 相同质量和显存限制下测 TTFT、每 token 延迟、命中率、H2D 字节、等待和抖动；证明计算/传输实际重叠，不只统计 promotion 次数 | +| P1 | 节点级服务与稳定客户端协议 | 提供版本化 acquire/release/describe 接口、多模型幂等 warmup、可查询操作状态、指标和能力探测；评估 FD 传递以减少对跨容器 /proc 路径的依赖 | 同节点多进程共享；不同 PID namespace 的可支持配置测试；daemon 崩溃、客户端存活映射及 COW 对租约和配额的影响明确 | +| P2 | 分布式分片迁移与重分布 | 现有 sharded store 绑定原 layout/rank；完整副本可本地重切片，但不提供旧分片到新分片的在线迁移。先做可信、可核验、有界传输，再考虑 RDMA | TP4→TP2/PP2→PP4 的真实模型迁移；校验所需 shard 完整后才发布；中断可恢复，重复请求不重复启动，不足内存时拒绝；GPU 拓扑与通信单独实测 | + +## 先解决真实加载链路中的复制 + +共享 CPU Tensor 并不意味着引擎会直接复用该 storage。常规 load_state_dict 可能复制到已分配参数，量化、packing 和 pinning 也会产生新存储。适配器应先测量这些步骤，明确参数及 storage 的所有权，再选择 meta 初始化和参数绑定方式。PyTorch 的官方加载说明展示了 mmap、meta 与 assign=True 的组合；这可作为适配切入点，并不代表所有模型和版本都能直接套用。[PyTorch 加载指南](https://docs.pytorch.org/tutorials/recipes/recipes/module_load_state_dict_tips.html) + +vLLM 已有按 worker 加载自身 TP/PP state 的原生 loader，因此接入要实现其期望的参数语义和存储格式。InfiniLoader 的显式 layout 并不是该格式的直接替代品。[vLLM ShardedStateLoader](https://docs.vllm.ai/en/latest/api/vllm/model_executor/model_loader/sharded_state_loader/) + +## 与 KTransformers 的合作边界 + +KT-Kernel 已有原生格式、CPU 内核、NUMA 配置、GPU 专家选择和动态放置相关接口，也公开了从 tensor 加载权重的入口。应优先适配和测量这些入口，让 InfiniLoader 提供可复用的权重及派生产物管理；不要将另一个参考 LRU/MoE 调度器视为主要产品收益。接入后仍需验证原生内核的布局、复制和异步使用契约，不能由 PyTorch 视图的共享性质推断原生内核零复制。[KT-Kernel 官方说明](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/README.md) + +## 分布式模式如何演进 + +固定布局且 RAM 紧张时优先分片；调试和频繁拓扑变更时完整 CPU 副本更方便。后续可以依据内存预算、预计切换频率与迁移成本选择模式，记录选择理由，但不应在缺少带宽和工作负载测量时宣称自动最优。 + +拓扑重分布需要参数语义、数据移动、协调和故障恢复共同工作。PyTorch Distributed Checkpoint 可作为状态分片与恢复接口的参考，不能把它的状态恢复能力等同于 InfiniLoader 已支持在线推理拓扑迁移。[PyTorch Distributed Checkpoint](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html) + +实验中的共享目录可见性问题还说明，节点上的权威操作状态和可查询结果比仅依赖目录 marker 更可靠。可继续接入既有算力调度平台的资源分配和租约，把权重准备作为作业阶段；InfiniLoader 保持权重层职责，不复制建设整套集群调度系统。 + +## 推荐实施顺序 + +1. **真实模型接入与基线。** 一个稠密模型、一个引擎、固定版本，完成结果正确性和原生/缓存加载对比，同时补齐 storage 复制与内存观测，并给出可复现安装、版本兼容和迁移说明。 +2. **减少重复准备与原生 MoE 接入。** 批量 catalog/export、派生格式缓存，再接一个 KT-Kernel 原生路径，量化收益与内存成本。 +3. **按瓶颈推进异步和集群能力。** 有性能证据后做预取、NUMA、热点策略；有可靠节点服务后再做在线重分布，最后评估 RDMA。 + +现阶段不建议同时承诺所有量化格式、所有推理引擎、多租户服务、训练状态缓存和通用 GPU 共享。它们具有不同的所有权与一致性要求，应由具体适配需求驱动。 From 8dc5f34887d5ea49d42451f64ec56aaaf6d73b31 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 01:31:02 +0800 Subject: [PATCH 12/29] Close cached descriptors across interrupted and partial loads --- src/infiniloader/cache.py | 52 ++++++++-------- tests/test_load_cleanup.py | 119 +++++++++++++++++++++++++++++++++++++ 2 files changed, 144 insertions(+), 27 deletions(-) create mode 100644 tests/test_load_cleanup.py diff --git a/src/infiniloader/cache.py b/src/infiniloader/cache.py index 7c66309..d1895eb 100644 --- a/src/infiniloader/cache.py +++ b/src/infiniloader/cache.py @@ -8,6 +8,7 @@ import stat import time import tempfile +import threading from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from pathlib import Path @@ -287,7 +288,7 @@ def copy_file_to_memfd(source: SourceFile, chunk_size: int = 64 * 1024 * 1024) - fcntl.F_SEAL_WRITE | fcntl.F_SEAL_GROW | fcntl.F_SEAL_SHRINK | fcntl.F_SEAL_SEAL) os.fchmod(dst_fd, 0o400) return CachedFile(source=source, fd=dst_fd, digest_prefix=digest.hexdigest()[:16]) - except Exception: + except BaseException: try: if dst_fd is not None: os.close(dst_fd) @@ -325,25 +326,31 @@ def load_model_into_memory( copied = 0 cached_by_index: list[CachedFile | None] = [None] * len(source_files) - futures = {} + ownership = threading.Lock() + aborted = False + + def load_one(index: int) -> CachedFile: + item = copy_file_to_memfd(source_files[index], chunk_size=chunk_size) + with ownership: + if aborted: + item.close() + else: + cached_by_index[index] = item + return item + try: if workers <= 1: for index, item in enumerate(source_files): - cached_item = copy_file_to_memfd(item, chunk_size=chunk_size) - cached_by_index[index] = cached_item + load_one(index) copied += item.size if progress: progress(copied, total, item.rel) else: with ThreadPoolExecutor(max_workers=workers) as pool: - futures = { - pool.submit(copy_file_to_memfd, item, chunk_size): (index, item) - for index, item in enumerate(source_files) - } + futures = [pool.submit(load_one, index) for index in range(len(source_files))] for future in as_completed(futures): - index, item = futures[future] cached_item = future.result() - cached_by_index[index] = cached_item + item = cached_item.source copied += item.size if progress: progress(copied, total, item.rel) @@ -358,23 +365,14 @@ def load_model_into_memory( exclude_patterns=exclude_patterns, workers=max(1, workers), ) - except Exception: - closed_fds = set() - for item in cached_by_index: - if item is None: - continue - item.close() - closed_fds.add(item.fd) - for future in futures: - if not future.done() or future.cancelled(): - continue - try: - item = future.result() - except Exception: - continue - if item.fd not in closed_fds: - item.close() - closed_fds.add(item.fd) + except BaseException: + # A second interrupt can also break the executor's shutdown wait. Workers + # finishing afterwards must release their own files instead of publishing. + with ownership: + aborted = True + for item in cached_by_index: + if item is not None: + item.close() raise diff --git a/tests/test_load_cleanup.py b/tests/test_load_cleanup.py new file mode 100644 index 0000000..4a7e3a4 --- /dev/null +++ b/tests/test_load_cleanup.py @@ -0,0 +1,119 @@ +"""Real descriptor ownership across interrupted loads and worker failures.""" +import gc +import os +from pathlib import Path +import threading + +import pytest + +from infiniloader import cache + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux memfd ownership") + + +def memfds(): + result = set() + for path in Path("/proc/self/fd").iterdir(): + try: + target = path.readlink() + except OSError: + continue + if "memfd:infiniloader-" in str(target): + result.add((path.name, str(target))) + return result + + +@pytest.fixture +def model(tmp_path): + root = tmp_path / "model" + root.mkdir() + for name in ("first.bin", "second.bin", "third.bin"): + (root / name).write_bytes(b"weights" * 1024) + return root + + +@pytest.mark.parametrize("workers", [1, 3]) +@pytest.mark.parametrize("error", [KeyboardInterrupt, SystemExit]) +def test_callback_interruption_releases_loaded_files(model, workers, error): + before = memfds() + + def interrupt(done, *_): + if done: + raise error("interrupt after copying weights") + + with pytest.raises(error): + cache.load_model_into_memory(model, workers=workers, progress=interrupt) + gc.collect() + assert memfds() == before + + +def test_copy_interruption_releases_raw_descriptor(model, monkeypatch): + before = memfds() + + def interrupt(*_): + raise KeyboardInterrupt("interrupt before ownership transfer") + + monkeypatch.setattr(cache, "_copy_with_copy_file_range", interrupt) + with pytest.raises(KeyboardInterrupt): + cache.copy_file_to_memfd(cache.collect_source_files(model)[0]) + gc.collect() + assert memfds() == before + + +@pytest.mark.parametrize("enqueue", [False, True]) +def test_partial_submit_failure_releases_worker_results(model, monkeypatch, enqueue): + before = memfds() + original = cache.ThreadPoolExecutor.submit + submitted = 0 + + def fail_second(pool, *args, **kwargs): + nonlocal submitted + submitted += 1 + if submitted != 2 or enqueue: + result = original(pool, *args, **kwargs) + if submitted == 2: + raise RuntimeError("cannot start another worker") + return result + + monkeypatch.setattr(cache.ThreadPoolExecutor, "submit", fail_second) + with pytest.raises(RuntimeError, match="cannot start"): + cache.load_model_into_memory(model, workers=2) + gc.collect() + assert memfds() == before + + +def test_worker_finishing_after_interrupted_shutdown_closes_its_file(model, monkeypatch): + before = memfds() + release = threading.Event() + waiting = threading.Event() + original_copy = cache.copy_file_to_memfd + original_shutdown = cache.ThreadPoolExecutor.shutdown + pools = [] + + def delayed_copy(source, **kwargs): + if source.rel.name == "second.bin": + waiting.set() + assert release.wait(10) + return original_copy(source, **kwargs) + + def interrupted_shutdown(pool, wait=True, **kwargs): + pools.append(pool) + original_shutdown(pool, wait=False, **kwargs) + raise KeyboardInterrupt("second interrupt while waiting for workers") + + def interrupt(done, *_): + if done: + assert waiting.wait(10) + raise KeyboardInterrupt("first interrupt") + + monkeypatch.setattr(cache, "copy_file_to_memfd", delayed_copy) + monkeypatch.setattr(cache.ThreadPoolExecutor, "shutdown", interrupted_shutdown) + try: + with pytest.raises(KeyboardInterrupt, match="second interrupt"): + cache.load_model_into_memory(model, workers=2, progress=interrupt) + finally: + release.set() + for pool in pools: + original_shutdown(pool, wait=True) + gc.collect() + assert memfds() == before From 2b6522d5975207abba6abae6e2202bed09c6a0f8 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 01:37:01 +0800 Subject: [PATCH 13/29] Normalize runtime paths and reject invalid cache locations --- src/infiniloader/cli.py | 12 +++-- src/infiniloader/daemon.py | 20 ++++---- src/infiniloader/local_cache.py | 13 +++-- src/infiniloader/paths.py | 21 +++++--- src/infiniloader/runtime.py | 6 +++ tests/test_runtime_paths.py | 90 +++++++++++++++++++++++++++++++++ 6 files changed, 135 insertions(+), 27 deletions(-) create mode 100644 tests/test_runtime_paths.py diff --git a/src/infiniloader/cli.py b/src/infiniloader/cli.py index ed2bc6a..b131ae2 100644 --- a/src/infiniloader/cli.py +++ b/src/infiniloader/cli.py @@ -14,13 +14,14 @@ from . import __version__ from .benchmark import benchmark_path from .cache import DEFAULT_EXCLUDES, build_view, human_bytes, inspect_source, load_model_into_memory, source_identity, _safe_remove_view -from .runtime import FileLock, process_matches, process_token, read_json, signal_owned_process, validate_name +from .runtime import FileLock, process_matches, process_token, read_json, signal_owned_process, validate_name, validate_socket_path from .memory import parse_bytes from .daemon import WeightLoaderDaemon from .distributed import build_tp_plan, create_rank_views, simulate_distributed_read from .ipc import LoaderConnectionError, request from .local_cache import cache_status, list_model_caches, materialize_model_cache, verify_cache, garbage_collect, lease_path, touch_cache from .paths import ( + absolute_path, default_log_path, default_pid_path, default_socket_path, @@ -41,19 +42,19 @@ def _print_json(data: object) -> None: def _socket(args: argparse.Namespace) -> Path: - return Path(args.socket).expanduser() if getattr(args, "socket", None) else default_socket_path(args.name, getattr(args, "base_dir", None)) + return absolute_path(args.socket) if getattr(args, "socket", None) else default_socket_path(args.name, getattr(args, "base_dir", None)) def _pid(args: argparse.Namespace) -> Path: - return Path(args.pid_file).expanduser() if getattr(args, "pid_file", None) else default_pid_path(args.name, getattr(args, "base_dir", None)) + return absolute_path(args.pid_file) if getattr(args, "pid_file", None) else default_pid_path(args.name, getattr(args, "base_dir", None)) def _state(args: argparse.Namespace) -> Path: - return Path(args.state_file).expanduser() if getattr(args, "state_file", None) else default_state_path(args.name, getattr(args, "base_dir", None)) + return absolute_path(args.state_file) if getattr(args, "state_file", None) else default_state_path(args.name, getattr(args, "base_dir", None)) def _view(args: argparse.Namespace) -> Path: - return Path(args.view).expanduser() if getattr(args, "view", None) else default_view_path(args.name, getattr(args, "base_dir", None)) + return absolute_path(args.view) if getattr(args, "view", None) else default_view_path(args.name, getattr(args, "base_dir", None)) def _resolve_name(args: argparse.Namespace) -> str: @@ -294,6 +295,7 @@ def _daemon_command(args: argparse.Namespace, model_path: Path | None = None, up def cmd_start(args: argparse.Namespace) -> int: + validate_socket_path(_socket(args)) args.identity = source_identity(Path(args.model), _patterns(args.pattern), tuple(args.exclude or DEFAULT_EXCLUDES)) with FileLock(_socket(args).with_suffix(".start.lock"), timeout=args.wait_timeout): return _start_locked(args) diff --git a/src/infiniloader/daemon.py b/src/infiniloader/daemon.py index 2f3a755..2463fd4 100644 --- a/src/infiniloader/daemon.py +++ b/src/infiniloader/daemon.py @@ -15,9 +15,10 @@ from typing import Any from .cache import CachedModel, build_view, human_bytes, load_model_into_memory, source_identity, _safe_remove_view, collect_source_files -from .runtime import FileLock, atomic_json, atomic_text, process_token, process_matches, read_json +from .runtime import FileLock, atomic_json, atomic_text, process_token, process_matches, read_json, validate_socket_path from .local_cache import lease_path, verify_cache from .memory import ResidentReservation +from .paths import absolute_path class WeightLoaderDaemon: @@ -47,13 +48,13 @@ def __init__( idle_ttl: float | None = None, ) -> None: self.name = name - self.source = source - self.view = view - self.socket_path = socket_path - self.pid_path = pid_path - self.state_path = state_path - self.upstream_source = upstream_source - self.local_cache_path = local_cache_path + self.source = absolute_path(source) + self.view = absolute_path(view) + self.socket_path = absolute_path(socket_path) + self.pid_path = absolute_path(pid_path) + self.state_path = absolute_path(state_path) + self.upstream_source = absolute_path(upstream_source) if upstream_source is not None else None + self.local_cache_path = absolute_path(local_cache_path) if local_cache_path is not None else None self.include_patterns = include_patterns self.exclude_patterns = exclude_patterns self.replace = replace @@ -301,7 +302,7 @@ def _claim_runtime_file(self, path: Path) -> None: def _check_runtime_paths(self) -> None: paths = (self.socket_path, self.pid_path, self.state_path) - if len({path.absolute() for path in paths}) != len(paths): + if len({path.resolve() for path in paths}) != len(paths): raise RuntimeError("socket, PID and state paths must be distinct") for path in paths: for source in (self.source, self.upstream_source): @@ -358,6 +359,7 @@ def _handle_signal(self, signum: int, _frame: object) -> None: self._stop = True def serve_forever(self) -> int: + validate_socket_path(self.socket_path) with ExitStack() as locks: locks.enter_context(FileLock(self.socket_path.with_suffix(".owner.lock"), timeout=0)) locks.enter_context(FileLock(self.view.parent / f".{self.view.name}.owner.lock", timeout=0)) diff --git a/src/infiniloader/local_cache.py b/src/infiniloader/local_cache.py index c2858e5..e77d45a 100644 --- a/src/infiniloader/local_cache.py +++ b/src/infiniloader/local_cache.py @@ -15,7 +15,7 @@ import math from .cache import DEFAULT_EXCLUDES, SourceFile, collect_source_files, human_bytes -from .paths import model_cache_root +from .paths import absolute_path, default_base_dir from .runtime import FileLock, atomic_json, read_json @@ -43,8 +43,11 @@ def cache_key(source_root: Path, name: str | None = None) -> str: return f"{_safe_name(name or source_root.name)}-{hashlib.sha256(str(source_root).encode()).hexdigest()[:12]}" -def _root(cache_root=None, base_dir=None) -> Path: - root = Path(cache_root).expanduser().resolve() if cache_root else model_cache_root(base_dir).resolve() +def _root(cache_root=None, base_dir=None, *, source_root: Path | None = None) -> Path: + base = absolute_path(base_dir) if base_dir else None + root = (Path(cache_root).expanduser() if cache_root else (base or default_base_dir()) / "model-cache").resolve() + if source_root is not None and root.is_relative_to(source_root.resolve()): + raise ValueError("disk cache directory must be outside the source model directory") root.mkdir(parents=True, exist_ok=True) return root @@ -165,7 +168,7 @@ def cache_path_for(source_root: Path, *, name=None, cache_root=None, base_dir=No include_patterns=(), exclude_patterns=DEFAULT_EXCLUDES) -> Path: source_root = source_root.expanduser().resolve() files = collect_source_files(source_root, include_patterns, exclude_patterns) - root = _root(cache_root, base_dir) + root = _root(cache_root, base_dir, source_root=source_root) prefix = f"{cache_key(source_root, name)}-{_revision(files, include_patterns, exclude_patterns)}" for candidate in sorted(root.glob(f"{prefix}*"), key=_created_at, reverse=True): @@ -214,7 +217,7 @@ def materialize_model_cache(source_root: Path, *, name=None, cache_root=None, ba raise ValueError("workers and chunk_size must be positive") start = time.perf_counter() source_root = source_root.expanduser().resolve() - root = _root(cache_root, base_dir) + root = _root(cache_root, base_dir, source_root=source_root) with FileLock(root / ".cache.lock", timeout=600): files = collect_source_files(source_root, include_patterns, exclude_patterns) if not files: diff --git a/src/infiniloader/paths.py b/src/infiniloader/paths.py index d5a811c..edae073 100644 --- a/src/infiniloader/paths.py +++ b/src/infiniloader/paths.py @@ -9,6 +9,11 @@ APP_DIR_NAME = "infiniloader" +def absolute_path(path: str | os.PathLike[str]) -> Path: + """Anchor paths to the caller's directory without following leaf symlinks.""" + return Path(path).expanduser().absolute() + + def _writable_dir(path: Path) -> bool: try: path.mkdir(parents=True, exist_ok=True) @@ -22,7 +27,7 @@ def _writable_dir(path: Path) -> bool: def default_base_dir() -> Path: configured = os.environ.get("ILOADER_BASE_DIR") if configured: - candidate = Path(configured).expanduser() + candidate = absolute_path(configured) if _writable_dir(candidate): return candidate raise RuntimeError(f"ILOADER_BASE_DIR is not writable: {candidate}") @@ -30,46 +35,46 @@ def default_base_dir() -> Path: candidates = [] xdg_cache_home = os.environ.get("XDG_CACHE_HOME") if xdg_cache_home: - candidates.append(Path(xdg_cache_home).expanduser() / APP_DIR_NAME) + candidates.append(absolute_path(xdg_cache_home) / APP_DIR_NAME) candidates.append(Path.home() / ".cache" / APP_DIR_NAME) for candidate in candidates: if _writable_dir(candidate): return candidate - fallback = Path(os.environ.get("TMPDIR", "/tmp")) / f"{APP_DIR_NAME}-{os.getuid()}" + fallback = absolute_path(os.environ.get("TMPDIR", "/tmp")) / f"{APP_DIR_NAME}-{os.getuid()}" fallback.mkdir(parents=True, exist_ok=True, mode=0o700) return fallback def runtime_dir(base_dir: str | os.PathLike[str] | None = None) -> Path: - base = Path(base_dir).expanduser() if base_dir else default_base_dir() + base = absolute_path(base_dir) if base_dir else default_base_dir() path = base / "run" path.mkdir(parents=True, exist_ok=True) return path def view_root(base_dir: str | os.PathLike[str] | None = None) -> Path: - base = Path(base_dir).expanduser() if base_dir else default_base_dir() + base = absolute_path(base_dir) if base_dir else default_base_dir() path = base / "views" path.mkdir(parents=True, exist_ok=True) return path def dist_view_root(base_dir: str | os.PathLike[str] | None = None) -> Path: - base = Path(base_dir).expanduser() if base_dir else default_base_dir() + base = absolute_path(base_dir) if base_dir else default_base_dir() path = base / "dist-views" path.mkdir(parents=True, exist_ok=True) return path def logs_dir(base_dir: str | os.PathLike[str] | None = None) -> Path: - base = Path(base_dir).expanduser() if base_dir else default_base_dir() + base = absolute_path(base_dir) if base_dir else default_base_dir() path = base / "logs" path.mkdir(parents=True, exist_ok=True) return path def model_cache_root(base_dir: str | os.PathLike[str] | None = None) -> Path: - base = Path(base_dir).expanduser() if base_dir else default_base_dir() + base = absolute_path(base_dir) if base_dir else default_base_dir() path = base / "model-cache" path.mkdir(parents=True, exist_ok=True) return path diff --git a/src/infiniloader/runtime.py b/src/infiniloader/runtime.py index e788040..816c332 100644 --- a/src/infiniloader/runtime.py +++ b/src/infiniloader/runtime.py @@ -17,6 +17,12 @@ def validate_name(name: str) -> str: return name +def validate_socket_path(path: Path) -> None: + encoded = os.fsencode(path) + if b"\0" in encoded or len(encoded) > 107: + raise ValueError("Unix socket path must fit in 107 bytes; use a shorter --base-dir or --socket") + + def read_json(path: Path) -> dict[str, Any] | None: try: value = json.loads(path.read_text(encoding="utf-8")) diff --git a/tests/test_runtime_paths.py b/tests/test_runtime_paths.py new file mode 100644 index 0000000..cbdb434 --- /dev/null +++ b/tests/test_runtime_paths.py @@ -0,0 +1,90 @@ +import os +from pathlib import Path +import sys +import tempfile + +import pytest + +from conftest import make_model_dir +from infiniloader import cli, paths +from infiniloader.local_cache import materialize_model_cache +from infiniloader.runtime import validate_socket_path + + +@pytest.mark.parametrize("option", ["base", "view"]) +@pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux daemon integration") +def test_relative_runtime_paths_survive_consumer_directory_change(monkeypatch, option): + with tempfile.TemporaryDirectory(prefix="ilp-") as directory: + root = Path(directory) + first, second = root / "first", root / "second" + first.mkdir() + second.mkdir() + model = make_model_dir(root / "model") + base = first / "state" + monkeypatch.setenv("PYTHONPATH", str(Path(cli.__file__).resolve().parents[1])) + monkeypatch.chdir(first) + args = ["start", str(model), "--name", "unit", "--base-dir", "state" if option == "base" else str(base)] + if option == "view": + args += ["--view", "relative-view"] + started = False + try: + assert cli.main(args) == 0 + started = True + monkeypatch.chdir(second) + script = ("import os; from pathlib import Path; p=Path(os.environ['ILOADER_MODEL_PATH']); " + "assert p.is_absolute(); assert (p/'model.safetensors').read_bytes()==b'0123456789abcdef'") + assert cli.main(["run", "--base-dir", str(base), "unit", "--", sys.executable, "-c", script]) == 0 + finally: + if started: + assert cli.main(["stop", "unit", "--base-dir", str(base), "--remove-view"]) == 0 + from test_lifecycle import wait_for + wait_for(lambda: not (base / "run/unit.pid").exists()) + + +def test_configured_relative_base_is_anchored(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("ILOADER_BASE_DIR", "state") + assert paths.default_base_dir() == tmp_path / "state" + + +def test_unix_socket_limit_counts_encoded_bytes(): + validate_socket_path(Path("/" + "a" * 106)) + with pytest.raises(ValueError, match="107 bytes"): + validate_socket_path(Path("/" + "a" * 107)) + with pytest.raises(ValueError, match="107 bytes"): + validate_socket_path(Path("/" + "模" * 36)) + + +def test_long_socket_is_rejected_before_cache_preparation(tmp_path, monkeypatch, capsys): + def forbidden(*_): + pytest.fail("must validate socket before preparing model data") + + monkeypatch.setattr(cli, "_prepare_model_source", forbidden) + model = make_model_dir(tmp_path / "model") + assert cli.main(["start", str(model), "--name", "x" * 96, "--base-dir", str(tmp_path / "state")]) == 1 + assert "107 bytes" in capsys.readouterr().err + + +@pytest.mark.parametrize("mode", ["explicit", "default", "symlink"]) +def test_nested_cache_root_is_rejected_before_copy(tmp_path, monkeypatch, mode): + model = make_model_dir(tmp_path / "model") + nested = model / "local-cache" + kwargs = {"cache_root": nested} + if mode == "default": + kwargs = {"base_dir": model / "state"} + elif mode == "symlink": + alias = tmp_path / "alias" + try: + alias.symlink_to(model, target_is_directory=True) + except OSError: + pytest.skip("symlink creation unavailable") + kwargs = {"cache_root": alias / "local-cache"} + + def forbidden(*_): + pytest.fail("must reject overlapping roots before copying") + + monkeypatch.setattr("infiniloader.local_cache._copy_one", forbidden) + with pytest.raises(ValueError, match="outside the source"): + materialize_model_cache(model, **kwargs) + assert not nested.exists() + assert not (model / "state").exists() From 5acb8b7a8ad2f9afa0c9224cfafc41f04dd27c4f Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 01:47:39 +0800 Subject: [PATCH 14/29] Publish runtime sockets safely and wait for cache shutdown --- src/infiniloader/cli.py | 73 +++++++++------ src/infiniloader/daemon.py | 89 +++++++++++++++--- src/infiniloader/diagnostics.py | 2 + src/infiniloader/ipc.py | 6 ++ tests/test_cli.py | 32 ++++++- tests/test_runtime_publication.py | 146 ++++++++++++++++++++++++++++++ 6 files changed, 309 insertions(+), 39 deletions(-) create mode 100644 tests/test_runtime_publication.py diff --git a/src/infiniloader/cli.py b/src/infiniloader/cli.py index b131ae2..62375b5 100644 --- a/src/infiniloader/cli.py +++ b/src/infiniloader/cli.py @@ -3,7 +3,6 @@ import argparse import json import os -import signal import shutil import subprocess import sys @@ -18,7 +17,7 @@ from .memory import parse_bytes from .daemon import WeightLoaderDaemon from .distributed import build_tp_plan, create_rank_views, simulate_distributed_read -from .ipc import LoaderConnectionError, request +from .ipc import LoaderConnectionError, LoaderStoppingError, request from .local_cache import cache_status, list_model_caches, materialize_model_cache, verify_cache, garbage_collect, lease_path, touch_cache from .paths import ( absolute_path, @@ -115,10 +114,7 @@ def _starting_status(args: argparse.Namespace, error: Exception) -> dict | None: def _load_state(path: Path) -> dict | None: - try: - return json.loads(path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError): - return None + return read_json(path) def _format_duration(seconds: float | None) -> str: @@ -318,17 +314,32 @@ def _require_identity(args: argparse.Namespace, status: dict) -> None: def _start_locked(args: argparse.Namespace) -> int: sock = _socket(args) - try: - status = request(sock, {"command": "status"}, timeout=1.0)["status"] - except LoaderConnectionError: - status = None - if status is not None: - _require_identity(args, status) - request(sock, {"command": "touch"}) - print(status["view"]) - print(f"already running: name={args.name} pid={status['pid']}", file=sys.stderr) - return 0 - starting = _starting_status(args, LoaderConnectionError("not ready")) + deadline = time.monotonic() + args.wait_timeout + while True: + try: + status = request(sock, {"command": "status"}, timeout=1.0)["status"] + except LoaderConnectionError: + status = None + starting = _starting_status(args, LoaderConnectionError("not ready")) if status is None else None + current = status or starting + if current and current.get("state") == "stopping": + if time.monotonic() >= deadline: + raise TimeoutError("previous daemon is still stopping; retry after it exits") + time.sleep(.05) + continue + if status is not None: + _require_identity(args, status) + try: + request(sock, {"command": "touch"}) + except (LoaderConnectionError, LoaderStoppingError) as exc: + if time.monotonic() >= deadline: + raise TimeoutError("daemon stopped while opening the cache; retry") from exc + time.sleep(.05) + continue + print(status["view"]) + print(f"already running: name={args.name} pid={status['pid']}", file=sys.stderr) + return 0 + break proc = None if starting: _require_identity(args, starting) @@ -336,6 +347,8 @@ def _start_locked(args: argparse.Namespace) -> int: print(starting["view"]) return 0 else: + with FileLock(sock.with_suffix(".owner.lock"), timeout=max(0, deadline - time.monotonic())): + pass # A stopping owner can still be finishing cleanup after removing its metadata. source, upstream_source, local_cache_path = _prepare_model_source(args) if source_identity(Path(args.model), _patterns(args.pattern), tuple(args.exclude or DEFAULT_EXCLUDES)) != args.identity: raise RuntimeError("source model changed while preparing local cache; retry") @@ -349,7 +362,6 @@ def _start_locked(args: argparse.Namespace) -> int: ) print(f"starting name={args.name} pid={proc.pid} log={log_path}", file=sys.stderr) - deadline = time.monotonic() + args.wait_timeout last_progress = None while time.monotonic() < deadline: if proc is not None and proc.poll() is not None: @@ -443,7 +455,7 @@ def _status_for_name(name: str, args: argparse.Namespace) -> dict: ) try: status = request(_socket(ns), {"command": "status"}, timeout=ns.timeout)["status"] - status["state"] = "ready" + status.setdefault("state", "ready") return status except LoaderConnectionError as exc: starting = _starting_status(ns, exc) @@ -494,27 +506,36 @@ def cmd_list(args: argparse.Namespace) -> int: return 0 +def _wait_for_exit(pid: int, token: str, timeout: float) -> None: + deadline = time.monotonic() + timeout + while process_matches(pid, token): + if time.monotonic() >= deadline: + raise TimeoutError("daemon is still stopping; view was preserved, retry after it exits") + time.sleep(.05) + + def _stop_one(args: argparse.Namespace) -> int: _resolve_name(args) try: response = request(_socket(args), {"command": "stop", "remove_view": args.remove_view, "force": getattr(args, "force", False)}, timeout=args.timeout) - print(f"{args.name}: {'stopping' if response.get('stopping') else 'ok'}") + pid, token = response.get("pid"), response.get("process_token") + if pid is not None and token is not None: + _wait_for_exit(int(pid), token, args.timeout) + print(f"{args.name}: stopped") + else: + print(f"{args.name}: {'stopping' if response.get('stopping') else 'ok'}") return 0 except LoaderConnectionError as exc: starting = _starting_status(args, exc) if not starting: print(f"{args.name}: {exc}", file=sys.stderr) return 1 - if starting.get("state") not in ("starting", "loading") and not getattr(args, "force", False): + if starting.get("state") not in ("starting", "loading", "loaded") and not getattr(args, "force", False): raise RuntimeError("ready daemon is unresponsive; no signal was sent, retry or use stop --force") signal_owned_process(int(starting["pid"]), starting.get("process_token")) + _wait_for_exit(int(starting["pid"]), starting.get("process_token"), args.timeout) if args.remove_view: - deadline = time.monotonic() + args.timeout - while process_matches(int(starting["pid"]), starting.get("process_token")): - if time.monotonic() >= deadline: - raise TimeoutError("daemon has not exited; view was preserved") - time.sleep(.05) view = Path(starting["view"]) if view.exists(): _safe_remove_view(view, expected_pid=int(starting["pid"])) diff --git a/src/infiniloader/daemon.py b/src/infiniloader/daemon.py index 2463fd4..5da576f 100644 --- a/src/infiniloader/daemon.py +++ b/src/infiniloader/daemon.py @@ -11,6 +11,7 @@ import threading import uuid import stat +import errno from pathlib import Path from typing import Any @@ -19,6 +20,7 @@ from .local_cache import lease_path, verify_cache from .memory import ResidentReservation from .paths import absolute_path +from .ipc import LoaderStoppingError class WeightLoaderDaemon: @@ -84,6 +86,7 @@ def __init__( self._connections = threading.BoundedSemaphore(16) self._claimed_runtime: dict[Path, tuple[int, int]] = {} self._socket_identity: dict | None = None + self._private_socket: Path | None = None def log(self, message: str) -> None: if not self.quiet: @@ -102,6 +105,7 @@ def _write_state(self, state: dict[str, Any]) -> None: state.setdefault("idle_ttl", self.idle_ttl) state.setdefault("memory_budget", self._reservation.info if self._reservation else None) state.setdefault("runtime_socket", self._socket_identity) + state.setdefault("runtime_private_socket", str(self._private_socket) if self._private_socket else None) atomic_json(self.state_path, state) self._claim_runtime_file(self.state_path) @@ -150,7 +154,7 @@ def load(self) -> None: load_seconds = self.load_finished_at - self.load_started_at self._write_state( { - "state": "ready", + "state": "loaded", "load_started_at": self.load_started_at, "load_finished_at": self.load_finished_at, "load_seconds": load_seconds, @@ -169,6 +173,7 @@ def load(self) -> None: def _status(self) -> dict[str, Any]: assert self.cached is not None return { + "state": "stopping" if self._stop else "ready", "name": self.name, "identity": self.identity, "process_token": self.process_token, @@ -211,7 +216,7 @@ def _touch(self) -> None: def _handle_locked(self, payload: dict[str, Any]) -> dict[str, Any]: command = payload.get("command") if self._stop and command not in ("ping", "status", "lease-release"): - raise RuntimeError("daemon is stopping; retry after it exits") + raise LoaderStoppingError("daemon is stopping; retry after it exits") if command == "ping": return {"ok": True, "pong": True} if command == "status": @@ -247,7 +252,8 @@ def _handle_locked(self, payload: dict[str, Any]) -> dict[str, Any]: raise RuntimeError("cache was used after collection was planned") self._remove_view_on_exit = bool(payload.get("remove_view", False)) self._stop = True - return {"ok": True, "stopping": True} + self._write_state({"state": "stopping", **self._status()}) + return {"ok": True, "stopping": True, "pid": os.getpid(), "process_token": self.process_token} return {"ok": False, "error": f"unknown command: {command!r}"} def _serve_connection(self, conn: socket.socket) -> None: @@ -278,6 +284,8 @@ def _serve_connection(self, conn: socket.socket) -> None: response = self._handle(payload) except (ValueError, TypeError, RuntimeError) as exc: response = {"ok": False, "error": str(exc)} + if isinstance(exc, LoaderStoppingError): + response["code"] = "stopping" conn.settimeout(1.0) conn.sendall(json.dumps(response).encode("utf-8") + b"\n") except OSError: @@ -301,6 +309,8 @@ def _claim_runtime_file(self, path: Path) -> None: self._claimed_runtime[path] = (st.st_dev, st.st_ino) def _check_runtime_paths(self) -> None: + stale_pid = None + stale_socket = None paths = (self.socket_path, self.pid_path, self.state_path) if len({path.resolve() for path in paths}) != len(paths): raise RuntimeError("socket, PID and state paths must be distinct") @@ -326,6 +336,7 @@ def _check_runtime_paths(self) -> None: if process_matches(state["pid"], state.get("process_token")): raise RuntimeError("runtime metadata belongs to a live daemon") if self.pid_path.exists(): + stale_pid = self.pid_path.lstat() try: pid = int(self.pid_path.read_text().strip()) except (OSError, ValueError) as exc: @@ -346,6 +357,69 @@ def _check_runtime_paths(self) -> None: raise RuntimeError(f"cannot establish that runtime socket is stale: {self.socket_path}") from exc else: raise RuntimeError(f"runtime socket still has a listener: {self.socket_path}") + stale_socket = st + # Validate every public resource before removing any stale one. Removing + # the old PID first avoids publishing a new state next to an old PID. + if state: + private = state.get("runtime_private_socket") + if isinstance(private, str): + path = absolute_path(private) + if path.parent == self.socket_path.parent and path not in paths: + try: + st = path.lstat() + expected = {"device": st.st_dev, "inode": st.st_ino} + if stat.S_ISSOCK(st.st_mode) and st.st_uid == os.getuid() and expected == state.get("runtime_socket"): + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: + probe.settimeout(.2) + try: + probe.connect(str(path)) + except ConnectionRefusedError: + self._remove_stale_file(path, st) + except OSError: + pass # Preserve anything whose ownership/liveness is uncertain. + if stale_pid is not None: + self._remove_stale_file(self.pid_path, stale_pid) + if stale_socket is not None: + self._remove_stale_file(self.socket_path, stale_socket) + + @staticmethod + def _remove_stale_file(path: Path, expected: os.stat_result) -> None: + current = path.lstat() + if (current.st_dev, current.st_ino) != (expected.st_dev, expected.st_ino): + raise RuntimeError(f"runtime resource changed during startup: {path}") + path.unlink() + + def _publish_socket(self, server: socket.socket) -> None: + # Keep private names no longer than the preflighted public basename. + width = min(12, len(os.fsencode(self.socket_path.name))) + for _ in range(64): + name = uuid.uuid4().hex[:width] + if width > 1: + name = "." + name[1:] + private = self.socket_path.with_name(name) + if private in (self.socket_path, self.pid_path, self.state_path): + continue + try: + server.bind(str(private)) + except OSError as exc: + if exc.errno == errno.EADDRINUSE: + continue + raise + break + else: + raise RuntimeError("cannot reserve a private runtime socket") + self._private_socket = private + self._claim_runtime_file(private) + st = private.lstat() + self._socket_identity = {"device": st.st_dev, "inode": st.st_ino} + os.chmod(private, 0o600) + server.listen(16) + # Persist the inode before making the socket visible. A crash can then + # leave either no public socket or one that can be verified on restart. + self._write_state({"state": "loaded"}) + self._claimed_runtime[self.socket_path] = (st.st_dev, st.st_ino) + os.link(private, self.socket_path, follow_symlinks=False) + private.unlink() def _install_signal_handlers(self) -> None: signal.signal(signal.SIGTERM, self._handle_signal) @@ -392,8 +466,6 @@ def _serve_owned(self) -> int: self.socket_path.parent.mkdir(parents=True, exist_ok=True) self.pid_path.parent.mkdir(parents=True, exist_ok=True) self.state_path.parent.mkdir(parents=True, exist_ok=True) - if self.socket_path.exists(): - self.socket_path.unlink() self._write_state({"state": "starting"}) atomic_text(self.pid_path, str(os.getpid())) self._claim_runtime_file(self.pid_path) @@ -410,12 +482,7 @@ def _serve_owned(self) -> int: server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) server.settimeout(1.0) - server.bind(str(self.socket_path)) - self._claim_runtime_file(self.socket_path) - st = self.socket_path.lstat() - self._socket_identity = {"device": st.st_dev, "inode": st.st_ino} - os.chmod(self.socket_path, 0o600) - server.listen(16) + self._publish_socket(server) self._write_state({"state": "ready", "progress": 1.0, "copied_bytes": self.cached.total_bytes, **self._status()}) self.log(f"[listen] socket={self.socket_path}") diff --git a/src/infiniloader/diagnostics.py b/src/infiniloader/diagnostics.py index 6a1b3be..63470eb 100644 --- a/src/infiniloader/diagnostics.py +++ b/src/infiniloader/diagnostics.py @@ -35,6 +35,8 @@ def record(name, ok, detail): with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: sock.bind(str(root / "probe.sock")) record("unix_socket", True, "socket paths under base directory are usable") + os.link(root / "probe.sock", root / "published.sock") + record("socket_publication", True, "exclusive socket hard-link publication works") fd = os.memfd_create("iloader-doctor", os.MFD_CLOEXEC | os.MFD_ALLOW_SEALING) try: import fcntl diff --git a/src/infiniloader/ipc.py b/src/infiniloader/ipc.py index f7d33a0..238b2bb 100644 --- a/src/infiniloader/ipc.py +++ b/src/infiniloader/ipc.py @@ -10,6 +10,10 @@ class LoaderConnectionError(RuntimeError): pass +class LoaderStoppingError(RuntimeError): + pass + + def request(socket_path: str | Path, payload: dict[str, Any], timeout: float = 5.0) -> dict[str, Any]: sock_path = str(socket_path) try: @@ -39,5 +43,7 @@ def request(socket_path: str | Path, payload: dict[str, Any], timeout: float = 5 except ValueError as exc: raise LoaderConnectionError(f"invalid response from {sock_path}: {exc}") from exc if not response.get("ok", False): + if response.get("code") == "stopping": + raise LoaderStoppingError(response.get("error", "daemon is stopping")) raise RuntimeError(response.get("error", "InfiniLoader daemon request failed")) return response diff --git a/tests/test_cli.py b/tests/test_cli.py index c7a796c..38ee815 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2,6 +2,7 @@ import json import os +import pytest from conftest import make_model_dir from infiniloader.cli import build_parser, main @@ -73,10 +74,15 @@ def test_list_reports_active_starting_cache_name(tmp_path, monkeypatch, capsys): def test_stop_all_can_stop_starting_daemon(tmp_path, monkeypatch, capsys): killed = [] + alive = True monkeypatch.setattr("infiniloader.cli._pid_is_running", lambda _pid: True) monkeypatch.setattr("infiniloader.cli.os.kill", lambda pid, sig: killed.append((pid, sig))) - monkeypatch.setattr("infiniloader.cli.process_matches", lambda pid, token: True) - monkeypatch.setattr("infiniloader.cli.signal_owned_process", lambda pid, token: killed.append((pid, token))) + monkeypatch.setattr("infiniloader.cli.process_matches", lambda pid, token: alive) + def signal_exit(pid, token): + nonlocal alive + killed.append((pid, token)) + alive = False + monkeypatch.setattr("infiniloader.cli.signal_owned_process", signal_exit) run_dir = tmp_path / "state" / "run" run_dir.mkdir(parents=True) (run_dir / "unit.pid").write_text("12345", encoding="utf-8") @@ -90,6 +96,28 @@ def test_stop_all_can_stop_starting_daemon(tmp_path, monkeypatch, capsys): assert "unit: stopping pid=12345" in captured.out +@pytest.mark.parametrize("state,force", [("loading", False), ("ready", True)]) +@pytest.mark.parametrize("remove_view", [False, True]) +def test_stop_fallback_waits_for_the_original_owner(tmp_path, monkeypatch, state, force, remove_view): + from infiniloader import cli + from infiniloader.ipc import LoaderConnectionError + events = [] + def offline(*args, **kwargs): + raise LoaderConnectionError("socket unavailable") + monkeypatch.setattr(cli, "request", offline) + monkeypatch.setattr(cli, "_starting_status", lambda *_: dict( + state=state, pid=12345, process_token="original", view=str(tmp_path / "missing"))) + monkeypatch.setattr(cli, "signal_owned_process", lambda pid, token: events.append(("signal", pid, token))) + monkeypatch.setattr(cli, "_wait_for_exit", lambda pid, token, timeout: events.append(("wait", pid, token))) + args = ["stop", "unit", "--base-dir", str(tmp_path)] + if force: + args.append("--force") + if remove_view: + args.append("--remove-view") + assert main(args) == 0 + assert events == [("signal", 12345, "original"), ("wait", 12345, "original")] + + def test_completion_bash_outputs_completion_function(capsys): code = main(["completion", "bash"]) captured = capsys.readouterr() diff --git a/tests/test_runtime_publication.py b/tests/test_runtime_publication.py new file mode 100644 index 0000000..3699c5b --- /dev/null +++ b/tests/test_runtime_publication.py @@ -0,0 +1,146 @@ +import json +import os +from pathlib import Path +import socket +import subprocess +import sys +import tempfile +from types import SimpleNamespace + +import pytest + +from conftest import make_model_dir +from infiniloader import cli +from infiniloader.daemon import WeightLoaderDaemon +from infiniloader.ipc import request + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux runtime publication") + + +@pytest.fixture +def runtime(monkeypatch): + with tempfile.TemporaryDirectory(prefix="ilr-") as directory: + root = Path(directory) + model = make_model_dir(root / "model") + base = root / "state" + monkeypatch.setenv("PYTHONPATH", str(Path(cli.__file__).resolve().parents[1])) + yield model, base + if (base / "run/unit.pid").exists(): + cli.main(["stop", "unit", "--base-dir", str(base), "--remove-view"]) + + +def start(model, base): + return cli.main(["start", str(model), "--name", "unit", "--base-dir", str(base), "--wait-timeout", "10"]) + + +@pytest.mark.parametrize("through_cli", [False, True]) +def test_stop_then_start_can_switch_model_without_retry(runtime, through_cli): + model, base = runtime + assert start(model, base) == 0 + sock = base / "run/unit.sock" + old_pid = request(sock, {"command": "status"})["status"]["pid"] + if through_cli: + assert cli.main(["stop", "unit", "--base-dir", str(base), "--remove-view"]) == 0 + assert not (base / "run/unit.pid").exists() + else: + assert request(sock, {"command": "stop", "remove_view": True})["stopping"] + replacement = make_model_dir(model.parent / "replacement", payload=b"new revision") + assert start(replacement, base) == 0 + status = request(sock, {"command": "status"})["status"] + assert status["pid"] != old_pid + assert (Path(status["view"]) / "model.safetensors").read_bytes() == b"new revision" + + +def crash_serve(model, base, point): + common = "from infiniloader.cli import main; from infiniloader.daemon import WeightLoaderDaemon; import os\n" + if point == "before_publication": + injection = ("original=WeightLoaderDaemon._write_state\n" + "def write(self,state):\n" + " if self._socket_identity is not None: os._exit(77)\n" + " return original(self,state)\n" + "WeightLoaderDaemon._write_state=write\n") + elif point == "after_publication": + injection = "original=os.link\ndef link(*a,**k):\n original(*a,**k)\n os._exit(77)\nos.link=link\n" + else: + injection = ("import infiniloader.daemon as daemon\n" + "daemon.atomic_text=lambda *a,**k: os._exit(77)\n") + result = subprocess.run([sys.executable, "-c", common + injection + "raise SystemExit(main())", + "_serve", str(model), "--name", "unit", "--base-dir", str(base), "--replace", "--quiet"], + capture_output=True, timeout=20) + assert result.returncode == 77, result.stderr.decode(errors="replace") + + +@pytest.mark.parametrize("point", ["before_publication", "after_publication"]) +def test_start_recovers_from_socket_publication_crash(runtime, point): + model, base = runtime + crash_serve(model, base, point) + if point == "after_publication": + state = json.loads((base / "run/unit.state.json").read_text()) + st = (base / "run/unit.sock").stat() + assert state["runtime_socket"] == {"device": st.st_dev, "inode": st.st_ino} + else: + assert not (base / "run/unit.sock").exists() + assert start(model, base) == 0 + + +def test_start_recovers_between_replacing_state_and_pid(runtime): + model, base = runtime + crash_serve(model, base, "after_publication") + crash_serve(model, base, "before_pid") + assert not (base / "run/unit.pid").exists() + assert start(model, base) == 0 + + +def test_publication_does_not_replace_a_concurrent_file(runtime, monkeypatch): + model, base = runtime + run = base / "run" + run.mkdir(parents=True) + daemon = WeightLoaderDaemon(name="unit", source=model, view=base / "view", socket_path=run / "unit.sock", + pid_path=run / "unit.pid", state_path=run / "unit.state.json", upstream_source=None, + local_cache_path=None, include_patterns=(), exclude_patterns=(), replace=True, + force=False, max_memory_fraction=.9, chunk_size=1024, workers=1) + original = os.link + + def race(source, destination, **kwargs): + Path(destination).write_text("preserve this file") + original(source, destination, **kwargs) + + monkeypatch.setattr(os, "link", race) + daemon._owns_runtime = True + try: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as server: + with pytest.raises(FileExistsError): + daemon._publish_socket(server) + finally: + daemon._cleanup_socket() + assert daemon.socket_path.read_text() == "preserve this file" + + +def test_private_socket_never_uses_the_public_name(runtime, monkeypatch): + model, base = runtime + run = base / "run" + run.mkdir(parents=True) + daemon = WeightLoaderDaemon(name="unit", source=model, view=base / "view", socket_path=run / ".abc", + pid_path=run / "unit.pid", state_path=run / "unit.state.json", upstream_source=None, + local_cache_path=None, include_patterns=(), exclude_patterns=(), replace=True, + force=False, max_memory_fraction=.9, chunk_size=1024, workers=1) + candidates = iter(["fabc", "fdef"]) + monkeypatch.setattr("infiniloader.daemon.uuid.uuid4", lambda: SimpleNamespace(hex=next(candidates))) + original = daemon._write_state + + def inspect(state): + assert not daemon.socket_path.exists() + assert daemon._private_socket != daemon.socket_path + original(state) + + monkeypatch.setattr(daemon, "_write_state", inspect) + daemon._owns_runtime = True + try: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as server: + daemon._publish_socket(server) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(str(daemon.socket_path)) + accepted, _ = server.accept() + accepted.close() + finally: + daemon._cleanup_socket() From 8fe11ecf8042662977b8f9ae6b38abdee4d00020 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 01:53:57 +0800 Subject: [PATCH 15/29] Share strict safetensors validation across planning commands --- src/infiniloader/distributed.py | 52 ++----------- src/infiniloader/formats.py | 125 ++++++++++++++++++++++++++++++++ src/infiniloader/shards.py | 95 ++++-------------------- tests/test_format_validation.py | 90 +++++++++++++++++++++++ 4 files changed, 234 insertions(+), 128 deletions(-) create mode 100644 src/infiniloader/formats.py create mode 100644 tests/test_format_validation.py diff --git a/src/infiniloader/distributed.py b/src/infiniloader/distributed.py index f1b8dd3..fcf6b04 100644 --- a/src/infiniloader/distributed.py +++ b/src/infiniloader/distributed.py @@ -4,7 +4,6 @@ import multiprocessing as mp import os import shutil -import struct import time import queue as queue_module import tempfile @@ -15,25 +14,7 @@ from .cache import DEFAULT_EXCLUDES, collect_source_files, human_bytes from .paths import dist_view_root from .runtime import validate_name, read_json - - -DTYPE_BYTES = { - "BOOL": 1, - "U8": 1, - "I8": 1, - "F8_E5M2": 1, - "F8_E4M3": 1, - "I16": 2, - "U16": 2, - "F16": 2, - "BF16": 2, - "I32": 4, - "U32": 4, - "F32": 4, - "F64": 8, - "I64": 8, - "U64": 8, -} +from .formats import DTYPE_BYTES, tensor_header, tensor_inventory @dataclass @@ -55,35 +36,14 @@ def nbytes(self) -> int: def parse_safetensors_header(path: Path) -> dict[str, Any]: - with path.open("rb") as f: - raw_len = f.read(8) - if len(raw_len) != 8: - raise ValueError(f"{path} is too small to be safetensors") - (header_len,) = struct.unpack(" 100 * 1024 * 1024 or header_len > path.stat().st_size - 8: - raise ValueError(f"invalid safetensors header length: {path}") - header = f.read(header_len) - return json.loads(header.decode("utf-8")) + return tensor_header(path, include_metadata=True)[0] def collect_safetensors_metadata(model_dir: Path) -> list[TensorInfo]: - model_dir = model_dir.expanduser().resolve() - rows: list[TensorInfo] = [] - for file_path in sorted(model_dir.glob("*.safetensors")): - header = parse_safetensors_header(file_path) - for name, item in header.items(): - if name == "__metadata__": - continue - rows.append( - TensorInfo( - file=file_path.relative_to(model_dir).as_posix(), - name=name, - dtype=item.get("dtype", ""), - shape=list(item.get("shape", [])), - data_offsets=list(item.get("data_offsets", [])), - ) - ) - return rows + _, tensors = tensor_inventory(model_dir) + return [TensorInfo(file=spec["file"], name=name, dtype=spec["dtype"], + shape=spec["shape"], data_offsets=spec["data_offsets"]) + for name, spec in tensors.items()] def infer_tp_split(tensor: TensorInfo, tp: int) -> dict[str, Any]: diff --git a/src/infiniloader/formats.py b/src/infiniloader/formats.py new file mode 100644 index 0000000..1f99667 --- /dev/null +++ b/src/infiniloader/formats.py @@ -0,0 +1,125 @@ +"""Shared safetensors metadata validation; payload hashing is a caller concern.""" +from __future__ import annotations + +import json +import math +import os +from pathlib import Path +import struct + +DTYPE_BYTES = { + "BOOL": 1, + "U8": 1, + "I8": 1, + "F8_E5M2": 1, + "F8_E4M3": 1, + "I16": 2, + "U16": 2, + "F16": 2, + "BF16": 2, + "I32": 4, + "U32": 4, + "F32": 4, + "F64": 8, + "I64": 8, + "U64": 8, +} + +MAX_HEADER = 100_000_000 + + +def _unique(pairs): + result = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate JSON key: {key}") + result[key] = value + return result + + +def read_json(path: Path): + return json.loads(path.read_text(encoding="utf-8"), object_pairs_hook=_unique) + + +def canonical(value) -> bytes: + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True, + allow_nan=False).encode() + + +def _signature(path: Path): + s = path.stat() + return s.st_dev, s.st_ino, s.st_size, s.st_mtime_ns, s.st_ctime_ns + + +def _int(value, *, minimum=0): + return type(value) is int and value >= minimum + + +def tensor_header(path: Path, *, include_metadata: bool = False) -> tuple[dict, int]: + """Reject holes, overlaps, truncated payloads and unsupported packed dtypes.""" + with path.open("rb") as stream: + size = os.fstat(stream.fileno()).st_size + raw = stream.read(8) + if len(raw) != 8: + raise ValueError(f"truncated safetensors: {path}") + length, = struct.unpack(" tuple[dict, dict]: + """Validate the complete tensor/index inventory without reading tensor data.""" + source = source.expanduser().resolve() + paths = sorted(source.rglob("*.safetensors")) + signatures = {path: _signature(path) for path in paths} + index_path = source / "model.safetensors.index.json" + index_signature = _signature(index_path) if index_path.exists() else None + files, tensors = {}, {} + for path in paths: + header, data_start = tensor_header(path) + relative = path.relative_to(source).as_posix() + files[relative] = {"size": signatures[path][2], "data_start": data_start} + for name, spec in header.items(): + if name in tensors: + raise ValueError(f"duplicate tensor across files: {name}") + tensors[name] = dict(spec, file=relative) + if not files or not tensors: + raise ValueError("no safetensors tensors found") + if index_signature is not None: + index = read_json(index_path) + if not isinstance(index, dict) or index.get("weight_map") != {name: spec["file"] for name, spec in tensors.items()}: + raise ValueError("safetensors index does not match tensor inventory") + if (paths != sorted(source.rglob("*.safetensors")) + or any(_signature(path) != signature for path, signature in signatures.items()) + or (_signature(index_path) if index_path.exists() else None) != index_signature): + raise RuntimeError("source inventory changed while reading metadata") + return files, tensors diff --git a/src/infiniloader/shards.py b/src/infiniloader/shards.py index 8a8269c..3a42b3b 100644 --- a/src/infiniloader/shards.py +++ b/src/infiniloader/shards.py @@ -17,28 +17,9 @@ import tempfile from pathlib import Path -from .distributed import DTYPE_BYTES +from .formats import DTYPE_BYTES, MAX_HEADER, _signature, canonical, read_json, tensor_header, tensor_inventory CHUNK = 1024 * 1024 -MAX_HEADER = 100_000_000 - - -def _unique(pairs): - result = {} - for key, value in pairs: - if key in result: - raise ValueError(f"duplicate JSON key: {key}") - result[key] = value - return result - - -def read_json(path: Path): - return json.loads(path.read_text(encoding="utf-8"), object_pairs_hook=_unique) - - -def canonical(value) -> bytes: - return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True, - allow_nan=False).encode() def digest(value) -> str: @@ -53,78 +34,28 @@ def file_hash(path: Path) -> str: return h.hexdigest() -def _signature(path: Path): - s = path.stat() - return s.st_dev, s.st_ino, s.st_size, s.st_mtime_ns, s.st_ctime_ns - - def _int(value, *, minimum=0): return type(value) is int and value >= minimum -def tensor_header(path: Path) -> tuple[dict, int]: - """Reject holes, overlaps, truncated payloads and unsupported packed dtypes.""" - with path.open("rb") as stream: - size = os.fstat(stream.fileno()).st_size - raw = stream.read(8) - if len(raw) != 8: - raise ValueError(f"truncated safetensors: {path}") - length, = struct.unpack(" dict: """Content identity is independent of hostname, inode and source directory.""" source = source.expanduser().resolve() - files, tensors = {}, {} paths = sorted(source.rglob("*.safetensors")) signatures = {p: _signature(p) for p in paths} - for path in paths: - before = signatures[path] - header, base = tensor_header(path) - rel = path.relative_to(source).as_posix() - files[rel] = {"sha256": file_hash(path), "size": before[2], "data_start": base} - if before != _signature(path): - raise RuntimeError(f"source changed while cataloging: {path}") - for name, spec in header.items(): - if name in tensors: - raise ValueError(f"duplicate tensor across files: {name}") - tensors[name] = dict(spec, file=rel) - if not files or not tensors: - raise ValueError("no safetensors tensors found") index_path = source / "model.safetensors.index.json" - if index_path.exists(): - index = read_json(index_path) - if not isinstance(index, dict) or index.get("weight_map") != {n: t["file"] for n, t in tensors.items()}: - raise ValueError("safetensors index does not match tensor inventory") - if paths != sorted(source.rglob("*.safetensors")) or any(_signature(p) != s for p, s in signatures.items()): + index_signature = _signature(index_path) if index_path.exists() else None + files, tensors = tensor_inventory(source) + if set(files) != {p.relative_to(source).as_posix() for p in paths}: + raise RuntimeError("source inventory changed while cataloging") + for relative, record in files.items(): + path = source / relative + record["sha256"] = file_hash(path) + if signatures[path] != _signature(path): + raise RuntimeError(f"source changed while cataloging: {path}") + if (paths != sorted(source.rglob("*.safetensors")) + or any(_signature(p) != s for p, s in signatures.items()) + or (_signature(index_path) if index_path.exists() else None) != index_signature): raise RuntimeError("source inventory changed while cataloging") result = {"format": 1, "files": files, "tensors": tensors} result["content_id"] = digest(result) diff --git a/tests/test_format_validation.py b/tests/test_format_validation.py new file mode 100644 index 0000000..8e74898 --- /dev/null +++ b/tests/test_format_validation.py @@ -0,0 +1,90 @@ +import json +import struct + +import pytest + +from infiniloader import distributed, shards +from infiniloader.formats import canonical, tensor_inventory + + +def write(path, header, payload=b"abcd"): + path.parent.mkdir(parents=True, exist_ok=True) + raw = canonical(header) if isinstance(header, dict) else header + path.write_bytes(struct.pack(" Date: Tue, 27 Oct 2026 02:02:43 +0800 Subject: [PATCH 16/29] Separate CLI arguments and derive completion from the parser --- src/infiniloader/arguments.py | 223 +++++++++++++++++++++++ src/infiniloader/cli.py | 323 +++------------------------------ src/infiniloader/completion.py | 89 +++++++++ tests/test_completion.py | 63 +++++++ 4 files changed, 401 insertions(+), 297 deletions(-) create mode 100644 src/infiniloader/arguments.py create mode 100644 src/infiniloader/completion.py create mode 100644 tests/test_completion.py diff --git a/src/infiniloader/arguments.py b/src/infiniloader/arguments.py new file mode 100644 index 0000000..1726cf5 --- /dev/null +++ b/src/infiniloader/arguments.py @@ -0,0 +1,223 @@ +"""Argument definitions for the public CLI and internal daemon entry point.""" +from __future__ import annotations + +import argparse +from collections.abc import Callable, Mapping +from pathlib import Path +import sys + +from . import __version__ +from .memory import parse_bytes + +CommandHandler = Callable[[argparse.Namespace], int] + + +def build_parser(handlers: Mapping[str, CommandHandler]) -> argparse.ArgumentParser: + prog = Path(sys.argv[0]).name + if prog in {"__main__.py", "python", "python3"}: + prog = "iloader" + parser = argparse.ArgumentParser(prog=prog, description="InfiniLoader: LLM weight loading and residency manager") + parser.add_argument("--version", action="version", version=f"InfiniLoader {__version__}") + sub = parser.add_subparsers(dest="command", required=True, metavar="COMMAND") + + inspect_p = sub.add_parser("inspect", help="inspect selected files and memory budget") + inspect_p.add_argument("model") + inspect_p.add_argument("--pattern", action="append", help="include glob; can be repeated") + inspect_p.add_argument("--exclude", action="append", help="exclude glob; replaces defaults when provided") + inspect_p.add_argument("--json", action="store_true") + inspect_p.set_defaults(func=handlers["cmd_inspect"]) + from .shards import configure_parser as configure_shards + configure_shards(sub) + + serve_p = sub.add_parser("_serve", help=argparse.SUPPRESS) + serve_p.add_argument("model") + serve_p.add_argument("--name", required=True) + serve_p.add_argument("--view") + serve_p.add_argument("--socket") + serve_p.add_argument("--pid-file") + serve_p.add_argument("--state-file", help=argparse.SUPPRESS) + serve_p.add_argument("--base-dir") + serve_p.add_argument("--upstream-source", help=argparse.SUPPRESS) + serve_p.add_argument("--local-cache-path", help=argparse.SUPPRESS) + serve_p.add_argument("--identity", help=argparse.SUPPRESS) + serve_p.add_argument("--pattern", action="append") + serve_p.add_argument("--exclude", action="append") + serve_p.add_argument("--replace", action="store_true") + serve_p.add_argument("--force", action="store_true") + serve_p.add_argument("--quiet", action="store_true") + serve_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) + serve_p.add_argument("--workers", type=int, default=1, help="parallel memfd preload workers") + serve_p.add_argument("--max-memory-fraction", type=float, default=0.90) + serve_p.add_argument("--local-cache", action="store_true", help="copy model to local cache before memfd preload") + serve_p.add_argument("--cache-root") + serve_p.add_argument("--cache-name") + serve_p.add_argument("--refresh-cache", action="store_true") + serve_p.add_argument("--cache-workers", type=int, default=4) + serve_p.set_defaults(func=handlers["cmd_serve"]) + sub._choices_actions = [action for action in sub._choices_actions if action.dest != "_serve"] + + start_p = sub.add_parser("start", help="start background daemon") + start_p.add_argument("model") + start_p.add_argument("--name", required=True) + start_p.add_argument("--view") + start_p.add_argument("--socket") + start_p.add_argument("--pid-file") + start_p.add_argument("--base-dir") + start_p.add_argument("--log-file") + start_p.add_argument("--pattern", action="append") + start_p.add_argument("--exclude", action="append") + start_p.add_argument("--force", action="store_true") + start_p.add_argument("--no-wait", action="store_true") + start_p.add_argument("--wait-timeout", type=float, default=600.0) + start_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) + start_p.add_argument("--workers", type=int, default=1, help="parallel memfd preload workers") + start_p.add_argument("--max-memory-fraction", type=float, default=0.90) + start_p.add_argument("--local-cache", action="store_true", help="copy model to local cache before memfd preload") + start_p.add_argument("--cache-root") + start_p.add_argument("--cache-name") + start_p.add_argument("--refresh-cache", action="store_true") + start_p.add_argument("--cache-workers", type=int, default=4) + start_p.set_defaults(func=handlers["cmd_start"]) + + for memory_parser in (serve_p, start_p): + memory_parser.add_argument("--max-memory-bytes", type=parse_bytes, help="per-cache preload cap including worker buffers") + memory_parser.add_argument("--reserve-bytes", type=parse_bytes, default=256 * 1024 * 1024, help="headroom left for the engine; default 256 MiB") + memory_parser.add_argument("--max-resident-bytes", type=parse_bytes, help="set persistent total resident quota for this base directory") + memory_parser.add_argument("--idle-ttl", type=float, help="evict after this many idle seconds; use run to protect readers") + + cache_p = sub.add_parser("cache", help="materialize a local model cache") + cache_p.add_argument("model") + cache_p.add_argument("--name") + cache_p.add_argument("--cache-name") + cache_p.add_argument("--cache-root") + cache_p.add_argument("--base-dir") + cache_p.add_argument("--pattern", action="append") + cache_p.add_argument("--exclude", action="append") + cache_p.add_argument("--refresh", action="store_true") + cache_p.add_argument("--cache-workers", type=int, default=4) + cache_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) + cache_p.add_argument("--quiet", action="store_true") + cache_p.add_argument("--json", action="store_true") + cache_p.set_defaults(func=handlers["cmd_cache"]) + + cache_status_p = sub.add_parser("cache-status", help="check whether a local model cache is valid") + cache_status_p.add_argument("model") + cache_status_p.add_argument("--name") + cache_status_p.add_argument("--cache-name") + cache_status_p.add_argument("--cache-root") + cache_status_p.add_argument("--base-dir") + cache_status_p.add_argument("--pattern", action="append") + cache_status_p.add_argument("--exclude", action="append") + cache_status_p.add_argument("--json", action="store_true") + cache_status_p.set_defaults(func=handlers["cmd_cache_status"]) + + cache_list_p = sub.add_parser("cache-list", help="list local model caches") + cache_list_p.add_argument("--cache-root") + cache_list_p.add_argument("--base-dir") + cache_list_p.add_argument("--json", action="store_true") + cache_list_p.set_defaults(func=handlers["cmd_cache_list"]) + + verify_p = sub.add_parser("verify", help="verify a versioned local cache using SHA-256") + verify_p.add_argument("path") + verify_p.add_argument("--quick", action="store_true", help="check file metadata without reading all bytes") + verify_p.add_argument("--json", action="store_true") + verify_p.set_defaults(func=handlers["cmd_verify"]) + + gc_p = sub.add_parser("gc", help="preview or apply LRU/TTL collection of unleased disk caches") + gc_p.add_argument("--base-dir") + gc_p.add_argument("--cache-root") + gc_p.add_argument("--max-bytes", type=parse_bytes, help="target total cache bytes; accepts GiB/MiB") + gc_p.add_argument("--resident", action="store_true", help="collect resident daemons instead of disk caches") + gc_p.add_argument("--ttl", type=float, help="collect caches unused for this many seconds") + gc_p.add_argument("--apply", action="store_true", help="delete the selected caches; default is dry run") + gc_p.add_argument("--json", action="store_true") + gc_p.set_defaults(func=handlers["cmd_gc"]) + + run_p = sub.add_parser("run", help="run a foreground command while protecting its cache from eviction") + run_p.add_argument("--base-dir") + run_p.add_argument("--cache-path") + run_p.add_argument("name_arg", nargs="?") + run_p.add_argument("program", nargs=argparse.REMAINDER) + run_p.set_defaults(func=handlers["cmd_run"]) + + doctor_p = sub.add_parser("doctor", help="check Linux, memfd, proc access, locks and resource limits") + doctor_p.add_argument("--base-dir") + doctor_p.add_argument("--json", action="store_true") + doctor_p.set_defaults(func=handlers["cmd_doctor"]) + + list_p = sub.add_parser("list", help="list active resident caches") + list_p.add_argument("--base-dir") + list_p.add_argument("--timeout", type=float, default=1.0) + list_p.add_argument("--json", action="store_true") + list_p.add_argument("--names-only", action="store_true", help=argparse.SUPPRESS) + list_p.set_defaults(func=handlers["cmd_list"]) + + dist_views_p = sub.add_parser("dist-views", help="create full HF-compatible model views for simulated ranks") + dist_views_p.add_argument("model") + dist_views_p.add_argument("--name", required=True) + dist_views_p.add_argument("--ranks", type=int, required=True) + dist_views_p.add_argument("--base-dir") + dist_views_p.add_argument("--replace", action="store_true") + dist_views_p.add_argument("--json", action="store_true") + dist_views_p.set_defaults(func=handlers["cmd_dist_views"]) + + dist_sim_p = sub.add_parser("dist-simulate", help="simulate concurrent multi-rank model reads on this host") + dist_sim_p.add_argument("model") + dist_sim_p.add_argument("--ranks", type=int, required=True) + dist_sim_p.add_argument("--rank-view-root", help="root containing rank0/rank1/... directories") + dist_sim_p.add_argument("--pattern", action="append") + dist_sim_p.add_argument("--exclude", action="append") + dist_sim_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) + dist_sim_p.add_argument("--timeout", type=float, default=60.0) + dist_sim_p.add_argument("--json", action="store_true") + dist_sim_p.set_defaults(func=handlers["cmd_dist_simulate"]) + + dist_plan_p = sub.add_parser("dist-plan", help="analyze safetensors metadata for TP split candidates") + dist_plan_p.add_argument("model") + dist_plan_p.add_argument("--tp", type=int, required=True) + dist_plan_p.add_argument("--limit", type=int, default=40) + dist_plan_p.add_argument("--json", action="store_true") + dist_plan_p.set_defaults(func=handlers["cmd_dist_plan"]) + + for name, help_text, func in [ + ("status", "show daemon status", handlers["cmd_status"]), + ("view", "print daemon view path", handlers["cmd_view"]), + ("manifest", "print daemon manifest", handlers["cmd_manifest"]), + ]: + p = sub.add_parser(name, help=help_text) + p.add_argument("name_arg", nargs="?") + p.add_argument("--name") + p.add_argument("--socket") + p.add_argument("--base-dir") + p.add_argument("--timeout", type=float, default=5.0) + if name in {"status"}: + p.add_argument("--json", action="store_true") + p.set_defaults(func=func) + + stop_p = sub.add_parser("stop", help="stop daemon") + stop_p.add_argument("name_arg", nargs="?") + stop_p.add_argument("--name") + stop_p.add_argument("--all", action="store_true", help="stop every active resident cache in the run directory") + stop_p.add_argument("--socket") + stop_p.add_argument("--base-dir") + stop_p.add_argument("--timeout", type=float, default=5.0) + stop_p.add_argument("--remove-view", action="store_true") + stop_p.add_argument("--force", action="store_true", help="stop even when commands hold active leases") + stop_p.set_defaults(func=handlers["cmd_stop"]) + + bench_p = sub.add_parser("bench", help="sequential read benchmark") + bench_p.add_argument("path") + bench_p.add_argument("--pattern", action="append") + bench_p.add_argument("--exclude", action="append") + bench_p.add_argument("--repeat", type=int, default=3) + bench_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) + bench_p.add_argument("--json", action="store_true") + bench_p.set_defaults(func=handlers["cmd_bench"]) + + selftest_p = sub.add_parser("selftest", help="verify memfd view with open/stat/mmap") + selftest_p.set_defaults(func=handlers["cmd_selftest"]) + + completion_p = sub.add_parser("completion", help="print shell completion script") + completion_p.add_argument("shell", choices=["bash"]) + completion_p.set_defaults(func=handlers["cmd_completion"]) + return parser diff --git a/src/infiniloader/cli.py b/src/infiniloader/cli.py index 62375b5..38bc2b5 100644 --- a/src/infiniloader/cli.py +++ b/src/infiniloader/cli.py @@ -10,11 +10,11 @@ import time from pathlib import Path -from . import __version__ +from .arguments import build_parser as argument_parser +from .completion import bash_completion from .benchmark import benchmark_path from .cache import DEFAULT_EXCLUDES, build_view, human_bytes, inspect_source, load_model_into_memory, source_identity, _safe_remove_view from .runtime import FileLock, process_matches, process_token, read_json, signal_owned_process, validate_name, validate_socket_path -from .memory import parse_bytes from .daemon import WeightLoaderDaemon from .distributed import build_tp_plan, create_rank_views, simulate_distributed_read from .ipc import LoaderConnectionError, LoaderStoppingError, request @@ -766,95 +766,9 @@ def _run_program(program, path, *, pass_fds=(), socket_path=None, lease=None) -> raise -BASH_COMPLETION = r'''_iloader_complete() { - local cur prev sub commands opts - COMPREPLY=() - cur="${COMP_WORDS[COMP_CWORD]}" - prev="${COMP_WORDS[COMP_CWORD-1]}" - commands="inspect start list status view manifest stop bench selftest cache cache-status cache-list verify gc run doctor dist-views dist-simulate dist-plan completion" - - if [[ ${COMP_CWORD} -eq 1 ]]; then - COMPREPLY=( $(compgen -W "${commands}" -- "${cur}") ) - return 0 - fi - - case "${prev}" in - --name) - COMPREPLY=( $(compgen -W "$(iloader list --names-only 2>/dev/null)" -- "${cur}") ) - return 0 - ;; - --device) - COMPREPLY=( $(compgen -W "cpu cuda cuda:0" -- "${cur}") ) - return 0 - ;; - --base-dir|--view|--cache-root|--log-file|--pid-file|--socket|--rank-view-root) - compopt -o filenames 2>/dev/null || true - COMPREPLY=( $(compgen -f -- "${cur}") ) - return 0 - ;; - esac - - sub="${COMP_WORDS[1]}" - case "${sub}" in - start) - opts="--name --view --socket --pid-file --base-dir --log-file --pattern --exclude --force --no-wait --wait-timeout --chunk-size --workers --max-memory-fraction --max-memory-bytes --max-resident-bytes --reserve-bytes --idle-ttl --local-cache --cache-root --cache-name --refresh-cache --cache-workers --help" - ;; - list) - opts="--base-dir --timeout --json --names-only --help" - ;; - status|view|manifest) - opts="--name --socket --base-dir --timeout --json --help" - ;; - stop) - opts="--name --all --socket --base-dir --timeout --remove-view --force --help" - ;; - cache|cache-status) - opts="--name --cache-name --cache-root --base-dir --pattern --exclude --refresh --cache-workers --chunk-size --quiet --json --help" - ;; - cache-list) - opts="--cache-root --base-dir --json --help" - ;; - verify) - opts="--quick --json --help" - ;; - gc) - opts="--base-dir --cache-root --max-bytes --ttl --resident --apply --json --help" - ;; - run) - opts="--base-dir --cache-path --help" - ;; - doctor) - opts="--base-dir --json --help" - ;; - dist-views) - opts="--name --ranks --base-dir --replace --json --help" - ;; - dist-simulate) - opts="--ranks --rank-view-root --pattern --exclude --chunk-size --timeout --json --help" - ;; - dist-plan) - opts="--tp --limit --json --help" - ;; - completion) - opts="bash --help" - ;; - *) - opts="--help --json --base-dir --name" - ;; - esac - - if [[ "${cur}" == -* ]]; then - COMPREPLY=( $(compgen -W "${opts}" -- "${cur}") ) - return 0 - fi -} -complete -F _iloader_complete iloader -''' - - def cmd_completion(args: argparse.Namespace) -> int: if args.shell == "bash": - print(BASH_COMPLETION) + print(bash_completion(build_parser())) return 0 raise ValueError(f"unsupported shell: {args.shell}") @@ -989,214 +903,29 @@ def cmd_selftest(_args: argparse.Namespace) -> int: def build_parser() -> argparse.ArgumentParser: - prog = Path(sys.argv[0]).name - if prog in {"__main__.py", "python", "python3"}: - prog = "iloader" - parser = argparse.ArgumentParser(prog=prog, description="InfiniLoader: LLM weight loading and residency manager") - parser.add_argument("--version", action="version", version=f"InfiniLoader {__version__}") - sub = parser.add_subparsers(dest="command", required=True, metavar="COMMAND") - - inspect_p = sub.add_parser("inspect", help="inspect selected files and memory budget") - inspect_p.add_argument("model") - inspect_p.add_argument("--pattern", action="append", help="include glob; can be repeated") - inspect_p.add_argument("--exclude", action="append", help="exclude glob; replaces defaults when provided") - inspect_p.add_argument("--json", action="store_true") - inspect_p.set_defaults(func=cmd_inspect) - from .shards import configure_parser as configure_shards - configure_shards(sub) - - serve_p = sub.add_parser("_serve", help=argparse.SUPPRESS) - serve_p.add_argument("model") - serve_p.add_argument("--name", required=True) - serve_p.add_argument("--view") - serve_p.add_argument("--socket") - serve_p.add_argument("--pid-file") - serve_p.add_argument("--state-file", help=argparse.SUPPRESS) - serve_p.add_argument("--base-dir") - serve_p.add_argument("--upstream-source", help=argparse.SUPPRESS) - serve_p.add_argument("--local-cache-path", help=argparse.SUPPRESS) - serve_p.add_argument("--identity", help=argparse.SUPPRESS) - serve_p.add_argument("--pattern", action="append") - serve_p.add_argument("--exclude", action="append") - serve_p.add_argument("--replace", action="store_true") - serve_p.add_argument("--force", action="store_true") - serve_p.add_argument("--quiet", action="store_true") - serve_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) - serve_p.add_argument("--workers", type=int, default=1, help="parallel memfd preload workers") - serve_p.add_argument("--max-memory-fraction", type=float, default=0.90) - serve_p.add_argument("--local-cache", action="store_true", help="copy model to local cache before memfd preload") - serve_p.add_argument("--cache-root") - serve_p.add_argument("--cache-name") - serve_p.add_argument("--refresh-cache", action="store_true") - serve_p.add_argument("--cache-workers", type=int, default=4) - serve_p.set_defaults(func=cmd_serve) - sub._choices_actions = [action for action in sub._choices_actions if action.dest != "_serve"] - - start_p = sub.add_parser("start", help="start background daemon") - start_p.add_argument("model") - start_p.add_argument("--name", required=True) - start_p.add_argument("--view") - start_p.add_argument("--socket") - start_p.add_argument("--pid-file") - start_p.add_argument("--base-dir") - start_p.add_argument("--log-file") - start_p.add_argument("--pattern", action="append") - start_p.add_argument("--exclude", action="append") - start_p.add_argument("--force", action="store_true") - start_p.add_argument("--no-wait", action="store_true") - start_p.add_argument("--wait-timeout", type=float, default=600.0) - start_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) - start_p.add_argument("--workers", type=int, default=1, help="parallel memfd preload workers") - start_p.add_argument("--max-memory-fraction", type=float, default=0.90) - start_p.add_argument("--local-cache", action="store_true", help="copy model to local cache before memfd preload") - start_p.add_argument("--cache-root") - start_p.add_argument("--cache-name") - start_p.add_argument("--refresh-cache", action="store_true") - start_p.add_argument("--cache-workers", type=int, default=4) - start_p.set_defaults(func=cmd_start) - - for memory_parser in (serve_p, start_p): - memory_parser.add_argument("--max-memory-bytes", type=parse_bytes, help="per-cache preload cap including worker buffers") - memory_parser.add_argument("--reserve-bytes", type=parse_bytes, default=256 * 1024 * 1024, help="headroom left for the engine; default 256 MiB") - memory_parser.add_argument("--max-resident-bytes", type=parse_bytes, help="set persistent total resident quota for this base directory") - memory_parser.add_argument("--idle-ttl", type=float, help="evict after this many idle seconds; use run to protect readers") - - cache_p = sub.add_parser("cache", help="materialize a local model cache") - cache_p.add_argument("model") - cache_p.add_argument("--name") - cache_p.add_argument("--cache-name") - cache_p.add_argument("--cache-root") - cache_p.add_argument("--base-dir") - cache_p.add_argument("--pattern", action="append") - cache_p.add_argument("--exclude", action="append") - cache_p.add_argument("--refresh", action="store_true") - cache_p.add_argument("--cache-workers", type=int, default=4) - cache_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) - cache_p.add_argument("--quiet", action="store_true") - cache_p.add_argument("--json", action="store_true") - cache_p.set_defaults(func=cmd_cache) - - cache_status_p = sub.add_parser("cache-status", help="check whether a local model cache is valid") - cache_status_p.add_argument("model") - cache_status_p.add_argument("--name") - cache_status_p.add_argument("--cache-name") - cache_status_p.add_argument("--cache-root") - cache_status_p.add_argument("--base-dir") - cache_status_p.add_argument("--pattern", action="append") - cache_status_p.add_argument("--exclude", action="append") - cache_status_p.add_argument("--json", action="store_true") - cache_status_p.set_defaults(func=cmd_cache_status) - - cache_list_p = sub.add_parser("cache-list", help="list local model caches") - cache_list_p.add_argument("--cache-root") - cache_list_p.add_argument("--base-dir") - cache_list_p.add_argument("--json", action="store_true") - cache_list_p.set_defaults(func=cmd_cache_list) - - verify_p = sub.add_parser("verify", help="verify a versioned local cache using SHA-256") - verify_p.add_argument("path") - verify_p.add_argument("--quick", action="store_true", help="check file metadata without reading all bytes") - verify_p.add_argument("--json", action="store_true") - verify_p.set_defaults(func=cmd_verify) - - gc_p = sub.add_parser("gc", help="preview or apply LRU/TTL collection of unleased disk caches") - gc_p.add_argument("--base-dir") - gc_p.add_argument("--cache-root") - gc_p.add_argument("--max-bytes", type=parse_bytes, help="target total cache bytes; accepts GiB/MiB") - gc_p.add_argument("--resident", action="store_true", help="collect resident daemons instead of disk caches") - gc_p.add_argument("--ttl", type=float, help="collect caches unused for this many seconds") - gc_p.add_argument("--apply", action="store_true", help="delete the selected caches; default is dry run") - gc_p.add_argument("--json", action="store_true") - gc_p.set_defaults(func=cmd_gc) - - run_p = sub.add_parser("run", help="run a foreground command while protecting its cache from eviction") - run_p.add_argument("--base-dir") - run_p.add_argument("--cache-path") - run_p.add_argument("name_arg", nargs="?") - run_p.add_argument("program", nargs=argparse.REMAINDER) - run_p.set_defaults(func=cmd_run) - - doctor_p = sub.add_parser("doctor", help="check Linux, memfd, proc access, locks and resource limits") - doctor_p.add_argument("--base-dir") - doctor_p.add_argument("--json", action="store_true") - doctor_p.set_defaults(func=cmd_doctor) - - list_p = sub.add_parser("list", help="list active resident caches") - list_p.add_argument("--base-dir") - list_p.add_argument("--timeout", type=float, default=1.0) - list_p.add_argument("--json", action="store_true") - list_p.add_argument("--names-only", action="store_true", help=argparse.SUPPRESS) - list_p.set_defaults(func=cmd_list) - - dist_views_p = sub.add_parser("dist-views", help="create full HF-compatible model views for simulated ranks") - dist_views_p.add_argument("model") - dist_views_p.add_argument("--name", required=True) - dist_views_p.add_argument("--ranks", type=int, required=True) - dist_views_p.add_argument("--base-dir") - dist_views_p.add_argument("--replace", action="store_true") - dist_views_p.add_argument("--json", action="store_true") - dist_views_p.set_defaults(func=cmd_dist_views) - - dist_sim_p = sub.add_parser("dist-simulate", help="simulate concurrent multi-rank model reads on this host") - dist_sim_p.add_argument("model") - dist_sim_p.add_argument("--ranks", type=int, required=True) - dist_sim_p.add_argument("--rank-view-root", help="root containing rank0/rank1/... directories") - dist_sim_p.add_argument("--pattern", action="append") - dist_sim_p.add_argument("--exclude", action="append") - dist_sim_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) - dist_sim_p.add_argument("--timeout", type=float, default=60.0) - dist_sim_p.add_argument("--json", action="store_true") - dist_sim_p.set_defaults(func=cmd_dist_simulate) - - dist_plan_p = sub.add_parser("dist-plan", help="analyze safetensors metadata for TP split candidates") - dist_plan_p.add_argument("model") - dist_plan_p.add_argument("--tp", type=int, required=True) - dist_plan_p.add_argument("--limit", type=int, default=40) - dist_plan_p.add_argument("--json", action="store_true") - dist_plan_p.set_defaults(func=cmd_dist_plan) - - for name, help_text, func in [ - ("status", "show daemon status", cmd_status), - ("view", "print daemon view path", cmd_view), - ("manifest", "print daemon manifest", cmd_manifest), - ]: - p = sub.add_parser(name, help=help_text) - p.add_argument("name_arg", nargs="?") - p.add_argument("--name") - p.add_argument("--socket") - p.add_argument("--base-dir") - p.add_argument("--timeout", type=float, default=5.0) - if name in {"status"}: - p.add_argument("--json", action="store_true") - p.set_defaults(func=func) - - stop_p = sub.add_parser("stop", help="stop daemon") - stop_p.add_argument("name_arg", nargs="?") - stop_p.add_argument("--name") - stop_p.add_argument("--all", action="store_true", help="stop every active resident cache in the run directory") - stop_p.add_argument("--socket") - stop_p.add_argument("--base-dir") - stop_p.add_argument("--timeout", type=float, default=5.0) - stop_p.add_argument("--remove-view", action="store_true") - stop_p.add_argument("--force", action="store_true", help="stop even when commands hold active leases") - stop_p.set_defaults(func=cmd_stop) - - bench_p = sub.add_parser("bench", help="sequential read benchmark") - bench_p.add_argument("path") - bench_p.add_argument("--pattern", action="append") - bench_p.add_argument("--exclude", action="append") - bench_p.add_argument("--repeat", type=int, default=3) - bench_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) - bench_p.add_argument("--json", action="store_true") - bench_p.set_defaults(func=cmd_bench) - - selftest_p = sub.add_parser("selftest", help="verify memfd view with open/stat/mmap") - selftest_p.set_defaults(func=cmd_selftest) - - completion_p = sub.add_parser("completion", help="print shell completion script") - completion_p.add_argument("shell", choices=["bash"]) - completion_p.set_defaults(func=cmd_completion) - return parser + return argument_parser({ + "cmd_bench": cmd_bench, + "cmd_cache": cmd_cache, + "cmd_cache_list": cmd_cache_list, + "cmd_cache_status": cmd_cache_status, + "cmd_completion": cmd_completion, + "cmd_dist_plan": cmd_dist_plan, + "cmd_dist_simulate": cmd_dist_simulate, + "cmd_dist_views": cmd_dist_views, + "cmd_doctor": cmd_doctor, + "cmd_gc": cmd_gc, + "cmd_inspect": cmd_inspect, + "cmd_list": cmd_list, + "cmd_manifest": cmd_manifest, + "cmd_run": cmd_run, + "cmd_selftest": cmd_selftest, + "cmd_serve": cmd_serve, + "cmd_start": cmd_start, + "cmd_status": cmd_status, + "cmd_stop": cmd_stop, + "cmd_verify": cmd_verify, + "cmd_view": cmd_view, + }) def main(argv: list[str] | None = None) -> int: diff --git a/src/infiniloader/completion.py b/src/infiniloader/completion.py new file mode 100644 index 0000000..0f25913 --- /dev/null +++ b/src/infiniloader/completion.py @@ -0,0 +1,89 @@ +"""Bash completion derived from the same parser used to execute commands.""" +from __future__ import annotations + +import argparse +import shlex + + +def bash_completion(parser: argparse.ArgumentParser) -> str: + transitions, cases, values = [], [], [] + + def visit(node: argparse.ArgumentParser, prefix: str) -> None: + options = [] + positional_choices = [] + children = {} + for action in node._actions: + if action.help == argparse.SUPPRESS: + continue + options.extend(action.option_strings) + if action.option_strings and action.choices is not None: + for option in action.option_strings: + choices = " ".join(str(choice) for choice in action.choices) + values.append(f" {shlex.quote(prefix + '|' + option)}) choices={shlex.quote(choices)} ;;") + if isinstance(action, argparse._SubParsersAction): + children.update({name: sub for name, sub in action.choices.items() if not name.startswith("_")}) + elif not action.option_strings and action.choices is not None: + positional_choices.extend(str(choice) for choice in action.choices) + commands = list(children) or positional_choices + cases.append(f" {shlex.quote(prefix)}) opts={shlex.quote(' '.join(options))}; commands={shlex.quote(' '.join(commands))} ;;") + for name, sub in children.items(): + child = f"{prefix} {name}".strip() + transitions.append(f" {shlex.quote(prefix + '|' + name)}) node={shlex.quote(child)} ;;") + visit(sub, child) + + visit(parser, "") + script = r'''_iloader_complete() { + COMPREPLY=() + local cur="${COMP_WORDS[COMP_CWORD]}" prev="${COMP_WORDS[COMP_CWORD-1]}" + local node="" word i opts="" commands="" choices="" base="" previous="" + local run_name=0 run_cache=0 run_skip=0 + for ((i=1; i/dev/null || true + mapfile -t COMPREPLY < <(compgen -f -- "$cur") + elif [[ "$prev" == --name ]] || { [[ "$cur" != -* ]] && [[ "$node" =~ ^(status|view|manifest|stop|run)$ ]] && [[ "$prev" != --base-dir ]]; }; then + local -a query=(list --names-only) + if [[ -n "$base" ]]; then query+=(--base-dir "$base"); fi + mapfile -t COMPREPLY < <(compgen -W "$(iloader "${query[@]}" 2>/dev/null)" -- "$cur") + elif [[ "$cur" == -* ]]; then + mapfile -t COMPREPLY < <(compgen -W "$opts" -- "$cur") + elif [[ -n "$commands" ]]; then + mapfile -t COMPREPLY < <(compgen -W "$commands" -- "$cur") + else + compopt -o filenames 2>/dev/null || true + mapfile -t COMPREPLY < <(compgen -f -- "$cur") + fi +} +complete -F _iloader_complete iloader +''' + return script.replace("@TRANSITIONS@", "\n".join(transitions)).replace("@CASES@", "\n".join(cases)).replace("@VALUES@", "\n".join(values)) diff --git a/tests/test_completion.py b/tests/test_completion.py new file mode 100644 index 0000000..41af4e7 --- /dev/null +++ b/tests/test_completion.py @@ -0,0 +1,63 @@ +import os +from pathlib import Path +import shlex +import shutil +import subprocess + +import pytest + +from infiniloader import cli +from infiniloader.completion import bash_completion + + +def complete(words): + script = bash_completion(cli.build_parser()) + script += "\nCOMP_WORDS=(" + " ".join(shlex.quote(word) for word in words) + ")\n" + script += f"COMP_CWORD={len(words) - 1}\n_iloader_complete\nprintf '%s\\n' \"${{COMPREPLY[@]}}\"\n" + result = subprocess.run(["bash", "-c", script], capture_output=True, text=True, check=True) + return set(result.stdout.splitlines()) - {""} + + +@pytest.mark.skipif(os.name != "posix" or not shutil.which("bash"), reason="Bash completion integration") +@pytest.mark.parametrize("words,expected", [ + (["iloader", "sh"], {"shard"}), + (["iloader", "_"], set()), + (["iloader", "shard", ""], {"plan", "export", "verify"}), + (["iloader", "shard", "plan", "--mode", "r"], {"replicated"}), + (["iloader", "run", "unit", "--", "--"], set()), + (["iloader", "run", "--cache-path", "/tmp/cache", "--"], set()), + (["iloader", "run", "--cache-path=/tmp/cache", "--"], set()), + (["iloader", "completion", "b"], {"bash"}), +]) +def test_actual_bash_completion_follows_command_context(words, expected): + assert complete(words) == expected + + +@pytest.mark.skipif(os.name != "posix" or not shutil.which("bash"), reason="Bash completion integration") +def test_nested_options_match_the_parser(): + options = complete(["iloader", "shard", "plan", "--"]) + assert {"--layout", "--output", "--mode", "--help"} <= options + assert "--rank" not in options + assert "--idle-ttl" in complete(["iloader", "start", "--"]) + + +@pytest.mark.skipif(os.name != "posix" or not shutil.which("bash"), reason="Bash completion integration") +def test_path_values_take_precedence_and_preserve_spaces(tmp_path): + directory = tmp_path / "cache directory" + directory.mkdir() + prefix = str(directory)[:-3] + assert str(directory) in complete(["iloader", "run", "--cache-path", prefix]) + assert str(directory) in complete(["iloader", "status", "--socket", prefix]) + + +def test_parser_handlers_and_remainder_contract_are_preserved(): + parser = cli.build_parser() + args = parser.parse_args(["start", "/model", "--name", "unit"]) + assert args.func is cli.cmd_start + args = parser.parse_args(["run", "--base-dir", "/state", "unit", "--", "engine", "--help"]) + assert args.func is cli.cmd_run + assert args.name_arg == "unit" and args.program == ["engine", "--help"] + args = parser.parse_args(["run", "--cache-path", "/cache", "engine", "--help"]) + assert args.name_arg == "engine" and args.program == ["--help"] + args = parser.parse_args(["shard", "plan", "/model", "--layout", "layout.json", "--output", "plan.json"]) + assert args.source == Path("/model") and args.shard_action == "plan" From f6ba2e569947ac65876f0f721f2704ebc929f466 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 02:17:54 +0800 Subject: [PATCH 17/29] Clarify cache usage and archive historical development records --- CONTRIBUTING.md | 49 +++ README.md | 341 ++++-------------- README.zh-CN.md | 273 +++----------- .../2026-10-18-distributed-results.json} | 0 .../2026-10-18-distributed-validation.md} | 6 +- .../2026-10-18-expert-review.md} | 0 .../2026-10-18-reliability-validation.md} | 2 +- .../2026-10-26-history-reorganization.json} | 0 .../2026-10-26-history-reorganization.md} | 6 +- docs/archive/2026-10-26-roadmap.zh-CN.md | 47 +++ docs/archive/README.md | 20 + docs/archive/original-engine-measurements.md | 27 ++ docs/distributed.md | 9 +- docs/operations.md | 41 ++- docs/roadmap.zh-CN.md | 50 +-- 15 files changed, 304 insertions(+), 567 deletions(-) create mode 100644 CONTRIBUTING.md rename docs/{distributed-results.json => archive/2026-10-18-distributed-results.json} (100%) rename docs/{distributed-validation.md => archive/2026-10-18-distributed-validation.md} (97%) rename docs/{expert-review.md => archive/2026-10-18-expert-review.md} (100%) rename docs/{validation.md => archive/2026-10-18-reliability-validation.md} (96%) rename docs/{history-reorganization.json => archive/2026-10-26-history-reorganization.json} (100%) rename docs/{history-reorganization.md => archive/2026-10-26-history-reorganization.md} (95%) create mode 100644 docs/archive/2026-10-26-roadmap.zh-CN.md create mode 100644 docs/archive/README.md create mode 100644 docs/archive/original-engine-measurements.md diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..2d5d729 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,49 @@ +# Contributing + +Keep changes focused on the cache's file interface, correctness and usability. +Optional tensor/device integrations should remain separate from the standard-library +core. Include a regression that reproduces a behavior failure rather than testing +only the implementation's current structure. + +## Core checks + +Use Linux with Python 3.10+: + +```bash +python3 -m pip install -e '.[test]' +python3 -m pytest -q +iloader selftest +iloader doctor +``` + +The file-cache tests intentionally use arbitrary file bytes; they test filesystem +and lifecycle behavior, not model semantics. Runtime tests use real subprocesses, +memfds and Unix sockets. Fault-injection tests use only their own temporary paths. + +## Optional PyTorch checks + +Install a PyTorch build appropriate for the target device, then: + +```bash +python3 -m pip install -e '.[test]' safetensors numpy +python3 -m pytest -q +python3 examples/hybrid_moe.py --device cpu --output hybrid-cpu.json +``` + +The optional tests cover tensor views, real local Gloo TP/PP execution and a small +reference MoE. Device tests skip when no compatible GPU is available; a CPU pass +does not validate CUDA/ROCm execution. Remote/device measurements must state their +commit, hardware, software versions and actual commands separately. + +## Review and documentation + +- Run the checks relevant to changed behavior and include their actual output. +- Preserve ordinary files and unknown runtime resources when testing cleanup. +- Keep CLI help, completion and both READMEs consistent. Completion is generated + from the argument parser and has executable Bash tests. +- Document memory ownership and lifetime when changing descriptors, mappings, + leases or device copies. +- Performance claims need a named baseline, reproducible commands and repeated + measurements; file-read bandwidth is not engine-startup latency. +- Historical maintenance records belong in [the archive](docs/archive/README.md), + while current operational contracts belong in [operations](docs/operations.md). diff --git a/README.md b/README.md index d211f07..6f460ec 100644 --- a/README.md +++ b/README.md @@ -1,309 +1,100 @@ -# InfiniLoader +# InfiniLoader -[中文](README.zh-CN.md) +[中文说明](README.zh-CN.md) -The development branch adds model identity checks, single-owner startup, sealed -memfd files, versioned disk caches, container-aware budgets, and lease-aware cache -management. See [operations and migration](docs/operations.md) for `doctor`, -`verify`, `run`, `gc`, and container guidance. +InfiniLoader keeps model files in a long-lived cache so inference-engine processes +can restart without fetching the same weights from their original storage again. +It is a Linux userspace cache: it does not implement model execution or request scheduling. -```bash -iloader doctor -iloader start "$MODEL" --name dev --local-cache \ - --max-memory-bytes 32GiB --reserve-bytes 4GiB --max-resident-bytes 64GiB -iloader run dev -- vllm serve '{model}' -iloader gc --resident --max-bytes 32GiB # preview -iloader gc --resident --max-bytes 32GiB --apply -``` - -Use `iloader cache` plus `iloader run --cache-path PATH -- COMMAND` for disk-only -caching. Existing direct-path readers remain supported; eviction protection requires -`run`, so leave automatic idle eviction disabled for unregistered readers. - -`InfiniLoader` is a userspace weight loading manager for LLM inference -framework development. It is designed for the loop where you repeatedly restart -vLLM, SGLang, InfiniLM, or a custom engine while debugging loader, kernel, -scheduler, or runtime code, and every restart pays the cost of reading model -weights from storage again. - -The default mode keeps a model directory resident in memory and exposes a normal -filesystem-looking Hugging Face model directory. Frameworks keep using their -existing file loaders; in the common path you only change the model path. - -## Quick Start - -The basic path is a memfd-backed, memfs-style in-memory view. If you usually run -a framework against `/path/to/huggingface-model`, add one command before the -framework launch and use the returned path instead. +The core uses Python's standard library. It copies selected files into sealed +`memfd` objects and exposes a directory of links to `/proc//fd/`. +Consumers keep their existing file loaders and use the published directory as the model path. -```bash -cd InfiniLoader -source env.sh - -MODEL=/path/to/huggingface-model -MODEL_VIEW=$(iloader start "$MODEL" --name dev-model) -``` +## Install and check the environment -Then pass `"$MODEL_VIEW"` to the inference engine as the model path. The -InfiniLoader command is the same; only the engine launch command changes. - -```bash -# InfiniLM -python python/infinilm/server/inference_server.py \ - --device cuda \ - --model="$MODEL_VIEW" \ - --tp=1 -``` +Linux and Python 3.10+ are required. Install inside the engine's environment, using +the same user and PID namespace as the consumer: ```bash -# vLLM -vllm serve "$MODEL_VIEW" --tensor-parallel-size 1 -``` - -```bash -# SGLang -python -m sglang.launch_server \ - --model-path "$MODEL_VIEW" \ - --tp 1 -``` - -When you are done: - -```bash -iloader stop dev-model --remove-view -``` - -That is the stable compatibility mode for 0.1.0: no framework loader code -changes, only the model path changes. vLLM, SGLang, InfiniLM, local cache, -distributed views, progress/status, and completion are covered below. - -## What It Solves - -Repeated LLM engine debugging often spends more time loading weights than -running the code under test: - -- network filesystems can be slow or bursty for large safetensors shards; -- hot OS page cache helps, but does not remove framework parsing and H2D costs; -- distributed runs may cause every rank to scan or read the same checkpoint; -- modifying framework code should not require editing model loading APIs. - -`InfiniLoader` provides a long-lived loader service that can outlive framework -processes. You can stop and restart vLLM/SGLang/InfiniLM while the weight view -stays warm. - -## How It Works - -The default `start/view` path is Linux-only and does not require root or FUSE: - -1. A background daemon reads selected files from a model directory. -2. Each file is copied into an anonymous `memfd_create` file. -3. The daemon keeps those file descriptors open. -4. A view directory mirrors the original model tree with symlinks to - `/proc//fd/`. -5. Frameworks use ordinary `open`, `stat`, and `mmap` calls against that view. - -Optional stable layers are available for local disk caching, distributed rank -views, and tensor-parallel split analysis. - -## Setup And Runtime State - -`env.sh` adds the repository `src` directory to `PYTHONPATH` and defines an -`iloader` shell function for the current terminal. It does not require a wrapper -script under `bin`. In bash, it also registers command completion for iloader -subcommands, options, and active cache names. - -You can also install InfiniLoader as a Python package: - -```bash -pip install -e . +python3 -m pip install . +iloader doctor iloader selftest ``` -Run the unit test suite with: - -```bash -pip install -e .[test] -python -m pytest -``` - -Runtime state is placed under `$ILOADER_BASE_DIR` when it is set, otherwise -under `$XDG_CACHE_HOME/infiniloader` or `~/.cache/infiniloader`. Every command -that creates runtime state also accepts `--base-dir`. - -## Command Surface - -The public 0.1.0 interface is intentionally small: - -- Primary workflow: `start`, `list`, `status`, `view`, `stop`. -- Utilities: `inspect`, `bench`, `completion`, `selftest`. -- Advanced but stable helpers: `cache`, `cache-status`, `cache-list`, - `dist-views`, `dist-simulate`, `dist-plan`. - -`serve` is an internal daemon entry used by `start` and is hidden from help and -completion. Transparent same-path interception, global mount management, and -GPU-resident tensor caches are not part of the 0.1.0 CLI; the supported contract -is to pass the view path printed by `iloader start` to the inference engine. - -## Detailed Usage - -### Basic Memfd Cache - -The default `start` command preloads selected model files into anonymous memory -files and prints a Hugging Face-compatible view path to stdout. Progress, -throughput, and ETA are printed to stderr while the daemon is loading. - -```bash -MODEL_VIEW=$(iloader start "$MODEL" --name dev-model) -``` +For development, `source env.sh` exposes the checkout without installing it and +enables Bash completion. The runtime directory must be node-local and support file +locks, Unix sockets and hard links; `doctor` checks these capabilities. -Use the view path with any framework command that accepts a model directory: +## Cache once and reuse ```bash -vllm serve "$MODEL_VIEW" --tensor-parallel-size 1 -``` +MODEL=/absolute/path/to/model +iloader start "$MODEL" --name dev +iloader status dev -```bash -python -m sglang.launch_server \ - --model-path "$MODEL_VIEW" \ - --tp 1 -``` +# Example when vLLM is already installed; validate your engine version. +iloader run dev -- vllm serve '{model}' --tensor-parallel-size 1 -```bash -python python/infinilm/server/inference_server.py \ - --device cuda \ - --model="$MODEL_VIEW" \ - --tp=1 +# After the engine exits, the cache remains available for another run. +iloader run dev -- vllm serve '{model}' --tensor-parallel-size 1 +iloader stop dev --remove-view ``` -From another terminal, you can inspect and manage active resident caches: +`run` replaces the literal `{model}` argument and sets `ILOADER_MODEL_PATH`. +It holds a lease while the foreground consumer runs. Put InfiniLoader options +before the cache name, and keep the consumer alive while its workers need the view. +`stop` waits for the original daemon to exit, or returns an error on timeout. -```bash -iloader list -iloader status --name dev-model -iloader stop dev-model --remove-view -``` +The file interface is the integration boundary. Engine-specific formats, individual +file-path resolution, containers and loader behavior can affect compatibility; +test your actual engine. Current CI exercises the cache service and reference +examples, not a versioned vLLM/SGLang/InfiniLM compatibility matrix. -Stopping a cache that is still `starting` or `loading` cancels the preload. +## Disk caching and memory limits -### Keep a Model Warm While Restarting a Framework +For a reusable local disk generation without another model-sized RAM allocation: ```bash -MODEL_VIEW=$(iloader start "$MODEL" --name dev-model) - -# Repeat this while editing framework code. -vllm serve "$MODEL_VIEW" --tensor-parallel-size 1 +LOCAL=$(iloader cache "$MODEL" --name dev) +iloader verify "$LOCAL" +iloader run --cache-path "$LOCAL" -- vllm serve '{model}' ``` -Stopping `vllm serve` does not stop the InfiniLoader daemon. The next framework -restart reuses the same warm view. If the daemon is already running, -`iloader start "$MODEL" --name dev-model` simply prints the existing view path. - -### Materialize a Local Cache First - -This is useful when the original model directory lives on remote or shared -storage and you want InfiniLoader to build a reusable local cache before -creating the in-memory view. +To preload into RAM from the local generation, use `start --local-cache`. +The disk-cache root must be outside the source model directory. ```bash -iloader cache "$MODEL" --name dev-model --cache-workers 8 -iloader cache-status "$MODEL" --name dev-model - -iloader start "$MODEL" \ - --name dev-model \ - --local-cache \ - --cache-workers 8 \ - --workers 8 -``` - -### Simulate Distributed Rank Reads - -`dist-views` creates one full Hugging Face-compatible view per rank. These are -not physically sharded checkpoints; every rank view still exposes the complete -model tree. This keeps framework loaders unchanged and is useful for launchers -that expect per-rank paths. - -```bash -RANK_VIEWS=$(iloader dist-views "$VIEW" --name smollm2-tp4 --ranks 4 --replace) - -iloader dist-simulate "$VIEW" \ - --ranks 4 \ - --rank-view-root "$RANK_VIEWS" \ - --pattern "*.safetensors" -``` - -To inspect which tensors could be split for a future engine-native -rank-sharded checkpoint: - -```bash -iloader dist-plan "$MODEL" --tp 4 --limit 20 -``` - -## Command Summary - -Explicit tensor sharding, live CPU layout selection and the CPU/GPU expert API -are documented in [Distributed and hybrid residency](docs/distributed.md). - -```bash -iloader inspect MODEL_DIR -iloader start MODEL_DIR --name NAME -iloader list -iloader status NAME -iloader view NAME -iloader manifest NAME -iloader stop NAME --remove-view -iloader stop --all --remove-view -iloader cache MODEL_DIR --name NAME -iloader cache-status MODEL_DIR --name NAME -iloader cache-list -iloader dist-views MODEL_OR_VIEW --name NAME --ranks N --replace -iloader dist-simulate MODEL_OR_VIEW --ranks N -iloader dist-plan MODEL_DIR --tp N -iloader bench MODEL_OR_VIEW --pattern "*.safetensors" --repeat 3 -iloader completion bash -iloader selftest +iloader start "$MODEL" --name dev --local-cache \ + --max-memory-bytes 32GiB --reserve-bytes 4GiB --max-resident-bytes 64GiB +iloader gc --resident --max-bytes 32GiB # preview only +iloader gc --resident --max-bytes 32GiB --apply ``` -## Measured Results - -Measurements below were collected on a single NVIDIA L20 server using remote -model storage plus local cache storage. They are meant as order-of-magnitude -guidance; exact results depend on model size, storage, framework version, JIT -state, and OS page cache state. - -| Scenario | Baseline | InfiniLoader path | Result | -| --- | ---: | ---: | --- | -| Bloom-560M safetensors read from NFS vs local cache | 0.91 GiB/s | 1.98 GiB/s | 2.18x read throughput | -| Multi-shard memfd preload with 1 vs 4 workers | 0.856 s | 0.240 s | 3.57x preload speedup | -| vLLM SmolLM2 cold-ish NFS weight stage vs memfd view | 33.82 s | 0.15 s | large cold-start improvement | -| vLLM SmolLM2 hot weight stage | 0.20 s | 0.15 s | small hot-cache improvement | -| SGLang 0.5.2 SmolLM2 hot engine load | 12.18 s | 10.37 s | faster hot reload | -| InfiniLM SmolLM2 internal load-only log | 902.7 ms | 219.5 ms | 4.1x weight-stage improvement | -| 4-rank simulated read from InfiniLoader-derived rank views | - | 2.70 GiB in 0.45 s | about 6.0 GiB/s aggregate | - -The most stable production-like gain comes from avoiding repeated NFS reads via -local cache plus memfd views. +The limits above are examples; size them for the model and engine. The cache holds +selected file bytes plus loading buffers. Engine allocations and client-held +mappings/COW copies are additional. Sealed memfd pages are immutable but are not +mlocked, so memory pressure and swap can affect latency. -## Compatibility Notes +Use immutable model revisions. A changed source requires a new cache name or an +explicit stop/rebuild. Readers launched directly against `iloader view dev` have +no automatic lease; keep idle eviction disabled for those readers. -- vLLM, SGLang, and InfiniLM work with the stable `start/view` path by changing - the model path. -- vLLM and SGLang have native `sharded_state` loaders in recent versions, which - are the preferred path for future physical rank-sharded checkpoints. -- InfiniLM currently fits the full-view path best. Rank-sharded checkpoint - loading would need an InfiniLM-specific adapter or native loader support. +## Documentation and optional APIs -## Limitations +- [Operations, verification, cleanup and containers](docs/operations.md) +- [Contributing and running tests](CONTRIBUTING.md) +- [Optional tensor sharding, TensorStore and expert examples](docs/distributed.md) +- [Cache-focused development scope](docs/roadmap.zh-CN.md) +- [Historical measurements and maintenance archive](docs/archive/README.md) -- Linux is required for `memfd_create` and `/proc//fd` views. -- The daemon must stay alive while framework processes read the view. -- The default view duplicates selected model bytes in RAM. -- Fully transparent same-path interception is not included in 0.1.0; pass the - printed view path to the engine. -- `shard plan/export/verify` supports explicit InfiniLoader TP/PP checkpoints. - `dist-plan` still only analyzes candidates. Loading these checkpoints into a - production model engine requires an adapter; see the reference PyTorch MLP. +TensorStore, explicit TP/PP layouts and the GPU expert pool are experimental APIs +with optional PyTorch dependencies. The reference MLP/MoE checks numerical behavior; +it does not provide a production model adapter, quantized kernel or serving engine. +They are not needed for the core `start/run` cache path. -## Development notes +Run `iloader --help` and `iloader COMMAND --help` for the complete CLI. Typical core +commands are `start`, `run`, `status`, `list`, `stop`, `cache`, `verify`, `gc`, `doctor` +and `bench`. `bench` measures file reads without clearing the OS page cache; +it does not measure end-to-end engine startup. -See [commit history and validation](docs/history-reorganization.md) and the -[development roadmap](docs/roadmap.zh-CN.md). +[MIT license](LICENSE). diff --git a/README.zh-CN.md b/README.zh-CN.md index 8357d52..9022d8f 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -1,266 +1,75 @@ -# InfiniLoader 中文说明 +# InfiniLoader -[English README](README.md) +[English](README.md) -开发分支新增了模型身份校验、daemon 启动互斥、memfd 写封印、不可变磁盘缓存版本、 -容器内存预算及租约保护。新增 `doctor`、`verify`、`run`、`gc` 命令，详见 -[运行管理与迁移说明](docs/operations.md)。 +InfiniLoader 把模型文件保存在长期运行的缓存进程中，让推理引擎反复重启时复用权重，减少对原始存储的重复读取。它是 Linux 用户态缓存，负责权重文件及其生命周期；模型计算和请求调度由推理引擎承担。 -```bash -iloader doctor -iloader start "$MODEL" --name dev --local-cache \ - --max-memory-bytes 32GiB --reserve-bytes 4GiB --max-resident-bytes 64GiB -iloader run dev -- vllm serve '{model}' -iloader gc --resident --max-bytes 32GiB # 只预览 -iloader gc --resident --max-bytes 32GiB --apply -``` - -内存紧张时，可只用 `iloader cache` 构建本地磁盘缓存，再用 -`iloader run --cache-path PATH -- COMMAND` 启动框架。原有直接传 view 路径的方式仍可使用； -自动清理无法识别未注册的读者，启用闲置过期或执行清理前应通过 `run` 保护任务。 - -`InfiniLoader` 是一个面向大模型推理框架开发调试的用户态权重加载管理器。它解决的核心问题是：调试 vLLM、SGLang、InfiniLM 或自研推理框架时，经常需要反复停止、修改代码、重新启动框架，而每次启动都要重新从存储读取模型权重。 - -默认模式下，InfiniLoader 会把模型目录常驻在内存里，并暴露出一个看起来像普通 Hugging Face 模型目录的路径。框架仍然使用原来的文件读取逻辑，通常只需要把模型路径换成 `iloader start/view` 输出的路径。 - -## 快速开始 - -最基础路径是 memfd 支撑的 memfs 风格内存 view。如果你原本会把框架直接指向 -`/path/to/huggingface-model`，现在只需要在启动框架前多执行一行命令，然后把模型路径换成它输出的 view。 +核心仅依赖 Python 标准库。选中的文件被复制到封存的 memfd，并通过指向 `/proc//fd/` 的目录链接供现有文件加载器读取。 -```bash -cd InfiniLoader -source env.sh - -MODEL=/path/to/huggingface-model -MODEL_VIEW=$(iloader start "$MODEL" --name dev-model) -``` +## 安装与环境检查 -之后把 `"$MODEL_VIEW"` 当成模型路径传给推理引擎。`iloader start` 是一样的，只有下面的引擎启动命令不同。 - -```bash -# InfiniLM -python python/infinilm/server/inference_server.py \ - --device cuda \ - --model="$MODEL_VIEW" \ - --tp=1 -``` +需要 Linux 和 Python 3.10+。建议安装在引擎所在环境，与消费者使用相同用户和 PID namespace： ```bash -# vLLM -vllm serve "$MODEL_VIEW" --tensor-parallel-size 1 -``` - -```bash -# SGLang -python -m sglang.launch_server \ - --model-path "$MODEL_VIEW" \ - --tp 1 -``` - -结束时： - -```bash -iloader stop dev-model --remove-view -``` - -这就是 0.1.0 的稳定兼容路径：不改框架权重加载代码，只换模型路径。vLLM、SGLang、InfiniLM、本地缓存、分布式 view、进度/status 和自动补全放在后面的详细教程里。 - -## 解决的问题 - -大模型推理框架调试中，权重加载常常成为启动慢的主要来源： - -- 远端或共享存储在读取大 safetensors shard 时容易慢或抖动； -- OS page cache 能改善热启动，但不能消除框架解析、tensor 构建和 CPU 到 GPU 拷贝开销； -- 分布式调试时，多 rank 可能重复扫描或读取同一个 checkpoint； -- 调试框架代码时，不希望为了加速加载而频繁修改框架 loader 接口。 - -InfiniLoader 的设计目标是让权重常驻服务和框架进程解耦。你可以反复停止和重启 vLLM/SGLang/InfiniLM，权重 view 仍然保持 warm。 - -## 工作原理 - -默认 `start/view` 路径只依赖 Linux 用户态能力，不需要 root，也不需要 FUSE： - -1. 后台 daemon 读取模型目录中选中的文件。 -2. 每个文件被复制到 `memfd_create` 创建的匿名内存文件中。 -3. daemon 持有这些文件描述符，使内容保持可读。 -4. view 目录镜像原始模型树，文件是指向 `/proc//fd/` 的符号链接。 -5. 推理框架继续使用普通 `open`、`stat`、`mmap` 读取模型文件。 - -此外，InfiniLoader 还支持本地二级缓存、分布式 rank view 和 TP 切分分析。 - -## 环境设置和运行目录 - -`env.sh` 会把仓库的 `src` 目录加入 `PYTHONPATH`，并在当前终端中定义 -`iloader` shell function；开发态不再依赖 `bin` 目录下的包装脚本。在 bash 中， -它也会自动注册 iloader 子命令、选项和活跃 cache 名称的补全。 - -也可以作为 Python 包安装： - -```bash -pip install -e . +python3 -m pip install . +iloader doctor iloader selftest ``` -运行单元测试： - -```bash -pip install -e .[test] -python -m pytest -``` +开发时可以使用 `source env.sh`，直接加载源码并启用 Bash 补全。运行目录应位于节点本地，并支持文件锁、Unix socket 和硬链接；`doctor` 会检查这些能力。 -运行状态会优先放在 `$ILOADER_BASE_DIR`；未设置时使用 -`$XDG_CACHE_HOME/infiniloader` 或 `~/.cache/infiniloader`。所有会创建运行状态的 -命令也都支持通过 `--base-dir` 显式指定目录。 - -## 命令边界 - -0.1.0 的公开接口刻意保持较小： - -- 主流程：`start`、`list`、`status`、`view`、`stop`。 -- 工具命令：`inspect`、`bench`、`completion`、`selftest`。 -- 稳定高级能力：`cache`、`cache-status`、`cache-list`、`dist-views`、`dist-simulate`、`dist-plan`。 - -`serve` 是 `start` 内部使用的 daemon 入口，不出现在 help 和自动补全里。透明同路径拦截、全局 mount 管理、GPU resident tensor cache 不属于 0.1.0 CLI；当前稳定契约是把 `iloader start` 输出的 view 路径传给推理引擎。 - -## 详细使用 - -### 基础 memfd 缓存 - -默认 `start` 命令会把模型文件预加载到匿名内存文件，并把 Hugging Face 兼容的 view 路径输出到 stdout。加载期间的进度、吞吐和 ETA 会输出到 stderr。 - -```bash -MODEL_VIEW=$(iloader start "$MODEL" --name dev-model) -``` - -任何接受模型目录的框架命令都可以使用这个 view： - -```bash -vllm serve "$MODEL_VIEW" --tensor-parallel-size 1 -``` - -```bash -python -m sglang.launch_server \ - --model-path "$MODEL_VIEW" \ - --tp 1 -``` - -```bash -python python/infinilm/server/inference_server.py \ - --device cuda \ - --model="$MODEL_VIEW" \ - --tp=1 -``` - -另一个终端里可以查看和管理当前活跃的 cache： +## 加载一次，重复使用 ```bash -iloader list -iloader status --name dev-model -iloader stop dev-model --remove-view -``` - -如果 cache 仍处于 `starting` 或 `loading`，`stop` 会取消这次 preload。 - -### 框架反复重启，权重保持 warm +MODEL=/absolute/path/to/model +iloader start "$MODEL" --name dev +iloader status dev -```bash -MODEL_VIEW=$(iloader start "$MODEL" --name dev-model) +# 示例：已安装 vLLM 时使用；仍需验证实际引擎版本。 +iloader run dev -- vllm serve '{model}' --tensor-parallel-size 1 -# 修改框架代码后反复运行这一条。 -vllm serve "$MODEL_VIEW" --tensor-parallel-size 1 +# 引擎退出后缓存继续保留，可重复运行。 +iloader run dev -- vllm serve '{model}' --tensor-parallel-size 1 +iloader stop dev --remove-view ``` -停止 `vllm serve` 不会停止 InfiniLoader daemon。下一次框架启动会复用同一个 warm view。如果 daemon 已经在运行，`iloader start "$MODEL" --name dev-model` 会直接输出已有 view 路径。 +`run` 替换参数中的字面量 `{model}`，并设置 `ILOADER_MODEL_PATH`。消费者前台运行期间有租约保护。InfiniLoader 选项放在缓存名之前；消费者应在其 worker 需要读取 view 时保持存活。`stop` 等待原 daemon 退出，超时则返回错误。 -### 先落本地缓存 +文件接口是集成边界。引擎格式、单文件路径解析、容器和加载器行为都可能影响兼容性，必须验证实际版本。目前 CI 覆盖缓存服务与参考示例，没有覆盖按版本维护的 vLLM/SGLang/InfiniLM 兼容矩阵。 -当原始模型目录位于远端或共享存储，而你希望先构建一个可复用的本地缓存时，建议开启 -`--local-cache`： +## 磁盘缓存与内存限制 -```bash -iloader cache "$MODEL" --name dev-model --cache-workers 8 -iloader cache-status "$MODEL" --name dev-model - -iloader start "$MODEL" \ - --name dev-model \ - --local-cache \ - --cache-workers 8 \ - --workers 8 -``` - -### 模拟分布式 rank 读取 - -`dist-views` 会为每个 rank 生成一个完整 Hugging Face 兼容 view。它不是物理 rank-sharded checkpoint；每个 rank view 仍然暴露完整模型目录。这样做的好处是保持框架默认 loader 不变。 +仅使用本地磁盘代际，不额外分配整份模型大小的 RAM： ```bash -RANK_VIEWS=$(iloader dist-views "$VIEW" --name smollm2-tp4 --ranks 4 --replace) - -iloader dist-simulate "$VIEW" \ - --ranks 4 \ - --rank-view-root "$RANK_VIEWS" \ - --pattern "*.safetensors" +LOCAL=$(iloader cache "$MODEL" --name dev) +iloader verify "$LOCAL" +iloader run --cache-path "$LOCAL" -- vllm serve '{model}' ``` -分析哪些 tensor 适合未来做框架原生的 rank-sharded checkpoint： +需要从本地代际继续预加载到 RAM 时，使用 `start --local-cache`。磁盘缓存根目录必须位于源模型目录之外。 ```bash -iloader dist-plan "$MODEL" --tp 4 --limit 20 -``` - -## 命令摘要 - -```bash -iloader inspect MODEL_DIR -iloader start MODEL_DIR --name NAME -iloader list -iloader status NAME -iloader view NAME -iloader manifest NAME -iloader stop NAME --remove-view -iloader stop --all --remove-view -iloader cache MODEL_DIR --name NAME -iloader cache-status MODEL_DIR --name NAME -iloader cache-list -iloader dist-views MODEL_OR_VIEW --name NAME --ranks N --replace -iloader dist-simulate MODEL_OR_VIEW --ranks N -iloader dist-plan MODEL_DIR --tp N -iloader bench MODEL_OR_VIEW --pattern "*.safetensors" --repeat 3 -iloader completion bash -iloader selftest +iloader start "$MODEL" --name dev --local-cache \ + --max-memory-bytes 32GiB --reserve-bytes 4GiB --max-resident-bytes 64GiB +iloader gc --resident --max-bytes 32GiB # 仅预览 +iloader gc --resident --max-bytes 32GiB --apply ``` -## 实测效果摘要 - -以下结果来自一台 NVIDIA L20 服务器，使用远端模型存储和本地缓存存储。结果用于说明量级，实际收益取决于模型大小、存储、框架版本、JIT 状态和 OS page cache 状态。 - -| 场景 | 基线 | InfiniLoader 路径 | 结果 | -| --- | ---: | ---: | --- | -| Bloom-560M safetensors 从 NFS 读 vs 本地 cache | 0.91 GiB/s | 1.98 GiB/s | 读吞吐约 2.18x | -| 多 shard memfd preload，1 worker vs 4 workers | 0.856 s | 0.240 s | preload 约 3.57x | -| vLLM SmolLM2 冷态 NFS 权重阶段 vs memfd view | 33.82 s | 0.15 s | 冷启动大幅改善 | -| vLLM SmolLM2 热态权重阶段 | 0.20 s | 0.15 s | 热态小幅改善 | -| SGLang 0.5.2 SmolLM2 热态 engine load | 12.18 s | 10.37 s | 热重启更快 | -| InfiniLM SmolLM2 load-only 内部日志 | 902.7 ms | 219.5 ms | 权重阶段约 4.1x | -| 4-rank 从 InfiniLoader 派生 rank views 模拟读取 | - | 2.70 GiB / 0.45 s | 聚合约 6.0 GiB/s | - -最稳定的收益来自本地 cache 加 memfd view，尤其适合绕过重复 NFS 读取。 +上述容量只是示例，应按模型和引擎实际占用设置。缓存计费包括文件内容与加载缓冲，引擎分配和客户端仍持有的映射/COW 副本另计。memfd 已封存但未 mlock，内存压力和 swap 仍会影响延迟。 -## 兼容性说明 +源模型应使用不可变版本。源变化后需更换缓存名或先停止再重建。直接使用 `iloader view dev` 路径的读者没有自动租约，应保持自动 idle eviction 关闭。 -- vLLM、SGLang、InfiniLM 都可以通过稳定的 `start/view` 路径使用，通常只需要改模型路径。 -- 新版本 vLLM 和 SGLang 具备原生 `sharded_state` loader，是未来支持物理 rank-sharded checkpoint 的优先方向。 -- InfiniLM 当前最适合完整 view 路径；rank-sharded checkpoint 需要 InfiniLM adapter 或原生 loader 支持。 +## 文档与可选接口 -## 限制 +- [操作、校验、回收与容器](docs/operations.md) +- [贡献与测试](CONTRIBUTING.md) +- [可选分片、TensorStore 和专家示例](docs/distributed.md) +- [围绕缓存目标的开发范围](docs/roadmap.zh-CN.md) +- [历史测量与维护资料](docs/archive/README.md) -- 需要 Linux；核心能力依赖 `memfd_create` 和 `/proc//fd`。 -- 框架读取 view 时，InfiniLoader daemon 必须保持运行。 -- 默认 view 会在内存中保留一份选中的模型文件内容。 -- 0.1.0 不包含完全透明的同路径拦截；需要把输出的 view 路径传给推理引擎。 -- `shard plan/export/verify` 已支持显式 TP/PP checkpoint；`dist-plan` 仍仅做候选分析。生产引擎加载这些分片需要适配器。 +TensorStore、显式 TP/PP layout 和 GPU 专家池是实验接口，使用可选 PyTorch 依赖。参考 MLP/MoE 验证基本数值行为，不提供生产模型适配器、量化内核或 serving engine；核心 `start/run` 不需要这些接口。 -## 开发说明 +完整 CLI 以 `iloader --help` 和 `iloader COMMAND --help` 为准。常用核心命令为 `start`、`run`、`status`、`list`、`stop`、`cache`、`verify`、`gc`、`doctor` 和 `bench`。`bench` 不清除 OS page cache，仅测文件读取，不能等同于引擎启动收益。 -显式分片、CPU TensorStore 与专家池接口见 [分布式与混合驻留](docs/distributed.md)。 -提交线及验证见 [提交重整记录](docs/history-reorganization.md)，后续能力见 [开发路线](docs/roadmap.zh-CN.md)。 +[MIT 许可](LICENSE)。 diff --git a/docs/distributed-results.json b/docs/archive/2026-10-18-distributed-results.json similarity index 100% rename from docs/distributed-results.json rename to docs/archive/2026-10-18-distributed-results.json diff --git a/docs/distributed-validation.md b/docs/archive/2026-10-18-distributed-validation.md similarity index 97% rename from docs/distributed-validation.md rename to docs/archive/2026-10-18-distributed-validation.md index b165353..e2b374d 100644 --- a/docs/distributed-validation.md +++ b/docs/archive/2026-10-18-distributed-validation.md @@ -2,7 +2,7 @@ > Historical 2026-10-18 measurements retain their original tested commit IDs. > The reorganized history preserves production code and examples byte-for-byte; -> see [history-reorganization.md](history-reorganization.md) for provenance +> see [history-reorganization.md](2026-10-26-history-reorganization.md) for provenance > and independently repeated CPU tests. No new four-node/device run is claimed. Date: 2026-10-18, Asia/Shanghai. Feature branch: `codex/distributed-residency`, based @@ -17,8 +17,8 @@ on reviewed reliability commit `7442271`. No merge into master was performed. | CPU weights plus GPU hot/CPU cold experts | Grouped, bounded GPU copies; stream-aware leases and LRU; real CPU/GPU ReLU MoE, repeated hot-set changes and failure injection on four DCUs | Sequential reference expert execution; no CPU/GPU overlap claim, distributed EP/MoE, quantized kernel or native KTransformers adapter. | The detailed protocol, commands, contracts and source references are in -[distributed.md](distributed.md). Sanitized numerical evidence is provided in -[distributed-results.json](distributed-results.json). +[distributed.md](../distributed.md). Sanitized numerical evidence is provided in +[distributed-results.json](2026-10-18-distributed-results.json). ## Environments and methodology diff --git a/docs/expert-review.md b/docs/archive/2026-10-18-expert-review.md similarity index 100% rename from docs/expert-review.md rename to docs/archive/2026-10-18-expert-review.md diff --git a/docs/validation.md b/docs/archive/2026-10-18-reliability-validation.md similarity index 96% rename from docs/validation.md rename to docs/archive/2026-10-18-reliability-validation.md index 55fa79f..6a386a5 100644 --- a/docs/validation.md +++ b/docs/archive/2026-10-18-reliability-validation.md @@ -38,7 +38,7 @@ TTL while a lease is active, then verifies eviction after lease release. The original branch HEAD `6a33d83` passed the complete remote Ubuntu Python 3.10-3.14 matrix in [Actions run 39101870656](https://github.com/InfiniTensor/InfiniLoader/actions/runs/39101870656). Check the newer branch workflow for the subsequent review fixes. Independent review -and per-commit test details are recorded in [expert-review.md](expert-review.md). +and per-commit test details are recorded in [expert-review.md](2026-10-18-expert-review.md). The Dockerfile is a non-root diagnostic image; Docker was unavailable locally, so the image was not built or exercised in this environment. diff --git a/docs/history-reorganization.json b/docs/archive/2026-10-26-history-reorganization.json similarity index 100% rename from docs/history-reorganization.json rename to docs/archive/2026-10-26-history-reorganization.json diff --git a/docs/history-reorganization.md b/docs/archive/2026-10-26-history-reorganization.md similarity index 95% rename from docs/history-reorganization.md rename to docs/archive/2026-10-26-history-reorganization.md index 0e3acdb..dc23ee7 100644 --- a/docs/history-reorganization.md +++ b/docs/archive/2026-10-26-history-reorganization.md @@ -29,7 +29,7 @@ Python 3.14.4、PyTorch 2.14.1+cpu、NumPy 2.5.4、safetensors 0.8.0、pytest 9. | 11 | 本文件所在提交 | Document validated history and development priorities | 文档与来源映射；最终 hash 的审核结果随交付报告给出 | 逐提交元数据、原始 pytest 日志 SHA-256 及完整映射见 -[history-reorganization.json](history-reorganization.json)。原始日志保存在本次任务的 +[history-reorganization.json](2026-10-26-history-reorganization.json)。原始日志保存在本次任务的 `InfiniLoader-commit-review-20261026` 证据目录；该本地目录不是公开下载链接。 生命周期专家审核了 1–5；缓存/内存专家审核了 1、3、4、7、8、10；CLI/验证专家审核 @@ -60,7 +60,7 @@ git diff 88d6bac improve-weight-residency -- tests 原四机 CPU/GPU 测量保留其真实代码版本（例如 `123c6e2`、`758f8c8`），不改写为新 hash。 源码相同说明实现被保留，不代表本次重新执行了 GPU、RCCL 或四机实验。两个真实设备 -测试在本地 CPU 环境跳过，原设备验证见 [distributed-validation.md](distributed-validation.md)。 +测试在本地 CPU 环境跳过，原设备验证见 [distributed-validation.md](2026-10-18-distributed-validation.md)。 当前分支的远端 CI 应查看该分支自己的 Actions；原报告中的旧 Actions 链接仍只证明旧提交。 ## 原提交归属 @@ -87,4 +87,4 @@ git diff 88d6bac improve-weight-residency -- tests | `758f8c8` | 9 | Startup barrier included in initial TP/PP harness | | `88d6bac` | 11 | Distributed contracts, historical numerical evidence and README | -后续能力优先级与验收条件见 [roadmap.zh-CN.md](roadmap.zh-CN.md)。 +后续能力优先级与验收条件见 [roadmap.zh-CN.md](2026-10-26-roadmap.zh-CN.md)。 diff --git a/docs/archive/2026-10-26-roadmap.zh-CN.md b/docs/archive/2026-10-26-roadmap.zh-CN.md new file mode 100644 index 0000000..973787e --- /dev/null +++ b/docs/archive/2026-10-26-roadmap.zh-CN.md @@ -0,0 +1,47 @@ +# InfiniLoader 后续能力路线 + +评估日期：2026-10-26。以下是基于当前实现和官方上游接口的开发建议，尚未实现，也不是性能承诺。 + +## 建议定位 + +InfiniLoader 适合继续发展成供推理引擎复用的权重存储与驻留层：管理内容身份、不可变代际、进程间共享、布局和设备副本生命周期。计算图、算子、KV cache 和请求调度仍由推理引擎负责。下一轮最有价值的交付是一个真实模型适配器和可信的端到端性能证据。 + +当前已经具备 sealed memfd、租约、磁盘代际、cgroup 内存预算、显式 TP/PP 切片、共享 CPU TensorStore 和同步 GPU ExpertPool。四机测试证明了小型 MLP 的 TP/PP 数值路径及参考 MoE 换入换出；它们没有验证生产 LLM、量化引擎兼容性、GPU 动态拓扑切换或分布式专家并行。 + +## 优先级与验收条件 + +| 优先级 | 能力 | 当前缺口与具体交付 | 验收依据 | +| --- | --- | --- | --- | +| P0 | 真实模型与引擎适配 | 固定一个模型族和引擎版本，从 PyTorch/Transformers 接入开始，再扩展 vLLM/InfiniLM；显式处理 GQA、packed QKV、bias、tied weights、dtype 和非连续 slice | 与原生加载比较 logits/生成结果；验证 TP/PP；重启和重复加载；实际模型的冷启动、热启动与内存结果 | +| P0 | 端到端性能与物理内存观测 | 将读盘、hash、映射、参数安装、转换、H2D、引擎初始化和首 token 分开计时；增加共享 PSS、cgroup、GPU allocated/reserved 与 NUMA 观测 | 同模型、同设备、同软件版本下比较原生 mmap、热 page cache、InfiniLoader；报告重复样本及 p50/p95，不用读文件带宽代替启动收益 | +| P1 | 批量 catalog、导出和复用 | 当前 export_rank 在导出前后分别 catalog 全量源；TensorStore 也会校验全量内容。设计一次 snapshot/catalog、多 rank 有界导出和代际级校验复用 | 记录总读字节与 hash 字节；rank 数增加时避免重复全量扫描；源变化、进程退出、损坏和重试测试仍通过 | +| P1 | 引擎原生格式与 KTransformers 适配 | 按内容 ID、格式、转换参数和内核版本缓存派生产物；提供原生 tensor/格式适配，明确 C++ 内核是否复制、持有还是借用存储 | 一个真实 MoE 端到端运行；精度检查、格式失效和重启复用；同时报告 CPU canonical、转换副本、GPU 权重与总显存 | +| P1 | 有界异步预取与冷热策略 | 当前 promote 同步，hot set 由调用方指定，参考执行按专家串行。先增加有界 staging、异步句柄、取消/失败状态，再对接路由统计和原生 CPU 内核 | 相同质量和显存限制下测 TTFT、每 token 延迟、命中率、H2D 字节、等待和抖动；证明计算/传输实际重叠，不只统计 promotion 次数 | +| P1 | 节点级服务与稳定客户端协议 | 提供版本化 acquire/release/describe 接口、多模型幂等 warmup、可查询操作状态、指标和能力探测；评估 FD 传递以减少对跨容器 /proc 路径的依赖 | 同节点多进程共享；不同 PID namespace 的可支持配置测试；daemon 崩溃、客户端存活映射及 COW 对租约和配额的影响明确 | +| P2 | 分布式分片迁移与重分布 | 现有 sharded store 绑定原 layout/rank；完整副本可本地重切片，但不提供旧分片到新分片的在线迁移。先做可信、可核验、有界传输，再考虑 RDMA | TP4→TP2/PP2→PP4 的真实模型迁移；校验所需 shard 完整后才发布；中断可恢复，重复请求不重复启动，不足内存时拒绝；GPU 拓扑与通信单独实测 | + +## 先解决真实加载链路中的复制 + +共享 CPU Tensor 并不意味着引擎会直接复用该 storage。常规 load_state_dict 可能复制到已分配参数，量化、packing 和 pinning 也会产生新存储。适配器应先测量这些步骤，明确参数及 storage 的所有权，再选择 meta 初始化和参数绑定方式。PyTorch 的官方加载说明展示了 mmap、meta 与 assign=True 的组合；这可作为适配切入点，并不代表所有模型和版本都能直接套用。[PyTorch 加载指南](https://docs.pytorch.org/tutorials/recipes/recipes/module_load_state_dict_tips.html) + +vLLM 已有按 worker 加载自身 TP/PP state 的原生 loader，因此接入要实现其期望的参数语义和存储格式。InfiniLoader 的显式 layout 并不是该格式的直接替代品。[vLLM ShardedStateLoader](https://docs.vllm.ai/en/latest/api/vllm/model_executor/model_loader/sharded_state_loader/) + +## 与 KTransformers 的合作边界 + +KT-Kernel 已有原生格式、CPU 内核、NUMA 配置、GPU 专家选择和动态放置相关接口，也公开了从 tensor 加载权重的入口。应优先适配和测量这些入口，让 InfiniLoader 提供可复用的权重及派生产物管理；不要将另一个参考 LRU/MoE 调度器视为主要产品收益。接入后仍需验证原生内核的布局、复制和异步使用契约，不能由 PyTorch 视图的共享性质推断原生内核零复制。[KT-Kernel 官方说明](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/README.md) + +## 分布式模式如何演进 + +固定布局且 RAM 紧张时优先分片；调试和频繁拓扑变更时完整 CPU 副本更方便。后续可以依据内存预算、预计切换频率与迁移成本选择模式，记录选择理由，但不应在缺少带宽和工作负载测量时宣称自动最优。 + +拓扑重分布需要参数语义、数据移动、协调和故障恢复共同工作。PyTorch Distributed Checkpoint 可作为状态分片与恢复接口的参考，不能把它的状态恢复能力等同于 InfiniLoader 已支持在线推理拓扑迁移。[PyTorch Distributed Checkpoint](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html) + +实验中的共享目录可见性问题还说明，节点上的权威操作状态和可查询结果比仅依赖目录 marker 更可靠。可继续接入既有算力调度平台的资源分配和租约，把权重准备作为作业阶段；InfiniLoader 保持权重层职责，不复制建设整套集群调度系统。 + +## 推荐实施顺序 + +1. **真实模型接入与基线。** 一个稠密模型、一个引擎、固定版本，完成结果正确性和原生/缓存加载对比，同时补齐 storage 复制与内存观测，并给出可复现安装、版本兼容和迁移说明。 +2. **减少重复准备与原生 MoE 接入。** 批量 catalog/export、派生格式缓存，再接一个 KT-Kernel 原生路径，量化收益与内存成本。 +3. **按瓶颈推进异步和集群能力。** 有性能证据后做预取、NUMA、热点策略；有可靠节点服务后再做在线重分布，最后评估 RDMA。 + +现阶段不建议同时承诺所有量化格式、所有推理引擎、多租户服务、训练状态缓存和通用 GPU 共享。它们具有不同的所有权与一致性要求，应由具体适配需求驱动。 diff --git a/docs/archive/README.md b/docs/archive/README.md new file mode 100644 index 0000000..509d6a3 --- /dev/null +++ b/docs/archive/README.md @@ -0,0 +1,20 @@ +# Historical measurements and maintenance records + +These files describe the versions and experiments named in each document. They are +retained for provenance, not as current engine compatibility or performance guarantees. +Some original raw logs are local-only; the archive says so rather than presenting +those paths as public download links. Original commit IDs remain unchanged. +The 2026-10-26 mapping records were copied from `33b176e`; references to their +containing commit describe that historical snapshot, not the current fixes. + +- [Original engine measurements](original-engine-measurements.md) +- [2026-10-18 reliability validation](2026-10-18-reliability-validation.md) +- [2026-10-18 review record](2026-10-18-expert-review.md) +- [2026-10-18 distributed validation](2026-10-18-distributed-validation.md) +- [Distributed numerical summary](2026-10-18-distributed-results.json) +- [2026-10-26 history mapping](2026-10-26-history-reorganization.md) +- [Original broader roadmap](2026-10-26-roadmap.zh-CN.md) +- [History mapping data](2026-10-26-history-reorganization.json) + +Current usage is in [operations](../operations.md); optional APIs are in +[distributed and hybrid residency](../distributed.md). diff --git a/docs/archive/original-engine-measurements.md b/docs/archive/original-engine-measurements.md new file mode 100644 index 0000000..3ba2b28 --- /dev/null +++ b/docs/archive/original-engine-measurements.md @@ -0,0 +1,27 @@ +# Original engine measurements + +Source: the README at original commit `5014a6a`. These observations were not +remeasured by the cache reliability fixes. The original table does not provide a +complete engine-version matrix, model revisions, raw samples or reproducible commands. +Retain the numbers as historical observations, not current release guarantees. +Weight-stage, engine-load and file-read measurements are different scopes and must +not be combined into an end-to-end speedup claim. + + +Measurements below were collected on a single NVIDIA L20 server using remote +model storage plus local cache storage. They are meant as order-of-magnitude +guidance; exact results depend on model size, storage, framework version, JIT +state, and OS page cache state. + +| Scenario | Baseline | InfiniLoader path | Result | +| --- | ---: | ---: | --- | +| Bloom-560M safetensors read from NFS vs local cache | 0.91 GiB/s | 1.98 GiB/s | 2.18x read throughput | +| Multi-shard memfd preload with 1 vs 4 workers | 0.856 s | 0.240 s | 3.57x preload speedup | +| vLLM SmolLM2 cold-ish NFS weight stage vs memfd view | 33.82 s | 0.15 s | large cold-start improvement | +| vLLM SmolLM2 hot weight stage | 0.20 s | 0.15 s | small hot-cache improvement | +| SGLang 0.5.2 SmolLM2 hot engine load | 12.18 s | 10.37 s | faster hot reload | +| InfiniLM SmolLM2 internal load-only log | 902.7 ms | 219.5 ms | 4.1x weight-stage improvement | +| 4-rank simulated read from InfiniLoader-derived rank views | - | 2.70 GiB in 0.45 s | about 6.0 GiB/s aggregate | + +The most stable production-like gain comes from avoiding repeated NFS reads via +local cache plus memfd views. diff --git a/docs/distributed.md b/docs/distributed.md index fad02f6..8743eba 100644 --- a/docs/distributed.md +++ b/docs/distributed.md @@ -1,5 +1,9 @@ # Distributed CPU residency and hybrid expert integration +These APIs and examples are experimental and optional. The core file-cache path +does not require PyTorch. The archived device/multi-node measurements identify +the exact historical versions; they do not establish production engine compatibility. + InfiniLoader now provides explicit safetensors partitioning, a PyTorch CPU tensor store, and a bounded GPU expert pool. The reference examples execute real TP/PP and a dense ReLU MoE. Framework-specific Llama/GQA, quantized kernels and @@ -30,7 +34,8 @@ that these engines consume the new InfiniLoader shard format directly. ## Explicit checkpoint preparation -The legacy `dist-plan` remains a heuristic report. The new executable path requires +The legacy `dist-plan` remains a heuristic report. It shares strict format and +index validation with `shard`, while reading headers rather than hashing payloads. The new executable path requires an exact tensor inventory and explicit TP dimension/PP stage for **every** tensor: ```json @@ -198,5 +203,5 @@ verify host/checkpoint identity, allocate devices, bound execution and collect a rank results. The example is a validation harness, not a fault-tolerant cluster controller, authenticated tensor-transfer service or inference serving framework. -See [distributed validation](distributed-validation.md) for the tested hardware, +See [distributed validation](archive/2026-10-18-distributed-validation.md) for the tested hardware, results, independent review and remaining limitations. diff --git a/docs/operations.md b/docs/operations.md index f9f36a8..672236e 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -1,8 +1,8 @@ -# Operating the development version +# Operating the cache service -This branch improves reliability and local usability. TP checkpoint conversion and -multi-node transport remain future work. Runtime dependencies remain limited to -Python's standard library. Linux and Python 3.10+ are required. +This page covers the core file-cache service. Runtime dependencies are limited to +Python's standard library; Linux and Python 3.10+ are required. Explicit tensor +sharding and CPU/GPU integration are [optional experimental APIs](distributed.md). ## Choose a backend @@ -33,6 +33,10 @@ weight transformation and CPU-to-device transfers remain the engine's responsibi ## Identity, verification and migration +Stop existing daemons before replacing the InfiniLoader package or checkout. The +CLI and daemon are expected to use the same revision. Engine restarts can reuse +the cache without upgrading or restarting InfiniLoader itself. + A resident name identifies one source path, file selection and source fingerprint. The fingerprint includes sizes, nanosecond timestamps, device and inode IDs. Reusing a name for another revision fails; stop the cache or choose another name. Changes @@ -49,6 +53,9 @@ NEW_LOCAL=$(iloader cache "$MODEL" --name dev --refresh) iloader verify "$NEW_LOCAL" --json ``` +The disk-cache root must be outside the source model directory, including through +symlink aliases. Overlapping roots are rejected before copying. + Refresh publishes a new generation, preserves old paths, and needs additional disk space. SHA-256 is computed during the copy and its CPU cost belongs in preparation benchmarks. Version-1 caches remain listed but are not reused; rebuild them, verify @@ -74,10 +81,17 @@ iloader start "$MODEL" --name dev --workers 4 \ caps remain enforced. Quotas account for InfiniLoader, not engine allocations. - Byte options accept integers and binary suffixes such as `256MiB` and `32GiB`. -Use a node-local `--base-dir`. A runtime registry shared between nodes is unsupported. +Use a node-local `--base-dir` supporting locks, Unix sockets and hard links. +Relative paths are anchored at startup. The encoded socket path must fit in 107 +bytes; use a shorter `--base-dir` or explicit `--socket` when needed. A runtime +registry shared between nodes is unsupported. ## Leases, LRU and TTL +`stop` waits for the original daemon PID/token to exit and returns an error if its +timeout is exceeded. Raw IPC stop requests are asynchronous. A subsequent `start` +waits for a stopping instance within its readiness timeout. + `run` protects a cache against ordinary stop and eviction. Local-cache daemons also hold a disk-generation lease. Foreground disk commands inherit its OS lock descriptor; resident commands register the child process's PID and start time. @@ -115,7 +129,8 @@ iloader list --json iloader dist-simulate "$VIEW" --ranks 4 --timeout 60 ``` -Doctor exercises locks, Unix sockets, memfd sealing and reopening in another process. +Doctor exercises locks, Unix sockets, exclusive socket hard-link publication, +memfd sealing and reopening in another process. It reports cgroup headroom, disk space, FD and memlock limits. Status includes identity, file bytes, the startup budget snapshot, active leases, reuse requests and last use. Distributed simulation is a bounded single-host read test, not a GPU benchmark. @@ -126,7 +141,11 @@ and pass proc and user-namespace access checks. A shared volume alone is insuffi Test an actual engine open in the deployment environment. Resolving individual view file symlinks to their final target can break reopening; use the published file paths. After a daemon exits, existing open descriptors/mappings may survive, but new opens -through its view fail. Automatic failover is not provided. +through its view fail. Automatic failover is not provided. New socket publication records its inode +before exposing the public path, allowing verified stale state to be reclaimed +after an interrupted startup. Unknown files, symlinks, listeners or sockets +without matching ownership are still refused. Preserve and inspect such paths +before manual recovery; do not indiscriminately delete a runtime directory. The non-root diagnostic image contains InfiniLoader, without an engine or model: @@ -137,6 +156,7 @@ docker run --rm --memory=2g infiniloader:dev ## Validation +See [contributing](../CONTRIBUTING.md) for core and optional test dependencies. Run `python -m pytest`, `iloader selftest`, and `iloader doctor`. CI tests Python 3.10-3.14 on Ubuntu. Tests exercise real daemon subprocesses, source mutations, disconnected clients, immutable generations, quotas, leases, and failed ranks. @@ -144,6 +164,7 @@ disconnected clients, immutable generations, quotas, leases, and failed ranks. `bench` reports each run, first-run and median time, and best throughput. It does not clear OS page cache. Compare cold/warm shared storage, cold/warm local disk, memfd and engine-native loaders separately. Include preparation, end-to-end engine readiness, -peak CPU memory, per-rank bytes and CPU-to-device transfers. Historical README numbers -are not new measurements for this branch; CPU tests do not establish GPU or multi-node -performance. +peak CPU memory, per-rank bytes and CPU-to-device transfers. The +[archived engine measurements](archive/original-engine-measurements.md) are +historical observations, not measurements for the current branch; CPU tests do +not establish GPU or multi-node performance. diff --git a/docs/roadmap.zh-CN.md b/docs/roadmap.zh-CN.md index 973787e..5ef4d41 100644 --- a/docs/roadmap.zh-CN.md +++ b/docs/roadmap.zh-CN.md @@ -1,47 +1,15 @@ -# InfiniLoader 后续能力路线 +# 开发范围与下一步 -评估日期：2026-10-26。以下是基于当前实现和官方上游接口的开发建议，尚未实现，也不是性能承诺。 +InfiniLoader 的目标是：模型文件加载一次，让后续进程快速、可靠地复用。核心提供文件缓存、内容与代际检查、使用期租约、内存预算和安全回收；模型执行与请求调度交给消费者。 -## 建议定位 +当前优先验收以下内容，不以增加功能数量为目标： -InfiniLoader 适合继续发展成供推理引擎复用的权重存储与驻留层：管理内容身份、不可变代际、进程间共享、布局和设备副本生命周期。计算图、算子、KV cache 和请求调度仍由推理引擎负责。下一轮最有价值的交付是一个真实模型适配器和可信的端到端性能证据。 +1. 用日常调试的一个真实模型和框架，通过已有 `start/run` 路径验证结果一致、重复加载耗时与内存占用。 +2. 连续重复加载、退出、停止与复用，检查 FD、内存及残留进程是否持续增长，并验证源变化与失败恢复。 +3. 保持安装、最小使用流程、支持环境与故障信息清楚可复现。公开性能结果应带模型 revision、引擎版本、命令、样本和测试范围。 -当前已经具备 sealed memfd、租约、磁盘代际、cgroup 内存预算、显式 TP/PP 切片、共享 CPU TensorStore 和同步 GPU ExpertPool。四机测试证明了小型 MLP 的 TP/PP 数值路径及参考 MoE 换入换出；它们没有验证生产 LLM、量化引擎兼容性、GPU 动态拓扑切换或分布式专家并行。 +后续优化应由这些测量发现的瓶颈驱动，例如重复文件扫描、校验或复制。先记录实际成本，再决定是否需要新的缓存策略或接口。 -## 优先级与验收条件 +显式分片、TensorStore 和 GPU 专家池保留为[实验接口](distributed.md)。它们不是核心 `start/run` 的依赖，也不作为前期缓存版本的生产引擎兼容保证。在线拓扑迁移、量化内核、路由调度和通用集群服务不属于当前交付范围。 -| 优先级 | 能力 | 当前缺口与具体交付 | 验收依据 | -| --- | --- | --- | --- | -| P0 | 真实模型与引擎适配 | 固定一个模型族和引擎版本，从 PyTorch/Transformers 接入开始，再扩展 vLLM/InfiniLM；显式处理 GQA、packed QKV、bias、tied weights、dtype 和非连续 slice | 与原生加载比较 logits/生成结果；验证 TP/PP；重启和重复加载；实际模型的冷启动、热启动与内存结果 | -| P0 | 端到端性能与物理内存观测 | 将读盘、hash、映射、参数安装、转换、H2D、引擎初始化和首 token 分开计时；增加共享 PSS、cgroup、GPU allocated/reserved 与 NUMA 观测 | 同模型、同设备、同软件版本下比较原生 mmap、热 page cache、InfiniLoader；报告重复样本及 p50/p95，不用读文件带宽代替启动收益 | -| P1 | 批量 catalog、导出和复用 | 当前 export_rank 在导出前后分别 catalog 全量源；TensorStore 也会校验全量内容。设计一次 snapshot/catalog、多 rank 有界导出和代际级校验复用 | 记录总读字节与 hash 字节；rank 数增加时避免重复全量扫描；源变化、进程退出、损坏和重试测试仍通过 | -| P1 | 引擎原生格式与 KTransformers 适配 | 按内容 ID、格式、转换参数和内核版本缓存派生产物；提供原生 tensor/格式适配，明确 C++ 内核是否复制、持有还是借用存储 | 一个真实 MoE 端到端运行；精度检查、格式失效和重启复用；同时报告 CPU canonical、转换副本、GPU 权重与总显存 | -| P1 | 有界异步预取与冷热策略 | 当前 promote 同步，hot set 由调用方指定，参考执行按专家串行。先增加有界 staging、异步句柄、取消/失败状态，再对接路由统计和原生 CPU 内核 | 相同质量和显存限制下测 TTFT、每 token 延迟、命中率、H2D 字节、等待和抖动；证明计算/传输实际重叠，不只统计 promotion 次数 | -| P1 | 节点级服务与稳定客户端协议 | 提供版本化 acquire/release/describe 接口、多模型幂等 warmup、可查询操作状态、指标和能力探测；评估 FD 传递以减少对跨容器 /proc 路径的依赖 | 同节点多进程共享；不同 PID namespace 的可支持配置测试；daemon 崩溃、客户端存活映射及 COW 对租约和配额的影响明确 | -| P2 | 分布式分片迁移与重分布 | 现有 sharded store 绑定原 layout/rank；完整副本可本地重切片，但不提供旧分片到新分片的在线迁移。先做可信、可核验、有界传输，再考虑 RDMA | TP4→TP2/PP2→PP4 的真实模型迁移；校验所需 shard 完整后才发布；中断可恢复，重复请求不重复启动，不足内存时拒绝；GPU 拓扑与通信单独实测 | - -## 先解决真实加载链路中的复制 - -共享 CPU Tensor 并不意味着引擎会直接复用该 storage。常规 load_state_dict 可能复制到已分配参数，量化、packing 和 pinning 也会产生新存储。适配器应先测量这些步骤，明确参数及 storage 的所有权，再选择 meta 初始化和参数绑定方式。PyTorch 的官方加载说明展示了 mmap、meta 与 assign=True 的组合；这可作为适配切入点，并不代表所有模型和版本都能直接套用。[PyTorch 加载指南](https://docs.pytorch.org/tutorials/recipes/recipes/module_load_state_dict_tips.html) - -vLLM 已有按 worker 加载自身 TP/PP state 的原生 loader，因此接入要实现其期望的参数语义和存储格式。InfiniLoader 的显式 layout 并不是该格式的直接替代品。[vLLM ShardedStateLoader](https://docs.vllm.ai/en/latest/api/vllm/model_executor/model_loader/sharded_state_loader/) - -## 与 KTransformers 的合作边界 - -KT-Kernel 已有原生格式、CPU 内核、NUMA 配置、GPU 专家选择和动态放置相关接口，也公开了从 tensor 加载权重的入口。应优先适配和测量这些入口，让 InfiniLoader 提供可复用的权重及派生产物管理；不要将另一个参考 LRU/MoE 调度器视为主要产品收益。接入后仍需验证原生内核的布局、复制和异步使用契约，不能由 PyTorch 视图的共享性质推断原生内核零复制。[KT-Kernel 官方说明](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/README.md) - -## 分布式模式如何演进 - -固定布局且 RAM 紧张时优先分片；调试和频繁拓扑变更时完整 CPU 副本更方便。后续可以依据内存预算、预计切换频率与迁移成本选择模式，记录选择理由，但不应在缺少带宽和工作负载测量时宣称自动最优。 - -拓扑重分布需要参数语义、数据移动、协调和故障恢复共同工作。PyTorch Distributed Checkpoint 可作为状态分片与恢复接口的参考，不能把它的状态恢复能力等同于 InfiniLoader 已支持在线推理拓扑迁移。[PyTorch Distributed Checkpoint](https://docs.pytorch.org/docs/stable/distributed.checkpoint.html) - -实验中的共享目录可见性问题还说明，节点上的权威操作状态和可查询结果比仅依赖目录 marker 更可靠。可继续接入既有算力调度平台的资源分配和租约，把权重准备作为作业阶段；InfiniLoader 保持权重层职责，不复制建设整套集群调度系统。 - -## 推荐实施顺序 - -1. **真实模型接入与基线。** 一个稠密模型、一个引擎、固定版本，完成结果正确性和原生/缓存加载对比，同时补齐 storage 复制与内存观测，并给出可复现安装、版本兼容和迁移说明。 -2. **减少重复准备与原生 MoE 接入。** 批量 catalog/export、派生格式缓存，再接一个 KT-Kernel 原生路径，量化收益与内存成本。 -3. **按瓶颈推进异步和集群能力。** 有性能证据后做预取、NUMA、热点策略；有可靠节点服务后再做在线重分布，最后评估 RDMA。 - -现阶段不建议同时承诺所有量化格式、所有推理引擎、多租户服务、训练状态缓存和通用 GPU 共享。它们具有不同的所有权与一致性要求，应由具体适配需求驱动。 +此前的较大范围建议保存在[历史路线](archive/2026-10-26-roadmap.zh-CN.md)，用于说明原始评估，不代表当前承诺。 From 92afa751ffec00c604d5efc6e6fde9804bc76151 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 02:30:52 +0800 Subject: [PATCH 18/29] Keep rank view metadata out of source file inventories --- src/infiniloader/cache.py | 2 ++ src/infiniloader/distributed.py | 7 ++++--- tests/test_distributed.py | 18 ++++++++++++++++++ 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/src/infiniloader/cache.py b/src/infiniloader/cache.py index d1895eb..a2e2d84 100644 --- a/src/infiniloader/cache.py +++ b/src/infiniloader/cache.py @@ -26,6 +26,8 @@ "*/.cache/*", ".iloader-cache/*", "*/.iloader-cache/*", + ".iloader-dist/*", + "*/.iloader-dist/*", ".iloader/*", "*/.iloader/*", "*.lock", diff --git a/src/infiniloader/distributed.py b/src/infiniloader/distributed.py index fcf6b04..5fe157a 100644 --- a/src/infiniloader/distributed.py +++ b/src/infiniloader/distributed.py @@ -196,11 +196,12 @@ def create_rank_views( "total_bytes": sum(item.size for item in files), } marker = rank_dir / ".iloader-dist" - marker.mkdir(exist_ok=True) - (marker / "rank.json").write_text(json.dumps(meta, indent=2, sort_keys=True), encoding="utf-8") + marker.mkdir() + with (marker / "rank.json").open("x", encoding="utf-8") as stream: + json.dump(meta, stream, indent=2, sort_keys=True) root_marker = tmp / ".iloader-dist" - root_marker.mkdir(exist_ok=True) + root_marker.mkdir() manifest = { "name": name, "source": str(source), diff --git a/tests/test_distributed.py b/tests/test_distributed.py index c12ef7a..312fc55 100644 --- a/tests/test_distributed.py +++ b/tests/test_distributed.py @@ -1,6 +1,8 @@ from __future__ import annotations import os +import json +from pathlib import Path import pytest @@ -8,6 +10,22 @@ from infiniloader.distributed import build_tp_plan, create_rank_views, simulate_distributed_read +@pytest.mark.skipif(os.name != "posix", reason="rank views use POSIX symlinks") +def test_rank_view_can_be_used_as_read_only_source(tmp_path): + model = make_model_dir(tmp_path / "model") + first = create_rank_views(model, name="first", ranks=2, base_dir=tmp_path / "state") + source = Path(first["rank_paths"][0]) + original = (source / ".iloader-dist/rank.json").read_bytes() + second = create_rank_views(source, name="second", ranks=3, base_dir=tmp_path / "state") + assert (source / ".iloader-dist/rank.json").read_bytes() == original + for rank, view in enumerate(second["rank_paths"]): + marker = Path(view) / ".iloader-dist/rank.json" + assert not marker.is_symlink() + assert json.loads(marker.read_text())["rank"] == rank + assert (Path(view) / "model.safetensors").read_bytes() == b"0123456789abcdef" + assert second["file_count"] == first["file_count"] + + @pytest.mark.skipif(os.name != "posix", reason="rank views use POSIX symlinks") def test_create_rank_views_and_simulate_read(tmp_path): model = make_model_dir(tmp_path / "model", payload=b"x" * 4096) From 6693f9b1c4fbad667e1bdf91d68e56c231bc6178 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 02:32:34 +0800 Subject: [PATCH 19/29] Recheck disk cache usage under the collection lease --- src/infiniloader/local_cache.py | 5 ++++ tests/test_local_versions.py | 42 +++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/src/infiniloader/local_cache.py b/src/infiniloader/local_cache.py index e77d45a..84c8162 100644 --- a/src/infiniloader/local_cache.py +++ b/src/infiniloader/local_cache.py @@ -312,6 +312,11 @@ def garbage_collect(*, cache_root=None, base_dir=None, max_bytes=None, ttl=None, with FileLock(lease_path(path), timeout=0): if path.is_symlink() or path.parent != root or not read_json(_manifest_path(path)): raise RuntimeError("cache ownership changed during collection") + usage = read_json(root / ".usage" / f"{path.name}.json") or {} + last_used = usage.get("last_used_at", _created_at(path)) + if last_used != row["last_used_at"]: + skipped.append(dict(cache_path=str(path), reason="used since collection scan")) + continue if apply: shutil.rmtree(path) (root / ".usage" / f"{path.name}.json").unlink(missing_ok=True) diff --git a/tests/test_local_versions.py b/tests/test_local_versions.py index bdea26e..e2a8f91 100644 --- a/tests/test_local_versions.py +++ b/tests/test_local_versions.py @@ -9,6 +9,7 @@ from infiniloader.cli import main from infiniloader.local_cache import garbage_collect, lease_path, materialize_model_cache, verify_cache from infiniloader.runtime import FileLock, atomic_json +from infiniloader import local_cache from test_lifecycle import running_daemon pytestmark = pytest.mark.skipif(os.name != "posix", reason="POSIX file locks") @@ -61,6 +62,47 @@ def test_gc_dry_run_lru_and_active_leases(tmp_path): assert garbage_collect(cache_root=root, ttl=0, apply=True)["remaining_bytes"] == 0 +@pytest.mark.parametrize("policy", ["ttl", "budget"]) +def test_gc_rechecks_usage_after_scanning(tmp_path, monkeypatch, policy): + root = tmp_path / "cache" + first = materialize_model_cache(make_model_dir(tmp_path / "a"), cache_root=root) + second = materialize_model_cache(make_model_dir(tmp_path / "b"), cache_root=root) + for index, cache in enumerate((first, second)): + atomic_json(root / ".usage" / f"{cache.cache_path.name}.json", {"last_used_at": index + 1}) + scan = local_cache.list_model_caches + + def scan_then_use(**kwargs): + rows = scan(**kwargs) + with FileLock(lease_path(first.cache_path), shared=True): + local_cache.touch_cache(first.cache_path) + return rows + + monkeypatch.setattr(local_cache, "list_model_caches", scan_then_use) + options = {"ttl": 3600} if policy == "ttl" else {"max_bytes": first.total_bytes} + result = garbage_collect(cache_root=root, apply=True, **options) + assert first.cache_path.exists() + assert not second.cache_path.exists() + assert result["remaining_bytes"] == first.total_bytes + assert result["skipped"] == [{"cache_path": str(first.cache_path), "reason": "used since collection scan"}] + assert result["budget_met"] + + +def test_gc_reports_unmet_budget_when_only_candidate_was_used(tmp_path, monkeypatch): + root = tmp_path / "cache" + cache = materialize_model_cache(make_model_dir(tmp_path / "model"), cache_root=root) + scan = local_cache.list_model_caches + + def scan_then_use(**kwargs): + rows = scan(**kwargs) + local_cache.touch_cache(cache.cache_path) + return rows + + monkeypatch.setattr(local_cache, "list_model_caches", scan_then_use) + result = garbage_collect(cache_root=root, max_bytes=0, apply=True) + assert cache.cache_path.exists() + assert not result["budget_met"] + + def test_daemon_pins_local_disk_generation(tmp_path): model = make_model_dir(tmp_path / "model") local = materialize_model_cache(model, base_dir=tmp_path / "state") From 9b0de02e355d68b84112e0007a0dd9198b9397d7 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 02:37:50 +0800 Subject: [PATCH 20/29] Forward run cancellation and support explicit daemon sockets --- docs/operations.md | 12 ++++++ src/infiniloader/arguments.py | 4 +- src/infiniloader/cli.py | 26 +----------- src/infiniloader/completion.py | 4 +- src/infiniloader/runner.py | 63 +++++++++++++++++++++++++++++ tests/test_completion.py | 2 + tests/test_lifecycle.py | 3 +- tests/test_management.py | 8 ++++ tests/test_runner.py | 73 ++++++++++++++++++++++++++++++++++ 9 files changed, 166 insertions(+), 29 deletions(-) create mode 100644 src/infiniloader/runner.py create mode 100644 tests/test_runner.py diff --git a/docs/operations.md b/docs/operations.md index 672236e..4cd0d35 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -18,6 +18,18 @@ iloader run dev -- vllm serve '{model}' --tensor-parallel-size 4 exit code. Put InfiniLoader options before the cache name. Use a foreground engine command that remains alive while its workers run. +For a daemon started with `--socket PATH`, pass the same override before the name: +`iloader run --socket PATH dev -- COMMAND '{model}'`. This option and `--cache-path` +are mutually exclusive. + +`run` forwards SIGTERM and SIGHUP to its direct child, waits for it to exit, and +kills it after a three-second cancellation grace period. SIGINT uses normal +foreground terminal handling. A signal-terminated child returns `128 + signal` +(for example 143 for SIGTERM); ordinary exit codes are unchanged. The lease lasts +until the child is reaped. Backgrounded descendants are outside this contract. +If the wrapper is killed by an uncatchable signal, a surviving direct child +continues to hold its lease. + For disk-only caching without another model-sized RAM allocation: ```bash diff --git a/src/infiniloader/arguments.py b/src/infiniloader/arguments.py index 1726cf5..19751d9 100644 --- a/src/infiniloader/arguments.py +++ b/src/infiniloader/arguments.py @@ -135,7 +135,9 @@ def build_parser(handlers: Mapping[str, CommandHandler]) -> argparse.ArgumentPar run_p = sub.add_parser("run", help="run a foreground command while protecting its cache from eviction") run_p.add_argument("--base-dir") - run_p.add_argument("--cache-path") + run_location = run_p.add_mutually_exclusive_group() + run_location.add_argument("--cache-path") + run_location.add_argument("--socket", help="socket of the resident cache; use the same override as start") run_p.add_argument("name_arg", nargs="?") run_p.add_argument("program", nargs=argparse.REMAINDER) run_p.set_defaults(func=handlers["cmd_run"]) diff --git a/src/infiniloader/cli.py b/src/infiniloader/cli.py index 38bc2b5..ebf700a 100644 --- a/src/infiniloader/cli.py +++ b/src/infiniloader/cli.py @@ -15,6 +15,7 @@ from .benchmark import benchmark_path from .cache import DEFAULT_EXCLUDES, build_view, human_bytes, inspect_source, load_model_into_memory, source_identity, _safe_remove_view from .runtime import FileLock, process_matches, process_token, read_json, signal_owned_process, validate_name, validate_socket_path +from .runner import run_program as _run_program from .daemon import WeightLoaderDaemon from .distributed import build_tp_plan, create_rank_views, simulate_distributed_read from .ipc import LoaderConnectionError, LoaderStoppingError, request @@ -741,31 +742,6 @@ def cmd_run(args: argparse.Namespace) -> int: pass -def _run_program(program, path, *, pass_fds=(), socket_path=None, lease=None) -> int: - env = dict(os.environ, ILOADER_MODEL_PATH=str(path)) - proc = subprocess.Popen([part.replace("{model}", str(path)) for part in program], env=env, pass_fds=pass_fds) - try: - if lease: - token = process_token(proc.pid) - if proc.poll() is None: - try: - request(socket_path, {"command": "lease-update", "lease": lease, "pid": proc.pid, - "process_token": token}) - except RuntimeError: - if proc.poll() is None: - raise - return proc.wait() - except BaseException: - if proc.poll() is None: - proc.terminate() - try: - proc.wait(timeout=3) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait() - raise - - def cmd_completion(args: argparse.Namespace) -> int: if args.shell == "bash": print(bash_completion(build_parser())) diff --git a/src/infiniloader/completion.py b/src/infiniloader/completion.py index 0f25913..13ce559 100644 --- a/src/infiniloader/completion.py +++ b/src/infiniloader/completion.py @@ -46,9 +46,9 @@ def visit(node: argparse.ArgumentParser, prefix: str) -> None: if (( run_skip )); then run_skip=0; continue; fi case "$word" in --) return 0 ;; - --base-dir) run_skip=1 ;; + --base-dir|--socket) run_skip=1 ;; --cache-path) run_cache=1; run_skip=1 ;; - --base-dir=*) ;; + --base-dir=*|--socket=*) ;; --cache-path=*) run_cache=1 ;; -*) ;; *) run_name=1 ;; diff --git a/src/infiniloader/runner.py b/src/infiniloader/runner.py new file mode 100644 index 0000000..36692ad --- /dev/null +++ b/src/infiniloader/runner.py @@ -0,0 +1,63 @@ +"""Run one foreground consumer while its caller holds the cache lease.""" +from __future__ import annotations + +import os +import signal +import subprocess +import time + +from .ipc import request +from .runtime import process_token + + +def run_program(program, path, *, pass_fds=(), socket_path=None, lease=None) -> int: + cancelled = None + proc = None + handlers = {} + + def cancel(signum, _frame): + nonlocal cancelled + cancelled = signum + + try: + # Install before spawn so cancellation cannot strand a just-created child. + # SIGINT retains Python's KeyboardInterrupt handling: terminal interrupts + # already reach both processes in the foreground process group. + for signum in (signal.SIGTERM, signal.SIGHUP): + handlers[signum] = signal.signal(signum, cancel) + env = dict(os.environ, ILOADER_MODEL_PATH=str(path)) + proc = subprocess.Popen([part.replace("{model}", str(path)) for part in program], + env=env, pass_fds=pass_fds) + if lease: + token = process_token(proc.pid) + if proc.poll() is None: + try: + request(socket_path, {"command": "lease-update", "lease": lease, "pid": proc.pid, + "process_token": token}) + except RuntimeError: + if proc.poll() is None: + raise + deadline = None + while proc.poll() is None: + if cancelled is not None and deadline is None: + proc.send_signal(cancelled) + deadline = time.monotonic() + 3 + if deadline is not None and time.monotonic() >= deadline: + proc.kill() + try: + proc.wait(timeout=.1) + except subprocess.TimeoutExpired: + continue + return proc.returncode if proc.returncode >= 0 else 128 - proc.returncode + finally: + try: + if proc is not None and proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + finally: + for signum, previous in handlers.items(): + signal.signal(signum, previous) diff --git a/tests/test_completion.py b/tests/test_completion.py index 41af4e7..8204508 100644 --- a/tests/test_completion.py +++ b/tests/test_completion.py @@ -39,6 +39,8 @@ def test_nested_options_match_the_parser(): assert {"--layout", "--output", "--mode", "--help"} <= options assert "--rank" not in options assert "--idle-ttl" in complete(["iloader", "start", "--"]) + assert "--socket" in complete(["iloader", "run", "--socket", "/tmp/custom.sock", "--"]) + assert "--socket" in complete(["iloader", "run", "--socket=/tmp/custom.sock", "--"]) @pytest.mark.skipif(os.name != "posix" or not shutil.which("bash"), reason="Bash completion integration") diff --git a/tests/test_lifecycle.py b/tests/test_lifecycle.py index c9230a3..adeb974 100644 --- a/tests/test_lifecycle.py +++ b/tests/test_lifecycle.py @@ -50,7 +50,8 @@ def registered(): return (state / "run/unit.state.json").exists() wait_for(registered, timeout=15) if not slow: - wait_for(lambda: (state / "run/unit.sock").exists()) + socket = Path(extra[extra.index("--socket") + 1]) if "--socket" in extra else state / "run/unit.sock" + wait_for(lambda: socket.exists()) yield proc, state finally: if proc.poll() is None: diff --git a/tests/test_management.py b/tests/test_management.py index 61da33c..8ede764 100644 --- a/tests/test_management.py +++ b/tests/test_management.py @@ -40,6 +40,14 @@ def test_run_substitutes_model_path_and_releases_lease(tmp_path): assert request(state / "run/unit.sock", {"command": "status"})["status"]["active_leases"] == 0 +def test_run_uses_the_explicit_start_socket(tmp_path): + socket = tmp_path / "custom.sock" + with running_daemon(tmp_path, make_model_dir(tmp_path / "model"), extra=["--socket", str(socket)]) as (_, state): + assert main(["run", "--base-dir", str(state), "--socket", str(socket), "unit", "--", + sys.executable, "-c", "import os,pathlib; assert pathlib.Path(os.environ['ILOADER_MODEL_PATH']).is_dir()"]) == 0 + assert request(socket, {"command": "status"})["status"]["active_leases"] == 0 + + def test_run_disk_cache_preserves_command_exit_code(tmp_path): local = materialize_model_cache(make_model_dir(tmp_path / "model"), base_dir=tmp_path / "state") assert main(["run", "--cache-path", str(local.cache_path), "--", sys.executable, "-c", "raise SystemExit(7)"]) == 7 diff --git a/tests/test_runner.py b/tests/test_runner.py new file mode 100644 index 0000000..f472882 --- /dev/null +++ b/tests/test_runner.py @@ -0,0 +1,73 @@ +import os +from pathlib import Path +import signal +import subprocess +import sys + +import pytest + +from conftest import make_model_dir +from infiniloader.local_cache import garbage_collect, materialize_model_cache +from infiniloader.runner import run_program +from infiniloader.runtime import process_matches, process_token, signal_owned_process +from test_lifecycle import wait_for + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux process and lease tests") + + +@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGHUP]) +def test_child_signal_exit_uses_shell_status(tmp_path, signum): + assert run_program([sys.executable, "-c", f"import os; os.kill(os.getpid(), {int(signum)})"], tmp_path) == 128 + signum + + +@pytest.mark.parametrize("signum,ignore,expected", [ + (signal.SIGTERM, False, 143), (signal.SIGHUP, False, 129), + (signal.SIGINT, False, 130), (signal.SIGTERM, True, 137), + (signal.SIGKILL, False, -signal.SIGKILL), +]) +def test_wrapper_cancellation_and_surviving_child_lease(tmp_path, signum, ignore, expected): + cache = materialize_model_cache(make_model_dir(tmp_path / "model"), base_dir=tmp_path / "state") + marker = tmp_path / "child.pid" + code = ( + "import os, pathlib, signal, sys, time; " + + ("signal.signal(signal.SIGTERM, signal.SIG_IGN); " if ignore else "") + + "pathlib.Path(sys.argv[1]).write_text(str(os.getpid())); time.sleep(60)" + ) + env = dict(os.environ, PYTHONPATH=str(Path(__file__).resolve().parents[1] / "src")) + log = (tmp_path / "wrapper.log").open("w") + wrapper = subprocess.Popen([sys.executable, "-m", "infiniloader", "run", "--cache-path", str(cache.cache_path), + "--", sys.executable, "-c", code, str(marker)], env=env, + stdout=log, stderr=log) + child, token = None, None + try: + def started(): + assert wrapper.poll() is None, (tmp_path / "wrapper.log").read_text() + return marker.exists() and marker.read_text() + wait_for(started, timeout=20) + child = int(marker.read_text()) + token = process_token(child) + assert token is not None + assert not garbage_collect(base_dir=tmp_path / "state", max_bytes=0, apply=True)["budget_met"] + wrapper.send_signal(signum) + assert wrapper.wait(timeout=8) == expected + if signum == signal.SIGKILL: + assert process_matches(child, token) + assert not garbage_collect(base_dir=tmp_path / "state", max_bytes=0, apply=True)["budget_met"] + else: + assert not process_matches(child, token) + assert garbage_collect(base_dir=tmp_path / "state", max_bytes=0, apply=True)["budget_met"] + finally: + if wrapper.poll() is None: + wrapper.kill() + wrapper.wait(timeout=3) + if child is not None and process_matches(child, token): + signal_owned_process(child, token) + wait_for(lambda: not process_matches(child, token)) + log.close() + + +def test_signal_handlers_are_restored_when_spawn_fails(tmp_path): + before = {s: signal.getsignal(s) for s in (signal.SIGTERM, signal.SIGHUP)} + with pytest.raises(FileNotFoundError): + run_program([str(tmp_path / "missing-command")], tmp_path) + assert {s: signal.getsignal(s) for s in before} == before From d51500199afc9632a6a84a321b121829186bb63a Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 02:38:03 +0800 Subject: [PATCH 21/29] Discount completed resident files from pending headroom reservations --- docs/operations.md | 4 +++ src/infiniloader/daemon.py | 2 ++ src/infiniloader/memory.py | 22 ++++++++++++-- tests/test_memory.py | 59 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 85 insertions(+), 2 deletions(-) diff --git a/docs/operations.md b/docs/operations.md index 4cd0d35..397de71 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -89,6 +89,10 @@ iloader start "$MODEL" --name dev --workers 4 \ directory. A subsequent start can explicitly update the cap. In-flight loads reserve model and buffer bytes; ready caches reserve model bytes. Dead reservations are pruned using boot ID and PID start time. +- Concurrent admissions deduct the outstanding part of pending reservations from + current headroom. Completed sealed files are already reflected in headroom and + are not deducted again. Incomplete files and worker buffers remain conservatively + reserved until completion; quotas are admission estimates, not a physical RAM guarantee. - `--force` bypasses the available-memory check, while explicit per-cache and shared caps remain enforced. Quotas account for InfiniLoader, not engine allocations. - Byte options accept integers and binary suffixes such as `256MiB` and `32GiB`. diff --git a/src/infiniloader/daemon.py b/src/infiniloader/daemon.py index 5da576f..433a5fb 100644 --- a/src/infiniloader/daemon.py +++ b/src/infiniloader/daemon.py @@ -110,6 +110,8 @@ def _write_state(self, state: dict[str, Any]) -> None: self._claim_runtime_file(self.state_path) def _progress(self, copied: int, total: int, rel: Path) -> None: + if self._reservation is not None: + self._reservation.materialized(copied) pct = copied / total * 100 if total else 100.0 elapsed = time.time() - self.load_started_at if self.load_started_at else 0.0 rate = copied / elapsed if elapsed > 0 and copied > 0 else None diff --git a/src/infiniloader/memory.py b/src/infiniloader/memory.py index 84a7e85..ec4444d 100644 --- a/src/infiniloader/memory.py +++ b/src/infiniloader/memory.py @@ -112,15 +112,33 @@ def __enter__(self): required = self.total + self.buffers if cap is not None and reserved + required > cap: raise MemoryError(f"resident quota exceeded: {reserved} reserved + {required} requested > {cap}") - pending = sum(row["reserved_bytes"] for row in data["entries"].values() if row["pending"]) + pending = 0 + for row in data["entries"].values(): + if row["pending"]: + credit = row.get("materialized_bytes", 0) + if type(credit) is not int or not 0 <= credit <= row["model_bytes"]: + credit = 0 + pending += row["reserved_bytes"] - credit self.info = check_memory(required, pending_bytes=pending, **self.options) self.info["max_resident_bytes"] = cap data["max_resident_bytes"] = cap data["entries"][self.key] = dict(name=self.name, pid=self.pid, process_token=self.token, - model_bytes=self.total, reserved_bytes=required, pending=True) + model_bytes=self.total, reserved_bytes=required, + materialized_bytes=0, pending=True) atomic_json(self.run_dir / "resident.json", data) return self + def materialized(self, copied_bytes: int): + """Credit only complete, sealed files still owned by this loader.""" + with FileLock(self.run_dir / "resident.lock"): + data = self._read() + row = data["entries"][self.key] + if (type(copied_bytes) is not int or not row["pending"] + or not row.get("materialized_bytes", 0) <= copied_bytes <= self.total): + raise ValueError("invalid materialized preload progress") + row["materialized_bytes"] = copied_bytes + atomic_json(self.run_dir / "resident.json", data) + def ready(self): with FileLock(self.run_dir / "resident.lock"): data = self._read() diff --git a/tests/test_memory.py b/tests/test_memory.py index d70ebd7..e06e9a1 100644 --- a/tests/test_memory.py +++ b/tests/test_memory.py @@ -63,6 +63,65 @@ def test_reservations_enforce_shared_quota_and_remove_stale_pids(tmp_path, monke assert not json.loads((tmp_path / "resident.json").read_text())["entries"] +@pytest.mark.skipif(os.name != "posix", reason="POSIX file locks") +@pytest.mark.parametrize("completed,ready,available,required,allowed", [ + (0, False, 100, 40, False), + (45, False, 55, 20, True), + (55, False, 45, 20, True), + (55, True, 45, 20, True), +]) +def test_pending_budget_deducts_only_outstanding_commitments(tmp_path, monkeypatch, completed, ready, available, required, allowed): + headroom = [100] + monkeypatch.setattr(memory, "memory_info", lambda: dict(available_bytes=headroom[0])) + with ResidentReservation(tmp_path, "a", 55, 5, reserve_bytes=0) as a: + a.materialized(completed) + if ready: + a.ready() + headroom[0] = available + if allowed: + with ResidentReservation(tmp_path, "b", required, 0, reserve_bytes=0) as b: + assert b.info["pending_bytes"] == (0 if ready else 60 - completed) + else: + with pytest.raises(MemoryError, match="effective budget"): + with ResidentReservation(tmp_path, "b", required, 0, reserve_bytes=0): + pass + + +@pytest.mark.skipif(os.name != "posix", reason="POSIX file locks") +def test_progress_does_not_discount_quota_and_rejects_invalid_credit(tmp_path, monkeypatch): + monkeypatch.setattr(memory, "memory_info", lambda: dict(available_bytes=10000)) + with ResidentReservation(tmp_path, "a", 55, 5, max_resident_bytes=70, reserve_bytes=0) as a: + a.materialized(45) + for invalid in (44, 56, True): + with pytest.raises(ValueError, match="progress"): + a.materialized(invalid) + with pytest.raises(MemoryError, match="quota"): + with ResidentReservation(tmp_path, "b", 11, 0, reserve_bytes=0): + pass + a.ready() + with ResidentReservation(tmp_path, "b", 15, 0, reserve_bytes=0): + pass + + +@pytest.mark.skipif(os.name != "posix", reason="POSIX file locks") +@pytest.mark.parametrize("credit", [None, 1000, -1, True]) +def test_missing_or_invalid_progress_retains_full_pending_reservation(tmp_path, monkeypatch, credit): + headroom = [100] + monkeypatch.setattr(memory, "memory_info", lambda: dict(available_bytes=headroom[0])) + with ResidentReservation(tmp_path, "a", 55, 5, reserve_bytes=0) as a: + path = tmp_path / "resident.json" + data = json.loads(path.read_text()) + if credit is None: + data["entries"][a.key].pop("materialized_bytes") + else: + data["entries"][a.key]["materialized_bytes"] = credit + atomic_json(path, data) + headroom[0] = 55 + with pytest.raises(MemoryError, match="effective budget"): + with ResidentReservation(tmp_path, "b", 20, 0, reserve_bytes=0): + pass + + @pytest.mark.parametrize("value, expected", [("0", 0), ("2GiB", 2 * 1024**3), ("256M", 256 * 1024**2)]) def test_human_readable_sizes(value, expected): assert parse_bytes(value) == expected From 53603385f2be71a63d49ed404d90950521ffe20d Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 02:41:42 +0800 Subject: [PATCH 22/29] Use one tensor layout contract for planning and runtime views --- src/infiniloader/layouts.py | 28 ++++++++++++++++++++++++++++ src/infiniloader/shards.py | 23 ++++------------------- src/infiniloader/torch_store.py | 18 +++++------------- tests/test_torch_store.py | 23 +++++++++++++++++++++++ 4 files changed, 60 insertions(+), 32 deletions(-) create mode 100644 src/infiniloader/layouts.py diff --git a/src/infiniloader/layouts.py b/src/infiniloader/layouts.py new file mode 100644 index 0000000..cd7d1d0 --- /dev/null +++ b/src/infiniloader/layouts.py @@ -0,0 +1,28 @@ +"""One explicit TP/PP layout contract for planning and canonical tensor views.""" +from __future__ import annotations + + +def validate_layout(tensors: dict, layout: dict) -> tuple[int, int, dict]: + if not isinstance(layout, dict) or layout.get("format") != 1: + raise ValueError("layout requires format=1") + tp, pp, rules = layout.get("tp"), layout.get("pp"), layout.get("tensors") + if (type(tp) is not int or tp < 1 or type(pp) is not int or pp < 1 + or not isinstance(rules, dict)): + raise ValueError("layout requires positive tp, pp and explicit tensor rules") + if set(rules) != set(tensors): + raise ValueError("layout must explicitly cover exactly all tensors") + stages = set() + for name, spec in tensors.items(): + rule = rules[name] + if not isinstance(rule, dict) or set(rule) != {"stage", "split_dim"}: + raise ValueError(f"rule requires stage and split_dim: {name}") + stage, dim = rule["stage"], rule["split_dim"] + if type(stage) is not int or not 0 <= stage < pp: + raise ValueError(f"invalid pipeline stage: {name}") + if dim is not None and (type(dim) is not int or not 0 <= dim < len(spec["shape"]) + or spec["shape"][dim] % tp): + raise ValueError(f"invalid or nondivisible split dimension: {name}") + stages.add(stage) + if len(stages) != pp: + raise ValueError("every pipeline stage must own at least one tensor") + return tp, pp, rules diff --git a/src/infiniloader/shards.py b/src/infiniloader/shards.py index 3a42b3b..8b33055 100644 --- a/src/infiniloader/shards.py +++ b/src/infiniloader/shards.py @@ -18,6 +18,7 @@ from pathlib import Path from .formats import DTYPE_BYTES, MAX_HEADER, _signature, canonical, read_json, tensor_header, tensor_inventory +from .layouts import validate_layout CHUNK = 1024 * 1024 @@ -65,23 +66,9 @@ def catalog(source: Path) -> dict: def build_plan(source: Path, layout: dict, *, mode="sharded", capacities: list[int] | None = None) -> dict: cat = catalog(source) - if not isinstance(layout, dict) or layout.get("format") != 1: - raise ValueError("layout requires format=1") - tp, pp = layout.get("tp"), layout.get("pp") - rules = layout.get("tensors") - if not _int(tp, minimum=1) or not _int(pp, minimum=1) or not isinstance(rules, dict): - raise ValueError("layout requires positive tp, pp and explicit tensor rules") - if mode not in ("sharded", "replicated") or set(rules) != set(cat["tensors"]): - raise ValueError("invalid mode or layout must explicitly cover exactly all tensors") - for name, spec in cat["tensors"].items(): - rule = rules[name] - if not isinstance(rule, dict) or set(rule) != {"stage", "split_dim"}: - raise ValueError(f"rule requires stage and split_dim: {name}") - stage, dim = rule["stage"], rule["split_dim"] - if not _int(stage) or stage >= pp: - raise ValueError(f"invalid pipeline stage: {name}") - if dim is not None and (not _int(dim) or dim >= len(spec["shape"]) or spec["shape"][dim] % tp): - raise ValueError(f"invalid or nondivisible split dimension: {name}") + tp, pp, rules = validate_layout(cat["tensors"], layout) + if mode not in ("sharded", "replicated"): + raise ValueError("invalid shard mode") ranks = [] for rank in range(tp * pp): stage, lane = divmod(rank, tp) @@ -100,8 +87,6 @@ def build_plan(source: Path, layout: dict, *, mode="sharded", capacities: list[i "bytes": math.prod(shape) * DTYPE_BYTES[spec["dtype"]]} ranks.append({"rank": rank, "stage": stage, "lane": lane, "tensors": parts, "tensor_bytes": sum(p["bytes"] for p in parts.values())}) - if any(not r["tensors"] for r in ranks): - raise ValueError("every pipeline stage must own at least one tensor") if capacities is not None: if len(capacities) != len(ranks) or not all(_int(v) for v in capacities): raise ValueError("capacities requires one nonnegative byte limit per rank") diff --git a/src/infiniloader/torch_store.py b/src/infiniloader/torch_store.py index bfc584e..15cc8cc 100644 --- a/src/infiniloader/torch_store.py +++ b/src/infiniloader/torch_store.py @@ -11,6 +11,7 @@ from pathlib import Path from .shards import catalog, verify_rank, read_json +from .layouts import validate_layout class TensorStore: @@ -75,28 +76,19 @@ def rank_tensors(self, layout: dict, rank: int) -> dict: """ if self._closed: raise RuntimeError("tensor store is closed") - tp, pp, rules = layout.get("tp"), layout.get("pp"), layout.get("tensors") - if (layout.get("format") != 1 or type(tp) is not int or tp < 1 - or type(pp) is not int or pp < 1 or type(rank) is not int or not 0 <= rank < tp * pp - or not isinstance(rules, dict)): + if type(rank) is not int: raise ValueError("invalid tensor layout/rank") if self.receipt and self.receipt["plan"]["mode"] == "sharded": if layout != self.receipt["plan"]["layout"] or rank != self.receipt["rank"]: raise ValueError("sharded store requires its exported layout and rank") return {n: self.tensor(n) for n in self.catalog["tensors"]} - if set(rules) != set(self.catalog["tensors"]): - raise ValueError("layout must cover exactly all canonical tensors") + tp, pp, rules = validate_layout(self.catalog["tensors"], layout) + if not 0 <= rank < tp * pp: + raise ValueError("invalid tensor layout/rank") result = {} for name, spec in self.catalog["tensors"].items(): rule = rules[name] - if not isinstance(rule, dict) or set(rule) != {"stage", "split_dim"}: - raise ValueError(f"invalid tensor rule: {name}") stage, dim = rule["stage"], rule["split_dim"] - if type(stage) is not int or not 0 <= stage < pp: - raise ValueError(f"invalid stage: {name}") - if dim is not None and (type(dim) is not int or not 0 <= dim < len(spec["shape"]) - or spec["shape"][dim] % tp): - raise ValueError(f"invalid split: {name}") if stage != rank // tp: continue tensor = self.tensor(name) diff --git a/tests/test_torch_store.py b/tests/test_torch_store.py index 1283f07..6d5e0f9 100644 --- a/tests/test_torch_store.py +++ b/tests/test_torch_store.py @@ -74,6 +74,29 @@ def test_sharded_store_enforces_exported_contract(tmp_path): store.rank_tensors(rules, 0) +def test_planner_and_canonical_store_reject_empty_pipeline_stages(tmp_path): + source = tmp_path / "source" + write_tensor(source) + rules = layout(0) + rules["pp"] = 2 + with resident(source, tmp_path / "view") as view, TensorStore(view) as store: + for mode in ("sharded", "replicated"): + with pytest.raises(ValueError, match="every pipeline stage"): + build_plan(source, rules, mode=mode) + for rank in range(4): + with pytest.raises(ValueError, match="every pipeline stage"): + store.rank_tensors(rules, rank) + + +def test_empty_tensor_is_a_valid_pipeline_stage_assignment(tmp_path): + source = tmp_path / "source" + write_tensor(source, shape=(0, 2)) + rules = layout(0) + assert build_plan(source, rules)["ranks"][0]["tensor_bytes"] == 0 + with resident(source, tmp_path / "view") as view, TensorStore(view) as store: + assert store.rank_tensors(rules, 0)["weight"].shape == (0, 2) + + def test_scalar_and_empty_cpu_view(tmp_path): scalar, empty = tmp_path / "scalar", tmp_path / "empty" write_tensor(scalar, shape=()) From 3bd87d599c974a7f4ea6babdbdcd33651addb00e Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 02:42:09 +0800 Subject: [PATCH 23/29] Keep the latest expert completion event per consumer stream --- src/infiniloader/experts.py | 10 ++-- tests/test_expert_events.py | 107 ++++++++++++++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 5 deletions(-) create mode 100644 tests/test_expert_events.py diff --git a/src/infiniloader/experts.py b/src/infiniloader/experts.py index 3c0af96..ab96ef6 100644 --- a/src/infiniloader/experts.py +++ b/src/infiniloader/experts.py @@ -19,7 +19,7 @@ class _Entry: weights: dict nbytes: int users: int = 0 - events: list = field(default_factory=list) + events: dict = field(default_factory=dict) @dataclass @@ -74,7 +74,7 @@ def _evict(self, expert): entry = self.entries[expert] if entry.users: raise RuntimeError("expert is leased") - for event in entry.events: + for event in entry.events.values(): event.synchronize() del self.entries[expert] self._stats["evictions"] += 1 @@ -150,11 +150,11 @@ def acquire(self, expert: str, *, stream=None): for consumer_stream in lease._streams: event = self.torch.cuda.Event() event.record(consumer_stream) - entry.events.append(event) + # Later work on this stream covers earlier leases. + # Retain the stream key as well as its latest event. + entry.events[consumer_stream] = event for tensor in entry.weights.values(): tensor.record_stream(consumer_stream) - # Completed events do not need to accumulate over hot hits. - entry.events[:] = [e for e in entry.events if not e.query()] except BaseException: self.torch.cuda.synchronize(self.device) raise diff --git a/tests/test_expert_events.py b/tests/test_expert_events.py new file mode 100644 index 0000000..5456464 --- /dev/null +++ b/tests/test_expert_events.py @@ -0,0 +1,107 @@ +"""CPU model of event bookkeeping; actual CUDA lifetime tests remain separate.""" +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest + +from infiniloader.experts import ExpertPool, _Entry + + +class Stream: + def __init__(self, device): + self.device = device + + +class Event: + def __init__(self): + self.synchronized = False + + def record(self, stream): + self.stream = stream + + def query(self): + raise AssertionError("hot release must not rescan historical events") + + def synchronize(self): + self.synchronized = True + + +class Cuda: + def __init__(self, stream): + self.current = stream + self.events = [] + self.synchronizations = 0 + + def Event(self): + event = Event() + self.events.append(event) + return event + + def current_stream(self, device): + return self.current + + def device(self, device): + return nullcontext() + + def stream(self, stream): + return nullcontext() + + def synchronize(self, device): + self.synchronizations += 1 + + +def pool_with_stream(): + device = object() + stream = Stream(device) + cuda = Cuda(stream) + recorded = [] + tensor = SimpleNamespace(record_stream=recorded.append) + store = SimpleNamespace(torch=SimpleNamespace(cuda=cuda), + catalog={"tensors": {"w": {"shape": [1], "dtype": "F32"}}}) + pool = ExpertPool(store, {"0": {"up": "w"}}, max_gpu_bytes=4) + pool.device = device + pool.entries["0"] = _Entry({"up": tensor}, 4) + return pool, cuda, stream, recorded + + +def test_repeated_hot_leases_keep_only_latest_event_per_stream(): + pool, cuda, stream, recorded = pool_with_stream() + for _ in range(1000): + with pool.acquire("0"): + pass + assert len(pool.entries["0"].events) == 1 + assert recorded == [stream] * 1000 + pool.evict("0") + assert cuda.events[-1].synchronized + assert not any(event.synchronized for event in cuda.events[:-1]) + + +def test_eviction_waits_for_latest_event_on_every_consumer_stream(): + pool, cuda, first, recorded = pool_with_stream() + second = Stream(pool.device) + with pool.acquire("0") as lease: + lease.record_stream(second) + initial = list(cuda.events) + with pool.acquire("0"): + pass + assert len(pool.entries["0"].events) == 2 + pool.close() + assert initial[1].synchronized and cuda.events[-1].synchronized + assert not initial[0].synchronized + assert recorded == [first, second, first] + + +def test_event_record_failure_synchronizes_device_and_releases_users(monkeypatch): + pool, cuda, stream, recorded = pool_with_stream() + + def fail(self, stream): + raise RuntimeError("injected event failure") + + monkeypatch.setattr(Event, "record", fail) + with pytest.raises(RuntimeError, match="injected"): + with pool.acquire("0") as lease: + pass + assert cuda.synchronizations == 1 + assert pool.status()["active_leases"] == 0 + assert pool.entries["0"].users == 0 + assert not lease.weights From 441202912ffc79e73463fae489bcb909ac7eca32 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 02:43:13 +0800 Subject: [PATCH 24/29] Document lease transfer limits and remaining preparation costs --- docs/distributed.md | 12 +++++++++++- docs/operations.md | 5 +++-- docs/roadmap.zh-CN.md | 4 ++++ 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/docs/distributed.md b/docs/distributed.md index 8743eba..4036db9 100644 --- a/docs/distributed.md +++ b/docs/distributed.md @@ -36,7 +36,13 @@ that these engines consume the new InfiniLoader shard format directly. The legacy `dist-plan` remains a heuristic report. It shares strict format and index validation with `shard`, while reading headers rather than hashing payloads. The new executable path requires -an exact tensor inventory and explicit TP dimension/PP stage for **every** tensor: +an exact tensor inventory and explicit TP dimension/PP stage for **every** tensor. + +Planning and canonical runtime slicing use the same layout validator. Every PP +stage must own at least one tensor, in both export modes; an empty tensor still +counts as an assignment. Replication changes storage placement, not layout validity. + +For example: ```json { @@ -137,6 +143,10 @@ weights default to the current consumer stream; use `acquire(..., stream=stream) and `lease.record_stream(other_stream)` for every additional consumer stream. Before eviction, the pool waits for recorded work and protects allocator reuse. Active leases cannot be evicted. Failed uploads never publish half an expert. +The pool retains only the latest completion event for each distinct consumer +stream per expert, until eviction; it does not scan a history of hot leases. +CPU event-model tests cover this bookkeeping. GPU ordering still requires the +separate real-device tests and is not established by the CPU test model. The budget counts pool-owned weight payload, including admission in progress; it does not cap activations, external tensor references, CUDA context or allocator diff --git a/docs/operations.md b/docs/operations.md index 397de71..f57a291 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -27,8 +27,9 @@ kills it after a three-second cancellation grace period. SIGINT uses normal foreground terminal handling. A signal-terminated child returns `128 + signal` (for example 143 for SIGTERM); ordinary exit codes are unchanged. The lease lasts until the child is reaped. Backgrounded descendants are outside this contract. -If the wrapper is killed by an uncatchable signal, a surviving direct child -continues to hold its lease. +After the child has inherited the disk lease or received the resident lease, +it remains protected if the wrapper is killed by an uncatchable signal. There is +a short spawn-to-transfer window for resident leases where this cannot be guaranteed. For disk-only caching without another model-sized RAM allocation: diff --git a/docs/roadmap.zh-CN.md b/docs/roadmap.zh-CN.md index 5ef4d41..ebc367f 100644 --- a/docs/roadmap.zh-CN.md +++ b/docs/roadmap.zh-CN.md @@ -10,6 +10,10 @@ InfiniLoader 的目标是：模型文件加载一次，让后续进程快速、 后续优化应由这些测量发现的瓶颈驱动，例如重复文件扫描、校验或复制。先记录实际成本，再决定是否需要新的缓存策略或接口。 +已有实现的成本需要单独量化并逐步收敛：磁盘构建目前按 cache root 串行；实验分片导出会按 rank 重复生成完整 catalog，TensorStore 初始化也有重复全量校验。优先考虑缩短构建锁范围、在一次导出任务中复用受校验的快照，并用稳定 FD 和 seals 减少重复读取。调整前后都必须保留源变化和错误发布的检测，不能直接删除校验来换取速度。 + +实验 MoE 的参考计算仍位于 `experts.py`；后续可将计算示例与驻留接口分离。真实引擎验收应分别记录文件读取、校验、映射、参数安装、格式转换、H2D 和初始化成本，避免把预热准备时间遗漏在比较之外。 + 显式分片、TensorStore 和 GPU 专家池保留为[实验接口](distributed.md)。它们不是核心 `start/run` 的依赖，也不作为前期缓存版本的生产引擎兼容保证。在线拓扑迁移、量化内核、路由调度和通用集群服务不属于当前交付范围。 此前的较大范围建议保存在[历史路线](archive/2026-10-26-roadmap.zh-CN.md)，用于说明原始评估，不代表当前承诺。 From 31fde8b0899b04d94721aa6e5962af2fb77b5c3a Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 03:45:46 +0800 Subject: [PATCH 25/29] Preserve inherited ignored signals in foreground consumers --- docs/operations.md | 3 +- src/infiniloader/runner.py | 2 ++ tests/test_runner.py | 56 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 1 deletion(-) diff --git a/docs/operations.md b/docs/operations.md index f57a291..44fc29c 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -22,7 +22,8 @@ For a daemon started with `--socket PATH`, pass the same override before the nam `iloader run --socket PATH dev -- COMMAND '{model}'`. This option and `--cache-path` are mutually exclusive. -`run` forwards SIGTERM and SIGHUP to its direct child, waits for it to exit, and +`run` preserves inherited `SIG_IGN` settings, including SIGHUP ignored by `nohup`. +For other settings, it forwards SIGTERM and SIGHUP to its direct child, waits for it to exit, and kills it after a three-second cancellation grace period. SIGINT uses normal foreground terminal handling. A signal-terminated child returns `128 + signal` (for example 143 for SIGTERM); ordinary exit codes are unchanged. The lease lasts diff --git a/src/infiniloader/runner.py b/src/infiniloader/runner.py index 36692ad..54022d8 100644 --- a/src/infiniloader/runner.py +++ b/src/infiniloader/runner.py @@ -24,6 +24,8 @@ def cancel(signum, _frame): # SIGINT retains Python's KeyboardInterrupt handling: terminal interrupts # already reach both processes in the foreground process group. for signum in (signal.SIGTERM, signal.SIGHUP): + if signal.getsignal(signum) == signal.SIG_IGN: + continue handlers[signum] = signal.signal(signum, cancel) env = dict(os.environ, ILOADER_MODEL_PATH=str(path)) proc = subprocess.Popen([part.replace("{model}", str(path)) for part in program], diff --git a/tests/test_runner.py b/tests/test_runner.py index f472882..0bf7db2 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1,5 +1,7 @@ import os +import json from pathlib import Path +import shutil import signal import subprocess import sys @@ -71,3 +73,57 @@ def test_signal_handlers_are_restored_when_spawn_fails(tmp_path): with pytest.raises(FileNotFoundError): run_program([str(tmp_path / "missing-command")], tmp_path) assert {s: signal.getsignal(s) for s in before} == before + + +@pytest.mark.parametrize("signum", [signal.SIGHUP, signal.SIGTERM]) +def test_wrapper_and_child_preserve_inherited_ignored_signals(tmp_path, signum): + if signum == signal.SIGHUP and not shutil.which("nohup"): + pytest.skip("requires nohup") + cache = materialize_model_cache(make_model_dir(tmp_path / "model"), base_dir=tmp_path / "state") + marker, finish = tmp_path / "consumer.json", tmp_path / "finish" + code = ( + "import json,os,pathlib,signal,sys,time\n" + "pathlib.Path(sys.argv[1]).write_text(json.dumps({'pid':os.getpid()," + f"'ignored':signal.getsignal({int(signum)}) == signal.SIG_IGN}}))\n" + "while not pathlib.Path(sys.argv[2]).exists(): time.sleep(.02)\n" + ) + command = [sys.executable, "-m", "infiniloader", "run", "--cache-path", str(cache.cache_path), + "--", sys.executable, "-c", code, str(marker), str(finish)] + if signum == signal.SIGHUP: + command = ["nohup", *command] + else: + command = [sys.executable, "-c", + "import os,signal,sys; signal.signal(signal.SIGTERM,signal.SIG_IGN); os.execv(sys.argv[1],sys.argv[1:])", + *command] + env = dict(os.environ, PYTHONPATH=str(Path(__file__).resolve().parents[1] / "src")) + child, token = None, None + with (tmp_path / "wrapper.log").open("w") as log: + wrapper = subprocess.Popen(command, env=env, stdin=subprocess.DEVNULL, stdout=log, stderr=log) + try: + def started(): + assert wrapper.poll() is None, (tmp_path / "wrapper.log").read_text() + return marker.exists() and marker.read_text() + wait_for(started, timeout=20) + data = json.loads(marker.read_text()) + child, token = data["pid"], process_token(data["pid"]) + assert data["ignored"] + status = Path(f"/proc/{wrapper.pid}/status").read_text() + mask = int(next(line.split()[1] for line in status.splitlines() if line.startswith("SigIgn:")), 16) + assert mask & (1 << (signum - 1)) + wrapper.send_signal(signum) + os.kill(child, signum) + assert process_matches(child, token) + assert not garbage_collect(base_dir=tmp_path / "state", max_bytes=0, apply=True)["budget_met"] + finish.touch() + assert wrapper.wait(timeout=8) == 0 + assert garbage_collect(base_dir=tmp_path / "state", max_bytes=0, apply=True)["budget_met"] + finally: + finish.touch() + if wrapper.poll() is None: + try: + wrapper.wait(timeout=3) + except subprocess.TimeoutExpired: + wrapper.kill() + wrapper.wait(timeout=3) + if child is not None: + wait_for(lambda: not process_matches(child, token)) From df177cfa4adf471aee93b36dc3d1701456da7db8 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 03:50:24 +0800 Subject: [PATCH 26/29] Retain recovery state until runtime cleanup completes --- docs/operations.md | 7 ++ src/infiniloader/daemon.py | 66 ++++++++++----- tests/test_runtime_cleanup.py | 150 ++++++++++++++++++++++++++++++++++ 3 files changed, 203 insertions(+), 20 deletions(-) create mode 100644 tests/test_runtime_cleanup.py diff --git a/docs/operations.md b/docs/operations.md index 44fc29c..e965d89 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -110,6 +110,13 @@ registry shared between nodes is unsupported. timeout is exceeded. Raw IPC stop requests are asynchronous. A subsequent `start` waits for a stopping instance within its readiness timeout. +Shutdown removes verified PID and socket resources before deleting their state. +An uncertain deletion or changed ownership retains the state and emits a cleanup +diagnostic, including in quiet mode. On the next start, unchanged owned stale +resources can be reclaimed; unrelated replacements are preserved for inspection. +If a recorded private socket cannot be verified or removed during recovery, +startup fails before replacing the old state, so a retry retains the same evidence. + `run` protects a cache against ordinary stop and eviction. Local-cache daemons also hold a disk-generation lease. Foreground disk commands inherit its OS lock descriptor; resident commands register the child process's PID and start time. diff --git a/src/infiniloader/daemon.py b/src/infiniloader/daemon.py index 433a5fb..027f208 100644 --- a/src/infiniloader/daemon.py +++ b/src/infiniloader/daemon.py @@ -12,6 +12,7 @@ import uuid import stat import errno +import sys from pathlib import Path from typing import Any @@ -298,13 +299,28 @@ def _serve_connection(self, conn: socket.socket) -> None: def _cleanup_socket(self) -> None: if not self._owns_runtime: return + complete = True for path, identity in self._claimed_runtime.items(): - try: - st = path.lstat() - if (st.st_dev, st.st_ino) == identity: - path.unlink() - except OSError: - pass + if path != self.state_path and not self._remove_claimed_runtime_file(path, identity): + complete = False + # State proves ownership on the next start. Keep it until all other + # claimed resources have been removed or confirmed absent. + if complete and self.state_path in self._claimed_runtime: + self._remove_claimed_runtime_file(self.state_path, self._claimed_runtime[self.state_path]) + + def _remove_claimed_runtime_file(self, path: Path, identity: tuple[int, int]) -> bool: + try: + st = path.lstat() + if (st.st_dev, st.st_ino) != identity or st.st_uid != os.getuid(): + print(f"[cleanup] retained {path}: runtime ownership changed", file=sys.stderr, flush=True) + return False + path.unlink() + return True + except FileNotFoundError: + return True + except OSError as exc: + print(f"[cleanup] retained {path}: {exc}", file=sys.stderr, flush=True) + return False def _claim_runtime_file(self, path: Path) -> None: st = path.lstat() @@ -364,21 +380,31 @@ def _check_runtime_paths(self) -> None: # the old PID first avoids publishing a new state next to an old PID. if state: private = state.get("runtime_private_socket") - if isinstance(private, str): + if private is not None: + if not isinstance(private, str): + raise RuntimeError("invalid private socket path in runtime state") path = absolute_path(private) - if path.parent == self.socket_path.parent and path not in paths: - try: - st = path.lstat() - expected = {"device": st.st_dev, "inode": st.st_ino} - if stat.S_ISSOCK(st.st_mode) and st.st_uid == os.getuid() and expected == state.get("runtime_socket"): - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: - probe.settimeout(.2) - try: - probe.connect(str(path)) - except ConnectionRefusedError: - self._remove_stale_file(path, st) - except OSError: - pass # Preserve anything whose ownership/liveness is uncertain. + if path.parent != self.socket_path.parent or path in paths: + raise RuntimeError("invalid private socket path in runtime state") + try: + st = path.lstat() + except FileNotFoundError: + st = None + if st is not None: + expected = {"device": st.st_dev, "inode": st.st_ino} + if (not stat.S_ISSOCK(st.st_mode) or st.st_uid != os.getuid() + or expected != state.get("runtime_socket")): + raise RuntimeError(f"private socket has no matching owned state: {path}") + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: + probe.settimeout(.2) + try: + probe.connect(str(path)) + except ConnectionRefusedError: + self._remove_stale_file(path, st) + except OSError as exc: + raise RuntimeError(f"cannot establish that private socket is stale: {path}") from exc + else: + raise RuntimeError(f"private socket still has a listener: {path}") if stale_pid is not None: self._remove_stale_file(self.pid_path, stale_pid) if stale_socket is not None: diff --git a/tests/test_runtime_cleanup.py b/tests/test_runtime_cleanup.py new file mode 100644 index 0000000..a68f94f --- /dev/null +++ b/tests/test_runtime_cleanup.py @@ -0,0 +1,150 @@ +import errno +import json +import os +from pathlib import Path +import socket +import subprocess +import sys +from unittest.mock import Mock + +import pytest + +from infiniloader.daemon import WeightLoaderDaemon +from infiniloader import cli +from test_runtime_publication import runtime, start, crash_serve + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux runtime cleanup") + + +def serve_with_cleanup_fault(model, base, fault, target): + code = f'''from pathlib import Path +import errno, os +from infiniloader.cli import main +from infiniloader.daemon import WeightLoaderDaemon +cleaning = False +publish = WeightLoaderDaemon._publish_socket +cleanup = WeightLoaderDaemon._cleanup_socket +unlink = Path.unlink +def publish_then_stop(self, server): + publish(self, server) + self._stop = True +def begin_cleanup(self): + global cleaning + cleaning = True + cleanup(self) +def injected_unlink(path, *args, **kwargs): + selected = cleaning and path.name == {target!r} + if selected and {fault!r} == 'unlink_error': + raise OSError(errno.EIO, 'injected cleanup failure') + result = unlink(path, *args, **kwargs) + if selected and {fault!r} == 'exit_after_unlink': + os._exit(77) + return result +WeightLoaderDaemon._publish_socket = publish_then_stop +WeightLoaderDaemon._cleanup_socket = begin_cleanup +Path.unlink = injected_unlink +raise SystemExit(main()) +''' + result = subprocess.run([sys.executable, "-c", code, "_serve", str(model), "--name", "unit", + "--base-dir", str(base), "--replace", "--quiet"], + capture_output=True, text=True, timeout=20) + expected = 77 if fault == "exit_after_unlink" else 0 + assert result.returncode == expected, result.stderr + return result + + +@pytest.mark.parametrize("target", ["unit.pid", "unit.sock", "unit.state.json"]) +def test_failed_unlink_keeps_state_and_allows_restart(runtime, target): + model, base = runtime + result = serve_with_cleanup_fault(model, base, "unlink_error", target) + assert (base / "run" / target).exists() + state = base / "run/unit.state.json" + assert json.loads(state.read_text())["runtime_format"] == 1 + assert "injected cleanup failure" in result.stderr + assert target in result.stderr + assert start(model, base) == 0 + + +@pytest.mark.parametrize("target", ["unit.pid", "unit.sock", "unit.state.json"]) +def test_cleanup_interruption_leaves_enough_evidence_to_restart(runtime, target): + model, base = runtime + serve_with_cleanup_fault(model, base, "exit_after_unlink", target) + if target == "unit.state.json": + assert not (base / "run/unit.pid").exists() + assert not (base / "run/unit.sock").exists() + else: + assert (base / "run/unit.state.json").exists() + assert start(model, base) == 0 + + +@pytest.mark.parametrize("fault", ["unlink_error", "already_missing", "disappeared_during_unlink", "replaced"]) +def test_private_socket_cleanup_preserves_state_on_uncertainty(runtime, monkeypatch, capsys, fault): + _, base = runtime + run = base / "run" + run.mkdir(parents=True) + daemon = WeightLoaderDaemon.__new__(WeightLoaderDaemon) + daemon._owns_runtime = True + daemon.state_path = run / "unit.state.json" + daemon._claimed_runtime = {} + daemon.state_path.write_text("owned state") + daemon._claim_runtime_file(daemon.state_path) + private = run / ".private" + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as server: + server.bind(str(private)) + daemon._claim_runtime_file(private) + real_unlink = Path.unlink + if fault == "already_missing": + private.unlink() + elif fault == "replaced": + replacement = run / "replacement" + replacement.write_text("keep this unrelated file") + replacement.replace(private) + else: + def fail(path, *args, **kwargs): + if path == private: + if fault == "disappeared_during_unlink": + real_unlink(path, *args, **kwargs) + raise FileNotFoundError(path) + raise OSError(errno.EIO, "injected private socket failure") + return real_unlink(path, *args, **kwargs) + monkeypatch.setattr(Path, "unlink", fail) + daemon._cleanup_socket() + if fault in ("already_missing", "disappeared_during_unlink"): + assert not daemon.state_path.exists() + else: + assert daemon.state_path.exists() and private.exists() + assert "[cleanup] retained" in capsys.readouterr().err + if fault == "replaced": + assert private.read_text() == "keep this unrelated file" + + +def test_private_unlink_error_during_restart_keeps_original_recovery_state(runtime, monkeypatch, capsys): + model, base = runtime + crash_serve(model, base, "after_publication") + state_path = base / "run/unit.state.json" + original_state = state_path.read_bytes() + state = json.loads(original_state) + private = Path(state["runtime_private_socket"]) + daemon = WeightLoaderDaemon.__new__(WeightLoaderDaemon) + daemon._owns_runtime, daemon.state_path, daemon._claimed_runtime = True, state_path, {} + for path in (state_path, base / "run/unit.pid", private, base / "run/unit.sock"): + daemon._claim_runtime_file(path) + real_unlink = Path.unlink + + def fail(path, *args, **kwargs): + if path == private: + raise OSError(errno.EIO, "persistent private socket failure") + return real_unlink(path, *args, **kwargs) + + with monkeypatch.context() as patch: + patch.setattr(Path, "unlink", fail) + daemon._cleanup_socket() + assert state_path.read_bytes() == original_state + serve = Mock(return_value=0) + patch.setattr(WeightLoaderDaemon, "_serve_owned", serve) + assert cli.main(["_serve", str(model), "--name", "unit", "--base-dir", str(base), "--replace", "--quiet"]) == 1 + serve.assert_not_called() + assert state_path.read_bytes() == original_state + assert private.exists() + assert "persistent private socket failure" in capsys.readouterr().err + assert start(model, base) == 0 From ca908bb03a342efd9b112b985e3dd9eb37d35619 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 03:50:37 +0800 Subject: [PATCH 27/29] Start daemon readiness waits after local cache preparation --- docs/operations.md | 7 +++++ src/infiniloader/arguments.py | 3 +- src/infiniloader/cli.py | 1 + tests/test_start_preparation.py | 54 +++++++++++++++++++++++++++++++++ 4 files changed, 64 insertions(+), 1 deletion(-) create mode 100644 tests/test_start_preparation.py diff --git a/docs/operations.md b/docs/operations.md index e965d89..c24a487 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -104,6 +104,13 @@ Relative paths are anchored at startup. The encoded socket path must fit in 107 bytes; use a shorter `--base-dir` or explicit `--socket` when needed. A runtime registry shared between nodes is unsupported. +`start --wait-timeout SECONDS` bounds runtime waits, not local disk-cache +preparation. A newly spawned daemon gets a fresh interval for readiness, or for +PID/state registration with `--no-wait`. Lock acquisition and waiting for an +existing instance still use bounded waits. Preparation failures return before a +daemon is spawned; a subsequent readiness timeout can leave the daemon loading, +so inspect `status` and the log before retrying. + ## Leases, LRU and TTL `stop` waits for the original daemon PID/token to exit and returns an error if its diff --git a/src/infiniloader/arguments.py b/src/infiniloader/arguments.py index 19751d9..20388ae 100644 --- a/src/infiniloader/arguments.py +++ b/src/infiniloader/arguments.py @@ -68,7 +68,8 @@ def build_parser(handlers: Mapping[str, CommandHandler]) -> argparse.ArgumentPar start_p.add_argument("--exclude", action="append") start_p.add_argument("--force", action="store_true") start_p.add_argument("--no-wait", action="store_true") - start_p.add_argument("--wait-timeout", type=float, default=600.0) + start_p.add_argument("--wait-timeout", type=float, default=600.0, + help="daemon readiness/registration wait in seconds; local cache preparation is excluded") start_p.add_argument("--chunk-size", type=int, default=64 * 1024 * 1024) start_p.add_argument("--workers", type=int, default=1, help="parallel memfd preload workers") start_p.add_argument("--max-memory-fraction", type=float, default=0.90) diff --git a/src/infiniloader/cli.py b/src/infiniloader/cli.py index ebf700a..5ca72f1 100644 --- a/src/infiniloader/cli.py +++ b/src/infiniloader/cli.py @@ -361,6 +361,7 @@ def _start_locked(args: argparse.Namespace) -> int: stdout=log, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, start_new_session=True, close_fds=True, ) + deadline = time.monotonic() + args.wait_timeout print(f"starting name={args.name} pid={proc.pid} log={log_path}", file=sys.stderr) last_progress = None diff --git a/tests/test_start_preparation.py b/tests/test_start_preparation.py new file mode 100644 index 0000000..992da71 --- /dev/null +++ b/tests/test_start_preparation.py @@ -0,0 +1,54 @@ +from pathlib import Path +import os +import time +from unittest.mock import Mock + +import pytest + +from infiniloader import cli +from infiniloader.ipc import request +from test_lifecycle import wait_for +from test_runtime_publication import runtime + +pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux daemon preparation") + + +@pytest.mark.parametrize("no_wait", [False, True]) +def test_slow_local_preparation_does_not_consume_daemon_wait(runtime, monkeypatch, capsys, no_wait): + model, base = runtime + real_clock = time.monotonic + elapsed = [0] + prepare = cli._prepare_model_source + + def slow_prepare(args): + result = prepare(args) + assert result[2] is not None + elapsed[0] += 20 + return result + + monkeypatch.setattr(cli.time, "monotonic", lambda: real_clock() + elapsed[0]) + monkeypatch.setattr(cli, "_prepare_model_source", slow_prepare) + args = ["start", str(model), "--name", "unit", "--base-dir", str(base), "--local-cache", "--wait-timeout", "10"] + if no_wait: + args.append("--no-wait") + assert cli.main(args) == 0 + assert capsys.readouterr().out.strip() == str(base / "views/unit") + wait_for(lambda: (base / "run/unit.sock").exists()) + status = request(base / "run/unit.sock", {"command": "status"})["status"] + assert status["state"] == "ready" + assert (Path(status["view"]) / "model.safetensors").read_bytes() == b"0123456789abcdef" + + +def test_preparation_failure_does_not_spawn_daemon(runtime, monkeypatch, capsys): + model, base = runtime + spawn = Mock() + + def fail(args): + raise OSError("injected preparation failure") + + monkeypatch.setattr(cli, "_prepare_model_source", fail) + monkeypatch.setattr(cli.subprocess, "Popen", spawn) + assert cli.main(["start", str(model), "--name", "unit", "--base-dir", str(base), "--local-cache"]) == 1 + result = capsys.readouterr() + assert not result.out and "injected preparation failure" in result.err + spawn.assert_not_called() From 146c2567e13444cb02d52feff5c4eccd737e9ba9 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 08:18:50 +0800 Subject: [PATCH 28/29] Release private socket ownership after successful publication --- docs/operations.md | 6 +- src/infiniloader/daemon.py | 32 ++++++---- tests/test_runtime_publication.py | 102 ++++++++++++++++++++++++++++++ 3 files changed, 125 insertions(+), 15 deletions(-) diff --git a/docs/operations.md b/docs/operations.md index c24a487..24c5cac 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -121,8 +121,10 @@ Shutdown removes verified PID and socket resources before deleting their state. An uncertain deletion or changed ownership retains the state and emits a cleanup diagnostic, including in quiet mode. On the next start, unchanged owned stale resources can be reclaimed; unrelated replacements are preserved for inspection. -If a recorded private socket cannot be verified or removed during recovery, -startup fails before replacing the old state, so a retry retains the same evidence. +Successfully removed private names are unregistered. If a released name is reused +by another inode, recovery leaves it untouched only after independently verifying +the old public socket as owned and stale. Uncertain lookup or removal errors still +stop recovery before replacing the old state, preserving evidence for a retry. `run` protects a cache against ordinary stop and eviction. Local-cache daemons also hold a disk-generation lease. Foreground disk commands inherit its OS lock descriptor; diff --git a/src/infiniloader/daemon.py b/src/infiniloader/daemon.py index 027f208..ff47216 100644 --- a/src/infiniloader/daemon.py +++ b/src/infiniloader/daemon.py @@ -392,19 +392,23 @@ def _check_runtime_paths(self) -> None: st = None if st is not None: expected = {"device": st.st_dev, "inode": st.st_ino} - if (not stat.S_ISSOCK(st.st_mode) or st.st_uid != os.getuid() - or expected != state.get("runtime_socket")): - raise RuntimeError(f"private socket has no matching owned state: {path}") - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: - probe.settimeout(.2) - try: - probe.connect(str(path)) - except ConnectionRefusedError: - self._remove_stale_file(path, st) - except OSError as exc: - raise RuntimeError(f"cannot establish that private socket is stale: {path}") from exc - else: - raise RuntimeError(f"private socket still has a listener: {path}") + # A verified stale public link proves publication completed. + # Another inode at the released private name is not ours. + reused = stale_socket is not None and expected != state.get("runtime_socket") + if not reused: + if (not stat.S_ISSOCK(st.st_mode) or st.st_uid != os.getuid() + or expected != state.get("runtime_socket")): + raise RuntimeError(f"private socket has no matching owned state: {path}") + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as probe: + probe.settimeout(.2) + try: + probe.connect(str(path)) + except ConnectionRefusedError: + self._remove_stale_file(path, st) + except OSError as exc: + raise RuntimeError(f"cannot establish that private socket is stale: {path}") from exc + else: + raise RuntimeError(f"private socket still has a listener: {path}") if stale_pid is not None: self._remove_stale_file(self.pid_path, stale_pid) if stale_socket is not None: @@ -448,6 +452,8 @@ def _publish_socket(self, server: socket.socket) -> None: self._claimed_runtime[self.socket_path] = (st.st_dev, st.st_ino) os.link(private, self.socket_path, follow_symlinks=False) private.unlink() + self._claimed_runtime.pop(private) + self._private_socket = None def _install_signal_handlers(self) -> None: signal.signal(signal.SIGTERM, self._handle_signal) diff --git a/tests/test_runtime_publication.py b/tests/test_runtime_publication.py index 3699c5b..ada831f 100644 --- a/tests/test_runtime_publication.py +++ b/tests/test_runtime_publication.py @@ -1,4 +1,5 @@ import json +import errno import os from pathlib import Path import socket @@ -138,9 +139,110 @@ def inspect(state): try: with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as server: daemon._publish_socket(server) + assert daemon._private_socket is None + assert set(daemon._claimed_runtime) == {daemon.socket_path, daemon.state_path} with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: client.connect(str(daemon.socket_path)) accepted, _ = server.accept() accepted.close() finally: daemon._cleanup_socket() + + +def test_failed_private_unlink_keeps_claim_and_recovery_state(runtime, monkeypatch): + model, base = runtime + run = base / "run" + run.mkdir(parents=True) + daemon = WeightLoaderDaemon(name="unit", source=model, view=base / "view", socket_path=run / "unit.sock", + pid_path=run / "unit.pid", state_path=run / "unit.state.json", upstream_source=None, + local_cache_path=None, include_patterns=(), exclude_patterns=(), replace=True, + force=False, max_memory_fraction=.9, chunk_size=1024, workers=1) + real_unlink = Path.unlink + + def fail(path, *args, **kwargs): + if path == daemon._private_socket: + raise OSError(errno.EIO, "injected publication unlink failure") + return real_unlink(path, *args, **kwargs) + + daemon._owns_runtime = True + try: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as server, monkeypatch.context() as patch: + patch.setattr(Path, "unlink", fail) + with pytest.raises(OSError, match="publication unlink failure"): + daemon._publish_socket(server) + private = daemon._private_socket + assert private is not None and private.exists() + assert private in daemon._claimed_runtime + state = json.loads(daemon.state_path.read_text()) + assert state["runtime_private_socket"] == str(private) + finally: + daemon._cleanup_socket() + + +@pytest.mark.parametrize("crash_after_release", [False, True]) +def test_released_private_name_can_be_reused_by_another_daemon(runtime, monkeypatch, crash_after_release): + model, base = runtime + sockets = model.parent / "sockets" + sockets.mkdir() + public, released = sockets / "a", sockets / "f" + neighbor_base = model.parent / "neighbor-state" + command = cli._daemon_command + launched = False + + def controlled_first_name(args, *positional, **kwargs): + nonlocal launched + result = command(args, *positional, **kwargs) + if args.name != "unit": + return result + crash = crash_after_release and not launched + launched = True + code = f'''from pathlib import Path +from types import SimpleNamespace +import os, uuid +from infiniloader.cli import main +original_uuid = uuid.uuid4 +first = True +def choose_uuid(): + global first + if first: + first = False + return SimpleNamespace(hex='f' * 32) + return original_uuid() +uuid.uuid4 = choose_uuid +if {crash!r}: + original_unlink = Path.unlink + def exit_after_release(path, *args, **kwargs): + result = original_unlink(path, *args, **kwargs) + if str(path) == {str(released)!r}: + os._exit(77) + return result + Path.unlink = exit_after_release +raise SystemExit(main()) +''' + return [sys.executable, "-c", code, *result[3:]] + + monkeypatch.setattr(cli, "_daemon_command", controlled_first_name) + start_args = ["start", str(model), "--name", "unit", "--base-dir", str(base), "--socket", str(public), "--wait-timeout", "10"] + stop_args = ["stop", "unit", "--base-dir", str(base), "--socket", str(public), "--remove-view"] + neighbor_args = ["--name", "neighbor", "--base-dir", str(neighbor_base), "--socket", str(released)] + try: + assert cli.main(start_args) == (1 if crash_after_release else 0) + assert not released.exists() + state_path = base / "run/unit.state.json" + if not crash_after_release: + assert json.loads(state_path.read_text())["runtime_private_socket"] is None + assert cli.main(["start", str(model), *neighbor_args]) == 0 + before = request(released, {"command": "status"})["status"] + inode = released.stat().st_ino + if not crash_after_release: + assert cli.main(stop_args) == 0 + assert not state_path.exists() + assert cli.main(start_args) == 0 + assert json.loads(state_path.read_text())["runtime_private_socket"] is None + after = request(released, {"command": "status"})["status"] + assert after["pid"] == before["pid"] and after["state"] == "ready" + assert released.stat().st_ino == inode + assert (Path(after["view"]) / "model.safetensors").read_bytes() == b"0123456789abcdef" + finally: + cli.main(["stop", *neighbor_args, "--remove-view"]) + cli.main(stop_args) From 8aabad2bb5a1e56ea2658877597761c4ba641fb6 Mon Sep 17 00:00:00 2001 From: whjthu Date: Tue, 27 Oct 2026 08:18:56 +0800 Subject: [PATCH 29/29] Isolate default signal tests from inherited dispositions --- tests/test_runner.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/test_runner.py b/tests/test_runner.py index 0bf7db2..8f27515 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -17,8 +17,22 @@ pytestmark = pytest.mark.skipif(not hasattr(os, "memfd_create"), reason="Linux process and lease tests") +@pytest.fixture +def default_signals(): + handlers = {signal.SIGTERM: signal.SIG_DFL, signal.SIGHUP: signal.SIG_DFL, + signal.SIGINT: signal.default_int_handler} + previous = {signum: signal.getsignal(signum) for signum in handlers} + try: + for signum, handler in handlers.items(): + signal.signal(signum, handler) + yield + finally: + for signum, handler in previous.items(): + signal.signal(signum, handler) + + @pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGHUP]) -def test_child_signal_exit_uses_shell_status(tmp_path, signum): +def test_child_signal_exit_uses_shell_status(tmp_path, signum, default_signals): assert run_program([sys.executable, "-c", f"import os; os.kill(os.getpid(), {int(signum)})"], tmp_path) == 128 + signum @@ -27,7 +41,7 @@ def test_child_signal_exit_uses_shell_status(tmp_path, signum): (signal.SIGINT, False, 130), (signal.SIGTERM, True, 137), (signal.SIGKILL, False, -signal.SIGKILL), ]) -def test_wrapper_cancellation_and_surviving_child_lease(tmp_path, signum, ignore, expected): +def test_wrapper_cancellation_and_surviving_child_lease(tmp_path, signum, ignore, expected, default_signals): cache = materialize_model_cache(make_model_dir(tmp_path / "model"), base_dir=tmp_path / "state") marker = tmp_path / "child.pid" code = (