heartbeat.py

# /// script
# requires-python = ">=3.10"
# dependencies = ["huggingface_hub", "requests"]
# ///
"""Heartbeat watchdog for the self-refreshing dataset-lineage pipeline.

Runs on its OWN schedule, independent of refresh_job.py, and fails LOUD — opens an
HF Discussion on the Space repo (the Hub emails watchers) and exits non-zero — the
moment any liveness signal goes stale. That turns the 2026-07 silent-rot failure
class (a dump 3 months stale, a Space serving a shadow db, a gpu phase that died
without finalizing) into an active, dashboard-red + emailed alert.

Checks (each independent; all are evaluated, then a single Discussion lists them):
  (a) scratch refresh/state.json cutoff older than MAX_STATE_AGE_DAYS
  (b) Space /stats reachable AND its built_at older than MAX_BUILT_AGE_DAYS
      (skipped when the live db predates the built_at stamp)
  (c) either cards dump (dataset + model) lastModified older than MAX_DUMP_AGE_DAYS
  (d) an inflight gpu marker older than MAX_INFLIGHT_AGE_HOURS

Schedule (self-updates from the Space repo, same pattern as refresh_job.py):
  hf jobs scheduled uv run "0 7 * * *" --flavor cpu-basic --timeout 15m \
      --secrets HF_TOKEN \
      https://huggingface.co/spaces/davanstrien/dataset-lineage-explorer/raw/main/heartbeat.py
"""

import datetime as dt
import json
import os

import requests
from huggingface_hub import HfApi, hf_hub_download

SCRATCH = "davanstrien/dataset-lineage-scratch"
SPACE = "davanstrien/dataset-lineage-explorer"
DATASET_DUMP = "librarian-bots/dataset_cards_with_metadata"
MODEL_DUMP = "librarian-bots/model_cards_with_metadata"
STATE_PATH = "refresh/state.json"
STATS_URL = f"https://{SPACE.replace('/', '-')}.hf.space/stats"

MAX_STATE_AGE_DAYS = 40      # a monthly refresh that stopped advancing the cutoff
MAX_BUILT_AGE_DAYS = 40      # a deployed db that stopped being rebuilt
MAX_DUMP_AGE_DAYS = 3        # upstream card-pipeline compile job stuck
MAX_INFLIGHT_AGE_HOURS = 36  # a gpu phase that died without finalizing

TOKEN = os.environ.get("HF_TOKEN")
api = HfApi(token=TOKEN)


def _parse(when):
    """Best-effort -> aware UTC datetime, or None."""
    if when is None:
        return None
    if isinstance(when, dt.datetime):
        return when if when.tzinfo else when.replace(tzinfo=dt.timezone.utc)
    try:
        d = dt.datetime.fromisoformat(str(when)[:19].replace(" ", "T"))
        return d.replace(tzinfo=dt.timezone.utc)
    except Exception:  # noqa: BLE001
        return None


def _age_days(when):
    d = _parse(when)
    if d is None:
        return None
    return (dt.datetime.now(dt.timezone.utc) - d).total_seconds() / 86400


def check_state(problems):
    try:
        p = hf_hub_download(SCRATCH, STATE_PATH, repo_type="dataset", token=TOKEN)
        state = json.load(open(p))
    except Exception as e:  # noqa: BLE001
        problems.append(f"(a) state.json unreadable in scratch ({e})")
        return
    age = _age_days(state.get("cutoff"))
    if age is None:
        problems.append(f"(a) state.json cutoff unparseable: {state.get('cutoff')!r}")
    elif age > MAX_STATE_AGE_DAYS:
        problems.append(f"(a) refresh cutoff is {age:.0f}d old (> {MAX_STATE_AGE_DAYS}) — "
                        "the monthly refresh has stopped advancing.")


def check_built(problems):
    try:
        stats = requests.get(STATS_URL, timeout=30).json()
    except (requests.RequestException, ValueError) as e:
        problems.append(f"(b) Space /stats unreachable/invalid ({e}) — the explorer may be down.")
        return
    built = stats.get("built_at")
    if not built:
        return  # older db without a built_at stamp — nothing to check
    age = _age_days(built)
    if age is None:
        problems.append(f"(b) /stats built_at unparseable: {built!r}")
    elif age > MAX_BUILT_AGE_DAYS:
        problems.append(f"(b) deployed db was built {age:.0f}d ago (> {MAX_BUILT_AGE_DAYS}) — "
                        "the Space is serving a stale graph.")


def check_dumps(problems):
    for repo in (DATASET_DUMP, MODEL_DUMP):
        try:
            age = _age_days(api.dataset_info(repo).last_modified)
        except Exception as e:  # noqa: BLE001
            problems.append(f"(c) dump {repo} info fetch failed ({e})")
            continue
        if age is None:
            problems.append(f"(c) dump {repo} lastModified unparseable")
        elif age > MAX_DUMP_AGE_DAYS:
            problems.append(f"(c) dump {repo} lastModified is {age:.1f}d old "
                            f"(> {MAX_DUMP_AGE_DAYS}) — upstream compile job likely stuck.")


def check_inflight(problems):
    try:
        files = [f for f in api.list_repo_files(SCRATCH, repo_type="dataset")
                 if f.startswith("refresh/inflight_") and f.endswith(".json")]
    except Exception as e:  # noqa: BLE001
        problems.append(f"(d) could not list scratch to check inflight markers ({e})")
        return
    for f in files:
        try:
            payload = json.load(open(hf_hub_download(SCRATCH, f, repo_type="dataset", token=TOKEN)))
        except Exception:  # noqa: BLE001
            continue
        age_h = None
        d = _parse(payload.get("submitted_at"))
        if d is not None:
            age_h = (dt.datetime.now(dt.timezone.utc) - d).total_seconds() / 3600
        if age_h is not None and age_h > MAX_INFLIGHT_AGE_HOURS:
            problems.append(f"(d) inflight marker {f} is {age_h:.0f}h old "
                            f"(> {MAX_INFLIGHT_AGE_HOURS}) — a gpu phase died without finalizing.")


def main():
    problems = []
    for check in (check_state, check_built, check_dumps, check_inflight):
        try:
            check(problems)
        except Exception as e:  # noqa: BLE001 — a check crash is itself a problem
            problems.append(f"heartbeat check {check.__name__} crashed: {e}")

    if not problems:
        print("heartbeat OK: all liveness signals fresh", flush=True)
        return

    body = ("The dataset-lineage refresh heartbeat found stale/failed signals:\n\n"
            + "\n".join(f"- {p}" for p in problems)
            + f"\n\nChecked at {dt.datetime.now(dt.timezone.utc):%Y-%m-%dT%H:%M:%SZ}. "
              "See jobs/refresh_job.py and the scratch repo refresh/ state.")
    print("HEARTBEAT PROBLEMS:\n" + body, flush=True)
    try:
        api.create_discussion(repo_id=SCRATCH, repo_type="dataset",
                              title="[refresh-heartbeat] stale/failed liveness signal",
                              description=body)
        print("posted heartbeat Discussion", flush=True)
    except Exception as e:  # noqa: BLE001
        print(f"could not post heartbeat Discussion ({e})", flush=True)
    raise SystemExit(1)


if __name__ == "__main__":
    main()