Get rid of FsspecUploader separate class

Author: aabmass

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

@@ -27,7 +27,6 @@ def fsspec_upload_hook() -> UploadHook:
     try:
         # pylint: disable=import-outside-toplevel
         from opentelemetry.util.genai._fsspec_upload.fsspec_hook import (
-            FsspecUploader,
             FsspecUploadHook,
         )
     except ImportError:
@@ -37,7 +36,4 @@ def fsspec_upload_hook() -> UploadHook:
     if not base_path:
         return _NoOpUploadHook()
 
-    return FsspecUploadHook(
-        uploader=FsspecUploader(),
-        base_path=base_path,
-    )
+    return FsspecUploadHook(base_path=base_path)

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

@@ -60,22 +60,6 @@ def fsspec_open(urlpath: str, mode: Literal["w"]) -> TextIO:
     return cast(TextIO, fsspec.open(urlpath, mode))  # pyright: ignore[reportUnknownMemberType]
 
 
-class FsspecUploader:
-    """Implements uploading GenAI completions to a generic backend using fsspec
-
-    This class is used by the `BatchUploadHook` to upload completions to an external
-    storage.
-    """
-
-    def upload(  # pylint: disable=no-self-use
-        self,
-        path: str,
-        json_encodeable: Callable[[], JsonEncodeable],
-    ) -> None:
-        with fsspec_open(path, "w") as file:
-            json.dump(json_encodeable(), file, separators=(",", ":"))
-
-
 class FsspecUploadHook(UploadHook):
     """An upload hook using ``fsspec`` to upload to external storage
 
@@ -93,12 +77,10 @@ class FsspecUploadHook(UploadHook):
     def __init__(
         self,
         *,
-        uploader: FsspecUploader,
         base_path: str,
         max_size: int = 20,
     ) -> None:
         self._base_path = base_path
-        self._uploader = uploader
         self._max_size = max_size
 
         # Use a ThreadPoolExecutor for its queueing and thread management. The semaphore
@@ -126,7 +108,7 @@ def done(future: Future[None]) -> None:
 
             try:
                 fut = self._executor.submit(
-                    self._uploader.upload, path, json_encodeable
+                    self._do_upload, path, json_encodeable
                 )
                 fut.add_done_callback(done)
             except RuntimeError:
@@ -152,6 +134,13 @@ def _calculate_ref_path(self) -> CompletionRefs:
             ),
         )
 
+    @staticmethod
+    def _do_upload(
+        path: str, json_encodeable: Callable[[], JsonEncodeable]
+    ) -> None:
+        with fsspec_open(path, "w") as file:
+            json.dump(json_encodeable(), file, separators=(",", ":"))
+
     def upload(
         self,
         *,

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

@@ -30,7 +30,6 @@
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.util.genai import types
 from opentelemetry.util.genai._fsspec_upload.fsspec_hook import (
-    FsspecUploader,
     FsspecUploadHook,
 )
 from opentelemetry.util.genai.upload_hook import (
@@ -88,15 +87,18 @@ def test_fsspec_entry_point_no_fsspec(self):
 
 class TestFsspecUploadHook(TestCase):
     def setUp(self):
-        self.mock_uploader = MagicMock(spec=FsspecUploader)
+        self._fsspec_patcher = patch(
+            "opentelemetry.util.genai._fsspec_upload.fsspec_hook.fsspec"
+        )
+        self.mock_fsspec = self._fsspec_patcher.start()
         self.hook = FsspecUploadHook(
-            uploader=self.mock_uploader,
             base_path=BASE_PATH,
             max_size=MAXSIZE,
         )
 
     def tearDown(self) -> None:
         self.hook.shutdown()
+        self._fsspec_patcher.stop()
 
     def test_shutdown_no_items(self):
         self.hook.shutdown()
@@ -111,7 +113,7 @@ def test_upload_then_shutdown(self):
         self.hook.shutdown()
 
         self.assertEqual(
-            self.mock_uploader.upload.call_count,
+            self.mock_fsspec.open.call_count,
             3,
             "should have uploaded 3 files",
         )
@@ -122,7 +124,7 @@ def test_upload_blocked(self):
         def blocked_upload(*args: Any) -> None:
             unblock_upload.wait()
 
-        self.mock_uploader.upload = MagicMock(wraps=blocked_upload)
+        self.mock_fsspec.open = MagicMock(wraps=blocked_upload)
 
         # fill the queue
         for _ in range(MAXSIZE):
@@ -133,7 +135,7 @@ def blocked_upload(*args: Any) -> None:
             )
 
         self.assertLessEqual(
-            self.mock_uploader.upload.call_count,
+            self.mock_fsspec.open.call_count,
             MAXSIZE,
             f"uploader should only be called {MAXSIZE=} times",
         )
@@ -155,7 +157,7 @@ def test_failed_upload_logs(self):
         def failing_upload(*args: Any) -> None:
             raise RuntimeError("failed to upload")
 
-        self.mock_uploader.upload = MagicMock(wraps=failing_upload)
+        self.mock_fsspec.open = MagicMock(wraps=failing_upload)
 
         with self.assertLogs(level=logging.ERROR) as logs:
             self.hook.upload(
@@ -184,8 +186,7 @@ def test_upload_after_shutdown_logs(self):
 
 class FsspecUploaderTest(TestCase):
     def test_upload(self):
-        uploader = FsspecUploader()
-        uploader.upload(
+        FsspecUploadHook._do_upload(
             "memory://my_path",
             lambda: [asdict(fake_input) for fake_input in FAKE_INPUTS],
         )
@@ -207,7 +208,6 @@ def assert_fsspec_equal(self, path: str, value: str) -> None:
 
     def test_upload_completions(self):
         hook = FsspecUploadHook(
-            uploader=FsspecUploader(),
             base_path=BASE_PATH,
         )
         hook.upload(