Repository navigation
Conversation
Replace fs::read (full file into memory) with a 64 KiB streaming loop so large shared libraries (.so, .gguf) do not need to fit in RAM during digest computation. Adds a test that verifies the streaming path produces the same SHA-256 and SHA-512 digests as a reference implementation on a 16 MiB synthetic file.
|
Hi @nrenvo, thanks for your interest in contributing! This project requires that pull request authors are vouched, and you are not in the list of vouched users. This PR will be closed automatically. See https://github.com/huggingface/kernels/blob/main/CONTRIBUTING.md for more details. |
danieldk
left a comment
danieldk
left a comment
There was a problem hiding this comment.
Thanks a lot! Added two small comments.
| loop { | ||
| let count = file | ||
| .read(&mut buffer) | ||
| .wrap_err_with(|| format!("Cannot read `{}` for hashing", path.display()))?; | ||
| if count == 0 { | ||
| break; | ||
| } |
There was a problem hiding this comment.
We should probably check for interrupted as well:
An error of the ErrorKind::Interrupted kind is non-fatal and the read operation should be retried if there is nothing else to do.
https://doc.rust-lang.org/beta/std/io/trait.Read.html
So something like:
let count = match file.read(&mut buffer) {
Ok(n) => n,
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(e) => return Err(e).wrap_err_with(|| format!("Cannot read `{}` for hashing", path.display())),
};|
|
||
| Ok(()) | ||
| } | ||
|
|
There was a problem hiding this comment.
The test uses multiples of 64KiB of the same value, we'd probably need an odd-sized tail and varying content. I also think it does not need to be very large, just 2-3 64KiB chunks and then something that's not 64KiB (and odd because the digest may work internally with chunks that are have a power-of-2 size)?
`Read::read` may return an error of kind `Interrupted`, which is non-fatal and should be retried. Handle it instead of failing digest computation. Also rework the streaming digest test to use varying content and an odd-sized tail after three full 64 KiB chunks, so it no longer only exercises buffer-aligned reads.
|
Thanks Daniel, good catches! I added the Interrupted retry, and reworked the test to use varying content with an odd-sized tail after three full 64 KiB chunks. Just pushed the update. |
Replace fs::read (entire file buffered into memory) with a 64 KiB streaming loop so large shared libraries or model weights do not need to fit in RAM during digest computation.
Previously, hashing a 1 GB .so file would spike the process's RSS by ~1 GB. With this change, peak memory during hashing is bounded by the 64 KiB buffer regardless of file size.
A new test (hash_variant_streams_large_file_without_changing_digest) verifies that the streaming path produces identical SHA-256 and SHA-512 digests as a reference implementation on a 16 MiB synthetic file.