fix: use temp files in push_to_hub to prevent OOM on large datasets #7895
+255
−16
Conversation
This file contains hidden or 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
When uploading large datasets with many shards, the previous implementation held all parquet shard bytes in memory via the `additions` list. This caused OOM for datasets with hundreds of shards (e.g., 902 shards × 300MB = 270GB). Changes: - Write parquet to temp file instead of BytesIO - Pass file path to CommitOperationAdd (huggingface_hub streams from disk) - Delete temp file after preupload_lfs_files completes - Remove unused BytesIO import Memory usage now stays constant (~1-2 shard sizes) instead of growing linearly with the number of shards. Fixes: #5
Wrap the fallible upload operations in try...finally to guarantee temp file cleanup even if preupload_lfs_files() throws (network timeout, 500 error, KeyboardInterrupt). Without this, a flaky upload of 900 shards could leave hundreds of ~300MB temp files in /tmp, potentially filling the disk.
Addresses CodeRabbit's CRITICAL review feedback: 1. Only delete temp file when _upload_mode == "lfs" - For LFS uploads, content is already on the Hub after preupload_lfs_files - For regular uploads, create_commit still needs to read from disk - Deleting temp files unconditionally broke non-LFS uploads 2. Replace XXXX placeholder with actual issue reference 3. Update tests to properly simulate huggingface_hub behavior: - Mock preupload_lfs_files to set _upload_mode attribute - Add test for regular (non-LFS) upload behavior - Mock upload_info to make size assertion meaningful
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Summary
Fixes memory accumulation in
_push_parquet_shards_to_hub_singlethat causes OOM when uploading large datasets with many shards.Root cause: The current implementation stores ALL parquet shard bytes in memory via
BytesIO, accumulating in theadditionslist. For N shards of ~300MB each, this requires N × 300MB RAM.Fix: Write parquet to temp file instead of
BytesIO, pass file path toCommitOperationAdd. Delete temp file afterpreupload_lfs_filescompletes (for LFS uploads only - regular uploads need the file untilcreate_commit).Changes
BytesIOwithtempfile.NamedTemporaryFilein_push_parquet_shards_to_hub_singleCommitOperationAdd.path_or_fileobjinstead of bytescreate_commit)try...finallyfor safe cleanup even on errorsBytesIOimportTest plan
tests/test_push_to_hub_memory.pywith 4 tests:test_push_to_hub_uses_file_path_not_bytes_in_commit_operationtest_push_to_hub_cleans_up_temp_files_for_lfs_uploadstest_push_to_hub_keeps_temp_files_for_regular_uploadstest_push_to_hub_uploaded_size_still_calculatedContext
This was discovered while uploading a 270GB neuroimaging dataset (ARC) with 902 shards. The process was killed by OOM after accumulating ~270GB in the
additionslist.Fixes #7893
Related: #5990, #7400