Add unitest

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

Author: Shunkang

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py

@@ -50,6 +50,7 @@ def __init__(self, dist: Distributed, enable_attention_dp: bool,
         self.waiting_queue: deque[RequestQueueItem] = deque()
         self.canceled_req_ids = []
         self.enable_attention_dp = enable_attention_dp
+        self.max_batch_size = max_batch_size
         self.max_beam_width = max_beam_width
         self.max_num_active_requests = max_num_active_requests
         self.is_disaggregated = is_disaggregated

tests/unittest/_torch/test_executor_request_queue.py

@@ -40,6 +40,7 @@ def executor_queue(mock_dist):
                                 max_beam_width=1,
                                 max_num_active_requests=16,
                                 enable_iter_perf_stats=True,
+                                batch_wait_timeout=0.0,
                                 is_disaggregated=False)
 
 
@@ -52,6 +53,7 @@ def integration_queue(mock_dist):
                                 max_beam_width=2,
                                 max_num_active_requests=8,
                                 enable_iter_perf_stats=True,
+                                batch_wait_timeout=0.0,
                                 is_disaggregated=False)
 
 
@@ -215,6 +217,75 @@ def test_get_from_request_queue_with_timeout(executor_queue):
     assert elapsed < 0.2  # Should finish within timeout
 
 
+def test_get_from_request_queue_async_behavior(executor_queue):
+    """Test asynchronous behavior where requests arrive over time."""
+    import threading
+
+    def add_requests_after_delay(delay, num_requests):
+        """Helper function to add requests after a delay."""
+        time.sleep(delay)
+        for i in range(num_requests):
+            item = RequestQueueItem(i + 10, Mock())
+            executor_queue.request_queue.put(item)
+
+    # Test 1: Without batch_wait_timeout (should only get initial requests)
+    executor_queue.batch_wait_timeout = 0.0
+
+    initial_requests = 3
+    for i in range(initial_requests):
+        item = RequestQueueItem(i, Mock())
+        executor_queue.request_queue.put(item)
+
+    thread = threading.Thread(target=add_requests_after_delay, args=(0.05, 2))
+    thread.start()
+
+    # Get requests immediately - should only get the initial ones
+    start_time = time.time()
+    items = executor_queue._get_from_request_queue(None)
+    elapsed = time.time() - start_time
+
+    assert len(items) == initial_requests
+    assert elapsed < 0.1
+    assert all(item.id < 10 for item in items)
+
+    thread.join()
+
+    # Test 2: With batch_wait_timeout (should wait and get all requests)
+    executor_queue.batch_wait_timeout = 0.2
+
+    # Clear the queue and add initial requests again
+    while not executor_queue.request_queue.empty():
+        try:
+            executor_queue.request_queue.get_nowait()
+        except queue.Empty:
+            break
+
+    initial_requests = 2
+    for i in range(initial_requests):
+        item = RequestQueueItem(i + 20, Mock())
+        executor_queue.request_queue.put(item)
+
+    thread = threading.Thread(target=add_requests_after_delay, args=(0.05, 3))
+    thread.start()
+
+    # Get requests with batch_wait_timeout - should wait and get all
+    start_time = time.time()
+    items = executor_queue._get_from_request_queue(None)
+    elapsed = time.time() - start_time
+
+    # Should wait and return all requests
+    assert len(items) == initial_requests + 3
+    assert elapsed >= 0.05
+    assert elapsed < 0.3
+
+    initial_ids = {item.id for item in items if 20 <= item.id < 30}
+    delayed_ids = {item.id for item in items if 10 <= item.id < 20}
+    assert len(initial_ids) == initial_requests
+    assert len(delayed_ids) == 3
+
+    thread.join()
+
+
 def test_get_from_waiting_queue(executor_queue):
     """Test getting items from waiting queue."""
     # Add items to waiting queue
@@ -371,6 +442,7 @@ def attention_dp_queue(mock_dist_attention_dp):
                                  max_beam_width=2,
                                  max_num_active_requests=8,
                                  enable_iter_perf_stats=True,
+                                 batch_wait_timeout=0.0,
                                  is_disaggregated=False)
     # Initialize all_ranks_num_active_requests
     return queue