Add environment variable for genai upload hook queue size

Author: aabmass

util/opentelemetry-util-genai/src/opentelemetry/util/genai/_upload/__init__.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import logging
 from os import environ
 
 from opentelemetry.util.genai.completion_hook import (
@@ -22,21 +23,65 @@
 )
 from opentelemetry.util.genai.environment_variables import (
     OTEL_INSTRUMENTATION_GENAI_UPLOAD_BASE_PATH,
+    OTEL_INSTRUMENTATION_GENAI_UPLOAD_FORMAT,
+    OTEL_INSTRUMENTATION_GENAI_UPLOAD_MAX_QUEUE_SIZE,
 )
 
+_logger = logging.getLogger(__name__)
+
 
 def upload_completion_hook() -> CompletionHook:
+    base_path = environ.get(OTEL_INSTRUMENTATION_GENAI_UPLOAD_BASE_PATH)
+    if not base_path:
+        _logger.warning(
+            "%s is required but not set, using no-op instead",
+            OTEL_INSTRUMENTATION_GENAI_UPLOAD_BASE_PATH,
+        )
+        return _NoOpCompletionHook()
+
     # If fsspec is not installed the hook will be a no-op.
     try:
         # pylint: disable=import-outside-toplevel
         from opentelemetry.util.genai._upload.completion_hook import (  # noqa: PLC0415
+            _DEFAULT_FORMAT,
+            _DEFAULT_MAX_QUEUE_SIZE,
+            _FORMATS,
             UploadCompletionHook,
         )
     except ImportError:
         return _NoOpCompletionHook()
 
-    base_path = environ.get(OTEL_INSTRUMENTATION_GENAI_UPLOAD_BASE_PATH)
-    if not base_path:
-        return _NoOpCompletionHook()
+    environ_max_queue_size = environ.get(
+        OTEL_INSTRUMENTATION_GENAI_UPLOAD_MAX_QUEUE_SIZE,
+        _DEFAULT_MAX_QUEUE_SIZE,
+    )
+    try:
+        environ_max_queue_size = int(environ_max_queue_size)
+    except ValueError:
+        _logger.warning(
+            "%s is not a valid integer for %s. Defaulting to %s.",
+            environ_max_queue_size,
+            OTEL_INSTRUMENTATION_GENAI_UPLOAD_MAX_QUEUE_SIZE,
+            _DEFAULT_MAX_QUEUE_SIZE,
+        )
+        environ_max_queue_size = _DEFAULT_MAX_QUEUE_SIZE
+
+    environ_format = environ.get(
+        OTEL_INSTRUMENTATION_GENAI_UPLOAD_FORMAT, _DEFAULT_FORMAT
+    ).lower()
+
+    if environ_format not in _FORMATS:
+        _logger.warning(
+            "%s is not a valid option for %s, should be be one of %s. Defaulting to %s.",
+            environ_format,
+            OTEL_INSTRUMENTATION_GENAI_UPLOAD_FORMAT,
+            _FORMATS,
+            _DEFAULT_FORMAT,
+        )
+        environ_format = _DEFAULT_FORMAT
 
-    return UploadCompletionHook(base_path=base_path)
+    return UploadCompletionHook(
+        base_path=base_path,
+        max_queue_size=environ_max_queue_size,
+        upload_format=environ_format,
+    )

util/opentelemetry-util-genai/src/opentelemetry/util/genai/_upload/completion_hook.py

@@ -24,7 +24,6 @@
 from contextlib import ExitStack
 from dataclasses import asdict, dataclass
 from functools import partial
-from os import environ
 from time import time
 from typing import Any, Callable, Final, Literal
 from uuid import uuid4
@@ -36,9 +35,6 @@
 from opentelemetry.trace import Span
 from opentelemetry.util.genai import types
 from opentelemetry.util.genai.completion_hook import CompletionHook
-from opentelemetry.util.genai.environment_variables import (
-    OTEL_INSTRUMENTATION_GENAI_UPLOAD_FORMAT,
-)
 from opentelemetry.util.genai.utils import gen_ai_json_dump
 
 GEN_AI_INPUT_MESSAGES_REF: Final = (
@@ -52,6 +48,9 @@
 )
 
 _MESSAGE_INDEX_KEY = "index"
+_DEFAULT_MAX_QUEUE_SIZE = 20
+_DEFAULT_FORMAT = "json"
+
 
 Format = Literal["json", "jsonl"]
 _FORMATS: tuple[Format, ...] = ("json", "jsonl")
@@ -105,36 +104,26 @@ def __init__(
         self,
         *,
         base_path: str,
-        max_size: int = 20,
-        upload_format: Format | None = None,
+        max_queue_size: int = _DEFAULT_MAX_QUEUE_SIZE,
+        upload_format: Format = _DEFAULT_FORMAT,
         lru_cache_max_size: int = 1024,
     ) -> None:
-        self._max_size = max_size
+        self._max_queue_size = max_queue_size
         self._fs, base_path = fsspec.url_to_fs(base_path)
         self._base_path = self._fs.unstrip_protocol(base_path)
         self.lru_dict: OrderedDict[str, bool] = OrderedDict()
         self.lru_cache_max_size = lru_cache_max_size
 
-        if upload_format not in _FORMATS + (None,):
+        if upload_format not in _FORMATS:
             raise ValueError(
                 f"Invalid {upload_format=}. Must be one of {_FORMATS}"
             )
-
-        if upload_format is None:
-            environ_format = environ.get(
-                OTEL_INSTRUMENTATION_GENAI_UPLOAD_FORMAT, "json"
-            ).lower()
-            if environ_format not in _FORMATS:
-                upload_format = "json"
-            else:
-                upload_format = environ_format
-
-        self._format: Final[Literal["json", "jsonl"]] = upload_format
+        self._format = upload_format
 
         # Use a ThreadPoolExecutor for its queueing and thread management. The semaphore
         # limits the number of queued tasks. If the queue is full, data will be dropped.
-        self._executor = ThreadPoolExecutor(max_workers=max_size)
-        self._semaphore = threading.BoundedSemaphore(max_size)
+        self._executor = ThreadPoolExecutor(max_workers=self._max_queue_size)
+        self._semaphore = threading.BoundedSemaphore(self._max_queue_size)
 
     def _submit_all(self, upload_data: UploadData) -> None:
         def done(future: Future[None]) -> None:
@@ -326,7 +315,7 @@ def shutdown(self, *, timeout_sec: float = 10.0) -> None:
 
         # Wait for all tasks to finish to flush the queue
         with ExitStack() as stack:
-            for _ in range(self._max_size):
+            for _ in range(self._max_queue_size):
                 remaining = deadline - time()
                 if not self._semaphore.acquire(timeout=remaining):  # pylint: disable=consider-using-with
                     # Couldn't finish flushing all uploads before timeout

util/opentelemetry-util-genai/src/opentelemetry/util/genai/environment_variables.py

@@ -50,3 +50,13 @@
 The format to use when uploading prompt and response data. Must be one of ``json`` or
 ``jsonl``. Defaults to ``json``.
 """
+
+OTEL_INSTRUMENTATION_GENAI_UPLOAD_MAX_QUEUE_SIZE = (
+    "OTEL_INSTRUMENTATION_GENAI_UPLOAD_MAX_QUEUE_SIZE"
+)
+"""
+.. envvar:: OTEL_INSTRUMENTATION_GENAI_UPLOAD_MAX_QUEUE_SIZE
+
+The maximum number of concurrent uploads to queue. New uploads will be dropped if the queue is
+full. Defaults to 20.
+"""

util/opentelemetry-util-genai/tests/test_upload.py

@@ -13,10 +13,8 @@
 # limitations under the License.
 
 
-# pylint: disable=import-outside-toplevel,no-name-in-module
-import importlib
+# pylint: disable=no-name-in-module
 import logging
-import sys
 import threading
 import time
 from contextlib import contextmanager
@@ -33,42 +31,10 @@
 from opentelemetry.util.genai._upload.completion_hook import (
     UploadCompletionHook,
 )
-from opentelemetry.util.genai.completion_hook import (
-    _NoOpCompletionHook,
-    load_completion_hook,
-)
 
 # Use MemoryFileSystem for testing
 # https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.implementations.memory.MemoryFileSystem
 BASE_PATH = "memory://"
-
-
-@patch.dict(
-    "os.environ",
-    {
-        "OTEL_INSTRUMENTATION_GENAI_COMPLETION_HOOK": "upload",
-        "OTEL_INSTRUMENTATION_GENAI_UPLOAD_BASE_PATH": BASE_PATH,
-    },
-    clear=True,
-)
-class TestUploadEntryPoint(TestCase):
-    def test_upload_entry_point(self):
-        self.assertIsInstance(load_completion_hook(), UploadCompletionHook)
-
-    def test_upload_entry_point_no_fsspec(self):
-        """Tests that the a no-op uploader is used when fsspec is not installed"""
-
-        from opentelemetry.util.genai import _upload  # noqa: PLC0415
-
-        # Simulate fsspec imports failing
-        with patch.dict(
-            sys.modules,
-            {"opentelemetry.util.genai._upload.completion_hook": None},
-        ):
-            importlib.reload(_upload)
-            self.assertIsInstance(load_completion_hook(), _NoOpCompletionHook)
-
-
 MAXSIZE = 5
 FAKE_INPUTS = [
     types.InputMessage(
@@ -125,7 +91,7 @@ def setUp(self):
         self.mock_fs.exists.return_value = False
 
         self.hook = UploadCompletionHook(
-            base_path=BASE_PATH, max_size=MAXSIZE, lru_cache_max_size=5
+            base_path=BASE_PATH, max_queue_size=MAXSIZE, lru_cache_max_size=5
         )
 
     def tearDown(self) -> None:
@@ -305,42 +271,6 @@ def test_upload_format_sets_content_type(self):
                 ANY, "w", content_type=expect_content_type
             )
 
-    def test_parse_upload_format_envvar(self):
-        for envvar_value, expect in (
-            ("", "json"),
-            ("json", "json"),
-            ("invalid", "json"),
-            ("jsonl", "jsonl"),
-            ("jSoNl", "jsonl"),
-        ):
-            with patch.dict(
-                "os.environ",
-                {"OTEL_INSTRUMENTATION_GENAI_UPLOAD_FORMAT": envvar_value},
-                clear=True,
-            ):
-                hook = UploadCompletionHook(base_path=BASE_PATH)
-                self.addCleanup(hook.shutdown)
-                self.assertEqual(
-                    hook._format,
-                    expect,
-                    f"expected upload format {expect=} with {envvar_value=} got {hook._format}",
-                )
-
-        with patch.dict(
-            "os.environ",
-            {"OTEL_INSTRUMENTATION_GENAI_UPLOAD_FORMAT": "json"},
-            clear=True,
-        ):
-            hook = UploadCompletionHook(
-                base_path=BASE_PATH, upload_format="jsonl"
-            )
-            self.addCleanup(hook.shutdown)
-            self.assertEqual(
-                hook._format,
-                "jsonl",
-                "upload_format kwarg should take precedence",
-            )
-
     def test_upload_after_shutdown_logs(self):
         self.hook.shutdown()
         with self.assertLogs(level=logging.INFO) as logs: