[Feature] Add native ModelExpress weight transfer backend (#58399)
Signed-off-by: Hyunjae Woo <hwoo@nvidia.com> Co-authored-by: aoshen02 <aoshen@inferact.ai>
Hyunjae Woo committed
Oct 1, 2026 at 19:29 UTC
5a661487e29433b009508c8b3045847eb11041fa
5 files changed
+558
docs/training/weight_transfer/README.md
+1
@@ -36,6 +36,7 @@ engine drives on your behalf:
36
| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Checkpoint-coordinate sparse weight patches |
37
| [sharded_rdt](sharded_rdt.md) | NIXL / Ray Direct Transport (pull-based) | Very large models where each worker needs only its own slice (MoE with expert parallelism) |
38
| [nccl_m2n](m2n.md) | NCCL M2N reshard | Trainer and inference use different sharding layouts |
39
+| [ModelExpress](modelexpress.md) | ModelExpress version-based transfer | Install immutable published weight versions, including object-storage checkpoints and deltas |
40
41
## Quickstart
42
docs/training/weight_transfer/modelexpress.md
new
+156
@@ -0,0 +1,156 @@
1
+# ModelExpress Engine
2
+
3
+The `modelexpress` weight transfer engine installs immutable weight versions
4
+published through [ModelExpress](https://github.com/ai-dynamo/modelexpress).
5
+Each update identifies the exact version to install, while ModelExpress handles
6
+source selection, transfer, and installation.
7
+
8
+## When to Use ModelExpress
9
+
10
+- Your workflow manages immutable weight versions through the ModelExpress
11
+ Refit API.
12
+- Inference workers need to pull a specific published version through
13
+ ModelExpress, using GPU peer transfer or object storage.
14
+- You want to load full checkpoints or replay checkpoint deltas from S3.
15
+
16
+## How It Works
17
+
18
+1. A publisher creates a weight-version resource through the ModelExpress Refit
19
+ API, publishes its payload, and marks the version READY.
20
+2. The orchestrator pauses generation and calls `start_weight_update` on the
21
+ inference engine.
22
+3. The orchestrator calls `update_weights` with the target `version_id`.
23
+ Each worker's `ModelExpressGeneratorClient` resolves and stages that version,
24
+ then applies it to the live vLLM model.
25
+4. After a successful update, `finish_weight_update` releases the staged handle.
26
+ The orchestrator can then resume inference.
27
+
28
+Initialize the backend before starting an update session, and send exactly one
29
+version per session. Invalid transitions must raise; an orchestrator must keep
30
+inference paused after an update error and retry with a new session only after
31
+cleanup. An error does not roll back weights that were already applied.
32
+
33
+`ModelExpressTrainerClient` is an optional publisher. Custom integrations can
34
+manage weight-version resources directly through the ModelExpress server's
35
+Refit API. For S3, create a STAGING version with the correct model, payload
36
+format, lineage, and index URI; upload the index and referenced shards in the
37
+format expected by `ModelExpressGeneratorClient`; then mark the version READY.
38
+See the [S3 Delta Weight Refit guide](https://github.com/ai-dynamo/modelexpress/blob/main/docs/S3_DELTA_WEIGHT_REFIT.md)
39
+for the publication sequence and generator artifact requirements.
40
+
41
+vLLM's `modelexpress_engine.py` re-exports the engine and its init/update
42
+dataclasses from
43
+`modelexpress_rl.inference.engines.vllm.weight_transfer_engine`.
44
+The implementation and configuration schema live in ModelExpress. Upgrading
45
+the installed ModelExpress package and restarting workers picks up compatible
46
+changes without a vLLM code change. The package must continue to implement
47
+vLLM's `WeightTransferEngine` interface.
48
+
49
+## Inference Side
50
+
51
+!!! warning "Development version"
52
+
53
+ This backend currently uses unreleased ModelExpress features from `main`,
54
+ which may introduce breaking changes. These features will be included in
55
+ the official ModelExpress v0.8.0 release.
56
+
57
+Install ModelExpress from the `main` branch on every worker:
58
+
59
+```bash
60
+uv pip install "git+https://github.com/ai-dynamo/modelexpress@main#subdirectory=modelexpress_client/python"
61
+```
62
+
63
+The backend is registered natively and loaded only when selected;
64
+`VLLM_PLUGINS=modelexpress` is not required for weight transfer. If the plugin
65
+is enabled, it preserves the native registration and uses the same engine
66
+implementation.
67
+
68
+```python
69
+from vllm import LLM
70
+from vllm.config import WeightTransferConfig
71
+
72
+llm = LLM(model="my-model", weight_transfer_config=WeightTransferConfig(backend="modelexpress"))
73
+```
74
+
75
+```bash
76
+vllm serve /models/launch \
77
+ --load-format safetensors \
78
+ --weight-transfer-config '{"backend":"modelexpress"}'
79
+```
80
+
81
+### Initialize the engine
82
+
83
+Initialize once by passing `init_info` to `init_weight_transfer_engine`.
84
+For HTTP administration, set `VLLM_SERVER_DEV_MODE=1` before starting vLLM and
85
+restrict these endpoints to trusted orchestrators.
86
+
87
+For S3 updates, send the following body to `POST /init_weight_transfer_engine`:
88
+
89
+```json
90
+{
91
+ "init_info": {
92
+ "model_name": "policy",
93
+ "server_url": "modelexpress:8001",
94
+ "object_storage_type": "S3",
95
+ "initial_base_version_id": "policy-v0",
96
+ "seed_checkpoint_path": "/models/launch",
97
+ "refit_checkpoint_dir": "/mxdelta/receiver",
98
+ "refit_checkpoint_max_size_gb": 200,
99
+ "object_storage_region_name": "us-west-2"
100
+ }
101
+}
102
+```
103
+
104
+### `ModelExpressWeightTransferInitInfo`
105
+
106
+These fields belong inside `init_info`, rather than
107
+`--weight-transfer-config`. The installed ModelExpress package defines the
108
+schema and defaults.
109
+
110
+| Field | Default | Description |
111
+| ----- | ------- | ----------- |
112
+| `model_name` | vLLM model name | Logical model identity used for published ModelExpress versions. Set this when it differs from the vLLM model path or name. |
113
+| `server_url` | ModelExpress configuration | ModelExpress server address, such as `modelexpress:8001`. If omitted, uses ModelExpress's environment settings, then `localhost:8001`. |
114
+| `initial_serving_version_id` | `None` | Version already loaded by the inference worker. If omitted, ModelExpress uses the initial version supplied by its runtime, when available. |
115
+| `registration_ttl_seconds` | Three MX heartbeat intervals | Worker registration lifetime in seconds, renewed by ModelExpress. |
116
+| `lease_ttl_seconds` | Registration lifetime | Weight-version lease lifetime in seconds, renewed while held. |
117
+| `max_transfer_attempts` | `3` | Maximum source-discovery and transfer attempts for one staged update. |
118
+| `max_replay_chain_length` | `64` | Maximum number of payload revisions replayed to reach the requested version. |
119
+| `rpc_timeout_seconds` | `30.0` | Deadline in seconds for each control-plane or manifest RPC. |
120
+
121
+### Object storage settings
122
+
123
+To enable object storage, provide `object_storage_type`,
124
+`initial_base_version_id`, `seed_checkpoint_path`, and `refit_checkpoint_dir`
125
+inside `init_info`. Endpoint and region settings also require these four fields.
126
+
127
+| Field | Default | Description |
128
+| ----- | ------- | ----------- |
129
+| `object_storage_type` | `None` | Object storage provider. Use `"S3"`; the generator currently supports S3. |
130
+| `initial_base_version_id` | `None` | Registered base version matching the seed checkpoint. Required for object storage. |
131
+| `seed_checkpoint_path` | `None` | Local seed checkpoint path accessible to each worker. Required for object storage. |
132
+| `refit_checkpoint_dir` | `None` | Local directory for cached checkpoints and delta replay. Required for object storage. |
133
+| `refit_checkpoint_max_size_gb` | ModelExpress default | Checkpoint cache quota in decimal GB. The example sets an explicit 200 GB quota; it does not reserve disk space. |
134
+| `object_storage_endpoint_url` | `None` | Optional endpoint override for an S3-compatible service. |
135
+| `object_storage_region_name` | `None` | Optional S3 region, such as `us-west-2`. |
136
+
137
+The seed checkpoint must match `initial_base_version_id`. Credentials use
138
+ModelExpress's normal credential chain. Omit the object storage settings when
139
+using only GPU peer transfer.
140
+
141
+### Apply a weight update
142
+
143
+After the target version is READY, pause generation and call
144
+`POST /start_weight_update`. Then send the exact immutable version ID to
145
+`POST /update_weights`:
146
+
147
+```json
148
+{"update_info": {"version_id": "policy-v1"}}
149
+```
150
+
151
+Call `POST /finish_weight_update` only after the update succeeds, then resume
152
+inference. The staged handle is retained until finish or shutdown. Staging and
153
+installation failures propagate; do not resume a worker with uncertain weights.
154
+
155
+See the ModelExpress documentation for publication, source selection, object
156
+storage, and deployment requirements.
tests/distributed/test_weight_transfer.py
+353
@@ -6,10 +6,14 @@ Unit tests for engine classes (parsing, validation, registry).
6
Integration tests for NCCL and IPC weight transfer between processes using Ray.
7
"""
8
9
+import builtins
10
import pickle
11
+import runpy
12
import threading
13
import time
14
from contextlib import nullcontext
15
+from pathlib import Path
16
+from types import SimpleNamespace
17
from unittest.mock import MagicMock
18
19
import pybase64 as base64
@@ -297,6 +301,36 @@ class TestEngineRegistry:
301
)
302
assert isinstance(engine, SparseNCCLWeightTransferEngine)
303
304
+ def test_modelexpress_registration_is_native_and_lazy(self, monkeypatch):
305
+ from vllm.distributed.weight_transfer import factory
306
+
307
+ imported = []
308
+ module = MagicMock()
309
+ module.ModelExpressWeightTransferEngine.__name__ = (
310
+ "ModelExpressWeightTransferEngine"
311
+ )
312
+
313
+ def import_module(name):
314
+ imported.append(name)
315
+ return module
316
+
317
+ monkeypatch.setattr(factory.importlib, "import_module", import_module)
318
+ config = WeightTransferConfig(backend="modelexpress")
319
+ vllm_config = create_mock_vllm_config()
320
+ device = torch.device("cpu")
321
+ model = MagicMock(spec=torch.nn.Module)
322
+
323
+ assert imported == []
324
+ engine = WeightTransferEngineFactory.create_engine(
325
+ config, vllm_config, device, model
326
+ )
327
+
328
+ assert imported == ["vllm.distributed.weight_transfer.modelexpress_engine"]
329
+ module.ModelExpressWeightTransferEngine.assert_called_once_with(
330
+ config, vllm_config, device, model
331
+ )
332
+ assert engine is module.ModelExpressWeightTransferEngine.return_value
333
+
334
def test_create_engine_invalid_backend(self):
335
config = WeightTransferConfig(backend="invalid")
336
with pytest.raises(ValueError, match="Invalid weight transfer backend"):
@@ -1976,3 +2010,322 @@ def test_sparse_nccl_trainer_non_sender_skips_client():
2010
engine.send_weights([_sparse_patch()])
2011
engine.send_weight_chunk([_sparse_patch()])
2012
assert client.order == []
2013
+
2014
+
2015
+# --- Optional ModelExpress Client Lifecycle ---
2016
+
2017
+
2018
+def _import_modelexpress_shim():
2019
+ path = (
2020
+ Path(__file__).parents[2]
2021
+ / "vllm/distributed/weight_transfer/modelexpress_engine.py"
2022
+ )
2023
+ runpy.run_path(str(path), run_name="_test_modelexpress_shim")
2024
+
2025
+
2026
+@pytest.mark.parametrize(
2027
+ "missing_module",
2028
+ [
2029
+ "modelexpress",
2030
+ "modelexpress_rl",
2031
+ "modelexpress_rl.inference",
2032
+ "modelexpress_rl.inference.engines",
2033
+ "modelexpress_rl.inference.engines.vllm",
2034
+ "modelexpress_rl.inference.engines.vllm.weight_transfer_engine",
2035
+ ],
2036
+)
2037
+def test_modelexpress_missing_backend_explains_source_install(
2038
+ monkeypatch, missing_module
2039
+):
2040
+ original_import = builtins.__import__
2041
+ error = ModuleNotFoundError(name=missing_module)
2042
+
2043
+ def import_without_backend(name, *args, **kwargs):
2044
+ if name == "modelexpress":
2045
+ raise error
2046
+ return original_import(name, *args, **kwargs)
2047
+
2048
+ monkeypatch.setattr(builtins, "__import__", import_without_backend)
2049
+ with pytest.raises(ImportError, match="modelexpress_client/python") as exc_info:
2050
+ _import_modelexpress_shim()
2051
+ assert exc_info.value.__cause__ is error
2052
+ assert "uv pip install" in str(exc_info.value)
2053
+
2054
+
2055
+@pytest.mark.parametrize("missing_module", ["grpc", "modelexpress.internal_dependency"])
2056
+def test_modelexpress_preserves_transitive_import_errors(monkeypatch, missing_module):
2057
+ original_import = builtins.__import__
2058
+ error = ModuleNotFoundError(name=missing_module)
2059
+
2060
+ def import_without_dependency(name, *args, **kwargs):
2061
+ if name == "modelexpress":
2062
+ raise error
2063
+ return original_import(name, *args, **kwargs)
2064
+
2065
+ monkeypatch.setattr(builtins, "__import__", import_without_dependency)
2066
+ with pytest.raises(ModuleNotFoundError) as exc_info:
2067
+ _import_modelexpress_shim()
2068
+ assert exc_info.value is error
2069
+
2070
+
2071
+@pytest.fixture
2072
+def mx_client(monkeypatch):
2073
+ client_module = pytest.importorskip("modelexpress_rl.inference.client")
2074
+ client = MagicMock()
2075
+ monkeypatch.setattr(
2076
+ client_module.ModelExpressGeneratorClient,
2077
+ "initialize",
2078
+ MagicMock(return_value=client),
2079
+ )
2080
+ monkeypatch.setattr(torch.accelerator, "synchronize", MagicMock())
2081
+ return client
2082
+
2083
+
2084
+def _make_modelexpress_engine():
2085
+ return WeightTransferEngineFactory.create_engine(
2086
+ WeightTransferConfig(backend="modelexpress"),
2087
+ SimpleNamespace(
2088
+ parallel_config=SimpleNamespace(),
2089
+ model_config=SimpleNamespace(model="test/model"),
2090
+ ),
2091
+ torch.device("cpu"),
2092
+ torch.nn.Linear(2, 2),
2093
+ )
2094
+
2095
+
2096
+@pytest.fixture
2097
+def mx_worker(mx_client, monkeypatch):
2098
+ from vllm.v1.worker import gpu_worker
2099
+
2100
+ monkeypatch.setattr(gpu_worker, "set_current_vllm_config", lambda _: nullcontext())
2101
+ worker = object.__new__(gpu_worker.Worker)
2102
+ worker.weight_transfer_engine = _make_modelexpress_engine()
2103
+ worker.vllm_config = worker.weight_transfer_engine.vllm_config
2104
+ worker._weight_update_active = False
2105
+ worker._weight_update_is_draft = False
2106
+ worker.model_runner = MagicMock()
2107
+ return worker
2108
+
2109
+
2110
+def test_modelexpress_worker_rejects_updates_before_initialization(
2111
+ mx_worker, mx_client
2112
+):
2113
+ with pytest.raises(RuntimeError, match="not initialized"):
2114
+ mx_worker.start_weight_update()
2115
+ with pytest.raises(RuntimeError, match="start_weight_update must be called"):
2116
+ mx_worker.update_weights({"version_id": "version-a"})
2117
+ with pytest.raises(RuntimeError, match="without a matching"):
2118
+ mx_worker.finish_weight_update()
2119
+ mx_client.stage_weight.assert_not_called()
2120
+
2121
+
2122
+@pytest.mark.parametrize(
2123
+ "second_update, error_type",
2124
+ [
2125
+ ({"version_id": "version-a"}, RuntimeError),
2126
+ ({"version_id": "version-b"}, RuntimeError),
2127
+ ({"version_id": ""}, ValueError),
2128
+ ({"version_id": None}, ValueError),
2129
+ ({}, ValueError),
2130
+ ],
2131
+)
2132
+def test_modelexpress_worker_recovers_after_rejected_update(
2133
+ mx_worker, mx_client, second_update, error_type
2134
+):
2135
+ """A rejected payload must release A so a fresh session can actually install B."""
2136
+ mx_worker.weight_transfer_engine.init_transfer_engine(
2137
+ mx_worker.weight_transfer_engine.parse_init_info({})
2138
+ )
2139
+ staged_a = SimpleNamespace(version_id="version-a", metrics={}, release=MagicMock())
2140
+ staged_b = SimpleNamespace(version_id="version-b", metrics={}, release=MagicMock())
2141
+ mx_client.stage_weight.side_effect = [staged_a, staged_b]
2142
+ mx_client.apply_weight.return_value = {}
2143
+ mx_worker.start_weight_update()
2144
+ mx_worker.update_weights({"version_id": "version-a"})
2145
+
2146
+ with pytest.raises(error_type):
2147
+ mx_worker.update_weights(second_update)
2148
+ staged_a.release.assert_called_once_with()
2149
+ mx_client.apply_weight.assert_called_once_with(staged_a)
2150
+
2151
+ mx_worker.start_weight_update()
2152
+ mx_worker.update_weights({"version_id": "version-b"})
2153
+ mx_worker.finish_weight_update()
2154
+
2155
+ assert [
2156
+ call.kwargs["version"].version_id
2157
+ for call in mx_client.stage_weight.call_args_list
2158
+ ] == ["version-a", "version-b"]
2159
+ mx_client.apply_weight.assert_called_with(staged_b)
2160
+ staged_b.release.assert_called_once_with()
2161
+
2162
+
2163
+def test_modelexpress_worker_can_retry_failed_finish(mx_worker, mx_client):
2164
+ mx_worker.weight_transfer_engine.init_transfer_engine(
2165
+ mx_worker.weight_transfer_engine.parse_init_info({})
2166
+ )
2167
+ staged = mx_client.stage_weight.return_value
2168
+ staged.release.side_effect = [RuntimeError("release failed"), None]
2169
+ mx_worker.start_weight_update()
2170
+ mx_worker.update_weights({"version_id": "version-a"})
2171
+
2172
+ with pytest.raises(RuntimeError, match="release failed"):
2173
+ mx_worker.finish_weight_update()
2174
+ mx_worker.finish_weight_update()
2175
+
2176
+ assert staged.release.call_count == 2
2177
+ assert not mx_worker._weight_update_active
2178
+ mx_worker.start_weight_update()
2179
+
2180
+
2181
+def test_native_engine_applies_and_releases_exact_versions(mx_client):
2182
+ from modelexpress_rl.inference.engines.vllm import (
2183
+ weight_transfer_engine as mx_engine,
2184
+ )
2185
+
2186
+ from vllm.distributed.weight_transfer import modelexpress_engine
2187
+
2188
+ engine = _make_modelexpress_engine()
2189
+ assert type(engine) is mx_engine.ModelExpressWeightTransferEngine
2190
+ assert (
2191
+ modelexpress_engine.ModelExpressWeightTransferInitInfo
2192
+ is mx_engine.ModelExpressWeightTransferInitInfo
2193
+ )
2194
+ assert (
2195
+ modelexpress_engine.ModelExpressWeightTransferUpdateInfo
2196
+ is mx_engine.ModelExpressWeightTransferUpdateInfo
2197
+ )
2198
+ assert not engine.supports_draft_weight_update
2199
+ engine.init_transfer_engine(engine.parse_init_info({}))
2200
+
2201
+ for version_id in ("version-a", "version-b"):
2202
+ staged = SimpleNamespace(version_id=version_id, metrics={}, release=MagicMock())
2203
+ mx_client.stage_weight.return_value = staged
2204
+ mx_client.apply_weight.return_value = {}
2205
+ engine.start_weight_update()
2206
+ engine.update_weights({"version_id": version_id})
2207
+
2208
+ assert (
2209
+ mx_client.stage_weight.call_args.kwargs["version"].version_id == version_id
2210
+ )
2211
+ mx_client.apply_weight.assert_called_with(staged)
2212
+ torch.accelerator.synchronize.assert_called()
2213
+ staged.release.assert_not_called()
2214
+ engine.finish_weight_update()
2215
+ staged.release.assert_called_once_with()
2216
+
2217
+ engine.shutdown()
2218
+ engine.shutdown()
2219
+ mx_client.close.assert_called_once_with()
2220
+
2221
+
2222
+def test_vime_init_config_reaches_mx_client(mx_client):
2223
+ from modelexpress_rl.inference.client import ModelExpressGeneratorClient
2224
+
2225
+ engine = _make_modelexpress_engine()
2226
+ ModelExpressGeneratorClient.initialize.assert_not_called()
2227
+ engine.init_transfer_engine(
2228
+ engine.parse_init_info(
2229
+ {
2230
+ "model_name": "policy",
2231
+ "server_url": "mx:8001",
2232
+ "initial_serving_version_id": "serving-a",
2233
+ "object_storage_type": "S3",
2234
+ "initial_base_version_id": "base-a",
2235
+ "seed_checkpoint_path": "/models/launch",
2236
+ "refit_checkpoint_dir": "/mxdelta/receiver",
2237
+ "refit_checkpoint_max_size_gb": 200,
2238
+ "object_storage_endpoint_url": "http://minio:9000",
2239
+ "object_storage_region_name": "us-west-2",
2240
+ "registration_ttl_seconds": 90,
2241
+ "lease_ttl_seconds": 60,
2242
+ "max_transfer_attempts": 4,
2243
+ "max_replay_chain_length": 17,
2244
+ "rpc_timeout_seconds": 12.5,
2245
+ }
2246
+ )
2247
+ )
2248
+ config = ModelExpressGeneratorClient.initialize.call_args.args[0]
2249
+ assert config.engine_context.model is engine.model
2250
+ assert config.engine_context.vllm_config is engine.vllm_config
2251
+ assert config.model_name == "policy"
2252
+ assert config.server_url == "mx:8001"
2253
+ assert config.initial_serving_version_id == "serving-a"
2254
+ assert config.registration_ttl_seconds == 90
2255
+ assert config.lease_ttl_seconds == 60
2256
+ assert config.max_transfer_attempts == 4
2257
+ assert config.max_replay_chain_length == 17
2258
+ assert config.rpc_timeout_seconds == 12.5
2259
+ storage = config.object_storage
2260
+ assert storage.storage_type.value == "S3"
2261
+ assert storage.initial_base_version_id == "base-a"
2262
+ assert storage.seed_checkpoint_path == "/models/launch"
2263
+ assert storage.refit_checkpoint_dir == "/mxdelta/receiver"
2264
+ assert storage.refit_checkpoint_max_size_gb == 200
2265
+ assert storage.endpoint_url == "http://minio:9000"
2266
+ assert storage.region_name == "us-west-2"
2267
+
2268
+
2269
+@pytest.mark.parametrize(
2270
+ "key",
2271
+ [
2272
+ "object_storage_type",
2273
+ "initial_base_version_id",
2274
+ "seed_checkpoint_path",
2275
+ "refit_checkpoint_dir",
2276
+ "object_storage_endpoint_url",
2277
+ "object_storage_region_name",
2278
+ ],
2279
+)
2280
+def test_incomplete_storage_config_rejected_before_init(mx_client, key):
2281
+ from modelexpress_rl.inference.client import ModelExpressGeneratorClient
2282
+
2283
+ engine = _make_modelexpress_engine()
2284
+ with pytest.raises(ValueError, match="object storage requires"):
2285
+ engine.init_transfer_engine(engine.parse_init_info({key: "value"}))
2286
+ ModelExpressGeneratorClient.initialize.assert_not_called()
2287
+
2288
+
2289
+@pytest.mark.parametrize("version_id", ["", " ", None, 1])
2290
+def test_invalid_version_rejected_before_staging(mx_client, version_id):
2291
+ engine = _make_modelexpress_engine()
2292
+ engine.init_transfer_engine(engine.parse_init_info({}))
2293
+ engine.start_weight_update()
2294
+ with pytest.raises(ValueError, match="version_id is required"):
2295
+ engine.update_weights({"version_id": version_id})
2296
+ mx_client.stage_weight.assert_not_called()
2297
+
2298
+
2299
+@pytest.mark.parametrize("failure", ["stage", "apply", "release"])
2300
+def test_update_failure_preserves_error_and_releases_handle(mx_client, failure):
2301
+ engine = _make_modelexpress_engine()
2302
+ engine.init_transfer_engine(engine.parse_init_info({}))
2303
+ staged = SimpleNamespace(version_id="version-a", metrics={}, release=MagicMock())
2304
+ mx_client.stage_weight.return_value = staged
2305
+ if failure == "stage":
2306
+ mx_client.stage_weight.side_effect = RuntimeError("stage failed")
2307
+ else:
2308
+ mx_client.apply_weight.side_effect = RuntimeError("apply failed")
2309
+ if failure == "release":
2310
+ staged.release.side_effect = RuntimeError("release failed")
2311
+ engine.start_weight_update()
2312
+ with pytest.raises(RuntimeError, match="stage failed|apply failed"):
2313
+ engine.update_weights({"version_id": "version-a"})
2314
+ assert staged.release.call_count == (0 if failure == "stage" else 1)
2315
+ torch.accelerator.synchronize.assert_not_called()
2316
+ engine.shutdown()
2317
+ mx_client.close.assert_called_once_with()
2318
+
2319
+
2320
+def test_plugin_does_not_replace_native_backend(mx_client):
2321
+ from modelexpress.engines.vllm.registration import (
2322
+ register_plugin_weight_transfer_engine,
2323
+ )
2324
+ from modelexpress_rl.inference.engines.vllm.weight_transfer_engine import (
2325
+ ModelExpressWeightTransferEngine,
2326
+ )
2327
+
2328
+ native_loader = WeightTransferEngineFactory._registry["modelexpress"]
2329
+ register_plugin_weight_transfer_engine()
2330
+ assert WeightTransferEngineFactory._registry["modelexpress"] is native_loader
2331
+ assert type(_make_modelexpress_engine()) is ModelExpressWeightTransferEngine
vllm/distributed/weight_transfer/factory.py
+7
@@ -250,6 +250,13 @@ WeightTransferEngineFactory.register_engine(
250
)
251
252
253
+WeightTransferEngineFactory.register_engine(
254
+ "modelexpress",
255
+ "vllm.distributed.weight_transfer.modelexpress_engine",
256
+ "ModelExpressWeightTransferEngine",
257
+)
258
+
259
+
260
# Trainer-side engines, parallel to the worker registry above.
261
WeightTransferTrainerFactory.register_engine(
262
"nccl",
vllm/distributed/weight_transfer/modelexpress_engine.py
new
+41
@@ -0,0 +1,41 @@
1
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+
5
+"""Re-export the weight-transfer backend maintained by ModelExpress.
6
+
7
+Install ModelExpress from the main branch's modelexpress_client/python
8
+subdirectory.
9
+"""
10
+
11
+try:
12
+ from modelexpress import configure_vllm_logging
13
+ from modelexpress_rl.inference.engines.vllm.weight_transfer_engine import (
14
+ ModelExpressWeightTransferEngine,
15
+ ModelExpressWeightTransferInitInfo,
16
+ ModelExpressWeightTransferUpdateInfo,
17
+ )
18
+except ModuleNotFoundError as exc:
19
+ if exc.name not in {
20
+ "modelexpress",
21
+ "modelexpress_rl",
22
+ "modelexpress_rl.inference",
23
+ "modelexpress_rl.inference.engines",
24
+ "modelexpress_rl.inference.engines.vllm",
25
+ "modelexpress_rl.inference.engines.vllm.weight_transfer_engine",
26
+ }:
27
+ raise
28
+ raise ImportError(
29
+ "The 'modelexpress' weight transfer backend requires ModelExpress. "
30
+ "Install it with `uv pip install "
31
+ "'git+https://github.com/ai-dynamo/modelexpress@main"
32
+ "#subdirectory=modelexpress_client/python'`."
33
+ ) from exc
34
+
35
+configure_vllm_logging()
36
+
37
+__all__ = [
38
+ "ModelExpressWeightTransferEngine",
39
+ "ModelExpressWeightTransferInitInfo",
40
+ "ModelExpressWeightTransferUpdateInfo",
41
+]