🎨 engine type hint

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>
vllm-project · Oct 31, 2024 · 28bee6d · 28bee6d
1 parent d38d052
commit 28bee6d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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()