src/eval_scheduler.py

"""Background scheduler: periodically runs live TS-Bench + TSFM.ai evaluation."""

from __future__ import annotations

import json
import logging
import os
import subprocess
import sys
import threading
import time
from datetime import datetime, timezone
from pathlib import Path

from src.benchmark_config import BENCHMARK_INTERVAL_SECONDS, LIVE_EVAL_ENABLED
from src.eval_schedule import compute_next_eval_at_unix, with_next_eval_fields

logger = logging.getLogger(__name__)

SPACE_ROOT = Path(__file__).resolve().parents[1]
EVAL_SCRIPT = SPACE_ROOT / "scripts" / "run_space_eval.py"
DEFAULT_EVAL_INTERVAL = BENCHMARK_INTERVAL_SECONDS


class LiveEvalScheduler:
    def __init__(
        self,
        interval_seconds: int = DEFAULT_EVAL_INTERVAL,
        run_on_startup: bool = True,
    ):
        self.interval_seconds = max(300, interval_seconds)
        self.run_on_startup = run_on_startup
        self._lock = threading.Lock()
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._next_run_at: float | None = None

    @property
    def enabled(self) -> bool:
        return LIVE_EVAL_ENABLED and bool(os.getenv("TSFM_API_KEY"))

    def start(self) -> None:
        if self._thread is not None:
            return
        if not LIVE_EVAL_ENABLED:
            logger.info(
                "UI-only mode (TSFM_LIVE_EVAL=0): scheduler disabled; "
                "leaderboard reads pushed results/."
            )
            return
        if not os.getenv("TSFM_API_KEY"):
            logger.warning(
                "Live eval disabled: TSFM_API_KEY not set. "
                "Leaderboard will only show static CSV results."
            )
            self._write_idle_state(
                "TSFM_API_KEY not set — configure it in HF Space Settings → Secrets."
            )
            return

        self._thread = threading.Thread(target=self._loop, name="live-eval-scheduler", daemon=True)
        self._thread.start()
        logger.info(
            "Live eval scheduler started (interval=%ss, script=%s)",
            self.interval_seconds,
            EVAL_SCRIPT,
        )

    def snapshot(self) -> dict:
        state_path = SPACE_ROOT / "results" / "eval_state.json"
        if state_path.exists():
            try:
                return json.loads(state_path.read_text())
            except json.JSONDecodeError:
                pass
        return {
            "status": "idle" if self.enabled else "disabled",
            "message": "Waiting for first evaluation run…" if self.enabled else "No API key",
        }

    def seconds_until_next_run(self) -> int | None:
        if self._next_run_at is None:
            return None
        return max(0, int(self._next_run_at - time.monotonic()))

    def next_eval_at_unix(self) -> int:
        state = self.snapshot()
        persisted = compute_next_eval_at_unix(state, self.interval_seconds)

        remaining = self.seconds_until_next_run()
        if remaining is None:
            return persisted

        in_memory = int(time.time()) + remaining
        return min(persisted, in_memory)

    def is_running(self) -> bool:
        return self.snapshot().get("status") == "running"

    def _loop(self) -> None:
        if self.run_on_startup:
            self._run_once()
        while not self._stop.wait(self.interval_seconds):
            self._run_once()

    def _run_once(self) -> None:
        if not self._lock.acquire(blocking=False):
            logger.info("Evaluation already in progress, skipping tick")
            return
        try:
            logger.info("Starting scheduled live evaluation")
            cmd = [sys.executable, str(EVAL_SCRIPT)]
            completed = subprocess.run(cmd, cwd=str(SPACE_ROOT), check=False)
            if completed.returncode != 0:
                logger.error("Evaluation exited with code %s", completed.returncode)
        finally:
            self._lock.release()
            self._next_run_at = time.monotonic() + self.interval_seconds
            self._persist_next_eval_after_run()

    def _persist_next_eval_after_run(self) -> None:
        state_path = SPACE_ROOT / "results" / "eval_state.json"
        if not state_path.exists():
            return
        try:
            state = json.loads(state_path.read_text())
        except json.JSONDecodeError:
            return
        if state.get("status") == "running":
            return
        anchor = state.get("finished_at") or state.get("updated_at")
        if not anchor:
            return
        enriched = with_next_eval_fields(state, anchor, self.interval_seconds)
        state_path.write_text(json.dumps(enriched, indent=4) + "\n")

    def _write_idle_state(self, message: str) -> None:
        results = SPACE_ROOT / "results"
        results.mkdir(parents=True, exist_ok=True)
        payload = {
            "status": "disabled",
            "message": message,
            "updated_at": datetime.now(timezone.utc).isoformat(),
        }
        (results / "eval_state.json").write_text(json.dumps(payload, indent=4) + "\n")


_scheduler: LiveEvalScheduler | None = None


def get_scheduler() -> LiveEvalScheduler:
    global _scheduler
    if _scheduler is None:
        _scheduler = LiveEvalScheduler()
    return _scheduler


def start_live_eval_scheduler() -> LiveEvalScheduler:
    scheduler = get_scheduler()
    scheduler.start()
    return scheduler