add async job polling API

README.md

@@ -13,13 +13,21 @@ short_description: High-fidelity 3D Generation from images
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
-API endpoint: `POST /v1/image-to-glb`
+API endpoints:
+
+- `POST /v1/image-to-glb` - create an async GLB generation job
+- `GET /v1/jobs/{job_id}` - poll job status
+- `GET /v1/jobs/{job_id}/result` - download the GLB when ready
 
 Example:
 
 ```bash
-curl -X POST "http://localhost:7860/v1/image-to-glb" \
+job_json=$(curl -s -X POST "http://localhost:7860/v1/image-to-glb" \
   -F 'image=@input.png' \
-  -F 'request={"generation":{"resolution":"512"},"export":{"remesh":true}}' \
-  --output output.glb
+  -F 'request={"generation":{"resolution":"512"},"export":{"remesh":true}}')
+
+job_id=$(python3 -c 'import json,sys; print(json.load(sys.stdin)["job_id"])' <<<"$job_json")
+
+curl -s "http://localhost:7860/v1/jobs/$job_id"
+curl -L "http://localhost:7860/v1/jobs/$job_id/result" --output output.glb
 ```

job_store.py

@@ -35,6 +35,21 @@ class JobStore:
         self.write_json(job_dir / "request.json", request_payload)
         return job_id, job_dir
 
+    def job_dir(self, job_id: str) -> Path:
+        return self.root_dir / job_id
+
+    def exists(self, job_id: str) -> bool:
+        return self.job_dir(job_id).is_dir()
+
+    def read_meta(self, job_id: str) -> dict[str, Any] | None:
+        return self.read_json(self.job_dir(job_id) / "meta.json")
+
+    def read_request(self, job_id: str) -> dict[str, Any] | None:
+        return self.read_json(self.job_dir(job_id) / "request.json")
+
+    def read_error(self, job_id: str) -> dict[str, Any] | None:
+        return self.read_json(self.job_dir(job_id) / "error.json")
+
     def update_status(self, job_dir: Path, status: str, **extra: Any) -> None:
         meta_path = job_dir / "meta.json"
         meta = self.read_json(meta_path, default={})

server.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import json
 import subprocess
 import sys
 import time
 from collections import deque
 from pathlib import Path
-from typing import Annotated
+from typing import Annotated, Any
 
-from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
+from fastapi import FastAPI, File, Form, HTTPException, UploadFile
 from fastapi.responses import FileResponse, JSONResponse
 from PIL import Image, UnidentifiedImageError
 
@@ -31,29 +32,37 @@ JOB_DIR = TMP_DIR / "jobs"
 EXPORT_TIMEOUT_SECONDS = 45
 
 
-class RequestQueue:
+class JobQueue:
     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:
+    def _queue_position_unlocked(self, job_id: str) -> int | None:
+        if self._active_job_id == job_id:
+            return 0
+        try:
+            index = self._waiting.index(job_id)
+        except ValueError:
+            return None
+        return index + (1 if self._active_job_id is not None else 0)
+
+    async def enqueue(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:
+            self._condition.notify_all()
+            queue_position = self._queue_position_unlocked(job_id)
+            return 0 if queue_position is None else queue_position
+
+    async def claim_next(self) -> str:
+        async with self._condition:
+            while not self._waiting:
+                await self._condition.wait()
+            job_id = self._waiting.popleft()
+            self._active_job_id = job_id
+            return job_id
+
+    async def complete(self, job_id: str) -> None:
         async with self._condition:
             if self._active_job_id == job_id:
                 self._active_job_id = None
@@ -61,10 +70,14 @@ class RequestQueue:
                 self._waiting.remove(job_id)
             self._condition.notify_all()
 
-    def snapshot(self) -> dict[str, object]:
+    def snapshot(self, job_id: str | None = None) -> dict[str, Any]:
+        queue_position = None
+        if job_id is not None:
+            queue_position = self._queue_position_unlocked(job_id)
         return {
             "active_job_id": self._active_job_id,
             "queued_jobs": len(self._waiting),
+            "queue_position": queue_position,
         }
 
 
@@ -72,11 +85,12 @@ class ServiceState:
     def __init__(self) -> None:
         self.runtime = TrellisRuntime()
         self.job_store = JobStore(JOB_DIR)
-        self.queue = RequestQueue()
+        self.queue = JobQueue()
+        self.worker_task: asyncio.Task[None] | None = None
 
 
 state = ServiceState()
-app = FastAPI(title="TRELLIS.2 API", version="1.0.0")
+app = FastAPI(title="TRELLIS.2 API", version="1.1.0")
 
 
 def _parse_request_payload(payload: str) -> ImageToGlbRequest:
@@ -105,6 +119,12 @@ def _open_image(contents: bytes) -> Image.Image:
         ) from error
 
 
+def _open_image_path(path: Path) -> Image.Image:
+    with Image.open(path) as image:
+        image.load()
+        return image.copy()
+
+
 def _run_export(job_dir: Path, request: ImageToGlbRequest) -> Path:
     output_path = job_dir / "result.glb"
     result_json = job_dir / "export_result.json"
@@ -162,111 +182,130 @@ def _run_export(job_dir: Path, request: ImageToGlbRequest) -> Path:
     )
 
 
-@app.on_event("startup")
-def startup() -> None:
-    TMP_DIR.mkdir(parents=True, exist_ok=True)
-    state.job_store.cleanup_expired()
-    state.runtime.load()
-
-
-@app.get("/healthz")
-async def healthz() -> dict[str, object]:
-    queue_state = state.queue.snapshot()
+def _job_urls(job_id: str) -> dict[str, str]:
     return {
-        "ok": state.runtime.is_healthy,
-        "runtime_healthy": state.runtime.is_healthy,
-        "busy": queue_state["active_job_id"] is not None,
-        "queued_jobs": queue_state["queued_jobs"],
-        "reason": state.runtime.unhealthy_reason,
+        "status_url": f"/v1/jobs/{job_id}",
+        "result_url": f"/v1/jobs/{job_id}/result",
     }
 
 
-@app.get("/")
-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,
+def _build_job_response(job_id: str) -> dict[str, Any]:
+    if not state.job_store.exists(job_id):
+        raise HTTPException(status_code=404, detail="Job not found")
+
+    meta = state.job_store.read_meta(job_id) or {}
+    failure = state.job_store.read_error(job_id)
+    queue_state = state.queue.snapshot(job_id)
+    response = {
+        "job_id": job_id,
+        "status": meta.get("status", "unknown"),
+        "created_at": meta.get("created_at"),
+        "updated_at": meta.get("updated_at"),
+        "input_filename": meta.get("input_filename"),
+        "runtime_healthy": state.runtime.is_healthy,
         "queued_jobs": queue_state["queued_jobs"],
+        "active_job_id": queue_state["active_job_id"],
+        **_job_urls(job_id),
     }
-
-
-@app.post("/v1/image-to-glb")
-async def image_to_glb(
-    http_request: Request,
-    image: Annotated[UploadFile, File(...)],
-    request_payload: str = Form("{}", alias="request"),
-):
-    job_id: str | None = None
-    job_dir: Path | None = None
-    queue_acquired = False
-    try:
-        started = time.perf_counter()
-        request_model = _parse_request_payload(request_payload)
-        state.job_store.cleanup_expired()
-        contents = await image.read()
-        pil_image = _open_image(contents)
-        job_id, job_dir = state.job_store.create(
-            request_model.model_dump(mode="json"), image.filename
+    queue_position = queue_state.get("queue_position")
+    if queue_position is not None:
+        response["queue_position"] = queue_position
+
+    for field in (
+        "queue_wait_ms",
+        "preprocess_ms",
+        "generate_ms",
+        "export_ms",
+        "duration_ms",
+        "queued_at",
+        "started_at",
+        "completed_at",
+        "output_path",
+    ):
+        if field in meta:
+            response[field] = meta[field]
+
+    if response["status"] == "succeeded" and meta.get("output_path"):
+        response["result_ready"] = Path(meta["output_path"]).exists()
+
+    if failure is not None:
+        response["failure"] = failure
+
+    return response
+
+
+def _process_job_sync(job_id: str) -> None:
+    job_dir = state.job_store.job_dir(job_id)
+    meta = state.job_store.read_meta(job_id) or {}
+    request_payload = state.job_store.read_request(job_id)
+    if request_payload is None:
+        state.job_store.record_failure(
+            job_dir,
+            stage="admission",
+            error_code="missing_request",
+            message="Job request payload is missing",
+            retryable=False,
         )
+        return
 
-        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)
-
-        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,
-            )
+    request_model = ImageToGlbRequest.model_validate(request_payload)
+    input_path_str = meta.get("input_path")
+    if not input_path_str:
+        state.job_store.record_failure(
+            job_dir,
+            stage="admission",
+            error_code="missing_input",
+            message="Uploaded image is missing",
+            retryable=False,
+        )
+        return
 
+    input_path = Path(input_path_str)
+    started = time.perf_counter()
+    queue_wait_ms = round(
+        (time.time() - float(meta.get("enqueued_at_ts", time.time()))) * 1000, 2
+    )
+    preprocess_started = time.perf_counter()
+    try:
         state.job_store.update_status(
             job_dir,
             "preprocessing",
+            started_at=time.time(),
             queue_wait_ms=queue_wait_ms,
-            queued_ahead=queued_ahead,
         )
-        preprocessed = state.runtime.preprocess(pil_image, request_model)
+        image = _open_image_path(input_path)
+        preprocessed = state.runtime.preprocess(image, request_model)
+        preprocess_ms = round((time.perf_counter() - preprocess_started) * 1000, 2)
 
-        state.job_store.update_status(job_dir, "generating")
+        generate_started = time.perf_counter()
+        state.job_store.update_status(
+            job_dir, "generating", preprocess_ms=preprocess_ms
+        )
         payload = state.runtime.generate_export_payload(preprocessed, request_model)
         save_export_payload(job_dir, payload)
+        generate_ms = round((time.perf_counter() - generate_started) * 1000, 2)
 
-        state.job_store.update_status(job_dir, "exporting")
+        export_started = time.perf_counter()
+        state.job_store.update_status(
+            job_dir,
+            "exporting",
+            preprocess_ms=preprocess_ms,
+            generate_ms=generate_ms,
+        )
         output_path = _run_export(job_dir, request_model)
+        export_ms = round((time.perf_counter() - export_started) * 1000, 2)
         duration_ms = round((time.perf_counter() - started) * 1000, 2)
         state.job_store.update_status(
             job_dir,
             "succeeded",
+            preprocess_ms=preprocess_ms,
+            generate_ms=generate_ms,
+            export_ms=export_ms,
             duration_ms=duration_ms,
+            completed_at=time.time(),
             output_path=str(output_path),
         )
-        return FileResponse(
-            path=output_path,
-            media_type="model/gltf-binary",
-            filename=f"{job_id}.glb",
-            headers={
-                "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:
-        raise
     except subprocess.TimeoutExpired as error:
-        assert job_id is not None and job_dir is not None
         state.job_store.record_failure(
             job_dir,
             stage="export",
@@ -275,24 +314,7 @@ async def image_to_glb(
             retryable=True,
             details={"timeout_seconds": error.timeout},
         )
-        return JSONResponse(
-            status_code=504,
-            content={
-                "job_id": job_id,
-                "stage": "export",
-                "error_code": "timeout",
-                "retryable": True,
-                "message": "GLB export timed out",
-                "details": {"timeout_seconds": error.timeout},
-            },
-        )
     except Exception as error:
-        if job_id is None or job_dir is None:
-            if isinstance(error, ServiceError):
-                raise HTTPException(
-                    status_code=error.status_code, detail=error.message
-                ) from error
-            raise
         service_error = classify_runtime_error("export", error)
         if isinstance(error, ServiceError):
             service_error = error
@@ -304,10 +326,115 @@ async def image_to_glb(
             retryable=service_error.retryable,
             details=service_error.details,
         )
-        return JSONResponse(
-            status_code=service_error.status_code,
-            content=service_error.to_dict(job_id),
+
+
+async def _worker_loop() -> None:
+    while True:
+        job_id = await state.queue.claim_next()
+        try:
+            await asyncio.to_thread(_process_job_sync, job_id)
+        finally:
+            await state.queue.complete(job_id)
+
+
+@app.on_event("startup")
+async def startup() -> None:
+    TMP_DIR.mkdir(parents=True, exist_ok=True)
+    state.job_store.cleanup_expired()
+    state.runtime.load()
+    if state.worker_task is None:
+        state.worker_task = asyncio.create_task(_worker_loop())
+
+
+@app.on_event("shutdown")
+async def shutdown() -> None:
+    if state.worker_task is None:
+        return
+    state.worker_task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await state.worker_task
+    state.worker_task = None
+
+
+@app.get("/healthz")
+async def healthz() -> dict[str, object]:
+    queue_state = state.queue.snapshot()
+    return {
+        "ok": state.runtime.is_healthy,
+        "runtime_healthy": state.runtime.is_healthy,
+        "busy": queue_state["active_job_id"] is not None,
+        "queued_jobs": queue_state["queued_jobs"],
+        "active_job_id": queue_state["active_job_id"],
+        "reason": state.runtime.unhealthy_reason,
+    }
+
+
+@app.get("/")
+async def root() -> dict[str, object]:
+    queue_state = state.queue.snapshot()
+    return {
+        "service": "TRELLIS.2 API",
+        "version": app.version,
+        "submit_endpoint": "/v1/image-to-glb",
+        "status_endpoint": "/v1/jobs/{job_id}",
+        "result_endpoint": "/v1/jobs/{job_id}/result",
+        "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(
+    image: Annotated[UploadFile, File(...)],
+    request_payload: str = Form("{}", alias="request"),
+):
+    request_model = _parse_request_payload(request_payload)
+    state.job_store.cleanup_expired()
+    contents = await image.read()
+    pil_image = _open_image(contents)
+
+    job_id, job_dir = state.job_store.create(
+        request_model.model_dump(mode="json"), image.filename
+    )
+    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)
+
+    queued_at = time.time()
+    queue_position = await state.queue.enqueue(job_id)
+    state.job_store.update_status(
+        job_dir,
+        "queued",
+        enqueued_at=queued_at,
+        enqueued_at_ts=queued_at,
+        input_path=str(input_path),
+        initial_queue_position=queue_position,
+    )
+
+    response = _build_job_response(job_id)
+    return JSONResponse(status_code=202, content=response)
+
+
+@app.get("/v1/jobs/{job_id}")
+async def get_job(job_id: str) -> dict[str, Any]:
+    return _build_job_response(job_id)
+
+
+@app.get("/v1/jobs/{job_id}/result")
+async def get_job_result(job_id: str):
+    response = _build_job_response(job_id)
+    status = response["status"]
+    if status == "succeeded":
+        output_path = response.get("output_path")
+        if not output_path or not Path(output_path).exists():
+            raise HTTPException(status_code=404, detail="Job result is missing")
+        return FileResponse(
+            path=output_path,
+            media_type="model/gltf-binary",
+            filename=f"{job_id}.glb",
+            headers={"X-Job-Id": job_id},
         )
-    finally:
-        if queue_acquired and job_id is not None:
-            await state.queue.release(job_id)
+    if status == "failed":
+        return JSONResponse(status_code=409, content=response)
+    return JSONResponse(status_code=202, content=response)