Batched OpenAI request worker async race (AssertionError: bopenai...)

[kentjhall]

In attempting to invoke an lmql query in an asynchronous context (specifically using asyncio.gather), I encounter the following assertion error every few runs or so:

AssertionError: bopenai requires the current event loop to be running in one of the worker threads

I did some digging, and what appears to be to be the issue is a race between async def complete and async def complete_request_worker. The latter is responsible for adding to the worker_loops set of the AsyncOpenAIAPI instance:

670         self.worker_loops.add(asyncio.get_running_loop())

And the former uses this set in an assertion:

741         assert loop in self.worker_loops, f"bopenai requires the current event loop to be running in one of the worker threads"

It seems complete_request_workers tasks are scheduled on construction of the AsyncOpenAIAPI instance:

579         self.complete_request_workers = [asyncio.create_task(self.complete_request_worker(self.complete_request_queue)) for i in range(5)]

And complete is invoked asynchronously on that instance right after construction, from __init__.py:

 95         if _api is None or not _api.is_available():
 96             _api = AsyncOpenAIAPI()
 97         return await _api.complete(*args, **kwargs)

I don't know how exactly tasks are scheduled in asynchronous Python land, but it seems to me that there's no guarantee as to the order of when these two tasks are executed—at least, this seems to be the case in my environment, as I found through some print statements that complete would sometimes execute before complete_request_workers had gotten the chance to add to the set, hence triggering the assert.