diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..5e19ae6 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,37 @@ +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 + 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 + - run: python examples/hybrid_moe.py --device cpu --output hybrid-cpu.json 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 c8ad4f9..6f460ec 100644 --- a/README.md +++ b/README.md @@ -1,282 +1,100 @@ -# InfiniLoader +# InfiniLoader -[中文](README.zh-CN.md) +[中文说明](README.zh-CN.md) -`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. +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. -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. +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. -## Quick Start +## Install and check the environment -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. +Linux and Python 3.10+ are required. Install inside the engine's environment, using +the same user and PID namespace as the consumer: ```bash -cd InfiniLoader -source env.sh - -MODEL=/path/to/huggingface-model -MODEL_VIEW=$(iloader start "$MODEL" --name dev-model) -``` - -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 -``` - -```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) -``` - -Use the view path with any framework command that accepts a model directory: +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. -```bash -vllm serve "$MODEL_VIEW" --tensor-parallel-size 1 -``` +## Cache once and reuse ```bash -python -m sglang.launch_server \ - --model-path "$MODEL_VIEW" \ - --tp 1 -``` +MODEL=/absolute/path/to/model +iloader start "$MODEL" --name dev +iloader status dev -```bash -python python/infinilm/server/inference_server.py \ - --device cuda \ - --model="$MODEL_VIEW" \ - --tp=1 -``` - -From another terminal, you can inspect and manage active resident caches: +# Example when vLLM is already installed; validate your engine version. +iloader run dev -- vllm serve '{model}' --tensor-parallel-size 1 -```bash -iloader list -iloader status --name dev-model -iloader stop dev-model --remove-view +# 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 ``` -Stopping a cache that is still `starting` or `loading` cancels the preload. +`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. -### Keep a Model Warm While Restarting a Framework +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. -```bash -MODEL_VIEW=$(iloader start "$MODEL" --name dev-model) +## Disk caching and memory limits -# Repeat this while editing framework code. -vllm serve "$MODEL_VIEW" --tensor-parallel-size 1 -``` - -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. +For a reusable local disk generation without another model-sized RAM allocation: ```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 +LOCAL=$(iloader cache "$MODEL" --name dev) +iloader verify "$LOCAL" +iloader run --cache-path "$LOCAL" -- vllm serve '{model}' ``` -### 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. +To preload into RAM from the local generation, use `start --local-cache`. +The disk-cache root must be outside the source model directory. ```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" +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 ``` -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 - -```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 -``` - -## 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. +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. -| 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 | +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. -The most stable production-like gain comes from avoiding repeated NFS reads via -local cache plus memfd views. +## Documentation and optional APIs -## Compatibility Notes +- [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) -- 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. +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. -## Limitations +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. -- 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. -- Physical rank-sharded checkpoint conversion is not included in this cleaned - release; `dist-plan` only analyzes candidate tensor splits. +[MIT license](LICENSE). diff --git a/README.zh-CN.md b/README.zh-CN.md index b24eebb..9022d8f 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -1,244 +1,75 @@ -# InfiniLoader 中文说明 +# InfiniLoader -[English README](README.md) +[English](README.md) -`InfiniLoader` 是一个面向大模型推理框架开发调试的用户态权重加载管理器。它解决的核心问题是：调试 vLLM、SGLang、InfiniLM 或自研推理框架时，经常需要反复停止、修改代码、重新启动框架，而每次启动都要重新从存储读取模型权重。 +InfiniLoader 把模型文件保存在长期运行的缓存进程中，让推理引擎反复重启时复用权重，减少对原始存储的重复读取。它是 Linux 用户态缓存，负责权重文件及其生命周期；模型计算和请求调度由推理引擎承担。 -默认模式下，InfiniLoader 会把模型目录常驻在内存里，并暴露出一个看起来像普通 Hugging Face 模型目录的路径。框架仍然使用原来的文件读取逻辑，通常只需要把模型路径换成 `iloader start/view` 输出的路径。 +核心仅依赖 Python 标准库。选中的文件被复制到封存的 memfd，并通过指向 `/proc//fd/` 的目录链接供现有文件加载器读取。 -## 快速开始 +## 安装与环境检查 -最基础路径是 memfd 支撑的 memfs 风格内存 view。如果你原本会把框架直接指向 -`/path/to/huggingface-model`，现在只需要在启动框架前多执行一行命令，然后把模型路径换成它输出的 view。 +需要 Linux 和 Python 3.10+。建议安装在引擎所在环境，与消费者使用相同用户和 PID namespace： ```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 -``` - -```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 -``` - -运行状态会优先放在 `$ILOADER_BASE_DIR`；未设置时使用 -`$XDG_CACHE_HOME/infiniloader` 或 `~/.cache/infiniloader`。所有会创建运行状态的 -命令也都支持通过 `--base-dir` 显式指定目录。 - -## 命令边界 +开发时可以使用 `source env.sh`，直接加载源码并启用 Bash 补全。运行目录应位于节点本地，并支持文件锁、Unix socket 和硬链接；`doctor` 会检查这些能力。 -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 -``` +MODEL=/absolute/path/to/model +iloader start "$MODEL" --name dev +iloader status dev -如果 cache 仍处于 `starting` 或 `loading`，`stop` 会取消这次 preload。 +# 示例：已安装 vLLM 时使用；仍需验证实际引擎版本。 +iloader run dev -- vllm serve '{model}' --tensor-parallel-size 1 -### 框架反复重启，权重保持 warm - -```bash -MODEL_VIEW=$(iloader start "$MODEL" --name dev-model) - -# 修改框架代码后反复运行这一条。 -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 退出，超时则返回错误。 -当原始模型目录位于远端或共享存储，而你希望先构建一个可复用的本地缓存时，建议开启 -`--local-cache`： +文件接口是集成边界。引擎格式、单文件路径解析、容器和加载器行为都可能影响兼容性，必须验证实际版本。目前 CI 覆盖缓存服务与参考示例，没有覆盖按版本维护的 vLLM/SGLang/InfiniLM 兼容矩阵。 -```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 不变。 - -```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" -``` - -分析哪些 tensor 适合未来做框架原生的 rank-sharded checkpoint： +仅使用本地磁盘代际，不额外分配整份模型大小的 RAM： ```bash -iloader dist-plan "$MODEL" --tp 4 --limit 20 +LOCAL=$(iloader cache "$MODEL" --name dev) +iloader verify "$LOCAL" +iloader run --cache-path "$LOCAL" -- vllm serve '{model}' ``` -## 命令摘要 +需要从本地代际继续预加载到 RAM 时，使用 `start --local-cache`。磁盘缓存根目录必须位于源模型目录之外。 ```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 状态。 +上述容量只是示例，应按模型和引擎实际占用设置。缓存计费包括文件内容与加载缓冲，引擎分配和客户端仍持有的映射/COW 副本另计。memfd 已封存但未 mlock，内存压力和 swap 仍会影响延迟。 -| 场景 | 基线 | 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 | +源模型应使用不可变版本。源变化后需更换缓存名或先停止再重建。直接使用 `iloader view dev` 路径的读者没有自动租约，应保持自动 idle eviction 关闭。 -最稳定的收益来自本地 cache 加 memfd view，尤其适合绕过重复 NFS 读取。 +## 文档与可选接口 -## 兼容性说明 +- [操作、校验、回收与容器](docs/operations.md) +- [贡献与测试](CONTRIBUTING.md) +- [可选分片、TensorStore 和专家示例](docs/distributed.md) +- [围绕缓存目标的开发范围](docs/roadmap.zh-CN.md) +- [历史测量与维护资料](docs/archive/README.md) -- vLLM、SGLang、InfiniLM 都可以通过稳定的 `start/view` 路径使用，通常只需要改模型路径。 -- 新版本 vLLM 和 SGLang 具备原生 `sharded_state` loader，是未来支持物理 rank-sharded checkpoint 的优先方向。 -- InfiniLM 当前最适合完整 view 路径；rank-sharded checkpoint 需要 InfiniLM adapter 或原生 loader 支持。 +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，仅测文件读取，不能等同于引擎启动收益。 -- 需要 Linux；核心能力依赖 `memfd_create` 和 `/proc//fd`。 -- 框架读取 view 时，InfiniLoader daemon 必须保持运行。 -- 默认 view 会在内存中保留一份选中的模型文件内容。 -- 0.1.0 不包含完全透明的同路径拦截；需要把输出的 view 路径传给推理引擎。 -- 当前清理版不包含物理 rank-sharded checkpoint 转换；`dist-plan` 只做候选 tensor 切分分析。 +[MIT 许可](LICENSE)。 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/archive/2026-10-18-distributed-results.json b/docs/archive/2026-10-18-distributed-results.json new file mode 100644 index 0000000..64a9fbc --- /dev/null +++ b/docs/archive/2026-10-18-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/archive/2026-10-18-distributed-validation.md b/docs/archive/2026-10-18-distributed-validation.md new file mode 100644 index 0000000..e2b374d --- /dev/null +++ b/docs/archive/2026-10-18-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](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 +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](2026-10-18-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/archive/2026-10-18-expert-review.md b/docs/archive/2026-10-18-expert-review.md new file mode 100644 index 0000000..a163eb9 --- /dev/null +++ b/docs/archive/2026-10-18-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/archive/2026-10-18-reliability-validation.md b/docs/archive/2026-10-18-reliability-validation.md new file mode 100644 index 0000000..6a386a5 --- /dev/null +++ b/docs/archive/2026-10-18-reliability-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](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. + +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. diff --git a/docs/archive/2026-10-26-history-reorganization.json b/docs/archive/2026-10-26-history-reorganization.json new file mode 100644 index 0000000..e8f2661 --- /dev/null +++ b/docs/archive/2026-10-26-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/archive/2026-10-26-history-reorganization.md b/docs/archive/2026-10-26-history-reorganization.md new file mode 100644 index 0000000..dc23ee7 --- /dev/null +++ b/docs/archive/2026-10-26-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](2026-10-26-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](2026-10-18-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](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 new file mode 100644 index 0000000..4036db9 --- /dev/null +++ b/docs/distributed.md @@ -0,0 +1,217 @@ +# 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 +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. 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. + +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 +{ + "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 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 +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](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 new file mode 100644 index 0000000..24c5cac --- /dev/null +++ b/docs/operations.md @@ -0,0 +1,204 @@ +# Operating the cache service + +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 + +```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 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` 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 +until the child is reaped. Backgrounded descendants are outside this contract. +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: + +```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 + +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 +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 +``` + +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 +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. +- 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`. + +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. + +`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 +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. +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; +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, 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. + +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. 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: + +```bash +docker build -f deploy/Dockerfile -t infiniloader:dev . +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. + +`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. 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 new file mode 100644 index 0000000..ebc367f --- /dev/null +++ b/docs/roadmap.zh-CN.md @@ -0,0 +1,19 @@ +# 开发范围与下一步 + +InfiniLoader 的目标是：模型文件加载一次，让后续进程快速、可靠地复用。核心提供文件缓存、内容与代际检查、使用期租约、内存预算和安全回收；模型执行与请求调度交给消费者。 + +当前优先验收以下内容，不以增加功能数量为目标： + +1. 用日常调试的一个真实模型和框架，通过已有 `start/run` 路径验证结果一致、重复加载耗时与内存占用。 +2. 连续重复加载、退出、停止与复用，检查 FD、内存及残留进程是否持续增长，并验证源变化与失败恢复。 +3. 保持安装、最小使用流程、支持环境与故障信息清楚可复现。公开性能结果应带模型 revision、引擎版本、命令、样本和测试范围。 + +后续优化应由这些测量发现的瓶颈驱动，例如重复文件扫描、校验或复制。先记录实际成本，再决定是否需要新的缓存策略或接口。 + +已有实现的成本需要单独量化并逐步收敛：磁盘构建目前按 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)，用于说明原始评估，不代表当前承诺。 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/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/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/arguments.py b/src/infiniloader/arguments.py new file mode 100644 index 0000000..20388ae --- /dev/null +++ b/src/infiniloader/arguments.py @@ -0,0 +1,226 @@ +"""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, + 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) + 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_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"]) + + 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/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/cache.py b/src/infiniloader/cache.py index 3dd6206..a2e2d84 100644 --- a/src/infiniloader/cache.py +++ b/src/infiniloader/cache.py @@ -7,10 +7,14 @@ import shutil import stat import time +import tempfile +import threading 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 +from .memory import memory_info, check_memory DEFAULT_EXCLUDES = ( @@ -22,6 +26,8 @@ "*/.cache/*", ".iloader-cache/*", "*/.iloader-cache/*", + ".iloader-dist/*", + "*/.iloader-dist/*", ".iloader/*", "*/.iloader/*", "*.lock", @@ -39,6 +45,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 +65,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 @@ -112,14 +122,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: @@ -148,16 +151,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 +186,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 +222,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 +231,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 +239,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 +269,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: + except BaseException: try: - os.close(dst_fd) + if dst_fd is not None: + os.close(dst_fd) finally: raise finally: @@ -270,7 +309,11 @@ 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]") 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) @@ -279,38 +322,42 @@ 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) - 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) + 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, @@ -320,58 +367,41 @@ 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 -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 +421,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..5ca72f1 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 @@ -11,14 +10,18 @@ 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 +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, request -from .local_cache import cache_status, list_model_caches, materialize_model_cache +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, default_log_path, default_pid_path, default_socket_path, @@ -39,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: @@ -106,14 +109,13 @@ 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 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: @@ -234,6 +236,11 @@ 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, + 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() @@ -274,80 +281,114 @@ 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)]) + 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))]) + if args.idle_ttl is not None: + command.extend(["--idle-ttl", str(args.idle_ttl)]) return command 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) + + +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") + 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: + raise RuntimeError(f"stop the running cache before changing --{option.replace('_', '-')}") + + +def _start_locked(args: argparse.Namespace) -> int: sock = _socket(args) - try: - status = request(sock, {"command": "status"}, timeout=1.0)["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 - - deadline = time.time() + 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 + deadline = time.monotonic() + args.wait_timeout + while True: try: status = request(sock, {"command": "status"}, timeout=1.0)["status"] - last_progress = _print_progress_update(args, last_progress) - if sys.stderr.isatty() and last_progress: - print(file=sys.stderr) + 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 - except Exception: + break + proc = None + if starting: + _require_identity(args, starting) + if args.no_wait: + 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") + 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, + ) + deadline = time.monotonic() + args.wait_timeout + print(f"starting name={args.name} pid={proc.pid} log={log_path}", file=sys.stderr) + + last_progress = None + 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) - 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) + request(sock, {"command": "touch"}) + 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: @@ -392,6 +433,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 @@ -414,7 +457,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) @@ -465,20 +508,39 @@ 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}, timeout=args.timeout) - print(f"{args.name}: {'stopping' if response.get('stopping') else 'ok'}") + response = request(_socket(args), {"command": "stop", "remove_view": args.remove_view, + "force": getattr(args, "force", False)}, timeout=args.timeout) + 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 - os.kill(int(starting["pid"]), signal.SIGTERM) + 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: - shutil.rmtree(_view(args), ignore_errors=True) + 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 @@ -504,8 +566,13 @@ def cmd_stop(args: argparse.Namespace) -> int: view=None, timeout=args.timeout, remove_view=args.remove_view, + force=getattr(args, "force", False), ) - 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) @@ -580,83 +647,105 @@ def cmd_cache_list(args: argparse.Namespace) -> int: return 0 -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" +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: + 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: + 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 + - 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 --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 --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" - ;; - dist-views) - opts="--name --ranks --base-dir --replace --json --help" - ;; - dist-simulate) - opts="--ranks --rank-view-root --pattern --exclude --chunk-size --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 _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 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}") @@ -690,6 +779,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) @@ -753,6 +843,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 @@ -789,181 +880,48 @@ 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) - - 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("--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) - - 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) - - 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("--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.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: 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 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/completion.py b/src/infiniloader/completion.py new file mode 100644 index 0000000..13ce559 --- /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/src/infiniloader/daemon.py b/src/infiniloader/daemon.py index 13c9e6b..ff47216 100644 --- a/src/infiniloader/daemon.py +++ b/src/infiniloader/daemon.py @@ -5,11 +5,23 @@ import signal import socket import time -import traceback +from contextlib import ExitStack +from concurrent.futures import ThreadPoolExecutor +import struct +import threading +import uuid +import stat +import errno +import sys 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, collect_source_files +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 +from .ipc import LoaderStoppingError class WeightLoaderDaemon: @@ -32,15 +44,20 @@ def __init__( chunk_size: int, 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, + 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 @@ -49,12 +66,28 @@ 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.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._private_socket: Path | None = None def log(self, message: str) -> None: if not self.quiet: @@ -62,17 +95,24 @@ 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("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) 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 @@ -109,13 +149,15 @@ 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) 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, @@ -134,7 +176,10 @@ 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, "pid": os.getpid(), "source": str(self.source), "upstream_source": str(self.upstream_source) if self.upstream_source else None, @@ -150,10 +195,31 @@ 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, + "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 LoaderStoppingError("daemon is stopping; retry after it exits") if command == "ping": return {"ok": True, "pong": True} if command == "status": @@ -161,28 +227,233 @@ 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} + 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 _cleanup_socket(self) -> None: - try: - if self.socket_path.exists() or self.socket_path.is_socket(): - self.socket_path.unlink() - except OSError: - pass + def _serve_connection(self, conn: socket.socket) -> None: try: - if self.pid_path.exists(): - self.pid_path.unlink() + 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)} + if isinstance(exc, LoaderStoppingError): + response["code"] = "stopping" + conn.settimeout(1.0) + conn.sendall(json.dumps(response).encode("utf-8") + b"\n") except OSError: - pass + 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 + complete = True + for path, identity in self._claimed_runtime.items(): + 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: - if self.state_path.exists(): - self.state_path.unlink() - except OSError: - pass + 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() + 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") + 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(): + stale_pid = self.pid_path.lstat() + 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}") + 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 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 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} + # 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: + 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() + self._claimed_runtime.pop(private) + self._private_socket = None def _install_signal_handlers(self) -> None: signal.signal(signal.SIGTERM, self._handle_signal) @@ -196,58 +467,84 @@ 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)) + 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") + 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() + 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 + clients = ThreadPoolExecutor(max_workers=4, thread_name_prefix="iloader-ipc") 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") server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) server.settimeout(1.0) - server.bind(str(self.socket_path)) - 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}") 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: 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() 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/diagnostics.py b/src/infiniloader/diagnostics.py new file mode 100644 index 0000000..63470eb --- /dev/null +++ b/src/infiniloader/diagnostics.py @@ -0,0 +1,64 @@ +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") + 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 + + 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/src/infiniloader/distributed.py b/src/infiniloader/distributed.py index 494997a..5fe157a 100644 --- a/src/infiniloader/distributed.py +++ b/src/infiniloader/distributed.py @@ -4,33 +4,17 @@ import multiprocessing as mp import os 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 - - -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 .runtime import validate_name, read_json +from .formats import DTYPE_BYTES, tensor_header, tensor_inventory @dataclass @@ -52,33 +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(" 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]: @@ -149,6 +114,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 +146,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 +168,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}" @@ -231,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), @@ -247,11 +213,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 +230,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 +241,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 +257,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 +277,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/experts.py b/src/infiniloader/experts.py new file mode 100644 index 0000000..ab96ef6 --- /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: dict = field(default_factory=dict) + + +@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.values(): + 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) + # 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) + 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/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/ipc.py b/src/infiniloader/ipc.py index f1684be..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: @@ -23,6 +27,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 +36,14 @@ 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): + 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/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/local_cache.py b/src/infiniloader/local_cache.py index 71fcaf3..84c8162 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 .paths import absolute_path, default_base_dir +from .runtime import FileLock, atomic_json, read_json @dataclass @@ -24,40 +29,27 @@ 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, *, 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 / cache_key(source_root, name) + return root def _manifest_path(path: Path) -> Path: @@ -65,236 +57,272 @@ 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, 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): + 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, source_root=source_root) + 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") + 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) + 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/src/infiniloader/memory.py b/src/infiniloader/memory.py new file mode 100644 index 0000000..ec4444d --- /dev/null +++ b/src/infiniloader/memory.py @@ -0,0 +1,152 @@ +"""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 = 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, + 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() + 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/src/infiniloader/paths.py b/src/infiniloader/paths.py index dda2340..edae073 100644 --- a/src/infiniloader/paths.py +++ b/src/infiniloader/paths.py @@ -1,19 +1,24 @@ from __future__ import annotations import os +import tempfile from pathlib import Path +from .runtime import validate_name 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) - 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 @@ -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,66 +35,71 @@ 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")) / APP_DIR_NAME - fallback.mkdir(parents=True, exist_ok=True) + 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 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/runner.py b/src/infiniloader/runner.py new file mode 100644 index 0000000..54022d8 --- /dev/null +++ b/src/infiniloader/runner.py @@ -0,0 +1,65 @@ +"""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): + 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], + 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/src/infiniloader/runtime.py b/src/infiniloader/runtime.py new file mode 100644 index 0000000..816c332 --- /dev/null +++ b/src/infiniloader/runtime.py @@ -0,0 +1,112 @@ +"""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 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")) + 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/src/infiniloader/shards.py b/src/infiniloader/shards.py new file mode 100644 index 0000000..8b33055 --- /dev/null +++ b/src/infiniloader/shards.py @@ -0,0 +1,263 @@ +"""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 .formats import DTYPE_BYTES, MAX_HEADER, _signature, canonical, read_json, tensor_header, tensor_inventory +from .layouts import validate_layout + +CHUNK = 1024 * 1024 + + +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 _int(value, *, minimum=0): + return type(value) is int and value >= minimum + + +def catalog(source: Path) -> dict: + """Content identity is independent of hostname, inode and source directory.""" + source = source.expanduser().resolve() + paths = sorted(source.rglob("*.safetensors")) + signatures = {p: _signature(p) for p in paths} + index_path = source / "model.safetensors.index.json" + 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) + 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) + 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) + 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 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/src/infiniloader/torch_store.py b/src/infiniloader/torch_store.py new file mode 100644 index 0000000..15cc8cc --- /dev/null +++ b/src/infiniloader/torch_store.py @@ -0,0 +1,111 @@ +"""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 +from .layouts import validate_layout + + +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") + 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"]} + 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] + stage, dim = rule["stage"], rule["split_dim"] + 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_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_cli.py b/tests/test_cli.py index 78e6015..38ee815 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2,9 +2,11 @@ import json import os +import pytest 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 +38,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 +63,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() @@ -70,11 +74,19 @@ 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: 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") + (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() @@ -84,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_completion.py b/tests/test_completion.py new file mode 100644 index 0000000..8204508 --- /dev/null +++ b/tests/test_completion.py @@ -0,0 +1,65 @@ +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", "--"]) + 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") +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" 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_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) 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 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"] 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(" 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("