diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 4c7084f37112..28d5a31f1cd6 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -151,6 +151,8 @@ def _discard_orphan_result(fut: "asyncio.Future[dict]") -> None: "weight_version", } ) +# An object sentinel cannot collide with a caller-supplied version. +_WEIGHT_VERSION_UNDECLARED: Final = object() def build_prompt_tokens_details( @@ -1185,7 +1187,7 @@ def __init__( # to prevent both bypassing the check before either inserts (atomicity). self._lora_capacity_guard = asyncio.Lock() self._paused: bool = False - self._weight_version: str = "initial" + self._weight_version: Any = _WEIGHT_VERSION_UNDECLARED embedding_loader = self.init_embedding_loader(config, encode_worker_client) @@ -2001,7 +2003,7 @@ async def abort_request(self, body: dict) -> dict: return {"status": "error", "message": str(e)} async def get_weight_version(self, body: dict) -> dict: - """Return the current weight version tag.""" + """Report the worker's current declared weight-version state.""" if body is None: body = {} elif not isinstance(body, dict): @@ -2009,7 +2011,29 @@ async def get_weight_version(self, body: dict) -> dict: "status": "error", "message": "request body must be a JSON object", } - return {"status": "ok", "version": getattr(self, "_weight_version", "initial")} + version = self._weight_version + is_declared = version is not _WEIGHT_VERSION_UNDECLARED + return { + "status": "ok", + "version": version if is_declared else None, + "version_declared": is_declared, + } + + async def set_weight_version(self, body: dict) -> dict: + """Validate and declare a weight version for subsequent requests.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + if "weight_version" not in body: + return {"status": "error", "message": "Missing 'weight_version' in body"} + version = body["weight_version"] + self._weight_version = version + logger.info("[RL] Weight version declared") + return {"status": "ok", "version": version} async def update_weights_from_disk(self, body: dict) -> dict: """Load weights from a shared filesystem checkpoint.""" @@ -2048,11 +2072,16 @@ async def update_weights_from_disk(self, body: dict) -> dict: # weights is now stale and must not be reused. Invalidate it # while still holding _pause_lock (generation is paused). await self.engine_client.reset_prefix_cache() - self._weight_version = version + if "weight_version" in body: + self._weight_version = version logger.info( f"[RL] Weights loaded from {path} (version={version}, rpc={rpc})" ) - return {"status": "ok", "version": version} + return { + "status": "ok", + "version": version, + "version_declared": "weight_version" in body, + } except EngineDeadError as e: self._shutdown_on_engine_dead(e) except Exception as e: @@ -2111,12 +2140,17 @@ async def update_weights_from_distributed(self, body: dict) -> dict: # Weights changed: stale prefix/KV cache must be invalidated # before resume so it is not reused under the new weights. await self.engine_client.reset_prefix_cache() - self._weight_version = version + if "weight_version" in body: + self._weight_version = version logger.info( f"[RL] Weights received via distributed " f"(version={version}, rpc={rpc})" ) - return {"status": "ok", "version": version} + return { + "status": "ok", + "version": version, + "version_declared": "weight_version" in body, + } except EngineDeadError as e: self._shutdown_on_engine_dead(e) except Exception as e: diff --git a/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py b/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py index ec415f1add45..6966155b4f3e 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py +++ b/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py @@ -98,6 +98,7 @@ def _make_handler( # BaseWorkerHandler.__init__ is bypassed above; the decode generate path # registers per-request deferred-abort guards here. handler._deferred_aborts = {} + handler._weight_version = mod._WEIGHT_VERSION_UNDECLARED return handler @@ -1924,6 +1925,8 @@ async def test_admin_rejects_non_dict_body(self): handler.resume_generation, handler.flush_cache, handler.abort_request, + handler.get_weight_version, + handler.set_weight_version, ): resp = await fn(body) assert resp["status"] == "error", (fn.__name__, body, resp) @@ -1948,7 +1951,7 @@ async def test_distributed_update_can_match_async_rl_semantics(self): } ) - assert resp == {"status": "ok", "version": 7} + assert resp == {"status": "ok", "version": 7, "version_declared": True} handler.engine_client.collective_rpc.assert_awaited_once_with( "update_weights", kwargs={"update_info": {"names": ["weight"]}}, @@ -1996,6 +1999,118 @@ async def test_distributed_update_preserves_safe_defaults(self): ) handler.engine_client.reset_prefix_cache.assert_awaited_once_with() + @staticmethod + def _make_rl_handler(): + handler = _make_handler() + handler._pause_lock = asyncio.Lock() + handler._paused = False + handler.engine_client = MagicMock() + handler.engine_client.collective_rpc = AsyncMock() + handler.engine_client.reset_prefix_cache = AsyncMock() + return handler + + @pytest.mark.asyncio + async def test_get_weight_version_reports_undeclared_before_any_update(self): + config = _make_config(enable_multimodal=False) + config.custom_encoder_class = None + with patch.object(mod, "VllmEngineMonitor"): + handler = mod.DecodeWorkerHandler( + runtime=MagicMock(), + config=config, + engine=MagicMock(), + default_sampling_params={}, + ) + + resp = await handler.get_weight_version({}) + + assert resp["status"] == "ok" + assert resp["version_declared"] is False + assert resp["version"] is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("version", [None, 7, "policy-43"]) + async def test_set_weight_version_declares_without_touching_the_engine( + self, version + ): + handler = self._make_rl_handler() + + resp = await handler.set_weight_version({"weight_version": version}) + + assert resp == {"status": "ok", "version": version} + assert await handler.get_weight_version({}) == { + "status": "ok", + "version": version, + "version_declared": True, + } + handler.engine_client.collective_rpc.assert_not_awaited() + handler.engine_client.reset_prefix_cache.assert_not_awaited() + + @pytest.mark.asyncio + async def test_set_weight_version_rejects_a_missing_version(self): + handler = self._make_rl_handler() + + resp = await handler.set_weight_version({}) + + assert resp["status"] == "error" + assert (await handler.get_weight_version({}))["version_declared"] is False + + @pytest.mark.asyncio + @pytest.mark.parametrize("version", [None, "initial", "unknown"]) + @pytest.mark.parametrize( + ("route", "body"), + [ + ("update_weights_from_disk", {"model_path": "/models/checkpoint-42"}), + ("update_weights_from_distributed", {"engine_rpc": "update_weights"}), + ], + ) + async def test_explicit_version_is_declared(self, route, body, version): + handler = self._make_rl_handler() + handler._paused = True + + response = await getattr(handler, route)({**body, "weight_version": version}) + + assert response == { + "status": "ok", + "version": version, + "version_declared": True, + } + assert await handler.get_weight_version({}) == { + "status": "ok", + "version": version, + "version_declared": True, + } + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "version", + [mod._WEIGHT_VERSION_UNDECLARED, None, "policy-42"], + ids=["undeclared", "declared-null", "declared-string"], + ) + @pytest.mark.parametrize( + ("route", "body"), + [ + ("update_weights_from_disk", {"model_path": "/models/checkpoint-43"}), + ("update_weights_from_distributed", {"engine_rpc": "update_weights"}), + ], + ) + async def test_update_without_a_version_preserves_declaration( + self, route, body, version + ): + handler = self._make_rl_handler() + handler._paused = True + if version is not mod._WEIGHT_VERSION_UNDECLARED: + await handler.set_weight_version({"weight_version": version}) + previous = await handler.get_weight_version({}) + + response = await getattr(handler, route)(body) + + assert response == { + "status": "ok", + "version": "unknown", + "version_declared": False, + } + assert await handler.get_weight_version({}) == previous + @pytest.mark.asyncio async def test_init_weights_update_group_succeeds_within_timeout(self): handler = _make_handler() diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index 8d5f1eaec356..26d53854712d 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/components/src/dynamo/vllm/worker_factory.py @@ -1826,6 +1826,7 @@ def register_engine_routes( "init_weights_update_group": handler.init_weights_update_group, "destroy_weights_update_group": handler.destroy_weights_update_group, "get_weight_version": handler.get_weight_version, + "set_weight_version": handler.set_weight_version, } if lora_enabled: diff --git a/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md b/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md index 2b7f6839752e..5cc70320eb95 100644 --- a/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md +++ b/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md @@ -122,6 +122,7 @@ A protocol version `1` response has this shape. Optional fields such as `model`, "liveness_probe", "pause_generation", "resume_generation", + "set_weight_version", "update_weights_from_disk" ] } @@ -167,12 +168,51 @@ The backends do not share one administration schema. Use the exact route returne | Stop and resume work | Python `pause_generation` / `resume_generation`; native-sidecar control routes reflect the installed vLLM RL API | `release_memory_occupation` unregisters and drains the worker; `resume_memory_occupation` restores it | | Clear KV state | Python `flush_cache`; native-sidecar behavior follows the advertised lifecycle route and vLLM configuration | `clear_kv_blocks` is a Dynamo worker request-plane endpoint, not a built-in `/engine/control/*` route, and rejects active requests | | Apply weights | Python disk or distributed update routes and group lifecycle; native sidecar init/start/update/finish routes when weight transfer is enabled | Built-in disk, tensor, distributed, or IPC update routes using the installed SGLang request schemas | -| Weight version | Python `get_weight_version` reads caller-supplied update metadata; native sidecar exposes get/update controls | `update_weight_version` changes metadata and can abort requests; it does not replace tensors | +| Weight version | Python `get_weight_version` reads the last declared version; `set_weight_version` declares metadata without loading weights; native sidecar exposes get/update controls | `update_weight_version` changes metadata and can abort requests; it does not replace tensors | | Custom controls | Python integrations can register and advertise trusted engine routes | Allowlist methods with `--engine-route` or `DYN_SGLANG_ENGINE_ROUTES` using path[=method][:engine|tm] | | Success body | Python RL routes commonly return a `status` field; native-sidecar routes follow vLLM's RL schemas | Built-in update routes commonly return `success` and `message`; check both HTTP status and body | Even two vLLM deployments can expose different route families because the Python worker and native sidecar adapt different backend control APIs. Never prepend, remove, or rename route segments returned in `routes`. +### Read and Declare the Weight Version + +A Python vLLM worker reports the last weight version declared to it through `get_weight_version`: + +```bash +curl http://10.0.0.12:8081/engine/get_weight_version \ + -H 'Content-Type: application/json' \ + -d '{}' +``` + +```json +{"status": "ok", "version": "42", "version_declared": true} +``` + +A worker tracks only the versions declared to it. `version_declared` is `false`, with `version` set to `null`, until something declares one, either through a `/engine/` weight-update route that carries `weight_version` or through `set_weight_version`. Branch on `version_declared` rather than comparing `version` against a placeholder string: any string, including `"initial"`, is a legal version tag that a caller can declare. + +For each Python vLLM worker, check that `routes` from `GET /v1/rl/workers` includes `set_weight_version` before relying on `version_declared`. Older workers omit both that route and the response field; treat a missing `version_declared` as unsupported declaration tracking, not as `false`, because their `"initial"` version cannot distinguish an undeclared worker from an explicit declaration. + +Version declarations accept any JSON value, including `null`. An explicit `{"weight_version": null}` is a declaration: `get_weight_version` then returns `"version": null` with `"version_declared": true`. + +A successful weight-update reply includes `version_declared` to indicate whether that update declared a version. If the request omits `weight_version`, the reply contains `"version": "unknown"` and `"version_declared": false`, and the previous declaration remains unchanged. An explicit `{"weight_version": "unknown"}` returns the same version with `"version_declared": true`. In an update reply, this flag describes the update; in `get_weight_version`, it describes the stored declaration. A worker with no previous declaration remains undeclared after an update that omits the version. Pass `weight_version` on every update whose version you want the worker to report. + +> [!WARNING] +> A worker reports the last version declared to it, not the weights loaded in its GPU memory. Dynamo observes only the weight updates that traverse its own `/engine/` routes. Loading weights by another path, such as calling `collective_rpc` on the engine object directly, leaves the reported version stale unless the loader declares the new version. + +When an RL framework loads weights outside Dynamo, declare the resulting version so the worker reports it. `set_weight_version` records the version without loading weights: + +```bash +curl http://10.0.0.12:8081/engine/set_weight_version \ + -H 'Content-Type: application/json' \ + -d '{"weight_version": "43"}' +``` + +```json +{"status": "ok", "version": "43"} +``` + +The route requires `weight_version` in the body and returns `{"status": "error"}` when it is missing. It neither pauses generation nor invalidates the prefix cache, so a caller that changed the weights must handle both itself. + ## Framework Compatibility | Framework | Status | Current Dynamo path | Key boundary |