Ensure Sync generators get closed if they are cancelled (#11396)

Author: freddyaboulton

.changeset/thick-hands-bet.md

@@ -0,0 +1,5 @@
+---
+"gradio": patch
+---
+
+fix:Ensure Sync generators get closed if they are cancelled

gradio/queueing.py

@@ -34,6 +34,7 @@
     LRUCache,
     error_payload,
     run_coro_in_background,
+    safe_aclose_iterator,
     safe_get_lock,
     set_task_name,
 )
@@ -808,6 +809,10 @@ async def reset_iterators(self, event_id: str):
             # Failure, but don't raise an error
             return
         async with app.lock:
+            try:
+                await safe_aclose_iterator(app.iterators[event_id])
+            except Exception:
+                pass
             del app.iterators[event_id]
             app.iterators_to_reset.add(event_id)
         return

gradio/utils.py

@@ -744,6 +744,9 @@ async def __anext__(self):
             run_sync_iterator_async, self.iterator, limiter=self.limiter
         )
 
+    def aclose(self):
+        self.iterator.close()
+
 
 async def async_iteration(iterator):
     return await anext(iterator)
@@ -1778,3 +1781,27 @@ def get_function_description(fn: Callable) -> tuple[str, dict[str, str], list[st
         pass
 
     return description, parameters, returns
+
+
+async def safe_aclose_iterator(iterator, timeout=60.0, retry_interval=0.05):
+    """
+    Safely close generators by calling the aclose method.
+    Sync generators are tricky because if you call `aclose` while the loop is running
+    then you get a ValueError and the generator will not shut down gracefully.
+    So the solution is to retry calling the aclose method until we succeed (with timeout).
+    """
+    start = time.monotonic()
+    if isinstance(iterator, SyncToAsyncIterator):
+        while True:
+            try:
+                iterator.aclose()
+                break
+            except ValueError as e:
+                if "already executing" in str(e):
+                    if time.monotonic() - start > timeout:
+                        raise
+                    await asyncio.sleep(retry_interval)
+                else:
+                    raise
+    else:
+        iterator.aclose()