[Test] Fix broken tests due to upstream change. (#1228)

Signed-off-by: Pooya Moradi <pooyam@google.com>
Co-authored-by: Pooya Moradi <pooyam@google.com>

Author: py4

tests/runner/test_kv_cache_manager.py

@@ -49,7 +49,9 @@ def setup_method(self):
                 swap_space=4,
                 cache_dtype="auto",
             )
-            scheduler_config = SchedulerConfig(max_num_seqs=16, )
+            scheduler_config = SchedulerConfig(max_num_seqs=16,
+                                               max_model_len=1024,
+                                               is_encoder_decoder=False)
             parallel_config = ParallelConfig(
                 pipeline_parallel_size=1,
                 tensor_parallel_size=1,

tests/runner/test_multimodal_manager.py

@@ -42,7 +42,9 @@ def setup_method(self):
                 swap_space=4,
                 cache_dtype="auto",
             )
-            scheduler_config = SchedulerConfig(max_num_seqs=16, )
+            scheduler_config = SchedulerConfig(max_num_seqs=16,
+                                               max_model_len=1024,
+                                               is_encoder_decoder=False)
             parallel_config = ParallelConfig(
                 pipeline_parallel_size=1,
                 tensor_parallel_size=1,

tests/runner/test_speculative_decoding_manager.py

@@ -41,7 +41,9 @@ def setup_method(self):
                 swap_space=4,
                 cache_dtype="auto",
             )
-            scheduler_config = SchedulerConfig(max_num_seqs=16, )
+            scheduler_config = SchedulerConfig(max_num_seqs=16,
+                                               max_model_len=1024,
+                                               is_encoder_decoder=False)
             parallel_config = ParallelConfig(
                 pipeline_parallel_size=1,
                 tensor_parallel_size=1,

tests/runner/test_structured_decoding_manager.py

@@ -39,7 +39,9 @@ def setup_method(self):
                 swap_space=4,
                 cache_dtype="auto",
             )
-            scheduler_config = SchedulerConfig(max_num_seqs=16, )
+            scheduler_config = SchedulerConfig(max_num_seqs=16,
+                                               max_model_len=1024,
+                                               is_encoder_decoder=False)
             parallel_config = ParallelConfig(
                 pipeline_parallel_size=1,
                 tensor_parallel_size=1,

tests/runner/test_tpu_runner.py

@@ -34,7 +34,9 @@ def setup_method(self):
                 swap_space=4,
                 cache_dtype="auto",
             )
-            scheduler_config = SchedulerConfig(max_num_seqs=16, )
+            scheduler_config = SchedulerConfig(max_num_seqs=16,
+                                               max_model_len=1024,
+                                               is_encoder_decoder=False)
             parallel_config = ParallelConfig(
                 pipeline_parallel_size=1,
                 tensor_parallel_size=1,
@@ -131,7 +133,9 @@ def setup_method(self):
                 swap_space=4,
                 cache_dtype="auto",
             )
-            scheduler_config = SchedulerConfig(max_num_seqs=16, )
+            scheduler_config = SchedulerConfig(max_num_seqs=16,
+                                               max_model_len=1024,
+                                               is_encoder_decoder=False)
             parallel_config = ParallelConfig(
                 pipeline_parallel_size=1,
                 tensor_parallel_size=1,

tests/spec_decode/test_eagle3.py

@@ -38,16 +38,19 @@ def _create_proposer(
         num_speculative_tokens=num_speculative_tokens,
     )
 
-    vllm_config = VllmConfig(
-        model_config=model_config,
-        cache_config=CacheConfig(block_size=16),
-        speculative_config=speculative_config,
-        device_config=DeviceConfig(device="tpu"),
-        parallel_config=ParallelConfig(pipeline_parallel_size=1,
-                                       tensor_parallel_size=1),
-        load_config=LoadConfig(),
-        scheduler_config=SchedulerConfig(max_num_batched_tokens=8192,
-                                         max_num_seqs=128))
+    vllm_config = VllmConfig(model_config=model_config,
+                             cache_config=CacheConfig(block_size=16),
+                             speculative_config=speculative_config,
+                             device_config=DeviceConfig(device="tpu"),
+                             parallel_config=ParallelConfig(
+                                 pipeline_parallel_size=1,
+                                 tensor_parallel_size=1),
+                             load_config=LoadConfig(),
+                             scheduler_config=SchedulerConfig(
+                                 max_num_batched_tokens=8192,
+                                 max_num_seqs=128,
+                                 max_model_len=model_config.max_model_len,
+                                 is_encoder_decoder=False))
 
     # Mock the runner, as the proposer needs it for initialization
     mock_runner = mock.MagicMock()