mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-17 10:42:36 +00:00
chore: fix flaky unit test and add proper shutdown for file batches
This commit is contained in:
parent
1970b4aa4b
commit
cf418f3394
9 changed files with 39 additions and 17 deletions
|
|
@ -293,6 +293,19 @@ class OpenAIVectorStoreMixin(ABC):
|
|||
await self._resume_incomplete_batches()
|
||||
self._last_file_batch_cleanup_time = 0
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
"""Clean up mixin resources including background tasks."""
|
||||
# Cancel any running file batch tasks gracefully
|
||||
if hasattr(self, "_file_batch_tasks"):
|
||||
tasks_to_cancel = list(self._file_batch_tasks.items())
|
||||
for _, task in tasks_to_cancel:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def delete_chunks(self, store_id: str, chunks_for_deletion: list[ChunkForDeletion]) -> None:
|
||||
"""Delete chunks from a vector store."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue