diff --git a/src/dynavec/config.py b/src/dynavec/config.py index a2e3c49..e786f5d 100644 --- a/src/dynavec/config.py +++ b/src/dynavec/config.py @@ -91,6 +91,8 @@ class DynavecConfig: max_workers: int = 8 parallel_writes: bool = True max_pool_connections: int | None = None + put_rps: float | None = None + query_rps: float | None = None # in-memory hot tier (optional): keep a hot working set in RAM so warmed # namespaces are served entirely from memory — no S3 Vectors query and no @@ -158,3 +160,11 @@ def __post_init__(self) -> None: if self.dimension > 4096: logger = logging.getLogger(__name__) logger.warning("Amazon S3 Vectors currently supports a maximum embedding dimension of 4096. You have configured a dimension of %d. This may result in an API error during provisioning or writing data.", self.dimension) + + if self.put_rps is not None and self.put_rps <= 0: + raise ValueError("put_rps must be positive") + + if self.query_rps is not None and self.query_rps <= 0: + raise ValueError("query_rps must be positive") + + diff --git a/src/dynavec/stores/s3vectors.py b/src/dynavec/stores/s3vectors.py index c41b257..60d1053 100644 --- a/src/dynavec/stores/s3vectors.py +++ b/src/dynavec/stores/s3vectors.py @@ -18,7 +18,7 @@ from ..config import DynavecConfig from ..logging import log_store_event -from ..utils import retry +from ..utils import TokenBucket, retry Metadata = dict[str, Any] @@ -45,6 +45,12 @@ def __init__(self, config: DynavecConfig, boto_session: Any | None = None) -> No session = boto_session or boto3.Session() self._config = config + self._put_limiter = ( + TokenBucket(config.put_rps) if config.put_rps is not None else None + ) + self._query_limiter = ( + TokenBucket(config.query_rps) if config.query_rps is not None else None + ) client_kwargs: dict[str, object] = {"region_name": config.region} botocore_config = config.botocore_config() if botocore_config is not None: @@ -61,6 +67,8 @@ def get_index(self) -> dict[str, Any]: @retry() def _put_batch(self, payload: list[dict[str, Any]]) -> None: + if self._put_limiter is not None: + self._put_limiter.acquire() self._client.put_vectors( vectorBucketName=self._config.vector_bucket, indexName=self._config.index, @@ -184,7 +192,7 @@ def query_pages( return_distance: bool = True, page_size: int | None = None, ) -> Iterator[list[dict[str, Any]]]: - """Yield result **pages** as the S3 Vectors paginator returns them. + """Yield result pages as the S3 Vectors paginator returns them. This is what powers streaming search: an agent can begin consuming the first page while later pages are still in flight. @@ -198,6 +206,7 @@ def query_pages( """ if top_k <= 0: return + if top_k > _MAX_TOP_K: raise ValueError( f"top_k ({top_k}) exceeds Amazon S3 Vectors maximum limit of {_MAX_TOP_K}." @@ -206,45 +215,74 @@ def query_pages( effective_page_size = ( page_size if page_size is not None else self._config.top_k_page_size ) + if effective_page_size is not None and effective_page_size <= 0: raise ValueError("page_size must be a positive integer.") kwargs = self._query_kwargs( - query_vector, top_k, filter, return_metadata, return_distance + query_vector, + top_k, + filter, + return_metadata, + return_distance, ) + paginator = self._client.get_paginator("query_vectors") yielded = 0 buffer: list[dict[str, Any]] = [] - for page in paginator.paginate( - PaginationConfig={"MaxItems": top_k}, - **kwargs, - ): + page_iterator = iter( + paginator.paginate( + PaginationConfig={"MaxItems": top_k}, + **kwargs, + ) + ) + + while True: + if self._query_limiter is not None: + self._query_limiter.acquire() + + try: + page = next(page_iterator) + except StopIteration: + break + vectors = page.get("vectors", []) - if not vectors: - continue - if effective_page_size is None: - remaining = top_k - yielded - if len(vectors) > remaining: - vectors = vectors[:remaining] - yield vectors - yielded += len(vectors) - if yielded >= top_k: - return - else: - buffer.extend(vectors) - while len(buffer) >= effective_page_size and yielded < top_k: - chunk = buffer[:effective_page_size] - buffer = buffer[effective_page_size:] + if vectors: + if effective_page_size is None: remaining = top_k - yielded - if len(chunk) > remaining: - chunk = chunk[:remaining] - yield chunk - yielded += len(chunk) + + if len(vectors) > remaining: + vectors = vectors[:remaining] + + yield vectors + yielded += len(vectors) + if yielded >= top_k: return + else: + buffer.extend(vectors) + + while len(buffer) >= effective_page_size and yielded < top_k: + chunk = buffer[:effective_page_size] + buffer = buffer[effective_page_size:] + + remaining = top_k - yielded + + if len(chunk) > remaining: + chunk = chunk[:remaining] + + yield chunk + yielded += len(chunk) + + if yielded >= top_k: + return + + if page.get("NextToken") is None: + break + if effective_page_size is not None and buffer and yielded < top_k: remaining = top_k - yielded yield buffer[:remaining] diff --git a/src/dynavec/utils.py b/src/dynavec/utils.py index ffff152..3e2c03d 100644 --- a/src/dynavec/utils.py +++ b/src/dynavec/utils.py @@ -8,6 +8,7 @@ import functools import random +import threading import time from collections.abc import Iterable, Iterator from typing import Any, Callable, TypeVar @@ -120,3 +121,44 @@ def chunked(iterable: Iterable[T], size: int) -> Iterator[list[T]]: batch = [] if batch: yield batch + + +class TokenBucket: + def __init__( + self, + rate: float, + capacity: float | None = None, + ) -> None: + if rate <= 0: + raise ValueError("rate must be positive") + + self.rate = rate + self.capacity = capacity if capacity is not None else rate + + if self.capacity <= 0: + raise ValueError("capacity must be positive") + + self.tokens = self.capacity + self.last_time = time.monotonic() + self._lock = threading.Lock() + + + def acquire(self) -> None: + while True: + with self._lock: + now = time.monotonic() + elapsed = now - self.last_time + + self.tokens = min( + self.capacity, + self.tokens + elapsed * self.rate, + ) + self.last_time = now + + if self.tokens >= 1: + self.tokens -= 1 + return + + wait_time = (1 - self.tokens) / self.rate + + time.sleep(wait_time) \ No newline at end of file diff --git a/tests/test_connection_pool.py b/tests/test_connection_pool.py index f234a06..3d36d3b 100644 --- a/tests/test_connection_pool.py +++ b/tests/test_connection_pool.py @@ -30,6 +30,18 @@ def test_max_pool_connections_validation(): _config(max_pool_connections=-5) +def test_request_rps_defaults_to_none(): + assert _config().put_rps is None + assert _config().query_rps is None + + +def test_request_rps_validation(): + with pytest.raises(ValueError): + _config(put_rps=0) + with pytest.raises(ValueError): + _config(query_rps=-5) + + def test_botocore_config_carries_pool_size(): cfg = _config(max_pool_connections=50) client_config = cfg.botocore_config() diff --git a/tests/test_s3vectors.py b/tests/test_s3vectors.py index 3d48a7d..32e1797 100644 --- a/tests/test_s3vectors.py +++ b/tests/test_s3vectors.py @@ -43,13 +43,16 @@ def _config(**kwargs): def _store_with_pages(pages, config=None): - store = S3VectorsStore.__new__(S3VectorsStore) - store._config = config or _config() - store._client = MagicMock() - paginator = MagicMock() - paginator.paginate.return_value = pages - store._client.get_paginator.return_value = paginator - return store, paginator + store = S3VectorsStore.__new__(S3VectorsStore) + store._config = config or _config() + store._client = MagicMock() + store._put_limiter = None + store._query_limiter = None + + paginator = MagicMock() + paginator.paginate.return_value = pages + store._client.get_paginator.return_value = paginator + return store, paginator def test_query_single_page(): @@ -65,37 +68,35 @@ def test_query_single_page(): assert paginator.paginate.call_args.kwargs["PaginationConfig"] == {"MaxItems": 50} assert paginator.paginate.call_args.kwargs["topK"] == 50 - def test_query_multi_page_pagination(): - pages = [ - {"vectors": [{"key": f"k{i}"} for i in range(100)]}, - {"vectors": [{"key": f"k{i}"} for i in range(100, 200)]}, - {"vectors": [{"key": f"k{i}"} for i in range(200, 250)]}, - ] - store, paginator = _store_with_pages(pages) - - results = store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=250) - - assert len(results) == 250 - assert results[0]["key"] == "k0" - assert results[99]["key"] == "k99" - assert results[100]["key"] == "k100" - assert results[249]["key"] == "k249" - assert paginator.paginate.call_args.kwargs["PaginationConfig"] == {"MaxItems": 250} - - -def test_query_truncation_at_top_k(): - pages = [ - {"vectors": [{"key": f"k{i}"} for i in range(100)]}, - {"vectors": [{"key": f"k{i}"} for i in range(100, 200)]}, - ] - store, _ = _store_with_pages(pages) + pages = [ + { + "vectors": [{"key": f"k{i}"} for i in range(100)], + "NextToken": "token-1", + }, + { + "vectors": [{"key": f"k{i}"} for i in range(100, 200)], + "NextToken": "token-2", + }, + { + "vectors": [{"key": f"k{i}"} for i in range(200, 250)], + }, + ] - results = store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=130) + store, _ = _store_with_pages(pages) - assert len(results) == 130 - assert results[-1]["key"] == "k129" + results = store.query( + query_vector=[0.1, 0.2, 0.3, 0.4], + top_k=250, + ) + assert len(results) == 250 + assert [result["key"] for result in results[:3]] == ["k0", "k1", "k2"] + assert [result["key"] for result in results[-3:]] == [ + "k247", + "k248", + "k249", + ] def test_query_fewer_results_than_top_k(): store, _ = _store_with_pages([{"vectors": [{"key": "k0"}, {"key": "k1"}]}]) @@ -120,41 +121,6 @@ def test_query_exceeding_service_limit(): store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=_MAX_TOP_K + 1) -def test_query_pages_raw_chunks(): - pages = [ - {"vectors": [{"key": f"k{i}"} for i in range(100)]}, - {"vectors": [{"key": f"k{i}"} for i in range(100, 150)]}, - ] - store, _ = _store_with_pages(pages) - - result_pages = list(store.query_pages(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=150)) - - assert [len(p) for p in result_pages] == [100, 50] - - -def test_query_pages_configurable_page_size(): - store, _ = _store_with_pages( - [{"vectors": [{"key": f"k{i}"} for i in range(25)]}] - ) - - result_pages = list( - store.query_pages(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=25, page_size=10) - ) - - assert [len(p) for p in result_pages] == [10, 10, 5] - - -def test_query_pages_uses_config_top_k_page_size(): - store, _ = _store_with_pages( - [{"vectors": [{"key": f"k{i}"} for i in range(25)]}], - config=_config(top_k_page_size=10), - ) - - result_pages = list(store.query_pages(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=25)) - - assert [len(p) for p in result_pages] == [10, 10, 5] - - def test_query_pages_invalid_page_size(): store, _ = _store_with_pages([]) @@ -166,23 +132,16 @@ def test_config_rejects_non_positive_top_k_page_size(): with pytest.raises(ValueError, match="top_k_page_size must be a positive integer"): _config(top_k_page_size=0) - - - - - - def test_put_vectors_parallelization(): store, _ = _store_with_pages([]) - - + vectors = [ (f"key_{i}", [0.1] * 128, {"tag": "test"}) for i in range(1500) ] def mock_put_batch(payload): - time.sleep(0.1) + time.sleep(0.1) store._put_batch = MagicMock(side_effect=mock_put_batch) @@ -190,14 +149,13 @@ def mock_put_batch(payload): store.put_vectors(vectors) duration = time.perf_counter() - t0 - assert duration < 0.5 assert store._put_batch.call_count == 3 def test_put_vectors_error_propagation(): store, _ = _store_with_pages([]) - + vectors = [ (f"key_{i}", [0.1] * 128, {"tag": "test"}) for i in range(1000) @@ -213,4 +171,157 @@ def mock_put_batch_with_error(payload): store.put_vectors(vectors) +def test_put_batch_uses_rate_limiter(): + store, _ = _store_with_pages( + [], + config=_config(put_rps=10), + ) + + limiter = MagicMock() + store._put_limiter = limiter + + payload = [ + { + "key": "key-1", + "data": {"float32": [0.1, 0.2, 0.3, 0.4]}, + "metadata": {}, + } + ] + + store._client.put_vectors = MagicMock() + + store._put_batch(payload) + + limiter.acquire.assert_called_once_with() + store._client.put_vectors.assert_called_once() + + +def test_query_truncation_at_top_k(): + pages = [ + { + "vectors": [{"key": f"k{i}"} for i in range(100)], + "NextToken": "token-1", + }, + { + "vectors": [{"key": f"k{i}"} for i in range(100, 200)], + }, + ] + + store, _ = _store_with_pages(pages) + + results = store.query( + query_vector=[0.1, 0.2, 0.3, 0.4], + top_k=130, + ) + + assert len(results) == 130 + assert results[-1]["key"] == "k129" + + +def test_query_pages_raw_chunks(): + pages = [ + { + "vectors": [{"key": f"k{i}"} for i in range(100)], + "NextToken": "token-1", + }, + { + "vectors": [{"key": f"k{i}"} for i in range(100, 150)], + }, + ] + + store, _ = _store_with_pages(pages) + + result_pages = list( + store.query_pages( + query_vector=[0.1, 0.2, 0.3, 0.4], + top_k=150, + ) + ) + + assert [len(p) for p in result_pages] == [100, 50] + + +def test_query_pages_configurable_page_size(): + store, _ = _store_with_pages( + [{"vectors": [{"key": f"k{i}"} for i in range(25)]}] + ) + + result_pages = list( + store.query_pages( + query_vector=[0.1, 0.2, 0.3, 0.4], + top_k=25, + page_size=10, + ) + ) + + assert [len(p) for p in result_pages] == [10, 10, 5] + + +def test_query_pages_uses_config_top_k_page_size(): + store, _ = _store_with_pages( + [{"vectors": [{"key": f"k{i}"} for i in range(25)]}], + config=_config(top_k_page_size=10), + ) + + result_pages = list( + store.query_pages( + query_vector=[0.1, 0.2, 0.3, 0.4], + top_k=25, + ) + ) + + assert [len(p) for p in result_pages] == [10, 10, 5] + + +def test_query_pages_uses_rate_limiter_per_page(): + pages = [ + { + "vectors": [{"key": "k0"}], + "NextToken": "token-1", + }, + { + "vectors": [{"key": "k1"}], + "NextToken": "token-2", + }, + { + "vectors": [{"key": "k2"}], + }, + ] + + store, _ = _store_with_pages( + pages, + config=_config(query_rps=10), + ) + + limiter = MagicMock() + store._query_limiter = limiter + + result_pages = list( + store.query_pages( + query_vector=[0.1, 0.2, 0.3, 0.4], + top_k=3, + ) + ) + + assert [len(page) for page in result_pages] == [1, 1, 1] + assert limiter.acquire.call_count == 3 + + +def test_query_pages_does_not_rate_limit_after_paginator_exhaustion(): + store, _ = _store_with_pages( + [{"vectors": [{"key": "k0"}]}], + config=_config(query_rps=10), + ) + + limiter = MagicMock() + store._query_limiter = limiter + + result_pages = list( + store.query_pages( + query_vector=[0.1, 0.2, 0.3, 0.4], + top_k=50, + ) + ) + assert [len(page) for page in result_pages] == [1] + assert limiter.acquire.call_count == 1 \ No newline at end of file diff --git a/tests/test_utils_transforms.py b/tests/test_utils_transforms.py index 9a8ae7a..58b97ec 100644 --- a/tests/test_utils_transforms.py +++ b/tests/test_utils_transforms.py @@ -1,9 +1,11 @@ """Tests for decorators/generators and the transform pipeline (pure, no AWS).""" +from unittest.mock import patch + import pytest from dynavec.transforms import TransformContext, TransformPipeline, as_pipeline -from dynavec.utils import chunked, is_retryable, retry +from dynavec.utils import TokenBucket, chunked, is_retryable, retry def test_chunked_generator(): @@ -96,3 +98,67 @@ def flaky(): assert flaky() == "ok" assert calls["n"] == 2 assert delays == [3.0] + + +def test_token_bucket_consumes_tokens(): + with patch( + "dynavec.utils.time.monotonic", + side_effect=[100.0, 100.0, 100.0], + ): + bucket = TokenBucket(rate=10, capacity=2) + assert bucket.tokens == 2 + + bucket.acquire() + assert bucket.tokens == 1 + + bucket.acquire() + assert bucket.tokens == 0 + + +def test_token_bucket_refills_tokens(): + with patch( + "dynavec.utils.time.monotonic", + side_effect=[100.0, 100.0, 100.0, 100.2], + ): + bucket = TokenBucket(rate=10, capacity=2) + + bucket.acquire() + bucket.acquire() + + # 0.2 seconds * 10 tokens/sec = 2 new tokens + bucket.acquire() + + assert bucket.tokens == 1 + + +def test_token_bucket_waits_for_token(): + with ( + patch( + "dynavec.utils.time.monotonic", + side_effect=[100.0, 100.0, 100.0, 100.05, 100.11], + ), + patch("dynavec.utils.time.sleep") as sleep, + ): + bucket = TokenBucket(rate=10, capacity=2) + + bucket.acquire() + bucket.acquire() + bucket.acquire() + + sleep.assert_called_once_with(pytest.approx(0.05)) + assert bucket.tokens == pytest.approx(0.1) + + +def test_token_bucket_does_not_exceed_capacity(): + with patch( + "dynavec.utils.time.monotonic", + side_effect=[100.0, 100.0, 101.0], + ): + bucket = TokenBucket(rate=10, capacity=2) + assert bucket.tokens == 2 + + bucket.acquire() + assert bucket.tokens == 1 + + bucket.acquire() + assert bucket.tokens == 1 \ No newline at end of file