feat(client): load models from the server cache when workers have no shared storage

[scydas]

What

Part of #569 — deliberately not a full close. The issue lists five things the fallback should cover:

Coverage required by #569	This PR
vLLM --load-format mx / modelexpress	done
SGLang remote_instance with the ModelExpress backend	partial, see below
config, tokenizer, weights and other repository files, not weights alone	done
pinned revisions and private Hugging Face repositories	not done, see below
atomic installation of a complete HF-compatible worker snapshot	done

Lets a vLLM worker get both the model metadata and the weights from the ModelExpress Server, so the fallback chain becomes P2P source → server cache/download → existing local/native fallback and a worker never needs direct Hugging Face access.

The metadata and the weights are fetched at different times, because the engine asks for them at different times:

When	Trigger	Needs
T0	EngineArgs.__post_init__ → get_model_path → snapshot_download	config/tokenizer + a resolvable HF cache layout
T1	tokenizer/processor load	metadata
T2	MxModelLoader.load_model → LoadStrategyChain	weights

Only T2 reaches the strategy chain. T0 happens while engine args are still being parsed — before any loader object exists — so a strategy-chain-only implementation cannot serve it. P2P does not help either: it transfers GPU tensors and never repository files, so every worker needs local metadata regardless of whether P2P later hits.

Hence the split:

• Metadata is pulled at T0, unconditionally, by patching snapshot_download. No weight moves on this path, so P2P keeps first refusal on the part that matters.
• Weights stay in the strategy chain. ServerCacheStrategy is inserted after RdmaStrategy, so it only runs on a P2P miss.

New modules (all under modelexpress_client/python/modelexpress/):

• model_client.py — ModelCacheClient, wrapping EnsureModelDownloaded / ListModelFiles / StreamModelFiles, with stream-protocol validation (chunk offset continuity, is_last_file semantics, commit-hash agreement) and two install entry points: metadata-only and weights-only.
• model_snapshot.py — the Hugging Face cache layout: models--<org>--<name>/{refs/main, snapshots/<commit>/}, flock-guarded staging, per-file rename on publish, and SnapshotPatch.rollback() so an interrupted weight install leaves no partial files.
• model_prefetch.py — T0 orchestration, the enable check, and repo-id recovery from a cache path.
• load_strategy/server_cache_strategy.py — the strategy itself; requires = (EngineAdapter.load_via_native,).
• engines/vllm/patches/patch_hf_snapshot_prefetch.py — takes over snapshot_download across huggingface_hub._snapshot_download / hf_api / huggingface_hub; a prefetch error is logged as a warning and delegates to the original, never blocking startup.

Also:

• generate_proto.sh now generates model_pb2 as well as p2p_pb2 (the Python client previously did not consume ModelService at all).
• Three env vars: MODEL_EXPRESS_NO_SHARED_STORAGE (the switch), MODEL_EXPRESS_CACHE_DIRECTORY, MODEL_EXPRESS_TRANSFER_CHUNK_SIZE.
• engines/vllm/loader.py: download_model returns early when the feature is on, since the snapshot is already in place.
• Docs: docs/ARCHITECTURE.md, docs/DEPLOYMENT.md.

Everything is behind MODEL_EXPRESS_NO_SHARED_STORAGE; with the switch off, none of this code runs (verified on hardware, see below).

model_snapshot.py carries its own weight-suffix list mirroring providers.rs, because the split has to happen client-side: the metadata phase needs everything except weights and the weight phase needs the inverse, and only the former is expressible server-side. #581 has since landed is_weight_file plus an ignore_weights flag on ModelFilesRequest, so the metadata half could stop duplicating the list — worth a follow-up issue rather than widening this PR into a filtering-protocol change.

SGLang is partly covered already: its nixl transport runs the same LoadStrategyChain, and SglangAdapter implements load_via_native, so ServerCacheStrategy is live there with no extra code. Two gaps remain for a separate PR — the transfer_engine transport falls straight to a native load without consulting the chain, and SGLang has no vllm.general_plugins equivalent, so T0 needs an explicit entry point.

The full list of requirements and operational limits lives in docs/DEPLOYMENT.md (Server-Backed Model Cache), and the phase split and stream/rollback contract in docs/ARCHITECTURE.md. Repeating the three that need a reviewer's judgement rather than a reader's awareness — the first two are the open half of #569's fourth bullet:

• Pinned revisions are not supported. The model RPC has no revision field, so the server always answers with what its cache holds for main. Reproduced on hardware: DGD pinned a338b55, server had 7ae5576 — the client warns and refuses to mix revisions rather than silently serving the wrong one. Fixing this needs a proto change plus a registry key change, which is why it is not folded in here.
• Private repositories are untested end to end. The server already supports HF_TOKEN and the worker never talks to Hugging Face, so the path should work by construction, but it has not been exercised on hardware and is not claimed.
• The weight stream is not deduped across pods. The server-side claim dedups the upstream fetch (experiments 6A/6B below), but each pod still streams its own copy, so N replicas cold-starting cost N x model size in server egress.

Note for reviewers of #581: its description states "the Python client is metadata-only and does not consume ModelService". That is no longer true after this PR — the Python client now calls all three ModelService RPCs.

Why

Without this, a worker with no shared storage still has to reach Hugging Face itself for config and tokenizer even when the ModelExpress Server already holds the model, and it has no way to get weights from the server when P2P misses (cold start, first replica, no rank match). The server cache existed but nothing in the Python engine path could use it.

The design constraint that shaped this is that P2P and repository files are disjoint. It is tempting to make the whole thing one strategy, but the engine resolves the model path long before any strategy runs, and it fails at parse_args with LocalEntryNotFoundError if the cache layout is not already valid. So the metadata half has to be earlier, and it has to be unconditional.

Testing

Unit: 1205 passed / 46 skipped for the full Python suite. New test files cover the client's stream protocol, the snapshot layout and rollback, the prefetch, the snapshot_download patch, and the strategy.

On real hardware — 2 nodes, 1× A800 80GB each, vLLM via Dynamo DGD, ModelExpress Server 0.4.0 with an internal HF mirror. Nine scenarios, every one re-run against this branch's single commit after the review fixes were folded in. "upstream" counts Hugging Face fetches by the server; "streams" counts StreamModelFiles calls served from the server's own cache, split into the metadata phase and the weight phase.

#	Scenario	upstream	metadata streams	weight streams	Result
1	Server cache cold	1	1	1	COMPLETE in 11.83s, READY
2	Server cache warm	0	1	1	metadata 0.157s
3	P2P available	0	1	0	RDMA 0.99 GB / 86ms / 92.5 Gbps, COMPLETE in 1.75s
4	Switch off	0	0	0	Server delta log 0 lines, none of this code executes
5	4 processes, same pod	0	1	1	4/4 ok; 8 calls collapsed to 2 streams by flock, zero leftovers
6A	4 independent pods	1	4	4	4/4 ok in 37s; claim dedups upstream, streams are per-pod
6B	DGD 2 replicas, simultaneous cold start	1	2	2	both released at the same instant, 12.48s / 9.19s
7	Server unreachable	0	0	0	degrades in 20.1s, no half-written files, failure semantics unchanged
8	7B / 4 shards / 15.24 GB	1	1 (8 files)	1 (4 files)	metadata 0.17s vs weights 114.27s, all shards present
9	Server killed mid-transfer	1	1	1, aborted	weights fully rolled back, metadata kept, retry resumes

Rows 6A and 6B show the dedup boundary: one upstream fetch no matter how many workers race, but one stream per worker in each phase. That is the cross-pod limit noted above, measured.

Row 3 is the one that carries the design claim. server-cache appears in Eligible loaders, so it could have run, but rdma is tried first and wins — and the server's own log for that window shows exactly one stream, Found 7 files, the metadata set. Not one weight byte came from the server while a P2P source was alive. Fetching metadata early does not weaken P2P-first.

Cross-process state loss

server-cache initially filtered itself out with No Hugging Face repo id for '/home/dynamo/.cache/.../snapshots/7ae5576...', skipping server cache. vLLM rewrites ModelConfig.model in place with the resolved path and loads weights in a separate EngineCore process, so the in-process path → repo id map recorded at T0 is empty where the strategy runs. Unit tests could not catch this because they run in one process. Fixed by deriving the repo id deterministically from the models--<org>--<name> directory name; TestRepoIdFromCachePath now pins that behaviour with the reasoning in its docstring.

Summary by CodeRabbit

• New Features

  • Added server-backed model caching for environments without shared storage.
  • Model metadata and weight files can now be fetched, validated, and installed locally.
  • Added prefetch support for Hugging Face model snapshots.
  • Added the public ModelCacheClient for managing cached models.
  • Added configuration options for cache location, storage mode, and transfer chunk size.

• Documentation

  • Documented server-backed caching, deployment configuration, fallback behavior, and operational considerations.

• Tests

  • Added extensive coverage for caching, streaming validation, snapshot management, prefetching, and loader integration.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Walkthrough

Changes

Server-backed model caching

Layer / File(s)	Summary
Model service contract and configuration modelexpress_client/python/generate_proto.sh, modelexpress_client/python/modelexpress/model_pb2*, modelexpress_client/python/modelexpress/envs.py	Added model download RPC bindings, cache configuration readers, package exports, and generation support for p2p and model protos.
Snapshot cache and streaming client modelexpress_client/python/modelexpress/model_snapshot.py, modelexpress_client/python/modelexpress/model_client.py, modelexpress_client/python/tests/test_model_{snapshot,client}.py	Added secure snapshot staging, atomic publication, transactional weight patching, validated file streaming, and rollback behavior.
Metadata prefetch and loader delegation modelexpress_client/python/modelexpress/model_prefetch.py, modelexpress_client/python/modelexpress/engines/vllm/..., modelexpress_client/python/tests/test_{model_prefetch,hf_snapshot_prefetch_patch,vllm_loader}.py	Added synchronized metadata prefetch, Hugging Face interception, transfer configuration, and vLLM download delegation for no-shared-storage mode.
Server cache load strategy modelexpress_client/python/modelexpress/load_strategy/..., modelexpress_client/python/tests/test_server_cache_strategy.py, docs/*	Added ServerCacheStrategy to the load chain and documented server-backed metadata and weight caching behavior.
Estimated code review effort: 5 (Critical)	~120 minutes

Poem

I’m a rabbit with snapshots tucked neat,
Streaming each weight in a filesystem suite.
Metadata hops, then the tensors arrive,
Atomic refs keep the cache alive.
P2P leads, server cache follows—
Hop, hop, hooray for cleaner tomorrows!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 22.92% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: server-cache model loading for workers without shared storage.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 5

🧹 Nitpick comments (3)

modelexpress_client/python/tests/test_model_prefetch.py (2)

65-76: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Rename the fake_client fixture to reflect its return value.

The fixture installs the fake client, but it returns the snapshot Path. Assertions then read ensure_metadata(REPO) == fake_client, which suggests a client comparison. A name such as installed_snapshot makes the assertions self-describing.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@modelexpress_client/python/tests/test_model_prefetch.py` around lines 65 -
76, Rename the fake_client fixture to installed_snapshot, reflecting that it
returns the snapshot Path rather than a client, and update all test parameters
and assertions that reference fake_client accordingly.

---

177-218: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Replace the sleep-based ordering with an explicit event.

The test depends on the second thread starting within the 0.3 s install window. On a loaded runner, that window can be missed, and installs == [REPO] can then fail. An threading.Event makes the interleaving deterministic and removes the fixed sleeps.

♻️ Proposed deterministic synchronization

         snapshot = tmp_path / "models--org--model" / "snapshots" / ("a" * 40)
         snapshot.mkdir(parents=True)
         installs = []
+        install_started = threading.Event()
+        release_install = threading.Event()
 
         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)
+                install_started.set()
+                release_install.wait(timeout=5)
                 return snapshot
 
         monkeypatch.setattr("modelexpress.model_client.ModelCacheClient", SlowClient)
 
         results = {}
 
-        def call(tag, delay):
-            time.sleep(delay)
-            results[tag] = model_prefetch.ensure_metadata(REPO)
+        def call(tag, wait_for_install):
+            if wait_for_install:
+                install_started.wait(timeout=5)
+            results[tag] = model_prefetch.ensure_metadata(REPO)
 
         threads = [
-            threading.Thread(target=call, args=("first", 0.0)),
-            threading.Thread(target=call, args=("second", 0.05)),
+            threading.Thread(target=call, args=("first", False)),
+            threading.Thread(target=call, args=("second", True)),
         ]
         for thread in threads:
             thread.start()
+        install_started.wait(timeout=5)
+        release_install.set()
         for thread in threads:
             thread.join()

Remove the now-unused time import if no other test needs it.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@modelexpress_client/python/tests/test_model_prefetch.py` around lines 177 -
218, Update test_second_caller_waits_and_gets_the_same_snapshot to coordinate
the two threads with a threading.Event: signal when install_metadata_snapshot
begins, then have the second caller wait for that signal before invoking
ensure_metadata. Remove the delay-based time.sleep calls and delete the time
import if unused, preserving the assertions that both callers receive the same
snapshot and only one install occurs.

modelexpress_client/python/tests/test_server_cache_strategy.py (1)

227-250: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Assert the chain order from behavior instead of source text.

The test inspects the source of LoadStrategyChain.run and compares substring positions. It fails on formatting-only edits, and it still passes if the ordering later moves out of run. Drive the chain instead and record the order in which strategies are attempted.

One approach: patch LoadStrategy.is_available to return True and LoadStrategy.load to append self.name and raise StrategyFailed(..., mutated=False), then assert the recorded names. The chain raises RuntimeError at the end, which the test can expect.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@modelexpress_client/python/tests/test_server_cache_strategy.py` around lines
227 - 250, The test test_sits_between_rdma_and_local_strategies should verify
runtime strategy-attempt order instead of inspecting LoadStrategyChain.run
source text. Patch LoadStrategy.is_available to allow every strategy, patch
LoadStrategy.load to record self.name and raise StrategyFailed with
mutated=False, execute the chain while expecting its final RuntimeError, then
assert the recorded names are RdmaStrategy, ServerCacheStrategy,
InstantTensorStrategy, and DefaultStrategy in that order.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/DEPLOYMENT.md`:
- Around line 742-746: Add `MODEL_EXPRESS_URL` and `MX_SERVER_ADDRESS` as
separate rows in the deployment environment-variable table, documenting their
roles as the server address options required alongside
`MODEL_EXPRESS_NO_SHARED_STORAGE`. Preserve the existing rows and the
`MODEL_EXPRESS_TRANSFER_CHUNK_SIZE` default.

In `@modelexpress_client/python/modelexpress/model_client.py`:
- Around line 87-90: Update the channel options in the client initialization to
include gRPC keepalive settings that detect silent or dropped peers without
imposing fixed deadlines on the streaming RPCs used by ensure_downloaded and
_stream_into. Add an appropriate bounded deadline to the unary list_files call,
while preserving the existing message-size limits and streaming behavior.
- Around line 86-96: Update the channel construction in the ModelCacheClient
stub initialization so the token added by auth.with_auth is transmitted over
encrypted TLS transport instead of grpc.insecure_channel. Configure and pass
appropriate TLS credentials, or enforce an mTLS/TLS proxy requirement for
authenticated deployments, while preserving the existing message-size options
and ModelServiceStub setup.

In `@modelexpress_client/python/modelexpress/model_snapshot.py`:
- Around line 272-292: Update the exception handler surrounding the staged
snapshot replacement in the snapshot commit method to catch BaseException rather
than Exception, while preserving the existing stale_path restoration condition
and re-raising behavior so interruptions also restore the displaced snapshot.

In `@modelexpress_client/python/tests/test_model_snapshot.py`:
- Around line 144-151: Update test_discard_removes_staging to capture
staging.path before calling staging.discard(), then assert that captured path no
longer exists afterward. Remove the post-discard _staging_path guard while
preserving the existing snapshots-directory assertion.

---

Nitpick comments:
In `@modelexpress_client/python/tests/test_model_prefetch.py`:
- Around line 65-76: Rename the fake_client fixture to installed_snapshot,
reflecting that it returns the snapshot Path rather than a client, and update
all test parameters and assertions that reference fake_client accordingly.
- Around line 177-218: Update
test_second_caller_waits_and_gets_the_same_snapshot to coordinate the two
threads with a threading.Event: signal when install_metadata_snapshot begins,
then have the second caller wait for that signal before invoking
ensure_metadata. Remove the delay-based time.sleep calls and delete the time
import if unused, preserving the assertions that both callers receive the same
snapshot and only one install occurs.

In `@modelexpress_client/python/tests/test_server_cache_strategy.py`:
- Around line 227-250: The test test_sits_between_rdma_and_local_strategies
should verify runtime strategy-attempt order instead of inspecting
LoadStrategyChain.run source text. Patch LoadStrategy.is_available to allow
every strategy, patch LoadStrategy.load to record self.name and raise
StrategyFailed with mutated=False, execute the chain while expecting its final
RuntimeError, then assert the recorded names are RdmaStrategy,
ServerCacheStrategy, InstantTensorStrategy, and DefaultStrategy in that order.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5203ad66-2fee-437f-8550-c44854f140d1

📥 Commits

Reviewing files that changed from the base of the PR and between a5982b1 and 782a3ac.

📒 Files selected for processing (22)

• docs/ARCHITECTURE.md
• docs/DEPLOYMENT.md
• modelexpress_client/python/generate_proto.sh
• modelexpress_client/python/modelexpress/__init__.py
• modelexpress_client/python/modelexpress/engines/vllm/loader.py
• modelexpress_client/python/modelexpress/engines/vllm/patches/__init__.py
• modelexpress_client/python/modelexpress/engines/vllm/patches/patch_hf_snapshot_prefetch.py
• modelexpress_client/python/modelexpress/envs.py
• modelexpress_client/python/modelexpress/load_strategy/__init__.py
• modelexpress_client/python/modelexpress/load_strategy/server_cache_strategy.py
• modelexpress_client/python/modelexpress/model_client.py
• modelexpress_client/python/modelexpress/model_pb2.py
• modelexpress_client/python/modelexpress/model_pb2_grpc.py
• modelexpress_client/python/modelexpress/model_prefetch.py
• modelexpress_client/python/modelexpress/model_snapshot.py
• modelexpress_client/python/tests/test_envs.py
• modelexpress_client/python/tests/test_hf_snapshot_prefetch_patch.py
• modelexpress_client/python/tests/test_model_client.py
• modelexpress_client/python/tests/test_model_prefetch.py
• modelexpress_client/python/tests/test_model_snapshot.py
• modelexpress_client/python/tests/test_server_cache_strategy.py
• modelexpress_client/python/tests/test_vllm_loader.py

[AndyDai-nv]

/ok to test dd3fcc2

[AndyDai-nv]

/ok to test 422df8b

[AndyDai-nv]

Thanks @scydas, lgtm, need your quick look @zhengluo-nv as well to see whether it is good to close 569

[zhengluo-nv]

/ok to test 645d969

[zhengluo-nv]

Approving as my understanding is the server_cache_strategy is disabled by default