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 tick() -> None:43 async with transaction() as conn:44 due = await fetch_all(conn, """select name from connectors where enabled and (next_run_at is null or next_run_at <= now())45 and (circuit_open_until is null or circuit_open_until <= now()) order by priority, coalesce(next_run_at, 'epoch') limit 6""")46 known = registry()47 names = [r["name"] for r in due if r["name"] in known and r["name"] not in _running]48 if names:49 log.info("due connectors", extra={"connectors": names})50 await asyncio.gather(*(run_connector(n) for n in names[:3]))51 await cache.heartbeat("scheduler", {"at": datetime.now(UTC).isoformat(), "running": sorted(_running), "due": names})525354async def hourly() -> None:55 from aiatlas.services.quality import recompute56 from aiatlas.services.stats import compute_stats5758 async with cache.lock("hourly", ttl_s=3000) as ok:59 if not ok:60 return61 try:62 await compute_stats()63 await recompute(limit=5000)64 async with transaction() as conn:65 n = await requeue_stale(conn)66 if n:67 log.info("requeued stale jobs", extra={"n": n})68 if settings.llm_available:69 from aiatlas.services.embeddings import pending_entity_ids7071 ids = await pending_entity_ids(limit=400)72 if ids:73 async with transaction() as conn:74 for i in range(0, len(ids), 50):75 await enqueue(conn, "embed_entity", {"entity_ids": ids[i:i + 50]}, priority=8, dedupe_key=f"embed:{ids[i]}")76 await cache.cache_invalidate()77 except Exception as exc: # noqa: BLE00178 log.warning("hourly maintenance failed", extra={"error": str(exc)})798081async def nightly_backup() -> None:82 from aiatlas.services.backup import backup_database8384 async with cache.lock("backup", ttl_s=3600) as ok:85 if ok:86 try:87 path = await asyncio.to_thread(backup_database)88 log.info("backup done", extra={"path": str(path)})89 except Exception as exc: # noqa: BLE00190 log.error("backup failed", extra={"error": str(exc)})919293async def main(*, with_worker: bool = True) -> None:94 settings.ensure_dirs()95 stop = asyncio.Event()96 loop = asyncio.get_running_loop()97 for sig in (signal.SIGINT, signal.SIGTERM):98 try:99 loop.add_signal_handler(sig, stop.set)100 except NotImplementedError:101 pass102 scheduler = AsyncIOScheduler(timezone="UTC")103 scheduler.add_job(tick, "interval", seconds=settings.scheduler_tick_s, max_instances=1, coalesce=True, id="tick")104 scheduler.add_job(hourly, "interval", minutes=60, max_instances=1, coalesce=True, id="hourly", next_run_time=datetime.now(UTC))105 minute, hour, *_ = settings.backup_cron.split()106 scheduler.add_job(nightly_backup, CronTrigger(minute=minute, hour=hour, timezone=settings.tz), id="backup")107 scheduler.start()108 log.info("scheduler started", extra={"tick_s": settings.scheduler_tick_s, "connectors": len(registry()), "worker": with_worker})109 async with transaction() as conn:110 await execute(conn, "update connectors set health = 'disabled' where not enabled")111 worker_task = asyncio.create_task(run_worker(stop=stop)) if with_worker else None112 await stop.wait()113 scheduler.shutdown(wait=False)114 if worker_task:115 await worker_task116 await cache.close()117118119__all__ = ["main", "tick", "run_connector", "hourly"]120