AsyncEngineDeadError when wrapped in asyncio.timeout

[benlebrun]

@gabegrand has identified an issue in which there is a large traceback of AsyncEngineDeadError errors at shutdown:

During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "<string>", line 35, in step
  File "/lingo/grandg/projects/disciple/disciple/base_models.py", line 123, in next_word
    token = await self.sample(self.context.next_token())
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lingo/grandg/projects/hfppl/hfppl/modeling.py", line 249, in sample
    x, _ = await dist.sample()
           ^^^^^^^^^^^^^^^^^^^
  File "/lingo/grandg/projects/hfppl/hfppl/distributions/lmcontext.py", line 37, in sample
    updated_logprobs = await self.ctx.lm.next_token_logprobs(self.ctx.tokens)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lingo/grandg/projects/hfppl/hfppl/llms.py", line 298, in next_token_logprobs
    logprobs = await self.model.next_token_logprobs(token_ids)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lingo/grandg/projects/disciple/.venv/lib/python3.11/site-packages/genlm_backend/llm/vllm.py", line 104, in next_token_logprobs
    result = await self._next_token_logprobs(key)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lingo/grandg/projects/disciple/.venv/lib/python3.11/site-packages/genlm_backend/llm/vllm.py", line 125, in _next_token_logprobs
    async for output in self.async_llm_engine.generate(
  File "/lingo/grandg/projects/disciple/.venv/lib/python3.11/site-packages/vllm/engine/async_llm_engine.py", line 1080, in generate
    await self.abort(request_id)
  File "/lingo/grandg/projects/disciple/.venv/lib/python3.11/site-packages/vllm/engine/async_llm_engine.py", line 1178, in abort
    raise AsyncEngineDeadError(
vllm.engine.async_llm_engine.AsyncEngineDeadError: Background loop is not running. If it was running, inspect the output to find the stacktrace of the error that caused the background loop to stop (AsyncEngineDeadError).

The particular use case is that the AsyncVirtualLM is wrapped in an asyncio.timeout. Current hypothesis is that once the timeout fires, the underlying AsyncLLMEngine gets deleted before it has had the chance to abort any active requests.