mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-12 12:06:04 +00:00
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
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:
parent
a6a600f845
commit
dd513449de
2 changed files with 582 additions and 561 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue