first
This commit is contained in:
valenti
@@ -0,0 +1,16 @@
|
||||
"""asf-worker — stateless RabbitMQ consumer ("gli addetti alle poste").
|
||||
|
||||
Each subscribed queue maps to a handler. Handlers look up the persistent
|
||||
Mongo ``Job`` record, update its status to running → completed|failed and
|
||||
optionally publish domain events.
|
||||
|
||||
Workers are horizontally scalable. All persistent state lives in MongoDB,
|
||||
never in process memory.
|
||||
"""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
|
||||
from .config import WorkerSettings, get_settings
|
||||
from .worker import run_worker, build_consumer, JobHandlerFn
|
||||
|
||||
__all__ = ["WorkerSettings", "get_settings", "run_worker", "build_consumer", "JobHandlerFn"]
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Allow running with `python -m app.worker`."""
|
||||
from .worker import run_worker
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_worker()
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Settings."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import lru_cache
|
||||
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class WorkerSettings(BaseSettings):
|
||||
model_config = SettingsConfigDict(
|
||||
env_file=".env", case_sensitive=False, extra="ignore"
|
||||
)
|
||||
asf_env: str = "development"
|
||||
asf_log_level: str = "INFO"
|
||||
rabbitmq_uri: str = "amqp://azurian:StrongPassword123!@172.19.40.201:5672/asf"
|
||||
mongo_uri: str = "mongodb://172.19.40.102:27017,172.19.40.103:27017/asf?replicaSet=rs0"
|
||||
mongo_db: str = "asf"
|
||||
ai_gateway_provider: str = "mock"
|
||||
|
||||
# Queue → routing key bindings consumed by this worker.
|
||||
# A single worker can subscribe to a subset; in production you'd run
|
||||
# separate worker deployments per queue pool.
|
||||
queues: list[str] = [
|
||||
"asf.ai.interactive",
|
||||
"asf.ai.background",
|
||||
"asf.indexing",
|
||||
"asf.code_analysis",
|
||||
"asf.testing",
|
||||
"asf.devops",
|
||||
"asf.notifications",
|
||||
]
|
||||
|
||||
# If true, AI jobs invoke the AI Gateway and write results back.
|
||||
enable_ai_handlers: bool = True
|
||||
|
||||
worker_id: str = "asf-worker-1"
|
||||
prefetch_count: int = 4
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_settings() -> WorkerSettings:
|
||||
return WorkerSettings()
|
||||
@@ -0,0 +1,317 @@
|
||||
"""Job handler dispatch + default handlers (AI + logging).
|
||||
|
||||
Handlers share the following signature::
|
||||
|
||||
async fn(
|
||||
envelope: JobRequestEnvelope,
|
||||
*,
|
||||
deps: WorkerDependencies,
|
||||
) -> HandlerResult
|
||||
|
||||
WorkerDependencies wraps Mongo collections, the AI Gateway, and logging.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Awaitable, Callable, Protocol
|
||||
from uuid import UUID
|
||||
|
||||
import structlog
|
||||
|
||||
from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase
|
||||
from asf_events.base import JobRequestEnvelope
|
||||
from asf_domain import JobStatus, JobUpdate, JobUsage
|
||||
|
||||
from app.config import WorkerSettings # type: ignore[attr-defined]
|
||||
|
||||
log = structlog.get_logger(__name__)
|
||||
|
||||
|
||||
JobHandlerFn = Callable[["JobRequestEnvelope", "WorkerDependencies"], Awaitable["HandlerResult"]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class HandlerResult:
|
||||
status: JobStatus
|
||||
result: dict[str, Any] = field(default_factory=dict)
|
||||
error: str | None = None
|
||||
error_detail: dict[str, Any] | None = None
|
||||
usage: JobUsage | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkerDependencies:
|
||||
db: AsyncIOMotorDatabase
|
||||
jobs: AsyncIOMotorCollection
|
||||
ai_executions: AsyncIOMotorCollection
|
||||
artifacts: AsyncIOMotorCollection
|
||||
settings: WorkerSettings
|
||||
ai_gateway: Any | None = None # will be set if enable_ai_handlers=True
|
||||
worker_id: str = "asf-worker"
|
||||
|
||||
|
||||
# ==========================================================================
|
||||
# Default handlers
|
||||
# ==========================================================================
|
||||
|
||||
|
||||
async def _mark_job_done(
|
||||
deps: WorkerDependencies,
|
||||
job_id: UUID,
|
||||
result: HandlerResult,
|
||||
) -> None:
|
||||
"""Persist the final handler outcome back to MongoDB jobs collection.
|
||||
|
||||
This is the ONLY place worker-state leaves process memory and hits the
|
||||
durable source of truth (MongoDB).
|
||||
"""
|
||||
updates: dict[str, Any] = {
|
||||
"status": result.status.value
|
||||
if isinstance(result.status, JobStatus)
|
||||
else str(result.status),
|
||||
"completed_at": __import__("datetime").datetime.now(
|
||||
__import__("datetime").timezone.utc
|
||||
),
|
||||
"updated_at": __import__("datetime").datetime.now(
|
||||
__import__("datetime").timezone.utc
|
||||
),
|
||||
"worker": deps.worker_id,
|
||||
}
|
||||
if result.result:
|
||||
updates["result"] = result.result
|
||||
if result.error:
|
||||
updates["error"] = result.error
|
||||
if result.error_detail:
|
||||
updates["error_detail"] = result.error_detail
|
||||
if result.usage is not None:
|
||||
updates["usage"] = (
|
||||
result.usage.model_dump()
|
||||
if hasattr(result.usage, "model_dump")
|
||||
else result.usage
|
||||
)
|
||||
await deps.jobs.update_one({"_id": job_id}, {"$set": updates})
|
||||
|
||||
|
||||
async def noop_handler(
|
||||
env: JobRequestEnvelope, deps: WorkerDependencies
|
||||
) -> HandlerResult:
|
||||
"""Catch-all handler: acknowledges the job and leaves a note.
|
||||
|
||||
Useful when a queue is not yet staffed with a real handler — ensures we
|
||||
never leave work stuck "running" with no outcome.
|
||||
"""
|
||||
log.warn(
|
||||
"no_handler_registered",
|
||||
job_id=str(env.job_id),
|
||||
queue=env.type,
|
||||
payload_keys=list(env.payload.keys()),
|
||||
)
|
||||
return HandlerResult(
|
||||
status=JobStatus.FAILED,
|
||||
error=f"No worker handler registered for job_type={env.type!r}",
|
||||
error_detail={"hint": "register a handler in app.worker.HANDLERS_BY_TYPE"},
|
||||
)
|
||||
|
||||
|
||||
async def echo_handler(
|
||||
env: JobRequestEnvelope, deps: WorkerDependencies
|
||||
) -> HandlerResult:
|
||||
"""Simple echo / dev handler: stores the payload back as result."""
|
||||
return HandlerResult(
|
||||
status=JobStatus.COMPLETED,
|
||||
result={"echo": env.payload, "handled_by": deps.worker_id},
|
||||
usage=JobUsage(cpu_seconds=0.001),
|
||||
)
|
||||
|
||||
|
||||
async def ai_handler(
|
||||
env: JobRequestEnvelope, deps: WorkerDependencies
|
||||
) -> HandlerResult:
|
||||
"""AI-interactive + AI-background handler.
|
||||
|
||||
Dispatches to the ASF AI Gateway (MockProvider in bootstrap) for:
|
||||
- "ask_asf" kind (chat completion)
|
||||
- "analyze_code"
|
||||
- "generate_tests"
|
||||
- "decompose_requirement"
|
||||
|
||||
In the bootstrap phase we safely fall back to echo if no gateway is
|
||||
attached (e.g. test runs with AI disabled).
|
||||
"""
|
||||
|
||||
kind = env.payload.get("kind", "unknown")
|
||||
log.info(
|
||||
"ai_handler_start",
|
||||
job_id=str(env.job_id),
|
||||
kind=kind,
|
||||
provider=deps.settings.ai_gateway_provider,
|
||||
)
|
||||
|
||||
if deps.ai_gateway is None:
|
||||
return HandlerResult(
|
||||
status=JobStatus.COMPLETED,
|
||||
result={
|
||||
"note": (
|
||||
"ai_gateway not attached to this worker. "
|
||||
"Install services/asf-ai and wire it in worker."
|
||||
),
|
||||
"payload": env.payload,
|
||||
},
|
||||
)
|
||||
|
||||
from asf_ai import ( # type: ignore[import-not-found]
|
||||
AIGateway,
|
||||
ChatRequest,
|
||||
ChatMessage,
|
||||
TraceContext,
|
||||
CodeAnalysisRequest,
|
||||
TestGenerationRequest,
|
||||
RequirementDecomposeRequest,
|
||||
)
|
||||
|
||||
gw: AIGateway = deps.ai_gateway
|
||||
trace = TraceContext(
|
||||
tenant_id=env.tenant_id,
|
||||
project_id=env.project_id,
|
||||
job_id=env.job_id,
|
||||
correlation_id=env.correlation_id,
|
||||
)
|
||||
|
||||
try:
|
||||
if kind == "ask_asf":
|
||||
question = env.payload.get("question", "")
|
||||
req = ChatRequest(
|
||||
trace=trace,
|
||||
messages=[ChatMessage(role="user", content=question)],
|
||||
context_summary=env.payload.get("context_summary"),
|
||||
)
|
||||
resp = await gw.chat(req)
|
||||
return HandlerResult(
|
||||
status=JobStatus.COMPLETED,
|
||||
result={
|
||||
"kind": "ask_asf",
|
||||
"question": question,
|
||||
"answer": resp.reply,
|
||||
"references": resp.references,
|
||||
},
|
||||
usage=JobUsage(
|
||||
input_tokens=resp.input_tokens,
|
||||
output_tokens=resp.output_tokens,
|
||||
estimated_cost=resp.estimated_cost,
|
||||
),
|
||||
)
|
||||
|
||||
if kind == "analyze_code":
|
||||
req = CodeAnalysisRequest(
|
||||
trace=trace,
|
||||
code_snippet=env.payload.get("code_snippet", ""),
|
||||
language=env.payload.get("language"),
|
||||
file_path=env.payload.get("file_path"),
|
||||
)
|
||||
resp = await gw.analyze_code(req)
|
||||
return HandlerResult(
|
||||
status=JobStatus.COMPLETED,
|
||||
result={
|
||||
"kind": "analyze_code",
|
||||
"summary": resp.summary,
|
||||
"findings": [f.model_dump(mode="json") for f in resp.findings],
|
||||
},
|
||||
usage=JobUsage(
|
||||
input_tokens=resp.input_tokens,
|
||||
output_tokens=resp.output_tokens,
|
||||
estimated_cost=resp.estimated_cost,
|
||||
),
|
||||
)
|
||||
|
||||
if kind == "generate_tests":
|
||||
req = TestGenerationRequest(
|
||||
trace=trace,
|
||||
code_snippet=env.payload.get("code_snippet", ""),
|
||||
language=env.payload.get("language", "python"),
|
||||
test_framework=env.payload.get("framework", "pytest"),
|
||||
goal=env.payload.get("goal", "Generate unit tests"),
|
||||
)
|
||||
resp = await gw.generate_tests(req)
|
||||
return HandlerResult(
|
||||
status=JobStatus.COMPLETED,
|
||||
result={
|
||||
"kind": "generate_tests",
|
||||
"test_cases": resp.test_cases,
|
||||
"fixtures_needed": resp.fixtures_needed,
|
||||
"coverage_notes": resp.coverage_notes,
|
||||
},
|
||||
usage=JobUsage(
|
||||
input_tokens=resp.input_tokens,
|
||||
output_tokens=resp.output_tokens,
|
||||
estimated_cost=resp.estimated_cost,
|
||||
),
|
||||
)
|
||||
|
||||
if kind == "decompose_requirement":
|
||||
req = RequirementDecomposeRequest(
|
||||
trace=trace,
|
||||
requirement_text=env.payload.get("requirement_text", ""),
|
||||
context=env.payload.get("context"),
|
||||
max_subtasks=int(env.payload.get("max_subtasks", 20)),
|
||||
)
|
||||
resp = await gw.decompose_requirement(req)
|
||||
return HandlerResult(
|
||||
status=JobStatus.COMPLETED,
|
||||
result={
|
||||
"kind": "decompose_requirement",
|
||||
"subtasks": [s.model_dump(mode="json") for s in resp.subtasks],
|
||||
"architecture_notes": resp.architecture_notes,
|
||||
"open_questions": resp.open_questions,
|
||||
},
|
||||
usage=JobUsage(
|
||||
input_tokens=resp.input_tokens,
|
||||
output_tokens=resp.output_tokens,
|
||||
estimated_cost=resp.estimated_cost,
|
||||
),
|
||||
)
|
||||
|
||||
return await echo_handler(env, deps)
|
||||
|
||||
except Exception as exc: # noqa: BLE001
|
||||
log.exception("ai_handler_failed", exc=str(exc))
|
||||
return HandlerResult(
|
||||
status=JobStatus.FAILED,
|
||||
error=f"AI handler raised: {type(exc).__name__}: {exc}",
|
||||
error_detail={"payload": env.payload},
|
||||
)
|
||||
|
||||
|
||||
async def indexing_handler(
|
||||
env: JobRequestEnvelope, deps: WorkerDependencies
|
||||
) -> HandlerResult:
|
||||
"""Stub handler for repository / doc indexing jobs."""
|
||||
log.info(
|
||||
"indexing_handler_stub",
|
||||
job_id=str(env.job_id),
|
||||
payload=env.payload,
|
||||
)
|
||||
return HandlerResult(
|
||||
status=JobStatus.COMPLETED,
|
||||
result={"stub": True, "indexed_items": 0, "source": env.payload.get("source")},
|
||||
usage=JobUsage(cpu_seconds=0.005),
|
||||
)
|
||||
|
||||
|
||||
async def notification_handler(
|
||||
env: JobRequestEnvelope, deps: WorkerDependencies
|
||||
) -> HandlerResult:
|
||||
"""Stub handler for email / webhook / in-app notifications."""
|
||||
log.info(
|
||||
"notification_stub",
|
||||
job_id=str(env.job_id),
|
||||
payload_keys=list(env.payload.keys()),
|
||||
)
|
||||
return HandlerResult(
|
||||
status=JobStatus.COMPLETED,
|
||||
result={"note": "notifications are a stub in bootstrap"},
|
||||
)
|
||||
|
||||
|
||||
# Final status persistence step, exported for reuse in the main worker loop.
|
||||
mark_job_done = _mark_job_done
|
||||
@@ -0,0 +1,22 @@
|
||||
"""ASF Worker job handlers — pluggable async handlers called from the generic consumer loop.
|
||||
|
||||
New capabilities are registered in ``REGISTRY`` below: mapping of ``job_type`` string
|
||||
to a ``JobHandlerFn``. The worker will route incoming messages by matching
|
||||
``JobRequestEnvelope.job_type`` against keys in ``REGISTRY``.
|
||||
|
||||
Currently stubbed; individual handler modules are added incrementally as features
|
||||
are plugged in (see bootstrap prompt §16).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from asf_events.base import JobRequestEnvelope
|
||||
|
||||
from ..handlers import HandlerResult, WorkerDependencies
|
||||
from ..worker import JobHandlerFn
|
||||
|
||||
REGISTRY: dict[str, JobHandlerFn] = {
|
||||
# "ai.analyze_code": run_analyze_code,
|
||||
# "ai.generate_tests": run_generate_tests,
|
||||
# "ai.decompose_requirements": run_decompose_requirements,
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
"""AI-related job handlers stubbed out for smoke import.
|
||||
|
||||
Actual execution routes through asf-ai.AIGateway once it is wired. The handler
|
||||
below produces a stable HandlerResult.ok(...) carrying a placeholder payload so
|
||||
that consumer-loop smoke tests do not need a live LLM.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from asf_events.base import JobRequestEnvelope
|
||||
|
||||
from ..handlers import HandlerResult
|
||||
from ..worker import WorkerDependencies
|
||||
|
||||
|
||||
async def run_ai_job(req: JobRequestEnvelope, deps: WorkerDependencies) -> HandlerResult:
|
||||
del deps
|
||||
payload = req.payload or {}
|
||||
handler = payload.get("handler") if isinstance(payload, dict) else None
|
||||
return HandlerResult.ok(job_id=req.job_id, output={"handler": handler or "ai.stub", "status": "ok"})
|
||||
@@ -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