diff --git a/entrypoint.sh b/entrypoint.sh index a13188c03..5fcdf4fb5 100644 --- a/entrypoint.sh +++ b/entrypoint.sh @@ -9,5 +9,5 @@ if [[ "${ENABLE_RAY_SERVE}" == "true" ]]; then uv run $ENV_ARG -m api.main else echo "🚀 Starting with Uvicorn..." - uv run --no-dev $ENV_ARG uvicorn api.main:app --host 0.0.0.0 --port ${APP_PORT:-8080} --reload --workers ${API_NUM_WORKERS:-1} + uv run --no-dev $ENV_ARG uvicorn api.main:app --host 0.0.0.0 --port ${APP_iPORT:-8080} --reload --workers ${API_NUM_WORKERS:-1} fi diff --git a/openrag/api/__init__.py b/openrag/api/__init__.py index e69de29bb..0c0061b07 100644 --- a/openrag/api/__init__.py +++ b/openrag/api/__init__.py @@ -0,0 +1,12 @@ +"""Compatibility entrypoint for ASGI loaders using ``api:app``.""" + + +def __getattr__(name: str): + if name == "app": + from api.main import app + + return app + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = ["app"] diff --git a/openrag/api/dependencies/auth.py b/openrag/api/dependencies/auth.py index d18289f1b..3d930ecba 100644 --- a/openrag/api/dependencies/auth.py +++ b/openrag/api/dependencies/auth.py @@ -1,16 +1,13 @@ import os -from config import load_config from core.utils.exceptions import OpenRAGError -from di.providers import get_auth_service, get_job_service, get_partition_service +from di.providers import get_auth_service, get_config, get_job_service, get_partition_service from fastapi import Depends, HTTPException, Request, status from utils.logger import get_logger -config = load_config() logger = get_logger() SUPER_ADMIN_MODE = os.getenv("SUPER_ADMIN_MODE", "false").lower() == "true" -DEFAULT_FILE_QUOTA = config.rdb.default_file_quota def current_user(request: Request): @@ -236,11 +233,13 @@ async def check_user_file_quota( user=Depends(current_user), auth_service=Depends(get_auth_service), job_service=Depends(get_job_service), + config=Depends(get_config), ): """Check if user has reached their file quota.""" + default_file_quota = config.rdb.default_file_quota if user.get("is_admin", False): return user - if DEFAULT_FILE_QUOTA < 0: + if default_file_quota < 0: return user user_quota = user.get("file_quota") if user_quota is not None and user_quota < 0: @@ -259,7 +258,7 @@ async def check_user_file_quota( auth_service.validate_file_quota( user, pending_task_count=pending_count, - default_quota=DEFAULT_FILE_QUOTA, + default_quota=default_file_quota, ) except OpenRAGError as exc: raise HTTPException( diff --git a/openrag/api/dependencies/files.py b/openrag/api/dependencies/files.py index b751042cd..a6e3f2c03 100644 --- a/openrag/api/dependencies/files.py +++ b/openrag/api/dependencies/files.py @@ -1,16 +1,10 @@ -from pathlib import Path from typing import Any -from config import load_config from core.indexing import validators as core_validators +from di.providers import get_config from fastapi import Depends, Form, UploadFile -config = load_config() - FORBIDDEN_CHARS_IN_FILE_ID = set("/") -LOG_FILE = Path(config.paths.log_dir or "logs") / "app.json" -ACCEPTED_FILE_FORMATS = config.loader.file_loaders.model_dump().keys() -DICT_MIMETYPES = config.loader.mimetypes.to_dict() async def validate_file_id(file_id: str): @@ -24,11 +18,14 @@ async def validate_metadata(metadata: Any | None = Form(None)): async def validate_file_format( file: UploadFile, metadata: dict = Depends(validate_metadata), + config=Depends(get_config), ): + accepted_file_formats = config.loader.file_loaders.model_dump().keys() + mimetypes = config.loader.mimetypes.to_dict() core_validators.validate_file_format( filename=file.filename, - accepted_formats=ACCEPTED_FILE_FORMATS, - accepted_mimetypes=DICT_MIMETYPES.keys(), + accepted_formats=accepted_file_formats, + accepted_mimetypes=mimetypes.keys(), mimetype=metadata.get("mimetype"), ) return file diff --git a/openrag/api/dependencies/llm.py b/openrag/api/dependencies/llm.py index a27853e68..6334257a3 100644 --- a/openrag/api/dependencies/llm.py +++ b/openrag/api/dependencies/llm.py @@ -1,12 +1,11 @@ import consts import openai from api.dependencies.auth import SUPER_ADMIN_MODE -from config import load_config -from fastapi import HTTPException, status +from di.providers import get_config +from fastapi import Depends, HTTPException, status from openai import AsyncOpenAI from utils.logger import get_logger -config = load_config() logger = get_logger() @@ -16,7 +15,7 @@ async def get_openai_models(base_url: str, api_key: str, timeout: int = 30): return models_response.data -async def check_llm_model_availability(): +async def check_llm_model_availability(config=Depends(get_config)): llm_param = config.llm base_url = llm_param.base_url model = llm_param.model diff --git a/openrag/api/dependencies/test_auth.py b/openrag/api/dependencies/test_auth.py index 68ef5a030..374e55c2e 100644 --- a/openrag/api/dependencies/test_auth.py +++ b/openrag/api/dependencies/test_auth.py @@ -1,5 +1,6 @@ +from types import SimpleNamespace + import pytest -from api.dependencies import auth as api_auth from api.dependencies.auth import check_user_file_quota, ensure_partition_role, require_task_owner from core.utils.exceptions import AuthError from fastapi import HTTPException @@ -85,6 +86,10 @@ async def get_user_pending_task_count(self, user_id: int | None) -> int: return self.pending_count +def _config(default_file_quota: int): + return SimpleNamespace(rdb=SimpleNamespace(default_file_quota=default_file_quota)) + + @pytest.mark.asyncio async def test_ensure_partition_role_allows_unknown_partition_without_membership(): partition_service = FakePartitionService(existing=set()) @@ -154,14 +159,14 @@ async def test_require_task_owner_reads_task_details_through_job_service(): @pytest.mark.asyncio -async def test_check_user_file_quota_reads_pending_count_through_job_service(monkeypatch): - monkeypatch.setattr(api_auth, "DEFAULT_FILE_QUOTA", 10) +async def test_check_user_file_quota_reads_pending_count_through_job_service(): job_service = FakeJobService(pending_count=2) user = await check_user_file_quota( user={"id": 7, "file_count": 1, "file_quota": 5}, auth_service=FakeAuthService, job_service=job_service, + config=_config(default_file_quota=10), ) assert user["id"] == 7 @@ -169,14 +174,13 @@ async def test_check_user_file_quota_reads_pending_count_through_job_service(mon @pytest.mark.asyncio -async def test_check_user_file_quota_default_zero_enforces_zero(monkeypatch): - monkeypatch.setattr(api_auth, "DEFAULT_FILE_QUOTA", 0) - +async def test_check_user_file_quota_default_zero_enforces_zero(): allowed_job_service = FakeJobService(pending_count=0) user = await check_user_file_quota( user={"id": 7, "file_count": 0, "file_quota": 1}, auth_service=EnforcingAuthService, job_service=allowed_job_service, + config=_config(default_file_quota=0), ) assert user["id"] == 7 assert allowed_job_service.pending_checks == [7] @@ -187,6 +191,7 @@ async def test_check_user_file_quota_default_zero_enforces_zero(monkeypatch): user={"id": 8, "file_count": 1, "file_quota": None}, auth_service=EnforcingAuthService, job_service=denied_job_service, + config=_config(default_file_quota=0), ) assert exc.value.status_code == 403 assert exc.value.detail == "File quota exceeded" diff --git a/openrag/api/main.py b/openrag/api/main.py index 70e82d796..bf1b79638 100644 --- a/openrag/api/main.py +++ b/openrag/api/main.py @@ -19,6 +19,7 @@ from __future__ import annotations +import asyncio import os import warnings from contextlib import asynccontextmanager @@ -52,6 +53,7 @@ from api.routers.user.search import router as search_router from config import load_config from di.container import ServiceContainer +from di.providers import set_container from di.workers import ensure_worker_bootstrap from dotenv import dotenv_values from fastapi import Depends, FastAPI @@ -70,8 +72,11 @@ # --------------------------------------------------------------------------- logger = get_logger() -config = load_config() -DATA_DIR = Path(config.paths.data_dir) +settings = load_config() +DATA_DIR = Path(settings.paths.data_dir) +CONTAINER_STARTUP_TIMEOUT = float( + os.getenv("OPENRAG_CONTAINER_STARTUP_TIMEOUT", max(60, settings.rdb.command_timeout * 4)) +) SHARED_ENV = os.environ.get("SHARED_ENV", None) env_vars = dotenv_values(SHARED_ENV) if SHARED_ENV else {} @@ -133,20 +138,29 @@ async def lifespan(app: FastAPI): container is absent. 4. ``prime_max_model_tokens`` — caches the vLLM ``max_model_len`` so per-request validation in the chat router stays synchronous. + 5. ``set_container`` — registers the resolved container as the + process-level singleton for callers that resolve it without a + request; cleared on shutdown. """ if not ray.is_initialized(): + logger.info("Startup: initializing Ray") ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True) + logger.info("Startup: Ray is initialized") # ``ensure_worker_bootstrap`` imports ``services.workers.bootstrap`` # for its side effect: creating the long-lived detached worker # actors (TaskStateManager, DocSerializer, MarkerPool, semaphores). # The indirection through :mod:`di.workers` keeps API code free of # direct ``services.workers`` imports. - ensure_worker_bootstrap() + logger.info("Startup: initializing worker bootstrap") + ensure_worker_bootstrap(settings) + logger.info("Startup: worker bootstrap initialized") container: ServiceContainer | None try: - container = ServiceContainer(config) + logger.info("Startup: wiring ServiceContainer") + container = ServiceContainer(settings) + logger.info("Startup: ServiceContainer wired") except Exception: # pragma: no cover - defensive boot guard logger.exception("ServiceContainer wiring skipped") container = None @@ -155,7 +169,13 @@ async def lifespan(app: FastAPI): if container is not None: try: - await container.initialize() + logger.info("Startup: initializing ServiceContainer", timeout=CONTAINER_STARTUP_TIMEOUT) + await asyncio.wait_for(container.initialize(), timeout=CONTAINER_STARTUP_TIMEOUT) + logger.info("Startup: ServiceContainer initialized") + except TimeoutError: # pragma: no cover - defensive boot guard + logger.exception("ServiceContainer.initialize timed out; serving degraded (503)") + app.state.container = None + container = None except Exception: # pragma: no cover - defensive boot guard # A half-initialised container (asyncpg pool never opened) # would route requests into broken repos and 500. Drop it so @@ -165,10 +185,20 @@ async def lifespan(app: FastAPI): container = None try: - await prime_max_model_tokens() + logger.info("Startup: priming max model token cache") + await prime_max_model_tokens(settings) + logger.info("Startup: max model token cache primed") except Exception: # pragma: no cover - defensive cache guard logger.exception("max_model_tokens cache priming failed; falling back to config") + # Mirror the resolved boot state into the process-level singleton so + # callers without a request (Chainlit, background tasks) resolve the + # same container the HTTP routes get via request.app.state. None on a + # degraded boot keeps di.providers serving the intended 503. + logger.info("Startup: registering process container", available=getattr(app.state, "container", None) is not None) + set_container(getattr(app.state, "container", None)) + logger.info("Startup: complete") + try: yield finally: @@ -178,6 +208,7 @@ async def lifespan(app: FastAPI): await live.shutdown() except Exception: # pragma: no cover - defensive shutdown guard logger.exception("ServiceContainer.shutdown skipped") + set_container(None) # --------------------------------------------------------------------------- @@ -266,7 +297,7 @@ def root_redirect(): @app.get("/config", summary="Get current configuration", tags=["Configuration"], dependencies=[Depends(require_admin)]) def get_config(): - return config + return settings # Router mounts. Phase 10F finished moving these into @@ -301,20 +332,20 @@ def get_config(): if __name__ == "__main__": - if config.ray.serve.enable: + if settings.ray.serve.enable: from ray import serve - @serve.deployment(num_replicas=config.ray.serve.num_replicas) + @serve.deployment(num_replicas=settings.ray.serve.num_replicas) @serve.ingress(app) class OpenRagAPI: pass - serve.start(http_options={"host": config.ray.serve.host, "port": config.ray.serve.port}) + serve.start(http_options={"host": settings.ray.serve.host, "port": settings.ray.serve.port}) if WITH_CHAINLIT_UI: from chainlit_api import app as chainlit_app serve.run(OpenRagAPI.bind(), route_prefix="/") - uvicorn.run(chainlit_app, host="0.0.0.0", port=config.ray.serve.chainlit_port) + uvicorn.run(chainlit_app, host="0.0.0.0", port=settings.ray.serve.chainlit_port) else: serve.run(OpenRagAPI.bind(), route_prefix="/", blocking=True) diff --git a/openrag/api/routers/admin/indexing.py b/openrag/api/routers/admin/indexing.py index b8e59f0ee..3a70b9839 100644 --- a/openrag/api/routers/admin/indexing.py +++ b/openrag/api/routers/admin/indexing.py @@ -29,8 +29,7 @@ ) from api.routers.admin.task_logs import collect_task_logs from components.indexer.utils.files import sanitize_filename, save_file_to_disk -from config import load_config -from di.providers import get_auth_service, get_indexing_service, get_partition_service +from di.providers import get_auth_service, get_config, get_indexing_service, get_partition_service from fastapi import ( APIRouter, Depends, @@ -46,22 +45,12 @@ logger = get_logger() -config = load_config() -DATA_DIR = config.paths.data_dir -LOG_FILE = Path(config.paths.log_dir or "logs") / "app.json" -# supported file formats or mimetypes -ACCEPTED_FILE_FORMATS = config.loader.file_loaders.model_dump().keys() -DICT_MIMETYPES = config.loader.mimetypes.to_dict() - -PREFERRED_URL_SCHEME = config.server.preferred_url_scheme - - -def build_url(request: Request, route_name: str, **path_params) -> str: +def build_url(request: Request, route_name: str, *, preferred_url_scheme: str | None = None, **path_params) -> str: """Build a URL using the preferred scheme if configured.""" url = request.url_for(route_name, **path_params) - if PREFERRED_URL_SCHEME: - url = url.replace(scheme=PREFERRED_URL_SCHEME) + if preferred_url_scheme: + url = url.replace(scheme=preferred_url_scheme) return str(url) @@ -76,7 +65,7 @@ def build_url(request: Request, route_name: str, **path_params) -> str: Returns a list of supported file extensions and MIME types that can be indexed by the system. """, ) -async def get_supported_types(): +async def get_supported_types(config=Depends(get_config)): """ Get a list of supported types for indexing. @@ -85,7 +74,9 @@ async def get_supported_types(): - `extensions`: List of supported file extensions. - `mimetypes`: List of supported MIME types. """ - resp = {"extensions": list(ACCEPTED_FILE_FORMATS), "mimetypes": list(DICT_MIMETYPES)} + accepted_file_formats = config.loader.file_loaders.model_dump().keys() + mimetypes = config.loader.mimetypes.to_dict() + resp = {"extensions": list(accepted_file_formats), "mimetypes": list(mimetypes)} return JSONResponse(content=resp) @@ -131,6 +122,7 @@ async def add_file( workspace_ids: str | None = Form(None, description="JSON array of workspace IDs to add the file to"), user=Depends(require_partition_editor), _quota_check=Depends(check_user_file_quota), + config=Depends(get_config), service=Depends(get_indexing_service), ): if await service.file_exists(file_id, partition): @@ -142,7 +134,7 @@ async def add_file( original_filename = file.filename file.filename = sanitize_filename(file.filename) try: - file_path = await save_file_to_disk(file, Path(DATA_DIR), with_random_prefix=True) + file_path = await save_file_to_disk(file, Path(config.paths.data_dir), with_random_prefix=True) except Exception as e: logger.exception("Failed to save file to disk.", error=str(e)) raise HTTPException( @@ -182,7 +174,14 @@ async def add_file( return JSONResponse( status_code=status.HTTP_201_CREATED, - content={"task_status_url": build_url(request, "get_task_status", task_id=task_id)}, + content={ + "task_status_url": build_url( + request, + "get_task_status", + preferred_url_scheme=config.server.preferred_url_scheme, + task_id=task_id, + ) + }, ) @@ -254,6 +253,7 @@ async def put_file( file: UploadFile = Depends(validate_file_format), metadata: dict = Depends(validate_metadata), user=Depends(require_partition_editor), + config=Depends(get_config), service=Depends(get_indexing_service), ): if not await service.file_exists(file_id, partition): @@ -267,7 +267,7 @@ async def put_file( # then deletes old ones — so the file is never left in a half-replaced state. original_filename = file.filename file.filename = sanitize_filename(file.filename) - file_path = await save_file_to_disk(file, Path(DATA_DIR), with_random_prefix=True) + file_path = await save_file_to_disk(file, Path(config.paths.data_dir), with_random_prefix=True) task_id = await service.add_file( file_path=str(file_path), @@ -282,7 +282,14 @@ async def put_file( return JSONResponse( status_code=status.HTTP_202_ACCEPTED, - content={"task_status_url": build_url(request, "get_task_status", task_id=task_id)}, + content={ + "task_status_url": build_url( + request, + "get_task_status", + preferred_url_scheme=config.server.preferred_url_scheme, + task_id=task_id, + ) + }, ) @@ -407,6 +414,7 @@ async def get_task_status( request: Request, task_id: str, task_details=Depends(require_task_owner), + config=Depends(get_config), service=Depends(get_indexing_service), ): state = await service.get_task_state(task_id) @@ -423,7 +431,12 @@ async def get_task_status( } if state == "FAILED": - content["error_url"] = build_url(request, "get_task_error", task_id=task_id) + content["error_url"] = build_url( + request, + "get_task_error", + preferred_url_scheme=config.server.preferred_url_scheme, + task_id=task_id, + ) return JSONResponse(status_code=status.HTTP_200_OK, content=content) @@ -473,12 +486,18 @@ async def get_task_error( **Note:** Logs are returned in chronological order (oldest first). """, ) -async def get_task_logs(task_id: str, max_lines: int = 100, task_details=Depends(require_task_owner)): - if not LOG_FILE.exists(): +async def get_task_logs( + task_id: str, + max_lines: int = 100, + task_details=Depends(require_task_owner), + config=Depends(get_config), +): + log_file = Path(config.paths.log_dir or "logs") / "app.json" + if not log_file.exists(): raise HTTPException(status_code=500, detail="Log file not found.") try: - logs = collect_task_logs(LOG_FILE, task_id, max_lines) + logs = collect_task_logs(log_file, task_id, max_lines) except ValueError as exc: raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, diff --git a/openrag/api/routers/admin/tools.py b/openrag/api/routers/admin/tools.py index 2ff97fb0b..e00222ea5 100644 --- a/openrag/api/routers/admin/tools.py +++ b/openrag/api/routers/admin/tools.py @@ -18,15 +18,12 @@ ) from api.schemas.admin.tools import ToolInfo from components.indexer.utils.files import save_file_to_disk -from config import load_config -from di.providers import get_conversion_service +from di.providers import get_config, get_conversion_service from fastapi import APIRouter, Depends, Form, HTTPException, UploadFile, status from fastapi.responses import JSONResponse from utils.logger import get_logger logger = get_logger() -config = load_config() -data_dir = config.paths.data_dir router = APIRouter() @@ -96,12 +93,13 @@ async def execute_tool( file: UploadFile = Depends(validate_file_format), tool: str = Depends(validate_tool), metadata: dict = Depends(validate_metadata), + config=Depends(get_config), service=Depends(get_conversion_service), ): file_path = None try: if tool["name"] == "extractText": - file_path = await save_file_to_disk(file, Path(data_dir), with_random_prefix=True) + file_path = await save_file_to_disk(file, Path(config.paths.data_dir), with_random_prefix=True) logger.debug(f"Execute tool extractText with file {file.filename}") sanitized_content = await service.serialize_file( diff --git a/openrag/api/routers/user/chat.py b/openrag/api/routers/user/chat.py index 2f3ba77b4..fd29157c8 100644 --- a/openrag/api/routers/user/chat.py +++ b/openrag/api/routers/user/chat.py @@ -13,6 +13,7 @@ import asyncio import json +from typing import TYPE_CHECKING from urllib.parse import urlparse import consts @@ -32,23 +33,29 @@ from components.indexer.utils.text_sanitizer import sanitize_text from components.utils import get_num_tokens from config import load_config -from di.providers import get_partition_service, get_query_service +from di.providers import get_config, get_partition_service, get_query_service from fastapi import APIRouter, Body, Depends, HTTPException, Request, status from fastapi.responses import JSONResponse, StreamingResponse from utils.exceptions.base import OpenRAGError from utils.logger import get_logger logger = get_logger() -config = load_config() router = APIRouter() +if TYPE_CHECKING: + from core.config.root import Settings + # Cached max model token limit. Populated by ``prime_max_model_tokens`` # which the application lifespan invokes during startup; ``get_max_model_tokens`` # falls back to ``config.llm_context.max_llm_context_size`` until then. _max_model_tokens: int | None = None -async def prime_max_model_tokens() -> None: +def _runtime_config(settings: "Settings | None" = None) -> "Settings": + return settings if settings is not None else load_config() + + +async def prime_max_model_tokens(settings: "Settings | None" = None) -> None: """Populate the cached max model token limit. Called once from the FastAPI lifespan in ``api/main.py`` (replaces the @@ -56,7 +63,7 @@ async def prime_max_model_tokens() -> None: it just refreshes the cache. """ global _max_model_tokens - _max_model_tokens = await _fetch_max_model_tokens() + _max_model_tokens = await _fetch_max_model_tokens(_runtime_config(settings)) def _make_sse_error(message: str, code: str) -> str: @@ -134,12 +141,14 @@ def chunk_url(extract_id) -> str: def is_direct_llm_model( request: OpenAIChatCompletionRequest | OpenAICompletionRequest, + settings: "Settings | None" = None, ) -> bool: """True if the request should use the LLM directly (no RAG partition).""" + config = _runtime_config(settings) return request.model is None or request.model == "" or request.model == config.llm.model -async def _fetch_max_model_tokens() -> int: +async def _fetch_max_model_tokens(config: "Settings") -> int: """Fetch the max model token limit from vLLM's OpenAI server. Falls back to ``config.llm_context.max_llm_context_size`` if unavailable. @@ -168,15 +177,18 @@ def get_max_model_tokens() -> int: """Return the cached max model token limit (populated at startup).""" if _max_model_tokens is not None: return _max_model_tokens + config = _runtime_config() return int(config.llm_context.max_llm_context_size) def validate_tokens_limit( request: OpenAIChatCompletionRequest | OpenAICompletionRequest, max_tokens_allowed: int, + settings: "Settings | None" = None, ) -> tuple[bool, str]: """Validate if the request respects the maximum token limit.""" try: + config = _runtime_config(settings) _length_function = get_num_tokens() if isinstance(request, OpenAIChatCompletionRequest): @@ -216,9 +228,14 @@ def validate_tokens_limit( def check_tokens_limit( request: OpenAIChatCompletionRequest | OpenAICompletionRequest, log, + settings: "Settings | None" = None, ): """Validate token limit and raise HTTPException(413) if exceeded.""" - is_valid, error_message = validate_tokens_limit(request, max_tokens_allowed=get_max_model_tokens()) + is_valid, error_message = validate_tokens_limit( + request, + max_tokens_allowed=get_max_model_tokens(), + settings=settings, + ) if not is_valid: log.info("Request exceeds token limit", detail=error_message) raise HTTPException( @@ -260,6 +277,7 @@ async def openai_chat_completion( _: None = Depends(check_llm_model_availability), service=Depends(get_query_service), partition_service=Depends(get_partition_service), + config=Depends(get_config), ): model_name = request.model or config.llm.model log = logger.bind(model=model_name, endpoint="/chat/completions") @@ -273,8 +291,8 @@ async def openai_chat_completion( log.debug("Received chat completion request with messages: {}", truncate(str(request.messages))) - if is_direct_llm_model(request): - check_tokens_limit(request, log) + if is_direct_llm_model(request, config): + check_tokens_limit(request, log, config) partitions = None else: partitions = await get_partition_name( @@ -352,6 +370,7 @@ async def openai_completion( _: None = Depends(check_llm_model_availability), service=Depends(get_query_service), partition_service=Depends(get_partition_service), + config=Depends(get_config), ): model_name = request.model or config.llm.model log = logger.bind(model=model_name, endpoint="/completions") @@ -367,8 +386,8 @@ async def openai_completion( detail="Streaming is not supported for this endpoint", ) - if is_direct_llm_model(request): - check_tokens_limit(request, log) + if is_direct_llm_model(request, config): + check_tokens_limit(request, log, config) partitions = None else: partitions = await get_partition_name( diff --git a/openrag/api/schemas/user/chat.py b/openrag/api/schemas/user/chat.py index 91cc5a055..f90e67a9e 100644 --- a/openrag/api/schemas/user/chat.py +++ b/openrag/api/schemas/user/chat.py @@ -1,10 +1,12 @@ from typing import Any, Literal -from config import load_config from pydantic import BaseModel, Field -config = load_config() -default_max_tokens = config.llm_context.max_output_tokens + +def default_max_tokens(): + from config import load_config + + return load_config().llm_context.max_output_tokens class OpenAIMessage(BaseModel): @@ -18,7 +20,7 @@ class OpenAIChatCompletionRequest(BaseModel): temperature: float | None = Field(0.3) top_p: float | None = Field(1.0) stream: bool | None = Field(False) - max_tokens: int | None = Field(default_max_tokens) + max_tokens: int | None = Field(default_factory=default_max_tokens) logprobs: int | None = Field(None) metadata: dict[str, Any] | None = Field( { @@ -39,7 +41,7 @@ class OpenAICompletionRequest(BaseModel): frequency_penalty: float | None = Field(0.0) logit_bias: dict | None = Field(None) logprobs: int | None = Field(None) - max_tokens: int | None = Field(default_max_tokens) + max_tokens: int | None = Field(default_factory=default_max_tokens) n: int | None = Field(1) presence_penalty: float | None = Field(0.0) seed: int | None = Field(None) diff --git a/openrag/api/test_main_proxy_headers.py b/openrag/api/test_main_proxy_headers.py index 1e3d975c9..4a09fa232 100644 --- a/openrag/api/test_main_proxy_headers.py +++ b/openrag/api/test_main_proxy_headers.py @@ -10,7 +10,10 @@ """ import ast +import importlib import os +import sys +from types import ModuleType _MAIN_PATH = os.path.dirname(__file__) + "/main.py" @@ -59,3 +62,15 @@ def test_default_forwarded_allow_ips_env_var_used(): with open(_MAIN_PATH) as f: src = f.read() assert "UVICORN_FORWARDED_ALLOW_IPS" in src + + +def test_api_package_exports_app_for_legacy_uvicorn_path(monkeypatch): + """Older images or overrides may still run ``uvicorn api:app``.""" + fake_app = object() + fake_main = ModuleType("api.main") + fake_main.app = fake_app + monkeypatch.setitem(sys.modules, "api.main", fake_main) + + api_module = importlib.import_module("api") + + assert api_module.app is fake_app diff --git a/openrag/components/indexer/chunker/chunker.py b/openrag/components/indexer/chunker/chunker.py index f4b794889..90c1e238f 100644 --- a/openrag/components/indexer/chunker/chunker.py +++ b/openrag/components/indexer/chunker/chunker.py @@ -30,10 +30,6 @@ from langchain_openai import ChatOpenAI logger = get_logger() -config = load_config() - -CONTEXTUALIZATION_TIMEOUT = config.chunker.contextualization_timeout -MAX_CONCURRENT_CONTEXTUALIZATION = config.chunker.max_concurrent_contextualization class _LangChainLLMAdapter(_CoreLLM): @@ -109,12 +105,14 @@ def __init__( self.contextual_retrieval = contextual_retrieval if contextual_retrieval: - _lc_llm = ChatOpenAI(**{**llm_config, "timeout": CONTEXTUALIZATION_TIMEOUT}) + config = load_config() + contextualization_timeout = config.chunker.contextualization_timeout + _lc_llm = ChatOpenAI(**{**llm_config, "timeout": contextualization_timeout}) self.contextualizer: _CoreChunkContextualizer | None = _CoreChunkContextualizer( llm=_LangChainLLMAdapter(_lc_llm), system_prompt=CHUNK_CONTEXTUALIZER_PROMPT, - timeout_seconds=CONTEXTUALIZATION_TIMEOUT, - max_concurrent=MAX_CONCURRENT_CONTEXTUALIZATION, + timeout_seconds=contextualization_timeout, + max_concurrent=config.chunker.max_concurrent_contextualization, semaphore=get_vlm_semaphore(), ) else: diff --git a/openrag/components/indexer/loaders/base.py b/openrag/components/indexer/loaders/base.py index 647e06fd8..ccc91ebac 100644 --- a/openrag/components/indexer/loaders/base.py +++ b/openrag/components/indexer/loaders/base.py @@ -28,7 +28,6 @@ from utils.logger import get_logger logger = get_logger() -config = load_config() class BaseLoader(ABC): @@ -41,7 +40,7 @@ class BaseLoader(ABC): def __init__(self, **kwargs) -> None: self.page_sep = "[PAGE_SEP]" - self.config = kwargs.get("config") + self.config = kwargs.get("config") or load_config() settings: dict = self.config.vlm.model_dump() model_settings = { "temperature": 0.2, diff --git a/openrag/components/prompts/prompts.py b/openrag/components/prompts/prompts.py index 9dc113c93..45c5e513b 100644 --- a/openrag/components/prompts/prompts.py +++ b/openrag/components/prompts/prompts.py @@ -16,17 +16,15 @@ from config import load_config from core.prompts.template_loader import load_template_by_key -config = load_config() - -prompts_dir: Path = config.paths.prompts_dir -prompt_mapping = config.prompts - def load_prompt( prompt_name: str, - prompts_dir: Path = prompts_dir, - prompt_mapping=prompt_mapping, + prompts_dir: Path | None = None, + prompt_mapping=None, ) -> str: + config = load_config() + prompts_dir = prompts_dir or config.paths.prompts_dir + prompt_mapping = prompt_mapping or config.prompts return load_template_by_key(prompts_dir, prompt_mapping, prompt_name) diff --git a/openrag/components/utils.py b/openrag/components/utils.py index f03a9f4d5..16529f6af 100644 --- a/openrag/components/utils.py +++ b/openrag/components/utils.py @@ -16,8 +16,6 @@ SOURCE_SEPARATOR = "-" * 10 + "\n\n" -# Global variables -config = load_config() logger = get_logger() @@ -43,6 +41,7 @@ def get_num_tokens(): try: from langchain_openai import ChatOpenAI + config = load_config() llm = ChatOpenAI(**config.llm.model_dump()) _cached_length_function = llm.get_num_tokens except Exception as exc: @@ -304,6 +303,7 @@ def detect_language(text: str): def get_llm_semaphore() -> DistributedSemaphore: + config = load_config() return DistributedSemaphore( name="llmSemaphore", max_concurrent_ops=config.semaphore.llm_semaphore, @@ -311,6 +311,7 @@ def get_llm_semaphore() -> DistributedSemaphore: def get_vlm_semaphore() -> DistributedSemaphore: + config = load_config() return DistributedSemaphore( name="vlmSemaphore", max_concurrent_ops=config.semaphore.vlm_semaphore, @@ -318,12 +319,8 @@ def get_vlm_semaphore() -> DistributedSemaphore: def get_audio_semaphore() -> DistributedSemaphore: + config = load_config() return DistributedSemaphore( name="audioSemaphore", max_concurrent_ops=config.loader.transcriber.max_concurrent_chunks, ) - - -get_llm_semaphore() -get_vlm_semaphore() -get_audio_semaphore() diff --git a/openrag/di/container.py b/openrag/di/container.py index 5b6186958..d48622b50 100644 --- a/openrag/di/container.py +++ b/openrag/di/container.py @@ -21,7 +21,7 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from core.embeddings import embedder_registry from core.llm import llm_registry @@ -33,6 +33,7 @@ from di.rerankers import register_rerankers from di.vector_stores import create_vector_store from di.vlms import register_vlms +from utils.logger import get_logger if TYPE_CHECKING: from core.config.root import Settings @@ -65,6 +66,8 @@ from services.orchestrators.workspace_service import WorkspaceService +logger = get_logger() + _NO_SETTINGS_MESSAGE = ( "ServiceContainer was constructed without a Settings instance — " "pass Settings to ServiceContainer(...) to wire storage adapters." @@ -90,6 +93,8 @@ def __init__(self, settings: Settings | None = None) -> None: self._oidc_config = OIDCConfig.from_env() self._settings = settings + self._initialized = False + self._inference_clients: list[Any] = [] self._catalog_store: CatalogStore | None = create_catalog_store(settings) if settings is not None else None self._vector_store: VectorStore | None = create_vector_store(settings) if settings is not None else None self._auth_service: AuthService | None = None @@ -121,13 +126,42 @@ def _require_settings(self) -> Settings: async def initialize(self) -> None: """Open the storage adapters (asyncpg pool + Alembic migrations).""" if self._catalog_store is not None: + logger.info("ServiceContainer.initialize: initializing catalog store") await self._catalog_store.initialize() + logger.info("ServiceContainer.initialize: ensuring admin user") await self.user_repo.ensure_admin_user(os.getenv("AUTH_TOKEN")) + logger.info("ServiceContainer.initialize: admin user ready") + self._initialized = True async def shutdown(self) -> None: - """Close the storage adapters cleanly.""" - if self._catalog_store is not None: - await self._catalog_store.shutdown() + """Close inference clients and storage adapters cleanly. + + Best-effort: a failure closing one client must not skip the + remaining clients, the database pool, or the state reset. + """ + try: + for client in self._inference_clients: + aclose = getattr(client, "aclose", None) + if aclose is not None: + try: + await aclose() + except Exception: + logger.exception("Failed to close inference client") + if self._catalog_store is not None: + await self._catalog_store.shutdown() + finally: + self._inference_clients.clear() + self._initialized = False + + @property + def is_initialized(self) -> bool: + """True once :meth:`initialize` has completed its async I/O.""" + return self._initialized + + @property + def config(self) -> Settings: + """The root settings this container was wired from.""" + return self._require_settings() # ------------------------------------------------------------------ # Storage adapters @@ -443,18 +477,23 @@ def conversion_service(self) -> ConversionService: # Registry-based inference factories (Phase 6) # ------------------------------------------------------------------ - @staticmethod - def create_embedder(name: str = "vllm", **kwargs): - return embedder_registry.create(name, **kwargs) + def create_embedder(self, name: str = "vllm", **kwargs): + """Build an embedder client, tracking it for shutdown cleanup.""" + return self._track(embedder_registry.create(name, **kwargs)) + + def create_llm(self, name: str = "vllm", **kwargs): + """Build an LLM client, tracking it for shutdown cleanup.""" + return self._track(llm_registry.create(name, **kwargs)) - @staticmethod - def create_llm(name: str = "vllm", **kwargs): - return llm_registry.create(name, **kwargs) + def create_reranker(self, name: str = "infinity", **kwargs): + """Build a reranker client, tracking it for shutdown cleanup.""" + return self._track(reranker_registry.create(name, **kwargs)) - @staticmethod - def create_reranker(name: str = "infinity", **kwargs): - return reranker_registry.create(name, **kwargs) + def create_vlm(self, name: str = "vllm", **kwargs): + """Build a VLM client, tracking it for shutdown cleanup.""" + return self._track(vlm_registry.create(name, **kwargs)) - @staticmethod - def create_vlm(name: str = "vllm", **kwargs): - return vlm_registry.create(name, **kwargs) + def _track(self, client: Any) -> Any: + """Register a built inference client so :meth:`shutdown` can close it.""" + self._inference_clients.append(client) + return client diff --git a/openrag/di/embedders.py b/openrag/di/embedders.py index 5f09ecda2..d78a87624 100644 --- a/openrag/di/embedders.py +++ b/openrag/di/embedders.py @@ -2,4 +2,6 @@ def register_embedders() -> None: + """Import embedder implementations so they register with the core registry.""" + import services.inference.ollama_client # noqa: F401 import services.inference.vllm_client # noqa: F401 diff --git a/openrag/di/factories.py b/openrag/di/factories.py new file mode 100644 index 000000000..507c4e24d --- /dev/null +++ b/openrag/di/factories.py @@ -0,0 +1,107 @@ +"""Generic cached component factory — bridges model config to registries. + +Phase 11A. :func:`make_component_factory` is the one pattern the composition +root reuses for all four inference component kinds (embedder, reranker, LLM, +VLM). Given a model name it looks up the matching entry in a config section, +resolves which registered implementation to build, instantiates it through the +:class:`~core.utils.registry.Registry`, and caches the result so subsequent +calls reuse the same client (and its underlying httpx connection pool). + +The factory returns a ``(factory_fn, cache)`` tuple. The cache dict is exposed +deliberately: + +* **Shutdown.** The container appends every cache to a shared + ``client_caches`` list and, on teardown, calls ``aclose()`` on each cached + instance to release httpx connections. +* **Invalidation.** Post-refactoring, when a model endpoint is renamed or + deleted at runtime, the owning service pops the stale entry + (``cache.pop(old_name, None)``) so the next call rebuilds against the new + config. + +``config_section`` is typed against the :class:`ModelEndpointConfig` protocol +below — the unified per-endpoint config shape that lands with the DB-backed +model registry. Until that config exists the protocol documents the contract +the factory depends on; any object exposing those attributes (including a test +double) is accepted. +""" + +from __future__ import annotations + +import threading +from collections.abc import Callable, Mapping +from typing import TYPE_CHECKING, Any, Protocol, TypeVar + +if TYPE_CHECKING: + from core.utils.registry import Registry + +T = TypeVar("T") + + +class ModelEndpointConfig(Protocol): + """Structural shape :func:`make_component_factory` reads off each entry. + + ``extra`` carries implementation-specific keyword arguments plus an + optional ``implementation`` control key that selects which registered + class to build (falling back to ``default_impl`` when absent). + """ + + endpoint: str + model_name: str + batch_size: int + timeout: float + extra: Mapping[str, Any] + + +def make_component_factory( + registry: Registry[T], + config_section: Mapping[str, ModelEndpointConfig], + default_impl: str, + client_caches: list[dict[str, T]], + extra_kwargs_fn: Callable[[ModelEndpointConfig], Mapping[str, Any]] | None = None, +) -> tuple[Callable[[str], T], dict[str, T]]: + """Build a cached ``(name) -> T`` factory from a registry and config section. + + Returns ``(factory_fn, cache)``. The cache is appended to ``client_caches`` + so the container can close every built client on shutdown. Instances are + created lazily on first request and reused thereafter; construction is + guarded by double-checked locking so concurrent first calls for the same + name build exactly one instance. + + ``extra_kwargs_fn``, when given, computes additional constructor kwargs + from the config entry (merged last, so it wins on key collisions) — the + seam for kwargs a kind needs but the unified config does not carry. + """ + cache: dict[str, T] = {} + lock = threading.Lock() + client_caches.append(cache) + + def factory(name: str = "default") -> T: + if name in cache: + return cache[name] + with lock: + if name in cache: + return cache[name] + model_cfg = config_section.get(name) + if model_cfg is None: + raise KeyError(f"Unknown model '{name}'. Available: {list(config_section)}") + # `implementation` is a control key (which class to build), not a + # constructor argument, so it is read out before splatting `extra`. + impl_kwargs = {k: v for k, v in model_cfg.extra.items() if k != "implementation"} + impl = model_cfg.extra.get("implementation", default_impl) + kwargs: dict[str, Any] = { + "endpoint": model_cfg.endpoint, + "model_name": model_cfg.model_name, + "batch_size": model_cfg.batch_size, + "timeout": model_cfg.timeout, + **impl_kwargs, + } + if extra_kwargs_fn is not None: + kwargs.update(extra_kwargs_fn(model_cfg)) + instance = registry.create(impl, **kwargs) + cache[name] = instance + return instance + + return factory, cache + + +__all__ = ["ModelEndpointConfig", "make_component_factory"] diff --git a/openrag/di/llms.py b/openrag/di/llms.py index 8620f7159..711d733c9 100644 --- a/openrag/di/llms.py +++ b/openrag/di/llms.py @@ -2,4 +2,6 @@ def register_llms() -> None: + """Import LLM implementations so they register with the core registry.""" + import services.inference.ollama_client # noqa: F401 import services.inference.vllm_client # noqa: F401 diff --git a/openrag/di/parsers.py b/openrag/di/parsers.py new file mode 100644 index 000000000..d33231595 --- /dev/null +++ b/openrag/di/parsers.py @@ -0,0 +1,19 @@ +"""Register document parser implementations with the core registry.""" + + +def register_parsers() -> None: + """Import parser implementations so they register with the core registry.""" + import core.indexing.parsers.audio.client_based # noqa: F401 + import core.indexing.parsers.audio.local_whisper # noqa: F401 + import core.indexing.parsers.doc_parser # noqa: F401 + import core.indexing.parsers.docx_parser # noqa: F401 + import core.indexing.parsers.eml_parser # noqa: F401 + import core.indexing.parsers.html_parser # noqa: F401 + import core.indexing.parsers.image_parser # noqa: F401 + import core.indexing.parsers.markdown_parser # noqa: F401 + import core.indexing.parsers.pdf.client_based # noqa: F401 + import core.indexing.parsers.pdf.docling # noqa: F401 + import core.indexing.parsers.pdf.marker # noqa: F401 + import core.indexing.parsers.pdf.pymupdf # noqa: F401 + import core.indexing.parsers.pptx_parser # noqa: F401 + import core.indexing.parsers.text_parser # noqa: F401 diff --git a/openrag/di/providers.py b/openrag/di/providers.py index fecd16e2c..1e6300a51 100644 --- a/openrag/di/providers.py +++ b/openrag/di/providers.py @@ -1,16 +1,16 @@ """FastAPI dependency providers. -Thin accessors over the request-scoped :class:`ServiceContainer` that -``main.py`` attaches at ``app.state.container``. Phase 8 keeps these as -one-liners — the container (``di/container.py``) is the composition -root. Phase 11 moves the attachment into a proper FastAPI lifespan and -wires ``container.initialize()``; until then the OIDC flow that needs -the asyncpg pool is dormant (token-mode auth routes already short-circuit -before reaching a service). +This module is the API layer's bridge to the composition root. During +the Phase 10 -> 11 transition it supports both access patterns: + +* request-scoped lookup from ``request.app.state.container``; +* process-level override via :func:`set_container` for tests and the + Phase 11 lifespan. """ from __future__ import annotations +import threading from typing import TYPE_CHECKING from fastapi import HTTPException, Request, status @@ -27,48 +27,117 @@ from services.orchestrators.user_service import UserService from services.orchestrators.workspace_service import WorkspaceService +_container: ServiceContainer | None = None +_container_lock = threading.Lock() + + +def set_container(container: ServiceContainer | None) -> None: + """Override the process-level container. + + Tests use this to install a fake container. The app lifespan can use + it after constructing the real container. Passing ``None`` clears the + override. + """ + global _container + with _container_lock: + _container = container -def get_container(request: Request) -> ServiceContainer: - container = getattr(request.app.state, "container", None) + +def get_container(request: Request = None) -> ServiceContainer: + """Resolve the active service container from request state or process state.""" + if request is not None and hasattr(request.app.state, "container"): + container = request.app.state.container + if container is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Service container is not available.", + ) + return container + + global _container + container = _container if container is None: + with _container_lock: + if _container is None: + from di.container import ServiceContainer + + _container = ServiceContainer() + container = _container + return container + + +def _require_initialized(request: Request = None) -> ServiceContainer: + """Return the active container only after its initialization guard passes.""" + container = get_container(request) + if not getattr(container, "is_initialized", True): raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, - detail="Service container is not available.", + detail="Service container has not been initialized.", ) return container -def get_auth_service(request: Request) -> AuthService: - return get_container(request).auth_service +def get_auth_service(request: Request = None) -> AuthService: + """Resolve the authentication orchestrator from the active container.""" + return _require_initialized(request).auth_service + + +def get_user_service(request: Request = None) -> UserService: + """Resolve the user orchestrator from the active container.""" + return _require_initialized(request).user_service + + +def get_partition_service(request: Request = None) -> PartitionService: + """Resolve the partition orchestrator from the active container.""" + return _require_initialized(request).partition_service -def get_user_service(request: Request) -> UserService: - return get_container(request).user_service +def get_workspace_service(request: Request = None) -> WorkspaceService: + """Resolve the workspace orchestrator from the active container.""" + return _require_initialized(request).workspace_service -def get_partition_service(request: Request) -> PartitionService: - return get_container(request).partition_service +def get_retrieval_service(request: Request = None) -> RetrievalService: + """Resolve the retrieval orchestrator from the active container.""" + return _require_initialized(request).retrieval_service -def get_workspace_service(request: Request) -> WorkspaceService: - return get_container(request).workspace_service +def get_query_service(request: Request = None) -> QueryService: + """Resolve the query orchestrator from the active container.""" + return _require_initialized(request).query_service -def get_retrieval_service(request: Request) -> RetrievalService: - return get_container(request).retrieval_service +def get_indexing_service(request: Request = None) -> IndexingService: + """Resolve the indexing orchestrator from the active container.""" + return _require_initialized(request).indexing_service -def get_query_service(request: Request) -> QueryService: - return get_container(request).query_service +def get_job_service(request: Request = None) -> JobService: + """Resolve the job orchestrator from the active container.""" + return _require_initialized(request).job_service -def get_indexing_service(request: Request) -> IndexingService: - return get_container(request).indexing_service +def get_conversion_service(request: Request = None) -> ConversionService: + """Resolve the conversion orchestrator from the active container.""" + return _require_initialized(request).conversion_service -def get_job_service(request: Request) -> JobService: - return get_container(request).job_service +def get_config(request: Request = None): + """Resolve application configuration from the active container.""" + return _require_initialized(request).config -def get_conversion_service(request: Request) -> ConversionService: - return get_container(request).conversion_service +__all__ = [ + "get_auth_service", + "get_config", + "get_container", + "get_conversion_service", + "get_indexing_service", + "get_job_service", + "get_partition_service", + "get_query_service", + "get_retrieval_service", + "get_user_service", + "get_workspace_service", + "set_container", +] diff --git a/openrag/di/test_container.py b/openrag/di/test_container.py index d9a9e7275..95678f5f0 100644 --- a/openrag/di/test_container.py +++ b/openrag/di/test_container.py @@ -26,9 +26,11 @@ from di.container import ServiceContainer from di.repositories import create_catalog_store from di.vector_stores import create_vector_store +from fastapi import HTTPException def _settings(database: str | None = None, collection: str = "vdb_test") -> Settings: + """Build minimal settings for container wiring tests.""" return Settings( rdb=RDBConfig(password="x", database=database), vectordb=VectorDBConfig(collection_name=collection), @@ -57,14 +59,17 @@ class TestLegacyContainerStillWorks: """The pre-Phase-7E callers do ``ServiceContainer()`` with no settings.""" def test_constructs_without_settings(self): + """Keep the legacy no-argument container constructor available.""" ServiceContainer() # must not raise def test_catalog_store_raises_when_unconfigured(self): + """Reject catalog store access when settings were not provided.""" c = ServiceContainer() with pytest.raises(RuntimeError, match="without a Settings instance"): _ = c.catalog_store def test_vector_store_raises_when_unconfigured(self): + """Reject vector store access when settings were not provided.""" c = ServiceContainer() with pytest.raises(RuntimeError, match="without a Settings instance"): _ = c.vector_store @@ -90,21 +95,27 @@ def test_vector_store_raises_when_unconfigured(self): ], ) def test_repo_properties_raise_when_unconfigured(self, name): + """Reject repository shortcuts when settings were not provided.""" c = ServiceContainer() with pytest.raises(RuntimeError, match="without a Settings instance"): getattr(c, name) class TestCatalogStoreWiring: + """Verify relational store wiring exposed by the service container.""" + def test_catalog_store_satisfies_port(self): + """Expose the catalog store through the expected core port.""" c = ServiceContainer(_settings()) assert isinstance(c.catalog_store, CatalogStore) def test_database_name_derived_from_collection(self): + """Derive the relational database name from the vector collection.""" c = ServiceContainer(_settings(database=None, collection="my_collection")) assert c.catalog_store._conn._conn_kwargs["database"] == "partitions_for_collection_my_collection" def test_explicit_database_overrides_fallback(self): + """Prefer an explicit relational database over the derived fallback.""" c = ServiceContainer(_settings(database="custom_db", collection="my_collection")) assert c.catalog_store._conn._conn_kwargs["database"] == "custom_db" @@ -129,6 +140,7 @@ def test_explicit_database_overrides_fallback(self): ], ) def test_repo_property_returns_port_typed_instance(self, name, port): + """Expose each repository shortcut as the matching core port.""" c = ServiceContainer(_settings()) repo = getattr(c, name) assert isinstance(repo, port) @@ -138,15 +150,20 @@ def test_repo_property_returns_port_typed_instance(self, name, port): @pytest.mark.asyncio async def test_initialize_seeds_admin_token(self, monkeypatch): + """Seed the configured admin token during container initialization.""" calls = [] async def ensure_admin_user(token): + """Record admin token seeding calls.""" calls.append(token) class FakeCatalogStore: + """Small catalog-store stand-in for initialization sequencing.""" + user_repo = SimpleNamespace(ensure_admin_user=ensure_admin_user) async def initialize(self): + """Record catalog initialization calls.""" calls.append("initialize") monkeypatch.setenv("AUTH_TOKEN", "admin-token") @@ -165,18 +182,21 @@ class TestVectorStoreWiring: ``tests/integration/test_milvus_store_integration.py``.""" def test_factory_returns_milvus_vector_store(self): + """Build Milvus vector stores from the DI factory.""" from services.storage.milvus_store import MilvusVectorStore store = create_vector_store(_settings()) assert isinstance(store, MilvusVectorStore) def test_container_property_returns_milvus_vector_store(self): + """Expose a Milvus vector store from the service container.""" from services.storage.milvus_store import MilvusVectorStore c = ServiceContainer(_settings()) assert isinstance(c.vector_store, MilvusVectorStore) def test_container_caches_vector_store(self): + """Cache the vector store so repeated reads reuse one client.""" c = ServiceContainer(_settings()) # Repeated property reads must return the same instance — every # construction opens a fresh pymilvus gRPC channel. @@ -184,14 +204,19 @@ def test_container_caches_vector_store(self): class TestRepositoriesFactory: + """Verify repository factory behavior independent of the container.""" + def test_returns_a_catalog_store(self): + """Build a catalog store through the repositories factory.""" assert isinstance(create_catalog_store(_settings()), CatalogStore) def test_run_migrations_flag_propagates(self): + """Pass the migration flag through to the catalog store.""" store = create_catalog_store(_settings(), run_migrations=False) assert store._run_migrations is False def test_does_not_mutate_input_settings(self): + """Leave caller-owned settings unchanged when deriving defaults.""" s = _settings(database=None, collection="abc") original_database = s.rdb.database create_catalog_store(s) @@ -220,6 +245,7 @@ class TestPhase8OrchestratorWiring: @pytest.mark.parametrize("prop,_provider", _ORCHESTRATORS) def test_property_is_lazy_and_cache_slot_starts_none(self, prop, _provider): + """Keep orchestrator properties lazy until the first access.""" # The public accessor is a property (lazy), not an eager attribute. assert isinstance(getattr(ServiceContainer, prop), property) # The cache slot exists and is None before first access (no @@ -228,6 +254,7 @@ def test_property_is_lazy_and_cache_slot_starts_none(self, prop, _provider): @pytest.mark.parametrize("prop,provider", _ORCHESTRATORS) def test_provider_delegates_to_container_property(self, prop, provider): + """Delegate each FastAPI provider to its container property.""" from types import SimpleNamespace from di import providers @@ -240,7 +267,135 @@ def test_provider_delegates_to_container_property(self, prop, provider): assert resolved is sentinel def test_no_orchestrator_is_missing_a_provider(self): + """Keep the provider surface aligned with container orchestrators.""" from di import providers wired = {name for name in vars(providers) if name.startswith("get_") and name.endswith("_service")} assert wired == {p for _, p in _ORCHESTRATORS} + + +class TestPhase11ProviderBridge: + """Verify Phase 11 request and process-level provider bridge behavior.""" + + def test_set_container_supports_no_request_lookup(self): + """Resolve services from the process-level override without a request.""" + from di import providers + + fake_container = SimpleNamespace(is_initialized=True, config=object()) + try: + providers.set_container(fake_container) + assert providers.get_container() is fake_container + assert providers.get_config() is fake_container.config + finally: + providers.set_container(None) + + def test_request_container_takes_precedence_over_process_container(self): + """Prefer request app-state containers over process-level overrides.""" + from di import providers + + process_container = SimpleNamespace(auth_service=object()) + request_container = SimpleNamespace(auth_service=object()) + request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(container=request_container))) + try: + providers.set_container(process_container) + assert providers.get_container(request) is request_container + assert providers.get_auth_service(request) is request_container.auth_service + finally: + providers.set_container(None) + + def test_request_container_none_returns_503(self): + """Return service-unavailable when request state has no container.""" + from di import providers + + request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(container=None))) + with pytest.raises(HTTPException) as exc: + providers.get_container(request) + + assert exc.value.status_code == 503 + + def test_uninitialized_container_returns_503_for_service_getter(self): + """Return service-unavailable until the active container is initialized.""" + from di import providers + + fake_container = SimpleNamespace(is_initialized=False, auth_service=object()) + try: + providers.set_container(fake_container) + with pytest.raises(HTTPException) as exc: + providers.get_auth_service() + finally: + providers.set_container(None) + + assert exc.value.status_code == 503 + + +class TestPhase11ContainerLifecycle: + """Phase 11 introspection + lifecycle the provider bridge relies on.""" + + def test_config_returns_wired_settings(self): + """Expose the settings the container was built from.""" + settings = _settings() + assert ServiceContainer(settings).config is settings + + def test_is_initialized_false_before_initialize(self): + """Report not-initialized until the async phase has run.""" + assert ServiceContainer(_settings()).is_initialized is False + + @pytest.mark.asyncio + async def test_initialize_sets_is_initialized(self): + """Flip the init guard once the async phase completes.""" + c = ServiceContainer() # no settings: initialize is a no-op but still flips the flag + await c.initialize() + assert c.is_initialized is True + + def test_create_tracks_inference_client(self): + """Record clients built through the factory for shutdown cleanup.""" + c = ServiceContainer() + client = c.create_llm(endpoint="http://vllm:8000/v1", model_name="m") + assert c._inference_clients == [client] + + @pytest.mark.asyncio + async def test_shutdown_closes_clients_and_resets_state(self): + """Close every tracked client, clear the list, and drop the init flag.""" + closed = [] + + class _FakeClient: + async def aclose(self): + """Record that the client was closed.""" + closed.append(self) + + c = ServiceContainer() + c._initialized = True + fake = _FakeClient() + c._inference_clients.append(fake) + + await c.shutdown() + + assert closed == [fake] + assert c._inference_clients == [] + assert c.is_initialized is False + + @pytest.mark.asyncio + async def test_shutdown_is_best_effort_when_a_client_fails(self): + """One client close failure must not skip the rest or the reset.""" + closed = [] + + class _BadClient: + async def aclose(self): + """Fail to close, exercising the best-effort path.""" + raise RuntimeError("boom") + + class _GoodClient: + async def aclose(self): + """Record a successful close after a prior failure.""" + closed.append(self) + + c = ServiceContainer() + c._initialized = True + good = _GoodClient() + c._inference_clients.extend([_BadClient(), good]) + + await c.shutdown() # must not raise + + assert closed == [good] + assert c._inference_clients == [] + assert c.is_initialized is False diff --git a/openrag/di/test_factories.py b/openrag/di/test_factories.py new file mode 100644 index 000000000..b472aaac0 --- /dev/null +++ b/openrag/di/test_factories.py @@ -0,0 +1,176 @@ +"""Tests for the generic cached component factory (Phase 11A).""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass, field +from typing import Any + +import pytest +from core.utils.registry import Registry +from di.factories import make_component_factory + + +@dataclass +class _FakeEndpoint: + """Stand-in for the forward-looking ``ModelEndpointConfig`` shape.""" + + endpoint: str = "http://host:8000/v1" + model_name: str = "m" + batch_size: int = 8 + timeout: float = 30.0 + extra: dict[str, Any] = field(default_factory=dict) + + +class _Dummy: + """Records the kwargs it was built with so tests can assert on them.""" + + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + +def _registry() -> Registry[_Dummy]: + reg: Registry[_Dummy] = Registry("dummy") + reg.register("vllm")(_Dummy) + reg.register("other")(type("_Other", (_Dummy,), {})) + return reg + + +class TestMakeComponentFactory: + def test_builds_and_caches_single_instance(self): + """Repeated calls for the same name reuse the first-built instance.""" + caches: list[dict[str, _Dummy]] = [] + factory, cache = make_component_factory( + _registry(), + {"default": _FakeEndpoint()}, + default_impl="vllm", + client_caches=caches, + ) + + first = factory("default") + second = factory("default") + + assert first is second + assert cache["default"] is first + + def test_cache_registered_in_client_caches(self): + """The returned cache is appended to ``client_caches`` for shutdown.""" + caches: list[dict[str, _Dummy]] = [] + _, cache = make_component_factory( + _registry(), + {"default": _FakeEndpoint()}, + default_impl="vllm", + client_caches=caches, + ) + + assert caches == [cache] + + def test_unknown_name_raises_key_error(self): + """A name absent from the config section raises ``KeyError``.""" + factory, _ = make_component_factory( + _registry(), + {"default": _FakeEndpoint()}, + default_impl="vllm", + client_caches=[], + ) + + with pytest.raises(KeyError, match="missing"): + factory("missing") + + def test_default_impl_used_when_unspecified(self): + """Entries without an ``implementation`` key build ``default_impl``.""" + factory, _ = make_component_factory( + _registry(), + {"default": _FakeEndpoint()}, + default_impl="vllm", + client_caches=[], + ) + + instance = factory("default") + + assert type(instance).__name__ == "_Dummy" + + def test_implementation_key_selects_impl_and_is_not_a_kwarg(self): + """``extra['implementation']`` selects the class but is not forwarded.""" + factory, _ = make_component_factory( + _registry(), + {"default": _FakeEndpoint(extra={"implementation": "other"})}, + default_impl="vllm", + client_caches=[], + ) + + instance = factory("default") + + assert type(instance).__name__ == "_Other" + assert "implementation" not in instance.kwargs + + def test_config_fields_and_extra_forwarded_as_kwargs(self): + """Endpoint fields plus impl-specific ``extra`` reach the constructor.""" + factory, _ = make_component_factory( + _registry(), + {"default": _FakeEndpoint(extra={"api_key": "secret"})}, + default_impl="vllm", + client_caches=[], + ) + + kwargs = factory("default").kwargs + + assert kwargs["endpoint"] == "http://host:8000/v1" + assert kwargs["model_name"] == "m" + assert kwargs["batch_size"] == 8 + assert kwargs["timeout"] == 30.0 + assert kwargs["api_key"] == "secret" + + def test_extra_kwargs_fn_merges_last(self): + """``extra_kwargs_fn`` output overrides config-derived kwargs.""" + factory, _ = make_component_factory( + _registry(), + {"default": _FakeEndpoint(model_name="from-config")}, + default_impl="vllm", + client_caches=[], + extra_kwargs_fn=lambda cfg: {"model_name": "overridden"}, + ) + + assert factory("default").kwargs["model_name"] == "overridden" + + def test_concurrent_first_calls_build_once(self): + """Double-checked locking yields exactly one instance under contention.""" + build_count = 0 + count_lock = threading.Lock() + + class _Counting: + def __init__(self, **_kwargs: Any) -> None: + nonlocal build_count + with count_lock: + build_count += 1 + + reg: Registry[_Counting] = Registry("counting") + reg.register("vllm")(_Counting) + + factory, _ = make_component_factory( + reg, + {"default": _FakeEndpoint()}, + default_impl="vllm", + client_caches=[], + ) + + n = 16 + barrier = threading.Barrier(n) + results: list[_Counting] = [] + results_lock = threading.Lock() + + def worker() -> None: + barrier.wait() + instance = factory("default") + with results_lock: + results.append(instance) + + threads = [threading.Thread(target=worker) for _ in range(n)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert build_count == 1 + assert len(results) == n + assert all(r is results[0] for r in results) diff --git a/openrag/di/test_inference.py b/openrag/di/test_inference.py index f38de06b2..2e332ce5b 100644 --- a/openrag/di/test_inference.py +++ b/openrag/di/test_inference.py @@ -1,54 +1,101 @@ from __future__ import annotations from core.embeddings import embedder_registry +from core.indexing.parsers import parser_registry from core.llm import llm_registry from core.rerankers import reranker_registry from core.vlm import vlm_registry from di.container import ServiceContainer from di.inference import register_inference +from di.parsers import register_parsers class TestRegisterInference: + """Verify inference implementation registration side effects.""" + def test_registries_populated(self): + """Register all supported inference backends in core registries.""" register_inference() assert "vllm" in llm_registry + assert "ollama" in llm_registry assert "vllm" in embedder_registry + assert "ollama" in embedder_registry assert "vllm" in vlm_registry assert "infinity" in reranker_registry assert "openai" in reranker_registry def test_idempotent(self): + """Allow repeated inference registration without duplicate failures.""" register_inference() register_inference() +class TestRegisterParsers: + """Verify parser implementation registration side effects.""" + + def test_parser_registry_populated(self): + """Register all supported parser implementations in the core registry.""" + register_parsers() + + assert { + "audio_client", + "doc", + "docling", + "docx", + "eml", + "html", + "image", + "local_whisper", + "markdown", + "marker", + "pdf_client", + "pptx", + "pymupdf", + "text", + }.issubset(set(parser_registry.list_registered())) + + def test_idempotent(self): + """Allow repeated parser registration without duplicate failures.""" + register_parsers() + register_parsers() + + class TestServiceContainer: + """Verify service container registry initialization and factories.""" + def test_container_populates_all_registries(self): + """Populate inference registries when a container is constructed.""" ServiceContainer() assert "vllm" in llm_registry + assert "ollama" in llm_registry assert "vllm" in embedder_registry + assert "ollama" in embedder_registry assert "vllm" in vlm_registry assert "infinity" in reranker_registry assert "openai" in reranker_registry def test_create_llm(self): + """Create an LLM client through the container factory.""" container = ServiceContainer() client = container.create_llm(endpoint="http://vllm:8000/v1", model_name="m") assert client is not None def test_create_embedder(self): + """Create an embedder client through the container factory.""" container = ServiceContainer() client = container.create_embedder(endpoint="http://vllm:8000/v1", model_name="m") assert client is not None def test_create_reranker(self): + """Create a reranker client through the container factory.""" container = ServiceContainer() client = container.create_reranker(endpoint="http://reranker:7997", model_name="m") assert client is not None def test_create_vlm(self): + """Create a VLM client through the container factory.""" container = ServiceContainer() client = container.create_vlm(endpoint="http://vllm:8000/v1", model_name="m") assert client is not None diff --git a/openrag/di/test_workers.py b/openrag/di/test_workers.py new file mode 100644 index 000000000..c4ee5b968 --- /dev/null +++ b/openrag/di/test_workers.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import sys +from importlib import import_module +from types import ModuleType +from unittest.mock import Mock + +from core.config.root import Settings +from di.workers import ensure_worker_bootstrap + + +def test_ensure_worker_bootstrap_initializes_explicitly() -> None: + """Startup calls the worker bootstrap function instead of relying on import side effects.""" + module = ModuleType("services.workers.bootstrap") + module.initialize_worker_bootstrap = Mock() + sys.modules["services.workers.bootstrap"] = module + settings = Settings() + + try: + ensure_worker_bootstrap(settings) + finally: + sys.modules.pop("services.workers.bootstrap", None) + + module.initialize_worker_bootstrap.assert_called_once_with(settings) + + +def test_worker_bootstrap_import_has_no_actor_side_effects() -> None: + """Importing worker bootstrap does not create detached actors.""" + sys.modules.pop("services.workers.bootstrap", None) + + module = import_module("services.workers.bootstrap") + + assert module.actor_creation_map == {} + assert not hasattr(module, "task_state_manager") + assert not hasattr(module, "serializer") diff --git a/openrag/di/workers.py b/openrag/di/workers.py index 1e6e69e23..20f83dc3f 100644 --- a/openrag/di/workers.py +++ b/openrag/di/workers.py @@ -22,9 +22,11 @@ def list_ray_actors() -> list[dict[str, str | None]]: ] -def ensure_worker_bootstrap() -> None: - """Import the worker bootstrap after Ray has been initialized.""" - import services.workers.bootstrap # noqa: F401 +def ensure_worker_bootstrap(settings: Any) -> None: + """Initialize the worker bootstrap after Ray has been initialized.""" + from services.workers.bootstrap import initialize_worker_bootstrap + + initialize_worker_bootstrap(settings) def get_actor_creation_map() -> Mapping[str, Callable[[], Any]]: diff --git a/openrag/services/workers/bootstrap.py b/openrag/services/workers/bootstrap.py index ff5a623f4..e9c12326b 100644 --- a/openrag/services/workers/bootstrap.py +++ b/openrag/services/workers/bootstrap.py @@ -19,21 +19,25 @@ """ from functools import wraps +from typing import TYPE_CHECKING import ray -from components.indexer.loaders.audio import WhisperActor, WhisperPool -from components.indexer.loaders.pdf_loaders.docling2 import DoclingPool -from components.indexer.loaders.pdf_loaders.marker import MarkerPool -from components.indexer.loaders.serializer import DocSerializer -from config import load_config -from services.inference.distributed_semaphore import DistributedSemaphoreActor -from services.workers.task_state import TaskStateManager from utils.logger import get_logger -config = load_config() +if TYPE_CHECKING: + from core.config.root import Settings + + logger = get_logger() actor_creation_map: dict[str, callable] = {} +_settings: "Settings | None" = None + + +def _require_settings() -> "Settings": + if _settings is None: + raise RuntimeError("Worker bootstrap has not been initialized.") + return _settings def _track_actor(func): @@ -56,14 +60,22 @@ def get_or_create_actor(name, cls, namespace="openrag", remote_args=(), **option def get_task_state_manager(): + from services.workers.task_state import TaskStateManager + return get_or_create_actor("TaskStateManager", TaskStateManager, lifetime="detached") def get_serializer(): + from components.indexer.loaders.serializer import DocSerializer + return get_or_create_actor("DocSerializer", DocSerializer, lifetime="detached") def get_marker_pool(): + from components.indexer.loaders.pdf_loaders.docling2 import DoclingPool + from components.indexer.loaders.pdf_loaders.marker import MarkerPool + + config = _require_settings() pdf_loader = config.loader.file_loaders.pdf match pdf_loader: case "DoclingLoader2": @@ -73,6 +85,9 @@ def get_marker_pool(): def init_audio_actor(): + from services.workers.parsers.whisper_workers import WhisperActor, WhisperPool, whisper_actor_options + + config = _require_settings() use_whisper_lang_detector = config.loader.transcriber.use_whisper_lang_detector file_loaders = config.loader.file_loaders loader_values = set(file_loaders.values()) if file_loaders else set() @@ -81,10 +96,13 @@ def init_audio_actor(): return get_or_create_actor("WhisperPool", WhisperPool, lifetime="detached") if "OpenAIAudioLoader" in loader_values and use_whisper_lang_detector: - return get_or_create_actor("WhisperActor", WhisperActor, lifetime="detached") + return get_or_create_actor("WhisperActor", WhisperActor, lifetime="detached", **whisper_actor_options(config)) def init_llm_semaphore(): + from services.inference.distributed_semaphore import DistributedSemaphoreActor + + config = _require_settings() return get_or_create_actor( "llmSemaphore", DistributedSemaphoreActor, @@ -94,6 +112,9 @@ def init_llm_semaphore(): def init_vlm_semaphore(): + from services.inference.distributed_semaphore import DistributedSemaphoreActor + + config = _require_settings() return get_or_create_actor( "vlmSemaphore", DistributedSemaphoreActor, @@ -103,6 +124,9 @@ def init_vlm_semaphore(): def init_audio_semaphore(): + from services.inference.distributed_semaphore import DistributedSemaphoreActor + + config = _require_settings() return get_or_create_actor( "audioSemaphore", DistributedSemaphoreActor, @@ -111,11 +135,15 @@ def init_audio_semaphore(): ) -init_llm_semaphore() -init_vlm_semaphore() -init_audio_semaphore() -init_audio_actor() -get_marker_pool() +def initialize_worker_bootstrap(settings: "Settings") -> None: + """Create the detached worker actors required by the request path.""" + global _settings + _settings = settings -task_state_manager = get_task_state_manager() -serializer = get_serializer() + init_llm_semaphore() + init_vlm_semaphore() + init_audio_semaphore() + init_audio_actor() + get_marker_pool() + get_task_state_manager() + get_serializer() diff --git a/openrag/services/workers/parsers/doc_serializer.py b/openrag/services/workers/parsers/doc_serializer.py index 6aca97d18..80c7d9e68 100644 --- a/openrag/services/workers/parsers/doc_serializer.py +++ b/openrag/services/workers/parsers/doc_serializer.py @@ -12,18 +12,8 @@ import ray import torch from components.indexer.loaders import get_loader_classes -from config import load_config from langchain_core.documents.base import Document -config = load_config() - -if torch.cuda.is_available(): - NUM_GPUS = config.ray.num_gpus -else: - NUM_GPUS = 0 - -DICT_MIMETYPES = config.loader.mimetypes.to_dict() - @ray.remote(max_restarts=5) class DocSerializer: @@ -61,10 +51,11 @@ async def serialize_document( p = Path(path) file_ext = p.suffix.lower() mimetype = metadata.get("mimetype", None) + mimetypes = self.config.loader.mimetypes.to_dict() if mimetype is None: loader_cls = self.loader_classes.get(file_ext) else: - loader_cls = self.loader_classes.get(DICT_MIMETYPES.get(mimetype)) + loader_cls = self.loader_classes.get(mimetypes.get(mimetype)) if loader_cls is None: log.warning(f"No loader available for {p.name}") diff --git a/openrag/services/workers/parsers/docling_workers.py b/openrag/services/workers/parsers/docling_workers.py index 7c0a9ef99..0a52e2ba2 100644 --- a/openrag/services/workers/parsers/docling_workers.py +++ b/openrag/services/workers/parsers/docling_workers.py @@ -31,20 +31,16 @@ from docling.document_converter import DocumentConverter, PdfFormatOption from utils.logger import get_logger -from ..ray_utils import call_ray_actor_with_timeout, retry_with_backoff, with_timeout +from ..ray_utils import call_ray_actor_with_timeout, retry_with_backoff logger = get_logger() -config = load_config() -if torch.cuda.is_available(): - DOCLING_NUM_GPUS = config.loader.docling_num_gpus -else: - DOCLING_NUM_GPUS = 0 -DOCLING_MAX_TASKS_PER_WORKER = config.loader.docling_max_tasks_per_worker +def _docling_num_gpus(config) -> float: + return config.loader.docling_num_gpus if torch.cuda.is_available() else 0 -@ray.remote(num_gpus=DOCLING_NUM_GPUS) +@ray.remote class DoclingWorker: def __init__(self): img_scale = 2 @@ -75,18 +71,20 @@ def __init__(self): self.logger = get_logger() self.config = load_config() self.pool_size = self.config.loader.docling_pool_size + self.max_tasks_per_worker = self.config.loader.docling_max_tasks_per_worker - self.actors = [DoclingWorker.remote() for _ in range(self.pool_size)] + self.actors = [ + DoclingWorker.options(num_gpus=_docling_num_gpus(self.config)).remote() for _ in range(self.pool_size) + ] self._queue: asyncio.Queue[ray.actor.ActorHandle] = asyncio.Queue() - for _ in range(DOCLING_MAX_TASKS_PER_WORKER): + for _ in range(self.max_tasks_per_worker): for actor in self.actors: self._queue.put_nowait(actor) - total_slots = self.pool_size * DOCLING_MAX_TASKS_PER_WORKER + total_slots = self.pool_size * self.max_tasks_per_worker self.logger.info( - f"Docling pool: {self.pool_size} actors × {DOCLING_MAX_TASKS_PER_WORKER} slots = " - f"{total_slots} PDF concurrency" + f"Docling pool: {self.pool_size} actors × {self.max_tasks_per_worker} slots = {total_slots} PDF concurrency" ) async def process_pdf(self, file_path: str) -> ConversionResult: @@ -121,6 +119,7 @@ class DoclingLoader(BasePooledParser): """ def __init__(self) -> None: + self.config = load_config() self.worker: DoclingPool = ray.get_actor("DoclingPool", namespace="openrag") def supported_types(self) -> list[str]: @@ -147,12 +146,12 @@ async def parse(self, document: Document) -> ProcessedDocument: page_count=len(result.pages), ) - @with_timeout( - seconds=config.loader.docling_timeout, - description="DoclingLoader PDF loading ({file_path})", - ) async def _dispatch(self, file_path: str) -> ConversionResult: - return self.worker.process_pdf.remote(file_path) + return await call_ray_actor_with_timeout( + self.worker.process_pdf.remote(file_path), + timeout=self.config.loader.docling_timeout, + task_description=f"DoclingLoader PDF loading ({file_path})", + ) @staticmethod def _build_text_blocks(result: ConversionResult) -> list[TextBlock]: diff --git a/openrag/services/workers/parsers/marker_workers.py b/openrag/services/workers/parsers/marker_workers.py index 3a495ca64..36417bfb5 100644 --- a/openrag/services/workers/parsers/marker_workers.py +++ b/openrag/services/workers/parsers/marker_workers.py @@ -19,18 +19,16 @@ from marker.converters.pdf import PdfConverter from utils.logger import get_logger -from ..ray_utils import with_retry, with_timeout +from ..ray_utils import call_ray_actor_with_timeout, retry_with_backoff logger = get_logger() -config = load_config() -if torch.cuda.is_available(): - MARKER_NUM_GPUS = config.loader.marker_num_gpus -else: # On CPU - MARKER_NUM_GPUS = 0 +def _marker_num_gpus(config) -> float: + return config.loader.marker_num_gpus if torch.cuda.is_available() else 0 -@ray.remote(num_gpus=MARKER_NUM_GPUS, max_restarts=5) + +@ray.remote class MarkerWorker: def __init__(self): import os @@ -167,9 +165,10 @@ def is_pool_broken(self): def __del__(self): """Clean up ProcessPoolExecutor on actor destruction""" - if self.executor: + executor = getattr(self, "executor", None) + if executor: try: - self.executor.shutdown(wait=False, cancel_futures=True) + executor.shutdown(wait=False, cancel_futures=True) except Exception: pass # Best effort cleanup @@ -184,7 +183,10 @@ def __init__(self): self.config = load_config() self.max_processes = self.config.loader.marker_max_processes self.pool_size = self.config.loader.marker_pool_size - self.actors = [MarkerWorker.remote() for _ in range(self.pool_size)] + self.actors = [ + MarkerWorker.options(num_gpus=_marker_num_gpus(self.config), max_restarts=5).remote() + for _ in range(self.pool_size) + ] self._queue: asyncio.Queue[ray.actor.ActorHandle] = asyncio.Queue() for _ in range(self.max_processes): @@ -216,52 +218,56 @@ def _create_chunks(page_count: int, chunk_size: int) -> list[tuple[list[int], st chunks.append((page_range, label)) return chunks - @with_timeout( - seconds=config.loader.marker_timeout, - description="MarkerWorker pool health check", - ) async def _check_pool_broken(self, worker): - return worker.is_pool_broken.remote() + return await call_ray_actor_with_timeout( + worker.is_pool_broken.remote(), + timeout=self.config.loader.marker_timeout, + task_description="MarkerWorker pool health check", + ) - @with_timeout( - seconds=config.loader.marker_timeout, - description="MarkerWorker pool reset", - ) async def _reset_worker_pool(self, worker): - return worker.setup_mp.remote() + return await call_ray_actor_with_timeout( + worker.setup_mp.remote(), + timeout=self.config.loader.marker_timeout, + task_description="MarkerWorker pool reset", + ) async def ensure_worker_pool_healthy(self, worker): if await self._check_pool_broken(worker): self.logger.warning("Worker ProcessPoolExecutor is broken. Reinitializing pool...") await self._reset_worker_pool(worker) - @with_timeout( - seconds=config.loader.marker_timeout, - description="MarkerPool PDF {label} ({file_path})", - ) async def _run_chunk(self, worker, file_path: str, page_range: list[int] | None, label: str): - return worker.process_pdf.remote(file_path, page_range=page_range) + return await call_ray_actor_with_timeout( + worker.process_pdf.remote(file_path, page_range=page_range), + timeout=self.config.loader.marker_timeout, + task_description=f"MarkerPool PDF {label} ({file_path})", + ) - @with_retry( - max_retries=config.loader.marker_max_task_retry, - base_delay=config.loader.marker_retry_base_delay, - description="MarkerPool PDF {label} ({file_path})", - ) async def _process_chunk(self, file_path: str, page_range: list[int] | None, label: str): """Acquire a worker slot, process a PDF chunk, and release the slot. A fresh worker is acquired per attempt so a flaky worker can be sidestepped and ``ensure_worker_pool_healthy`` re-runs each time. - Retries are handled by ``@with_retry``. + Retries are handled by ``retry_with_backoff``. """ - worker = await self._queue.get() - try: - self.logger.info(f"MarkerWorker allocated for {label}") - await self.ensure_worker_pool_healthy(worker) - return await self._run_chunk(worker, file_path, page_range, label) - finally: - await self._queue.put(worker) - self.logger.debug(f"MarkerWorker returned to pool for {label}") + + async def attempt(_i: int): + worker = await self._queue.get() + try: + self.logger.info(f"MarkerWorker allocated for {label}") + await self.ensure_worker_pool_healthy(worker) + return await self._run_chunk(worker, file_path, page_range, label) + finally: + await self._queue.put(worker) + self.logger.debug(f"MarkerWorker returned to pool for {label}") + + return await retry_with_backoff( + attempt, + max_retries=self.config.loader.marker_max_task_retry, + base_delay=self.config.loader.marker_retry_base_delay, + task_description=f"MarkerPool PDF {label} ({file_path})", + ) async def process_pdf(self, file_path: str): chunk_size = self.config.loader.marker_chunk_size @@ -340,6 +346,7 @@ class MarkerLoader(BasePooledParser): _PAGE_MARKER_RE = re.compile(r"\{(\d+)\}" + re.escape(PAGE_SEP)) def __init__(self) -> None: + self.config = load_config() self.worker = ray.get_actor("MarkerPool", namespace="openrag") def supported_types(self) -> list[str]: @@ -369,12 +376,12 @@ async def parse(self, document: Document) -> ProcessedDocument: # ----- helpers ----- - @with_timeout( - seconds=config.loader.marker_timeout, - description="MarkerLoader PDF loading ({file_path})", - ) async def _convert_pdf(self, file_path: str): - return self.worker.process_pdf.remote(file_path) + return await call_ray_actor_with_timeout( + self.worker.process_pdf.remote(file_path), + timeout=self.config.loader.marker_timeout, + task_description=f"MarkerLoader PDF loading ({file_path})", + ) async def _dispatch(self, file_path: str) -> tuple[str, dict]: start = time.time() diff --git a/openrag/services/workers/parsers/whisper_workers.py b/openrag/services/workers/parsers/whisper_workers.py index 1893d88bd..697f2fee2 100644 --- a/openrag/services/workers/parsers/whisper_workers.py +++ b/openrag/services/workers/parsers/whisper_workers.py @@ -14,27 +14,28 @@ from faster_whisper import WhisperModel from utils.logger import get_logger -from ..ray_utils import with_retry, with_timeout +from ..ray_utils import call_ray_actor_with_timeout, retry_with_backoff logger = get_logger() -config = load_config() -if torch.cuda.is_available(): - WHISPER_NUM_GPUS = config.loader.local_whisper.whisper_num_gpus -else: # On CPU - WHISPER_NUM_GPUS = 0 +def _whisper_num_gpus(config) -> float: + return config.loader.local_whisper.whisper_num_gpus if torch.cuda.is_available() else 0 -WHISPER_CONCURRENCY_PER_WORKER = config.loader.local_whisper.whisper_concurrency_per_worker + +def whisper_actor_options(config) -> dict[str, float | int]: + return { + "num_gpus": _whisper_num_gpus(config), + "max_restarts": 5, + "max_concurrency": config.loader.local_whisper.whisper_concurrency_per_worker, + } # Duration of the audio sample used for language detection LANG_DETECT_SAMPLE_MS = 30_000 # 30 s -@ray.remote( - num_gpus=WHISPER_NUM_GPUS, max_restarts=5, max_concurrency=WHISPER_CONCURRENCY_PER_WORKER -) # Ensure each worker processes one file at a time +@ray.remote class WhisperActor: def __init__(self): import torch @@ -101,34 +102,39 @@ class WhisperPool: """ def __init__(self): + from config import load_config from utils.logger import get_logger self.logger = get_logger() + self.config = load_config() - n_workers = config.loader.local_whisper.whisper_n_workers + n_workers = self.config.loader.local_whisper.whisper_n_workers self.logger.info(f"Starting WhisperPool with {n_workers} workers") - self.workers = [WhisperActor.remote() for _ in range(n_workers)] + self.workers = [WhisperActor.options(**whisper_actor_options(self.config)).remote() for _ in range(n_workers)] self._pending = [0] * n_workers - @with_timeout( - seconds=config.loader.local_whisper.whisper_timeout, - description="WhisperPool transcribe ({path})", - ) async def _transcribe_chunk(self, idx: int, path): - return self.workers[idx].transcribe.remote(path) + return await call_ray_actor_with_timeout( + self.workers[idx].transcribe.remote(path), + timeout=self.config.loader.local_whisper.whisper_timeout, + task_description=f"WhisperPool transcribe ({path})", + ) - @with_retry( - max_retries=config.loader.local_whisper.whisper_max_task_retry, - base_delay=config.loader.local_whisper.whisper_retry_base_delay, - description="WhisperPool transcribe ({path})", - ) async def transcribe(self, path): - idx = min(range(len(self._pending)), key=lambda j: self._pending[j]) - self._pending[idx] += 1 - try: - return await self._transcribe_chunk(idx, path) - finally: - self._pending[idx] -= 1 + async def attempt(_i: int): + idx = min(range(len(self._pending)), key=lambda j: self._pending[j]) + self._pending[idx] += 1 + try: + return await self._transcribe_chunk(idx, path) + finally: + self._pending[idx] -= 1 + + return await retry_with_backoff( + attempt, + max_retries=self.config.loader.local_whisper.whisper_max_task_retry, + base_delay=self.config.loader.local_whisper.whisper_retry_base_delay, + task_description=f"WhisperPool transcribe ({path})", + ) async def detect_language_via_actor( @@ -143,10 +149,14 @@ async def detect_language_via_actor( loader's optional language detector) can stay Ray-free. Returns ``None`` on failure so callers can fall back to a default behaviour. """ - from ..ray_utils import call_ray_actor_with_timeout - + config = load_config() try: - actor = WhisperActor.options(name="WhisperActor", namespace="openrag", get_if_exists=True).remote() + actor = WhisperActor.options( + name="WhisperActor", + namespace="openrag", + get_if_exists=True, + **whisper_actor_options(config), + ).remote() except Exception: logger.exception("Error getting WhisperActor") return None @@ -173,6 +183,7 @@ class LocalWhisperLoader(BasePooledParser): """ def __init__(self): + self.config = load_config() self.whisper_actor: WhisperPool = ray.get_actor("WhisperPool", namespace="openrag") def supported_types(self) -> list[str]: diff --git a/openrag/utils/logger.py b/openrag/utils/logger.py index a794e2b84..64771b9e8 100644 --- a/openrag/utils/logger.py +++ b/openrag/utils/logger.py @@ -4,8 +4,6 @@ from config import load_config from loguru import logger -config = load_config() - def escape_markup(s: str) -> str: return s.replace("\\", "\\\\").replace("<", "\\<").replace(">", "\\>") @@ -27,7 +25,9 @@ def mask_email(email: str | None) -> str: return f"{masked_local}@{domain}" -def get_logger(): +def get_logger(config=None): + config = config or load_config() + def formatter(record): level = record["level"].name mod = record["name"]