Stop caching LoRA requests and query vllm api server cache to check if it contains incoming LoRA requests

Author: jberkhahn

src/vllm_tgis_adapter/__main__.py

@@ -19,7 +19,7 @@
 from vllm_tgis_adapter.tgis_utils.logs import add_logging_wrappers
 
 from .grpc import run_grpc_server
-from .http import run_http_server
+from .http import build_http_server, run_http_server
 from .logging import DEFAULT_LOGGER_NAME, init_logger
 from .tgis_utils.args import EnvVarArgumentParser, add_tgis_args, postprocess_tgis_args
 from .utils import check_for_failed_tasks, write_termination_log
@@ -43,15 +43,16 @@ async def start_servers(args: argparse.Namespace) -> None:
     async with build_async_engine_client(args) as engine:
         add_logging_wrappers(engine)
 
+        vllm_server = await build_http_server(args, engine)
         http_server_task = loop.create_task(
-            run_http_server(args, engine, sock),
+            run_http_server(args, vllm_server, sock),
             name="http_server",
         )
         # The http server task will catch interrupt signals for us
         tasks.append(http_server_task)
 
         grpc_server_task = loop.create_task(
-            run_grpc_server(args, engine),
+            run_grpc_server(args, engine, vllm_server),
             name="grpc_server",
         )
         tasks.append(grpc_server_task)

src/vllm_tgis_adapter/grpc/adapters.py

@@ -15,20 +15,29 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from vllm.lora.request import LoRARequest
+from vllm.entrypoints.openai.protocol import ErrorResponse
 from vllm.prompt_adapter.request import PromptAdapterRequest
 
 from vllm_tgis_adapter.logging import init_logger
 from vllm_tgis_adapter.tgis_utils.convert_pt_to_prompt import convert_pt_to_peft
 
 from .validation import TGISValidationError
 
+try:
+    from vllm.entrypoints.openai.protocol import LoadLoRAAdapterRequest
+except ImportError:
+    from vllm.entrypoints.openai.protocol import (
+        LoadLoraAdapterRequest as LoadLoRAAdapterRequest,
+    )
+
 if TYPE_CHECKING:
     from vllm.entrypoints.grpc.pb.generation_pb2 import (
         BatchedGenerationRequest,
         BatchedTokenizeRequest,
         SingleGenerationRequest,
     )
+    from vllm.entrypoints.openai.serving_models import OpenAIServingModels
+    from vllm.lora.request import LoRARequest
 
 global_thread_pool = None  # used for loading adapter files from disk
 
@@ -49,7 +58,8 @@ class AdapterMetadata:
 class AdapterStore:
     cache_path: str  # Path to local store of adapters to load from
     adapters: dict[str, AdapterMetadata]
-    next_unique_id: int = 1
+    # Pick a large number to avoid colliding with vllm's adapter IDs
+    next_unique_id: int = 1000001
     load_locks: dict[str, asyncio.Lock] = dataclasses.field(default_factory=dict)
 
 
@@ -58,6 +68,7 @@ async def validate_adapters(
     | BatchedGenerationRequest
     | BatchedTokenizeRequest,
     adapter_store: AdapterStore | None,
+    vllm_model_handler: OpenAIServingModels,
 ) -> dict[str, LoRARequest | PromptAdapterRequest]:
     """Validate the adapters.
 
@@ -81,6 +92,12 @@ async def validate_adapters(
 
     # Guard against concurrent access for the same adapter
     async with adapter_store.load_locks.setdefault(adapter_id, asyncio.Lock()):
+        # Check VLLM server lora cache if this request matches an existing
+        # LoRA adapter
+        for existing_lora_request in vllm_model_handler.lora_requests:
+            if existing_lora_request.lora_name == adapter_id:
+                return {"lora_request": existing_lora_request}
+
         # If not already cached, we need to validate that files exist and
         # grab the type out of the adapter_config.json file
         if (adapter_metadata := adapter_store.adapters.get(adapter_id)) is None:
@@ -107,16 +124,19 @@ async def validate_adapters(
             )
 
             # Add to cache
+            # Query vllm's cache for lora requests
+            if adapter_metadata.adapter_type == "LORA":
+                lora_request = await _load_lora_adapter(
+                    request,
+                    adapter_id,
+                    adapter_metadata,
+                    vllm_model_handler,
+                )
+                return {"lora_request": lora_request}
+            # Use our cache for everything else
             adapter_store.adapters[adapter_id] = adapter_metadata
 
     # Build the proper vllm request object
-    if adapter_metadata.adapter_type == "LORA":
-        lora_request = LoRARequest(
-            lora_name=adapter_id,
-            lora_int_id=adapter_metadata.unique_id,
-            lora_path=adapter_metadata.full_path,
-        )
-        return {"lora_request": lora_request}
     if adapter_metadata.adapter_type == "PROMPT_TUNING":
         prompt_adapter_request = PromptAdapterRequest(
             prompt_adapter_id=adapter_metadata.unique_id,
@@ -126,12 +146,36 @@ async def validate_adapters(
                 "num_virtual_tokens", 0
             ),
         )
-        return {"prompt_adapter_request": prompt_adapter_request}
+    return {"prompt_adapter_request": prompt_adapter_request}
 
     # All other types unsupported
     TGISValidationError.AdapterUnsupported.error(adapter_metadata.adapter_type)  # noqa: RET503
 
 
+async def _load_lora_adapter(
+    request: SingleGenerationRequest
+    | BatchedGenerationRequest
+    | BatchedTokenizeRequest,
+    adapter_id: str,
+    adapter_metadata: AdapterMetadata,
+    vllm_model_handler: OpenAIServingModels,
+) -> LoRARequest:
+    load_request = LoadLoRAAdapterRequest(
+        lora_path=adapter_metadata.full_path,
+        lora_name=adapter_id,
+    )
+    load_result = await vllm_model_handler.load_lora_adapter(
+        request=load_request,
+        base_model_name=request.model_id,
+    )
+    if isinstance(load_result, ErrorResponse):
+        raise ValueError(load_result.message)  ## noqa: TRY004
+    for existing_lora_request in vllm_model_handler.lora_requests:
+        if existing_lora_request.lora_name == adapter_id:
+            return existing_lora_request
+    raise RuntimeError("vllm server failed to load LoRA adapter")
+
+
 def _load_adapter_metadata(adapter_id: str, adapter_path: str, unique_id: int) -> dict:
     """Get adapter metadata from files.
 

src/vllm_tgis_adapter/grpc/grpc_server.py

@@ -59,10 +59,12 @@
     import argparse
     from collections.abc import AsyncIterator, MutableSequence
 
+    from fastapi import FastAPI
     from grpc.aio import ServicerContext
     from vllm import CompletionOutput, RequestOutput
     from vllm.config import ModelConfig
     from vllm.engine.protocol import EngineClient
+    from vllm.entrypoints.openai.serving_models import OpenAIServingModels
     from vllm.lora.request import LoRARequest
     from vllm.sequence import Logprob
     from vllm.transformers_utils.tokenizer import AnyTokenizer
@@ -167,9 +169,11 @@ def __init__(
         args: argparse.Namespace,
         health_servicer: health.HealthServicer,
         stop_event: asyncio.Event,
+        vllm_server: FastAPI,
     ):
         self.engine: EngineClient = engine
         self.stop_event = stop_event
+        self.vllm_server = vllm_server
 
         # This is set in post_init()
         self.config: ModelConfig | None = None
@@ -218,7 +222,11 @@ async def Generate(
         start_time = time.time()
         service_metrics.count_generate_request(len(request.requests))
         request_id = self.request_id(context)
-        kwargs = await self._validate_adapters(request, context)
+        kwargs = await self._validate_adapters(
+            request,
+            context,
+            self.vllm_server.state.openai_serving_models,
+        )
         tokenizer = await self._get_tokenizer(kwargs)
 
         sampling_params, deadline = await self._validate_and_convert_params(
@@ -308,7 +316,11 @@ async def GenerateStream(  # noqa: PLR0915, C901
         start_time = time.time()
         service_metrics.count_generate_request()
         request_id = self.request_id(context)
-        adapter_kwargs = await self._validate_adapters(request, context)
+        adapter_kwargs = await self._validate_adapters(
+            request,
+            context,
+            self.vllm_server.state.openai_serving_models,
+        )
         tokenizer = await self._get_tokenizer(adapter_kwargs)
 
         sampling_params, deadline = await self._validate_and_convert_params(
@@ -628,10 +640,13 @@ async def _validate_adapters(
         | BatchedGenerationRequest
         | BatchedTokenizeRequest,
         context: ServicerContext,
+        vllm_model_handler: OpenAIServingModels,
     ) -> dict[str, LoRARequest | PromptAdapterRequest]:
         try:
             adapters = await validate_adapters(
-                request=request, adapter_store=self.adapter_store
+                request=request,
+                adapter_store=self.adapter_store,
+                vllm_model_handler=vllm_model_handler,
             )
         except ValueError as e:
             service_metrics.count_request_failure(FailureReasonLabel.VALIDATION)
@@ -812,7 +827,11 @@ async def Tokenize(
         service_metrics.count_tokenization_request(request)
 
         # TODO simplify to only check for lora adapter
-        adapter_kwargs = await self._validate_adapters(request, context)
+        adapter_kwargs = await self._validate_adapters(
+            request,
+            context,
+            self.vllm_server.state.openai_serving_models,
+        )
         tokenizer = await self._get_tokenizer(adapter_kwargs)
 
         responses: list[TokenizeResponse] = []
@@ -886,13 +905,20 @@ async def start_grpc_server(
     args: argparse.Namespace,
     engine: EngineClient,
     stop_event: asyncio.Event,
+    vllm_server: FastAPI,
 ) -> aio.Server:
     server = aio.server()
 
     health_servicer = health.HealthServicer()
     health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
 
-    generation = TextGenerationService(engine, args, health_servicer, stop_event)
+    generation = TextGenerationService(
+        engine,
+        args,
+        health_servicer,
+        stop_event,
+        vllm_server,
+    )
     await generation.post_init()
     generation_pb2_grpc.add_GenerationServiceServicer_to_server(generation, server)
 
@@ -951,9 +977,10 @@ async def start_grpc_server(
 async def run_grpc_server(
     args: argparse.Namespace,
     engine: EngineClient,
+    vllm_server: FastAPI,
 ) -> None:
     stop_event = asyncio.Event()
-    server = await start_grpc_server(args, engine, stop_event)
+    server = await start_grpc_server(args, engine, stop_event, vllm_server)
 
     # Add a task to watch for the stop event, so that the server can kill
     # itself from within its own handlers

src/vllm_tgis_adapter/http.py

@@ -13,7 +13,7 @@
     import argparse
     import socket
 
-    from fastapi import Request, Response
+    from fastapi import FastAPI, Request, Response
     from vllm.engine.async_llm_engine import AsyncLLMEngine
     from vllm.engine.protocol import AsyncEngineClient
 
@@ -22,14 +22,12 @@
 logger = init_logger(__name__)
 
 
-async def run_http_server(
+async def build_http_server(
     args: argparse.Namespace,
     engine: AsyncLLMEngine | AsyncEngineClient,
-    sock: socket.socket | None = None,
-    **uvicorn_kwargs,  # noqa: ANN003
-) -> None:
-    # modified copy of vllm.entrypoints.openai.api_server.run_server that
-    # allows passing of the engine
+) -> FastAPI:
+    # builds the vllm api server so we can pass reference to it
+    # within the tgis adapter
 
     app = build_app(args)
 
@@ -53,6 +51,18 @@ async def set_correlation_id(request: Request, call_next: Callable) -> Response:
     if inspect.isawaitable(maybe_coroutine):
         await maybe_coroutine
 
+    return app
+
+
+async def run_http_server(
+    args: argparse.Namespace,
+    app: FastAPI,
+    sock: socket.socket | None = None,
+    **uvicorn_kwargs,  # noqa: ANN003
+) -> None:
+    # modified copy of vllm.entrypoints.openai.api_server.run_server that
+    # allows passing of the engine
+
     serve_kwargs = {
         "host": args.host,
         "port": args.port,