Add optional Colab GPU analysis worker

.env.example

@@ -4,6 +4,8 @@ NODE_ENV=development
 FRONTEND_URL=http://localhost:8080
 BACKEND_URL=http://localhost:4000
 ANALYSIS_PROXY_SECRET=replace-with-a-random-secret-of-at-least-32-characters
+# Optional: use the same random secret in the Colab notebook to register its worker.
+# COLAB_WORKER_SECRET=replace-with-a-different-random-secret-of-at-least-32-characters
 SUPABASE_URL=https://your-project.supabase.co
 SUPABASE_ANON_KEY=your-publishable-or-anon-key
 SUPABASE_SERVICE_ROLE_KEY=your-service-role-or-secret-key

README.md

@@ -37,6 +37,17 @@ Set these Space variables/secrets from `.env.example`:
 
 Store tokens, the service-role key, and proxy secret as Space secrets. `HF_TOKEN` requires read access to Core and write access to the private storage dataset. Use a dedicated storage dataset for this application. `start.py` sets `ANALYSIS_API_URL` and `BACKEND_INTERNAL_URL` automatically; do not configure a separate public AI URL.
 
+## Optional Colab GPU worker
+
+Analysis automatically uses Colab for inference when its worker lease is healthy. If no worker is registered, Colab is on CPU, or its connection is lost, the backend falls back to its existing local analysis pipeline. Colab returns analysis reports only; the backend still prepares playback videos, writes the normal `sessions/{session_id}` dataset folder, and updates Supabase.
+
+1. Generate a separate random secret with at least 32 characters. Add it to the Hugging Face Space as `COLAB_WORKER_SECRET`. Do not reuse `ANALYSIS_PROXY_SECRET`.
+2. Open [`Colab Server/worker.ipynb`](../Colab%20Server/worker.ipynb) in Colab and add these Colab Secrets: `HF_TOKEN`, `STORAGE_REPO_ID`, `CORE_REPO_ID`, `SPORALIZE_BACKEND_URL`, and `SPORALIZE_COLAB_WORKER_SECRET`. `HF_TOKEN` needs read access to Core and read/write access to the private storage dataset; `STORAGE_REPO_ID` and `CORE_REPO_ID` must match the backend values. Set `CORE_REVISION` too if the backend pins one. The backend URL is the public Space origin, such as `https://shoraky-sporalize-api.hf.space`; the Colab secret must match `COLAB_WORKER_SECRET` exactly.
+3. Select a GPU runtime and run the notebook. It starts the worker and Cloudflare quick tunnel, waits for model initialization, then registers the generated URL automatically. Keep the notebook cell running while analyses should use Colab.
+4. The notebook refreshes the worker lease every 20 seconds. The backend treats it as offline after 75 seconds without a heartbeat. Stopping or disconnecting Colab returns new jobs to the normal backend pipeline.
+
+The worker accepts one analysis at a time. The backend stages inputs and receives reports through temporary folders in the private Hugging Face dataset; large files and report payloads do not cross the tunnel. A job already running on Colab can resume on the backend if the tunnel is lost; Colab does not write final session records or artifacts, so fallback does not create duplicate saved sessions. Temporary handoff folders are deleted when the job completes or falls back. The backend still retains its usual CPU processing requirements for fallback. No Supabase schema change is needed.
+
 For a local Docker check, build with `docker build -t sporalize-backend .`, then run `docker run --rm --env-file .env -e PORT=7860 -p 7860:7860 sporalize-backend`. The `.env` NODE_ENV setting controls whether production HTTPS checks apply.
 
 Node health: `GET /health`. AI health is internal at `http://127.0.0.1:8000/healthz` and becomes available after model initialization. Node health alone does not indicate model readiness. Video uploads go through Node directly to Python without Vercel. Validate hosting timeouts and representative large uploads before release.

app.py

@@ -67,6 +67,123 @@ def backend(path, user_id, method="GET", body=None, is_admin=False):
             time.sleep(2 ** attempt)
 
 
+def registered_colab_worker():
+    secret = os.environ.get("COLAB_WORKER_SECRET", "")
+    if len(secret) < 32:
+        return None
+    try:
+        response = httpx.get(
+            f'{os.environ["BACKEND_INTERNAL_URL"].rstrip("/")}/internal/colab-worker/current',
+            timeout=3,
+            headers={"x-sporalize-proxy-secret": os.environ["ANALYSIS_PROXY_SECRET"]},
+        )
+        response.raise_for_status()
+        registration = response.json()
+        if not registration.get("active") or not registration.get("url"):
+            return None
+        base_url = registration["url"].rstrip("/")
+        health = httpx.get(
+            f"{base_url}/worker/health", timeout=5,
+            headers={"x-sporalize-colab-secret": secret},
+        )
+        health.raise_for_status()
+        status = health.json()
+        return base_url if status.get("ready") and status.get("cuda") else None
+    except (httpx.HTTPError, ValueError, KeyError) as error:
+        logger.info("Colab GPU worker is not available: %s", error)
+        return None
+
+
+def run_pipeline_on_colab(worker_url, client_id, work, cameras, job):
+    """Use HF as a temporary handoff so large videos never cross the tunnel."""
+    secret = os.environ["COLAB_WORKER_SECRET"]
+    headers = {"x-sporalize-colab-secret": secret}
+    input_path = f"colab-staging/{client_id}"
+    result_path = f"colab-results/{client_id}/reports.json"
+    job.update(status="running", progress=5, phase="Preparing analysis", stage="analysis")
+
+    try:
+        hub().upload_folder(
+            repo_id=os.environ["STORAGE_REPO_ID"], repo_type="dataset", folder_path=str(work),
+            path_in_repo=input_path, allow_patterns=["actions.json", "calibration.npz", "videos/*"],
+            commit_message=f"Stage analysis inputs {client_id}", token=os.environ["HF_TOKEN"],
+        )
+        video_files = [{"order": order, "filename": Path(path).name} for order, path in cameras.items()]
+        response = httpx.post(
+            f"{worker_url}/worker/jobs",
+            json={"jobId": client_id, "sessionId": job["sessionId"], "inputPath": input_path, "videos": video_files},
+            headers=headers, timeout=20,
+        )
+        response.raise_for_status()
+
+        failures = 0
+        while True:
+            if job.get("cancel"):
+                try:
+                    httpx.post(f"{worker_url}/worker/jobs/{client_id}/cancel", headers=headers, timeout=5)
+                except httpx.HTTPError:
+                    pass
+                raise InterruptedError("Processing cancelled by user")
+
+            try:
+                status_response = httpx.get(f"{worker_url}/worker/jobs/{client_id}", headers=headers, timeout=10)
+                status_response.raise_for_status()
+                remote = status_response.json()
+                failures = 0
+            except (httpx.HTTPError, ValueError) as error:
+                failures += 1
+                if failures >= 4:
+                    logger.warning("Colab worker connection was lost; using the primary analysis server: %s", error)
+                    return None
+                time.sleep(3)
+                continue
+
+            status = remote.get("status")
+            if status == "completed":
+                try:
+                    downloaded = hf_hub_download(
+                        repo_id=os.environ["STORAGE_REPO_ID"], repo_type="dataset", filename=result_path,
+                        token=os.environ["HF_TOKEN"], force_download=True,
+                    )
+                    reports = json.loads(Path(downloaded).read_text(encoding="utf-8"))
+                    return reports if isinstance(reports, list) else None
+                except Exception as error:
+                    logger.warning("Colab result could not be read; using the primary analysis server: %s", error)
+                    return None
+            if status in {"failed", "missing"}:
+                logger.warning("Colab analysis failed; using the primary analysis server: %s", remote.get("error", status))
+                return None
+            if status == "cancelled":
+                raise InterruptedError("Processing cancelled by user")
+
+            try:
+                percentage = min(100.0, max(0.0, float(remote.get("progress", 0))))
+            except (TypeError, ValueError):
+                percentage = 0.0
+            job.update(
+                status="running", stage="analysis", progress=5 + 0.85 * percentage,
+                phase=remote.get("phase") or "Analyzing session",
+                step=remote.get("step", 0), total=remote.get("total", 0),
+                action=remote.get("action", 0), actionsTotal=remote.get("actionsTotal", 0),
+                updatedAt=int(time.time() * 1000),
+            )
+            time.sleep(2)
+    except InterruptedError:
+        raise
+    except Exception as error:
+        logger.warning("Colab handoff failed; using the primary analysis server: %s", error)
+        return None
+    finally:
+        for remote_path in (input_path, f"colab-results/{client_id}"):
+            try:
+                hub().delete_folder(
+                    repo_id=os.environ["STORAGE_REPO_ID"], repo_type="dataset", path_in_repo=remote_path,
+                    commit_message=f"Clean up analysis handoff {client_id}", token=os.environ["HF_TOKEN"],
+                )
+            except Exception as error:
+                logger.info("Temporary Colab handoff cleanup skipped for %s: %s", remote_path, error)
+
+
 @lru_cache(maxsize=1)
 def runtime():
     root = Path(__file__).resolve().parent.parent / "Core"
@@ -375,18 +492,24 @@ def run_job(client_id, work, cameras):
         if job.get("cancel"):
             raise InterruptedError()
         job.update(status="running", progress=5, phase="Analyzing", stage="analysis")
-        pipeline, weights = runtime()
-
-        def progress(action, total, frame, frames):
-            action_number = action + 1
-            frame_total = max(1, frames)
-            job.update(progress=5 + 85 * (action + (frame + 1) / frame_total) / max(total, 1),
-                       phase=f"Analyzing action {action_number}/{total}", step=frame + 1,
-                       total=frame_total, action=action_number, actionsTotal=total,
-                       updatedAt=int(time.time() * 1000))
-            return not job.get("cancel")
-
-        reports = pipeline(cameras, {**weights, "CALIBRATION_PATH": str(work / "calibration.npz"), "ACTIONS_PATH": str(work / "actions.json")}, {"YOLO_IMGSZ": YOLO_IMGSZ}, progress)
+        worker_url = registered_colab_worker()
+        reports = run_pipeline_on_colab(worker_url, client_id, work, cameras, job) if worker_url else None
+        if reports is None:
+            if job.get("cancel"):
+                raise InterruptedError()
+            job.update(status="running", progress=5, phase="Analyzing session", stage="analysis")
+            pipeline, weights = runtime()
+
+            def progress(action, total, frame, frames):
+                action_number = action + 1
+                frame_total = max(1, frames)
+                job.update(progress=5 + 85 * (action + (frame + 1) / frame_total) / max(total, 1),
+                           phase=f"Analyzing action {action_number}/{total}", step=frame + 1,
+                           total=frame_total, action=action_number, actionsTotal=total,
+                           updatedAt=int(time.time() * 1000))
+                return not job.get("cancel")
+
+            reports = pipeline(cameras, {**weights, "CALIBRATION_PATH": str(work / "calibration.npz"), "ACTIONS_PATH": str(work / "actions.json")}, {"YOLO_IMGSZ": YOLO_IMGSZ}, progress)
         if job.get("cancel"):
             raise InterruptedError()
         (work / "results.json").write_text(json.dumps(json_value(reports), allow_nan=False), encoding="utf-8")

src/colab-worker.ts

@@ -0,0 +1,61 @@
+import { timingSafeEqual } from "node:crypto";
+import { Router, type Request } from "express";
+import { z } from "zod";
+
+const registrationSchema = z.object({ active: z.boolean().default(true), url: z.string().url().optional() }).strict();
+const HEARTBEAT_TTL_MS = 75_000;
+
+let worker: { url: string; lastSeenAt: number } | null = null;
+
+function secretMatches(expected: string | undefined, received: string | undefined) {
+  if (!expected || !received) return false;
+  const expectedBytes = Buffer.from(expected);
+  const receivedBytes = Buffer.from(received);
+  return expectedBytes.length === receivedBytes.length && timingSafeEqual(expectedBytes, receivedBytes);
+}
+
+function validTunnelUrl(value: string) {
+  try {
+    const url = new URL(value);
+    return url.protocol === "https:" && url.hostname.endsWith(".trycloudflare.com") &&
+      !url.username && !url.password && url.pathname === "/" && !url.search && !url.hash;
+  } catch {
+    return false;
+  }
+}
+
+export function colabWorkerRouter(analysisSecret: string, workerSecret?: string) {
+  const router = Router();
+
+  // The Colab notebook refreshes this lease while its tunnel is alive.
+  router.post("/register", (req, res) => {
+    if (!secretMatches(workerSecret, req.header("x-sporalize-colab-secret"))) {
+      return res.status(401).json({ error: "Invalid worker registration secret." });
+    }
+    const parsed = registrationSchema.safeParse(req.body);
+    if (!parsed.success) return res.status(400).json({ error: "Invalid worker registration." });
+    if (!parsed.data.active) {
+      worker = null;
+      return res.json({ active: false });
+    }
+    if (!parsed.data.url || !validTunnelUrl(parsed.data.url)) {
+      return res.status(400).json({ error: "Worker URL must be an HTTPS Cloudflare quick tunnel URL." });
+    }
+    worker = { url: parsed.data.url.replace(/\/$/, ""), lastSeenAt: Date.now() };
+    return res.json({ active: true, expiresInSeconds: HEARTBEAT_TTL_MS / 1000 });
+  });
+
+  // Only the private Python service can ask Node for the current worker lease.
+  router.get("/current", (req: Request, res) => {
+    if (!secretMatches(analysisSecret, req.header("x-sporalize-proxy-secret"))) {
+      return res.status(401).json({ error: "Unauthorized." });
+    }
+    if (!worker || Date.now() - worker.lastSeenAt > HEARTBEAT_TTL_MS) {
+      worker = null;
+      return res.json({ active: false });
+    }
+    return res.json({ active: true, url: worker.url });
+  });
+
+  return router;
+}

src/config.ts

@@ -9,6 +9,7 @@ const envSchema = z.object({
   BACKEND_URL: z.string().url(),
   ANALYSIS_API_URL: z.string().url().default("http://127.0.0.1:8000"),
   ANALYSIS_PROXY_SECRET: z.string().min(32),
+  COLAB_WORKER_SECRET: z.string().optional(),
   SUPABASE_URL: z.string().url(),
   SUPABASE_ANON_KEY: z.string().min(1),
   SUPABASE_SERVICE_ROLE_KEY: z.string().min(1),
@@ -21,6 +22,10 @@ if (!parsed.success) {
 }
 
 export const env = parsed.data;
+if (env.COLAB_WORKER_SECRET && env.COLAB_WORKER_SECRET.length < 32) {
+  console.error("COLAB_WORKER_SECRET must contain at least 32 characters when configured");
+  process.exit(1);
+}
 export const isProduction = env.NODE_ENV === "production";
 // Production browsers reach auth through Cloudflare Pages, including OAuth.
 export const authPublicOrigin = new URL(isProduction ? env.FRONTEND_URL : env.BACKEND_URL).origin;

src/server.ts

@@ -13,6 +13,7 @@ import { authenticatedUser, clearPasswordRecoveryCookie, clearSessionCookies, ha
 import { hashPassword, oauthClient, supabaseAdmin, supabaseAuth, syncProfile, userHasPassword, userPayload } from "./supabase.js";
 import { registerProfileImageRoutes } from "./profile-images.js";
 import { registerAdminRoutes } from "./admin.js";
+import { colabWorkerRouter } from "./colab-worker.js";
 
 const app = express();
 app.disable("x-powered-by");
@@ -25,6 +26,7 @@ app.use(cors({
 }));
 app.use("/internal/analysis", analysisRouter(supabaseAdmin, env.ANALYSIS_PROXY_SECRET));
 app.use(express.json({ limit: "32kb", type: ["application/json", "text/plain"] }));
+app.use("/internal/colab-worker", colabWorkerRouter(env.ANALYSIS_PROXY_SECRET, env.COLAB_WORKER_SECRET));
 app.use(cookieParser());
 app.use((req, res, next) => {
   if (["GET", "HEAD", "OPTIONS"].includes(req.method)) return next();