From 73b623f2f7db092053c1c86fe796bed89eb3dc71 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Fri, 21 Aug 2026 08:50:59 +0800 Subject: [PATCH] [Core] Add Host Weight Runtime foundation (#6419) Signed-off-by: Hongsheng Liu --- docs/.nav.yml | 2 + docs/design/feature/host_weight_runtime.md | 302 +++++ docs/design/index.md | 14 +- docs/design/module/host_weight_runtime.md | 264 ++++ tests/host_weight_runtime/test_contracts.py | 232 ++++ .../test_filesystem_store.py | 1106 +++++++++++++++ tests/host_weight_runtime/test_runtime.py | 364 +++++ vllm_omni/host_weight_runtime/__init__.py | 139 ++ vllm_omni/host_weight_runtime/config.py | 183 +++ vllm_omni/host_weight_runtime/errors.py | 106 ++ .../filesystem/__init__.py | 11 + .../host_weight_runtime/filesystem/locks.py | 115 ++ .../host_weight_runtime/filesystem/store.py | 1189 +++++++++++++++++ .../host_weight_runtime/filesystem/writer.py | 425 ++++++ vllm_omni/host_weight_runtime/identity.py | 407 ++++++ vllm_omni/host_weight_runtime/inspection.py | 80 ++ vllm_omni/host_weight_runtime/lease.py | 143 ++ vllm_omni/host_weight_runtime/manifest.py | 353 +++++ vllm_omni/host_weight_runtime/outcomes.py | 119 ++ vllm_omni/host_weight_runtime/protocols.py | 228 ++++ vllm_omni/host_weight_runtime/runtime.py | 332 +++++ 21 files changed, 6108 insertions(+), 6 deletions(-) create mode 100644 docs/design/feature/host_weight_runtime.md create mode 100644 docs/design/module/host_weight_runtime.md create mode 100644 tests/host_weight_runtime/test_contracts.py create mode 100644 tests/host_weight_runtime/test_filesystem_store.py create mode 100644 tests/host_weight_runtime/test_runtime.py create mode 100644 vllm_omni/host_weight_runtime/__init__.py create mode 100644 vllm_omni/host_weight_runtime/config.py create mode 100644 vllm_omni/host_weight_runtime/errors.py create mode 100644 vllm_omni/host_weight_runtime/filesystem/__init__.py create mode 100644 vllm_omni/host_weight_runtime/filesystem/locks.py create mode 100644 vllm_omni/host_weight_runtime/filesystem/store.py create mode 100644 vllm_omni/host_weight_runtime/filesystem/writer.py create mode 100644 vllm_omni/host_weight_runtime/identity.py create mode 100644 vllm_omni/host_weight_runtime/inspection.py create mode 100644 vllm_omni/host_weight_runtime/lease.py create mode 100644 vllm_omni/host_weight_runtime/manifest.py create mode 100644 vllm_omni/host_weight_runtime/outcomes.py create mode 100644 vllm_omni/host_weight_runtime/protocols.py create mode 100644 vllm_omni/host_weight_runtime/runtime.py diff --git a/docs/.nav.yml b/docs/.nav.yml index a2c93d534..ded87fe9c 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -113,6 +113,7 @@ nav: - Feature Design: - Runtime and Stage Execution: - Disaggregated Inference: design/feature/disaggregated_inference.md + - Host Weight Runtime: design/feature/host_weight_runtime.md - Async Chunk: design/feature/async_chunk.md - Async Diffusion Output: design/feature/async_diffusion_output.md - Async Omni Output Materialization: design/feature/omni_async_output_materialization.md @@ -169,6 +170,7 @@ nav: - Offloader: design/module/diffusion/offloader.md - Execution Platforms: design/module/execution_platforms.md - Cache Management: design/module/cache_management.md + - Host Weight Runtime: design/module/host_weight_runtime.md - Quantization: design/module/quantization.md - Observability: design/module/observability.md - Profiling: design/module/profiling.md diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md new file mode 100644 index 000000000..d26aa6202 --- /dev/null +++ b/docs/design/feature/host_weight_runtime.md @@ -0,0 +1,302 @@ +# Host Weight Runtime + +The Host Weight Runtime is a loader-facing feature for reusing immutable, +runtime-ready model weights across workers on one host. It lets a model +integration replace repeated checkpoint loading and transformation with an +exact lookup of a previously published host representation, while keeping the +canonical model source authoritative. + +This feature is shared infrastructure. Diffusion, autoregressive, speech, and +future model integrations may depend on the top-level +`vllm_omni.host_weight_runtime` package without depending on one another's +loader or execution stack. The internal component and filesystem contracts are +specified in the [Host Weight Runtime module design](../module/host_weight_runtime.md). + +## Status + +The first implementation provides contracts and a CPU local-filesystem store. +It does not enable cached loading for a model by itself. A consumer must still +provide an exact identity, a representation producer, and a model restorer. + +V1 includes: + +- exact runtime-weight identity and immutable manifests; +- coordinated local lookup and one-producer publication; +- descriptor-backed safetensors mmap leases; +- preferred and required resolution policy; +- validation, deny, quarantine, cleanup, and capacity controls; and +- typed reports for every terminal resolution outcome. + +V1 does not include: + +- a public CLI or default-on model integration; +- a generic BF16, FP8, quantized, or model-specific producer; +- CUDA registration, pinned staging, H2D scheduling, or GPU kernels; +- a remote artifact provider or cross-node coordination; +- post-load publication (V1 accepts only pre-load-safe producers); +- automatic eviction; or +- a change to DLO AllGather or no-AllGather behavior. + +## Motivation and use cases + +Model loading can create the same final host representation repeatedly. This +is especially expensive when loading performs checkpoint decoding, tensor +renaming, TP slicing, quantization, packing, or scale construction before GPU +transfer. Independent workers may then retain private CPU copies of identical +runtime weights. + +The Host Weight Runtime is useful when: + +- multiple same-node workers request semantically identical final weights; +- constructing those weights is materially more expensive than validating and + mapping a local artifact; +- the runtime layout can be identified exactly and reproduced deterministically; +- the workers can share immutable file-backed pages through the OS page cache; + and +- canonical loading remains available when policy permits fallback. + +Typical consumers include same-node diffusion replicas, DLO host-weight +sources, and future transformed or quantized model loaders. Request routing and +replica orchestration remain outside this feature. + +This feature is not zero-copy GPU execution. A lease exposes stable CPU tensor +views and mapped ranges. A separate transport still decides whether to use +registered mmap, private pinned staging, synchronous copies, or asynchronous +H2D transfer. + +## Resolution behavior + +The loader resolves the immutable canonical source and computes the exact +requested identity before asking the runtime for a representation. + +```mermaid +flowchart TD + L["Loader resolves canonical source"] --> I["Compute exact WeightArtifactIdentity"] + I --> M{"Runtime mode"} + M -->|"disabled"| C["Canonical loader"] + M -->|"preferred or required"| A{"Validated local artifact?"} + A -->|"yes"| R["Return HostWeightLease"] + A -->|"no"| X["Remote exact artifact (future)"] + X --> P{"Registered producer available?"} + P -->|"yes"| B["Build, validate, and publish atomically"] + B --> R + P -->|"no or recoverable failure"| F{"Policy permits fallback?"} + F -->|"preferred"| C + F -->|"required"| E["Fail startup"] + P -->|"nonretryable failure"| E +``` + +The modes are: + +| Mode | Observable behavior | +| --- | --- | +| `disabled` | Use canonical loading directly without probing storage, identity, credentials, or topology. | +| `preferred` | Return an exact lease when available; otherwise use canonical fallback only for a miss, invalid cache entry, or typed retryable failure. | +| `required` | Fail when an exact lease cannot be acquired. | + +An unsupported capability, semantic identity collision, producer failure, or +publication failure is nonretryable and remains visible even in preferred +mode. Storage policy must never disguise a semantic or configuration error as +a cache miss. + +## Loader and restoration sequence + +The model integration owns the end-to-end loading transaction. The runtime +does not construct a model or invoke the canonical loader. + +```mermaid +sequenceDiagram + participant L as Loader adapter + participant R as HostWeightRuntime + participant S as HostWeightStore + participant P as WeightProducer + participant X as WeightRestorer + participant T as GPU transport + + L->>L: Resolve canonical revision and exact identity + L->>R: resolve(identity, producer) + R->>S: lookup(exact identity) + alt validated hit + S-->>R: HostWeightLease + else miss with allowed producer + R->>S: get_or_build(identity, producer) + S->>P: produce(store-scoped writer) + P-->>S: final-layout tensors and metadata + S-->>R: validated HostWeightLease + else policy permits canonical fallback + R-->>L: CANONICAL_FALLBACK + L->>L: Run canonical loader + else required or nonretryable failure + R-->>L: FAILED report + end + R-->>L: lease and terminal report + L->>X: plan_restore(model, lease) + X-->>L: validation-only restore plan + L->>X: commit() once + L->>T: consume restored CPU tensors + T-->>L: transfer teardown complete + L->>L: close lease +``` + +`plan_restore()` must not mutate the model or lease. `commit() -> None` is the +sole one-shot model mutation. If planning fails, canonical fallback may reuse +the untouched model. If commit begins and fails, the partially hydrated model +must be discarded and canonical fallback must construct a fresh model. + +## Exact representation identity + +The store performs exact lookup. It never silently converts one representation +into another or substitutes the first backing that responds. + +Identity includes: + +- immutable model revision and source fingerprint; +- model component and ownership boundary; +- representation name, dtype, and format metadata; +- final tensor layout and semantic parallel coordinates; +- static adaptation identity; and +- producer implementation, manifest, and restorer schema versions. + +The requested representation is selected by the loader. The store chooses +where to obtain that exact representation: validated local artifact, future +remote materialization, or a registered producer. + +### Parallelism + +Parallel coordinates are included only when they change bytes, shape, layout, +or component ownership. + +| Dimension | Identity and sharing rule | +| --- | --- | +| DP | Exclude DP rank for replicated weights so same-node replicas can acquire the same artifact. Include a coordinate only if DP changes weight ownership. | +| TP | Include TP size and rank when each rank owns different tensor slices or packed layouts. Each distinct TP shard is a separate exact artifact. | +| SP | Include SP size and backend when they change the weight layout. Exclude SP rank when all SP ranks consume identical bytes. | +| PP | Encode component ownership or stage-local layout when PP changes which weights belong to the consumer. | +| EP | Encode expert ownership and layout whenever ranks own different expert weights. | + +The runtime does not infer these rules from process-group topology. The model +adapter is responsible for constructing a correct identity. + +### Quantization and adaptations + +Runtime BF16, FP8, packed quantized weights, and future formats are distinct +representations. Conversion is performed by an explicitly versioned producer, +never by cache coercion. A producer must publish the final layout consumed by +its matching restorer. + +Dynamic LoRA overlays are not part of a reusable base-weight artifact. A +statically merged adapter is cacheable only as a separate identity containing +the adapter fingerprint and merge semantics. + +## Host sharing and GPU transport + +Every process receives its own virtual mappings and `HostWeightLease`, but +workers mapping the same local artifact can share physical file-backed pages +through the OS page cache. This avoids independent private tensor allocations; +it does not guarantee that process PSS reports exactly one model copy. + +The ownership boundary is: + +- the store owns immutable artifact files, publication, and lifecycle; +- the kernel owns page-cache residency and placement; +- the lease owns process-local mappings and the shared artifact lock; and +- transport owns page registration, page locking, private staging, H2D copies, + streams, device buffers, and lease release ordering. + +DLO AllGather and no-AllGather remain transport and execution choices. A DLO +integration may consume a lease, but the Host Weight Runtime does not choose a +parallel collective or orchestrate DP requests. See the +[DLO feature design](offloader/distributed_layerwise_offload.md). + +## Locality and NUMA policy + +The V1 backend is node-local. It accepts an allowlist of known local disk +filesystems plus tmpfs/ramfs, records the detected kernel filesystem type, and +rejects known remote or unknown filesystems. NFS, CIFS, Lustre, Ceph, and other +cross-node mounts cannot silently satisfy the local backend because their +page-cache and advisory-lock behavior does not provide the required node-local +contract. + +A NUMA-specific root creates a separate storage domain but does not by itself +place page-cache pages on that NUMA node. A topology-aware integration must +also pin workers, establish local first touch or prefault, perform transport +registration in the same domain, and collect residency evidence. + +Tmpfs artifacts consume host memory and may consume swap. They must be +accounted as memory-backed storage rather than ordinary disk capacity. + +## Concurrency and timeouts + +One process per exact identity owns a build; other workers wait and then acquire +leases for the published artifact. Publication is invisible until all payloads +and metadata are validated, hashed, fsynced, and atomically renamed. + +`coordination_timeout_seconds` bounds filesystem lock acquisition. It does not +cancel synchronous validation, a producer that has already started, or atomic +publication. A hung in-process producer therefore blocks its owning process and +must be handled by external process supervision. Enforceable producer +cancellation requires a future process-isolated producer contract. + +Leases keep a shared artifact lock for their lifetime. A forked child may close +its inherited descriptors but cannot unlock the parent process's lease. Cleanup +uses a nonblocking exclusive lock and reports an active lease instead of +removing live mappings. + +## Failure and fallback contract + +| Condition | Preferred mode | Required mode | +| --- | --- | --- | +| Exact local hit | Return lease | Return lease | +| Local miss | Try later backing or canonical fallback | Try later backing or fail | +| Invalid or denied artifact | Quarantine/rebuild when possible; otherwise canonical fallback | Quarantine/rebuild or fail | +| Retryable lock, domain, or capacity failure | Canonical fallback | Fail | +| Unsupported producer or backend | Fail | Fail | +| Identity collision | Fail | Fail | +| Producer or publication failure | Fail | Fail | +| Restore planning failure | Canonical fallback may reuse untouched model | Fail | +| Restore commit failure | Discard model; fallback requires a fresh instance | Fail | + +Cache failures should normally reduce performance rather than availability, +but only when the typed failure explicitly permits fallback. Canonical source, +authentication, configuration, and semantic errors remain visible. + +## Consumer integration requirements + +A consumer PR must: + +1. Resolve an immutable canonical model source before cache lookup. +2. Construct the exact final representation identity, including relevant + parallel, quantization, and adaptation semantics. +3. Register a deterministic producer for that identity or support lookup-only + operation. +4. Provide a validation-only restorer with a one-shot commit. +5. Keep GPU transport outside the producer, store, and restorer contracts. +6. Handle preferred fallback and required failure without reusing a model after + a failed restore commit. +7. Retain the lease until every transport operation that may access mapped + memory has completed. +8. Emit and test the terminal resolution report. + +Consumer validation must prove: + +- output parity with canonical loading; +- exact warm-hit work avoidance; +- correct TP/SP/PP/EP ownership and layout identity; +- correct BF16, FP8, quantized, or adaptation semantics; +- shared-backing evidence when host-memory savings are claimed; +- clean unmapping, unregistration, and artifact-lock release; and +- startup, host-memory, and transport effects for the claimed deployment. + +## Rollout + +The rollout is intentionally staged: + +1. Land the neutral contracts and local filesystem store. +2. Add one model-specific producer/restorer integration with parity evidence. +3. Connect eligible DLO or other transport paths without moving transport + ownership into the runtime. +4. Add additional final-layout and quantized producers independently. +5. Introduce remote materialization only through the same local lease contract. + +The architectural decisions and deferred work are tracked in +[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414). diff --git a/docs/design/index.md b/docs/design/index.md index acac90458..066d257d1 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -18,6 +18,7 @@ implementation contract; it is not, by itself, a general support claim. ### Runtime and stage execution - [Disaggregated Inference](feature/disaggregated_inference.md) +- [Host Weight Runtime](feature/host_weight_runtime.md) - [Async Chunk](feature/async_chunk.md) - [Async Diffusion Output](feature/async_diffusion_output.md) - [Async Omni Output Materialization](feature/omni_async_output_materialization.md) @@ -88,14 +89,15 @@ The design contracts separate selection mechanics from backend algorithms: - [Model Integration](module/model_integration.md) - [Autoregressive Runtime](module/ar_runtime.md) - Diffusion - - [Overview](module/diffusion/index.md) - - [Runtime](module/diffusion/diffusion_runtime.md) - - [Model Integration](module/diffusion/diffusion_model_integration.md) - - [Continuous Batching](module/diffusion/continuous_batching.md) - - [Parallelism](module/diffusion/parallelism.md) - - [Offloader](module/diffusion/offloader.md) + - [Overview](module/diffusion/index.md) + - [Runtime](module/diffusion/diffusion_runtime.md) + - [Model Integration](module/diffusion/diffusion_model_integration.md) + - [Continuous Batching](module/diffusion/continuous_batching.md) + - [Parallelism](module/diffusion/parallelism.md) + - [Offloader](module/diffusion/offloader.md) - [Execution Platforms](module/execution_platforms.md) - [Cache Management](module/cache_management.md) +- [Host Weight Runtime](module/host_weight_runtime.md) - [Quantization](module/quantization.md) - [Observability](module/observability.md) - [Profiling](module/profiling.md) diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md new file mode 100644 index 000000000..441d2d07e --- /dev/null +++ b/docs/design/module/host_weight_runtime.md @@ -0,0 +1,264 @@ +--- +title: Host Weight Runtime +kind: module +status: draft +owners: + - "@hsliuustc0106" +primary_code_paths: + - vllm_omni/host_weight_runtime/** +related_code_paths: + - vllm_omni/diffusion/model_loader/** + - vllm_omni/diffusion/offloader/** +depends_on: + - model_integration.md + - quantization.md +validation_paths: + - tests/host_weight_runtime/** +upstream_refs: + - https://github.com/vllm-project/vllm-omni/issues/6414 +last_reviewed: 2026-08-21 +--- + +# Host Weight Runtime + +The Host Weight Runtime is a loader-adjacent, transport-independent foundation +for reusing immutable, runtime-ready model weights. It separates the semantic +question—"are these exactly the weights this loader requested?"—from the +physical question—"where is the validated host representation stored?" + +The runtime can save repeated checkpoint loading, transformation, and peak host +memory when several workers need the same final representation. It also creates +a second, versioned path for reconstructing a valid model, so a cache miss or +recoverable store failure normally falls back to the canonical loader. + +This module does not enable host-weight caching for any model by itself. A model +integration must still provide an exact identity, a producer, and a restorer. +For feature behavior and consumer integration requirements, see the +[Host Weight Runtime feature design](../feature/host_weight_runtime.md). + +## Component boundary + +```mermaid +flowchart LR + LA["Model loader adapter"] -->|"exact identity + policy"| RT["HostWeightRuntime"] + RT -->|"lookup / get_or_build"| ST["HostWeightStore"] + PR["WeightProducer"] -->|"runtime-ready tensors"| ST + ST -->|"validated mmap views"| LE["HostWeightLease"] + LE --> RE["WeightRestorer"] + RE --> MO["Restored host model"] + MO --> TR["DLO or another GPU transport"] + + CS["Canonical model source"] --> PR + LA -. "fallback" .-> CS +``` + +The contracts have deliberately narrow ownership: + +- `HostWeightRuntime` applies resolution and fallback policy and emits one + terminal report. It does not own a model or call the canonical loader. +- `HostWeightStore` owns artifact lookup, coordinated construction, + publication, validation, quarantine, and lifecycle. +- `HostWeightLease` owns process-local tensor views, mapped-file resources, and + the shared artifact lock that stabilizes their lifetime. +- `WeightProducer` creates one declared final-layout representation. It writes + only through a store-scoped writer and cannot publish paths directly. +- `WeightRestorer` validates a lease-to-model plan without mutation. Its + one-shot `commit() -> None` is the sole model-mutating phase. Concrete + restorers remain with the model or quantization integration. +- GPU transport owns page registration, private staging, H2D scheduling, and + lease release. DLO must not parse manifests or manage store artifacts. + +## Exact semantic identity + +Storage lookup is exact; the store never coerces one representation into +another. `WeightArtifactIdentity` includes: + +- immutable model revision and source-file fingerprint; +- component and ownership boundary; +- runtime representation, dtype, and representation metadata; +- final tensor layout and semantic parallel coordinates; +- static adaptation identity; and +- producer version, implementation fingerprint, manifest schema, and restorer + schema. + +The identity excludes transfer policy such as registered mmap versus private +pinned staging. Those paths move the same host representation and must not +produce duplicate artifacts. + +Parallel coordinates follow semantics rather than process topology: + +- TP size and rank are included when they change tensor contents or shape. +- DP rank is excluded for replicated weights. +- SP size and backend are included when they affect layout. SP rank is excluded + only when all SP ranks consume an identical representation. +- PP, EP, or future parallel modes must be added by the model adapter whenever + they change component ownership or tensor layout. + +Dynamic LoRA overlays are not cached as base weights. A statically merged +adapter is a distinct artifact identity with its own fingerprint. + +## Deterministic resolution + +The loader first resolves the canonical source and computes the exact requested +identity. The runtime then follows this order: + +```text +validated local artifact + -> exact remote artifact materialized locally (future) + -> registered local producer + -> canonical loader fallback, if policy permits +``` + +An already published local artifact wins because it avoids network transfer and +is already in node-accessible storage. "Wins" means the highest-priority valid +representation, not whichever backing responds first. + +Configuration modes are: + +- `disabled`: return the canonical-direct outcome without probing a root, + credentials, topology, or identity; +- `preferred`: recoverable store failures may return canonical fallback; and +- `required`: fail when no exact lease can be acquired. + +Failure-to-action policy is centralized in `HostWeightRuntime`. A clean miss or +invalid cache artifact may fall back in preferred mode. A typed failure may do +so only when it is marked retryable; unsupported capabilities, identity +collisions, producer failures, and publication failures remain visible as +startup failures. + +`coordination_timeout_seconds` bounds acquisition of lookup and build locks; it +is not a hard wall-clock deadline for in-process validation, production, or +atomic publication. Once this V1 implementation becomes the producer, the +synchronous producer and publication run to completion. A hung producer blocks +its owning process and requires external process supervision; enforceable +producer cancellation requires a future process-isolated producer contract. A +waiter timing out never cancels another process's valid build. + +Remote providers are represented by protocols but are explicitly unsupported +in the first implementation. V1 also invokes only `PRE_LOAD_SAFE` producers; +enabling post-load publication is rejected during configuration until a +separate post-load entry point and model-lifetime contract exist. + +## Restoration transaction boundary + +`WeightRestorer.plan_restore()` is validation-only: it must not modify the model +or lease. The returned plan may be committed exactly once, `commit()` returns +`None`, and commit is the sole phase allowed to mutate the model. Registration, +staging, and GPU transfer remain transport responsibilities rather than restore +side effects. + +If planning fails, preferred-mode canonical fallback may reuse the untouched +model. If commit begins and then fails, the model is considered partially +hydrated and must be discarded; canonical fallback must construct a fresh model +instance instead of reusing it. + +## Filesystem domain and artifact lifecycle + +A storage domain describes physical locality and policy; it is not part of +semantic identity. The first implementation supports node- or NUMA-scoped +roots on verified local disk filesystems or tmpfs. The detected kernel +filesystem type is preserved in `domain.json` and inspection output. Known +remote filesystems and unknown filesystem types are rejected rather than being +silently treated as local disk, because cross-node page cache and advisory-lock +semantics do not satisfy this backend's contract. A NUMA root alone does not +guarantee page placement; CPU pinning, local first-touch/prefault, and residency +evidence belong to a later topology-aware integration. + +```text +/ +├── domain.json +├── domain-policy.json +├── locks/ +│ ├── .build.lock +│ └── .artifact.lock +├── artifacts// +│ ├── READY.json +│ ├── manifest.json +│ └── *.safetensors +├── tmp/ +├── quarantine/ +└── deny/ +``` + +Lock files live outside replaceable artifact directories. Every cooperative +mutation follows one lock order: build lock, then exclusive artifact lock. +Lookup acquires a shared artifact lock, validates `READY` and the manifest, +opens payloads with non-following descriptor-based access, creates mmap-backed +views, and transfers that shared lock to the lease. + +Lease acquisition rechecks the producer and restorer schemas against the exact +identity. Malformed safetensors payloads are reported as typed invalid +artifacts and denied rather than escaping the resolution policy as raw parser +exceptions. + +Each `FileLock` records its creator PID. A forked child closes an inherited lock +descriptor without issuing `LOCK_UN`, because `flock` state belongs to the +shared open-file description and unlocking in the child would also release the +parent's active lease lock. + +Publication follows this lifecycle: + +```text +ABSENT -> invisible temporary build -> fully hashed and fsynced + -> atomic rename -> READY + +READY -> denied -> quarantined -> explicit removal +``` + +`READY.json` is written last in the temporary directory. Readers therefore see +either a complete validated artifact or a miss, never an in-progress build. +Failed process-local builds are left invisible and a later build owner removes +their stale temporary directories. Explicit cleanup is nonblocking and reports +an active builder or lease instead of unlinking live mappings. + +## Validation and capacity + +Every publication computes full file and tensor hashes. Local lookup supports: + +- `manifest_and_metadata`, which validates identity, immutable publication + metadata, file descriptors, sizes, tensor keys, shapes, dtypes, and strides; + and +- `full_checksum`, which additionally reads and hashes every payload file. + +Filesystem verity is modeled but unsupported. Invalid artifacts receive an +external deny marker before their shared lookup lock is released; a build owner +then quarantines the immutable directory under the exclusive lock. +Lookup checks the marker both before validation and after constructing the +lease so a weaker overlapping lookup cannot return a lease after a stronger +validation has denied the artifact. + +`inspect_domain()` and `inspect_artifact()` return typed, read-only snapshots. +Inspection validation does not create a deny marker. Reported lock activity is +point-in-time evidence only; kernel locks are not a persistent owner registry. + +Capacity policy covers all store-owned bytes, including ready, temporary, and +quarantined data. The local writer preflights `max_artifact_bytes`, +`max_store_bytes`, and `min_free_bytes`, and preallocates payloads where the +filesystem supports it. `ENOSPC` remains a normal typed store failure because +concurrent preflight is inherently racy. Automatic eviction and strict +concurrent reservations are not implemented. + +For tmpfs, artifact bytes are host-memory consumption and may consume swap; +they must not be reported operationally as ordinary disk capacity. The store +owns files and reports capacity. The kernel controls page-cache residency, and +transport controls registration or page locking. + +## First implementation scope + +The foundation implements a CPU local filesystem store and +`SINGLE_PROCESS` production for one exact artifact. It intentionally adds no: + +- public CLI or DLO configuration; +- BF16, FP8, or model-specific producer/restorer; +- change to DLO AllGather or no-AllGather behavior; +- CUDA registration, pinned staging, or H2D scheduling; +- remote provider, group-cohort producer, artifact bundle, eviction, or + cross-node coordination. + +Consumer PRs must prove ordinary-loader parity, warm-hit work avoidance, +parallel-layout correctness, clean teardown, shared-backing evidence, and +representation-specific startup and memory behavior. Generic foundation tests +remain CPU-only and make no production performance claim. + +The complete rationale and staged rollout are tracked in +[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414). diff --git a/tests/host_weight_runtime/test_contracts.py b/tests/host_weight_runtime/test_contracts.py new file mode 100644 index 000000000..d831863da --- /dev/null +++ b/tests/host_weight_runtime/test_contracts.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Contract tests for semantic identity and immutable artifact schemas.""" + +from __future__ import annotations + +import ast +import math +from pathlib import Path +from typing import get_type_hints + +import pytest + +from vllm_omni.host_weight_runtime import ( + AdaptationIdentity, + CanonicalJson, + ComponentIdentity, + FileManifestEntry, + IdentityValidationError, + ManifestValidationError, + ProducerIdentity, + ProductionPolicy, + PublicationMarker, + RuntimeWeightLayout, + TensorKind, + TensorManifestEntry, + WeightArtifactIdentity, + WeightManifest, + WeightRepresentation, + WeightRestorePlan, + WeightRestorer, + WeightSourceIdentity, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_SHA_A = "a" * 64 +_SHA_B = "b" * 64 +_SHA_C = "c" * 64 + + +def _identity( + *, + tp_size: int = 1, + tp_rank: int = 0, + sp_size: int = 1, + adaptation: AdaptationIdentity | None = None, + source_metadata: CanonicalJson | None = None, +) -> WeightArtifactIdentity: + return WeightArtifactIdentity( + schema_version=1, + source=WeightSourceIdentity( + model_id="test-org/test-model", + revision="0123456789abcdef", + fingerprint="source-files-sha256", + metadata=source_metadata or CanonicalJson.empty(), + ), + component=ComponentIdentity(name="transformer", ownership="diffusion.dit"), + representation=WeightRepresentation(name="runtime-bf16", dtype="torch.bfloat16"), + layout=RuntimeWeightLayout( + name="final-module-layout", + tensor_parallel_size=tp_size, + tensor_parallel_rank=tp_rank, + sequence_parallel_size=sp_size, + sequence_parallel_backend="ulysses" if sp_size > 1 else "none", + ), + adaptation=adaptation or AdaptationIdentity(), + producer=ProducerIdentity( + producer_id="test.final-layout", + version="1", + implementation_fingerprint="implementation-sha256", + manifest_schema="test-producer-v1", + restorer_schema="test-restorer-v1", + ), + ) + + +def test_canonical_metadata_is_a_stable_immutable_snapshot() -> None: + items: list[object] = [1, {"value": 2}] + mutable = {"z": items, "a": True} + snapshot = CanonicalJson.from_value(mutable) + items.append(3) + + assert snapshot.encoded == b'{"a":true,"z":[1,{"value":2}]}' + assert snapshot.to_value() == {"a": True, "z": [1, {"value": 2}]} + assert CanonicalJson.from_value({"a": True, "z": [1, {"value": 2}]}) == snapshot + + +@pytest.mark.parametrize("value", [math.inf, math.nan, object(), {1: "non-string-key"}]) +def test_canonical_metadata_rejects_process_specific_or_non_json_values(value: object) -> None: + with pytest.raises(IdentityValidationError): + CanonicalJson.from_value(value) + + +@pytest.mark.parametrize("encoded", [b"not-json", b'{"b":1, "a":2}', b'{"value":NaN}']) +def test_canonical_metadata_cannot_bypass_stable_encoding(encoded: bytes) -> None: + with pytest.raises(IdentityValidationError): + CanonicalJson(encoded) + + +def test_identity_round_trip_and_hash_are_deterministic() -> None: + first = _identity(source_metadata=CanonicalJson.from_value({"files": ["b", "a"], "size": 3})) + second = _identity(source_metadata=CanonicalJson.from_value({"size": 3, "files": ["b", "a"]})) + + assert first == second + assert first.key == second.key + assert WeightArtifactIdentity.from_dict(first.to_dict()) == first + assert len(first.key) == 64 + + +def test_identity_separates_semantic_parallel_layout_and_static_adaptation() -> None: + base = _identity() + tp0 = _identity(tp_size=2, tp_rank=0) + tp1 = _identity(tp_size=2, tp_rank=1) + sp2 = _identity(sp_size=2) + merged_lora = _identity(adaptation=AdaptationIdentity(kind="merged-lora", fingerprint="adapter-sha256")) + + assert len({item.key for item in (base, tp0, tp1, sp2, merged_lora)}) == 5 + assert "data_parallel" not in base.layout.to_dict() + assert "sequence_parallel_rank" not in sp2.layout.to_dict() + + +def test_identity_schema_rejects_unknown_fields() -> None: + document = _identity().to_dict() + document["transfer_mode"] = "registered" + + with pytest.raises(IdentityValidationError, match="unexpected"): + WeightArtifactIdentity.from_dict(document) + + +def test_identity_schema_does_not_coerce_field_types() -> None: + document = _identity().to_dict() + document["schema_version"] = "1" + + with pytest.raises(IdentityValidationError, match="must be an integer"): + WeightArtifactIdentity.from_dict(document) + + +def test_restoration_contract_has_a_validation_only_one_shot_commit_boundary() -> None: + assert get_type_hints(WeightRestorePlan.commit)["return"] is type(None) + assert "one-shot" in (WeightRestorePlan.__doc__ or "") + assert "without mutating" in (WeightRestorer.__doc__ or "") + + +def test_post_load_publication_policy_is_explicitly_unsupported() -> None: + with pytest.raises(ValueError, match="post-load publication is not implemented"): + ProductionPolicy(allow_post_load_publish=True) + + +def test_manifest_and_publication_marker_round_trip_exactly() -> None: + identity = _identity() + manifest = WeightManifest( + schema_version=1, + artifact_key=identity.key, + identity=identity, + producer_schema="test-producer-v1", + restorer_schema="test-restorer-v1", + tensors=( + TensorManifestEntry( + name="layer.weight", + kind=TensorKind.PARAMETER, + role="weight", + file_name="weights.safetensors", + storage_key="layer.weight", + shape=(2, 4), + dtype="torch.bfloat16", + stride=(4, 1), + nbytes=16, + sha256=_SHA_A, + ), + ), + files=(FileManifestEntry(file_name="weights.safetensors", size=128, sha256=_SHA_B),), + artifact_content_sha256=_SHA_C, + format_metadata=CanonicalJson.from_value({"quantization": "none"}), + ) + marker = PublicationMarker( + schema_version=1, + publication_id="publication-id", + manifest_sha256=_SHA_A, + artifact_content_sha256=_SHA_C, + integrity_mode="full_checksum", + ) + + assert WeightManifest.from_dict(manifest.to_dict()) == manifest + assert PublicationMarker.from_dict(marker.to_dict()) == marker + + document = manifest.to_dict() + document["unknown"] = True + with pytest.raises(ManifestValidationError, match="unexpected"): + WeightManifest.from_dict(document) + + malformed = manifest.to_dict() + tensors = malformed["tensors"] + assert isinstance(tensors, list) and isinstance(tensors[0], dict) + tensors[0]["nbytes"] = "16" + with pytest.raises(ManifestValidationError, match="must be an integer"): + WeightManifest.from_dict(malformed) + + +@pytest.mark.parametrize("file_name", ["", ".", "..", "../weights.safetensors", "nested/weights.safetensors"]) +def test_manifest_rejects_unsafe_payload_names(file_name: str) -> None: + with pytest.raises(ManifestValidationError, match="safe relative basename"): + FileManifestEntry(file_name=file_name, size=1, sha256=_SHA_A) + + +def test_package_source_keeps_the_neutral_dependency_boundary() -> None: + package_root = Path(__file__).resolve().parents[2] / "vllm_omni" / "host_weight_runtime" + assert package_root.is_dir() + forbidden = ( + "vllm_omni.diffusion", + "vllm_omni.distributed", + "torch.cuda", + "torch.distributed", + ) + violations: list[str] = [] + for source_path in package_root.rglob("*.py"): + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + violations.extend( + f"{source_path.name}:{node.lineno}:{alias.name}" + for alias in node.names + if alias.name.startswith(forbidden) + ) + elif isinstance(node, ast.ImportFrom) and node.level == 0: + module = node.module or "" + if module.startswith(forbidden) or ( + module == "torch" and any(alias.name in {"cuda", "distributed"} for alias in node.names) + ): + violations.append(f"{source_path.name}:{node.lineno}:{module}") + + assert violations == [] diff --git a/tests/host_weight_runtime/test_filesystem_store.py b/tests/host_weight_runtime/test_filesystem_store.py new file mode 100644 index 000000000..763d2622a --- /dev/null +++ b/tests/host_weight_runtime/test_filesystem_store.py @@ -0,0 +1,1106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Lifecycle and concurrency tests for the local filesystem store.""" + +from __future__ import annotations + +import contextlib +import errno +import hashlib +import json +import multiprocessing as mp +import os +import shutil +import threading +import time +import warnings +from collections.abc import Sequence +from dataclasses import replace +from multiprocessing.process import BaseProcess +from pathlib import Path + +import pytest +import torch + +from vllm_omni.host_weight_runtime import ( + AdaptationIdentity, + ArtifactInventoryState, + BuildRequest, + CapacityPolicy, + ComponentIdentity, + CoordinationScope, + FailureCode, + HostWeightError, + HostWeightLease, + IntegrityPolicy, + ProducerIdentity, + ProductionMetadata, + ProductionSourceMode, + ResolutionStage, + RuntimeWeightLayout, + StorageClass, + StorageDomainPolicy, + StorageScope, + StoreResult, + StoreStatus, + TensorKind, + TensorWriteSpec, + ValidationLevel, + WeightArtifactIdentity, + WeightProductionSpec, + WeightRepresentation, + WeightSourceIdentity, +) +from vllm_omni.host_weight_runtime.filesystem import ( + FilesystemHostWeightStore, + detect_filesystem_type, + detect_storage_class, +) +from vllm_omni.host_weight_runtime.filesystem import store as filesystem_store_module +from vllm_omni.host_weight_runtime.filesystem.locks import FileLock + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Spawning imports the complete vLLM package and exceeded 10 seconds on a loaded +# CI worker. Keep process-bootstrap bounds separate from the short store +# deadlines exercised by these tests. +_PROCESS_START_TIMEOUT_SECONDS = 60.0 +_PROCESS_COMPLETION_TIMEOUT_SECONDS = 60.0 +_PROCESS_CLEANUP_TIMEOUT_SECONDS = 5.0 + + +def _identity(*, tp_rank: int = 0) -> WeightArtifactIdentity: + return WeightArtifactIdentity( + schema_version=1, + source=WeightSourceIdentity("test-org/test-model", "0123456789abcdef", "source-files-sha256"), + component=ComponentIdentity("transformer", "diffusion.dit"), + representation=WeightRepresentation("runtime-bf16", "torch.bfloat16"), + layout=RuntimeWeightLayout( + "final-module-layout", + tensor_parallel_size=2, + tensor_parallel_rank=tp_rank, + ), + adaptation=AdaptationIdentity(), + producer=ProducerIdentity( + "test.final-layout", + "1", + "implementation-sha256", + "test-producer-v1", + "test-restorer-v1", + ), + ) + + +def _weight(offset: int = 0) -> torch.Tensor: + return torch.arange(offset, offset + 12, dtype=torch.float32).to(torch.bfloat16).reshape(4, 3) + + +def _specs() -> tuple[TensorWriteSpec, ...]: + return ( + TensorWriteSpec("layer.weight", (4, 3), torch.bfloat16), + TensorWriteSpec("layer.scale", (2,), torch.float32, kind=TensorKind.BUFFER, role="scale"), + ) + + +def _open_fd_targets() -> tuple[str, ...]: + targets = [] + for path in Path("/proc/self/fd").iterdir(): + with contextlib.suppress(OSError): + targets.append(os.readlink(path)) + return tuple(targets) + + +class FakeProducer: + def __init__( + self, + identity: WeightArtifactIdentity, + *, + counter_path: Path | None = None, + delay_seconds: float = 0.0, + started: object | None = None, + release: object | None = None, + write_mode: str = "rows", + weight_offset: int = 0, + coordination_scope: CoordinationScope = CoordinationScope.SINGLE_PROCESS, + ) -> None: + self._spec = WeightProductionSpec( + producer_id="test.final-layout", + outputs=(identity,), + source_mode=ProductionSourceMode.FINALIZED_MODEL, + coordination_scope=coordination_scope, + ) + self.counter_path = counter_path + self.delay_seconds = delay_seconds + self.started = started + self.release = release + self.write_mode = write_mode + self.weight_offset = weight_offset + + @property + def spec(self) -> WeightProductionSpec: + return self._spec + + def produce(self, writer: object) -> ProductionMetadata: + if self.counter_path is not None: + fd = os.open(self.counter_path, os.O_CREAT | os.O_APPEND | os.O_WRONLY | os.O_CLOEXEC, 0o600) + try: + os.write(fd, f"{os.getpid()}\n".encode()) + os.fsync(fd) + finally: + os.close(fd) + if self.started is not None: + self.started.set() # type: ignore[attr-defined] + if self.release is not None and not self.release.wait(timeout=10): # type: ignore[attr-defined] + raise TimeoutError("test producer was not released") + if self.delay_seconds: + time.sleep(self.delay_seconds) + + with writer.open_tensor_file("weights.safetensors", _specs()) as output: # type: ignore[attr-defined] + if self.write_mode == "rows": + output.write_rows("layer.weight", 0, _weight(self.weight_offset)[:2]) + output.write_rows("layer.weight", 2, _weight(self.weight_offset)[2:]) + elif self.write_mode == "complete": + output.write_tensor("layer.weight", _weight(self.weight_offset)) + elif self.write_mode == "incomplete": + output.write_rows("layer.weight", 0, _weight(self.weight_offset)[:2]) + elif self.write_mode == "overlap": + output.write_rows("layer.weight", 0, _weight(self.weight_offset)[:3]) + output.write_rows("layer.weight", 2, _weight(self.weight_offset)[2:]) + else: # pragma: no cover - test helper guard + raise AssertionError(f"unknown write mode {self.write_mode}") + output.write_tensor("layer.scale", torch.tensor([1.0, 2.0])) + return ProductionMetadata("test-producer-v1", "test-restorer-v1") + + +class CrashProducer(FakeProducer): + def produce(self, writer: object) -> ProductionMetadata: + output = writer.open_tensor_file("weights.safetensors", _specs()) # type: ignore[attr-defined] + output.write_rows("layer.weight", 0, _weight()[:2]) + os._exit(17) + + +class NonCooperativeStore(FilesystemHostWeightStore): + """Test double representing a publisher that does not join the build lock.""" + + def _build_lock_path(self, key: str) -> Path: + return self.locks_dir / f"{key}.noncooperative-build.lock" + + def _remove_stale_temps_locked(self, _key: str) -> None: + return + + +class InitialLookupBarrierStore(FilesystemHostWeightStore): + """Make both test workers observe the initial miss before either builds.""" + + def __init__( + self, + domain: StorageDomainPolicy, + capacity: CapacityPolicy, + integrity: IntegrityPolicy, + *, + initial_lookup_barrier: object, + ) -> None: + super().__init__(domain, capacity, integrity) + self._initial_lookup_barrier = initial_lookup_barrier + self._lookup_count = 0 + + def lookup( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + deadline: float | None = None, + ) -> StoreResult: + result = super().lookup(identity, validation=validation, deadline=deadline) + self._lookup_count += 1 + if self._lookup_count == 1: + self._initial_lookup_barrier.wait(timeout=10) # type: ignore[attr-defined] + return result + + +class PausingOpenLeaseStore(FilesystemHostWeightStore): + """Pause after opening a lease so another lookup can publish a deny marker.""" + + def __init__( + self, + domain: StorageDomainPolicy, + capacity: CapacityPolicy, + integrity: IntegrityPolicy, + *, + lease_opened: threading.Event, + resume_lookup: threading.Event, + ) -> None: + super().__init__(domain, capacity, integrity) + self._lease_opened = lease_opened + self._resume_lookup = resume_lookup + + def _open_lease( + self, + identity: WeightArtifactIdentity, + entry_path: Path, + artifact_lock: FileLock, + validation: ValidationLevel, + ) -> HostWeightLease: + lease = super()._open_lease(identity, entry_path, artifact_lock, validation) + self._lease_opened.set() + if not self._resume_lookup.wait(5): + lease.close() + raise TimeoutError("test did not resume paused lookup") + return lease + + +def _make_store( + root: Path, + *, + capacity: CapacityPolicy | None = None, + integrity: IntegrityPolicy | None = None, +) -> FilesystemHostWeightStore: + root.mkdir(parents=True, exist_ok=True, mode=0o700) + domain = StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)) + return FilesystemHostWeightStore(domain, capacity or CapacityPolicy(), integrity or IntegrityPolicy()) + + +def _build_process(root: str, counter: str, initial_lookup_barrier: object, result_queue: object) -> None: + try: + identity = _identity() + root_path = Path(root) + root_path.mkdir(parents=True, exist_ok=True, mode=0o700) + store = InitialLookupBarrierStore( + StorageDomainPolicy(root=root_path, storage_class=detect_storage_class(root_path)), + CapacityPolicy(), + IntegrityPolicy(), + initial_lookup_barrier=initial_lookup_barrier, + ) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, counter_path=Path(counter), delay_seconds=0.25), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 15, + ) + fingerprint = result.lease.provenance.backing_fingerprint if result.lease is not None else None + if result.lease is not None: + result.lease.close() + result_queue.put((result.status.value, fingerprint, None)) # type: ignore[attr-defined] + except BaseException as exc: + result_queue.put(("error", None, repr(exc))) # type: ignore[attr-defined] + + +def _controlled_build_process(root: str, started: object, release: object, result_queue: object) -> None: + try: + identity = _identity() + store = _make_store(Path(root)) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, started=started, release=release), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 15, + ) + if result.lease is not None: + result.lease.close() + result_queue.put((result.status.value, None)) # type: ignore[attr-defined] + except BaseException as exc: + result_queue.put(("error", repr(exc))) # type: ignore[attr-defined] + + +def _crash_build_process(root: str) -> None: + identity = _identity() + store = _make_store(Path(root)) + store.get_or_build( + BuildRequest(identity), + CrashProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 15, + ) + + +def _join_processes( + processes: Sequence[BaseProcess], + timeout: float = _PROCESS_COMPLETION_TIMEOUT_SECONDS, +) -> None: + deadline = time.monotonic() + timeout + for process in processes: + process.join(max(0.0, deadline - time.monotonic())) + timed_out = [process for process in processes if process.is_alive()] + if not timed_out: + return + timed_out_workers = ", ".join(f"{process.name}(pid={process.pid})" for process in timed_out) + for process in timed_out: + process.terminate() + for process in timed_out: + process.join(_PROCESS_CLEANUP_TIMEOUT_SECONDS) + survivors = [process for process in timed_out if process.is_alive()] + for process in survivors: + process.kill() + process.join(_PROCESS_CLEANUP_TIMEOUT_SECONDS) + unstoppable = [process for process in survivors if process.is_alive()] + assert not unstoppable, "timed-out test worker processes could not be stopped" + pytest.fail(f"test worker processes exceeded the {timeout:g}-second group deadline: {timed_out_workers}") + + +def test_build_lookup_share_one_published_backing_and_lease_blocks_cleanup(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.status is StoreStatus.BUILT + assert built.lease is not None + assert torch.equal(built.lease.tensors["layer.weight"], _weight()) + assert torch.equal(built.lease.tensors["layer.scale"], torch.tensor([1.0, 2.0])) + assert built.lease.mapped_regions + assert all(region.address > 0 and region.length > 0 and region.read_only for region in built.lease.mapped_regions) + with pytest.raises(TypeError, match="process-local"): + built.lease.__reduce_ex__(4) + + second = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert second.status is StoreStatus.HIT + assert second.lease is not None + assert second.lease.provenance.publication_id == built.lease.provenance.publication_id + assert second.lease.provenance.backing_fingerprint == built.lease.provenance.backing_fingerprint + payload = store.artifacts_dir / identity.key / "weights.safetensors" + assert any(str(payload) in target for target in _open_fd_targets()) + + artifact = store.inspect_artifact(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert artifact.status is StoreStatus.HIT + assert artifact.manifest is not None and artifact.manifest.identity == identity + assert artifact.publication_id == built.lease.provenance.publication_id + assert artifact.artifact_lock_active + domain = store.inspect_domain() + assert domain.domain_uuid == built.lease.provenance.domain_uuid + assert domain.filesystem_type == detect_filesystem_type(store.root) + assert any( + item.artifact_key == identity.key and item.state is ArtifactInventoryState.READY and item.size_bytes > 0 + for item in domain.inventory + ) + + active = store.cleanup(identity) + assert active is not None and active.code is FailureCode.ACTIVE_LEASE + second.lease.close() + built.lease.close() + built.lease.close() + assert not built.lease.tensors + assert not any(str(payload) in target for target in _open_fd_targets()) + + assert store.cleanup(identity) is None + assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.MISS + + +def test_concurrent_lease_close_waits_for_resource_teardown(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + lease = built.lease + teardown_started = threading.Event() + allow_teardown = threading.Event() + second_started = threading.Event() + second_returned = threading.Event() + errors: list[BaseException] = [] + + def blocking_closer() -> None: + teardown_started.set() + if not allow_teardown.wait(5): + raise TimeoutError("test did not release lease teardown") + + def close_lease(*, started: threading.Event | None = None, returned: threading.Event | None = None) -> None: + if started is not None: + started.set() + try: + lease.close() + except BaseException as exc: + errors.append(exc) + finally: + if returned is not None: + returned.set() + + lease._resource_closers = (*lease._resource_closers, blocking_closer) + first = threading.Thread(target=close_lease) + second = threading.Thread(target=close_lease, kwargs={"started": second_started, "returned": second_returned}) + first.start() + assert teardown_started.wait(5) + second.start() + assert second_started.wait(5) + assert not second_returned.wait(0.2) + + allow_teardown.set() + first.join(5) + second.join(5) + assert not first.is_alive() and not second.is_alive() + assert second_returned.is_set() + assert not errors + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork-specific flock ownership regression") +def test_forked_child_close_does_not_release_parent_lease_lock(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + child_pid = os.fork() + if child_pid == 0: # pragma: no cover - assertions run in the parent + try: + built.lease.close() + except BaseException: + os._exit(1) + os._exit(0) + _, child_status = os.waitpid(child_pid, 0) + assert os.waitstatus_to_exitcode(child_status) == 0 + + active = store.cleanup(identity) + assert active is not None and active.code is FailureCode.ACTIVE_LEASE + built.lease.close() + assert store.cleanup(identity) is None + + +@pytest.mark.parametrize("write_mode", ["incomplete", "overlap"]) +def test_failed_streaming_writes_never_publish_partial_artifacts(tmp_path: Path, write_mode: str) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, write_mode=write_mode), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert result.status is StoreStatus.FAILED + assert result.failure is not None and result.failure.code is FailureCode.PRODUCER_FAILED + assert not list(store.tmp_dir.iterdir()) + assert not list(store.artifacts_dir.iterdir()) + + +@pytest.mark.parametrize( + "dtype", + [ + dtype + for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz") + if (dtype := getattr(torch, name, None)) is not None + ], +) +def test_available_float8_formats_round_trip_through_streaming_store(tmp_path: Path, dtype: torch.dtype) -> None: + identity = replace( + _identity(), + representation=WeightRepresentation(name="runtime-fp8", dtype=str(dtype)), + ) + expected = torch.zeros((2, 4), dtype=dtype) + + class Float8Producer: + spec = WeightProductionSpec( + producer_id="test.final-layout", + outputs=(identity,), + source_mode=ProductionSourceMode.FINALIZED_MODEL, + ) + + def produce(self, writer: object) -> ProductionMetadata: + spec = TensorWriteSpec("weight", tuple(expected.shape), dtype) + with writer.open_tensor_file("weights.safetensors", (spec,)) as output: # type: ignore[attr-defined] + output.write_tensor("weight", expected) + return ProductionMetadata("test-producer-v1", "test-restorer-v1") + + store = _make_store(tmp_path / "store") + result = store.get_or_build( + BuildRequest(identity), + Float8Producer(), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert result.status is StoreStatus.BUILT + assert result.lease is not None + actual = result.lease.tensors["weight"] + assert actual.dtype is dtype + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + result.lease.close() + + +def test_lookup_treats_orphan_deny_marker_as_miss(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + deny = store.deny_dir / f"{identity.key}.json" + deny.write_text("{}", encoding="utf-8") + + result = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + + assert result.status is StoreStatus.MISS + + +def test_overlapping_weaker_lookup_cannot_return_after_artifact_is_denied(tmp_path: Path) -> None: + identity = _identity() + root = tmp_path / "store" + store = _make_store(root) + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + built.lease.close() + + artifact_dir = store.artifacts_dir / identity.key + payload = artifact_dir / "weights.safetensors" + os.chmod(artifact_dir, 0o755) + os.chmod(payload, 0o644) + with payload.open("r+b") as handle: + handle.seek(-1, os.SEEK_END) + original = handle.read(1) + handle.seek(-1, os.SEEK_END) + handle.write(bytes([original[0] ^ 0xFF])) + handle.flush() + os.fsync(handle.fileno()) + os.chmod(payload, 0o444) + os.chmod(artifact_dir, 0o555) + + lease_opened = threading.Event() + resume_lookup = threading.Event() + overlapping = PausingOpenLeaseStore( + StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)), + CapacityPolicy(), + IntegrityPolicy(), + lease_opened=lease_opened, + resume_lookup=resume_lookup, + ) + weak_results: list[StoreResult] = [] + + def weak_lookup() -> None: + weak_results.append(overlapping.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA)) + + thread = threading.Thread(target=weak_lookup) + thread.start() + try: + assert lease_opened.wait(5) + strong = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert strong.status is StoreStatus.INVALID + assert (store.deny_dir / f"{identity.key}.json").is_file() + finally: + resume_lookup.set() + thread.join(5) + + assert not thread.is_alive() + assert len(weak_results) == 1 + weak = weak_results[0] + assert weak.status is StoreStatus.INVALID + assert weak.failure is not None and weak.failure.code is FailureCode.ARTIFACT_DENIED + + +def test_full_checksum_denies_quarantines_and_rebuilds_corrupt_payload(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + built.lease.close() + + artifact_dir = store.artifacts_dir / identity.key + payload = artifact_dir / "weights.safetensors" + os.chmod(artifact_dir, 0o755) + os.chmod(payload, 0o644) + with payload.open("r+b") as handle: + handle.seek(-1, os.SEEK_END) + original = handle.read(1) + handle.seek(-1, os.SEEK_END) + handle.write(bytes([original[0] ^ 0xFF])) + handle.flush() + os.fsync(handle.fileno()) + os.chmod(payload, 0o444) + os.chmod(artifact_dir, 0o555) + + metadata_only = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA) + assert metadata_only.status is StoreStatus.HIT + assert metadata_only.lease is not None + metadata_only.lease.close() + + inspected = store.inspect_artifact(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert inspected.status is StoreStatus.INVALID + assert inspected.failure is not None and inspected.failure.code is FailureCode.CHECKSUM_MISMATCH + assert not (store.deny_dir / f"{identity.key}.json").exists() + + invalid = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert invalid.status is StoreStatus.INVALID + assert invalid.failure is not None and invalid.failure.code is FailureCode.CHECKSUM_MISMATCH + assert (store.deny_dir / f"{identity.key}.json").is_file() + + rebuilt = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert rebuilt.status is StoreStatus.BUILT + assert rebuilt.lease is not None + assert torch.equal(rebuilt.lease.tensors["layer.weight"], _weight()) + rebuilt.lease.close() + assert not (store.deny_dir / f"{identity.key}.json").exists() + assert len(list(store.quarantine_dir.iterdir())) == 1 + + +def test_malformed_safetensors_is_a_typed_invalid_artifact(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + built.lease.close() + + artifact_dir = store.artifacts_dir / identity.key + payload = artifact_dir / "weights.safetensors" + os.chmod(artifact_dir, 0o755) + os.chmod(payload, 0o644) + with payload.open("r+b") as handle: + handle.write(b"\xff" * 8) + handle.flush() + os.fsync(handle.fileno()) + os.chmod(payload, 0o444) + os.chmod(artifact_dir, 0o555) + + result = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA) + + assert result.status is StoreStatus.INVALID + assert result.failure is not None and result.failure.code is FailureCode.MANIFEST_CORRUPT + assert (store.deny_dir / f"{identity.key}.json").is_file() + + +@pytest.mark.parametrize("schema_field", ["producer_schema", "restorer_schema"]) +def test_lookup_rejects_manifest_schema_that_differs_from_identity( + tmp_path: Path, + schema_field: str, +) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + built.lease.close() + + artifact_dir = store.artifacts_dir / identity.key + manifest_path = artifact_dir / "manifest.json" + marker_path = artifact_dir / "READY.json" + manifest_document = json.loads(manifest_path.read_bytes()) + marker_document = json.loads(marker_path.read_bytes()) + manifest_document[schema_field] = "incompatible-schema" + manifest_bytes = json.dumps( + manifest_document, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + marker_document["manifest_sha256"] = hashlib.sha256(manifest_bytes).hexdigest() + marker_bytes = json.dumps( + marker_document, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + + os.chmod(artifact_dir, 0o755) + for path, contents in ((manifest_path, manifest_bytes), (marker_path, marker_bytes)): + os.chmod(path, 0o644) + path.write_bytes(contents) + os.chmod(path, 0o444) + os.chmod(artifact_dir, 0o555) + + result = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA) + + assert result.status is StoreStatus.INVALID + assert result.failure is not None and result.failure.code is FailureCode.MANIFEST_INCOMPATIBLE + assert (store.deny_dir / f"{identity.key}.json").is_file() + + +def test_noncooperative_conflicting_publication_is_reported_without_replacement(tmp_path: Path) -> None: + identity = _identity() + root = tmp_path / "store" + store = _make_store(root) + started = threading.Event() + release = threading.Event() + thread_result: list[StoreResult] = [] + + def build_expected_content() -> None: + thread_result.append( + store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, started=started, release=release), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + ) + + builder = threading.Thread(target=build_expected_content) + builder.start() + try: + assert started.wait(timeout=5) + competing = NonCooperativeStore(store.domain_policy, store.capacity_policy, store.integrity_policy) + published = competing.get_or_build( + BuildRequest(identity), + FakeProducer(identity, weight_offset=100), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert published.status is StoreStatus.BUILT + assert published.lease is not None + published.lease.close() + finally: + release.set() + builder.join(timeout=10) + assert not builder.is_alive() + assert len(thread_result) == 1 + collision = thread_result[0] + assert collision.status is StoreStatus.FAILED + assert collision.failure is not None and collision.failure.code is FailureCode.IDENTITY_COLLISION + + winner = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert winner.status is StoreStatus.HIT + assert winner.lease is not None + assert torch.equal(winner.lease.tensors["layer.weight"], _weight(100)) + winner.lease.close() + assert len(list(store.quarantine_dir.iterdir())) == 1 + + +def test_payload_symlink_is_rejected_without_following_it(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + built.lease.close() + + artifact_dir = store.artifacts_dir / identity.key + payload = artifact_dir / "weights.safetensors" + outside = tmp_path / "outside.safetensors" + os.chmod(artifact_dir, 0o755) + payload.rename(outside) + payload.symlink_to(outside) + os.chmod(artifact_dir, 0o555) + + result = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA) + assert result.status is StoreStatus.INVALID + assert result.failure is not None and result.failure.code is FailureCode.MANIFEST_CORRUPT + + +def test_coordination_deadline_does_not_claim_to_cancel_an_active_producer(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + started = time.monotonic() + + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, delay_seconds=0.15), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 0.02, + ) + + assert time.monotonic() - started >= 0.1 + assert result.status is StoreStatus.BUILT + assert result.lease is not None + result.lease.close() + + +@pytest.mark.parallel +def test_multiprocess_contention_runs_one_builder_and_joins_same_artifact(tmp_path: Path) -> None: + ctx = mp.get_context("spawn") + root = tmp_path / "store" + _make_store(root) + counter = tmp_path / "producer-count" + initial_lookup_barrier = ctx.Barrier(2) + result_queue = ctx.Queue() + processes = [ + ctx.Process( + target=_build_process, + args=(str(root), str(counter), initial_lookup_barrier, result_queue), + ) + for _ in range(2) + ] + for process in processes: + process.start() + _join_processes(processes) + + results = [result_queue.get(timeout=3) for _ in processes] + assert all(process.exitcode == 0 for process in processes) + assert sorted(item[0] for item in results) == [StoreStatus.BUILT.value, StoreStatus.JOINED.value] + assert all(item[2] is None for item in results) + assert len({item[1] for item in results}) == 1 + assert counter.read_text(encoding="utf-8").count("\n") == 1 + + +@pytest.mark.parallel +def test_reader_sees_miss_during_build_and_waiter_has_bounded_timeout(tmp_path: Path) -> None: + ctx = mp.get_context("spawn") + root = tmp_path / "store" + store = _make_store(root) + identity = _identity() + started = ctx.Event() + release = ctx.Event() + result_queue = ctx.Queue() + process = ctx.Process( + target=_controlled_build_process, + args=(str(root), started, release, result_queue), + ) + process.start() + try: + assert started.wait(timeout=_PROCESS_START_TIMEOUT_SECONDS) + assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.MISS + timed_out = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 0.1, + ) + assert timed_out.status is StoreStatus.TIMEOUT + assert timed_out.failure is not None and timed_out.failure.code is FailureCode.ACTIVE_BUILD_TIMEOUT + finally: + release.set() + _join_processes([process]) + + status, error = result_queue.get(timeout=3) + assert process.exitcode == 0 + assert (status, error) == (StoreStatus.BUILT.value, None) + + +@pytest.mark.parallel +def test_crashed_builder_leaves_only_invisible_temp_and_next_builder_recovers(tmp_path: Path) -> None: + ctx = mp.get_context("spawn") + root = tmp_path / "store" + store = _make_store(root) + process = ctx.Process(target=_crash_build_process, args=(str(root),)) + process.start() + _join_processes([process]) + + assert process.exitcode == 17 + assert list(store.tmp_dir.iterdir()) + assert not list(store.artifacts_dir.iterdir()) + + identity = _identity() + recovered = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert recovered.status is StoreStatus.BUILT + assert recovered.lease is not None + recovered.lease.close() + assert not list(store.tmp_dir.iterdir()) + + +def test_capacity_and_enospc_failures_are_typed_and_clean_staging( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + identity = _identity() + limited = _make_store(tmp_path / "limited", capacity=CapacityPolicy(max_artifact_bytes=64)) + too_large = limited.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert too_large.status is StoreStatus.FAILED + assert too_large.failure is not None and too_large.failure.code is FailureCode.ARTIFACT_TOO_LARGE + assert not list(limited.tmp_dir.iterdir()) + + no_space = _make_store(tmp_path / "enospc") + + def raise_enospc(_fd: int, _offset: int, _length: int) -> None: + raise OSError(errno.ENOSPC, "injected no space") + + monkeypatch.setattr(os, "posix_fallocate", raise_enospc) + failed = no_space.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert failed.status is StoreStatus.FAILED + assert failed.failure is not None and failed.failure.code is FailureCode.ENOSPC + assert not list(no_space.tmp_dir.iterdir()) + assert not list(no_space.artifacts_dir.iterdir()) + + +@pytest.mark.parametrize( + ("failure_errno", "expected_stage", "expected_code"), + [ + (errno.ENOSPC, ResolutionStage.CAPACITY, FailureCode.ENOSPC), + (errno.EDQUOT, ResolutionStage.CAPACITY, FailureCode.ENOSPC), + (errno.EACCES, ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE), + ], +) +def test_staging_directory_failures_are_typed_and_clean( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + failure_errno: int, + expected_stage: ResolutionStage, + expected_code: FailureCode, +) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + original_mkdir = Path.mkdir + + def fail_staging_mkdir( + path: Path, + mode: int = 0o777, + parents: bool = False, + exist_ok: bool = False, + ) -> None: + if path.parent == store.tmp_dir: + raise OSError(failure_errno, "injected staging directory failure") + original_mkdir(path, mode=mode, parents=parents, exist_ok=exist_ok) + + monkeypatch.setattr(Path, "mkdir", fail_staging_mkdir) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert result.status is StoreStatus.FAILED + assert result.failure is not None + assert result.failure.stage is expected_stage + assert result.failure.code is expected_code + assert not list(store.tmp_dir.iterdir()) + assert not list(store.artifacts_dir.iterdir()) + + +def test_store_and_free_space_capacity_limits_fail_before_publication(tmp_path: Path) -> None: + identity = _identity() + store_limited = _make_store(tmp_path / "store-limit", capacity=CapacityPolicy(max_store_bytes=1)) + store_result = store_limited.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert store_result.status is StoreStatus.FAILED + assert store_result.failure is not None and store_result.failure.code is FailureCode.STORE_LIMIT_EXCEEDED + assert not list(store_limited.artifacts_dir.iterdir()) + + free_root = tmp_path / "free-limit" + free_root.mkdir(mode=0o700) + unavailable_reserve = shutil.disk_usage(free_root).free + 1 + free_limited = _make_store(free_root, capacity=CapacityPolicy(min_free_bytes=unavailable_reserve)) + free_result = free_limited.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert free_result.status is StoreStatus.FAILED + assert free_result.failure is not None and free_result.failure.code is FailureCode.INSUFFICIENT_CAPACITY + assert not list(free_limited.artifacts_dir.iterdir()) + + +def test_domain_capacity_policy_is_authoritative(tmp_path: Path) -> None: + root = tmp_path / "store" + _make_store(root, capacity=CapacityPolicy(max_store_bytes=4096)) + + with pytest.raises(HostWeightError, match="capacity policy differs"): + _make_store(root, capacity=CapacityPolicy(max_store_bytes=8192)) + + +@pytest.mark.parametrize("filesystem_type", ["nfs4", "cifs", "lustre", "ceph", "unknownfs"]) +def test_remote_and_unknown_filesystems_cannot_back_the_local_store( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + filesystem_type: str, +) -> None: + monkeypatch.setattr(filesystem_store_module, "detect_filesystem_type", lambda _path: filesystem_type) + root = tmp_path / "store" + + with pytest.raises(ValueError, match="not node-local|not recognized"): + detect_storage_class(root) + with pytest.raises(HostWeightError) as error: + FilesystemHostWeightStore( + StorageDomainPolicy(root=root, storage_class=StorageClass.DISK), + CapacityPolicy(), + IntegrityPolicy(), + ) + + assert error.value.failure.code is FailureCode.DOMAIN_POLICY_MISMATCH + assert not root.exists() + + +def test_trusted_store_rejects_group_or_world_writable_root(tmp_path: Path) -> None: + root = tmp_path / "untrusted" + root.mkdir(mode=0o700) + root.chmod(0o777) + + with pytest.raises(HostWeightError) as error: + FilesystemHostWeightStore( + StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)), + CapacityPolicy(), + IntegrityPolicy(), + ) + + assert error.value.failure.code is FailureCode.DOMAIN_UNAVAILABLE + + +def test_unimplemented_coordination_and_fs_verity_return_typed_failures(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + + unsupported_producer = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, coordination_scope=CoordinationScope.GROUP_COHORT), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert unsupported_producer.status is StoreStatus.FAILED + assert unsupported_producer.failure is not None + assert unsupported_producer.failure.code is FailureCode.PRODUCER_UNSUPPORTED + + fs_verity = store.lookup(identity, validation=ValidationLevel.FS_VERITY) + assert fs_verity.status is StoreStatus.FAILED + assert fs_verity.failure is not None and fs_verity.failure.code is FailureCode.UNSUPPORTED + + +def test_unimplemented_numa_enforcement_and_background_scrub_fail_explicitly(tmp_path: Path) -> None: + numa_root = tmp_path / "numa" + numa_root.mkdir(mode=0o700) + with pytest.raises(HostWeightError) as numa_error: + FilesystemHostWeightStore( + StorageDomainPolicy( + root=numa_root, + scope=StorageScope.NUMA, + storage_class=detect_storage_class(numa_root), + numa_prefault=True, + ), + CapacityPolicy(), + IntegrityPolicy(), + ) + assert numa_error.value.failure.code is FailureCode.UNSUPPORTED + + with pytest.raises(HostWeightError) as scrub_error: + _make_store(tmp_path / "scrub", integrity=IntegrityPolicy(background_scrub=True)) + assert scrub_error.value.failure.code is FailureCode.UNSUPPORTED diff --git a/tests/host_weight_runtime/test_runtime.py b/tests/host_weight_runtime/test_runtime.py new file mode 100644 index 000000000..afd7c943a --- /dev/null +++ b/tests/host_weight_runtime/test_runtime.py @@ -0,0 +1,364 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Resolution-policy tests for the loader-adjacent Host Weight Runtime.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +from vllm_omni.host_weight_runtime import ( + AdaptationIdentity, + CanonicalJson, + ComponentIdentity, + FailureCode, + HostWeightFailure, + HostWeightRuntime, + HostWeightRuntimeConfig, + IntegrityPolicy, + LookupPhase, + ProducerIdentity, + ProductionMetadata, + ProductionSourceMode, + RemoteImportPolicy, + RemoteOnMiss, + ResolutionOutcome, + ResolutionReport, + ResolutionStage, + RuntimeMode, + RuntimeWeightLayout, + StorageDomainPolicy, + StoreResult, + StoreStatus, + TensorWriteSpec, + ValidationLevel, + WeightArtifactIdentity, + WeightProductionSpec, + WeightRepresentation, + WeightSourceIdentity, +) +from vllm_omni.host_weight_runtime.filesystem import detect_storage_class + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _identity() -> WeightArtifactIdentity: + return WeightArtifactIdentity( + schema_version=1, + source=WeightSourceIdentity("test-org/test-model", "0123456789abcdef", "source-files-sha256"), + component=ComponentIdentity("transformer", "diffusion.dit"), + representation=WeightRepresentation("runtime-bf16", "torch.bfloat16"), + layout=RuntimeWeightLayout("final-module-layout"), + adaptation=AdaptationIdentity(), + producer=ProducerIdentity( + "test.final-layout", + "1", + "implementation-sha256", + "test-producer-v1", + "test-restorer-v1", + ), + ) + + +class CountingProducer: + def __init__( + self, + identity: WeightArtifactIdentity, + *, + lookup_phase: LookupPhase = LookupPhase.PRE_LOAD_SAFE, + ) -> None: + self.calls = 0 + self._spec = WeightProductionSpec( + producer_id="test.final-layout", + outputs=(identity,), + source_mode=ProductionSourceMode.FINALIZED_MODEL, + lookup_phase=lookup_phase, + ) + + @property + def spec(self) -> WeightProductionSpec: + return self._spec + + def produce(self, writer: object) -> ProductionMetadata: + self.calls += 1 + spec = TensorWriteSpec("weight", (2, 2), torch.bfloat16) + with writer.open_tensor_file("weights.safetensors", (spec,)) as output: # type: ignore[attr-defined] + output.write_tensor("weight", torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)) + return ProductionMetadata( + producer_schema="test-producer-v1", + restorer_schema="test-restorer-v1", + format_metadata=CanonicalJson.from_value({"source": "test"}), + ) + + +def _domain(root: Path) -> StorageDomainPolicy: + return StorageDomainPolicy(root=root, storage_class=detect_storage_class(root.parent)) + + +def test_disabled_mode_does_not_probe_or_create_configured_root(tmp_path: Path) -> None: + root = tmp_path / "must-not-exist" + runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=RuntimeMode.DISABLED, domain=_domain(root))) + + resolution = runtime.resolve() + + assert resolution.report.outcome is ResolutionOutcome.CANONICAL_DIRECT + assert resolution.report.attempts == () + assert resolution.lease is None + assert not root.exists() + + +@pytest.mark.parametrize( + ("mode", "expected"), + [ + (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK), + (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED), + ], +) +def test_local_miss_uses_configured_fallback_policy( + tmp_path: Path, + mode: RuntimeMode, + expected: ResolutionOutcome, +) -> None: + runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=mode, domain=_domain(tmp_path / mode.value))) + + resolution = runtime.resolve(_identity()) + + assert resolution.report.outcome is expected + assert resolution.lease is None + assert [attempt.result.value for attempt in resolution.report.attempts] == ["miss", "skipped"] + expected_action = "canonical_fallback" if mode is RuntimeMode.PREFERRED else "fail_startup" + assert [attempt.action.value for attempt in resolution.report.attempts] == [expected_action, expected_action] + + +def test_local_production_then_exact_warm_hit_emits_one_terminal_report_each(tmp_path: Path) -> None: + reports: list[ResolutionReport] = [] + identity = _identity() + producer = CountingProducer(identity) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")), + observer=reports.append, + ) + + cold = runtime.resolve(identity, producer=producer) + assert cold.report.outcome is ResolutionOutcome.LOCAL_PRODUCTION + assert cold.lease is not None + assert torch.equal( + cold.lease.tensors["weight"], + torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2), + ) + cold.lease.close() + assert cold.report.attempts[0].action.value == "try_producer" + + warm = runtime.resolve(identity) + assert warm.report.outcome is ResolutionOutcome.LOCAL_HIT + assert warm.lease is not None + warm.lease.close() + + assert producer.calls == 1 + assert reports == [cold.report, warm.report] + assert all(report.resolution_id for report in reports) + + +@pytest.mark.parametrize( + ("mode", "expected"), + [ + (RuntimeMode.PREFERRED, ResolutionOutcome.FAILED), + (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED), + ], +) +def test_remote_required_is_explicitly_unsupported_and_skips_local_production( + tmp_path: Path, + mode: RuntimeMode, + expected: ResolutionOutcome, +) -> None: + identity = _identity() + producer = CountingProducer(identity) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=mode, + domain=_domain(tmp_path / mode.value), + remote=RemoteImportPolicy( + on_local_miss=RemoteOnMiss.REQUIRE, + providers=("future-mooncake",), + ), + ) + ) + + resolution = runtime.resolve(identity, producer=producer) + + assert resolution.report.outcome is expected + assert producer.calls == 0 + assert [attempt.stage.value for attempt in resolution.report.attempts] == ["lookup", "remote"] + assert resolution.report.attempts[-1].failure is not None + assert resolution.report.attempts[-1].failure.code.value == "unsupported" + assert resolution.report.attempts[-1].action.value == "fail_startup" + + +def test_observer_failure_cannot_change_resolution_outcome(tmp_path: Path) -> None: + def broken_observer(_report: object) -> None: + raise RuntimeError("metrics backend unavailable") + + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")), + observer=broken_observer, + ) + + resolution = runtime.resolve(_identity()) + + assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + + +def test_preload_resolution_does_not_run_postload_only_producer(tmp_path: Path) -> None: + identity = _identity() + producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")) + ) + + resolution = runtime.resolve(identity, producer=producer) + + assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + assert producer.calls == 0 + assert resolution.report.attempts[0].action.value == "canonical_fallback" + + +@pytest.mark.parametrize( + ("mode", "expected"), + [ + (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK), + (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED), + ], +) +def test_retryable_store_initialization_failure_obeys_runtime_mode( + tmp_path: Path, + mode: RuntimeMode, + expected: ResolutionOutcome, +) -> None: + root = tmp_path / "not-a-directory" + root.write_text("occupied", encoding="utf-8") + runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=mode, domain=_domain(root))) + + resolution = runtime.resolve(_identity()) + + assert resolution.report.outcome is expected + assert resolution.lease is None + assert len(resolution.report.attempts) == 1 + assert resolution.report.attempts[0].failure is not None + assert resolution.report.attempts[0].failure.code.value == "domain_unavailable" + assert resolution.report.attempts[0].failure.retryable + + +@pytest.mark.parametrize( + ("stage", "code"), + [ + (ResolutionStage.PRODUCTION, FailureCode.PRODUCER_UNSUPPORTED), + (ResolutionStage.VALIDATION, FailureCode.IDENTITY_COLLISION), + (ResolutionStage.LIFECYCLE, FailureCode.PUBLICATION_FAILED), + ], +) +def test_nonretryable_post_initialization_failure_cannot_fall_back_in_preferred_mode( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + stage: ResolutionStage, + code: FailureCode, +) -> None: + identity = _identity() + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / code.value)) + ) + assert runtime.store is not None + + def fail_production( + _request: object, + _producer: object, + *, + validation: object, + deadline: float, + ) -> StoreResult: + del validation, deadline + return StoreResult( + StoreStatus.FAILED, + failure=HostWeightFailure( + stage=stage, + code=code, + retryable=False, + message="injected nonretryable production failure", + details=CanonicalJson.empty(), + ), + ) + + monkeypatch.setattr(runtime.store, "get_or_build", fail_production) + resolution = runtime.resolve(identity, producer=CountingProducer(identity)) + + assert resolution.report.outcome is ResolutionOutcome.FAILED + assert resolution.report.attempts[-1].action.value == "fail_startup" + assert resolution.report.attempts[-1].failure is not None + assert resolution.report.attempts[-1].failure.code is code + + +def test_retryable_post_initialization_failure_can_fall_back_in_preferred_mode( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + identity = _identity() + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "retryable")) + ) + assert runtime.store is not None + + def fail_production( + _request: object, + _producer: object, + *, + validation: object, + deadline: float, + ) -> StoreResult: + del validation, deadline + return StoreResult( + StoreStatus.FAILED, + failure=HostWeightFailure( + stage=ResolutionStage.CAPACITY, + code=FailureCode.ENOSPC, + retryable=True, + message="injected recoverable capacity failure", + details=CanonicalJson.empty(), + ), + ) + + monkeypatch.setattr(runtime.store, "get_or_build", fail_production) + resolution = runtime.resolve(identity, producer=CountingProducer(identity)) + + assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + assert resolution.report.attempts[-1].action.value == "canonical_fallback" + + +def test_nonretryable_store_configuration_failure_remains_visible_in_preferred_mode(tmp_path: Path) -> None: + root = tmp_path / "untrusted" + root.mkdir(mode=0o700) + root.chmod(0o777) + runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(root))) + + resolution = runtime.resolve(_identity()) + + assert resolution.report.outcome is ResolutionOutcome.FAILED + assert resolution.report.attempts[0].action.value == "fail_startup" + assert resolution.report.attempts[0].failure is not None + assert not resolution.report.attempts[0].failure.retryable + + +def test_unsupported_integrity_policy_fails_instead_of_disguising_configuration_error(tmp_path: Path) -> None: + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.PREFERRED, + domain=_domain(tmp_path / "store"), + integrity=IntegrityPolicy(local_lookup=ValidationLevel.FS_VERITY), + ) + ) + + resolution = runtime.resolve(_identity()) + + assert resolution.report.outcome is ResolutionOutcome.FAILED + assert resolution.report.attempts[0].failure is not None + assert resolution.report.attempts[0].failure.code.value == "unsupported" diff --git a/vllm_omni/host_weight_runtime/__init__.py b/vllm_omni/host_weight_runtime/__init__.py new file mode 100644 index 000000000..c7bbf51dc --- /dev/null +++ b/vllm_omni/host_weight_runtime/__init__.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Neutral contracts for immutable runtime-ready host weight artifacts.""" + +from .config import ( + CapacityPolicy, + HostWeightRuntimeConfig, + IntegrityPolicy, + ProductionPolicy, + RemoteImportPolicy, + RemoteOnMiss, + RuntimeMode, + StorageClass, + StorageDomainPolicy, + StorageScope, + ValidationLevel, + WaitPolicy, +) +from .errors import FailureCode, HostWeightError, HostWeightFailure, ResolutionAction, ResolutionStage +from .identity import ( + AdaptationIdentity, + CanonicalJson, + ComponentIdentity, + IdentityValidationError, + ProducerIdentity, + RuntimeWeightLayout, + WeightArtifactIdentity, + WeightRepresentation, + WeightSourceIdentity, +) +from .inspection import ( + ArtifactInspection, + ArtifactInventoryEntry, + ArtifactInventoryState, + DomainInspection, + HostWeightInspector, +) +from .lease import HostWeightLease, LeaseProvenance, MappedHostRegion +from .manifest import ( + FileManifestEntry, + ManifestValidationError, + PublicationMarker, + TensorKind, + TensorManifestEntry, + WeightManifest, +) +from .outcomes import ( + AttemptResult, + HostWeightResolution, + ResolutionAttempt, + ResolutionOutcome, + ResolutionReport, + StoreResult, + StoreStatus, +) +from .protocols import ( + ArtifactWriter, + BuildRequest, + CoordinationScope, + HostWeightStore, + LookupPhase, + ProductionMetadata, + ProductionSourceMode, + RemoteArtifactDescriptor, + RemoteArtifactProvider, + TensorFileWriter, + TensorWriteSpec, + WeightProducer, + WeightProductionSpec, + WeightRestorePlan, + WeightRestorer, +) +from .runtime import HostWeightRuntime + +__all__ = [ + "AdaptationIdentity", + "ArtifactInspection", + "ArtifactInventoryEntry", + "ArtifactInventoryState", + "ArtifactWriter", + "AttemptResult", + "BuildRequest", + "CanonicalJson", + "CapacityPolicy", + "ComponentIdentity", + "CoordinationScope", + "DomainInspection", + "FailureCode", + "FileManifestEntry", + "HostWeightError", + "HostWeightFailure", + "HostWeightLease", + "HostWeightResolution", + "HostWeightInspector", + "HostWeightRuntime", + "HostWeightRuntimeConfig", + "HostWeightStore", + "IdentityValidationError", + "IntegrityPolicy", + "LeaseProvenance", + "LookupPhase", + "ManifestValidationError", + "MappedHostRegion", + "ProducerIdentity", + "ProductionMetadata", + "ProductionPolicy", + "ProductionSourceMode", + "PublicationMarker", + "RemoteArtifactDescriptor", + "RemoteArtifactProvider", + "RemoteImportPolicy", + "RemoteOnMiss", + "ResolutionAction", + "ResolutionAttempt", + "ResolutionOutcome", + "ResolutionReport", + "ResolutionStage", + "RuntimeMode", + "RuntimeWeightLayout", + "StorageClass", + "StorageDomainPolicy", + "StorageScope", + "StoreResult", + "StoreStatus", + "TensorFileWriter", + "TensorKind", + "TensorManifestEntry", + "TensorWriteSpec", + "ValidationLevel", + "WaitPolicy", + "WeightArtifactIdentity", + "WeightManifest", + "WeightProducer", + "WeightProductionSpec", + "WeightRepresentation", + "WeightRestorePlan", + "WeightRestorer", + "WeightSourceIdentity", +] diff --git a/vllm_omni/host_weight_runtime/config.py b/vllm_omni/host_weight_runtime/config.py new file mode 100644 index 000000000..f2422899e --- /dev/null +++ b/vllm_omni/host_weight_runtime/config.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Typed policy for the loader-adjacent Host Weight Runtime.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path + + +class RuntimeMode(str, Enum): + DISABLED = "disabled" + PREFERRED = "preferred" + REQUIRED = "required" + + +class StorageScope(str, Enum): + NODE = "node" + NUMA = "numa" + + +class StorageClass(str, Enum): + DISK = "disk" + TMPFS = "tmpfs" + + +class ValidationLevel(str, Enum): + MANIFEST_AND_METADATA = "manifest_and_metadata" + FULL_CHECKSUM = "full_checksum" + FS_VERITY = "fs_verity" + + +class RemoteOnMiss(str, Enum): + DISABLED = "disabled" + TRY = "try" + REQUIRE = "require" + + +@dataclass(frozen=True) +class StorageDomainPolicy: + root: Path + scope: StorageScope = StorageScope.NODE + domain_id: str = "node" + storage_class: StorageClass = StorageClass.DISK + numa_prefault: bool = False + verify_numa_residency: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.root, Path) or not str(self.root): + raise ValueError("host weight store root must not be empty") + if not isinstance(self.scope, StorageScope) or not isinstance(self.storage_class, StorageClass): + raise ValueError("storage scope and class must use their typed enum values") + if not isinstance(self.domain_id, str) or not self.domain_id: + raise ValueError("host weight store domain_id must not be empty") + if not isinstance(self.numa_prefault, bool) or not isinstance(self.verify_numa_residency, bool): + raise ValueError("NUMA placement options must be booleans") + if self.scope is StorageScope.NODE and (self.numa_prefault or self.verify_numa_residency): + raise ValueError("NUMA placement options require scope='numa'") + + +@dataclass(frozen=True) +class CapacityPolicy: + max_artifact_bytes: int | None = None + max_store_bytes: int | None = None + min_free_bytes: int = 0 + eviction: str = "none" + + def __post_init__(self) -> None: + for name, value in ( + ("max_artifact_bytes", self.max_artifact_bytes), + ("max_store_bytes", self.max_store_bytes), + ): + if value is not None and (isinstance(value, bool) or not isinstance(value, int) or value <= 0): + raise ValueError(f"{name} must be positive when configured") + if isinstance(self.min_free_bytes, bool) or not isinstance(self.min_free_bytes, int) or self.min_free_bytes < 0: + raise ValueError("min_free_bytes must not be negative") + if not isinstance(self.eviction, str) or self.eviction != "none": + raise ValueError("automatic host weight eviction is not implemented") + + +@dataclass(frozen=True) +class IntegrityPolicy: + local_lookup: ValidationLevel = ValidationLevel.MANIFEST_AND_METADATA + require_trusted_permissions: bool = True + background_scrub: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.local_lookup, ValidationLevel): + raise ValueError("local lookup validation must use ValidationLevel") + if not isinstance(self.require_trusted_permissions, bool) or not isinstance(self.background_scrub, bool): + raise ValueError("integrity switches must be booleans") + + +@dataclass(frozen=True) +class RemoteImportPolicy: + on_local_miss: RemoteOnMiss = RemoteOnMiss.DISABLED + providers: tuple[str, ...] = () + + def __post_init__(self) -> None: + if not isinstance(self.on_local_miss, RemoteOnMiss): + raise ValueError("remote on-miss policy must use RemoteOnMiss") + if not isinstance(self.providers, tuple): + raise ValueError("remote providers must be an immutable tuple") + if self.on_local_miss is not RemoteOnMiss.DISABLED and not self.providers: + raise ValueError("remote providers are required when remote import is enabled") + if any(not isinstance(provider, str) or not provider for provider in self.providers): + raise ValueError("remote provider identifiers must not be empty") + + +@dataclass(frozen=True) +class ProductionPolicy: + allow_local_build: bool = True + allow_post_load_publish: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.allow_local_build, bool) or not isinstance(self.allow_post_load_publish, bool): + raise ValueError("production switches must be booleans") + if self.allow_post_load_publish: + raise ValueError("post-load publication is not implemented") + + +@dataclass(frozen=True) +class WaitPolicy: + coordination_timeout_seconds: float = 600.0 + + def __post_init__(self) -> None: + timeout = self.coordination_timeout_seconds + if ( + isinstance(timeout, bool) + or not isinstance(timeout, int | float) + or not math.isfinite(timeout) + or timeout <= 0 + ): + raise ValueError("coordination_timeout_seconds must be positive") + + +@dataclass(frozen=True) +class HostWeightRuntimeConfig: + mode: RuntimeMode = RuntimeMode.DISABLED + domain: StorageDomainPolicy | None = None + remote: RemoteImportPolicy = field(default_factory=RemoteImportPolicy) + production: ProductionPolicy = field(default_factory=ProductionPolicy) + capacity: CapacityPolicy = field(default_factory=CapacityPolicy) + wait: WaitPolicy = field(default_factory=WaitPolicy) + integrity: IntegrityPolicy = field(default_factory=IntegrityPolicy) + schema_version: int = 1 + + def __post_init__(self) -> None: + if isinstance(self.schema_version, bool) or self.schema_version != 1: + raise ValueError("only Host Weight Runtime config schema_version=1 is supported") + if not isinstance(self.mode, RuntimeMode): + raise ValueError("runtime mode must use RuntimeMode") + if self.domain is not None and not isinstance(self.domain, StorageDomainPolicy): + raise ValueError("runtime domain must use StorageDomainPolicy") + for name, value, expected in ( + ("remote", self.remote, RemoteImportPolicy), + ("production", self.production, ProductionPolicy), + ("capacity", self.capacity, CapacityPolicy), + ("wait", self.wait, WaitPolicy), + ("integrity", self.integrity, IntegrityPolicy), + ): + if not isinstance(value, expected): + raise ValueError(f"runtime {name} policy has the wrong type") + if self.mode is not RuntimeMode.DISABLED and self.domain is None: + raise ValueError("enabled Host Weight Runtime requires a storage domain") + + +__all__ = [ + "CapacityPolicy", + "HostWeightRuntimeConfig", + "IntegrityPolicy", + "ProductionPolicy", + "RemoteImportPolicy", + "RemoteOnMiss", + "RuntimeMode", + "StorageClass", + "StorageDomainPolicy", + "StorageScope", + "ValidationLevel", + "WaitPolicy", +] diff --git a/vllm_omni/host_weight_runtime/errors.py b/vllm_omni/host_weight_runtime/errors.py new file mode 100644 index 000000000..d3ab4f3c2 --- /dev/null +++ b/vllm_omni/host_weight_runtime/errors.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stable Host Weight Runtime failures and actions.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum + +from .identity import CanonicalJson + + +class ResolutionStage(str, Enum): + CONFIGURATION = "configuration" + DOMAIN = "domain" + LOOKUP = "lookup" + VALIDATION = "validation" + REMOTE = "remote" + PRODUCTION = "production" + CAPACITY = "capacity" + RESTORATION = "restoration" + LIFECYCLE = "lifecycle" + CANONICAL_LOADING = "canonical_loading" + + +class FailureCode(str, Enum): + INVALID_CONFIG = "invalid_config" + COHORT_POLICY_MISMATCH = "cohort_policy_mismatch" + DOMAIN_POLICY_MISMATCH = "domain_policy_mismatch" + DOMAIN_UNAVAILABLE = "domain_unavailable" + STORAGE_CLASS_MISMATCH = "storage_class_mismatch" + NUMA_PLACEMENT_INVALID = "numa_placement_invalid" + LOCAL_MISS = "local_miss" + ARTIFACT_DENIED = "artifact_denied" + ACTIVE_BUILD_TIMEOUT = "active_build_timeout" + MANIFEST_INCOMPATIBLE = "manifest_incompatible" + MANIFEST_CORRUPT = "manifest_corrupt" + CHECKSUM_MISMATCH = "checksum_mismatch" + IDENTITY_COLLISION = "identity_collision" + REMOTE_MISS = "remote_miss" + REMOTE_UNAVAILABLE = "remote_unavailable" + REMOTE_AUTH_FAILED = "remote_auth_failed" + REMOTE_TIMEOUT = "remote_timeout" + REMOTE_ARTIFACT_INVALID = "remote_artifact_invalid" + PRODUCER_UNAVAILABLE = "producer_unavailable" + PRODUCER_UNSUPPORTED = "producer_unsupported" + PRODUCER_FAILED = "producer_failed" + PRODUCTION_TIMEOUT = "production_timeout" + ARTIFACT_TOO_LARGE = "artifact_too_large" + STORE_LIMIT_EXCEEDED = "store_limit_exceeded" + INSUFFICIENT_CAPACITY = "insufficient_capacity" + ENOSPC = "enospc" + RESTORER_UNAVAILABLE = "restorer_unavailable" + RESTORE_PREFLIGHT_FAILED = "restore_preflight_failed" + RESTORE_COMMIT_FAILED = "restore_commit_failed" + ACTIVE_LEASE = "active_lease" + PUBLICATION_FAILED = "publication_failed" + QUARANTINE_FAILED = "quarantine_failed" + CANONICAL_SOURCE_FAILED = "canonical_source_failed" + UNSUPPORTED = "unsupported" + + +class ResolutionAction(str, Enum): + RETURN_LEASE = "return_lease" + TRY_REMOTE = "try_remote" + TRY_PRODUCER = "try_producer" + QUARANTINE_THEN_PRODUCE = "quarantine_then_produce" + CANONICAL_FALLBACK = "canonical_fallback" + FAIL_STARTUP = "fail_startup" + SKIP = "skip" + + +@dataclass(frozen=True) +class HostWeightFailure: + stage: ResolutionStage + code: FailureCode + retryable: bool + message: str + details: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + if not isinstance(self.stage, ResolutionStage) or not isinstance(self.code, FailureCode): + raise ValueError("host weight failure stage and code must use their typed enums") + if not isinstance(self.retryable, bool): + raise ValueError("host weight failure retryable must be a boolean") + if not isinstance(self.message, str) or not self.message: + raise ValueError("host weight failure requires a message") + if not isinstance(self.details, CanonicalJson): + raise ValueError("host weight failure details must use CanonicalJson") + + +class HostWeightError(RuntimeError): + """Internal exception carrying a stable, policy-neutral failure.""" + + def __init__(self, failure: HostWeightFailure): + super().__init__(failure.message) + self.failure = failure + + +__all__ = [ + "FailureCode", + "HostWeightError", + "HostWeightFailure", + "ResolutionAction", + "ResolutionStage", +] diff --git a/vllm_omni/host_weight_runtime/filesystem/__init__.py b/vllm_omni/host_weight_runtime/filesystem/__init__.py new file mode 100644 index 000000000..2083c1e85 --- /dev/null +++ b/vllm_omni/host_weight_runtime/filesystem/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Local filesystem Host Weight Store implementation.""" + +from .store import FilesystemHostWeightStore, detect_filesystem_type, detect_storage_class + +__all__ = [ + "FilesystemHostWeightStore", + "detect_filesystem_type", + "detect_storage_class", +] diff --git a/vllm_omni/host_weight_runtime/filesystem/locks.py b/vllm_omni/host_weight_runtime/filesystem/locks.py new file mode 100644 index 000000000..70a400dac --- /dev/null +++ b/vllm_omni/host_weight_runtime/filesystem/locks.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stable POSIX advisory locks used by the filesystem store.""" + +from __future__ import annotations + +import errno +import fcntl +import os +import time +from pathlib import Path + + +class FileLockTimeoutError(TimeoutError): + pass + + +class FileLock: + """One held flock whose inode is stable outside artifact directories.""" + + def __init__(self, path: Path, *, exclusive: bool, deadline: float | None, nonblocking: bool = False) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + self.path = path + self._fd = os.open(path, os.O_CREAT | os.O_RDWR | os.O_CLOEXEC | os.O_NOFOLLOW, 0o600) + self._creator_pid = os.getpid() + self._closed = False + operation = fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH + if nonblocking: + operation |= fcntl.LOCK_NB + try: + if nonblocking: + self._acquire_once(operation) + elif deadline is None: + self._acquire_blocking(operation) + else: + self._acquire_until(operation | fcntl.LOCK_NB, deadline) + except BaseException: + os.close(self._fd) + self._closed = True + raise + + def _acquire_once(self, operation: int) -> None: + try: + fcntl.flock(self._fd, operation) + except OSError as exc: + if exc.errno in (errno.EACCES, errno.EAGAIN): + raise FileLockTimeoutError(f"lock is active: {self.path}") from exc + raise + + def _acquire_blocking(self, operation: int) -> None: + while True: + try: + fcntl.flock(self._fd, operation) + return + except OSError as exc: + if exc.errno != errno.EINTR: + raise + + def _acquire_until(self, operation: int, deadline: float) -> None: + while True: + try: + fcntl.flock(self._fd, operation) + return + except OSError as exc: + if exc.errno == errno.EINTR: + continue + if exc.errno not in (errno.EACCES, errno.EAGAIN): + raise + remaining = deadline - time.monotonic() + if remaining <= 0: + raise FileLockTimeoutError(f"timed out waiting for {self.path}") from exc + time.sleep(min(0.05, remaining)) + + def close(self) -> None: + if self._closed: + return + self._closed = True + if os.getpid() != self._creator_pid: + # flock state belongs to the inherited open-file description. A + # child must only drop its descriptor reference; LOCK_UN would + # release the creator process's active lock as well. + os.close(self._fd) + return + try: + fcntl.flock(self._fd, fcntl.LOCK_UN) + finally: + os.close(self._fd) + + def __enter__(self) -> FileLock: + return self + + def __exit__(self, exc_type: object, exc: object, _traceback: object) -> None: + self.close() + + +def lock_is_active(path: Path) -> bool: + """Return a point-in-time contention snapshot without creating a lock file.""" + try: + fd = os.open(path, os.O_RDWR | os.O_CLOEXEC | os.O_NOFOLLOW) + except FileNotFoundError: + return False + try: + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as exc: + if exc.errno in (errno.EACCES, errno.EAGAIN): + return True + raise + fcntl.flock(fd, fcntl.LOCK_UN) + return False + finally: + os.close(fd) + + +__all__ = ["FileLock", "FileLockTimeoutError", "lock_is_active"] diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py new file mode 100644 index 000000000..9c81d2972 --- /dev/null +++ b/vllm_omni/host_weight_runtime/filesystem/store.py @@ -0,0 +1,1189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Crash-safe node-local filesystem implementation of HostWeightStore.""" + +from __future__ import annotations + +import contextlib +import errno +import hashlib +import json +import os +import shutil +import stat +import time +import uuid +from collections import defaultdict +from collections.abc import Callable +from pathlib import Path +from typing import Protocol + +import regex as re +import torch +from safetensors import SafetensorError, safe_open + +from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel +from ..errors import FailureCode, HostWeightError, HostWeightFailure, ResolutionStage +from ..identity import CanonicalJson, WeightArtifactIdentity, canonical_json +from ..inspection import ( + ArtifactInspection, + ArtifactInventoryEntry, + ArtifactInventoryState, + DomainInspection, +) +from ..lease import HostWeightLease, LeaseProvenance, MappedHostRegion +from ..manifest import ManifestValidationError, PublicationMarker, TensorManifestEntry, WeightManifest +from ..outcomes import StoreResult, StoreStatus +from ..protocols import BuildRequest, CoordinationScope, WeightProducer +from .locks import FileLock, FileLockTimeoutError, lock_is_active +from .writer import FilesystemArtifactWriter + +_DOMAIN_FILE = "domain.json" +_DOMAIN_POLICY_FILE = "domain-policy.json" +_MANIFEST_FILE = "manifest.json" +_READY_FILE = "READY.json" +_MAX_METADATA_BYTES = 64 * 1024**2 +_FILE_HASH_CHUNK_BYTES = 8 * 1024**2 +_ARTIFACT_KEY_RE = re.compile(r"[0-9a-f]{64}") +_LOCAL_DISK_FILESYSTEM_TYPES = frozenset( + { + "aufs", + "bcachefs", + "btrfs", + "erofs", + "ext2", + "ext3", + "ext4", + "f2fs", + "fuseblk", + "jfs", + "ntfs3", + "overlay", + "reiserfs", + "rootfs", + "xfs", + "zfs", + } +) +_REMOTE_FILESYSTEM_TYPES = frozenset( + { + "9p", + "afs", + "beegfs", + "ceph", + "cifs", + "fuse.gcsfuse", + "fuse.s3fs", + "fuse.sshfs", + "gcsfuse", + "glusterfs", + "gpfs", + "lustre", + "nfs", + "nfs4", + "orangefs", + "panfs", + "s3fs", + "smb3", + "sshfs", + "virtiofs", + "wekafs", + } +) + + +class _ExitContext(Protocol): + def __exit__(self, exc_type: None, exc: None, _traceback: None) -> object: ... + + +def _make_fd_closer(fd: int) -> Callable[[], None]: + def close() -> None: + os.close(fd) + + return close + + +def _make_context_closer(context: _ExitContext) -> Callable[[], None]: + def close() -> None: + context.__exit__(None, None, None) + + return close + + +def _failure( + stage: ResolutionStage, + code: FailureCode, + message: str, + *, + retryable: bool = False, + details: object | None = None, +) -> HostWeightFailure: + return HostWeightFailure( + stage=stage, + code=code, + retryable=retryable, + message=message, + details=CanonicalJson.from_value({} if details is None else details), + ) + + +def _fsync_directory(path: Path) -> None: + fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def _write_atomic_json(path: Path, value: object, *, mode: int = 0o444) -> None: + data = canonical_json(value) + temporary = path.parent / f".{path.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp" + fd = os.open(temporary, os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_CLOEXEC, 0o600) + try: + offset = 0 + while offset < len(data): + count = os.write(fd, data[offset:]) + if count <= 0: + raise OSError(errno.EIO, f"short write for {temporary}") + offset += count + os.fsync(fd) + os.fchmod(fd, mode) + os.fsync(fd) + except BaseException: + os.close(fd) + temporary.unlink(missing_ok=True) + raise + else: + os.close(fd) + os.replace(temporary, path) + _fsync_directory(path.parent) + + +def _read_json_file(path: Path) -> dict[str, object]: + fd = os.open(path, os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW) + try: + return _read_json_fd(fd, path.name) + finally: + os.close(fd) + + +def _read_json_at(directory_fd: int, name: str, *, require_read_only: bool = False) -> dict[str, object]: + fd = os.open(name, os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, dir_fd=directory_fd) + try: + return _read_json_fd(fd, name, require_read_only=require_read_only) + finally: + os.close(fd) + + +def _read_json_fd(fd: int, name: str, *, require_read_only: bool = False) -> dict[str, object]: + file_stat = os.fstat(fd) + if not stat.S_ISREG(file_stat.st_mode) or file_stat.st_nlink != 1: + raise ManifestValidationError(f"metadata file {name!r} is not one regular file") + if file_stat.st_size > _MAX_METADATA_BYTES: + raise ManifestValidationError(f"metadata file {name!r} exceeds {_MAX_METADATA_BYTES} bytes") + if require_read_only and (file_stat.st_uid != os.geteuid() or file_stat.st_mode & 0o222): + raise ManifestValidationError(f"published metadata file {name!r} is not immutable") + chunks: list[bytes] = [] + offset = 0 + while offset < file_stat.st_size: + chunk = os.pread(fd, min(1024**2, file_stat.st_size - offset), offset) + if not chunk: + raise ManifestValidationError(f"metadata file {name!r} ended unexpectedly") + chunks.append(chunk) + offset += len(chunk) + try: + value = json.loads(b"".join(chunks)) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ManifestValidationError(f"metadata file {name!r} is not valid JSON") from exc + if not isinstance(value, dict): + raise ManifestValidationError(f"metadata file {name!r} must contain an object") + return value + + +def _sha256_fd(fd: int, size: int) -> str: + digest = hashlib.sha256() + offset = 0 + while offset < size: + chunk = os.pread(fd, min(_FILE_HASH_CHUNK_BYTES, size - offset), offset) + if not chunk: + raise OSError(errno.EIO, "artifact file ended while hashing") + digest.update(chunk) + offset += len(chunk) + return digest.hexdigest() + + +def _mount_unescape(value: str) -> str: + return value.replace("\\040", " ").replace("\\011", "\t").replace("\\012", "\n").replace("\\134", "\\") + + +def detect_filesystem_type(path: Path) -> str: + """Return the kernel filesystem type for the most specific mount.""" + resolved = path.resolve() + mountinfo = Path("/proc/self/mountinfo") + if not mountinfo.is_file(): # pragma: no cover - vLLM-Omni production is Linux + raise ValueError("cannot verify filesystem locality without /proc/self/mountinfo") + best_length = -1 + best_type = "" + for line in mountinfo.read_text(encoding="utf-8").splitlines(): + left, separator, right = line.partition(" - ") + if not separator: + continue + left_fields = left.split() + right_fields = right.split() + if len(left_fields) < 5 or not right_fields: + continue + mount_point = Path(_mount_unescape(left_fields[4])) + try: + resolved.relative_to(mount_point) + except ValueError: + continue + if len(str(mount_point)) > best_length: + best_length = len(str(mount_point)) + best_type = right_fields[0] + if not best_type: + raise ValueError(f"cannot identify the filesystem containing {resolved}") + return best_type + + +def _storage_class_for_filesystem_type(filesystem_type: str) -> StorageClass: + if filesystem_type in {"tmpfs", "ramfs"}: + return StorageClass.TMPFS + if filesystem_type in _LOCAL_DISK_FILESYSTEM_TYPES: + return StorageClass.DISK + if filesystem_type in _REMOTE_FILESYSTEM_TYPES: + raise ValueError(f"filesystem type {filesystem_type!r} is not node-local") + raise ValueError(f"filesystem type {filesystem_type!r} is not recognized as a local filesystem") + + +def detect_storage_class(path: Path) -> StorageClass: + """Classify a verified node-local mount without spawning platform tools.""" + return _storage_class_for_filesystem_type(detect_filesystem_type(path)) + + +class FilesystemHostWeightStore: + """A cooperative local store with stable locks and immutable publication.""" + + def __init__( + self, + domain: StorageDomainPolicy, + capacity: CapacityPolicy, + integrity: IntegrityPolicy, + ) -> None: + self.domain_policy = domain + self.capacity_policy = capacity + self.integrity_policy = integrity + self.root = domain.root.expanduser().resolve() + self.locks_dir = self.root / "locks" + self.artifacts_dir = self.root / "artifacts" + self.tmp_dir = self.root / "tmp" + self.quarantine_dir = self.root / "quarantine" + self.deny_dir = self.root / "deny" + if integrity.local_lookup is ValidationLevel.FS_VERITY: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.UNSUPPORTED, + "filesystem-verity validation is not implemented", + ) + ) + if domain.numa_prefault or domain.verify_numa_residency: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.UNSUPPORTED, + "enforced NUMA placement and residency verification are not implemented", + ) + ) + if integrity.background_scrub: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.UNSUPPORTED, + "background integrity scrubbing is not implemented", + ) + ) + try: + self._initialize_domain() + except HostWeightError: + raise + except OSError as exc: + raise HostWeightError( + _failure( + ResolutionStage.DOMAIN, + FailureCode.DOMAIN_UNAVAILABLE, + f"cannot initialize storage domain {self.root}: {exc}", + retryable=True, + ) + ) from exc + except (ValueError, ManifestValidationError) as exc: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"storage domain metadata is invalid under {self.root}: {exc}", + ) + ) from exc + + def _initialize_domain(self) -> None: + try: + self.filesystem_type = detect_filesystem_type(self.root) + detected_storage_class = _storage_class_for_filesystem_type(self.filesystem_type) + except ValueError as exc: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"local filesystem storage is unavailable for {self.root}: {exc}", + ) + ) from exc + if detected_storage_class is not self.domain_policy.storage_class: + raise HostWeightError( + _failure( + ResolutionStage.DOMAIN, + FailureCode.STORAGE_CLASS_MISMATCH, + "configured storage class " + f"{self.domain_policy.storage_class.value} differs from detected {detected_storage_class.value}", + ) + ) + + try: + self.root.mkdir(parents=True, exist_ok=True, mode=0o700) + self.locks_dir.mkdir(exist_ok=True, mode=0o700) + for directory in (self.artifacts_dir, self.tmp_dir, self.quarantine_dir, self.deny_dir): + directory.mkdir(exist_ok=True, mode=0o700) + except OSError as exc: + raise HostWeightError( + _failure( + ResolutionStage.DOMAIN, + FailureCode.DOMAIN_UNAVAILABLE, + f"cannot initialize {self.root}: {exc}", + retryable=True, + ) + ) from exc + + if self.integrity_policy.require_trusted_permissions: + for directory in ( + self.root, + self.locks_dir, + self.artifacts_dir, + self.tmp_dir, + self.quarantine_dir, + self.deny_dir, + ): + directory_stat = directory.lstat() + if ( + not stat.S_ISDIR(directory_stat.st_mode) + or directory_stat.st_uid != os.geteuid() + or directory_stat.st_mode & 0o022 + ): + raise HostWeightError( + _failure( + ResolutionStage.DOMAIN, + FailureCode.DOMAIN_UNAVAILABLE, + f"host weight store directory is not exclusively controlled by this user: {directory}", + ) + ) + + with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=None): + domain_path = self.root / _DOMAIN_FILE + if domain_path.exists(): + domain_metadata = _read_json_file(domain_path) + expected = { + "schema_version": 1, + "scope": self.domain_policy.scope.value, + "domain_id": self.domain_policy.domain_id, + "backend": "local_filesystem", + "storage_class": self.domain_policy.storage_class.value, + "filesystem_type": self.filesystem_type, + } + if set(domain_metadata) != {*expected, "domain_uuid"}: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"storage domain metadata has an incompatible schema under {self.root}", + ) + ) + actual = {key: domain_metadata.get(key) for key in expected} + domain_uuid = domain_metadata.get("domain_uuid") + if actual != expected or not isinstance(domain_uuid, str) or not domain_uuid: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"storage domain metadata does not match configured policy under {self.root}", + ) + ) + else: + domain_metadata = { + "schema_version": 1, + "domain_uuid": uuid.uuid4().hex, + "scope": self.domain_policy.scope.value, + "domain_id": self.domain_policy.domain_id, + "backend": "local_filesystem", + "storage_class": self.domain_policy.storage_class.value, + "filesystem_type": self.filesystem_type, + } + _write_atomic_json(domain_path, domain_metadata) + self.domain_uuid = str(domain_metadata["domain_uuid"]) + + policy_path = self.root / _DOMAIN_POLICY_FILE + requested_policy = self._capacity_document(policy_version=1) + if policy_path.exists(): + actual_policy = _read_json_file(policy_path) + if set(actual_policy) != set(requested_policy): + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"capacity policy has an incompatible schema under {self.root}", + ) + ) + comparable_actual = {key: actual_policy.get(key) for key in requested_policy if key != "policy_version"} + comparable_requested = { + key: value for key, value in requested_policy.items() if key != "policy_version" + } + if comparable_actual != comparable_requested: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"capacity policy differs from the authoritative policy under {self.root}", + ) + ) + policy_version = actual_policy.get("policy_version") + if isinstance(policy_version, bool) or not isinstance(policy_version, int) or policy_version < 1: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"capacity policy has an invalid version under {self.root}", + ) + ) + self.policy_version = policy_version + else: + _write_atomic_json(policy_path, requested_policy) + self.policy_version = 1 + + def _capacity_document(self, *, policy_version: int) -> dict[str, object]: + return { + "schema_version": 1, + "policy_version": policy_version, + "max_artifact_bytes": self.capacity_policy.max_artifact_bytes, + "max_store_bytes": self.capacity_policy.max_store_bytes, + "min_free_bytes": self.capacity_policy.min_free_bytes, + "eviction": self.capacity_policy.eviction, + } + + def _build_lock_path(self, key: str) -> Path: + return self.locks_dir / f"{key}.build.lock" + + def _artifact_lock_path(self, key: str) -> Path: + return self.locks_dir / f"{key}.artifact.lock" + + def _entry_path(self, key: str) -> Path: + return self.artifacts_dir / key + + def _deny_path(self, key: str) -> Path: + return self.deny_dir / f"{key}.json" + + def lookup( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + deadline: float | None = None, + ) -> StoreResult: + return self._lookup(identity, validation=validation, deadline=deadline, record_invalid=True) + + def _lookup( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + deadline: float | None, + record_invalid: bool, + ) -> StoreResult: + if validation is ValidationLevel.FS_VERITY: + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.VALIDATION, + FailureCode.UNSUPPORTED, + "filesystem-verity validation is modeled but not implemented", + ), + ) + try: + artifact_lock = FileLock(self._artifact_lock_path(identity.key), exclusive=False, deadline=deadline) + except FileLockTimeoutError as exc: + return StoreResult( + StoreStatus.TIMEOUT, + failure=_failure( + ResolutionStage.LOOKUP, + FailureCode.ACTIVE_BUILD_TIMEOUT, + str(exc), + retryable=True, + ), + ) + except OSError as exc: + return StoreResult( + StoreStatus.FAILED, + failure=_failure(ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE, str(exc), retryable=True), + ) + + entry_path = self._entry_path(identity.key) + if not entry_path.exists(): + artifact_lock.close() + return StoreResult(StoreStatus.MISS) + if self._deny_path(identity.key).exists(): + artifact_lock.close() + return StoreResult( + StoreStatus.INVALID, + failure=_failure( + ResolutionStage.LOOKUP, + FailureCode.ARTIFACT_DENIED, + f"artifact {identity.key} is denied pending quarantine", + ), + ) + try: + lease = self._open_lease(identity, entry_path, artifact_lock, validation) + if self._deny_path(identity.key).exists(): + lease.close() + return StoreResult( + StoreStatus.INVALID, + failure=_failure( + ResolutionStage.LOOKUP, + FailureCode.ARTIFACT_DENIED, + f"artifact {identity.key} was denied during lookup", + ), + ) + return StoreResult(StoreStatus.HIT, lease=lease) + except HostWeightError as exc: + try: + if record_invalid: + with contextlib.suppress(OSError): + self._write_deny(identity.key, exc.failure) + finally: + artifact_lock.close() + return StoreResult(StoreStatus.INVALID, failure=exc.failure) + except (OSError, SafetensorError, ValueError, ManifestValidationError) as exc: + failure = _failure( + ResolutionStage.VALIDATION, + FailureCode.MANIFEST_CORRUPT, + f"artifact {identity.key} failed validation: {exc}", + ) + try: + if record_invalid: + with contextlib.suppress(OSError): + self._write_deny(identity.key, failure) + finally: + artifact_lock.close() + return StoreResult(StoreStatus.INVALID, failure=failure) + + def _open_lease( + self, + identity: WeightArtifactIdentity, + entry_path: Path, + artifact_lock: FileLock, + validation: ValidationLevel, + ) -> HostWeightLease: + resources: list[Callable[[], None]] = [] + try: + entry_fd = os.open(entry_path, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW) + resources.append(_make_fd_closer(entry_fd)) + entry_stat = os.fstat(entry_fd) + if self.integrity_policy.require_trusted_permissions and ( + entry_stat.st_uid != os.geteuid() or entry_stat.st_mode & 0o222 + ): + raise ManifestValidationError("published artifact directory is not immutable") + names = set(os.listdir(entry_fd)) + marker_document = _read_json_at( + entry_fd, + _READY_FILE, + require_read_only=self.integrity_policy.require_trusted_permissions, + ) + manifest_document = _read_json_at( + entry_fd, + _MANIFEST_FILE, + require_read_only=self.integrity_policy.require_trusted_permissions, + ) + marker = PublicationMarker.from_dict(marker_document) + manifest = WeightManifest.from_dict(manifest_document) + if marker.schema_version != 1 or manifest.schema_version != 1 or marker.integrity_mode != "full_checksum": + raise HostWeightError( + _failure( + ResolutionStage.VALIDATION, + FailureCode.MANIFEST_INCOMPATIBLE, + "artifact publication schema or integrity mode is unsupported", + ) + ) + expected_names = {_READY_FILE, _MANIFEST_FILE, *(entry.file_name for entry in manifest.files)} + if names != expected_names: + raise ManifestValidationError( + f"artifact directory entries differ: missing={sorted(expected_names - names)}, " + f"unexpected={sorted(names - expected_names)}" + ) + manifest_digest = hashlib.sha256(manifest.canonical_bytes).hexdigest() + if marker.manifest_sha256 != manifest_digest: + raise HostWeightError( + _failure(ResolutionStage.VALIDATION, FailureCode.CHECKSUM_MISMATCH, "manifest digest mismatch") + ) + if marker.artifact_content_sha256 != manifest.artifact_content_sha256: + raise HostWeightError( + _failure(ResolutionStage.VALIDATION, FailureCode.CHECKSUM_MISMATCH, "content digest mismatch") + ) + if manifest.artifact_key != identity.key or manifest.identity.canonical_bytes != identity.canonical_bytes: + raise HostWeightError( + _failure( + ResolutionStage.VALIDATION, + FailureCode.MANIFEST_INCOMPATIBLE, + "artifact semantic identity differs from the exact request", + ) + ) + if ( + manifest.producer_schema != identity.producer.manifest_schema + or manifest.restorer_schema != identity.producer.restorer_schema + ): + raise HostWeightError( + _failure( + ResolutionStage.VALIDATION, + FailureCode.MANIFEST_INCOMPATIBLE, + "artifact producer or restorer schema differs from the exact request", + ) + ) + + tensors_by_file: dict[str, list[TensorManifestEntry]] = defaultdict(list) + for tensor_entry in manifest.tensors: + tensors_by_file[tensor_entry.file_name].append(tensor_entry) + tensors: dict[str, torch.Tensor] = {} + regions: list[MappedHostRegion] = [] + backing_digest = hashlib.sha256() + for file_entry in manifest.files: + fd = os.open( + file_entry.file_name, + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=entry_fd, + ) + resources.append(_make_fd_closer(fd)) + file_stat = os.fstat(fd) + if not stat.S_ISREG(file_stat.st_mode) or file_stat.st_nlink != 1: + raise ManifestValidationError(f"payload {file_entry.file_name!r} is not one regular file") + if file_stat.st_size != file_entry.size: + raise HostWeightError( + _failure( + ResolutionStage.VALIDATION, + FailureCode.CHECKSUM_MISMATCH, + f"payload size differs for {file_entry.file_name!r}", + ) + ) + if self.integrity_policy.require_trusted_permissions and ( + file_stat.st_uid != os.geteuid() or file_stat.st_mode & 0o222 + ): + raise ManifestValidationError(f"published payload {file_entry.file_name!r} is not immutable") + if ( + validation is ValidationLevel.FULL_CHECKSUM + and _sha256_fd(fd, file_stat.st_size) != file_entry.sha256 + ): + raise HostWeightError( + _failure( + ResolutionStage.VALIDATION, + FailureCode.CHECKSUM_MISMATCH, + f"payload digest differs for {file_entry.file_name!r}", + ) + ) + backing_digest.update( + canonical_json( + { + "file": file_entry.file_name, + "device": file_stat.st_dev, + "inode": file_stat.st_ino, + "size": file_stat.st_size, + } + ) + ) + handle_context = safe_open(f"/proc/self/fd/{fd}", framework="pt", device="cpu") + handle = handle_context.__enter__() + resources.append(_make_context_closer(handle_context)) + expected_entries = tensors_by_file[file_entry.file_name] + if set(handle.keys()) != {entry.storage_key for entry in expected_entries}: + raise ManifestValidationError(f"payload tensor keys differ for {file_entry.file_name!r}") + addresses: list[tuple[int, int]] = [] + for tensor_entry in expected_entries: + tensor = handle.get_tensor(tensor_entry.storage_key) + actual_nbytes = tensor.numel() * tensor.element_size() + if ( + tuple(tensor.shape) != tensor_entry.shape + or str(tensor.dtype) != tensor_entry.dtype + or tuple(tensor.stride()) != tensor_entry.stride + or actual_nbytes != tensor_entry.nbytes + or tensor.device.type != "cpu" + ): + raise ManifestValidationError(f"mapped tensor metadata differs for {tensor_entry.name!r}") + tensors[tensor_entry.name] = tensor + if actual_nbytes: + addresses.append((tensor.data_ptr(), tensor.data_ptr() + actual_nbytes)) + if addresses: + merged: list[list[int]] = [] + for start, end in sorted(addresses): + if merged and start <= merged[-1][1]: + merged[-1][1] = max(merged[-1][1], end) + else: + merged.append([start, end]) + regions.extend( + MappedHostRegion( + file_name=file_entry.file_name, + address=start, + length=end - start, + ) + for start, end in merged + ) + + backing_digest.update(self.domain_uuid.encode("ascii")) + backing_digest.update(marker.publication_id.encode("ascii")) + provenance = LeaseProvenance( + resolution_id=uuid.uuid4().hex, + identity_digest=identity.key, + artifact_content_sha256=manifest.artifact_content_sha256, + domain_uuid=self.domain_uuid, + publication_id=marker.publication_id, + resolution_source="local", + artifact_origin="producer", + producer_id=identity.producer.producer_id, + provider_id=None, + validation_level=validation.value, + backing_fingerprint=backing_digest.hexdigest(), + ) + return HostWeightLease( + identity=identity, + manifest=manifest, + tensors=tensors, + mapped_regions=tuple(regions), + provenance=provenance, + resource_closers=tuple(resources), + release_artifact_lock=artifact_lock.close, + ) + except BaseException: + for closer in reversed(resources): + with contextlib.suppress(Exception): + closer() + raise + + def get_or_build( + self, + request: BuildRequest, + producer: WeightProducer, + *, + validation: ValidationLevel, + deadline: float, + ) -> StoreResult: + identity = request.identity + if producer.spec.coordination_scope is not CoordinationScope.SINGLE_PROCESS: + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.PRODUCTION, + FailureCode.PRODUCER_UNSUPPORTED, + f"coordination scope {producer.spec.coordination_scope.value} is not implemented", + ), + ) + if tuple(output.key for output in producer.spec.outputs) != (identity.key,): + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.PRODUCTION, + FailureCode.PRODUCER_UNAVAILABLE, + "producer does not declare exactly the requested artifact identity", + ), + ) + + initial = self.lookup(identity, validation=validation, deadline=deadline) + if initial.status is StoreStatus.HIT: + return initial + if initial.status in {StoreStatus.TIMEOUT, StoreStatus.FAILED}: + return initial + + try: + build_lock = FileLock(self._build_lock_path(identity.key), exclusive=True, deadline=deadline) + except FileLockTimeoutError as exc: + return StoreResult( + StoreStatus.TIMEOUT, + failure=_failure( + ResolutionStage.PRODUCTION, + FailureCode.ACTIVE_BUILD_TIMEOUT, + str(exc), + retryable=True, + ), + ) + except OSError as exc: + return StoreResult( + StoreStatus.FAILED, + failure=_failure(ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE, str(exc), retryable=True), + ) + + with build_lock: + recheck = self.lookup(identity, validation=validation, deadline=deadline) + if recheck.status is StoreStatus.HIT: + return StoreResult(StoreStatus.JOINED, lease=recheck.lease) + if recheck.status in {StoreStatus.TIMEOUT, StoreStatus.FAILED}: + return recheck + + try: + with FileLock(self._artifact_lock_path(identity.key), exclusive=True, deadline=deadline): + if self._entry_path(identity.key).exists(): + self._quarantine_locked(identity.key, reason="validation_failure") + self._remove_deny_locked(identity.key) + self._remove_stale_temps_locked(identity.key) + except FileLockTimeoutError as exc: + return StoreResult( + StoreStatus.TIMEOUT, + failure=_failure( + ResolutionStage.PRODUCTION, + FailureCode.ACTIVE_BUILD_TIMEOUT, + str(exc), + retryable=True, + ), + ) + except OSError as exc: + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.LIFECYCLE, + FailureCode.QUARANTINE_FAILED, + f"failed to prepare artifact rebuild: {exc}", + ), + ) + + staging = self.tmp_dir / f"{identity.key}.{os.getpid()}.{uuid.uuid4().hex}" + writer: FilesystemArtifactWriter | None = None + manifest: WeightManifest | None = None + try: + staging.mkdir(mode=0o700) + writer = FilesystemArtifactWriter(staging, capacity_check=self._check_capacity) + metadata = producer.produce(writer) + manifest, _ = writer.finalize(identity, metadata) + artifact_bytes = self._tree_bytes(staging) + self._check_capacity(artifact_bytes, 0) + os.chmod(staging, 0o555) + _fsync_directory(staging) + except HostWeightError as exc: + if writer is not None: + writer.abort() + self._discard_temp(identity.key, staging) + return StoreResult(StoreStatus.FAILED, failure=exc.failure) + except OSError as exc: + if writer is not None: + writer.abort() + self._discard_temp(identity.key, staging) + if exc.errno in (errno.ENOSPC, errno.EDQUOT): + code = FailureCode.ENOSPC + stage = ResolutionStage.CAPACITY + elif writer is None: + code = FailureCode.DOMAIN_UNAVAILABLE + stage = ResolutionStage.DOMAIN + else: + code = FailureCode.PRODUCER_FAILED + stage = ResolutionStage.PRODUCTION + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + stage, + code, + f"artifact production failed: {exc}", + retryable=stage is ResolutionStage.CAPACITY, + ), + ) + except Exception as exc: + if writer is not None: + writer.abort() + self._discard_temp(identity.key, staging) + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.PRODUCTION, + FailureCode.PRODUCER_FAILED, + f"artifact producer failed: {exc}", + ), + ) + + assert manifest is not None + try: + with FileLock(self._artifact_lock_path(identity.key), exclusive=True, deadline=None): + entry_path = self._entry_path(identity.key) + if entry_path.exists(): + existing = WeightManifest.from_dict(_read_json_file(entry_path / _MANIFEST_FILE)) + same_semantic_content = ( + existing.identity.canonical_bytes == manifest.identity.canonical_bytes + and existing.artifact_content_sha256 == manifest.artifact_content_sha256 + and existing.producer_schema == manifest.producer_schema + and existing.restorer_schema == manifest.restorer_schema + and existing.format_metadata == manifest.format_metadata + ) + if not same_semantic_content: + collision = self.quarantine_dir / f"{identity.key}.identity-collision.{uuid.uuid4().hex}" + os.replace(staging, collision) + _fsync_directory(self.quarantine_dir) + _fsync_directory(self.tmp_dir) + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.VALIDATION, + FailureCode.IDENTITY_COLLISION, + "the same semantic identity produced different artifact content", + ), + ) + self._remove_tree(staging) + else: + os.replace(staging, entry_path) + _fsync_directory(self.artifacts_dir) + _fsync_directory(self.tmp_dir) + self._remove_deny_locked(identity.key) + except (OSError, ValueError, ManifestValidationError) as exc: + self._discard_temp(identity.key, staging) + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.LIFECYCLE, + FailureCode.PUBLICATION_FAILED, + f"atomic artifact publication failed: {exc}", + ), + ) + + opened = self.lookup(identity, validation=validation, deadline=deadline) + if opened.status is not StoreStatus.HIT: + return opened + return StoreResult(StoreStatus.BUILT, lease=opened.lease) + + def _check_capacity(self, projected_artifact_bytes: int, additional_bytes: int) -> None: + policy = self.capacity_policy + if policy.max_artifact_bytes is not None and projected_artifact_bytes > policy.max_artifact_bytes: + raise HostWeightError( + _failure( + ResolutionStage.CAPACITY, + FailureCode.ARTIFACT_TOO_LARGE, + f"projected artifact size {projected_artifact_bytes} exceeds {policy.max_artifact_bytes}", + retryable=True, + ) + ) + store_bytes = self._tree_bytes(self.root) + if policy.max_store_bytes is not None and store_bytes + additional_bytes > policy.max_store_bytes: + raise HostWeightError( + _failure( + ResolutionStage.CAPACITY, + FailureCode.STORE_LIMIT_EXCEEDED, + f"store size plus allocation would exceed {policy.max_store_bytes}", + retryable=True, + ) + ) + free_bytes = shutil.disk_usage(self.root).free + if free_bytes - additional_bytes < policy.min_free_bytes: + raise HostWeightError( + _failure( + ResolutionStage.CAPACITY, + FailureCode.INSUFFICIENT_CAPACITY, + f"allocation would leave less than {policy.min_free_bytes} free bytes", + retryable=True, + ) + ) + + def _write_deny(self, key: str, failure: HostWeightFailure) -> None: + path = self._deny_path(key) + document = { + "schema_version": 1, + "stage": failure.stage.value, + "reason": failure.code.value, + "message": failure.message, + } + try: + fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_CLOEXEC, 0o444) + except FileExistsError: + return + try: + data = canonical_json(document) + offset = 0 + while offset < len(data): + written = os.write(fd, data[offset:]) + if written == 0: + raise OSError("short write while publishing deny marker") + offset += written + os.fsync(fd) + finally: + os.close(fd) + _fsync_directory(self.deny_dir) + + def _remove_deny_locked(self, key: str) -> None: + path = self._deny_path(key) + if path.exists(): + path.unlink() + _fsync_directory(self.deny_dir) + + def _quarantine_locked(self, key: str, *, reason: str) -> Path | None: + entry = self._entry_path(key) + if not entry.exists(): + return None + destination = self.quarantine_dir / f"{key}.{reason}.{uuid.uuid4().hex}" + os.replace(entry, destination) + _fsync_directory(self.artifacts_dir) + _fsync_directory(self.quarantine_dir) + return destination + + def _remove_stale_temps_locked(self, key: str) -> None: + for path in self.tmp_dir.glob(f"{key}.*"): + self._remove_tree(path) + _fsync_directory(self.tmp_dir) + + def _discard_temp(self, key: str, path: Path) -> None: + with contextlib.suppress(OSError, FileLockTimeoutError): + with FileLock(self._artifact_lock_path(key), exclusive=True, deadline=None): + self._remove_tree(path) + _fsync_directory(self.tmp_dir) + + @staticmethod + def _remove_tree(path: Path) -> None: + if path.is_symlink() or path.is_file(): + path.unlink(missing_ok=True) + elif path.exists(): + for directory, subdirectories, _ in os.walk(path, topdown=False, followlinks=False): + for name in subdirectories: + child = Path(directory) / name + if not child.is_symlink(): + os.chmod(child, 0o700) + os.chmod(directory, 0o700) + shutil.rmtree(path) + + @staticmethod + def _tree_bytes(root: Path) -> int: + with contextlib.suppress(OSError): + root_stat = root.lstat() + if stat.S_ISREG(root_stat.st_mode): + return root_stat.st_size + total = 0 + for directory, _, files in os.walk(root, followlinks=False): + for name in files: + path = Path(directory) / name + with contextlib.suppress(OSError): + file_stat = path.lstat() + if stat.S_ISREG(file_stat.st_mode): + total += file_stat.st_size + return total + + def cleanup(self, identity: WeightArtifactIdentity) -> HostWeightFailure | None: + """Explicitly remove an inactive artifact without waiting for leases.""" + try: + build_lock = FileLock( + self._build_lock_path(identity.key), + exclusive=True, + deadline=None, + nonblocking=True, + ) + except FileLockTimeoutError: + return _failure( + ResolutionStage.LIFECYCLE, + FailureCode.ACTIVE_BUILD_TIMEOUT, + f"artifact {identity.key} has an active builder", + retryable=True, + ) + with build_lock: + try: + artifact_lock = FileLock( + self._artifact_lock_path(identity.key), + exclusive=True, + deadline=None, + nonblocking=True, + ) + except FileLockTimeoutError: + return _failure( + ResolutionStage.LIFECYCLE, + FailureCode.ACTIVE_LEASE, + f"artifact {identity.key} has an active lease", + retryable=True, + ) + with artifact_lock: + entry = self._entry_path(identity.key) + if entry.exists(): + removal = self.quarantine_dir / f"{identity.key}.cleanup.{uuid.uuid4().hex}" + os.replace(entry, removal) + _fsync_directory(self.artifacts_dir) + self._remove_tree(removal) + _fsync_directory(self.quarantine_dir) + self._remove_deny_locked(identity.key) + return None + + def inspect_domain(self) -> DomainInspection: + usage = shutil.disk_usage(self.root) + inventory: list[ArtifactInventoryEntry] = [] + for state, directory in ( + (ArtifactInventoryState.READY, self.artifacts_dir), + (ArtifactInventoryState.DENIED, self.deny_dir), + (ArtifactInventoryState.QUARANTINED, self.quarantine_dir), + (ArtifactInventoryState.TEMPORARY, self.tmp_dir), + ): + for path in directory.iterdir(): + artifact_key = path.name[:64] + if _ARTIFACT_KEY_RE.fullmatch(artifact_key) is None: + continue + reason = None + if state is ArtifactInventoryState.DENIED: + with contextlib.suppress(OSError, ValueError, ManifestValidationError): + reason_value = _read_json_file(path).get("reason") + if isinstance(reason_value, str): + reason = reason_value + inventory.append( + ArtifactInventoryEntry( + artifact_key=artifact_key, + state=state, + storage_name=path.name, + size_bytes=self._tree_bytes(path), + reason=reason, + ) + ) + inventory.sort(key=lambda item: (item.artifact_key, item.state.value, item.storage_name)) + return DomainInspection( + domain_uuid=self.domain_uuid, + scope=self.domain_policy.scope.value, + domain_id=self.domain_policy.domain_id, + storage_class=self.domain_policy.storage_class.value, + filesystem_type=self.filesystem_type, + policy_version=self.policy_version, + store_bytes=self._tree_bytes(self.root), + filesystem_free_bytes=usage.free, + max_artifact_bytes=self.capacity_policy.max_artifact_bytes, + max_store_bytes=self.capacity_policy.max_store_bytes, + min_free_bytes=self.capacity_policy.min_free_bytes, + inventory=tuple(inventory), + ) + + def inspect_artifact( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + ) -> ArtifactInspection: + result = self._lookup( + identity, + validation=validation, + deadline=time.monotonic(), + record_invalid=False, + ) + manifest = None + publication_id = None + backing_fingerprint = None + if result.lease is not None: + manifest = result.lease.manifest + publication_id = result.lease.provenance.publication_id + backing_fingerprint = result.lease.provenance.backing_fingerprint + result.lease.close() + entry = self._entry_path(identity.key) + return ArtifactInspection( + artifact_key=identity.key, + status=result.status, + size_bytes=self._tree_bytes(entry), + manifest=manifest, + publication_id=publication_id, + backing_fingerprint=backing_fingerprint, + failure=result.failure, + build_lock_active=lock_is_active(self._build_lock_path(identity.key)), + artifact_lock_active=lock_is_active(self._artifact_lock_path(identity.key)), + ) + + +__all__ = ["FilesystemHostWeightStore", "detect_storage_class"] diff --git a/vllm_omni/host_weight_runtime/filesystem/writer.py b/vllm_omni/host_weight_runtime/filesystem/writer.py new file mode 100644 index 000000000..8ecd17799 --- /dev/null +++ b/vllm_omni/host_weight_runtime/filesystem/writer.py @@ -0,0 +1,425 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Store-scoped streaming safetensors publication.""" + +from __future__ import annotations + +import contextlib +import errno +import hashlib +import json +import os +import struct +import uuid +from collections.abc import Callable +from pathlib import Path, PurePosixPath +from types import TracebackType + +import torch +from safetensors import safe_open + +from ..identity import WeightArtifactIdentity, canonical_json +from ..manifest import FileManifestEntry, PublicationMarker, TensorManifestEntry, WeightManifest +from ..protocols import ProductionMetadata, TensorWriteSpec + +_HASH_CHUNK_BYTES = 8 * 1024**2 +_MANIFEST_FILE = "manifest.json" +_READY_FILE = "READY.json" + +_SAFETENSORS_DTYPES: dict[torch.dtype, str] = { + torch.bool: "BOOL", + torch.uint8: "U8", + torch.int8: "I8", + torch.int16: "I16", + torch.int32: "I32", + torch.int64: "I64", + torch.float16: "F16", + torch.bfloat16: "BF16", + torch.float32: "F32", + torch.float64: "F64", +} +for _torch_name, _safe_name in ( + ("float8_e4m3fn", "F8_E4M3"), + ("float8_e5m2", "F8_E5M2"), + ("float8_e4m3fnuz", "F8_E4M3FNUZ"), + ("float8_e5m2fnuz", "F8_E5M2FNUZ"), +): + if (_dtype := getattr(torch, _torch_name, None)) is not None: + _SAFETENSORS_DTYPES[_dtype] = _safe_name + + +class ArtifactWriterError(RuntimeError): + pass + + +def _safe_relative_name(value: str) -> str: + path = PurePosixPath(value) + if not value or path.is_absolute() or len(path.parts) != 1 or path.name != value or value in {".", ".."}: + raise ArtifactWriterError(f"artifact payload name must be one safe relative basename: {value!r}") + if not value.endswith(".safetensors"): + raise ArtifactWriterError(f"artifact payload must use the .safetensors format: {value!r}") + return value + + +def _pwrite_all(fd: int, data: bytes | bytearray | memoryview, offset: int) -> None: + view = memoryview(data).cast("B") + written = 0 + while written < len(view): + count = os.pwrite(fd, view[written:], offset + written) + if count <= 0: + raise OSError(errno.EIO, "short pwrite while publishing host weights") + written += count + + +def _fsync_directory(path: Path) -> None: + fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(_HASH_CHUNK_BYTES): + digest.update(chunk) + return digest.hexdigest() + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + digest = hashlib.sha256() + byte_view = tensor.detach().reshape(-1).view(torch.uint8) + for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES): + chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES] + digest.update(memoryview(chunk.numpy())) + return digest.hexdigest() + + +class _Coverage: + def __init__(self, size: int) -> None: + self._size = size + self._ranges: list[tuple[int, int]] = [] + + def add(self, start: int, end: int) -> None: + if start < 0 or end < start or end > self._size: + raise ArtifactWriterError(f"write range [{start}, {end}) exceeds tensor storage {self._size}") + if start == end: + return + for existing_start, existing_end in self._ranges: + if start < existing_end and existing_start < end: + raise ArtifactWriterError(f"overlapping tensor write ranges [{start}, {end}) and {self._ranges}") + self._ranges.append((start, end)) + + def complete(self) -> bool: + if self._size == 0: + return True + cursor = 0 + for start, end in sorted(self._ranges): + if start != cursor: + return False + cursor = end + return cursor == self._size + + +class StreamingSafeTensorFile: + """One preallocated safetensors file filled through bounded CPU writes.""" + + def __init__( + self, + owner: FilesystemArtifactWriter, + relative_name: str, + specs: tuple[TensorWriteSpec, ...], + ) -> None: + self._owner = owner + self.relative_name = _safe_relative_name(relative_name) + if not specs: + raise ArtifactWriterError("a tensor payload file must contain at least one tensor") + by_name = {spec.name: spec for spec in specs} + if len(by_name) != len(specs): + raise ArtifactWriterError("tensor payload contains duplicate storage keys") + unsupported = [spec.dtype for spec in specs if spec.dtype not in _SAFETENSORS_DTYPES] + if unsupported: + raise ArtifactWriterError(f"unsupported safetensors dtypes: {unsupported}") + self.specs = tuple(sorted(specs, key=lambda spec: spec.name)) + self._specs = {spec.name: spec for spec in self.specs} + self._coverage = {spec.name: _Coverage(spec.nbytes) for spec in self.specs} + self._closed = False + + offsets: dict[str, int] = {} + cursor = 0 + for spec in self.specs: + offsets[spec.name] = cursor + cursor += spec.nbytes + header: dict[str, object] = {"__metadata__": {"format": "pt", "host_weight_runtime": "1"}} + for spec in self.specs: + start = offsets[spec.name] + header[spec.name] = { + "dtype": _SAFETENSORS_DTYPES[spec.dtype], + "shape": list(spec.shape), + "data_offsets": [start, start + spec.nbytes], + } + encoded_header = json.dumps(header, sort_keys=True, separators=(",", ":")).encode("utf-8") + encoded_header += b" " * (-len(encoded_header) % 8) + self._payload_offset = 8 + len(encoded_header) + self._offsets = offsets + self._path = owner.staging_dir / self.relative_name + total_size = self._payload_offset + cursor + owner._reserve_file(total_size) + self._fd = os.open(self._path, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_CLOEXEC, 0o600) + try: + if hasattr(os, "posix_fallocate"): + try: + os.posix_fallocate(self._fd, 0, total_size) + except OSError as exc: + if exc.errno not in (errno.EINVAL, errno.ENOSYS, errno.EOPNOTSUPP): + raise + os.ftruncate(self._fd, total_size) + else: # pragma: no cover - Linux production path has posix_fallocate + os.ftruncate(self._fd, total_size) + _pwrite_all(self._fd, struct.pack(" memoryview: + if tensor.device.type != "cpu" or tensor.is_meta: + raise ArtifactWriterError("artifact producers must provide materialized CPU tensor chunks") + contiguous = tensor.detach().contiguous() + return memoryview(contiguous.view(torch.uint8).numpy()).cast("B") + + def write_tensor(self, name: str, tensor: torch.Tensor) -> None: + spec = self._get_spec(name) + if tuple(tensor.shape) != spec.shape or tensor.dtype != spec.dtype: + raise ArtifactWriterError( + f"tensor {name!r} differs from its declared shape/dtype: " + f"got {tuple(tensor.shape)}/{tensor.dtype}, expected {spec.shape}/{spec.dtype}" + ) + self._write(name, 0, self._bytes(tensor)) + + def write_rows(self, name: str, start_row: int, tensor: torch.Tensor) -> None: + spec = self._get_spec(name) + if not spec.shape: + if start_row != 0 or tuple(tensor.shape) != (): + raise ArtifactWriterError(f"scalar tensor {name!r} only accepts one complete scalar write") + self.write_tensor(name, tensor) + return + if tensor.dtype != spec.dtype or tuple(tensor.shape[1:]) != spec.shape[1:]: + raise ArtifactWriterError(f"row chunk for {name!r} differs from its declared dtype or trailing shape") + row_count = tensor.shape[0] + if start_row < 0 or start_row + row_count > spec.shape[0]: + raise ArtifactWriterError(f"row chunk for {name!r} exceeds its declared first dimension") + row_bytes = spec.nbytes // spec.shape[0] if spec.shape[0] else 0 + self._write(name, start_row * row_bytes, self._bytes(tensor)) + + def _get_spec(self, name: str) -> TensorWriteSpec: + if self._closed: + raise ArtifactWriterError("cannot write through a closed tensor file") + try: + return self._specs[name] + except KeyError as exc: + raise ArtifactWriterError(f"tensor {name!r} was not declared in this payload file") from exc + + def _write(self, name: str, byte_offset: int, data: memoryview) -> None: + coverage = self._coverage[name] + coverage.add(byte_offset, byte_offset + len(data)) + _pwrite_all(self._fd, data, self._payload_offset + self._offsets[name] + byte_offset) + + def close(self) -> None: + if self._closed: + return + missing = [name for name, coverage in self._coverage.items() if not coverage.complete()] + if missing: + self.abort() + raise ArtifactWriterError(f"tensor payload closed with incomplete writes: {missing}") + self._closed = True + try: + os.fsync(self._fd) + finally: + os.close(self._fd) + self._owner._closed(self) + + def abort(self) -> None: + if self._closed: + return + self._closed = True + os.close(self._fd) + size = self._path.stat().st_size if self._path.exists() else 0 + self._path.unlink(missing_ok=True) + self._owner._aborted(self, size) + + def __enter__(self) -> StreamingSafeTensorFile: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + _traceback: TracebackType | None, + ) -> None: + if exc_type is None: + self.close() + else: + self.abort() + + +class FilesystemArtifactWriter: + """Scoped writer that cannot publish outside one invisible staging dir.""" + + def __init__( + self, + staging_dir: Path, + *, + capacity_check: Callable[[int, int], None], + ) -> None: + self.staging_dir = staging_dir + self._capacity_check = capacity_check + self._active: set[StreamingSafeTensorFile] = set() + self._files: dict[str, tuple[TensorWriteSpec, ...]] = {} + self._reserved_bytes = 0 + self._finalized = False + + def open_tensor_file( + self, + relative_name: str, + specs: tuple[TensorWriteSpec, ...], + ) -> StreamingSafeTensorFile: + if self._finalized: + raise ArtifactWriterError("cannot add payloads after artifact finalization") + relative_name = _safe_relative_name(relative_name) + if relative_name in self._files or any(item.relative_name == relative_name for item in self._active): + raise ArtifactWriterError(f"artifact payload {relative_name!r} was declared more than once") + return StreamingSafeTensorFile(self, relative_name, specs) + + def _reserve_file(self, size: int) -> None: + self._capacity_check(self._reserved_bytes + size, size) + self._reserved_bytes += size + + def _forget_reserved_file(self, size: int) -> None: + self._reserved_bytes = max(0, self._reserved_bytes - size) + + def _opened(self, writer: StreamingSafeTensorFile) -> None: + self._active.add(writer) + + def _closed(self, writer: StreamingSafeTensorFile) -> None: + self._active.remove(writer) + self._files[writer.relative_name] = writer.specs + + def _aborted(self, writer: StreamingSafeTensorFile, size: int) -> None: + self._active.discard(writer) + self._forget_reserved_file(size) + + def abort(self) -> None: + for writer in tuple(self._active): + with contextlib.suppress(Exception): + writer.abort() + + def finalize( + self, + identity: WeightArtifactIdentity, + metadata: ProductionMetadata, + ) -> tuple[WeightManifest, PublicationMarker]: + if self._finalized: + raise ArtifactWriterError("artifact writer was finalized more than once") + if self._active: + raise ArtifactWriterError("artifact producer left tensor payload files open") + if not self._files: + raise ArtifactWriterError("artifact producer emitted no tensor payload files") + if metadata.producer_schema != identity.producer.manifest_schema: + raise ArtifactWriterError("production metadata does not match the identity's manifest schema") + if metadata.restorer_schema != identity.producer.restorer_schema: + raise ArtifactWriterError("production metadata does not match the identity's restorer schema") + + tensor_entries: list[TensorManifestEntry] = [] + file_entries: list[FileManifestEntry] = [] + seen_tensors: set[str] = set() + for file_name in sorted(self._files): + path = self.staging_dir / file_name + os.chmod(path, 0o444) + with path.open("rb") as handle: + os.fsync(handle.fileno()) + file_entries.append( + FileManifestEntry( + file_name=file_name, + size=path.stat().st_size, + sha256=_file_sha256(path), + ) + ) + expected = {spec.name: spec for spec in self._files[file_name]} + with safe_open(str(path), framework="pt", device="cpu") as handle: + if set(handle.keys()) != set(expected): + raise ArtifactWriterError(f"published tensor keys differ in {file_name!r}") + for name in sorted(expected): + if name in seen_tensors: + raise ArtifactWriterError(f"artifact tensor {name!r} occurs in multiple payload files") + seen_tensors.add(name) + spec = expected[name] + tensor = handle.get_tensor(name) + if tuple(tensor.shape) != spec.shape or tensor.dtype != spec.dtype: + raise ArtifactWriterError(f"published tensor {name!r} differs from its declared metadata") + entry = TensorManifestEntry( + name=name, + kind=spec.kind, + role=spec.role, + file_name=file_name, + storage_key=name, + shape=tuple(tensor.shape), + dtype=str(tensor.dtype), + stride=tuple(tensor.stride()), + nbytes=tensor.numel() * tensor.element_size(), + sha256=_tensor_sha256(tensor), + ) + tensor_entries.append(entry) + + ordered_tensor_entries = tuple(sorted(tensor_entries, key=lambda entry: entry.name)) + content_digest = hashlib.sha256() + for entry in ordered_tensor_entries: + content_digest.update(canonical_json(entry.semantic_dict())) + + manifest = WeightManifest( + schema_version=1, + artifact_key=identity.key, + identity=identity, + producer_schema=metadata.producer_schema, + restorer_schema=metadata.restorer_schema, + tensors=ordered_tensor_entries, + files=tuple(file_entries), + artifact_content_sha256=content_digest.hexdigest(), + format_metadata=metadata.format_metadata, + ) + manifest_bytes = manifest.canonical_bytes + self._write_metadata_file(_MANIFEST_FILE, manifest_bytes) + marker = PublicationMarker( + schema_version=1, + publication_id=uuid.uuid4().hex, + manifest_sha256=hashlib.sha256(manifest_bytes).hexdigest(), + artifact_content_sha256=manifest.artifact_content_sha256, + integrity_mode="full_checksum", + ) + self._write_metadata_file(_READY_FILE, marker.canonical_bytes) + _fsync_directory(self.staging_dir) + self._finalized = True + return manifest, marker + + def _write_metadata_file(self, name: str, data: bytes) -> None: + path = self.staging_dir / name + fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_CLOEXEC, 0o600) + try: + _pwrite_all(fd, data, 0) + os.fsync(fd) + os.fchmod(fd, 0o444) + os.fsync(fd) + finally: + os.close(fd) + + +__all__ = [ + "ArtifactWriterError", + "FilesystemArtifactWriter", + "StreamingSafeTensorFile", +] diff --git a/vllm_omni/host_weight_runtime/identity.py b/vllm_omni/host_weight_runtime/identity.py new file mode 100644 index 000000000..77b769728 --- /dev/null +++ b/vllm_omni/host_weight_runtime/identity.py @@ -0,0 +1,407 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Deterministic semantic identity for runtime-ready host weights.""" + +from __future__ import annotations + +import hashlib +import json +import math +from dataclasses import dataclass, field +from typing import TypeAlias, cast + +JsonScalar: TypeAlias = None | bool | int | float | str +JsonValue: TypeAlias = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"] + + +class IdentityValidationError(ValueError): + """Raised when an artifact identity is not deterministic JSON.""" + + +def _validate_json(value: object, path: str = "$") -> JsonValue: + if value is None or isinstance(value, (bool, int, str)): + return value + if isinstance(value, float): + if not math.isfinite(value): + raise IdentityValidationError(f"{path} contains a non-finite float") + return value + if isinstance(value, list | tuple): + return [_validate_json(item, f"{path}[{index}]") for index, item in enumerate(value)] + if isinstance(value, dict): + normalized: dict[str, JsonValue] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise IdentityValidationError(f"{path} contains a non-string mapping key") + normalized[key] = _validate_json(item, f"{path}.{key}") + return {key: normalized[key] for key in sorted(normalized)} + raise IdentityValidationError(f"{path} contains unsupported value type {type(value).__qualname__}") + + +def canonical_json(value: object) -> bytes: + """Return one stable JSON encoding, rejecting process-specific objects.""" + return json.dumps( + _validate_json(value), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode("utf-8") + + +@dataclass(frozen=True) +class CanonicalJson: + """Immutable snapshot of deterministic JSON metadata.""" + + encoded: bytes + + def __post_init__(self) -> None: + if not isinstance(self.encoded, bytes): + raise IdentityValidationError("canonical metadata must be immutable bytes") + try: + value = json.loads(self.encoded) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise IdentityValidationError("canonical metadata is not valid JSON") from exc + if canonical_json(value) != self.encoded: + raise IdentityValidationError("canonical metadata does not use the required stable encoding") + + @classmethod + def from_value(cls, value: object) -> CanonicalJson: + return cls(canonical_json(value)) + + @classmethod + def empty(cls) -> CanonicalJson: + return cls(b"{}") + + def to_value(self) -> JsonValue: + return json.loads(self.encoded) + + +def _require_text(name: str, value: object) -> str: + if not isinstance(value, str) or not value: + raise IdentityValidationError(f"{name} must not be empty") + return value + + +def _require_int(name: str, value: object) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise IdentityValidationError(f"{name} must be an integer") + return value + + +def _require_metadata(name: str, value: object) -> None: + if not isinstance(value, CanonicalJson): + raise IdentityValidationError(f"{name} must use CanonicalJson") + + +def _require_exact_keys(value: dict[str, object], expected: set[str], kind: str) -> None: + keys = set(value) + if keys != expected: + raise IdentityValidationError( + f"{kind} fields differ: missing={sorted(expected - keys)}, unexpected={sorted(keys - expected)}" + ) + + +@dataclass(frozen=True) +class WeightSourceIdentity: + model_id: str + revision: str + fingerprint: str + metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + _require_text("source.model_id", self.model_id) + _require_text("source.revision", self.revision) + _require_text("source.fingerprint", self.fingerprint) + _require_metadata("source.metadata", self.metadata) + + def to_dict(self) -> dict[str, JsonValue]: + return { + "model_id": self.model_id, + "revision": self.revision, + "fingerprint": self.fingerprint, + "metadata": self.metadata.to_value(), + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> WeightSourceIdentity: + _require_exact_keys(value, {"model_id", "revision", "fingerprint", "metadata"}, "source identity") + return cls( + model_id=_require_text("source.model_id", value["model_id"]), + revision=_require_text("source.revision", value["revision"]), + fingerprint=_require_text("source.fingerprint", value["fingerprint"]), + metadata=CanonicalJson.from_value(value["metadata"]), + ) + + +@dataclass(frozen=True) +class ComponentIdentity: + name: str + ownership: str + metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + _require_text("component.name", self.name) + _require_text("component.ownership", self.ownership) + _require_metadata("component.metadata", self.metadata) + + def to_dict(self) -> dict[str, JsonValue]: + return {"name": self.name, "ownership": self.ownership, "metadata": self.metadata.to_value()} + + @classmethod + def from_dict(cls, value: dict[str, object]) -> ComponentIdentity: + _require_exact_keys(value, {"name", "ownership", "metadata"}, "component identity") + return cls( + name=_require_text("component.name", value["name"]), + ownership=_require_text("component.ownership", value["ownership"]), + metadata=CanonicalJson.from_value(value["metadata"]), + ) + + +@dataclass(frozen=True) +class WeightRepresentation: + name: str + dtype: str + metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + _require_text("representation.name", self.name) + _require_text("representation.dtype", self.dtype) + _require_metadata("representation.metadata", self.metadata) + + def to_dict(self) -> dict[str, JsonValue]: + return {"name": self.name, "dtype": self.dtype, "metadata": self.metadata.to_value()} + + @classmethod + def from_dict(cls, value: dict[str, object]) -> WeightRepresentation: + _require_exact_keys(value, {"name", "dtype", "metadata"}, "weight representation") + return cls( + name=_require_text("representation.name", value["name"]), + dtype=_require_text("representation.dtype", value["dtype"]), + metadata=CanonicalJson.from_value(value["metadata"]), + ) + + +@dataclass(frozen=True) +class RuntimeWeightLayout: + name: str + tensor_parallel_size: int = 1 + tensor_parallel_rank: int = 0 + sequence_parallel_size: int = 1 + sequence_parallel_backend: str = "none" + metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + _require_text("layout.name", self.name) + tensor_parallel_size = _require_int("layout.tensor_parallel_size", self.tensor_parallel_size) + tensor_parallel_rank = _require_int("layout.tensor_parallel_rank", self.tensor_parallel_rank) + sequence_parallel_size = _require_int("layout.sequence_parallel_size", self.sequence_parallel_size) + if tensor_parallel_size < 1: + raise IdentityValidationError("tensor_parallel_size must be positive") + if not 0 <= tensor_parallel_rank < tensor_parallel_size: + raise IdentityValidationError("tensor_parallel_rank must be within tensor_parallel_size") + if sequence_parallel_size < 1: + raise IdentityValidationError("sequence_parallel_size must be positive") + _require_text("layout.sequence_parallel_backend", self.sequence_parallel_backend) + _require_metadata("layout.metadata", self.metadata) + + def to_dict(self) -> dict[str, JsonValue]: + return { + "name": self.name, + "tensor_parallel_size": self.tensor_parallel_size, + "tensor_parallel_rank": self.tensor_parallel_rank, + "sequence_parallel_size": self.sequence_parallel_size, + "sequence_parallel_backend": self.sequence_parallel_backend, + "metadata": self.metadata.to_value(), + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> RuntimeWeightLayout: + expected = { + "name", + "tensor_parallel_size", + "tensor_parallel_rank", + "sequence_parallel_size", + "sequence_parallel_backend", + "metadata", + } + _require_exact_keys(value, expected, "runtime weight layout") + return cls( + name=_require_text("layout.name", value["name"]), + tensor_parallel_size=_require_int("layout.tensor_parallel_size", value["tensor_parallel_size"]), + tensor_parallel_rank=_require_int("layout.tensor_parallel_rank", value["tensor_parallel_rank"]), + sequence_parallel_size=_require_int("layout.sequence_parallel_size", value["sequence_parallel_size"]), + sequence_parallel_backend=_require_text( + "layout.sequence_parallel_backend", value["sequence_parallel_backend"] + ), + metadata=CanonicalJson.from_value(value["metadata"]), + ) + + +@dataclass(frozen=True) +class AdaptationIdentity: + kind: str = "base" + fingerprint: str | None = None + metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + _require_text("adaptation.kind", self.kind) + if self.kind != "base" and not self.fingerprint: + raise IdentityValidationError("non-base adaptation requires a fingerprint") + if self.fingerprint is not None and (not isinstance(self.fingerprint, str) or not self.fingerprint): + raise IdentityValidationError("adaptation fingerprint must be a non-empty string or null") + _require_metadata("adaptation.metadata", self.metadata) + + def to_dict(self) -> dict[str, JsonValue]: + return {"kind": self.kind, "fingerprint": self.fingerprint, "metadata": self.metadata.to_value()} + + @classmethod + def from_dict(cls, value: dict[str, object]) -> AdaptationIdentity: + _require_exact_keys(value, {"kind", "fingerprint", "metadata"}, "adaptation identity") + fingerprint = value["fingerprint"] + if fingerprint is not None and not isinstance(fingerprint, str): + raise IdentityValidationError("adaptation fingerprint must be a string or null") + return cls( + kind=_require_text("adaptation.kind", value["kind"]), + fingerprint=fingerprint, + metadata=CanonicalJson.from_value(value["metadata"]), + ) + + +@dataclass(frozen=True) +class ProducerIdentity: + producer_id: str + version: str + implementation_fingerprint: str + manifest_schema: str + restorer_schema: str + + def __post_init__(self) -> None: + for name, value in ( + ("producer.producer_id", self.producer_id), + ("producer.version", self.version), + ("producer.implementation_fingerprint", self.implementation_fingerprint), + ("producer.manifest_schema", self.manifest_schema), + ("producer.restorer_schema", self.restorer_schema), + ): + _require_text(name, value) + + def to_dict(self) -> dict[str, JsonValue]: + return { + "producer_id": self.producer_id, + "version": self.version, + "implementation_fingerprint": self.implementation_fingerprint, + "manifest_schema": self.manifest_schema, + "restorer_schema": self.restorer_schema, + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> ProducerIdentity: + expected = { + "producer_id", + "version", + "implementation_fingerprint", + "manifest_schema", + "restorer_schema", + } + _require_exact_keys(value, expected, "producer identity") + return cls( + producer_id=_require_text("producer.producer_id", value["producer_id"]), + version=_require_text("producer.version", value["version"]), + implementation_fingerprint=_require_text( + "producer.implementation_fingerprint", value["implementation_fingerprint"] + ), + manifest_schema=_require_text("producer.manifest_schema", value["manifest_schema"]), + restorer_schema=_require_text("producer.restorer_schema", value["restorer_schema"]), + ) + + +@dataclass(frozen=True) +class WeightArtifactIdentity: + schema_version: int + source: WeightSourceIdentity + component: ComponentIdentity + representation: WeightRepresentation + layout: RuntimeWeightLayout + adaptation: AdaptationIdentity + producer: ProducerIdentity + + def __post_init__(self) -> None: + if _require_int("identity.schema_version", self.schema_version) < 1: + raise IdentityValidationError("identity schema_version must be positive") + for name, value, expected in ( + ("source", self.source, WeightSourceIdentity), + ("component", self.component, ComponentIdentity), + ("representation", self.representation, WeightRepresentation), + ("layout", self.layout, RuntimeWeightLayout), + ("adaptation", self.adaptation, AdaptationIdentity), + ("producer", self.producer, ProducerIdentity), + ): + if not isinstance(value, expected): + raise IdentityValidationError(f"identity {name} has the wrong contract type") + + def to_dict(self) -> dict[str, JsonValue]: + return { + "schema_version": self.schema_version, + "source": self.source.to_dict(), + "component": self.component.to_dict(), + "representation": self.representation.to_dict(), + "layout": self.layout.to_dict(), + "adaptation": self.adaptation.to_dict(), + "producer": self.producer.to_dict(), + } + + @property + def canonical_bytes(self) -> bytes: + return canonical_json(self.to_dict()) + + @property + def key(self) -> str: + return hashlib.sha256(self.canonical_bytes).hexdigest() + + @classmethod + def from_dict(cls, value: dict[str, object]) -> WeightArtifactIdentity: + expected = { + "schema_version", + "source", + "component", + "representation", + "layout", + "adaptation", + "producer", + } + _require_exact_keys(value, expected, "weight artifact identity") + source = value["source"] + component = value["component"] + representation = value["representation"] + layout = value["layout"] + adaptation = value["adaptation"] + producer = value["producer"] + if any( + not isinstance(section, dict) + for section in (source, component, representation, layout, adaptation, producer) + ): + raise IdentityValidationError("identity sections must be JSON objects") + return cls( + schema_version=_require_int("identity.schema_version", value["schema_version"]), + source=WeightSourceIdentity.from_dict(cast(dict[str, object], source)), + component=ComponentIdentity.from_dict(cast(dict[str, object], component)), + representation=WeightRepresentation.from_dict(cast(dict[str, object], representation)), + layout=RuntimeWeightLayout.from_dict(cast(dict[str, object], layout)), + adaptation=AdaptationIdentity.from_dict(cast(dict[str, object], adaptation)), + producer=ProducerIdentity.from_dict(cast(dict[str, object], producer)), + ) + + +__all__ = [ + "AdaptationIdentity", + "CanonicalJson", + "ComponentIdentity", + "IdentityValidationError", + "JsonValue", + "ProducerIdentity", + "RuntimeWeightLayout", + "WeightArtifactIdentity", + "WeightRepresentation", + "WeightSourceIdentity", + "canonical_json", +] diff --git a/vllm_omni/host_weight_runtime/inspection.py b/vllm_omni/host_weight_runtime/inspection.py new file mode 100644 index 000000000..c61760502 --- /dev/null +++ b/vllm_omni/host_weight_runtime/inspection.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Typed, read-only snapshots for local Host Weight Store inspection.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Protocol + +from .config import ValidationLevel +from .errors import HostWeightFailure +from .identity import WeightArtifactIdentity +from .manifest import WeightManifest +from .outcomes import StoreStatus + + +class ArtifactInventoryState(str, Enum): + READY = "ready" + DENIED = "denied" + QUARANTINED = "quarantined" + TEMPORARY = "temporary" + + +@dataclass(frozen=True) +class ArtifactInventoryEntry: + artifact_key: str + state: ArtifactInventoryState + storage_name: str + size_bytes: int + reason: str | None = None + + +@dataclass(frozen=True) +class DomainInspection: + domain_uuid: str + scope: str + domain_id: str + storage_class: str + filesystem_type: str + policy_version: int + store_bytes: int + filesystem_free_bytes: int + max_artifact_bytes: int | None + max_store_bytes: int | None + min_free_bytes: int + inventory: tuple[ArtifactInventoryEntry, ...] + + +@dataclass(frozen=True) +class ArtifactInspection: + artifact_key: str + status: StoreStatus + size_bytes: int + manifest: WeightManifest | None + publication_id: str | None + backing_fingerprint: str | None + failure: HostWeightFailure | None + build_lock_active: bool + artifact_lock_active: bool + + +class HostWeightInspector(Protocol): + def inspect_domain(self) -> DomainInspection: ... + + def inspect_artifact( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + ) -> ArtifactInspection: ... + + +__all__ = [ + "ArtifactInspection", + "ArtifactInventoryEntry", + "ArtifactInventoryState", + "DomainInspection", + "HostWeightInspector", +] diff --git a/vllm_omni/host_weight_runtime/lease.py b/vllm_omni/host_weight_runtime/lease.py new file mode 100644 index 000000000..7446f0bbe --- /dev/null +++ b/vllm_omni/host_weight_runtime/lease.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Process-local ownership for immutable mmap-backed host tensors.""" + +from __future__ import annotations + +import contextlib +import threading +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import NoReturn, SupportsIndex + +import torch + +from .identity import WeightArtifactIdentity +from .manifest import WeightManifest + + +@dataclass(frozen=True) +class MappedHostRegion: + """One process-local virtual range belonging to an artifact payload file.""" + + file_name: str + address: int + length: int + read_only: bool = True + + def __post_init__(self) -> None: + if not self.file_name or self.address < 0 or self.length < 0: + raise ValueError("mapped host region contains invalid fields") + + +@dataclass(frozen=True) +class LeaseProvenance: + resolution_id: str + identity_digest: str + artifact_content_sha256: str + domain_uuid: str + publication_id: str + resolution_source: str + artifact_origin: str + producer_id: str | None + provider_id: str | None + validation_level: str + backing_fingerprint: str + + +class HostWeightLease: + """Stable tensor views and resources retained until transport teardown.""" + + def __init__( + self, + *, + identity: WeightArtifactIdentity, + manifest: WeightManifest, + tensors: Mapping[str, torch.Tensor], + mapped_regions: tuple[MappedHostRegion, ...], + provenance: LeaseProvenance, + resource_closers: tuple[Callable[[], None], ...], + release_artifact_lock: Callable[[], None], + ) -> None: + if manifest.identity.canonical_bytes != identity.canonical_bytes: + raise ValueError("lease manifest identity differs from the requested identity") + if set(tensors) != {entry.name for entry in manifest.tensors}: + raise ValueError("lease tensor names differ from the validated manifest") + entries = {entry.name: entry for entry in manifest.tensors} + for name, tensor in tensors.items(): + entry = entries[name] + if ( + not isinstance(tensor, torch.Tensor) + or tensor.device.type != "cpu" + or tuple(tensor.shape) != entry.shape + or str(tensor.dtype) != entry.dtype + or tuple(tensor.stride()) != entry.stride + or tensor.numel() * tensor.element_size() != entry.nbytes + ): + raise ValueError(f"lease tensor {name!r} differs from the validated host metadata") + if provenance.identity_digest != identity.key: + raise ValueError("lease provenance identity digest differs from the requested identity") + if provenance.artifact_content_sha256 != manifest.artifact_content_sha256: + raise ValueError("lease provenance content digest differs from the validated manifest") + if not isinstance(mapped_regions, tuple) or any( + not isinstance(region, MappedHostRegion) for region in mapped_regions + ): + raise ValueError("lease mapped regions must be an immutable tuple") + if not isinstance(resource_closers, tuple) or any(not callable(closer) for closer in resource_closers): + raise ValueError("lease resource closers must be an immutable tuple of callables") + if not callable(release_artifact_lock): + raise ValueError("lease artifact lock release must be callable") + self.identity = identity + self.manifest = manifest + self._tensor_storage = dict(tensors) + self.tensors = MappingProxyType(self._tensor_storage) + self.mapped_regions = mapped_regions + self.provenance = provenance + self._resource_closers = resource_closers + self._release_artifact_lock = release_artifact_lock + self._state_lock = threading.RLock() + self._closed = False + + @property + def closed(self) -> bool: + with self._state_lock: + return self._closed + + def close(self) -> None: + with self._state_lock: + if self._closed: + return + self._closed = True + + errors: list[BaseException] = [] + self._tensor_storage.clear() + for closer in reversed(self._resource_closers): + try: + closer() + except BaseException as exc: # pragma: no cover - best-effort teardown aggregation + errors.append(exc) + try: + self._release_artifact_lock() + except BaseException as exc: # pragma: no cover - best-effort teardown aggregation + errors.append(exc) + if errors: + raise RuntimeError(f"failed to close {len(errors)} host weight lease resources") from errors[0] + + def __enter__(self) -> HostWeightLease: + if self.closed: + raise RuntimeError("cannot re-enter a closed HostWeightLease") + return self + + def __exit__(self, exc_type: object, exc: object, _traceback: object) -> None: + self.close() + + def __reduce_ex__(self, _protocol: SupportsIndex) -> NoReturn: + raise TypeError("HostWeightLease is process-local and cannot be serialized") + + def __del__(self) -> None: + with contextlib.suppress(Exception): + self.close() + + +__all__ = ["HostWeightLease", "LeaseProvenance", "MappedHostRegion"] diff --git a/vllm_omni/host_weight_runtime/manifest.py b/vllm_omni/host_weight_runtime/manifest.py new file mode 100644 index 000000000..dfeaf2b95 --- /dev/null +++ b/vllm_omni/host_weight_runtime/manifest.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Versioned immutable manifests for runtime-ready host weights.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from pathlib import PurePosixPath +from typing import cast + +import regex as re + +from .identity import CanonicalJson, JsonValue, WeightArtifactIdentity, canonical_json + +_SHA256_RE = re.compile(r"[0-9a-f]{64}") + + +class ManifestValidationError(ValueError): + """Raised when a published manifest violates its schema.""" + + +class TensorKind(str, Enum): + PARAMETER = "parameter" + BUFFER = "buffer" + + +def _require_text(name: str, value: object) -> str: + if not isinstance(value, str) or not value: + raise ManifestValidationError(f"{name} must be a non-empty string") + return value + + +def _require_int(name: str, value: object) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise ManifestValidationError(f"{name} must be an integer") + return value + + +def _require_sha256(name: str, value: object) -> str: + if not isinstance(value, str) or _SHA256_RE.fullmatch(value) is None: + raise ManifestValidationError(f"{name} must be a lowercase SHA-256 digest") + return value + + +def _require_safe_file_name(value: object) -> str: + if not isinstance(value, str): + raise ManifestValidationError("artifact file name must be a string") + path = PurePosixPath(value) + if not value or path.is_absolute() or len(path.parts) != 1 or path.name != value or value in {".", ".."}: + raise ManifestValidationError(f"artifact file name must be one safe relative basename: {value!r}") + return value + + +def _require_int_tuple(name: str, value: object) -> tuple[int, ...]: + if not isinstance(value, list): + raise ManifestValidationError(f"{name} must be an array") + return tuple(_require_int(f"{name}[{index}]", item) for index, item in enumerate(value)) + + +def _exact_keys(value: dict[str, object], expected: set[str], kind: str) -> None: + keys = set(value) + if keys != expected: + raise ManifestValidationError( + f"{kind} fields differ: missing={sorted(expected - keys)}, unexpected={sorted(keys - expected)}" + ) + + +@dataclass(frozen=True) +class TensorManifestEntry: + name: str + kind: TensorKind + role: str + file_name: str + storage_key: str + shape: tuple[int, ...] + dtype: str + stride: tuple[int, ...] + nbytes: int + sha256: str + + def __post_init__(self) -> None: + _require_text("tensor.name", self.name) + _require_text("tensor.storage_key", self.storage_key) + _require_text("tensor.role", self.role) + _require_text("tensor.dtype", self.dtype) + if not isinstance(self.kind, TensorKind): + raise ManifestValidationError("tensor kind must use TensorKind") + _require_safe_file_name(self.file_name) + if not isinstance(self.shape, tuple) or any( + isinstance(dimension, bool) or not isinstance(dimension, int) or dimension < 0 for dimension in self.shape + ): + raise ManifestValidationError(f"tensor {self.name!r} has a negative shape dimension") + if ( + not isinstance(self.stride, tuple) + or len(self.stride) != len(self.shape) + or any(isinstance(value, bool) or not isinstance(value, int) or value < 0 for value in self.stride) + ): + raise ManifestValidationError(f"tensor {self.name!r} has an invalid stride") + if _require_int(f"tensors[{self.name!r}].nbytes", self.nbytes) < 0: + raise ManifestValidationError(f"tensor {self.name!r} has negative nbytes") + _require_sha256(f"tensors[{self.name!r}].sha256", self.sha256) + + def semantic_dict(self) -> dict[str, JsonValue]: + """Return content identity independent of physical file sharding.""" + return { + "name": self.name, + "kind": self.kind.value, + "role": self.role, + "shape": list(self.shape), + "dtype": self.dtype, + "stride": list(self.stride), + "nbytes": self.nbytes, + "sha256": self.sha256, + } + + def to_dict(self) -> dict[str, JsonValue]: + return { + "name": self.name, + "kind": self.kind.value, + "role": self.role, + "file_name": self.file_name, + "storage_key": self.storage_key, + "shape": list(self.shape), + "dtype": self.dtype, + "stride": list(self.stride), + "nbytes": self.nbytes, + "sha256": self.sha256, + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> TensorManifestEntry: + expected = { + "name", + "kind", + "role", + "file_name", + "storage_key", + "shape", + "dtype", + "stride", + "nbytes", + "sha256", + } + _exact_keys(value, expected, "tensor manifest entry") + kind_value = _require_text("tensor.kind", value["kind"]) + try: + kind = TensorKind(kind_value) + except ValueError as exc: + raise ManifestValidationError(f"unknown tensor kind {value['kind']!r}") from exc + return cls( + name=_require_text("tensor.name", value["name"]), + kind=kind, + role=_require_text("tensor.role", value["role"]), + file_name=_require_safe_file_name(value["file_name"]), + storage_key=_require_text("tensor.storage_key", value["storage_key"]), + shape=_require_int_tuple("tensor.shape", value["shape"]), + dtype=_require_text("tensor.dtype", value["dtype"]), + stride=_require_int_tuple("tensor.stride", value["stride"]), + nbytes=_require_int("tensor.nbytes", value["nbytes"]), + sha256=_require_sha256("tensor.sha256", value["sha256"]), + ) + + +@dataclass(frozen=True) +class FileManifestEntry: + file_name: str + size: int + sha256: str + + def __post_init__(self) -> None: + _require_safe_file_name(self.file_name) + if _require_int(f"files[{self.file_name!r}].size", self.size) < 0: + raise ManifestValidationError(f"artifact file {self.file_name!r} has negative size") + _require_sha256(f"files[{self.file_name!r}].sha256", self.sha256) + + def to_dict(self) -> dict[str, JsonValue]: + return {"file_name": self.file_name, "size": self.size, "sha256": self.sha256} + + @classmethod + def from_dict(cls, value: dict[str, object]) -> FileManifestEntry: + _exact_keys(value, {"file_name", "size", "sha256"}, "file manifest entry") + return cls( + file_name=_require_safe_file_name(value["file_name"]), + size=_require_int("file.size", value["size"]), + sha256=_require_sha256("file.sha256", value["sha256"]), + ) + + +@dataclass(frozen=True) +class WeightManifest: + schema_version: int + artifact_key: str + identity: WeightArtifactIdentity + producer_schema: str + restorer_schema: str + tensors: tuple[TensorManifestEntry, ...] + files: tuple[FileManifestEntry, ...] + artifact_content_sha256: str + format_metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + if _require_int("manifest.schema_version", self.schema_version) < 1: + raise ManifestValidationError("manifest schema_version must be positive") + _require_sha256("artifact_key", self.artifact_key) + _require_sha256("artifact_content_sha256", self.artifact_content_sha256) + if not isinstance(self.identity, WeightArtifactIdentity): + raise ManifestValidationError("manifest identity must use WeightArtifactIdentity") + if self.artifact_key != self.identity.key: + raise ManifestValidationError("manifest artifact_key does not match its semantic identity") + _require_text("manifest.producer_schema", self.producer_schema) + _require_text("manifest.restorer_schema", self.restorer_schema) + if ( + not isinstance(self.tensors, tuple) + or not isinstance(self.files, tuple) + or not self.tensors + or not self.files + or any(not isinstance(entry, TensorManifestEntry) for entry in self.tensors) + or any(not isinstance(entry, FileManifestEntry) for entry in self.files) + ): + raise ManifestValidationError("an artifact must contain at least one tensor and one file") + if not isinstance(self.format_metadata, CanonicalJson): + raise ManifestValidationError("manifest format_metadata must use CanonicalJson") + tensor_names = [entry.name for entry in self.tensors] + if tensor_names != sorted(tensor_names) or len(set(tensor_names)) != len(tensor_names): + raise ManifestValidationError("manifest tensors must have unique names in canonical order") + file_names = [entry.file_name for entry in self.files] + if file_names != sorted(file_names) or len(set(file_names)) != len(file_names): + raise ManifestValidationError("manifest files must have unique names in canonical order") + available_files = set(file_names) + if any(entry.file_name not in available_files for entry in self.tensors): + raise ManifestValidationError("tensor manifest references an undeclared artifact file") + storage_bindings = [(entry.file_name, entry.storage_key) for entry in self.tensors] + if len(set(storage_bindings)) != len(storage_bindings): + raise ManifestValidationError("manifest tensor storage bindings must be unique") + + @property + def canonical_bytes(self) -> bytes: + return canonical_json(self.to_dict()) + + @property + def total_tensor_bytes(self) -> int: + return sum(entry.nbytes for entry in self.tensors) + + def to_dict(self) -> dict[str, JsonValue]: + return { + "schema_version": self.schema_version, + "artifact_key": self.artifact_key, + "identity": self.identity.to_dict(), + "producer_schema": self.producer_schema, + "restorer_schema": self.restorer_schema, + "tensors": [entry.to_dict() for entry in self.tensors], + "files": [entry.to_dict() for entry in self.files], + "artifact_content_sha256": self.artifact_content_sha256, + "format_metadata": self.format_metadata.to_value(), + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> WeightManifest: + expected = { + "schema_version", + "artifact_key", + "identity", + "producer_schema", + "restorer_schema", + "tensors", + "files", + "artifact_content_sha256", + "format_metadata", + } + _exact_keys(value, expected, "weight manifest") + identity = value["identity"] + if not isinstance(identity, dict): + raise ManifestValidationError("manifest identity must be an object") + tensors = value["tensors"] + files = value["files"] + if not isinstance(tensors, list) or not all(isinstance(item, dict) for item in tensors): + raise ManifestValidationError("manifest tensors must be an array of objects") + if not isinstance(files, list) or not all(isinstance(item, dict) for item in files): + raise ManifestValidationError("manifest files must be an array of objects") + return cls( + schema_version=_require_int("manifest.schema_version", value["schema_version"]), + artifact_key=_require_sha256("manifest.artifact_key", value["artifact_key"]), + identity=WeightArtifactIdentity.from_dict(cast(dict[str, object], identity)), + producer_schema=_require_text("manifest.producer_schema", value["producer_schema"]), + restorer_schema=_require_text("manifest.restorer_schema", value["restorer_schema"]), + tensors=tuple(TensorManifestEntry.from_dict(cast(dict[str, object], item)) for item in tensors), + files=tuple(FileManifestEntry.from_dict(cast(dict[str, object], item)) for item in files), + artifact_content_sha256=_require_sha256( + "manifest.artifact_content_sha256", value["artifact_content_sha256"] + ), + format_metadata=CanonicalJson.from_value(value["format_metadata"]), + ) + + +@dataclass(frozen=True) +class PublicationMarker: + schema_version: int + publication_id: str + manifest_sha256: str + artifact_content_sha256: str + integrity_mode: str + + def __post_init__(self) -> None: + if _require_int("publication.schema_version", self.schema_version) < 1: + raise ManifestValidationError("publication schema_version must be positive") + _require_text("publication.publication_id", self.publication_id) + _require_text("publication.integrity_mode", self.integrity_mode) + _require_sha256("publication.manifest_sha256", self.manifest_sha256) + _require_sha256("publication.artifact_content_sha256", self.artifact_content_sha256) + + @property + def canonical_bytes(self) -> bytes: + return canonical_json(self.to_dict()) + + def to_dict(self) -> dict[str, JsonValue]: + return { + "schema_version": self.schema_version, + "publication_id": self.publication_id, + "manifest_sha256": self.manifest_sha256, + "artifact_content_sha256": self.artifact_content_sha256, + "integrity_mode": self.integrity_mode, + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> PublicationMarker: + expected = { + "schema_version", + "publication_id", + "manifest_sha256", + "artifact_content_sha256", + "integrity_mode", + } + _exact_keys(value, expected, "publication marker") + return cls( + schema_version=_require_int("publication.schema_version", value["schema_version"]), + publication_id=_require_text("publication.publication_id", value["publication_id"]), + manifest_sha256=_require_sha256("publication.manifest_sha256", value["manifest_sha256"]), + artifact_content_sha256=_require_sha256( + "publication.artifact_content_sha256", value["artifact_content_sha256"] + ), + integrity_mode=_require_text("publication.integrity_mode", value["integrity_mode"]), + ) + + +__all__ = [ + "FileManifestEntry", + "ManifestValidationError", + "PublicationMarker", + "TensorKind", + "TensorManifestEntry", + "WeightManifest", +] diff --git a/vllm_omni/host_weight_runtime/outcomes.py b/vllm_omni/host_weight_runtime/outcomes.py new file mode 100644 index 000000000..7cae2aaad --- /dev/null +++ b/vllm_omni/host_weight_runtime/outcomes.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Typed store and end-to-end resolution outcomes.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING + +from .errors import HostWeightFailure, ResolutionAction, ResolutionStage + +if TYPE_CHECKING: + from .lease import HostWeightLease + + +class StoreStatus(str, Enum): + HIT = "hit" + MISS = "miss" + BUILT = "built" + JOINED = "joined" + INVALID = "invalid" + TIMEOUT = "timeout" + FAILED = "failed" + + +@dataclass(frozen=True) +class StoreResult: + status: StoreStatus + lease: HostWeightLease | None = None + failure: HostWeightFailure | None = None + + def __post_init__(self) -> None: + has_lease = self.status in {StoreStatus.HIT, StoreStatus.BUILT, StoreStatus.JOINED} + if has_lease != (self.lease is not None): + raise ValueError(f"store result {self.status.value} has inconsistent lease ownership") + if self.status in {StoreStatus.HIT, StoreStatus.MISS, StoreStatus.BUILT, StoreStatus.JOINED}: + if self.failure is not None: + raise ValueError(f"successful store result {self.status.value} must not carry a failure") + if self.status in {StoreStatus.INVALID, StoreStatus.TIMEOUT, StoreStatus.FAILED} and self.failure is None: + raise ValueError(f"store result {self.status.value} requires a typed failure") + + +class AttemptResult(str, Enum): + HIT = "hit" + MISS = "miss" + SUCCESS = "success" + TIMEOUT = "timeout" + ERROR = "error" + SKIPPED = "skipped" + + +class ResolutionOutcome(str, Enum): + LOCAL_HIT = "local_hit" + REMOTE_IMPORT = "remote_import" + LOCAL_PRODUCTION = "local_production" + CANONICAL_DIRECT = "canonical_direct" + CANONICAL_FALLBACK = "canonical_fallback" + FAILED = "failed" + + +@dataclass(frozen=True) +class ResolutionAttempt: + stage: ResolutionStage + result: AttemptResult + action: ResolutionAction + elapsed_seconds: float + failure: HostWeightFailure | None = None + + def __post_init__(self) -> None: + if not math.isfinite(self.elapsed_seconds) or self.elapsed_seconds < 0: + raise ValueError("resolution attempt elapsed_seconds must be finite and non-negative") + if self.result in {AttemptResult.HIT, AttemptResult.MISS, AttemptResult.SUCCESS} and self.failure is not None: + raise ValueError(f"resolution attempt {self.result.value} must not carry a failure") + if self.result in {AttemptResult.TIMEOUT, AttemptResult.ERROR} and self.failure is None: + raise ValueError(f"resolution attempt {self.result.value} requires a typed failure") + + +@dataclass(frozen=True) +class ResolutionReport: + resolution_id: str + outcome: ResolutionOutcome + attempts: tuple[ResolutionAttempt, ...] + elapsed_seconds: float + + def __post_init__(self) -> None: + if not self.resolution_id: + raise ValueError("resolution report requires a resolution_id") + if not isinstance(self.attempts, tuple): + raise ValueError("resolution report attempts must be an immutable tuple") + if not math.isfinite(self.elapsed_seconds) or self.elapsed_seconds < 0: + raise ValueError("resolution report elapsed_seconds must be finite and non-negative") + + +@dataclass(frozen=True) +class HostWeightResolution: + report: ResolutionReport + lease: HostWeightLease | None = None + + def __post_init__(self) -> None: + lease_outcomes = { + ResolutionOutcome.LOCAL_HIT, + ResolutionOutcome.REMOTE_IMPORT, + ResolutionOutcome.LOCAL_PRODUCTION, + } + if (self.report.outcome in lease_outcomes) != (self.lease is not None): + raise ValueError("resolution outcome has inconsistent lease ownership") + + +__all__ = [ + "AttemptResult", + "HostWeightResolution", + "ResolutionAttempt", + "ResolutionOutcome", + "ResolutionReport", + "StoreResult", + "StoreStatus", +] diff --git a/vllm_omni/host_weight_runtime/protocols.py b/vllm_omni/host_weight_runtime/protocols.py new file mode 100644 index 000000000..c5572c9fb --- /dev/null +++ b/vllm_omni/host_weight_runtime/protocols.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Provider and consumer protocols for Host Weight Runtime extensions.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from enum import Enum +from types import TracebackType +from typing import TYPE_CHECKING, Protocol + +import torch + +from .identity import CanonicalJson, WeightArtifactIdentity +from .manifest import TensorKind + +if TYPE_CHECKING: + from .config import ValidationLevel + from .lease import HostWeightLease + from .outcomes import StoreResult + + +class ProductionSourceMode(str, Enum): + FINALIZED_MODEL = "finalized_model" + CHECKPOINT_STREAM = "checkpoint_stream" + + +class CoordinationScope(str, Enum): + SINGLE_PROCESS = "single_process" + GROUP_COHORT = "group_cohort" + ARTIFACT_BUNDLE = "artifact_bundle" + + +class LookupPhase(str, Enum): + PRE_LOAD_SAFE = "pre_load_safe" + POST_LOAD_ONLY = "post_load_only" + + +@dataclass(frozen=True) +class TensorWriteSpec: + name: str + shape: tuple[int, ...] + dtype: torch.dtype + kind: TensorKind = TensorKind.PARAMETER + role: str = "weight" + + def __post_init__(self) -> None: + if not isinstance(self.name, str) or not self.name or not isinstance(self.role, str) or not self.role: + raise ValueError("tensor write name and role must not be empty") + if not isinstance(self.shape, tuple) or any( + isinstance(dimension, bool) or not isinstance(dimension, int) or dimension < 0 for dimension in self.shape + ): + raise ValueError(f"tensor {self.name!r} has a negative shape dimension") + if not isinstance(self.dtype, torch.dtype): + raise ValueError(f"tensor {self.name!r} must declare one torch dtype") + if not isinstance(self.kind, TensorKind): + raise ValueError(f"tensor {self.name!r} must use TensorKind") + + @property + def numel(self) -> int: + return math.prod(self.shape) if self.shape else 1 + + @property + def nbytes(self) -> int: + return self.numel * torch.empty((), dtype=self.dtype).element_size() + + +class TensorFileWriter(Protocol): + def write_tensor(self, name: str, tensor: torch.Tensor) -> None: ... + + def write_rows(self, name: str, start_row: int, tensor: torch.Tensor) -> None: ... + + def close(self) -> None: ... + + def __enter__(self) -> TensorFileWriter: ... + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + _traceback: TracebackType | None, + ) -> bool | None: ... + + +class ArtifactWriter(Protocol): + def open_tensor_file( + self, + relative_name: str, + specs: tuple[TensorWriteSpec, ...], + ) -> TensorFileWriter: ... + + +@dataclass(frozen=True) +class WeightProductionSpec: + producer_id: str + outputs: tuple[WeightArtifactIdentity, ...] + source_mode: ProductionSourceMode + coordination_scope: CoordinationScope = CoordinationScope.SINGLE_PROCESS + lookup_phase: LookupPhase = LookupPhase.PRE_LOAD_SAFE + + def __post_init__(self) -> None: + if not isinstance(self.producer_id, str) or not self.producer_id or not self.outputs: + raise ValueError("producer_id and outputs must not be empty") + if not isinstance(self.outputs, tuple) or any( + not isinstance(output, WeightArtifactIdentity) for output in self.outputs + ): + raise ValueError("producer outputs must be an immutable tuple of artifact identities") + if not isinstance(self.source_mode, ProductionSourceMode): + raise ValueError("producer source mode must use ProductionSourceMode") + if not isinstance(self.coordination_scope, CoordinationScope): + raise ValueError("producer coordination scope must use CoordinationScope") + if not isinstance(self.lookup_phase, LookupPhase): + raise ValueError("producer lookup phase must use LookupPhase") + if len({output.key for output in self.outputs}) != len(self.outputs): + raise ValueError("producer outputs must have unique semantic identities") + if any(output.producer.producer_id != self.producer_id for output in self.outputs): + raise ValueError("producer outputs must name the declaring producer_id") + + +@dataclass(frozen=True) +class ProductionMetadata: + producer_schema: str + restorer_schema: str + format_metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + if ( + not isinstance(self.producer_schema, str) + or not self.producer_schema + or not isinstance(self.restorer_schema, str) + or not self.restorer_schema + ): + raise ValueError("producer and restorer schema identifiers must not be empty") + if not isinstance(self.format_metadata, CanonicalJson): + raise ValueError("production format metadata must use CanonicalJson") + + +class WeightProducer(Protocol): + """Synchronous V1 producer; callers do not preempt an active build.""" + + @property + def spec(self) -> WeightProductionSpec: ... + + def produce(self, writer: ArtifactWriter) -> ProductionMetadata: + """Produce one exact artifact synchronously or raise.""" + ... + + +@dataclass(frozen=True) +class BuildRequest: + identity: WeightArtifactIdentity + + +class WeightRestorePlan(Protocol): + """A validated, one-shot model mutation plan.""" + + def commit(self) -> None: + """Mutate the target model exactly once or raise.""" + ... + + +class WeightRestorer(Protocol): + """Validate a lease-to-model restore without mutating either input.""" + + @property + def schema(self) -> str: ... + + def plan_restore(self, model: object, lease: HostWeightLease) -> WeightRestorePlan: + """Return a validated plan; all model mutation is deferred to commit().""" + ... + + +class HostWeightStore(Protocol): + """Exact artifact access whose deadlines bound lock acquisition only.""" + + def lookup( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + deadline: float | None = None, + ) -> StoreResult: ... + + def get_or_build( + self, + request: BuildRequest, + producer: WeightProducer, + *, + validation: ValidationLevel, + deadline: float, + ) -> StoreResult: ... + + +@dataclass(frozen=True) +class RemoteArtifactDescriptor: + identity: WeightArtifactIdentity + provider_id: str + manifest_sha256: str + content_sha256: str + + +class RemoteArtifactProvider(Protocol): + @property + def provider_id(self) -> str: ... + + def probe(self, identity: WeightArtifactIdentity) -> RemoteArtifactDescriptor | None: ... + + def fetch(self, descriptor: RemoteArtifactDescriptor, writer: ArtifactWriter) -> ProductionMetadata: ... + + +__all__ = [ + "ArtifactWriter", + "BuildRequest", + "CoordinationScope", + "HostWeightStore", + "LookupPhase", + "ProductionMetadata", + "ProductionSourceMode", + "RemoteArtifactDescriptor", + "RemoteArtifactProvider", + "TensorFileWriter", + "TensorWriteSpec", + "WeightProducer", + "WeightProductionSpec", + "WeightRestorePlan", + "WeightRestorer", +] diff --git a/vllm_omni/host_weight_runtime/runtime.py b/vllm_omni/host_weight_runtime/runtime.py new file mode 100644 index 000000000..4e911e04a --- /dev/null +++ b/vllm_omni/host_weight_runtime/runtime.py @@ -0,0 +1,332 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Policy-driven exact resolution over HostWeightStore contracts.""" + +from __future__ import annotations + +import logging +import time +import uuid +from collections.abc import Callable + +from .config import HostWeightRuntimeConfig, RemoteOnMiss, RuntimeMode +from .errors import FailureCode, HostWeightError, HostWeightFailure, ResolutionAction, ResolutionStage +from .identity import CanonicalJson, WeightArtifactIdentity +from .lease import HostWeightLease +from .outcomes import ( + AttemptResult, + HostWeightResolution, + ResolutionAttempt, + ResolutionOutcome, + ResolutionReport, + StoreResult, + StoreStatus, +) +from .protocols import BuildRequest, HostWeightStore, LookupPhase, WeightProducer + +logger = logging.getLogger(__name__) + + +class HostWeightRuntime: + """Resolve one exact representation without owning model restoration.""" + + def __init__( + self, + config: HostWeightRuntimeConfig, + store: HostWeightStore | None = None, + *, + observer: Callable[[ResolutionReport], None] | None = None, + _initialization_failure: HostWeightFailure | None = None, + ) -> None: + if observer is not None and not callable(observer): + raise ValueError("Host Weight Runtime observer must be callable") + if config.mode is RuntimeMode.DISABLED and (store is not None or _initialization_failure is not None): + raise ValueError("disabled Host Weight Runtime must not initialize a store") + if config.mode is not RuntimeMode.DISABLED and (store is None) == (_initialization_failure is None): + raise ValueError("enabled Host Weight Runtime requires exactly one store or initialization failure") + self.config = config + self.store = store + self._initialization_failure = _initialization_failure + self._observer = observer + + @classmethod + def from_config( + cls, + config: HostWeightRuntimeConfig, + *, + observer: Callable[[ResolutionReport], None] | None = None, + ) -> HostWeightRuntime: + """Construct lazily so disabled mode performs no filesystem probe.""" + if config.mode is RuntimeMode.DISABLED: + return cls(config, observer=observer) + assert config.domain is not None + from .filesystem import FilesystemHostWeightStore # noqa: PLC0415 + + try: + store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity) + except HostWeightError as exc: + return cls(config, observer=observer, _initialization_failure=exc.failure) + return cls(config, store, observer=observer) + + @property + def enabled(self) -> bool: + return self.config.mode is not RuntimeMode.DISABLED + + def resolve( + self, + identity: WeightArtifactIdentity | None = None, + *, + producer: WeightProducer | None = None, + ) -> HostWeightResolution: + started = time.monotonic() + if not self.enabled: + return self._finish( + ResolutionReport( + resolution_id=uuid.uuid4().hex, + outcome=ResolutionOutcome.CANONICAL_DIRECT, + attempts=(), + elapsed_seconds=time.monotonic() - started, + ) + ) + if identity is None: + raise ValueError("enabled Host Weight Runtime requires an exact artifact identity") + if self._initialization_failure is not None: + failure = self._initialization_failure + action = self._action_for_failure(failure) + return self._finish_terminal( + ( + ResolutionAttempt( + stage=failure.stage, + result=AttemptResult.ERROR, + action=action, + elapsed_seconds=time.monotonic() - started, + failure=failure, + ), + ), + started, + ) + assert self.store is not None + deadline = started + self.config.wait.coordination_timeout_seconds + attempts: list[ResolutionAttempt] = [] + + lookup_started = time.monotonic() + local = self.store.lookup( + identity, + validation=self.config.integrity.local_lookup, + deadline=deadline, + ) + if local.status is StoreStatus.HIT: + assert local.lease is not None + attempts.append( + ResolutionAttempt( + stage=ResolutionStage.LOOKUP, + result=AttemptResult.HIT, + action=ResolutionAction.RETURN_LEASE, + elapsed_seconds=time.monotonic() - lookup_started, + ) + ) + return self._finish( + ResolutionReport( + resolution_id=local.lease.provenance.resolution_id, + outcome=ResolutionOutcome.LOCAL_HIT, + attempts=tuple(attempts), + elapsed_seconds=time.monotonic() - started, + ), + lease=local.lease, + ) + + local_action = self._next_after_local(local, producer) + attempts.append( + self._store_attempt( + ResolutionStage.LOOKUP, + local, + time.monotonic() - lookup_started, + action=local_action, + ) + ) + if local.status in {StoreStatus.TIMEOUT, StoreStatus.FAILED}: + return self._finish_terminal(tuple(attempts), started) + + producer_allowed = self._producer_allowed(producer) + if self.config.remote.on_local_miss is not RemoteOnMiss.DISABLED: + remote_failure = HostWeightFailure( + stage=ResolutionStage.REMOTE, + code=FailureCode.UNSUPPORTED, + retryable=False, + message="remote artifact providers are modeled but not implemented", + details=CanonicalJson.empty(), + ) + attempts.append( + ResolutionAttempt( + stage=ResolutionStage.REMOTE, + result=AttemptResult.SKIPPED, + action=self._action_for_failure(remote_failure), + elapsed_seconds=0.0, + failure=remote_failure, + ) + ) + return self._finish_terminal(tuple(attempts), started) + + if producer_allowed: + assert producer is not None + production_started = time.monotonic() + produced = self.store.get_or_build( + BuildRequest(identity), + producer, + validation=self.config.integrity.local_lookup, + deadline=deadline, + ) + if produced.status in {StoreStatus.HIT, StoreStatus.BUILT, StoreStatus.JOINED}: + assert produced.lease is not None + attempts.append( + ResolutionAttempt( + stage=ResolutionStage.PRODUCTION, + result=AttemptResult.SUCCESS, + action=ResolutionAction.RETURN_LEASE, + elapsed_seconds=time.monotonic() - production_started, + ) + ) + return self._finish( + ResolutionReport( + resolution_id=produced.lease.provenance.resolution_id, + outcome=ResolutionOutcome.LOCAL_PRODUCTION, + attempts=tuple(attempts), + elapsed_seconds=time.monotonic() - started, + ), + lease=produced.lease, + ) + attempts.append( + self._store_attempt( + ResolutionStage.PRODUCTION, + produced, + time.monotonic() - production_started, + action=self._terminal_action_for_store_result(produced), + ) + ) + else: + unavailable = HostWeightFailure( + stage=ResolutionStage.PRODUCTION, + code=FailureCode.PRODUCER_UNAVAILABLE, + retryable=True, + message="no allowed producer can build the exact requested representation", + details=CanonicalJson.empty(), + ) + attempts.append( + ResolutionAttempt( + stage=ResolutionStage.PRODUCTION, + result=AttemptResult.SKIPPED, + action=self._action_for_failure(unavailable), + elapsed_seconds=0.0, + failure=unavailable, + ) + ) + + return self._finish_terminal(tuple(attempts), started) + + def _mode_terminal_action(self) -> ResolutionAction: + if self.config.mode is RuntimeMode.PREFERRED: + return ResolutionAction.CANONICAL_FALLBACK + return ResolutionAction.FAIL_STARTUP + + def _action_for_failure(self, failure: HostWeightFailure | None) -> ResolutionAction: + if failure is None: + raise ValueError("failed store results require a typed failure") + if self.config.mode is RuntimeMode.PREFERRED and failure.retryable: + return ResolutionAction.CANONICAL_FALLBACK + return ResolutionAction.FAIL_STARTUP + + def _terminal_action_for_store_result(self, result: StoreResult) -> ResolutionAction: + if result.status in {StoreStatus.MISS, StoreStatus.INVALID}: + return self._mode_terminal_action() + return self._action_for_failure(result.failure) + + def _next_after_local( + self, + result: StoreResult, + producer: WeightProducer | None, + ) -> ResolutionAction: + if result.status in {StoreStatus.TIMEOUT, StoreStatus.FAILED}: + return self._terminal_action_for_store_result(result) + if self.config.remote.on_local_miss is not RemoteOnMiss.DISABLED: + return ResolutionAction.TRY_REMOTE + if self._producer_allowed(producer): + if result.status is StoreStatus.INVALID: + return ResolutionAction.QUARANTINE_THEN_PRODUCE + return ResolutionAction.TRY_PRODUCER + return self._mode_terminal_action() + + def _producer_allowed(self, producer: WeightProducer | None) -> bool: + return ( + producer is not None + and self.config.production.allow_local_build + and producer.spec.lookup_phase is LookupPhase.PRE_LOAD_SAFE + ) + + @staticmethod + def _store_attempt( + stage: ResolutionStage, + result: StoreResult, + elapsed: float, + *, + action: ResolutionAction, + ) -> ResolutionAttempt: + if result.status is StoreStatus.MISS: + return ResolutionAttempt( + stage=stage, + result=AttemptResult.MISS, + action=action, + elapsed_seconds=elapsed, + ) + if result.status is StoreStatus.TIMEOUT: + attempt_result = AttemptResult.TIMEOUT + elif result.status is StoreStatus.INVALID: + attempt_result = AttemptResult.ERROR + else: + attempt_result = AttemptResult.ERROR + return ResolutionAttempt( + stage=stage, + result=attempt_result, + action=action, + elapsed_seconds=elapsed, + failure=result.failure, + ) + + def _finish_terminal( + self, + attempts: tuple[ResolutionAttempt, ...], + started: float, + ) -> HostWeightResolution: + action = attempts[-1].action + if action is ResolutionAction.CANONICAL_FALLBACK: + outcome = ResolutionOutcome.CANONICAL_FALLBACK + elif action is ResolutionAction.FAIL_STARTUP: + outcome = ResolutionOutcome.FAILED + else: + raise ValueError(f"resolution cannot terminate with action {action.value}") + return self._finish( + ResolutionReport( + resolution_id=uuid.uuid4().hex, + outcome=outcome, + attempts=attempts, + elapsed_seconds=time.monotonic() - started, + ) + ) + + def _finish( + self, + report: ResolutionReport, + *, + lease: HostWeightLease | None = None, + ) -> HostWeightResolution: + if self._observer is not None: + try: + self._observer(report) + except Exception: + logger.exception( + "Host Weight Runtime observer failed", + extra={"resolution_id": report.resolution_id}, + ) + return HostWeightResolution(report=report, lease=lease) + + +__all__ = ["HostWeightRuntime"]