Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[PERF]: Improved server-side serialization with orjson + async #1938

Merged
merged 12 commits into from
Apr 4, 2024

Conversation

tazarov
Copy link
Contributor

@tazarov tazarov commented Mar 28, 2024

This is PR #1695 migrated to Chroma repo for stacking the outstanding changes.

Description of changes

Summarize the changes made by this PR.

  • Improvements & Bug fixes
    • Added orjson serialization improving serialization performance especially on large batches 100+ docs 2x faster (tested with locust)
      • Added async body serialization further improving performance (tested with locust)
      • ⚠️ Fixed an issue with SQLite lack of case_sensitive_like for FastAPI Server thread pool connections

Test plan

How are these changes tested?

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

For performance tests, 5 concurrent clients were querying, and 1 was adding pre-computed OpenAI (ada-02) embeddings (runtime ~1m)

Batch size: 100 (1000 also attached in a zip)

Tests with existing main codebase:

Screenshot 2024-02-03 at 17 52 39

Tests with orjson + async:

Screenshot 2024-02-03 at 17 53 17

1m-100batch-5concur.zip
1m-1000batch-5concur.zip

Documentation Changes

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

Refs

Copy link

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 (Readability, Modularity, Intuitiveness)

@tazarov tazarov mentioned this pull request Mar 28, 2024
1 task
@tazarov
Copy link
Contributor Author

tazarov commented Mar 28, 2024

Picking up from #1695 (comment):

tazarov added 6 commits March 28, 2024 21:19
- Added orjson serialization improving serialization performance especially on large batches 100+ docs 2x faster (tested with locust)
- Added async body serialization further improving performance (tested with locust)
- Added async handling with AnyIO of the more impactful server queries further reducing concurrent request response times

Future work:  Add orjson serialization at client-side
- Removed commented out code
- Added clarification for to_thread usage
- Configurable via `chroma_server_thread_pool_size` setting
…f the main loop

TODO: Telemetry/Authz has the potentially to negatively impact things - needs async refactor
@tazarov tazarov mentioned this pull request Mar 28, 2024
1 task
@tazarov tazarov force-pushed the feature/orjson-parser-async-endpoints branch from 76492a2 to ed53f6a Compare March 28, 2024 19:25
tazarov added 3 commits March 28, 2024 22:00
- Body serialization is off the main thread
- Get collection did not have enough args
- Body was not awaited in modify collection
- Minor cleanups to make code more readable
Copy link
Contributor

@beggers beggers left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A few questions and requests but this is looking pretty good!

@@ -789,7 +791,7 @@ def reset(self) -> bool:
def get_settings(self) -> Settings:
return self._settings

@property
@cached_property
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is there a reason for this? it adds a small amount of cognitive overhead ("what's a cached property?") for no benefit that i can see?

@@ -34,6 +34,7 @@ def __init__(self, conn_pool: Pool, stack: local):
@override
def __enter__(self) -> base.Cursor:
if len(self._tx_stack.stack) == 0:
self._conn.execute("PRAGMA case_sensitive_like = ON")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Was this meant to be included? It seems unrelated to the rest of the changes.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I found this bug while testing the PR, so I added it here. I can split it up in another PR, as it is indeed unrelated.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Move to PR or leave it here?

chromadb/test/client/test_multiple_clients_concurrency.py Outdated Show resolved Hide resolved

return wrapper
if asyncio.iscoroutinefunction(f):

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we fix the spacing here? It's funky to have an empty line after an if-statement

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Interesting, I am sure I have fixed this and ran pre-commit run which must have resulted in this formatting again.

chromadb/server/fastapi/__init__.py Outdated Show resolved Hide resolved
chromadb/server/fastapi/__init__.py Outdated Show resolved Hide resolved
chromadb/server/fastapi/__init__.py Outdated Show resolved Hide resolved
chromadb/server/fastapi/__init__.py Show resolved Hide resolved
Copy link
Contributor

@beggers beggers left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ty for answering my questions. If we clean up the style and import nits I'd be very happy to check this in.

chromadb/server/fastapi/__init__.py Outdated Show resolved Hide resolved
chromadb/server/fastapi/__init__.py Show resolved Hide resolved
@beggers beggers merged commit 0f8b5b6 into main Apr 4, 2024
119 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants