fixes from review comments

neuralmagic · Jul 3, 2024 · 407079a · 407079a
1 parent 6a3ada2
commit 407079a
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 7 deletions.
diff --git a/src/guidellm/request/base.py b/src/guidellm/request/base.py
@@ -8,6 +8,7 @@
 from transformers import AutoTokenizer, PreTrainedTokenizer
 
 from guidellm.core.request import TextGenerationRequest
+from guidellm.utils import STANDARD_SLEEP_INTERVAL
 
 __all__ = ["RequestGenerator"]
 
@@ -146,13 +147,13 @@ def _populate_queue(self):
             try:
                 if self._queue.qsize() < self._async_queue_size:
                     item = self.create_item()
-                    self._queue.put(item, timeout=0.1)
+                    self._queue.put(item, timeout=STANDARD_SLEEP_INTERVAL)
                     logger.debug(
                         "Item added to queue. Current queue size: {}",
                         self._queue.qsize(),
                     )
                 else:
-                    time.sleep(0.1)
+                    time.sleep(STANDARD_SLEEP_INTERVAL)
             except Full:
                 continue
         logger.info("RequestGenerator stopped populating queue")
diff --git a/src/guidellm/utils/__init__.py b/src/guidellm/utils/__init__.py
@@ -1,3 +1,7 @@
-from .constants import PREFERRED_DATA_COLUMNS, PREFERRED_DATA_SPLITS
+from .constants import (
+    PREFERRED_DATA_COLUMNS,
+    PREFERRED_DATA_SPLITS,
+    STANDARD_SLEEP_INTERVAL,
+)
 
-__all__ = ["PREFERRED_DATA_COLUMNS", "PREFERRED_DATA_SPLITS"]
+__all__ = ["PREFERRED_DATA_COLUMNS", "PREFERRED_DATA_SPLITS", "STANDARD_SLEEP_INTERVAL"]
diff --git a/src/guidellm/utils/constants.py b/src/guidellm/utils/constants.py
@@ -1,4 +1,4 @@
-__all__ = ["PREFERRED_DATA_COLUMNS", "PREFERRED_DATA_SPLITS"]
+__all__ = ["PREFERRED_DATA_COLUMNS", "PREFERRED_DATA_SPLITS", "STANDARD_SLEEP_INTERVAL"]
 
 
 PREFERRED_DATA_COLUMNS = [
@@ -15,3 +15,5 @@
 ]
 
 PREFERRED_DATA_SPLITS = ["test", "validation", "train"]
+
+STANDARD_SLEEP_INTERVAL = 0.1
diff --git a/tests/integration/request/test_base.py b/tests/integration/request/test_base.py
@@ -1,5 +1,6 @@
 import pytest
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
 from guidellm.core.request import TextGenerationRequest
 from guidellm.request.base import RequestGenerator
 

diff --git a/tests/unit/request/test_base.py b/tests/unit/request/test_base.py
@@ -1,3 +1,4 @@
+import time
 from unittest.mock import Mock, patch
 
 import pytest
@@ -79,6 +80,7 @@ def test_request_generator_repr():
 @pytest.mark.regression
 def test_request_generator_create_item_not_implemented():
     with pytest.raises(TypeError):
+
         class IncompleteRequestGenerator(RequestGenerator):
             pass
 
@@ -106,7 +108,7 @@ def test_request_generator_iter_calls_create_item():
         if len(items) == 5:
             break
 
-    assert len(items) == 5
+    assert generator._queue.qsize() == 0
     generator.create_item.assert_called()
 
 
@@ -124,5 +126,7 @@ def test_request_generator_async_iter_calls_create_item():
             break
 
     generator.stop()
-    assert len(items) == 5
+    stop_size = generator._queue.qsize()
+    time.sleep(0.1)
+    assert generator._queue.qsize() == stop_size
     generator.create_item.assert_called()