queue API requests serially

server.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
+import asyncio
 import json
 import subprocess
 import sys
-import threading
 import time
+from collections import deque
 from pathlib import Path
 from typing import Annotated
 
-from fastapi import FastAPI, File, Form, HTTPException, UploadFile
+from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
 from fastapi.responses import FileResponse, JSONResponse
 from PIL import Image, UnidentifiedImageError
 
@@ -27,15 +28,51 @@ from service_runtime import (
 ROOT_DIR = Path(__file__).resolve().parent
 TMP_DIR = ROOT_DIR / "tmp"
 JOB_DIR = TMP_DIR / "jobs"
-BUSY_RETRY_AFTER_SECONDS = 15
 EXPORT_TIMEOUT_SECONDS = 45
 
 
+class RequestQueue:
+    def __init__(self) -> None:
+        self._condition = asyncio.Condition()
+        self._waiting: deque[str] = deque()
+        self._active_job_id: str | None = None
+
+    async def acquire(self, job_id: str) -> int:
+        async with self._condition:
+            self._waiting.append(job_id)
+            queued_ahead = len(self._waiting) - 1
+            try:
+                while self._waiting[0] != job_id or self._active_job_id is not None:
+                    await self._condition.wait()
+                self._waiting.popleft()
+                self._active_job_id = job_id
+                return queued_ahead
+            except BaseException:
+                if job_id in self._waiting:
+                    self._waiting.remove(job_id)
+                self._condition.notify_all()
+                raise
+
+    async def release(self, job_id: str) -> None:
+        async with self._condition:
+            if self._active_job_id == job_id:
+                self._active_job_id = None
+            elif job_id in self._waiting:
+                self._waiting.remove(job_id)
+            self._condition.notify_all()
+
+    def snapshot(self) -> dict[str, object]:
+        return {
+            "active_job_id": self._active_job_id,
+            "queued_jobs": len(self._waiting),
+        }
+
+
 class ServiceState:
     def __init__(self) -> None:
         self.runtime = TrellisRuntime()
         self.job_store = JobStore(JOB_DIR)
-        self.lock = threading.Lock()
+        self.queue = RequestQueue()
 
 
 state = ServiceState()
@@ -68,19 +105,6 @@ def _open_image(contents: bytes) -> Image.Image:
         ) from error
 
 
-def _busy_response() -> JSONResponse:
-    return JSONResponse(
-        status_code=429,
-        headers={"Retry-After": str(BUSY_RETRY_AFTER_SECONDS)},
-        content={
-            "stage": "admission",
-            "error_code": "busy",
-            "retryable": True,
-            "message": "Another request is already running. Retry shortly.",
-        },
-    )
-
-
 def _run_export(job_dir: Path, request: ImageToGlbRequest) -> Path:
     output_path = job_dir / "result.glb"
     result_json = job_dir / "export_result.json"
@@ -146,38 +170,42 @@ def startup() -> None:
 
 
 @app.get("/healthz")
-def healthz() -> dict[str, object]:
+async def healthz() -> dict[str, object]:
+    queue_state = state.queue.snapshot()
     return {
         "ok": state.runtime.is_healthy,
         "runtime_healthy": state.runtime.is_healthy,
-        "busy": state.lock.locked(),
+        "busy": queue_state["active_job_id"] is not None,
+        "queued_jobs": queue_state["queued_jobs"],
         "reason": state.runtime.unhealthy_reason,
     }
 
 
 @app.get("/")
-def root() -> dict[str, object]:
+async def root() -> dict[str, object]:
+    queue_state = state.queue.snapshot()
     return {
         "service": "TRELLIS.2 API",
         "version": app.version,
         "endpoint": "/v1/image-to-glb",
         "healthy": state.runtime.is_healthy,
+        "busy": queue_state["active_job_id"] is not None,
+        "queued_jobs": queue_state["queued_jobs"],
     }
 
 
 @app.post("/v1/image-to-glb")
 async def image_to_glb(
+    http_request: Request,
     image: Annotated[UploadFile, File(...)],
-    request: str = Form("{}"),
+    request_payload: str = Form("{}", alias="request"),
 ):
-    if not state.lock.acquire(blocking=False):
-        return _busy_response()
-
     job_id: str | None = None
     job_dir: Path | None = None
+    queue_acquired = False
     try:
         started = time.perf_counter()
-        request_model = _parse_request_payload(request)
+        request_model = _parse_request_payload(request_payload)
         state.job_store.cleanup_expired()
         contents = await image.read()
         pil_image = _open_image(contents)
@@ -188,7 +216,27 @@ async def image_to_glb(
         input_suffix = Path(image.filename or "input.png").suffix or ".png"
         input_path = job_dir / f"input{input_suffix}"
         save_input_image(pil_image, input_path)
-        state.job_store.update_status(job_dir, "preprocessing")
+
+        queued_at = time.perf_counter()
+        state.job_store.update_status(job_dir, "queued")
+        queued_ahead = await state.queue.acquire(job_id)
+        queue_acquired = True
+        queue_wait_ms = round((time.perf_counter() - queued_at) * 1000, 2)
+        if await http_request.is_disconnected():
+            raise ServiceError(
+                stage="admission",
+                error_code="client_disconnected",
+                message="Client disconnected while waiting in queue",
+                retryable=False,
+                status_code=499,
+            )
+
+        state.job_store.update_status(
+            job_dir,
+            "preprocessing",
+            queue_wait_ms=queue_wait_ms,
+            queued_ahead=queued_ahead,
+        )
         preprocessed = state.runtime.preprocess(pil_image, request_model)
 
         state.job_store.update_status(job_dir, "generating")
@@ -212,6 +260,7 @@ async def image_to_glb(
                 "X-Job-Id": job_id,
                 "X-Request-Id": job_id,
                 "X-Duration-Ms": str(duration_ms),
+                "X-Queue-Wait-Ms": str(queue_wait_ms),
             },
         )
     except HTTPException:
@@ -260,4 +309,5 @@ async def image_to_glb(
             content=service_error.to_dict(job_id),
         )
     finally:
-        state.lock.release()
+        if queue_acquired and job_id is not None:
+            await state.queue.release(job_id)