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, 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 await cache.cache_invalidate()38 finally:39 _running.discard(name)404142async def _pop_run_now() -> list[str]:43 """Admin `POST /admin/connectors/{name}/run` sets `aia:run-now:<name>`; consume (delete) every such key."""44 names: list[str] = []45 try:46 r = cache.redis()47 async for key in r.scan_iter(match="aia:run-now:*", count=100):48 if await r.delete(key):49 names.append(key.decode().rsplit(":", 1)[-1])50 except Exception as exc: # noqa: BLE00151 log.debug("run-now scan skipped", extra={"error": str(exc)})52 return names535455async def tick() -> None:56 known = registry()57 forced = [n for n in await _pop_run_now() if n in known and n not in _running]58 if forced:59 log.info("run-now connectors", extra={"connectors": forced})60 await asyncio.gather(*(run_connector(n, force=True) for n in forced[:3]))61 async with transaction() as conn:62 due = await fetch_all(conn, """select name from connectors where enabled and (next_run_at is null or next_run_at <= now())63 and (circuit_open_until is null or circuit_open_until <= now()) order by priority, coalesce(next_run_at, 'epoch') limit 6""")64 names = [r["name"] for r in due if r["name"] in known and r["name"] not in _running and r["name"] not in forced]65 if names:66 log.info("due connectors", extra={"connectors": names})67 await asyncio.gather(*(run_connector(n) for n in names[:3]))68 await cache.heartbeat("scheduler", {"at": datetime.now(UTC).isoformat(), "running": sorted(_running), "due": names, "forced": forced})697071async def hourly() -> None:72 from aiatlas.services.quality import recompute73 from aiatlas.services.stats import compute_stats7475 async with cache.lock("hourly", ttl_s=3000) as ok:76 if not ok:77 return78 try:79 await compute_stats()80 await recompute(limit=5000)81 async with transaction() as conn:82 n = await requeue_stale(conn)83 if n:84 log.info("requeued stale jobs", extra={"n": n})85 if settings.llm_available:86 from aiatlas.services.embeddings import pending_entity_ids8788 ids = await pending_entity_ids(limit=400)89 if ids:90 async with transaction() as conn:91 for i in range(0, len(ids), 50):92 await enqueue(conn, "embed_entity", {"entity_ids": ids[i:i + 50]}, priority=8, dedupe_key=f"embed:{ids[i]}")93 await cache.cache_invalidate()94 except Exception as exc: # noqa: BLE00195 log.warning("hourly maintenance failed", extra={"error": str(exc)})969798async def nightly_backup() -> None:99 from aiatlas.services.backup import backup_database100101 async with cache.lock("backup", ttl_s=3600) as ok:102 if ok:103 try:104 path = await asyncio.to_thread(backup_database)105 log.info("backup done", extra={"path": str(path)})106 except Exception as exc: # noqa: BLE001107 log.error("backup failed", extra={"error": str(exc)})108109110async def main(*, with_worker: bool = True) -> None:111 settings.ensure_dirs()112 stop = asyncio.Event()113 loop = asyncio.get_running_loop()114 for sig in (signal.SIGINT, signal.SIGTERM):115 try:116 loop.add_signal_handler(sig, stop.set)117 except NotImplementedError:118 pass119 scheduler = AsyncIOScheduler(timezone="UTC")120 scheduler.add_job(tick, "interval", seconds=settings.scheduler_tick_s, max_instances=1, coalesce=True, id="tick")121 scheduler.add_job(hourly, "interval", minutes=60, max_instances=1, coalesce=True, id="hourly", next_run_time=datetime.now(UTC))122 minute, hour, *_ = settings.backup_cron.split()123 scheduler.add_job(nightly_backup, CronTrigger(minute=minute, hour=hour, timezone=settings.tz), id="backup")124 scheduler.start()125 log.info("scheduler started", extra={"tick_s": settings.scheduler_tick_s, "connectors": len(registry()), "worker": with_worker})126 async with transaction() as conn:127 await execute(conn, "update connectors set health = 'disabled' where not enabled")128 worker_task = asyncio.create_task(run_worker(stop=stop)) if with_worker else None129 await stop.wait()130 scheduler.shutdown(wait=False)131 if worker_task:132 await worker_task133 await cache.close()134135136__all__ = ["hourly", "main", "run_connector", "tick"]137