"""Throwaway ZeroGPU transport spike. Validates whether custom `/v1` HTTP routes, custom SSE, DELETE cancellation, custom `/v1` `@spaces.GPU` page routes, and page-scoped Gradio `@spaces.GPU` API functions can coexist in one HF Space. This app is deliberately small and model-free: the GPU function sleeps for a short duration to exercise ZeroGPU queue allocation without requiring weights. """ from __future__ import annotations import asyncio import json import tempfile import time import uuid from dataclasses import dataclass, field from pathlib import Path from typing import Any import spaces from gradio import Server from pydantic import BaseModel, Field from starlette.responses import JSONResponse, StreamingResponse app = Server() class CreateJobRequest(BaseModel): pages: int = Field(default=1, ge=1, le=10) @dataclass class Job: job_id: str total_pages: int status: str = "created" completed_pages: set[int] = field(default_factory=set) cancelled: bool = False events: list[dict[str, Any]] = field(default_factory=list) next_sequence: int = 1 def emit(self, event: str, data: dict[str, Any] | None = None) -> dict[str, Any]: envelope = { "event": event, "sequence": self.next_sequence, "job_id": self.job_id, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "data": data or {}, } self.next_sequence += 1 self.events.append(envelope) return envelope JOBS: dict[str, Job] = {} STATE_DIR = Path(tempfile.gettempdir()) / "zerogpu_transport_spike_state" STATE_DIR.mkdir(parents=True, exist_ok=True) def _state_path(job_id: str) -> Path: return STATE_DIR / f"{job_id}.json" def _serialize_job(job: Job) -> dict[str, Any]: return { "job_id": job.job_id, "total_pages": job.total_pages, "status": job.status, "completed_pages": sorted(job.completed_pages), "cancelled": job.cancelled, "events": job.events, "next_sequence": job.next_sequence, } def _deserialize_job(data: dict[str, Any]) -> Job: return Job( job_id=data["job_id"], total_pages=int(data["total_pages"]), status=data.get("status", "created"), completed_pages=set(data.get("completed_pages", [])), cancelled=bool(data.get("cancelled", False)), events=list(data.get("events", [])), next_sequence=int(data.get("next_sequence", 1)), ) def _save_job(job: Job) -> None: tmp = _state_path(job.job_id).with_suffix(".tmp") tmp.write_text(json.dumps(_serialize_job(job), separators=(",", ":")), encoding="utf-8") tmp.replace(_state_path(job.job_id)) def _get_job(job_id: str) -> Job | None: path = _state_path(job_id) if path.exists(): job = _deserialize_job(json.loads(path.read_text(encoding="utf-8"))) JOBS[job_id] = job return job return JOBS.get(job_id) def _job_json(job: Job) -> dict[str, Any]: return { "job_id": job.job_id, "status": job.status, "total_pages": job.total_pages, "completed_pages": sorted(job.completed_pages), "cancelled": job.cancelled, "event_count": len(job.events), } @app.get("/") async def root(): return {"name": "zerogpu-transport-spike", "ok": True} @app.get("/v1/health/live") async def live(): return {"ok": True, "kind": "live"} @app.get("/v1/health/ready") async def ready(): return {"ok": True, "kind": "ready"} @app.post("/v1/jobs") async def create_job(request: CreateJobRequest): job_id = f"job_{uuid.uuid4().hex[:16]}" job = Job(job_id=job_id, total_pages=request.pages) job.emit("job.created", {"total_pages": request.pages}) JOBS[job_id] = job _save_job(job) return JSONResponse(_job_json(job), status_code=201) @app.get("/v1/jobs/{job_id}") async def get_job(job_id: str): job = _get_job(job_id) if not job: return JSONResponse({"error": {"code": "JOB_NOT_FOUND"}}, status_code=404) return _job_json(job) @app.delete("/v1/jobs/{job_id}") async def cancel_job(job_id: str): job = _get_job(job_id) if not job: return JSONResponse({"error": {"code": "JOB_NOT_FOUND"}}, status_code=404) if job.status not in {"completed", "cancelled"}: job.status = "cancelled" job.cancelled = True job.emit("job.cancelled", {}) _save_job(job) return _job_json(job) @app.get("/v1/jobs/{job_id}/events") async def job_events(job_id: str): job = _get_job(job_id) if not job: return JSONResponse({"error": {"code": "JOB_NOT_FOUND"}}, status_code=404) async def stream(): cursor = 0 while True: refreshed = _get_job(job_id) if refreshed: job.events = refreshed.events job.status = refreshed.status while cursor < len(job.events): event = job.events[cursor] cursor += 1 yield f"id: {event['sequence']}\n" yield f"event: {event['event']}\n" yield f"data: {json.dumps(event)}\n\n" if job.status in {"completed", "cancelled"}: break await asyncio.sleep(0.25) return StreamingResponse(stream(), media_type="text/event-stream") def _run_page_impl(job_id: str, page_number: int, *, transport: str) -> dict[str, Any]: job = _get_job(job_id) if not job: return {"ok": False, "error": {"code": "JOB_NOT_FOUND"}} if job.cancelled: job.emit("page.failed", {"page_number": page_number, "code": "JOB_CANCELLED", "transport": transport}) _save_job(job) return {"ok": False, "error": {"code": "JOB_CANCELLED"}} if page_number < 1 or page_number > job.total_pages: return {"ok": False, "error": {"code": "PAGE_OUT_OF_RANGE"}} if job.status == "created": job.status = "processing_pages" job.emit("job.started", {"transport": transport}) job.emit("page.started", {"page_number": page_number, "transport": transport}) time.sleep(1.0) job.completed_pages.add(page_number) job.emit( "page.completed", { "page_number": page_number, "completed_pages": len(job.completed_pages), "total_pages": job.total_pages, "transport": transport, }, ) if len(job.completed_pages) >= job.total_pages: job.status = "completed" job.emit("job.completed", {"transport": transport}) _save_job(job) return { "ok": True, "job_id": job_id, "page_number": page_number, "completed_pages": sorted(job.completed_pages), "transport": transport, } @app.post("/v1/jobs/{job_id}/pages/{page_number}/run") @spaces.GPU(duration=5) def run_page_route(job_id: str, page_number: int): """Custom REST page-scoped GPU route. This matches the intended Marker `/v1` route shape. Passing this route is the primary proof that FastAPI/Gradio custom routes can themselves be ZeroGPU-decorated in the deployed Space. """ result = _run_page_impl(job_id, page_number, transport="rest") status = 200 if result.get("ok") else 400 return JSONResponse(result, status_code=status) @app.api() @spaces.GPU(duration=5) def run_page(job_id: str, page_number: int) -> dict[str, Any]: """Page-scoped Gradio GPU function. This fallback validates the alternative design if custom REST GPU routes do not work on ZeroGPU. """ return _run_page_impl(job_id, page_number, transport="gradio") app.launch(show_error=True)