[PERF] prefetch blocks for fulltext index writer by codetheweb · Pull Request #3370 · chroma-core/chroma

codetheweb · 2024-12-30T19:36:42Z

Description of changes

Prefetchs blocks in parallel when writing the full text index rather than fetching the blocks sequentially on-demand. Fetching in parallel saves us quite a bit of time and having prefetch as a distinct operator also allows us to mark it as I/O.

Test plan

How are these changes tested?

Tests pass locally with pytest for python, yarn test for js, cargo test for rust

Tested by disabling the write-through cache and observing that the prefetch step pulled in the correct blocks.

Documentation Changes

Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the docs repository?

n/a

github-actions · 2024-12-30T19:36:52Z

rust/index/src/fulltext/types.rs

rust/blockstore/src/arrow/ordered_blockfile_writer.rs

HammadB

I think @sanketkedia should weigh in on this before merge, as we have operators for prefetch and established patterns for it

codetheweb · 2025-01-02T22:30:33Z

I think @sanketkedia should weigh in on this before merge, as we have operators for prefetch and established patterns for it

I will ask
I originally sent more context to Robert in a DM, should have included it here:

I didn't implement prefetching in an operator in this first pass because it doesn't fit well into the upstream SegmentWriter abstraction & associated operators

rust/blockstore/src/types/writer.rs

rust/blockstore/src/arrow/ordered_blockfile_writer.rs

rust/index/src/fulltext/types.rs

rust/blockstore/src/arrow/blockfile.rs

codetheweb · 2025-01-13T22:36:13Z

This stack of pull requests is managed by Graphite. Learn more about stacking.

rust/worker/src/execution/orchestration/compact.rs

rust/index/src/fulltext/types.rs

rust/worker/src/execution/operators/prefetch_keys.rs

HammadB · 2025-01-21T21:27:42Z

rust/blockstore/src/arrow/ordered_blockfile_writer.rs

+        let num_block_cache_hits = self
+            .num_block_cache_hits
+            .load(std::sync::atomic::Ordering::Relaxed);
+        let num_block_cache_misses = self


Foyer already exposes the underlying hit rate and miss rate - i think this is redundant?

chroma/rust/cache/src/foyer.rs

Line 251 in 8f7a0ff

cache_hit: opentelemetry::metrics::Counter<u64>,

Also if we are doing it - i i think it should be a metric not a trace

ah yep missed that
I think it's actually kinda useful having it on the trace; if all you have is a metric point with a low cache hit rate it could be pretty hard to debug
will leave out for now

ah makes sense

rust/index/src/fulltext/types.rs

codetheweb · 2025-01-22T19:25:41Z

rust/worker/src/execution/orchestration/compact.rs

        let task = wrap(operator, input, ctx.receiver());
        self.send(task, ctx).await;
+
+        let segments = self.get_all_segments().await.unwrap();


(Not putting this in the initial_tasks() method from the Orchestrator trait as that would require making initial_tasks() async and my understanding is that it's intended to be a static list of tasks. A cleaner solution would perhaps be adding a GetSegment operator and dispatching that from initial_tasks().)

codetheweb · 2025-01-22T20:52:18Z

rust/blockstore/src/provider.rs

        Ok(())
    }
+
+    pub async fn prefetch(&self, id: &uuid::Uuid) -> Result<usize, Box<dyn ChromaError>> {


this is a new top-level API, not attached to a reader or writer
I think this makes more sense semantically since we're not fetching the blocks for any particular reader or writer instance and it lets us avoid K & V generics

HammadB · 2025-01-23T04:32:54Z

rust/blockstore/src/arrow/provider.rs

    }

+    pub async fn prefetch(&self, id: &Uuid) -> Result<usize, ArrowBlockfileProviderPrefetchError> {
+        let block_ids = self


Why not just fetch the root and then do this? The abstractions feel a bit off to me here.

Fetching the root requires K to be known, which would break a general-purpose prefetch operator that has no knowledge of the concrete K/V types for each file path.

ah that sucks, might be nice to comment that

HammadB · 2025-01-23T06:11:45Z

Barring any fallout from my question above this makes sense to me. Thanks for the attention to quality.

## Description of changes These changes were originally part of #3370 but decided to break them out as they're no longer needed for that PR. Makes the interface slightly more flexible and puts the burden of owned data on the caller. ## Test plan *How are these changes tested?* - [x] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust ## Documentation Changes *Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?* n/a

## Description of changes These changes were originally part of chroma-core#3370 but decided to break them out as they're no longer needed for that PR. Makes the interface slightly more flexible and puts the burden of owned data on the caller. ## Test plan *How are these changes tested?* - [x] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust ## Documentation Changes *Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?* n/a

rescrv reviewed Jan 2, 2025

View reviewed changes

rust/index/src/fulltext/types.rs Outdated Show resolved Hide resolved

rust/blockstore/src/arrow/ordered_blockfile_writer.rs Outdated Show resolved Hide resolved

HammadB reviewed Jan 2, 2025

View reviewed changes

codetheweb commented Jan 10, 2025

View reviewed changes

rust/blockstore/src/types/writer.rs Outdated Show resolved Hide resolved

codetheweb commented Jan 10, 2025

View reviewed changes

rust/blockstore/src/arrow/ordered_blockfile_writer.rs Outdated Show resolved Hide resolved

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from 8b144d3 to 154c080 Compare January 10, 2025 23:36

codetheweb commented Jan 10, 2025

View reviewed changes

rust/index/src/fulltext/types.rs Outdated Show resolved Hide resolved

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch 2 times, most recently from 8ed8d3b to 931a5f2 Compare January 10, 2025 23:54

codetheweb commented Jan 10, 2025

View reviewed changes

rust/blockstore/src/arrow/blockfile.rs Outdated Show resolved Hide resolved

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from 931a5f2 to 323b15e Compare January 13, 2025 22:34

codetheweb changed the base branch from main to fix-arrow-cached-check-bug January 13, 2025 22:35

codetheweb mentioned this pull request Jan 13, 2025

[BUG]: BlockManager::cached() should not always return true #3480

Merged

1 task

codetheweb force-pushed the fix-arrow-cached-check-bug branch from ce5521e to 2c3884c Compare January 14, 2025 18:30

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from 323b15e to 31ee004 Compare January 14, 2025 18:31

codetheweb marked this pull request as ready for review January 14, 2025 18:35

codetheweb requested a review from sanketkedia January 14, 2025 18:35

codetheweb changed the base branch from fix-arrow-cached-check-bug to graphite-base/3370 January 15, 2025 02:09

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from 31ee004 to f7ecfd9 Compare January 15, 2025 02:09

codetheweb force-pushed the graphite-base/3370 branch from 2c3884c to 6ae644d Compare January 15, 2025 02:09

codetheweb changed the base branch from graphite-base/3370 to main January 15, 2025 02:10

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from f7ecfd9 to a13d405 Compare January 15, 2025 02:10

sanketkedia reviewed Jan 15, 2025

View reviewed changes

rust/worker/src/execution/orchestration/compact.rs Outdated Show resolved Hide resolved

sanketkedia reviewed Jan 15, 2025

View reviewed changes

rust/index/src/fulltext/types.rs Outdated Show resolved Hide resolved

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from a13d405 to 8c7288e Compare January 16, 2025 00:53

codetheweb commented Jan 16, 2025

View reviewed changes

rust/worker/src/execution/operators/prefetch_keys.rs Outdated Show resolved Hide resolved

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from e0f7533 to 577da0e Compare January 16, 2025 18:35

codetheweb requested a review from sanketkedia January 16, 2025 18:37

codetheweb requested a review from HammadB January 17, 2025 00:51

HammadB reviewed Jan 21, 2025

View reviewed changes

rust/index/src/fulltext/types.rs Outdated Show resolved Hide resolved

codetheweb added 2 commits January 22, 2025 10:41

[PERF] prefetch blocks for fulltext index writer

f311949

Pass keys in prefetch input

afb02fa

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch 3 times, most recently from 77c6a76 to f7c73c4 Compare January 22, 2025 19:21

codetheweb mentioned this pull request Jan 22, 2025

[ENH]: improve interface for load_blocks_for_keys() #3534

Merged

1 task

codetheweb commented Jan 22, 2025

View reviewed changes

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch 3 times, most recently from c547c2f to a264c8f Compare January 22, 2025 20:51

codetheweb commented Jan 22, 2025

View reviewed changes

Prefetch entire segment

8241a65

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch from a264c8f to 8241a65 Compare January 22, 2025 21:01

codetheweb requested a review from HammadB January 22, 2025 21:02

HammadB reviewed Jan 23, 2025

View reviewed changes

HammadB approved these changes Jan 23, 2025

View reviewed changes

codetheweb force-pushed the perf-prefetch-fulltext-index-writer-blocks branch 2 times, most recently from 3440b66 to 8241a65 Compare January 23, 2025 21:46

Add comment

61d49f0

codetheweb enabled auto-merge (squash) January 23, 2025 21:49

codetheweb added 2 commits January 23, 2025 14:07

Merge branch 'main' into perf-prefetch-fulltext-index-writer-blocks

daf9a7d

Fix imports

c9f7fa0

codetheweb merged commit b848326 into main Jan 23, 2025
79 checks passed

codetheweb deleted the perf-prefetch-fulltext-index-writer-blocks branch January 23, 2025 23:07

Conversation

codetheweb commented Dec 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description of changes

Test plan

Documentation Changes

Uh oh!

github-actions bot commented Dec 30, 2024

Reviewer Checklist

Testing, Bugs, Errors, Logs, Documentation

System Compatibility

Quality

Uh oh!

Uh oh!

Uh oh!

HammadB left a comment

Choose a reason for hiding this comment

Uh oh!

codetheweb commented Jan 2, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codetheweb commented Jan 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codetheweb Jan 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HammadB commented Jan 23, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

codetheweb commented Dec 30, 2024 •

edited

Loading

codetheweb commented Jan 13, 2025 •

edited

Loading

codetheweb Jan 22, 2025 •

edited

Loading