[Bug fix] vLLM upstream compatibility. Fix DP scheduler (#1057)

Signed-off-by: wenxindongwork <wenxindong@google.com>

Author: wenxindongwork

tests/core/test_dp_scheduler.py

@@ -5,6 +5,7 @@
 from vllm.config import VllmConfig
 from vllm.v1.core.sched.output import (CachedRequestData, GrammarOutput,
                                        SchedulerOutput)
+from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.engine import EngineCoreOutputs
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats
@@ -24,9 +25,10 @@ def mock_vllm_config(self):
         config.sharding_config = MagicMock()
         config.sharding_config.total_dp_size = 2
         config.scheduler_config = MagicMock()
-        config.scheduler_config._original_scheduler_cls = "vllm.v1.core.sched.scheduler.Scheduler"
+        config.scheduler_config._original_scheduler_cls = Scheduler
         config.scheduler_config.max_num_seqs = 8
         config.scheduler_config.max_num_batched_tokens = 1024
+        config.scheduler_config.async_scheduling = False
         return config
 
     @pytest.fixture
@@ -46,18 +48,14 @@ def _create_dp_scheduler_with_mocks(self, mock_vllm_config,
                                         mock_structured_output_manager,
                                         **kwargs):
         """Helper to create a DPScheduler with properly mocked schedulers."""
-        with patch(
-                "tpu_inference.core.sched.dp_scheduler.resolve_obj_by_qualname"
-        ) as mock_resolve:
-            # Create individual mock scheduler instances
-            mock_scheduler_0 = MagicMock()
-            mock_scheduler_1 = MagicMock()
-
-            # Set up the mock class to return these instances
-            mock_scheduler_cls = MagicMock(
-                side_effect=[mock_scheduler_0, mock_scheduler_1])
-            mock_resolve.return_value = mock_scheduler_cls
-
+        # Create individual mock scheduler instances
+        mock_scheduler_0 = MagicMock()
+        mock_scheduler_1 = MagicMock()
+
+        # Patch the Scheduler class to return our mock instances
+        with patch.object(
+                mock_vllm_config.scheduler_config, '_original_scheduler_cls',
+                MagicMock(side_effect=[mock_scheduler_0, mock_scheduler_1])):
             scheduler = DPScheduler(
                 vllm_config=mock_vllm_config,
                 kv_cache_config=mock_kv_cache_config,
@@ -67,38 +65,36 @@ def _create_dp_scheduler_with_mocks(self, mock_vllm_config,
 
             return scheduler
 
-    @patch("tpu_inference.core.sched.dp_scheduler.resolve_obj_by_qualname")
     def test_init_creates_per_rank_schedulers(
         self,
-        mock_resolve,
         mock_vllm_config,
         mock_kv_cache_config,
         mock_structured_output_manager,
     ):
         """Test Initialization creates schedulers for each DP rank."""
         # Mock the scheduler class
-        mock_scheduler_cls = MagicMock()
         mock_scheduler_instance = MagicMock()
-        mock_scheduler_cls.return_value = mock_scheduler_instance
-        mock_resolve.return_value = mock_scheduler_cls
-
-        scheduler = DPScheduler(
-            vllm_config=mock_vllm_config,
-            kv_cache_config=mock_kv_cache_config,
-            structured_output_manager=mock_structured_output_manager,
-            block_size=16,
-            log_stats=True,
-        )
+        mock_scheduler_cls = MagicMock(return_value=mock_scheduler_instance)
+
+        with patch.object(mock_vllm_config.scheduler_config,
+                          '_original_scheduler_cls', mock_scheduler_cls):
+            scheduler = DPScheduler(
+                vllm_config=mock_vllm_config,
+                kv_cache_config=mock_kv_cache_config,
+                structured_output_manager=mock_structured_output_manager,
+                block_size=16,
+                log_stats=True,
+            )
 
-        # Verify schedulers were created
-        assert len(scheduler.schedulers) == 2
-        assert scheduler.dp_size == 2
-        assert scheduler.log_stats is True
-        assert len(scheduler.per_rank_kv_cache_configs) == 2
+            # Verify schedulers were created
+            assert len(scheduler.schedulers) == 2
+            assert scheduler.dp_size == 2
+            assert scheduler.log_stats is True
+            assert len(scheduler.per_rank_kv_cache_configs) == 2
 
-        # Verify each rank got the correct config
-        for rank_config in scheduler.per_rank_kv_cache_configs:
-            assert rank_config.num_blocks == 50  # 100 / 2
+            # Verify each rank got the correct config
+            for rank_config in scheduler.per_rank_kv_cache_configs:
+                assert rank_config.num_blocks == 50  # 100 / 2
 
     def test_get_rank_token_counts(self, mock_vllm_config,
                                    mock_kv_cache_config,
@@ -296,9 +292,9 @@ def test_combine_cached_request_data(self, mock_vllm_config,
                                          mock_kv_cache_config,
                                          mock_structured_output_manager):
         """Test _combine_cached_request_data combines data from all ranks."""
-        with patch(
-                "tpu_inference.core.sched.dp_scheduler.resolve_obj_by_qualname"
-        ):
+        mock_scheduler_cls = MagicMock(return_value=MagicMock())
+        with patch.object(mock_vllm_config.scheduler_config,
+                          '_original_scheduler_cls', mock_scheduler_cls):
             scheduler = DPScheduler(
                 vllm_config=mock_vllm_config,
                 kv_cache_config=mock_kv_cache_config,
@@ -403,9 +399,9 @@ def test_get_grammar_bitmask_no_structured_output(
             self, mock_vllm_config, mock_kv_cache_config,
             mock_structured_output_manager):
         """Test get_grammar_bitmask returns None when no structured output."""
-        with patch(
-                "tpu_inference.core.sched.dp_scheduler.resolve_obj_by_qualname"
-        ):
+        mock_scheduler_cls = MagicMock(return_value=MagicMock())
+        with patch.object(mock_vllm_config.scheduler_config,
+                          '_original_scheduler_cls', mock_scheduler_cls):
             scheduler = DPScheduler(
                 vllm_config=mock_vllm_config,
                 kv_cache_config=mock_kv_cache_config,
@@ -452,9 +448,9 @@ def test_update_from_output_routes_to_schedulers(
             self, mock_vllm_config, mock_kv_cache_config,
             mock_structured_output_manager):
         """Test update_from_output splits output and updates each scheduler."""
-        with patch(
-                "tpu_inference.core.sched.dp_scheduler.resolve_obj_by_qualname"
-        ):
+        mock_scheduler_cls = MagicMock(return_value=MagicMock())
+        with patch.object(mock_vllm_config.scheduler_config,
+                          '_original_scheduler_cls', mock_scheduler_cls):
             scheduler = DPScheduler(
                 vllm_config=mock_vllm_config,
                 kv_cache_config=mock_kv_cache_config,
@@ -551,9 +547,9 @@ def test_split_model_output_by_rank(self, mock_vllm_config,
                                         mock_kv_cache_config,
                                         mock_structured_output_manager):
         """Test _split_model_output_by_rank distributes output correctly."""
-        with patch(
-                "tpu_inference.core.sched.dp_scheduler.resolve_obj_by_qualname"
-        ):
+        mock_scheduler_cls = MagicMock(return_value=MagicMock())
+        with patch.object(mock_vllm_config.scheduler_config,
+                          '_original_scheduler_cls', mock_scheduler_cls):
             scheduler = DPScheduler(
                 vllm_config=mock_vllm_config,
                 kv_cache_config=mock_kv_cache_config,
@@ -597,9 +593,9 @@ def test_cleanup_finished_requests(self, mock_vllm_config,
                                        mock_kv_cache_config,
                                        mock_structured_output_manager):
         """Test _cleanup_finished_requests removes finished requests."""
-        with patch(
-                "tpu_inference.core.sched.dp_scheduler.resolve_obj_by_qualname"
-        ):
+        mock_scheduler_cls = MagicMock(return_value=MagicMock())
+        with patch.object(mock_vllm_config.scheduler_config,
+                          '_original_scheduler_cls', mock_scheduler_cls):
             scheduler = DPScheduler(
                 vllm_config=mock_vllm_config,
                 kv_cache_config=mock_kv_cache_config,
@@ -669,9 +665,9 @@ def test_has_finished_requests(self, mock_vllm_config,
                                    mock_kv_cache_config,
                                    mock_structured_output_manager):
         """Test has_finished_requests checks all ranks."""
-        with patch(
-                "tpu_inference.core.sched.dp_scheduler.resolve_obj_by_qualname"
-        ):
+        mock_scheduler_cls = MagicMock(return_value=MagicMock())
+        with patch.object(mock_vllm_config.scheduler_config,
+                          '_original_scheduler_cls', mock_scheduler_cls):
             scheduler = DPScheduler(
                 vllm_config=mock_vllm_config,
                 kv_cache_config=mock_kv_cache_config,
@@ -798,9 +794,9 @@ def test_make_stats_with_logging_disabled(self, mock_vllm_config,
                                               mock_kv_cache_config,
                                               mock_structured_output_manager):
         """Test make_stats returns None when logging is disabled."""
-        with patch(
-                "tpu_inference.core.sched.dp_scheduler.resolve_obj_by_qualname"
-        ):
+        mock_scheduler_cls = MagicMock(return_value=MagicMock())
+        with patch.object(mock_vllm_config.scheduler_config,
+                          '_original_scheduler_cls', mock_scheduler_cls):
             scheduler = DPScheduler(
                 vllm_config=mock_vllm_config,
                 kv_cache_config=mock_kv_cache_config,
@@ -878,11 +874,12 @@ def test_update_config_with_dp_size_greater_than_one(self):
         mock_config.sharding_config.total_dp_size = 2
         mock_config.scheduler_config._original_scheduler_cls = None
         mock_config.scheduler_config.scheduler_cls = "vllm.v1.core.sched.scheduler.Scheduler"
+        mock_config.scheduler_config.async_scheduling = False
 
         update_vllm_config_for_dp_scheduler(mock_config)
 
         # Verify config was updated
-        assert mock_config.scheduler_config._original_scheduler_cls == "vllm.v1.core.sched.scheduler.Scheduler"
+        assert mock_config.scheduler_config._original_scheduler_cls == Scheduler
         assert mock_config.scheduler_config.scheduler_cls == DPScheduler
 
     def test_update_config_with_dp_size_one(self):

tpu_inference/core/sched/dp_scheduler.py

@@ -6,7 +6,7 @@
 import torch
 from vllm.config import VllmConfig
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm.utils.import_utils import resolve_obj_by_qualname
+from vllm.v1.core.sched.async_scheduler import AsyncScheduler
 from vllm.v1.core.sched.interface import SchedulerInterface
 from vllm.v1.core.sched.output import (CachedRequestData, GrammarOutput,
                                        SchedulerOutput)
@@ -76,8 +76,7 @@ def __init__(
         self._create_per_rank_configs(kv_cache_config)
 
         # The original scheduler class could be Scheduler or AsyncScheduler
-        original_scheduler_cls = resolve_obj_by_qualname(
-            vllm_config.scheduler_config._original_scheduler_cls)
+        original_scheduler_cls = vllm_config.scheduler_config._original_scheduler_cls
         self.schedulers: List[Scheduler] = []
         for rank in range(self.dp_size):
             scheduler = original_scheduler_cls(
@@ -92,7 +91,8 @@ def __init__(
             self.schedulers.append(scheduler)
 
         logger.info(
-            f"DPScheduler per-rank limits: max_seqs={self.vllm_config.scheduler_config.max_num_seqs}, "
+            f"DPScheduler (Async = {self.vllm_config.scheduler_config.async_scheduling}) "
+            f"per-rank limits: max_seqs={self.vllm_config.scheduler_config.max_num_seqs}, "
             f"max_tokens={self.vllm_config.scheduler_config.max_num_batched_tokens}"
         )
 
@@ -515,5 +515,9 @@ def update_vllm_config_for_dp_scheduler(vllm_config: Any) -> None:
     dp_size = vllm_config.sharding_config.total_dp_size
 
     if dp_size > 1:
-        vllm_config.scheduler_config._original_scheduler_cls = vllm_config.scheduler_config.scheduler_cls
+        if vllm_config.scheduler_config.async_scheduling:
+            vllm_config.scheduler_config._original_scheduler_cls = AsyncScheduler
+        else:
+            vllm_config.scheduler_config._original_scheduler_cls = Scheduler
+
         vllm_config.scheduler_config.scheduler_cls = DPScheduler