first
This commit is contained in:
valenti
@@ -0,0 +1,237 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user