test_abort (#6743)

tests/output/test_process_batch_output.py

@@ -53,6 +53,7 @@ class MockConfig:
 
     class SchedulerConfig:
         name = "default"
+        splitwise_role = "decode"
 
     class CacheConfig:
         enable_prefix_caching = False
@@ -104,15 +105,15 @@ class MockResourceManager:
         self.stop_flags = [False]
         self.tasks_list = [MockTask()]
         self.to_be_rescheduled_request_id_set = set()
+        self.to_be_aborted_req_id_set = set()
         self.abort_req_ids_set = set()
         self.req_dict = {}
+        self.recycle_abort_task = MagicMock(side_effect=lambda rid: self.to_be_aborted_req_id_set.discard(rid))
+        self.reschedule_preempt_task = MagicMock()
 
     def info(self):
         return "Mock resource manager info"
 
-    def reschedule_preempt_task(self, task_id):
-        pass
-
 
 class MockCachedGeneratedTokens:
     def __init__(self):
@@ -323,7 +324,7 @@ class TestTokenProcessorProcessBatchOutput(unittest.TestCase):
         task_id = "test_aborted_request"
         task = processor.resource_manager.tasks_list[0]
         task.request_id = task_id
-        processor.resource_manager.abort_req_ids_set = {task_id}
+        processor.resource_manager.to_be_aborted_req_id_set = {task_id}
 
         # Add the task to req_dict to prevent _recycle_aborted_task from processing it early
         # batch_id should be < batch_size - 1 to avoid early recycling
@@ -339,8 +340,8 @@ class TestTokenProcessorProcessBatchOutput(unittest.TestCase):
         # Set up output tokens with negative token
         # batch = 2 (so batch_id=0 is < batch_size-1=1)
         processor.output_tokens[1, 0].set_tensor(paddle.to_tensor(2))
-        # Set negative token for first task (batch_id=0)
-        processor.output_tokens[2, 0].set_tensor(paddle.to_tensor(-1))
+        # Set negative token (PREEMPTED_TOKEN_ID) for first task (batch_id=0)
+        processor.output_tokens[2, 0].set_tensor(paddle.to_tensor(-9))
         # Set positive token for second task (batch_id=1)
         processor.output_tokens[3, 0].set_tensor(paddle.to_tensor(100))
 
@@ -356,15 +357,15 @@ class TestTokenProcessorProcessBatchOutput(unittest.TestCase):
         # Mock llm_logger to capture the log message and envs.ENABLE_V1_KVCACHE_SCHEDULER
         with (
             patch("fastdeploy.output.token_processor.llm_logger") as mock_logger,
-            patch("fastdeploy.output.token_processor.envs.ENABLE_V1_KVCACHE_SCHEDULER", 0),
+            patch("fastdeploy.output.token_processor.envs.ENABLE_V1_KVCACHE_SCHEDULER", 1),
         ):
             # Call the method
             processor._process_batch_output()
+            print(mock_logger)
 
             # Verify the recycling logic was triggered
-            mock_logger.info.assert_any_call(f"Aborted task {task_id} received negative token. Recycling.")
-            processor._recycle_resources.assert_called_once_with(task_id, 0, task)
-            self.assertNotIn(task_id, processor.resource_manager.abort_req_ids_set)
+            processor.resource_manager.recycle_abort_task.assert_called_once_with(task_id)
+            self.assertNotIn(task_id, processor.resource_manager.to_be_aborted_req_id_set)
 
     def test_process_batch_output_non_aborted_task_negative_token(self):
         """Test non-aborted task receiving negative token does not trigger recycling"""

tests/output/test_process_batch_output_use_zmq.py

@@ -158,11 +158,14 @@ class TestTokenProcessorLogprobs(unittest.TestCase):
         # Set up task as aborted
         task_id = "test_aborted_request"
         self.task_mock.request_id = task_id
-        self.processor.resource_manager.abort_req_ids_set = {task_id}
+        self.processor.resource_manager.to_be_aborted_req_id_set = {task_id}
+        self.processor.resource_manager.recycle_abort_task = MagicMock(
+            side_effect=lambda rid: self.processor.resource_manager.to_be_aborted_req_id_set.discard(rid)
+        )
 
-        # Create stream data with negative token
+        # Create stream data with negative token (PREEMPTED_TOKEN_ID = -9)
         stream_data = MagicMock()
-        stream_data.tokens = np.array([1, 2, -1])  # Last token is negative
+        stream_data.tokens = np.array([1, 2, -9])  # Last token is PREEMPTED_TOKEN_ID
         stream_data.batch_id = 0
 
         # Mock _recycle_resources to track if it's called
@@ -171,19 +174,16 @@ class TestTokenProcessorLogprobs(unittest.TestCase):
         # Mock the llm_logger module and envs.ENABLE_V1_KVCACHE_SCHEDULER
         with (
             patch("fastdeploy.output.token_processor.llm_logger") as mock_logger,
-            patch("fastdeploy.output.token_processor.envs.ENABLE_V1_KVCACHE_SCHEDULER", 0),
+            patch("fastdeploy.output.token_processor.envs.ENABLE_V1_KVCACHE_SCHEDULER", 1),
         ):
             # Call the method
             result = self.processor._process_batch_output_use_zmq([stream_data])
 
             # Verify the recycling logic was triggered
-            mock_logger.info.assert_any_call(f"Aborted task {task_id} received negative token. Recycling.")
-            self.processor._recycle_resources.assert_called_once_with(task_id, 0, self.task_mock)
-            self.assertNotIn(task_id, self.processor.resource_manager.abort_req_ids_set)
-            self.assertEqual(len(result), 1)  # Should return abort result
-            self.assertEqual(result[0].finished, True)
-            self.assertEqual(result[0].error_code, 499)
-            self.assertIn("aborted", result[0].error_msg.lower())
+            mock_logger.info.assert_any_call(f"start to recycle abort request_id {task_id}")
+            self.processor.resource_manager.recycle_abort_task.assert_called_once_with(task_id)
+            self.assertNotIn(task_id, self.processor.resource_manager.to_be_aborted_req_id_set)
+            self.assertEqual(len(result), 0)  # Aborted task is skipped (continue)
 
     def test_process_batch_output_use_zmq_non_aborted_task_negative_token(self):
         """Test non-aborted task receiving negative token does not trigger recycling"""

tests/output/test_token_processor.py

@@ -74,6 +74,7 @@ class _DummyResourceManager:
         self.req_dict = {}
         self.requests = {}
         self.to_be_rescheduled_request_id_set = set()
+        self.to_be_aborted_req_id_set = set()
         self.abort_req_ids_set = set()
         self.recycled = []
         self.cached_tasks = []
@@ -85,6 +86,10 @@ class _DummyResourceManager:
     def reschedule_preempt_task(self, request_id):
         self.recycled.append(f"reschedule-{request_id}")
 
+    def recycle_abort_task(self, request_id):
+        self.recycled.append(f"recycle-abort-{request_id}")
+        self.to_be_aborted_req_id_set.discard(request_id)
+
     def finish_requests_async(self, request_id):
         self.recycled.append(f"finish-{request_id}")