diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index ba079d7a4637..4f497787e143 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -555,6 +555,15 @@ class CompletionRequest(OpenAIBaseModel): default=None, description=("Parameters for disaggregated serving"), ) + priority: Optional[float] = Field( + default=None, + ge=0.0, + le=1.0, + description= + ("Scheduling priority in [0.0, 1.0]; higher is served first. Only honored " + "when the engine runs with scheduler_config.waiting_queue_policy=priority. " + "If unset, the engine default (0.5) is used."), + ) conversation_params: Optional[ConversationParams] = Field( default=None, description=("Parameters for multi-turn conversation routing"), @@ -954,6 +963,15 @@ class ChatCompletionRequest(OpenAIBaseModel): ("If specified, KV cache will be salted with the provided string " "to limit the kv cache reuse on with the requests having the same string." )) + priority: Optional[float] = Field( + default=None, + ge=0.0, + le=1.0, + description= + ("Scheduling priority in [0.0, 1.0]; higher is served first. Only honored " + "when the engine runs with scheduler_config.waiting_queue_policy=priority. " + "If unset, the engine default (0.5) is used."), + ) agent_hierarchy: Optional[AgentHierarchy] = Field( default=None, description="Agent hierarchy ") diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 3e1aa0d644db..6e91986774f0 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -35,6 +35,7 @@ # yapf: disable from tensorrt_llm.executor import CppExecutorError from tensorrt_llm.executor.postproc_worker import PostprocParams +from tensorrt_llm.executor.request import DEFAULT_REQUEST_PRIORITY from tensorrt_llm.inputs import prompt_inputs from tensorrt_llm.inputs.data import TokensPrompt from tensorrt_llm.inputs.media_io import BaseMediaIO @@ -1623,6 +1624,8 @@ async def chat_stream_generator( cache_salt=request.cache_salt, trace_headers=trace_headers, scheduling_params=scheduling_params, + priority=request.priority + if request.priority is not None else DEFAULT_REQUEST_PRIORITY, ) asyncio.create_task(self.await_disconnected(raw_request, promise)) if not self.postproc_worker_enabled: @@ -1951,7 +1954,9 @@ async def generator_wrapper(generator: AsyncIterator[Any]): lora_request=request.lora_request, disaggregated_params=disaggregated_params, conversation_params=conversation_params, - trace_headers=trace_headers) + trace_headers=trace_headers, + priority=request.priority if request.priority is not None + else DEFAULT_REQUEST_PRIORITY) asyncio.create_task( self.await_disconnected(raw_request, promise)) if not self.postproc_worker_enabled: @@ -2112,6 +2117,8 @@ async def create_streaming_generator(promise: RequestOutput, disaggregated_params=disaggregated_params, conversation_params=conversation_params, trace_headers=trace_headers, + priority=request.priority + if request.priority is not None else DEFAULT_REQUEST_PRIORITY, ) if not self.postproc_worker_enabled: postproc_args.num_prompt_tokens = len(promise.prompt_token_ids) diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index 471491aaa0b9..99d263914f05 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -301,6 +301,12 @@ models: default: null status: prototype required: false + priority: + kind: extension + type: Optional[float] + default: null + status: prototype + required: false CompletionResponseChoice: fields: @@ -802,6 +808,12 @@ models: default: null status: beta required: false + priority: + kind: extension + type: Optional[float] + default: null + status: prototype + required: false ChatCompletionResponseChoice: fields: