diff --git a/CHANGELOG.md b/CHANGELOG.md index bbfe8308..c63d1446 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,8 +5,8 @@ - `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata. - `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` / `mtrag_clapnq_live_uncompacted` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, live arms with and without `EvidenceCompactionCapability`, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, per-turn tool-traffic attributes, and `citation_status` / `turn_citation_status` eval attributes. - - BTree indexes on `chunks.id`, `chunks.document_id` and `documents.id`, and a Bitmap index on `document_items.label`. Existing databases need `haiku-rag migrate`. +- `lancedb.read_consistency_interval_seconds` (default 30), `lancedb.index_cache_size_bytes` and `lancedb.metadata_cache_size_bytes`. The LanceDB session is shared across connections in a process, so its index and metadata caches survive a connection being closed. ### Changed diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index 4d86b215..57ab3aa4 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -122,6 +122,18 @@ The `storage_options` keys are case-insensitive and passed directly to the under **Note:** Table optimization is automatically handled by LanceDB Cloud (`db://` URIs) and is disabled for better performance. For object storage backends (S3, Azure, GCS), optimization and vector indexing are still performed normally. +### Caching and Read Consistency + +```yaml +lancedb: + read_consistency_interval_seconds: 30 # null to never re-check + index_cache_size_bytes: 536870912 # null for the LanceDB default + metadata_cache_size_bytes: 268435456 +``` + +- **read_consistency_interval_seconds**: how often a connection checks for writes from another process. `null` never checks, so a long-lived reader never sees the ingester's writes. `0` checks on every read. +- **index_cache_size_bytes** / **metadata_cache_size_bytes**: sizes for the caches held by the LanceDB session, which is shared across every connection in the process. The first vector query loads the index into it, so on object storage the cache is what stops the next connection refetching it. Size it for the total set of indexes a process keeps warm, against the memory available to it. + ### Deployment Pattern: One Writer, Many Readers LanceDB on S3 supports **exactly one writer + N readers** per database URI. Multiple writers against the same URI can race on the manifest commit and corrupt state. This is a LanceDB property, not something `haiku.rag` enforces. diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 0bfbfec7..2fa58a8e 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -62,10 +62,21 @@ class StorageConfig(BaseModel): class LanceDBConfig(BaseModel): + """LanceDB connection settings. + + read_consistency_interval_seconds bounds how stale a reader may be. None + never re-checks, so a long-lived reader never sees another process's writes. + The cache sizes are per process, since the session is shared across + connections. + """ + uri: str = "" api_key: str = "" region: str = "" storage_options: dict[str, str] = Field(default_factory=dict) + read_consistency_interval_seconds: float | None = Field(default=30, ge=0) + index_cache_size_bytes: int | None = Field(default=None, ge=0) + metadata_cache_size_bytes: int | None = Field(default=None, ge=0) class EmbeddingsConfig(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 002283a1..f47a5d82 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -55,25 +55,56 @@ class ConnectionMode(Enum): return ConnectionMode.OBJECT_STORAGE +_sessions: dict[tuple[int | None, int | None], lancedb.Session] = {} + + +def _session(config: AppConfig) -> lancedb.Session: + """The process's session for these cache sizes. + + Sessions hold the index and metadata caches. Sharing one across connections + is what keeps a cached index from being refetched per connection, which on + object storage is the dominant cost of the first query. + """ + key = ( + config.lancedb.index_cache_size_bytes, + config.lancedb.metadata_cache_size_bytes, + ) + if key not in _sessions: + kwargs = {} + if key[0] is not None: + kwargs["index_cache_size_bytes"] = key[0] + if key[1] is not None: + kwargs["metadata_cache_size_bytes"] = key[1] + _sessions[key] = lancedb.Session(**kwargs) + return _sessions[key] + + async def connect_lancedb( config: AppConfig, db_path: Path | None = None ) -> lancedb.AsyncConnection: + interval = config.lancedb.read_consistency_interval_seconds + kwargs: dict[str, Any] = { + "session": _session(config), + "read_consistency_interval": ( + timedelta(seconds=interval) if interval is not None else None + ), + } mode = ConnectionMode.from_config(config) if mode == ConnectionMode.CLOUD: return await lancedb.connect_async( uri=config.lancedb.uri, api_key=config.lancedb.api_key, region=config.lancedb.region, + **kwargs, ) elif mode == ConnectionMode.OBJECT_STORAGE: - kwargs: dict[str, Any] = {"uri": config.lancedb.uri} if config.lancedb.storage_options: kwargs["storage_options"] = config.lancedb.storage_options - return await lancedb.connect_async(**kwargs) + return await lancedb.connect_async(uri=config.lancedb.uri, **kwargs) else: if db_path is None: raise ValueError("No lancedb.uri configured and no db_path provided") - return await lancedb.connect_async(db_path.absolute()) + return await lancedb.connect_async(db_path.absolute(), **kwargs) class DocumentRecord(LanceModel): diff --git a/tests/test_lancedb_connection.py b/tests/test_lancedb_connection.py index 8e624250..33fa349c 100644 --- a/tests/test_lancedb_connection.py +++ b/tests/test_lancedb_connection.py @@ -1,6 +1,8 @@ +from datetime import timedelta from unittest.mock import AsyncMock, patch import pytest +from pydantic import ValidationError from haiku.rag.config import Config from haiku.rag.config.models import AppConfig, LanceDBConfig @@ -49,7 +51,8 @@ class TestConnectLancedb: "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: await connect_lancedb(config, db_path=temp_db_path) - mock_connect.assert_called_once_with(temp_db_path.absolute()) + mock_connect.assert_awaited_once() + assert mock_connect.call_args.args == (temp_db_path.absolute(),) @pytest.mark.asyncio async def test_local_resolves_relative_db_path(self, tmp_path, monkeypatch): @@ -62,7 +65,8 @@ class TestConnectLancedb: "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: await connect_lancedb(config, db_path=relative) - mock_connect.assert_called_once_with(relative.absolute()) + mock_connect.assert_awaited_once() + assert mock_connect.call_args.args == (relative.absolute(),) @pytest.mark.asyncio async def test_cloud_passes_uri_api_key_region(self): @@ -75,9 +79,11 @@ class TestConnectLancedb: "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: await connect_lancedb(config) - mock_connect.assert_called_once_with( - uri="db://my-database", api_key="test-key", region="us-west-2" - ) + mock_connect.assert_awaited_once() + kwargs = mock_connect.call_args.kwargs + assert kwargs["uri"] == "db://my-database" + assert kwargs["api_key"] == "test-key" + assert kwargs["region"] == "us-west-2" @pytest.mark.asyncio async def test_object_storage_passes_uri_and_storage_options(self): @@ -94,13 +100,13 @@ class TestConnectLancedb: "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: await connect_lancedb(config) - mock_connect.assert_called_once_with( - uri="s3://bucket/path", - storage_options={ - "endpoint": "http://minio:9000", - "region": "us-east-1", - }, - ) + mock_connect.assert_awaited_once() + kwargs = mock_connect.call_args.kwargs + assert kwargs["uri"] == "s3://bucket/path" + assert kwargs["storage_options"] == { + "endpoint": "http://minio:9000", + "region": "us-east-1", + } @pytest.mark.asyncio async def test_object_storage_without_storage_options(self): @@ -109,7 +115,10 @@ class TestConnectLancedb: "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock ) as mock_connect: await connect_lancedb(config) - mock_connect.assert_called_once_with(uri="s3://bucket/path") + mock_connect.assert_awaited_once() + kwargs = mock_connect.call_args.kwargs + assert kwargs["uri"] == "s3://bucket/path" + assert "storage_options" not in kwargs @pytest.mark.asyncio async def test_local_without_db_path_raises(self): @@ -398,3 +407,119 @@ class TestStoreMiscellany: async with Store(temp_db_path, create=True) as store: with pytest.raises(ValueError, match="Unknown table"): await store.list_table_versions("not_a_table") + + +class TestSessionAndConsistency: + @pytest.mark.asyncio + async def test_session_is_shared_across_connections(self): + config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path")) + with patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ) as mock_connect: + await connect_lancedb(config) + await connect_lancedb(config) + + sessions = [c.kwargs["session"] for c in mock_connect.call_args_list] + assert sessions[0] is sessions[1] + + @pytest.mark.asyncio + async def test_cache_sizes_select_distinct_sessions(self): + small = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", index_cache_size_bytes=1 << 20 + ) + ) + large = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", index_cache_size_bytes=1 << 30 + ) + ) + with patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ) as mock_connect: + await connect_lancedb(small) + await connect_lancedb(large) + + sessions = [c.kwargs["session"] for c in mock_connect.call_args_list] + assert sessions[0] is not sessions[1] + + @pytest.mark.asyncio + async def test_both_cache_sizes_are_applied(self): + config = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", + index_cache_size_bytes=2 << 20, + metadata_cache_size_bytes=4 << 20, + ) + ) + with ( + patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ), + patch("haiku.rag.store.engine.lancedb.Session") as mock_session, + ): + await connect_lancedb(config) + + mock_session.assert_called_once_with( + index_cache_size_bytes=2 << 20, metadata_cache_size_bytes=4 << 20 + ) + + @pytest.mark.asyncio + async def test_read_consistency_interval_is_forwarded(self): + config = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", read_consistency_interval_seconds=5 + ) + ) + with patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ) as mock_connect: + await connect_lancedb(config) + + assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta( + seconds=5 + ) + + @pytest.mark.asyncio + async def test_read_consistency_interval_omitted_when_disabled(self): + config = AppConfig( + lancedb=LanceDBConfig( + uri="s3://bucket/path", read_consistency_interval_seconds=None + ) + ) + with patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ) as mock_connect: + await connect_lancedb(config) + + assert mock_connect.call_args.kwargs["read_consistency_interval"] is None + + @pytest.mark.asyncio + async def test_local_connection_also_gets_session_and_consistency(self, tmp_path): + config = AppConfig() + with patch( + "haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock + ) as mock_connect: + await connect_lancedb(config, tmp_path / "db.lancedb") + + assert mock_connect.call_args.kwargs["session"] is not None + assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta( + seconds=30 + ) + + +class TestLanceDBConfigValidation: + def test_negative_values_are_rejected(self): + """Negatives overflow or panic inside Lance, so reject them here.""" + with pytest.raises(ValidationError): + LanceDBConfig(read_consistency_interval_seconds=-1) + with pytest.raises(ValidationError): + LanceDBConfig(index_cache_size_bytes=-1) + with pytest.raises(ValidationError): + LanceDBConfig(metadata_cache_size_bytes=-1) + + def test_zero_is_allowed(self): + config = LanceDBConfig( + read_consistency_interval_seconds=0, index_cache_size_bytes=0 + ) + assert config.read_consistency_interval_seconds == 0