Adjusting self ping for datfid master for not to sleep

main.py

@@ -11,6 +11,8 @@ UPSTREAM_URL = os.environ.get("hf_url", "").rstrip("/")             # url to acc
 HF_TOKEN     = os.environ.get("hf_token")                           # HF access token to access Private Space
 DEMO_FORWARD_URL   = os.getenv("DEMO_FORWARD_URL", "").rstrip("/")  # url to acces demo space 
 DATFID_DEMO_TOKEN  = os.getenv("DATFID_DEMO_TOKEN", "")             # token to access demo space
+# This Space's public URL (used to ping self while waiting so HF does not put this Space to sleep). Override with SELF_URL env if different.
+SELF_URL = os.getenv("SELF_URL", "https://datfid-org-datfid-master.hf.space").rstrip("/")
 
 if not HF_TOKEN:
     raise RuntimeError("Missing secret 'hf_token' in public Space.")
@@ -39,8 +41,9 @@ SDK_MAX_BODY_BYTES_extended = int(os.getenv("SDK_MAX_BODY_BYTES_extended", "1250
 
 # How long to wait for upstream (API) response; long runs may need 30+ min (1800+)
 UPSTREAM_TIMEOUT = float(os.getenv("UPSTREAM_TIMEOUT", "900"))  # 15 minutes
-# While waiting for a response, ping upstream every N seconds so the backend Space does not sleep
+# While waiting for a response: send secure_ping to upstream every N seconds; if SELF_URL set, also ping self every SELF_PING_INTERVAL so this Space stays awake
 PING_UPSTREAM_INTERVAL = float(os.getenv("PING_UPSTREAM_INTERVAL", "270"))
+SELF_PING_INTERVAL = float(os.getenv("SELF_PING_INTERVAL", "120"))
 
 # Early global body-size guard (runs before routes)
 @app.middleware("http")
@@ -79,25 +82,31 @@ 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."""
+async def _ping_upstream_loop(ping_url: str, headers: dict, interval: float, self_url: str = ""):
+    """Every `interval` seconds: if self_url set, GET self_url/keep-alive (keep this Space awake), then GET ping_url (backend). Stops when cancelled."""
     if interval <= 0:
         return
     while True:
         await asyncio.sleep(interval)
         try:
             async with httpx.AsyncClient(timeout=10.0) as client:
+                if self_url:
+                    try:
+                        await client.get(f"{self_url}/keep-alive")
+                    except Exception:
+                        pass
                 await client.get(ping_url, headers=headers)
         except asyncio.CancelledError:
             break
         except Exception:
-            pass  # ignore ping errors
+            pass
 
-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:
+def _start_ping_task(ping_url: str, headers: dict, self_url: str = ""):
+    """Start background ping task. If SELF_URL set, uses SELF_PING_INTERVAL and pings self + backend; else PING_UPSTREAM_INTERVAL, backend only."""
+    interval = (SELF_PING_INTERVAL if self_url and SELF_PING_INTERVAL > 0 else PING_UPSTREAM_INTERVAL) or PING_UPSTREAM_INTERVAL
+    if interval <= 0:
         return None
-    return asyncio.create_task(_ping_upstream_loop(ping_url, headers, PING_UPSTREAM_INTERVAL))
+    return asyncio.create_task(_ping_upstream_loop(ping_url, headers, interval, self_url))
 
 async def _cancel_ping_task(task: asyncio.Task | None):
     if task is None:
@@ -133,8 +142,8 @@ async def _forward(path: str, method: str = "GET", json_body=None, user_token: s
     if user_token:
         headers["X-API-Key"] = user_token
 
-    # While waiting for upstream, call backend secure_ping every PING_UPSTREAM_INTERVAL seconds
-    ping_url = f"{UPSTREAM_URL}secure-ping/"
+    # While waiting: ping backend (secure_ping) and, if SELF_URL set, ping self so this Space stays awake
+    ping_url = f"{UPSTREAM_URL}/secure-ping/"
     ping_headers = {"Authorization": f"Bearer {HF_TOKEN}"}
     if user_token:
         ping_headers["X-API-Key"] = user_token
@@ -149,12 +158,13 @@ async def _forward(path: str, method: str = "GET", json_body=None, user_token: s
 
     request_task = asyncio.create_task(do_request())
     interval = PING_UPSTREAM_INTERVAL if PING_UPSTREAM_INTERVAL > 0 else 0.0
+    wait_sec = (SELF_PING_INTERVAL if SELF_URL and SELF_PING_INTERVAL > 0 else interval) or interval
     try:
         while not request_task.done():
-            if interval <= 0:
+            if wait_sec <= 0:
                 await request_task
                 break
-            ping_sleep = asyncio.create_task(asyncio.sleep(interval))
+            ping_sleep = asyncio.create_task(asyncio.sleep(wait_sec))
             done, pending = await asyncio.wait(
                 {request_task, ping_sleep},
                 return_when=asyncio.FIRST_COMPLETED,
@@ -169,11 +179,16 @@ async def _forward(path: str, method: str = "GET", json_body=None, user_token: s
                         pass
             if request_task in done:
                 break
-            try:
-                async with httpx.AsyncClient(timeout=10.0) as client:
-                    await client.get(ping_url, headers=ping_headers)
-            except Exception:
-                pass
+            async with httpx.AsyncClient(timeout=10.0) as c:
+                if SELF_URL:
+                    try:
+                        await c.get(f"{SELF_URL}/keep-alive")
+                    except Exception:
+                        pass
+                try:
+                    await c.get(ping_url, headers=ping_headers)
+                except Exception:
+                    pass
         if not request_task.done():
             request_task.cancel()
             try:
@@ -216,7 +231,7 @@ async def _forward_stream(path: str, files=None, data=None, user_token: str | No
     ping_headers = {"Authorization": f"Bearer {HF_TOKEN}"}
     if user_token:
         ping_headers["X-API-Key"] = user_token
-    ping_task = _start_ping_task(f"{UPSTREAM_URL}secure-ping/", ping_headers)
+    ping_task = _start_ping_task(f"{UPSTREAM_URL}/secure-ping/", ping_headers, SELF_URL)
     timeout = httpx.Timeout(UPSTREAM_TIMEOUT)
     client = httpx.AsyncClient(timeout=timeout, follow_redirects=True)
 
@@ -280,7 +295,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 to upstream and return JSON. Sends secure_ping every PING_UPSTREAM_INTERVAL while waiting."""
+    """POST multipart to upstream and return JSON. While waiting: secure_ping to backend; if SELF_URL set, ping self."""
     url = f"{UPSTREAM_URL}{path}"
     headers = {
         "Authorization": f"Bearer {HF_TOKEN}",
@@ -288,7 +303,7 @@ async def _forward_multipart_json(path: str, files=None, data=None, user_token:
     }
     if user_token:
         headers["X-API-Key"] = user_token
-    ping_url = f"{UPSTREAM_URL}secure-ping/"
+    ping_url = f"{UPSTREAM_URL}/secure-ping/"
     ping_headers = {"Authorization": f"Bearer {HF_TOKEN}"}
     if user_token:
         ping_headers["X-API-Key"] = user_token
@@ -302,12 +317,13 @@ async def _forward_multipart_json(path: str, files=None, data=None, user_token:
 
     request_task = asyncio.create_task(do_request())
     interval = PING_UPSTREAM_INTERVAL if PING_UPSTREAM_INTERVAL > 0 else 0.0
+    wait_sec = (SELF_PING_INTERVAL if SELF_URL and SELF_PING_INTERVAL > 0 else interval) or interval
     try:
         while not request_task.done():
-            if interval <= 0:
+            if wait_sec <= 0:
                 await request_task
                 break
-            ping_sleep = asyncio.create_task(asyncio.sleep(interval))
+            ping_sleep = asyncio.create_task(asyncio.sleep(wait_sec))
             done, pending = await asyncio.wait(
                 {request_task, ping_sleep},
                 return_when=asyncio.FIRST_COMPLETED,
@@ -322,11 +338,16 @@ async def _forward_multipart_json(path: str, files=None, data=None, user_token:
                         pass
             if request_task in done:
                 break
-            try:
-                async with httpx.AsyncClient(timeout=10.0) as client:
-                    await client.get(ping_url, headers=ping_headers)
-            except Exception:
-                pass
+            async with httpx.AsyncClient(timeout=10.0) as c:
+                if SELF_URL:
+                    try:
+                        await c.get(f"{SELF_URL}/keep-alive")
+                    except Exception:
+                        pass
+                try:
+                    await c.get(ping_url, headers=ping_headers)
+                except Exception:
+                    pass
         if not request_task.done():
             request_task.cancel()
             try: