diff --git a/vectordb_bench/backend/clients/antfly/antfly.py b/vectordb_bench/backend/clients/antfly/antfly.py index 2c39f2b6e..96b82f935 100644 --- a/vectordb_bench/backend/clients/antfly/antfly.py +++ b/vectordb_bench/backend/clients/antfly/antfly.py @@ -41,6 +41,36 @@ def _make_client(base_url: str, timeout: float) -> httpx.Client: ) +def _looks_like_json_response(response: httpx.Response) -> bool: + content_type = response.headers.get("content-type", "") + body = response.content.strip() + return ( + response.is_success + and ("json" in content_type or body.startswith((b"{", b"["))) + ) + + +def _detect_metadata_base_url(host: str, port: int) -> str: + root = f"http://{_httpx_host(host)}:{port}" + candidates = (f"{root}/db/v1", f"{root}/api/v1") + for base_url in candidates: + try: + with httpx.Client(base_url=base_url, timeout=5) as client: + r = client.get("/status") + if _looks_like_json_response(r): + return base_url + except Exception: + pass + try: + with httpx.Client(base_url=root, timeout=5) as client: + r = client.get("/api/v1/tables") + if _looks_like_json_response(r): + return f"{root}/api/v1" + except Exception: + pass + return candidates[0] + + class Antfly(VectorDB): def __init__( self, @@ -56,13 +86,14 @@ def __init__( self.collection_name = collection_name self.dim = dim - self._metadata_base_url = ( - f"http://{_httpx_host(db_config['host'])}:{db_config['port']}/api/v1" + self._metadata_base_url = _detect_metadata_base_url( + db_config["host"], db_config["port"] ) self._store_host = _httpx_host(db_config.get("store_host") or db_config["host"]) self._store_port = db_config.get("store_port") self._use_direct_store_search = bool(db_config.get("use_direct_store_search")) self._pack_query_vectors = bool(db_config.get("pack_query_vectors")) + self._legacy_wire = not self._pack_query_vectors self._direct_shard_id: str | None = None num_shards = db_config.get("num_shards", 1) @@ -77,16 +108,24 @@ def __init__( r = client.delete(f"/tables/{self.collection_name}") log.info(f"Drop table response: {r.status_code}") - table = self._get_table_status_or_none(client) - if table is None: + def create_table_if_needed() -> None: + table = self._get_table_status_or_none(client) + if table is not None: + log.info("Reusing existing table: %s", self.collection_name) + return r = client.post( f"/tables/{self.collection_name}", json={"num_shards": num_shards} ) log.info(f"Create table response: {r.status_code}") r.raise_for_status() - else: - log.info("Reusing existing table: %s", self.collection_name) + def reset_table() -> None: + r = client.delete(f"/tables/{self.collection_name}") + log.info(f"Reset table response: {r.status_code}") + create_table_if_needed() + self._wait_for_shard_ready(client) + + create_table_if_needed() self._wait_for_shard_ready(client) if self._get_index_status(client) is None: @@ -97,28 +136,48 @@ def __init__( **self.case_config.index_param(), } index_error = None - # Try each index type, with and without field, to handle - # both old binaries (require field) and new source (reject field with external). - for index_type in INDEX_TYPES: - for extra in ({}, {"field": SOURCE_FIELD}): - r = client.post( - f"/tables/{self.collection_name}/indexes/{INDEX_NAME}", - json={"type": index_type, **index_def, **extra}, - ) - log.info( - f"Add embeddings index response ({index_type}, field={'field' in extra}): {r.status_code}" - ) - if r.is_success: - index_error = None - break + index_created = False + extras = ({"field": SOURCE_FIELD}, {}) if self._legacy_wire else ({}, {"field": SOURCE_FIELD}) + candidates = [ + (index_type, extra) + for index_type in INDEX_TYPES + for extra in extras + ] + # Treat a create-index HTTP success as provisional. Older stable + # binaries can accept a metadata change that the shard later + # rejects, so require a ready status before loading vectors. + for idx, (index_type, extra) in enumerate(candidates): + if idx > 0: + reset_table() + r = client.post( + f"/tables/{self.collection_name}/indexes/{INDEX_NAME}", + json={"type": index_type, **index_def, **extra}, + ) + log.info( + f"Add embeddings index response ({index_type}, field={'field' in extra}): {r.status_code}" + ) + if not r.is_success: index_error = r - if index_error is None: + continue + if self._legacy_wire: + index_created = True break - if index_error is not None: + if self._wait_for_index_ready( + client, + expected_total=0, + timeout=TABLE_READY_TIMEOUT, + ): + index_created = True + break + index_error = r + if not index_created and index_error is not None: index_error.raise_for_status() + if not index_created: + raise RuntimeError("Antfly index was created but never became ready") else: log.info("Reusing existing embeddings index: %s", INDEX_NAME) - self._wait_for_index_ready(client, expected_total=0) + if not self._legacy_wire: + self._wait_for_index_ready(client, expected_total=0) self._refresh_direct_search_routing(client) finally: client.close() @@ -151,21 +210,59 @@ def _wait_for_shard_ready(self, client: httpx.Client): def _get_index_status(self, client: httpx.Client) -> dict | None: r = client.get(f"/tables/{self.collection_name}/indexes/{INDEX_NAME}") if r.status_code == 404: - return None + return self._get_legacy_index_status(client) r.raise_for_status() + body = r.content.strip() + if not body or not body.startswith(b"{"): + return self._get_legacy_index_status(client) return r.json() def _get_table_status(self, client: httpx.Client) -> dict: + table = self._get_table_status_or_none(client) + if table is None: + raise RuntimeError(f"Antfly table not found: {self.collection_name}") + return table + + def _get_table_status_or_none(self, client: httpx.Client) -> dict | None: r = client.get(f"/tables/{self.collection_name}") + if r.status_code == 404: + return self._get_legacy_table_status(client) r.raise_for_status() + body = r.content.strip() + if not body or not body.startswith(b"{"): + return self._get_legacy_table_status(client) return r.json() - def _get_table_status_or_none(self, client: httpx.Client) -> dict | None: - r = client.get(f"/tables/{self.collection_name}") + def _get_legacy_table_status(self, client: httpx.Client) -> dict | None: + r = client.get("/tables") if r.status_code == 404: return None r.raise_for_status() - return r.json() + body = r.content.strip() + if not body or not body.startswith(b"["): + return None + for table in r.json(): + if isinstance(table, dict) and table.get("name") == self.collection_name: + return table + return None + + def _get_legacy_index_status(self, client: httpx.Client) -> dict | None: + r = client.get("/status") + if r.status_code == 404: + return None + r.raise_for_status() + body = r.content.strip() + if not body or not body.startswith(b"{"): + return None + statuses = (((r.json().get("shards") or {}).get("statuses")) or {}) + for shard in statuses.values(): + if not isinstance(shard, dict) or shard.get("table") != self.collection_name: + continue + indexes = ((((shard.get("info") or {}).get("shard_stats") or {}).get("indexes")) or {}) + status = indexes.get(INDEX_NAME) + if isinstance(status, dict): + return {"status": status} + return None def _refresh_direct_search_routing(self, client: httpx.Client): if not self._use_direct_store_search: @@ -186,21 +283,28 @@ def _index_status_is_ready( if payload is None: return False if status is None: - return expected_total == 0 + return False rebuilding = bool(status.get("rebuilding")) wal_backlog = int(status.get("wal_backlog", 0) or 0) total_indexed = int(status.get("total_indexed", 0) or 0) + doc_count = int(status.get("doc_count", 0) or 0) has_error = bool(status.get("error")) + if expected_total == 0 and total_indexed == 0 and doc_count == 0: + return not has_error and wal_backlog == 0 + if has_error or rebuilding or wal_backlog > 0: return False return expected_total is None or total_indexed >= expected_total def _wait_for_index_ready( - self, client: httpx.Client, expected_total: int | None = None - ): - deadline = time.monotonic() + INDEX_READY_TIMEOUT + self, + client: httpx.Client, + expected_total: int | None = None, + timeout: int = INDEX_READY_TIMEOUT, + ) -> bool: + deadline = time.monotonic() + timeout last_status = None while time.monotonic() < deadline: @@ -210,17 +314,18 @@ def _wait_for_index_ready( last_status = status if self._index_status_is_ready(payload, status, expected_total): log.info(f"Embeddings index is ready: {status}") - return + return True except Exception as e: last_status = {"error": str(e)} time.sleep(INDEX_READY_POLL_INTERVAL) log.warning( "Embeddings index readiness timeout after %ss, expected_total=%s, last_status=%s", - INDEX_READY_TIMEOUT, + timeout, expected_total, last_status, ) + return False @contextmanager def init(self): @@ -271,8 +376,10 @@ def _serialize_query_vector(self, vector: list[float]) -> list[float] | str: return self._pack_vector(vector) return vector - def _serialize_insert_vector(self, vector: list[float]) -> str: - return self._pack_vector(vector) + def _serialize_insert_vector(self, vector: list[float]) -> list[float] | str: + if self._pack_query_vectors or os.environ.get("ANTFLY_PACK_VECTORS") == "1": + return self._pack_vector(vector) + return vector def _metadata_query_body(self, query: list[float], k: int) -> dict[str, Any]: return {