Adjustment non-sleep of datfid master space while datfid api is still working

main.py

@@ -38,7 +38,7 @@ SDK_MAX_BODY_BYTES = int(os.getenv("SDK_MAX_BODY_BYTES", "25000000"))  # 25MB de
 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", "900")) # 15 minutes
+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
 PING_UPSTREAM_INTERVAL = float(os.getenv("PING_UPSTREAM_INTERVAL", "270"))
 
@@ -124,6 +124,7 @@ async def _forward(path: str, method: str = "GET", json_body=None, user_token: s
       - '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.
+    Uses explicit wait (response OR ping timer) so pings run even when the event loop doesn't preempt the request.
     """
     url = f"{UPSTREAM_URL}{path}"
     headers = {
@@ -133,13 +134,63 @@ async def _forward(path: str, method: str = "GET", json_body=None, user_token: s
     if user_token:
         headers["X-API-Key"] = user_token
 
-    ping_task = _start_ping_task(f"{UPSTREAM_URL}/", {"Authorization": f"Bearer {HF_TOKEN}"})
-    try:
-        timeout = httpx.Timeout(UPSTREAM_TIMEOUT)
+    ping_url = f"{UPSTREAM_URL}/"
+    ping_headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+    timeout = httpx.Timeout(UPSTREAM_TIMEOUT)
+    r = None
+
+    async def do_request():
+        nonlocal r
         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)
+
+    request_task = asyncio.create_task(do_request())
+    interval = PING_UPSTREAM_INTERVAL if PING_UPSTREAM_INTERVAL > 0 else 0.0
+    try:
+        while not request_task.done():
+            if interval <= 0:
+                await request_task
+                break
+            ping_sleep = asyncio.create_task(asyncio.sleep(interval))
+            done, pending = await asyncio.wait(
+                {request_task, ping_sleep},
+                return_when=asyncio.FIRST_COMPLETED,
+                timeout=UPSTREAM_TIMEOUT + 10,
+            )
+            for t in pending:
+                if t is not request_task:
+                    t.cancel()
+                    try:
+                        await t
+                    except asyncio.CancelledError:
+                        pass
+            if request_task in done:
+                break
+            # Ping timer fired: hit upstream so backend does not sleep
+            try:
+                async with httpx.AsyncClient(timeout=10.0) as client:
+                    await client.get(ping_url, headers=ping_headers)
+            except Exception:
+                pass
+        if not request_task.done():
+            request_task.cancel()
+            try:
+                await request_task
+            except asyncio.CancelledError:
+                pass
+        await request_task
+    except asyncio.CancelledError:
+        request_task.cancel()
+        try:
+            await request_task
+        except asyncio.CancelledError:
+            pass
+        raise
+    exc = request_task.exception()
+    if exc is not None:
+        raise exc
+    if r is None:
+        raise RuntimeError("Upstream request did not complete")
 
     ct = r.headers.get("content-type", "")
 
@@ -224,7 +275,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. Pings upstream every PING_UPSTREAM_INTERVAL while waiting."""
+    """POST multipart (files + data) to upstream and return JSON. Explicit ping-while-wait (same as _forward)."""
     url = f"{UPSTREAM_URL}{path}"
     headers = {
         "Authorization": f"Bearer {HF_TOKEN}",
@@ -232,13 +283,62 @@ async def _forward_multipart_json(path: str, files=None, data=None, user_token:
     }
     if user_token:
         headers["X-API-Key"] = user_token
-    ping_task = _start_ping_task(f"{UPSTREAM_URL}/", {"Authorization": f"Bearer {HF_TOKEN}"})
-    try:
-        timeout = httpx.Timeout(UPSTREAM_TIMEOUT)
+    ping_url = f"{UPSTREAM_URL}/"
+    ping_headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+    timeout = httpx.Timeout(UPSTREAM_TIMEOUT)
+    r = None
+
+    async def do_request():
+        nonlocal r
         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)
+
+    request_task = asyncio.create_task(do_request())
+    interval = PING_UPSTREAM_INTERVAL if PING_UPSTREAM_INTERVAL > 0 else 0.0
+    try:
+        while not request_task.done():
+            if interval <= 0:
+                await request_task
+                break
+            ping_sleep = asyncio.create_task(asyncio.sleep(interval))
+            done, pending = await asyncio.wait(
+                {request_task, ping_sleep},
+                return_when=asyncio.FIRST_COMPLETED,
+                timeout=UPSTREAM_TIMEOUT + 10,
+            )
+            for t in pending:
+                if t is not request_task:
+                    t.cancel()
+                    try:
+                        await t
+                    except asyncio.CancelledError:
+                        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
+        if not request_task.done():
+            request_task.cancel()
+            try:
+                await request_task
+            except asyncio.CancelledError:
+                pass
+        await request_task
+    except asyncio.CancelledError:
+        request_task.cancel()
+        try:
+            await request_task
+        except asyncio.CancelledError:
+            pass
+        raise
+    exc = request_task.exception()
+    if exc is not None:
+        raise exc
+    if r is None:
+        raise RuntimeError("Upstream request did not complete")
     ct = r.headers.get("content-type", "")
     if "application/json" in ct:
         try: