From 90e110ea53c6dfc785b8a0abaf7a02750a6c9ad3 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 25 Mar 2025 13:24:12 -0700 Subject: [PATCH] [V1][Minor] Use SchedulerInterface type for engine scheduler field So that IDEs know what the methods are. Signed-off-by: Nick Hill --- vllm/v1/engine/core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index f4bb4583bea4..42511777feeb 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -23,6 +23,7 @@ zmq_socket_ctx) from vllm.v1.core.kv_cache_utils import (get_kv_cache_config, unify_kv_cache_configs) +from vllm.v1.core.sched.interface import SchedulerInterface from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler as V1Scheduler from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest, @@ -84,7 +85,7 @@ def __init__( "compatibility may not be maintained.", vllm_config.scheduler_config.scheduler_cls) - self.scheduler = Scheduler( + self.scheduler: SchedulerInterface = Scheduler( scheduler_config=vllm_config.scheduler_config, model_config=vllm_config.model_config, cache_config=vllm_config.cache_config,