[ENH]: Load HNSW index without disk intermediary #5159
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
Reviewer ChecklistPlease leverage this checklist to ensure your code review is thorough before approving Testing, Bugs, Errors, Logs, Documentation
System Compatibility
Quality
|
|
Enable In-Memory Loading of HNSW Index from S3 (No Disk Intermediary) This PR refactors the HNSW index loading pipeline to support direct in-memory loading of index data from S3, rather than writing and reading intermediate files to/from disk. It updates interface contracts, adds a new code path for loading buffers directly, introduces relevant changes throughout the index provider and index types, and bumps the hnswlib dependency to a specific commit supporting this feature. Key Changes• Added PersistentIndex::load_from_hnsw_data API (with trait and implementations) to allow direct in-memory index loading. Affected Areas• rust/index/src/hnsw_provider.rs This summary was automatically generated by @propel-code-bot |
|
8 Jobs Failed: PR checks / LintPR checks / Rust tests / test-benches (blacksmith-16vcpu-ubuntu-2204, --bench get)PR checks / Python tests / test-cluster-rust-frontend (3.9, chromadb/test/property/test_embeddings.py)PR checks / Rust tests / Integration test ci_k8s_integration 2PR checks / Rust tests / test (blacksmith-8vcpu-ubuntu-2204)PR checks / Rust tests / test-longPR checks / all-required-pr-checks-passed1 job failed running on non-Blacksmith runners. Summary: 1 successful workflow, 1 failed workflow
Last updated: 2025-07-29 15:10:29 UTC |
![propel-code-bot[bot]](https://avatars.githubusercontent.com/in/1179004?s=60&v=4){kind=small}
Comment on lines
+223
to
+230
| None => match HnswIndex::load_from_hnsw_data( | ||
| self.fetch_hnsw_segment(&new_id, prefix_path) | ||
| .await | ||
| .map_err(|e| Box::new(HnswIndexProviderForkError::FileError(*e)))?, | ||
| &index_config, | ||
| ef_search, | ||
| new_id, | ||
| ) { |
There was a problem hiding this comment.
[CriticalError]
The logic for loading the index within the fork method appears to be incorrect. It attempts to fetch the segment from remote storage using new_id, but the index for new_id doesn't exist in storage yet. The files for source_id have just been copied to a local directory.
The previous implementation using HnswIndex::load(storage_path_str, ...) correctly loaded the index from this new local directory. Since fork is intended to create a mutable, file-backed copy of an index, it seems the original approach of loading from the local path should be restored.
Comment on lines
+287
to
+301
| let hnsw_data = self.fetch_hnsw_segment(source_id, prefix_path).await?; | ||
| let getters = [ | ||
| |hnsw_data: &hnswlib::HnswData| Arc::new(Vec::from(hnsw_data.header_buffer())), | ||
| |hnsw_data: &hnswlib::HnswData| Arc::new(Vec::from(hnsw_data.data_level0_buffer())), | ||
| |hnsw_data: &hnswlib::HnswData| Arc::new(Vec::from(hnsw_data.length_buffer())), | ||
| |hnsw_data: &hnswlib::HnswData| Arc::new(Vec::from(hnsw_data.link_list_buffer())), | ||
| ]; | ||
|
|
||
| for (file, getter) in FILES.iter().zip(getters) { | ||
| let file_path = index_storage_path.join(file); | ||
| self.copy_bytes_to_local_file(&file_path, getter(&hnsw_data)) | ||
| .await?; | ||
| } | ||
| Ok(()) | ||
| } |
There was a problem hiding this comment.
[PerformanceOptimization]
This function now fetches the entire HNSW segment into an in-memory HnswData object before writing the individual files to disk. The previous implementation streamed each file directly. For large indexes, this change could significantly increase memory usage during the fork operation. Was this change intentional? If fork still needs to write to disk, perhaps restoring the previous file-by-file download logic for this function would be more memory-efficient.
Comment on lines
+408
to
+415
| None => match HnswIndex::load_from_hnsw_data( | ||
| self.fetch_hnsw_segment(id, prefix_path) | ||
| .await | ||
| .map_err(|e| Box::new(HnswIndexProviderOpenError::FileError(*e)))?, | ||
| &index_config, | ||
| ef_search, | ||
| *id, | ||
| ) { |
There was a problem hiding this comment.
[PerformanceOptimization]
This change successfully loads the index from memory, which aligns with the PR's goal. However, the open function still contains calls that write the index to a temporary directory on disk (create_dir_all, load_hnsw_segment_into_directory) before this memory-based loading occurs. These disk operations now seem redundant.
To fully load without a disk intermediary and improve efficiency, you could remove the calls to create_dir_all, load_hnsw_segment_into_directory, and purge_one_id from this function.
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.
Description of changes
(WIP)
Test plan
How are these changes tested?
pytestfor python,yarn testfor js,cargo testfor rustMigration plan
Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?
Observability plan
What is the plan to instrument and monitor this change?
Documentation Changes
Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the docs section?