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 re-poll its subscription Source.

    Mirrors the Layer-2 import refetch: with the bad file gone, the source's
    next gallery-dl run re-fetches a fresh copy, which re-imports as a new
    record and re-enters the GPU pipeline. The record delete cascades the
    error tombstones with it. 'no_source' when no enabled, real-URL Source is
    reachable via the image's provenance — manual remediation 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"}
    path = rec.path
    summary = delete_images(session, image_ids=[image_id], images_root=images_root)
    # 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) and queued a "
        "re-check of source %s to re-fetch it", image_id, path, src_id,
    )
    return {"status": "refetch_queued", "source_id": src_id, **summary}