diff --git a/components/src/dynamo/common/backend/README.md b/components/src/dynamo/common/backend/README.md index 0acf4e8715b9..37815d2be03c 100644 --- a/components/src/dynamo/common/backend/README.md +++ b/components/src/dynamo/common/backend/README.md @@ -5,7 +5,8 @@ inference, metrics + Prometheus bridging, KV event publishing, KV-aware (DP-rank) routing, health-check canaries, OpenTelemetry tracing, and request-side guided decoding / structural tag. -> **Work in progress.** Multimodal, diffusion (image/video/DLLM), LoRA, +> **Work in progress.** Multimodal, diffusion (image/video/DLLM), +> LoRA (SGLang / TRT-LLM — vLLM is supported), > engine routes (pause/resume, profiling, weight updates), > text-in-text-out, and snapshot/CRIU are still on the non-unified > path. See [Feature Gaps](#feature-gaps) for the per-engine matrix. @@ -440,6 +441,31 @@ Lifecycle and runtime: - `DynamoException` error chain wrapping - Finish reason normalization handled by the Rust layer - Engine control plumbing, with per-backend profiling, pause/resume, and supported weight-update controls +- **Dynamic LoRA (vLLM)** — load / unload / list adapters at runtime, + with ModelDeploymentCard publishing for frontend discovery, + per-adapter serialization locks, and per-request routing. Gated on + `--enable-lora` **and** `DYN_LORA_ENABLED=true`; SGLang / TRT-LLM + advertise no LoRA updates yet. Because LoRA ops mutate engine-managed + adapters rather than the serving lifecycle, they ride the generic + **engine-update** mechanism (a sibling of engine controls, kept separate + so the control surface isn't inflated): + - **Canonical API** (unified backend): `POST /engine/update/load_lora` + `{lora_name, source:{uri}}`, `POST /engine/update/unload_lora` + `{lora_name}`, `POST /engine/update/list_loras` `{}` (uniform `POST` + + JSON body). Engine updates return **HTTP 200** with a + `{"status": "error", ...}` body on semantic failure (5xx only when the + handler raises). + - **`/v1/loras` compatibility alias** — for the unified backend, the + legacy surface forwards to the engine updates (`POST /v1/loras` → + `load_lora`, `GET /v1/loras` → `list_loras`, + `DELETE /v1/loras/{name}` → `unload_lora`), preserving legacy HTTP + semantics (a `status:"error"` response maps to **HTTP 500** on + load/unload). When LoRA is unsupported, these return an explicit + "LoRA management not available" error rather than failing opaquely. + - **Legacy (non-unified) vLLM** continues to serve `/v1/loras` + unchanged. + - Loaded adapters appear in `GET /v1/models`; inference selects an + adapter by sending `"model": ""`. - **Disaggregated serving** (`agg`/`prefill`/`decode`) — KV transfer uses NIXL across all three engines; SGLang exchanges a Dynamo-level bootstrap address, vLLM and TRT-LLM use an engine-internal handshake. @@ -498,7 +524,7 @@ Request handling: | Text-in-text-out mode | OpenAI-compatible chat/completion with engine-side tokenization. Unified hardcodes `ModelInput.Tokens`. | | Multimodal | Images / video / embeddings, NIXL embedding transfer, encode workers. `worker.py:_to_rust_disaggregation_mode` rejects the `ENCODE` role. | | Diffusion | Image (FLUX), video (Wan2.1), LLM diffusion (DLLM) workers; no diffusion engine, MediaOutput, or media scheduling on the unified path. | -| LoRA adapters | Dynamic load / unload / list, ModelDeploymentCard publishing, per-adapter serialization locks, per-request adapter threading on prefill. | +| LoRA adapters (SGLang / TRT-LLM) | Dynamic load / unload / list, ModelDeploymentCard publishing, per-adapter serialization locks, per-request adapter threading. **vLLM is supported on the unified path** — see [What works today](#what-works-today); SGLang and TRT-LLM advertise no LoRA updates yet. | | Snapshot / checkpoint | CRIU-based engine state save/restore + identity reload. | ### vLLM-specific gaps @@ -517,7 +543,6 @@ Request handling: | `--benchmark-mode` family | The `--benchmark-*` flag family (mode, prefill/decode granularities, warmup, output path, timeout) injects into `vllm_config.additional_config` | | "Omni" alternative entry point | `dynamo.vllm.omni.*` parallel mode for alternative tensor workflows | | Multimodal (vLLM) | NIXL embedding transfer (`EmbeddingTransferMode`, `--embedding-transfer-mode`), embedding LRU cache (`--multimodal-embedding-cache-capacity-gb`), Qwen VL mRoPE, `EncodeWorkerHandler`, `--route-to-encoder` | -| LoRA (vLLM) | Three endpoints (`load_lora`, `unload_lora`, `list_loras`); also: unified prefill doesn't thread per-request LoRA adapters into the engine call | ### SGLang-specific gaps @@ -563,9 +588,10 @@ For users picking what to land next on the unified path: 1. **Text-in-text-out** (`ModelInput.Text`) — common ask; needs engine-side tokenization + chat templating path. -2. **LoRA dynamic load/unload + MDC publishing** — production-visible - feature with concrete API surface (three endpoints on vLLM - `handlers.py`). +2. **LoRA dynamic load/unload + MDC publishing** — **done for vLLM** + (engine updates `/engine/update/load_lora|unload_lora|list_loras` + a + `/v1/loras` compatibility alias; see [What works today](#what-works-today)). + Remaining: SGLang and TRT-LLM, which advertise no LoRA updates yet. 3. **Engine routes / lifecycle endpoints** — sleep/wake, profile start/stop, weight updates, KV block clearing, prefix cache reset. Visible in operator workflows. diff --git a/components/src/dynamo/common/backend/engine.py b/components/src/dynamo/common/backend/engine.py index dae051f185b1..824a1b8b9f54 100644 --- a/components/src/dynamo/common/backend/engine.py +++ b/components/src/dynamo/common/backend/engine.py @@ -42,9 +42,14 @@ class GenerateRequest(TypedDict, total=False): Disaggregated-serving keys (``prefill_result``, ``bootstrap_info``) are set by the frontend's PrefillRouter on decode requests; engines read them via ``dynamo.common.backend.disagg`` helpers. + + ``model`` carries the requested model name (set by the Rust + preprocessor). Engines that support dynamic LoRA read it to route a + request to a loaded adapter. """ token_ids: Required[list[int]] + model: str sampling_options: dict[str, Any] stop_conditions: dict[str, Any] output_options: dict[str, Any] @@ -273,11 +278,10 @@ async def health_check_payload(self) -> Optional[dict[str, Any]]: return None def supported_controls(self) -> set[str]: - """Engine-control capability keys this engine supports. + """Return the set of engine-control capability keys this engine supports. - The unified backend maps these keys to runtime endpoints. Engines only - advertise and implement semantic controls; they do not own transport or - route registration details. + Controls are semantic operations on the engine's serving lifecycle. + Engines advertise the keys they implement. """ return set() @@ -290,6 +294,30 @@ async def engine_control( "message": f"unsupported engine control: {control}", } + def supported_updates(self) -> set[str]: + """Return the set of engine-update capability keys this engine supports. + + Updates are a sibling surface to :meth:`supported_controls` for + operations that mutate engine-managed assets rather than the engine's + serving lifecycle. Engines advertise the keys they implement. + """ + return set() + + async def engine_update(self, update: str, body: dict[str, Any]) -> dict[str, Any]: + """Handle one advertised engine-update request.""" + return { + "status": "error", + "message": f"unsupported engine update: {update}", + } + + async def on_endpoint_ready(self, endpoint) -> None: + """Receive the runtime serving ``Endpoint`` once, before serving begins. + + Default no-op. Engines that publish their own discovery records stash + it for use from :meth:`engine_update`. ``Worker`` calls this exactly + once; a raised exception is fatal to startup.""" + return None + class LLMEngine(BaseEngine): """Abstract base for token-based inference engines (vLLM, SGLang, TRT-LLM). diff --git a/components/src/dynamo/sglang/request_handlers/handler_base.py b/components/src/dynamo/sglang/request_handlers/handler_base.py index 66aae108e0ba..96d6567a9efd 100644 --- a/components/src/dynamo/sglang/request_handlers/handler_base.py +++ b/components/src/dynamo/sglang/request_handlers/handler_base.py @@ -975,28 +975,29 @@ def register_engine_routes(self, runtime: DistributedRuntime) -> None: Args: runtime: The DistributedRuntime instance to register routes on. """ - runtime.register_engine_route("start_profile", self.start_profile) - runtime.register_engine_route("stop_profile", self.stop_profile) + runtime.register_engine_route("control/start_profile", self.start_profile) + runtime.register_engine_route("control/stop_profile", self.stop_profile) runtime.register_engine_route( - "release_memory_occupation", self.release_memory_occupation + "control/release_memory_occupation", self.release_memory_occupation ) runtime.register_engine_route( - "resume_memory_occupation", self.resume_memory_occupation + "control/resume_memory_occupation", self.resume_memory_occupation ) runtime.register_engine_route( - "update_weights_from_disk", self.update_weights_from_disk + "control/update_weights_from_disk", self.update_weights_from_disk ) runtime.register_engine_route( - "update_weights_from_tensor", self.update_weights_from_tensor + "control/update_weights_from_tensor", self.update_weights_from_tensor ) runtime.register_engine_route( - "update_weights_from_distributed", self.update_weights_from_distributed + "control/update_weights_from_distributed", + self.update_weights_from_distributed, ) runtime.register_engine_route( - "update_weights_from_ipc", self.update_weights_from_ipc + "control/update_weights_from_ipc", self.update_weights_from_ipc ) runtime.register_engine_route( - "update_weight_version", self.update_weight_version + "control/update_weight_version", self.update_weight_version ) if getattr(self.config, "dynamo_args", None) and getattr( self.config.dynamo_args, "enable_rl", False diff --git a/components/src/dynamo/trtllm/workers/llm_worker.py b/components/src/dynamo/trtllm/workers/llm_worker.py index 8036158bff49..66aa1026989f 100644 --- a/components/src/dynamo/trtllm/workers/llm_worker.py +++ b/components/src/dynamo/trtllm/workers/llm_worker.py @@ -135,16 +135,16 @@ def _sync_config_from_engine_args(config: Config, engine_args: dict) -> None: def _register_memory_routes(runtime, handler) -> None: runtime.register_engine_route( - "release_memory_occupation", + "control/release_memory_occupation", handler.release_memory_occupation, ) runtime.register_engine_route( - "resume_memory_occupation", + "control/resume_memory_occupation", handler.resume_memory_occupation, ) logging.info( "Registered engine routes: " - "/engine/release_memory_occupation, /engine/resume_memory_occupation" + "/engine/control/release_memory_occupation, /engine/control/resume_memory_occupation" ) diff --git a/components/src/dynamo/vllm/llm_engine.py b/components/src/dynamo/vllm/llm_engine.py index 6133494c83e6..57c03d28ff8d 100644 --- a/components/src/dynamo/vllm/llm_engine.py +++ b/components/src/dynamo/vllm/llm_engine.py @@ -19,6 +19,7 @@ from vllm.config import VllmConfig from vllm.distributed.kv_events import ZmqEventPublisher from vllm.inputs import TokensPrompt +from vllm.lora.request import LoRARequest from vllm.usage.usage_lib import UsageContext from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.metrics.loggers import StatLoggerBase @@ -52,7 +53,17 @@ from dynamo.common.backend.publisher import ComponentSnapshot, KvEventSource, ZmqSource from dynamo.common.backend.worker import WorkerConfig from dynamo.common.constants import DisaggregationMode -from dynamo.llm import ModelInput +from dynamo.common.lora.manager import LoRAInfo, get_lora_manager +from dynamo.llm import ( + ModelInput, + ModelRuntimeConfig, + ModelType, + WorkerType, + lora_name_to_id, + register_model, + unregister_model, +) +from dynamo.runtime import Endpoint from dynamo.vllm.args import configure_rl_logprobs_mode, parse_args from dynamo.vllm.cache_info import ( configure_kv_event_block_size, @@ -144,6 +155,13 @@ def __call__(self, vllm_config: VllmConfig, dp_rank: int) -> StatLoggerBase: return _UnifiedStatLogger(self, dp_rank) +# Number of stripe locks serializing per-adapter LoRA load/unload. Fixed so +# lock memory stays bounded no matter how many distinct adapter names are seen; +# distinct names may share a stripe (harmless extra serialization on this +# control-plane path). +_LORA_LOCK_STRIPES = 32 + + class VllmLLMEngine(LLMEngine): # Class-level default so ``__new__``-built instances (tests skipping # ``__init__``) still expose what ``generate()`` reads; ``start()`` sets it. @@ -155,12 +173,16 @@ def __init__( disaggregation_mode: DisaggregationMode, served_model_name: str, component: str, + dyn_tool_call_parser: Optional[str] = None, + dyn_reasoning_parser: Optional[str] = None, enable_rl: bool = False, ): self.engine_args = engine_args self.disaggregation_mode = disaggregation_mode self._served_model_name = served_model_name self._component = component + self._dyn_tool_call_parser = dyn_tool_call_parser + self._dyn_reasoning_parser = dyn_reasoning_parser self.enable_rl = enable_rl self.engine_client: AsyncLLM | None = None self._vllm_config: Any = None @@ -168,6 +190,13 @@ def __init__( self._prometheus_temp_dir: tempfile.TemporaryDirectory[str] | None = None self._model_max_len: int | None = None self._dp_range: Optional[tuple[int, int]] = None + # Effective KV-event block size, computed in start(). LoRA MDCs must + # publish this (not engine_args.block_size) so LoRA block hashes match + # vLLM's emitted KV events for routing. + self._kv_event_block_size: int | None = None + # Per-rank KV block count, computed in start() and published on both the + # base-model and LoRA MDCs so the router sees the worker's real capacity. + self._total_kv_blocks: int | None = None # Constructed in start() before AsyncLLM init so vLLM's stat-logger # factory call sees a valid object. `num_gpu_blocks` is patched # after KV profiling finishes. @@ -177,6 +206,26 @@ def __init__( self._pause_lock = asyncio.Lock() self._scale_ep_lock = asyncio.Lock() self._scale_ep_in_progress = False + # Dynamic LoRA state. `_endpoint` is set by `on_endpoint_ready` before + # serving begins; LoRA discovery (register_model/unregister_model) + # publishes against it. + self._endpoint: Optional[Endpoint] = None + self.loaded_loras: dict[str, LoRAInfo] = {} + # Adapters whose discovery ModelDeploymentCard is currently published. + # Tracked separately from `loaded_loras` because the engine load and the + # discovery publish can diverge on partial failure: an adapter may be + # loaded into vLLM yet have no card (publish failed and engine-side + # rollback also failed), or have a stale card with no engine load + # (unregister failed and re-add rollback also failed). Keeping the two + # states apart lets a retried load/unload reconcile the divergence + # instead of short-circuiting as "already loaded" / "not found". + self._published_loras: set[str] = set() + # Striped locks serialize concurrent load/unload of the same adapter. + # A fixed array keyed by hash bounds lock memory (no per-name growth) + # and preserves the "same name -> same lock" invariant by construction, + # so there is no lock-eviction race. Relies only on per-process hash + # stability, which is all we need within a single worker. + self._lora_load_locks = [asyncio.Lock() for _ in range(_LORA_LOCK_STRIPES)] @classmethod async def from_args( @@ -207,6 +256,8 @@ async def from_args( mode, served_model_name=config.served_model_name or config.model, component=config.component, + dyn_tool_call_parser=config.dyn_tool_call_parser, + dyn_reasoning_parser=config.dyn_reasoning_parser, enable_rl=config.enable_rl, ) worker_config = WorkerConfig.from_runtime_config( @@ -217,6 +268,10 @@ async def from_args( ) return engine, worker_config + async def on_endpoint_ready(self, endpoint: Endpoint) -> None: + """Stash the serving endpoint for dynamic-LoRA discovery publishing.""" + self._endpoint = endpoint + async def start(self, worker_id: int) -> EngineConfig: """Start vLLM and return normalized metadata for runtime registration.""" del worker_id # vLLM's NixlConnector handles its own per-worker IDs @@ -257,6 +312,7 @@ async def start(self, worker_id: int) -> EngineConfig: if per_rank_num_gpu_blocks is None: raise RuntimeError("per-rank KV block count is not set") self._stat_logger_factory.num_gpu_blocks = per_rank_num_gpu_blocks + self._total_kv_blocks = per_rank_num_gpu_blocks self._model_max_len = getattr( getattr(vllm_config, "model_config", None), "max_model_len", None ) @@ -268,6 +324,7 @@ async def start(self, worker_id: int) -> EngineConfig: # both to the runtime via `EngineConfig` and to any future readers. await configure_kv_event_block_size(self.engine_client, vllm_config) block_size = get_configured_kv_event_block_size(vllm_config) + self._kv_event_block_size = block_size return EngineConfig( model=self.engine_args.model, @@ -366,11 +423,17 @@ async def generate( ) local_dp_rank = None if rank is None else rank - dp_start + # Route to a loaded LoRA adapter when the request names one; the base + # model resolves to None. With LoRA enabled, an unknown adapter name + # raises rather than silently falling back to the base model. + lora_request = self._resolve_lora_request(request.get("model")) + gen = self.engine_client.generate( prompt, sampling_params, request_id, data_parallel_rank=local_dp_rank, + lora_request=lora_request, **telemetry.engine_trace_kwargs(context), ) @@ -545,6 +608,14 @@ async def register_prometheus(self, metrics: "EngineMetrics") -> None: multiproc_only_prefixes=["lmcache:"], ) + def _lora_enabled(self) -> bool: + """Dynamic-LoRA updates are available only when the engine was built + with ``--enable-lora`` AND the LoRA manager is initialized + (``DYN_LORA_ENABLED=true``).""" + return bool(getattr(self.engine_args, "enable_lora", False)) and ( + get_lora_manager() is not None + ) + def supported_controls(self) -> set[str]: controls = {"start_profile", "stop_profile", "sleep", "wake_up"} if self.engine_client is not None and hasattr( @@ -573,6 +644,506 @@ async def engine_control(self, control: str, body: dict) -> dict: } return await handler(body or {}) + def supported_updates(self) -> set[str]: + # LoRA lifecycle ops mutate engine-managed adapters, so they ride the + # engine-update surface (/engine/update/) rather than inflating + # the engine-control surface. + if self._lora_enabled(): + return {"load_lora", "unload_lora", "list_loras"} + return set() + + async def engine_update(self, update: str, body: dict) -> dict: + handlers = {} + if self._lora_enabled(): + handlers["load_lora"] = self.load_lora + handlers["unload_lora"] = self.unload_lora + handlers["list_loras"] = self.list_loras + + handler = handlers.get(update) + if handler is None: + return { + "status": "error", + "message": f"unsupported engine update: {update}", + } + return await handler(body or {}) + + def _resolve_lora_request(self, model_name: str | None) -> LoRARequest | None: + """Return a LoRARequest for a loaded adapter, or None for the base model. + + Raises ValueError when LoRA is enabled and ``model_name`` is a non-base + name with no loaded adapter, so an unknown or just-unloaded adapter + fails loudly instead of being silently served by the base model. When + LoRA is disabled there are no adapters, so a non-base name is left to + the engine (current behavior) rather than rejected here. + """ + if not model_name or model_name in ( + self._served_model_name, + self.engine_args.model, + ): + return None + lora = self.loaded_loras.get(model_name) + if lora is not None: + return LoRARequest( + lora_name=model_name, + lora_int_id=lora.id, + lora_path=lora.path, + ) + if self._lora_enabled(): + raise ValueError(f"unknown model or LoRA adapter: '{model_name}'") + return None + + def _get_lora_lock(self, lora_name: str) -> asyncio.Lock: + """Return the stripe lock that serializes load/unload for ``lora_name``. + + A name always maps to the same stripe within the process, so every + load/unload for a given ``lora_name`` serializes on the *same* lock. + Because the stripe set is fixed, there is no per-name lock to evict and + thus no eviction race; the cost is that two distinct names sharing a + stripe serialize against each other (harmless on this control path). + """ + return self._lora_load_locks[hash(lora_name) % _LORA_LOCK_STRIPES] + + async def _publish_lora_card(self, lora_name: str, lora_id: int) -> None: + """Publish a LoRA adapter as a ModelDeploymentCard for discovery. + + Assumes ``self._endpoint`` is set (callers gate on it). Raises on + failure so callers can roll back or retry; on success the caller is + responsible for recording ``lora_name`` in ``self._published_loras``. + """ + assert self._endpoint is not None + + user_data = { + "lora_adapter": True, + "lora_id": lora_id, + } + + # Match the base-model registration topology (see main.py + # register_vllm_model + worker_factory) so the frontend builds the LoRA + # pipeline against the right component. Without this, a prefill worker + # would publish the adapter as a decode-capable chat/completions model + # and the frontend would route chat traffic straight to prefill, which + # then waits forever for a KV transfer. + model_type, worker_type, needs = self._lora_registration_topology() + + runtime_config = ModelRuntimeConfig() + # Prefill workers don't run tool/reasoning parsing (mirrors the base + # model registration in main.py:register_vllm_model). + if model_type != ModelType.Prefill: + runtime_config.tool_call_parser = self._dyn_tool_call_parser + runtime_config.reasoning_parser = self._dyn_reasoning_parser + + # Carry the worker's DP-rank range and capacity metadata (the same + # effective vLLM values the base-model MDC publishes via EngineConfig in + # start()), so multi-DP LoRA requests are routed/attributed per rank + # instead of as if every worker only served rank 0. start() always runs + # before a load; guard in case a load somehow races ahead of it. + if self._vllm_config is not None and self._dp_range is not None: + scheduler_config = self._vllm_config.scheduler_config + if self._total_kv_blocks is not None: + runtime_config.total_kv_blocks = self._total_kv_blocks + runtime_config.max_num_seqs = scheduler_config.max_num_seqs + runtime_config.max_num_batched_tokens = ( + scheduler_config.max_num_batched_tokens + ) + runtime_config.data_parallel_start_rank = self._dp_range[0] + runtime_config.data_parallel_size = self._dp_range[1] + + # Publish the effective KV-event block size (computed in start() and + # used by the base-model MDC) so LoRA block hashes match vLLM's emitted + # KV events. start() always runs before a load, but fall back to the + # engine arg if it somehow hasn't. + kv_cache_block_size = ( + self._kv_event_block_size + if self._kv_event_block_size is not None + else self.engine_args.block_size + ) + + await register_model( + model_input=ModelInput.Tokens, + model_type=model_type, + endpoint=self._endpoint, + model_path=self.engine_args.model, + kv_cache_block_size=kv_cache_block_size, + runtime_config=runtime_config, + user_data=user_data, + lora_name=lora_name, + base_model_path=self.engine_args.model, + worker_type=worker_type, + needs=needs, + ) + + def _lora_registration_topology( + self, + ) -> tuple[ModelType, WorkerType, list[list[WorkerType]]]: + """Map the worker's disaggregation role to the LoRA MDC topology. + + Returns ``(model_type, worker_type, needs)`` matching how the base + model registers (main.py:register_vllm_model + worker_factory). + """ + if self.disaggregation_mode == DisaggregationMode.PREFILL: + return ModelType.Prefill, WorkerType.Prefill, [[WorkerType.Decode]] + if self.disaggregation_mode == DisaggregationMode.DECODE: + return ( + ModelType.Chat | ModelType.Completions, + WorkerType.Decode, + [[WorkerType.Prefill]], + ) + return ModelType.Chat | ModelType.Completions, WorkerType.Aggregated, [] + + async def load_lora(self, body: dict) -> dict: + """Load a LoRA adapter dynamically into vLLM's AsyncLLM engine. + + Request body: ``{"lora_name": str, "source": {"uri": str}}``. + + Idempotent: concurrent loads of the same name are serialized and only + one load operation happens. + """ + request = body or {} + if self.engine_client is None: + return {"status": "error", "message": "Engine is not running"} + try: + lora_name = request.get("lora_name") + if not lora_name: + return { + "status": "error", + "message": "'lora_name' is required in request", + } + + # Reject names that collide with the base model. A LoRA card shares + # the frontend model key with its name, so an adapter named after the + # base model would shadow it and make _resolve_lora_request route + # plain base-model requests through the adapter. + if lora_name in (self._served_model_name, self.engine_args.model): + return { + "status": "error", + "message": ( + f"LoRA name '{lora_name}' collides with the base model; " + "choose a different adapter name" + ), + } + + logger.debug("load_lora request keys: %s", list(request.keys())) + + source = request.get("source") + if not source or not isinstance(source, dict): + return { + "status": "error", + "message": "'source' object is required in request", + } + + lora_uri = source.get("uri") + if not lora_uri: + return { + "status": "error", + "message": "'source.uri' is required in request", + } + + lora_manager = get_lora_manager() + if lora_manager is None: + return { + "status": "error", + "message": "LoRAManager not initialized. Set DYN_LORA_ENABLED=true to enable URI-based LoRA loading.", + } + + # Serialize load/unload operations per lora_name. + lock = self._get_lora_lock(lora_name) + async with lock: + # Idempotency check after acquiring the lock: a 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 + # The adapter is loaded into the engine, but its + # discovery card may be missing (a prior publish failed + # and the engine-side rollback also failed). Reconcile by + # retrying the publish instead of reporting early success. + if ( + self._endpoint is not None + and lora_name not in self._published_loras + ): + logger.info( + "LoRA '%s' loaded but unpublished; " + "retrying discovery publish", + lora_name, + ) + try: + await self._publish_lora_card(lora_name, lora_id) + self._published_loras.add(lora_name) + except Exception as e: + logger.exception( + "Failed to publish LoRA %s ModelDeploymentCard: %s", + lora_name, + e, + ) + return { + "status": "error", + "message": f"LoRA '{lora_name}' is loaded but discovery publish failed: {str(e)}", + "lora_name": lora_name, + } + logger.info( + "LoRA adapter already loaded (concurrent request completed): " + "%s with ID %s", + lora_name, + lora_id, + ) + return { + "status": "success", + "message": f"LoRA adapter '{lora_name}' already loaded", + "lora_name": lora_name, + "lora_id": lora_id, + } + + logger.info("Downloading LoRA adapter: %s from %s", lora_name, lora_uri) + download_result = await lora_manager.download_lora(lora_uri) + + if download_result["status"] != "success": + return { + "status": "error", + "message": f"Failed to download LoRA: {download_result.get('message', 'Unknown error')}", + } + + lora_path = download_result["local_path"] + logger.debug("LoRA downloaded to: %s", lora_path) + + # Deterministic ID from lora_name before using it. + lora_id = lora_name_to_id(lora_name) + + await self.engine_client.add_lora( + LoRARequest( + lora_name=lora_name, + lora_int_id=lora_id, + lora_path=lora_path, + ) + ) + + self.loaded_loras[lora_name] = LoRAInfo(id=lora_id, path=lora_path) + logger.info( + "Successfully loaded LoRA adapter: %s with ID %s", + lora_name, + lora_id, + ) + + # Publish the LoRA as a ModelDeploymentCard so the frontend + # can discover it and route to this worker instance. + if self._endpoint is not None: + logger.debug( + "Publishing LoRA '%s' ModelDeploymentCard to %s", + lora_name, + self._endpoint, + ) + try: + await self._publish_lora_card(lora_name, lora_id) + self._published_loras.add(lora_name) + logger.info( + "Successfully published LoRA '%s' ModelDeploymentCard", + lora_name, + ) + except Exception as e: + logger.exception( + "Failed to publish LoRA %s ModelDeploymentCard: %s", + lora_name, + e, + ) + + # Rollback: remove the LoRA from the engine to keep + # engine state and discovery consistent. If the + # rollback itself fails, the entry stays in + # `loaded_loras` but absent from `_published_loras`, + # so a retried load reconciles the publish. + try: + logger.debug( + "Rolling back: removing LoRA '%s' from engine", + lora_name, + ) + await self.engine_client.remove_lora(lora_id) + self.loaded_loras.pop(lora_name, None) + logger.debug( + "Successfully rolled back LoRA '%s'", lora_name + ) + except Exception as rollback_error: + logger.exception( + "Failed to rollback LoRA %s: %s", + lora_name, + rollback_error, + ) + self._published_loras.discard(lora_name) + + return { + "status": "error", + "message": f"Failed to register LoRA '{lora_name}' in discovery registry: {str(e)}", + "lora_name": lora_name, + } + else: + logger.debug( + "Cannot publish LoRA '%s': serving endpoint not ready", + lora_name, + ) + + return { + "status": "success", + "message": f"LoRA adapter '{lora_name}' loaded successfully", + "lora_name": lora_name, + "lora_id": lora_id, + } + except Exception as e: + logger.exception("Failed to load LoRA adapter: %s", e) + return {"status": "error", "message": str(e)} + + async def unload_lora(self, body: dict) -> dict: + """Unload a LoRA adapter dynamically from vLLM's AsyncLLM engine. + + Request body: ``{"lora_name": str}``. + """ + request = body or {} + if self.engine_client is None: + return {"status": "error", "message": "Engine is not running"} + try: + lora_name = request.get("lora_name") + if not lora_name: + return { + "status": "error", + "message": "'lora_name' is required in request", + } + + # Serialize load/unload operations per lora_name. + lock = self._get_lora_lock(lora_name) + async with lock: + # Check existence *after* waiting for any in-progress load. + lora = self.loaded_loras.get(lora_name) + if lora is None: + # The adapter is gone from the engine but may still have + # a stale discovery card (a prior unload's unregister + # failed and the re-add rollback also failed). Reconcile + # by retrying the unregister so discovery converges. + if ( + self._endpoint is not None + and lora_name in self._published_loras + ): + logger.info( + "LoRA '%s' not loaded but still published; " + "retrying discovery unregister", + lora_name, + ) + try: + await unregister_model( + endpoint=self._endpoint, + lora_name=lora_name, + ) + self._published_loras.discard(lora_name) + return { + "status": "success", + "message": f"LoRA adapter '{lora_name}' discovery card removed", + "lora_name": lora_name, + } + except Exception as e: + logger.exception( + "Failed to unregister stale LoRA %s ModelDeploymentCard: %s", + lora_name, + e, + ) + return { + "status": "error", + "message": f"Failed to unregister stale LoRA '{lora_name}' from discovery registry: {str(e)}", + "lora_name": lora_name, + } + return { + "status": "error", + "message": f"LoRA adapter '{lora_name}' not found. Available LoRAs: {list(self.loaded_loras.keys())}", + } + + logger.debug("Unloading LoRA adapter: %s", lora_name) + lora_id = lora.id + + # Stop advertising the adapter *before* removing it from the + # engine, so the frontend stops routing LoRA traffic here + # while the adapter still exists. Removing it first would + # leave a window where requests route to a worker that no + # longer has the adapter (falling back to base or failing). + if self._endpoint is not None and lora_name in self._published_loras: + logger.debug( + "Unregistering LoRA '%s' ModelDeploymentCard", + lora_name, + ) + try: + await unregister_model( + endpoint=self._endpoint, + lora_name=lora_name, + ) + self._published_loras.discard(lora_name) + logger.info( + "Successfully unregistered LoRA '%s' ModelDeploymentCard", + lora_name, + ) + except Exception as e: + # Nothing mutated yet: the engine still has the + # adapter and discovery still advertises it + # (consistent and still routable). Surface the error + # and leave state intact for a retry. + logger.exception( + "Failed to unregister LoRA %s ModelDeploymentCard: %s", + lora_name, + e, + ) + return { + "status": "error", + "message": f"Failed to unregister LoRA '{lora_name}' from discovery registry: {str(e)}", + "lora_name": lora_name, + } + elif self._endpoint is None: + logger.debug( + "Cannot unregister LoRA '%s': serving endpoint not ready", + lora_name, + ) + + # Discovery no longer routes to this adapter; remove it from + # the engine. + try: + await self.engine_client.remove_lora(lora_id) + except Exception as e: + # The discovery card is already gone but the engine still + # holds the adapter (loaded-but-unpublished). Leave it in + # loaded_loras so a retried unload skips the unregister + # and retries only the engine removal. + logger.exception( + "Failed to remove LoRA %s from engine: %s", + lora_name, + e, + ) + return { + "status": "error", + "message": f"Failed to remove LoRA '{lora_name}' from engine: {str(e)}", + "lora_name": lora_name, + } + + del self.loaded_loras[lora_name] + + logger.info( + "Successfully unloaded LoRA adapter: %s with ID %s", + lora_name, + lora_id, + ) + return { + "status": "success", + "message": f"LoRA adapter '{lora_name}' unloaded successfully", + "lora_name": lora_name, + "lora_id": lora_id, + } + except Exception as e: + logger.exception("Failed to unload LoRA adapter: %s", e) + return {"status": "error", "message": str(e)} + + async def list_loras(self, body: dict) -> dict: + """List all loaded LoRA adapters as a lora_name -> lora_id mapping.""" + try: + loras = {name: lora.id for name, lora in self.loaded_loras.items()} + return { + "status": "success", + "loras": loras, + "count": len(loras), + } + except Exception as e: + logger.error("Failed to list LoRA adapters: %s", e) + return {"status": "error", "message": str(e)} + async def sleep(self, body: dict) -> dict: body = body or {} level = body.get("level", 1) @@ -726,6 +1297,15 @@ async def cleanup(self) -> None: finally: self.engine_client = None self._pause_controller = None + # Drop the serving endpoint and dynamic-LoRA bookkeeping so a + # shut-down engine holds no dangling endpoint reference and no + # stale adapter state. Discovery cards published for the worker are + # reclaimed when the endpoint's lease expires on process exit. The + # stripe locks are fixed process state, not per-adapter, so they + # are left intact. + self._endpoint = None + self.loaded_loras.clear() + self._published_loras.clear() if self._prometheus_temp_dir is not None: if ( os.environ.get("PROMETHEUS_MULTIPROC_DIR") diff --git a/components/src/dynamo/vllm/tests/test_vllm_lora.py b/components/src/dynamo/vllm/tests/test_vllm_lora.py new file mode 100644 index 000000000000..b9890f8531e4 --- /dev/null +++ b/components/src/dynamo/vllm/tests/test_vllm_lora.py @@ -0,0 +1,726 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for unified-vLLM dynamic-LoRA support. + +Covers engine-control gating, per-request adapter resolution, the +load/unload/list lifecycle (including discovery publish/unpublish and the +add_lora<->register_model rollback couplings), and the on_endpoint_ready +handoff. Everything is mocked: no GPU, no real AsyncLLM, no real discovery. +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +pytest.importorskip("vllm.lora.request") +pytest.importorskip("vllm.usage.usage_lib") +pytest.importorskip("vllm.v1.engine.async_llm") + +from dynamo.common.constants import DisaggregationMode # noqa: E402 +from dynamo.common.lora.manager import LoRAInfo # noqa: E402 +from dynamo.llm import ModelType, WorkerType # noqa: E402 +from dynamo.vllm import llm_engine as llm_engine_mod # noqa: E402 +from dynamo.vllm.llm_engine import VllmLLMEngine # noqa: E402 + +pytestmark = [ + pytest.mark.unit, + pytest.mark.vllm, + pytest.mark.gpu_0, + pytest.mark.pre_merge, +] + + +def _make_lora_engine(enable_lora: bool = True, endpoint=None) -> VllmLLMEngine: + """Build a VllmLLMEngine with only the LoRA-relevant state populated. + + Calls the real ``__init__`` (side-effect-free: only attribute assignment + and lock creation; AsyncLLM is built later in ``start()``), so new engine + attributes get their real defaults automatically and this helper does not + silently drift from the constructor. Only the fields ``__init__`` leaves + None/unset that the LoRA paths read are overridden below. + """ + engine = VllmLLMEngine( + engine_args=SimpleNamespace( + enable_lora=enable_lora, + model="/models/base", + block_size=16, + ), + disaggregation_mode=DisaggregationMode.AGGREGATED, + served_model_name="base-model", + component="test", + ) + engine.engine_client = SimpleNamespace( + add_lora=AsyncMock(), + remove_lora=AsyncMock(), + ) + engine._kv_event_block_size = 16 + engine._endpoint = endpoint + return engine + + +def _patch_discovery(monkeypatch, *, manager=None, name_to_id=None): + """Patch the discovery + LoRA-manager symbols imported into llm_engine. + + Returns the (register_model, unregister_model) AsyncMocks for assertions. + """ + if manager is None: + manager = SimpleNamespace( + download_lora=AsyncMock( + return_value={"status": "success", "local_path": "/cache/adapter"} + ) + ) + register = AsyncMock() + unregister = AsyncMock() + monkeypatch.setattr(llm_engine_mod, "get_lora_manager", lambda: manager) + monkeypatch.setattr(llm_engine_mod, "register_model", register) + monkeypatch.setattr(llm_engine_mod, "unregister_model", unregister) + monkeypatch.setattr( + llm_engine_mod, "lora_name_to_id", name_to_id or (lambda name: 123) + ) + monkeypatch.setattr(llm_engine_mod, "ModelRuntimeConfig", MagicMock()) + return register, unregister + + +# --------------------------------------------------------------------------- # +# Engine-update gating +# +# LoRA lifecycle ops ride the engine-*update* surface (supported_updates / +# engine_update), not engine controls, so they don't inflate the control set. +# --------------------------------------------------------------------------- # + + +def test_lora_updates_not_advertised_without_manager(monkeypatch): + monkeypatch.setattr(llm_engine_mod, "get_lora_manager", lambda: None) + engine = _make_lora_engine(enable_lora=True) + + updates = engine.supported_updates() + + assert "load_lora" not in updates + assert "unload_lora" not in updates + assert "list_loras" not in updates + # LoRA must never leak back into the control surface. + assert {"load_lora", "unload_lora", "list_loras"}.isdisjoint( + engine.supported_controls() + ) + + +def test_lora_updates_not_advertised_without_enable_lora(monkeypatch): + monkeypatch.setattr(llm_engine_mod, "get_lora_manager", lambda: MagicMock()) + engine = _make_lora_engine(enable_lora=False) + + updates = engine.supported_updates() + + assert {"load_lora", "unload_lora", "list_loras"}.isdisjoint(updates) + + +@pytest.mark.asyncio +async def test_lora_updates_advertised_and_dispatchable(monkeypatch): + monkeypatch.setattr(llm_engine_mod, "get_lora_manager", lambda: MagicMock()) + engine = _make_lora_engine(enable_lora=True) + + updates = engine.supported_updates() + assert {"load_lora", "unload_lora", "list_loras"} <= updates + # LoRA must not appear among controls. + assert {"load_lora", "unload_lora", "list_loras"}.isdisjoint( + engine.supported_controls() + ) + + # The dispatcher routes the update to the real method. + result = await engine.engine_update("list_loras", {}) + assert result["status"] == "success" + + +@pytest.mark.asyncio +async def test_disabled_lora_update_is_rejected_by_dispatcher(monkeypatch): + monkeypatch.setattr(llm_engine_mod, "get_lora_manager", lambda: None) + engine = _make_lora_engine(enable_lora=True) + + result = await engine.engine_update("load_lora", {"lora_name": "x"}) + + assert result["status"] == "error" + assert "unsupported engine update" in result["message"] + engine.engine_client.add_lora.assert_not_awaited() + + +# --------------------------------------------------------------------------- # +# Per-request adapter resolution +# --------------------------------------------------------------------------- # + + +def test_resolve_lora_request_for_loaded_adapter(): + engine = _make_lora_engine() + engine.loaded_loras = {"adapterA": LoRAInfo(id=7, path="/path/a")} + + lora_request = engine._resolve_lora_request("adapterA") + + assert lora_request is not None + assert lora_request.lora_name == "adapterA" + assert lora_request.lora_int_id == 7 + assert lora_request.lora_path == "/path/a" + + +def test_resolve_lora_request_for_base_is_none(monkeypatch): + monkeypatch.setattr(llm_engine_mod, "get_lora_manager", lambda: MagicMock()) + engine = _make_lora_engine() + engine.loaded_loras = {"adapterA": LoRAInfo(id=7, path="/path/a")} + + # Base-model name (served name and the engine_args.model path) and the + # absent-model case both resolve to the base model (None), even with LoRA + # enabled. + assert engine._resolve_lora_request("base-model") is None + assert engine._resolve_lora_request("/models/base") is None + assert engine._resolve_lora_request(None) is None + + +def test_resolve_lora_request_unknown_adapter_raises_when_lora_enabled(monkeypatch): + monkeypatch.setattr(llm_engine_mod, "get_lora_manager", lambda: MagicMock()) + engine = _make_lora_engine(enable_lora=True) + engine.loaded_loras = {"adapterA": LoRAInfo(id=7, path="/path/a")} + + # An unknown or just-unloaded adapter name must fail loudly rather than be + # silently served by the base model. + with pytest.raises(ValueError, match="unknown model or LoRA adapter"): + engine._resolve_lora_request("ghost-adapter") + + +def test_resolve_lora_request_unknown_adapter_is_none_when_lora_disabled(monkeypatch): + monkeypatch.setattr(llm_engine_mod, "get_lora_manager", lambda: None) + engine = _make_lora_engine(enable_lora=False) + + # With LoRA disabled there are no adapters, so a non-base name is left to + # the engine instead of being rejected here. + assert engine._resolve_lora_request("ghost-adapter") is None + + +@pytest.mark.asyncio +async def test_generate_passes_resolved_lora_request(monkeypatch): + engine = _make_lora_engine() + engine.loaded_loras = {"adapterA": LoRAInfo(id=9, path="/p/a")} + engine._default_sampling_params = SimpleNamespace() + engine._model_max_len = None + engine._dp_range = None + + captured: dict = {} + + async def _empty_gen(): + return + yield # pragma: no cover - marks this as an async generator + + def _fake_generate(*args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + return _empty_gen() + + engine.engine_client.generate = _fake_generate + monkeypatch.setattr( + llm_engine_mod, + "build_sampling_params", + lambda *a, **k: SimpleNamespace(extra_args=None, max_tokens=10, min_tokens=0), + ) + monkeypatch.setattr( + llm_engine_mod.telemetry, "engine_trace_kwargs", lambda context: {} + ) + + context = SimpleNamespace(id=lambda: "req-1") + + # Adapter request -> resolved LoRARequest. + _ = [ + c + async for c in engine.generate( + {"token_ids": [1, 2], "model": "adapterA"}, context + ) + ] + assert captured["kwargs"]["lora_request"] is not None + assert captured["kwargs"]["lora_request"].lora_name == "adapterA" + + # Base-model request -> None. + _ = [ + c + async for c in engine.generate({"token_ids": [1, 2], "model": "base"}, context) + ] + assert captured["kwargs"]["lora_request"] is None + + +# --------------------------------------------------------------------------- # +# load_lora +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_load_lora_happy_path(monkeypatch): + engine = _make_lora_engine(endpoint=object()) + register, _ = _patch_discovery(monkeypatch) + + result = await engine.load_lora( + {"lora_name": "adapterA", "source": {"uri": "file:///x"}} + ) + + assert result["status"] == "success" + assert result["lora_id"] == 123 + engine.engine_client.add_lora.assert_awaited_once() + register.assert_awaited_once() + assert register.await_args.kwargs["lora_name"] == "adapterA" + assert engine.loaded_loras["adapterA"].id == 123 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("collision", ["base-model", "/models/base"]) +async def test_load_lora_rejects_base_model_name_collision(monkeypatch, collision): + # An adapter named after the base model would shadow its frontend model key + # and route plain base-model requests through the adapter. Reject it before + # touching the download/engine/discovery path. + engine = _make_lora_engine(endpoint=object()) + manager = SimpleNamespace(download_lora=AsyncMock()) + register, _ = _patch_discovery(monkeypatch, manager=manager) + + result = await engine.load_lora( + {"lora_name": collision, "source": {"uri": "file:///x"}} + ) + + assert result["status"] == "error" + assert "collides" in result["message"] + manager.download_lora.assert_not_awaited() + engine.engine_client.add_lora.assert_not_awaited() + register.assert_not_awaited() + assert collision not in engine.loaded_loras + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mode, expected_model_type, expected_worker_type, expected_needs", + [ + ( + DisaggregationMode.AGGREGATED, + ModelType.Chat | ModelType.Completions, + WorkerType.Aggregated, + [], + ), + ( + DisaggregationMode.DECODE, + ModelType.Chat | ModelType.Completions, + WorkerType.Decode, + [[WorkerType.Prefill]], + ), + ( + DisaggregationMode.PREFILL, + ModelType.Prefill, + WorkerType.Prefill, + [[WorkerType.Decode]], + ), + ], +) +async def test_load_lora_publishes_disagg_topology( + monkeypatch, mode, expected_model_type, expected_worker_type, expected_needs +): + # The LoRA MDC must match the base-model registration topology so the + # frontend builds the pipeline against the right component. A prefill worker + # publishing the adapter as decode-capable chat/completions would make the + # frontend route chat traffic straight to prefill. + engine = _make_lora_engine(endpoint=object()) + engine.disaggregation_mode = mode + engine._dyn_tool_call_parser = "hermes" + engine._dyn_reasoning_parser = "deepseek_r1" + register, _ = _patch_discovery(monkeypatch) + + result = await engine.load_lora( + {"lora_name": "adapterA", "source": {"uri": "file:///x"}} + ) + + assert result["status"] == "success" + kwargs = register.await_args.kwargs + # ModelType is a Rust bitflags pyclass without Python __eq__, so combined + # flags (Chat | Completions) compare by identity and never match a freshly + # built copy. Compare the deterministic str() form instead. + assert str(kwargs["model_type"]) == str(expected_model_type) + assert kwargs["worker_type"] == expected_worker_type + assert kwargs["needs"] == expected_needs + + +@pytest.mark.asyncio +async def test_load_lora_publishes_dp_and_capacity_metadata(monkeypatch): + # The LoRA MDC must carry the worker's real DP-rank range and capacity (the + # same effective vLLM metadata the base-model card publishes), so multi-DP + # LoRA requests are routed/attributed per rank instead of as if every worker + # only served rank 0. + engine = _make_lora_engine(endpoint=object()) + engine._vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace( + max_num_seqs=256, + max_num_batched_tokens=8192, + ) + ) + engine._dp_range = (2, 4) + engine._total_kv_blocks = 1024 + register, _ = _patch_discovery(monkeypatch) + + result = await engine.load_lora( + {"lora_name": "adapterA", "source": {"uri": "file:///x"}} + ) + + assert result["status"] == "success" + runtime_config = register.await_args.kwargs["runtime_config"] + assert runtime_config.total_kv_blocks == 1024 + assert runtime_config.max_num_seqs == 256 + assert runtime_config.max_num_batched_tokens == 8192 + assert runtime_config.data_parallel_start_rank == 2 + assert runtime_config.data_parallel_size == 4 + + +@pytest.mark.asyncio +async def test_load_lora_idempotent_reload(monkeypatch): + engine = _make_lora_engine(endpoint=object()) + manager = SimpleNamespace(download_lora=AsyncMock()) + register, _ = _patch_discovery(monkeypatch, manager=manager) + # A healthy adapter is both loaded into the engine AND published to + # discovery; re-loading it must short-circuit without re-publishing. + engine.loaded_loras = {"adapterA": LoRAInfo(id=123, path="/cache/adapter")} + engine._published_loras = {"adapterA"} + + result = await engine.load_lora( + {"lora_name": "adapterA", "source": {"uri": "file:///x"}} + ) + + assert result["status"] == "success" + assert "already loaded" in result["message"] + manager.download_lora.assert_not_awaited() + engine.engine_client.add_lora.assert_not_awaited() + register.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_load_lora_reconciles_loaded_but_unpublished(monkeypatch): + # Simulate a sticky partial-failure state: the adapter is loaded into the + # engine but its discovery card was never published (a prior publish + its + # rollback both failed). A retried load must re-publish rather than report + # early success. + engine = _make_lora_engine(endpoint=object()) + manager = SimpleNamespace(download_lora=AsyncMock()) + register, _ = _patch_discovery(monkeypatch, manager=manager) + engine.loaded_loras = {"adapterA": LoRAInfo(id=123, path="/cache/adapter")} + # _published_loras intentionally left empty -> loaded-but-unpublished. + + result = await engine.load_lora( + {"lora_name": "adapterA", "source": {"uri": "file:///x"}} + ) + + assert result["status"] == "success" + # Engine load is not repeated, but discovery publication is reconciled. + manager.download_lora.assert_not_awaited() + engine.engine_client.add_lora.assert_not_awaited() + register.assert_awaited_once() + assert "adapterA" in engine._published_loras + + +@pytest.mark.asyncio +async def test_load_lora_reconcile_publish_failure_surfaces_error(monkeypatch): + engine = _make_lora_engine(endpoint=object()) + register, _ = _patch_discovery(monkeypatch) + register.side_effect = RuntimeError("discovery is down") + engine.loaded_loras = {"adapterA": LoRAInfo(id=123, path="/cache/adapter")} + + result = await engine.load_lora( + {"lora_name": "adapterA", "source": {"uri": "file:///x"}} + ) + + # Reconciliation tried and failed: report error (not a false "already + # loaded" success) and leave the adapter unpublished so a later retry can + # reconcile again. + assert result["status"] == "error" + assert "discovery publish failed" in result["message"] + assert "adapterA" not in engine._published_loras + + +@pytest.mark.asyncio +async def test_load_lora_marks_adapter_published(monkeypatch): + engine = _make_lora_engine(endpoint=object()) + _patch_discovery(monkeypatch) + + await engine.load_lora({"lora_name": "adapterA", "source": {"uri": "file:///x"}}) + + assert "adapterA" in engine._published_loras + + +@pytest.mark.asyncio +async def test_load_lora_rollback_failure_leaves_adapter_unpublished(monkeypatch): + # register fails AND the engine-side remove_lora rollback also fails: the + # adapter stays in loaded_loras but must NOT be marked published, so a + # subsequent load reconciles instead of short-circuiting. + engine = _make_lora_engine(endpoint=object()) + register, _ = _patch_discovery(monkeypatch) + register.side_effect = RuntimeError("discovery is down") + engine.engine_client.remove_lora.side_effect = RuntimeError("engine wedged") + + result = await engine.load_lora( + {"lora_name": "adapterA", "source": {"uri": "file:///x"}} + ) + + assert result["status"] == "error" + assert "adapterA" not in engine._published_loras + + +@pytest.mark.asyncio +async def test_load_lora_errors_when_manager_missing(monkeypatch): + engine = _make_lora_engine(endpoint=object()) + monkeypatch.setattr(llm_engine_mod, "get_lora_manager", lambda: None) + + result = await engine.load_lora( + {"lora_name": "adapterA", "source": {"uri": "file:///x"}} + ) + + assert result["status"] == "error" + assert "LoRAManager not initialized" in result["message"] + engine.engine_client.add_lora.assert_not_awaited() + assert "adapterA" not in engine.loaded_loras + + +@pytest.mark.asyncio +async def test_load_lora_rolls_back_when_register_fails(monkeypatch): + engine = _make_lora_engine(endpoint=object()) + register, _ = _patch_discovery(monkeypatch) + register.side_effect = RuntimeError("discovery is down") + + result = await engine.load_lora( + {"lora_name": "adapterA", "source": {"uri": "file:///x"}} + ) + + assert result["status"] == "error" + assert "Failed to register" in result["message"] + # Rollback removes the adapter from the engine and tracking. + engine.engine_client.remove_lora.assert_awaited_once_with(123) + assert "adapterA" not in engine.loaded_loras + + +# --------------------------------------------------------------------------- # +# unload_lora +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_unload_lora_happy_path(monkeypatch): + engine = _make_lora_engine(endpoint=object()) + _, unregister = _patch_discovery(monkeypatch) + # Healthy adapter: loaded into the engine AND published to discovery. + engine.loaded_loras = {"adapterA": LoRAInfo(id=123, path="/cache/adapter")} + engine._published_loras = {"adapterA"} + + result = await engine.unload_lora({"lora_name": "adapterA"}) + + assert result["status"] == "success" + engine.engine_client.remove_lora.assert_awaited_once_with(123) + unregister.assert_awaited_once() + assert "adapterA" not in engine.loaded_loras + assert "adapterA" not in engine._published_loras + + +@pytest.mark.asyncio +async def test_unload_lora_unregisters_before_engine_removal(monkeypatch): + # Discovery must be unpublished before the engine drops the adapter, so the + # frontend stops routing LoRA traffic before the adapter disappears. + engine = _make_lora_engine(endpoint=object()) + _, unregister = _patch_discovery(monkeypatch) + engine.loaded_loras = {"adapterA": LoRAInfo(id=123, path="/cache/adapter")} + engine._published_loras = {"adapterA"} + + order: list[str] = [] + unregister.side_effect = lambda **_: order.append("unregister") + engine.engine_client.remove_lora.side_effect = lambda *_: order.append("remove") + + result = await engine.unload_lora({"lora_name": "adapterA"}) + + assert result["status"] == "success" + assert order == ["unregister", "remove"] + + +@pytest.mark.asyncio +async def test_unload_lora_loaded_but_unpublished_skips_unregister(monkeypatch): + # Loaded-but-unpublished adapter (a prior load's publish failed): unload + # should not attempt an unregister, just drop it from the engine. + engine = _make_lora_engine(endpoint=object()) + _, unregister = _patch_discovery(monkeypatch) + engine.loaded_loras = {"adapterA": LoRAInfo(id=123, path="/cache/adapter")} + # _published_loras intentionally empty. + + result = await engine.unload_lora({"lora_name": "adapterA"}) + + assert result["status"] == "success" + unregister.assert_not_awaited() + engine.engine_client.remove_lora.assert_awaited_once_with(123) + assert "adapterA" not in engine.loaded_loras + + +@pytest.mark.asyncio +async def test_unload_lora_not_found(monkeypatch): + engine = _make_lora_engine(endpoint=object()) + _patch_discovery(monkeypatch) + + result = await engine.unload_lora({"lora_name": "nope"}) + + assert result["status"] == "error" + assert "not found" in result["message"] + engine.engine_client.remove_lora.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_unload_lora_happy_path_clears_published(monkeypatch): + engine = _make_lora_engine(endpoint=object()) + _, unregister = _patch_discovery(monkeypatch) + engine.loaded_loras = {"adapterA": LoRAInfo(id=123, path="/cache/adapter")} + engine._published_loras = {"adapterA"} + + result = await engine.unload_lora({"lora_name": "adapterA"}) + + assert result["status"] == "success" + unregister.assert_awaited_once() + assert "adapterA" not in engine._published_loras + + +@pytest.mark.asyncio +async def test_unload_lora_reconciles_stale_published_card(monkeypatch): + # Adapter is gone from the engine but still has a stale discovery card (a + # prior unload's unregister + re-add rollback both failed). A retried + # unload must retry the unregister so discovery converges. + engine = _make_lora_engine(endpoint=object()) + _, unregister = _patch_discovery(monkeypatch) + engine.loaded_loras = {} + engine._published_loras = {"adapterA"} + + result = await engine.unload_lora({"lora_name": "adapterA"}) + + assert result["status"] == "success" + unregister.assert_awaited_once() + engine.engine_client.remove_lora.assert_not_awaited() + assert "adapterA" not in engine._published_loras + + +@pytest.mark.asyncio +async def test_unload_lora_unregister_failure_leaves_state_intact(monkeypatch): + # Unregister runs first; if it fails, nothing has been mutated yet, so the + # adapter stays both loaded and published (consistent and still routable) + # and the engine is never touched. + engine = _make_lora_engine(endpoint=object()) + _, unregister = _patch_discovery(monkeypatch) + unregister.side_effect = RuntimeError("discovery is down") + engine.loaded_loras = {"adapterA": LoRAInfo(id=123, path="/cache/adapter")} + engine._published_loras = {"adapterA"} + + result = await engine.unload_lora({"lora_name": "adapterA"}) + + assert result["status"] == "error" + assert "Failed to unregister" in result["message"] + engine.engine_client.remove_lora.assert_not_awaited() + assert "adapterA" in engine.loaded_loras + assert "adapterA" in engine._published_loras + + +@pytest.mark.asyncio +async def test_unload_lora_engine_removal_failure_after_unregister(monkeypatch): + # Unregister succeeds (card removed, discarded from published) but the + # engine removal fails: the adapter stays loaded-but-unpublished so a + # retried unload skips the unregister and retries only the engine removal. + engine = _make_lora_engine(endpoint=object()) + _, unregister = _patch_discovery(monkeypatch) + engine.engine_client.remove_lora.side_effect = RuntimeError("engine wedged") + engine.loaded_loras = {"adapterA": LoRAInfo(id=123, path="/cache/adapter")} + engine._published_loras = {"adapterA"} + + result = await engine.unload_lora({"lora_name": "adapterA"}) + + assert result["status"] == "error" + assert "Failed to remove" in result["message"] + unregister.assert_awaited_once() + assert "adapterA" not in engine._published_loras + assert "adapterA" in engine.loaded_loras + + +# --------------------------------------------------------------------------- # +# list_loras +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_list_loras_reports_loaded_adapters(): + engine = _make_lora_engine() + engine.loaded_loras = { + "a": LoRAInfo(id=1, path="/a"), + "b": LoRAInfo(id=2, path="/b"), + } + + result = await engine.list_loras({}) + + assert result["status"] == "success" + assert result["loras"] == {"a": 1, "b": 2} + assert result["count"] == 2 + + +# --------------------------------------------------------------------------- # +# on_endpoint_ready handoff +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_on_endpoint_ready_stashes_endpoint_for_publishing(monkeypatch): + engine = _make_lora_engine(endpoint=None) + register, _ = _patch_discovery(monkeypatch) + + sentinel = object() + await engine.on_endpoint_ready(sentinel) + assert engine._endpoint is sentinel + + # load_lora now publishes against the stashed endpoint. + result = await engine.load_lora( + {"lora_name": "adapterA", "source": {"uri": "file:///x"}} + ) + assert result["status"] == "success" + register.assert_awaited_once() + assert register.await_args.kwargs["endpoint"] is sentinel + + +@pytest.mark.asyncio +async def test_load_lora_skips_publish_without_endpoint(monkeypatch): + engine = _make_lora_engine(endpoint=None) + register, _ = _patch_discovery(monkeypatch) + + result = await engine.load_lora( + {"lora_name": "adapterA", "source": {"uri": "file:///x"}} + ) + + # Adapter still loads into the engine; discovery publish is skipped. + assert result["status"] == "success" + engine.engine_client.add_lora.assert_awaited_once() + register.assert_not_awaited() + assert engine.loaded_loras["adapterA"].id == 123 + + +@pytest.mark.asyncio +async def test_cleanup_clears_endpoint_and_lora_state(): + engine = _make_lora_engine(endpoint=object()) + engine.engine_client.shutdown = MagicMock() + engine.loaded_loras = {"adapterA": LoRAInfo(id=7, path="/path/a")} + engine._published_loras = {"adapterA"} + + await engine.cleanup() + + # A shut-down engine must not retain a dangling endpoint reference or any + # stale adapter bookkeeping. The fixed stripe locks are process state, not + # per-adapter, so they are left intact. + assert engine._endpoint is None + assert engine.loaded_loras == {} + assert engine._published_loras == set() + + +def test_get_lora_lock_is_stable_and_bounded(): + from dynamo.vllm.llm_engine import _LORA_LOCK_STRIPES + + engine = _make_lora_engine() + + # The same name always maps to the same stripe lock: this is the + # serialization invariant load/unload depends on. + assert engine._get_lora_lock("adapterA") is engine._get_lora_lock("adapterA") + + # The lock store is a fixed set of stripes, so it does not grow per distinct + # adapter name (bounded memory, no eviction needed). + for i in range(_LORA_LOCK_STRIPES * 4): + engine._get_lora_lock(f"adapter-{i}") + assert len(engine._lora_load_locks) == _LORA_LOCK_STRIPES diff --git a/components/src/dynamo/vllm/tests/test_vllm_unit.py b/components/src/dynamo/vllm/tests/test_vllm_unit.py index ef23dab3c4e1..2721c189fe5f 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_unit.py +++ b/components/src/dynamo/vllm/tests/test_vllm_unit.py @@ -345,6 +345,8 @@ def test_unified_from_args_applies_rl_logprobs_default(monkeypatch): model="Qwen/Qwen3-0.6B", disaggregation_mode=CommonDisaggregationMode.AGGREGATED, component="backend", + dyn_tool_call_parser=None, + dyn_reasoning_parser=None, ) worker_config = object() diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index 1900c695583c..fe90dd46a494 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/components/src/dynamo/vllm/worker_factory.py @@ -892,11 +892,13 @@ def register_engine_routes( Args: runtime: The DistributedRuntime instance to register routes on. """ - runtime.register_engine_route("start_profile", handler.start_profile) - runtime.register_engine_route("stop_profile", handler.stop_profile) - runtime.register_engine_route("sleep", handler.sleep) - runtime.register_engine_route("wake_up", handler.wake_up) - runtime.register_engine_route("scale_elastic_ep", handler.scale_elastic_ep) + runtime.register_engine_route("control/start_profile", handler.start_profile) + runtime.register_engine_route("control/stop_profile", handler.stop_profile) + runtime.register_engine_route("control/sleep", handler.sleep) + runtime.register_engine_route("control/wake_up", handler.wake_up) + runtime.register_engine_route( + "control/scale_elastic_ep", handler.scale_elastic_ep + ) rl_routes: dict = { "liveness_probe": handler.liveness_probe, @@ -931,8 +933,9 @@ async def unload_lora(body: dict) -> dict: ) logger.info( - "Registered engine routes: sleep, wake_up, scale_elastic_ep, " - "start_profile, stop_profile, and RL admin routes: %s%s", + "Registered engine routes: control/sleep, control/wake_up, " + "control/scale_elastic_ep, control/start_profile, control/stop_profile, " + "and RL admin routes: %s%s", ", ".join(sorted(rl_routes)), " (LoRA routes: load_lora, unload_lora)" if lora_enabled else "", ) diff --git a/docs/backends/sglang/sglang-reference-guide.md b/docs/backends/sglang/sglang-reference-guide.md index 19f071cfc0e9..bb7cf5548924 100644 --- a/docs/backends/sglang/sglang-reference-guide.md +++ b/docs/backends/sglang/sglang-reference-guide.md @@ -127,15 +127,15 @@ SGLang workers expose operational endpoints via Dynamo's system server: | Route | Description | |-------|-------------| -| `/engine/start_profile` | Start PyTorch profiling | -| `/engine/stop_profile` | Stop profiling and save traces | -| `/engine/release_memory_occupation` | Release GPU memory for maintenance | -| `/engine/resume_memory_occupation` | Resume GPU memory after release | -| `/engine/update_weights_from_disk` | Update model weights from disk | -| `/engine/update_weights_from_tensor` | Update model weights from tensor payload | -| `/engine/update_weights_from_distributed` | Update model weights from distributed source | -| `/engine/update_weights_from_ipc` | Update model weights from IPC payload | -| `/engine/update_weight_version` | Update weight version metadata | +| `/engine/control/start_profile` | Start PyTorch profiling | +| `/engine/control/stop_profile` | Stop profiling and save traces | +| `/engine/control/release_memory_occupation` | Release GPU memory for maintenance | +| `/engine/control/resume_memory_occupation` | Resume GPU memory after release | +| `/engine/control/update_weights_from_disk` | Update model weights from disk | +| `/engine/control/update_weights_from_tensor` | Update model weights from tensor payload | +| `/engine/control/update_weights_from_distributed` | Update model weights from distributed source | +| `/engine/control/update_weights_from_ipc` | Update model weights from IPC payload | +| `/engine/control/update_weight_version` | Update weight version metadata | ## See Also diff --git a/docs/components/profiler/profiler-examples.md b/docs/components/profiler/profiler-examples.md index ac57bd3b7e64..d786b7c7a2c3 100644 --- a/docs/components/profiler/profiler-examples.md +++ b/docs/components/profiler/profiler-examples.md @@ -189,14 +189,14 @@ Profile SGLang workers at runtime via HTTP endpoints: ```bash # Start profiling -curl -X POST http://localhost:9090/engine/start_profile \ +curl -X POST http://localhost:9090/engine/control/start_profile \ -H "Content-Type: application/json" \ -d '{"output_dir": "/tmp/profiler_output"}' # Run inference requests to generate profiling data... # Stop profiling -curl -X POST http://localhost:9090/engine/stop_profile +curl -X POST http://localhost:9090/engine/control/stop_profile ``` A test script is provided at `examples/backends/sglang/test_sglang_profile.py`: diff --git a/docs/components/profiler/profiler-guide.md b/docs/components/profiler/profiler-guide.md index c184a0ba51e6..41cdb98398ed 100644 --- a/docs/components/profiler/profiler-guide.md +++ b/docs/components/profiler/profiler-guide.md @@ -641,14 +641,14 @@ SGLang workers expose profiling endpoints for runtime performance analysis: ```bash # Start profiling -curl -X POST http://localhost:9090/engine/start_profile \ +curl -X POST http://localhost:9090/engine/control/start_profile \ -H "Content-Type: application/json" \ -d '{"output_dir": "/tmp/profiler_output"}' # Run inference requests... # Stop profiling -curl -X POST http://localhost:9090/engine/stop_profile +curl -X POST http://localhost:9090/engine/control/stop_profile ``` View traces using Chrome's `chrome://tracing`, [Perfetto UI](https://ui.perfetto.dev/), or TensorBoard. diff --git a/examples/backends/sglang/test_sglang_profile.py b/examples/backends/sglang/test_sglang_profile.py index 75ce367ee9b0..923744c1f088 100644 --- a/examples/backends/sglang/test_sglang_profile.py +++ b/examples/backends/sglang/test_sglang_profile.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """ -Test script for /engine/start_profile and /engine/stop_profile routes. +Test script for /engine/control/start_profile and /engine/control/stop_profile routes. This script demonstrates the new custom engine route registration feature. It starts a simple sglang server with dynamo and tests the profiling endpoints. @@ -146,17 +146,17 @@ def start_sglang_backend(): def test_profiling_endpoints(): - """Test the /engine/start_profile and /engine/stop_profile endpoints""" + """Test the /engine/control/start_profile and /engine/control/stop_profile endpoints""" base_url = f"http://{HOST}:{SYSTEM_PORT}" print("\n" + "=" * 60) - print("Testing /engine/start_profile and /engine/stop_profile") + print("Testing /engine/control/start_profile and /engine/control/stop_profile") print("=" * 60) # Test 1: Start profiling with parameters (no num_steps so we control stop manually) print("\n1. Starting profiling with parameters...") response = requests.post( - f"{base_url}/engine/start_profile", + f"{base_url}/engine/control/start_profile", json={ "output_dir": PROFILER_OUTPUT_DIR, "activities": ["CPU", "GPU"], @@ -196,7 +196,7 @@ def test_profiling_endpoints(): # Test 2: Stop profiling print("\n4. Stopping profiling...") - response = requests.post(f"{base_url}/engine/stop_profile") + response = requests.post(f"{base_url}/engine/control/stop_profile") print(f" Status: {response.status_code}") print(f" Response: {response.json()}") assert response.status_code == 200, f"Expected 200, got {response.status_code}" @@ -204,7 +204,7 @@ def test_profiling_endpoints(): # Test 3: Test with empty body (GET-like POST) print("\n5. Starting profiling with empty body...") - response = requests.post(f"{base_url}/engine/start_profile") + response = requests.post(f"{base_url}/engine/control/start_profile") print(f" Status: {response.status_code}") print(f" Response: {response.json()}") assert response.status_code == 200, f"Expected 200, got {response.status_code}" @@ -217,7 +217,7 @@ def test_profiling_endpoints(): assert response.status_code == 404, f"Expected 404, got {response.status_code}" # Stop profiling again - response = requests.post(f"{base_url}/engine/stop_profile") + response = requests.post(f"{base_url}/engine/control/stop_profile") print("\n" + "=" * 60) print("✓ All tests passed!") diff --git a/fern/components/profiler/profiler_guide.md b/fern/components/profiler/profiler_guide.md index 23c4268e108f..8ef9c4111b8c 100644 --- a/fern/components/profiler/profiler_guide.md +++ b/fern/components/profiler/profiler_guide.md @@ -488,14 +488,14 @@ SGLang workers expose profiling endpoints for runtime performance analysis: ```bash # Start profiling -curl -X POST http://localhost:9090/engine/start_profile \ +curl -X POST http://localhost:9090/engine/control/start_profile \ -H "Content-Type: application/json" \ -d '{"output_dir": "/tmp/profiler_output"}' # Run inference requests... # Stop profiling -curl -X POST http://localhost:9090/engine/stop_profile +curl -X POST http://localhost:9090/engine/control/stop_profile ``` View traces using Chrome's `chrome://tracing`, [Perfetto UI](https://ui.perfetto.dev/), or TensorBoard. diff --git a/lib/backend-common/CLAUDE.md b/lib/backend-common/CLAUDE.md index 327d3335c276..85809da15ab4 100644 --- a/lib/backend-common/CLAUDE.md +++ b/lib/backend-common/CLAUDE.md @@ -29,7 +29,7 @@ parse args, start engine, wire Prometheus serve requests pre-cleanup sh return engine return metadata (optional) (concurrent) drain release ``` -The trait has six methods. `from_args` is NOT on the trait — each +The trait has twelve methods. `from_args` is NOT on the trait — each backend exposes a backend-specific constructor (typically a sync `from_args(argv) -> Result<(Self, WorkerConfig)>` inherent method). This keeps the trait fully object-safe without a `where Self: Sized` @@ -87,6 +87,38 @@ opt-out and lets `run.rs` stay non-generic. successful first returns `Ok(())` without re-entering teardown. The conformance kit pins both — `CleanupWithoutStartFailed` and `SecondCleanupFailed`. +- `health_check_payload(&self) -> Result, DynamoError>` — + optional, default `Ok(None)`. Canary payload the runtime sends through + `generate` to actively probe an idle endpoint; `None` disables active + probing. Operator overrides (`DYN_HEALTH_CHECK_PAYLOAD` / `WorkerConfig`) + take precedence. +- `supported_controls(&self) -> Result, DynamoError>` — + optional, default empty. Semantic engine-control keys this engine + advertises (e.g. `start_profile`, `sleep`, `wake_up`). The Worker maps + each onto a `/engine/control/{key}` route via `register_engine_controls`. +- `engine_control(&self, control: String, body: Value) -> Result` — + optional, default returns a `status:"error"` body. Dispatches one + advertised control. Returning a `status:"error"` body is HTTP 200 at the + `/engine/*` layer (it 5xx's only when this *raises*). +- `supported_updates(&self) -> Result, DynamoError>` — + optional, default empty. A sibling surface to `supported_controls` for + ops that mutate engine-managed assets rather than the serving lifecycle + (e.g. vLLM dynamic LoRA `load_lora` / `unload_lora` / `list_loras`). Kept + separate so LoRA doesn't inflate the control surface. The Worker maps + each onto a `/engine/update/{key}` route via `register_engine_updates` + (no quiesce/resume policy — updates never toggle discovery). +- `engine_update(&self, update: String, body: Value) -> Result` — + optional, default returns a `status:"error"` body. Dispatches one + advertised update; same HTTP-200-on-`status:"error"` semantics as + `engine_control`. +- `on_endpoint_ready(&self, endpoint: Endpoint) -> Result<(), DynamoError>` — + optional, default no-op. The Worker hands the engine its serving + `Endpoint` exactly once, after it exists and **before** + `register_engine_controls` / `register_engine_updates` (so `/engine/*` + can't fire before the engine has the endpoint). A failure is **fatal to + startup**. Engines that publish their own discovery records stash it + (e.g. vLLM dynamic LoRA via `register_model` / `unregister_model`). + Mirrors the `on_publisher_ready` handoff idiom. ## Contract for `generate` diff --git a/lib/backend-common/Cargo.toml b/lib/backend-common/Cargo.toml index f8034a51b2ac..5460a4a70ac4 100644 --- a/lib/backend-common/Cargo.toml +++ b/lib/backend-common/Cargo.toml @@ -17,6 +17,9 @@ autoexamples = false # Enables the conformance test kit (testing.rs). Intended for `[dev-dependencies]` # in backend crates — adds no weight to production binaries. testing = [] +# Enables NATS-backed integration tests (e.g. the on_endpoint_ready handoff). +# Forwards to dynamo-runtime's integration feature for `create_test_drt_async`. +integration = ["dynamo-runtime/integration"] [dependencies] dynamo-runtime = { workspace = true } diff --git a/lib/backend-common/src/engine.rs b/lib/backend-common/src/engine.rs index 3ef581e18fb7..d9a5904752ec 100644 --- a/lib/backend-common/src/engine.rs +++ b/lib/backend-common/src/engine.rs @@ -343,6 +343,45 @@ pub trait LLMEngine: Send + Sync + 'static { "message": format!("unsupported engine control: {control}"), })) } + + /// Semantic engine updates this engine supports. Empty by default. + /// + /// Updates are a sibling surface to [`supported_controls`](LLMEngine::supported_controls) + /// for operations that mutate engine-managed assets (e.g. vLLM dynamic + /// LoRA load/unload/list) rather than the engine's serving lifecycle. + /// Keeping them separate avoids inflating the control surface. Engines + /// advertise update keys and implement them via [`LLMEngine::engine_update`]; + /// the unified backend maps each key onto an `/engine/update/{key}` route. + async fn supported_updates(&self) -> Result, DynamoError> { + Ok(Vec::new()) + } + + /// Handle one semantic engine-update request. + async fn engine_update( + &self, + update: String, + _body: serde_json::Value, + ) -> Result { + Ok(serde_json::json!({ + "status": "error", + "message": format!("unsupported engine update: {update}"), + })) + } + + /// Hand the engine its runtime serving [`Endpoint`](dynamo_runtime::component::Endpoint), + /// exactly once, after it exists and before serving begins. Default no-op. + /// + /// Engines that publish their own discovery records (e.g. vLLM dynamic + /// LoRA via `register_model`) stash it here for later use from + /// [`engine_update`](LLMEngine::engine_update). Mirrors the + /// [`on_publisher_ready`](MetricsBindings::on_publisher_ready) handoff idiom. + /// Errors abort startup; `cleanup` runs on the partial state. + async fn on_endpoint_ready( + &self, + _endpoint: dynamo_runtime::component::Endpoint, + ) -> Result<(), DynamoError> { + Ok(()) + } } /// Raw media-generation engine trait — the non-token sibling of [`LLMEngine`]. diff --git a/lib/backend-common/src/worker.rs b/lib/backend-common/src/worker.rs index 10d2ae527be8..64a6d9237158 100644 --- a/lib/backend-common/src/worker.rs +++ b/lib/backend-common/src/worker.rs @@ -295,6 +295,39 @@ impl EngineKind { } } + async fn supported_updates(&self) -> Result, DynamoError> { + match self { + EngineKind::Llm(e) => e.supported_updates().await, + // Raw media engines advertise no semantic engine updates. + EngineKind::Raw(_) => Ok(Vec::new()), + } + } + + async fn engine_update( + &self, + update: String, + body: serde_json::Value, + ) -> Result { + match self { + EngineKind::Llm(e) => e.engine_update(update, body).await, + EngineKind::Raw(_) => Ok(serde_json::json!({ + "status": "error", + "message": format!("unsupported engine update: {update}"), + })), + } + } + + async fn on_endpoint_ready( + &self, + endpoint: dynamo_runtime::component::Endpoint, + ) -> Result<(), DynamoError> { + match self { + EngineKind::Llm(e) => e.on_endpoint_ready(endpoint).await, + // Raw media engines publish no discovery records of their own. + EngineKind::Raw(_) => Ok(()), + } + } + /// Raw media engines (image/video/audio) register name-only — the engine /// loads the model itself and the model has no LLM artifacts (tokenizer / /// chat template / config.json) for Dynamo to fetch. @@ -639,12 +672,41 @@ impl Worker { endpoint.clone(), control_lock.clone(), ); - registry.register(&control_name, callback); + // Namespace control routes under `/engine/control/` so they + // share the `/engine/{*path}` route without colliding with updates. + registry.register(&format!("control/{control_name}"), callback); } tracing::info!(control_count, "registered engine management controls"); Ok(()) } + /// Register advertised engine updates on the runtime system server. + /// + /// Updates are a sibling surface to controls for operations that mutate + /// engine-managed assets (e.g. vLLM dynamic LoRA). They register under + /// `/engine/update/` and, unlike controls, never toggle discovery + /// registration — so no quiesce/resume policy wrapper or serialization lock. + async fn register_engine_updates( + &self, + endpoint: &dynamo_runtime::component::Endpoint, + ) -> Result<(), DynamoError> { + let updates = self.engine.supported_updates().await?; + if updates.is_empty() { + tracing::debug!("engine returned no management updates"); + return Ok(()); + } + + let registry = endpoint.drt().engine_routes(); + let update_count = updates.len(); + for update_name in updates { + let callback = engine_update_callback(update_name.clone(), self.engine.clone()); + // Namespace update routes under `/engine/update/`. + registry.register(&format!("update/{update_name}"), callback); + } + tracing::info!(update_count, "registered engine management updates"); + Ok(()) + } + /// Full graceful-shutdown orchestrator: discovery unregister → /// grace period → engine drain → cleanup. Shared by every shutdown path — /// pre-serve (mid-start signal) and the serve loop's signal arm. @@ -734,6 +796,16 @@ impl Worker { let mut local_model = build_local_model(&self.config, engine_config, self.engine.is_raw()).await?; tracing::debug!("local model built"); + + // Hand the engine its serving endpoint before registering the model + // with discovery. on_endpoint_ready is a fatal handoff: doing it first + // means a failure leaves nothing published, so there is no stale + // discovery entry to reclaim. Engines that publish their own discovery + // records (e.g. vLLM dynamic LoRA) stash the endpoint here, and this + // still runs before `register_engine_controls`, so `/engine/*` cannot + // fire before the engine has the endpoint. + self.engine.on_endpoint_ready(endpoint.clone()).await?; + local_model .attach( &endpoint, @@ -751,7 +823,9 @@ impl Worker { ) })?; tracing::debug!("model registered with discovery"); + self.register_engine_controls(&endpoint).await?; + self.register_engine_updates(&endpoint).await?; let served = resolve_served_name(&self.config, engine_config) .unwrap_or_else(|| engine_config.model.clone()); @@ -1096,6 +1170,16 @@ fn control_request_body_error(body: &serde_json::Value) -> Option Option { + if body.is_object() { + None + } else { + Some(control_error_response( + "engine update request body must be a JSON object", + )) + } +} + fn engine_control_callback(control_name: String, engine: EngineKind) -> EngineRouteCallback { Arc::new(move |body| { let engine = engine.clone(); @@ -1109,6 +1193,22 @@ fn engine_control_callback(control_name: String, engine: EngineKind) -> EngineRo }) } +fn engine_update_callback(update_name: String, engine: EngineKind) -> EngineRouteCallback { + Arc::new(move |body| { + let engine = engine.clone(); + let update_name = update_name.clone(); + Box::pin(async move { + if let Some(response) = update_request_body_error(&body) { + return Ok(response); + } + engine + .engine_update(update_name, body) + .await + .map_err(|e| anyhow::anyhow!(e.to_string())) + }) + }) +} + fn wrap_engine_control_callback( control_name: String, callback: EngineRouteCallback, @@ -1135,7 +1235,7 @@ fn wrap_engine_control_callback( if let Err(e) = endpoint.unregister_endpoint_instance().await { return Ok(control_error_response(format!( - "failed to unregister endpoint before /engine/{control_name}: {e}" + "failed to unregister endpoint before /engine/control/{control_name}: {e}" ))); } @@ -1169,12 +1269,12 @@ fn wrap_engine_control_callback( && let Err(e) = endpoint.register_endpoint_instance().await { // The engine is serving-safe but absent from discovery. The - // operation is idempotent: retrying /engine/{control_name} + // operation is idempotent: retrying /engine/control/{control_name} // re-registers without repeating the wake/resume work (the // controller short-circuits "already awake/resumed"), so surface // that it is safe to retry. return Ok(control_error_response(format!( - "engine resumed but re-registration failed after /engine/{control_name}: {e}; retry /engine/{control_name} to rejoin discovery" + "engine resumed but re-registration failed after /engine/control/{control_name}: {e}; retry /engine/control/{control_name} to rejoin discovery" ))); } Ok(response) @@ -1532,6 +1632,26 @@ mod tests { } } + #[test] + fn update_request_body_validation_requires_json_object() { + assert!(update_request_body_error(&serde_json::json!({})).is_none()); + assert!(update_request_body_error(&serde_json::json!({"lora_name": "a"})).is_none()); + + for body in [ + serde_json::json!(null), + serde_json::json!(true), + serde_json::json!("bad"), + serde_json::json!(["lora_name"]), + ] { + let response = update_request_body_error(&body).unwrap(); + assert!(control_response_is_error(&response)); + assert_eq!( + response.get("message").and_then(|value| value.as_str()), + Some("engine update request body must be a JSON object") + ); + } + } + #[test] fn control_response_error_detection_matches_backend_conventions() { assert!(control_response_is_error(&serde_json::json!({ @@ -2292,3 +2412,239 @@ mod tests { ); } } + +// Integration tests for the `on_endpoint_ready` handoff. These need a real +// `DistributedRuntime`/`Endpoint` (NATS-backed), so they live behind the +// `integration` feature: +// cargo test -p dynamo-backend-common --features integration on_endpoint_ready +#[cfg(all(test, feature = "integration"))] +mod handoff_integration_tests { + use super::*; + use crate::engine::PreprocessedRequest; + use async_trait::async_trait; + use dynamo_runtime::distributed_test_utils::create_test_drt_async; + use futures::stream::BoxStream; + use std::sync::Mutex as StdMutex; + + /// Build a real serving `Endpoint` from a test DRT, mirroring how + /// `run_inner` resolves namespace → component → endpoint. + async fn test_endpoint() -> dynamo_runtime::component::Endpoint { + let drt = create_test_drt_async().await; + drt.namespace("handoff_ns") + .unwrap() + .component("handoff_comp") + .unwrap() + .endpoint("generate") + } + + /// Mock engine that records the order of `on_endpoint_ready`, + /// `supported_controls`, and `supported_updates` calls, lets a test force + /// `on_endpoint_ready` to fail, and advertises configurable control/update + /// sets. + struct HandoffMockEngine { + log: Arc>>, + endpoint_ready_should_fail: bool, + controls: Vec, + updates: Vec, + } + + impl HandoffMockEngine { + fn new( + endpoint_ready_should_fail: bool, + controls: Vec, + updates: Vec, + ) -> (Arc, Arc>>) { + let log = Arc::new(StdMutex::new(Vec::new())); + let eng = Arc::new(Self { + log: log.clone(), + endpoint_ready_should_fail, + controls, + updates, + }); + (eng, log) + } + } + + #[async_trait] + impl LLMEngine for HandoffMockEngine { + async fn start(&self, _worker_id: u64) -> Result { + Ok(EngineConfig { + model: "mock".to_string(), + ..EngineConfig::default() + }) + } + + async fn generate( + &self, + _request: PreprocessedRequest, + _ctx: crate::engine::GenerateContext, + ) -> Result< + BoxStream<'static, Result>, + DynamoError, + > { + unreachable!("not used in handoff tests") + } + + async fn cleanup(&self) -> Result<(), DynamoError> { + Ok(()) + } + + async fn supported_controls(&self) -> Result, DynamoError> { + self.log.lock().unwrap().push("supported_controls"); + Ok(self.controls.clone()) + } + + async fn supported_updates(&self) -> Result, DynamoError> { + self.log.lock().unwrap().push("supported_updates"); + Ok(self.updates.clone()) + } + + async fn on_endpoint_ready( + &self, + _endpoint: dynamo_runtime::component::Endpoint, + ) -> Result<(), DynamoError> { + self.log.lock().unwrap().push("on_endpoint_ready"); + if self.endpoint_ready_should_fail { + Err(err( + ErrorType::Backend(BackendError::Unknown), + "synthetic on_endpoint_ready failure", + )) + } else { + Ok(()) + } + } + } + + /// Engine that overrides only the required methods, so it inherits the + /// trait-default `on_endpoint_ready` / `supported_controls`. + struct DefaultsEngine; + + #[async_trait] + impl LLMEngine for DefaultsEngine { + async fn start(&self, _worker_id: u64) -> Result { + Ok(EngineConfig::default()) + } + + async fn generate( + &self, + _request: PreprocessedRequest, + _ctx: crate::engine::GenerateContext, + ) -> Result< + BoxStream<'static, Result>, + DynamoError, + > { + unreachable!("not used in handoff tests") + } + + async fn cleanup(&self) -> Result<(), DynamoError> { + Ok(()) + } + } + + /// The trait default `on_endpoint_ready` is a no-op that succeeds against a + /// real `Endpoint`. + #[tokio::test] + async fn default_on_endpoint_ready_is_noop() { + let endpoint = test_endpoint().await; + let engine = Arc::new(DefaultsEngine); + engine + .on_endpoint_ready(endpoint) + .await + .expect("default on_endpoint_ready must succeed"); + } + + /// `serve_with_orchestrator` runs `on_endpoint_ready` before + /// `register_engine_controls` and `register_engine_updates`. Drive the same + /// three production calls in that order and assert: (1) the handoff is + /// observed before the engine is asked for its controls/updates, and (2) the + /// advertised control lands under `control/` and the advertised update + /// under `update/` in the DRT's engine-route registry, so + /// `/engine/control/` and `/engine/update/` become routable. + #[tokio::test] + async fn handoff_precedes_registration_and_populates_namespaced_registry() { + let endpoint = test_endpoint().await; + let (engine, log) = HandoffMockEngine::new( + false, + vec!["start_profile".to_string()], + vec!["load_lora".to_string()], + ); + let worker = Worker::new(engine, WorkerConfig::default()); + + // Mirror serve_with_orchestrator's handoff + registration calls exactly. + worker + .engine + .on_endpoint_ready(endpoint.clone()) + .await + .expect("handoff should succeed"); + worker + .register_engine_controls(&endpoint) + .await + .expect("control registration should succeed"); + worker + .register_engine_updates(&endpoint) + .await + .expect("update registration should succeed"); + + let recorded = log.lock().unwrap().clone(); + assert_eq!( + recorded, + vec![ + "on_endpoint_ready", + "supported_controls", + "supported_updates" + ], + "endpoint handoff must happen before controls/updates are enumerated/registered" + ); + let routes = endpoint.drt().engine_routes(); + assert!( + routes.get("control/start_profile").is_some(), + "advertised control must be registered under control/" + ); + assert!( + routes.get("update/load_lora").is_some(), + "advertised update must be registered under update/" + ); + // Bare (unprefixed) keys must NOT be registered by the unified Worker. + assert!( + routes.get("start_profile").is_none(), + "control must not be registered under its bare name" + ); + assert!( + routes.get("load_lora").is_none(), + "update must not be registered under its bare name" + ); + } + + /// A failing `on_endpoint_ready` aborts startup: the `?` in + /// `serve_with_orchestrator` propagates the error before + /// `register_engine_controls`/`register_engine_updates` run, so nothing is + /// registered. + #[tokio::test] + async fn failed_handoff_is_fatal_and_skips_registration() { + let endpoint = test_endpoint().await; + let (engine, log) = HandoffMockEngine::new( + true, + vec!["start_profile".to_string()], + vec!["load_lora".to_string()], + ); + let worker = Worker::new(engine, WorkerConfig::default()); + + let result = worker.engine.on_endpoint_ready(endpoint.clone()).await; + assert!(result.is_err(), "failed handoff must surface as an error"); + + // Production code returns here via `?`; we do NOT call + // register_engine_controls/register_engine_updates. Confirm nothing + // was registered. + let recorded = log.lock().unwrap().clone(); + assert_eq!(recorded, vec!["on_endpoint_ready"]); + let routes = endpoint.drt().engine_routes(); + assert!( + routes.get("control/start_profile").is_none(), + "no controls should be registered after a fatal handoff" + ); + assert!( + routes.get("update/load_lora").is_none(), + "no updates should be registered after a fatal handoff" + ); + } +} diff --git a/lib/bindings/python/rust/backend.rs b/lib/bindings/python/rust/backend.rs index 460231321b2b..fb21698a43c9 100644 --- a/lib/bindings/python/rust/backend.rs +++ b/lib/bindings/python/rust/backend.rs @@ -1024,6 +1024,119 @@ impl PyEngineCore { })? .map_err(py_err_to_dynamo) } + + async fn supported_updates(&self) -> Result, DynamoError> { + let engine = self.engine.clone(); + let join = tokio::task::spawn_blocking(move || { + Python::with_gil(|py| -> PyResult> { + let result = engine.bind(py).call_method0("supported_updates")?; + let mut updates = Vec::new(); + for item in result.try_iter()? { + updates.push(item?.extract()?); + } + Ok(updates) + }) + }) + .await; + + match join { + Ok(Ok(v)) => Ok(v), + Ok(Err(e)) => Err(py_err_to_dynamo(e)), + Err(join_err) => Err(DynamoError::builder() + .error_type(ErrorType::Backend(BackendError::Unknown)) + .message(format!( + "supported_updates spawn_blocking join failed: {join_err}" + )) + .build()), + } + } + + async fn engine_update( + &self, + update: String, + body: serde_json::Value, + ) -> Result { + let engine = self.engine.clone(); + let event_loop = self.event_loop.clone(); + let py_future = tokio::task::spawn_blocking(move || { + Python::with_gil(|py| { + let py_body = pythonize(py, &body).map_err(|e| { + PyErr::new::(format!( + "Failed to convert engine update request body to Python: {e}" + )) + })?; + let coroutine = engine + .bind(py) + .call_method1("engine_update", (update, py_body))?; + let locals = TaskLocals::new(event_loop.bind(py).clone()); + pyo3_async_runtimes::into_future_with_locals(&locals, coroutine) + }) + }) + .await + .map_err(|e| { + DynamoError::builder() + .error_type(ErrorType::Backend(BackendError::Unknown)) + .message(format!("engine_update offload error: {e}")) + .build() + })? + .map_err(py_err_to_dynamo)?; + + let py_result = py_future.await.map_err(py_err_to_dynamo)?; + + tokio::task::spawn_blocking(move || { + Python::with_gil(|py| { + depythonize::(py_result.bind(py)).map_err(|e| { + PyErr::new::(format!( + "Failed to serialize engine update response: {e}" + )) + }) + }) + }) + .await + .map_err(|e| { + DynamoError::builder() + .error_type(ErrorType::Backend(BackendError::Unknown)) + .message(format!("engine_update response offload error: {e}")) + .build() + })? + .map_err(py_err_to_dynamo) + } + + async fn on_endpoint_ready( + &self, + endpoint: rs::component::Endpoint, + ) -> Result<(), DynamoError> { + let engine = self.engine.clone(); + let event_loop = self.event_loop.clone(); + + let py_future = tokio::task::spawn_blocking(move || { + Python::with_gil(|py| -> PyResult<_> { + let py_endpoint = Py::new( + py, + crate::Endpoint { + inner: endpoint, + event_loop: event_loop.bind(py).clone().unbind(), + }, + )?; + let coroutine = engine + .bind(py) + .call_method1("on_endpoint_ready", (py_endpoint,))?; + let locals = TaskLocals::new(event_loop.bind(py).clone()); + pyo3_async_runtimes::into_future_with_locals(&locals, coroutine) + }) + }) + .await + .map_err(|e| { + DynamoError::builder() + .error_type(ErrorType::Backend(BackendError::Unknown)) + .message(format!("on_endpoint_ready offload error: {e}")) + .build() + })? + .map_err(py_err_to_dynamo)?; + + py_future.await.map_err(py_err_to_dynamo)?; + Ok(()) + } } impl PyEngineCore { @@ -1216,6 +1329,25 @@ impl LLMEngine for PyLLMEngine { ) -> Result { self.core.engine_control(control, body).await } + + async fn supported_updates(&self) -> Result, DynamoError> { + self.core.supported_updates().await + } + + async fn engine_update( + &self, + update: String, + body: serde_json::Value, + ) -> Result { + self.core.engine_update(update, body).await + } + + async fn on_endpoint_ready( + &self, + endpoint: rs::component::Endpoint, + ) -> Result<(), DynamoError> { + self.core.on_endpoint_ready(endpoint).await + } } // --------------------------------------------------------------------------- diff --git a/lib/bindings/python/rust/lib.rs b/lib/bindings/python/rust/lib.rs index 162faee7d816..46d133132a16 100644 --- a/lib/bindings/python/rust/lib.rs +++ b/lib/bindings/python/rust/lib.rs @@ -1007,7 +1007,7 @@ impl DistributedRuntime { /// Register an async Python callback for /engine/{route_name} /// /// Args: - /// route_name: Route path (e.g., "start_profile" → /engine/start_profile) + /// route_name: Route path (e.g., "control/start_profile" → /engine/control/start_profile) /// callback: Async function with signature: async def(body: dict) -> dict /// /// Example: @@ -1016,7 +1016,7 @@ impl DistributedRuntime { /// await engine.start_profile(**body) /// return {"status": "ok"} /// - /// runtime.register_engine_route("start_profile", start_profile) + /// runtime.register_engine_route("control/start_profile", start_profile) /// ``` #[pyo3(signature = (route_name, callback))] fn register_engine_route( diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index f5f7d88d2c6a..616c2a7786d3 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -189,7 +189,7 @@ class DistributedRuntime: Register an async callback for /engine/{route_name} on the system status server. Args: - route_name: The route path (e.g., "start_profile" creates /engine/start_profile) + route_name: The route path (e.g., "control/start_profile" creates /engine/control/start_profile) callback: Async function with signature: async def(body: dict) -> dict Example: @@ -197,7 +197,7 @@ class DistributedRuntime: await engine.start_profile(**body) return {"status": "ok", "message": "Profiling started"} - runtime.register_engine_route("start_profile", start_profile) + runtime.register_engine_route("control/start_profile", start_profile) The callback receives the JSON request body as a dict and should return a dict that will be serialized as the JSON response. diff --git a/lib/gpu_memory_service/README.md b/lib/gpu_memory_service/README.md index 85747241414a..1852712d13cb 100644 --- a/lib/gpu_memory_service/README.md +++ b/lib/gpu_memory_service/README.md @@ -583,7 +583,7 @@ The integration patches `torch_memory_saver` to route both weight and KV-cache o Both integrations support releasing and reclaiming GPU memory for shadow engine patterns. The API names differ by framework: -- **vLLM**: `sleep` / `wake_up` (via `/engine/sleep` and `/engine/wake_up` HTTP endpoints) +- **vLLM**: `sleep` / `wake_up` (via `/engine/control/sleep` and `/engine/control/wake_up` HTTP endpoints) - **SGLang**: `release_memory_occupation` / `resume_memory_occupation` (via the corresponding HTTP endpoints) Under the hood, pausing calls `unmap_all_vas()` + `abort()` to release GPU memory while preserving VA reservations. Resuming is tag-specific: diff --git a/lib/runtime/src/engine_routes.rs b/lib/runtime/src/engine_routes.rs index c5020b75bf1b..eacb7a7b6d0b 100644 --- a/lib/runtime/src/engine_routes.rs +++ b/lib/runtime/src/engine_routes.rs @@ -33,7 +33,7 @@ impl EngineRouteRegistry { } } - /// Register a callback for a route (e.g., "start_profile" for /engine/start_profile) + /// Register a callback for a route (e.g., "control/start_profile" for /engine/control/start_profile) /// /// A route name is expected to be registered exactly once. Re-registering an /// existing name overwrites the previous callback and emits a warning, since diff --git a/lib/runtime/src/system_status_server.rs b/lib/runtime/src/system_status_server.rs index e7e6aca203a3..f0c6db7dcc66 100644 --- a/lib/runtime/src/system_status_server.rs +++ b/lib/runtime/src/system_status_server.rs @@ -526,10 +526,21 @@ async fn metadata_file_handler( } } -/// Helper function to call a LoRA management endpoint locally via in-process registry +/// Helper function to call a LoRA management endpoint for the local worker. /// -/// This function ONLY uses the local endpoint registry for direct in-process calls. -/// It does NOT fall back to network discovery if the endpoint is not found. +/// Resolution order (both are in-process, never network discovery): +/// 1. The legacy local endpoint registry, populated by non-unified workers via +/// `.register_local_engine()`. +/// 2. The generic engine-route registry (`/engine/*`). Unified-backend workers +/// advertise LoRA lifecycle ops (`load_lora`/`unload_lora`/`list_loras`) as +/// engine *updates*, registered under `update/`; this fallback maps the +/// bare LoRA name onto that key so the legacy `/v1/loras` surface forwards to +/// them (the `/v1/loras` compatibility shim). +/// +/// Because legacy workers populate the local registry, they never reach the +/// fallback — their `/v1/loras` behavior is unchanged. If neither registry +/// holds the name, returns an explicit "LoRA management not available" error +/// rather than an opaque "endpoint not found". async fn call_lora_endpoint( drt: &crate::DistributedRuntime, endpoint_name: &str, @@ -537,40 +548,50 @@ async fn call_lora_endpoint( ) -> anyhow::Result { use crate::engine::AsyncEngine; - tracing::debug!("Calling local endpoint: '{endpoint_name}'"); + tracing::debug!("Calling LoRA endpoint: '{endpoint_name}'"); - // Get the endpoint from the local registry (in-process call only) - let local_registry = drt.local_endpoint_registry(); - let engine = local_registry - .get(endpoint_name) - .ok_or_else(|| { - anyhow::anyhow!( - "Endpoint '{}' not found in local registry. Make sure it's registered with .register_local_engine()", - endpoint_name - ) - })?; + // 1. Legacy local registry (in-process call only). + if let Some(engine) = drt.local_endpoint_registry().get(endpoint_name) { + tracing::debug!( + "Found endpoint '{}' in local registry, calling directly", + endpoint_name + ); - tracing::debug!( - "Found endpoint '{}' in local registry, calling directly", - endpoint_name - ); + let request = crate::pipeline::SingleIn::new(request_body); + let mut stream = engine.generate(request).await?; - // Call the engine directly without going through the network stack - let request = crate::pipeline::SingleIn::new(request_body); - let mut stream = engine.generate(request).await?; + if let Some(response) = stream.next().await { + let response_data = response.data.unwrap_or_default(); + let lora_response = serde_json::from_value::(response_data.clone()) + .unwrap_or_else(|_| parse_lora_response(&response_data)); + return Ok(lora_response); + } - // Get the first response - if let Some(response) = stream.next().await { - let response_data = response.data.unwrap_or_default(); + anyhow::bail!("No response received from endpoint '{}'", endpoint_name) + } - // Try structured deserialization first, fall back to manual field extraction + // 2. Unified-backend engine-update registry fallback. The unified Worker + // registers LoRA ops as engine updates under `update/`, so map the + // bare LoRA endpoint name onto that namespaced key. + let update_key = format!("update/{endpoint_name}"); + if let Some(callback) = drt.engine_routes().get(&update_key) { + tracing::debug!( + "Found '{}' in engine routes registry, invoking update callback", + update_key + ); + let response_data = callback(request_body).await?; let lora_response = serde_json::from_value::(response_data.clone()) .unwrap_or_else(|_| parse_lora_response(&response_data)); - return Ok(lora_response); } - anyhow::bail!("No response received from endpoint '{}'", endpoint_name) + anyhow::bail!( + "LoRA management not available: no '{}' handler is registered \ + (neither a local LoRA endpoint nor an engine update). This worker \ + either has LoRA disabled or its backend does not support LoRA \ + management.", + endpoint_name + ) } /// Helper to parse response data into LoraResponse @@ -1243,4 +1264,94 @@ mod integration_tests { ) .await; } + + /// `/v1/loras` compat shim: with the legacy local registry empty, a LoRA + /// update registered in `engine_routes()` under `update/` resolves via + /// the fallback and its JSON response is parsed into a `LoraResponse`. This + /// is the path unified-backend workers take for `/v1/loras`. + #[tokio::test] + async fn test_call_lora_endpoint_resolves_via_engine_routes() { + temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async { + let drt = create_test_drt_async().await; + + let callback: crate::engine_routes::EngineRouteCallback = Arc::new(|_body| { + Box::pin(async move { + Ok(serde_json::json!({ + "status": "success", + "lora_name": "adapterA", + "lora_id": 42, + })) + }) + }); + // Unified Worker registers LoRA ops under the `update/` namespace. + drt.engine_routes().register("update/load_lora", callback); + + // Local registry is empty, so resolution must fall through to + // engine_routes. + assert!(drt.local_endpoint_registry().get("load_lora").is_none()); + + let response = call_lora_endpoint( + &drt, + "load_lora", + serde_json::json!({"lora_name": "adapterA"}), + ) + .await + .expect("engine_routes fallback should resolve the control"); + + assert_eq!(response.status, "success"); + assert_eq!(response.lora_name.as_deref(), Some("adapterA")); + assert_eq!(response.lora_id, Some(42)); + }) + .await; + } + + /// When neither the local registry nor `engine_routes()` holds the name, + /// the caller gets an explicit "LoRA management not available" error + /// rather than an opaque "endpoint not found". + #[tokio::test] + async fn test_call_lora_endpoint_missing_returns_clean_error() { + temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async { + let drt = create_test_drt_async().await; + + let err = call_lora_endpoint(&drt, "load_lora", serde_json::json!({})) + .await + .expect_err("missing handler must error"); + + assert!( + err.to_string().contains("LoRA management not available"), + "expected explicit unavailable message, got: {err}" + ); + }) + .await; + } + + /// An update callback that returns `{"status":"error",...}` (rather than + /// raising) surfaces as a `LoraResponse{status:"error"}`. The `/v1/loras` + /// load/unload handlers map this to HTTP 500, preserving legacy semantics; + /// the direct `/engine/*` route would instead return HTTP 200 + this JSON. + #[tokio::test] + async fn test_call_lora_endpoint_propagates_error_status() { + temp_env::async_with_vars([(env_system::DYN_SYSTEM_PORT, None::<&str>)], async { + let drt = create_test_drt_async().await; + + let callback: crate::engine_routes::EngineRouteCallback = Arc::new(|_body| { + Box::pin(async move { + Ok(serde_json::json!({ + "status": "error", + "message": "adapter not found", + })) + }) + }); + // Unified Worker registers LoRA ops under the `update/` namespace. + drt.engine_routes().register("update/unload_lora", callback); + + let response = call_lora_endpoint(&drt, "unload_lora", serde_json::json!({})) + .await + .expect("a non-raising callback returns Ok even on logical error"); + + assert_eq!(response.status, "error"); + assert_eq!(response.message.as_deref(), Some("adapter not found")); + }) + .await; + } } diff --git a/tests/fault_tolerance/deploy/templates/vllm/run_bare_multinode_elastic_ep_scale_test.sh b/tests/fault_tolerance/deploy/templates/vllm/run_bare_multinode_elastic_ep_scale_test.sh index a40ace03b6b3..540732806031 100755 --- a/tests/fault_tolerance/deploy/templates/vllm/run_bare_multinode_elastic_ep_scale_test.sh +++ b/tests/fault_tolerance/deploy/templates/vllm/run_bare_multinode_elastic_ep_scale_test.sh @@ -163,7 +163,7 @@ print('text:', repr(text), ' usage:', tokens) } # Calls vLLM's native scale endpoint on port 8000 (same port as inference). -# NOTE: bare vLLM uses /scale_elastic_ep directly, NOT /engine/scale_elastic_ep. +# NOTE: bare vLLM uses /scale_elastic_ep directly, NOT /engine/control/scale_elastic_ep. scale() { local from_dp="$1" local to_dp="$2" diff --git a/tests/fault_tolerance/deploy/templates/vllm/run_elastic_ep_scale_test.sh b/tests/fault_tolerance/deploy/templates/vllm/run_elastic_ep_scale_test.sh index cff366d9dde5..96b260e17948 100755 --- a/tests/fault_tolerance/deploy/templates/vllm/run_elastic_ep_scale_test.sh +++ b/tests/fault_tolerance/deploy/templates/vllm/run_elastic_ep_scale_test.sh @@ -139,8 +139,8 @@ scale() { echo "SCALE dp=$from_dp → dp=$to_dp at $(date -u +%Y-%m-%dT%H:%M:%SZ)" echo " worker pod: $(worker_pod)" echo "==========================================" - echo "--- request: POST /engine/scale_elastic_ep {\"new_data_parallel_size\": $to_dp} ---" - RESP=$(curl -s -X POST http://localhost:8001/engine/scale_elastic_ep \ + echo "--- request: POST /engine/control/scale_elastic_ep {\"new_data_parallel_size\": $to_dp} ---" + RESP=$(curl -s -X POST http://localhost:8001/engine/control/scale_elastic_ep \ -H "Content-Type: application/json" \ -d "{\"new_data_parallel_size\": $to_dp}" \ --max-time "$timeout") diff --git a/tests/fault_tolerance/deploy/templates/vllm/run_multinode_elastic_ep_scale_test.sh b/tests/fault_tolerance/deploy/templates/vllm/run_multinode_elastic_ep_scale_test.sh index 631f45de6d17..1f0c55817fe4 100755 --- a/tests/fault_tolerance/deploy/templates/vllm/run_multinode_elastic_ep_scale_test.sh +++ b/tests/fault_tolerance/deploy/templates/vllm/run_multinode_elastic_ep_scale_test.sh @@ -163,7 +163,7 @@ scale() { local lpod lpod=$(head_pod) RESP=$(kubectl exec "$lpod" -n "$NS" -- \ - curl -s -X POST http://localhost:9090/engine/scale_elastic_ep \ + curl -s -X POST http://localhost:9090/engine/control/scale_elastic_ep \ -H "Content-Type: application/json" \ -d "{\"new_data_parallel_size\": $to_dp}" \ --max-time "$timeout" \ diff --git a/tests/gpu_memory_service/common/runtime.py b/tests/gpu_memory_service/common/runtime.py index d52853a5cfaa..a623b83a6aac 100644 --- a/tests/gpu_memory_service/common/runtime.py +++ b/tests/gpu_memory_service/common/runtime.py @@ -201,7 +201,7 @@ def _request_engine( action: str, ) -> dict: response = requests.post( - f"http://localhost:{self.system_port}/engine/{route}", + f"http://localhost:{self.system_port}/engine/control/{route}", json=payload, timeout=timeout, ) diff --git a/tests/runtime/test_engine_controls_e2e.py b/tests/runtime/test_engine_controls_e2e.py index c6d4f551afbb..13c70b5e3e41 100644 --- a/tests/runtime/test_engine_controls_e2e.py +++ b/tests/runtime/test_engine_controls_e2e.py @@ -53,13 +53,13 @@ async def sleep_control(body: dict[str, Any]) -> dict[str, Any]: calls.append(body) return {"status": "ok", "control": "sleep", "body": body} - runtime.register_engine_route("sleep", sleep_control) + runtime.register_engine_route("control/sleep", sleep_control) try: async with httpx.AsyncClient(timeout=5.0) as client: response = await _post_with_retry( client, - f"http://127.0.0.1:{system_port}/engine/sleep", + f"http://127.0.0.1:{system_port}/engine/control/sleep", {"level": 1}, )