diff --git a/python/openai/README.md b/python/openai/README.md index 945b050be0..538d086d73 100644 --- a/python/openai/README.md +++ b/python/openai/README.md @@ -542,6 +542,125 @@ available arguments and default values. For more information on the `tritonfrontend` python bindings, see the docs [here](https://github.com/triton-inference-server/server/blob/main/docs/customization_guide/tritonfrontend.md). +## Model Management + +The OpenAI-compatible frontend supports explicit model control, allowing you to +dynamically load and unload models at runtime without restarting the server. +This is particularly useful when hosting multiple large models on a shared GPU +cluster and needing to swap models on demand. + +### Model Control Mode + +Use `--model-control-mode` to specify how models are managed at startup and +runtime. The default is `none`. + +| Mode | Behavior | +|------|----------| +| `none` (default) | All models in the repository are loaded at startup. Load/unload APIs are not available. | +| `explicit` | No models are loaded at startup unless specified with `--load-model`. Load and unload are controlled via the management API. | + +> [!NOTE] +> This matches the native `tritonserver --model-control-mode` behavior. +> See [Triton Model Management](../../docs/user_guide/model_management.md) for +> more details. + +### Load Models at Startup (Explicit Mode) + +When using `--model-control-mode=explicit`, use `--load-model` to specify which +models should be loaded at startup. It may be specified multiple times to load +multiple models. + +
+Example + +```bash +# Start in explicit mode with no models loaded +python3 openai_frontend/main.py \ + --model-repository /path/to/models \ + --tokenizer meta-llama/Meta-Llama-3.1-8B-Instruct \ + --model-control-mode explicit + +# Start in explicit mode and load a specific model at startup +python3 openai_frontend/main.py \ + --model-repository /path/to/models \ + --tokenizer meta-llama/Meta-Llama-3.1-8B-Instruct \ + --model-control-mode explicit \ + --load-model llama-3.1-8b-instruct + +# Load multiple models at startup +python3 openai_frontend/main.py \ + --model-repository /path/to/models \ + --tokenizer meta-llama/Meta-Llama-3.1-8B-Instruct \ + --model-control-mode explicit \ + --load-model model-a \ + --load-model model-b + +# Load ALL models in the repository at startup +python3 openai_frontend/main.py \ + --model-repository /path/to/models \ + --tokenizer meta-llama/Meta-Llama-3.1-8B-Instruct \ + --model-control-mode explicit \ + --load-model '*' +``` + +
+ +> [!IMPORTANT] +> - `--load-model` requires `--model-control-mode=explicit`. +> - `--load-model=*` can not be used together with loading specific models. + +### Dynamic Load / Unload API + +Once the server is running in `explicit` mode, models can be loaded and unloaded +at runtime via the following endpoints: + +| Method | Endpoint | Description | +|--------|----------|-------------| +| `POST` | `/v1/models/{model_name}/load` | Load a model. Blocks until model is fully loaded and ready. | +| `POST` | `/v1/models/{model_name}/unload` | Unload a model. Blocks until fully unloaded. In-flight requests complete before removal. | + +Both endpoints return an error if `--model-control-mode` is not `explicit`. + +#### Load a model + +```bash +MODEL="llama-3.1-8b-instruct" +curl -s -X POST http://localhost:9000/v1/models/${MODEL}/load | jq +``` + +
+Example output + +```json +{ + "id": "llama-3.1-8b-instruct", + "object": "model", + "created": 1750000000, + "owned_by": "Triton Inference Server" +} +``` + +
+ +#### Unload a model + +```bash +MODEL="llama-3.1-8b-instruct" +curl -s -X POST http://localhost:9000/v1/models/${MODEL}/unload | jq +``` + +
+Example output + +```json +{ + "status": "success", + "model": "llama-3.1-8b-instruct" +} +``` + +
+ ## Model Parallelism Support - [x] vLLM ([EngineArgs](https://github.com/triton-inference-server/vllm_backend/blob/main/README.md#using-the-vllm-backend)) @@ -794,9 +913,11 @@ Use the `--openai-restricted-api` command-line argument to configure endpoint re - `POST /v1/completions` - **embedding**: Embedding endpoint - `POST /v1/embeddings` - - **model-repository**: Model listing and information endpoints + - **model-repository**: Model listing, information, and dynamic load/unload endpoints - `GET /v1/models` - `GET /v1/models/{model_name}` + - `POST /v1/models/{model_name}/load` + - `POST /v1/models/{model_name}/unload` - **metrics**: Server metrics endpoint - `GET /metrics` - **health**: Health check endpoint diff --git a/python/openai/openai_frontend/engine/engine.py b/python/openai/openai_frontend/engine/engine.py index 2dfeafb1db..da68ac024b 100644 --- a/python/openai/openai_frontend/engine/engine.py +++ b/python/openai/openai_frontend/engine/engine.py @@ -1,4 +1,4 @@ -# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2024-2026, 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 @@ -100,3 +100,19 @@ def embedding(self, request: CreateEmbeddingRequest) -> CreateEmbeddingResponse: Returns a CreateEmbeddingResponse. """ pass + + async def load_model(self, model_name: str) -> Model: + """ + Loads a model by name. Only available in EXPLICIT model control mode. + Blocks until the model is fully loaded and ready, matching standard + Triton server load behavior. + """ + pass + + async def unload_model(self, model_name: str) -> None: + """ + Unloads a model by name. Only available in EXPLICIT model control mode. + Blocks until the model is fully unloaded, matching standard Triton + server unload behavior. In-flight requests complete before unload. + """ + pass diff --git a/python/openai/openai_frontend/engine/triton_engine.py b/python/openai/openai_frontend/engine/triton_engine.py index d194fc2e1f..19a5b2f157 100644 --- a/python/openai/openai_frontend/engine/triton_engine.py +++ b/python/openai/openai_frontend/engine/triton_engine.py @@ -27,6 +27,7 @@ from __future__ import annotations +import asyncio import base64 import json import time @@ -137,10 +138,8 @@ def __init__( self.lora_separator = lora_separator self.default_max_tokens = default_max_tokens - # NOTE: Creation time and model metadata will be static at startup for - # now, and won't account for dynamically loading/unloading models. - self.create_time = int(time.time()) self.model_metadata = self._get_model_metadata() + self._metadata_lock = asyncio.Lock() self.tool_call_parser = ( ToolParserManager.get_tool_parser_cls(tool_call_parser) if tool_call_parser @@ -493,53 +492,114 @@ def _get_tokenizer(self, tokenizer_name: str): return tokenizer + def _build_model_metadata(self, name: str) -> TritonModelMetadata: + model = self.server.model(name) + backend = model.config()["backend"] + if not backend and model.config()["platform"] == "ensemble": + backend = "ensemble" + print(f"Found model: {name=}, {backend=}") + + lora_configs = _parse_lora_configs( + self.server.options.model_repository, + name, + model.version, + backend if self.backend is None else self.backend, + ) + + echo_tensor_name = None + for input in model.config()["input"]: + if input["name"] in [ + "exclude_input_in_output", + "sampling_param_exclude_input_from_output", + ]: + echo_tensor_name = input["name"] + break + + return TritonModelMetadata( + name=name, + backend=backend, + model=model, + tokenizer=self.tokenizer, + lora_configs=lora_configs, + echo_tensor_name=echo_tensor_name, + create_time=int(time.time()), + inference_request_converter=self._determine_request_converter( + backend, RequestKind.GENERATION + ), + embedding_request_converter=self._determine_request_converter( + backend, RequestKind.EMBEDDING + ), + ) + def _get_model_metadata(self) -> Dict[str, TritonModelMetadata]: - # One tokenizer and creation time shared for all loaded models for now. + # One tokenizer is shared for all loaded models; creation time is per model. model_metadata = {} + for name, _ in self.server.models(exclude_not_ready=True).keys(): + model_metadata[name] = self._build_model_metadata(name) + return model_metadata - # Read all triton models and store the necessary metadata for each - for name, _ in self.server.models().keys(): - model = self.server.model(name) - backend = model.config()["backend"] - # Explicitly handle ensembles to avoid any runtime validation errors - if not backend and model.config()["platform"] == "ensemble": - backend = "ensemble" - print(f"Found model: {name=}, {backend=}") - - lora_configs = _parse_lora_configs( - self.server.options.model_repository, - name, - model.version, - backend if self.backend is None else self.backend, + async def load_model(self, model_name: str) -> Model: + if ( + self.server.options.model_control_mode + != tritonserver.ModelControlMode.EXPLICIT + ): + raise ClientError( + "Model load/unload requires --model-control-mode=explicit" ) - echo_tensor_name = None - for input in model.config()["input"]: - if input["name"] in [ - "exclude_input_in_output", - "sampling_param_exclude_input_from_output", - ]: - echo_tensor_name = input["name"] - break - - metadata = TritonModelMetadata( - name=name, - backend=backend, - model=model, - tokenizer=self.tokenizer, - lora_configs=lora_configs, - echo_tensor_name=echo_tensor_name, - create_time=self.create_time, - inference_request_converter=self._determine_request_converter( - backend, RequestKind.GENERATION - ), - embedding_request_converter=self._determine_request_converter( - backend, RequestKind.EMBEDDING - ), + async with self._metadata_lock: + if model_name in self.model_metadata: + raise ClientError(f"Model '{model_name}' is already loaded") + + # Blocking C API call dispatched to thread pool to avoid blocking + # the event loop. The C API blocks until model is fully loaded and + # ready, matching standard Triton server behavior. + try: + metadata = await asyncio.to_thread(self._load_model_sync, model_name) + except tritonserver.InvalidArgumentError as e: + raise ClientError(f"Failed to load model '{model_name}': {e}") + except tritonserver.TritonError as e: + raise ServerError(f"Failed to load model '{model_name}': {e}") + + self.model_metadata[model_name] = metadata + + return Model( + id=model_name, + created=metadata.create_time, + object=ObjectType.model, + owned_by="Triton Inference Server", + ) + + def _load_model_sync(self, model_name: str) -> TritonModelMetadata: + self.server.load(model_name) + return self._build_model_metadata(model_name) + + async def unload_model(self, model_name: str) -> None: + if ( + self.server.options.model_control_mode + != tritonserver.ModelControlMode.EXPLICIT + ): + raise ClientError( + "Model load/unload requires --model-control-mode=explicit" ) - model_metadata[name] = metadata - return model_metadata + async with self._metadata_lock: + if model_name not in self.model_metadata: + raise ClientError(f"Unknown model: {model_name}") + + # Blocking C API call dispatched to thread pool. The C API handles + # in-flight request draining and conflict resolution internally. + try: + await asyncio.to_thread(self._unload_model_sync, model_name) + except tritonserver.InvalidArgumentError as e: + raise ClientError(f"Failed to unload model '{model_name}': {e}") + except tritonserver.TritonError as e: + raise ServerError(f"Failed to unload model '{model_name}': {e}") + + del self.model_metadata[model_name] + + def _unload_model_sync(self, model_name: str) -> None: + self.server.unload(model_name) def _get_streaming_chat_response_chunk( self, diff --git a/python/openai/openai_frontend/frontend/fastapi/middleware/api_restriction.py b/python/openai/openai_frontend/frontend/fastapi/middleware/api_restriction.py index c2216f9e1d..085434ad46 100644 --- a/python/openai/openai_frontend/frontend/fastapi/middleware/api_restriction.py +++ b/python/openai/openai_frontend/frontend/fastapi/middleware/api_restriction.py @@ -1,4 +1,4 @@ -# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2025-2026, 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 @@ -36,7 +36,10 @@ "POST /v1/completions", "POST /v1/embeddings", ], - "model-repository": ["GET /v1/models"], + "model-repository": [ + "GET /v1/models", + "POST /v1/models/", + ], "metrics": ["GET /metrics"], "health": ["GET /health/ready"], } diff --git a/python/openai/openai_frontend/frontend/fastapi/routers/model_management.py b/python/openai/openai_frontend/frontend/fastapi/routers/model_management.py new file mode 100644 index 0000000000..2ddecb8883 --- /dev/null +++ b/python/openai/openai_frontend/frontend/fastapi/routers/model_management.py @@ -0,0 +1,90 @@ +# Copyright 2026, 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 +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import traceback + +from fastapi import APIRouter, HTTPException, Request +from schemas.openai import Model +from utils.utils import ClientError, ServerError, StatusCode + +router = APIRouter() + + +@router.post( + "/v1/models/{model_name}/load", + response_model=Model, + tags=["Model Management"], +) +async def load_model(model_name: str, raw_request: Request) -> Model: + """ + Loads a model by name. Only available in EXPLICIT model control mode. + Blocks until the model is fully loaded and ready. + """ + if not raw_request.app.engine: + raise HTTPException( + status_code=StatusCode.SERVER_ERROR, + detail="No attached inference engine", + ) + + try: + return await raw_request.app.engine.load_model(model_name) + except ClientError as e: + raise HTTPException(status_code=StatusCode.CLIENT_ERROR, detail=f"{e}") + except ServerError as e: + print(traceback.format_exc()) + raise HTTPException(status_code=StatusCode.SERVER_ERROR, detail=f"{e}") + except Exception as e: + print(traceback.format_exc()) + raise HTTPException(status_code=StatusCode.SERVER_ERROR, detail=f"{e}") + + +@router.post( + "/v1/models/{model_name}/unload", + tags=["Model Management"], +) +async def unload_model(model_name: str, raw_request: Request) -> dict: + """ + Unloads a model by name. Only available in EXPLICIT model control mode. + Blocks until the model is fully unloaded. In-flight requests are allowed + to complete before the model is removed. + """ + if not raw_request.app.engine: + raise HTTPException( + status_code=StatusCode.SERVER_ERROR, + detail="No attached inference engine", + ) + + try: + await raw_request.app.engine.unload_model(model_name) + return {"status": "success", "model": model_name} + except ClientError as e: + raise HTTPException(status_code=StatusCode.CLIENT_ERROR, detail=f"{e}") + except ServerError as e: + print(traceback.format_exc()) + raise HTTPException(status_code=StatusCode.SERVER_ERROR, detail=f"{e}") + except Exception as e: + print(traceback.format_exc()) + raise HTTPException(status_code=StatusCode.SERVER_ERROR, detail=f"{e}") diff --git a/python/openai/openai_frontend/frontend/fastapi/routers/models.py b/python/openai/openai_frontend/frontend/fastapi/routers/models.py index 39e2f26d53..871e300276 100644 --- a/python/openai/openai_frontend/frontend/fastapi/routers/models.py +++ b/python/openai/openai_frontend/frontend/fastapi/routers/models.py @@ -1,4 +1,4 @@ -# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2024-2026, 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 @@ -36,7 +36,7 @@ @router.get("/v1/models", response_model=ListModelsResponse, tags=["Models"]) -def list_models(request: Request) -> ListModelsResponse: +async def list_models(request: Request) -> ListModelsResponse: """ Lists the currently available models, and provides basic information about each one such as the owner and availability. """ @@ -50,7 +50,7 @@ def list_models(request: Request) -> ListModelsResponse: @router.get("/v1/models/{model_name}", response_model=Model, tags=["Models"]) -def retrieve_model(request: Request, model_name: str) -> Model: +async def retrieve_model(request: Request, model_name: str) -> Model: """ Retrieves a model instance, providing basic information about the model such as the owner and permissioning. """ diff --git a/python/openai/openai_frontend/frontend/fastapi_frontend.py b/python/openai/openai_frontend/frontend/fastapi_frontend.py index e5fb01deae..f7fa8aab3c 100644 --- a/python/openai/openai_frontend/frontend/fastapi_frontend.py +++ b/python/openai/openai_frontend/frontend/fastapi_frontend.py @@ -1,4 +1,4 @@ -# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2024-2026, 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 @@ -38,6 +38,7 @@ chat, completions, embeddings, + model_management, models, observability, ) @@ -101,6 +102,7 @@ def _create_app(self): app.include_router(observability.router) app.include_router(models.router) + app.include_router(model_management.router) app.include_router(completions.router) app.include_router(chat.router) app.include_router(embeddings.router) diff --git a/python/openai/openai_frontend/main.py b/python/openai/openai_frontend/main.py index 88aed600d4..ac310e662c 100755 --- a/python/openai/openai_frontend/main.py +++ b/python/openai/openai_frontend/main.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2024-2026, 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 @@ -150,6 +150,31 @@ def parse_args(): default=16, help="The default maximum number of tokens to generate if not specified in the request. The default is 16.", ) + triton_group.add_argument( + "--model-control-mode", + type=str, + default="none", + choices=["none", "explicit"], + help="Specify the mode for model management. Options are 'none', and 'explicit'. " + "The default is 'none'. For 'none', the server will load all models in the model " + "repository at startup and will not make any changes to the loaded " + "models after that. For 'explicit', model load and unload are initiated by using the " + "model control APIs, and only models specified with --load-model will " + "be loaded at startup.", + ) + triton_group.add_argument( + "--load-model", + type=str, + action="append", + default=None, + help="Name of the model to be loaded on server startup. It may be specified " + "multiple times to add multiple models. To load ALL models at startup, " + "specify '*' as the model name with --load-model=* as the ONLY " + "--load-model argument, this does not imply any pattern matching. " + "Specifying --load-model=* in conjunction with another --load-model " + "argument will result in error. Note that this option will only take " + "effect if --model-control-mode is set to 'explicit'.", + ) # OpenAI-Compatible Frontend (FastAPI) openai_group = parser.add_argument_group("Triton OpenAI-Compatible Frontend") @@ -200,8 +225,25 @@ def main(): args = parse_args() # Initialize a Triton Inference Server pointing at LLM models + model_control_mode = ( + tritonserver.ModelControlMode.EXPLICIT + if args.model_control_mode == "explicit" + else tritonserver.ModelControlMode.NONE + ) + + load_models = args.load_model or [] + if load_models and model_control_mode != tritonserver.ModelControlMode.EXPLICIT: + print( + "Error: Use of '--load-model' requires setting " + "'--model-control-mode=explicit' as well.", + file=sys.stderr, + ) + sys.exit(1) + server: tritonserver.Server = tritonserver.Server( model_repository=args.model_repository, + model_control_mode=model_control_mode, + startup_models=load_models, log_verbose=args.tritonserver_log_verbose_level, log_info=True, log_warn=True, diff --git a/python/openai/tests/test_model_management.py b/python/openai/tests/test_model_management.py new file mode 100644 index 0000000000..7cf420eed1 --- /dev/null +++ b/python/openai/tests/test_model_management.py @@ -0,0 +1,613 @@ +# Copyright 2026, 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 +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import concurrent.futures +import http.client +import os +import random +import time +from pathlib import Path +from urllib.parse import urlparse + +import pytest +import requests +from tests.utils import OpenAIServer + +TEST_MODEL_REPOSITORY = str(Path(__file__).parent / "test_models") +TEST_MODEL = "mock_llm" +TEST_MODEL_2 = "identity_py" + + +# Test "--load-model" and "--model-control-mode" CLI options +@pytest.mark.openai +class TestModelControlCLIOptions: + @staticmethod + def _assert_server_launch_fails(args, expected_error: str): + """Helper: verify server fails to start and stderr contains expected_error.""" + with pytest.raises(Exception) as exc_info: + with OpenAIServer(args): + pass + assert expected_error in str(exc_info.value) + + def test_non_explicit_mode_load_model_error(self): + """--load-model without --model-control-mode=explicit must exit with error. + Error message matches native tritonserver exactly.""" + self._assert_server_launch_fails( + args=[ + "--model-repository", + TEST_MODEL_REPOSITORY, + "--load-model", + TEST_MODEL, + ], + expected_error="Error: Use of '--load-model' requires setting '--model-control-mode=explicit' as well.", + ) + + def test_explicit_mode_load_zero_model(self): + args = [ + "--model-repository", + TEST_MODEL_REPOSITORY, + "--model-control-mode", + "explicit", + ] + with OpenAIServer(args) as openai_server: + r = requests.get(openai_server.url_for("v1", "models"), timeout=10) + assert r.status_code == 200 + assert len(r.json()["data"]) == 0 + + def test_explicit_mode_load_one_model(self): + args = [ + "--model-repository", + TEST_MODEL_REPOSITORY, + "--model-control-mode", + "explicit", + "--load-model", + TEST_MODEL, + ] + with OpenAIServer(args) as openai_server: + r = requests.get(openai_server.url_for("v1", "models"), timeout=10) + names = [m["id"] for m in r.json()["data"]] + assert TEST_MODEL in names + assert TEST_MODEL_2 not in names + + def test_explicit_mode_load_multiple_models(self): + args = [ + "--model-repository", + TEST_MODEL_REPOSITORY, + "--model-control-mode", + "explicit", + "--load-model", + TEST_MODEL, + "--load-model", + TEST_MODEL_2, + ] + with OpenAIServer(args) as openai_server: + r = requests.get(openai_server.url_for("v1", "models"), timeout=10) + names = [m["id"] for m in r.json()["data"]] + assert TEST_MODEL in names + assert TEST_MODEL_2 in names + + def test_explicit_mode_load_all_models(self): + args = [ + "--model-repository", + TEST_MODEL_REPOSITORY, + "--model-control-mode", + "explicit", + "--load-model", + "*", + ] + with OpenAIServer(args) as openai_server: + r = requests.get(openai_server.url_for("v1", "models"), timeout=10) + names = [m["id"] for m in r.json()["data"]] + assert TEST_MODEL in names + assert TEST_MODEL_2 in names + + def test_explicit_mode_load_all_models_and_specific_model_error(self): + self._assert_server_launch_fails( + args=[ + "--model-repository", + TEST_MODEL_REPOSITORY, + "--model-control-mode", + "explicit", + "--load-model", + "*", + "--load-model", + TEST_MODEL, + ], + expected_error="Wildcard model name '*' must be the ONLY startup model if specified at all.", + ) + + def test_explicit_mode_load_nonexistent_model_error(self): + self._assert_server_launch_fails( + args=[ + "--model-repository", + TEST_MODEL_REPOSITORY, + "--model-control-mode", + "explicit", + "--load-model", + "nonexistent_model", + ], + expected_error="failed to poll model 'nonexistent_model': model not found in any model repository", + ) + + def test_explicit_mode_load_invalid_model_name_error(self): + invalid_model_names = [ + ( + os.path.relpath("/etc", TEST_MODEL_REPOSITORY), + "at least one version must be available under the version policy", + ), + ( + os.path.relpath("/etc/passwd", TEST_MODEL_REPOSITORY), + "Poll failed for model directory", + ), + ("model/..", "model not found in any model repository"), + ("..", "at least one version must be available under the version policy"), + ("/etc/passwd", "model not found in any model repository"), + ("", "Invalid model name"), + (" ", "model not found in any model repository"), + ("\n\t", "model not found in any model repository"), + ] + for model_name, expected_error in invalid_model_names: + self._assert_server_launch_fails( + args=[ + "--model-repository", + TEST_MODEL_REPOSITORY, + "--model-control-mode", + "explicit", + "--load-model", + model_name, + ], + expected_error=expected_error, + ) + + +class _ModelManagementBase: + @pytest.fixture(scope="class") + def base_url(self, server: OpenAIServer): + return server.url_root + + @pytest.fixture + def all_models(self) -> list[str]: + return [TEST_MODEL, TEST_MODEL_2] + + @pytest.fixture(scope="class") + def load_model(self) -> list[str]: + return [] + + @pytest.fixture(scope="class") + def server( + self, + model_repository: str, + tokenizer_model: str, + backend: str, + model_control_mode: str, + load_model: list[str], + ): + args = [ + "--model-repository", + model_repository, + ] + if tokenizer_model: + args += ["--tokenizer", tokenizer_model] + if backend: + args += ["--backend", backend] + if model_control_mode: + args += ["--model-control-mode", model_control_mode] + if load_model: + for model in load_model: + args += ["--load-model", model] + + with OpenAIServer(args) as openai_server: + yield openai_server + + @pytest.fixture(autouse=True) + def _cleanup(self, base_url, all_models: list[str]): + """Ensure clean state before and after every test by unloading the models.""" + for name in all_models: + requests.post(f"{base_url}/v1/models/{name}/unload") + yield + for name in all_models: + requests.post(f"{base_url}/v1/models/{name}/unload") + + @staticmethod + def _list_available_models(base_url: str) -> list[str]: + response = requests.get(f"{base_url}/v1/models") + assert response.status_code == 200 + return [m["id"] for m in response.json()["data"]] + + @staticmethod + def _assert_unknown_model(response: requests.Response): + assert response.status_code == 400 + assert "unknown model" in response.json()["detail"].lower() + + +class TestModelControlModeNone(_ModelManagementBase): + """Test NONE mode rejects load/unload API calls.""" + + @pytest.fixture(scope="class") + def model_repository(self): + return TEST_MODEL_REPOSITORY + + @pytest.fixture(scope="class") + def model_control_mode(self, request): + return request.param + + @pytest.mark.parametrize("model_control_mode", [None, "none"], indirect=True) + def test_load_and_unload_rejected(self, base_url): + """Test NONE mode rejects load/unload API calls.""" + + for api in ["load", "unload"]: + response = requests.post(f"{base_url}/v1/models/{TEST_MODEL}/{api}") + assert response.status_code == 400 + assert ( + "model load/unload requires --model-control-mode=explicit" + in response.json()["detail"].lower() + ) + + +class TestModelManagement(_ModelManagementBase): + """Test load/unload operations in EXPLICIT mode.""" + + @pytest.fixture(scope="class") + def model_control_mode(self): + return "explicit" + + @pytest.fixture(scope="class") + def model_repository(self): + return TEST_MODEL_REPOSITORY + + @staticmethod + def _assert_model_metadata(model_data: dict): + assert model_data["id"] + assert model_data["object"] == "model" + assert model_data["created"] > 0 + assert model_data["owned_by"] == "Triton Inference Server" + + def test_load_model(self, base_url): + # Server should start with no models loaded + assert self._list_available_models(base_url) == [] + + response = requests.post(f"{base_url}/v1/models/{TEST_MODEL}/load") + assert response.status_code == 200 + self._assert_model_metadata(response.json()) + assert response.json()["id"] == TEST_MODEL + assert TEST_MODEL in self._list_available_models(base_url) + + response = requests.get(f"{base_url}/v1/models/{TEST_MODEL}") + assert response.status_code == 200 + self._assert_model_metadata(response.json()) + + def test_unload_model(self, base_url): + assert ( + requests.post(f"{base_url}/v1/models/{TEST_MODEL}/load").status_code == 200 + ) + + response = requests.post(f"{base_url}/v1/models/{TEST_MODEL}/unload") + assert response.status_code == 200 + body = response.json() + assert body["status"] == "success" + assert body["model"] == TEST_MODEL + + assert TEST_MODEL not in self._list_available_models(base_url) + assert requests.get(f"{base_url}/v1/models/{TEST_MODEL}").status_code == 404 + + def test_load_rejects_duplicate(self, base_url): + assert ( + requests.post(f"{base_url}/v1/models/{TEST_MODEL}/load").status_code == 200 + ) + + response = requests.post(f"{base_url}/v1/models/{TEST_MODEL}/load") + assert response.status_code == 400 + assert "already loaded" in response.json()["detail"].lower() + + def test_load_unload_unknown_model(self, base_url): + response = requests.post(f"{base_url}/v1/models/unknown_model/load") + assert response.status_code == 500 + assert "failed to poll from model repository" in response.json()["detail"] + + response = requests.post(f"{base_url}/v1/models/unknown_model/unload") + self._assert_unknown_model(response) + + def test_load_unload_invalid_model_name(self, base_url): + invalid_model_names = [ + (os.path.relpath("/etc", TEST_MODEL_REPOSITORY), 404), + (os.path.relpath("/etc/passwd", TEST_MODEL_REPOSITORY), 404), + ("model/..", 404), + ("..", 400), + ("/etc/passwd", 404), + ("model/subdir", 404), + ("model/", 404), + ("", 404), + ("%20%20", 400), + ("%0A%09", 400), + ] + parsed = urlparse(base_url) + for model_name, expected_status in invalid_model_names: + for endpoint in ["load", "unload"]: + conn = http.client.HTTPConnection(parsed.hostname, parsed.port) + conn.request("POST", f"/v1/models/{model_name}/{endpoint}") + response = conn.getresponse() + + assert response.status == expected_status, ( + f"Expected {expected_status} for model name {model_name!r}, " + f"got {response.status} {response.read().decode()}" + ) + + def test_load_unload_reload(self, base_url): + assert ( + requests.post(f"{base_url}/v1/models/{TEST_MODEL}/load").status_code == 200 + ) + assert TEST_MODEL in self._list_available_models(base_url) + + assert ( + requests.post(f"{base_url}/v1/models/{TEST_MODEL}/unload").status_code + == 200 + ) + assert TEST_MODEL not in self._list_available_models(base_url) + + assert ( + requests.post(f"{base_url}/v1/models/{TEST_MODEL}/load").status_code == 200 + ) + assert TEST_MODEL in self._list_available_models(base_url) + + def test_load_multiple_models(self, base_url): + assert ( + requests.post(f"{base_url}/v1/models/{TEST_MODEL}/load").status_code == 200 + ) + assert ( + requests.post(f"{base_url}/v1/models/{TEST_MODEL_2}/load").status_code + == 200 + ) + + names = self._list_available_models(base_url) + assert TEST_MODEL in names + assert TEST_MODEL_2 in names + + # Unload the first model + assert ( + requests.post(f"{base_url}/v1/models/{TEST_MODEL}/unload").status_code + == 200 + ) + names = self._list_available_models(base_url) + assert TEST_MODEL not in names + assert TEST_MODEL_2 in names + + # Unload the second model + assert ( + requests.post(f"{base_url}/v1/models/{TEST_MODEL_2}/unload").status_code + == 200 + ) + names = self._list_available_models(base_url) + assert TEST_MODEL not in names + assert TEST_MODEL_2 not in names + + +class TestConcurrentModelManagement(_ModelManagementBase): + """Test sequential and concurrent load/unload.""" + + @pytest.fixture(scope="class") + def model_repository(self): + return TEST_MODEL_REPOSITORY + + @pytest.fixture(scope="class") + def model_control_mode(self): + return "explicit" + + def test_concurrent_load_model(self, base_url): + futures = [] + concurrency = 5 + with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool: + for _ in range(concurrency): + futures.append( + pool.submit( + requests.post, f"{base_url}/v1/models/{TEST_MODEL}/load" + ) + ) + + codes = sorted([future.result().status_code for future in futures]) + assert codes == [200] + [400] * ( + concurrency - 1 + ), f"Expected one 200 and the rest 400 for concurrent loads, got {codes}" + assert TEST_MODEL in self._list_available_models(base_url) + + def test_concurrent_unload_model(self, base_url): + # Load the model + requests.post(f"{base_url}/v1/models/{TEST_MODEL}/load") + assert TEST_MODEL in self._list_available_models(base_url) + + futures = [] + concurrency = 5 + with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool: + for _ in range(concurrency): + futures.append( + pool.submit( + requests.post, f"{base_url}/v1/models/{TEST_MODEL}/unload" + ) + ) + + codes = sorted([future.result().status_code for future in futures]) + assert codes == [200] + [400] * ( + concurrency - 1 + ), f"Expected one 200 and the rest 400 for concurrent unloads, got {codes}" + assert TEST_MODEL not in self._list_available_models(base_url) + + def test_concurrent_load_multiple_models(self, base_url): + concurrency = 10 + tasks = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool: + for _ in range(concurrency // 2): + for model_name in [TEST_MODEL, TEST_MODEL_2]: + tasks.append( + pool.submit( + requests.post, f"{base_url}/v1/models/{model_name}/load" + ) + ) + + codes = sorted([task.result().status_code for task in tasks]) + assert codes == [200, 200] + [400] * ( + concurrency - 2 + ), f"Expected two 200s and the rest 400 for concurrent loads, got {codes}" + available_models = self._list_available_models(base_url) + assert TEST_MODEL in available_models + assert TEST_MODEL_2 in available_models + + def test_concurrent_unload_multiple_models(self, base_url): + # Load both models + for model_name in [TEST_MODEL, TEST_MODEL_2]: + requests.post(f"{base_url}/v1/models/{model_name}/load") + assert model_name in self._list_available_models(base_url) + + concurrency = 10 + tasks = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool: + for _ in range(concurrency // 2): + for model_name in [TEST_MODEL, TEST_MODEL_2]: + tasks.append( + pool.submit( + requests.post, f"{base_url}/v1/models/{model_name}/unload" + ) + ) + + codes = sorted([task.result().status_code for task in tasks]) + assert codes == [200, 200] + [400] * ( + concurrency - 2 + ), f"Expected two 200s and the rest 400 for concurrent unloads, got {codes}" + available_models = self._list_available_models(base_url) + for model_name in [TEST_MODEL, TEST_MODEL_2]: + assert model_name not in available_models + + def test_concurrent_load_unload_stress(self, base_url): + futures = [] + concurrency = 50 + with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool: + for _ in range(concurrency): + action = random.choice(["load", "unload"]) + model_name = random.choice([TEST_MODEL, TEST_MODEL_2]) + futures.append( + pool.submit( + requests.post, f"{base_url}/v1/models/{model_name}/{action}" + ) + ) + done, _ = concurrent.futures.wait( + futures, return_when=concurrent.futures.ALL_COMPLETED + ) + assert ( + len(done) == concurrency + ), f"Expected {concurrency} requests to be completed, got {len(done)}" + for future in done: + response = future.result() + assert response.status_code in ( + 200, + 400, + ), f"Unexpected status code: {response.status_code}" + + # Wait for server to be ready + for _ in range(10): + response = requests.get(f"{base_url}/health/ready") + if response.status_code == 200: + break + time.sleep(1) + assert response.status_code == 200 + + +# Test Inference with real LLM backend (vLLM / TRT-LLM) after load/unload +@pytest.mark.openai +class TestModelManagementInference(_ModelManagementBase): + @pytest.fixture(scope="class") + def model_control_mode(self): + return "explicit" + + @pytest.fixture(scope="class") + def load_model(self, model: str) -> list[str]: + # For tensorrt_llm_bls, we need to load all the dependent models + if model == "tensorrt_llm_bls": + return ["postprocessing", "preprocessing", "tensorrt_llm"] + + @staticmethod + def _assert_load(base_url, model_name: str): + r = requests.post( + f"{base_url}/v1/models/{model_name}/load", + ) + assert r.status_code == 200, f"Load {model_name} failed: {r.text}" + + @staticmethod + def _assert_unload(base_url, model_name: str): + r = requests.post( + f"{base_url}/v1/models/{model_name}/unload", + ) + assert r.status_code == 200, f"Unload {model_name} failed: {r.text}" + + @staticmethod + def _assert_usage(usage): + assert usage is not None + assert usage["prompt_tokens"] > 0 + assert usage["completion_tokens"] > 0 + assert ( + usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"] + ) + + @staticmethod + def _completions(base_url, model_name: str, **kwargs): + return requests.post( + f"{base_url}/v1/completions", + json={ + "model": model_name, + "prompt": "What is machine learning?", + "max_tokens": 10, + **kwargs, + }, + ) + + def test_load_completions(self, base_url, model: str): + self._assert_unknown_model(self._completions(base_url, model)) + + self._assert_load(base_url, model) + r = self._completions(base_url, model) + assert r.status_code == 200 + data = r.json() + assert data["choices"][0]["text"].strip() + assert data["choices"][0]["finish_reason"] == "stop" + self._assert_usage(data["usage"]) + + self._assert_unload(base_url, model) + + def test_unload_rejects_inference(self, base_url, model: str): + self._assert_load(base_url, model) + assert self._completions(base_url, model).status_code == 200 + + self._assert_unload(base_url, model) + self._assert_unknown_model(self._completions(base_url, model)) + + def test_reload_inference(self, base_url, model: str): + self._assert_load(base_url, model) + assert self._completions(base_url, model).status_code == 200 + + self._assert_unload(base_url, model) + self._assert_unknown_model(self._completions(base_url, model)) + + self._assert_load(base_url, model) + r = self._completions(base_url, model) + assert r.status_code == 200 + assert r.json()["choices"][0]["text"].strip() diff --git a/python/openai/tests/test_openai_restricted_apis.py b/python/openai/tests/test_openai_restricted_apis.py index 8e26addb11..200412df34 100755 --- a/python/openai/tests/test_openai_restricted_apis.py +++ b/python/openai/tests/test_openai_restricted_apis.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2025-2026, 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 @@ -40,7 +40,7 @@ def assert_response_success( """Assert that a response was successful.""" assert ( response.status_code == expected_status - ), f"{description} should return {expected_status}, got {response.status_code}" + ), f"{description} should return {expected_status}, got {response.status_code} {response.text}" def assert_response_unauthorized( @@ -49,7 +49,7 @@ def assert_response_unauthorized( """Assert that a response was unauthorized.""" assert ( response.status_code == expected_status - ), f"{description} should be unauthorized with {expected_status}, got {response.status_code}" + ), f"{description} should be unauthorized with {expected_status}, got {response.status_code} {response.text}" def make_get_request( @@ -129,6 +129,7 @@ def make_completion_request( def verify_model_repository_endpoints( base_url, model, headers, expected_success, description_prefix ): + # Verify model repository endpoints response = make_get_request(base_url, "/v1/models", headers=headers) if expected_success: assert_response_success( @@ -149,6 +150,22 @@ def verify_model_repository_endpoints( response, description=f"{description_prefix} Specific model endpoint" ) + # Verify model management endpoints + for endpoint in ["unload", "load"]: + response = requests.post( + f"{base_url}/v1/models/{model}/{endpoint}", headers=headers + ) + if expected_success: + assert_response_success( + response, + description=f"{description_prefix} - Model {endpoint} endpoint", + ) + else: + assert_response_unauthorized( + response, + description=f"{description_prefix} - Model {endpoint} endpoint", + ) + def verify_metrics_endpoint(base_url, headers, expected_success, description_prefix): # Test metrics endpoint @@ -288,6 +305,10 @@ def server_with_restrictions(self, model_repository, tokenizer_model, backend): tokenizer_model, "--backend", backend, + "--model-control-mode", + "explicit", + "--load-model", + "*", "--openai-restricted-api", "inference,model-repository", "admin-key", @@ -345,6 +366,10 @@ def server_multiple_restrictions(self, model_repository, tokenizer_model, backen tokenizer_model, "--backend", backend, + "--model-control-mode", + "explicit", + "--load-model", + "*", "--openai-restricted-api", "model-repository", "model-key", diff --git a/python/openai/tests/utils.py b/python/openai/tests/utils.py index 0088fc0571..c73453ee90 100644 --- a/python/openai/tests/utils.py +++ b/python/openai/tests/utils.py @@ -42,9 +42,15 @@ # TODO: Cleanup, refactor, mock, etc. -def setup_server(model_repository: str): +def setup_server( + model_repository: str, + model_control_mode: tritonserver.ModelControlMode = tritonserver.ModelControlMode.NONE, + load_models: Optional[List[str]] = None, +): server: tritonserver.Server = tritonserver.Server( model_repository=model_repository, + model_control_mode=model_control_mode, + startup_models=load_models or [], log_verbose=0, log_info=True, log_warn=True, diff --git a/qa/L0_openai/generate_engine.py b/qa/L0_openai/generate_engine.py index b454896cfb..22c21d2209 100644 --- a/qa/L0_openai/generate_engine.py +++ b/qa/L0_openai/generate_engine.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -32,7 +32,7 @@ def generate_model_engine(model: str, engines_path: str): - config = BuildConfig(plugin_config=PluginConfig.from_dict({"_gemm_plugin": "auto"})) + config = BuildConfig(plugin_config=PluginConfig(gemm_plugin="auto")) lora_config = LoraConfig( lora_target_modules=["attn_q", "attn_k", "attn_v"], diff --git a/src/command_line_parser.cc b/src/command_line_parser.cc index cc11803ae8..359a02ea6e 100644 --- a/src/command_line_parser.cc +++ b/src/command_line_parser.cc @@ -1,4 +1,4 @@ -// Copyright 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2022-2026, 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 @@ -424,11 +424,11 @@ TritonParser::SetupOptions() "Specify the mode for model management. Options are \"none\", \"poll\" " "and \"explicit\". The default is \"none\". " "For \"none\", the server will load all models in the model " - "repository(s) at startup and will not make any changes to the load " + "repository(s) at startup and will not make any changes to the loaded " "models after that. For \"poll\", the server will poll the model " "repository(s) to detect changes and will load/unload models based on " "those changes. The poll rate is controlled by 'repository-poll-secs'. " - "For \"explicit\", model load and unload is initiated by using the " + "For \"explicit\", model load and unload are initiated by using the " "model control APIs, and only models specified with --load-model will " "be loaded at startup."}); model_repo_options_.push_back(