[CI] Add unit tests about concurrent partial prefills

Signed-off-by: Csrayz <[email protected]>

Author: Csrayz

tests/ut/core/test_schedule_config.py

@@ -51,6 +51,8 @@ def test_initialize_from_config_with_override(self):
                 num_scheduler_steps=1,
                 scheduler_cls="vllm_ascend.core.scheduler.AscendScheduler",
                 max_num_batched_tokens=2048,
+                max_long_partial_prefills=1,
+                long_prefill_token_threshold=512,
             ),
         )
         self.assertEqual(ascend_config.enable_chunked_prefill, False)
@@ -60,7 +62,8 @@ def test_initialize_from_config_with_override(self):
                          "vllm_ascend.core.scheduler.AscendScheduler")
         self.assertEqual(ascend_config.max_num_batched_tokens, 2048)
         self.assertEqual(ascend_config.encoder_cache_size, 2048)
-
+        self.assertEqual(ascend_config.max_long_partial_prefills, 1)
+        self.assertEqual(ascend_config.long_prefill_token_threshold, 512)
     def test_not_implemented_policy(self):
         with self.assertRaises(NotImplementedError) as context:
             AscendSchedulerConfig.initialize_from_config(

tests/ut/core/test_scheduler.py

@@ -226,7 +226,7 @@ def test_get_num_unfinished_requests(self):
                              len(requests) - i - 1)
 
     def test_schedule(self):
-        '''Test scheduling. 
+        '''Test scheduling.
         Two cases: default APC/no prompt logprobs; APC=True + prompt logprobs
         '''
         scheduler = self.create_scheduler()
@@ -251,6 +251,27 @@ def test_schedule(self):
         for i, request in enumerate(requests):
             self.assertEqual(scheduler.running[i], request)
 
+    def test_concurrent_partial_prefills_schedule(self):
+        '''Test concurrent partial prefills scheduling.
+        total requests = 10, every request has 10 token.
+        while set long_prefill_token_threshold = 1, scheduler can
+        only schedule max_long_partial_prefills long request.
+        '''
+        scheduler = self.create_scheduler()
+        scheduler.scheduler_config.chunked_prefill_enabled = False
+        scheduler.scheduler_config.max_long_partial_prefills = 2
+        scheduler.scheduler_config.long_prefill_token_threshold = 1
+        requests = create_requests(num_requests=10, num_tokens=20)
+        for request in requests:
+            scheduler.add_request(request)
+
+        # Test initial scheduling
+        output = scheduler.schedule()
+        self.assertEqual(len(output.scheduled_new_reqs),
+            scheduler.scheduler_config.max_long_partial_prefills)
+        self.assertEqual(output.scheduled_cached_reqs.num_reqs, 0)
+        self.assertEqual(len(output.finished_req_ids), 0)
+
     def test_schedule_enable_prefix_caching(self):
         '''Test scheduling.
         Two cases: default APC/no prompt logprobs; APC=True + prompt logprobs