github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2026-07-16 17:00:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: worker config

Browse Source
This commit is contained in:
Javier Martinez
2026-07-14 17:58:45 +02:00
parent 8487812af0
commit 1371c275a1
38 changed files with 627 additions and 414 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
Makefile
View File

@@ -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)

24
fern/docs/pages/configuration/cli.mdx
View File

@@ -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
<Tabs>
<Tab title="macOS / Linux">
```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
</Tab>
<Tab title="Windows (PowerShell)">
```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
</Tab>
<Tab title="Windows (CMD)">
```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

3
private_gpt/arq/__init__.py
View File

@@ -0,0 +1,3 @@
from private_gpt.arq.runner import run_arq_worker
__all__ = ["run_arq_worker"]

4
private_gpt/arq/__main__.py Normal file
View File

@@ -0,0 +1,4 @@
from private_gpt.arq import run_arq_worker
if __name__ == "__main__":
run_arq_worker()

1
private_gpt/arq/chat/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""ARQ chat worker support."""

2
private_gpt/arq/iteration_state.py → private_gpt/arq/chat/iteration_state.py
View File

@@ -1,3 +1,5 @@
"""Redis-backed state for resumable ARQ chat execution."""
from __future__ import annotations
import json

8
private_gpt/arq/enqueue.py
View File

@@ -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:

22
private_gpt/arq/healthcheck.py
View File

@@ -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

5
private_gpt/arq/lifecycle.py
View File

@@ -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

63
private_gpt/arq/runner.py
View File

@@ -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())

15
private_gpt/arq/settings.py
View File

@@ -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:

6
private_gpt/arq/task_registry.py
View File

@@ -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)

12
private_gpt/arq/tasks/chat/__init__.py
View File

@@ -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)

9
private_gpt/arq/tasks/chat/resume.py
View File

@@ -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,
)

11
private_gpt/arq/tasks/chat/settings.py Normal file
View File

@@ -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)

4
private_gpt/arq/tasks/chat/start.py
View File

@@ -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,

10
private_gpt/arq/worker.py
View File

@@ -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()

6
private_gpt/celery/base.py
View File

@@ -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

14
private_gpt/celery/celery.py
View File

@@ -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,

10
private_gpt/celery/healthcheck.py
View File

@@ -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 '<empty>'}"
return status

11
private_gpt/celery/task_registry.py
View File

@@ -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)

5
private_gpt/celery/tasks/__init__.py
View File

@@ -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)

180
private_gpt/cli/commands/worker.py
View File

@@ -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()

2
private_gpt/components/engines/chat/checkpoint_store.py
View File

@@ -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,

12
private_gpt/worker/__init__.py Normal file
View File

@@ -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"]

6
private_gpt/worker/__main__.py Normal file
View File

@@ -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:])

177
private_gpt/worker/modes.py Normal file
View File

@@ -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)

31
private_gpt/worker/registry.py Normal file
View File

@@ -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)

121
scripts/worker_entrypoint
View File

@@ -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" "$@"

7
tests/arq/tasks/chat/test_start.py
View File

@@ -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",

9
tests/arq/test_healthcheck.py Normal file
View File

@@ -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"},
}

37
tests/arq/test_runner.py Normal file
View File

@@ -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()

20
tests/arq/test_task_registry.py Normal file
View File

@@ -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",
}

30
tests/celery/test_chat_worker.py
View File

@@ -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()

32
tests/celery/test_healthcheck.py Normal file
View File

@@ -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

52
tests/celery/test_task_registry.py Normal file
View File

@@ -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

18
tests/worker/test_modes.py Normal file
View File

@@ -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"

33
tests/worker/test_registry.py Normal file
View File

@@ -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")