[Core] Add Host Weight Runtime foundation (#6419)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu authored and GitHub committed
1 parent
5730af0a95
commit
73b623f2f7
21 files changed
+6108
-6
No files matched your search
@@ -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
|
||||
|
||||
@@ -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).
|
||||
@@ -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)
|
||||
|
||||
@@ -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>/
|
||||
├── domain.json
|
||||
├── domain-policy.json
|
||||
├── locks/
|
||||
│ ├── <identity>.build.lock
|
||||
│ └── <identity>.artifact.lock
|
||||
├── artifacts/<identity>/
|
||||
│ ├── 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).
|
||||
@@ -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 == []
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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"
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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"]
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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("<Q", len(encoded_header)), 0)
|
||||
_pwrite_all(self._fd, encoded_header, 8)
|
||||
except BaseException:
|
||||
os.close(self._fd)
|
||||
self._path.unlink(missing_ok=True)
|
||||
owner._forget_reserved_file(total_size)
|
||||
raise
|
||||
owner._opened(self)
|
||||
|
||||
@staticmethod
|
||||
def _bytes(tensor: torch.Tensor) -> 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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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"]
|
||||
Reference in new issue
Block a user