Add backend progress polling for job predictions

web/main.py

@@ -5,7 +5,9 @@ import io
 import json
 import os
 import re
+import threading
 import time
+import uuid
 from pathlib import Path
 from typing import Any, Optional
 
@@ -131,6 +133,20 @@ class VQAServerState:
         self.preload_models = os.getenv("WEB_PRELOAD_MODELS", "0") == "1"
         # Chạy lần lượt và giải phóng model sau mỗi lượt để giảm đỉnh RAM/VRAM.
         self.release_after_predict = os.getenv("WEB_RELEASE_AFTER_PREDICT", "1") == "1"
+        self.progress_state: dict[str, Any] = {
+            "job_id": "",
+            "active": False,
+            "status": "idle",
+            "current_variant": "",
+            "current_index": 0,
+            "total": 0,
+            "completed": 0,
+            "message": "Idle",
+            "updated_at": time.time(),
+        }
+        self.latest_result: dict[str, Any] | None = None
+        self.latest_error: str = ""
+        self.progress_lock = threading.Lock()
 
     @property
     def phobert_model(self) -> str:
@@ -149,6 +165,31 @@ def _artifact_exists(path: Path) -> bool:
     return path.exists()
 
 
+def _set_progress(
+    *,
+    job_id: str = "",
+    active: bool,
+    status: str,
+    message: str,
+    current_variant: str = "",
+    current_index: int = 0,
+    total: int = 0,
+    completed: int = 0,
+) -> None:
+    with state.progress_lock:
+        state.progress_state = {
+            "job_id": job_id,
+            "active": active,
+            "status": status,
+            "current_variant": current_variant,
+            "current_index": current_index,
+            "total": total,
+            "completed": completed,
+            "message": message,
+            "updated_at": time.time(),
+        }
+
+
 def _release_variant_cache(variant: str) -> None:
     if variant in {"A1", "A2"}:
         bundle = state.a_models.pop(variant, None)
@@ -927,6 +968,81 @@ async def predict_variant(variant: str, question: str, image: Image.Image) -> di
             _release_variant_cache(variant)
 
 
+async def _predict_models(
+    selected_models: list[str],
+    question: str,
+    pil_img: Image.Image,
+    job_id: str = "",
+) -> dict[str, Any]:
+    results = []
+    total = len(selected_models)
+    _set_progress(job_id=job_id, active=True, status="running", message="Starting comparison...", total=total, completed=0)
+    async with load_lock:
+        for index, variant in enumerate(selected_models, start=1):
+            _set_progress(
+                job_id=job_id,
+                active=True,
+                status="running",
+                message=f"Running {variant} ({index}/{total})",
+                current_variant=variant,
+                current_index=index,
+                total=total,
+                completed=index - 1,
+            )
+            result = await predict_variant(variant, question, pil_img)
+            results.append(result)
+            _set_progress(
+                job_id=job_id,
+                active=True,
+                status="running",
+                message=f"Finished {variant} ({index}/{total})",
+                current_variant=variant,
+                current_index=index,
+                total=total,
+                completed=index,
+            )
+
+    predictions = {item["variant"]: item["prediction"] for item in results if item.get("status") == "ok"}
+    summary = {
+        "majority_vote": majority_answer(list(predictions.values())) if predictions else "",
+        "success_count": sum(1 for item in results if item.get("status") == "ok"),
+        "error_count": sum(1 for item in results if item.get("status", "").startswith("error")),
+    }
+    payload = {
+        "question": question,
+        "selected_models": selected_models,
+        "results": results,
+        "summary": summary,
+    }
+    _set_progress(
+        job_id=job_id,
+        active=False,
+        status="done",
+        message=f"Finished {total}/{total} models.",
+        total=total,
+        completed=total,
+    )
+    return payload
+
+
+def _run_predict_job(job_id: str, selected_models: list[str], question: str, image_bytes: bytes) -> None:
+    try:
+        pil_img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        payload = asyncio.run(_predict_models(selected_models, question, pil_img, job_id=job_id))
+        with state.progress_lock:
+            state.latest_result = {"job_id": job_id, "payload": payload, "status": "done"}
+            state.latest_error = ""
+    except Exception as exc:
+        with state.progress_lock:
+            state.latest_result = None
+            state.latest_error = str(exc)
+        _set_progress(job_id=job_id, active=False, status="error", message=f"Failed: {exc}")
+    finally:
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+
 def _parse_model_selection(raw_model_name: Optional[str], raw_model_names: Optional[str]) -> list[str]:
     if raw_model_names:
         try:
@@ -1009,28 +1125,65 @@ async def predict(
         raise HTTPException(status_code=400, detail=f"Failed to read image file: {exc}") from exc
 
     selected_models = _parse_model_selection(model_name, model_names)
-    results = []
-    async with load_lock:
-        for variant in selected_models:
-            if state.release_after_predict:
-                _release_variant_cache(variant)
-            results.append(await predict_variant(variant, question, pil_img))
+    payload = await _predict_models(selected_models, question, pil_img)
+    return JSONResponse(payload)
 
-    predictions = {item["variant"]: item["prediction"] for item in results if item.get("status") == "ok"}
-    summary = {
-        "majority_vote": majority_answer(list(predictions.values())) if predictions else "",
-        "success_count": sum(1 for item in results if item.get("status") == "ok"),
-        "error_count": sum(1 for item in results if item.get("status", "").startswith("error")),
-    }
 
-    return JSONResponse(
-        {
-            "question": question,
-            "selected_models": selected_models,
-            "results": results,
-            "summary": summary,
+@app.post("/v1/predict-job")
+async def predict_job(
+    question: str = Form(..., description="Question for VQA"),
+    model_name: Optional[str] = Form(None, description="Legacy single model name"),
+    model_names: Optional[str] = Form(None, description="Comma-separated or JSON list of models"),
+    image: UploadFile = File(..., description="Image input (JPEG/PNG)"),
+) -> JSONResponse:
+    if not question.strip():
+        raise HTTPException(status_code=400, detail="Question is required.")
+
+    try:
+        img_bytes = await image.read()
+    except Exception as exc:
+        raise HTTPException(status_code=400, detail=f"Failed to read image file: {exc}") from exc
+
+    selected_models = _parse_model_selection(model_name, model_names)
+    job_id = uuid.uuid4().hex
+    with state.progress_lock:
+        state.latest_result = None
+        state.latest_error = ""
+        state.progress_state = {
+            "job_id": job_id,
+            "active": True,
+            "status": "queued",
+            "current_variant": "",
+            "current_index": 0,
+            "total": len(selected_models),
+            "completed": 0,
+            "message": "Queued for prediction...",
+            "updated_at": time.time(),
         }
+
+    thread = threading.Thread(
+        target=_run_predict_job,
+        args=(job_id, selected_models, question, img_bytes),
+        daemon=True,
     )
+    thread.start()
+
+    return JSONResponse({"job_id": job_id, "status": "queued", "selected_models": selected_models}, status_code=202)
+
+
+@app.get("/v1/progress")
+def predict_progress() -> JSONResponse:
+    return JSONResponse(state.progress_state)
+
+
+@app.get("/v1/result")
+def predict_result() -> JSONResponse:
+    with state.progress_lock:
+        if state.latest_result is not None:
+            return JSONResponse(state.latest_result)
+        if state.latest_error:
+            return JSONResponse({"status": "error", "error": state.latest_error}, status_code=500)
+        return JSONResponse({"status": "pending"}, status_code=202)
 
 
 @app.get("/v1/question-suggestions")

web/static/index.html

@@ -269,6 +269,16 @@ Reset
 </div>
 
 <div class="space-y-5 pt-2">
+<div class="space-y-2">
+  <div class="flex items-center justify-between text-[12px] uppercase tracking-[0.22em] text-china-gold font-bold">
+    <span>Backend Progress</span>
+    <span id="progress-label">Idle</span>
+  </div>
+  <div class="h-3 rounded-full bg-[#E7E1D6] overflow-hidden border border-china-gold/25">
+    <div id="progress-bar" class="h-full w-0 bg-gradient-to-r from-imperial-red via-china-gold to-gold-light transition-[width] duration-300 ease-out"></div>
+  </div>
+  <div id="progress-detail" class="text-[12px] italic font-serif text-ink-black/60">Waiting for a request.</div>
+</div>
 <div class="flex items-center gap-3">
 <span class="text-xs font-bold uppercase tracking-widest text-china-gold">Model set:</span>
 <div class="flex gap-2 overflow-x-auto pb-1 no-scrollbar">
@@ -383,11 +393,15 @@ Medical VQA web demo for six-model comparison.
     resetBtn: document.getElementById("reset-btn"),
     statusText: document.getElementById("status-text"),
     resultsGrid: document.getElementById("results-grid"),
+    progressBar: document.getElementById("progress-bar"),
+    progressLabel: document.getElementById("progress-label"),
+    progressDetail: document.getElementById("progress-detail"),
   };
 
   let currentImageFile = null;
   let selectedModels = new Set(MODEL_ORDER);
   let questionSuggestions = [];
+  let progressTimer = null;
 
   function escapeHtml(value) {
     return String(value ?? "")
@@ -405,6 +419,56 @@ Medical VQA web demo for six-model comparison.
     el.statusText.textContent = message;
   }
 
+  function setProgressUI(state) {
+    const total = Number(state?.total || 0);
+    const completed = Number(state?.completed || 0);
+    const pct = total > 0 ? Math.max(0, Math.min(100, Math.round((completed / total) * 100))) : 0;
+    el.progressBar.style.width = `${pct}%`;
+    el.progressLabel.textContent = state?.active ? (state?.status || "running").toUpperCase() : "IDLE";
+    el.progressDetail.textContent = state?.message || "Waiting for a request.";
+  }
+
+  async function refreshProgress() {
+    try {
+      const res = await fetch("/v1/progress", { cache: "no-store" });
+      if (!res.ok) return;
+      const data = await res.json();
+      setProgressUI(data);
+      if (!data?.active && progressTimer) {
+        clearInterval(progressTimer);
+        progressTimer = null;
+      }
+      return data;
+    } catch (err) {
+      // ignore polling noise
+    }
+    return null;
+  }
+
+  function startProgressPolling() {
+    if (progressTimer) return;
+    refreshProgress();
+    progressTimer = setInterval(refreshProgress, 750);
+  }
+
+  function stopProgressPolling() {
+    if (progressTimer) {
+      clearInterval(progressTimer);
+      progressTimer = null;
+    }
+    refreshProgress();
+  }
+
+  async function waitForJobCompletion() {
+    while (true) {
+      const data = await refreshProgress();
+      if (data?.status === "done" || data?.status === "error") {
+        return data;
+      }
+      await new Promise((resolve) => setTimeout(resolve, 750));
+    }
+  }
+
   function setPreview(file) {
     currentImageFile = file || null;
     if (!file) {
@@ -641,6 +705,7 @@ Medical VQA web demo for six-model comparison.
     setStatus("Running all selected models...");
     renderRunningModelGrid();
     applyTiltEffect(".tilt-card", 5);
+    startProgressPolling();
 
     try {
       const formData = new FormData();
@@ -648,19 +713,30 @@ Medical VQA web demo for six-model comparison.
       formData.append("model_names", JSON.stringify(Array.from(selectedModels)));
       formData.append("image", currentImageFile);
 
-      const res = await fetch("/v1/predict", { method: "POST", body: formData });
+      const res = await fetch("/v1/predict-job", { method: "POST", body: formData });
       const data = await res.json();
       if (!res.ok) {
         throw new Error(data?.detail || "Prediction failed");
       }
-      renderModelGrid(data.results || []);
+
+      setStatus(`Job queued: ${data.job_id}`);
+      await waitForJobCompletion();
+
+      const resultRes = await fetch("/v1/result", { cache: "no-store" });
+      const resultData = await resultRes.json();
+      if (!resultRes.ok) {
+        throw new Error(resultData?.error || "Prediction failed");
+      }
+
+      renderModelGrid(resultData?.payload?.results || []);
       applyTiltEffect(".tilt-card", 5);
-      setStatus(`Done. ${data.summary?.success_count ?? 0} models succeeded.`);
+      setStatus(`Done. ${resultData?.payload?.summary?.success_count ?? 0} models succeeded.`);
     } catch (err) {
       setStatus(err.message || "Prediction failed");
     } finally {
       el.runBtn.disabled = false;
       el.runBtn.querySelector("span").textContent = "Run Comparison";
+      stopProgressPolling();
     }
   });
 
@@ -669,6 +745,7 @@ Medical VQA web demo for six-model comparison.
   loadModels();
   loadQuestionSuggestions();
   renderModelGrid([], "", null);
+  refreshProgress();
   applyTiltEffect(".tilt-card", 5);
 </script>