diff --git a/components/src/dynamo/common/rl/__init__.py b/components/src/dynamo/common/rl/__init__.py new file mode 100644 index 000000000000..4c8af36db383 --- /dev/null +++ b/components/src/dynamo/common/rl/__init__.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared RL admin utilities.""" + +from .admin import ( + RLAdminValidationError, + RLRouteHandler, + RLRouteRegistry, + env_bool, + first_endpoint_response, + register_rl_routes, + require_lora_load_request, + require_lora_unload_request, +) + +__all__ = [ + "RLAdminValidationError", + "RLRouteHandler", + "RLRouteRegistry", + "env_bool", + "first_endpoint_response", + "register_rl_routes", + "require_lora_load_request", + "require_lora_unload_request", +] diff --git a/components/src/dynamo/common/rl/admin.py b/components/src/dynamo/common/rl/admin.py new file mode 100644 index 000000000000..e99533fb0358 --- /dev/null +++ b/components/src/dynamo/common/rl/admin.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared helpers for RL admin request-plane endpoints.""" + +from __future__ import annotations + +import logging +import os +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from typing import Any + +logger = logging.getLogger(__name__) + +TRUE_ENV_VALUES = {"1", "true", "yes", "on"} + +RLRouteHandler = Callable[[dict[str, Any]], Awaitable[dict[str, Any] | None]] +EndpointGenerator = Callable[[dict[str, Any]], AsyncIterator[dict[str, Any] | None]] + + +class RLAdminValidationError(ValueError): + """Validation error whose message can be returned directly to RL clients.""" + + +def env_bool(name: str, default: bool = False) -> bool: + """Parse a boolean environment variable using Dynamo's common true values.""" + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() in TRUE_ENV_VALUES + + +async def first_endpoint_response( + endpoint_handler: EndpointGenerator, + body: dict[str, Any], +) -> dict[str, Any]: + """Return the first response from an async-generator endpoint handler. + + The generator is explicitly closed before returning so handlers that hold + resources across their yield (e.g. load_lora/unload_lora holding a per-LoRA + lock) release them promptly rather than waiting for garbage collection. + """ + gen = endpoint_handler(body) + try: + async for response in gen: + return response or {"status": "ok"} + return {"status": "ok"} + finally: + aclose = getattr(gen, "aclose", None) + if aclose is not None: + await aclose() + + +def require_lora_load_request(request: Mapping[str, Any] | None) -> tuple[str, str]: + """Validate the shared URI-based LoRA load request shape.""" + if request is None or not isinstance(request, Mapping): + raise RLAdminValidationError( + "Request is required with 'lora_name' and 'source.uri'" + ) + + lora_name = request.get("lora_name") + if not isinstance(lora_name, str) or not lora_name: + raise RLAdminValidationError("'lora_name' is required and must be a string") + + source = request.get("source") + if not source or not isinstance(source, Mapping): + raise RLAdminValidationError("'source' object is required in request") + + lora_uri = source.get("uri") + if not isinstance(lora_uri, str) or not lora_uri: + raise RLAdminValidationError("'source.uri' is required and must be a string") + + return lora_name, lora_uri + + +def require_lora_unload_request(request: Mapping[str, Any] | None) -> str: + """Validate the shared LoRA unload request shape.""" + if request is None or not isinstance(request, Mapping): + raise RLAdminValidationError("Request is required with 'lora_name' field") + + lora_name = request.get("lora_name") + if not isinstance(lora_name, str) or not lora_name: + raise RLAdminValidationError("'lora_name' is required and must be a string") + + return lora_name + + +class RLRouteRegistry: + """Registry for worker RL admin route descriptors.""" + + def __init__( + self, + runtime: Any, + *, + logger_: logging.Logger | None = None, + ) -> None: + self._runtime = runtime + self._logger = logger_ or logger + self.routes: dict[str, RLRouteHandler] = {} + + def add_route(self, name: str, handler: RLRouteHandler) -> None: + self.routes[name] = handler + + def add_routes(self, routes: Mapping[str, RLRouteHandler]) -> None: + for name, handler in routes.items(): + self.add_route(name, handler) + + def describe(self) -> dict[str, Any]: + response: dict[str, Any] = { + "status": "ok", + "routes": sorted(self.routes), + } + + system_url_fn = getattr(self._runtime, "system_status_server_url", None) + if callable(system_url_fn): + system_url = system_url_fn() + if system_url: + response["system_url"] = system_url + + return response + + async def dispatch( + self, request: Mapping[str, Any] | None = None + ) -> dict[str, Any]: + if request is None or not isinstance(request, Mapping): + return {"status": "error", "message": "rl_dispatch: request required"} + + method = request.get("method") + + if not isinstance(method, str) or not method: + return {"status": "error", "message": "rl_dispatch: missing 'method' (str)"} + + if method != "routes": + return { + "status": "error", + "method": method, + "message": "rl request-plane endpoint only supports method='routes'", + } + + if "kwargs" in request and not isinstance(request.get("kwargs"), Mapping): + return { + "status": "error", + "method": method, + "message": "rl_dispatch: 'kwargs' must be an object", + } + + return self.describe() + + async def dispatch_stream( + self, request: Mapping[str, Any] | None = None + ) -> AsyncIterator[dict[str, Any]]: + yield await self.dispatch(request) + + +def register_rl_routes( + runtime: Any, + registry: RLRouteRegistry, + routes: Mapping[str, RLRouteHandler], + *, + enable_dispatch: bool, +) -> None: + """Register worker system routes and optionally expose route descriptors.""" + for name, handler in routes.items(): + runtime.register_engine_route(name, handler) + if enable_dispatch: + registry.add_route(name, handler) diff --git a/components/src/dynamo/common/tests/test_rl_admin.py b/components/src/dynamo/common/tests/test_rl_admin.py new file mode 100644 index 000000000000..176005cd2de5 --- /dev/null +++ b/components/src/dynamo/common/tests/test_rl_admin.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio + +import pytest + +from dynamo.common.rl import ( + RLAdminValidationError, + RLRouteRegistry, + first_endpoint_response, + register_rl_routes, + require_lora_load_request, + require_lora_unload_request, +) + +pytestmark = [pytest.mark.pre_merge, pytest.mark.unit, pytest.mark.gpu_0] + + +class _Runtime: + def __init__(self, system_url: str | None = None) -> None: + self.system_url = system_url + self.registered: list[tuple[str, object]] = [] + + def system_status_server_url(self) -> str | None: + return self.system_url + + def register_engine_route(self, name: str, handler: object) -> None: + self.registered.append((name, handler)) + + +def test_route_registry_describes_routes() -> None: + runtime = _Runtime("http://worker:8081") + registry = RLRouteRegistry(runtime) + + async def ping(body: dict) -> dict: + return {"status": "ok", "body": body} + + registry.add_route("ping", ping) + + routes = asyncio.run(registry.dispatch({"method": "routes"})) + assert routes == { + "status": "ok", + "routes": ["ping"], + "system_url": "http://worker:8081", + } + + routes_with_kwargs = asyncio.run( + registry.dispatch({"method": "routes", "kwargs": {}}) + ) + assert routes_with_kwargs == routes + + +def test_route_registry_rejects_request_plane_admin_execution() -> None: + registry = RLRouteRegistry(_Runtime()) + + response = asyncio.run(registry.dispatch({"method": "missing"})) + + assert response["status"] == "error" + assert response["method"] == "missing" + assert ( + response["message"] == "rl request-plane endpoint only supports method='routes'" + ) + + +def test_route_registry_rejects_non_object_kwargs_for_routes() -> None: + registry = RLRouteRegistry(_Runtime()) + + response = asyncio.run(registry.dispatch({"method": "routes", "kwargs": []})) + + assert response == { + "status": "error", + "method": "routes", + "message": "rl_dispatch: 'kwargs' must be an object", + } + + +def test_register_rl_routes_always_registers_engine_route() -> None: + runtime = _Runtime() + registry = RLRouteRegistry(runtime) + + async def ping(body: dict) -> dict: + return {"status": "ok", "body": body} + + register_rl_routes(runtime, registry, {"ping": ping}, enable_dispatch=False) + + assert runtime.registered == [("ping", ping)] + assert registry.routes == {} + + register_rl_routes(runtime, registry, {"ping": ping}, enable_dispatch=True) + + assert registry.routes == {"ping": ping} + + +def test_first_endpoint_response_returns_first_chunk() -> None: + async def endpoint(_body: dict): + yield {"status": "ok", "value": 1} + yield {"status": "ok", "value": 2} + + response = asyncio.run(first_endpoint_response(endpoint, {})) + + assert response == {"status": "ok", "value": 1} + + +def test_lora_load_request_validation() -> None: + assert require_lora_load_request( + {"lora_name": "adapter", "source": {"uri": "file:///tmp/adapter"}} + ) == ("adapter", "file:///tmp/adapter") + + try: + require_lora_load_request({"lora_name": "adapter"}) + except RLAdminValidationError as exc: + assert str(exc) == "'source' object is required in request" + else: + raise AssertionError("expected validation error") + + +def test_lora_unload_request_validation() -> None: + assert require_lora_unload_request({"lora_name": "adapter"}) == "adapter" + + try: + require_lora_unload_request({}) + except RLAdminValidationError as exc: + assert str(exc) == "'lora_name' is required and must be a string" + else: + raise AssertionError("expected validation error") + + # Non-string scalars must be rejected, not str()-coerced. + for bad in ([], {}, 123, ["adapter"]): + try: + require_lora_unload_request({"lora_name": bad}) + except RLAdminValidationError: + pass + else: + raise AssertionError(f"expected validation error for lora_name={bad!r}") + + +def test_lora_load_request_rejects_non_string_fields() -> None: + # lora_name / source.uri must be strings (no str() coercion of lists/dicts). + for req in ( + {"lora_name": ["a"], "source": {"uri": "file:///x"}}, + {"lora_name": "a", "source": {"uri": {}}}, + {"lora_name": "a", "source": {"uri": ["file:///x"]}}, + ): + try: + require_lora_load_request(req) + except RLAdminValidationError: + pass + else: + raise AssertionError(f"expected validation error for {req!r}") diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index 170306c44f73..33c77685bb2a 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -17,7 +17,17 @@ from abc import ABC, abstractmethod from collections import deque from contextlib import asynccontextmanager -from typing import Any, AsyncIterator, Dict, Final, Generic, Iterator, Optional, TypeVar +from typing import ( + Any, + AsyncIterator, + Dict, + Final, + Generic, + Iterator, + NoReturn, + Optional, + TypeVar, +) import torch from vllm.config import ModelConfig, VllmConfig @@ -54,6 +64,13 @@ MmKwargsTransferMetadata, ) from dynamo.common.multimodal.video_loader import VideoLoader +from dynamo.common.rl import ( + RLAdminValidationError, + RLRouteRegistry, + env_bool, + require_lora_load_request, + require_lora_unload_request, +) from dynamo.common.utils import nvtx_utils as _nvtx from dynamo.common.utils.engine_response import normalize_finish_reason from dynamo.common.utils.input_params import InputParamManager @@ -120,15 +137,28 @@ class _DeferredAbort: abort. """ - def __init__(self, engine_client: Any, request_id: str): + def __init__( + self, + engine_client: Any, + request_id: str, + on_engine_dead: Optional[Any] = None, + ): self._engine_client = engine_client self._request_id = request_id + # Escalation hook invoked if the (possibly deferred/background) engine + # abort hits EngineDeadError, so engine death shuts the runtime down even + # when the abort runs outside the request's own error handling (the + # disconnect-monitor or deferred-after-first-token path). + self._on_engine_dead = on_engine_dead self._first_token_received = False self._first_token_event = asyncio.Event() # Strong reference to the deferred-abort background task so it is not # garbage collected mid-execution (asyncio.create_task only holds a # weak reference via the event loop). self._abort_task: Optional[asyncio.Task] = None + # Exception the real engine abort raised (if it has run), so the admin + # abort_request route can report failure instead of a false "ok". + self._abort_exc: Optional[BaseException] = None def signal_first_token(self) -> None: """Called when the first engine output for the request is received.""" @@ -153,8 +183,20 @@ async def abort(self) -> None: f"{self._request_id}, spawning background task" ) self._abort_task = asyncio.create_task(self._wait_and_abort()) + # Only block on completion when the abort runs immediately (post first + # token). A pre-first-token deferred abort fires in the background when + # the first token arrives; awaiting it here would hang the caller (admin + # route or disconnect monitor) until — or unless — generation produces + # output. _abort_task keeps the background task alive; close() reaps it. + if not self._first_token_received: + return try: - await self._abort_task + # shield() so that if the caller (e.g. a cancelled system-route + # request or a disconnected client) is cancelled while awaiting, the + # cancellation is NOT propagated into the abort task — it really does + # continue in the background. A bare `await self._abort_task` would + # cancel the task too, silently dropping the abort. + await asyncio.shield(self._abort_task) except asyncio.CancelledError: logger.debug( f"Deferred abort: shielded from cancellation for request " @@ -167,10 +209,17 @@ async def _run_abort(self) -> None: await self._engine_client.abort(self._request_id) logger.debug(f"Aborted Request ID: {self._request_id}") except Exception as e: + # Record so abort_request can report the failure rather than a false + # success. Also escalate engine death here, since a deferred or + # disconnect-monitor abort runs in the background with no caller + # awaiting the result to handle EngineDeadError. + self._abort_exc = e logger.warning( f"Deferred abort: engine abort raised for request " f"{self._request_id}: {e}" ) + if isinstance(e, EngineDeadError) and self._on_engine_dead is not None: + self._on_engine_dead(e) async def _wait_and_abort(self) -> None: """Background task: wait for first token, then abort.""" @@ -202,7 +251,10 @@ async def close(self) -> None: self._abort_task.cancel() try: - await self._abort_task + # shield so that if cleanup is awaiting a real post-first-token abort + # and the caller is cancelled, the abort still completes (the + # pre-first-token path was cancelled just above and resolves here). + await asyncio.shield(self._abort_task) except asyncio.CancelledError: pass except Exception as e: @@ -214,7 +266,11 @@ async def close(self) -> None: @asynccontextmanager async def _deferred_abort_guard( - engine_client: Any, request_id: str, is_decode_only: bool + engine_client: Any, + request_id: str, + is_decode_only: bool, + registry: Optional[dict[str, "_DeferredAbort"]] = None, + on_engine_dead: Optional[Any] = None, ) -> AsyncIterator[Optional[_DeferredAbort]]: """Own the _DeferredAbort lifecycle for a single request. @@ -223,13 +279,32 @@ async def _deferred_abort_guard( when generation finishes without producing output (case 1b). close() is specifically designed not to call engine_client.abort() in the unsafe pre-first-token window. + + When `registry` is provided, the guard registers itself under `request_id` + for the request's lifetime so out-of-band callers (the admin abort_request + route) can route their abort through this same deferred path instead of + calling engine_client.abort() directly in the unsafe window. """ - guard = _DeferredAbort(engine_client, request_id) if is_decode_only else None + guard = ( + _DeferredAbort(engine_client, request_id, on_engine_dead) + if is_decode_only + else None + ) + if guard is not None and registry is not None: + registry[request_id] = guard try: yield guard finally: if guard is not None: - await guard.close() + # Keep the guard registered until close() finishes: close() may + # await a deferred abort, and an out-of-band admin abort_request + # during that window must still find the guard and route through + # the deferred path instead of taking the unsafe direct abort. + try: + await guard.close() + finally: + if registry is not None: + registry.pop(request_id, None) class VllmEnginePauseController: @@ -920,6 +995,8 @@ def __init__( self._lora_load_locks: dict[str, asyncio.Lock] = {} # Guard lock-map access in case handlers are invoked from multiple threads. self._lora_load_locks_guard = threading.Lock() + self._paused: bool = False + self._weight_version: str = "initial" self.image_loader = ImageLoader( enable_frontend_decoding=enable_frontend_decoding @@ -937,6 +1014,11 @@ def __init__( self.dp_range = get_dp_range_for_worker(self.engine_client.vllm_config) self._pause_controller = VllmEnginePauseController(self.engine_client) self._pause_lock = asyncio.Lock() + # Maps request_id -> _DeferredAbort for in-flight decode-only requests so + # admin abort_request can route through the deferred-abort path instead + # of calling engine_client.abort() during the unsafe pre-first-token + # NIXL-KV-transfer window. + self._deferred_aborts: dict[str, _DeferredAbort] = {} self._mm_kwargs_receiver: MmKwargsNixlReceiver | None = None # Some models (Kimi-K2.5) declare their image modality as @@ -974,6 +1056,15 @@ def __init__( # Store shutdown event for graceful shutdown monitoring self.shutdown_event = shutdown_event + # Request-plane RL method map served by rl_dispatch on + # dyn://..rl when --enable-rl / DYN_ENABLE_RL is set. + self.rl_route_registry = RLRouteRegistry(self.runtime, logger_=logger) + + def _shutdown_on_engine_dead(self, e: EngineDeadError) -> NoReturn: + logger.error(f"vLLM EngineDeadError: {e}") + logger.warning("Initiating Dynamo Runtime shutdown.") + self.runtime.shutdown() + os._exit(1) def init_embedding_loader( self, config: Config, encode_worker_client: Optional[Client] = None @@ -1037,7 +1128,13 @@ async def sleep(self, body: dict) -> dict: 2. Abort and drain in-flight requests 3. Sleep engine - safe once generation has stopped """ - body = body or {} + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } level = body.get("level", 1) async with self._pause_lock: if self._pause_controller.is_paused: @@ -1108,7 +1205,13 @@ async def scale_elastic_ep(self, body: dict) -> dict: already reserved by the pod, then hot-swap the expert routing table. No pod restart is needed. """ - body = body or {} + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } new_dp_size = body.get("new_data_parallel_size") if new_dp_size is None: return { @@ -1208,7 +1311,13 @@ async def wake_up(self, body: dict) -> dict: 1. Wake engine - restore GPU memory 2. Re-register endpoint instance - allow frontend to route requests here again """ - body = body or {} + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } tags = body.get("tags") async with self._pause_lock: needs_recovery = self._pause_controller.needs_resume_recovery @@ -1261,6 +1370,323 @@ async def stop_profile(self, body: dict) -> dict: logger.error(f"Failed to stop profiling: {e}") return {"status": "error", "message": str(e)} + async def rl_dispatch(self, request=None): + """Request-plane dispatcher for the worker's ``rl`` endpoint.""" + async for response in self.rl_route_registry.dispatch_stream(request): + yield response + + async def liveness_probe(self, body: dict) -> dict: + """Engine event-loop liveness probe.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + try: + if hasattr(self.engine_client, "check_health"): + await self.engine_client.check_health() + else: + await self.engine_client.collective_rpc("liveness_probe", kwargs={}) + return {"status": "ok", "alive": True} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.warning(f"[RL] liveness_probe failed: {e}") + return {"status": "error", "alive": False, "message": str(e)} + + async def pause_generation(self, body: dict) -> dict: + """Pause generation before a weight update.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + mode = body.get("mode", "keep") + clear_cache = bool(body.get("clear_cache", False)) + if mode not in ("keep", "wait", "abort"): + return { + "status": "error", + "message": f"Invalid mode '{mode}'; expected keep|wait|abort", + } + async with self._pause_lock: + try: + try: + await self.engine_client.pause_generation( + mode=mode, clear_cache=clear_cache + ) + except TypeError: + await self.engine_client.pause_generation() + if clear_cache: + await self.engine_client.reset_prefix_cache() + self._paused = True + logger.info( + f"[RL] Engine paused (mode={mode}, clear_cache={clear_cache})" + ) + return { + "status": "ok", + "message": "Engine paused", + "mode": mode, + "clear_cache": clear_cache, + } + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] Failed to pause: {e}") + return {"status": "error", "message": str(e)} + + async def resume_generation(self, body: dict) -> dict: + """Resume generation after a weight update.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + # Serialize pause / resume / weight-update so a concurrent resume cannot + # re-enable generation while an update's collective_rpc is still + # mutating weights (dynamo-ops). + async with self._pause_lock: + try: + await self.engine_client.resume_generation() + self._paused = False + logger.info("[RL] Engine resumed") + return {"status": "ok", "message": "Engine resumed"} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] Failed to resume: {e}") + return {"status": "error", "message": str(e)} + + async def flush_cache(self, body: dict) -> dict: + """Invalidate prefix / KV cache.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + # Serialize under _pause_lock so a flush cannot race with a locked + # weight-update / pause / resume mutating engine cache state. + async with self._pause_lock: + try: + await self.engine_client.reset_prefix_cache() + logger.debug("[RL] Prefix cache flushed") + return {"status": "ok", "message": "Cache flushed"} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] Failed to flush cache: {e}") + return {"status": "error", "message": str(e)} + + async def abort_request(self, body: dict) -> dict: + """Abort a single in-flight request by request_id.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + request_id = body.get("request_id") + if not request_id: + return {"status": "error", "message": "Missing 'request_id' in body"} + try: + guard = self._deferred_aborts.get(request_id) + if guard is not None: + # Route through the per-request deferred-abort guard so that in + # disaggregated decode mode the real engine abort is deferred + # until the first token, never firing during an in-flight NIXL + # KV transfer (which can crash EngineCore). + await guard.abort() + # If the abort already ran (post-first-token) and failed, report + # it instead of a false "ok"; escalate engine death like the + # direct path does. (Pre-first-token aborts are queued and have + # no result yet, so they correctly report accepted/ok.) + abort_exc = guard._abort_exc + if abort_exc is not None: + if isinstance(abort_exc, EngineDeadError): + self._shutdown_on_engine_dead(abort_exc) + return { + "status": "error", + "request_id": request_id, + "message": f"abort failed: {abort_exc}", + } + else: + await self.engine_client.abort(request_id) + logger.debug(f"[RL] Aborted request {request_id}") + return {"status": "ok", "request_id": request_id} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] Failed to abort request {request_id}: {e}") + return {"status": "error", "message": str(e)} + + async def get_weight_version(self, body: dict) -> dict: + """Return the current weight version tag.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + return {"status": "ok", "version": getattr(self, "_weight_version", "initial")} + + async def update_weights_from_disk(self, body: dict) -> dict: + """Load weights from a shared filesystem checkpoint.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + # Hold _pause_lock across the paused-state check and the weight RPC so a + # concurrent resume cannot re-enable generation mid-update (dynamo-ops). + async with self._pause_lock: + if not getattr(self, "_paused", False): + return { + "status": "error", + "message": ( + "Worker must be paused via pause_generation() before " + "updating weights. Call pause_generation() first, then " + "update, then resume_generation()." + ), + } + path = body.get("model_path") + if not path: + return {"status": "error", "message": "Missing 'model_path' in body"} + version = body.get("weight_version", "unknown") + rpc = body.get("engine_rpc", "reload_weights") + kwargs = ( + {"weights_path": path} + if rpc == "reload_weights" + else {"weight_path": path} + ) + try: + await self.engine_client.collective_rpc(rpc, kwargs=kwargs) + # Weights changed: any prefix/KV cache computed under the old + # 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 + logger.info( + f"[RL] Weights loaded from {path} (version={version}, rpc={rpc})" + ) + return {"status": "ok", "version": version} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] update_weights_from_disk failed: {e}") + return {"status": "error", "message": str(e)} + + async def update_weights_from_distributed(self, body: dict) -> dict: + """Receive weights via a distributed transport such as NCCL.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + async with self._pause_lock: + if not getattr(self, "_paused", False): + return { + "status": "error", + "message": ( + "Worker must be paused via pause_generation() before " + "updating weights. Call pause_generation() first, then " + "update, then resume_generation()." + ), + } + version = body.get("weight_version", "unknown") + rpc = body.get("engine_rpc", "update_weights_from_path") + rpc_kwargs = { + k: v + for k, v in body.items() + if k not in ("engine_rpc", "weight_version") + } + try: + await self.engine_client.collective_rpc(rpc, kwargs=rpc_kwargs) + # 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 + logger.info( + f"[RL] Weights received via distributed " + f"(version={version}, rpc={rpc})" + ) + return {"status": "ok", "version": version} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] update_weights_from_distributed failed: {e}") + return {"status": "error", "message": str(e)} + + async def update_weights_from_tensor(self, body: dict) -> dict: + """Not implemented: in-process tensor transfer is not yet supported.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + return { + "status": "error", + "message": "update_weights_from_tensor is not implemented", + } + + async def init_weights_update_group(self, body: dict) -> dict: + """Initialize the distributed weight-update communication group.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + rpc = body.get("engine_rpc", "init_broadcaster") + kwargs = {k: v for k, v in body.items() if k != "engine_rpc"} + async with self._pause_lock: + try: + await self.engine_client.collective_rpc(rpc, kwargs=kwargs) + logger.info(f"[RL] Weight update group initialized (rpc={rpc})") + return {"status": "ok", "message": "Weight update group initialized"} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] init_weights_update_group failed: {e}") + return {"status": "error", "message": str(e)} + + async def destroy_weights_update_group(self, body: dict) -> dict: + """Tear down the distributed weight-update communication group.""" + if body is None: + body = {} + elif not isinstance(body, dict): + return { + "status": "error", + "message": "request body must be a JSON object", + } + rpc = body.get("engine_rpc", "destroy_broadcaster") + kwargs = {k: v for k, v in body.items() if k != "engine_rpc"} + async with self._pause_lock: + try: + await self.engine_client.collective_rpc(rpc, kwargs=kwargs) + logger.info(f"[RL] Weight update group destroyed (rpc={rpc})") + return {"status": "ok", "message": "Weight update group destroyed"} + except EngineDeadError as e: + self._shutdown_on_engine_dead(e) + except Exception as e: + logger.error(f"[RL] destroy_weights_update_group failed: {e}") + return {"status": "error", "message": str(e)} + @abstractmethod def generate(self, request: RequestT, context: Context) -> AsyncIterator[ResponseT]: raise NotImplementedError @@ -1431,46 +1857,22 @@ async def load_lora(self, request=None): } } - This method is idempotent - concurrent calls for the same LoRA will be - serialized and only one load operation will happen. + Concurrent calls for the same LoRA are serialized. Re-loading an already + loaded LoRA is idempotent by default. Set + ``DYN_LORA_HOTSWAP_ENABLED=true`` to replace an already loaded LoRA with + a new URI. """ try: - if request is None: - yield { - "status": "error", - "message": "Request is required with 'lora_name' and 'source.uri'", - } - return - - lora_name = request.get("lora_name") - if not lora_name: - yield { - "status": "error", - "message": "'lora_name' is required in request", - } + try: + lora_name, lora_uri = require_lora_load_request(request) + except RLAdminValidationError as e: + yield {"status": "error", "message": str(e)} return # Debug: Log the incoming request logger.debug(f"load_lora request keys: {list(request.keys())}") logger.debug(f"load_lora request: {request}") - # Check for URI-based API format (source.uri) - source = request.get("source") - if not source or not isinstance(source, dict): - yield { - "status": "error", - "message": "'source' object is required in request", - } - return - - lora_uri = source.get("uri") - if not lora_uri: - yield { - "status": "error", - "message": "'source.uri' is required in request", - } - return - # Use LoRAManager to download from URI lora_manager = get_lora_manager() if lora_manager is None: @@ -1484,19 +1886,21 @@ async def load_lora(self, request=None): lock = self._get_lora_lock(lora_name) async with lock: try: - # Check if already loaded (idempotency check after acquiring lock). - # Another concurrent request may have loaded this LoRA while we waited. - if lora_name in self.loaded_loras: - lora_id = self.loaded_loras[lora_name].id + old_info = self.loaded_loras.get(lora_name) + hot_swap_enabled = env_bool("DYN_LORA_HOTSWAP_ENABLED") + is_hot_swap = old_info is not None and hot_swap_enabled + + if old_info is not None and not hot_swap_enabled: logger.info( - f"LoRA adapter already loaded (concurrent request completed): " - f"{lora_name} with ID {lora_id}" + f"LoRA adapter already loaded: {lora_name} " + f"with ID {old_info.id}" ) yield { "status": "success", "message": f"LoRA adapter '{lora_name}' already loaded", "lora_name": lora_name, - "lora_id": lora_id, + "lora_id": old_info.id, + "hot_swap": False, } return @@ -1518,25 +1922,115 @@ async def load_lora(self, request=None): # Generate deterministic ID from lora_name before using it lora_id = lora_name_to_id(lora_name) - # Add the LoRA to the engine - await self.engine_client.add_lora( - LoRARequest( - lora_name=lora_name, - lora_int_id=lora_id, - lora_path=lora_path, + if is_hot_swap and old_info is not None: + try: + await self.engine_client.remove_lora(old_info.id) + except Exception as e: + logger.error( + f"Failed to remove existing LoRA '{lora_name}' " + f"before hot-swap: {e}" + ) + yield { + "status": "error", + "message": ( + f"Failed to remove existing LoRA '{lora_name}' " + f"before hot-swap: {e}" + ), + "lora_name": lora_name, + } + return + + try: + await self.engine_client.add_lora( + LoRARequest( + lora_name=lora_name, + lora_int_id=lora_id, + lora_path=lora_path, + ) ) - ) + except Exception as e: + if is_hot_swap and old_info is not None: + try: + await self.engine_client.add_lora( + LoRARequest( + lora_name=lora_name, + lora_int_id=old_info.id, + lora_path=old_info.path, + ) + ) + except Exception as rollback_error: + self.loaded_loras.pop(lora_name, None) + logger.exception( + f"Rollback failed for LoRA {lora_name}: " + f"{rollback_error}" + ) + yield { + "status": "error", + "message": f"Failed to add LoRA '{lora_name}': {e}", + "lora_name": lora_name, + } + return # Track the LoRA self.loaded_loras[lora_name] = LoRAInfo(id=lora_id, path=lora_path) logger.info( - f"Successfully loaded LoRA adapter: {lora_name} with ID {lora_id}" + f"Successfully {'hot-swapped' if is_hot_swap else 'loaded'} " + f"LoRA adapter: {lora_name} with ID {lora_id}" ) + if is_hot_swap: + try: + await self.engine_client.reset_prefix_cache() + except Exception as e: + # The new adapter is already active in the engine, but + # the prefix cache still holds entries computed under + # the old adapter and could be reused incorrectly. + # Roll the ENGINE back to old_info (remove new, re-add + # old) so engine state and our tracking stay consistent + # — a metadata-only rollback would leave the new adapter + # live while we report/route the old one (codex). + rolled_back = "tracking only" + if old_info is not None: + try: + await self.engine_client.remove_lora(lora_id) + await self.engine_client.add_lora( + LoRARequest( + lora_name=lora_name, + lora_int_id=old_info.id, + lora_path=old_info.path, + ) + ) + self.loaded_loras[lora_name] = old_info + rolled_back = "engine+tracking" + except Exception as rollback_error: + # Engine is in an indeterminate adapter state; + # drop tracking so we never claim a clean swap. + self.loaded_loras.pop(lora_name, None) + logger.exception( + f"LoRA '{lora_name}' hot-swap engine " + f"rollback failed: {rollback_error}" + ) + else: + self.loaded_loras.pop(lora_name, None) + logger.error( + f"LoRA '{lora_name}' hot-swap rolled back " + f"({rolled_back}): prefix cache reset failed: {e}" + ) + yield { + "status": "error", + "message": ( + f"LoRA '{lora_name}' hot-swap aborted; prefix " + f"cache reset failed: {e}" + ), + "lora_name": lora_name, + "lora_id": lora_id, + } + return + # Publish LoRA as a ModelDeploymentCard with format: # v1/mdc/{namespace}/{component}/{endpoint}/{instance_id}/{lora_slug} # This allows the frontend to discover it and route correctly to the worker instance - if self.generate_endpoint is not None: + if not is_hot_swap and self.generate_endpoint is not None: logger.debug( f"Publishing LoRA '{lora_name}' ModelDeploymentCard to {self.generate_endpoint}" ) @@ -1637,16 +2131,20 @@ async def load_lora(self, request=None): "lora_name": lora_name, } return - else: + elif not is_hot_swap: logger.debug( f"Cannot publish LoRA '{lora_name}': generate_endpoint={self.generate_endpoint}, config={self.config}" ) yield { "status": "success", - "message": f"LoRA adapter '{lora_name}' loaded successfully", + "message": ( + f"LoRA adapter '{lora_name}' " + f"{'hot-swapped' if is_hot_swap else 'loaded'} successfully" + ), "lora_name": lora_name, "lora_id": lora_id, + "hot_swap": is_hot_swap, } finally: # Avoid lock-map growth on failed loads: if this attempt did not leave the LoRA @@ -1670,18 +2168,10 @@ async def unload_lora(self, request=None): } """ try: - if request is None: - yield { - "status": "error", - "message": "Request is required with 'lora_name' field", - } - return - lora_name = request.get("lora_name") - if not lora_name: - yield { - "status": "error", - "message": "'lora_name' is required in request", - } + try: + lora_name = require_lora_unload_request(request) + except RLAdminValidationError as e: + yield {"status": "error", "message": str(e)} return # Serialize load/unload operations per lora_name. @@ -2774,7 +3264,11 @@ async def _generate_token_mode(self, request, context, request_id): # any deferred-abort waiter spawned by the monitor is in a stable # state when close() is awaited. async with _deferred_abort_guard( - self.engine_client, request_id, is_decode_only + self.engine_client, + request_id, + is_decode_only, + self._deferred_aborts, + self._shutdown_on_engine_dead, ) as abort_guard: async with self._abort_monitor( context, request_id, abort_guard=abort_guard @@ -2859,7 +3353,20 @@ async def _generate_text_mode(self, request, context, request_id): trace_headers = context.trace_headers() - async with self._abort_monitor(context, request_id): + # Mirror _generate_token_mode: in disagg decode mode route aborts through + # the per-request deferred guard so engine_client.abort() never fires in + # the unsafe pre-first-token window, and the admin abort_request route can + # reach this request via self._deferred_aborts. + is_decode_only = self.config.disaggregation_mode == DisaggregationMode.DECODE + async with _deferred_abort_guard( + self.engine_client, + request_id, + is_decode_only, + self._deferred_aborts, + self._shutdown_on_engine_dead, + ) as abort_guard, self._abort_monitor( + context, request_id, abort_guard=abort_guard + ): try: gen = self.engine_client.generate( prompt, @@ -2888,6 +3395,8 @@ async def _generate_text_mode(self, request, context, request_id): break for output in res.outputs: + if abort_guard is not None: + abort_guard.signal_first_token() output_idx = getattr(output, "index", 0) or 0 previous_text = previous_text_per_choice.get(output_idx, "") # Calculate the delta text (new text since last chunk) 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 bf8c4cc4cb38..eca8d44dd6b6 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py +++ b/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py @@ -29,6 +29,7 @@ pytestmark = [ pytest.mark.pre_merge, + pytest.mark.unit, pytest.mark.vllm, pytest.mark.gpu_0, pytest.mark.multimodal, @@ -89,6 +90,9 @@ def _make_handler( encode_worker_client=encode_worker_client, ) handler.model_config = model_config + # BaseWorkerHandler.__init__ is bypassed above; the decode generate path + # registers per-request deferred-abort guards here. + handler._deferred_aborts = {} return handler @@ -555,6 +559,7 @@ def _make_decode_handler( handler.otel_tracing_enabled = False handler.input_param_manager = MagicMock() handler.input_param_manager.get_extra_params.return_value = {} + handler._deferred_aborts = {} return handler @@ -820,16 +825,15 @@ async def test_abort_before_first_token_does_not_fire_immediately(self): engine_client.abort = AsyncMock() guard = mod._DeferredAbort(engine_client, "req-1") - abort_task = asyncio.create_task(guard.abort()) - # Yield so the deferred waiter is scheduled and parks on - # _first_token_event.wait(). - await asyncio.sleep(0) + # abort() before first token returns promptly; the real abort is + # deferred to a background task and engine.abort is NOT called yet. + await asyncio.wait_for(guard.abort(), timeout=1.0) engine_client.abort.assert_not_called() - assert not abort_task.done() + assert guard._abort_task is not None + assert not guard._abort_task.done() - # Cleanup: close() cancels the deferred waiter, which unblocks abort_task. + # Cleanup: close() cancels the parked deferred waiter without firing abort. await guard.close() - await abort_task engine_client.abort.assert_not_called() @pytest.mark.asyncio @@ -852,15 +856,15 @@ async def test_deferred_background_task_fires_after_first_token(self): engine_client.abort = AsyncMock() guard = mod._DeferredAbort(engine_client, "req-3") - abort_task = asyncio.create_task(guard.abort()) - await asyncio.sleep(0) + # abort() returns promptly (deferred); engine.abort not called yet. + await asyncio.wait_for(guard.abort(), timeout=1.0) engine_client.abort.assert_not_called() - assert not abort_task.done() + assert guard._abort_task is not None + assert not guard._abort_task.done() - # Signalling first token wakes the deferred waiter, which then runs - # engine.abort() and unblocks abort_task. + # Signalling first token wakes the deferred waiter, which runs abort(). guard.signal_first_token() - await abort_task + await guard._abort_task engine_client.abort.assert_awaited_once_with("req-3") @@ -999,10 +1003,9 @@ async def test_close_observes_already_completed_deferred_abort(self): engine_client.abort = AsyncMock() guard = mod._DeferredAbort(engine_client, "req-close-done") - abort_task = asyncio.create_task(guard.abort()) - await asyncio.sleep(0) + await asyncio.wait_for(guard.abort(), timeout=1.0) guard.signal_first_token() - await abort_task + await guard._abort_task assert guard._abort_task is not None assert guard._abort_task.done() @@ -1044,8 +1047,8 @@ async def test_generate_token_mode_closes_guard_on_no_output(self): created_guards: list[mod._DeferredAbort] = [] real_deferred_abort = mod._DeferredAbort - def _capture(engine_client, request_id): - g = real_deferred_abort(engine_client, request_id) + def _capture(engine_client, request_id, on_engine_dead=None): + g = real_deferred_abort(engine_client, request_id, on_engine_dead) g.close = AsyncMock(wraps=g.close) created_guards.append(g) return g @@ -1303,3 +1306,84 @@ def test_mixed_and_empty(self): h64 = "f" * 64 assert mod._pad_mm_hashes_to_64([]) == [] assert mod._pad_mm_hashes_to_64(["abc", h64]) == ["abc" + "0" * 61, h64] + + +class TestRLAdminRouteHardening: + """Regressions for the codex round-2 RL admin fixes.""" + + @pytest.mark.asyncio + async def test_admin_rejects_non_dict_body(self): + handler = _make_handler() + handler.engine_client = MagicMock() + for body in ([], "x", 5, ["pause"]): + for fn in ( + handler.pause_generation, + handler.resume_generation, + handler.flush_cache, + handler.abort_request, + ): + resp = await fn(body) + assert resp["status"] == "error", (fn.__name__, body, resp) + assert "JSON object" in resp["message"] + + @pytest.mark.asyncio + async def test_abort_request_surfaces_deferred_abort_failure(self): + handler = _make_handler() + handler.engine_client = MagicMock() + + class _FailingGuard: + def __init__(self): + self._abort_exc = RuntimeError("engine abort boom") + + async def abort(self): + return None + + handler._deferred_aborts = {"req-x": _FailingGuard()} + resp = await handler.abort_request({"request_id": "req-x"}) + assert resp["status"] == "error" + assert "boom" in resp["message"] + + @pytest.mark.asyncio + async def test_abort_request_ok_when_deferred_clean(self): + handler = _make_handler() + handler.engine_client = MagicMock() + + class _CleanGuard: + def __init__(self): + self._abort_exc = None + + async def abort(self): + return None + + handler._deferred_aborts = {"req-y": _CleanGuard()} + resp = await handler.abort_request({"request_id": "req-y"}) + assert resp["status"] == "ok" + assert resp["request_id"] == "req-y" + + @pytest.mark.asyncio + async def test_deferred_abort_does_not_block_before_first_token(self): + # abort() before the first token must return promptly (the real abort is + # deferred to a background task), not hang on the first-token event. + guard = mod._DeferredAbort(MagicMock(), "req-z") + await asyncio.wait_for(guard.abort(), timeout=1.0) + assert guard._abort_exc is None + await guard.close() + + @pytest.mark.asyncio + async def test_deferred_abort_escalates_engine_dead(self): + from vllm.v1.engine.exceptions import EngineDeadError + + escalated = [] + + async def boom(_request_id): + raise EngineDeadError("engine dead") + + engine = MagicMock() + engine.abort = boom + guard = mod._DeferredAbort( + engine, "req-d", on_engine_dead=lambda e: escalated.append(e) + ) + guard.signal_first_token() # post-first-token -> immediate abort path + await guard.abort() + assert len(escalated) == 1 + assert isinstance(escalated[0], EngineDeadError) diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index 1f8497c657fa..1900c695583c 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/components/src/dynamo/vllm/worker_factory.py @@ -16,6 +16,7 @@ from vllm.v1.engine.async_llm import AsyncLLM from dynamo import prometheus_names +from dynamo.common.rl import first_endpoint_response, register_rl_routes from dynamo.common.utils.endpoint_types import parse_endpoint_types from dynamo.common.utils.prometheus import ( LLMBackendMetrics, @@ -388,11 +389,18 @@ async def _create_decode_worker( clear_endpoint = runtime.endpoint( f"{config.namespace}.{config.component}.clear_kv_blocks" ) + rl_endpoint = ( + runtime.endpoint(f"{config.namespace}.{config.component}.rl") + if config.enable_rl + else None + ) shutdown_endpoints[:] = [ generate_endpoint, clear_endpoint, ] + if rl_endpoint is not None: + shutdown_endpoints.append(rl_endpoint) lora_enabled = config.engine_args.enable_lora if lora_enabled: @@ -521,7 +529,7 @@ async def _create_decode_worker( ) # Register engine routes - self.register_engine_routes(runtime, handler) + self.register_engine_routes(runtime, handler, lora_enabled=lora_enabled) # Parse endpoint types from --endpoint-types flag model_type = parse_endpoint_types(config.endpoint_types) @@ -616,6 +624,14 @@ async def _create_decode_worker( ), ] + if rl_endpoint is not None: + serve_tasks.append( + rl_endpoint.serve_endpoint( + handler.rl_dispatch, + metrics_labels=model_metrics_labels, + ) + ) + if lora_enabled: serve_tasks.extend( [ @@ -661,6 +677,11 @@ async def _create_prefill_worker( clear_endpoint = runtime.endpoint( f"{config.namespace}.{config.component}.clear_kv_blocks" ) + rl_endpoint = ( + runtime.endpoint(f"{config.namespace}.{config.component}.rl") + if config.enable_rl + else None + ) # Use pre-created engine if provided (checkpoint mode), otherwise create new fpm_worker_id = str(generate_endpoint.connection_id()) @@ -750,7 +771,9 @@ async def _create_prefill_worker( ) # Register engine routes - self.register_engine_routes(runtime, handler) + self.register_engine_routes( + runtime, handler, lora_enabled=config.engine_args.enable_lora + ) await self._maybe_wait_for_failover_lock(handler, runtime, config) @@ -765,6 +788,8 @@ async def _create_prefill_worker( f"{config.namespace}.{config.component}.get_perf_metrics" ) shutdown_endpoints[:] = [generate_endpoint, clear_endpoint, perf_endpoint] + if rl_endpoint is not None: + shutdown_endpoints.append(rl_endpoint) # Prefill workers expose no OpenAI surface — the role is carried by # `worker_type=Prefill`. We register the legacy `ModelType.Prefill` @@ -809,7 +834,7 @@ async def _create_prefill_worker( try: logger.debug("Starting serve_endpoint for prefill worker") - await asyncio.gather( + serve_tasks = [ generate_endpoint.serve_endpoint( handler.generate, # type: ignore graceful_shutdown=True, @@ -824,7 +849,15 @@ async def _create_prefill_worker( handler.get_perf_metrics, metrics_labels=prefill_metrics_labels, ), - ) + ] + if rl_endpoint is not None: + serve_tasks.append( + rl_endpoint.serve_endpoint( + handler.rl_dispatch, + metrics_labels=prefill_metrics_labels, + ) + ) + await asyncio.gather(*serve_tasks) logger.debug("serve_endpoint completed for prefill worker") except Exception as e: logger.error(f"Failed to serve endpoints: {e}") @@ -849,7 +882,10 @@ async def _maybe_get_encode_worker_client( return None def register_engine_routes( - self, runtime: DistributedRuntime, handler: BaseWorkerHandler + self, + runtime: DistributedRuntime, + handler: BaseWorkerHandler, + lora_enabled: bool = False, ) -> None: """Register all engine routes for this handler. @@ -862,6 +898,41 @@ def register_engine_routes( runtime.register_engine_route("wake_up", handler.wake_up) runtime.register_engine_route("scale_elastic_ep", handler.scale_elastic_ep) + rl_routes: dict = { + "liveness_probe": handler.liveness_probe, + "pause_generation": handler.pause_generation, + "resume_generation": handler.resume_generation, + "flush_cache": handler.flush_cache, + "abort_request": handler.abort_request, + "update_weights_from_disk": handler.update_weights_from_disk, + "update_weights_from_distributed": handler.update_weights_from_distributed, + "update_weights_from_tensor": handler.update_weights_from_tensor, + "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, + } + + if lora_enabled: + + async def load_lora(body: dict) -> dict: + return await first_endpoint_response(handler.load_lora, body) + + async def unload_lora(body: dict) -> dict: + return await first_endpoint_response(handler.unload_lora, body) + + rl_routes["load_lora"] = load_lora + rl_routes["unload_lora"] = unload_lora + + register_rl_routes( + runtime, + handler.rl_route_registry, + rl_routes, + enable_dispatch=handler.config.enable_rl, + ) + logger.info( - "Registered engine routes: /engine/sleep, /engine/wake_up, /engine/scale_elastic_ep, /engine/start_profile, /engine/stop_profile" + "Registered engine routes: sleep, wake_up, scale_elastic_ep, " + "start_profile, stop_profile, and RL admin routes: %s%s", + ", ".join(sorted(rl_routes)), + " (LoRA routes: load_lora, unload_lora)" if lora_enabled else "", )