112 lines
3.1 KiB
Python
112 lines
3.1 KiB
Python
 
|
"""Worker entrypoint (refactor §1).
|
||
|
|
|
||
|
|
One image, two roles selected by ``WORKER_ROLE``:
|
||
|
|
|
||
|
|
- ``jobs`` — arq consumer for the ``interactive`` and ``bulk`` queues (both in one
|
||
|
|
process, per-queue concurrency/timeout).
|
||
|
|
- ``sync`` — runs the account-sync supervisor (long-lived per-account loops).
|
||
|
|
- ``all`` — both, for single-process development (the default).
|
||
|
|
|
||
|
|
Run with: ``python -m workers.run``.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
import signal
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from arq.worker import Worker
|
||
|
|
|
||
|
|
from app.config import get_settings
|
||
|
|
from app.queue import QUEUE_BULK, QUEUE_INTERACTIVE, redis_settings
|
||
|
|
from workers.worker import (
|
||
|
|
BULK_FUNCTIONS,
|
||
|
|
INTERACTIVE_FUNCTIONS,
|
||
|
|
sync_on_shutdown,
|
||
|
|
sync_on_startup,
|
||
|
|
)
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
# arq requires a non-empty function set; the sync worker consumes a queue nothing
|
||
|
|
# is ever enqueued to, so it just keeps the event loop alive for the supervisor.
|
||
|
|
_SYNC_CONTROL_QUEUE = "sync:control"
|
||
|
|
|
||
|
|
|
||
|
|
async def _noop(ctx: dict[str, Any]) -> None:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _worker(
|
||
|
|
functions: list[Any],
|
||
|
|
queue_name: str,
|
||
|
|
max_jobs: int,
|
||
|
|
job_timeout: int,
|
||
|
|
on_startup: Any = None,
|
||
|
|
on_shutdown: Any = None,
|
||
|
|
) -> Worker:
|
||
|
|
return Worker(
|
||
|
|
functions=functions,
|
||
|
|
queue_name=queue_name,
|
||
|
|
redis_settings=redis_settings(),
|
||
|
|
max_jobs=max_jobs,
|
||
|
|
job_timeout=job_timeout,
|
||
|
|
on_startup=on_startup,
|
||
|
|
on_shutdown=on_shutdown,
|
||
|
|
handle_signals=False, # the runner owns signal handling (graceful shutdown)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def build_workers(role: str) -> list[Worker]:
|
||
|
|
"""Construct the arq Worker instances for the given role."""
|
||
|
|
s = get_settings()
|
||
|
|
workers: list[Worker] = []
|
||
|
|
if role in ("sync", "all"):
|
||
|
|
workers.append(
|
||
|
|
_worker([_noop], _SYNC_CONTROL_QUEUE, 1, 3600, sync_on_startup, sync_on_shutdown)
|
||
|
|
)
|
||
|
|
if role in ("jobs", "all"):
|
||
|
|
workers.append(
|
||
|
|
_worker(INTERACTIVE_FUNCTIONS, QUEUE_INTERACTIVE, s.interactive_concurrency, 120)
|
||
|
|
)
|
||
|
|
workers.append(_worker(BULK_FUNCTIONS, QUEUE_BULK, s.bulk_concurrency, 3600))
|
||
|
|
return workers
|
||
|
|
|
||
|
|
|
||
|
|
async def _run() -> None:
|
||
|
|
logging.basicConfig(level=get_settings().log_level)
|
||
|
|
role = get_settings().worker_role
|
||
|
|
workers = build_workers(role)
|
||
|
|
logger.info("worker starting: role=%s (%d arq worker(s))", role, len(workers))
|
||
|
|
|
||
|
|
tasks = [asyncio.create_task(w.async_run()) for w in workers]
|
||
|
|
stop = asyncio.Event()
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
||
|
|
loop.add_signal_handler(sig, stop.set)
|
||
|
|
|
||
|
|
await stop.wait()
|
||
|
|
logger.info("worker stopping: cancelling %d worker(s)", len(tasks))
|
||
|
|
for task in tasks:
|
||
|
|
task.cancel() # arq runs each worker's on_shutdown in its finally
|
||
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
||
|
|
|
||
|
|
from app.db import dispose_engine
|
||
|
|
from app.queue import close_arq_pool
|
||
|
|
from app.redis_client import close_redis
|
||
|
|
|
||
|
|
await dispose_engine()
|
||
|
|
await close_redis()
|
||
|
|
await close_arq_pool()
|
||
|
|
logger.info("worker stopped")
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
asyncio.run(_run())
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|