Merge pull request #932 from cbcoutinho/fix/embedded-qdrant-thread-offload-926

fix(vector): offload embedded Qdrant ops to a worker thread (#926)
This commit is contained in:
Chris Coutinho
2026-06-19 10:44:25 +02:00
committed by GitHub
2 changed files with 265 additions and 3 deletions
+108 -1
View File
@@ -1,9 +1,14 @@
"""Qdrant client wrapper.""" """Qdrant client wrapper."""
import functools
import inspect
import logging import logging
from typing import Any from collections.abc import Coroutine
from typing import Any, cast
import anyio import anyio
import anyio.to_thread
from anyio import CapacityLimiter
from qdrant_client import AsyncQdrantClient, models from qdrant_client import AsyncQdrantClient, models
from qdrant_client.http.exceptions import UnexpectedResponse from qdrant_client.http.exceptions import UnexpectedResponse
from qdrant_client.models import ( from qdrant_client.models import (
@@ -111,6 +116,90 @@ _qdrant_client: AsyncQdrantClient | None = None
_qdrant_init_lock: anyio.Lock | None = None _qdrant_init_lock: anyio.Lock | None = None
def _drive_local_coroutine(coro: Coroutine[Any, Any, Any]) -> Any:
"""Run an embedded-Qdrant coroutine to completion on the current thread.
qdrant-client's local/in-memory backend (``AsyncQdrantLocal``) exposes an
``async def`` surface over purely synchronous, CPU-bound work: its
coroutines never await real I/O, so they complete on the first ``send``
(verified against qdrant-client 1.17). Driving one here — inside the worker
thread dispatched by :func:`anyio.to_thread.run_sync` — is what moves that
CPU work off the event loop.
A genuine suspension would mean the backend started doing real async I/O,
which a plain thread cannot drive correctly. Any real awaitable yields a
non-``None`` object (an asyncio ``Future`` / a trio checkpoint), so the
``yielded is not None`` guard catches it and fails loudly rather than
spinning or silently mis-driving the coroutine. Only a literal bare
``yield`` / ``yield None`` is treated as a non-suspension and re-driven —
consistent with the local backend being purely synchronous. If a qdrant
adapter ever starts inserting real checkpoints, prefer wrapping it in
network mode over relaxing this guard.
"""
try:
while True:
yielded = coro.send(None)
if yielded is not None:
coro.close()
raise RuntimeError(
"embedded Qdrant coroutine suspended unexpectedly "
f"(yielded {yielded!r}); the thread-offload shim in "
"vector/qdrant_client.py assumes the local backend is "
"synchronous"
)
except StopIteration as exc:
return exc.value
class _ThreadOffloadingQdrantClient:
"""Run embedded-Qdrant operations on a worker thread to free the event loop.
qdrant-client's embedded backends — ``:memory:`` and ``path=`` local mode —
execute every operation synchronously on the calling thread despite the
``AsyncQdrantClient`` surface (``AsyncQdrantLocal`` contains no
``to_thread``/executor offload). A single background scan of thousands of
tagged files therefore pins the event loop on a CPU-constrained host, so
``/health/live`` / ``/health/ready`` probes and in-flight MCP requests stall
until the scan finishes — the failure mode reported in issue #926.
This proxy forwards attribute access to the wrapped client and offloads
every coroutine-returning call to a worker thread, keeping the event loop
responsive. A private ``CapacityLimiter(1)`` serialises those offloads:
``AsyncQdrantLocal`` is not safe for concurrent access, and today the
single-threaded event loop is what implicitly serialises it — the limiter
preserves that invariant once the work runs off-loop, while still letting
the loop interleave health checks and other requests between operations.
Network mode (``QDRANT_URL``) is never wrapped: there the client performs
real non-blocking I/O and already yields to the event loop.
"""
def __init__(self, inner: AsyncQdrantClient) -> None:
self._inner = inner
self._limiter = CapacityLimiter(1)
def __getattr__(self, name: str) -> Any:
# Only reached for names absent from the instance __dict__; ``_inner``
# and ``_limiter`` resolve normally and never recurse through here.
attr = getattr(self._inner, name)
if not callable(attr):
return attr
@functools.wraps(attr)
def _maybe_offload(*args: Any, **kwargs: Any) -> Any:
result = attr(*args, **kwargs)
if inspect.iscoroutine(result):
return self._offload(result)
return result
return _maybe_offload
async def _offload(self, coro: Coroutine[Any, Any, Any]) -> Any:
return await anyio.to_thread.run_sync(
_drive_local_coroutine, coro, limiter=self._limiter
)
async def _create_one_payload_index( async def _create_one_payload_index(
client: AsyncQdrantClient, client: AsyncQdrantClient,
collection_name: str, collection_name: str,
@@ -596,6 +685,24 @@ async def get_qdrant_client() -> AsyncQdrantClient:
logger.warning("No Qdrant mode configured, defaulting to :memory:") logger.warning("No Qdrant mode configured, defaulting to :memory:")
provisional = AsyncQdrantClient(":memory:") provisional = AsyncQdrantClient(":memory:")
# Embedded Qdrant (``:memory:`` / ``path=``) runs every operation
# synchronously on the calling thread, so without this each call —
# including the O(N) startup backfill scroll below and every later
# scan/search query — would block the event loop and stall health
# checks on a CPU-constrained host (issue #926). Offload local-mode
# work to a worker thread. Network mode (mirrors the ``if
# settings.qdrant_url`` branch above) already does non-blocking I/O
# and is left untouched.
#
# Wrap here, before the ``_backfill_doc_id_to_string`` /
# ``_ensure_payload_indexes`` migrations below, so that the O(N)
# startup scroll runs off the event loop too — not just steady-state
# scan/search queries.
if not settings.qdrant_url:
provisional = cast(
AsyncQdrantClient, _ThreadOffloadingQdrantClient(provisional)
)
# Get collection name (auto-generated from deployment ID + model) # Get collection name (auto-generated from deployment ID + model)
collection_name = settings.get_collection_name() collection_name = settings.get_collection_name()
+157 -2
View File
@@ -18,6 +18,7 @@ from types import SimpleNamespace
from unittest.mock import call from unittest.mock import call
import anyio import anyio
import anyio.lowlevel
import httpx import httpx
import pytest import pytest
from qdrant_client.http.exceptions import UnexpectedResponse from qdrant_client.http.exceptions import UnexpectedResponse
@@ -28,8 +29,10 @@ from nextcloud_mcp_server.vector.qdrant_client import (
_DOC_ID_BACKFILL_SENTINEL_ID, _DOC_ID_BACKFILL_SENTINEL_ID,
_PAYLOAD_INDEX_FIELDS, _PAYLOAD_INDEX_FIELDS,
_backfill_doc_id_to_string, _backfill_doc_id_to_string,
_drive_local_coroutine,
_ensure_payload_indexes, _ensure_payload_indexes,
_group_int_doc_ids, _group_int_doc_ids,
_ThreadOffloadingQdrantClient,
get_qdrant_client, get_qdrant_client,
) )
@@ -976,7 +979,11 @@ async def test_get_qdrant_client_creates_collection_on_local_mode_value_error(
client = await get_qdrant_client() client = await get_qdrant_client()
assert client is provisional # Local/in-memory mode is wrapped in the thread-offload proxy (issue #926),
# so the returned client is the proxy over ``provisional`` rather than the
# bare mock; the proxy forwards every awaited call straight through to it.
assert isinstance(client, _ThreadOffloadingQdrantClient)
assert client._inner is provisional
provisional.create_collection.assert_awaited_once() provisional.create_collection.assert_awaited_once()
# The created collection should be the auto-generated name from # The created collection should be the auto-generated name from
# settings — guards against accidental collection-name drift. # settings — guards against accidental collection-name drift.
@@ -1009,4 +1016,152 @@ async def test_get_qdrant_client_propagates_unrelated_value_error(
with pytest.raises(ValueError, match="Bad collection_name"): with pytest.raises(ValueError, match="Bad collection_name"):
await get_qdrant_client() await get_qdrant_client()
provisional.create_collection.assert_not_awaited()
# ---------------------------------------------------------------------------
# Thread-offload proxy for embedded Qdrant (issue #926)
# ---------------------------------------------------------------------------
@pytest.mark.unit
def test_drive_local_coroutine_returns_value():
"""A non-suspending coroutine is driven to completion and its value returned."""
async def _coro():
return 42
assert _drive_local_coroutine(_coro()) == 42
@pytest.mark.unit
def test_drive_local_coroutine_propagates_exception():
"""An exception raised inside the coroutine surfaces to the driver."""
async def _coro():
raise ValueError("boom")
with pytest.raises(ValueError, match="boom"):
_drive_local_coroutine(_coro())
@pytest.mark.unit
def test_drive_local_coroutine_rejects_real_suspension():
"""A coroutine that genuinely suspends (yields non-None) fails loudly.
The embedded backend is assumed synchronous; a real ``await`` of I/O would
mean this shim can no longer drive it from a plain thread.
"""
class _Suspends:
# Mimics a real future/coroutine that yields control to the event loop.
def __await__(self):
yield "pending-future"
async def _coro():
await _Suspends()
return 1
with pytest.raises(RuntimeError, match="suspended unexpectedly"):
_drive_local_coroutine(_coro())
@pytest.mark.unit
async def test_offloading_proxy_runs_off_event_loop_thread(mocker):
"""The proxy executes the wrapped coroutine on a worker thread, not the loop.
This is the whole point of issue #926: embedded-Qdrant CPU work must not run
on the event loop thread, where it would stall health checks.
"""
import threading
loop_thread = threading.get_ident()
worker_threads: list[int] = []
inner = mocker.AsyncMock()
async def _record_thread(*args, **kwargs):
worker_threads.append(threading.get_ident())
return "ok"
inner.scroll.side_effect = _record_thread
proxy = _ThreadOffloadingQdrantClient(inner)
result = await proxy.scroll("collection")
assert result == "ok"
assert worker_threads, "wrapped coroutine never ran"
assert worker_threads[0] != loop_thread
inner.scroll.assert_awaited_once_with("collection")
@pytest.mark.unit
async def test_offloading_proxy_passes_through_non_callables(mocker):
"""Non-callable attributes resolve straight to the wrapped client."""
inner = mocker.AsyncMock()
inner.collection_name = "my-collection"
proxy = _ThreadOffloadingQdrantClient(inner)
assert proxy.collection_name == "my-collection"
@pytest.mark.unit
async def test_offloading_proxy_forwards_sync_callables_directly(mocker):
"""A callable returning a non-coroutine is forwarded without thread offload.
``_maybe_offload`` only routes through the worker thread when the call
produces a coroutine; plain sync methods return their value straight
through.
"""
inner = mocker.Mock() # Mock (not AsyncMock): calls return plain values.
inner.get_fastembed_vector_params.return_value = {"size": 384}
proxy = _ThreadOffloadingQdrantClient(inner)
result = proxy.get_fastembed_vector_params()
assert result == {"size": 384}
inner.get_fastembed_vector_params.assert_called_once_with()
@pytest.mark.unit
async def test_get_qdrant_client_does_not_wrap_network_mode(
mocker, monkeypatch, reset_qdrant_singleton
):
"""Network mode (``QDRANT_URL``) returns the bare client, never the proxy.
Remote Qdrant already does non-blocking I/O; wrapping it would needlessly
funnel every call through the serialising worker-thread limiter.
"""
from nextcloud_mcp_server.config import Settings
settings = Settings(
qdrant_url="https://qdrant:6333",
ollama_embedding_model="nomic-embed-text",
vector_sync_enabled=False,
)
monkeypatch.setattr(
"nextcloud_mcp_server.vector.qdrant_client.get_settings", lambda: settings
)
embedding_service = mocker.Mock()
embedding_service.provider = mocker.Mock(spec_set=[])
embedding_service.get_dimension = lambda: 4
monkeypatch.setattr(
"nextcloud_mcp_server.embedding.get_embedding_service",
lambda: embedding_service,
)
# Drive the create path via the network-mode "not found" signal — an
# UnexpectedResponse(404), as the real HTTP client raises (local mode's
# ValueError("not found") is the other branch). This skips the
# dimension-validation branch that would need a fully-shaped CollectionInfo.
provisional = _stub_provisional(
mocker, _make_unexpected(404, b'{"status":{"error":"Not found"}}')
)
monkeypatch.setattr(
"nextcloud_mcp_server.vector.qdrant_client.AsyncQdrantClient",
lambda *a, **kw: provisional,
)
client = await get_qdrant_client()
assert client is provisional
assert not isinstance(client, _ThreadOffloadingQdrantClient)