create event loop before predictor setup

[technillogue]

The current implementation (prior to #1350) throws an error:

Traceback (most recent call last):
File "/usr/local/lib/python3.11/site-packages/cog/server/worker.py", line 226, in _predict
for r in result:
...
File "/src/src/download.py", line 118, in sync_download_file
return self.loop.run_until_complete(self.download_file(url))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/asyncio/base_events.py", line 629, in run_until_complete
self._check_running()
File "/usr/local/lib/python3.11/asyncio/base_events.py", line 590, in _check_running
raise RuntimeError(
RuntimeError: Cannot run the event loop while another loop is running

The problem is that llama's setup() creates a Downloader, which can't find an event loop and creates one, then tries to use that loop in predict, and by that time _ChildWorker.run has used asyncio.run

This solves that specific problem, and since we didn't see any other errors besides llama maybe few other models use an event loop. However, we may want/need to consider moving the check where we inspect if predict() is async earlier, and have more separate code paths so that we don't need to be async if the user isn't using the async feature. On the other hand, doing that could get messy and involve a lot of duplicated code, so if we're sure our users use asyncio.get_running_loop and not asyncio.create_event_loop, it would be best to not do that.

[nickstenning]

IIUC this is the fix for cog v0.9.0b9 being broken for some models? If so would you have a look at what's broken the tests here so we can get this merged? It's not great for main to be be unreleasable.

[nickstenning]

The fact that is_async_predictor checks both setup and predict makes this a little hard to read IMO.

I think it would be clearer if the setup path did:

self._predictor = load_predictor_from_ref(self._predictor_ref)
# Could be a function or a class
if hasattr(self._predictor, "setup"):
    if is_async(self._predictor.setup):
        get_loop().run_until_complete(run_setup_async(self._predictor))
    else:
        run_setup(self._predictor)

and then similarly, later on

def _loop(self) -> None:
    if is_async(get_predict(self._predictor)):
        return get_loop().run_until_complete(self._loop_async())
    return self._loop_sync()

[technillogue]

I think the rationale here is that is predict is async, the user has opted into the async world entirely, and we want to create an event loop before running setup, so get_running_loop in the users' setup code will find the same event loop we're going to use to run _loop_async. that may not be the right approach though

[nickstenning]

Ah, ok, that makes sense. That wasn't clear to me, but the comment you added is very helpful.

[technillogue]

actually now that I think about it just using new_event_loop doesn't mean that get_running_loop will work, I might be completely wrong...

the crux of the issue is def setup() creating a Downloader, which does get_running_loop/new_event_loop, creates a ClientSession implicitly attached to that loop, uses sync_download_files_to_disk, and then (in the future), async def predict() might use await downloader.download_file() to download a finetune, and ClientSession will error because it's not the same event loop as when it was created. my specific downloader code tries to guard against this but a different user might not do that

worse, after setup() finishes doing new_event_loop().run_until_complete() and returns, our subsequent get_event_loop() won't find anything because the loop from setup() isn't running anymore

maybe if predict is async we should just always wrap setup() in a coroutine so that any get_running_loop will do the right thing?

[nickstenning]

I think what you're doing here should work.

In [24]: async def who_is_running():
    ...:     print(id(asyncio.get_running_loop()))
    ...:

In [25]: await who_is_running()
4338068384

In [26]: loop = asyncio.new_event_loop()

In [27]: id(loop)
Out[27]: 4342403856

In [28]: loop.run_until_complete(who_is_running())
4342403856

[technillogue]

it's more like this

def get_loop():
    try:
        return asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.new_event_loop()


async def who_is_running(msg):
    print(msg, id(asyncio.get_running_loop()))


def setup():
    user_loop = get_loop()
    user_loop.run_until_complete(who_is_running("user setup loop"))


async def predict():
    await who_is_running("predict loop")


loop = get_loop()
setup()
loop.run_until_complete(predict())

user setup loop 139844438211888
predict loop 139844449272592

if that's IPython with autoawait you can get confusing results

[yorickvP]

We don't have an existing contract on async def setup functions, so we're free to break people who are spawning their own event loops there. They'll discover and fix that during testing.

We do have an established behavior of allowing event loops in non-async def setup or predict functions, disallowing that now will probably break user code.
Maybe that means they should run in another thread so they're free to run their own loop?
(note that we should only call 1 at a time)

[technillogue]

I think #1366 (comment) addresses the motivation and the current code in this PR doesn't break people who spawn event loops in setup; running it in a thread shouldn't be needed

[technillogue]

if neither setup or predict is async, no event loop is created; if predict is async but setup isn't, then an event loop is created but not running, so asyncio.run will work, get_running_loop won't, and people will fall through to new_event_loop. this will be a different event loop than the one that will run an async predict, but that's tolerable

[nickstenning]

Just a couple of mostly superficial requests to make things easier to follow. Mostly it looks good.

[nickstenning]

I still think we should probably use is_async here?

[technillogue]

setup shouldn't be allowed to be a generator

[nickstenning]

Another option that might (?) make things less icky, if I've finally understood the problem here.

def run(self) -> None:
    ...
 
    self._stream_redirector.start()
    asyncio.run(self._run())
    self._stream_redirector.shutdown()

async def _run(self) -> Awaitable[None]:
    await self._setup()
    await self._loop()

async def _setup(self) -> Awaitable[None]:
    ...

[technillogue]

we don't know if we're async or not until we load_predictor_from_ref, which might throw an error, that should be handled in the same block as errors from calling setup

[technillogue]

we could duplicate the setup error handling with a similar pattern as _handle_predict_error and then this could work, I was just a little confused what was happening with Done and _cancelable

[nickstenning]

Maybe I'm missing something, but why does the predictor need to be async for these methods to be async?

[technillogue]

if the predictor isn't async, they haven't opted into the async world, and if we make these methods async we would get the same problems we were having before with there being a running event loop that the user isn't expecting, breaking unsuspecting code

[technillogue]

we could have a kind of silly check like "async" in open("predict.py").read()

[technillogue]

after reflecting on it, I think the current code has the correct behavior, but not for the reasons I initially thought. here are some some scenarios:

1. you use the same event loop from sync setup and from sync predict. nothing changes for you.
2. you have sync setup that uses asyncio.run, and an async predict. your setup and predict will use different event loops, but if setup is only concerned with downloading and predict is only concerned with the inference engine (no downloading), that's fine, and we'd rather not make people rewrite everything here. in this case there is no difference if we new_event_loop before the setup code or not.
3. you have both async setup and async predict. you can't use asyncio.run/run_until_complete, but you will be in the same running loop
4. if you really want, you can have async setup and sync predict, which is silly but won't break anything

for comparison, the other approach we paired on is here d68a147