[ENH]: CIP-5: Large Batch Handling Improvements Proposal

[tazarov]

• Including only CIP for review.

Refs: #1049

Description of changes

Summarize the changes made by this PR.

• Improvements & Bug fixes
  • New proposal to handle large batches of embeddings gracefully

Test plan

How are these changes tested?

• Tests pass locally with pytest for python, yarn test for js

Documentation Changes

TBD

[github-actions]

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

• Can you think of any use case in which the code does not behave as intended? Have they been tested?
• Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
• If appropriate, are there adequate property based tests?
• If appropriate, are there adequate unit tests?
• Should any logging, debugging, tracing information be added or removed?
• Are error messages user-friendly?
• Have all documentation changes needed been made?
• Have all non-obvious changes been commented?

System Compatibility

• Are there any potential impacts on other parts of the system or backward compatibility?
• Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

• Is this code of a unexpectedly high quality (Readbility, Modularity, Intuitiveness)

[HammadB]

is the assumption that these are fetched once on client creation?

[tazarov]

Haven't thought about it but actual impl caches the values

[tazarov]

https://github.com/amikos-tech/chroma-core/blob/93f0eda92ff1061a42cb6b3cf219ce6291c6f4ef/chromadb/api/fastapi.py#L400

[HammadB]

My main feedback here is we should not gracefully break up the batches for you, as partial failures are harder to reason about. It's easier to delegate the batching responsibility to the client and have you be responsible for handling partial failures in a way that makes sense to you. Otherwise this LGTM!

[tazarov]

@HammadB, true that we cannot guarantee ACID-like behaviour this way, but even without, things seem to go wrong on very large batches. Do you think we can introduce tenacity (Apache 2 lic) for some added bonus of retries? I think this can be helpful further down the road, where a retry can resolve the problem.

In my experience, in situations where transactionality cannot be guaranteed, we can implement mechanisms to fail gracefully - e.g. store the failed batch in a temp file and allow the user to resubmit later. wdyt?

[imartinez]

I think this discussion brought by @HammadB and the last proposed solution by @tazarov requires more time and analysis (the last solution also has its implications).

I'd suggest to reduce the scope of this change to just exposing the max_batch_size in the API so it can be merged and clients can handle batches themselves (and add it to the docs) -in the case of privateGPT it is causing errors, I guess there are more clients impacted- and discuss the error management in a different CIP. WDYT?

[tazarov]

@imartinez, let me do that now.

[tazarov]

@imartinez @HammadB, Refactored. The new impl no longer splits batches automatically. We do ship a utility library for that though chromadb/utils/batch_utils.py. The API now validates the batch size (client-side), prior to trying to persist it. Everything else is left as-is:

• pre-flight-check endpoint is still there
• max_batch_size property is still available at API level for clients to use to break up batches

Added a couple more tests to ensure correctness.

CIP doc has also been updated to reflect above changes.

[tazarov]

@HammadB,

Another unrelated failure:

chromadb/test/property/test_persist.py .F                                [100%]

=================================== FAILURES ===================================
___________________ test_persist_embeddings_state[settings0] ___________________

caplog = <_pytest.logging.LogCaptureFixture object at 0x7f496ee70550>
settings = Settings(environment='', chroma_db_impl=None, chroma_api_impl='chromadb.api.segment.SegmentAPI', chroma_telemetry_impl...=None, chroma_server_auth_token_transport_header=None, anonymized_telemetry=True, allow_reset=True, migrations='apply')

    def test_persist_embeddings_state(
        caplog: pytest.LogCaptureFixture, settings: Settings
    ) -> None:
        caplog.set_level(logging.ERROR)
        api = chromadb.Client(settings)
        run_state_machine_as_test(
>           lambda: PersistEmbeddingsStateMachine(settings=settings, api=api)
        )  # type: ignore

chromadb/test/property/test_persist.py:226: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/opt/hostedtoolcache/Python/3.7.17/x64/lib/python3.7/site-packages/hypothesis/stateful.py:216: in run_state_machine_as_test
    run_state_machine(state_machine_factory)
/opt/hostedtoolcache/Python/3.7.17/x64/lib/python3.7/site-packages/hypothesis/stateful.py:112: in run_state_machine
    @given(st.data())
chromadb/test/property/test_embeddings.py:210: in ann_accuracy
    embedding_function=self.embedding_function,
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

collection = Collection(name=A01
)
record_set = {'documents': [None], 'embeddings': [[0.09765625, 0.430419921875, 0.20556640625, 0.08978271484375, -0.1527099609375, 0.2[917](https://github.com/chroma-core/chroma/actions/runs/6082759956/job/16515250514?pr=1077#step:5:918)48046875, ...]], 'ids': ['1'], 'metadatas': [{'1': '10'}]}
n_results = 1, min_recall = 0.95
embedding_function = <chromadb.test.property.strategies.hashing_embedding_function object at 0x7f4956ddcc50>

[tazarov]

@HammadB, can you rerun the failed tests just to see if we can rule out intermittent failure? If it keeps failing, I'll debug it further

[HammadB]

@tazarov forsure, I've rerun them a handful of times but happy to do so again

[tazarov]

@HammadB, Tests in my fork are passing, so this is ready for merge after your review.