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:
@@ -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()
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user