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

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?*

- [x] Tests pass locally with `pytest` for python, `yarn test` for js
- [x] 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](https://github.com/chroma-core/chroma/assets/1157440/5d87b4d5-4dae-48fe-908c-7c09db2a5abc)


Tests with orjson + async:

![Screenshot 2024-02-03 at 17 53
17](https://github.com/chroma-core/chroma/assets/1157440/d9818fdd-11c3-45c9-81dd-8baecbb638cf)


[1m-100batch-5concur.zip](https://github.com/chroma-core/chroma/files/14152062/1m-100batch-5concur.zip)

[1m-1000batch-5concur.zip](https://github.com/chroma-core/chroma/files/14152063/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](https://github.com/chroma-core/docs)?*


## Refs

- https://showmax.engineering/articles/json-python-libraries-overview
- https://github.com/ijl/orjson
-
https://catnotfoundnear.github.io/finding-the-fastest-python-json-library-on-all-python-versions-8-compared.html

---------

Co-authored-by: Ben Eggers <[email protected]>

chromadb/api/segment.py

@@ -1,3 +1,5 @@
+from functools import cached_property
+
 from chromadb.api import ServerAPI
 from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Settings, System
 from chromadb.db.system import SysDB
@@ -789,7 +791,7 @@ def reset(self) -> bool:
     def get_settings(self) -> Settings:
         return self._settings
 
-    @property
+    @cached_property
     @override
     def max_batch_size(self) -> int:
         return self._producer.max_batch_size

chromadb/auth/fastapi.py

@@ -1,3 +1,6 @@
+import asyncio
+
+import chromadb
 from contextvars import ContextVar
 from functools import wraps
 import logging
@@ -173,7 +176,7 @@ def authz_context(
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
     def decorator(f: Callable[..., Any]) -> Callable[..., Any]:
         @wraps(f)
-        def wrapped(*args: Any, **kwargs: Dict[Any, Any]) -> Any:
+        async def wrapped(*args: Any, **kwargs: Dict[Any, Any]) -> Any:
             _dynamic_kwargs = {
                 "api": args[0]._api,
                 "function": f,
@@ -213,6 +216,7 @@ def wrapped(*args: Any, **kwargs: Dict[Any, Any]) -> Any:
                     )
 
                     if _provider:
+                        # TODO this will block the event loop if it takes too long - refactor for async
                         a_authz_responses.append(_provider.authorize(_context))
                 if not any(a_authz_responses):
                     raise AuthorizationError("Unauthorized")
@@ -239,6 +243,8 @@ def wrapped(*args: Any, **kwargs: Dict[Any, Any]) -> Any:
                         ):
                             kwargs["database"].name = desired_database
 
+            if asyncio.iscoroutinefunction(f):
+                return await f(*args, **kwargs)
             return f(*args, **kwargs)
 
         return wrapped

chromadb/config.py

@@ -144,6 +144,8 @@ def empty_str_to_none(cls, v: str) -> Optional[str]:
         return v
 
     chroma_server_nofile: Optional[int] = None
+    # the number of maximum threads to handle synchronous tasks in the FastAPI server
+    chroma_server_thread_pool_size: int = 40
 
     chroma_server_auth_provider: Optional[str] = None
 

chromadb/db/impl/sqlite.py

@@ -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")
             self._conn.execute("BEGIN;")
         self._tx_stack.stack.append(self)
         return self._conn.cursor()  # type: ignore