Implement zstd compression for multipart payloads on the Rust side. #1401
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
obi1kenobi
added a commit
that referenced
this pull request
Jan 9, 2025
Requires #1401 to merge first, and will need a re-lock of poetry dependencies when that happens. Tested end-to-end by sending zstd-compressed traces to LangSmith servers.
I've tested this with both the Python and Rust changes applied, and I've successfully sent zstd-compressed traces to LangSmith servers. This introduces breaking changes in the APIs for both `langsmith-tracing-client` and `langsmith-pyo3`, so I'm bumping the left-most non-zero version number in each of their manifests. In Rust, the left-most non-zero number is considered a "major" version -- in other words, leading zeroes are ignored for SemVer purposes. This will require us to publish a new `langsmith-pyo3` Python package version with the new changes. I'd like to trigger the publish workflow after this is merged, so I can make the corresponding Python changes to enable zstd end-to-end. Lmk if you have concerns or want me to hold off on publishing a new version.
obi1kenobi
force-pushed
the
pg/zstd-compression-in-tracing
branch
from
c007aed to
2446ce9
Compare
|
Something to consider in the code review: from the server's perspective, this approach is equivalent to the current pure-Python zstd implementation, but has an internal caveat that may prove to be a problem down the line. Due to how the In an ideal world, we'd do the compression in a streaming fashion. This would either require some changes to |
obi1kenobi
commented
Jan 9, 2025
Comment on lines
+351
to
+380
| // `reqwest` doesn't have a public method for getting *just* the multipart form data: | ||
| // the `Form::reader()` method isn't public. Instead, we pretend to be preparing a request, | ||
| // place the multipart data, and ask for the body to be buffered and made available as `&[u8]`. | ||
| // *This does not send the request!* We merely prepare a request in memory as a workaround. | ||
| // | ||
| // Just in case, we use `example.com` as the URL here, which is explicitly reserved for | ||
| // example use in the standards and is guaranteed to not be taken. This means that | ||
| // under no circumstances will we send multipart data to any other place, | ||
| // even if `reqwest` were to suddenly change the API to send data on `.build()`. | ||
| let builder = reqwest::blocking::Client::new().get("http://example.com/").multipart(form); | ||
| let mut request = builder.build().expect("failed to construct request"); | ||
| let multipart_form_bytes = request | ||
| .body_mut() | ||
| .as_mut() | ||
| .expect("multipart request had no body, this should never happen") | ||
| .buffer() | ||
| .map_err(|e| TracingClientError::IoError(e.to_string()))?; | ||
|
|
||
| let mut buffer = Vec::with_capacity(multipart_form_bytes.len()); | ||
| let mut encoder = zstd::Encoder::new(&mut buffer, compression_level) | ||
| .and_then(|mut encoder| { | ||
| encoder.multithread(n_workers)?; | ||
| Ok(encoder) | ||
| }) | ||
| .map_err(|e| TracingClientError::IoError(e.to_string()))?; | ||
|
|
||
| encoder | ||
| .write_all(multipart_form_bytes) | ||
| .map_err(|e| TracingClientError::IoError(e.to_string()))?; | ||
| encoder.finish().map_err(|e| TracingClientError::IoError(e.to_string()))?; |
There was a problem hiding this comment.
This workaround will not be required anymore once the next reqwest version is released. This code will be replaced in #1403 which makes use of new reqwest functionality to do this properly.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
I've tested this with both the Python and Rust changes applied, and I've successfully sent zstd-compressed traces to LangSmith servers.
This introduces breaking changes in the APIs for both
langsmith-tracing-clientandlangsmith-pyo3, so I'm bumping the left-most non-zero version number in each of their manifests. In Rust, the left-most non-zero number is considered a "major" version -- in other words, leading zeroes are ignored for SemVer purposes.This will require us to publish a new
langsmith-pyo3Python package version with the new changes. I'd like to trigger the publish workflow after this is merged, so I can make the corresponding Python changes to enable zstd end-to-end. Lmk if you have concerns or want me to hold off on publishing a new version.