feat: Add multi-LoRA support to OpenAI frontend (#8038)

Author: kthui

python/openai/README.md

@@ -230,6 +230,69 @@ pip install -r requirements-test.txt
 pytest -v tests/
 ```
 
+### LoRA Adapters
+
+If the command line argument `--lora-separator=<separator_string>` is provided
+when starting the OpenAI Frontend, a vLLM LoRA adaptor listed on the
+`multi_lora.json` may be selected by appending the LoRA name to the model name,
+separated by the LoRA separator, on the inference request in
+`<model_name><separator_string><lora_name>` format.
+
+<details>
+<summary>For example</summary>
+
+```bash
+# start server with model named gemma-2b
+python3 openai_frontend/main.py --lora-separator=_lora_ ...
+
+# inference without LoRA
+curl -s http://localhost:9000/v1/completions -H 'Content-Type: application/json' -d '{
+  "model": "gemma-2b",
+  "temperature": 0,
+  "prompt": "When was the wheel invented?"
+}'
+{
+  ...
+  "choices":[{..."text":"\n\nThe wheel was invented by the Sumerians in Mesopotamia around 350"}],
+  ...
+}
+
+# inference with LoRA named doll
+curl -s http://localhost:9000/v1/completions -H 'Content-Type: application/json' -d '{
+  "model": "gemma-2b_lora_doll",
+  "temperature": 0,
+  "prompt": "When was the wheel invented?"
+}'
+{
+  ...
+  "choices":[{..."text":"\n\nThe wheel was invented in Mesopotamia around 3500 BC.\n\n"}],
+  ...
+}
+
+# inference with LoRA named sheep
+curl -s http://localhost:9000/v1/completions -H 'Content-Type: application/json' -d '{
+  "model": "gemma-2b_lora_sheep",
+  "temperature": 0,
+  "prompt": "When was the wheel invented?"
+}'
+{
+  ...
+  "choices":[{..."text":"\n\nThe wheel was invented around 3000 BC in Mesopotamia.\n\n"}],
+  ...
+}
+```
+
+</details>
+
+When listing or retrieving model(s), the model id will include the LoRA name in
+the same `<model_name><separator_string><lora_name>` format for each LoRA
+adapter listed on the `multi_lora.json`. Note: The LoRA name inclusion is
+limited to locally stored models, inference requests are not limited though.
+
+See the
+[vLLM documentation](https://github.com/triton-inference-server/vllm_backend/blob/main/docs/llama_multi_lora_tutorial.md)
+on how to serve a model with LoRA adapters.
+
 ## TensorRT-LLM
 
 0. Prepare your model repository for a TensorRT-LLM model, build the engine, etc. You can try any of the following options:

python/openai/openai_frontend/engine/triton_engine.py

@@ -1,4 +1,4 @@
-# Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -39,6 +39,7 @@
     _create_trtllm_inference_request,
     _create_vllm_inference_request,
     _get_output,
+    _get_vllm_lora_names,
     _validate_triton_responses_non_streaming,
 )
 from schemas.openai import (
@@ -70,6 +71,8 @@ class TritonModelMetadata:
     model: tritonserver.Model
     # Tokenizers used for chat templates
     tokenizer: Optional[Any]
+    # LoRA names supported by the backend
+    lora_names: Optional[List[str]]
     # Time that model was loaded by Triton
     create_time: int
     # Conversion format between OpenAI and Triton requests
@@ -78,13 +81,18 @@ class TritonModelMetadata:
 
 class TritonLLMEngine(LLMEngine):
     def __init__(
-        self, server: tritonserver.Server, tokenizer: str, backend: Optional[str] = None
+        self,
+        server: tritonserver.Server,
+        tokenizer: str,
+        backend: Optional[str] = None,
+        lora_separator: Optional[str] = None,
     ):
         # Assume an already configured and started server
         self.server = server
         self.tokenizer = self._get_tokenizer(tokenizer)
         # TODO: Reconsider name of "backend" vs. something like "request_format"
         self.backend = backend
+        self.lora_separator = lora_separator
 
         # NOTE: Creation time and model metadata will be static at startup for
         # now, and won't account for dynamically loading/unloading models.
@@ -100,22 +108,35 @@ def metrics(self) -> str:
     def models(self) -> List[Model]:
         models = []
         for metadata in self.model_metadata.values():
-            models.append(
-                Model(
-                    id=metadata.name,
-                    created=metadata.create_time,
-                    object=ObjectType.model,
-                    owned_by="Triton Inference Server",
-                ),
-            )
+            model_names = [metadata.name]
+            if (
+                self.lora_separator is not None
+                and len(self.lora_separator) > 0
+                and metadata.lora_names is not None
+            ):
+                for lora_name in metadata.lora_names:
+                    model_names.append(
+                        f"{metadata.name}{self.lora_separator}{lora_name}"
+                    )
+
+            for model_name in model_names:
+                models.append(
+                    Model(
+                        id=model_name,
+                        created=metadata.create_time,
+                        object=ObjectType.model,
+                        owned_by="Triton Inference Server",
+                    ),
+                )
 
         return models
 
     async def chat(
         self, request: CreateChatCompletionRequest
     ) -> CreateChatCompletionResponse | AsyncIterator[str]:
-        metadata = self.model_metadata.get(request.model)
-        self._validate_chat_request(request, metadata)
+        model_name, lora_name = self._get_model_and_lora_name(request.model)
+        metadata = self.model_metadata.get(model_name)
+        self._validate_chat_request(request, metadata, lora_name)
 
         conversation = [
             message.model_dump(exclude_none=True) for message in request.messages
@@ -130,7 +151,7 @@ async def chat(
 
         # Convert to Triton request format and perform inference
         responses = metadata.model.async_infer(
-            metadata.request_converter(metadata.model, prompt, request)
+            metadata.request_converter(metadata.model, prompt, request, lora_name)
         )
 
         # Prepare and send responses back to client in OpenAI format
@@ -174,20 +195,23 @@ async def completion(
         self, request: CreateCompletionRequest
     ) -> CreateCompletionResponse | AsyncIterator[str]:
         # Validate request and convert to Triton format
-        metadata = self.model_metadata.get(request.model)
-        self._validate_completion_request(request, metadata)
+        model_name, lora_name = self._get_model_and_lora_name(request.model)
+        metadata = self.model_metadata.get(model_name)
+        self._validate_completion_request(request, metadata, lora_name)
 
         # Convert to Triton request format and perform inference
         responses = metadata.model.async_infer(
-            metadata.request_converter(metadata.model, request.prompt, request)
+            metadata.request_converter(
+                metadata.model, request.prompt, request, lora_name
+            )
         )
 
         # Prepare and send responses back to client in OpenAI format
         request_id = f"cmpl-{uuid.uuid1()}"
         created = int(time.time())
         if request.stream:
             return self._streaming_completion_iterator(
-                request_id, created, metadata.name, responses
+                request_id, created, request.model, responses
             )
 
         # Response validation with decoupled models in mind
@@ -208,7 +232,7 @@ async def completion(
             system_fingerprint=None,
             object=ObjectType.text_completion,
             created=created,
-            model=metadata.name,
+            model=request.model,
         )
 
     # TODO: This behavior should be tested further
@@ -234,6 +258,16 @@ def _determine_request_converter(self, backend: str):
         # an ensemble, a python or BLS model, a TRT-LLM backend model, etc.
         return _create_trtllm_inference_request
 
+    def _get_model_and_lora_name(self, request_model_name: str):
+        if self.lora_separator is None or len(self.lora_separator) == 0:
+            return request_model_name, None
+
+        names = request_model_name.split(self.lora_separator)
+        if len(names) != 2:
+            return request_model_name, None
+
+        return names[0], names[1]
+
     def _get_tokenizer(self, tokenizer_name: str):
         tokenizer = None
         if tokenizer_name:
@@ -254,11 +288,18 @@ def _get_model_metadata(self) -> Dict[str, TritonModelMetadata]:
                 backend = "ensemble"
             print(f"Found model: {name=}, {backend=}")
 
+            lora_names = None
+            if self.backend == "vllm" or backend == "vllm":
+                lora_names = _get_vllm_lora_names(
+                    self.server.options.model_repository, name, model.version
+                )
+
             metadata = TritonModelMetadata(
                 name=name,
                 backend=backend,
                 model=model,
                 tokenizer=self.tokenizer,
+                lora_names=lora_names,
                 create_time=self.create_time,
                 request_converter=self._determine_request_converter(backend),
             )
@@ -343,7 +384,10 @@ async def _streaming_chat_iterator(
         yield "data: [DONE]\n\n"
 
     def _validate_chat_request(
-        self, request: CreateChatCompletionRequest, metadata: TritonModelMetadata
+        self,
+        request: CreateChatCompletionRequest,
+        metadata: TritonModelMetadata,
+        lora_name: str | None,
     ):
         """
         Validates a chat request to align with currently supported features.
@@ -362,6 +406,13 @@ def _validate_chat_request(
         if not metadata.request_converter:
             raise Exception(f"Unknown request format for model: {request.model}")
 
+        if (
+            metadata.lora_names is not None
+            and lora_name is not None
+            and lora_name not in metadata.lora_names
+        ):
+            raise Exception(f"Unknown LoRA: {lora_name}; for model: {request.model}")
+
         # Reject unsupported features if requested
         if request.n and request.n > 1:
             raise Exception(
@@ -396,7 +447,10 @@ async def _streaming_completion_iterator(
         yield "data: [DONE]\n\n"
 
     def _validate_completion_request(
-        self, request: CreateCompletionRequest, metadata: TritonModelMetadata
+        self,
+        request: CreateCompletionRequest,
+        metadata: TritonModelMetadata,
+        lora_name: str | None,
     ):
         """
         Validates a completions request to align with currently supported features.
@@ -411,6 +465,13 @@ def _validate_completion_request(
         if not metadata.request_converter:
             raise Exception(f"Unknown request format for model: {request.model}")
 
+        if (
+            metadata.lora_names is not None
+            and lora_name is not None
+            and lora_name not in metadata.lora_names
+        ):
+            raise Exception(f"Unknown LoRA: {lora_name}; for model: {request.model}")
+
         # Reject unsupported features if requested
         if request.suffix is not None:
             raise Exception("suffix is not currently supported")