FabledCurator/backend/app/services/ml/gpu_triage.py

"""GPU-failure triage (#125): classify errored jobs, PROBE the file, recover.

An errored GPU job is a tombstone with a stored reason, but the reason alone is
a suspicion, not a verdict — a timeout can hit a perfectly fine file, and
"moov atom not found" can mean a truncated download OR a one-off transfer
fault. So triage EVALUATES: it runs the real integrity probe (sha256 recompute
+ PIL/ffprobe — verify_integrity's own machinery) on each errored image ONCE
and records both verdicts:

  ImageRecord.integrity_status  <- file-level verdict (ok / corrupt / ...)
  GpuJob.triage_status          <- 'defect'  (file is bad: recovery material,
                                              excluded from /retry_errors)
                                   'file_ok' (file passes: the failure was
                                              operational, safe to retry)

Recovery reuses established primitives: delete the defective copy + record
(cleanup_service.delete_images — full cascade) and re-poll the image's
subscription Source (the Layer-2 refetch pattern: gallery-dl re-fetches the
now-absent file on the next source check). Images without a pollable Source
report 'no_source' — manual remediation. Every classification is logged at
WARNING so the operator notices in Logs / System Activity.
"""
import logging
import time
from datetime import UTC, datetime
from pathlib import Path

from sqlalchemy import select, update
from sqlalchemy.orm import Session

from ...models import GpuJob, ImageProvenance, ImageRecord, Source
from ..cleanup_service import delete_images

log = logging.getLogger(__name__)

# Reason buckets for the triage overview (reporting only — the PROBE decides
# 'defect', never the string). Ordered: first match wins.
_REASON_BUCKETS = (
    ("poisoned", ("poisoned:",)),
    ("transient", ("gave up after repeated transient", "curator unreachable",
                   "connection", "read timed out")),
    ("timeout", ("timed out", "timeout")),
    ("truncated_or_corrupt", ("moov atom", "invalid data", "end of file",
                              "header missing", "error reading header",
                              "truncated", "premature", "corrupt",
                              "no frames sampled")),
    ("decode", ("cannot identify", "decompression", "broken data stream",
                "unrecognized data")),
)


def classify_reason(error: str | None) -> str:
    """Bucket a stored job-error string for the overview table."""
    text = (error or "").lower()
    if not text:
        return "other"
    for bucket, needles in _REASON_BUCKETS:
        if any(n in text for n in needles):
            return bucket
    return "other"


def triage_errored_jobs(
    session: Session, *, time_budget_seconds: float = 300.0,
) -> dict:
    """Probe every not-yet-triaged errored image and write both verdicts.

    Time-boxed (sha256 of a large original over NFS can take tens of seconds)
    and inherently resumable: rows are selected by `triage_status IS NULL`, so
    the next sweep continues exactly where a budget cut stopped. Commits per
    image so a mid-run crash keeps completed verdicts."""
    image_ids = session.execute(
        select(GpuJob.image_record_id)
        .where(GpuJob.status == "error", GpuJob.triage_status.is_(None))
        .group_by(GpuJob.image_record_id)
        .order_by(GpuJob.image_record_id)
    ).scalars().all()
    counts = {"probed": 0, "defect": 0, "file_ok": 0, "partial": False}
    if not image_ids:
        return counts
    # Lazy imports: the probe helper lives in the maintenance task module and
    # the hasher in the importer — importing either at module load would pull
    # celery into every service consumer.
    from ...tasks.maintenance import _verify_one
    from ..importer import _sha256_of

    started = time.monotonic()
    for image_id in image_ids:
        if time.monotonic() - started > time_budget_seconds:
            counts["partial"] = True
            break
        rec = session.get(ImageRecord, image_id)
        if rec is None:      # record deleted since the job errored
            continue
        verdict = _verify_one(Path(rec.path), rec.sha256, rec.mime, _sha256_of)
        # 'ok' means the failure was operational; anything else (corrupt /
        # failed_verification = missing/unreadable) makes the file itself the
        # problem — recovery material.
        triage = "file_ok" if verdict == "ok" else "defect"
        reason = session.execute(
            select(GpuJob.error)
            .where(GpuJob.image_record_id == image_id, GpuJob.status == "error")
            .limit(1)
        ).scalar_one_or_none()
        rec.integrity_status = verdict
        session.execute(
            update(GpuJob)
            .where(GpuJob.image_record_id == image_id, GpuJob.status == "error")
            .values(triage_status=triage, updated_at=datetime.now(UTC))
        )
        session.commit()
        counts["probed"] += 1
        counts[triage] += 1
        log.warning(
            "gpu triage: image %s (%s) job error %r -> integrity probe %r -> %s",
            image_id, rec.path, (reason or "")[:120], verdict, triage,
        )
    return counts


def recover_defective_image(
    session: Session, image_id: int, *, images_root: Path,
) -> dict:
    """Delete the defective copy + record and queue its re-fetch — surgically
    where possible.

    Two re-fetch layers (operator 2026-07-03: deep back-catalogue items are
    NEVER re-walked by the normal cadence, so recovery can't rely on it):
    1. SURGICAL: any ExternalLink rows on the image's post(s) are reset +
       re-dispatched — this is how external-host files (the common defect
       case: big videos) come back regardless of post age. Sha-dedupe at
       import discards payload files that still exist.
    2. BROAD: a source re-check, which re-fetches gallery-dl-NATIVE files the
       walk still reaches (recent posts). A native file deeper than the walk
       needs a per-source backfill/deep scan — reported via links_reset=0 so
       the caller can say so.

    The record delete cascades the error tombstones with it. 'no_source' when
    no enabled, real-URL Source resolves via provenance — manual there."""
    rec = session.get(ImageRecord, image_id)
    if rec is None:
        return {"status": "not_found"}
    src_id = session.execute(
        select(Source.id)
        .join(ImageProvenance, ImageProvenance.source_id == Source.id)
        .where(
            ImageProvenance.image_record_id == image_id,
            Source.enabled.is_(True),
            ~Source.url.like("sidecar:%"),   # synthetic anchor — not pollable
        )
        .order_by(Source.id.asc())
    ).scalars().first()
    if src_id is None:
        return {"status": "no_source"}
    # Capture the post linkage BEFORE the delete cascades provenance away.
    post_ids = session.execute(
        select(ImageProvenance.post_id)
        .where(ImageProvenance.image_record_id == image_id)
    ).scalars().all()
    path = rec.path
    summary = delete_images(session, image_ids=[image_id], images_root=images_root)
    from ..external_links import refetch_links_for_post

    links_reset = 0
    for pid in post_ids:
        links_reset += refetch_links_for_post(session, pid)["links_reset"]
    # Lazy import (services -> tasks would cycle at module load).
    from ...tasks.download import download_source

    download_source.delay(src_id)
    log.warning(
        "gpu triage recovery: deleted defective image %s (%s); reset %d "
        "external link(s) and queued a re-check of source %s",
        image_id, path, links_reset, src_id,
    )
    return {
        "status": "refetch_queued", "source_id": src_id,
        "links_reset": links_reset, **summary,
    }