diff --git a/conf/config.yaml b/conf/config.yaml index 5567f4e27..d286b86f5 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -55,7 +55,7 @@ embedder: # --- Vector Database (Milvus) --- # Env: VDB_HOST, VDB_PORT, VDB_CONNECTOR_NAME, VDB_COLLECTION_NAME, VDB_HYBRID_SEARCH, -# VDB_ENABLE_INSERTION +# VDB_ENABLE_INSERTION, VDB_TIMEOUT vectordb: host: milvus port: 19530 @@ -63,6 +63,7 @@ vectordb: collection_name: vdb_test hybrid_search: true enable: true + timeout: 120.0 # per-request timeout (s) for the Milvus sync/async clients schema_version: 1 # Increment when the collection schema changes and a migration is required # --- Relational Database (PostgreSQL) --- @@ -257,29 +258,11 @@ loader: # --- Ray --- ray: - # Env: RAY_NUM_GPUS, RAY_POOL_SIZE, RAY_MAX_TASKS_PER_WORKER - num_gpus: 0.01 - pool_size: 1 - max_tasks_per_worker: 8 - indexer: - # Env: RAY_MAX_TASK_RETRIES, INDEXER_SERIALIZE_TIMEOUT, VECTORDB_TIMEOUT - max_task_retries: 2 - serialize_timeout: 3600 - vectordb_timeout: 30 - concurrency_groups: - # Env: INDEXER_DEFAULT_CONCURRENCY, INDEXER_UPDATE_CONCURRENCY, etc. - default: 1000 - update: 100 - search: 100 - delete: 100 - serialize: 50 - chunk: 1000 - insert: 100 - - semaphore: - # Env: RAY_SEMAPHORE_CONCURRENCY - concurrency: 100000 + # Env: RAY_POOL_SIZE, RAY_MAX_TASKS_PER_WORKER + # Indexing capacity = pool_size (worker actors) × max_tasks_per_worker (files per worker). + pool_size: 1 + max_tasks_per_worker: 50 serve: # Env: ENABLE_RAY_SERVE, RAY_SERVE_NUM_REPLICAS, RAY_SERVE_HOST, diff --git a/docs/content/docs/documentation/deploy_ray_cluster.md b/docs/content/docs/documentation/deploy_ray_cluster.md index c5ad01f7e..5738b376f 100644 --- a/docs/content/docs/documentation/deploy_ray_cluster.md +++ b/docs/content/docs/documentation/deploy_ray_cluster.md @@ -16,7 +16,6 @@ Ensure your `.env` file includes the standard app variables **plus Ray-specific // .env # Ray # Resources for all files -RAY_NUM_GPUS=0.1 RAY_POOL_SIZE=1 RAY_MAX_TASKS_PER_WORKER=5 @@ -54,15 +53,12 @@ UV_CACHE_DIR=/tmp/uv-cache ``` :::tip[**Tips**] -- `RAY_NUM_GPUS` defines **per-actor resource requirements**. Ray will not start a task until these resources are available on one of the nodes. -For example, if one indexation consumes ~1GB of VRAM and your GPU has 4GB, setting `RAY_NUM_GPUS=0.25` allows you to run **4 indexers per node**. In a 2-node cluster, that means up to **8 concurrent indexation tasks**. - -- `RAY_POOL_SIZE` defines the number of worker actors that will be created to handle indexation tasks. It acts like a **maximum concurrency limit**. -Using the previous example, you can set `POOL_SIZE=8` to fully utilize your cluster capacity. +- `RAY_POOL_SIZE` defines the number of indexer worker actors created to handle indexation tasks, and `RAY_MAX_TASKS_PER_WORKER` the number of files each worker processes concurrently. +Total indexing concurrency is `RAY_POOL_SIZE × RAY_MAX_TASKS_PER_WORKER` — e.g. `RAY_POOL_SIZE=8` with `RAY_MAX_TASKS_PER_WORKER=5` allows up to **40 concurrent indexation tasks**. Size these to your cluster capacity. ::: :::caution -If other GPU-intensive services are running on your nodes (e.g. vLLM, the RAG API), make sure to **reserve enough GPU memory** for them and subtract that from your total when calculating the safe pool size. +`RAY_POOL_SIZE` and `RAY_MAX_TASKS_PER_WORKER` size **indexing throughput**, not GPU reservation — the indexer workers don't claim GPU memory. GPU on each node is consumed by the model servers (e.g. vLLM) and by the GPU-accelerated parsers, so budget GPU memory through **`MARKER_NUM_GPUS`** / **`MARKER_POOL_SIZE`** (and the Docling equivalents) and the model-server settings rather than through the indexer pool size. ::: --- diff --git a/docs/content/docs/documentation/env_vars.md b/docs/content/docs/documentation/env_vars.md index de5c1527d..68c4832b1 100644 --- a/docs/content/docs/documentation/env_vars.md +++ b/docs/content/docs/documentation/env_vars.md @@ -180,6 +180,7 @@ The vector database stores embeddings and is configured using the following envi | `VDB_COLLECTION_NAME` | str | vdb_test | Name of the collection storing embeddings | |`VDB_HYBRID_SEARCH`| `bool` | true |To activate hybrid search (semantic similarity + Keyword search)| | `VDB_ENABLE_INSERTION` | bool | true | Enable or disable vector database insertion. When disabled, documents are processed but not inserted into Milvus. Useful for testing. | +| `VDB_TIMEOUT` | float | 120.0 | Per-request timeout (seconds) applied to the Milvus sync and async clients | These variables can be overridden when using an external vector database service. @@ -375,8 +376,8 @@ Ray is used for distributed task processing and parallel execution in the RAG pi | Variable | Type | Default | Description | |----------|------|---------|-------------| -| `RAY_POOL_SIZE` | `int` | 1 | Number of serializer actor instances (typically 1 actor per cluster node) | -| `RAY_MAX_TASKS_PER_WORKER` | `int` | 8 | Maximum number of concurrent tasks (serialization tasks) per serializer actor instance | +| `RAY_POOL_SIZE` | `int` | 1 | Number of indexer worker actors in the pool. Total indexing capacity = `RAY_POOL_SIZE` × `RAY_MAX_TASKS_PER_WORKER`. | +| `RAY_MAX_TASKS_PER_WORKER` | `int` | 50 | Maximum number of files processed concurrently per indexer worker actor | | `RAY_DASHBOARD_PORT` | `int` | 8265 | Ray Dashboard port used for monitoring. In production, [comment out this line](https://github.com/linagora/openrag/blob/ee732ea8e080dcde0107d62d12703a7525f810cd/docker-compose.yaml#L21C1-L22C1) to avoid exposing the port, as it may introduce security vulnerabilities. | :::danger[Attention] @@ -391,33 +392,6 @@ The following environment variables control Ray's logging behavior, task retry s | `RAY_ENABLE_UV_RUN_RUNTIME_ENV` | `number` | `0` | Controls UV runtime environment integration. **Critical**: Must be set to `0` when using the newest version of UV to avoid compatibility issues. | |`RAY_memory_monitor_refresh_ms`| `number` | 250 ms | To control the frequency of memory usage checks and task or actor termination if needed. If you set this value to 0, task killing is disabled. | -#### Indexer Configuration - -| Variable | Type | Default | Description | -|----------|------|---------|-------------| -| `RAY_MAX_TASK_RETRIES` | int | 2 | Number of retry attempts for failed tasks | -| `INDEXER_SERIALIZE_TIMEOUT` | int | 36000 | Timeout in seconds for serialization operations (10 hours) | - -#### Indexer Concurrency Groups - -Controls the maximum number of concurrent operations for different indexer tasks: - -| Variable | Type | Default | Description | -|----------|------|---------|-------------| -| `INDEXER_DEFAULT_CONCURRENCY` | int | 1000 | Default concurrency limit for general operations | -| `INDEXER_UPDATE_CONCURRENCY` | int | 100 | Maximum concurrent document update operations | -| `INDEXER_SERIALIZE_CONCURRENCY` | int | 50 | Maximum concurrent serialization operations | -| `INDEXER_SEARCH_CONCURRENCY` | int | 100 | Maximum concurrent search/retrieval operations | -| `INDEXER_DELETE_CONCURRENCY` | int | 100 | Maximum concurrent document deletion operations | -| `INDEXER_CHUNK_CONCURRENCY` | int | 1000 | Maximum concurrent document chunking operations | -| `INDEXER_INSERT_CONCURRENCY` | int | 10 | Maximum concurrent document insertion operations | - -#### Semaphore Configuration - -| Variable | Type | Default | Description | -|----------|------|---------|-------------| -| `RAY_SEMAPHORE_CONCURRENCY` | int | 100000 | Global concurrency limit for Ray semaphore operations | - #### Ray Serve Configuration Ray Serve enables deployment of the FastAPI as a scalable service. For simple deployment, without the intend to scale, one can usage the [uvicorn deployment mode](/openrag/documentation/env_vars/#ray-serve-configuration) diff --git a/infra/charts/openrag-stack/values.yaml b/infra/charts/openrag-stack/values.yaml index f7e4d4ac4..a89edb243 100644 --- a/infra/charts/openrag-stack/values.yaml +++ b/infra/charts/openrag-stack/values.yaml @@ -358,7 +358,6 @@ env: TRANSCRIBER_BASE_URL: "http://{{ .Release.Name }}-whisper-engine-service/v1" # Ray - RAY_NUM_GPUS: "0.1" RAY_POOL_SIZE: "3" RAY_MAX_TASKS_PER_WORKER: "50" RAY_DASHBOARD_PORT: "8265" diff --git a/infra/compose/.env.ollama b/infra/compose/.env.ollama index feaf6504e..92cb94e81 100644 --- a/infra/compose/.env.ollama +++ b/infra/compose/.env.ollama @@ -38,7 +38,6 @@ CHAINLIT_AUTH_SECRET=openrag_local_dev_secret_2026 # RAY & System RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 -RAY_NUM_GPUS=0 RAY_DASHBOARD_PORT=8265 RAY_memory_usage_threshold=0.99 RAY_memory_monitor_refresh_ms=0 diff --git a/openrag/core/config/infrastructure.py b/openrag/core/config/infrastructure.py index 32001a85f..ecc4052dd 100644 --- a/openrag/core/config/infrastructure.py +++ b/openrag/core/config/infrastructure.py @@ -20,6 +20,8 @@ class VectorDBConfig(ConfigMixin): collection_name: str = "vdb_test" hybrid_search: bool = True enable: bool = True + # Per-request timeout (s) applied to the Milvus sync and async clients. + timeout: float = Field(default=120.0, gt=0) schema_version: int = 1 @@ -47,31 +49,14 @@ class RDBConfig(ConfigMixin): # --------------------------------------------------------------------------- -# Ray — concurrency groups, serve config +# Ray — worker pool, serve config # --------------------------------------------------------------------------- -class IndexerConcurrencyGroupsConfig(ConfigMixin): - default: int = 1000 - update: int = 100 - search: int = 100 - delete: int = 100 - serialize: int = 50 - chunk: int = 1000 - insert: int = 100 - - class RayIndexerConfig(ConfigMixin): - max_task_retries: int = 2 - serialize_timeout: int = 3600 - vectordb_timeout: int = 30 - concurrency_groups: IndexerConcurrencyGroupsConfig = Field( - default_factory=IndexerConcurrencyGroupsConfig, - ) - - -class RaySemaphoreConfig(ConfigMixin): - concurrency: int = 100000 + # Indexing capacity = pool_size (worker actors) × max_tasks_per_worker (files per worker). + pool_size: int = Field(default=1, ge=1) + max_tasks_per_worker: int = Field(default=50, ge=1) class RayServeConfig(ConfigMixin): @@ -83,11 +68,7 @@ class RayServeConfig(ConfigMixin): class RayConfig(ConfigMixin): - num_gpus: float = 0.01 - pool_size: int = 1 - max_tasks_per_worker: int = 8 indexer: RayIndexerConfig = Field(default_factory=RayIndexerConfig) - semaphore: RaySemaphoreConfig = Field(default_factory=RaySemaphoreConfig) serve: RayServeConfig = Field(default_factory=RayServeConfig) diff --git a/openrag/core/config/loader.py b/openrag/core/config/loader.py index c5d8beb61..6febadb4c 100644 --- a/openrag/core/config/loader.py +++ b/openrag/core/config/loader.py @@ -52,6 +52,7 @@ ("VDB_COLLECTION_NAME", "vectordb.collection_name", str), ("VDB_HYBRID_SEARCH", "vectordb.hybrid_search", bool), ("VDB_ENABLE_INSERTION", "vectordb.enable", bool), + ("VDB_TIMEOUT", "vectordb.timeout", float), # RDB (Postgres) ("POSTGRES_HOST", "rdb.host", str), ("POSTGRES_PORT", "rdb.port", int), @@ -127,20 +128,8 @@ ("OPENAI_LOADER_CONCURRENCY_LIMIT", "loader.openai.concurrency_limit", int), ("OPENAI_LOADER_ENABLE_THINKING", "loader.openai.enable_thinking", bool), # Ray - ("RAY_NUM_GPUS", "ray.num_gpus", float), - ("RAY_POOL_SIZE", "ray.pool_size", int), - ("RAY_MAX_TASKS_PER_WORKER", "ray.max_tasks_per_worker", int), - ("RAY_MAX_TASK_RETRIES", "ray.indexer.max_task_retries", int), - ("INDEXER_SERIALIZE_TIMEOUT", "ray.indexer.serialize_timeout", int), - ("VECTORDB_TIMEOUT", "ray.indexer.vectordb_timeout", int), - ("INDEXER_DEFAULT_CONCURRENCY", "ray.indexer.concurrency_groups.default", int), - ("INDEXER_UPDATE_CONCURRENCY", "ray.indexer.concurrency_groups.update", int), - ("INDEXER_SEARCH_CONCURRENCY", "ray.indexer.concurrency_groups.search", int), - ("INDEXER_DELETE_CONCURRENCY", "ray.indexer.concurrency_groups.delete", int), - ("INDEXER_SERIALIZE_CONCURRENCY", "ray.indexer.concurrency_groups.serialize", int), - ("INDEXER_CHUNK_CONCURRENCY", "ray.indexer.concurrency_groups.chunk", int), - ("INDEXER_INSERT_CONCURRENCY", "ray.indexer.concurrency_groups.insert", int), - ("RAY_SEMAPHORE_CONCURRENCY", "ray.semaphore.concurrency", int), + ("RAY_POOL_SIZE", "ray.indexer.pool_size", int), + ("RAY_MAX_TASKS_PER_WORKER", "ray.indexer.max_tasks_per_worker", int), ("ENABLE_RAY_SERVE", "ray.serve.enable", bool), ("RAY_SERVE_NUM_REPLICAS", "ray.serve.num_replicas", int), ("RAY_SERVE_HOST", "ray.serve.host", str), diff --git a/openrag/services/orchestrators/model_endpoint_service.py b/openrag/services/orchestrators/model_endpoint_service.py index f416cd39b..41bc0a01e 100644 --- a/openrag/services/orchestrators/model_endpoint_service.py +++ b/openrag/services/orchestrators/model_endpoint_service.py @@ -98,7 +98,16 @@ async def seed_defaults(self) -> None: logger.info(f"Seeded default {model_type} endpoint '{row.name}'.") def _build_default_seeds(self) -> dict[str, dict[str, Any]]: - """Build seed data from env overrides + existing Settings fallbacks.""" + """Build seed data from env overrides + existing Settings fallbacks. + + The ``*_ENDPOINT`` / ``*_MODEL`` env vars below are seed-specific names + the config loader does NOT map onto ``Settings``, so they are read here + directly. The api-key env vars (``API_KEY``, ``EMBEDDER_API_KEY``, ...) + ARE mapped by the loader (loader.py), so ``s..api_key`` already + reflects any env override — reading them via ``os.getenv`` again would + be redundant double-handling (and non-deterministic when a local .env is + loaded into the process via ``load_dotenv``). + """ s = self._config return { "embedder": { @@ -106,20 +115,14 @@ def _build_default_seeds(self) -> dict[str, dict[str, Any]]: "model_name": os.getenv("EMBEDDING_MODEL", s.embedder.model_name), "batch_size": s.embedder.batch_size, "timeout": s.embedder.timeout, - "extra": _with_api_key( - {"implementation": "vllm"}, - os.getenv("EMBEDDER_API_KEY", s.embedder.api_key), - ), + "extra": _with_api_key({"implementation": "vllm"}, s.embedder.api_key), }, "llm": { "endpoint": os.getenv("LLM_ENDPOINT", s.llm.base_url), "model_name": os.getenv("LLM_MODEL", s.llm.model), "timeout": s.llm.timeout, "extra": _with_enable_thinking( - _with_api_key( - {"implementation": "vllm"}, - os.getenv("API_KEY", s.llm.api_key), - ), + _with_api_key({"implementation": "vllm"}, s.llm.api_key), s.llm.enable_thinking, ), }, @@ -128,10 +131,7 @@ def _build_default_seeds(self) -> dict[str, dict[str, Any]]: "model_name": os.getenv("VLM_MODEL", s.vlm.model), "timeout": s.vlm.timeout, "extra": _with_enable_thinking( - _with_api_key( - {"implementation": "vllm"}, - os.getenv("VLM_API_KEY", s.vlm.api_key), - ), + _with_api_key({"implementation": "vllm"}, s.vlm.api_key), s.vlm.enable_thinking, ), }, @@ -145,10 +145,7 @@ def _build_default_seeds(self) -> dict[str, dict[str, Any]]: "endpoint": os.getenv("RERANKER_ENDPOINT", s.reranker.base_url), "model_name": os.getenv("RERANKER_MODEL", s.reranker.model_name), "timeout": s.reranker.timeout, - "extra": _with_api_key( - {"implementation": s.reranker.provider}, - os.getenv("RERANKER_API_KEY", s.reranker.api_key), - ), + "extra": _with_api_key({"implementation": s.reranker.provider}, s.reranker.api_key), }, } diff --git a/openrag/services/storage/milvus_store.py b/openrag/services/storage/milvus_store.py index 2ff4c0cf4..dbf426e45 100644 --- a/openrag/services/storage/milvus_store.py +++ b/openrag/services/storage/milvus_store.py @@ -144,7 +144,7 @@ def __init__(self, config: VectorDBConfig) -> None: self._collection_name = config.collection_name self._hybrid = config.hybrid_search self._uri = f"http://{config.host}:{config.port}" - self._timeout = 60 + self._timeout = config.timeout try: self._client = MilvusClient(uri=self._uri, timeout=self._timeout) self._async_client = AsyncMilvusClient(uri=self._uri, timeout=self._timeout) diff --git a/openrag/services/workers/dispatcher.py b/openrag/services/workers/dispatcher.py index 80d329b2b..8a38c99dc 100644 --- a/openrag/services/workers/dispatcher.py +++ b/openrag/services/workers/dispatcher.py @@ -90,17 +90,25 @@ async def dispatch_indexing( task_description=f"set_details({task_id})", ) - task = self._pool.process_file.remote( - task_id=task_id, - path=path, - metadata=metadata, - partition=partition, - user=user, - workspace_ids=workspace_ids, - replace=replace, - indexation_config=indexation_config, - embedder_name=embedder_name, + # ``IndexerPool`` is a Ray actor; ``submit`` returns ``[worker_ref]`` + # (wrapped so Ray doesn't auto-dereference and block on the worker task). + # Awaiting the submit call yields that list; element 0 is the worker ref + # that ``cancel_task``/``ray.cancel`` must target. + submitted = await self._call( + self._pool.submit.remote( + task_id=task_id, + path=path, + metadata=metadata, + partition=partition, + user=user, + workspace_ids=workspace_ids, + replace=replace, + indexation_config=indexation_config, + embedder_name=embedder_name, + ), + task_description=f"submit({task_id})", ) + task = submitted[0] await self._call( self._tsm.set_object_ref.remote(task_id, {"ref": task}), diff --git a/openrag/services/workers/indexer_pool.py b/openrag/services/workers/indexer_pool.py index f82643612..7b66e16d1 100644 --- a/openrag/services/workers/indexer_pool.py +++ b/openrag/services/workers/indexer_pool.py @@ -15,11 +15,17 @@ # this window (and on a miss), bounding both staleness and DB load regardless # of indexing throughput. _MODEL_REGISTRY_TTL_SECONDS = 60.0 +_INDEXER_POOL_DISPATCHER_ACTOR_NAME = "IndexerPoolDispatcher" @ray.remote -class IndexerPool: - """Thin Ray actor wrapper around ``IndexerWorker``.""" +class IndexerWorkerActor: + """Thin Ray actor wrapping ``IndexerWorker`` — one instance per pool slot. + + A worker runs up to ``ray.indexer.max_tasks_per_worker`` files concurrently + (its Ray ``max_concurrency``). ``IndexerPool`` holds ``ray.indexer.pool_size`` + of these and load-balances across them. + """ def __init__(self) -> None: import services.inference.ollama_client # noqa: F401 @@ -223,18 +229,107 @@ async def process_file( return result +@ray.remote +class IndexerPool: + """Single detached dispatcher actor over a fleet of ``IndexerWorkerActor``. + + Exactly **one** instance exists cluster-wide — it is created named, detached + and with ``get_if_exists=True`` (see :func:`build_indexer_pool`). Every API + process and every Ray Serve replica therefore shares the *same* dispatcher, + so the least-loaded view is global. A per-replica client object would keep + its own ``_inflight`` counters, and since Ray Serve runs each replica as a + separate actor/process, those views would diverge and unbalance dispatch + across the shared workers under bursts. + + It holds ``ray.indexer.pool_size`` worker actors and dispatches each file to + the least-loaded one (fewest in-flight files). ``submit`` returns the + worker's Ray ``ObjectRef`` wrapped in a one-element list, so the caller keeps + per-task cancellation (``ray.cancel``) and task-state tracking. The wrapper + matters: a *bare* returned ``ObjectRef`` may be auto-dereferenced by Ray + (blocking ``ray.get`` until the file finishes indexing), whereas a ref nested + inside a container is returned unresolved — see Ray's nested-task semantics. + + This is an async Ray actor, so ``submit`` and the release callbacks run on a + single event loop; ``_inflight`` is mutated from that one thread only and + needs no lock. + """ + + def __init__(self, pool_size: int, max_tasks_per_worker: int, namespace: str = "openrag") -> None: + if pool_size < 1: + raise ValueError("IndexerPool requires pool_size >= 1") + # The workers are themselves named + detached + get_if_exists, so they + # are shared singletons too; only this dispatcher creates them. + self._workers = [ + IndexerWorkerActor.options( # type: ignore[attr-defined] + name=f"IndexerWorker-{i}", + namespace=namespace, + get_if_exists=True, + lifetime="detached", + max_concurrency=max_tasks_per_worker, + ).remote() + for i in range(pool_size) + ] + self._inflight = [0] * len(self._workers) + self._release_tasks: set[asyncio.Task[Any]] = set() + + async def size(self) -> int: + return len(self._workers) + + async def submit(self, **kwargs: Any) -> list[Any]: + """Dispatch ``process_file`` to the least-loaded worker. + + Returns ``[worker_ref]`` (the worker's Ray ``ObjectRef`` in a + one-element list — see the class docstring); in-flight bookkeeping is + released when the task settles (success, failure, or cancellation). + """ + idx = min(range(len(self._workers)), key=self._inflight.__getitem__) + self._inflight[idx] += 1 + try: + ref = self._workers[idx].process_file.remote(**kwargs) + except Exception: + # Submission failed before a ref exists (e.g. unserializable args or + # a dead actor); roll back so load balancing stays accurate. + self._inflight[idx] -= 1 + raise + task = asyncio.get_running_loop().create_task(self._release(idx, ref)) + # Keep a strong ref so the tracker isn't GC'd mid-flight (asyncio docs). + self._release_tasks.add(task) + task.add_done_callback(self._release_tasks.discard) + return [ref] + + async def _release(self, idx: int, ref: Any) -> None: + try: + # return_exceptions=True so a failed/cancelled task still decrements. + await asyncio.gather(ref, return_exceptions=True) + finally: + self._inflight[idx] -= 1 + + def build_indexer_pool(namespace: str = "openrag") -> Any: from core.config import load_config cfg = load_config() - max_concurrency = max(1, cfg.ray.max_tasks_per_worker) + pool_size = cfg.ray.indexer.pool_size + max_tasks_per_worker = cfg.ray.indexer.max_tasks_per_worker + # One detached dispatcher actor shared by all API / Serve replicas via + # get_if_exists. Its own max_concurrency only bounds concurrent submit() + # calls (each returns promptly without awaiting the worker), so size it to + # the whole fleet's capacity. The constructor args are honoured only on the + # first creation; later get_if_exists calls reuse the existing dispatcher + # and ignore them — which is correct, since every replica loads the same cfg. + # Do not reuse the old "IndexerPool" name: that detached actor exposed + # process_file(), while this dispatcher exposes submit(). return IndexerPool.options( # type: ignore[attr-defined] - name="IndexerPool", + name=_INDEXER_POOL_DISPATCHER_ACTOR_NAME, namespace=namespace, get_if_exists=True, lifetime="detached", - max_concurrency=max_concurrency, - ).remote() + max_concurrency=max(1, pool_size * max_tasks_per_worker), + ).remote( + pool_size=pool_size, + max_tasks_per_worker=max_tasks_per_worker, + namespace=namespace, + ) def _required_llm_names(indexation_config: dict[str, Any] | None) -> list[str]: @@ -648,4 +743,4 @@ def _global_vlm_endpoint_config(cfg: Any) -> Any | None: ) -__all__ = ["IndexerPool", "build_indexer_pool"] +__all__ = ["IndexerPool", "IndexerWorkerActor", "build_indexer_pool"] diff --git a/openrag/services/workers/task_state.py b/openrag/services/workers/task_state.py index 8f0584b70..7e5bc6329 100644 --- a/openrag/services/workers/task_state.py +++ b/openrag/services/workers/task_state.py @@ -10,8 +10,8 @@ from core.config import load_config as _load_config _cfg = _load_config() - _POOL_SIZE: int = _cfg.ray.pool_size - _MAX_TASKS_PER_WORKER: int = _cfg.ray.max_tasks_per_worker + _POOL_SIZE: int = _cfg.ray.indexer.pool_size + _MAX_TASKS_PER_WORKER: int = _cfg.ray.indexer.max_tasks_per_worker except (ImportError, AttributeError) as _cfg_err: import logging as _logging diff --git a/tests/unit/core/config/test_ray_config.py b/tests/unit/core/config/test_ray_config.py new file mode 100644 index 000000000..020279574 --- /dev/null +++ b/tests/unit/core/config/test_ray_config.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +import pytest +from core.config.infrastructure import RayIndexerConfig +from pydantic import ValidationError + + +def test_ray_indexer_config_defaults() -> None: + cfg = RayIndexerConfig() + assert cfg.pool_size == 1 + assert cfg.max_tasks_per_worker == 50 + + +@pytest.mark.parametrize("field", ["pool_size", "max_tasks_per_worker"]) +def test_ray_indexer_config_rejects_sub_one(field: str) -> None: + # The runtime max(1, ...) floor was replaced by ge=1 validation at the boundary. + with pytest.raises(ValidationError): + RayIndexerConfig(**{field: 0}) diff --git a/tests/unit/core/config/test_vectordb_config.py b/tests/unit/core/config/test_vectordb_config.py new file mode 100644 index 000000000..fd83a1d77 --- /dev/null +++ b/tests/unit/core/config/test_vectordb_config.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +import pytest +from core.config import load_config +from core.config.infrastructure import VectorDBConfig +from pydantic import ValidationError + + +def test_vectordb_config_timeout_default() -> None: + # Replaces the value formerly hardcoded as ``MilvusVectorStore._timeout``. + cfg = VectorDBConfig() + assert cfg.timeout == 120.0 + assert isinstance(cfg.timeout, float) + + +@pytest.mark.parametrize("bad_timeout", [0, -1, -0.5]) +def test_vectordb_config_rejects_non_positive_timeout(bad_timeout: float) -> None: + # A non-positive timeout must fail at config load time, not at client usage. + with pytest.raises(ValidationError): + VectorDBConfig(timeout=bad_timeout) + + +def test_vdb_timeout_can_be_overridden_from_env(monkeypatch, tmp_path) -> None: + (tmp_path / "config.yaml").write_text("retriever:\n type: single\n", encoding="utf-8") + monkeypatch.setenv("VDB_TIMEOUT", "200") + + settings = load_config(config_path=tmp_path) + + assert settings.vectordb.timeout == 200.0 diff --git a/tests/unit/services/workers/test_dispatcher.py b/tests/unit/services/workers/test_dispatcher.py index 87bc976bc..39e3db6f5 100644 --- a/tests/unit/services/workers/test_dispatcher.py +++ b/tests/unit/services/workers/test_dispatcher.py @@ -14,8 +14,10 @@ def _remote_mock(return_value: Any = None) -> MagicMock: def _pool_with_ref(ref: object) -> MagicMock: pool = MagicMock() - pool.process_file = MagicMock() - pool.process_file.remote = MagicMock(return_value=ref) + # IndexerPool is a Ray actor; submit.remote() picks the least-loaded worker + # and returns its ObjectRef wrapped in a one-element list (the dispatcher + # awaits the call and takes element 0). + pool.submit = _remote_mock([ref]) return pool @@ -131,7 +133,7 @@ async def test_dispatch_indexing_queues_worker_pool_task_and_records_ref() -> No metadata={"filename": "report.txt"}, user_id=42, ) - pool.process_file.remote.assert_called_once_with( + pool.submit.remote.assert_called_once_with( task_id="task-1", path="/data/report.txt", metadata={"file_id": "file-1", "source": "/data/report.txt", "filename": "report.txt"}, diff --git a/tests/unit/services/workers/test_indexer_pool.py b/tests/unit/services/workers/test_indexer_pool.py index d888108d8..8b9809e97 100644 --- a/tests/unit/services/workers/test_indexer_pool.py +++ b/tests/unit/services/workers/test_indexer_pool.py @@ -52,9 +52,9 @@ def test_build_chunker_rejects_non_callable_chunk_attr(monkeypatch: pytest.Monke @pytest.mark.asyncio async def test_catalog_initialization_is_single_flight() -> None: - from services.workers.indexer_pool import IndexerPool + from services.workers.indexer_pool import IndexerWorkerActor - actor_class = IndexerPool.__ray_metadata__.modified_class + actor_class = IndexerWorkerActor.__ray_metadata__.modified_class pool = actor_class.__new__(actor_class) class Store: @@ -76,29 +76,80 @@ async def initialize(self) -> None: assert pool._catalog_initialized is True -def test_build_indexer_pool_uses_detached_actor_with_configured_concurrency( +def test_build_indexer_pool_uses_new_detached_dispatcher_name( monkeypatch: pytest.MonkeyPatch, ) -> None: import core.config import services.workers.indexer_pool as module - calls = {} + options_calls = [] + remote_calls = [] class Options: - def remote(self): - return "actor" + def __init__(self, kwargs): + self._kwargs = kwargs + + def remote(self, **rkwargs): + remote_calls.append(rkwargs) + return "dispatcher-actor" def fake_options(**kwargs): - calls.update(kwargs) - return Options() + options_calls.append(kwargs) + return Options(kwargs) - cfg = SimpleNamespace(ray=SimpleNamespace(max_tasks_per_worker=4)) + cfg = SimpleNamespace(ray=SimpleNamespace(indexer=SimpleNamespace(pool_size=3, max_tasks_per_worker=4))) monkeypatch.setattr(core.config, "load_config", lambda: cfg) monkeypatch.setattr(module.IndexerPool, "options", fake_options) - assert module.build_indexer_pool() == "actor" - assert calls["lifetime"] == "detached" - assert calls["max_concurrency"] == 4 + pool = module.build_indexer_pool() + + # A single shared dispatcher actor — not one client object per replica. + assert pool == "dispatcher-actor" + assert len(options_calls) == 1 + opts = options_calls[0] + # The dispatcher has a different public interface from the old detached + # IndexerPool actor, so it must not reuse that actor name during upgrades. + assert opts["name"] == "IndexerPoolDispatcher" + assert opts["namespace"] == "openrag" + assert opts["get_if_exists"] is True + assert opts["lifetime"] == "detached" + # max_concurrency bounds concurrent submit() calls → whole-fleet capacity. + assert opts["max_concurrency"] == 12 + # pool_size / max_tasks_per_worker are passed to the actor constructor. + assert remote_calls == [{"pool_size": 3, "max_tasks_per_worker": 4, "namespace": "openrag"}] + + +def test_indexer_pool_actor_spawns_pool_size_detached_workers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import services.workers.indexer_pool as module + + calls = [] + + class Options: + def __init__(self, kwargs): + self._kwargs = kwargs + + def remote(self): + return f"actor-{self._kwargs['name']}" + + def fake_options(**kwargs): + calls.append(kwargs) + return Options(kwargs) + + monkeypatch.setattr(module.IndexerWorkerActor, "options", fake_options) + + actor_class = module.IndexerPool.__ray_metadata__.modified_class + pool = actor_class(pool_size=3, max_tasks_per_worker=4) + + # One detached worker actor per pool_size slot, each capped at max_tasks_per_worker. + assert len(pool._workers) == 3 + assert {c["name"] for c in calls} == {"IndexerWorker-0", "IndexerWorker-1", "IndexerWorker-2"} + for c in calls: + assert c["lifetime"] == "detached" + assert c["max_concurrency"] == 4 + assert c["get_if_exists"] is True + assert c["namespace"] == "openrag" def test_build_topic_tagger_factory_resolves_named_llm(monkeypatch: pytest.MonkeyPatch) -> None: @@ -453,9 +504,9 @@ def test_reload_decision_treats_default_global_fallback_as_resolvable() -> None: # block-reload every window forever without converging. import time as _time - from services.workers.indexer_pool import IndexerPool + from services.workers.indexer_pool import IndexerWorkerActor - actor_class = IndexerPool.__ray_metadata__.modified_class + actor_class = IndexerWorkerActor.__ray_metadata__.modified_class def _pool(has_fallback: bool): pool = actor_class.__new__(actor_class) @@ -475,9 +526,9 @@ def _pool(has_fallback: bool): @pytest.mark.asyncio async def test_ensure_registry_fresh_is_single_flight() -> None: - from services.workers.indexer_pool import IndexerPool + from services.workers.indexer_pool import IndexerWorkerActor - actor_class = IndexerPool.__ray_metadata__.modified_class + actor_class = IndexerWorkerActor.__ray_metadata__.modified_class pool = actor_class.__new__(actor_class) class Service: @@ -508,9 +559,9 @@ async def test_ttl_refresh_runs_in_background_without_blocking() -> None: # the reload runs in the background. import time as _time - from services.workers.indexer_pool import _MODEL_REGISTRY_TTL_SECONDS, IndexerPool + from services.workers.indexer_pool import _MODEL_REGISTRY_TTL_SECONDS, IndexerWorkerActor - actor_class = IndexerPool.__ray_metadata__.modified_class + actor_class = IndexerWorkerActor.__ray_metadata__.modified_class pool = actor_class.__new__(actor_class) started = asyncio.Event() @@ -729,6 +780,128 @@ async def chat(self, messages, **kwargs): llm_registry._registry.pop("test-contextualizer-llm", None) +class _FakeWorker: + """Stand-in for an ``IndexerWorkerActor`` handle. + + ``process_file.remote(**kwargs)`` returns an ``asyncio.Future`` that + plays the role of a Ray ``ObjectRef`` (``asyncio.gather`` accepts both), + so tests can drive task completion deterministically. + """ + + def __init__(self) -> None: + self.calls: list[dict] = [] + self.futures: list[asyncio.Future] = [] + self.process_file = SimpleNamespace(remote=self._remote) + + def _remote(self, **kwargs): + fut = asyncio.get_running_loop().create_future() + self.calls.append(kwargs) + self.futures.append(fut) + return fut + + +def _bare_pool(workers: list) -> object: + """An ``IndexerPool`` actor instance with ``__init__`` bypassed. + + The dispatch/release logic under test lives on the actor class; we set the + fields it touches directly so tests can inject fake workers instead of + spawning real Ray actors. + """ + from services.workers.indexer_pool import IndexerPool + + actor_class = IndexerPool.__ray_metadata__.modified_class + pool = actor_class.__new__(actor_class) + pool._workers = list(workers) + pool._inflight = [0] * len(workers) + pool._release_tasks = set() + return pool + + +async def _settle_pool_release_tasks(pool: object, *futures: asyncio.Future[object]) -> None: + for fut in futures: + if not fut.done(): + fut.set_result(None) + release_tasks = list(getattr(pool, "_release_tasks")) + if release_tasks: + await asyncio.gather(*release_tasks) + + +def test_pool_requires_positive_pool_size() -> None: + from services.workers.indexer_pool import IndexerPool + + actor_class = IndexerPool.__ray_metadata__.modified_class + with pytest.raises(ValueError): + actor_class(pool_size=0, max_tasks_per_worker=4) + + +@pytest.mark.asyncio +async def test_pool_dispatches_to_least_loaded_and_passes_ref_through() -> None: + workers = [_FakeWorker(), _FakeWorker()] + pool = _bare_pool(workers) + + ref0 = await pool.submit(task_id="a") # tie -> worker 0 + await pool.submit(task_id="b") # worker 0 busy -> worker 1 + await pool.submit(task_id="c") # tie (1 each) -> worker 0 + + assert len(workers[0].calls) == 2 + assert len(workers[1].calls) == 1 + # The ObjectRef is passed through wrapped in a one-element list (the + # dispatcher unwraps it; the wrapper stops Ray auto-dereferencing the ref). + assert ref0 == [workers[0].futures[0]] + await _settle_pool_release_tasks( + pool, + workers[0].futures[0], + workers[1].futures[0], + workers[0].futures[1], + ) + + +@pytest.mark.asyncio +async def test_pool_releases_inflight_when_task_settles() -> None: + workers = [_FakeWorker(), _FakeWorker()] + pool = _bare_pool(workers) + + await pool.submit(task_id="a") # worker 0 + await pool.submit(task_id="b") # worker 1 + assert pool._inflight == [1, 1] + + # One success, one failure — both must decrement the in-flight count. + workers[0].futures[0].set_result({"ok": True}) + workers[1].futures[0].set_exception(RuntimeError("boom")) + + for _ in range(20): + await asyncio.sleep(0) + if pool._inflight == [0, 0]: + break + assert pool._inflight == [0, 0] + + # Freed workers are eligible again on the next dispatch. + await pool.submit(task_id="c") + assert pool._inflight[0] == 1 + await _settle_pool_release_tasks(pool, workers[0].futures[1]) + + +@pytest.mark.asyncio +async def test_pool_rolls_back_inflight_when_submission_raises() -> None: + # If process_file.remote raises (e.g. unserializable args or a dead actor), + # the in-flight count must be rolled back so the worker isn't seen as busy. + class _RaisingWorker: + def __init__(self) -> None: + def _boom(**_kwargs): + raise RuntimeError("remote submission failed") + + self.process_file = SimpleNamespace(remote=_boom) + + pool = _bare_pool([_RaisingWorker()]) + + with pytest.raises(RuntimeError, match="remote submission failed"): + await pool.submit(task_id="a") + + assert pool._inflight == [0] + + assert pool._inflight == [0] + + def test_indexer_pool_wires_contextualizer_factory(monkeypatch: pytest.MonkeyPatch) -> None: import core.config import core.embeddings @@ -795,7 +968,7 @@ def fake_get_actor(*args, **kwargs): monkeypatch.setattr(module.ray, "get_actor", fake_get_actor) monkeypatch.setattr(module, "IndexerWorker", Worker) - actor_class = module.IndexerPool.__ray_metadata__.modified_class + actor_class = module.IndexerWorkerActor.__ray_metadata__.modified_class actor_class() assert actor_calls