fix: InferenceStore workers being cancelled on event loop change (#4373)
Some checks failed
Unit Tests / unit-tests (3.13) (push) Failing after 2m9s
Integration Auth Tests / test-matrix (oauth2_token) (push) Failing after 2s
SqlStore Integration Tests / test-postgres (3.13) (push) Failing after 1s
Integration Tests (Replay) / generate-matrix (push) Successful in 5s
Vector IO Integration Tests / test-matrix (push) Failing after 51s
SqlStore Integration Tests / test-postgres (3.12) (push) Failing after 53s
Unit Tests / unit-tests (3.12) (push) Failing after 2m20s
Pre-commit / pre-commit (push) Successful in 2m46s
Integration Tests (Replay) / Integration Tests (, , , client=, ) (push) Failing after 2m51s

# What does this PR do?

The InferenceStore write queue workers were created during initialize()
which runs in the app factory's event loop. When uvicorn starts, it
creates its own event loop, causing the original worker tasks to be
cancelled. This resulted in chat completions being queued but never
written to the database.

Root cause: asyncio tasks are bound to the event loop they were created
in. When the event loop changes (app factory -> uvicorn server), tasks
from the old loop are cancelled. However, the task references remained
in self._worker_tasks, so _ensure_workers_started() thought workers
were still running.

The fix removes eager worker creation from initialize(). Workers must
not be started during app factory execution because that event loop is
temporary. Instead, _ensure_workers_started() creates workers lazily
during store_chat_completion(), which runs during request handling and
is guaranteed to be in uvicorn's event loop.

Additionally, _ensure_workers_started() now checks for active (non-done)
tasks rather than just checking if the list is non-empty. This makes
the code resilient to any future event loop transitions.

This follows the standard async Python pattern: defer background task
creation until the server is actually running, not during app init.

## Test Plan

Run the server with postgres or sqlite, I tested with postgres:

```
uv run llama stack run llama_stack/distributions/starter/run-with-postgres-store.yaml
```

Execute an inference call:

```
export INFERENCE_MODEL=gpt-4o-mini && curl -fsS http://127.0.0.1:8321/v1/chat/completions -H "Content-Type: application/json" -d "{\"model\": \"openai/$INFERENCE_MODEL\",\"messages\": [{\"role\": \"user\", \"content\": \"What color is grass?\"}], \"max_tokens\": 128, \"temperature\": 0.0}"
```

Check the DB being populated:

```
podman exec postgres psql -U llamastack -d llamastack -t -c "SELECT COUNT(*) FROM chat_completions;" 
    1
```

**AGAINST STABLE BRANCH 0.3.X SINCE MAIN DOES NOT HAVE THE ISSUE.**

---------

Signed-off-by: Sébastien Han <seb@redhat.com>
This commit is contained in:
Sébastien Han 2025-12-11 17:32:04 +01:00 committed by GitHub
parent a6a600f845
commit dd513449de
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 582 additions and 561 deletions

View file

@ -66,14 +66,6 @@ class InferenceStore:
},
)
if self.enable_write_queue:
self._queue = asyncio.Queue(maxsize=self._max_write_queue_size)
for _ in range(self._num_writers):
self._worker_tasks.append(asyncio.create_task(self._worker_loop()))
logger.debug(
f"Inference store write queue enabled with {self._num_writers} writers, max queue size {self._max_write_queue_size}"
)
async def shutdown(self) -> None:
if not self._worker_tasks:
return

1135
uv.lock generated

File diff suppressed because it is too large Load diff