diff --git a/.github/workflows/_python-tests.yml b/.github/workflows/_python-tests.yml index 588627e34ba..ff1c2a178e8 100644 --- a/.github/workflows/_python-tests.yml +++ b/.github/workflows/_python-tests.yml @@ -157,7 +157,7 @@ jobs: - test-glob: "chromadb/test/api" - test-glob: "chromadb/test/property/test_collections_with_database_tenant.py" - test-glob: "chromadb/test/property/test_collections_with_database_tenant_overwrite.py" - runs-on: blacksmith-8vcpu-ubuntu-2404 + runs-on: blacksmith-16vcpu-ubuntu-2404 # OIDC token auth for AWS permissions: contents: read diff --git a/chromadb/test/api/test_schema_e2e.py b/chromadb/test/api/test_schema_e2e.py index 8781964d3e7..331ab795b6f 100644 --- a/chromadb/test/api/test_schema_e2e.py +++ b/chromadb/test/api/test_schema_e2e.py @@ -203,10 +203,17 @@ def test_schema_vector_config_persistence_with_ef( if not is_spann_disabled_mode: assert vector_index.config.spann is not None spann_config = vector_index.config.spann - assert spann_config.search_nprobe == 16 - assert spann_config.write_nprobe == 32 - assert spann_config.ef_construction == 120 - assert spann_config.max_neighbors == 24 + # Capture server-returned values. Quantization may override the requested + # values with its own tuned defaults, so we verify persistence (values + # survive a client reload) rather than asserting specific numbers. + initial_search_nprobe = spann_config.search_nprobe + initial_write_nprobe = spann_config.write_nprobe + initial_ef_construction = spann_config.ef_construction + initial_max_neighbors = spann_config.max_neighbors + assert initial_search_nprobe is not None + assert initial_write_nprobe is not None + assert initial_ef_construction is not None + assert initial_max_neighbors is not None else: assert vector_index.config.spann is None assert vector_index.config.hnsw is not None @@ -226,8 +233,8 @@ def test_schema_vector_config_persistence_with_ef( spann_json = persisted_json["keys"]["#embedding"]["float_list"]["vector_index"][ "config" ]["spann"] - assert spann_json["search_nprobe"] == 16 - assert spann_json["write_nprobe"] == 32 + assert spann_json["search_nprobe"] == initial_search_nprobe + assert spann_json["write_nprobe"] == initial_write_nprobe else: hnsw_json = persisted_json["keys"]["#embedding"]["float_list"]["vector_index"][ "config" @@ -252,8 +259,8 @@ def test_schema_vector_config_persistence_with_ef( assert reloaded_vector_index.config.space == "cosine" if not is_spann_disabled_mode: assert reloaded_vector_index.config.spann is not None - assert reloaded_vector_index.config.spann.search_nprobe == 16 - assert reloaded_vector_index.config.spann.write_nprobe == 32 + assert reloaded_vector_index.config.spann.search_nprobe == initial_search_nprobe + assert reloaded_vector_index.config.spann.write_nprobe == initial_write_nprobe else: assert reloaded_vector_index.config.hnsw is not None assert reloaded_vector_index.config.hnsw.ef_construction == 100 @@ -291,7 +298,7 @@ def get_config(self) -> Dict[str, Any]: @staticmethod def build_from_config( - config: Dict[str, Any] + config: Dict[str, Any], ) -> "DeterministicSparseEmbeddingFunction": return DeterministicSparseEmbeddingFunction(config.get("label", "det_sparse")) @@ -2344,14 +2351,24 @@ def test_modify_collection_no_initial_config_creates_default_schema( # Verify schema was updated updated_schema = collection.schema assert updated_schema is not None - assert updated_schema.defaults.float_list.vector_index.config.spann.search_nprobe == 32 # type: ignore - assert updated_schema.keys["#embedding"].float_list.vector_index.config.spann.search_nprobe == 32 # type: ignore + assert ( + updated_schema.defaults.float_list.vector_index.config.spann.search_nprobe == 32 + ) # type: ignore + assert ( + updated_schema.keys[ + "#embedding" + ].float_list.vector_index.config.spann.search_nprobe + == 32 + ) # type: ignore # Re-fetch from server collection_refreshed = client.get_collection(collection_name) refreshed_schema = collection_refreshed.schema assert refreshed_schema is not None - assert refreshed_schema.defaults.float_list.vector_index.config.spann.search_nprobe == 32 # type: ignore + assert ( + refreshed_schema.defaults.float_list.vector_index.config.spann.search_nprobe + == 32 + ) # type: ignore @pytest.mark.skipif(not is_spann_disabled_mode, reason="SPANN is disabled") @@ -2379,13 +2396,18 @@ def test_modify_collection_no_initial_config_creates_default_schema_local( updated_schema = collection.schema assert updated_schema is not None assert updated_schema.defaults.float_list.vector_index.config.hnsw.ef_search == 100 # type: ignore - assert updated_schema.keys["#embedding"].float_list.vector_index.config.hnsw.ef_search == 100 # type: ignore + assert ( + updated_schema.keys["#embedding"].float_list.vector_index.config.hnsw.ef_search + == 100 + ) # type: ignore # Re-fetch from server collection_refreshed = client.get_collection(collection_name) refreshed_schema = collection_refreshed.schema assert refreshed_schema is not None - assert refreshed_schema.defaults.float_list.vector_index.config.hnsw.ef_search == 100 # type: ignore + assert ( + refreshed_schema.defaults.float_list.vector_index.config.hnsw.ef_search == 100 + ) # type: ignore @pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled) @@ -2404,12 +2426,20 @@ def test_modify_collection_with_initial_spann_schema(client: ClientAPI) -> None: }, ) - # Verify initial schema has the specified config + # Verify initial schema has a SPANN config. Quantization may override the + # requested values with tuned defaults, so capture the server-returned values + # from each location and verify they are preserved after a modify. schema = collection.schema assert schema is not None assert schema.defaults.float_list.vector_index.config.spann is not None # type: ignore - assert schema.defaults.float_list.vector_index.config.spann.search_nprobe == 10 # type: ignore - assert schema.defaults.float_list.vector_index.config.spann.ef_search == 50 # type: ignore + initial_defaults_ef_search = ( + schema.defaults.float_list.vector_index.config.spann.ef_search + ) # type: ignore + initial_embedding_ef_search = schema.keys[ + "#embedding" + ].float_list.vector_index.config.spann.ef_search # type: ignore + assert initial_defaults_ef_search is not None + assert initial_embedding_ef_search is not None # Modify to update search_nprobe to a different value within limits collection.modify(configuration={"spann": {"search_nprobe": 20}}) @@ -2417,20 +2447,39 @@ def test_modify_collection_with_initial_spann_schema(client: ClientAPI) -> None: # Verify update updated_schema = collection.schema assert updated_schema is not None - assert updated_schema.defaults.float_list.vector_index.config.spann.search_nprobe == 20 # type: ignore - # ef_search should remain unchanged - assert updated_schema.defaults.float_list.vector_index.config.spann.ef_search == 50 # type: ignore + assert ( + updated_schema.defaults.float_list.vector_index.config.spann.search_nprobe == 20 + ) # type: ignore + # ef_search should remain unchanged from the initial server-returned value + assert ( + updated_schema.defaults.float_list.vector_index.config.spann.ef_search + == initial_defaults_ef_search + ) # type: ignore # Verify both locations updated - assert updated_schema.keys["#embedding"].float_list.vector_index.config.spann.search_nprobe == 20 # type: ignore - assert updated_schema.keys["#embedding"].float_list.vector_index.config.spann.ef_search == 50 # type: ignore + assert ( + updated_schema.keys[ + "#embedding" + ].float_list.vector_index.config.spann.search_nprobe + == 20 + ) # type: ignore + assert ( + updated_schema.keys["#embedding"].float_list.vector_index.config.spann.ef_search + == initial_embedding_ef_search + ) # type: ignore # Re-fetch and verify collection_refreshed = client.get_collection(collection_name) refreshed_schema = collection_refreshed.schema assert refreshed_schema is not None - assert refreshed_schema.defaults.float_list.vector_index.config.spann.search_nprobe == 20 # type: ignore - assert refreshed_schema.defaults.float_list.vector_index.config.spann.ef_search == 50 # type: ignore + assert ( + refreshed_schema.defaults.float_list.vector_index.config.spann.search_nprobe + == 20 + ) # type: ignore + assert ( + refreshed_schema.defaults.float_list.vector_index.config.spann.ef_search + == initial_defaults_ef_search + ) # type: ignore @pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled) @@ -2446,6 +2495,14 @@ def test_modify_collection_updates_schema_spann_multiple_fields( configuration={"spann": {"search_nprobe": 64, "ef_search": 100}}, ) + # Capture the server-assigned default for write_nprobe before modifying. + # This value depends on whether quantization is enabled (32 without, 64 with). + initial_schema = collection.schema + assert initial_schema is not None + initial_write_nprobe = ( + initial_schema.defaults.float_list.vector_index.config.spann.write_nprobe + ) # type: ignore + # Modify multiple fields collection.modify( configuration={ @@ -2462,20 +2519,36 @@ def test_modify_collection_updates_schema_spann_multiple_fields( assert schema.defaults.float_list.vector_index.config.spann.search_nprobe == 128 # type: ignore assert schema.defaults.float_list.vector_index.config.spann.ef_search == 200 # type: ignore - # Verify other fields were preserved - assert schema.defaults.float_list.vector_index.config.spann.write_nprobe == 32 # type: ignore + # Verify other fields were preserved (not changed by the modify) + assert ( + schema.defaults.float_list.vector_index.config.spann.write_nprobe + == initial_write_nprobe + ) # type: ignore # Verify in both locations - assert schema.keys["#embedding"].float_list.vector_index.config.spann.search_nprobe == 128 # type: ignore - assert schema.keys["#embedding"].float_list.vector_index.config.spann.ef_search == 200 # type: ignore + assert ( + schema.keys["#embedding"].float_list.vector_index.config.spann.search_nprobe + == 128 + ) # type: ignore + assert ( + schema.keys["#embedding"].float_list.vector_index.config.spann.ef_search == 200 + ) # type: ignore # Re-fetch from server collection_refreshed = client.get_collection(collection_name) refreshed_schema = collection_refreshed.schema assert refreshed_schema is not None - assert refreshed_schema.defaults.float_list.vector_index.config.spann.search_nprobe == 128 # type: ignore - assert refreshed_schema.defaults.float_list.vector_index.config.spann.ef_search == 200 # type: ignore - assert refreshed_schema.defaults.float_list.vector_index.config.spann.write_nprobe == 32 # type: ignore + assert ( + refreshed_schema.defaults.float_list.vector_index.config.spann.search_nprobe + == 128 + ) # type: ignore + assert ( + refreshed_schema.defaults.float_list.vector_index.config.spann.ef_search == 200 + ) # type: ignore + assert ( + refreshed_schema.defaults.float_list.vector_index.config.spann.write_nprobe + == initial_write_nprobe + ) # type: ignore @pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled) @@ -2530,7 +2603,10 @@ def test_modify_collection_preserves_other_schema_fields(client: ClientAPI) -> N assert updated_schema.keys["#document"].string.fts_index.enabled is True # Verify vector index WAS updated - assert updated_schema.defaults.float_list.vector_index.config.spann.search_nprobe == 128 # type: ignore + assert ( + updated_schema.defaults.float_list.vector_index.config.spann.search_nprobe + == 128 + ) # type: ignore # Re-fetch from server to verify persistence collection_refreshed = client.get_collection(collection_name) @@ -2538,7 +2614,10 @@ def test_modify_collection_preserves_other_schema_fields(client: ClientAPI) -> N assert refreshed_schema is not None # Verify vector index was updated on server - assert refreshed_schema.defaults.float_list.vector_index.config.spann.search_nprobe == 128 # type: ignore + assert ( + refreshed_schema.defaults.float_list.vector_index.config.spann.search_nprobe + == 128 + ) # type: ignore # Verify other value types are still intact on server assert refreshed_schema.defaults.string is not None @@ -2837,7 +2916,5 @@ def test_fts_disabled_search_api_blocks_document_filter( ) with pytest.raises(InvalidArgumentError) as exc_info: - collection.search( - Search(where=Key.DOCUMENT.contains("alpha")) - ) + collection.search(Search(where=Key.DOCUMENT.contains("alpha"))) assert "fts" in str(exc_info.value).lower() diff --git a/chromadb/test/property/invariants.py b/chromadb/test/property/invariants.py index f1509863ab4..4f0b03c7df7 100644 --- a/chromadb/test/property/invariants.py +++ b/chromadb/test/property/invariants.py @@ -319,6 +319,25 @@ def ann_accuracy( dim = len(embeddings[0]) accuracy_threshold = accuracy_threshold * math.pow(10, int(math.log10(dim))) + # Detect whether quantization is active on this collection so the + # per-result distance equality check below can use a wider tolerance + # derived from the RaBitQ paper's theoretical error bound. The + # quantize field lives in the schema (under #embedding key), not in + # configuration_json. + quantization_active = False + serialized_schema = collection._model.serialized_schema + if serialized_schema is not None: + embedding_spann_cfg = ( + serialized_schema.get("keys", {}) + .get("#embedding", {}) + .get("float_list", {}) + .get("vector_index", {}) + .get("config", {}) + .get("spann", {}) + ) + if embedding_spann_cfg.get("quantize") not in (None, "none"): + quantization_active = True + # Perform exact distance computation if query_embeddings is None: query_embeddings = ( @@ -398,11 +417,39 @@ def ann_accuracy( unexpected_id = id not in expected_ids index = id_to_index[id] - correct_distance = np.allclose( - distances_i[index], - query_results["distances"][i][j], - atol=accuracy_threshold, - ) + # For quantized SPANN (4-bit RaBitQ) the paper's Theorem 3.2 + # bounds the distance-estimator error by O(1/sqrt(D)) with a + # constant that scales with the norms of the data and query + # vectors. We observed on test_add_mcmr that 100% of failing + # quantized comparisons fall within 5 * ||q|| * ||d|| / sqrt(D) + # (P95 at ~9% of that bound, max at 16%). The corresponding + # *relative* error can be large when the true distance value + # happens to be near zero (notably for inner product, whose + # value has no non-zero floor), which is why a simple rtol + # cannot cover quantized IP. + # + # For quantized collections we therefore use an absolute + # tolerance derived from the RaBitQ bound in addition to the + # standard accuracy_threshold floor. For non-quantized + # collections we omit rtol so numpy's default rtol=1e-5 + # applies, matching pre-branch behavior (e.g. + # test_cross_version_persist's local HNSW relies on it for + # small numerical drift). + if quantization_active: + q_norm = float(np.linalg.norm(query_embeddings[i])) + d_norm = float(np.linalg.norm(embeddings[index])) + rabitq_atol = 5.0 * q_norm * d_norm / math.sqrt(dim) + correct_distance = np.allclose( + distances_i[index], + query_results["distances"][i][j], + atol=max(accuracy_threshold, rabitq_atol), + ) + else: + correct_distance = np.allclose( + distances_i[index], + query_results["distances"][i][j], + atol=accuracy_threshold, + ) if unexpected_id: # If the ID is unexpcted, but the distance is correct, then we # have a duplicate in the data. In this case, we should not reduce recall. diff --git a/chromadb/test/property/strategies.py b/chromadb/test/property/strategies.py index 2fd0575f9d6..22c8b6dccb3 100644 --- a/chromadb/test/property/strategies.py +++ b/chromadb/test/property/strategies.py @@ -694,7 +694,15 @@ def collections( name = draw(collection_name()) metadata = draw(collection_metadata) - dimension = draw(st.integers(min_value=2, max_value=2048)) + # Dimension bounds tuned for 4-bit RaBitQ quantization (enabled broadly + # in the distributed test tenant configs): + # - min 128: below this, RaBitQ's concentration-of-measure bound + # (~1/sqrt(d)) produces relative error >10% on cosine/L2 distances, + # and hypothesis eagerly shrinks to pathological low-dim cases. + # - max 512: quantized compaction's per-record cost is dominated by + # a dense O(dim^2) random rotation matmul; capping dimension keeps + # per-compaction CPU bounded so the test fits in COMPACTION_SLEEP. + dimension = draw(st.integers(min_value=128, max_value=512)) dtype = draw(st.sampled_from(float_types)) use_persistent_hnsw_params = draw(with_persistent_hnsw_params) diff --git a/k8s/distributed-chroma/values.dev.yaml b/k8s/distributed-chroma/values.dev.yaml index e72a35e1b9d..8b5717c681f 100644 --- a/k8s/distributed-chroma/values.dev.yaml +++ b/k8s/distributed-chroma/values.dev.yaml @@ -35,9 +35,9 @@ queryService: jemallocConfig: "prof:true,prof_active:true,lg_prof_sample:19" resources: limits: - cpu: 100m + cpu: 300m requests: - cpu: 100m + cpu: 300m compactionService: env: @@ -46,17 +46,17 @@ compactionService: jemallocConfig: "prof:true,prof_active:true,lg_prof_sample:19" resources: limits: - cpu: 200m + cpu: 1500m requests: - cpu: 200m + cpu: 1500m rustLogService: replicaCount: 1 resources: limits: - cpu: 200m + cpu: 500m requests: - cpu: 200m + cpu: 500m garbageCollector: jemallocConfig: "prof:true,prof_active:true,lg_prof_sample:19" diff --git a/k8s/distributed-chroma/values2.dev.yaml b/k8s/distributed-chroma/values2.dev.yaml index af8021ad1fe..28bc3bad2e0 100644 --- a/k8s/distributed-chroma/values2.dev.yaml +++ b/k8s/distributed-chroma/values2.dev.yaml @@ -35,9 +35,9 @@ queryService: jemallocConfig: "prof:true,prof_active:true,lg_prof_sample:19" resources: limits: - cpu: 200m + cpu: 300m requests: - cpu: 200m + cpu: 300m replicaCount: 1 compactionService: @@ -47,17 +47,17 @@ compactionService: jemallocConfig: "prof:true,prof_active:true,lg_prof_sample:19" resources: limits: - cpu: 200m + cpu: 1500m requests: - cpu: 200m + cpu: 1500m rustLogService: replicaCount: 1 resources: limits: - cpu: 200m + cpu: 500m requests: - cpu: 200m + cpu: 500m garbageCollector: jemallocConfig: "prof:true,prof_active:true,lg_prof_sample:19" diff --git a/rust/frontend/sample_configs/distributed.yaml b/rust/frontend/sample_configs/distributed.yaml index 981da516172..5985d16dd34 100644 --- a/rust/frontend/sample_configs/distributed.yaml +++ b/rust/frontend/sample_configs/distributed.yaml @@ -84,6 +84,7 @@ enable_schema: true tenants_to_migrate_immediately: - "default_tenant" tenants_to_migrate_immediately_threshold: "ffffffff-ffff-ffff-ffff-ffffffffffff" -tenants_with_quantization_enabled: [] +tenants_with_quantization_enabled: +- "*" tenants_with_maxscore_enabled: [] enable_log_scouting: true diff --git a/rust/frontend/sample_configs/distributed2.yaml b/rust/frontend/sample_configs/distributed2.yaml index b3992b5c954..e081cb0330f 100644 --- a/rust/frontend/sample_configs/distributed2.yaml +++ b/rust/frontend/sample_configs/distributed2.yaml @@ -74,6 +74,7 @@ enable_schema: true tenants_to_migrate_immediately: - "default_tenant" tenants_to_migrate_immediately_threshold: "ffffffff-ffff-ffff-ffff-ffffffffffff" -tenants_with_quantization_enabled: [] +tenants_with_quantization_enabled: +- "*" tenants_with_maxscore_enabled: [] enable_log_scouting: true diff --git a/rust/frontend/sample_configs/distributed_mcmr.yaml b/rust/frontend/sample_configs/distributed_mcmr.yaml index d5a96335e2e..71960e0d4f9 100644 --- a/rust/frontend/sample_configs/distributed_mcmr.yaml +++ b/rust/frontend/sample_configs/distributed_mcmr.yaml @@ -74,6 +74,7 @@ enable_schema: true tenants_to_migrate_immediately: - "default_tenant" tenants_to_migrate_immediately_threshold: "ffffffff-ffff-ffff-ffff-ffffffffffff" -tenants_with_quantization_enabled: [] +tenants_with_quantization_enabled: +- "*" tenants_with_maxscore_enabled: [] enable_log_scouting: true diff --git a/rust/frontend/src/impls/service_based_frontend.rs b/rust/frontend/src/impls/service_based_frontend.rs index c4d1a79dc59..84a8f7bb1ae 100644 --- a/rust/frontend/src/impls/service_based_frontend.rs +++ b/rust/frontend/src/impls/service_based_frontend.rs @@ -3012,6 +3012,10 @@ mod tests { ) || segments .iter() .any(|s| s.r#type == SegmentType::Spann && s.scope == SegmentScope::VECTOR) + || segments + .iter() + .any(|s| s.r#type == SegmentType::QuantizedSpann + && s.scope == SegmentScope::VECTOR) ); assert!(segments .iter()