Draft of zstd compression in the PyO3 client implementation.

python/langsmith/client.py

@@ -1,4 +1,4 @@
 """Client for interacting with the LangSmith API.
 
 Use the client to customize API keys / workspace ocnnections, SSl certs,
 etc. for tracing.
@@ -556,6 +556,9 @@
                 batch_size = 100
                 batch_timeout_millis = 1000
                 worker_threads = 1
+                compression_level = int(
+                    ls_utils.get_env_var("RUN_COMPRESSION_LEVEL", 3)
+                )
 
                 try:
                     self._pyo3_client = langsmith_pyo3.BlockingTracingClient(
@@ -565,6 +568,7 @@
                         batch_size,
                         batch_timeout_millis,
                         worker_threads,
+                        compression_level,
                     )
                 except Exception as e:
                     logger.warning(

rust/Cargo.toml

@@ -11,6 +11,7 @@ resolver = "2"
 chrono = "0.4.38"
 flate2 = "1.0.34"
 futures = "0.3.31"
+http = "1.2.0"
 rayon = "1.10.0"
 serde = { version = "1.0.210", features = ["derive"] }
 serde_json = "1.0.128"
@@ -20,6 +21,7 @@ tokio = { version = "1", features = ["full"] }
 tokio-util = "0.7.12"
 ureq = "2.10.1"
 uuid = { version = "1.11.0", features = ["v4"] }
+zstd = { version = "0.13.2", features = ["zstdmt"] }
 
 # Use rustls instead of OpenSSL, because OpenSSL is a nightmare when compiling across platforms.
 # OpenSSL is a default feature, so we have to disable all default features, then re-add

rust/crates/langsmith-pyo3/src/blocking_tracing_client.rs

@@ -27,6 +27,7 @@ impl BlockingTracingClient {
         batch_size: usize,
         batch_timeout_millis: u64,
         worker_threads: usize,
+        compression_level: i32,
     ) -> PyResult<Self> {
         let config = langsmith_tracing_client::client::blocking::ClientConfig {
             endpoint,
@@ -39,6 +40,7 @@ impl BlockingTracingClient {
 
             headers: None, // TODO: support custom headers
             num_worker_threads: worker_threads,
+            compression_level,
         };
 
         let client = RustTracingClient::new(config)

rust/crates/langsmith-tracing-client/Cargo.toml

@@ -17,6 +17,8 @@ futures = { workspace = true }
 rayon = { workspace = true }
 ureq = { workspace = true }
 flate2 = { workspace = true }
+zstd = { workspace = true }
+http = { workspace = true }
 
 [dev-dependencies]
 multer = "3.1.0"

rust/crates/langsmith-tracing-client/src/client/blocking/processor.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::sync::mpsc::{Receiver, Sender};
 use std::sync::{mpsc, Arc, Mutex};
+use std::thread::available_parallelism;
 use std::time::{Duration, Instant};
 
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
@@ -266,14 +267,34 @@ impl RunProcessor {
         for (part_name, part) in json_parts.into_iter().chain(attachment_parts) {
             form = form.part(part_name, part);
         }
+        let content_type = format!("multipart/form-data; boundary={}", form.boundary());
+
+        let compressed_data = {
+            let mut buffer = Vec::with_capacity(4096);
+            let mut encoder = zstd::Encoder::new(&mut buffer, self.config.compression_level)
+                .and_then(|mut encoder| {
+                    encoder.multithread(
+                        available_parallelism().map(|p| p.get() as u32).unwrap_or_default(),
+                    )?;
+                    Ok(encoder)
+                })
+                .map_err(|e| TracingClientError::IoError(format!("{e}")))?;
+            std::io::copy(&mut form.reader(), &mut encoder)
+                .map_err(|e| TracingClientError::IoError(format!("{e}")))?;
+            encoder.finish().map_err(|e| TracingClientError::IoError(format!("{e}")))?;
+
+            buffer
+        };
 
         // send the multipart POST request
         let start_send_batch = Instant::now();
         let response = self
             .http_client
             .post(format!("{}/runs/multipart", self.config.endpoint))
-            .multipart(form)
             .headers(self.config.headers.as_ref().cloned().unwrap_or_default())
+            .header(http::header::CONTENT_TYPE, content_type)
+            .header(http::header::CONTENT_ENCODING, "zstd")
+            .body(compressed_data)
             .send()?;
         // println!("Sending batch took {:?}", start_send_batch.elapsed());
 

rust/crates/langsmith-tracing-client/src/client/blocking/tracing_client.rs

@@ -19,6 +19,7 @@ pub struct ClientConfig {
     pub batch_timeout: Duration,
     pub headers: Option<HeaderMap>,
     pub num_worker_threads: usize,
+    pub compression_level: i32,
 }
 
 pub struct TracingClient {