GarminHealthLab/backend/services/jobs.py

"""
The coach's producer/consumer queue.

Producing one insight costs 40s to several minutes against the gateway, so it
can never happen inside a request. Screens *enqueue*; a worker thread consumes.

Priority is the whole point of the queue rather than a plain background thread:
after a sync the backfill enqueues every scope at low priority, and those jobs
may take half an hour to work through — but the moment the user opens a screen,
that screen's job is promoted to the front and runs next. What they are looking
at is always what the queue is working on.

The queue lives in the database, not in memory, because Gunicorn runs several
workers: a job is claimed with a holder id and re-read to confirm, the same way
`scheduler.py` claims its tick, so exactly one worker runs a given job.
"""
import datetime
import hashlib
import os
import threading
import time

from config import DB_TYPE
from db import execute, query_one, query_all

# Lower runs first.
PRIORITY_INTERACTIVE = 0   # a screen the user has open right now
PRIORITY_PREFETCH = 10     # backfill after a sync

# A claim older than this is treated as abandoned: the worker holding it died
# mid-generation, and without expiry that job would never run again.
CLAIM_TIMEOUT_SECONDS = int(os.environ.get("AI_JOB_CLAIM_TIMEOUT") or 1800)

# Generations are slow, not frequent; polling this often costs nothing and
# keeps an interactive job's wait to a couple of seconds.
POLL_SECONDS = float(os.environ.get("AI_JOB_POLL_SECONDS") or 2)

MAX_ATTEMPTS = int(os.environ.get("AI_JOB_MAX_ATTEMPTS") or 3)

# The current gateway at ai.zichuan.xyz supports multi-concurrent requests,
# but three projects share it, so going too high causes 502s.
# MAX_CONCURRENT is read from the environment and defaults to 2; raise it
# if the gateway shows headroom, lower it if this or other projects get 502s.
MAX_CONCURRENT = int(os.environ.get("AI_JOB_CONCURRENCY") or 2)
GAP_SECONDS = float(os.environ.get("AI_JOB_GAP_SECONDS") or 5)

# How long a job that exhausted its attempts stays given up on before it is
# tried again. Without this a transient upstream outage is permanent: three
# quick failures while the gateway is unreachable would retire that screen's
# insight until its underlying data happened to change, which for a screen the
# user is not syncing could be days.
FAILED_RETRY_SECONDS = int(os.environ.get("AI_JOB_RETRY_AFTER") or 1800)

ENABLED = (os.environ.get("AI_JOBS") or "true").lower() not in ("0", "false", "no")

_started = False
_start_lock = threading.Lock()

# Set by analysis.py at import time. Injected rather than imported so this
# module stays free of the feature logic it schedules — and so the two do not
# import each other in a cycle.
_runner = None

# Optional refiller, called when the queue is idle to keep the backlog
# populated with daily/activity items that have no insight yet.
_refiller = None


def set_runner(fn):
    """Register `fn(user_id, kind, subject) -> None`, called for each job."""
    global _runner
    _runner = fn


def set_refiller(fn):
    """Register `fn() -> int`, called when the queue is idle to refill backlog."""
    global _refiller
    _refiller = fn


def _now():
    return datetime.datetime.utcnow()


def _iso(dt):
    return dt.isoformat(timespec="seconds")


def _parse(value):
    if not value:
        return None
    try:
        return datetime.datetime.fromisoformat(str(value).replace(" ", "T"))
    except ValueError:
        return None


def job_id(user_id, kind, subject):
    return hashlib.sha256(
        f"{user_id}|{kind}|{subject}".encode("utf-8")
    ).hexdigest()[:64]


def enqueue(user_id, kind, subject, fingerprint=None,
            priority=PRIORITY_PREFETCH):
    """Queue one generation, or promote it if it is already queued.

    Returns the job's current status. Idempotent by design: the screen polls
    every few seconds while it waits, and every one of those polls calls this.

    A finished job is re-queued only when the data it was derived from has
    changed — that is what `fingerprint` is for, and it is why a poll on
    unchanged data does not restart the work that just completed.
    """
    jid = job_id(user_id, kind, subject)
    now = _iso(_now())
    row = query_one("SELECT * FROM ai_jobs WHERE id = ?", [jid])

    if row:
        if row["status"] == "running":
            claimed = _parse(row.get("claimed_at"))
            if claimed and (_now() - claimed).total_seconds() < CLAIM_TIMEOUT_SECONDS:
                # Already being generated. Promoting it now would not make the
                # in-flight call any faster.
                return "running"

        stale = fingerprint and row.get("fingerprint") != fingerprint
        if row["status"] == "done" and not stale:
            return "done"
        if row["status"] == "failed" and row["attempts"] >= MAX_ATTEMPTS and not stale:
            gave_up = _parse(row.get("updated_at"))
            if gave_up and (_now() - gave_up).total_seconds() < FAILED_RETRY_SECONDS:
                return "failed"
            # Past the cooldown: reset the attempt count so the outage that
            # exhausted it does not count against the retry.
            execute(
                "UPDATE ai_jobs SET status = 'pending', attempts = 0, error = NULL, "
                "holder = NULL, claimed_at = NULL, priority = ?, updated_at = ? "
                "WHERE id = ?",
                [min(priority, row["priority"]), now, jid],
            )
            return "pending"

        # Promote (never demote): a screen the user just opened must not be
        # pushed back by the prefetch entry that was already sitting there.
        execute(
            "UPDATE ai_jobs SET status = 'pending', priority = ?, "
            "fingerprint = ?, holder = NULL, claimed_at = NULL, "
            "attempts = ?, updated_at = ? WHERE id = ?",
            [
                min(priority, row["priority"]),
                fingerprint or row.get("fingerprint"),
                0 if stale else row["attempts"],
                now, jid,
            ],
        )
        return "pending"

    cols = ["id", "user_id", "kind", "subject", "fingerprint", "priority",
            "status", "attempts", "created_at", "updated_at"]
    execute(
        f"INSERT INTO ai_jobs ({', '.join(cols)}) "
        f"VALUES ({', '.join(['?'] * len(cols))})",
        [jid, user_id, kind, subject, fingerprint, priority, "pending", 0,
         now, now],
    )
    return "pending"


def supersede(user_id, kind, subject):
    """Drop queued work of the same kind for a different subject.

    A single-entry scope has one live subject; anything else queued under that
    kind is about a snapshot that no longer exists, and running it would spend
    a gateway call on an answer nothing will read.

    This is a safety net, not the mechanism: subjects are supposed to be stable
    (see scopes.py). It exists because they were not — a row count in the key
    made every poll mint a new `trends` job, and production had 36 of them
    queued before anyone noticed.
    """
    execute(
        "DELETE FROM ai_jobs WHERE user_id = ? AND kind = ? AND subject != ? "
        "AND status = 'pending'",
        [user_id, kind, subject],
    )


def status_of(user_id, kind, subject):
    row = query_one("SELECT * FROM ai_jobs WHERE id = ?",
                    [job_id(user_id, kind, subject)])
    if not row:
        return None
    return {
        "status": row["status"],
        "priority": row["priority"],
        "attempts": row["attempts"],
        "error": row.get("error"),
        "updatedAt": row.get("updated_at"),
    }


def list_jobs(user_id, limit=60):
    """The queue as it stands, for the 设置 screen.

    Ordered the way the worker will actually take them — priority, then age —
    so the list reads as "what happens next" rather than as a table of rows.
    Finished jobs come last: they are history, not queue.
    """
    rows = query_all(
        "SELECT * FROM ai_jobs WHERE user_id = ? "
        "ORDER BY CASE status WHEN 'running' THEN 0 WHEN 'pending' THEN 1 "
        "WHEN 'failed' THEN 2 ELSE 3 END, priority ASC, created_at ASC",
        [user_id],
    )
    return [{
        "kind": r["kind"],
        "subject": r["subject"],
        "status": r["status"],
        "priority": r["priority"],
        # The queue only distinguishes "the user is looking at this" from
        # "backfill"; showing the raw number would mean explaining the scale.
        "interactive": r["priority"] <= PRIORITY_INTERACTIVE,
        "attempts": r["attempts"],
        "maxAttempts": MAX_ATTEMPTS,
        "error": r.get("error"),
        "updatedAt": r.get("updated_at"),
    } for r in rows[:limit]]


def settings():
    """What the consumer is configured to do, for the same screen."""
    return {
        "enabled": ENABLED,
        "concurrency": MAX_CONCURRENT,
        "gapSeconds": GAP_SECONDS,
        "maxAttempts": MAX_ATTEMPTS,
        "retryAfterSeconds": FAILED_RETRY_SECONDS,
    }


def retry_failed(user_id):
    """Put every given-up job back in the queue, now, at the user's request.

    The automatic retry waits out a cooldown so a flapping upstream is not
    hammered; a person pressing 重试 has decided it is worth trying again.
    """
    execute(
        "UPDATE ai_jobs SET status = 'pending', attempts = 0, error = NULL, "
        "holder = NULL, claimed_at = NULL, updated_at = ? "
        "WHERE user_id = ? AND status = 'failed'",
        [_iso(_now()), user_id],
    )
    return pending_count(user_id)


def pending_count(user_id=None):
    sql = "SELECT COUNT(*) AS n FROM ai_jobs WHERE status IN ('pending', 'running')"
    params = []
    if user_id:
        sql += " AND user_id = ?"
        params.append(user_id)
    row = query_one(sql, params)
    return (row or {}).get("n") or 0


def _claim_next():
    """Take the highest-priority runnable job, or None.

    Ordered by priority then age so the interactive job wins and, among equals,
    the one that has waited longest goes first.
    """
    cutoff = _iso(_now() - datetime.timedelta(seconds=CLAIM_TIMEOUT_SECONDS))

    # Concurrency is counted across the deployment, not per process: two
    # Gunicorn workers each running "one job" is two concurrent calls into a
    # gateway that has four threads for three projects.
    running = query_one(
        "SELECT COUNT(*) AS n FROM ai_jobs WHERE status = 'running' "
        "AND claimed_at IS NOT NULL AND claimed_at >= ?",
        [cutoff],
    )
    if (running or {}).get("n", 0) >= MAX_CONCURRENT:
        return None

    rows = query_all(
        "SELECT * FROM ai_jobs WHERE status = 'pending' "
        "OR (status = 'running' AND (claimed_at IS NULL OR claimed_at < ?)) "
        "ORDER BY priority ASC, created_at ASC",
        [cutoff],
    )
    holder = f"{os.getpid()}-{threading.get_ident()}"

    for row in rows:
        if row["attempts"] >= MAX_ATTEMPTS:
            continue
        execute(
            "UPDATE ai_jobs SET status = 'running', holder = ?, claimed_at = ?, "
            "attempts = ?, updated_at = ? WHERE id = ? AND status = ?",
            [holder, _iso(_now()), row["attempts"] + 1, _iso(_now()),
             row["id"], row["status"]],
        )
        # Re-read: another worker may have claimed it between the SELECT and
        # the UPDATE, in which case its holder is the one now recorded.
        check = query_one("SELECT holder FROM ai_jobs WHERE id = ?", [row["id"]])
        if check and check.get("holder") == holder:
            return row
    return None


def _finish(jid, error=None):
    execute(
        "UPDATE ai_jobs SET status = ?, error = ?, holder = NULL, "
        "claimed_at = NULL, updated_at = ? WHERE id = ?",
        ["failed" if error else "done", (error or "")[:500] if error else None,
         _iso(_now()), jid],
    )


def run_once():
    """Claim and run one job. Returns True when something was run."""
    if _runner is None:
        return False
    row = _claim_next()
    if not row:
        return False
    try:
        _runner(row["user_id"], row["kind"], row["subject"])
    except Exception as e:  # noqa: BLE001 - one bad job must not stop the queue
        _finish(row["id"], f"{type(e).__name__}: {e}")
        print(f"[ai-jobs] {row['kind']}:{row['subject']} failed: {e}")
        return True
    _finish(row["id"])
    return True


def _loop():
    idle_cycles = 0
    while True:
        try:
            # A gap after each job, not just when idle. Running them
            # back-to-back is what saturates the shared gateway: a backfill
            # finishing twenty minutes later costs nobody anything, a gateway
            # that stops answering costs all three projects.
            ran = run_once()
            if ran:
                idle_cycles = 0
                time.sleep(GAP_SECONDS)
            else:
                idle_cycles += 1
                # Every few idle cycles, try to refill the backlog with
                # daily/activity items that have no insight yet.
                if idle_cycles >= 5 and _refiller is not None:
                    try:
                        n = _refiller()
                        if n:
                            print(f"[ai-jobs] refilled {n} backlog items")
                    except Exception as e:  # noqa: BLE001
                        print(f"[ai-jobs] refill error: {e}")
                    idle_cycles = 0
                time.sleep(POLL_SECONDS)
        except Exception as e:  # noqa: BLE001 - the loop must outlive any failure
            print(f"[ai-jobs] worker error: {e}")
            time.sleep(POLL_SECONDS)


def start():
    """Start one consumer per process."""
    global _started
    if not ENABLED:
        print("[ai-jobs] disabled by AI_JOBS")
        return
    with _start_lock:
        if _started:
            return
        _started = True
    threading.Thread(target=_loop, daemon=True, name="ai-jobs").start()
    print("[ai-jobs] worker started")


def reset_stale_claims():
    """Release jobs a previous process was running when it stopped.

    Without this they sit in `running` until the claim expires, which for the
    screen waiting on one looks exactly like a generation that never finishes.
    """
    execute(
        "UPDATE ai_jobs SET status = 'pending', holder = NULL, claimed_at = NULL "
        "WHERE status = 'running'"
    )