Added ping every 270 seconds to not fall asleep if the core space is still counting

main.py

@@ -1,4 +1,4 @@
-import os, httpx
+import os, asyncio, httpx
 from fastapi import FastAPI, Request, HTTPException, UploadFile, File, Form
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
@@ -37,6 +37,11 @@ async def ratelimit_handler(request: Request, exc: RateLimitExceeded):
 SDK_MAX_BODY_BYTES = int(os.getenv("SDK_MAX_BODY_BYTES", "25000000"))  # 25MB default for demo
 SDK_MAX_BODY_BYTES_extended = int(os.getenv("SDK_MAX_BODY_BYTES_extended", "125000000")) # 125MB default for prod 
 
+# How long to wait for upstream (API) response; long runs may need 30+ min (1800+)
+UPSTREAM_TIMEOUT = float(os.getenv("UPSTREAM_TIMEOUT", "1800"))
+# While waiting for a response, ping upstream every N seconds so the backend Space does not sleep
+PING_UPSTREAM_INTERVAL = float(os.getenv("PING_UPSTREAM_INTERVAL", "270"))
+
 # Early global body-size guard (runs before routes)
 @app.middleware("http")
 async def limit_body_size(request: Request, call_next):
@@ -74,6 +79,35 @@ def _filter_resp_headers(headers: dict) -> dict:
     allowed = {"content-disposition"}
     return {k: v for k, v in headers.items() if k.lower() in allowed}
 
+async def _ping_upstream_loop(ping_url: str, headers: dict, interval: float):
+    """Background task: every `interval` seconds, GET ping_url to keep the backend Space awake. Stops when cancelled."""
+    if interval <= 0:
+        return
+    while True:
+        await asyncio.sleep(interval)
+        try:
+            async with httpx.AsyncClient(timeout=10.0) as client:
+                await client.get(ping_url, headers=headers)
+        except asyncio.CancelledError:
+            break
+        except Exception:
+            pass  # ignore ping errors
+
+def _start_ping_task(ping_url: str, headers: dict):
+    """Start a background task that pings the given URL every PING_UPSTREAM_INTERVAL seconds. Returns the task (or None); cancel it when the main request finishes."""
+    if PING_UPSTREAM_INTERVAL <= 0:
+        return None
+    return asyncio.create_task(_ping_upstream_loop(ping_url, headers, PING_UPSTREAM_INTERVAL))
+
+async def _cancel_ping_task(task: asyncio.Task | None):
+    if task is None:
+        return
+    task.cancel()
+    try:
+        await task
+    except asyncio.CancelledError:
+        pass
+
 def _extract_user_token(req: Request) -> str | None:
     """
     Read user's DATFID token from Authorization: Bearer <dt+...>.
@@ -89,6 +123,7 @@ async def _forward(path: str, method: str = "GET", json_body=None, user_token: s
     Forward request to the PRIVATE Space:
       - 'Authorization: Bearer <HF_TOKEN>' to pass HF private gate
       - 'X-API-Key: <dt+...>' so your private app can validate the user token
+    While waiting, pings upstream every PING_UPSTREAM_INTERVAL seconds so the backend Space does not sleep.
     """
     url = f"{UPSTREAM_URL}{path}"
     headers = {
@@ -98,9 +133,13 @@ async def _forward(path: str, method: str = "GET", json_body=None, user_token: s
     if user_token:
         headers["X-API-Key"] = user_token
 
-    timeout = httpx.Timeout(600.0)
-    async with httpx.AsyncClient(timeout=timeout) as client:
-        r = await client.request(method, url, headers=headers, json=json_body)
+    ping_task = _start_ping_task(f"{UPSTREAM_URL}/", {"Authorization": f"Bearer {HF_TOKEN}"})
+    try:
+        timeout = httpx.Timeout(UPSTREAM_TIMEOUT)
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            r = await client.request(method, url, headers=headers, json=json_body)
+    finally:
+        await _cancel_ping_task(ping_task)
 
     ct = r.headers.get("content-type", "")
 
@@ -121,8 +160,8 @@ async def _forward_stream(path: str, files=None, data=None, user_token: str | No
     if user_token:
         headers["X-API-Key"] = user_token
 
-    timeout = httpx.Timeout(600.0)
-
+    ping_task = _start_ping_task(f"{UPSTREAM_URL}/", {"Authorization": f"Bearer {HF_TOKEN}"})
+    timeout = httpx.Timeout(UPSTREAM_TIMEOUT)
     client = httpx.AsyncClient(timeout=timeout, follow_redirects=True)
 
     # Don't open the context yet; the iterator must own the context lifetime.
@@ -150,6 +189,7 @@ async def _forward_stream(path: str, files=None, data=None, user_token: str | No
                 async for chunk in resp.aiter_raw():
                     yield chunk
         finally:
+            await _cancel_ping_task(ping_task)
             await client.aclose()
 
     response = StreamingResponse(
@@ -162,7 +202,7 @@ async def _forward_stream(path: str, files=None, data=None, user_token: str | No
 
 
 async def _forward_multipart_json(path: str, files=None, data=None, user_token: str | None = None, method: str = "POST"):
-    """POST multipart (files + data) to upstream and return JSON response."""
+    """POST multipart (files + data) to upstream and return JSON response. Pings upstream every PING_UPSTREAM_INTERVAL while waiting."""
     url = f"{UPSTREAM_URL}{path}"
     headers = {
         "Authorization": f"Bearer {HF_TOKEN}",
@@ -170,9 +210,13 @@ async def _forward_multipart_json(path: str, files=None, data=None, user_token:
     }
     if user_token:
         headers["X-API-Key"] = user_token
-    timeout = httpx.Timeout(600.0)
-    async with httpx.AsyncClient(timeout=timeout) as client:
-        r = await client.request(method, url, headers=headers, files=files, data=data)
+    ping_task = _start_ping_task(f"{UPSTREAM_URL}/", {"Authorization": f"Bearer {HF_TOKEN}"})
+    try:
+        timeout = httpx.Timeout(UPSTREAM_TIMEOUT)
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            r = await client.request(method, url, headers=headers, files=files, data=data)
+    finally:
+        await _cancel_ping_task(ping_task)
     ct = r.headers.get("content-type", "")
     if "application/json" in ct:
         try:
@@ -195,6 +239,8 @@ async def _forward_demo_stream(path: str, *, files: dict | None, data: dict | No
         "X-DATFID-Token": DATFID_DEMO_TOKEN,
         "Accept": "*/*",
     }
+    ping_headers = {"Authorization": f"Bearer {HF_TOKEN}", "X-DATFID-Token": DATFID_DEMO_TOKEN}
+    ping_task = _start_ping_task(f"{DEMO_FORWARD_URL.rstrip('/')}/", ping_headers)
 
     timeout = httpx.Timeout(120.0)
     client = httpx.AsyncClient(timeout=timeout, follow_redirects=True)
@@ -219,6 +265,7 @@ async def _forward_demo_stream(path: str, *, files: dict | None, data: dict | No
                 async for chunk in resp.aiter_raw():
                     yield chunk
         finally:
+            await _cancel_ping_task(ping_task)
             await client.aclose()
 
     response = StreamingResponse(
@@ -235,6 +282,23 @@ async def root(req: Request):
     user_token = _extract_user_token(req)  # optional here
     return await _forward("/", "GET", user_token=user_token)
 
+
+@app.get("/keep-alive")
+async def keep_alive():
+    """
+    Hit this from an external cron (e.g. every 4 min) to keep this Space and the backend awake.
+    Sleep timeout is configured on Hugging Face (Space settings), not in code; pinging more
+    often than that (e.g. every 4 min if HF sleep is 5 min) prevents both spaces from sleeping.
+    Pings the upstream API so both see activity. No auth; use UptimeRobot, cron-job.org, etc.
+    """
+    try:
+        url = f"{UPSTREAM_URL}/"
+        async with httpx.AsyncClient(timeout=10.0) as client:
+            r = await client.get(url, headers={"Authorization": f"Bearer {HF_TOKEN}"})
+        return JSONResponse(content={"ok": True, "upstream_status": r.status_code})
+    except Exception as e:
+        return JSONResponse(status_code=502, content={"ok": False, "error": str(e)[:200]})
+
 @app.get("/secure-ping/")
 async def secure_ping(req: Request):
     # Require user's DATFID token in Authorization: Bearer <dt+...>