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) -> None:
        """Take the job out of the queue before doing any work on it.

        Without this the next poll would pick the same row up again while the
        first execution is still running on the GPU.
        """

        job.status = "running"
        db.add(job)
        db.commit()

    @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"
            ]
            for job in rows:
                try:
                    AnalysisJobWorker._claim(session, job)
                    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 len(rows)
        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