gen ai uploader timeout and fix flaky bugs

Author: aabmass

util/opentelemetry-util-genai/src/opentelemetry/util/genai/_fsspec_upload/fsspec_hook.py

@@ -20,8 +20,10 @@
 import posixpath
 import threading
 from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import ExitStack
 from dataclasses import asdict, dataclass
 from functools import partial
+from time import time
 from typing import Any, Callable, Final, Literal, TextIO, cast
 from uuid import uuid4
 
@@ -102,12 +104,12 @@ def __init__(
 
     def _submit_all(self, upload_data: UploadData) -> None:
         def done(future: Future[None]) -> None:
-            self._semaphore.release()
-
             try:
                 future.result()
             except Exception:  # pylint: disable=broad-except
                 _logger.exception("fsspec uploader failed")
+            finally:
+                self._semaphore.release()
 
         for path, json_encodeable in upload_data.items():
             # could not acquire, drop data
@@ -127,7 +129,7 @@ def done(future: Future[None]) -> None:
                 _logger.info(
                     "attempting to upload file after FsspecUploadHook.shutdown() was already called"
                 )
-                break
+                self._semaphore.release()
 
     def _calculate_ref_path(self) -> CompletionRefs:
         # TODO: experimental with using the trace_id and span_id, or fetching
@@ -203,6 +205,17 @@ def to_dict(
                 **references,
             }
 
-    def shutdown(self) -> None:
-        # TODO: support timeout
-        self._executor.shutdown()
+    def shutdown(self, timeout_sec: float = 10.0) -> None:
+        deadline = time() + timeout_sec
+
+        # Wait for all tasks to finish to flush the queue
+        with ExitStack() as stack:
+            for _ in range(self._max_size):
+                if not self._semaphore.acquire(timeout=deadline - time()):  # pylint: disable=consider-using-with
+                    # Couldn't finish flushing all uploads before timeout
+                    break
+
+                stack.callback(self._semaphore.release)
+
+            # Queue is flushed and blocked, start shutdown
+            self._executor.shutdown(wait=False)

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

@@ -85,10 +85,21 @@ def test_fsspec_entry_point_no_fsspec(self):
 FAKE_SYSTEM_INSTRUCTION = [types.Text(content="You are a helpful assistant.")]
 
 
+class ThreadSafeMagicMock(MagicMock):
+    def __init__(self, *args, **kwargs) -> None:
+        self.__dict__["_lock"] = threading.Lock()
+        super().__init__(*args, **kwargs)
+
+    def _increment_mock_call(self, /, *args, **kwargs):
+        with self.__dict__["_lock"]:
+            super()._increment_mock_call(*args, **kwargs)
+
+
 class TestFsspecUploadHook(TestCase):
     def setUp(self):
         self._fsspec_patcher = patch(
-            "opentelemetry.util.genai._fsspec_upload.fsspec_hook.fsspec"
+            "opentelemetry.util.genai._fsspec_upload.fsspec_hook.fsspec",
+            new=ThreadSafeMagicMock(),
         )
         self.mock_fsspec = self._fsspec_patcher.start()
         self.hook = FsspecUploadHook(
@@ -155,10 +166,7 @@ def blocked_upload(*args: Any):
         unblock_upload.set()
 
     def test_failed_upload_logs(self):
-        def failing_upload(*args: Any) -> None:
-            raise RuntimeError("failed to upload")
-
-        self.mock_fsspec.open = MagicMock(wraps=failing_upload)
+        self.mock_fsspec.open.side_effect = RuntimeError("failed to upload")
 
         with self.assertLogs(level=logging.ERROR) as logs:
             self.hook.upload(
@@ -178,7 +186,7 @@ def test_upload_after_shutdown_logs(self):
                 outputs=FAKE_OUTPUTS,
                 system_instruction=FAKE_SYSTEM_INSTRUCTION,
             )
-        self.assertEqual(len(logs.output), 1)
+        self.assertEqual(len(logs.output), 3)
         self.assertIn(
             "attempting to upload file after FsspecUploadHook.shutdown() was already called",
             logs.output[0],