diff --git a/tests/test_aws_opensearch.py b/tests/test_aws_opensearch.py index 2e2930338..b099e3028 100644 --- a/tests/test_aws_opensearch.py +++ b/tests/test_aws_opensearch.py @@ -1,16 +1,32 @@ from types import SimpleNamespace import pytest +from opensearchpy import ConnectionTimeout, TransportError from vectordb_bench import config -from vectordb_bench.backend.clients.aws_opensearch.aws_opensearch import AWSOpenSearch +from vectordb_bench.backend.clients.aws_opensearch.aws_opensearch import ( + BULK_MAX_ATTEMPTS, + AWSOpenSearch, + OpenSearchBulkInsertError, +) -def test_serverless_insert_uses_configured_batch_size(monkeypatch) -> None: - bulk_requests = [] +def _bulk_response(*statuses: int) -> dict[str, object]: + items = [] + for position, status in enumerate(statuses): + result = {"_id": str(position), "status": status} + if status >= 300: + result["error"] = {"type": "rejected", "reason": "test failure"} + items.append({"index": result}) + return {"errors": any(status >= 300 for status in statuses), "items": items} - def bulk(*, body): + +def test_serverless_insert_uses_configured_batch_size(monkeypatch: pytest.MonkeyPatch) -> None: + bulk_requests: list[list[dict[str, object]]] = [] + + def bulk(*, body: list[dict[str, object]]) -> dict[str, object]: bulk_requests.append(body) + return _bulk_response(*(201 for _ in body[::2])) monkeypatch.setattr(config, "NUM_PER_BATCH", 2) @@ -33,7 +49,10 @@ def bulk(*, body): @pytest.mark.parametrize("batch_size", [0, -1]) -def test_serverless_insert_rejects_non_positive_batch_size(monkeypatch, batch_size: int) -> None: +def test_serverless_insert_rejects_non_positive_batch_size( + monkeypatch: pytest.MonkeyPatch, + batch_size: int, +) -> None: monkeypatch.setattr(config, "NUM_PER_BATCH", batch_size) db = object.__new__(AWSOpenSearch) @@ -44,3 +63,203 @@ def test_serverless_insert_rejects_non_positive_batch_size(monkeypatch, batch_si embeddings=[[0.1]], metadata=[1], ) + + +def test_serverless_insert_retries_only_failed_documents(monkeypatch: pytest.MonkeyPatch) -> None: + bulk_requests: list[list[dict[str, object]]] = [] + retry_delays: list[int] = [] + responses = iter([_bulk_response(201, 429, 201), _bulk_response(201)]) + + def bulk(*, body: list[dict[str, object]]) -> dict[str, object]: + bulk_requests.append(body) + return next(responses) + + monkeypatch.setattr(config, "NUM_PER_BATCH", 3) + monkeypatch.setattr( + "vectordb_bench.backend.clients.aws_opensearch.aws_opensearch.time.sleep", + retry_delays.append, + ) + + db = object.__new__(AWSOpenSearch) + db.client = SimpleNamespace(bulk=bulk) + db._is_serverless = True + db.index_name = "test-index" + db.vector_col_name = "embedding" + db.with_scalar_labels = False + + inserted, error = db._insert_with_single_client( + embeddings=[[0.1], [0.2], [0.3]], + metadata=[1, 2, 3], + ) + + assert inserted == 3 + assert error is None + assert [len(request) // 2 for request in bulk_requests] == [3, 1] + assert bulk_requests[1][1]["id"] == 2 + assert retry_delays == [2] + + +def test_serverless_insert_fails_after_partial_failure_exhausts_attempts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + bulk_requests: list[list[dict[str, object]]] = [] + retry_delays: list[int] = [] + errors: list[str] = [] + + def bulk(*, body: list[dict[str, object]]) -> dict[str, object]: + bulk_requests.append(body) + if len(bulk_requests) == 1: + return _bulk_response(201, 429) + return _bulk_response(429) + + monkeypatch.setattr(config, "NUM_PER_BATCH", 2) + monkeypatch.setattr( + "vectordb_bench.backend.clients.aws_opensearch.aws_opensearch.time.sleep", + retry_delays.append, + ) + monkeypatch.setattr( + "vectordb_bench.backend.clients.aws_opensearch.aws_opensearch.log.error", + errors.append, + ) + + db = object.__new__(AWSOpenSearch) + db.client = SimpleNamespace(bulk=bulk) + db._is_serverless = True + db.index_name = "test-index" + db.vector_col_name = "embedding" + db.with_scalar_labels = False + + inserted, error = db._insert_with_single_client( + embeddings=[[0.1], [0.2]], + metadata=[1, 2], + ) + + assert inserted == 1 + assert isinstance(error, OpenSearchBulkInsertError) + assert error.non_retryable is True + assert "left 1 documents uninserted after 30 attempts; successful=1" in str(error) + assert len(bulk_requests) == BULK_MAX_ATTEMPTS + assert [len(request) // 2 for request in bulk_requests] == [2] + [1] * (BULK_MAX_ATTEMPTS - 1) + assert retry_delays == [2, 4, 8, 16, 32] + [60] * 24 + assert sum(retry_delays) == 1502 + assert any("left 1 documents uninserted after 30 attempts; successful=1" in message for message in errors) + + +def test_serverless_insert_retries_request_level_429(monkeypatch: pytest.MonkeyPatch) -> None: + bulk_requests: list[list[dict[str, object]]] = [] + retry_delays: list[int] = [] + + def bulk(*, body: list[dict[str, object]]) -> dict[str, object]: + bulk_requests.append(body) + if len(bulk_requests) <= 2: + raise TransportError(429, "too many requests") + return _bulk_response(201) + + monkeypatch.setattr(config, "NUM_PER_BATCH", 1) + monkeypatch.setattr( + "vectordb_bench.backend.clients.aws_opensearch.aws_opensearch.time.sleep", + retry_delays.append, + ) + + db = object.__new__(AWSOpenSearch) + db.client = SimpleNamespace(bulk=bulk) + db._is_serverless = True + db.index_name = "test-index" + db.vector_col_name = "embedding" + db.with_scalar_labels = False + + inserted, error = db._insert_with_single_client( + embeddings=[[0.1]], + metadata=[1], + ) + + assert inserted == 1 + assert error is None + assert len(bulk_requests) == 3 + assert retry_delays == [2, 4] + + +def test_serverless_insert_does_not_retry_ambiguous_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + bulk_requests: list[list[dict[str, object]]] = [] + retry_delays: list[int] = [] + + def bulk(*, body: list[dict[str, object]]) -> dict[str, object]: + bulk_requests.append(body) + raise ConnectionTimeout(None, "timed out", None) + + monkeypatch.setattr(config, "NUM_PER_BATCH", 1) + monkeypatch.setattr( + "vectordb_bench.backend.clients.aws_opensearch.aws_opensearch.time.sleep", + retry_delays.append, + ) + + db = object.__new__(AWSOpenSearch) + db.client = SimpleNamespace(bulk=bulk) + db._is_serverless = True + db.index_name = "test-index" + db.vector_col_name = "embedding" + db.with_scalar_labels = False + + inserted, error = db._insert_with_single_client( + embeddings=[[0.1]], + metadata=[1], + ) + + assert inserted == 0 + assert isinstance(error, OpenSearchBulkInsertError) + assert error.non_retryable is True + assert "ambiguous outcome" in str(error) + assert len(bulk_requests) == 1 + assert retry_delays == [] + + +def test_multiple_clients_fallback_preserves_labels(monkeypatch: pytest.MonkeyPatch) -> None: + class BulkClient: + def __init__(self, response: dict[str, object]) -> None: + self.response = response + + def bulk(self, *, body: list[dict[str, object]]) -> dict[str, object]: + return self.response + + def close(self) -> None: + return None + + clients = iter([BulkClient(_bulk_response(201)), BulkClient(_bulk_response(429))]) + monkeypatch.setattr( + "vectordb_bench.backend.clients.aws_opensearch.aws_opensearch.OpenSearch", + lambda **_: next(clients), + ) + monkeypatch.setattr("vectordb_bench.backend.clients.aws_opensearch.aws_opensearch.time.sleep", lambda _: None) + + fallback_requests: list[list[dict[str, object]]] = [] + + def fallback_bulk(*, body: list[dict[str, object]]) -> dict[str, object]: + fallback_requests.append(body) + return _bulk_response(*(201 for _ in body[::2])) + + db = object.__new__(AWSOpenSearch) + db.client = SimpleNamespace( + bulk=fallback_bulk, + indices=SimpleNamespace( + stats=lambda **_: {"_all": {"primaries": {"indexing": {"index_total": 1}}}}, + ), + ) + db.db_config = {} + db.case_config = SimpleNamespace(use_routing=False) + db.index_name = "test-index" + db.id_col_name = "_id" + db.vector_col_name = "embedding" + db.label_col_name = "label" + db.with_scalar_labels = True + db._is_serverless = False + + inserted, error = db._insert_with_multiple_clients( + embeddings=[[0.1], [0.2]], + metadata=[1, 2], + num_clients=2, + labels_data=["first", "second"], + ) + + assert inserted == 2 + assert error is None + assert [document["label"] for document in fallback_requests[0][1::2]] == ["first", "second"] diff --git a/vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py b/vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py index eb0195b88..833e198b5 100644 --- a/vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py +++ b/vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py @@ -2,6 +2,7 @@ import time from collections.abc import Iterable from contextlib import contextmanager +from typing import Any from opensearchpy import OpenSearch @@ -16,6 +17,13 @@ WAITING_FOR_REFRESH_SEC = 30 WAITING_FOR_FORCE_MERGE_SEC = 30 SECONDS_WAITING_FOR_REPLICAS_TO_BE_ENABLED_SEC = 30 +BULK_MAX_ATTEMPTS = 30 +BULK_INITIAL_RETRY_DELAY_SEC = 2 +BULK_MAX_RETRY_DELAY_SEC = 60 + + +class OpenSearchBulkInsertError(RuntimeError): + non_retryable = True class AWSOpenSearch(VectorDB): @@ -285,20 +293,132 @@ def _insert_with_single_client( other_data[self.label_col_name] = batch_labels[j] insert_data.append(other_data) + inserted, error = self._execute_bulk_with_retries( + self.client, + insert_data, + f"index {self.index_name}", + retry_ambiguous_request_errors=not self._is_serverless, + ) + total_inserted += inserted + if error is not None: + return total_inserted, error + + return total_inserted, None + + @staticmethod + def _parse_bulk_response( + response: dict[str, Any], + insert_data: list[dict[str, Any]], + ) -> tuple[int, list[dict[str, Any]], list[str]]: + expected_count = len(insert_data) // 2 + if not response.get("errors"): + return expected_count, [], [] + + items = response.get("items") + if not isinstance(items, list): + return 0, insert_data, ["response did not contain an items list"] + + success_count = 0 + failed_data = [] + failure_samples = [] + for position in range(expected_count): + item = items[position] if position < len(items) else None + if isinstance(item, dict) and len(item) == 1: + operation, result = next(iter(item.items())) + if isinstance(result, dict): + status = result.get("status") + if isinstance(status, int) and 200 <= status < 300 and "error" not in result: + success_count += 1 + continue + error = result.get("error", "unknown") + failure_samples.append( + f"item[{position}] {operation} id={result.get('_id', 'unknown')} " + f"status={status or 'unknown'} error={error}" + ) + else: + failure_samples.append(f"item[{position}] {operation}=malformed") + else: + failure_samples.append(f"item[{position}]=malformed") + failed_data.extend(insert_data[position * 2 : position * 2 + 2]) + + return success_count, failed_data, failure_samples + + def _execute_bulk_with_retries( + self, + client: OpenSearch, + insert_data: list[dict[str, Any]], + context: str, + retry_ambiguous_request_errors: bool = True, + ) -> tuple[int, Exception | None]: + pending_data = insert_data + total_inserted = 0 + final_error = None + + for attempt in range(1, BULK_MAX_ATTEMPTS + 1): + response: Any = None + request_error = None try: - self.client.bulk(body=insert_data) - total_inserted += len(batch_embeddings) + response = client.bulk(body=pending_data) except Exception as e: - log.warning(f"Failed to insert batch: {self.index_name} error: {e!s}") - time.sleep(10) - try: - self.client.bulk(body=insert_data) - total_inserted += len(batch_embeddings) - except Exception as retry_e: - log.warning(f"Retry failed for batch: {retry_e!s}") - return total_inserted, retry_e + request_error = e + + if request_error is None and not isinstance(response, dict): + request_error = TypeError("OpenSearch bulk response was not an object") + + if request_error is not None: + status_code = getattr(request_error, "status_code", None) + request_is_retryable = status_code == 429 or retry_ambiguous_request_errors + if request_is_retryable and attempt < BULK_MAX_ATTEMPTS: + retry_delay = min( + BULK_INITIAL_RETRY_DELAY_SEC * (2 ** (attempt - 1)), + BULK_MAX_RETRY_DELAY_SEC, + ) + log.warning( + f"Bulk request failed for {context}; next attempt {attempt + 1}/{BULK_MAX_ATTEMPTS} " + f"in {retry_delay}s: {request_error!s}" + ) + time.sleep(retry_delay) + continue + if request_is_retryable: + message = f"Bulk request failed for {context} after {BULK_MAX_ATTEMPTS} attempts" + else: + message = ( + f"Bulk request failed for {context} with an ambiguous outcome; " + "not retrying an auto-ID request" + ) + final_error = OpenSearchBulkInsertError(f"{message}; successful={total_inserted}: {request_error!s}") + break - return total_inserted, None + inserted, failed_data, failure_samples = self._parse_bulk_response(response, pending_data) + total_inserted += inserted + if not failed_data: + return total_inserted, None + + failed_count = len(failed_data) // 2 + sample_summary = "; ".join(failure_samples[:3]) + if attempt < BULK_MAX_ATTEMPTS: + retry_delay = min( + BULK_INITIAL_RETRY_DELAY_SEC * (2 ** (attempt - 1)), + BULK_MAX_RETRY_DELAY_SEC, + ) + log.warning( + f"Bulk response for {context} rejected {failed_count} documents; " + f"next attempt {attempt + 1}/{BULK_MAX_ATTEMPTS} in {retry_delay}s; " + f"{sample_summary}" + ) + pending_data = failed_data + time.sleep(retry_delay) + continue + + final_error = OpenSearchBulkInsertError( + f"Bulk insert for {context} left {failed_count} documents uninserted after " + f"{BULK_MAX_ATTEMPTS} attempts; successful={total_inserted}; {sample_summary}" + ) + break + + assert final_error is not None + log.error(str(final_error)) + return total_inserted, final_error def _insert_with_multiple_clients( self, @@ -342,19 +462,7 @@ def insert_chunk(client_idx: int, chunk_idx: int): other_data[self.label_col_name] = chunk_labels_data[i] insert_data.append(other_data) - max_retries = 10 - for attempt in range(max_retries): - try: - client.bulk(body=insert_data) - return len(chunk_embeddings), None - except Exception as e: - if "429" in str(e) and attempt < max_retries - 1: - log.warning(f"Client {client_idx} got 429 error, retry {attempt + 1}/{max_retries} after 10s") - time.sleep(10) - else: - log.warning(f"Client {client_idx} failed to insert data: {e!s}") - return 0, e - return 0, Exception("Max retries exceeded") + return self._execute_bulk_with_retries(client, insert_data, f"client {client_idx}") results = [] with ThreadPoolExecutor(max_workers=len(clients)) as executor: @@ -380,7 +488,7 @@ def insert_chunk(client_idx: int, chunk_idx: int): if errors: log.warning("Some clients failed to insert data, retrying with single client") time.sleep(10) - return self._insert_with_single_client(embeddings, metadata) + return self._insert_with_single_client(embeddings, metadata, labels_data) resp = self.client.indices.stats(index=self.index_name) log.info(