update the cv modal inference proxy server with optimization

workspace/sprint_1_2/CODEBASE/backend/services/celery_app.py

@@ -0,0 +1,22 @@
+from celery import Celery
+
+from backend.implementation import config
+
+celery_app = Celery(
+    "cv-inference",
+    broker=f"redis://{config.REDIS_HOST}:{config.REDIS_PORT}/{config.REDIS_DB}",
+    backend=f"redis://{config.REDIS_HOST}:{config.REDIS_PORT}/{config.REDIS_DB}",
+    # Explicit include — autodiscover looks for tasks.py, not cv_tasks.py
+    include=["backend.implementation.tasks.cv_tasks"],
+)
+
+# Must match the @task(name=...) value, not the Python module path
+celery_app.conf.task_routes = {
+    "cv_inference.run_chunk": {"queue": "cv-inference"},
+}
+
+celery_app.conf.task_serializer = "json"
+celery_app.conf.result_serializer = "json"
+celery_app.conf.accept_content = ["json"]
+celery_app.conf.result_expires = 3600
+celery_app.conf.task_default_queue = "cv-inference"

workspace/sprint_1_2/CODEBASE/backend/services/cv_celery_service.py

@@ -0,0 +1,130 @@
+import base64
+import dataclasses
+import logging
+import os
+import time
+from typing import Any
+
+from backend.services.cv_inference_service import CvInferenceOptions, _encode_image_to_bytes
+from backend.implementation.tasks.cv_tasks import run_cv_chunk
+from backend.services.celery_app import celery_app
+
+logger = logging.getLogger(__name__)
+
+CELERY_CHUNK_SIZE = int(os.getenv("CELERY_CHUNK_SIZE", "4"))
+CELERY_BATCH_POLL_CACHE_TTL_MS = float(os.getenv("CELERY_BATCH_POLL_CACHE_TTL_MS", "2000"))
+
+_batch_poll_cache: dict[str, tuple[float, dict[str, Any]]] = {}
+_batch_timings: dict[str, float] = {}
+
+
+def submit_celery_batch(images, frame_ids, options):
+    if not images:
+        raise ValueError("images must not be empty")
+    if len(frame_ids) != len(images):
+        raise ValueError("frame_ids length must match images length")
+
+    from celery import group
+
+    chunks = []
+    for i in range(0, len(images), CELERY_CHUNK_SIZE):
+        chunk_imgs = images[i : i + CELERY_CHUNK_SIZE]
+        chunk_fids = frame_ids[i : i + CELERY_CHUNK_SIZE]
+        b64_images = [
+            base64.b64encode(_encode_image_to_bytes(img)).decode() for img in chunk_imgs
+        ]
+        chunk_payload = {
+            "images_b64": b64_images,
+            "frame_ids": chunk_fids,
+            "calibration": dataclasses.asdict(options.calibration) if options.calibration else {},
+            "model_versions": options.model_versions,
+        }
+        chunks.append(chunk_payload)
+
+    job = group(run_cv_chunk.s(chunk) for chunk in chunks)
+    result = job.apply_async()
+    try:
+        result.save()
+    except Exception:
+        logger.exception("Failed to persist Celery GroupResult for job %s", result.id)
+
+    _batch_timings[result.id] = time.monotonic()
+    logger.info(
+        "Celery batch submitted job_id=%s chunks=%d images=%d",
+        result.id,
+        len(chunks),
+        len(images),
+    )
+    return result.id
+
+
+async def get_celery_batch_result(job_id):
+    from celery.result import GroupResult
+
+    now = time.monotonic()
+    cached = _batch_poll_cache.get(job_id)
+    if cached and cached[0] > now:
+        return cached[1]
+
+    result = GroupResult.restore(job_id, app=celery_app)
+    if result is None:
+        payload = {
+            "status": "unknown",
+            "detail": "Job ID not found in result backend",
+        }
+        _batch_poll_cache[job_id] = (now + CELERY_BATCH_POLL_CACHE_TTL_MS / 1000, payload)
+        return payload
+
+    completed = result.completed_count()
+    total = len(result.results)
+
+    if not result.ready():
+        payload = {
+            "status": "pending",
+            "completed": completed,
+            "total": total,
+            "progress": round(completed / total, 2) if total > 0 else 0,
+        }
+        _batch_poll_cache[job_id] = (now + CELERY_BATCH_POLL_CACHE_TTL_MS / 1000, payload)
+        return payload
+
+    if result.failed():
+        errors = []
+        for r in result.results:
+            if r.failed():
+                errors.append(str(r.result))
+        payload = {
+            "status": "failed",
+            "errors": errors,
+        }
+        _batch_poll_cache[job_id] = (now + CELERY_BATCH_POLL_CACHE_TTL_MS / 1000, payload)
+        _log_batch_timing(job_id, "failed")
+        return payload
+
+    all_results = []
+    for r in result.results:
+        chunk_results = r.get()
+        all_results.extend(chunk_results)
+
+    payload = {
+        "status": "completed",
+        "image_count": len(all_results),
+        "results": all_results,
+    }
+    _batch_poll_cache[job_id] = (now + CELERY_BATCH_POLL_CACHE_TTL_MS / 1000, payload)
+    _log_batch_timing(job_id, "completed")
+    return payload
+
+
+def _log_batch_timing(job_id: str, status: str) -> None:
+    start = _batch_timings.pop(job_id, None)
+    if start is not None:
+        duration_ms = (time.monotonic() - start) * 1000
+        logger.info(
+            "Celery batch %s job_id=%s duration_ms=%.0f",
+            status,
+            job_id,
+            duration_ms,
+        )
+    else:
+        logger.info("Celery batch %s job_id=%s (no start time recorded)", status, job_id)

workspace/sprint_1_2/CODEBASE/backend/services/cv_inference_service.py

@@ -5,6 +5,7 @@ import asyncio
 import base64
 import io
 import logging
+import os
 from dataclasses import dataclass
 from typing import Any
 
@@ -52,8 +53,6 @@ SEGMENT_CLASSES_POST = {
     6: "baker's cyst",
 }
 
-_triton_pipeline_lock = asyncio.Lock()
-
 
 @dataclass
 class CvInferenceOptions:
@@ -76,6 +75,12 @@ def _encode_image_to_data_url(image_pil: Image.Image) -> str:
     return f"data:image/png;base64,{encoded}"
 
 
+def _encode_image_to_bytes(image_pil: Image.Image) -> bytes:
+    buffered = io.BytesIO()
+    image_pil.save(buffered, format="PNG")
+    return buffered.getvalue()
+
+
 def _build_inflammation_payload(logits_row: np.ndarray, config: CalibrationConfig) -> dict:
     interpreted = interpret_inflammation_logits(logits_row, config)
     return {
@@ -147,6 +152,9 @@ def _build_segmentation_result(
     classes_detected = [name for name, mask in masks.items() if np.sum(mask) > 0 and name != "background"]
     color_legend = _build_color_legend(classes_detected, angle_type)
 
+    segmented_png_bytes = _encode_image_to_bytes(overlay)
+    segmented_b64 = base64.b64encode(segmented_png_bytes).decode()
+
     result: dict[str, Any] = {
         "success": True,
         "angle": angle_payload,
@@ -161,7 +169,7 @@ def _build_segmentation_result(
         },
         "images": {
             "enhanced": enhanced_data_url,
-            "segmented": _encode_image_to_data_url(overlay),
+            "segmented": f"data:image/png;base64,{segmented_b64}",
         },
         "models_used": {
             "angle": angle_model,
@@ -281,19 +289,30 @@ async def _run_batch_uncached(
     if len(frame_ids) != len(images):
         raise ValueError("frame_ids length must match images length")
 
-    async with _triton_pipeline_lock:
-        results: list[dict[str, Any]] = []
-        infer_modes: list[str] = []
-        triton_call_count = 0
-        for image_pil, fid in zip(images, frame_ids, strict=True):
-            item, mode, calls = await _run_spec_cv_pipeline_single(
+    concurrency = int(os.getenv("CV_BATCH_CONCURRENCY", "2"))
+    semaphore = asyncio.Semaphore(concurrency)
+
+    async def process_one(image_pil: Image.Image, fid: str, index: int):
+        async with semaphore:
+            if index > 0:
+                await asyncio.sleep(min(index * 0.15, 1.0))
+            return await _run_spec_cv_pipeline_single(
                 image_pil,
                 frame_id=fid,
                 options=options,
             )
-            results.append(item)
-            infer_modes.append(mode)
-            triton_call_count += calls
+
+    outcomes = await asyncio.gather(
+        *[process_one(img, fid, idx) for idx, (img, fid) in enumerate(zip(images, frame_ids))],
+    )
+
+    results: list[dict[str, Any]] = []
+    infer_modes: list[str] = []
+    triton_call_count = 0
+    for item, mode, calls in outcomes:
+        results.append(item)
+        infer_modes.append(mode)
+        triton_call_count += calls
     return CvBatchResult(
         results=results,
         triton_infer_calls=triton_call_count,

workspace/sprint_1_2/CODEBASE/backend/services/triton_runtime_service.py

@@ -25,12 +25,11 @@ INPUT_NAME = "input_image"
 OUTPUT_NAME = "logits"
 
 TRITON_INFER_TIMEOUT = float(os.getenv("TRITON_INFER_TIMEOUT", "120"))
-TRITON_INFER_RETRIES = int(os.getenv("TRITON_INFER_RETRIES", "5"))
-TRITON_RETRY_BASE_S = float(os.getenv("TRITON_RETRY_BASE_S", "4"))
+TRITON_INFER_RETRIES = int(os.getenv("TRITON_INFER_RETRIES", "2"))
+TRITON_RETRY_BASE_S = float(os.getenv("TRITON_RETRY_BASE_S", "2.0"))
 TRITON_USE_BATCH_INFER = os.getenv("TRITON_USE_BATCH_INFER", "true").lower()
 RETRYABLE_STATUS = {429, 502, 503, 504}
 
-_triton_infer_lock = asyncio.Lock()
 _adapter: TritonAdapter | None = None
 _adapter_endpoint: str | None = None
 
@@ -43,6 +42,11 @@ def _get_adapter() -> TritonAdapter:
     global _adapter, _adapter_endpoint
     endpoint = get_triton_endpoint()
     if _adapter is None or _adapter_endpoint != endpoint:
+        if _adapter is not None:
+            try:
+                asyncio.get_event_loop().run_until_complete(_adapter.close())
+            except RuntimeError:
+                pass
         _adapter = TritonAdapter(endpoint_url=endpoint, timeout=TRITON_INFER_TIMEOUT)
         _adapter_endpoint = endpoint
     return _adapter
@@ -201,7 +205,7 @@ def _infer_angle_logits_chunk(images: list[Image.Image], model_name: str) -> tup
     if not _should_try_batched_infer(len(images)):
         return _infer_angle_logits_sequential(images, model_name), "sequential"
     try:
-        return _infer_angle_logits_batch(images, model_name, max_retries=1), "batched"
+        return _infer_angle_logits_batch(images, model_name), "batched"
     except Exception as exc:
         logger.warning(
             "Batched angle infer×%s failed (%s); falling back to sequential single-image calls",
@@ -249,7 +253,7 @@ def _infer_inflammation_logits_chunk(
     if not _should_try_batched_infer(len(images)):
         return _infer_inflammation_logits_sequential(images, model_name), "sequential"
     try:
-        return _infer_inflammation_logits_batch(images, model_name, max_retries=1), "batched"
+        return _infer_inflammation_logits_batch(images, model_name), "batched"
     except Exception as exc:
         logger.warning(
             "Batched inflammation infer×%s failed (%s); falling back to sequential",
@@ -297,7 +301,7 @@ def _infer_segmentation_logits_chunk(
     if not _should_try_batched_infer(len(images)):
         return _infer_segmentation_logits_sequential(images, model_name), "sequential"
     try:
-        return _infer_segmentation_logits_batch(images, model_name, max_retries=1), "batched"
+        return _infer_segmentation_logits_batch(images, model_name), "batched"
     except Exception as exc:
         logger.warning(
             "Batched segmentation infer×%s failed (%s); falling back to sequential",
@@ -311,61 +315,57 @@ async def infer_angle_logits(
     images: list[Image.Image],
     model_name: str,
 ) -> tuple[np.ndarray, Literal["batched", "sequential"], int]:
-    """Run angle classification under the global Triton lock. Returns (logits, mode, call_count)."""
-    async with _triton_infer_lock:
-        all_logits: list[np.ndarray] = []
-        modes: list[str] = []
-        call_count = 0
-        for chunk in chunk_sequence(images):
-            logits, mode = await asyncio.to_thread(_infer_angle_logits_chunk, chunk, model_name)
-            all_logits.append(logits)
-            modes.append(mode)
-            call_count += 1 if mode == "batched" else len(chunk)
-        combined = np.concatenate(all_logits, axis=0) if len(all_logits) > 1 else all_logits[0]
-        infer_mode: Literal["batched", "sequential"] = (
-            "sequential" if any(m == "sequential" for m in modes) else "batched"
-        )
-        return combined, infer_mode, call_count
+    all_logits: list[np.ndarray] = []
+    modes: list[str] = []
+    call_count = 0
+    for chunk in chunk_sequence(images):
+        logits, mode = await asyncio.to_thread(_infer_angle_logits_chunk, chunk, model_name)
+        all_logits.append(logits)
+        modes.append(mode)
+        call_count += 1 if mode == "batched" else len(chunk)
+    combined = np.concatenate(all_logits, axis=0) if len(all_logits) > 1 else all_logits[0]
+    infer_mode: Literal["batched", "sequential"] = (
+        "sequential" if any(m == "sequential" for m in modes) else "batched"
+    )
+    return combined, infer_mode, call_count
 
 
 async def infer_inflammation_logits(
     images: list[Image.Image],
     model_name: str,
 ) -> tuple[np.ndarray, Literal["batched", "sequential"], int]:
-    async with _triton_infer_lock:
-        all_logits: list[np.ndarray] = []
-        modes: list[str] = []
-        call_count = 0
-        for chunk in chunk_sequence(images):
-            logits, mode = await asyncio.to_thread(_infer_inflammation_logits_chunk, chunk, model_name)
-            all_logits.append(logits)
-            modes.append(mode)
-            call_count += 1 if mode == "batched" else len(chunk)
-        combined = np.concatenate(all_logits, axis=0) if len(all_logits) > 1 else all_logits[0]
-        infer_mode: Literal["batched", "sequential"] = (
-            "sequential" if any(m == "sequential" for m in modes) else "batched"
-        )
-        return combined, infer_mode, call_count
+    all_logits: list[np.ndarray] = []
+    modes: list[str] = []
+    call_count = 0
+    for chunk in chunk_sequence(images):
+        logits, mode = await asyncio.to_thread(_infer_inflammation_logits_chunk, chunk, model_name)
+        all_logits.append(logits)
+        modes.append(mode)
+        call_count += 1 if mode == "batched" else len(chunk)
+    combined = np.concatenate(all_logits, axis=0) if len(all_logits) > 1 else all_logits[0]
+    infer_mode: Literal["batched", "sequential"] = (
+        "sequential" if any(m == "sequential" for m in modes) else "batched"
+    )
+    return combined, infer_mode, call_count
 
 
 async def infer_segmentation_logits(
     images: list[Image.Image],
     model_name: str,
 ) -> tuple[np.ndarray, Literal["batched", "sequential"], int]:
-    async with _triton_infer_lock:
-        all_logits: list[np.ndarray] = []
-        modes: list[str] = []
-        call_count = 0
-        for chunk in chunk_sequence(images):
-            logits, mode = await asyncio.to_thread(_infer_segmentation_logits_chunk, chunk, model_name)
-            all_logits.append(logits)
-            modes.append(mode)
-            call_count += 1 if mode == "batched" else len(chunk)
-        combined = np.concatenate(all_logits, axis=0) if len(all_logits) > 1 else all_logits[0]
-        infer_mode: Literal["batched", "sequential"] = (
-            "sequential" if any(m == "sequential" for m in modes) else "batched"
-        )
-        return combined, infer_mode, call_count
+    all_logits: list[np.ndarray] = []
+    modes: list[str] = []
+    call_count = 0
+    for chunk in chunk_sequence(images):
+        logits, mode = await asyncio.to_thread(_infer_segmentation_logits_chunk, chunk, model_name)
+        all_logits.append(logits)
+        modes.append(mode)
+        call_count += 1 if mode == "batched" else len(chunk)
+    combined = np.concatenate(all_logits, axis=0) if len(all_logits) > 1 else all_logits[0]
+    infer_mode: Literal["batched", "sequential"] = (
+        "sequential" if any(m == "sequential" for m in modes) else "batched"
+    )
+    return combined, infer_mode, call_count
 
 
 async def infer_angle_logits_single(image: Image.Image, model_name: str) -> tuple[np.ndarray, str, int]:

workspace/sprint_1_2/CODEBASE/backend/services/triton_warmup.py

@@ -1,6 +1,7 @@
 """Warm up Modal Triton on CV service startup — reduces cold-start 502s and first-frame latency."""
 from __future__ import annotations
 
+import asyncio
 import logging
 import os
 
@@ -23,6 +24,20 @@ def _warmup_model_versions() -> dict[str, str]:
     return versions
 
 
+async def _warmup_one(
+    name: str,
+    coro,
+    timeout: float,
+) -> None:
+    try:
+        await asyncio.wait_for(coro, timeout=timeout)
+        logger.info("Triton warmup %s complete", name)
+    except asyncio.TimeoutError:
+        logger.warning("Triton warmup %s timed out after %.1fs", name, timeout)
+    except Exception as exc:
+        logger.warning("Triton warmup %s failed: %s", name, exc)
+
+
 async def warmup_triton_models() -> None:
     if os.getenv("CV_SKIP_TRITON_WARMUP", "").lower() in {"1", "true", "yes"}:
         logger.info("Triton warmup skipped (CV_SKIP_TRITON_WARMUP)")
@@ -36,13 +51,29 @@ async def warmup_triton_models() -> None:
     img224 = Image.new("RGB", (224, 224), color=(128, 128, 128))
     img512 = Image.new("RGB", (512, 512), color=(128, 128, 128))
 
+    warmup_timeout = float(os.getenv("TRITON_WARMUP_TIMEOUT", "15"))
+
     logger.info(
-        "Warming up Triton (angle=%s, inflammation=%s, segmentation=%s)…",
+        "Warming up Triton (angle=%s, inflammation=%s, segmentation=%s, timeout=%.1fs)…",
         angle_model,
         inflam_model,
         seg_model,
+        warmup_timeout,
+    )
+
+    await _warmup_one(
+        "angle",
+        triton_runtime.infer_angle_logits_single(img224, angle_model),
+        warmup_timeout,
+    )
+    await _warmup_one(
+        "inflammation",
+        triton_runtime.infer_inflammation_logits_single(img224, inflam_model),
+        warmup_timeout,
+    )
+    await _warmup_one(
+        "segmentation",
+        triton_runtime.infer_segmentation_logits_single(img512, seg_model),
+        warmup_timeout,
     )
-    await triton_runtime.infer_angle_logits_single(img224, angle_model)
-    await triton_runtime.infer_inflammation_logits_single(img224, inflam_model)
-    await triton_runtime.infer_segmentation_logits_single(img512, seg_model)
     logger.info("Triton warmup complete")