[ENH] Add quota component and test for static #1720

nicolasgere · 2024-02-14T23:25:15Z

Description of changes

Summarize the changes made by this PR.

New functionality
- Add quota check, it will be use to be able to rate limit, apply static check to payload etc.

Test plan

How are these changes tested?

Tests pass locally with pytest, added unit test

Documentation Changes

CIP doc

github-actions · 2024-02-14T23:25:29Z

HammadB

Thanks , some nits and observations but is great overall!

chromadb/quota/__init__.py

HammadB · 2024-02-15T19:19:17Z

chromadb/quota/__init__.py

+        self.resource = resource
+
+
+class QuotaProvider(Component):


nit - add docstrings about this class

chromadb/quota/__init__.py

tazarov

As a general comment, I think that quota or any limit additions to the codebase are unnecessary technical debt better moved to the Proxy/Cache/Coordinator communication pattern. However, that is just my two cents on the long-term impacts.

tazarov · 2024-02-16T11:53:32Z

chromadb/quota/__init__.py

+            self.should_enforce = True
+        self.system = system
+
+    def payload_static_check(self, metadatas: Optional[Metadatas] = None, documents: Optional[Documents] = None, embeddings: Optional[Embeddings]= None, collection_id: Optional[str]= None):


We should also check ids are those are user-passed strings and in the DB we store them as text. That can get pretty large.

I will add them in another PR

tazarov · 2024-02-17T07:15:38Z

chromadb/api/segment.py

@@ -101,6 +101,7 @@ def __init__(self, system: System):
        self._settings = system.settings
        self._sysdb = self.require(SysDB)
        self._manager = self.require(SegmentManager)
+        self._quota = self.require(QuotaEnforcer)


Should we enforce metadata checks on the collection level?

Good question, I will ask with the team.

tazarov · 2024-02-17T07:23:49Z

chromadb/quota/__init__.py

+    def payload_static_check(self, metadatas: Optional[Metadatas] = None, documents: Optional[Documents] = None, embeddings: Optional[Embeddings]= None, collection_id: Optional[str]= None):
+        if not self.should_enforce:
+            return
+        metadata_key_length_quota = self._quota_provider.get_for_subject(resource="METADATA_KEY_LENGTH", subject=collection_id)


Does this imply that all checks can only be performed against specific collections? How about quotas at user, tenant, and DB levels? Does this approach force Chroma to maintain per-collection quotas? Why are per-collection quotas more interesting than quotas at other levels of object inheritance?

This is just for static one, but yes, later, tenant based quota would also be possible

tazarov · 2024-02-17T07:25:54Z

chromadb/server/fastapi/__init__.py

@@ -291,6 +293,12 @@ def app(self) -> fastapi.FastAPI:
    def root(self) -> Dict[str, int]:
        return {"nanosecond heartbeat": self._api.heartbeat()}

+    async def quota_exception_handler(request: Request, exc: QuotaError):
+        return JSONResponse(
+            status_code=429,


Why is 429 (Too Many Requests) more appropriate than 413 (Payload Too Large)?

we want something generic, and 429 is also use for resource exhausted. (gcp, aws)

tazarov · 2024-02-17T07:27:32Z

chromadb/quota/__init__.py

+                    raise QuotaError(resource="DOCUMENT_SIZE", actual=len(document) , quota=document_size_quota)
+        embedding_dimension_quota = self._quota_provider.get_for_subject(resource="EMBEDDINGS_DIMENSION", subject=collection_id)
+        if embedding_dimension_quota and embeddings:
+            for embedding in embeddings:


using any() might be a bit more python and codebase idiomatic:

if any(len(embedding) > embedding_dimension_quota for embedding in embeddings): raise QuotaError("EMBEDDINGS_DIMENSION", actual=max(len(embedding) for embedding in embeddings), quota=embedding_dimension_quota)

Same applies for the other checks

tazarov · 2024-02-17T07:41:29Z

chromadb/server/fastapi/__init__.py

@@ -140,6 +141,7 @@ def __init__(self, settings: Settings):
            allow_origins=settings.chroma_server_cors_allow_origins,
            allow_methods=["*"],
        )
+        self._app.add_exception_handler(QuotaError, self.quota_exception_handler)


Why is a dedicated exception handler better than using the existing catch_exceptions_middleware to handle the error?

I did not see it, can be part of another refactoring. Or maybe we should use add_exception_handler which is the way of doing it in fastapi

tazarov · 2024-02-17T07:44:15Z

chromadb/test/quota/test_static_quota_enforcer.py

@@ -0,0 +1,78 @@
+import random


Do you think it makes sense that limits-related tests are subject to the rigor of property tests?

as there is no state, I think unit test is ok. wdyt?

## Description of changes *Summarize the changes made by this PR.* - New functionality - Add quota check, it will be use to be able to rate limit, apply static check to payload etc. ## Test plan *How are these changes tested?* - [ ] Tests pass locally with `pytest`, added unit test --------- Co-authored-by: nicolas <[email protected]>

add quota component and test for static

b081553

fix test

8dfab64

nicolasgere changed the title ~~WIP: add quota component and test for static~~ Add quota component and test for static Feb 15, 2024

nicolas added 2 commits February 15, 2024 10:09

refactor test

bfeaebb

remote print

f98a2ad

nicolasgere requested review from HammadB and beggers February 15, 2024 18:10

HammadB changed the title ~~Add quota component and test for static~~ [ENH] Add quota component and test for static Feb 15, 2024

HammadB approved these changes Feb 15, 2024

View reviewed changes

add comments, refactor resource with enum

8cf2633

nicolasgere requested a review from tazarov February 15, 2024 21:11

nicolasgere merged commit e6ceeee into main Feb 16, 2024
98 checks passed

tazarov reviewed Feb 17, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ENH] Add quota component and test for static #1720

[ENH] Add quota component and test for static #1720

nicolasgere commented Feb 14, 2024 •

edited

Loading

github-actions bot commented Feb 14, 2024 •

edited by nicolasgere

Loading

HammadB left a comment

HammadB Feb 15, 2024

tazarov left a comment

tazarov Feb 16, 2024

nicolasgere Feb 20, 2024

tazarov Feb 17, 2024

nicolasgere Feb 20, 2024

tazarov Feb 17, 2024

nicolasgere Feb 20, 2024

tazarov Feb 17, 2024

nicolasgere Feb 20, 2024

tazarov Feb 17, 2024

tazarov Feb 17, 2024

nicolasgere Feb 20, 2024

tazarov Feb 17, 2024

nicolasgere Feb 20, 2024

[ENH] Add quota component and test for static #1720

[ENH] Add quota component and test for static #1720

Conversation

nicolasgere commented Feb 14, 2024 • edited Loading

Description of changes

Test plan

Documentation Changes

github-actions bot commented Feb 14, 2024 • edited by nicolasgere Loading

Reviewer Checklist

Testing, Bugs, Errors, Logs, Documentation

System Compatibility

Quality

HammadB left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tazarov left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nicolasgere commented Feb 14, 2024 •

edited

Loading

github-actions bot commented Feb 14, 2024 •

edited by nicolasgere

Loading