Initial public release

backend/app/services/analysis_job_worker.py

@@ -0,0 +1,172 @@
+"""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