signal bg threads data is available instead of sleeping

langchain-ai · Dec 13, 2024 · 0e06bde · 0e06bde
1 parent efa4bd6
commit 0e06bde
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 18 deletions.
diff --git a/python/langsmith/_internal/_background_thread.py b/python/langsmith/_internal/_background_thread.py
@@ -6,7 +6,6 @@
 import logging
 import sys
 import threading
-import time
 import weakref
 from multiprocessing import cpu_count
 from queue import Empty, Queue
@@ -276,23 +275,25 @@ def keep_thread_active() -> bool:
             # for now, keep thread alive
             return True
 
-    while keep_thread_active():
-        try:
-            data_stream = _tracing_thread_drain_compressed_buffer(
-                client, size_limit, size_limit_bytes
-            )
-            if data_stream is not None:
-                try:
-                    HTTP_REQUEST_THREAD_POOL.submit(
-                        client._send_compressed_multipart_req, data_stream
-                    )
-                except RuntimeError:
-                    client._send_compressed_multipart_req(data_stream)
-            else:
-                time.sleep(0.05)
-        except Exception:
-            logger.error("Error in tracing compression thread", exc_info=True)
-            time.sleep(0.1)  # Wait before retrying on error
+    while True:
+        triggered = client._data_available_event.wait(timeout=0.05)
+        if not keep_thread_active():
+            break
+        if not triggered:
+            continue
+        client._data_available_event.clear()
+
+        data_stream = _tracing_thread_drain_compressed_buffer(
+            client, size_limit, size_limit_bytes
+        )
+
+        if data_stream is not None:
+            try:
+                HTTP_REQUEST_THREAD_POOL.submit(
+                    client._send_compressed_multipart_req, data_stream
+                )
+            except RuntimeError:
+                client._send_compressed_multipart_req(data_stream)
 
     # Drain the buffer on exit
     try:

diff --git a/python/langsmith/client.py b/python/langsmith/client.py
@@ -402,6 +402,7 @@ class Client:
         "_run_count",
         "_buffer_lock",
         "compressed_runs_buffer",
+        "_data_available_event",
     ]
 
     def __init__(
@@ -511,6 +512,7 @@ def __init__(
                 level=3, threads=-1
             ).stream_writer(self.compressed_runs_buffer, closefd=False)
             self._buffer_lock: threading.Lock = threading.Lock()
+            self._data_available_event = threading.Event()
             self._run_count: int = 0
         else:
             self.compressed_runs_buffer = None
@@ -1334,6 +1336,7 @@ def create_run(
                         multipart_form, self.compressor_writer, self.boundary
                     )
                     self._run_count += 1
+                    self._data_available_event.set()
             elif self.tracing_queue is not None:
                 serialized_op = serialize_run_dict("post", run_create)
                 self.tracing_queue.put(
@@ -1863,6 +1866,7 @@ def update_run(
                         multipart_form, self.compressor_writer, self.boundary
                     )
                     self._run_count += 1
+                    self._data_available_event.set()
             elif self.tracing_queue is not None:
                 self.tracing_queue.put(
                     TracingQueueItem(data["dotted_order"], serialized_op)