perf: [LS-2561] Stream compress multipart runs

python/langsmith/_internal/_background_thread.py

@@ -1,14 +1,18 @@
 from __future__ import annotations
 
 import functools
+import io
 import logging
 import sys
 import threading
+import time
 import weakref
+import zstandard as zstd
 from queue import Empty, Queue
 from typing import (
     TYPE_CHECKING,
     List,
+    Optional,
     Union,
     cast,
 )
@@ -88,6 +92,34 @@ def _tracing_thread_drain_queue(
     return next_batch
 
 
+def _tracing_thread_drain_compressed_buffer(
+    client: Client,
+    runs_limit: int = 100,
+    max_buffer_size: int = 50 * 1024 * 1024
+) -> Optional[bytes]:
+    with client._buffer_lock:
+        current_size = client.tracing_queue.tell()
+
+        # Check if we should send now
+        if not (client._run_count >= runs_limit or current_size >= max_buffer_size):
+            return None
+
+        # Write final boundary and close compression stream
+        client.compressor_writer.write(f'--{client.boundary}--\r\n'.encode())
+        client.compressor_writer.flush()
+        client.compressor_writer.close()
+
+        client.tracing_queue.seek(0)
+        data = client.tracing_queue.getvalue()
+
+        # Reinitialize for next batch
+        client.tracing_queue = io.BytesIO()
+        client.compressor = zstd.ZstdCompressor()
+        client.compressor_writer = client.compressor.stream_writer(
+            client.tracing_queue, closefd=False)
+        client._run_count = 0
+        return data
+
 def _tracing_thread_handle_batch(
     client: Client,
     tracing_queue: Queue,
@@ -199,6 +231,19 @@ def keep_thread_active() -> bool:
     ):
         _tracing_thread_handle_batch(client, tracing_queue, next_batch, use_multipart)
 
+def tracing_control_thread_func_compress(client_ref: weakref.ref[Client]) -> None:
+    client = client_ref()
+    if client is None:
+        return
+
+    while True:
+        result = _tracing_thread_drain_compressed_buffer(client)
+        if result is not None:
+            time.sleep(0.150)  # Simulate call to backend
+        else:
+            time.sleep(0.1)  # Avoid busy-waiting if no data ready
+
+
 
 def _tracing_sub_thread_func(
     client_ref: weakref.ref[Client],

python/langsmith/_internal/_operations.py

@@ -3,6 +3,8 @@
 import itertools
 import logging
 import uuid
+import io
+import zstandard
 from typing import Literal, Optional, Union, cast
 
 from langsmith import schemas as ls_schemas
@@ -271,3 +273,36 @@ def serialized_run_operation_to_multipart_parts_and_context(
         acc_parts,
         f"trace={op.trace_id},id={op.id}",
     )
+
+
+def compress_multipart_parts_and_context(
+    parts_and_context: MultipartPartsAndContext, 
+    compressor_writer: zstandard.ZstdCompressorWriter,
+    boundary: str
+) -> None:
+    for part_name, (filename, data, content_type, headers) in parts_and_context.parts:
+        part_header = f'--{boundary}\r\n'
+        part_header += f'Content-Disposition: form-data; name="{part_name}"'
+
+        if filename:
+            part_header += f'; filename="{filename}"'
+
+        part_header += f'\r\nContent-Type: {content_type}\r\n'
+
+        for header_name, header_value in headers.items():
+            part_header += f'{header_name}: {header_value}\r\n'
+
+        part_header += '\r\n'
+        compressor_writer.write(part_header.encode())
+
+        if isinstance(data, (bytes, bytearray)):
+            with memoryview(data) as view:
+                chunk_size = 1024 * 1024  # 1MB chunks
+                for i in range(0, len(view), chunk_size):
+                    chunk = view[i:i + chunk_size]
+                    compressor_writer.write(chunk)
+        else:
+            compressor_writer.write(str(data).encode())
+
+        # Write part terminator
+        compressor_writer.write(b'\r\n')

python/langsmith/client.py

@@ -33,6 +33,7 @@
 import uuid
 import warnings
 import weakref
+import zstandard
 from inspect import signature
 from queue import PriorityQueue
 from typing import (
@@ -75,6 +76,7 @@
 )
 from langsmith._internal._background_thread import (
     tracing_control_thread_func as _tracing_control_thread_func,
+    tracing_control_thread_func_compress as _tracing_control_thread_func_compress
 )
 from langsmith._internal._beta_decorator import warn_beta
 from langsmith._internal._constants import (
@@ -94,6 +96,7 @@
     serialize_run_dict,
     serialized_feedback_operation_to_multipart_parts_and_context,
     serialized_run_operation_to_multipart_parts_and_context,
+    compress_multipart_parts_and_context,
 )
 from langsmith._internal._serde import dumps_json as _dumps_json
 
@@ -489,6 +492,16 @@ def __init__(
         # Create a session and register a finalizer to close it
         session_ = session if session else requests.Session()
         self.session = session_
+        self.compress_traces = os.getenv("LANGSMITH_COMPRESS_TRACES") == "true"
+        if self.compress_traces:
+            self.boundary = BOUNDARY
+            self.compressor = zstandard.ZstdCompressor()
+            self.compressor_writer = self.compressor.stream_writer(
+                self.tracing_queue, closefd=False)
+            self.tracing_queue = io.BytesIO()
+            self._buffer_lock = threading.Lock()
+            self._run_count = 0
+
         self._info = (
             info
             if info is None or isinstance(info, ls_schemas.LangSmithInfo)
@@ -497,7 +510,14 @@ def __init__(
         weakref.finalize(self, close_session, self.session)
         atexit.register(close_session, session_)
         # Initialize auto batching
-        if auto_batch_tracing:
+        if auto_batch_tracing and self.compress_traces:
+            threading.Thread(
+                target=_tracing_control_thread_func_compress,
+                # arg must be a weakref to self to avoid the Thread object
+                # preventing garbage collection of the Client object
+                args=(weakref.ref(self),),
+            ).start()
+        elif auto_batch_tracing:
             self.tracing_queue: Optional[PriorityQueue] = PriorityQueue()
 
             threading.Thread(
@@ -1291,9 +1311,17 @@ def create_run(
                 self._pyo3_client.create_run(run_create)
             elif self.tracing_queue is not None:
                 serialized_op = serialize_run_dict("post", run_create)
-                self.tracing_queue.put(
-                    TracingQueueItem(run_create["dotted_order"], serialized_op)
-                )
+                if self.compress_traces:
+                    multipart_form = self.serialized_run_operation_to_multipart_parts_and_context(
+                        serialized_op)
+                    with self._buffer_lock:
+                        compress_multipart_parts_and_context(
+                            multipart_form, self.compressor_writer, self.boundary)
+                        self._run_count += 1
+                else:
+                    self.tracing_queue.put(
+                        TracingQueueItem(run_create["dotted_order"], serialized_op)
+                    )
             else:
                 # Neither Rust nor Python batch ingestion is configured,
                 # fall back to the non-batch approach.
@@ -1755,9 +1783,16 @@ def update_run(
         if use_multipart and self.tracing_queue is not None:
             # not collecting attachments currently, use empty dict
             serialized_op = serialize_run_dict(operation="patch", payload=data)
-            self.tracing_queue.put(
-                TracingQueueItem(data["dotted_order"], serialized_op)
-            )
+            if self.compress_traces:
+                multipart_form = serialized_run_operation_to_multipart_parts_and_context(serialized_op)
+                with self._buffer_lock:
+                    compress_multipart_parts_and_context(
+                        multipart_form, self.compressor_writer, self.boundary)
+                    self._run_count += 1
+            else:
+                self.tracing_queue.put(
+                    TracingQueueItem(data["dotted_order"], serialized_op)
+                )
         else:
             self._update_run(data)