From 407079a62fcc707be92dd2ceb151ae99f770bdf8 Mon Sep 17 00:00:00 2001
From: Mark Kurtz <mark.kurtz@neuralmagic.com>
Date: Wed, 3 Jul 2024 01:56:31 -0400
Subject: [PATCH] fixes from review comments

---
 src/guidellm/request/base.py           | 5 +++--
 src/guidellm/utils/__init__.py         | 8 ++++++--
 src/guidellm/utils/constants.py        | 4 +++-
 tests/integration/request/test_base.py | 1 +
 tests/unit/request/test_base.py        | 8 ++++++--
 5 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/src/guidellm/request/base.py b/src/guidellm/request/base.py
index d575d42..3e84819 100644
--- 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
index 92c5d75..fefad35 100644
--- 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
index 39596c2..8d1e8a0 100644
--- 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
index a631909..ffb1fd0 100644
--- 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
index c0afdcd..1871d1c 100644
--- 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()