Adjusted ping every 270 seconds for large computations up to 900 seconds

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", "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
 PING_UPSTREAM_INTERVAL = float(os.getenv("PING_UPSTREAM_INTERVAL", "270"))
 
@@ -173,22 +173,44 @@ async def _forward_stream(path: str, files=None, data=None, user_token: str | No
     headers_holder = {}
 
     async def body_iter():
+        chunk_queue: asyncio.Queue = asyncio.Queue()
+        stream_done = {"done": False, "exc": None}
+
+        async def stream_reader():
+            try:
+                async with stream_ctx as resp:
+                    status_holder["code"] = resp.status_code
+                    media_type_holder["ct"] = resp.headers.get("content-type", "application/octet-stream")
+                    headers_holder.update(_filter_resp_headers(resp.headers))
+                    if resp.status_code >= 400:
+                        chunk = await resp.aread()
+                        await chunk_queue.put(chunk)
+                        await chunk_queue.put(None)
+                        return
+                    async for chunk in resp.aiter_raw():
+                        await chunk_queue.put(chunk)
+                    await chunk_queue.put(None)
+            except Exception as e:
+                stream_done["exc"] = e
+                await chunk_queue.put(None)
+            finally:
+                stream_done["done"] = True
+
+        reader_task = asyncio.create_task(stream_reader())
         try:
-            async with stream_ctx as resp:
-                status_holder["code"] = resp.status_code
-                media_type_holder["ct"] = resp.headers.get("content-type", "application/octet-stream")
-                headers_holder.update(_filter_resp_headers(resp.headers))
-
-                # If upstream already returned an error, buffer it (short text/json)
-                if resp.status_code >= 400:
-                    # Buffer entire payload and yield once
-                    chunk = await resp.aread()
-                    yield chunk
-                    return
-
-                async for chunk in resp.aiter_raw():
-                    yield chunk
+            while True:
+                chunk = await chunk_queue.get()
+                if chunk is None:
+                    if stream_done["exc"]:
+                        raise stream_done["exc"]
+                    break
+                yield chunk
         finally:
+            reader_task.cancel()
+            try:
+                await reader_task
+            except asyncio.CancelledError:
+                pass
             await _cancel_ping_task(ping_task)
             await client.aclose()
 
@@ -251,20 +273,44 @@ async def _forward_demo_stream(path: str, *, files: dict | None, data: dict | No
     headers_holder = {}
 
     async def body_iter():
+        chunk_queue: asyncio.Queue = asyncio.Queue()
+        stream_done = {"done": False, "exc": None}
+
+        async def stream_reader():
+            try:
+                async with stream_ctx as resp:
+                    status_holder["code"] = resp.status_code
+                    media_type_holder["ct"] = resp.headers.get("content-type", "application/octet-stream")
+                    headers_holder.update(_filter_resp_headers(resp.headers))
+                    if resp.status_code >= 400:
+                        chunk = await resp.aread()
+                        await chunk_queue.put(chunk)
+                        await chunk_queue.put(None)
+                        return
+                    async for chunk in resp.aiter_raw():
+                        await chunk_queue.put(chunk)
+                    await chunk_queue.put(None)
+            except Exception as e:
+                stream_done["exc"] = e
+                await chunk_queue.put(None)
+            finally:
+                stream_done["done"] = True
+
+        reader_task = asyncio.create_task(stream_reader())
         try:
-            async with stream_ctx as resp:
-                status_holder["code"] = resp.status_code
-                media_type_holder["ct"] = resp.headers.get("content-type", "application/octet-stream")
-                headers_holder.update(_filter_resp_headers(resp.headers))
-
-                if resp.status_code >= 400:
-                    chunk = await resp.aread()
-                    yield chunk
-                    return
-
-                async for chunk in resp.aiter_raw():
-                    yield chunk
+            while True:
+                chunk = await chunk_queue.get()
+                if chunk is None:
+                    if stream_done["exc"]:
+                        raise stream_done["exc"]
+                    break
+                yield chunk
         finally:
+            reader_task.cancel()
+            try:
+                await reader_task
+            except asyncio.CancelledError:
+                pass
             await _cancel_ping_task(ping_task)
             await client.aclose()