From 1371c275a1d53ed54cde107d6c80244565c62d53 Mon Sep 17 00:00:00 2001 From: Javier Martinez Date: Tue, 14 Jul 2026 17:58:45 +0200 Subject: [PATCH] fix: worker config --- Makefile | 29 ++- fern/docs/pages/configuration/cli.mdx | 24 +-- private_gpt/arq/__init__.py | 3 + private_gpt/arq/__main__.py | 4 + private_gpt/arq/chat/__init__.py | 1 + private_gpt/arq/{ => chat}/iteration_state.py | 2 + private_gpt/arq/enqueue.py | 8 +- private_gpt/arq/healthcheck.py | 22 +-- private_gpt/arq/lifecycle.py | 5 +- private_gpt/arq/runner.py | 63 +++--- private_gpt/arq/settings.py | 15 +- private_gpt/arq/task_registry.py | 6 +- private_gpt/arq/tasks/chat/__init__.py | 12 +- private_gpt/arq/tasks/chat/resume.py | 9 +- private_gpt/arq/tasks/chat/settings.py | 11 ++ private_gpt/arq/tasks/chat/start.py | 4 +- private_gpt/arq/worker.py | 10 - private_gpt/celery/base.py | 6 +- private_gpt/celery/celery.py | 14 +- private_gpt/celery/healthcheck.py | 10 +- private_gpt/celery/task_registry.py | 11 +- private_gpt/celery/tasks/__init__.py | 5 - private_gpt/cli/commands/worker.py | 180 +----------------- .../engines/chat/checkpoint_store.py | 2 +- private_gpt/worker/__init__.py | 12 ++ private_gpt/worker/__main__.py | 6 + private_gpt/worker/modes.py | 177 +++++++++++++++++ private_gpt/worker/registry.py | 31 +++ scripts/worker_entrypoint | 121 +----------- tests/arq/tasks/chat/test_start.py | 7 +- tests/arq/test_healthcheck.py | 9 + tests/arq/test_runner.py | 37 ++++ tests/arq/test_task_registry.py | 20 ++ tests/celery/test_chat_worker.py | 30 +++ tests/celery/test_healthcheck.py | 32 ++++ tests/celery/test_task_registry.py | 52 +++++ tests/worker/test_modes.py | 18 ++ tests/worker/test_registry.py | 33 ++++ 38 files changed, 627 insertions(+), 414 deletions(-) create mode 100644 private_gpt/arq/__main__.py create mode 100644 private_gpt/arq/chat/__init__.py rename private_gpt/arq/{ => chat}/iteration_state.py (98%) create mode 100644 private_gpt/arq/tasks/chat/settings.py delete mode 100644 private_gpt/arq/worker.py create mode 100644 private_gpt/worker/__init__.py create mode 100644 private_gpt/worker/__main__.py create mode 100644 private_gpt/worker/modes.py create mode 100644 private_gpt/worker/registry.py create mode 100644 tests/arq/test_healthcheck.py create mode 100644 tests/arq/test_runner.py create mode 100644 tests/arq/test_task_registry.py create mode 100644 tests/celery/test_healthcheck.py create mode 100644 tests/celery/test_task_registry.py create mode 100644 tests/worker/test_modes.py create mode 100644 tests/worker/test_registry.py diff --git a/Makefile b/Makefile index 3b01f968..50d80dac 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,7 @@ # Any args passed to the make script, use with $(call args, default_value) args = `arg="$(filter-out $@,$(MAKECMDGOALS))" && echo $${arg:-${1}}` AUTO_DISCOVER_ARGS ?= +WORKER_ARGS ?= PROD_DIST_DIR ?= .dist PROD_PYTHON ?= .venv/bin/python PROD_BINARY ?= .venv/bin/private-gpt @@ -11,7 +12,7 @@ TEST_LOCAL_DATA_DIR ?= $(TEST_PGPT_HOME)/local_data/tests WIPE_PGPT_HOME := $(if $(PGPT_HOME),$(PGPT_HOME),$(HOME)/.local/share/private-gpt) WIPE_LOCAL_DATA_DIR := $(WIPE_PGPT_HOME)/local_data -.PHONY: test test-coverage black ruff format mypy check auto-discover-models update-openapi-spec run dev-windows dev prod-run api-docs docs ingest wipe celery flower celery-worker arq-worker +.PHONY: test test-coverage black ruff format mypy check auto-discover-models update-openapi-spec run dev-windows dev prod-run api-docs docs ingest wipe celery flower celery-worker arq-worker chat-worker tools-worker ######################################################################################################################## # Quality checks @@ -103,13 +104,31 @@ wipe: ######################################################################################################################## celery: - PGPT_WORKER_MODE=celery uv run private-gpt worker + PGPT_WORKER_MODE=celery ./scripts/worker_entrypoint $(WORKER_ARGS) flower: - PGPT_WORKER_MODE=flower uv run private-gpt worker + PGPT_WORKER_MODE=flower ./scripts/worker_entrypoint $(WORKER_ARGS) celery-worker: - PGPT_WORKER_MODE=celery uv run private-gpt worker + PGPT_WORKER_MODE=celery ./scripts/worker_entrypoint $(WORKER_ARGS) arq-worker: - PGPT_WORKER_MODE=arq uv run private-gpt worker + PGPT_WORKER_MODE=arq ./scripts/worker_entrypoint $(WORKER_ARGS) + +chat-worker: + PGPT_WORKER_MODE=arq \ + PGPT_WORKER_APP_MODULE=private_gpt \ + PGPT_ARQ_QUEUE=chat \ + PGPT_ARQ_TASK_PACKAGES=private_gpt.arq.tasks.chat \ + PGPT_STATEFUL_WORKER_TYPE=chat \ + PGPT_WORKER_WARM_PROFILE=chat \ + ./scripts/worker_entrypoint $(WORKER_ARGS) + +tools-worker: + PGPT_WORKER_MODE=celery \ + PGPT_WORKER_APP_MODULE=private_gpt \ + PGPT_CELERY_QUEUES=tools \ + PGPT_CELERY_TASK_PACKAGES=private_gpt.celery.tasks.tools \ + PGPT_STATEFUL_WORKER_TYPE=tools \ + PGPT_WORKER_WARM_PROFILE=tools \ + ./scripts/worker_entrypoint $(WORKER_ARGS) diff --git a/fern/docs/pages/configuration/cli.mdx b/fern/docs/pages/configuration/cli.mdx index 038abc02..9c0394a7 100644 --- a/fern/docs/pages/configuration/cli.mdx +++ b/fern/docs/pages/configuration/cli.mdx @@ -74,12 +74,12 @@ Starts the Celery worker stack. All behaviour is controlled through environment private-gpt worker ``` -The worker process manages up to three sub-processes depending on `PGPT_WORKER_MODE`: +`PGPT_WORKER_MODE` is required and selects one registered worker runtime: | `PGPT_WORKER_MODE` | What starts | |---|---| -| `mixed` (default) | Celery worker + Flower UI | -| `worker` | Celery worker only | +| `arq` | ARQ worker | +| `celery` | Celery worker | | `flower` | Flower UI only | **Examples:** @@ -87,11 +87,7 @@ The worker process manages up to three sub-processes depending on `PGPT_WORKER_M ```bash - # Default: worker + Flower UI - private-gpt worker - - # Worker only, no Flower - PGPT_WORKER_MODE=worker private-gpt worker + PGPT_WORKER_MODE=celery private-gpt worker # Custom concurrency and queues PGPT_CELERY_QUEUES=ingest,chat PGPT_CELERY_CONCURRENCY=4 private-gpt worker @@ -99,11 +95,7 @@ The worker process manages up to three sub-processes depending on `PGPT_WORKER_M ```powershell - # Default: worker + Flower UI - private-gpt worker - - # Worker only, no Flower - $env:PGPT_WORKER_MODE = "worker" + $env:PGPT_WORKER_MODE = "celery" private-gpt worker # Custom concurrency and queues @@ -114,11 +106,7 @@ The worker process manages up to three sub-processes depending on `PGPT_WORKER_M ```cmd - :: Default: worker + Flower UI - private-gpt worker - - :: Worker only, no Flower - set PGPT_WORKER_MODE=worker + set PGPT_WORKER_MODE=celery private-gpt worker :: Custom concurrency and queues diff --git a/private_gpt/arq/__init__.py b/private_gpt/arq/__init__.py index e69de29b..876b4748 100644 --- a/private_gpt/arq/__init__.py +++ b/private_gpt/arq/__init__.py @@ -0,0 +1,3 @@ +from private_gpt.arq.runner import run_arq_worker + +__all__ = ["run_arq_worker"] diff --git a/private_gpt/arq/__main__.py b/private_gpt/arq/__main__.py new file mode 100644 index 00000000..4300e27e --- /dev/null +++ b/private_gpt/arq/__main__.py @@ -0,0 +1,4 @@ +from private_gpt.arq import run_arq_worker + +if __name__ == "__main__": + run_arq_worker() diff --git a/private_gpt/arq/chat/__init__.py b/private_gpt/arq/chat/__init__.py new file mode 100644 index 00000000..5f896047 --- /dev/null +++ b/private_gpt/arq/chat/__init__.py @@ -0,0 +1 @@ +"""ARQ chat worker support.""" diff --git a/private_gpt/arq/iteration_state.py b/private_gpt/arq/chat/iteration_state.py similarity index 98% rename from private_gpt/arq/iteration_state.py rename to private_gpt/arq/chat/iteration_state.py index b1df9cc0..e05c2c9c 100644 --- a/private_gpt/arq/iteration_state.py +++ b/private_gpt/arq/chat/iteration_state.py @@ -1,3 +1,5 @@ +"""Redis-backed state for resumable ARQ chat execution.""" + from __future__ import annotations import json diff --git a/private_gpt/arq/enqueue.py b/private_gpt/arq/enqueue.py index b2a7a86d..4b9da379 100644 --- a/private_gpt/arq/enqueue.py +++ b/private_gpt/arq/enqueue.py @@ -4,7 +4,7 @@ from typing import Any from arq import create_pool from arq.jobs import Job -from private_gpt.arq.settings import get_queue_name, get_redis_settings +from private_gpt.arq.settings import get_redis_settings from private_gpt.settings.settings import settings as _settings logger = logging.getLogger(__name__) @@ -32,13 +32,13 @@ def _log_dispatch( async def enqueue_job( *, task_name: str, + queue_name: str, args: tuple[Any, ...] = (), correlation_id: str, job_id: str | None = None, defer_seconds: int | None = None, ) -> None: current_settings = _settings() - queue_name = get_queue_name(current_settings) _log_dispatch( task_name=task_name, queue_name=queue_name, @@ -58,14 +58,14 @@ async def enqueue_job( await redis.aclose() -async def abort_job(*, job_id: str) -> bool: +async def abort_job(*, job_id: str, queue_name: str) -> bool: current_settings = _settings() redis = await create_pool(get_redis_settings(current_settings)) try: job = Job( job_id, redis=redis, - _queue_name=get_queue_name(current_settings), + _queue_name=queue_name, ) return await job.abort(timeout=2) finally: diff --git a/private_gpt/arq/healthcheck.py b/private_gpt/arq/healthcheck.py index d69f89d3..85b6d768 100644 --- a/private_gpt/arq/healthcheck.py +++ b/private_gpt/arq/healthcheck.py @@ -1,28 +1,14 @@ from typing import Any -from arq import create_pool -from fastapi import FastAPI, HTTPException - -from private_gpt.arq.settings import get_health_check_key, get_redis_settings -from private_gpt.settings.settings import settings +from fastapi import FastAPI app = FastAPI() @app.get("/health") async def health() -> dict[str, Any]: - current_settings = settings() - redis = await create_pool(get_redis_settings(current_settings)) - try: - healthy = bool(await redis.exists(get_health_check_key(current_settings))) - finally: - await redis.aclose() - - status = { - "status": "healthy" if healthy else "unhealthy", + return { + "status": "healthy", "mode": "arq-worker", - "services": {"worker": "healthy" if healthy else "unhealthy"}, + "services": {"worker": "healthy"}, } - if not healthy: - raise HTTPException(status_code=503, detail=status) - return status diff --git a/private_gpt/arq/lifecycle.py b/private_gpt/arq/lifecycle.py index 0e4aad45..2a294131 100644 --- a/private_gpt/arq/lifecycle.py +++ b/private_gpt/arq/lifecycle.py @@ -1,3 +1,4 @@ +import os from typing import Any import nest_asyncio # type: ignore @@ -19,7 +20,9 @@ async def startup(ctx: dict[Any, Any]) -> None: nest_asyncio.apply() injector = get_global_injector(allow_to_generate_new_injectors=True) set_global_injector(injector) - warm(injector, profile="chat") + warm_profile = os.environ.get("PGPT_WORKER_WARM_PROFILE", "").strip() + if warm_profile: + warm(injector, profile=warm_profile) ctx["injector"] = injector diff --git a/private_gpt/arq/runner.py b/private_gpt/arq/runner.py index 49cf85a0..198753ad 100644 --- a/private_gpt/arq/runner.py +++ b/private_gpt/arq/runner.py @@ -1,23 +1,19 @@ import asyncio import contextlib -import importlib import os import signal import subprocess import sys -from typing import TYPE_CHECKING +from collections.abc import Callable +from arq.typing import StartupShutdown from arq.worker import Worker -from private_gpt.arq.settings import ( - get_health_check_key, - get_queue_name, - get_redis_settings, -) -from private_gpt.settings.settings import settings - -if TYPE_CHECKING: - from private_gpt.settings.settings import Settings +from private_gpt.arq.hooks import on_job_end +from private_gpt.arq.lifecycle import shutdown, startup +from private_gpt.arq.settings import get_queue_name, get_redis_settings +from private_gpt.arq.tasks import autodiscover_registered_tasks +from private_gpt.settings.settings import Settings, settings def _default_concurrency() -> int: @@ -27,12 +23,33 @@ def _default_concurrency() -> int: return 1 << (cpu_count.bit_length() - 1) -def run_arq_worker() -> None: +def _task_packages() -> tuple[str, ...]: + configured = os.environ.get("PGPT_ARQ_TASK_PACKAGES", "") + task_packages = tuple( + package.strip() for package in configured.split(",") if package.strip() + ) + if not task_packages: + raise ValueError("PGPT_ARQ_TASK_PACKAGES must configure at least one package") + return task_packages + + +def _queue_name() -> str: + queue = os.environ.get("PGPT_ARQ_QUEUE", "").strip() + if not queue: + raise ValueError("PGPT_ARQ_QUEUE must configure a queue") + return get_queue_name(queue) + + +def run_arq_worker( + *, + settings_resolver: Callable[[], Settings] = settings, + startup_hook: StartupShutdown = startup, + shutdown_hook: StartupShutdown = shutdown, +) -> None: app_module = os.environ.get("PGPT_WORKER_APP_MODULE", "private_gpt") - worker_module = importlib.import_module(f"{app_module}.arq.worker") - settings_module = importlib.import_module(f"{app_module}.settings.settings") - settings_resolver = getattr(settings_module, "settings", settings) - current_settings: Settings = settings_resolver() + current_settings = settings_resolver() + task_packages = _task_packages() + queue_name = _queue_name() max_jobs = int(os.environ.get("PGPT_ARQ_MAX_JOBS", str(_default_concurrency()))) job_timeout = int(os.environ.get("PGPT_ARQ_JOB_TIMEOUT", "21600")) api_enabled = os.environ.get("API_ENABLED", "true").lower() == "true" @@ -77,12 +94,12 @@ def run_arq_worker() -> None: async def _main() -> None: worker = Worker( - functions=worker_module.functions, - queue_name=get_queue_name(current_settings), + functions=autodiscover_registered_tasks(*task_packages), + queue_name=queue_name, redis_settings=get_redis_settings(current_settings), - on_startup=worker_module.startup, - on_shutdown=worker_module.shutdown, - on_job_end=worker_module.on_job_end, + on_startup=startup_hook, + on_shutdown=shutdown_hook, + on_job_end=on_job_end, handle_signals=False, allow_abort_jobs=True, max_jobs=max_jobs, @@ -90,7 +107,6 @@ def run_arq_worker() -> None: retry_jobs=False, keep_result=0, job_timeout=job_timeout, - health_check_key=get_health_check_key(current_settings), health_check_interval=30, job_completion_wait=5, ) @@ -111,7 +127,8 @@ def run_arq_worker() -> None: task.cancel() print( - f"Starting arq worker queue={get_queue_name(current_settings)} max_jobs={max_jobs} job_timeout={job_timeout}" + f"Starting arq worker queue={queue_name} task_packages={','.join(task_packages)} " + f"max_jobs={max_jobs} job_timeout={job_timeout}" ) try: asyncio.run(_main()) diff --git a/private_gpt/arq/settings.py b/private_gpt/arq/settings.py index c7627b5f..24ef0b9e 100644 --- a/private_gpt/arq/settings.py +++ b/private_gpt/arq/settings.py @@ -2,20 +2,11 @@ from arq.connections import RedisSettings from private_gpt.settings.settings import Settings -CHAT_HEALTH_CHECK_KEY_PREFIX = "private_gpt:arq:health" -START_CHAT_TASK_NAME = "private_gpt.chat.start" -RESUME_ITERATION_TASK_NAME = "private_gpt.chat.resume_iteration" -TOOL_RESUME_TASK_NAME = "private_gpt.tool.resume" -CHAT_TIMEOUT_TASK_NAME = "private_gpt.chat.timeout" -CHAT_QUEUE_PREFIX = "private_gpt:arq:queue" +QUEUE_PREFIX = "private_gpt:arq:queue" -def get_queue_name(settings: Settings) -> str: - return f"{CHAT_QUEUE_PREFIX}:{settings.scheduler.chat.celery_queue}" - - -def get_health_check_key(settings: Settings) -> str: - return f"{CHAT_HEALTH_CHECK_KEY_PREFIX}:{settings.scheduler.chat.celery_queue}" +def get_queue_name(queue: str) -> str: + return f"{QUEUE_PREFIX}:{queue}" def get_redis_settings(settings: Settings) -> RedisSettings: diff --git a/private_gpt/arq/task_registry.py b/private_gpt/arq/task_registry.py index 6f584b15..dc9b66be 100644 --- a/private_gpt/arq/task_registry.py +++ b/private_gpt/arq/task_registry.py @@ -10,14 +10,10 @@ def register_task_packages(*task_packages: str) -> None: def get_task_packages(*task_packages: str) -> tuple[str, ...]: - packages = list(BASE_TASK_PACKAGES) + packages = list(task_packages or BASE_TASK_PACKAGES) for task_package in _EXTERNAL_TASK_PACKAGES: if task_package not in packages: packages.append(task_package) - for task_package in task_packages: - if task_package not in packages: - packages.append(task_package) - return tuple(packages) diff --git a/private_gpt/arq/tasks/chat/__init__.py b/private_gpt/arq/tasks/chat/__init__.py index 00185aaa..ac4aeabb 100644 --- a/private_gpt/arq/tasks/chat/__init__.py +++ b/private_gpt/arq/tasks/chat/__init__.py @@ -2,6 +2,8 @@ import asyncio from typing import TYPE_CHECKING, Any from private_gpt.arq.enqueue import abort_job +from private_gpt.arq.tasks.chat.settings import get_queue_name +from private_gpt.settings.settings import settings if TYPE_CHECKING: from private_gpt.arq.tasks.chat.resume import ( @@ -14,8 +16,14 @@ if TYPE_CHECKING: async def abort_chat_job(*, correlation_id: str) -> bool: results = await asyncio.gather( - abort_job(job_id=f"{correlation_id}:start"), - abort_job(job_id=f"{correlation_id}:resume"), + abort_job( + job_id=f"{correlation_id}:start", + queue_name=get_queue_name(settings()), + ), + abort_job( + job_id=f"{correlation_id}:resume", + queue_name=get_queue_name(settings()), + ), ) return any(results) diff --git a/private_gpt/arq/tasks/chat/resume.py b/private_gpt/arq/tasks/chat/resume.py index a9bea291..dd7a364d 100644 --- a/private_gpt/arq/tasks/chat/resume.py +++ b/private_gpt/arq/tasks/chat/resume.py @@ -1,15 +1,17 @@ from typing import Any from private_gpt.arq.enqueue import enqueue_job -from private_gpt.arq.settings import ( +from private_gpt.arq.tasks import arq_task +from private_gpt.arq.tasks.chat.settings import ( CHAT_TIMEOUT_TASK_NAME, RESUME_ITERATION_TASK_NAME, TOOL_RESUME_TASK_NAME, + get_queue_name, ) -from private_gpt.arq.tasks import arq_task from private_gpt.components.engines.chat.async_chat_engine import AsyncChatEngine from private_gpt.di import get_global_injector from private_gpt.server.chat.chat_service import ChatService +from private_gpt.settings.settings import settings async def enqueue_resume_iteration_job( @@ -17,6 +19,7 @@ async def enqueue_resume_iteration_job( ) -> None: await enqueue_job( task_name=RESUME_ITERATION_TASK_NAME, + queue_name=get_queue_name(settings()), args=(correlation_id,), job_id=job_id or f"{correlation_id}:resume", correlation_id=correlation_id, @@ -32,6 +35,7 @@ async def enqueue_chat_timeout_job( ) -> None: await enqueue_job( task_name=CHAT_TIMEOUT_TASK_NAME, + queue_name=get_queue_name(settings()), args=(correlation_id, checkpoint_id), job_id=job_id, correlation_id=correlation_id, @@ -44,6 +48,7 @@ async def enqueue_tool_resume_job( ) -> None: await enqueue_job( task_name=TOOL_RESUME_TASK_NAME, + queue_name=get_queue_name(settings()), args=(correlation_id, tool_id, result), correlation_id=correlation_id, ) diff --git a/private_gpt/arq/tasks/chat/settings.py b/private_gpt/arq/tasks/chat/settings.py new file mode 100644 index 00000000..55b53346 --- /dev/null +++ b/private_gpt/arq/tasks/chat/settings.py @@ -0,0 +1,11 @@ +from private_gpt.arq.settings import get_queue_name as build_queue_name +from private_gpt.settings.settings import Settings + +START_CHAT_TASK_NAME = "private_gpt.chat.start" +RESUME_ITERATION_TASK_NAME = "private_gpt.chat.resume_iteration" +TOOL_RESUME_TASK_NAME = "private_gpt.tool.resume" +CHAT_TIMEOUT_TASK_NAME = "private_gpt.chat.timeout" + + +def get_queue_name(settings: Settings) -> str: + return build_queue_name(settings.scheduler.chat.celery_queue) diff --git a/private_gpt/arq/tasks/chat/start.py b/private_gpt/arq/tasks/chat/start.py index 3bd89497..b2140c86 100644 --- a/private_gpt/arq/tasks/chat/start.py +++ b/private_gpt/arq/tasks/chat/start.py @@ -1,10 +1,11 @@ from typing import Any from private_gpt.arq.enqueue import enqueue_job -from private_gpt.arq.settings import START_CHAT_TASK_NAME from private_gpt.arq.tasks import arq_task +from private_gpt.arq.tasks.chat.settings import START_CHAT_TASK_NAME, get_queue_name from private_gpt.di import get_global_injector from private_gpt.server.chat.chat_service import ChatService +from private_gpt.settings.settings import settings async def enqueue_start_chat_job( @@ -17,6 +18,7 @@ async def enqueue_start_chat_job( ) -> None: await enqueue_job( task_name=START_CHAT_TASK_NAME, + queue_name=get_queue_name(settings()), args=(request_data, correlation_id, stream_type, metadata), job_id=job_id or f"{correlation_id}:start", correlation_id=correlation_id, diff --git a/private_gpt/arq/worker.py b/private_gpt/arq/worker.py deleted file mode 100644 index 5638cf61..00000000 --- a/private_gpt/arq/worker.py +++ /dev/null @@ -1,10 +0,0 @@ -from private_gpt.arq.hooks import ( - on_job_end, # noqa: F401 — used by runner via worker_module.on_job_end -) -from private_gpt.arq.lifecycle import ( # noqa: F401 — used by runner via worker_module.* - shutdown, - startup, -) -from private_gpt.arq.tasks import autodiscover_tasks - -functions = autodiscover_tasks() diff --git a/private_gpt/celery/base.py b/private_gpt/celery/base.py index f91e1782..757d4c72 100644 --- a/private_gpt/celery/base.py +++ b/private_gpt/celery/base.py @@ -336,7 +336,11 @@ class StatefulBackgroundTask(_BackgroundTask): from private_gpt.eager_loading import warm injector = get_global_injector(allow_to_generate_new_injectors=True) - profile = os.environ.get("PGPT_STATEFUL_WORKER_TYPE", "chat").strip() + profile = os.environ.get("PGPT_WORKER_WARM_PROFILE", "").strip() + if not profile: + raise ValueError( + "PGPT_WORKER_WARM_PROFILE is required for stateful workers" + ) warm(injector, profile=profile) @classmethod diff --git a/private_gpt/celery/celery.py b/private_gpt/celery/celery.py index bf163896..f5b8c32d 100644 --- a/private_gpt/celery/celery.py +++ b/private_gpt/celery/celery.py @@ -1,3 +1,4 @@ +import os from collections.abc import Callable, Sequence from typing import Any @@ -10,6 +11,15 @@ from private_gpt.initialize import initialize_globals, initialize_observability from private_gpt.settings.settings import Settings, settings +def _configured_task_packages(task_packages: Sequence[str]) -> tuple[str, ...]: + if task_packages: + return tuple(task_packages) + configured = os.environ.get("PGPT_CELERY_TASK_PACKAGES", "") + return tuple( + package.strip() for package in configured.split(",") if package.strip() + ) + + def _autodiscover_tasks( celery_app: Celery, task_packages: Sequence[str], *, force: bool = False ) -> None: @@ -52,7 +62,9 @@ def create_app( if before_setup is not None: before_setup() - resolved_task_packages = get_task_packages(*task_packages) + resolved_task_packages = get_task_packages( + *_configured_task_packages(task_packages) + ) celery_app = Celery( app_name, diff --git a/private_gpt/celery/healthcheck.py b/private_gpt/celery/healthcheck.py index 858fb22e..fe42038d 100644 --- a/private_gpt/celery/healthcheck.py +++ b/private_gpt/celery/healthcheck.py @@ -45,21 +45,25 @@ async def check_worker() -> bool: async def health_check() -> dict[str, Any]: - mode = os.getenv("PGPT_WORKER_MODE", "mixed") + mode = os.getenv("PGPT_WORKER_MODE", "").strip().lower() status: dict[str, Any] = {"status": "healthy", "mode": mode, "services": {}} - if mode in ["mixed", "flower"]: + if mode == "flower": flower_health = await check_flower() status["services"]["flower"] = "healthy" if flower_health else "unhealthy" if not flower_health: status["status"] = "unhealthy" - if mode in ["mixed", "worker", "celery"]: + elif mode == "celery": worker_health = await check_worker() status["services"]["worker"] = "healthy" if worker_health else "unhealthy" if not worker_health: status["status"] = "unhealthy" + else: + status["status"] = "unhealthy" + status["error"] = f"Unsupported worker mode: {mode or ''}" + return status diff --git a/private_gpt/celery/task_registry.py b/private_gpt/celery/task_registry.py index a33b113d..1d66b024 100644 --- a/private_gpt/celery/task_registry.py +++ b/private_gpt/celery/task_registry.py @@ -1,4 +1,7 @@ -BASE_TASK_PACKAGES = ("private_gpt.celery.tasks",) +BASE_TASK_PACKAGES = ( + "private_gpt.celery.tasks.ingestion", + "private_gpt.celery.tasks.tools", +) _EXTERNAL_TASK_PACKAGES: list[str] = [] @@ -10,14 +13,10 @@ def register_task_packages(*task_packages: str) -> None: def get_task_packages(*task_packages: str) -> tuple[str, ...]: - packages = list(BASE_TASK_PACKAGES) + packages = list(task_packages or BASE_TASK_PACKAGES) for task_package in _EXTERNAL_TASK_PACKAGES: if task_package not in packages: packages.append(task_package) - for task_package in task_packages: - if task_package not in packages: - packages.append(task_package) - return tuple(packages) diff --git a/private_gpt/celery/tasks/__init__.py b/private_gpt/celery/tasks/__init__.py index 373b2099..c9c2ef67 100644 --- a/private_gpt/celery/tasks/__init__.py +++ b/private_gpt/celery/tasks/__init__.py @@ -1,6 +1 @@ -from private_gpt.celery.tasks.ingestion import __all__ as ingestion_all -from private_gpt.celery.tasks.tools import __all__ as tools_all - __all__: list[str] = [] -__all__.extend(ingestion_all) -__all__.extend(tools_all) diff --git a/private_gpt/cli/commands/worker.py b/private_gpt/cli/commands/worker.py index 956a58de..6df9b07a 100644 --- a/private_gpt/cli/commands/worker.py +++ b/private_gpt/cli/commands/worker.py @@ -1,180 +1,6 @@ -import os -import signal -import subprocess -import sys - -import typer - -from private_gpt.arq.runner import run_arq_worker -from private_gpt.settings.settings import settings - - -def _resolve_worker_mode() -> str: - env_mode = os.environ.get("PGPT_WORKER_MODE", "").strip().lower() - if env_mode: - return env_mode - - current_settings = settings() - if current_settings.scheduler.chat.mode == "arq": - return "arq" - return "celery" - - -def _build_flower_args(mode: str) -> list[str]: - args: list[str] = [] - url_prefix = os.environ.get("PGPT_FLOWER_URL_PREFIX", "") - port = os.environ.get("PGPT_FLOWER_PORT", "5555") - password = os.environ.get("PGPT_FLOWER_PASSWORD", "flower") - persistent = os.environ.get("PGPT_FLOWER_PERSISTENT", "true") - persistent_db = os.environ.get("PGPT_FLOWER_PERSISTENT_DB", "celery_flower.db") - max_tasks = os.environ.get("PGPT_FLOWER_MAXIMUM_TASKS", "1000") - - if url_prefix: - args.append(f"--url_prefix={url_prefix}") - if port: - args.append(f"--port={port}") - if password: - args.append(f"--basic_auth=flower:{password}") - if mode == "mixed": - args.append("--logging=none") - if persistent.lower() == "true": - args.append("--persistent=True") - if persistent_db: - args.append(f"--db={persistent_db}") - if max_tasks: - args.append(f"--max-tasks={max_tasks}") - return args - - -def _build_worker_args(max_tasks_per_child: int = 1000) -> list[str]: - args: list[str] = [] - log_level = os.environ.get("PGPT_CELERY_LOG_LEVEL", "info") - queues = os.environ.get("PGPT_CELERY_QUEUES", "") - hostname = os.environ.get("PGPT_CELERY_HOSTNAME", "default") - pool = os.environ.get("PGPT_CELERY_POOL", "prefork") - concurrency = os.environ.get("PGPT_CELERY_CONCURRENCY", "") - time_limit = os.environ.get("PGPT_CELERY_TIME_LIMIT", "") - stateful_type = os.environ.get("PGPT_STATEFUL_WORKER_TYPE", "").strip() - - if log_level: - args.append(f"--loglevel={log_level}") - - if stateful_type: - queues = queues or stateful_type - hostname = hostname or stateful_type - args.append(f"--max-tasks-per-child={max_tasks_per_child}") - - if queues: - args.append(f"--queues={queues}") - if hostname: - args.append(f"--hostname={hostname}%h") - if pool: - args.append(f"--pool={pool}") - if concurrency: - args.append(f"--concurrency={concurrency}") - if time_limit: - args.append(f"--time-limit={time_limit}") - return args +from private_gpt.worker import run_private_gpt_worker def worker_command() -> None: - """Start a background worker process. - - Behaviour is fully controlled through environment variables. - """ - mode = _resolve_worker_mode() - - if mode == "arq": - run_arq_worker() - return - - app_module = os.environ.get("PGPT_WORKER_APP_MODULE", "private_gpt") - celery_app = f"{app_module}.celery" - healthcheck_app = f"{app_module}.celery.healthcheck:app" - procs: list[subprocess.Popen[bytes]] = [] - - def _cleanup(signum: int = 0, frame: object = None) -> None: - typer.echo("Shutting down services...") - for p in procs: - p.terminate() - for p in procs: - try: - p.wait(timeout=10) - except subprocess.TimeoutExpired: - p.kill() - sys.exit(0) - - signal.signal(signal.SIGTERM, _cleanup) - signal.signal(signal.SIGINT, _cleanup) - - if mode in ("flower", "mixed"): - flower_args = _build_flower_args(mode) - typer.echo( - f"Starting flower on port {os.environ.get('PGPT_FLOWER_PORT', '5555')}" - ) - procs.append( - subprocess.Popen( - [ - sys.executable, - "-m", - "celery", - "--app", - celery_app, - "flower", - *flower_args, - ] - ) - ) - - if mode in ("celery", "mixed"): - worker_args = _build_worker_args(settings().celery.max_tasks_per_child) - typer.echo(f"Starting celery worker with args: {' '.join(worker_args)}") - # The PGPT_STATEFUL_WORKER_TYPE env var triggers eager warm-up in the - # worker process via bootsteps.py. - procs.append( - subprocess.Popen( - [ - sys.executable, - "-m", - "celery", - "--app", - celery_app, - "worker", - *worker_args, - ] - ) - ) - - if os.environ.get("API_ENABLED", "true").lower() == "true": - api_port = os.environ.get("API_PORT", "8090") - typer.echo(f"Starting healthcheck server on port {api_port}") - procs.append( - subprocess.Popen( - [ - sys.executable, - "-m", - "uvicorn", - healthcheck_app, - "--host", - "0.0.0.0", - "--port", - api_port, - "--no-access-log", - "--log-level", - "critical", - ] - ) - ) - - if not procs: - typer.echo( - f"No processes started. Check PGPT_WORKER_MODE={mode!r}. Supported: mixed, arq, celery", - err=True, - ) - raise SystemExit(1) - - try: - for p in procs: - p.wait() - except KeyboardInterrupt: - _cleanup() + """Start the worker mode selected by ``PGPT_WORKER_MODE``.""" + run_private_gpt_worker() diff --git a/private_gpt/components/engines/chat/checkpoint_store.py b/private_gpt/components/engines/chat/checkpoint_store.py index 6c40192d..eedfb61a 100644 --- a/private_gpt/components/engines/chat/checkpoint_store.py +++ b/private_gpt/components/engines/chat/checkpoint_store.py @@ -120,7 +120,7 @@ class ChatCheckpointStoreFactory: def get(self) -> ChatCheckpointStore: if self._settings.scheduler.chat.mode == "arq": - from private_gpt.arq.iteration_state import RedisChatCheckpointStore + from private_gpt.arq.chat.iteration_state import RedisChatCheckpointStore return cast( ChatCheckpointStore, diff --git a/private_gpt/worker/__init__.py b/private_gpt/worker/__init__.py new file mode 100644 index 00000000..c483949f --- /dev/null +++ b/private_gpt/worker/__init__.py @@ -0,0 +1,12 @@ +from collections.abc import Sequence + +from private_gpt.worker.modes import register_private_gpt_worker_modes +from private_gpt.worker.registry import register_worker_mode, run_worker + + +def run_private_gpt_worker(args: Sequence[str] = ()) -> None: + register_private_gpt_worker_modes() + run_worker(args) + + +__all__ = ["register_worker_mode", "run_private_gpt_worker", "run_worker"] diff --git a/private_gpt/worker/__main__.py b/private_gpt/worker/__main__.py new file mode 100644 index 00000000..ea3422a2 --- /dev/null +++ b/private_gpt/worker/__main__.py @@ -0,0 +1,6 @@ +import sys + +from private_gpt.worker import run_private_gpt_worker + +if __name__ == "__main__": + run_private_gpt_worker(sys.argv[1:]) diff --git a/private_gpt/worker/modes.py b/private_gpt/worker/modes.py new file mode 100644 index 00000000..37c73a24 --- /dev/null +++ b/private_gpt/worker/modes.py @@ -0,0 +1,177 @@ +import importlib +import os +import signal +import subprocess +import sys +from collections.abc import Callable, Sequence +from functools import partial + +import typer + +from private_gpt.settings.settings import settings +from private_gpt.worker.registry import register_worker_mode + + +def _app_module() -> str: + return os.environ.get("PGPT_WORKER_APP_MODULE", "private_gpt") + + +def _build_flower_args() -> list[str]: + args: list[str] = [] + url_prefix = os.environ.get("PGPT_FLOWER_URL_PREFIX", "") + port = os.environ.get("PGPT_FLOWER_PORT", "5555") + password = os.environ.get("PGPT_FLOWER_PASSWORD", "flower") + persistent = os.environ.get("PGPT_FLOWER_PERSISTENT", "true") + persistent_db = os.environ.get("PGPT_FLOWER_PERSISTENT_DB", "celery_flower.db") + max_tasks = os.environ.get("PGPT_FLOWER_MAXIMUM_TASKS", "1000") + + if url_prefix: + args.append(f"--url_prefix={url_prefix}") + if port: + args.append(f"--port={port}") + if password: + args.append(f"--basic_auth=flower:{password}") + if persistent.lower() == "true": + args.append("--persistent=True") + if persistent_db: + args.append(f"--db={persistent_db}") + if max_tasks: + args.append(f"--max-tasks={max_tasks}") + return args + + +def _build_celery_args(max_tasks_per_child: int) -> list[str]: + args: list[str] = [] + log_level = os.environ.get("PGPT_CELERY_LOG_LEVEL", "info") + queues = os.environ.get("PGPT_CELERY_QUEUES", "") + hostname = os.environ.get("PGPT_CELERY_HOSTNAME", "default") + pool = os.environ.get("PGPT_CELERY_POOL", "prefork") + concurrency = os.environ.get("PGPT_CELERY_CONCURRENCY", "") + time_limit = os.environ.get("PGPT_CELERY_TIME_LIMIT", "") + stateful_type = os.environ.get("PGPT_STATEFUL_WORKER_TYPE", "").strip() + + if log_level: + args.append(f"--loglevel={log_level}") + if stateful_type: + queues = queues or stateful_type + hostname = hostname or stateful_type + args.append(f"--max-tasks-per-child={max_tasks_per_child}") + if queues: + args.append(f"--queues={queues}") + if hostname: + args.append(f"--hostname={hostname}%h") + if pool: + args.append(f"--pool={pool}") + if concurrency: + args.append(f"--concurrency={concurrency}") + if time_limit: + args.append(f"--time-limit={time_limit}") + return args + + +def _run_processes(commands: Sequence[Sequence[str]]) -> None: + processes = [subprocess.Popen(command) for command in commands] + + def cleanup(signum: int = 0, frame: object = None) -> None: + del signum, frame + typer.echo("Shutting down services...") + for process in processes: + process.terminate() + for process in processes: + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + + signal.signal(signal.SIGTERM, cleanup) + signal.signal(signal.SIGINT, cleanup) + + try: + for process in processes: + process.wait() + except KeyboardInterrupt: + cleanup() + + +def _healthcheck_command() -> list[str]: + return [ + sys.executable, + "-m", + "uvicorn", + f"{_app_module()}.celery.healthcheck:app", + "--host", + "0.0.0.0", + "--port", + os.environ.get("API_PORT", "8090"), + "--no-access-log", + "--log-level", + "critical", + ] + + +def _with_healthcheck(command: list[str]) -> list[list[str]]: + commands = [command] + if os.environ.get("API_ENABLED", "true").lower() == "true": + commands.append(_healthcheck_command()) + return commands + + +def run_arq(args: Sequence[str]) -> None: + if args: + raise ValueError("The arq worker mode does not support arguments yet") + arq_app = importlib.import_module(f"{_app_module()}.arq") + arq_app.run_arq_worker() + + +def run_celery( + args: Sequence[str], + *, + max_tasks_per_child_resolver: Callable[[], int] = ( + lambda: settings().celery.max_tasks_per_child + ), +) -> None: + celery_args = _build_celery_args(max_tasks_per_child_resolver()) + typer.echo(f"Starting celery worker with args: {' '.join(celery_args)}") + command = [ + sys.executable, + "-m", + "celery", + "--app", + f"{_app_module()}.celery", + "worker", + *celery_args, + *args, + ] + _run_processes(_with_healthcheck(command)) + + +def run_flower(args: Sequence[str]) -> None: + flower_args = _build_flower_args() + typer.echo(f"Starting flower on port {os.environ.get('PGPT_FLOWER_PORT', '5555')}") + command = [ + sys.executable, + "-m", + "celery", + "--app", + f"{_app_module()}.celery", + "flower", + *flower_args, + *args, + ] + _run_processes(_with_healthcheck(command)) + + +def register_private_gpt_worker_modes( + max_tasks_per_child_resolver: Callable[[], int] = ( + lambda: settings().celery.max_tasks_per_child + ), +) -> None: + register_worker_mode("arq", run_arq) + register_worker_mode( + "celery", + partial( + run_celery, + max_tasks_per_child_resolver=max_tasks_per_child_resolver, + ), + ) + register_worker_mode("flower", run_flower) diff --git a/private_gpt/worker/registry.py b/private_gpt/worker/registry.py new file mode 100644 index 00000000..4405479a --- /dev/null +++ b/private_gpt/worker/registry.py @@ -0,0 +1,31 @@ +import os +from collections.abc import Callable, Sequence + +WorkerModeHandler = Callable[[Sequence[str]], None] + +_worker_modes: dict[str, WorkerModeHandler] = {} + + +def register_worker_mode(name: str, handler: WorkerModeHandler) -> None: + normalized_name = name.strip().lower() + if not normalized_name: + raise ValueError("Worker mode name cannot be empty") + _worker_modes[normalized_name] = handler + + +def get_worker_mode(name: str) -> WorkerModeHandler: + normalized_name = name.strip().lower() + try: + return _worker_modes[normalized_name] + except KeyError as exc: + supported_modes = ", ".join(sorted(_worker_modes)) + raise ValueError( + f"Unsupported PGPT_WORKER_MODE={name!r}. Registered modes: {supported_modes}" + ) from exc + + +def run_worker(args: Sequence[str] = ()) -> None: + mode = os.environ.get("PGPT_WORKER_MODE", "").strip() + if not mode: + raise ValueError("PGPT_WORKER_MODE is required") + get_worker_mode(mode)(args) diff --git a/scripts/worker_entrypoint b/scripts/worker_entrypoint index 972d7a1f..c4d91ffc 100755 --- a/scripts/worker_entrypoint +++ b/scripts/worker_entrypoint @@ -1,122 +1,7 @@ #!/bin/bash +set -euo pipefail -: "${PGPT_WORKER_MODE:=mixed}" # mixed, celery, flower, arq -: "${PGPT_FLOWER_PORT:=5555}" -: "${PGPT_FLOWER_PASSWORD:=flower}" -: "${PGPT_FLOWER_URL_PREFIX:=}" -: "${PGPT_FLOWER_PERSISTENT:=true}" -: "${PGPT_FLOWER_PERSISTENT_DB:=celery_flower.db}" -: "${PGPT_FLOWER_MAXIMUM_TASKS:=1000}" -: "${PGPT_CELERY_LOG_LEVEL:=info}" -: "${PGPT_CELERY_QUEUES:=}" -: "${PGPT_CELERY_HOSTNAME:=default}" -: "${PGPT_CELERY_POOL:=prefork}" -: "${PGPT_CELERY_CONCURRENCY:=}" -: "${PGPT_CELERY_TIME_LIMIT:=}" +: "${PGPT_WORKER_MODE:?PGPT_WORKER_MODE is required}" : "${PGPT_WORKER_APP_MODULE:=private_gpt}" -: "${PGPT_STATEFUL_WORKER_TYPE:=}" -: "${PGPT_MAX_WORKERS:=}" -: "${PGPT_CELERY_MAX_TASKS_PER_CHILD:=1000}" -: "${API_ENABLED:=true}" -: "${API_PORT:=8090}" -CELERY_APP_MODULE="${PGPT_WORKER_APP_MODULE}.celery" -HEALTHCHECK_APP_MODULE="${PGPT_WORKER_APP_MODULE}.celery.healthcheck:app" -ARQ_HEALTHCHECK_APP_MODULE="${PGPT_WORKER_APP_MODULE}.arq.healthcheck:app" - -if [[ "$PGPT_WORKER_MODE" == "arq" ]]; then - exec python -m private_gpt.arq.runner -fi - -# Create directory for PID files -mkdir -p /tmp/private_gpt_pids - -# Add signal handling -cleanup() { - echo "Shutting down services..." - jobs -p | xargs -r kill - exit 0 -} -trap cleanup EXIT INT TERM - -# Run flower on the specified port -if [[ "$PGPT_WORKER_MODE" == "flower" || "$PGPT_WORKER_MODE" == "mixed" ]]; then - echo "Starting flower on port $PGPT_FLOWER_PORT" - FLOWER_ARGS=() - if [[ -n "$PGPT_FLOWER_URL_PREFIX" ]]; then - FLOWER_ARGS+=("--url_prefix=$PGPT_FLOWER_URL_PREFIX") - fi - if [[ -n "$PGPT_FLOWER_PORT" ]]; then - FLOWER_ARGS+=("--port=$PGPT_FLOWER_PORT") - fi - if [[ -n "$PGPT_FLOWER_PASSWORD" ]]; then - FLOWER_ARGS+=("--basic_auth=flower:$PGPT_FLOWER_PASSWORD") - fi - if [[ "$PGPT_WORKER_MODE" == "mixed" ]]; then - FLOWER_ARGS+=("--logging=none") - fi - if [[ "$PGPT_FLOWER_PERSISTENT" == "true" ]]; then - FLOWER_ARGS+=("--persistent=True") - if [[ -n "$PGPT_FLOWER_PERSISTENT_DB" ]]; then - FLOWER_ARGS+=("--db=$PGPT_FLOWER_PERSISTENT_DB") - fi - if [[ -n "$PGPT_FLOWER_MAXIMUM_TASKS" ]]; then - FLOWER_ARGS+=("--max-tasks=$PGPT_FLOWER_MAXIMUM_TASKS") - fi - fi - - # print all args except password - args=("${FLOWER_ARGS[@]}") - args=("${args[@]//:$PGPT_FLOWER_PASSWORD/}") - - echo "Starting flower with args: ${args[*]}" - python -m celery --app "$CELERY_APP_MODULE" flower "${FLOWER_ARGS[@]}" & -fi - -# Run celery worker using the config. -# Stateful workers (PGPT_STATEFUL_WORKER_TYPE set) keep warm DI/runtime state -# and recycle children after the configured task count. -if [[ "$PGPT_WORKER_MODE" == "celery" || "$PGPT_WORKER_MODE" == "mixed" ]]; then - WORKER_ARGS=() - if [[ -n "$PGPT_CELERY_LOG_LEVEL" ]]; then - WORKER_ARGS+=("--loglevel=$PGPT_CELERY_LOG_LEVEL") - fi - - QUEUES="$PGPT_CELERY_QUEUES" - HOSTNAME="${PGPT_CELERY_HOSTNAME}" - - # If stateful worker type is set, derive defaults. - if [[ -n "$PGPT_STATEFUL_WORKER_TYPE" ]]; then - QUEUES="${QUEUES:-$PGPT_STATEFUL_WORKER_TYPE}" - HOSTNAME="${HOSTNAME:-$PGPT_STATEFUL_WORKER_TYPE}" - WORKER_ARGS+=("--max-tasks-per-child=$PGPT_CELERY_MAX_TASKS_PER_CHILD") - fi - - if [[ -n "$QUEUES" ]]; then - WORKER_ARGS+=("--queues=$QUEUES") - fi - if [[ -n "$HOSTNAME" ]]; then - WORKER_ARGS+=("--hostname=${HOSTNAME}%h") - fi - if [[ -n "$PGPT_CELERY_POOL" ]]; then - WORKER_ARGS+=("--pool=$PGPT_CELERY_POOL") - fi - if [[ -n "$PGPT_CELERY_CONCURRENCY" ]]; then - WORKER_ARGS+=("--concurrency=$PGPT_CELERY_CONCURRENCY") - fi - if [[ -n "$PGPT_CELERY_TIME_LIMIT" ]]; then - WORKER_ARGS+=("--time-limit=$PGPT_CELERY_TIME_LIMIT") - fi - - echo "Starting celery worker with args: ${WORKER_ARGS[*]}" - python -m celery --app "$CELERY_APP_MODULE" worker "${WORKER_ARGS[@]}" & -fi - -# Start the Healthcheck server -if [[ "$API_ENABLED" == "true" ]]; then - echo "Starting healthcheck server on port $API_PORT" - python -m uvicorn "$HEALTHCHECK_APP_MODULE" --host 0.0.0.0 --port $API_PORT --no-access-log --log-level critical & -fi - -# Wait for all background processes -wait +exec python -m "${PGPT_WORKER_APP_MODULE}.worker" "$@" diff --git a/tests/arq/tasks/chat/test_start.py b/tests/arq/tasks/chat/test_start.py index 912d30d9..1f6cca46 100644 --- a/tests/arq/tasks/chat/test_start.py +++ b/tests/arq/tasks/chat/test_start.py @@ -2,8 +2,12 @@ from unittest.mock import AsyncMock import pytest -from private_gpt.arq.settings import START_CHAT_TASK_NAME from private_gpt.arq.tasks.chat import enqueue_start_chat_job +from private_gpt.arq.tasks.chat.settings import ( + START_CHAT_TASK_NAME, + get_queue_name, +) +from private_gpt.settings.settings import settings @pytest.mark.parametrize( @@ -37,6 +41,7 @@ async def test_enqueue_start_chat_job_dispatches_generic_arq_job( enqueue_job.assert_awaited_once_with( task_name=START_CHAT_TASK_NAME, + queue_name=get_queue_name(settings()), args=( request_data, "execution-id", diff --git a/tests/arq/test_healthcheck.py b/tests/arq/test_healthcheck.py new file mode 100644 index 00000000..35638246 --- /dev/null +++ b/tests/arq/test_healthcheck.py @@ -0,0 +1,9 @@ +from private_gpt.arq.healthcheck import health + + +async def test_healthcheck_does_not_require_redis() -> None: + assert await health() == { + "status": "healthy", + "mode": "arq-worker", + "services": {"worker": "healthy"}, + } diff --git a/tests/arq/test_runner.py b/tests/arq/test_runner.py new file mode 100644 index 00000000..33bb1589 --- /dev/null +++ b/tests/arq/test_runner.py @@ -0,0 +1,37 @@ +import pytest + +from private_gpt.arq.runner import _queue_name, _task_packages + + +def test_task_packages_are_loaded_from_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "PGPT_ARQ_TASK_PACKAGES", + "private_gpt.arq.tasks.chat, custom.tasks ", + ) + + assert _task_packages() == ( + "private_gpt.arq.tasks.chat", + "custom.tasks", + ) + + +def test_task_packages_are_required(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PGPT_ARQ_TASK_PACKAGES", raising=False) + + with pytest.raises(ValueError, match="PGPT_ARQ_TASK_PACKAGES"): + _task_packages() + + +def test_queue_is_loaded_from_environment(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PGPT_ARQ_QUEUE", "chat") + + assert _queue_name() == "private_gpt:arq:queue:chat" + + +def test_queue_is_required(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PGPT_ARQ_QUEUE", raising=False) + + with pytest.raises(ValueError, match="PGPT_ARQ_QUEUE"): + _queue_name() diff --git a/tests/arq/test_task_registry.py b/tests/arq/test_task_registry.py new file mode 100644 index 00000000..2207135f --- /dev/null +++ b/tests/arq/test_task_registry.py @@ -0,0 +1,20 @@ +from private_gpt.arq.task_registry import get_task_packages +from private_gpt.arq.tasks import autodiscover_registered_tasks + + +def test_explicit_task_packages_replace_defaults() -> None: + assert get_task_packages("private_gpt.arq.tasks.chat") == ( + "private_gpt.arq.tasks.chat", + ) + + +def test_chat_package_discovers_only_chat_tasks() -> None: + assert { + task.name + for task in autodiscover_registered_tasks("private_gpt.arq.tasks.chat") + } == { + "private_gpt.chat.resume_iteration", + "private_gpt.chat.start", + "private_gpt.chat.timeout", + "private_gpt.tool.resume", + } diff --git a/tests/celery/test_chat_worker.py b/tests/celery/test_chat_worker.py index f51cc563..f4f57e69 100644 --- a/tests/celery/test_chat_worker.py +++ b/tests/celery/test_chat_worker.py @@ -1,5 +1,8 @@ import asyncio from typing import Any +from unittest.mock import Mock + +import pytest from private_gpt.celery.base import StatefulBackgroundTask @@ -25,3 +28,30 @@ def test_stateful_task_reuses_one_event_loop() -> None: assert first_loop_id == second_loop_id finally: LoopIdTask.shutdown_runtime() + + +async def test_stateful_worker_uses_explicit_warm_profile( + monkeypatch: pytest.MonkeyPatch, +) -> None: + injector = Mock() + warm = Mock() + monkeypatch.setenv("PGPT_STATEFUL_WORKER_TYPE", "stateful-type") + monkeypatch.setenv("PGPT_WORKER_WARM_PROFILE", "tools") + monkeypatch.setattr( + "private_gpt.celery.base.get_global_injector", + Mock(return_value=injector), + ) + monkeypatch.setattr("private_gpt.eager_loading.warm", warm) + + await StatefulBackgroundTask._warm_async() + + warm.assert_called_once_with(injector, profile="tools") + + +async def test_stateful_worker_requires_explicit_warm_profile( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("PGPT_WORKER_WARM_PROFILE", raising=False) + + with pytest.raises(ValueError, match="PGPT_WORKER_WARM_PROFILE"): + await StatefulBackgroundTask._warm_async() diff --git a/tests/celery/test_healthcheck.py b/tests/celery/test_healthcheck.py new file mode 100644 index 00000000..286c7c4d --- /dev/null +++ b/tests/celery/test_healthcheck.py @@ -0,0 +1,32 @@ +import pytest + +from private_gpt.celery import healthcheck + + +@pytest.mark.parametrize("mode", ["", "mixed", "worker", "unknown"]) +async def test_healthcheck_rejects_unregistered_modes( + monkeypatch: pytest.MonkeyPatch, + mode: str, +) -> None: + monkeypatch.setenv("PGPT_WORKER_MODE", mode) + + status = await healthcheck.health_check() + + assert status["status"] == "unhealthy" + assert status["services"] == {} + + +async def test_healthcheck_checks_only_celery_worker( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("PGPT_WORKER_MODE", "celery") + monkeypatch.setattr(healthcheck, "check_worker", lambda: _async_result(True)) + + status = await healthcheck.health_check() + + assert status["status"] == "healthy" + assert status["services"] == {"worker": "healthy"} + + +async def _async_result(value: bool) -> bool: + return value diff --git a/tests/celery/test_task_registry.py b/tests/celery/test_task_registry.py new file mode 100644 index 00000000..3c5b03e6 --- /dev/null +++ b/tests/celery/test_task_registry.py @@ -0,0 +1,52 @@ +import ast +import os +import subprocess +import sys + +import pytest + +from private_gpt.celery.task_registry import get_task_packages + + +def test_explicit_task_packages_replace_defaults() -> None: + assert get_task_packages("private_gpt.celery.tasks.tools") == ( + "private_gpt.celery.tasks.tools", + ) + + +@pytest.mark.parametrize( + ("task_package", "expected_tasks"), + [ + ( + "private_gpt.celery.tasks.ingestion", + { + "private_gpt.ingestion.delete", + "private_gpt.ingestion.vector_index", + }, + ), + ( + "private_gpt.celery.tasks.tools", + {"private_gpt.tools.run"}, + ), + ], +) +def test_worker_registers_only_configured_task_package( + task_package: str, + expected_tasks: set[str], +) -> None: + env = os.environ.copy() + env["PGPT_CELERY_TASK_PACKAGES"] = task_package + output = subprocess.check_output( + [ + sys.executable, + "-c", + "from private_gpt.celery.celery import celery_app; " + "celery_app.loader.import_default_modules(); " + "print(sorted(name for name in celery_app.tasks " + "if name.startswith('private_gpt.')))", + ], + env=env, + text=True, + ) + + assert set(ast.literal_eval(output.strip())) == expected_tasks diff --git a/tests/worker/test_modes.py b/tests/worker/test_modes.py new file mode 100644 index 00000000..f42dcabc --- /dev/null +++ b/tests/worker/test_modes.py @@ -0,0 +1,18 @@ +import pytest + +from private_gpt.worker import modes + + +def test_celery_mode_forwards_worker_arguments( + monkeypatch: pytest.MonkeyPatch, +) -> None: + commands: list[list[list[str]]] = [] + monkeypatch.setenv("API_ENABLED", "false") + monkeypatch.setattr(modes, "_run_processes", lambda value: commands.append(value)) + + modes.run_celery( + ["--without-gossip"], + max_tasks_per_child_resolver=lambda: 1000, + ) + + assert commands[0][0][-1] == "--without-gossip" diff --git a/tests/worker/test_registry.py b/tests/worker/test_registry.py new file mode 100644 index 00000000..d5c66146 --- /dev/null +++ b/tests/worker/test_registry.py @@ -0,0 +1,33 @@ +import pytest + +from private_gpt.worker.registry import ( + get_worker_mode, + register_worker_mode, + run_worker, +) + + +def test_registered_application_mode_can_be_dispatched( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + register_worker_mode("custom", lambda args: calls.extend(args)) + monkeypatch.setenv("PGPT_WORKER_MODE", "CUSTOM") + + run_worker(["custom"]) + + assert calls == ["custom"] + + +def test_worker_mode_is_required(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PGPT_WORKER_MODE", raising=False) + + with pytest.raises(ValueError, match="PGPT_WORKER_MODE is required"): + run_worker() + + +def test_unknown_worker_mode_lists_registered_modes() -> None: + register_worker_mode("known", lambda args: None) + + with pytest.raises(ValueError, match=r"Registered modes: .*known"): + get_worker_mode("unknown")