Livrare LOT 1 - Didi

ai_platform/modules/didi_brain/scheduler/main.py

@@ -0,0 +1,146 @@
+"""Scheduler entry point — orchestrates feeder + auditor + watcher.
+
+Runs as a single container with three independent asyncio tasks. Each task
+is supervisor-style: catches its own exceptions and continues, so a failure
+in one (e.g., RSS feed temporarily down) doesn't kill the others.
+
+Health check: writes ``/tmp/scheduler.healthy`` periodically. Docker
+healthcheck probes the file's mtime to detect a stuck loop.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import signal
+import time
+from pathlib import Path
+
+from brain_api.db import db
+from scheduler.auditor import run_auditor
+from scheduler.brain_client import BrainClient
+from scheduler.config import settings
+from scheduler.feeder import run_feeder
+from scheduler.watcher import run_watcher
+from shared.llm_client import LlmClient
+from shared.logging import setup_logging, get_logger
+
+log = get_logger(__name__)
+
+HEALTH_FILE = Path("/tmp/scheduler.healthy")
+HEALTH_INTERVAL_S = 30.0
+
+
+async def _heartbeat() -> None:
+    """Periodically touch the health file so healthcheck sees activity."""
+    while True:
+        try:
+            HEALTH_FILE.write_text(str(time.time()))
+        except Exception:  # noqa: BLE001
+            pass
+        await asyncio.sleep(HEALTH_INTERVAL_S)
+
+
+async def _wait_brain_ready(brain: BrainClient, *, max_attempts: int = 60) -> None:
+    """Spin until brain-api answers /health, with bounded retries.
+
+    Compose dependencies don't always order brain-api before scheduler at
+    runtime; we'd rather wait than crash on first call.
+    """
+    for attempt in range(max_attempts):
+        if await brain.health():
+            log.info("brain_api_reachable", attempts=attempt + 1)
+            return
+        await asyncio.sleep(2.0)
+    log.warning("brain_api_unreachable_after_retries", attempts=max_attempts)
+
+
+async def _supervised(name: str, coro_factory) -> None:
+    """Wrap a long-running task so a crash inside doesn't unwind the rest.
+
+    The task never returns under normal operation; if it raises, we log and
+    restart after a short backoff.
+    """
+    backoff = 5.0
+    while True:
+        try:
+            await coro_factory()
+            log.warning("supervised_task_returned", name=name)
+        except asyncio.CancelledError:
+            log.info("supervised_task_cancelled", name=name)
+            raise
+        except Exception:  # noqa: BLE001
+            log.exception("supervised_task_crashed", name=name)
+        await asyncio.sleep(backoff)
+        backoff = min(backoff * 2, 300.0)  # cap at 5 min
+
+
+async def main() -> None:
+    setup_logging()
+    log.info(
+        "scheduler_starting",
+        feeder=settings.feeder_enabled,
+        auditor=settings.auditor_enabled,
+        watcher=settings.watcher_enabled,
+        brain_url=settings.brain_api_url,
+    )
+
+    # Connect to PG (auditor needs direct DB access for SELECT/UPDATE).
+    try:
+        await db.connect()
+    except Exception:  # noqa: BLE001
+        log.exception("scheduler_db_connect_failed")
+        # Without DB the auditor can't run, but feeder + watcher only need
+        # HTTP — degrade gracefully rather than exit.
+        pass
+
+    brain = BrainClient()
+    llm = LlmClient()
+
+    await _wait_brain_ready(brain)
+
+    # Graceful shutdown: cancel all tasks on SIGTERM/SIGINT.
+    loop = asyncio.get_running_loop()
+    stop_event = asyncio.Event()
+
+    def _trigger_stop() -> None:
+        stop_event.set()
+
+    for sig in (signal.SIGTERM, signal.SIGINT):
+        try:
+            loop.add_signal_handler(sig, _trigger_stop)
+        except NotImplementedError:
+            # Windows / restricted env — ignore.
+            pass
+
+    tasks = [
+        asyncio.create_task(_heartbeat(), name="heartbeat"),
+        asyncio.create_task(
+            _supervised("feeder", lambda: run_feeder(brain)),
+            name="feeder",
+        ),
+        asyncio.create_task(
+            _supervised("auditor", lambda: run_auditor(brain, llm)),
+            name="auditor",
+        ),
+        asyncio.create_task(
+            _supervised("watcher", lambda: run_watcher(brain, llm)),
+            name="watcher",
+        ),
+    ]
+
+    try:
+        await stop_event.wait()
+    finally:
+        log.info("scheduler_stopping")
+        for t in tasks:
+            t.cancel()
+        # Give tasks a moment to clean up.
+        await asyncio.gather(*tasks, return_exceptions=True)
+        await brain.aclose()
+        await llm.aclose()
+        await db.close()
+        log.info("scheduler_stopped")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())