Add support for async generators

[freddyaboulton]

Description

Closes: #3618 . @deep-diver would appreciate if you could test this to see if it fixes your issue. Click on the link on the gradio-pr-bot message, Go to Files in the top right, requirements.txt will have a url hosting a wheel built for this PR.

To test

Blocks case

import asyncio
import gradio as gr
import time


async def async_iteration(count: int):
    for i in range(count):
        yield i
        await asyncio.sleep(1)


def iteration(count: int):
    for i in range(count):
        yield i
        time.sleep(1)


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            num1 = gr.Number(value=4, precision=0)
            o1 = gr.Number()
            async_iterate = gr.Button(value="Async Iteration")
            async_iterate.click(async_iteration, num1, o1)
        with gr.Column():
            num2 = gr.Number(value=4, precision=0)
            o2 = gr.Number()
            iterate = gr.Button(value="Iterate")
            iterate.click(iteration, num2, o2)

demo.queue(concurrency_count=20).launch()

Interface case

app1 = gr.Interface(async_iteration, gr.Number(precision=0, value=4), gr.Number())
app2 = gr.Interface(iteration, gr.Number(precision=0, value=4), gr.Number())

with gr.Blocks() as demo:
    with gr.Row():
        app1.render()
    with gr.Row():
        app2.render()

demo.queue(concurrency_count=20).launch()

Checklist:

• I have performed a self-review of my own code
• I have added a short summary of my change to the CHANGELOG.md
• My code follows the style guidelines of this project
• I have commented my code in hard-to-understand areas
• I have made corresponding changes to the documentation
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

A note about the CHANGELOG

Hello 👋 and thank you for contributing to Gradio!

All pull requests must update the change log located in CHANGELOG.md, unless the pull request is labeled with the "no-changelog-update" label.

Please add a brief summary of the change to the Upcoming Release > Full Changelog section of the CHANGELOG.md file and include
a link to the PR (formatted in markdown) and a link to your github profile (if you like). For example, "* Added a cool new feature by [@myusername](link-to-your-github-profile) in [PR 11111](https://github.com/gradio-app/gradio/pull/11111)".

If you would like to elaborate on your change further, feel free to include a longer explanation in the other sections.
If you would like an image/gif/video showcasing your feature, it may be best to edit the CHANGELOG file using the
GitHub web UI since that lets you upload files directly via drag-and-drop.

[gradio-pr-bot]

All the demos for this PR have been deployed at https://huggingface.co/spaces/gradio-pr-deploys/pr-3821-all-demos

[deep-diver]

Thanks for doing this @freddyaboulton! I will test and get back as soon as possible!

[freddyaboulton]

Just checking in @deep-diver :)

[davidrios]

I can confirm this works. Here's my app using it: https://github.com/davidrios/speech2text-server, the response comes streaming when using this branch.
But I wasn't able to catch the stop action from the form. Do you know how to do that in this case?

[freddyaboulton]

Thanks for testing and verifying it works! Can you elaborate a bit more on what you couldn't do @davidrios ?

[davidrios]

Oh, sure. The Gradio form has a Submit button, when the user clicks it the process starts and the results come streaming, meanwhile the Submit turns to a Stop button. I would like to catch that Stop action to stop the underlying processing, because it's a heavy process running in another thread. Right now when I click the Stop button, in the UI it seems like it stopped, but nothing happens in my code and the processing keeps going until the end. Some times when I click submit again it just continues from where it left off and some times it really stops but I get this in the log:

INFO:     127.0.0.1:52450 - "POST /api/predict HTTP/1.1" 200 OK
Executing <Task finished name='Task-51' coro=<RequestResponseCycle.run_asgi() done, defined at /home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/uvicorn/protocols/http/httptools_impl.py:433> result=None created at /home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/uvicorn/protocols/http/httptools_impl.py:312> took 0.984 seconds
/home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/gradio/processing_utils.py:180: UserWarning: Trying to convert audio automatically from int32 to 16-bit int format.
  warnings.warn(warning.format(data.dtype))
INFO:     127.0.0.1:52450 - "POST /api/predict HTTP/1.1" 200 OK
Executing <Task finished name='Task-53' coro=<RequestResponseCycle.run_asgi() done, defined at /home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/uvicorn/protocols/http/httptools_impl.py:433> result=None created at /home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/uvicorn/protocols/http/httptools_impl.py:312> took 0.948 seconds
/home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/gradio/processing_utils.py:180: UserWarning: Trying to convert audio automatically from int32 to 16-bit int format.
  warnings.warn(warning.format(data.dtype))
INFO:     127.0.0.1:52450 - "POST /api/predict HTTP/1.1" 200 OK
Executing <Task finished name='Task-55' coro=<RequestResponseCycle.run_asgi() done, defined at /home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/uvicorn/protocols/http/httptools_impl.py:433> result=None created at /home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/uvicorn/protocols/http/httptools_impl.py:312> took 0.979 seconds
INFO:     127.0.0.1:47866 - "POST /run/predict HTTP/1.1" 200 OK
INFO:     127.0.0.1:52450 - "POST /reset HTTP/1.1" 200 OK
Task exception was never retrieved
future: <Task finished name='Task-60' coro=<<async_generator_athrow without __name__>()> exception=RuntimeError('async generator ignored GeneratorExit') created at /usr/lib/python3.11/asyncio/runners.py:118>
source_traceback: Object created at (most recent call last):
  File "<string>", line 1, in <module>
  File "/usr/lib/python3.11/multiprocessing/spawn.py", line 120, in spawn_main
    exitcode = _main(fd, parent_sentinel)
  File "/usr/lib/python3.11/multiprocessing/spawn.py", line 133, in _main
    return self._bootstrap(parent_sentinel)
  File "/usr/lib/python3.11/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/usr/lib/python3.11/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/uvicorn/_subprocess.py", line 76, in subprocess_started
    target(sockets=sockets)
  File "/home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/uvicorn/server.py", line 61, in run
    return asyncio.run(self.serve(sockets=sockets))
  File "/usr/lib/python3.11/asyncio/runners.py", line 190, in run
    return runner.run(main)
  File "/usr/lib/python3.11/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
RuntimeError: async generator ignored GeneratorExit
INFO:     127.0.0.1:47876 - "POST /run/predict HTTP/1.1" 200 OK

Also, if I try to stop the server, it stops fine if I have never clicked Stop, but after that it goes:

^CINFO:     Shutting down
INFO:     Waiting for background tasks to complete. (CTRL+C to force quit)

But it never stops and if I force:

.... lots of lines
  File "/home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/fastapi/routing.py", line 289, in app
    await dependant.call(**values)
  File "/home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/gradio/routes.py", line 535, in join_queue
    await asyncio.sleep(1)
  File "/usr/lib/python3.11/asyncio/tasks.py", line 639, in sleep
    return await future
           ^^^^^^^^^^^^
asyncio.exceptions.CancelledError

... lots of lines
ERROR:    Traceback (most recent call last):
  File "/home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/starlette/routing.py", line 686, in lifespan
    await receive()
  File "/home/david/work/speech2text/server/.venv/lib/python3.11/site-packages/uvicorn/lifespan/on.py", line 139, in receive
    return await self.receive_queue.get()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/asyncio/queues.py", line 158, in get
    await getter
asyncio.exceptions.CancelledError

[davidrios]

So, after improving the robustness of my code and adding a lot of logging, I'm almost certain the issue is that my async generator passed to Gradio is never cancelled or awaited when I click the Stop button. Something inside Gradio stops but it doesn't clean up the generator, so I can't run my clean up code either.

[davidrios]

Here's a minimal example to reproduce the issue:

import asyncio

import gradio as gr

execution = 0


async def gradio_output_iterator(text):
    try:
        output = "* Counting to 30... *\n"
        yield output
        for i in range(1, 31):
            await asyncio.sleep(1)
            output += f"{i}... "
            yield output
            if i % 10 == 0:
                output += "\n"
                yield output
        output += "\ndone!"
        yield output
    finally:
        global execution
        execution += 1
        print(
            f"*** {execution} finished gradio output. this should appear in the log "
            "when the user clicks Stop or the counter finishes"
        )


demo = gr.Interface(
    gradio_output_iterator,
    inputs=[
        gr.Text(info="no need to put anything"),
    ],
    outputs=gr.Text(label="Results"),
    allow_flagging="never",
)

demo.queue(concurrency_count=20).launch()

[freddyaboulton]

Refactored this test to use a websocket connection

[freddyaboulton]

Refactored the Interface cancel button to use then as opposed to a finally block. Helps prevent the issues with the generator not being cleaned up

[freddyaboulton]

The idea here is to wrap an sync iterator into an async one so that we can have one api for yielding outputs from the iterator

[freddyaboulton]

Think I figured out the issue @davidrios - the finally block in your repro prints for me and there are no scary warnings printed to the console. Thanks for the help!

[abidlabs]

Elegant solution @freddyaboulton! I was able to fine just one edge case where it seems to get stuck, which is when the async generator is started and stopped, then the sync generator is started and stopped, and and then the async generator is started and stopped Interface example. In the frontend, I get this error:

runtime.esm.js:1  Uncaught (in promise) TypeError: Cannot read properties of null (reading 'id')
    at Array.Wl (runtime.esm.js:1:5521)
    at Object.p (Button.svelte:24:3)
    at Vt (index.mjs:100:14)
    at Object.p (Button.svelte:25:9)
    at Mn (index.mjs:1273:36)
    at en (index.mjs:1233:17)

I think that even though the async generator is stopped, the frontend still expects some payload returned? Here's a GIF showing the behavior happening:

To reproduce, run:

import asyncio
import gradio as gr
import time


async def async_iteration(count: int):
    for i in range(count):
        yield i
        await asyncio.sleep(1)


def iteration(count: int):
    for i in range(count):
        yield i
        time.sleep(1)


app1 = gr.Interface(async_iteration, gr.Number(precision=0, value=4), gr.Number())
app2 = gr.Interface(iteration, gr.Number(precision=0, value=4), gr.Number())

with gr.Blocks() as demo:
    with gr.Row():
        app1.render()
    with gr.Row():
        app2.render()

demo.queue(concurrency_count=20).launch()

Then:

• Start the async Interface and stop it before prediction is complete
• Start the sync Interface and stop it before prediction is complete
• Start the async Interface and stop it before prediction is complete --> error is thrown

It's definitely an edge case, but it's happening consistently so might be worth investigating if there's a deeper issue here?

[freddyaboulton]

@abidlabs Thanks for the review! The app was crashing because we were retuning None as the value of the stop button. This was because we were also returning update values for the buttons from the prediction function. Long story short, I refactored the submit -> stop -> submit button update flow to use then api entirely. This is much nicer, as the generator no longer outputs values for the buttons so view_api for the client is much nicer. The one downside is that because of #4113, if the generator fails the stop button won't get converted back to submit button. The user can click the stop button again to "reset" it though. This is a pre-existing bug so I'm ok merging this in but we can also fix #4113 beforehand as not changing the stop button back is technically a change of behavior.

Also, there can still be a weird race condition if you stop and start separate iterators in the example you tried. The app will not crash however.

[aliabid94]

Working on fixing #4113 now

[abidlabs]

Looks great!! Thanks @freddyaboulton for the great PR and @aliabid94 for the quick fix for .then(). Merging in

[deepankarm]

Do we know when is the next release planned with this PR?

[freddyaboulton]

@deepankarm This got released yesterday with gradio 3.29.0!