geointel/backend/app/services/analysis_job_worker.py

"""Background execution for queued analysis runs.

Tiled GPU inference is minutes of work. Running it inside the HTTP request
holds a worker thread for the whole duration, times the client out and leaves
the operator without progress. Queued ``detection.run`` and
``segmentation.run`` jobs are picked up here instead, mirroring the polling
worker the AOI operations already use so the runtime keeps one job model.
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any
from uuid import UUID

from app.db.session import SessionLocal
from app.models import Job

logger = logging.getLogger("geointel.analysis_worker")


class AnalysisJobWorker:
    HANDLED_JOB_TYPES = ("detection.run", "segmentation.run")
    BATCH_SIZE = 4

    @staticmethod
    def _uuid(value: Any) -> UUID | None:
        if isinstance(value, UUID):
            return value
        try:
            return UUID(str(value))
        except (TypeError, ValueError):
            return None

    @staticmethod
    def _dispatch(db, job: Job) -> Any:
        # Imported lazily: both services import each other's helpers, and the
        # worker must not add a third edge to that cycle at module load.
        from app.services.detection_service import DetectionService
        from app.services.segmentation_service import SegmentationService

        parameters = job.parameters_json if isinstance(job.parameters_json, dict) else {}
        project_id = AnalysisJobWorker._uuid(parameters.get("project_id"))
        dataset_id = AnalysisJobWorker._uuid(parameters.get("dataset_id"))
        if project_id is None or dataset_id is None:
            raise ValueError("Queued analysis job is missing project_id or dataset_id")

        common = {
            "db": db,
            "project_id": project_id,
            "dataset_id": dataset_id,
            "model_id": parameters.get("model_id"),
            "confidence_threshold": float(parameters.get("confidence_threshold") or 0.0),
            "class_filter": parameters.get("class_filter") or [],
            "tile_manifest_path": parameters.get("tile_manifest_path"),
            "parameters_json": parameters.get("parameters_json") or {},
            "existing_job": job,
        }
        if job.job_type == "detection.run":
            return DetectionService.run_detection(
                model_asset_id=parameters.get("model_asset_id"),
                **common,
            )
        return SegmentationService.run_segmentation(**common)

    @staticmethod
    def claim(db, job: Job) -> bool:
        """Take the job out of the queue, atomically. Returns whether we won.

        Selecting and then updating in a second statement lets two workers —
        a restarted process overlapping the previous one, or a second replica —
        both start tiled GPU inference on the same row. The conditional update
        makes exactly one caller see a row count of 1; the AOI worker beside
        this one already claims with FOR UPDATE SKIP LOCKED for the same reason.
        """

        claimed = (
            db.query(Job)
            .filter(Job.id == job.id, Job.status == "queued")
            .update({Job.status: "running"}, synchronize_session=False)
        )
        db.commit()
        if not claimed:
            return False
        job.status = "running"
        return True

    @staticmethod
    def _finalize(db, job: Job, result: Any) -> None:
        """Close a job the handler left open.

        The analysis services normally set the terminal status themselves.
        If one returns without doing so, recording the outcome here is what
        keeps the job from sitting in "running" for ever.
        """

        if job.status != "running":
            return
        status = getattr(result, "status", None)
        if status == "success":
            job.status = "success"
            job.result_json = {
                "detection_count": getattr(result, "detection_count", None),
                "segmentation_count": getattr(result, "segmentation_count", None),
            }
        else:
            job.status = "failed"
            job.error_message = getattr(result, "message", None) or "Analysis run did not complete"
            job.result_json = {"error_code": getattr(result, "error_code", None) or "ANALYSIS_JOB_INCOMPLETE"}
        db.add(job)
        db.commit()

    @staticmethod
    def _mark_failed(db, job: Job, *, code: str, message: str) -> None:
        try:
            db.rollback()
        except Exception:
            pass
        job.status = "failed"
        job.error_message = message
        job.result_json = {"error_code": code, "message": message}
        db.add(job)
        db.commit()

    @staticmethod
    def run_once(db=None) -> int:
        """Execute one batch of queued analysis jobs. Returns the batch size."""

        owns_session = db is None
        session = db if db is not None else SessionLocal()
        try:
            rows = [
                job
                for job in (
                    session.query(Job)
                    .filter(Job.status == "queued")
                    .filter(Job.job_type.in_(AnalysisJobWorker.HANDLED_JOB_TYPES))
                    .order_by(Job.created_at)
                    .limit(AnalysisJobWorker.BATCH_SIZE)
                    .all()
                )
                if job.job_type in AnalysisJobWorker.HANDLED_JOB_TYPES and job.status == "queued"
            ]
            claimed_count = 0
            for job in rows:
                if not AnalysisJobWorker.claim(session, job):
                    # Another worker took it between the select and the claim.
                    continue
                claimed_count += 1
                try:
                    result = AnalysisJobWorker._dispatch(session, job)
                    AnalysisJobWorker._finalize(session, job, result)
                except Exception as exc:
                    code = getattr(exc, "code", None) or "ANALYSIS_JOB_INTERNAL_ERROR"
                    message = getattr(exc, "message", None) or str(exc) or "Unexpected analysis job failure"
                    AnalysisJobWorker._mark_failed(session, job, code=str(code), message=str(message))
                    logger.exception("Analysis job failed job_id=%s job_type=%s", job.id, job.job_type)
            return claimed_count
        finally:
            if owns_session:
                session.close()

    @staticmethod
    async def run(stop_event: asyncio.Event, poll_seconds: float) -> None:
        while not stop_event.is_set():
            processed = await asyncio.to_thread(AnalysisJobWorker.run_once)
            if processed == 0:
                try:
                    await asyncio.wait_for(stop_event.wait(), timeout=poll_seconds)
                except TimeoutError:
                    pass