HTML 77.2%
TypeScript 10.5%
Python 9.6%
JavaScript 2.5%
1"""Scheduler process (`aia schedule`): runs due connectors (adaptive intervals, Redis lock per connector), the job worker,2hourly stats/quality, embedding backlog, nightly backup. One process per node is enough; several nodes cooperate via locks."""3from __future__ import annotations45import asyncio6import logging7import signal8from datetime import UTC, datetime910from apscheduler.schedulers.asyncio import AsyncIOScheduler11from apscheduler.triggers.cron import CronTrigger1213from aiatlas.config import settings14from aiatlas.connectors import get, registry15from aiatlas.db import execute, fetch_all, fetch_one, transaction16from aiatlas.services import cache17from aiatlas.services.jobs import enqueue, requeue_stale, run_worker1819log = logging.getLogger(__name__)2021_running: set[str] = set()222324async def run_connector(name: str, *, force: bool = False) -> None:25 if name in _running:26 return27 _running.add(name)28 try:29 async with cache.lock(f"connector:{name}", ttl_s=4 * 3600) as ok:30 if not ok:31 log.info("connector locked elsewhere", extra={"connector": name})32 return33 try:34 await get(name).run(force=force)35 except Exception as exc: # noqa: BLE00136 log.warning("connector run failed", extra={"connector": name, "error": str(exc)})37 # a connector run can touch every read model (stats, changes, benchmarks, prices, models, frontier, pulse, families…):38 # flush the whole `aia:api:` namespace so new routes are covered without maintaining a prefix list39 await cache.cache_invalidate()40 finally:41 _running.discard(name)424344LIGHT_CANONICALIZE_STEPS = ["variants", "artifacts", "families", "events", "anomalies"]454647async def canonicalize_job(*, light: bool = False) -> None:48 """Nightly full `canonicalize --apply` (03:30 local) and a light pass every 6 h; one line of summary in the log."""49 from aiatlas.services.canonical import canonicalize5051 async with cache.lock("canonicalize", ttl_s=2 * 3600) as ok:52 if not ok:53 log.info("canonicalize locked elsewhere, skipping", extra={"light": light})54 return55 try:56 report = await canonicalize(apply=True, steps=LIGHT_CANONICALIZE_STEPS if light else None)57 summary = ", ".join(f"{s.name}={s.changes}" for s in report.steps if s.changes) or "no changes"58 log.info("canonicalize %s pass: %d change(s) — %s", "light" if light else "full", report.changes, summary,59 extra={"light": light, "changes": report.changes, "steps": {s.name: s.changes for s in report.steps}})60 if report.changes:61 await cache.cache_invalidate()62 except Exception as exc: # noqa: BLE00163 log.warning("canonicalize pass failed", extra={"light": light, "error": str(exc)})646566async def _pop_run_now() -> list[str]:67 """Admin `POST /admin/connectors/{name}/run` sets `aia:run-now:<name>`; consume (delete) every such key."""68 names: list[str] = []69 try:70 r = cache.redis()71 async for key in r.scan_iter(match="aia:run-now:*", count=100):72 if await r.delete(key):73 names.append(key.decode().rsplit(":", 1)[-1])74 except Exception as exc: # noqa: BLE00175 log.debug("run-now scan skipped", extra={"error": str(exc)})76 return names777879async def tick() -> None:80 known = registry()81 forced = [n for n in await _pop_run_now() if n in known and n not in _running]82 if forced:83 log.info("run-now connectors", extra={"connectors": forced})84 for n in forced[:3]: # one at a time: concurrent runs touch the same organisation rows and can deadlock Postgres85 await run_connector(n, force=True)86 async with transaction() as conn:87 due = await fetch_all(conn, """select name from connectors where enabled and (next_run_at is null or next_run_at <= now())88 and (circuit_open_until is null or circuit_open_until <= now()) order by priority, coalesce(next_run_at, 'epoch') limit 6""")89 names = [r["name"] for r in due if r["name"] in known and r["name"] not in _running and r["name"] not in forced]90 if names:91 log.info("due connectors", extra={"connectors": names})92 for n in names[:2]:93 await run_connector(n)94 await cache.heartbeat("scheduler", {"at": datetime.now(UTC).isoformat(), "running": sorted(_running), "due": names, "forced": forced})959697async def hourly() -> None:98 from aiatlas.services.quality import recompute99 from aiatlas.services.stats import compute_stats100101 async with cache.lock("hourly", ttl_s=3000) as ok:102 if not ok:103 return104 try:105 await compute_stats()106 await recompute(limit=5000)107 async with transaction() as conn:108 n = await requeue_stale(conn)109 if n:110 log.info("requeued stale jobs", extra={"n": n})111 async with transaction() as conn:112 extraction_pending = await fetch_one(conn, "select 1 from jobs where kind = 'llm_extract' and status in ('queued','running') limit 1")113 if settings.llm_available and not extraction_pending: # never interleave embeddings with extraction: the server would swap models114 from aiatlas.services.embeddings import pending_entity_ids115116 ids = await pending_entity_ids(limit=400)117 if ids:118 async with transaction() as conn:119 for i in range(0, len(ids), 50):120 await enqueue(conn, "embed_entity", {"entity_ids": ids[i:i + 50]}, priority=8, dedupe_key=f"embed:{ids[i]}")121 await cache.cache_invalidate()122 except Exception as exc: # noqa: BLE001123 log.warning("hourly maintenance failed", extra={"error": str(exc)})124125126async def nightly_backup() -> None:127 from aiatlas.services.backup import backup_database128129 async with cache.lock("backup", ttl_s=3600) as ok:130 if ok:131 try:132 path = await asyncio.to_thread(backup_database)133 log.info("backup done", extra={"path": str(path)})134 except Exception as exc: # noqa: BLE001135 log.error("backup failed", extra={"error": str(exc)})136137138async def main(*, with_worker: bool = True) -> None:139 settings.ensure_dirs()140 stop = asyncio.Event()141 loop = asyncio.get_running_loop()142 for sig in (signal.SIGINT, signal.SIGTERM):143 try:144 loop.add_signal_handler(sig, stop.set)145 except NotImplementedError:146 pass147 scheduler = AsyncIOScheduler(timezone="UTC")148 scheduler.add_job(tick, "interval", seconds=settings.scheduler_tick_s, max_instances=1, coalesce=True, id="tick")149 scheduler.add_job(hourly, "interval", minutes=60, max_instances=1, coalesce=True, id="hourly", next_run_time=datetime.now(UTC))150 minute, hour, *_ = settings.backup_cron.split()151 scheduler.add_job(nightly_backup, CronTrigger(minute=minute, hour=hour, timezone=settings.tz), id="backup")152 scheduler.add_job(canonicalize_job, CronTrigger(minute=30, hour=3, timezone=settings.tz), id="canonicalize-full", max_instances=1, coalesce=True)153 scheduler.add_job(canonicalize_job, "interval", hours=6, kwargs={"light": True}, id="canonicalize-light", max_instances=1, coalesce=True)154 scheduler.start()155 log.info("scheduler started", extra={"tick_s": settings.scheduler_tick_s, "connectors": len(registry()), "worker": with_worker})156 async with transaction() as conn:157 await execute(conn, "update connectors set health = 'disabled' where not enabled")158 worker_task = asyncio.create_task(run_worker(stop=stop)) if with_worker else None159 await stop.wait()160 scheduler.shutdown(wait=False)161 if worker_task:162 await worker_task163 await cache.close()164165166__all__ = ["LIGHT_CANONICALIZE_STEPS", "canonicalize_job", "hourly", "main", "run_connector", "tick"]167