diff --git a/docs/user_guide/how_to_guides/mcp.md b/docs/user_guide/how_to_guides/mcp.md index 6afd2282d..9b37fd598 100644 --- a/docs/user_guide/how_to_guides/mcp.md +++ b/docs/user_guide/how_to_guides/mcp.md @@ -195,6 +195,7 @@ indexes: - `runtime.default_embed_text_field` is only required when the server should generate embeddings during upsert - `vectorizer` is required for query embedding and server-side embedding, but optional for fulltext-only configs - `runtime.max_result_window` caps deep paging by limiting the maximum `offset + limit` +- `runtime.startup_timeout_seconds` bounds three startup steps for the binding: connecting to Redis together with the first `INFO`, which share one deadline, then reading the index, then building the vectorizer - `schema_overrides` is only for patching incomplete field attrs discovered from Redis ### Fulltext-Only Config diff --git a/docs/user_guide/how_to_guides/migrate-indexes.md b/docs/user_guide/how_to_guides/migrate-indexes.md index f9cd2df52..1a4e2a64c 100644 --- a/docs/user_guide/how_to_guides/migrate-indexes.md +++ b/docs/user_guide/how_to_guides/migrate-indexes.md @@ -887,7 +887,10 @@ backup is written *before* the index is dropped, so a crash at any point leaves a complete backup on disk. In multi-worker mode, dump and quantize are fused (each worker reads, backs up, and converts its shard in one pass *after* the index drop). A crash during this fused phase may leave partial -backup shards. Re-running detects and resumes from partial state. +backup shards. Re-running detects and resumes from partial state. With +`--async` (`AsyncMigrationExecutor`), the first worker to fail cancels the +others; the sync executor lets them finish their shards. Either way, no shard +is still being written when the migration reports the failure. ``` ### Rollback diff --git a/redisvl/mcp/server.py b/redisvl/mcp/server.py index 4568482ce..66130920c 100644 --- a/redisvl/mcp/server.py +++ b/redisvl/mcp/server.py @@ -979,13 +979,20 @@ async def _connect_redis_client(self, timeout: int) -> Any: if self.config is None: raise RuntimeError("MCP server config not loaded") - client = await asyncio.wait_for( - RedisConnectionFactory._get_aredis_connection( - redis_url=self.config.server.redis_url - ), - timeout=timeout, - ) - await asyncio.wait_for(client.info("server"), timeout=timeout) + # Connecting and the first INFO share one deadline, so a slow server + # fails startup within ``timeout`` rather than up to twice that. + client: Any = None + try: + async with asyncio.timeout(timeout): + client = await RedisConnectionFactory._get_aredis_connection( + redis_url=self.config.server.redis_url + ) + await client.info("server") + except BaseException: + # The caller only takes ownership of a client that answered INFO. + if client is not None: + await client.aclose() + raise return client async def _load_effective_schema( diff --git a/redisvl/migration/async_executor.py b/redisvl/migration/async_executor.py index 24237656a..a5fc19776 100644 --- a/redisvl/migration/async_executor.py +++ b/redisvl/migration/async_executor.py @@ -564,8 +564,9 @@ async def apply( Async counterpart of :meth:`MigrationExecutor.apply`. Uses ``await`` for Redis I/O so the event loop remains responsive during - large quantization jobs. Multi-worker quantization uses - ``asyncio.gather`` with independent connections. + large quantization jobs. Multi-worker quantization runs in an + ``asyncio.TaskGroup`` with independent connections, so one worker's + failure cancels the rest. Args: plan: The migration plan to apply (from diff --git a/redisvl/migration/quantize.py b/redisvl/migration/quantize.py index 02de574ed..a35b83cff 100644 --- a/redisvl/migration/quantize.py +++ b/redisvl/migration/quantize.py @@ -4,7 +4,7 @@ the per-key HGET loop with batched pipeline operations. Also provides multi-worker orchestration for parallel quantization -using ThreadPoolExecutor (sync) or asyncio.gather (async). +using ThreadPoolExecutor (sync) or asyncio.TaskGroup (async). """ import hashlib @@ -538,9 +538,19 @@ async def async_multi_worker_quantize( progress_callback: Optional[Callable[[str, int, int], None]] = None, worker_backup_paths: Optional[List[str]] = None, ) -> MultiWorkerResult: - """Orchestrate async multi-worker quantization via asyncio.gather. + """Orchestrate async multi-worker quantization via asyncio.TaskGroup. Each worker gets its own async Redis connection and backup file shard. + If a worker fails, its siblings are cancelled, and each closes its + connection, before the failure is raised, so no worker is still writing + once the caller sees the error. A cancelled worker's shard records only + whole batches, so a later run resumes it like any interrupted shard. + + Raises: + Exception: The failing worker's own exception, not an + ``ExceptionGroup``. If several workers fail before the + cancellation reaches them, the first is raised and the rest + are logged. """ import asyncio @@ -562,22 +572,34 @@ async def async_multi_worker_quantize( f"({len(worker_backup_paths)} != {actual_workers})" ) - coroutines = [ - _async_worker_quantize( - worker_id=i, - redis_url=redis_url, - keys=slices[i], - datatype_changes=datatype_changes, - backup_path=worker_backup_paths[i], - index_name=index_name, - batch_size=batch_size, - progress_callback=progress_callback, - ) - for i in range(actual_workers) - ] - - results = await asyncio.gather(*coroutines) - worker_results = list(results) + failure: Exception | None = None + try: + async with asyncio.TaskGroup() as group: + tasks = [ + group.create_task( + _async_worker_quantize( + worker_id=i, + redis_url=redis_url, + keys=slices[i], + datatype_changes=datatype_changes, + backup_path=worker_backup_paths[i], + index_name=index_name, + batch_size=batch_size, + progress_callback=progress_callback, + ) + ) + for i in range(actual_workers) + ] + except ExceptionGroup as group_error: + failure = group_error.exceptions[0] + for other in group_error.exceptions[1:]: + logger.error("Another quantization worker also failed", exc_info=other) + if failure is not None: + # Raised outside the except block so the worker's traceback and cause + # reach the caller unchanged, as they did under asyncio.gather. + raise failure + + worker_results = [task.result() for task in tasks] total_docs = sum(r["docs"] for r in worker_results) return MultiWorkerResult( diff --git a/tests/unit/test_mcp/test_server.py b/tests/unit/test_mcp/test_server.py index 394ffe656..e9918c8ab 100644 --- a/tests/unit/test_mcp/test_server.py +++ b/tests/unit/test_mcp/test_server.py @@ -478,3 +478,52 @@ async def test_shutdown_after_stop_is_idempotent(): await server.shutdown() assert server._lifecycle_state.name == "STOPPED" + + +class _SlowInfoClient: + def __init__(self, info_delay: float): + self.info_delay = info_delay + self.aclose_calls = 0 + + async def info(self, section): + await asyncio.sleep(self.info_delay) + return {"redis_version": "8.0.0"} + + async def aclose(self): + self.aclose_calls += 1 + + +def _patch_slow_connect(monkeypatch, client, connect_delay: float): + async def slow_connect(redis_url=None, **kwargs): + await asyncio.sleep(connect_delay) + return client + + monkeypatch.setattr( + "redisvl.mcp.server.RedisConnectionFactory._get_aredis_connection", + slow_connect, + ) + + +@pytest.mark.asyncio +async def test_connect_redis_client_shares_one_deadline(monkeypatch): + # Each step fits the timeout alone, but not both together. + client = _SlowInfoClient(info_delay=0.15) + _patch_slow_connect(monkeypatch, client, connect_delay=0.15) + server = RedisVLMCPServer(_dummy_settings()) + server.config = _startup_config() + + with pytest.raises(TimeoutError): + await server._connect_redis_client(0.2) + + assert client.aclose_calls == 1 + + +@pytest.mark.asyncio +async def test_connect_redis_client_returns_client_within_deadline(monkeypatch): + client = _SlowInfoClient(info_delay=0.05) + _patch_slow_connect(monkeypatch, client, connect_delay=0.05) + server = RedisVLMCPServer(_dummy_settings()) + server.config = _startup_config() + + assert await server._connect_redis_client(1) is client + assert client.aclose_calls == 0 diff --git a/tests/unit/test_multi_worker_quantize.py b/tests/unit/test_multi_worker_quantize.py index 8a3808b30..0a68ba6ee 100644 --- a/tests/unit/test_multi_worker_quantize.py +++ b/tests/unit/test_multi_worker_quantize.py @@ -9,10 +9,13 @@ - Progress aggregation """ +import asyncio +import logging import struct from unittest.mock import MagicMock, patch import pytest +from redis.exceptions import ConnectionError as RedisConnectionError def _make_float32_vector(dims: int = 4, seed: float = 0.0) -> bytes: @@ -364,3 +367,171 @@ async def test_async_worker_loads_existing_backup(self, tmp_path): fields=dt_changes, batch_size=2, ) + + +class _FakeAsyncRedis: + """Hash store shared by every worker; ``on_execute`` may fail or stall.""" + + def __init__(self, store, on_execute): + self.store = store + self.on_execute = on_execute + self.closed = False + + def pipeline(self, transaction=False): + return _FakePipeline(self) + + async def aclose(self): + self.closed = True + + +class _FakePipeline: + def __init__(self, client): + self.client = client + self.ops = [] + + def hget(self, key, field): + self.ops.append(("hget", key, field)) + + def hset(self, key, mapping): + self.ops.append(("hset", key, mapping)) + + async def execute(self): + results = [] + for op, key, arg in self.ops: + if op == "hget": + results.append(self.client.store.get(key, {}).get(arg)) + else: + self.client.store.setdefault(key, {}).update(arg) + results.append(len(arg)) + # Runs after the commands land, like a reply lost after Redis applied them. + await self.client.on_execute(self.ops) + return results + + +class _WorkerFailure: + """Worker 0 fails once workers 1 and 2 are stalled inside their second batch. + + Keys are ``doc:0``..``doc:11``: three workers of four keys each, in + batches of two. + """ + + def __init__(self, op): + self.op = op + self.stalled = 0 + self.siblings_stalled = asyncio.Event() + self.cancelled = [] + + async def on_execute(self, ops): + op, key, _ = ops[0] + if op != self.op: + return + doc = int(key.split(":")[1]) + worker = doc // 4 + if worker == 0: + await self.siblings_stalled.wait() + raise RedisConnectionError("worker 0 lost its connection") + if doc % 4 == 2: + self.stalled += 1 + if self.stalled == 2: + self.siblings_stalled.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancelled.append(worker) + raise + + +async def _healthy(ops): + return None + + +class TestAsyncMultiWorkerFailure: + """One worker's failure cancels its siblings, and the shards still resume.""" + + DATATYPE_CHANGES = { + "embedding": {"source": "float32", "target": "float16", "dims": 4} + } + + def _run(self, store, on_execute, tmp_path, clients): + from redisvl.migration.quantize import async_multi_worker_quantize + + def from_url(url): + client = _FakeAsyncRedis(store, on_execute) + clients.append(client) + return client + + async def run(): + with patch("redis.asyncio.from_url", from_url): + async with asyncio.timeout(5): + return await async_multi_worker_quantize( + redis_url="redis://localhost:6379", + keys=[f"doc:{i}" for i in range(12)], + datatype_changes=self.DATATYPE_CHANGES, + backup_dir=str(tmp_path), + index_name="idx", + num_workers=3, + batch_size=2, + ) + + return run() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "op, dump_batches, quantized_batches", + [("hget", 1, 0), ("hset", 2, 1)], + ids=["during-dump", "during-quantize"], + ) + async def test_failure_cancels_siblings_and_backups_resume( + self, tmp_path, op, dump_batches, quantized_batches + ): + from redisvl.migration.backup import VectorBackup + from redisvl.migration.quantize import ( + build_worker_backup_paths, + convert_vectors, + ) + + originals = { + f"doc:{i}": {"embedding": _make_float32_vector(4, seed=float(i))} + for i in range(12) + } + store = {key: dict(fields) for key, fields in originals.items()} + scenario = _WorkerFailure(op) + clients = [] + + with pytest.raises(RedisConnectionError, match="worker 0"): + await self._run(store, scenario.on_execute, tmp_path, clients) + + # The siblings were cancelled before the error surfaced, not left running. + assert sorted(scenario.cancelled) == [1, 2] + assert len(clients) == 3 + assert all(client.closed for client in clients) + + # Each cancelled shard recorded only the batches it finished. + paths = build_worker_backup_paths(str(tmp_path), "idx", 3) + for path in paths[1:]: + header = VectorBackup.load(path).header + assert header.dump_completed_batches == dump_batches + assert header.quantize_completed_batches == quantized_batches + + result = await self._run(store, _healthy, tmp_path, []) + + assert result.total_docs_quantized == 12 + assert all(VectorBackup.load(p).header.phase == "completed" for p in paths) + assert store == convert_vectors(originals, self.DATATYPE_CHANGES) + + @pytest.mark.asyncio + async def test_simultaneous_failures_raise_the_first_not_a_group( + self, tmp_path, caplog + ): + store = {f"doc:{i}": {"embedding": _make_float32_vector(4)} for i in range(12)} + + async def every_worker_fails(ops): + raise RedisConnectionError(f"failed at {ops[0][1]}") + + with ( + caplog.at_level(logging.ERROR, logger="redisvl.migration.quantize"), + pytest.raises(RedisConnectionError, match="failed at doc:0"), + ): + await self._run(store, every_worker_fails, tmp_path, []) + + assert caplog.text.count("Another quantization worker also failed") == 2