From 2cab52a2c80a971539dac4209730c03854f87dc1 Mon Sep 17 00:00:00 2001 From: Vishal Bala Date: Fri, 30 Oct 2026 16:20:59 +0100 Subject: [PATCH 1/4] fix(mcp): share one startup deadline across connect and INFO _connect_redis_client wrapped the connection and the first INFO in separate asyncio.wait_for calls, so a slow server could hold startup for up to twice runtime.startup_timeout_seconds before failing. Both now run under one asyncio.timeout block. The raised exception is still TimeoutError, which map_exception already handles. A client that connected but then failed or timed out on INFO was never closed; it is now closed before the error propagates. --- docs/user_guide/how_to_guides/mcp.md | 1 + redisvl/mcp/server.py | 21 +++++++---- tests/unit/test_mcp/test_server.py | 52 ++++++++++++++++++++++++++++ 3 files changed, 67 insertions(+), 7 deletions(-) diff --git a/docs/user_guide/how_to_guides/mcp.md b/docs/user_guide/how_to_guides/mcp.md index 6afd2282d..529824ed3 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 each startup step for the binding: connecting to Redis and the first `INFO` share one deadline, then reading the index and building the vectorizer each get their own - `schema_overrides` is only for patching incomplete field attrs discovered from Redis ### Fulltext-Only Config 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/tests/unit/test_mcp/test_server.py b/tests/unit/test_mcp/test_server.py index 394ffe656..b18960351 100644 --- a/tests/unit/test_mcp/test_server.py +++ b/tests/unit/test_mcp/test_server.py @@ -478,3 +478,55 @@ 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() + + loop = asyncio.get_running_loop() + started = loop.time() + with pytest.raises(TimeoutError): + await server._connect_redis_client(0.2) + + assert loop.time() - started < 0.3 + 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 From 3ccd2126dee6089b7440c3dbb80f897ffc57a9f7 Mon Sep 17 00:00:00 2001 From: Vishal Bala Date: Fri, 30 Oct 2026 16:24:41 +0100 Subject: [PATCH 2/4] fix(migration): cancel sibling quantize workers when one fails async_multi_worker_quantize ran its workers under asyncio.gather, which raised the first worker's exception while the others kept reading from and writing to Redis after the caller had been told the run failed. It now runs them in an asyncio.TaskGroup, so the siblings are cancelled and have closed their connections before the error propagates. Callers keep the old contract. The TaskGroup's ExceptionGroup is unwrapped and the first worker exception is raised as is; any others that failed in the same instant are logged. Without that, the async executor's report would read 'unhandled errors in a TaskGroup' instead of the worker's own message. A cancelled worker stops at an awaited pipeline call, and every backup write between those awaits is synchronous, so its shard records only whole batches. Resume re-reads an unrecorded dump batch and re-converts an unrecorded quantize batch from the stored originals, which is idempotent. Tests cover both cancellation points and the resume. --- .../how_to_guides/migrate-indexes.md | 4 +- redisvl/migration/async_executor.py | 5 +- redisvl/migration/quantize.py | 58 ++++-- tests/unit/test_multi_worker_quantize.py | 167 ++++++++++++++++++ 4 files changed, 213 insertions(+), 21 deletions(-) diff --git a/docs/user_guide/how_to_guides/migrate-indexes.md b/docs/user_guide/how_to_guides/migrate-indexes.md index f9cd2df52..91d750399 100644 --- a/docs/user_guide/how_to_guides/migrate-indexes.md +++ b/docs/user_guide/how_to_guides/migrate-indexes.md @@ -887,7 +887,9 @@ 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 +`AsyncMigrationExecutor`, the first worker to fail cancels the others, so no +shard is still being written when the migration reports the failure. ``` ### Rollback 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_multi_worker_quantize.py b/tests/unit/test_multi_worker_quantize.py index 8a3808b30..87674fc56 100644 --- a/tests/unit/test_multi_worker_quantize.py +++ b/tests/unit/test_multi_worker_quantize.py @@ -9,10 +9,12 @@ - Progress aggregation """ +import asyncio 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 +366,168 @@ 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 pytest.raises(RedisConnectionError, match="failed at"): + await self._run(store, every_worker_fails, tmp_path, []) + + assert caplog.text.count("Another quantization worker also failed") == 2 From a178129a3c747507735535abb4a744a048492986 Mon Sep 17 00:00:00 2001 From: Vishal Bala Date: Fri, 30 Oct 2026 16:35:08 +0100 Subject: [PATCH 3/4] docs(mcp,migration): narrow the timeout and cancellation claims The MCP guide now names the three startup steps that runtime.startup_timeout_seconds bounds, since the native-hybrid probe runs without one. The migration guide says cancellation applies with --async and that the sync executor lets sibling workers finish. The shared-deadline test drops a wall-clock bound that the TimeoutError already proves, and the simultaneous-failure test pins that the first worker's error is the one raised. --- docs/user_guide/how_to_guides/mcp.md | 2 +- docs/user_guide/how_to_guides/migrate-indexes.md | 5 +++-- tests/unit/test_mcp/test_server.py | 3 --- tests/unit/test_multi_worker_quantize.py | 2 +- 4 files changed, 5 insertions(+), 7 deletions(-) diff --git a/docs/user_guide/how_to_guides/mcp.md b/docs/user_guide/how_to_guides/mcp.md index 529824ed3..9b37fd598 100644 --- a/docs/user_guide/how_to_guides/mcp.md +++ b/docs/user_guide/how_to_guides/mcp.md @@ -195,7 +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 each startup step for the binding: connecting to Redis and the first `INFO` share one deadline, then reading the index and building the vectorizer each get their own +- `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 91d750399..1a4e2a64c 100644 --- a/docs/user_guide/how_to_guides/migrate-indexes.md +++ b/docs/user_guide/how_to_guides/migrate-indexes.md @@ -888,8 +888,9 @@ 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. With -`AsyncMigrationExecutor`, the first worker to fail cancels the others, so no -shard is still being written when the migration reports the failure. +`--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/tests/unit/test_mcp/test_server.py b/tests/unit/test_mcp/test_server.py index b18960351..e9918c8ab 100644 --- a/tests/unit/test_mcp/test_server.py +++ b/tests/unit/test_mcp/test_server.py @@ -512,12 +512,9 @@ async def test_connect_redis_client_shares_one_deadline(monkeypatch): server = RedisVLMCPServer(_dummy_settings()) server.config = _startup_config() - loop = asyncio.get_running_loop() - started = loop.time() with pytest.raises(TimeoutError): await server._connect_redis_client(0.2) - assert loop.time() - started < 0.3 assert client.aclose_calls == 1 diff --git a/tests/unit/test_multi_worker_quantize.py b/tests/unit/test_multi_worker_quantize.py index 87674fc56..d8e294289 100644 --- a/tests/unit/test_multi_worker_quantize.py +++ b/tests/unit/test_multi_worker_quantize.py @@ -527,7 +527,7 @@ async def test_simultaneous_failures_raise_the_first_not_a_group( async def every_worker_fails(ops): raise RedisConnectionError(f"failed at {ops[0][1]}") - with pytest.raises(RedisConnectionError, match="failed at"): + with 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 From 84430e9acdc24bfda23556caad42ade7b44d3565 Mon Sep 17 00:00:00 2001 From: Vishal Bala Date: Fri, 30 Oct 2026 16:36:01 +0100 Subject: [PATCH 4/4] test(migration): capture worker-failure logs under --log-level=CRITICAL make test passes --log-level=CRITICAL, which hid the ERROR records the simultaneous-failure test counts. The test now sets the level itself. --- tests/unit/test_multi_worker_quantize.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_multi_worker_quantize.py b/tests/unit/test_multi_worker_quantize.py index d8e294289..0a68ba6ee 100644 --- a/tests/unit/test_multi_worker_quantize.py +++ b/tests/unit/test_multi_worker_quantize.py @@ -10,6 +10,7 @@ """ import asyncio +import logging import struct from unittest.mock import MagicMock, patch @@ -527,7 +528,10 @@ async def test_simultaneous_failures_raise_the_first_not_a_group( async def every_worker_fails(ops): raise RedisConnectionError(f"failed at {ops[0][1]}") - with pytest.raises(RedisConnectionError, match="failed at doc:0"): + 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