patch executor queue mocks

Signed-off-by: Olya Kozlova <[email protected]>

Author: evezhier

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py

@@ -163,23 +163,18 @@ def _can_process_attention_dp_request(
 
     def _get_request_id(self):
         # (next_request_id + 1) % UINT64_MAX
+        current_id = self.next_request_id
         self.next_request_id = (self.next_request_id + 1) & ((1 << 64) - 1)
-        return self.next_request_id
+        return current_id
 
     def _generate_child_request_ids(
             self, request: ExecutorRequest) -> List[int] | None:
         """ Generate child request IDs if needed. """
         child_req_ids = None
-        sampling_config = request.sampling_config
-        beam_width = (sampling_config.beam_width
-                      if sampling_config.beam_width else 1)
-        num_return_sequences = (sampling_config.num_return_sequences
-                                if sampling_config.num_return_sequences else 1)
-
-        # Create child requests if beam_width == 1 and num_return_sequences > 1.
-        if beam_width == 1 and num_return_sequences > 1:
+        num_children = self._get_num_child_requests(request)
+        if num_children > 0:
             child_req_ids = []
-            for _ in range(num_return_sequences - 1):
+            for _ in range(num_children):
                 child_req_id = self._get_request_id()
                 if self.enable_iter_perf_stats:
                     self.start_times[child_req_id] = time.time()
@@ -599,8 +594,8 @@ def _merge_requests(self, new_requests: list[RequestQueueItem]):
                     req_item.id, req_item.request, req_item.child_req_ids,
                     self._should_exclude_last_generation_logits())
                 req_with_children.append(req)
-                if req.children:
-                    req_with_children.extend(req.children)
+                if req.child_requests:
+                    req_with_children.extend(req.child_requests)
             return req_with_children
 
     def _merge_star_attention_requests(self,

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -334,7 +334,7 @@ def __init__(
                                   return_log_probs, return_context_logits,
                                   return_generation_logits,
                                   exclude_last_generation_logits)
-        self.children = []
+        self.child_requests = []
 
     def is_generation_only_request(self):
         return self.py_llm_request_type == LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY
@@ -377,14 +377,14 @@ def create_child_request(self, child_id):
         py_request.py_batch_idx = None
         py_request.py_seq_slot = None
 
-        py_request.children = []
+        py_request.child_requests = []
 
         assert py_request.is_child
         assert py_request.request_id == child.request_id
         assert py_request.parent_request_id == self.request_id
         assert py_request.sampling_config.random_seed != self.sampling_config.random_seed
 
-        self.children.append(py_request)
+        self.child_requests.append(py_request)
 
 
 def convert_wordlist(word_list) -> List[List[int]]:

tests/unittest/_torch/test_best_of_n.py

@@ -66,10 +66,10 @@ def test_create_child_request(n: int):
             parent.request_id + parent.sampling_config.num_return_sequences):
         parent.create_child_request(child_id)
 
-    assert len(
-        parent.children) == parent.sampling_config.num_return_sequences - 1
+    assert len(parent.child_requests
+               ) == parent.sampling_config.num_return_sequences - 1
 
-    for ind, child in enumerate(parent.children):
+    for ind, child in enumerate(parent.child_requests):
         assert child.request_id == ind + parent.request_id + 1
         assert child.py_request_id == child.request_id
         assert child.parent_request_id == parent.request_id
@@ -88,7 +88,7 @@ def test_create_child_request(n: int):
         assert child.get_tokens() == parent.get_tokens()
         assert child.get_tokens() is not parent.get_tokens()
 
-        assert child.children == []
+        assert child.child_requests == []
 
 
 @force_ampere  # Save H100 resource

tests/unittest/_torch/test_executor_request_queue.py

@@ -75,7 +75,8 @@ def test_enqueue_requests(executor_queue):
     """Test enqueuing multiple requests."""
     mock_requests = [Mock(), Mock(), Mock()]
 
-    with patch('time.time', return_value=1234.5):
+    with (patch('time.time', return_value=1234.5),
+          patch.object(executor_queue, '_generate_child_request_ids')):
         req_ids = executor_queue.enqueue_requests(mock_requests)  # type: ignore
 
     assert len(req_ids) == 3
@@ -92,7 +93,8 @@ def test_enqueue_request_single(executor_queue):
     """Test enqueuing a single request."""
     mock_request = Mock()
 
-    with patch('time.time', return_value=1234.5):
+    with (patch('time.time', return_value=1234.5),
+          patch.object(executor_queue, '_generate_child_request_ids')):
         req_id = executor_queue.enqueue_request(mock_request)
 
     assert req_id == 8
@@ -104,8 +106,8 @@ def test_enqueue_request_with_query(executor_queue):
     """Test enqueuing a request with query data."""
     mock_request = Mock()
     query_data = [1, 2, 3, 4]
-
-    req_id = executor_queue.enqueue_request(mock_request, query=query_data)
+    with patch.object(executor_queue, '_generate_child_request_ids'):
+        req_id = executor_queue.enqueue_request(mock_request, query=query_data)
 
     assert req_id == 8
 
@@ -115,6 +117,31 @@ def test_enqueue_request_with_query(executor_queue):
     assert item.request == mock_request
 
 
+@pytest.mark.parametrize("n_children", [0, 1, 2])
+def test_enqueue_request_with_child_ids(executor_queue, n_children):
+    """Test enqueuing a request with query data."""
+    mock_request = Mock()
+    query_data = [1, 2, 3, 4]
+    with patch.object(executor_queue,
+                      '_get_num_child_requests') as mock_children:
+        mock_children.return_value = n_children
+        req_id = executor_queue.enqueue_request(mock_request, query=query_data)
+
+    assert req_id == 8
+
+    # Verify the item was enqueued with child ids
+    item = executor_queue.request_queue.get_nowait()
+    assert item.id == req_id
+    assert item.request == mock_request
+    if n_children == 0:
+        assert item.child_req_ids is None
+    else:
+        assert item.child_req_ids is not None
+        assert len(item.child_req_ids) == n_children
+        assert item.child_req_ids == list(
+            range(1 + req_id, 1 + req_id + n_children))
+
+
 def test_enqueue_cancel_request(executor_queue):
     """Test enqueuing a cancel request."""
     req_id = 42
@@ -253,11 +280,10 @@ def test_validate_and_filter_requests(executor_queue):
 )
 def test_merge_requests_default(mock_convert, executor_queue):
     """Test merging requests with default configuration."""
-    mock_llm_request = Mock()
+    mock_llm_request = Mock(child_requests=[])
     mock_convert.return_value = mock_llm_request
 
     requests = [RequestQueueItem(1, Mock()), RequestQueueItem(2, Mock())]
-
     result = executor_queue._merge_requests(requests)
 
     assert len(result) == 2