From 28bee6dc888d497a4f31cbf28bbfad701e638018 Mon Sep 17 00:00:00 2001 From: Joe Runde Date: Thu, 31 Oct 2024 14:09:04 -0600 Subject: [PATCH] :art: engine type hint Signed-off-by: Joe Runde --- vllm/entrypoints/llm.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index d1d140a4b4a68..ea55d142bea58 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -204,7 +204,9 @@ def __init__( ) # Logic to switch between engines is done at runtime instead of import # to avoid import order issues - self.engine_class = V1LLMEngine if envs.VLLM_USE_V1 else LLMEngine + self.engine_class: Union[ + V1LLMEngine, + LLMEngine] = V1LLMEngine if envs.VLLM_USE_V1 else LLMEngine self.llm_engine = self.engine_class.from_engine_args( engine_args, usage_context=UsageContext.LLM_CLASS) self.request_counter = Counter()