diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index d62c2434c..0446bd9b4 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -159,6 +159,7 @@ ModelExpress/ │ │ ├── context.py # LoadContext and LoadResult │ │ ├── base.py # LoadStrategy ABC and shared helpers │ │ ├── rdma_strategy.py # RdmaStrategy (P2P GPU transfer via NIXL) +│ │ ├── server_cache_strategy.py # ServerCacheStrategy (weights from MX Server) │ │ ├── instant_tensor_strategy.py # InstantTensorStrategy (fast local safetensors) │ │ ├── model_streamer_strategy.py # ModelStreamerStrategy (S3/GCS/Azure/local) │ │ ├── gds_strategy.py # GdsStrategy (GPUDirect Storage) @@ -646,7 +647,10 @@ Loading precedence: CLI args > environment variables > config file > defaults. | `adapter.py` | `EngineAdapter` lifecycle hooks and strategy retry errors | | `vllm_loader.py` | Compatibility shim for `modelexpress.engines.vllm.loader` | | `metadata/` | Metadata publishing, source identity, heartbeat, worker manifest serving, metadata client selection, and engine-agnostic cache-artifact transfer | -| `load_strategy/` | Engine-neutral loading strategy chain: `RdmaStrategy`, `InstantTensorStrategy` (fast local safetensors), `ModelStreamerStrategy` (S3/GCS/Azure/local), `GdsStrategy`, `DefaultStrategy` | +| `load_strategy/` | Engine-neutral loading strategy chain: `RdmaStrategy`, `ServerCacheStrategy` (weights streamed from MX Server), `InstantTensorStrategy` (fast local safetensors), `ModelStreamerStrategy` (S3/GCS/Azure/local), `GdsStrategy`, `DefaultStrategy` | +| `model_client.py` | `ModelCacheClient` - `ModelService` RPCs plus stream validation for server-cached models | +| `model_snapshot.py` | Hugging Face cache layout: path validation, atomic snapshot publication, `refs/main` | +| `model_prefetch.py` | Pre-engine metadata prefetch and repo-id resolution for server-backed loading | | `engines/vllm/` | `VllmAdapter` and `MxModelLoader` map strategy hooks to vLLM loader APIs; `refit/` contains the vLLM-specific MDL installer and geometry-capture/PWAL receiver | | `engines/sglang/` | `SglangAdapter` and `MxModelLoader` - maps strategy hooks to SGLang's `remote_instance` backend | | `tensor_utils.py` | Tensor collection, checksums, storage views, `capture_tensor_attrs` | @@ -774,13 +778,35 @@ Auto-detects the best loading strategy with a prioritized chain. Each strategy i | Priority | Strategy | `is_available()` | Behavior | |---|---|---|---| | p0 | `RdmaStrategy` | NIXL available | `ListSources(READY)`, filter by `worker_rank` and runtime `accelerator`, order the survivors via the configured `SourceSelector` (`MX_P2P_SOURCE_SELECTOR`: `random` default or `rendezvous_hash`), then try candidates (max 3). Filtering before the retry slice prevents incompatible sources from exhausting the retry budget; a post-`GetMetadata` accelerator check remains as defense-in-depth. Before preparing target tensors, P2P sources must serve a manifest for the selected runtime `worker_id`; generation mismatches and transfer failures retry the next candidate, reinitializing the target first when it may have been mutated. | -| p1 | `InstantTensorStrategy` | `MX_INSTANT_TENSOR` enabled (default) + `instanttensor` installed + CUDA device + adapter implements `build_instanttensor_weight_iter` (and `apply_weight_iter`) | Load the model's own safetensors directly onto CUDA via the `instanttensor` library (distributed loading, pipelined prefetch, direct I/O, GDS when available). Reuses vLLM's built-in `--load-format instanttensor` path, so it needs no `MX_MODEL_URI`; the engine resolves and (if needed) downloads the weight files. Falls through on failure. | -| p2 | `ModelStreamerStrategy` | `MX_MODEL_URI` set + `runai_model_streamer` installed | Stream safetensors to GPU via CPU staging buffer. `MX_MODEL_URI` accepts remote URIs (`s3://`, `gs://`, `az://`), absolute local paths, or HF model IDs (resolved via `HF_HUB_CACHE`). All storage backends (S3, GCS, Azure) included by default. | -| p3 | `GdsStrategy` | Active accelerator backend supports GDS and GDS hardware is available | Load via `MxGdsLoader` (direct file-to-GPU). Falls through on failure. Reads full checkpoint tensors and slices for TP downstream — see [GDS Reads Full Checkpoint Tensors Under TP](#gds-reads-full-checkpoint-tensors-under-tp). | -| p4 | `DefaultStrategy` | Engine native fallback loader available | Native loader fallback (for vLLM, `DefaultModelLoader`, CPU-staged, auto-downloads from HF Hub). | +| p1 | `ServerCacheStrategy` | `MODEL_EXPRESS_NO_SHARED_STORAGE` enabled + server address configured + adapter implements `load_via_native` | Stream the model's weight files from ModelExpress Server into the snapshot the engine already resolved, then hand off to the engine's native loader. The cold-miss path for workers with no route to Hugging Face: the server downloads and caches the model once, and every later worker is served from that cache. Non-weight files arrive earlier, before the engine starts — see [Server-Backed Model Cache](#server-backed-model-cache). Falls through on failure. | +| p2 | `InstantTensorStrategy` | `MX_INSTANT_TENSOR` enabled (default) + `instanttensor` installed + CUDA device + adapter implements `build_instanttensor_weight_iter` (and `apply_weight_iter`) | Load the model's own safetensors directly onto CUDA via the `instanttensor` library (distributed loading, pipelined prefetch, direct I/O, GDS when available). Reuses vLLM's built-in `--load-format instanttensor` path, so it needs no `MX_MODEL_URI`; the engine resolves and (if needed) downloads the weight files. Falls through on failure. | +| p3 | `ModelStreamerStrategy` | `MX_MODEL_URI` set + `runai_model_streamer` installed | Stream safetensors to GPU via CPU staging buffer. `MX_MODEL_URI` accepts remote URIs (`s3://`, `gs://`, `az://`), absolute local paths, or HF model IDs (resolved via `HF_HUB_CACHE`). All storage backends (S3, GCS, Azure) included by default. | +| p4 | `GdsStrategy` | Active accelerator backend supports GDS and GDS hardware is available | Load via `MxGdsLoader` (direct file-to-GPU). Falls through on failure. Reads full checkpoint tensors and slices for TP downstream — see [GDS Reads Full Checkpoint Tensors Under TP](#gds-reads-full-checkpoint-tensors-under-tp). | +| p5 | `DefaultStrategy` | Engine native fallback loader available | Native loader fallback (for vLLM, `DefaultModelLoader`, CPU-staged, auto-downloads from HF Hub). | See [ModelExpress Benchmarks](BENCHMARKS.md) for measured loading-path, NIXL registration, and artifact-transfer results with explicit timing boundaries. +### Server-Backed Model Cache + +Workers without shared storage need repository files at two different moments, and only one of them is late enough for the strategy chain. + +An engine resolves the model long before it loads weights: vLLM calls `snapshot_download` while parsing engine args and rewrites `ModelConfig.model` with the resolved local path, and the tokenizer follows immediately after. Under `HF_HUB_OFFLINE=1` with an empty cache, that call fails before any loader exists. P2P cannot cover it either — it transfers GPU tensors and never repository files, so even a worker that will end up loading over RDMA still needs config and tokenizer on local disk first. + +So the two halves are fetched separately: + +| Phase | What | When | Where | +|---|---|---|---| +| Metadata | Everything except weights | Before the engine resolves the model | `model_prefetch.ensure_metadata()`, invoked from a `snapshot_download` patch | +| Weights | Weight files only | After `RdmaStrategy` finds no source | `ServerCacheStrategy` | + +Fetching metadata unconditionally does not weaken P2P-first, because no weight moves on that path — a live source still serves every byte of the weights. + +`model_client.py` wraps the `ModelService` RPCs (`EnsureModelDownloaded`, `ListModelFiles`, `StreamModelFiles`) and validates the stream: chunk offsets must be contiguous, sizes must match the manifest, and every listed file must arrive before the final marker. `model_snapshot.py` owns the local layout, writing `refs/main` alongside `snapshots//` — without that ref, `snapshot_download(local_files_only=True)` cannot resolve a repo id no matter how complete the snapshot is. Metadata is published by renaming a staging directory; weights are added to the live snapshot one atomic rename at a time, and a failure part-way through rolls back the files it already published rather than leaving a partial weight set the engine would load as complete. + +Each phase opens with `EnsureModelDownloaded` and pins everything after it to the revision that call reported, so the manifest and the stream come from one commit and a default revision that moves mid-phase cannot mix two; when the server names no revision, later calls go unpinned and the stream's commit-hash validation is what refuses a mid-phase change. Reuse of a local snapshot is gated on the same value and fails closed without it: a manifest carries only paths and sizes, so a revision that changed neither is indistinguishable from the copy on disk, and reusing it would skip the stream that would have caught the difference. The metadata phase asks with `ignore_weights=true`; the server keys its registry entry on the weight mode, so that claim cannot satisfy the weight phase's later full-weight request, and a cold server does not fetch weights before `RdmaStrategy` has had its chance at them. + +Two limits are worth knowing. The client never asks for a particular revision, only for whichever one the server resolves by default, so a worker that pinned a revision of its own gets a logged mismatch rather than the revision it wanted. And a server that already holds an unpinned model reports no revision at all, so the metadata phase restreams instead of reusing what is on disk; the files are small and the stream carries the commit, which makes restreaming the cheap way to stay correct. + Strategies handle the loading path and NIXL tensor registration. `LoadContext.accelerator_backend` centralizes accelerator-specific torch operations and capability gates for fast paths such as pool registration, VMM arena registration, and GDS. Backends that do not support those CUDA-specific paths, such as XPU, leave the gates disabled and use the generic fallback path. XPU transfer deployments still require a UCX/NIXL runtime that can register XPU device memory. Adapter hooks handle engine lifecycle such as vLLM `process_weights_after_loading`, and the chain performs best-effort metadata publication after a successful strategy. New strategies can be added by creating a new file in `load_strategy/` and registering it in `LoadStrategyChain.run()`. ### Source Selection diff --git a/docs/DEPLOYMENT.md b/docs/DEPLOYMENT.md index b89034ceb..4d27609f3 100644 --- a/docs/DEPLOYMENT.md +++ b/docs/DEPLOYMENT.md @@ -781,6 +781,41 @@ artifact transfer only within a trusted deployment, and network-isolate the MX server and worker gRPC endpoints from untrusted clients. ModelExpress does not currently sign cache artifacts. +### Server-Backed Model Cache (No Shared Storage) + +For workers that cannot reach the Hugging Face Hub themselves, ModelExpress Server can act as the only route to the model. The worker asks the server for repository files; the server downloads the model once on a cold miss and serves every later worker from its own cache. + +Weights and everything else are fetched at different times. Config, tokenizer, and index files arrive before the engine starts, because the engine resolves the model path (and fails offline) long before any weight loader runs. Weights stay behind the strategy chain, so a live P2P source is still the first choice and the server is only asked after `RdmaStrategy` finds nothing. Fetching metadata early does not weaken P2P-first — no weight moves on that path. + +Enable it on the worker: + +```yaml +MODEL_EXPRESS_URL: http://:8001 +MODEL_EXPRESS_NO_SHARED_STORAGE: "1" +MODEL_EXPRESS_CACHE_DIRECTORY: /home/dynamo/.cache/huggingface/hub +HF_HUB_CACHE: /home/dynamo/.cache/huggingface/hub +HF_HUB_OFFLINE: "1" +``` + +| Variable | Default | Description | +|----------|---------|-------------| +| `MODEL_EXPRESS_NO_SHARED_STORAGE` | `0` | Fetch repository files from ModelExpress Server. When unset, nothing changes: no extra RPCs and no change to P2P or local loading. | +| `MODEL_EXPRESS_URL` | unset | ModelExpress Server address. Required together with the switch above; without an address the feature stays off. | +| `MX_SERVER_ADDRESS` | unset | Alternative spelling of the server address, accepted for parity with the P2P client. Either variable satisfies the requirement. | +| `MODEL_EXPRESS_CACHE_DIRECTORY` | `HF_HUB_CACHE` | Where the worker installs snapshots. Point it at the same path as `HF_HUB_CACHE` so the engine reads what ModelExpress wrote. | +| `MODEL_EXPRESS_TRANSFER_CHUNK_SIZE` | `1048576` | gRPC file-stream chunk size in bytes. Values outside 1..`MAX_CHUNK_SIZE` fall back to the default rather than failing startup. | + +Requirements and limits: + +- ModelExpress Server needs a writable cache directory, egress to Hugging Face, and `HF_TOKEN` for private repositories. The worker needs none of these. +- The server must be from a release newer than v0.5.0 — the first generation that keys registry entries on the weight mode. The metadata phase claims a metadata-only download; an older server records that claim against the model name alone, which marks the model complete, so the weight phase finds nothing left to fetch and the worker falls through to a native load that an offline pod cannot perform. +- Mount the cache path as a volume shared by every container that touches it. Without a volume the snapshot lands in the container's writable layer, invisible to other containers and lost on restart. +- Requesting a pinned revision is not supported yet: the client never sets `revision` on the first call of a phase, so the server picks its default revision. Every call after that is pinned to the revision the server reported — when it named none, later calls stay unpinned and the stream's own commit-hash validation is the guard, which still refuses a mid-stream change. A worker that asked for a specific revision gets the mismatch logged, and weights whose commit differs from the local snapshot directory are refused rather than mixed in. +- Reusing a local snapshot requires the server to name its revision. A server that already holds an unpinned model answers without naming one, so the worker restreams the metadata rather than assume the copy on disk is current. Metadata is small — well under a second — and the stream carries the commit, which makes restreaming the cheap way to stay correct. +- On a cold server the metadata phase waits only for the non-weight files. The weights are downloaded later, and only if P2P found no source. The server keys its registry entry on the weight mode, so the metadata-only request does not mark the model complete. +- The server dedups the upstream download but not the per-worker stream. Concurrent workers on a cold model all wait on one Hugging Face fetch, then each streams its own copy, so N replicas starting together cost N x model size in server egress. Size the server's network accordingly, or stagger large rollouts. +- An unreachable server costs about 20 seconds per worker before loading falls through to the next strategy. That is the TCP connect timeout; a shorter deadline would abort legitimate cold-cache downloads, which can take minutes. + ### InstantTensor (Fast Local Safetensors) InstantTensor loads the model's own safetensors directly onto CUDA using distributed loading, pipelined prefetching, and direct I/O, with GPUDirect Storage when the hardware supports it. It sits right after P2P RDMA in the loading chain: when no peer source is already serving, it is the fastest local-disk path before falling back to ModelStreamer, GDS, or the native loader. Unlike ModelStreamer it needs no `MX_MODEL_URI`; it reuses vLLM's built-in `--load-format instanttensor` path, so the engine resolves the model's weight files (downloading from the Hugging Face Hub into the local cache first if they are not already local). diff --git a/modelexpress_client/python/generate_proto.sh b/modelexpress_client/python/generate_proto.sh index d01257821..6b7cf14b6 100755 --- a/modelexpress_client/python/generate_proto.sh +++ b/modelexpress_client/python/generate_proto.sh @@ -13,27 +13,32 @@ SPDX_HEADER="# SPDX-FileCopyrightText: Copyright (c) 2025-${YEAR} NVIDIA CORPORA # SPDX-License-Identifier: Apache-2.0 #" -# Generate protobuf files -echo "Generating protobuf files from ${PROTO_DIR}/p2p.proto..." -python -m grpc_tools.protoc \ - "-I${PROTO_DIR}" \ - "--python_out=${OUT_DIR}" \ - "--grpc_python_out=${OUT_DIR}" \ - "${PROTO_DIR}/p2p.proto" - -# Fix relative import in grpc file -echo "Fixing imports in p2p_pb2_grpc.py..." -tmp_file="$(mktemp)" -sed 's/^import p2p_pb2 as/from . import p2p_pb2 as/' "${OUT_DIR}/p2p_pb2_grpc.py" > "${tmp_file}" -mv "${tmp_file}" "${OUT_DIR}/p2p_pb2_grpc.py" - -# Add SPDX header to generated files -for file in "${OUT_DIR}/p2p_pb2.py" "${OUT_DIR}/p2p_pb2_grpc.py"; do - echo "Adding SPDX header to ${file}..." - tmp_file=$(mktemp) - echo "${SPDX_HEADER}" > "${tmp_file}" - cat "${file}" >> "${tmp_file}" - mv "${tmp_file}" "${file}" +PROTOS=(p2p model) + +for proto in "${PROTOS[@]}"; do + # Generate protobuf files + echo "Generating protobuf files from ${PROTO_DIR}/${proto}.proto..." + python -m grpc_tools.protoc \ + "-I${PROTO_DIR}" \ + "--python_out=${OUT_DIR}" \ + "--grpc_python_out=${OUT_DIR}" \ + "${PROTO_DIR}/${proto}.proto" + + # Fix relative import in grpc file + echo "Fixing imports in ${proto}_pb2_grpc.py..." + tmp_file="$(mktemp)" + sed "s/^import ${proto}_pb2 as/from . import ${proto}_pb2 as/" \ + "${OUT_DIR}/${proto}_pb2_grpc.py" > "${tmp_file}" + mv "${tmp_file}" "${OUT_DIR}/${proto}_pb2_grpc.py" + + # Add SPDX header to generated files + for file in "${OUT_DIR}/${proto}_pb2.py" "${OUT_DIR}/${proto}_pb2_grpc.py"; do + echo "Adding SPDX header to ${file}..." + tmp_file=$(mktemp) + echo "${SPDX_HEADER}" > "${tmp_file}" + cat "${file}" >> "${tmp_file}" + mv "${tmp_file}" "${file}" + done done echo "Done." diff --git a/modelexpress_client/python/modelexpress/__init__.py b/modelexpress_client/python/modelexpress/__init__.py index 3a553f5ea..d657be3a1 100644 --- a/modelexpress_client/python/modelexpress/__init__.py +++ b/modelexpress_client/python/modelexpress/__init__.py @@ -84,9 +84,11 @@ def register_modelexpress_loaders(): from .gds_loader import MxGdsLoader # noqa: F401 from .gds_transfer import GdsTransferManager # noqa: F401 from .metadata.publisher import PublisherThread # noqa: F401 +from .model_client import ModelCacheClient # noqa: F401 __all__ = [ "GdsTransferManager", + "ModelCacheClient", "MxClient", "MxGdsLoader", "PublisherThread", diff --git a/modelexpress_client/python/modelexpress/engines/vllm/loader.py b/modelexpress_client/python/modelexpress/engines/vllm/loader.py index 2ad3afec2..ec29f6ba4 100644 --- a/modelexpress_client/python/modelexpress/engines/vllm/loader.py +++ b/modelexpress_client/python/modelexpress/engines/vllm/loader.py @@ -13,10 +13,11 @@ Uses LoadStrategyChain to auto-detect the best loading strategy: 1. RDMA (P2P GPU transfer via NIXL) - if a source is already serving - 2. InstantTensor (fast local safetensors, direct I/O + GDS) - set MX_INSTANT_TENSOR=0 to disable - 3. ModelStreamer (S3/GCS/Azure/local via runai-model-streamer) - set MX_MODEL_URI - 4. GDS (GPUDirect Storage) - direct file-to-GPU, bypassing CPU - 5. Default (vLLM DefaultModelLoader) - standard CPU-staged loading + 2. ServerCache (stream weights from ModelExpress Server) - set MODEL_EXPRESS_NO_SHARED_STORAGE=1 + 3. InstantTensor (fast local safetensors, direct I/O + GDS) - set MX_INSTANT_TENSOR=0 to disable + 4. ModelStreamer (S3/GCS/Azure/local via runai-model-streamer) - set MX_MODEL_URI + 5. GDS (GPUDirect Storage) - direct file-to-GPU, bypassing CPU + 6. Default (vLLM DefaultModelLoader) - standard CPU-staged loading Usage: --load-format modelexpress @@ -31,7 +32,7 @@ import torch import torch.nn as nn -from ... import configure_vllm_logging, envs +from ... import configure_vllm_logging, envs, model_prefetch from ...load_strategy import LoadContext, LoadStrategyChain from ...nixl_transfer import NixlTransferManager from ...vmm.runtime import log_arena_post_load, maybe_enter_vmm_arena @@ -130,6 +131,17 @@ def load_model( def download_model(self, model_config: ModelConfig) -> None: """Download the model so it can be loaded immediately.""" + if model_prefetch.is_enabled(): + # Without shared storage this would pull the full weight set from + # Hugging Face before any strategy runs, defeating P2P-first and + # failing outright when the worker is offline. The strategy chain + # decides where the weights come from. + logger.info( + "MODEL_EXPRESS_NO_SHARED_STORAGE is set; leaving weight " + "acquisition to the ModelExpress strategy chain" + ) + return + import copy disk_config = copy.copy(self.load_config) diff --git a/modelexpress_client/python/modelexpress/engines/vllm/patches/__init__.py b/modelexpress_client/python/modelexpress/engines/vllm/patches/__init__.py index 82c67b95a..90e0613ec 100644 --- a/modelexpress_client/python/modelexpress/engines/vllm/patches/__init__.py +++ b/modelexpress_client/python/modelexpress/engines/vllm/patches/__init__.py @@ -3,12 +3,14 @@ """Runtime compatibility patches for vLLM.""" +from .patch_hf_snapshot_prefetch import patch_hf_snapshot_prefetch from .patch_humming_regex_ignore import patch_humming_regex_ignore from .patch_object_storage_format_check import patch_object_storage_format_check PATCHES = ( patch_object_storage_format_check, patch_humming_regex_ignore, + patch_hf_snapshot_prefetch, ) __all__ = ["PATCHES"] diff --git a/modelexpress_client/python/modelexpress/engines/vllm/patches/patch_hf_snapshot_prefetch.py b/modelexpress_client/python/modelexpress/engines/vllm/patches/patch_hf_snapshot_prefetch.py new file mode 100644 index 000000000..ca6ec578e --- /dev/null +++ b/modelexpress_client/python/modelexpress/engines/vllm/patches/patch_hf_snapshot_prefetch.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fetch model metadata from ModelExpress Server before the engine resolves it.""" + +from __future__ import annotations + +import logging + +from .... import model_prefetch + +logger = logging.getLogger(__name__) + +# snapshot_download is re-exported into several huggingface_hub namespaces, and +# callers bind it from whichever one they imported. vLLM reaches it through +# HfApi.snapshot_download, which calls the name bound in hf_api, so patching the +# package attribute alone would miss the path that actually fails. +_PATCH_TARGETS = ( + "huggingface_hub._snapshot_download", + "huggingface_hub.hf_api", + "huggingface_hub", +) + + +def patch_hf_snapshot_prefetch() -> bool: + """Route offline snapshot resolution through the ModelExpress model cache. + + Without shared storage the engine has no way to resolve a model before its + weight loader runs: it calls snapshot_download while parsing engine args, + finds nothing on disk, and fails under HF_HUB_OFFLINE. This hook fills the + local cache from the server first, then lets the original call proceed. + Weights are not fetched here -- P2P keeps first refusal on those. + """ + if not model_prefetch.is_enabled(): + return False + + try: + import importlib + + from huggingface_hub import _snapshot_download + except ImportError: + return False + + original = _snapshot_download.snapshot_download + if getattr(original, "__modelexpress_patched__", False): + return False + + def patched(*args, **kwargs): + """Install the snapshot from the server, then run the original call. + + The original always runs, so the engine keeps the behaviour and the + errors it already handles. Prefetch failures are logged and swallowed: + this sits on the startup path, and a ModelExpress error here would + replace the diagnostic the engine expects with an unfamiliar one. + """ + repo_id = kwargs["repo_id"] if "repo_id" in kwargs else (args[0] if args else None) + repo_type = kwargs.get("repo_type") or "model" + if isinstance(repo_id, str) and repo_type == "model": + try: + model_prefetch.ensure_metadata(repo_id, kwargs.get("revision")) + except Exception as exc: + # Let the original call report the failure the engine expects + # instead of replacing it with a ModelExpress error. + logger.warning( + "ModelExpress metadata prefetch failed for %s: %s", repo_id, exc + ) + return original(*args, **kwargs) + + patched.__modelexpress_patched__ = True + patched.__wrapped__ = original + + for module_name in _PATCH_TARGETS: + try: + module = importlib.import_module(module_name) + except ImportError: + continue + if getattr(module, "snapshot_download", None) is original: + module.snapshot_download = patched + return True diff --git a/modelexpress_client/python/modelexpress/envs.py b/modelexpress_client/python/modelexpress/envs.py index 56d4fc589..34a5cb9db 100644 --- a/modelexpress_client/python/modelexpress/envs.py +++ b/modelexpress_client/python/modelexpress/envs.py @@ -42,6 +42,9 @@ # ModelExpress server address / logging MODEL_EXPRESS_URL: Optional[str] MX_SERVER_ADDRESS: Optional[str] + MODEL_EXPRESS_CACHE_DIRECTORY: Optional[str] + MODEL_EXPRESS_NO_SHARED_STORAGE: bool + MODEL_EXPRESS_TRANSFER_CHUNK_SIZE: Optional[str] MODEL_EXPRESS_LOG_LEVEL: str MODEL_NAME: Optional[str] # Auth (client) @@ -209,6 +212,11 @@ def _env_positive_float(name: str, default: float) -> float: # Site-varying defaults: return raw (None when unset), callers add defaults. "MODEL_EXPRESS_URL": lambda: os.environ.get("MODEL_EXPRESS_URL"), "MX_SERVER_ADDRESS": lambda: os.environ.get("MX_SERVER_ADDRESS"), + "MODEL_EXPRESS_CACHE_DIRECTORY": lambda: os.environ.get("MODEL_EXPRESS_CACHE_DIRECTORY"), + "MODEL_EXPRESS_NO_SHARED_STORAGE": lambda: _env_bool("MODEL_EXPRESS_NO_SHARED_STORAGE", False), + "MODEL_EXPRESS_TRANSFER_CHUNK_SIZE": lambda: os.environ.get( + "MODEL_EXPRESS_TRANSFER_CHUNK_SIZE" + ), "MODEL_EXPRESS_LOG_LEVEL": lambda: os.environ.get("MODEL_EXPRESS_LOG_LEVEL", "").upper(), "MODEL_NAME": lambda: os.environ.get("MODEL_NAME"), # ── Auth (client) ────────────────────────────────────────────────────── diff --git a/modelexpress_client/python/modelexpress/load_strategy/__init__.py b/modelexpress_client/python/modelexpress/load_strategy/__init__.py index 485494e1a..0db43375f 100644 --- a/modelexpress_client/python/modelexpress/load_strategy/__init__.py +++ b/modelexpress_client/python/modelexpress/load_strategy/__init__.py @@ -62,6 +62,7 @@ def run(model: nn.Module, ctx: LoadContext) -> nn.Module: Raises RuntimeError if no strategy succeeds. """ from .rdma_strategy import RdmaStrategy + from .server_cache_strategy import ServerCacheStrategy from .instant_tensor_strategy import InstantTensorStrategy from .model_streamer_strategy import ModelStreamerStrategy from .gds_strategy import GdsStrategy @@ -69,6 +70,7 @@ def run(model: nn.Module, ctx: LoadContext) -> nn.Module: all_strategies: list[LoadStrategy] = [ RdmaStrategy(), + ServerCacheStrategy(), InstantTensorStrategy(), ModelStreamerStrategy(), GdsStrategy(), diff --git a/modelexpress_client/python/modelexpress/load_strategy/server_cache_strategy.py b/modelexpress_client/python/modelexpress/load_strategy/server_cache_strategy.py new file mode 100644 index 000000000..95c83b7b8 --- /dev/null +++ b/modelexpress_client/python/modelexpress/load_strategy/server_cache_strategy.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Server cache loading strategy: fetch weights from ModelExpress Server. + +Runs after RdmaStrategy, so a live P2P source is always preferred. This is the +cold-miss path: no peer is serving the model, the local cache holds only the +metadata installed before the engine started, and the worker has no route to +Hugging Face of its own. +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +from .. import model_prefetch +from ..adapter import EngineAdapter, StrategyFailed +from .base import LoadContext, LoadStrategy, _as_load_result, register_tensors +from .context import LoadResult + +logger = logging.getLogger("modelexpress.strategy_server_cache") + + +class ServerCacheStrategy(LoadStrategy): + """Install weights from the server, then load them with the engine's loader.""" + + name = "server-cache" + requires = (EngineAdapter.load_via_native,) + + def is_available(self, ctx: LoadContext) -> bool: + """Return whether the server can supply weights for this model. + + Needs the no-shared-storage switch, a server address, and a Hugging + Face repo id. The repo id may have to be recovered from the resolved + cache path, because the engine rewrites the model name in place and + loads weights in a process that never ran the prefetch. + """ + if not super().is_available(ctx): + return False + if not model_prefetch.is_enabled(): + return False + if _repo_id(ctx) is None: + logger.info( + f"[Worker {ctx.global_rank}] No Hugging Face repo id for " + f"{ctx.identity.model_name!r}, skipping server cache" + ) + return False + return True + + def load(self, result: LoadResult, ctx: LoadContext) -> LoadResult: + """Stream the weights into the resolved snapshot, then load natively. + + Raises :class:`StrategyFailed` with ``mutated=False`` while the model + is still untouched, so the chain can try the next strategy, and with + ``mutated=True`` once the engine's own loader has started writing into + it and only a reinit can recover. + """ + result = _as_load_result(result) + if ctx.adapter is None: + raise StrategyFailed( + "ModelExpress Server cache requires an engine adapter", mutated=False + ) + + repo_id = _repo_id(ctx) + if repo_id is None: + raise StrategyFailed("No Hugging Face repo id for this model", mutated=False) + + try: + snapshot_path = self._snapshot_path(ctx, repo_id) + logger.info( + f"[Worker {ctx.global_rank}] Fetching {repo_id} weights from " + f"ModelExpress Server into {snapshot_path}" + ) + self._install_weights(repo_id, snapshot_path) + except StrategyFailed: + raise + except Exception as exc: + raise StrategyFailed( + f"ModelExpress Server cache failed: {exc}", mutated=False + ) from exc + + try: + result = ctx.adapter.load_via_native(result) + result = ctx.adapter.after_native_load(result) + except Exception as exc: + raise StrategyFailed(str(exc), mutated=True) from exc + + register_tensors(result, ctx) + return result + + def _snapshot_path(self, ctx: LoadContext, repo_id: str) -> Path: + """Return the snapshot the engine resolved, installing it if needed.""" + engine_path = getattr(ctx.model_config, "model", None) + if engine_path: + candidate = Path(str(engine_path)) + if candidate.is_dir(): + return candidate + + revision = getattr(ctx.model_config, "revision", None) + snapshot_path = model_prefetch.ensure_metadata(repo_id, revision) + if snapshot_path is None: + raise StrategyFailed( + f"No local snapshot for {repo_id} and metadata prefetch did not apply", + mutated=False, + ) + return snapshot_path + + @staticmethod + def _install_weights(repo_id: str, snapshot_path: Path) -> None: + from ..model_client import ModelCacheClient + + with ModelCacheClient( + chunk_size=model_prefetch.configured_chunk_size() + ) as client: + client.install_weight_files(repo_id, snapshot_path) + + +def _repo_id(ctx: LoadContext) -> str | None: + """Resolve the repo id to ask the server for. + + ``identity.model_name`` is whatever the engine put in ModelConfig, which + vLLM overwrites with the resolved local path while parsing engine args. + model_prefetch keeps the mapping back to the original repo id. + """ + return model_prefetch.repo_id_for(ctx.identity.model_name) diff --git a/modelexpress_client/python/modelexpress/model_client.py b/modelexpress_client/python/modelexpress/model_client.py new file mode 100644 index 000000000..8c44cd248 --- /dev/null +++ b/modelexpress_client/python/modelexpress/model_client.py @@ -0,0 +1,463 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Client for the ModelExpress Server model cache. + +Wraps the ``ModelService`` RPCs and installs what they return into the local +Hugging Face cache. Two entry points, matching the two moments a worker needs +files from the server: + +- :meth:`ModelCacheClient.install_metadata_snapshot` runs before the engine + starts. It fetches everything except weights, which is enough for config and + tokenizer resolution, and leaves P2P as the first choice for the weights. +- :meth:`ModelCacheClient.install_weight_files` runs after a P2P miss and adds + the weights to the snapshot the engine already resolved. + +The server streams files relative to its own snapshot directory; layout and +atomicity live in :mod:`modelexpress.model_snapshot`. +""" + +from __future__ import annotations + +import logging +import os +from pathlib import Path +from typing import Mapping, Sequence + +import grpc + +from . import auth +from . import model_pb2 +from . import model_pb2_grpc +from .client import _get_server_url +from .model_snapshot import ( + ModelSnapshotCache, + ModelSnapshotError, + SnapshotSink, + safe_commit_hash, + split_by_weight, +) + +logger = logging.getLogger("modelexpress.model_client") + +DEFAULT_MAX_MESSAGE_SIZE = 100 * 1024 * 1024 +# The server default is 32 KiB, which costs one round trip per 32 KiB of a +# multi-GB weight file. Stay an order of magnitude below gRPC's usual 4 MiB +# message ceiling so the server never has to raise its encoding limit. +DEFAULT_CHUNK_SIZE = 1024 * 1024 +MAX_CHUNK_SIZE = (1 << 32) - 1 + + +class ModelCacheError(RuntimeError): + """Raised when the server's download or file stream cannot be trusted.""" + + +class ModelCacheClient: + """Synchronous client for the ModelExpress ``ModelService`` RPCs.""" + + def __init__( + self, + server_url: str | None = None, + cache_directory: str | os.PathLike[str] | None = None, + chunk_size: int | None = None, + max_message_size: int = DEFAULT_MAX_MESSAGE_SIZE, + ): + self.server_url = _get_server_url(server_url) + self.cache_directory = cache_directory + self.chunk_size = DEFAULT_CHUNK_SIZE if chunk_size is None else chunk_size + if not 0 < self.chunk_size <= MAX_CHUNK_SIZE: + raise ValueError(f"chunk_size must be between 1 and {MAX_CHUNK_SIZE} bytes") + if max_message_size <= 0: + raise ValueError("max_message_size must be positive") + + self._max_message_size = max_message_size + self._channel: grpc.Channel | None = None + self._stub: model_pb2_grpc.ModelServiceStub | None = None + + def __enter__(self) -> "ModelCacheClient": + return self + + def __exit__(self, *exc_info) -> None: + self.close() + + @property + def stub(self) -> model_pb2_grpc.ModelServiceStub: + """Return (and lazily create) the model-service stub.""" + if self._stub is None: + options = [ + ("grpc.max_send_message_length", self._max_message_size), + ("grpc.max_receive_message_length", self._max_message_size), + # No RPC here carries a deadline, because a cold-cache download + # is legitimately slow. Keepalive is what separates "slow" from + # "dead": without it a silently dropped connection blocks the + # engine's startup path forever, so the pod neither becomes + # ready nor crash-loops. + ("grpc.keepalive_time_ms", 30_000), + ("grpc.keepalive_timeout_ms", 10_000), + ("grpc.keepalive_permit_without_calls", 1), + ("grpc.http2.max_pings_without_data", 0), + ] + self._channel = auth.with_auth( + grpc.insecure_channel(self.server_url, options=options) + ) + self._stub = model_pb2_grpc.ModelServiceStub(self._channel) + logger.debug("ModelCacheClient connected to %s", self.server_url) + return self._stub + + def close(self) -> None: + """Close the underlying gRPC channel.""" + if self._channel is not None: + self._channel.close() + self._channel = None + self._stub = None + + # -- RPCs ----------------------------------------------------------------- + + def ensure_downloaded( + self, + model_name: str, + provider: int = model_pb2.HUGGING_FACE, + ignore_weights: bool = False, + ) -> str | None: + """Block until the server reports the model as downloaded. + + ``ignore_weights`` asks for a metadata-only download. Servers that + key their registry entries on the weight mode keep that claim separate + from a later full-weight request, so the weight phase asks again and + gets its own download. Servers older than that key on the model name + alone: the metadata-only claim registers the model as complete there, + no later weight fetch happens, and the weights become unreachable + through the cache -- see the server requirement in DEPLOYMENT.md. + + Returns the commit the server resolved the request to, or ``None`` + when it named none. A server that already holds an unpinned model + answers without naming a revision, so ``None`` is ordinary and means + only that the caller learned nothing -- not that the local cache is + current. + """ + request = model_pb2.ModelDownloadRequest( + model_name=model_name, + provider=provider, + ignore_weights=ignore_weights, + ) + for update in self.stub.EnsureModelDownloaded(request): + if update.message: + logger.info("Model %s: %s", model_name, update.message) + if update.status == model_pb2.DOWNLOADED: + return _reported_revision(model_name, update) + if update.status == model_pb2.ERROR: + raise ModelCacheError( + f"ModelExpress failed to download {model_name}: " + f"{update.message or 'unknown server error'}" + ) + if update.status != model_pb2.DOWNLOADING: + raise ModelCacheError( + f"ModelExpress reported unknown status {update.status} for {model_name}" + ) + raise ModelCacheError( + f"ModelExpress status stream ended before {model_name} was downloaded" + ) + + def list_files( + self, + model_name: str, + provider: int = model_pb2.HUGGING_FACE, + ignore_weights: bool = False, + revision: str | None = None, + ) -> dict[str, int]: + """Return the server's file manifest as {relative_path: size}. + + ``revision`` pins the manifest to one snapshot, so a server whose + default revision moves between this call and the stream cannot answer + the two from different commits. + """ + request = model_pb2.ModelFilesRequest( + model_name=model_name, + provider=provider, + ignore_weights=ignore_weights, + ) + if revision is not None: + request.revision = revision + response = self.stub.ListModelFiles(request) + return _manifest_to_dict(response) + + # -- Installation --------------------------------------------------------- + + def install_metadata_snapshot( + self, + model_name: str, + provider: int = model_pb2.HUGGING_FACE, + ) -> Path: + """Install every non-weight file as a resolvable Hugging Face snapshot. + + Asks the server for a metadata-only download, so a cold server does + not fetch the weights before ``RdmaStrategy`` has had its chance at + them. Returns the snapshot directory, reusing an existing local + snapshot only when the server named the revision it is holding and + that revision is the one already on disk. + """ + revision = self.ensure_downloaded(model_name, provider, ignore_weights=True) + manifest = self.list_files( + model_name, provider, ignore_weights=True, revision=revision + ) + # Still split: an older server ignores ignore_weights and answers with + # the whole repository. + metadata_paths, _ = split_by_weight(manifest.keys()) + if not metadata_paths: + raise ModelCacheError( + f"ModelExpress returned no non-weight files for {model_name}" + ) + expected = {path: manifest[path] for path in metadata_paths} + + cache = ModelSnapshotCache(model_name, self.cache_directory) + with cache.lock(): + existing = cache.resolve_snapshot(expected, revision) + if existing is not None: + logger.info("Reusing local snapshot for %s at %s", model_name, existing) + return existing + + staging = cache.staging() + try: + commit_hash = self._stream_into( + model_name, + provider, + metadata_paths, + expected, + staging, + expected_commit=revision, + revision=revision, + ) + snapshot_path = staging.publish(commit_hash, expected) + except BaseException: + staging.discard() + raise + + logger.info( + "Installed %d metadata files for %s at %s", + len(expected), + model_name, + snapshot_path, + ) + return snapshot_path + + def install_weight_files( + self, + model_name: str, + snapshot_path: Path, + provider: int = model_pb2.HUGGING_FACE, + ) -> None: + """Add the model's weight files to an already published snapshot. + + Refuses to write weights the server reports under a different commit + than the snapshot is named after: the engine addresses pinned + revisions by directory name, so mixing commits would hand it weights + from one revision under the name of another. + """ + revision = self.ensure_downloaded(model_name, provider) + if revision is not None and revision != snapshot_path.name: + raise ModelCacheError( + f"Server resolved {model_name} to commit {revision} but the local " + f"snapshot is {snapshot_path.name}; refusing to mix revisions" + ) + manifest = self.list_files(model_name, provider, revision=revision) + _, weight_paths = split_by_weight(manifest.keys()) + if not weight_paths: + raise ModelCacheError( + f"ModelExpress returned no weight files for {model_name}" + ) + expected = {path: manifest[path] for path in weight_paths} + + cache = ModelSnapshotCache(model_name, self.cache_directory) + with cache.lock(): + if cache.has_files(snapshot_path, expected): + logger.info("Weights for %s already present at %s", model_name, snapshot_path) + return + + patch = cache.patch(snapshot_path) + try: + self._stream_into( + model_name, + provider, + weight_paths, + expected, + patch, + expected_commit=snapshot_path.name, + revision=revision, + ) + patch.commit() + except BaseException: + patch.rollback() + raise + finally: + patch.close() + + logger.info( + "Installed %d weight files for %s at %s", + len(expected), + model_name, + snapshot_path, + ) + + # -- Stream handling ------------------------------------------------------ + + def _stream_into( + self, + model_name: str, + provider: int, + paths: Sequence[str], + expected: Mapping[str, int], + sink: SnapshotSink, + expected_commit: str | None = None, + revision: str | None = None, + ) -> str: + """Stream ``paths`` into ``sink``, validating the protocol as it goes. + + Returns the commit hash the server reported for the snapshot. When + ``expected_commit`` is given, the stream is rejected on the very first + chunk if the commit differs -- the alternative is transferring the + whole model before noticing, which for a sharded checkpoint means tens + of gigabytes thrown away. ``revision`` pins the server to one snapshot + so the manifest and the stream cannot come from different commits. + """ + request = model_pb2.ModelFilesRequest( + model_name=model_name, + provider=provider, + chunk_size=self.chunk_size, + file_selector=model_pb2.ModelFileSelector(paths=list(paths)), + ) + if revision is not None: + request.revision = revision + + commit_hash: str | None = None + received: dict[str, int] = {} + current_path: str | None = None + current_size = 0 + current_total = 0 + saw_final_marker = False + + for chunk in self.stub.StreamModelFiles(request): + if saw_final_marker: + raise ModelCacheError("Server sent data after the final stream marker") + + if commit_hash is None: + if not chunk.HasField("commit_hash"): + raise ModelCacheError("First file chunk did not carry a commit hash") + commit_hash = _require_commit_hash(chunk.commit_hash) + if expected_commit is not None and commit_hash != expected_commit: + raise ModelCacheError( + f"Server streamed commit {commit_hash} but the local snapshot " + f"is {expected_commit}; refusing to mix revisions" + ) + elif chunk.HasField("commit_hash") and chunk.commit_hash != commit_hash: + raise ModelCacheError("Server changed the commit hash mid-stream") + + relative_path = chunk.relative_path + if relative_path not in expected: + raise ModelCacheError(f"Server streamed an unrequested file: {relative_path!r}") + + if current_path != relative_path: + if current_path is not None: + raise ModelCacheError( + f"Server started {relative_path!r} before {current_path!r} finished" + ) + if relative_path in received: + raise ModelCacheError(f"Server streamed {relative_path!r} twice") + if chunk.offset != 0: + raise ModelCacheError( + f"First chunk of {relative_path!r} has offset {chunk.offset}" + ) + if chunk.total_size != expected[relative_path]: + raise ModelCacheError( + f"Size mismatch for {relative_path!r}: manifest " + f"{expected[relative_path]}, stream {chunk.total_size}" + ) + sink.begin_file(relative_path) + current_path = relative_path + current_size = 0 + current_total = chunk.total_size + elif chunk.total_size != current_total: + raise ModelCacheError(f"Server changed the size of {relative_path!r} mid-file") + + if chunk.offset != current_size: + raise ModelCacheError( + f"Unexpected offset {chunk.offset} for {relative_path!r}, " + f"expected {current_size}" + ) + if current_size + len(chunk.data) > current_total: + raise ModelCacheError(f"Data for {relative_path!r} exceeds its advertised size") + if chunk.is_last_file and not chunk.is_last_chunk: + raise ModelCacheError("Final-file marker set before the file's final chunk") + + sink.write(chunk.data) + current_size += len(chunk.data) + + if chunk.is_last_chunk: + if current_size != current_total: + raise ModelCacheError( + f"Incomplete file {relative_path!r}: received {current_size}, " + f"expected {current_total}" + ) + sink.end_file() + received[relative_path] = current_size + current_path = None + saw_final_marker = chunk.is_last_file + elif current_size == current_total: + raise ModelCacheError(f"File {relative_path!r} completed without a final chunk") + + if commit_hash is None: + raise ModelCacheError("Server streamed no model files") + if not saw_final_marker: + raise ModelCacheError("Stream ended before the final file marker") + if received != dict(expected): + missing = sorted(set(expected) - set(received)) + raise ModelCacheError(f"Stream did not match the manifest; missing files: {missing}") + return commit_hash + + +def _manifest_to_dict(manifest: model_pb2.ModelFileList) -> dict[str, int]: + """Validate a file manifest and return it keyed by relative path.""" + files: dict[str, int] = {} + for file_info in manifest.files: + relative_path = file_info.relative_path + if relative_path in files: + raise ModelCacheError(f"Server returned duplicate file path: {relative_path!r}") + files[relative_path] = file_info.size + + if not files: + raise ModelCacheError("Server returned an empty model file manifest") + total = sum(files.values()) + if total != manifest.total_size: + raise ModelCacheError( + f"Manifest total mismatch: files add up to {total} bytes, " + f"manifest advertises {manifest.total_size} bytes" + ) + return files + + +def _reported_revision(model_name: str, update: model_pb2.ModelStatusUpdate) -> str | None: + """Return the commit an update names, or None when it names none. + + An unusable value degrades to None rather than raising. This runs on the + engine's startup path, and the only thing the value buys is the reuse + shortcut: dropping it costs one metadata stream, while raising would cost + the worker its start. Nothing downstream trusts it unchecked either -- + the stream reports the commit again and validates it there. + """ + if not update.HasField("resolved_revision"): + return None + try: + return safe_commit_hash(update.resolved_revision) + except ModelSnapshotError: + logger.warning( + "ModelExpress reported an unusable revision %r for %s; ignoring it", + update.resolved_revision, + model_name, + ) + return None + + +def _require_commit_hash(commit_hash: str) -> str: + try: + return safe_commit_hash(commit_hash) + except ModelSnapshotError as exc: + raise ModelCacheError(str(exc)) from exc diff --git a/modelexpress_client/python/modelexpress/model_pb2.py b/modelexpress_client/python/modelexpress/model_pb2.py new file mode 100644 index 000000000..ccf8616fa --- /dev/null +++ b/modelexpress_client/python/modelexpress/model_pb2.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: model.proto +# Protobuf Python Version: 5.27.2 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 5, + 27, + 2, + '', + 'model.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0bmodel.proto\x12\x13model_express.model\"^\n\x12\x44\x65leteModelRequest\x12\x12\n\nmodel_name\x18\x01 \x01(\t\x12\x34\n\x08provider\x18\x02 \x01(\x0e\x32\".model_express.model.ModelProvider\"H\n\x13\x44\x65leteModelResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x14\n\x07message\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\n\n\x08_message\"\x9c\x01\n\x14ModelDownloadRequest\x12\x12\n\nmodel_name\x18\x01 \x01(\t\x12\x34\n\x08provider\x18\x02 \x01(\x0e\x32\".model_express.model.ModelProvider\x12\x16\n\x0eignore_weights\x18\x03 \x01(\x08\x12\x15\n\x08revision\x18\x04 \x01(\tH\x00\x88\x01\x01\x42\x0b\n\t_revision\"\xe7\x01\n\x11ModelStatusUpdate\x12\x12\n\nmodel_name\x18\x01 \x01(\t\x12\x30\n\x06status\x18\x02 \x01(\x0e\x32 .model_express.model.ModelStatus\x12\x14\n\x07message\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x34\n\x08provider\x18\x04 \x01(\x0e\x32\".model_express.model.ModelProvider\x12\x1e\n\x11resolved_revision\x18\x05 \x01(\tH\x01\x88\x01\x01\x42\n\n\x08_messageB\x14\n\x12_resolved_revision\"\xec\x01\n\x11ModelFilesRequest\x12\x12\n\nmodel_name\x18\x01 \x01(\t\x12\x34\n\x08provider\x18\x02 \x01(\x0e\x32\".model_express.model.ModelProvider\x12\x12\n\nchunk_size\x18\x03 \x01(\r\x12=\n\rfile_selector\x18\x04 \x01(\x0b\x32&.model_express.model.ModelFileSelector\x12\x16\n\x0eignore_weights\x18\x05 \x01(\x08\x12\x15\n\x08revision\x18\x06 \x01(\tH\x00\x88\x01\x01\x42\x0b\n\t_revision\"\"\n\x11ModelFileSelector\x12\r\n\x05paths\x18\x01 \x03(\t\"\xab\x01\n\tFileChunk\x12\x15\n\rrelative_path\x18\x01 \x01(\t\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\x12\x0e\n\x06offset\x18\x03 \x01(\x04\x12\x12\n\ntotal_size\x18\x04 \x01(\x04\x12\x15\n\ris_last_chunk\x18\x05 \x01(\x08\x12\x14\n\x0cis_last_file\x18\x06 \x01(\x08\x12\x18\n\x0b\x63ommit_hash\x18\x07 \x01(\tH\x00\x88\x01\x01\x42\x0e\n\x0c_commit_hash\"j\n\rModelFileList\x12\x12\n\nmodel_name\x18\x01 \x01(\t\x12\x31\n\x05\x66iles\x18\x02 \x03(\x0b\x32\".model_express.model.ModelFileInfo\x12\x12\n\ntotal_size\x18\x03 \x01(\x04\"4\n\rModelFileInfo\x12\x15\n\rrelative_path\x18\x01 \x01(\t\x12\x0c\n\x04size\x18\x02 \x01(\x04*9\n\x0bModelStatus\x12\x0f\n\x0b\x44OWNLOADING\x10\x00\x12\x0e\n\nDOWNLOADED\x10\x01\x12\t\n\x05\x45RROR\x10\x02*3\n\rModelProvider\x12\x10\n\x0cHUGGING_FACE\x10\x00\x12\x07\n\x03NGC\x10\x01\x12\x07\n\x03GCS\x10\x02\x32\x9a\x03\n\x0cModelService\x12l\n\x15\x45nsureModelDownloaded\x12).model_express.model.ModelDownloadRequest\x1a&.model_express.model.ModelStatusUpdate0\x01\x12\\\n\x10StreamModelFiles\x12&.model_express.model.ModelFilesRequest\x1a\x1e.model_express.model.FileChunk0\x01\x12\\\n\x0eListModelFiles\x12&.model_express.model.ModelFilesRequest\x1a\".model_express.model.ModelFileList\x12`\n\x0b\x44\x65leteModel\x12\'.model_express.model.DeleteModelRequest\x1a(.model_express.model.DeleteModelResponseb\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'model_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals['_MODELSTATUS']._serialized_start=1210 + _globals['_MODELSTATUS']._serialized_end=1267 + _globals['_MODELPROVIDER']._serialized_start=1269 + _globals['_MODELPROVIDER']._serialized_end=1320 + _globals['_DELETEMODELREQUEST']._serialized_start=36 + _globals['_DELETEMODELREQUEST']._serialized_end=130 + _globals['_DELETEMODELRESPONSE']._serialized_start=132 + _globals['_DELETEMODELRESPONSE']._serialized_end=204 + _globals['_MODELDOWNLOADREQUEST']._serialized_start=207 + _globals['_MODELDOWNLOADREQUEST']._serialized_end=363 + _globals['_MODELSTATUSUPDATE']._serialized_start=366 + _globals['_MODELSTATUSUPDATE']._serialized_end=597 + _globals['_MODELFILESREQUEST']._serialized_start=600 + _globals['_MODELFILESREQUEST']._serialized_end=836 + _globals['_MODELFILESELECTOR']._serialized_start=838 + _globals['_MODELFILESELECTOR']._serialized_end=872 + _globals['_FILECHUNK']._serialized_start=875 + _globals['_FILECHUNK']._serialized_end=1046 + _globals['_MODELFILELIST']._serialized_start=1048 + _globals['_MODELFILELIST']._serialized_end=1154 + _globals['_MODELFILEINFO']._serialized_start=1156 + _globals['_MODELFILEINFO']._serialized_end=1208 + _globals['_MODELSERVICE']._serialized_start=1323 + _globals['_MODELSERVICE']._serialized_end=1733 +# @@protoc_insertion_point(module_scope) diff --git a/modelexpress_client/python/modelexpress/model_pb2_grpc.py b/modelexpress_client/python/modelexpress/model_pb2_grpc.py new file mode 100644 index 000000000..7e125833d --- /dev/null +++ b/modelexpress_client/python/modelexpress/model_pb2_grpc.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc +import warnings + +from . import model_pb2 as model__pb2 + +GRPC_GENERATED_VERSION = '1.66.2' +GRPC_VERSION = grpc.__version__ +_version_not_supported = False + +try: + from grpc._utilities import first_version_is_lower + _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION) +except ImportError: + _version_not_supported = True + +if _version_not_supported: + raise RuntimeError( + f'The grpc package installed is at version {GRPC_VERSION},' + + f' but the generated code in model_pb2_grpc.py depends on' + + f' grpcio>={GRPC_GENERATED_VERSION}.' + + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' + ) + + +class ModelServiceStub(object): + """Model service for handling model downloads and status + """ + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.EnsureModelDownloaded = channel.unary_stream( + '/model_express.model.ModelService/EnsureModelDownloaded', + request_serializer=model__pb2.ModelDownloadRequest.SerializeToString, + response_deserializer=model__pb2.ModelStatusUpdate.FromString, + _registered_method=True) + self.StreamModelFiles = channel.unary_stream( + '/model_express.model.ModelService/StreamModelFiles', + request_serializer=model__pb2.ModelFilesRequest.SerializeToString, + response_deserializer=model__pb2.FileChunk.FromString, + _registered_method=True) + self.ListModelFiles = channel.unary_unary( + '/model_express.model.ModelService/ListModelFiles', + request_serializer=model__pb2.ModelFilesRequest.SerializeToString, + response_deserializer=model__pb2.ModelFileList.FromString, + _registered_method=True) + self.DeleteModel = channel.unary_unary( + '/model_express.model.ModelService/DeleteModel', + request_serializer=model__pb2.DeleteModelRequest.SerializeToString, + response_deserializer=model__pb2.DeleteModelResponse.FromString, + _registered_method=True) + + +class ModelServiceServicer(object): + """Model service for handling model downloads and status + """ + + def EnsureModelDownloaded(self, request, context): + """Ensure a model is downloaded and stream status updates until completion + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def StreamModelFiles(self, request, context): + """Stream model files from server to client (used when shared storage is disabled) + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def ListModelFiles(self, request, context): + """Get list of files for a model (useful for resumable transfers) + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def DeleteModel(self, request, context): + """Delete a model record from the server-side registry (used by `model clear`) + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_ModelServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'EnsureModelDownloaded': grpc.unary_stream_rpc_method_handler( + servicer.EnsureModelDownloaded, + request_deserializer=model__pb2.ModelDownloadRequest.FromString, + response_serializer=model__pb2.ModelStatusUpdate.SerializeToString, + ), + 'StreamModelFiles': grpc.unary_stream_rpc_method_handler( + servicer.StreamModelFiles, + request_deserializer=model__pb2.ModelFilesRequest.FromString, + response_serializer=model__pb2.FileChunk.SerializeToString, + ), + 'ListModelFiles': grpc.unary_unary_rpc_method_handler( + servicer.ListModelFiles, + request_deserializer=model__pb2.ModelFilesRequest.FromString, + response_serializer=model__pb2.ModelFileList.SerializeToString, + ), + 'DeleteModel': grpc.unary_unary_rpc_method_handler( + servicer.DeleteModel, + request_deserializer=model__pb2.DeleteModelRequest.FromString, + response_serializer=model__pb2.DeleteModelResponse.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'model_express.model.ModelService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('model_express.model.ModelService', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class ModelService(object): + """Model service for handling model downloads and status + """ + + @staticmethod + def EnsureModelDownloaded(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream( + request, + target, + '/model_express.model.ModelService/EnsureModelDownloaded', + model__pb2.ModelDownloadRequest.SerializeToString, + model__pb2.ModelStatusUpdate.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def StreamModelFiles(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream( + request, + target, + '/model_express.model.ModelService/StreamModelFiles', + model__pb2.ModelFilesRequest.SerializeToString, + model__pb2.FileChunk.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def ListModelFiles(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.model.ModelService/ListModelFiles', + model__pb2.ModelFilesRequest.SerializeToString, + model__pb2.ModelFileList.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def DeleteModel(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/model_express.model.ModelService/DeleteModel', + model__pb2.DeleteModelRequest.SerializeToString, + model__pb2.DeleteModelResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) diff --git a/modelexpress_client/python/modelexpress/model_prefetch.py b/modelexpress_client/python/modelexpress/model_prefetch.py new file mode 100644 index 000000000..d3a50d212 --- /dev/null +++ b/modelexpress_client/python/modelexpress/model_prefetch.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Metadata prefetch for workers without shared model storage. + +An engine needs a resolvable local snapshot long before it loads weights: +vLLM calls ``snapshot_download`` while parsing engine args, and the tokenizer +follows right after. Neither can wait for the weight loader, and neither is +served by P2P, which transfers GPU tensors and never repository files. + +So the two halves are fetched at different times. Everything except the +weights is pulled here, unconditionally, before the engine resolves the model. +The weights stay with the strategy chain, where P2P keeps first refusal and +:mod:`modelexpress.load_strategy.server_cache_strategy` only steps in on a +miss. Fetching metadata early costs one small transfer and does not weaken +P2P-first, because no weight ever moves on this path. + +This module also remembers which repo id produced which snapshot directory. +vLLM rewrites ``ModelConfig.model`` in place with the resolved local path, so +by the time a strategy runs, the repo id the server needs is gone unless +something recorded it. +""" + +from __future__ import annotations + +import logging +import os +import re +import threading +from pathlib import Path + +from . import envs + +logger = logging.getLogger("modelexpress.model_prefetch") + +COMMIT_HASH_PATTERN = re.compile(r"^[0-9a-fA-F]{40}$") + +_REPO_DIR_PREFIX = "models--" + +# Reentrant: ensure_metadata holds this across the whole install and calls +# helpers that take it again. +_lock = threading.RLock() +_snapshot_to_repo_id: dict[str, str] = {} +_installed: set[str] = set() + + +def is_enabled() -> bool: + """Return whether server-backed model fetching is configured.""" + if not envs.MODEL_EXPRESS_NO_SHARED_STORAGE: + return False + return bool(envs.MODEL_EXPRESS_URL or envs.MX_SERVER_ADDRESS) + + +def is_repo_id(model: str) -> bool: + """Return whether ``model`` looks like a Hugging Face repo id we can fetch.""" + if not model or os.path.isabs(model) or os.path.sep in os.path.dirname(model or ""): + return False + try: + from huggingface_hub.utils import validate_repo_id + except ImportError: + return False + try: + validate_repo_id(model) + except Exception: + return False + return not Path(model).exists() + + +def ensure_metadata(repo_id: str, revision: str | None = None) -> Path | None: + """Install the model's non-weight files locally, once per process. + + Returns the snapshot directory, or None when the prefetch does not apply. + Errors from the server propagate; callers on the engine's critical path + decide whether to fail or fall through. + + The install runs under the lock so that a caller arriving mid-install waits + for that result instead of walking away with None -- the engine resolves + the model right after this returns, and a None would send it looking for a + snapshot that is still being written. + """ + if not is_enabled() or not is_repo_id(repo_id): + return None + + from .model_client import ModelCacheClient + + with _lock: + if repo_id in _installed: + # Later calls in the same process (tokenizer, processor) resolve + # from the snapshot the first call installed. + return _known_snapshot(repo_id) + + with ModelCacheClient(chunk_size=configured_chunk_size()) as client: + snapshot_path = client.install_metadata_snapshot(repo_id) + + _warn_on_revision_mismatch(repo_id, revision, snapshot_path) + _installed.add(repo_id) + _snapshot_to_repo_id[_normalize(snapshot_path)] = repo_id + return snapshot_path + + +def repo_id_for(model: str | os.PathLike[str]) -> str | None: + """Map a resolved snapshot path (or a repo id) back to its repo id. + + The in-process record only covers the process that ran the prefetch. vLLM + loads weights in a separate EngineCore process, which never sees it, so the + cache layout itself has to be authoritative: a snapshot path carries the + repo id in its ``models----`` directory. + """ + model = str(model) + with _lock: + recorded = _snapshot_to_repo_id.get(_normalize(model)) + if recorded is not None: + return recorded + if is_repo_id(model): + return model + return repo_id_from_cache_path(model) + + +def repo_id_from_cache_path(path: str | os.PathLike[str]) -> str | None: + """Recover a repo id from a Hugging Face cache path, or None.""" + try: + from huggingface_hub.utils import validate_repo_id + except ImportError: + return None + + candidate = Path(str(path)) + for part in (candidate, *candidate.parents): + name = part.name + if not name.startswith(_REPO_DIR_PREFIX): + continue + repo_id = name[len(_REPO_DIR_PREFIX):].replace("--", "/") + try: + validate_repo_id(repo_id) + except Exception: + return None + return repo_id + return None + + +def reset() -> None: + """Forget prefetch state. For tests.""" + with _lock: + _snapshot_to_repo_id.clear() + _installed.clear() + + +def _known_snapshot(repo_id: str) -> Path | None: + with _lock: + for snapshot_path, known_repo_id in _snapshot_to_repo_id.items(): + if known_repo_id == repo_id: + return Path(snapshot_path) + return None + + +def _normalize(path: str | os.PathLike[str]) -> str: + return os.path.normpath(str(path)) + + +def configured_chunk_size() -> int | None: + """Return the configured transfer chunk size, or None to use the default. + + Every rejection falls back to the default rather than raising: this runs on + the engine's startup path, and a bad env var should not be the reason a + worker fails to start. + """ + from .model_client import MAX_CHUNK_SIZE + + raw = envs.MODEL_EXPRESS_TRANSFER_CHUNK_SIZE + if not raw: + return None + try: + value = int(raw) + except ValueError: + logger.warning("Invalid MODEL_EXPRESS_TRANSFER_CHUNK_SIZE=%r; using default", raw) + return None + if not 0 < value <= MAX_CHUNK_SIZE: + logger.warning( + "MODEL_EXPRESS_TRANSFER_CHUNK_SIZE=%r must be between 1 and %d; using default", + raw, + MAX_CHUNK_SIZE, + ) + return None + return value + + +def _warn_on_revision_mismatch(repo_id: str, revision: str | None, snapshot_path: Path) -> None: + """Warn when the server's snapshot is not the revision the engine asked for. + + The client pins its own follow-up calls to the revision the server + reported, but it never asks for a particular one, so the engine's choice + does not reach the server and the answer is whatever the default revision + resolves to. A mismatch is not silently unsafe -- Hugging Face addresses + pinned revisions by directory name, so the engine simply will not see this + snapshot -- but the failure that follows is unhelpful unless the real + reason is logged here. + """ + if not revision or revision == "main": + return + if COMMIT_HASH_PATTERN.fullmatch(revision) and snapshot_path.name == revision: + return + logger.warning( + "ModelExpress Server returned commit %s for %s, but this worker asked for " + "revision %r. This client does not pin a revision on the model RPCs yet, so " + "the engine will not resolve this snapshot.", + snapshot_path.name, + repo_id, + revision, + ) diff --git a/modelexpress_client/python/modelexpress/model_snapshot.py b/modelexpress_client/python/modelexpress/model_snapshot.py new file mode 100644 index 000000000..364a8b8cb --- /dev/null +++ b/modelexpress_client/python/modelexpress/model_snapshot.py @@ -0,0 +1,538 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hugging Face cache layout for model files streamed from ModelExpress Server. + +Streamed files carry paths relative to the server's snapshot directory. This +module turns them into a cache an engine can resolve while offline:: + + /models----/ + refs/main commit hash of the published snapshot + snapshots// the files themselves + +``refs/main`` is what lets ``snapshot_download(local_files_only=True)`` resolve +a repo id. Without it the engine raises ``LocalEntryNotFoundError`` even when +every file is already on disk. + +There are two write paths because their atomicity requirements differ: + +- :class:`SnapshotStaging` builds a snapshot out of band and publishes the + whole directory with a single rename. Use it before the engine starts. +- :class:`SnapshotPatch` adds files to a snapshot the engine has already + resolved, renaming one file at a time so the directory is never swapped + out from under it. +""" + +from __future__ import annotations + +import logging +import os +import shutil +import tempfile +import uuid +from contextlib import contextmanager +from fcntl import LOCK_EX, LOCK_UN, flock +from pathlib import Path +from typing import Iterator, Mapping + +from huggingface_hub.constants import HF_HUB_CACHE + +from . import envs + +logger = logging.getLogger("modelexpress.model_snapshot") + +# Mirrors ModelProviderExt::is_weight_file in +# modelexpress_common/src/providers.rs. The server uses that list to decide +# what `ignore_weights` skips, so the two must stay in sync. +WEIGHT_FILE_SUFFIXES = ( + ".bin", + ".safetensors", + ".h5", + ".msgpack", + ".ckpt.index", + ".iop", + ".gas", +) + +MAIN_REF = "main" + +_LOCK_FILE = ".modelexpress.lock" +_STAGING_PREFIX = ".modelexpress-staging-" +_STALE_PREFIX = ".modelexpress-stale-" +_TEMP_PREFIX = ".modelexpress-tmp-" +_BACKUP_PREFIX = ".modelexpress-backup-" + + +class ModelSnapshotError(RuntimeError): + """Raised when server-provided paths or the local cache cannot be trusted.""" + + +def is_weight_file(relative_path: str) -> bool: + """Return whether a repo-relative path holds model weights.""" + return relative_path.endswith(WEIGHT_FILE_SUFFIXES) + + +def split_by_weight(paths) -> tuple[list[str], list[str]]: + """Split repo-relative paths into (metadata, weights), preserving order.""" + metadata: list[str] = [] + weights: list[str] = [] + for path in paths: + (weights if is_weight_file(path) else metadata).append(path) + return metadata, weights + + +def safe_relative_path(relative_path: str) -> Path: + """Validate a server-provided path and return it as a relative Path.""" + if ( + not relative_path + or "\x00" in relative_path + or "\\" in relative_path + or relative_path.startswith("/") + ): + raise ModelSnapshotError(f"Unsafe model file path: {relative_path!r}") + parts = relative_path.split("/") + if any(part in ("", ".", "..") for part in parts): + raise ModelSnapshotError(f"Unsafe model file path: {relative_path!r}") + return Path(*parts) + + +def safe_commit_hash(commit_hash: str) -> str: + """Validate a server-provided commit hash used as a directory name.""" + if ( + not commit_hash + or commit_hash in (".", "..") + or "\x00" in commit_hash + or "/" in commit_hash + or "\\" in commit_hash + ): + raise ModelSnapshotError(f"Unsafe commit hash: {commit_hash!r}") + return commit_hash + + +def repo_dir_name(model_name: str) -> str: + """Return the Hugging Face cache directory name for a model id.""" + if ( + not model_name + or "\x00" in model_name + or "\\" in model_name + or model_name.startswith("/") + ): + raise ValueError(f"Invalid Hugging Face model name: {model_name!r}") + parts = model_name.split("/") + if any(part in ("", ".", "..") for part in parts): + raise ValueError(f"Invalid Hugging Face model name: {model_name!r}") + return f"models--{'--'.join(parts)}" + + +def resolve_cache_root(explicit: str | os.PathLike[str] | None = None) -> Path: + """Resolve the local cache root. + + Priority: explicit argument, ``MODEL_EXPRESS_CACHE_DIRECTORY``, then + huggingface_hub's own ``HF_HUB_CACHE``. + """ + if explicit is not None: + return Path(explicit).expanduser() + configured = envs.MODEL_EXPRESS_CACHE_DIRECTORY + if configured: + return Path(configured).expanduser() + return Path(HF_HUB_CACHE).expanduser() + + +def _fsync_directory(directory: Path) -> None: + descriptor = os.open(directory, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def _is_contained(path: Path, root: Path) -> bool: + try: + return path.resolve().is_relative_to(root) + except OSError: + return False + + +def _ensure_directory(directory: Path, cache_root: Path) -> None: + if directory.is_symlink(): + raise ModelSnapshotError(f"Refusing to use symlinked cache directory: {directory}") + directory.mkdir(parents=True, exist_ok=True) + if not _is_contained(directory, cache_root): + raise ModelSnapshotError(f"Cache directory resolves outside the cache root: {directory}") + + +class SnapshotSink: + """Writes one streamed file at a time below ``root``.""" + + def __init__(self, root: Path, cache_root: Path): + self._root = root + self._cache_root = cache_root + self._handle = None + self._target: Path | None = None + self._relative_path: str | None = None + + @property + def current_file(self) -> str | None: + """Repo-relative path of the file currently open, if any.""" + return self._relative_path + + def begin_file(self, relative_path: str) -> None: + """Open ``relative_path`` for writing.""" + if self._handle is not None: + raise ModelSnapshotError( + f"Cannot start {relative_path!r} while {self._relative_path!r} is open" + ) + target = self._root / safe_relative_path(relative_path) + target.parent.mkdir(parents=True, exist_ok=True) + if not _is_contained(target.parent, self._cache_root): + raise ModelSnapshotError(f"File path resolves outside the cache root: {target}") + self._target = target + self._relative_path = relative_path + self._handle = self._open(target) + + def write(self, data: bytes) -> None: + """Append a chunk to the open file.""" + if self._handle is None: + raise ModelSnapshotError("No model file is open for writing") + written = self._handle.write(data) + if written != len(data): + raise ModelSnapshotError( + f"Short local write for {self._relative_path!r}: " + f"wrote {written} of {len(data)} bytes" + ) + + def end_file(self) -> None: + """Flush, sync and finalize the open file.""" + if self._handle is None or self._target is None: + raise ModelSnapshotError("No model file is open for writing") + self._handle.flush() + os.fsync(self._handle.fileno()) + self._handle.close() + self._handle = None + self._finalize(self._target) + self._target = None + self._relative_path = None + + def close(self) -> None: + """Drop a partially written file. Safe to call more than once.""" + if self._handle is None: + return + self._handle.close() + self._handle = None + if self._target is not None: + self._discard(self._target) + self._target = None + self._relative_path = None + + def _open(self, target: Path): + raise NotImplementedError + + def _finalize(self, target: Path) -> None: + raise NotImplementedError + + def _discard(self, target: Path) -> None: + raise NotImplementedError + + +class SnapshotStaging(SnapshotSink): + """Collects a snapshot in a staging directory, then publishes it atomically.""" + + def __init__(self, cache: "ModelSnapshotCache"): + _ensure_directory(cache.repo_root, cache.cache_root) + staging_path = Path( + tempfile.mkdtemp(prefix=_STAGING_PREFIX, dir=cache.repo_root) + ) + super().__init__(staging_path, cache.cache_root) + self._cache = cache + self._staging_path: Path | None = staging_path + + @property + def path(self) -> Path: + """Staging directory backing this snapshot.""" + if self._staging_path is None: + raise ModelSnapshotError("Staging directory has already been consumed") + return self._staging_path + + def publish(self, commit_hash: str, expected_files: Mapping[str, int]) -> Path: + """Move the staged files into ``snapshots/`` and update refs/main.""" + staging_path = self.path + commit_hash = safe_commit_hash(commit_hash) + snapshots_root = self._cache.repo_root / "snapshots" + _ensure_directory(snapshots_root, self._cache.cache_root) + snapshot_path = snapshots_root / commit_hash + + if self._cache.has_files(snapshot_path, expected_files): + logger.info( + "Snapshot %s already complete, discarding staged copy", snapshot_path + ) + shutil.rmtree(staging_path, ignore_errors=True) + self._staging_path = None + self._cache.write_main_ref(commit_hash) + return snapshot_path + + if snapshot_path.is_dir() and not snapshot_path.is_symlink(): + # Same commit means same content, so the directory already on disk + # holds files this manifest never mentions -- weights, above all. + # Merge into it rather than swapping it out, or installing metadata + # would delete a weight set nothing here checks for. + self._merge_into(snapshot_path) + self._cache.write_main_ref(commit_hash) + self._staging_path = None + return snapshot_path + + stale_path: Path | None = None + if snapshot_path.exists() or snapshot_path.is_symlink(): + stale_path = self._cache.repo_root / f"{_STALE_PREFIX}{uuid.uuid4().hex}" + os.replace(snapshot_path, stale_path) + + try: + os.replace(staging_path, snapshot_path) + _fsync_directory(snapshots_root) + self._cache.write_main_ref(commit_hash) + except BaseException: + # BaseException, not Exception: a KeyboardInterrupt here would + # otherwise strand the moved-aside directory with no owner and no + # cleanup path, leaking a partial model's worth of disk. + if stale_path is not None and not snapshot_path.exists(): + os.replace(stale_path, snapshot_path) + raise + self._staging_path = None + + if stale_path is not None: + try: + shutil.rmtree(stale_path) + except OSError: + logger.warning("Failed to clean up stale snapshot %s", stale_path) + return snapshot_path + + def _merge_into(self, snapshot_path: Path) -> None: + """Move every staged file into an existing snapshot, one rename at a time. + + Staging and the snapshot share a filesystem, so each rename is atomic: + a reader sees either the old file or the new one, never a partial write. + """ + staging_path = self.path + touched_dirs: set[Path] = set() + for source in sorted(staging_path.rglob("*")): + if source.is_dir(): + continue + target = snapshot_path / source.relative_to(staging_path) + target.parent.mkdir(parents=True, exist_ok=True) + if not _is_contained(target.parent, self._cache.cache_root): + raise ModelSnapshotError( + f"Staged file resolves outside the cache root: {target}" + ) + os.replace(source, target) + touched_dirs.add(target.parent) + for directory in touched_dirs: + _fsync_directory(directory) + shutil.rmtree(staging_path, ignore_errors=True) + + def discard(self) -> None: + """Remove the staging directory if it was never published.""" + self.close() + if self._staging_path is not None: + shutil.rmtree(self._staging_path, ignore_errors=True) + self._staging_path = None + + def _open(self, target: Path): + return target.open("xb") + + def _finalize(self, target: Path) -> None: + return None + + def _discard(self, target: Path) -> None: + target.unlink(missing_ok=True) + + +class SnapshotPatch(SnapshotSink): + """Adds files to a published snapshot one atomic rename at a time.""" + + def __init__(self, cache: "ModelSnapshotCache", snapshot_path: Path): + if not snapshot_path.is_dir(): + raise ModelSnapshotError(f"Snapshot directory does not exist: {snapshot_path}") + if not _is_contained(snapshot_path, cache.cache_root): + raise ModelSnapshotError( + f"Snapshot resolves outside the cache root: {snapshot_path}" + ) + super().__init__(snapshot_path, cache.cache_root) + self._temp_paths: dict[Path, Path] = {} + self._published: list[Path] = [] + self._backups: dict[Path, Path] = {} + + def commit(self) -> None: + """Make this patch final, dropping the backups it took. + + Call this only once every file has arrived. Afterwards + :meth:`rollback` has nothing to undo, so a caller that commits and + then fails cannot delete the files it just published. + """ + touched: set[Path] = set() + for backup in self._backups.values(): + backup.unlink(missing_ok=True) + touched.add(backup.parent) + self._backups.clear() + self._published.clear() + for directory in touched: + _fsync_directory(directory) + + def rollback(self) -> None: + """Undo this patch, leaving the snapshot as it was before it started. + + A half-applied patch is worse than none: the engine would see a subset + of the weights and load it as if it were complete. A file this patch + replaced is restored from its backup rather than left deleted -- + rolling back a refresh must not cost the snapshot a shard it already + had. + """ + self.close() + touched: set[Path] = set() + while self._published: + target = self._published.pop() + target.unlink(missing_ok=True) + backup = self._backups.pop(target, None) + if backup is not None: + os.replace(backup, target) + touched.add(target.parent) + # A backup with no published file means the rename never landed; the + # original is the copy to put back, not the one to drop. + for target, backup in self._backups.items(): + if target.exists(): + backup.unlink(missing_ok=True) + else: + os.replace(backup, target) + touched.add(target.parent) + self._backups.clear() + for directory in touched: + _fsync_directory(directory) + + def _open(self, target: Path): + temp_path = target.parent / f"{_TEMP_PREFIX}{uuid.uuid4().hex}-{target.name}" + self._temp_paths[target] = temp_path + return temp_path.open("xb") + + def _finalize(self, target: Path) -> None: + temp_path = self._temp_paths.pop(target) + # os.replace overwrites, so an existing file is gone the moment the + # rename lands. Move it aside first or rollback has nothing to restore. + if target.exists() or target.is_symlink(): + backup = target.parent / f"{_BACKUP_PREFIX}{uuid.uuid4().hex}-{target.name}" + os.replace(target, backup) + self._backups[target] = backup + os.replace(temp_path, target) + _fsync_directory(target.parent) + self._published.append(target) + + def _discard(self, target: Path) -> None: + temp_path = self._temp_paths.pop(target, None) + if temp_path is not None: + temp_path.unlink(missing_ok=True) + + +class ModelSnapshotCache: + """One Hugging Face repo directory inside a local cache root.""" + + def __init__( + self, + model_name: str, + cache_root: str | os.PathLike[str] | None = None, + ): + self.model_name = model_name + root = resolve_cache_root(cache_root) + root.mkdir(parents=True, exist_ok=True) + self.cache_root = root.resolve() + self.repo_root = self.cache_root / repo_dir_name(model_name) + + @contextmanager + def lock(self) -> Iterator[None]: + """Serialize cache writes across the workers sharing this directory.""" + _ensure_directory(self.repo_root, self.cache_root) + lock_path = self.repo_root / _LOCK_FILE + with lock_path.open("a", encoding="utf-8") as handle: + flock(handle.fileno(), LOCK_EX) + try: + yield + finally: + flock(handle.fileno(), LOCK_UN) + + def snapshot_path(self, commit_hash: str) -> Path: + """Return the directory a given commit's snapshot lives in.""" + return self.repo_root / "snapshots" / safe_commit_hash(commit_hash) + + def read_main_ref(self) -> str | None: + """Return the commit hash refs/main points at, or None.""" + ref_path = self.repo_root / "refs" / MAIN_REF + if not ref_path.is_file() or ref_path.is_symlink(): + return None + try: + return safe_commit_hash(ref_path.read_text(encoding="utf-8").strip()) + except (OSError, UnicodeError, ModelSnapshotError): + return None + + def write_main_ref(self, commit_hash: str) -> None: + """Point refs/main at ``commit_hash``, replacing any previous value.""" + commit_hash = safe_commit_hash(commit_hash) + refs_root = self.repo_root / "refs" + _ensure_directory(refs_root, self.cache_root) + temp_ref = refs_root / f"{_TEMP_PREFIX}{uuid.uuid4().hex}" + try: + with temp_ref.open("x", encoding="utf-8") as ref_file: + ref_file.write(commit_hash) + ref_file.flush() + os.fsync(ref_file.fileno()) + os.replace(temp_ref, refs_root / MAIN_REF) + _fsync_directory(refs_root) + finally: + temp_ref.unlink(missing_ok=True) + + def has_files(self, snapshot_path: Path, expected_files: Mapping[str, int]) -> bool: + """Return whether every expected file is present at its expected size.""" + if not snapshot_path.is_dir(): + return False + try: + if not _is_contained(snapshot_path, self.cache_root): + return False + for relative_path, expected_size in expected_files.items(): + file_path = snapshot_path / safe_relative_path(relative_path) + if not file_path.is_file(): + return False + if not _is_contained(file_path, self.cache_root): + return False + if file_path.stat().st_size != expected_size: + return False + except (OSError, ModelSnapshotError): + return False + return True + + def resolve_snapshot( + self, + expected_files: Mapping[str, int], + expected_commit: str | None, + ) -> Path | None: + """Return the snapshot refs/main points at when it holds every file. + + ``expected_commit`` is the revision the server reported for this + request. Reuse fails closed unless it matches: a file manifest carries + only paths and sizes, so a revision that changed neither is + indistinguishable from the one already on disk, and reusing it would + hand the engine stale files without ever opening a stream to notice. + ``None`` means the server named no revision, which is not proof of + anything and so never justifies reuse. + """ + if expected_commit is None: + return None + commit_hash = self.read_main_ref() + if commit_hash is None or commit_hash != expected_commit: + return None + snapshot_path = self.snapshot_path(commit_hash) + if self.has_files(snapshot_path, expected_files): + return snapshot_path + return None + + def staging(self) -> SnapshotStaging: + """Open a staging directory for a fresh snapshot.""" + return SnapshotStaging(self) + + def patch(self, snapshot_path: Path) -> SnapshotPatch: + """Open a writer that adds files to an already published snapshot.""" + return SnapshotPatch(self, snapshot_path) diff --git a/modelexpress_client/python/tests/test_envs.py b/modelexpress_client/python/tests/test_envs.py index b87a8370a..c8778913a 100644 --- a/modelexpress_client/python/tests/test_envs.py +++ b/modelexpress_client/python/tests/test_envs.py @@ -22,6 +22,9 @@ def test_defaults_when_unset(monkeypatch): "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", "MODEL_EXPRESS_URL", "MX_SERVER_ADDRESS", + "MODEL_EXPRESS_CACHE_DIRECTORY", + "MODEL_EXPRESS_NO_SHARED_STORAGE", + "MODEL_EXPRESS_TRANSFER_CHUNK_SIZE", "MX_GDS_TIMEOUT", "MX_HEARTBEAT_INTERVAL_SECS", "MX_RESHARD_FUSED_WIRE", @@ -40,6 +43,9 @@ def test_defaults_when_unset(monkeypatch): assert envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR is None assert envs.MODEL_EXPRESS_URL is None assert envs.MX_SERVER_ADDRESS is None + assert envs.MODEL_EXPRESS_CACHE_DIRECTORY is None + assert envs.MODEL_EXPRESS_NO_SHARED_STORAGE is False + assert envs.MODEL_EXPRESS_TRANSFER_CHUNK_SIZE is None assert envs.MX_GDS_TIMEOUT == pytest.approx(120.0) assert envs.MX_HEARTBEAT_INTERVAL_SECS == 30 assert envs.MX_RESHARD_FUSED_WIRE is True @@ -78,6 +84,15 @@ def test_bool_parsing(monkeypatch, caplog): monkeypatch.setenv("MX_VMM_ARENA", "1") assert envs.MX_VMM_ARENA is True + for truthy in ("1", "TRUE", "yes", "On"): + monkeypatch.setenv("MODEL_EXPRESS_NO_SHARED_STORAGE", truthy) + assert envs.MODEL_EXPRESS_NO_SHARED_STORAGE is True + for falsy in ("0", "FALSE", "no", "Off"): + monkeypatch.setenv("MODEL_EXPRESS_NO_SHARED_STORAGE", falsy) + assert envs.MODEL_EXPRESS_NO_SHARED_STORAGE is False + monkeypatch.setenv("MODEL_EXPRESS_NO_SHARED_STORAGE", "maybe") + assert envs.MODEL_EXPRESS_NO_SHARED_STORAGE is False + for truthy in ("1", "TRUE", "yes", "On"): monkeypatch.setenv("MX_ARTIFACT_TRANSFER", truthy) assert envs.MX_ARTIFACT_TRANSFER is True diff --git a/modelexpress_client/python/tests/test_hf_snapshot_prefetch_patch.py b/modelexpress_client/python/tests/test_hf_snapshot_prefetch_patch.py new file mode 100644 index 000000000..3393d3222 --- /dev/null +++ b/modelexpress_client/python/tests/test_hf_snapshot_prefetch_patch.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the snapshot_download prefetch hook.""" + +import importlib + +import pytest + +from modelexpress import model_prefetch +from modelexpress.engines.vllm.patches.patch_hf_snapshot_prefetch import ( + _PATCH_TARGETS, + patch_hf_snapshot_prefetch, +) + +REPO = "org/model" + + +# Captured at import time, before anything here can patch it. Restoring from a +# fixture-local snapshot is not enough: the patch mutates huggingface_hub +# process-wide, and a stub that survives teardown silently reroutes every later +# test that resolves a model. +_PRISTINE = importlib.import_module("huggingface_hub._snapshot_download").snapshot_download + +# Which namespaces bind the name varies by version: 1.8 has HfApi import it +# inside the method body, older releases bind it at hf_api module level. +_BOUND_MODULES = [ + module + for module in (importlib.import_module(name) for name in _PATCH_TARGETS) + if getattr(module, "snapshot_download", None) is _PRISTINE +] + + +def _modules(): + return _BOUND_MODULES + + +@pytest.fixture(autouse=True) +def restore_snapshot_download(): + yield + for module in _BOUND_MODULES: + module.snapshot_download = _PRISTINE + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + for name in ("MODEL_EXPRESS_NO_SHARED_STORAGE", "MODEL_EXPRESS_URL", "MX_SERVER_ADDRESS"): + monkeypatch.delenv(name, raising=False) + monkeypatch.delenv("MX_DISABLE_PATCHES", raising=False) + model_prefetch.reset() + + +@pytest.fixture +def enabled(monkeypatch): + monkeypatch.setenv("MODEL_EXPRESS_NO_SHARED_STORAGE", "1") + monkeypatch.setenv("MODEL_EXPRESS_URL", "http://mx:8001") + + +@pytest.fixture +def recorded(monkeypatch): + calls = [] + + def fake_ensure(repo_id, revision=None): + calls.append((repo_id, revision)) + return None + + monkeypatch.setattr(model_prefetch, "ensure_metadata", fake_ensure) + return calls + + +@pytest.fixture +def stub_download(): + """Replace the real downloader so patched calls stay offline. + + Deliberately not monkeypatch: restore_snapshot_download is the single + owner of this global, and two undo mechanisms racing over the same + attribute is what let a stub escape into the rest of the suite. + """ + calls = [] + + def fake_snapshot_download(*args, **kwargs): + calls.append((args, kwargs)) + return "/local/snapshot" + + for module in _BOUND_MODULES: + module.snapshot_download = fake_snapshot_download + return calls + + +class TestPatchGating: + def test_no_op_when_disabled(self, stub_download): + assert patch_hf_snapshot_prefetch() is False + for module in _modules(): + assert getattr(module.snapshot_download, "__modelexpress_patched__", False) is False + + def test_applies_when_enabled(self, enabled, stub_download): + assert patch_hf_snapshot_prefetch() is True + for module in _modules(): + assert module.snapshot_download.__modelexpress_patched__ is True + + def test_is_idempotent(self, enabled, stub_download): + assert patch_hf_snapshot_prefetch() is True + assert patch_hf_snapshot_prefetch() is False + + def test_disable_patches_env_skips_it(self, enabled, stub_download, monkeypatch): + from modelexpress.patches import apply_patches + + monkeypatch.setenv("MX_DISABLE_PATCHES", "1") + apply_patches([patch_hf_snapshot_prefetch]) + for module in _modules(): + assert getattr(module.snapshot_download, "__modelexpress_patched__", False) is False + + +class TestPatchedBehavior: + def test_prefetches_then_delegates(self, enabled, stub_download, recorded): + patch_hf_snapshot_prefetch() + import huggingface_hub + + result = huggingface_hub.snapshot_download(REPO, revision="abc") + + assert recorded == [(REPO, "abc")] + assert result == "/local/snapshot" + + def test_reaches_the_hf_api_call_path(self, enabled, stub_download, recorded): + """The path vLLM actually takes: HfApi().snapshot_download. + + This is the call in the issue's traceback, and it resolves the name + differently across huggingface_hub versions -- module-level in older + releases, a local import inside the method in 1.8. Patching the source + module has to cover both. + """ + patch_hf_snapshot_prefetch() + from huggingface_hub import HfApi + + HfApi().snapshot_download(repo_id=REPO) + + assert recorded == [(REPO, None)] + + def test_repo_id_as_keyword(self, enabled, stub_download, recorded): + patch_hf_snapshot_prefetch() + import huggingface_hub + + huggingface_hub.snapshot_download(repo_id=REPO) + + assert recorded == [(REPO, None)] + + def test_skips_non_model_repos(self, enabled, stub_download, recorded): + patch_hf_snapshot_prefetch() + import huggingface_hub + + huggingface_hub.snapshot_download(REPO, repo_type="dataset") + + assert recorded == [] + + def test_prefetch_failure_keeps_original_semantics(self, enabled, stub_download, monkeypatch): + def failing_ensure(repo_id, revision=None): + raise RuntimeError("server unreachable") + + monkeypatch.setattr(model_prefetch, "ensure_metadata", failing_ensure) + patch_hf_snapshot_prefetch() + import huggingface_hub + + assert huggingface_hub.snapshot_download(REPO) == "/local/snapshot" + assert len(stub_download) == 1 + + def test_arguments_are_forwarded_untouched(self, enabled, stub_download, recorded): + patch_hf_snapshot_prefetch() + import huggingface_hub + + huggingface_hub.snapshot_download( + REPO, revision="abc", local_files_only=True, allow_patterns=["*.json"] + ) + + args, kwargs = stub_download[0] + assert args == (REPO,) + assert kwargs["revision"] == "abc" + assert kwargs["local_files_only"] is True + assert kwargs["allow_patterns"] == ["*.json"] diff --git a/modelexpress_client/python/tests/test_model_client.py b/modelexpress_client/python/tests/test_model_client.py new file mode 100644 index 000000000..3bf801088 --- /dev/null +++ b/modelexpress_client/python/tests/test_model_client.py @@ -0,0 +1,703 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the ModelExpress model-cache client and its stream validation.""" + +import pytest + +from modelexpress import model_pb2 +from modelexpress.model_client import ModelCacheClient, ModelCacheError +from modelexpress.model_snapshot import ModelSnapshotCache, ModelSnapshotError + +COMMIT = "c" * 40 +MODEL = "org/model" + + +def chunk( + relative_path, + data, + *, + offset=0, + total_size=None, + is_last_chunk=True, + is_last_file=False, + commit_hash=None, +): + payload = model_pb2.FileChunk( + relative_path=relative_path, + data=data, + offset=offset, + total_size=len(data) if total_size is None else total_size, + is_last_chunk=is_last_chunk, + is_last_file=is_last_file, + ) + if commit_hash is not None: + payload.commit_hash = commit_hash + return payload + + +def whole_file(relative_path, data, *, is_last_file=False, commit_hash=None): + return chunk( + relative_path, + data, + is_last_chunk=True, + is_last_file=is_last_file, + commit_hash=commit_hash, + ) + + +class FakeStub: + """Records requests and replays canned ModelService responses.""" + + def __init__(self, *, files=None, chunks=None, updates=None, resolved_revision=None): + self.files = files or {} + self.chunks = chunks or [] + self.updates = updates + self.resolved_revision = resolved_revision + self.stream_requests = [] + self.list_requests = [] + self.download_requests = [] + + def EnsureModelDownloaded(self, request): + self.download_requests.append(request) + updates = self.updates + if updates is None: + update = model_pb2.ModelStatusUpdate( + model_name=request.model_name, status=model_pb2.DOWNLOADED + ) + # A server that already holds an unpinned model names no revision, + # so leaving this unset is the ordinary warm-cache answer. + if self.resolved_revision is not None: + update.resolved_revision = self.resolved_revision + updates = [update] + return iter(updates) + + def ListModelFiles(self, request): + self.list_requests.append(request) + return model_pb2.ModelFileList( + model_name=request.model_name, + files=[ + model_pb2.ModelFileInfo(relative_path=path, size=size) + for path, size in self.files.items() + ], + total_size=sum(self.files.values()), + ) + + def StreamModelFiles(self, request): + self.stream_requests.append(request) + return iter(self.chunks) + + +def make_client(tmp_path, stub, **kwargs): + client = ModelCacheClient(server_url="localhost:1", cache_directory=tmp_path, **kwargs) + client._stub = stub + return client + + +@pytest.fixture(autouse=True) +def no_cache_env(monkeypatch): + monkeypatch.delenv("MODEL_EXPRESS_CACHE_DIRECTORY", raising=False) + + +class TestConstruction: + def test_rejects_zero_chunk_size(self, tmp_path): + with pytest.raises(ValueError): + ModelCacheClient(cache_directory=tmp_path, chunk_size=0) + + def test_rejects_zero_max_message_size(self, tmp_path): + with pytest.raises(ValueError): + ModelCacheClient(cache_directory=tmp_path, max_message_size=0) + + +class TestEnsureDownloaded: + def test_returns_on_downloaded(self, tmp_path): + stub = FakeStub( + updates=[ + model_pb2.ModelStatusUpdate(model_name=MODEL, status=model_pb2.DOWNLOADING), + model_pb2.ModelStatusUpdate(model_name=MODEL, status=model_pb2.DOWNLOADED), + ] + ) + make_client(tmp_path, stub).ensure_downloaded(MODEL) + + assert stub.download_requests[0].ignore_weights is False + + def test_reports_the_resolved_revision(self, tmp_path): + stub = FakeStub(resolved_revision=COMMIT) + assert make_client(tmp_path, stub).ensure_downloaded(MODEL) == COMMIT + + def test_none_when_the_server_names_no_revision(self, tmp_path): + stub = FakeStub() + assert make_client(tmp_path, stub).ensure_downloaded(MODEL) is None + + def test_unusable_revision_degrades_to_none(self, tmp_path): + """A bad value costs the reuse shortcut, not the worker's start.""" + update = model_pb2.ModelStatusUpdate(model_name=MODEL, status=model_pb2.DOWNLOADED) + update.resolved_revision = "../escape" + stub = FakeStub(updates=[update]) + + assert make_client(tmp_path, stub).ensure_downloaded(MODEL) is None + + def test_raises_on_error_status(self, tmp_path): + stub = FakeStub( + updates=[ + model_pb2.ModelStatusUpdate( + model_name=MODEL, status=model_pb2.ERROR, message="no disk" + ) + ] + ) + with pytest.raises(ModelCacheError, match="no disk"): + make_client(tmp_path, stub).ensure_downloaded(MODEL) + + def test_raises_when_stream_ends_early(self, tmp_path): + stub = FakeStub( + updates=[ + model_pb2.ModelStatusUpdate(model_name=MODEL, status=model_pb2.DOWNLOADING) + ] + ) + with pytest.raises(ModelCacheError, match="ended before"): + make_client(tmp_path, stub).ensure_downloaded(MODEL) + + +class TestListFiles: + def test_returns_manifest(self, tmp_path): + stub = FakeStub(files={"config.json": 2, "model.safetensors": 7}) + assert make_client(tmp_path, stub).list_files(MODEL) == { + "config.json": 2, + "model.safetensors": 7, + } + + def test_rejects_empty_manifest(self, tmp_path): + stub = FakeStub(files={}) + with pytest.raises(ModelCacheError, match="empty model file manifest"): + make_client(tmp_path, stub).list_files(MODEL) + + def test_rejects_total_size_mismatch(self, tmp_path, monkeypatch): + stub = FakeStub(files={"config.json": 2}) + original = stub.ListModelFiles + + def lying_list(request): + response = original(request) + response.total_size = 999 + return response + + stub.ListModelFiles = lying_list + with pytest.raises(ModelCacheError, match="total mismatch"): + make_client(tmp_path, stub).list_files(MODEL) + + +class TestInstallMetadataSnapshot: + def test_requests_only_non_weight_files(self, tmp_path): + stub = FakeStub( + files={"config.json": 2, "model.safetensors": 7, "tokenizer.json": 2}, + chunks=[ + whole_file("config.json", b"{}", commit_hash=COMMIT), + whole_file("tokenizer.json", b"[]", is_last_file=True), + ], + ) + snapshot = make_client(tmp_path, stub).install_metadata_snapshot(MODEL) + + assert list(stub.stream_requests[0].file_selector.paths) == [ + "config.json", + "tokenizer.json", + ] + assert (snapshot / "config.json").read_bytes() == b"{}" + assert not (snapshot / "model.safetensors").exists() + assert snapshot.name == COMMIT + + def test_writes_main_ref(self, tmp_path): + stub = FakeStub( + files={"config.json": 2}, + chunks=[whole_file("config.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + ) + make_client(tmp_path, stub).install_metadata_snapshot(MODEL) + + cache = ModelSnapshotCache(MODEL, tmp_path) + assert cache.read_main_ref() == COMMIT + + def test_reuses_existing_snapshot_when_the_revision_matches(self, tmp_path): + stub = FakeStub( + files={"config.json": 2}, + chunks=[whole_file("config.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + resolved_revision=COMMIT, + ) + client = make_client(tmp_path, stub) + first = client.install_metadata_snapshot(MODEL) + second = client.install_metadata_snapshot(MODEL) + + assert first == second + assert len(stub.stream_requests) == 1 + + def test_restreams_when_the_server_names_no_revision(self, tmp_path): + """Reuse fails closed: a server that names no revision proves nothing.""" + stub = FakeStub( + files={"config.json": 2}, + chunks=[whole_file("config.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + ) + client = make_client(tmp_path, stub) + first = client.install_metadata_snapshot(MODEL) + second = client.install_metadata_snapshot(MODEL) + + assert first == second + assert len(stub.stream_requests) == 2 + + def test_does_not_reuse_a_stale_snapshot_behind_an_advanced_main(self, tmp_path): + """Same file names and sizes, new commit -- the manifest cannot tell. + + The server's default revision moves while the local snapshot keeps a + matching manifest. Reuse must not hand the engine the old files. + """ + stub = FakeStub( + files={"config.json": 2}, + chunks=[whole_file("config.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + resolved_revision=COMMIT, + ) + client = make_client(tmp_path, stub) + first = client.install_metadata_snapshot(MODEL) + + moved = "d" * 40 + stub.resolved_revision = moved + stub.chunks = [whole_file("config.json", b"{}", is_last_file=True, commit_hash=moved)] + second = client.install_metadata_snapshot(MODEL) + + assert first != second + assert second.name == moved + assert len(stub.stream_requests) == 2 + + def test_pins_the_manifest_and_stream_to_the_reported_revision(self, tmp_path): + stub = FakeStub( + files={"config.json": 2}, + chunks=[whole_file("config.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + resolved_revision=COMMIT, + ) + make_client(tmp_path, stub).install_metadata_snapshot(MODEL) + + assert stub.list_requests[0].revision == COMMIT + assert stub.stream_requests[0].revision == COMMIT + + def test_metadata_phase_asks_for_a_metadata_only_download(self, tmp_path): + """A cold server must not fetch the weights before RdmaStrategy runs. + + The server keys its registry entry on the weight mode, so this claim + does not satisfy the weight phase's later full-weight request. + """ + stub = FakeStub( + files={"config.json": 2}, + chunks=[whole_file("config.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + ) + make_client(tmp_path, stub).install_metadata_snapshot(MODEL) + + assert [r.ignore_weights for r in stub.download_requests] == [True] + assert [r.ignore_weights for r in stub.list_requests] == [True] + + def test_rejects_manifest_without_metadata(self, tmp_path): + stub = FakeStub(files={"model.safetensors": 7}) + with pytest.raises(ModelCacheError, match="no non-weight files"): + make_client(tmp_path, stub).install_metadata_snapshot(MODEL) + + def test_leaves_no_snapshot_when_stream_fails(self, tmp_path): + stub = FakeStub( + files={"config.json": 2, "tokenizer.json": 2}, + chunks=[whole_file("config.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + ) + with pytest.raises(ModelCacheError, match="missing files"): + make_client(tmp_path, stub).install_metadata_snapshot(MODEL) + + cache = ModelSnapshotCache(MODEL, tmp_path) + assert cache.read_main_ref() is None + leftovers = [ + p.name for p in cache.repo_root.iterdir() if p.name.startswith(".modelexpress-") + ] + assert leftovers == [] + + + def test_installed_snapshot_resolves_offline(self, tmp_path): + """What the engine actually does with the snapshot, end to end. + + vLLM resolves the model through snapshot_download(local_files_only=True) + while parsing engine args, well before the weight loader runs. + """ + from huggingface_hub import snapshot_download + + stub = FakeStub( + files={"config.json": 2, "model.safetensors": 7}, + chunks=[whole_file("config.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + ) + snapshot = make_client(tmp_path, stub).install_metadata_snapshot(MODEL) + + resolved = snapshot_download(MODEL, cache_dir=str(tmp_path), local_files_only=True) + + assert resolved == str(snapshot) + + +class TestInstallWeightFiles: + def _snapshot(self, tmp_path): + stub = FakeStub( + files={"config.json": 2}, + chunks=[whole_file("config.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + ) + return make_client(tmp_path, stub).install_metadata_snapshot(MODEL) + + def test_requests_only_weight_files(self, tmp_path): + snapshot = self._snapshot(tmp_path) + stub = FakeStub( + files={"config.json": 2, "model.safetensors": 7}, + chunks=[ + whole_file("model.safetensors", b"weights", is_last_file=True, commit_hash=COMMIT) + ], + ) + make_client(tmp_path, stub).install_weight_files(MODEL, snapshot) + + assert list(stub.stream_requests[0].file_selector.paths) == ["model.safetensors"] + assert (snapshot / "model.safetensors").read_bytes() == b"weights" + assert (snapshot / "config.json").read_bytes() == b"{}" + + def test_skips_when_weights_present(self, tmp_path): + snapshot = self._snapshot(tmp_path) + (snapshot / "model.safetensors").write_bytes(b"weights") + stub = FakeStub(files={"config.json": 2, "model.safetensors": 7}) + + make_client(tmp_path, stub).install_weight_files(MODEL, snapshot) + + assert stub.stream_requests == [] + + def test_rejects_an_advanced_revision_before_the_present_weights_shortcut(self, tmp_path): + """The mirror of the metadata reuse hole, on the weight path. + + ``has_files`` compares names and sizes only, so weights left by an + earlier revision satisfy it. Without the revision check the call + returns happily and the engine loads the older checkpoint. + """ + snapshot = self._snapshot(tmp_path) + (snapshot / "model.safetensors").write_bytes(b"weights") + stub = FakeStub( + files={"config.json": 2, "model.safetensors": 7}, + resolved_revision="d" * 40, + ) + + with pytest.raises(ModelCacheError, match="refusing to mix revisions"): + make_client(tmp_path, stub).install_weight_files(MODEL, snapshot) + + def test_weight_phase_asks_for_the_weights(self, tmp_path): + snapshot = self._snapshot(tmp_path) + stub = FakeStub( + files={"config.json": 2, "model.safetensors": 7}, + chunks=[ + whole_file("model.safetensors", b"weights", is_last_file=True, commit_hash=COMMIT) + ], + ) + make_client(tmp_path, stub).install_weight_files(MODEL, snapshot) + + assert [r.ignore_weights for r in stub.download_requests] == [False] + + def test_wrong_commit_aborts_before_transferring(self, tmp_path): + """Reject on the first chunk, not after the whole checkpoint arrives. + + A sharded model is tens of gigabytes; noticing the mismatch only at the + end means throwing all of it away. + """ + snapshot = self._snapshot(tmp_path) + produced = [] + + def counting_stream(request): + for payload in ( + whole_file("a.safetensors", b"A", commit_hash="d" * 40), + whole_file("b.safetensors", b"B"), + whole_file("c.safetensors", b"C", is_last_file=True), + ): + produced.append(payload.relative_path) + yield payload + + stub = FakeStub( + files={"config.json": 2, "a.safetensors": 1, "b.safetensors": 1, "c.safetensors": 1} + ) + stub.StreamModelFiles = counting_stream + + with pytest.raises(ModelCacheError, match="refusing to mix revisions"): + make_client(tmp_path, stub).install_weight_files(MODEL, snapshot) + + assert produced == ["a.safetensors"] + assert list(snapshot.iterdir()) == [snapshot / "config.json"] + + def test_refuses_weights_from_a_different_commit(self, tmp_path): + """Pinned revisions are addressed by directory name, so commits must match.""" + snapshot = self._snapshot(tmp_path) + stub = FakeStub( + files={"config.json": 2, "model.safetensors": 7}, + chunks=[ + whole_file( + "model.safetensors", b"weights", is_last_file=True, commit_hash="d" * 40 + ) + ], + ) + with pytest.raises(ModelCacheError, match="refusing to mix revisions"): + make_client(tmp_path, stub).install_weight_files(MODEL, snapshot) + + assert list(snapshot.iterdir()) == [snapshot / "config.json"] + + def test_rejects_manifest_without_weights(self, tmp_path): + snapshot = self._snapshot(tmp_path) + stub = FakeStub(files={"config.json": 2}) + with pytest.raises(ModelCacheError, match="no weight files"): + make_client(tmp_path, stub).install_weight_files(MODEL, snapshot) + + def test_leaves_no_partial_file_when_stream_fails(self, tmp_path): + snapshot = self._snapshot(tmp_path) + stub = FakeStub( + files={"config.json": 2, "model.safetensors": 7}, + chunks=[ + chunk( + "model.safetensors", + b"weig", + total_size=7, + is_last_chunk=False, + commit_hash=COMMIT, + ) + ], + ) + with pytest.raises(ModelCacheError, match="final file marker"): + make_client(tmp_path, stub).install_weight_files(MODEL, snapshot) + + assert list(snapshot.iterdir()) == [snapshot / "config.json"] + + def test_rolls_back_completed_files_when_a_later_file_fails(self, tmp_path): + """A half-applied weight set would load as if it were complete.""" + snapshot = self._snapshot(tmp_path) + stub = FakeStub( + files={"config.json": 2, "a.safetensors": 1, "b.safetensors": 1}, + chunks=[whole_file("a.safetensors", b"A", commit_hash=COMMIT)], + ) + with pytest.raises(ModelCacheError, match="final file marker"): + make_client(tmp_path, stub).install_weight_files(MODEL, snapshot) + + assert list(snapshot.iterdir()) == [snapshot / "config.json"] + + def test_ensures_the_server_has_the_model(self, tmp_path): + snapshot = self._snapshot(tmp_path) + stub = FakeStub( + files={"config.json": 2, "model.safetensors": 7}, + chunks=[ + whole_file("model.safetensors", b"weights", is_last_file=True, commit_hash=COMMIT) + ], + ) + make_client(tmp_path, stub).install_weight_files(MODEL, snapshot) + + assert len(stub.download_requests) == 1 + + +class TestStreamValidation: + """One canned bad stream per protocol rule the client has to enforce.""" + + def _install(self, tmp_path, chunks, files=None): + stub = FakeStub(files=files or {"config.json": 2}, chunks=chunks) + return make_client(tmp_path, stub).install_metadata_snapshot(MODEL) + + def test_first_chunk_must_carry_commit_hash(self, tmp_path): + with pytest.raises(ModelCacheError, match="commit hash"): + self._install(tmp_path, [whole_file("config.json", b"{}", is_last_file=True)]) + + def test_commit_hash_may_not_change(self, tmp_path): + with pytest.raises(ModelCacheError, match="changed the commit hash"): + self._install( + tmp_path, + [ + whole_file("config.json", b"{}", commit_hash=COMMIT), + whole_file( + "tokenizer.json", b"[]", is_last_file=True, commit_hash="d" * 40 + ), + ], + files={"config.json": 2, "tokenizer.json": 2}, + ) + + def test_rejects_unrequested_file(self, tmp_path): + with pytest.raises(ModelCacheError, match="unrequested file"): + self._install( + tmp_path, + [whole_file("secret.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + ) + + def test_rejects_weight_file_in_metadata_stream(self, tmp_path): + with pytest.raises(ModelCacheError, match="unrequested file"): + self._install( + tmp_path, + [ + whole_file( + "model.safetensors", b"weights", is_last_file=True, commit_hash=COMMIT + ) + ], + files={"config.json": 2, "model.safetensors": 7}, + ) + + def test_rejects_size_mismatch_against_manifest(self, tmp_path): + with pytest.raises(ModelCacheError, match="Size mismatch"): + self._install( + tmp_path, + [ + chunk( + "config.json", + b"{}", + total_size=99, + is_last_chunk=False, + commit_hash=COMMIT, + ) + ], + ) + + def test_rejects_non_zero_first_offset(self, tmp_path): + with pytest.raises(ModelCacheError, match="offset"): + self._install( + tmp_path, + [ + chunk( + "config.json", + b"{}", + offset=1, + total_size=2, + is_last_file=True, + commit_hash=COMMIT, + ) + ], + ) + + def test_rejects_offset_gap(self, tmp_path): + with pytest.raises(ModelCacheError, match="Unexpected offset"): + self._install( + tmp_path, + [ + chunk( + "config.json", b"{", total_size=4, is_last_chunk=False, commit_hash=COMMIT + ), + chunk("config.json", b"}", offset=3, total_size=4, is_last_file=True), + ], + files={"config.json": 4}, + ) + + def test_rejects_data_beyond_total_size(self, tmp_path): + with pytest.raises(ModelCacheError, match="exceeds its advertised size"): + self._install( + tmp_path, + [ + chunk( + "config.json", + b"{oversized}", + total_size=2, + is_last_file=True, + commit_hash=COMMIT, + ) + ], + ) + + def test_rejects_interleaved_files(self, tmp_path): + with pytest.raises(ModelCacheError, match="before"): + self._install( + tmp_path, + [ + chunk( + "config.json", b"{", total_size=2, is_last_chunk=False, commit_hash=COMMIT + ), + whole_file("tokenizer.json", b"[]"), + ], + files={"config.json": 2, "tokenizer.json": 2}, + ) + + def test_rejects_duplicate_file(self, tmp_path): + with pytest.raises(ModelCacheError, match="twice"): + self._install( + tmp_path, + [ + whole_file("config.json", b"{}", commit_hash=COMMIT), + whole_file("config.json", b"{}", is_last_file=True), + ], + ) + + def test_rejects_final_file_marker_before_final_chunk(self, tmp_path): + with pytest.raises(ModelCacheError, match="Final-file marker"): + self._install( + tmp_path, + [ + chunk( + "config.json", + b"{", + total_size=2, + is_last_chunk=False, + is_last_file=True, + commit_hash=COMMIT, + ) + ], + ) + + def test_rejects_data_after_final_marker(self, tmp_path): + with pytest.raises(ModelCacheError, match="after the final stream marker"): + self._install( + tmp_path, + [ + whole_file("config.json", b"{}", is_last_file=True, commit_hash=COMMIT), + whole_file("tokenizer.json", b"[]"), + ], + files={"config.json": 2, "tokenizer.json": 2}, + ) + + def test_rejects_missing_final_marker(self, tmp_path): + with pytest.raises(ModelCacheError, match="final file marker"): + self._install( + tmp_path, [whole_file("config.json", b"{}", commit_hash=COMMIT)] + ) + + def test_rejects_empty_stream(self, tmp_path): + with pytest.raises(ModelCacheError, match="no model files"): + self._install(tmp_path, []) + + def test_rejects_incomplete_file(self, tmp_path): + with pytest.raises(ModelCacheError, match="Incomplete file"): + self._install( + tmp_path, + [ + chunk( + "config.json", + b"{", + total_size=2, + is_last_chunk=True, + is_last_file=True, + commit_hash=COMMIT, + ) + ], + ) + + def test_unsafe_path_is_caught_by_the_manifest_check(self, tmp_path): + with pytest.raises(ModelCacheError, match="unrequested file"): + self._install( + tmp_path, + [whole_file("../escape.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + ) + + def test_unsafe_path_in_the_manifest_never_reaches_the_filesystem(self, tmp_path): + """Defense in depth: a compromised manifest must not steer the writer.""" + with pytest.raises(ModelSnapshotError, match="Unsafe model file path"): + self._install( + tmp_path, + [whole_file("../escape.json", b"{}", is_last_file=True, commit_hash=COMMIT)], + files={"../escape.json": 2}, + ) + assert not (tmp_path / "escape.json").exists() + + def test_accepts_multi_chunk_file(self, tmp_path): + snapshot = self._install( + tmp_path, + [ + chunk("config.json", b"{", total_size=4, is_last_chunk=False, commit_hash=COMMIT), + chunk("config.json", b"a", offset=1, total_size=4, is_last_chunk=False), + chunk("config.json", b"b}", offset=2, total_size=4, is_last_file=True), + ], + files={"config.json": 4}, + ) + assert (snapshot / "config.json").read_bytes() == b"{ab}" + + def test_accepts_empty_file(self, tmp_path): + snapshot = self._install( + tmp_path, + [ + whole_file("config.json", b"{}", commit_hash=COMMIT), + whole_file(".gitattributes", b"", is_last_file=True), + ], + files={"config.json": 2, ".gitattributes": 0}, + ) + assert (snapshot / ".gitattributes").read_bytes() == b"" diff --git a/modelexpress_client/python/tests/test_model_prefetch.py b/modelexpress_client/python/tests/test_model_prefetch.py new file mode 100644 index 000000000..ebca08738 --- /dev/null +++ b/modelexpress_client/python/tests/test_model_prefetch.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the pre-engine metadata prefetch.""" + +import logging +import threading +import time +from pathlib import Path + +import pytest + +from modelexpress import model_prefetch + +REPO = "org/model" +COMMIT = "e" * 40 + + +class FakeClient: + """Stands in for ModelCacheClient, recording what the prefetch asked for.""" + + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.calls = [] + self.snapshot = None + self.error = None + FakeClient.instances.append(self) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return None + + def install_metadata_snapshot(self, repo_id, *args, **kwargs): + self.calls.append(repo_id) + if self.error is not None: + raise self.error + return self.snapshot + + +@pytest.fixture(autouse=True) +def clean_state(monkeypatch, tmp_path): + model_prefetch.reset() + FakeClient.instances = [] + for name in ( + "MODEL_EXPRESS_NO_SHARED_STORAGE", + "MODEL_EXPRESS_URL", + "MX_SERVER_ADDRESS", + "MODEL_EXPRESS_TRANSFER_CHUNK_SIZE", + ): + monkeypatch.delenv(name, raising=False) + yield + model_prefetch.reset() + + +@pytest.fixture +def enabled(monkeypatch): + monkeypatch.setenv("MODEL_EXPRESS_NO_SHARED_STORAGE", "1") + monkeypatch.setenv("MODEL_EXPRESS_URL", "http://mx:8001") + + +@pytest.fixture +def fake_client(monkeypatch, tmp_path): + snapshot = tmp_path / "models--org--model" / "snapshots" / COMMIT + snapshot.mkdir(parents=True) + + def factory(**kwargs): + client = FakeClient(**kwargs) + client.snapshot = snapshot + return client + + monkeypatch.setattr("modelexpress.model_client.ModelCacheClient", factory) + return snapshot + + +class TestIsEnabled: + def test_off_by_default(self): + assert model_prefetch.is_enabled() is False + + def test_needs_a_server_address(self, monkeypatch): + monkeypatch.setenv("MODEL_EXPRESS_NO_SHARED_STORAGE", "1") + assert model_prefetch.is_enabled() is False + + def test_needs_the_switch(self, monkeypatch): + monkeypatch.setenv("MODEL_EXPRESS_URL", "http://mx:8001") + assert model_prefetch.is_enabled() is False + + def test_enabled_with_both(self, enabled): + assert model_prefetch.is_enabled() is True + + def test_mx_server_address_also_counts(self, monkeypatch): + monkeypatch.setenv("MODEL_EXPRESS_NO_SHARED_STORAGE", "1") + monkeypatch.setenv("MX_SERVER_ADDRESS", "mx:8001") + assert model_prefetch.is_enabled() is True + + +class TestIsRepoId: + @pytest.mark.parametrize("model", ["org/model", "model", "org/model-v2.5"]) + def test_repo_ids(self, model): + assert model_prefetch.is_repo_id(model) is True + + @pytest.mark.parametrize( + "model", + ["", "/abs/path", "/home/dynamo/.cache/huggingface/hub/models--org--model", "a/b/c"], + ) + def test_not_repo_ids(self, model): + assert model_prefetch.is_repo_id(model) is False + + def test_existing_local_directory_is_not_a_repo_id(self, tmp_path, monkeypatch): + (tmp_path / "local-model").mkdir() + monkeypatch.chdir(tmp_path) + assert model_prefetch.is_repo_id("local-model") is False + + +class TestEnsureMetadata: + def test_no_op_when_disabled(self, fake_client): + assert model_prefetch.ensure_metadata(REPO) is None + assert FakeClient.instances == [] + + def test_no_op_for_local_path(self, enabled, fake_client, tmp_path): + assert model_prefetch.ensure_metadata(str(tmp_path)) is None + assert FakeClient.instances == [] + + def test_installs_and_returns_snapshot(self, enabled, fake_client): + assert model_prefetch.ensure_metadata(REPO) == fake_client + assert FakeClient.instances[0].calls == [REPO] + + def test_second_call_does_not_hit_the_server(self, enabled, fake_client): + first = model_prefetch.ensure_metadata(REPO) + second = model_prefetch.ensure_metadata(REPO) + + assert first == second + assert len(FakeClient.instances) == 1 + + def test_failure_is_retryable(self, enabled, fake_client, monkeypatch): + def failing_factory(**kwargs): + client = FakeClient(**kwargs) + client.error = RuntimeError("server down") + return client + + monkeypatch.setattr("modelexpress.model_client.ModelCacheClient", failing_factory) + with pytest.raises(RuntimeError, match="server down"): + model_prefetch.ensure_metadata(REPO) + + def working_factory(**kwargs): + client = FakeClient(**kwargs) + client.snapshot = fake_client + return client + + monkeypatch.setattr("modelexpress.model_client.ModelCacheClient", working_factory) + assert model_prefetch.ensure_metadata(REPO) == fake_client + + def test_passes_configured_chunk_size(self, enabled, fake_client, monkeypatch): + monkeypatch.setenv("MODEL_EXPRESS_TRANSFER_CHUNK_SIZE", "65536") + model_prefetch.ensure_metadata(REPO) + assert FakeClient.instances[0].kwargs["chunk_size"] == 65536 + + @pytest.mark.parametrize("raw", ["not-a-number", "0", "-1", "99999999999999"]) + def test_bad_chunk_size_falls_back(self, enabled, fake_client, monkeypatch, raw): + """A bad env var must not be the reason a worker fails to start.""" + monkeypatch.setenv("MODEL_EXPRESS_TRANSFER_CHUNK_SIZE", raw) + model_prefetch.ensure_metadata(REPO) + assert FakeClient.instances[0].kwargs["chunk_size"] is None + + +class TestConcurrentEnsureMetadata: + """A second caller arriving mid-install must get the snapshot, not None. + + The engine resolves the model immediately after ensure_metadata returns, so + handing back None while another thread is still writing the snapshot sends + it looking for files that do not exist yet. + """ + + def test_second_caller_waits_and_gets_the_same_snapshot( + self, enabled, monkeypatch, tmp_path + ): + snapshot = tmp_path / "models--org--model" / "snapshots" / ("a" * 40) + snapshot.mkdir(parents=True) + installs = [] + + class SlowClient: + def __init__(self, **kwargs): + pass + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return None + + def install_metadata_snapshot(self, repo_id, *args, **kwargs): + installs.append(repo_id) + time.sleep(0.3) + return snapshot + + monkeypatch.setattr("modelexpress.model_client.ModelCacheClient", SlowClient) + + results = {} + + def call(tag, delay): + time.sleep(delay) + results[tag] = model_prefetch.ensure_metadata(REPO) + + threads = [ + threading.Thread(target=call, args=("first", 0.0)), + threading.Thread(target=call, args=("second", 0.05)), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert results["first"] == snapshot + assert results["second"] == snapshot + assert installs == [REPO] + + +class TestRevisionMismatch: + def test_matching_commit_is_quiet(self, enabled, fake_client, caplog): + with caplog.at_level(logging.WARNING): + model_prefetch.ensure_metadata(REPO, COMMIT) + assert "does not pin a revision" not in caplog.text + + def test_main_is_quiet(self, enabled, fake_client, caplog): + with caplog.at_level(logging.WARNING): + model_prefetch.ensure_metadata(REPO, "main") + assert "does not pin a revision" not in caplog.text + + def test_pinned_mismatch_explains_itself(self, enabled, fake_client, caplog): + with caplog.at_level(logging.WARNING): + model_prefetch.ensure_metadata(REPO, "f" * 40) + assert "does not pin a revision" in caplog.text + + def test_branch_name_warns(self, enabled, fake_client, caplog): + with caplog.at_level(logging.WARNING): + model_prefetch.ensure_metadata(REPO, "refs/pr/1") + assert "does not pin a revision" in caplog.text + + +class TestRepoIdFor: + def test_maps_snapshot_path_back(self, enabled, fake_client): + snapshot = model_prefetch.ensure_metadata(REPO) + assert model_prefetch.repo_id_for(snapshot) == REPO + assert model_prefetch.repo_id_for(str(snapshot) + "/") == REPO + + def test_passes_through_repo_id(self): + assert model_prefetch.repo_id_for(REPO) == REPO + + def test_none_for_unknown_path(self, tmp_path): + assert model_prefetch.repo_id_for(tmp_path / "unknown") is None + + def test_none_for_unregistered_local_model_dir(self, tmp_path): + assert model_prefetch.repo_id_for("/opt/models/llama") is None + + def test_accepts_path_objects(self, enabled, fake_client): + snapshot = model_prefetch.ensure_metadata(REPO) + assert model_prefetch.repo_id_for(Path(snapshot)) == REPO + + +class TestRepoIdFromCachePath: + """vLLM loads weights in a separate EngineCore process. + + That process never runs the prefetch, so the in-process record is empty + there and the cache layout has to carry the repo id on its own. + """ + + def test_recovers_from_snapshot_path_without_any_record(self, tmp_path): + path = tmp_path / "models--Qwen--Qwen2.5-0.5B-Instruct" / "snapshots" / ("a" * 40) + assert model_prefetch.repo_id_for(path) == "Qwen/Qwen2.5-0.5B-Instruct" + assert model_prefetch._snapshot_to_repo_id == {} + + def test_recovers_from_repo_root(self, tmp_path): + path = tmp_path / "models--org--model" + assert model_prefetch.repo_id_from_cache_path(path) == "org/model" + + def test_recovers_from_file_inside_snapshot(self, tmp_path): + path = ( + tmp_path / "models--org--model" / "snapshots" / ("b" * 40) / "config.json" + ) + assert model_prefetch.repo_id_from_cache_path(path) == "org/model" + + def test_single_segment_repo(self, tmp_path): + assert model_prefetch.repo_id_from_cache_path(tmp_path / "models--gpt2") == "gpt2" + + @pytest.mark.parametrize( + "path", ["/opt/models/llama", "/home/dynamo/.cache/huggingface/hub", "/"] + ) + def test_none_for_non_cache_paths(self, path): + assert model_prefetch.repo_id_from_cache_path(path) is None diff --git a/modelexpress_client/python/tests/test_model_snapshot.py b/modelexpress_client/python/tests/test_model_snapshot.py new file mode 100644 index 000000000..c07ab4fb1 --- /dev/null +++ b/modelexpress_client/python/tests/test_model_snapshot.py @@ -0,0 +1,392 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the Hugging Face cache layout used by server-streamed models.""" + +import pytest + +from modelexpress.model_snapshot import ( + MAIN_REF, + ModelSnapshotCache, + ModelSnapshotError, + is_weight_file, + repo_dir_name, + resolve_cache_root, + safe_commit_hash, + safe_relative_path, + split_by_weight, +) + +COMMIT = "a" * 40 +OTHER_COMMIT = "b" * 40 + + +@pytest.fixture +def cache(tmp_path, monkeypatch): + monkeypatch.delenv("MODEL_EXPRESS_CACHE_DIRECTORY", raising=False) + return ModelSnapshotCache("org/model", tmp_path) + + +def _write(cache, files, commit=COMMIT): + """Publish ``files`` ({path: bytes}) as a snapshot and return its path.""" + staging = cache.staging() + for relative_path, payload in files.items(): + staging.begin_file(relative_path) + staging.write(payload) + staging.end_file() + expected = {path: len(payload) for path, payload in files.items()} + return staging.publish(commit, expected) + + +class TestWeightClassification: + @pytest.mark.parametrize( + "path", + [ + "model.safetensors", + "pytorch_model-00001-of-00002.bin", + "sub/dir/model.safetensors", + "tf_model.h5", + "flax_model.msgpack", + ], + ) + def test_weight_files(self, path): + assert is_weight_file(path) is True + + @pytest.mark.parametrize( + "path", + [ + "config.json", + "tokenizer.json", + "model.safetensors.index.json", + "README.md", + ], + ) + def test_metadata_files(self, path): + assert is_weight_file(path) is False + + def test_split_preserves_order(self): + metadata, weights = split_by_weight( + ["config.json", "a.safetensors", "tokenizer.json", "b.bin"] + ) + assert metadata == ["config.json", "tokenizer.json"] + assert weights == ["a.safetensors", "b.bin"] + + +class TestPathValidation: + @pytest.mark.parametrize( + "path", + [ + "", + "/etc/passwd", + "../escape", + "sub/../../escape", + "sub/./file", + "back\\slash", + "nul\x00byte", + "trailing/", + ], + ) + def test_rejects_unsafe_paths(self, path): + with pytest.raises(ModelSnapshotError): + safe_relative_path(path) + + def test_accepts_nested_path(self): + assert safe_relative_path("sub/dir/file.json").parts == ("sub", "dir", "file.json") + + @pytest.mark.parametrize("commit", ["", ".", "..", "a/b", "a\\b", "a\x00b"]) + def test_rejects_unsafe_commit(self, commit): + with pytest.raises(ModelSnapshotError): + safe_commit_hash(commit) + + def test_repo_dir_name(self): + assert repo_dir_name("org/model") == "models--org--model" + assert repo_dir_name("model") == "models--model" + + @pytest.mark.parametrize("name", ["", "/abs", "org/../model", "back\\slash"]) + def test_repo_dir_name_rejects_unsafe(self, name): + with pytest.raises(ValueError): + repo_dir_name(name) + + +class TestCacheRootResolution: + def test_explicit_wins(self, tmp_path, monkeypatch): + monkeypatch.setenv("MODEL_EXPRESS_CACHE_DIRECTORY", str(tmp_path / "env")) + assert resolve_cache_root(tmp_path / "explicit") == tmp_path / "explicit" + + def test_env_used_when_no_explicit(self, tmp_path, monkeypatch): + monkeypatch.setenv("MODEL_EXPRESS_CACHE_DIRECTORY", str(tmp_path / "env")) + assert resolve_cache_root() == tmp_path / "env" + + def test_falls_back_to_hf_hub_cache(self, monkeypatch): + from huggingface_hub.constants import HF_HUB_CACHE + + monkeypatch.delenv("MODEL_EXPRESS_CACHE_DIRECTORY", raising=False) + assert str(resolve_cache_root()) == str(HF_HUB_CACHE) + + +class TestPublish: + def test_layout_and_ref(self, cache): + snapshot = _write(cache, {"config.json": b"{}", "sub/tok.json": b"[]"}) + + assert snapshot == cache.repo_root / "snapshots" / COMMIT + assert (snapshot / "config.json").read_bytes() == b"{}" + assert (snapshot / "sub" / "tok.json").read_bytes() == b"[]" + assert (cache.repo_root / "refs" / MAIN_REF).read_text() == COMMIT + assert cache.read_main_ref() == COMMIT + + def test_no_staging_directory_left_behind(self, cache): + _write(cache, {"config.json": b"{}"}) + leftovers = [ + p.name for p in cache.repo_root.iterdir() if p.name.startswith(".modelexpress-") + ] + assert leftovers == [] + + def test_discard_removes_staging(self, cache): + staging = cache.staging() + staging_path = staging.path + staging.begin_file("config.json") + staging.write(b"{}") + staging.discard() + + assert not staging_path.exists() + assert not (cache.repo_root / "snapshots").exists() + + def test_republish_same_commit_reuses_complete_snapshot(self, cache): + snapshot = _write(cache, {"config.json": b"{}"}) + (snapshot / "extra.json").write_text("kept") + + again = _write(cache, {"config.json": b"{}"}) + + assert again == snapshot + assert (snapshot / "extra.json").read_text() == "kept" + + def test_republish_updates_incomplete_snapshot(self, cache): + snapshot = _write(cache, {"config.json": b"{}"}) + (snapshot / "config.json").unlink() + + again = _write(cache, {"config.json": b"{'v': 2}"}) + + assert again == snapshot + assert (snapshot / "config.json").read_bytes() == b"{'v': 2}" + + def test_republish_keeps_files_the_manifest_does_not_mention(self, cache): + """Installing metadata must not delete an already-installed weight set. + + The commit hash comes from the server resolving ``main``, so a second + install targets the same ``snapshots//``. The manifest passed + here covers metadata only, so replacing the directory wholesale would + drop weights that no expected-file check ever looks at. + """ + snapshot = _write(cache, {"config.json": b"{}"}) + weights = snapshot / "model.safetensors" + weights.write_bytes(b"W" * 64) + (snapshot / "shards" / "extra").mkdir(parents=True) + (snapshot / "shards" / "extra" / "part.safetensors").write_bytes(b"S" * 16) + + again = _write(cache, {"config.json": b"{}", "chat_template.jinja": b"tpl"}) + + assert again == snapshot + assert weights.read_bytes() == b"W" * 64 + assert (snapshot / "shards" / "extra" / "part.safetensors").read_bytes() == b"S" * 16 + assert (snapshot / "chat_template.jinja").read_bytes() == b"tpl" + + def test_republish_leaves_no_staging_or_stale_directories(self, cache): + snapshot = _write(cache, {"config.json": b"{}"}) + (snapshot / "model.safetensors").write_bytes(b"W") + + _write(cache, {"config.json": b"{}", "chat_template.jinja": b"tpl"}) + + leftovers = [ + entry.name + for entry in cache.repo_root.iterdir() + if entry.name.startswith((".modelexpress-stale-", ".modelexpress-staging-")) + ] + assert leftovers == [] + + def test_second_commit_moves_ref(self, cache): + _write(cache, {"config.json": b"{}"}, commit=COMMIT) + _write(cache, {"config.json": b"{}"}, commit=OTHER_COMMIT) + + assert cache.read_main_ref() == OTHER_COMMIT + assert (cache.repo_root / "snapshots" / COMMIT).is_dir() + + def test_rejects_unsafe_streamed_path(self, cache): + staging = cache.staging() + with pytest.raises(ModelSnapshotError): + staging.begin_file("../escape.json") + staging.discard() + + def test_rejects_overlapping_files(self, cache): + staging = cache.staging() + staging.begin_file("a.json") + with pytest.raises(ModelSnapshotError): + staging.begin_file("b.json") + staging.discard() + + +class TestResolveSnapshot: + def test_returns_snapshot_when_files_present(self, cache): + snapshot = _write(cache, {"config.json": b"{}"}) + assert cache.resolve_snapshot({"config.json": 2}, COMMIT) == snapshot + + def test_none_when_file_missing(self, cache): + _write(cache, {"config.json": b"{}"}) + assert cache.resolve_snapshot({"config.json": 2, "tokenizer.json": 5}, COMMIT) is None + + def test_none_when_size_differs(self, cache): + _write(cache, {"config.json": b"{}"}) + assert cache.resolve_snapshot({"config.json": 99}, COMMIT) is None + + def test_none_without_ref(self, cache): + _write(cache, {"config.json": b"{}"}) + (cache.repo_root / "refs" / MAIN_REF).unlink() + assert cache.resolve_snapshot({"config.json": 2}, COMMIT) is None + + def test_none_when_server_named_no_revision(self, cache): + """No reported revision is not proof the local snapshot is current.""" + _write(cache, {"config.json": b"{}"}) + assert cache.resolve_snapshot({"config.json": 2}, None) is None + + def test_none_when_revision_differs(self, cache): + """The case a manifest cannot catch: same names and sizes, new commit. + + This is what makes size-only reuse unsafe. Without the commit check + the stale snapshot is returned and no stream ever opens to notice. + """ + _write(cache, {"config.json": b"{}"}) + assert cache.resolve_snapshot({"config.json": 2}, OTHER_COMMIT) is None + + +class TestPatch: + def test_adds_files_to_published_snapshot(self, cache): + snapshot = _write(cache, {"config.json": b"{}"}) + + patch = cache.patch(snapshot) + patch.begin_file("model.safetensors") + patch.write(b"weights") + patch.end_file() + patch.close() + + assert (snapshot / "model.safetensors").read_bytes() == b"weights" + assert (snapshot / "config.json").read_bytes() == b"{}" + assert cache.read_main_ref() == COMMIT + + def test_leaves_no_temp_file_on_abort(self, cache): + snapshot = _write(cache, {"config.json": b"{}"}) + + patch = cache.patch(snapshot) + patch.begin_file("model.safetensors") + patch.write(b"partial") + patch.close() + + assert not (snapshot / "model.safetensors").exists() + assert list(snapshot.iterdir()) == [snapshot / "config.json"] + + def test_rejects_missing_snapshot(self, cache): + with pytest.raises(ModelSnapshotError): + cache.patch(cache.repo_root / "snapshots" / COMMIT) + + def test_rollback_restores_a_replaced_shard(self, cache): + """A refresh that fails part-way must not cost the snapshot a shard. + + Publishing shard A overwrites the copy already on disk. If the patch + then fails on shard B, rolling back by deleting what it published + would leave the snapshot short of a shard it had before the patch + started -- worse than the partial write the rollback exists to avoid. + """ + snapshot = _write( + cache, + {"shard-1.safetensors": b"old-one", "shard-2.safetensors": b"old-two"}, + ) + + patch = cache.patch(snapshot) + patch.begin_file("shard-1.safetensors") + patch.write(b"new-one") + patch.end_file() + patch.rollback() + + assert (snapshot / "shard-1.safetensors").read_bytes() == b"old-one" + assert (snapshot / "shard-2.safetensors").read_bytes() == b"old-two" + assert sorted(p.name for p in snapshot.iterdir()) == [ + "shard-1.safetensors", + "shard-2.safetensors", + ] + + def test_rollback_removes_a_newly_added_file(self, cache): + snapshot = _write(cache, {"config.json": b"{}"}) + + patch = cache.patch(snapshot) + patch.begin_file("model.safetensors") + patch.write(b"weights") + patch.end_file() + patch.rollback() + + assert not (snapshot / "model.safetensors").exists() + assert list(snapshot.iterdir()) == [snapshot / "config.json"] + + def test_commit_drops_the_backups(self, cache): + snapshot = _write(cache, {"shard-1.safetensors": b"old-one"}) + + patch = cache.patch(snapshot) + patch.begin_file("shard-1.safetensors") + patch.write(b"new-one") + patch.end_file() + patch.commit() + patch.close() + + assert (snapshot / "shard-1.safetensors").read_bytes() == b"new-one" + assert list(snapshot.iterdir()) == [snapshot / "shard-1.safetensors"] + + def test_rollback_after_commit_keeps_the_published_files(self, cache): + snapshot = _write(cache, {"shard-1.safetensors": b"old-one"}) + + patch = cache.patch(snapshot) + patch.begin_file("shard-1.safetensors") + patch.write(b"new-one") + patch.end_file() + patch.commit() + patch.rollback() + + assert (snapshot / "shard-1.safetensors").read_bytes() == b"new-one" + + +class TestLock: + def test_released_after_context(self, cache): + with cache.lock(): + pass + with cache.lock(): + pass + assert (cache.repo_root / ".modelexpress.lock").is_file() + + +def test_published_snapshot_resolves_offline(cache, monkeypatch): + """huggingface_hub must resolve the published layout with no network. + + Regression guard for issue #569: the engine resolves the model through + ``snapshot_download(local_files_only=True)`` long before the weight loader + runs, and that call fails with LocalEntryNotFoundError unless refs/main + points at a snapshot directory. + """ + from huggingface_hub import snapshot_download + + snapshot = _write(cache, {"config.json": b"{}", "tokenizer.json": b"[]"}) + + resolved = snapshot_download( + "org/model", cache_dir=str(cache.cache_root), local_files_only=True + ) + + assert resolved == str(snapshot) + + +def test_snapshot_without_ref_is_unresolvable(cache): + """The failure mode from the issue, pinned so the ref write cannot regress.""" + from huggingface_hub import snapshot_download + from huggingface_hub.errors import LocalEntryNotFoundError + + _write(cache, {"config.json": b"{}"}) + (cache.repo_root / "refs" / MAIN_REF).unlink() + + with pytest.raises(LocalEntryNotFoundError): + snapshot_download( + "org/model", cache_dir=str(cache.cache_root), local_files_only=True + ) diff --git a/modelexpress_client/python/tests/test_server_cache_strategy.py b/modelexpress_client/python/tests/test_server_cache_strategy.py new file mode 100644 index 000000000..b7993e58e --- /dev/null +++ b/modelexpress_client/python/tests/test_server_cache_strategy.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for ServerCacheStrategy.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from modelexpress import model_prefetch, p2p_pb2 +from modelexpress.adapter import EngineAdapter, StrategyFailed +from modelexpress.load_strategy.context import LoadResult +from modelexpress.load_strategy.server_cache_strategy import ServerCacheStrategy + +REPO = "org/model" +COMMIT = "a" * 40 + + +class _FakeAdapter(EngineAdapter): + """Adapter implementing the native-load capability the strategy requires.""" + + def __init__(self, *, native_error=None): + self.native_error = native_error + self.native_calls = 0 + self.post_calls = 0 + + def discover_tensors(self, result: LoadResult): + return {} + + def load_via_native(self, result: LoadResult) -> LoadResult: + self.native_calls += 1 + if self.native_error is not None: + raise self.native_error + return result + + def after_native_load(self, result: LoadResult) -> LoadResult: + self.post_calls += 1 + return result + + +class _NoNativeAdapter(EngineAdapter): + """Adapter without load_via_native, so the strategy must be ineligible.""" + + def discover_tensors(self, result: LoadResult): + return {} + + +class FakeClient: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.calls = [] + self.error = None + FakeClient.instances.append(self) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return None + + def install_weight_files(self, repo_id, snapshot_path, *args, **kwargs): + self.calls.append((repo_id, snapshot_path)) + if self.error is not None: + raise self.error + + +@pytest.fixture(autouse=True) +def clean_state(monkeypatch): + model_prefetch.reset() + FakeClient.instances = [] + for name in ("MODEL_EXPRESS_NO_SHARED_STORAGE", "MODEL_EXPRESS_URL", "MX_SERVER_ADDRESS"): + monkeypatch.delenv(name, raising=False) + yield + model_prefetch.reset() + + +@pytest.fixture +def enabled(monkeypatch): + monkeypatch.setenv("MODEL_EXPRESS_NO_SHARED_STORAGE", "1") + monkeypatch.setenv("MODEL_EXPRESS_URL", "http://mx:8001") + + +@pytest.fixture +def snapshot(tmp_path): + path = tmp_path / "models--org--model" / "snapshots" / COMMIT + path.mkdir(parents=True) + (path / "config.json").write_text("{}") + return path + + +@pytest.fixture +def fake_client(monkeypatch): + monkeypatch.setattr("modelexpress.model_client.ModelCacheClient", FakeClient) + return FakeClient + + +def _make_context(model_name, *, adapter=None, model_path=None, revision=None): + from modelexpress.load_strategy import LoadContext + + return LoadContext( + model_config=SimpleNamespace(model=model_path, revision=revision), + load_config=MagicMock(), + target_device=torch.device("cpu"), + global_rank=0, + worker_rank=0, + device_id=0, + identity=p2p_pb2.SourceIdentity(model_name=model_name, tensor_parallel_size=1), + mx_client=MagicMock(), + worker_id="test-worker", + adapter=adapter if adapter is not None else _FakeAdapter(), + ) + + +class TestIsAvailable: + def test_unavailable_when_switch_is_off(self): + assert ServerCacheStrategy().is_available(_make_context(REPO)) is False + + def test_unavailable_without_server_address(self, monkeypatch): + monkeypatch.setenv("MODEL_EXPRESS_NO_SHARED_STORAGE", "1") + assert ServerCacheStrategy().is_available(_make_context(REPO)) is False + + def test_unavailable_without_native_load_capability(self, enabled): + ctx = _make_context(REPO, adapter=_NoNativeAdapter()) + assert ServerCacheStrategy().is_available(ctx) is False + + def test_available_for_a_repo_id(self, enabled): + assert ServerCacheStrategy().is_available(_make_context(REPO)) is True + + def test_available_for_a_registered_snapshot_path(self, enabled, snapshot): + model_prefetch._snapshot_to_repo_id[str(snapshot)] = REPO + assert ServerCacheStrategy().is_available(_make_context(str(snapshot))) is True + + def test_unavailable_for_an_unknown_local_path(self, enabled): + ctx = _make_context("/opt/models/llama") + assert ServerCacheStrategy().is_available(ctx) is False + + def test_available_in_a_process_that_never_ran_the_prefetch(self, enabled, snapshot): + """The EngineCore process has no prefetch record; the path must suffice. + + vLLM rewrites ModelConfig.model with the resolved snapshot path and + loads weights in a separate process, so is_available() has to recover + the repo id from the cache layout alone. + """ + model_prefetch.reset() + ctx = _make_context(str(snapshot)) + assert ServerCacheStrategy().is_available(ctx) is True + + +class TestLoad: + def test_installs_weights_then_loads_natively(self, enabled, snapshot, fake_client): + adapter = _FakeAdapter() + ctx = _make_context(REPO, adapter=adapter, model_path=str(snapshot)) + result = LoadResult(value=MagicMock(), model=MagicMock()) + + with patch( + "modelexpress.load_strategy.server_cache_strategy.register_tensors" + ) as register: + out = ServerCacheStrategy().load(result, ctx) + + assert FakeClient.instances[0].calls == [(REPO, snapshot)] + assert adapter.native_calls == 1 + assert adapter.post_calls == 1 + assert register.call_count == 1 + assert out is result + + def test_uses_the_snapshot_the_engine_resolved(self, enabled, snapshot, fake_client): + """model_config.model is the path the engine is already reading from.""" + model_prefetch._snapshot_to_repo_id[str(snapshot)] = REPO + ctx = _make_context(str(snapshot), model_path=str(snapshot)) + + with patch("modelexpress.load_strategy.server_cache_strategy.register_tensors"): + ServerCacheStrategy().load(LoadResult(value=MagicMock()), ctx) + + assert FakeClient.instances[0].calls == [(REPO, snapshot)] + + def test_installs_metadata_when_no_snapshot_exists(self, enabled, snapshot, fake_client): + ctx = _make_context(REPO, model_path=None) + + with patch.object(model_prefetch, "ensure_metadata", return_value=snapshot) as ensure: + with patch("modelexpress.load_strategy.server_cache_strategy.register_tensors"): + ServerCacheStrategy().load(LoadResult(value=MagicMock()), ctx) + + assert ensure.call_count == 1 + assert FakeClient.instances[0].calls == [(REPO, snapshot)] + + def test_server_failure_is_a_clean_miss(self, enabled, snapshot, monkeypatch): + def failing_factory(**kwargs): + client = FakeClient(**kwargs) + client.error = RuntimeError("server unreachable") + return client + + monkeypatch.setattr("modelexpress.model_client.ModelCacheClient", failing_factory) + adapter = _FakeAdapter() + ctx = _make_context(REPO, adapter=adapter, model_path=str(snapshot)) + + with pytest.raises(StrategyFailed) as excinfo: + ServerCacheStrategy().load(LoadResult(value=MagicMock()), ctx) + + assert excinfo.value.mutated is False + assert adapter.native_calls == 0 + + def test_native_load_failure_reports_a_mutated_model(self, enabled, snapshot, fake_client): + adapter = _FakeAdapter(native_error=RuntimeError("bad checkpoint")) + ctx = _make_context(REPO, adapter=adapter, model_path=str(snapshot)) + + with pytest.raises(StrategyFailed) as excinfo: + ServerCacheStrategy().load(LoadResult(value=MagicMock()), ctx) + + assert excinfo.value.mutated is True + + def test_missing_snapshot_is_a_clean_miss(self, enabled, fake_client): + ctx = _make_context(REPO, model_path=None) + + with patch.object(model_prefetch, "ensure_metadata", return_value=None): + with pytest.raises(StrategyFailed) as excinfo: + ServerCacheStrategy().load(LoadResult(value=MagicMock()), ctx) + + assert excinfo.value.mutated is False + assert FakeClient.instances == [] + + +class TestChainOrder: + def test_sits_between_rdma_and_local_strategies(self): + import inspect + + from modelexpress.load_strategy import LoadStrategyChain + + source = inspect.getsource(LoadStrategyChain.run) + order = [ + name + for name in ( + "RdmaStrategy()", + "ServerCacheStrategy()", + "InstantTensorStrategy()", + "DefaultStrategy()", + ) + if name in source + ] + assert order == [ + "RdmaStrategy()", + "ServerCacheStrategy()", + "InstantTensorStrategy()", + "DefaultStrategy()", + ] + assert source.index("RdmaStrategy()") < source.index("ServerCacheStrategy()") + assert source.index("ServerCacheStrategy()") < source.index("InstantTensorStrategy()") diff --git a/modelexpress_client/python/tests/test_vllm_loader.py b/modelexpress_client/python/tests/test_vllm_loader.py index 573cf7c38..eaef75ab7 100644 --- a/modelexpress_client/python/tests/test_vllm_loader.py +++ b/modelexpress_client/python/tests/test_vllm_loader.py @@ -251,9 +251,23 @@ def test_no_remaining_abstract_methods(self): def test_download_model_delegates(self): loader = _make_loader() cfg = MagicMock() - with patch("modelexpress.engines.vllm.loader.DefaultModelLoader") as mock_cls: - loader.download_model(cfg) - mock_cls.return_value.download_model.assert_called_once_with(cfg) + with patch.dict("os.environ", {}, clear=True): + with patch("modelexpress.engines.vllm.loader.DefaultModelLoader") as mock_cls: + loader.download_model(cfg) + mock_cls.return_value.download_model.assert_called_once_with(cfg) + + def test_download_model_defers_without_shared_storage(self): + """A full pre-download here would pull weights before P2P gets a turn.""" + loader = _make_loader() + cfg = MagicMock() + env = { + "MODEL_EXPRESS_NO_SHARED_STORAGE": "1", + "MODEL_EXPRESS_URL": "http://mx:8001", + } + with patch.dict("os.environ", env, clear=True): + with patch("modelexpress.engines.vllm.loader.DefaultModelLoader") as mock_cls: + loader.download_model(cfg) + mock_cls.return_value.download_model.assert_not_called() def test_load_weights_delegates(self): loader = _make_loader()