238 lines
7.7 KiB
Python
238 lines
7.7 KiB
Python
"""RabbitMQ consumer loop — main entry point for the asf-worker process.
|
|
|
|
Run::
|
|
python -m app.worker
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import signal
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from typing import Awaitable, Callable
|
|
|
|
import aio_pika
|
|
import aio_pika.abc
|
|
import structlog
|
|
|
|
from motor.motor_asyncio import AsyncIOMotorClient
|
|
|
|
from .config import WorkerSettings, get_settings
|
|
from .handlers import (
|
|
WorkerDependencies,
|
|
noop_handler,
|
|
echo_handler,
|
|
ai_handler,
|
|
indexing_handler,
|
|
notification_handler,
|
|
mark_job_done,
|
|
HandlerResult,
|
|
)
|
|
from asf_events.base import JobRequestEnvelope
|
|
from asf_domain import JobStatus, JobUpdate
|
|
|
|
log = structlog.get_logger(__name__)
|
|
|
|
JobHandlerFn = Callable[[JobRequestEnvelope, WorkerDependencies], Awaitable[HandlerResult]]
|
|
"""Type alias for a queue handler: (envelope, deps) → coroutine returning HandlerResult."""
|
|
|
|
|
|
# Queue name → handler function
|
|
HANDLERS_BY_QUEUE: dict[str, Callable] = {
|
|
"asf.ai.interactive": ai_handler,
|
|
"asf.ai.background": ai_handler,
|
|
"asf.indexing": indexing_handler,
|
|
"asf.code_analysis": ai_handler, # use analyze_code via AI gateway
|
|
"asf.testing": ai_handler, # use generate_tests via AI gateway
|
|
"asf.devops": echo_handler, # stub
|
|
"asf.notifications": notification_handler,
|
|
}
|
|
|
|
|
|
async def _claim_running(deps: WorkerDependencies, envelope: JobRequestEnvelope) -> None:
|
|
"""Atomically mark the job as running + claim worker ownership.
|
|
|
|
Guards against duplicate-delivery races from RabbitMQ (at-least-once).
|
|
If the job is already in a terminal state we bail out early.
|
|
"""
|
|
from datetime import datetime, timezone
|
|
|
|
now = datetime.now(tz=timezone.utc)
|
|
res = await deps.jobs.update_one(
|
|
{
|
|
"_id": envelope.job_id,
|
|
"tenant_id": envelope.tenant_id,
|
|
"status": {"$in": [JobStatus.QUEUED.value]},
|
|
},
|
|
{
|
|
"$set": {
|
|
"status": JobStatus.RUNNING.value,
|
|
"started_at": now,
|
|
"updated_at": now,
|
|
"worker": deps.worker_id,
|
|
},
|
|
"$inc": {"retries": 1},
|
|
},
|
|
)
|
|
if res.matched_count == 0:
|
|
|
|
log.warn(
|
|
"job_already_claimed_or_terminal",
|
|
job_id=str(envelope.job_id),
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class Consumer:
|
|
settings: WorkerSettings
|
|
connection: aio_pika.abc.AbstractConnection | None = None
|
|
channel: aio_pika.abc.AbstractChannel | None = None
|
|
mongo_client: AsyncIOMotorClient | None = None
|
|
deps: WorkerDependencies | None = None
|
|
stop_event: asyncio.Event | None = None
|
|
|
|
async def setup(self) -> None:
|
|
self.stop_event = asyncio.Event()
|
|
|
|
# ---- Mongo ----
|
|
self.mongo_client = AsyncIOMotorClient(
|
|
self.settings.mongo_uri,
|
|
uuidRepresentation="standard",
|
|
tz_aware=True,
|
|
)
|
|
db = self.mongo_client[self.settings.mongo_db]
|
|
self.deps = WorkerDependencies(
|
|
db=db,
|
|
jobs=db["jobs"],
|
|
ai_executions=db["ai_executions"],
|
|
artifacts=db["artifacts"],
|
|
settings=self.settings,
|
|
worker_id=self.settings.worker_id,
|
|
)
|
|
|
|
# ---- Optional AI gateway wiring ----
|
|
if self.settings.enable_ai_handlers:
|
|
try:
|
|
from asf_ai.app.gateway import AIGateway, AIGatewayConfig # type: ignore[import-not-found]
|
|
self.deps.ai_gateway = AIGateway(config=AIGatewayConfig(
|
|
default_provider=self.settings.ai_gateway_provider,
|
|
))
|
|
log.info(
|
|
"ai_gateway_attached", provider=self.settings.ai_gateway_provider
|
|
)
|
|
except Exception as exc: # noqa: BLE001
|
|
log.warn("ai_gateway_unavailable", exc=str(exc))
|
|
|
|
# ---- RabbitMQ ----
|
|
self.connection = await aio_pika.connect_robust(self.settings.rabbitmq_uri)
|
|
self.channel = await self.connection.channel()
|
|
await self.channel.set_qos(prefetch_count=self.settings.prefetch_count)
|
|
|
|
jobs_exchange = await self.channel.declare_exchange(
|
|
"asf.jobs", type=aio_pika.ExchangeType.DIRECT, durable=True
|
|
)
|
|
await self.channel.declare_exchange(
|
|
"asf.events", type=aio_pika.ExchangeType.TOPIC, durable=True
|
|
)
|
|
|
|
for qname in self.settings.queues:
|
|
q = await self.channel.declare_queue(qname, durable=True)
|
|
await q.bind(jobs_exchange, routing_key=qname)
|
|
await q.consume(self._make_consumer_callback(qname))
|
|
log.info("subscribed", queue=qname)
|
|
|
|
def _make_consumer_callback(self, qname: str):
|
|
handler = HANDLERS_BY_QUEUE.get(qname, noop_handler)
|
|
|
|
async def _cb(message: aio_pika.abc.AbstractIncomingMessage) -> None:
|
|
async with message.process(requeue=True, reject_on_redelivered=False):
|
|
assert self.deps is not None
|
|
try:
|
|
envelope = JobRequestEnvelope.model_validate_json(message.body)
|
|
except Exception: # noqa: BLE001
|
|
log.exception("bad_message_body", queue=qname)
|
|
return # ack & drop poison
|
|
|
|
# 1. claim ownership in Mongo
|
|
await _claim_running(self.deps, envelope)
|
|
# 2. run the handler
|
|
try:
|
|
result: HandlerResult = await handler(envelope, self.deps)
|
|
except Exception as exc: # noqa: BLE001
|
|
log.exception("handler_unhandled_exception", queue=qname)
|
|
result = HandlerResult(
|
|
status=JobStatus.FAILED,
|
|
error=f"{type(exc).__name__}: {exc}",
|
|
)
|
|
# 3. persist final state
|
|
try:
|
|
await mark_job_done(self.deps, envelope.job_id, result)
|
|
except Exception: # noqa: BLE001
|
|
log.exception("job_persistence_failed", job_id=str(envelope.job_id))
|
|
|
|
return _cb
|
|
|
|
async def run_forever(self) -> None:
|
|
assert self.stop_event is not None
|
|
log.info(
|
|
"asf_worker_started",
|
|
queues=self.settings.queues,
|
|
worker_id=self.settings.worker_id,
|
|
)
|
|
await self.stop_event.wait()
|
|
|
|
async def close(self) -> None:
|
|
log.info("asf_worker_shutting_down")
|
|
if self.channel is not None:
|
|
try:
|
|
await self.channel.close()
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
if self.connection is not None:
|
|
try:
|
|
await self.connection.close()
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
if self.mongo_client is not None:
|
|
self.mongo_client.close()
|
|
|
|
|
|
def build_consumer(settings: WorkerSettings | None = None) -> Consumer:
|
|
return Consumer(settings=settings or get_settings())
|
|
|
|
|
|
async def run_worker_async() -> None:
|
|
settings = get_settings()
|
|
consumer = build_consumer(settings)
|
|
await consumer.setup()
|
|
loop = asyncio.get_running_loop()
|
|
|
|
def _graceful() -> None:
|
|
if consumer.stop_event is not None:
|
|
consumer.stop_event.set()
|
|
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
try:
|
|
loop.add_signal_handler(sig, _graceful)
|
|
except NotImplementedError:
|
|
# Windows
|
|
pass
|
|
|
|
try:
|
|
await consumer.run_forever()
|
|
finally:
|
|
await consumer.close()
|
|
|
|
|
|
def run_worker() -> None:
|
|
"""Synchronous entry point (e.g. python -m app.worker)."""
|
|
try:
|
|
asyncio.run(run_worker_async())
|
|
except KeyboardInterrupt:
|
|
print("Interrupted. Goodbye.", file=sys.stderr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_worker()
|