mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-06 10:37:22 +00:00
fix: harden storage semantics (#4118)
Fixes issues in the storage system by guaranteeing immediate durability for responses and ensuring background writers stay alive. Three related fixes: * Responses to the OpenAI-compatible API now write directly to Postgres/SQLite inside the request instead of detouring through an async queue that might never drain; this restores the expected read-after-write behavior and removes the "response not found" races reported by users. * The access-control shim was stamping owner_principal/access_attributes as SQL NULL, which Postgres interprets as non-public rows; fixing it to use the empty-string/JSON-null pattern means conversations and responses stored without an authenticated user stay queryable (matching SQLite). * The inference-store queue remains for batching, but its worker tasks now start lazily on the live event loop so server startup doesn't cancel them—writes keep flowing even when the stack is launched via llama stack run. Closes #4115 Added a matrix entry to test our "base" suite against Postgres as the store.
This commit is contained in:
parent
56d87f5133
commit
81e44b06ff
27 changed files with 1195 additions and 160 deletions
|
|
@ -66,13 +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()))
|
||||
else:
|
||||
logger.info("Write queue disabled for SQLite to avoid concurrency issues")
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
if not self._worker_tasks:
|
||||
return
|
||||
|
|
@ -93,10 +86,29 @@ class InferenceStore:
|
|||
if self.enable_write_queue and self._queue is not None:
|
||||
await self._queue.join()
|
||||
|
||||
async def _ensure_workers_started(self) -> None:
|
||||
"""Ensure the async write queue workers run on the current loop."""
|
||||
if not self.enable_write_queue:
|
||||
return
|
||||
|
||||
if self._queue is None:
|
||||
self._queue = asyncio.Queue(maxsize=self._max_write_queue_size)
|
||||
logger.debug(
|
||||
f"Inference store write queue created with max size {self._max_write_queue_size} "
|
||||
f"and {self._num_writers} writers"
|
||||
)
|
||||
|
||||
if not self._worker_tasks:
|
||||
loop = asyncio.get_running_loop()
|
||||
for _ in range(self._num_writers):
|
||||
task = loop.create_task(self._worker_loop())
|
||||
self._worker_tasks.append(task)
|
||||
|
||||
async def store_chat_completion(
|
||||
self, chat_completion: OpenAIChatCompletion, input_messages: list[OpenAIMessageParam]
|
||||
) -> None:
|
||||
if self.enable_write_queue:
|
||||
await self._ensure_workers_started()
|
||||
if self._queue is None:
|
||||
raise ValueError("Inference store is not initialized")
|
||||
try:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue