phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-16 22:29:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

chore: fix flaky unit test and add proper shutdown for file batches

Browse source
This commit is contained in:
Swapna Lekkala 2025-10-07 14:09:28 -07:00
parent 1970b4aa4b
commit cf418f3394
9 changed files with 39 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

3
llama_stack/providers/remote/vector_io/chroma/chroma.py
View file

@ -167,7 +167,8 @@ class ChromaVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
self.openai_vector_stores = await self._load_openai_vector_stores()
async def shutdown(self) -> None:
pass
# Clean up mixin resources (file batch tasks)
await super().shutdown()
async def register_vector_db(
self,

2
llama_stack/providers/remote/vector_io/milvus/milvus.py
View file

@ -349,6 +349,8 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
async def shutdown(self) -> None:
self.client.close()
# Clean up mixin resources (file batch tasks)
await super().shutdown()
async def register_vector_db(
self,

2
llama_stack/providers/remote/vector_io/pgvector/pgvector.py
View file

@ -390,6 +390,8 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
if self.conn is not None:
self.conn.close()
log.info("Connection to PGVector database server closed")
# Clean up mixin resources (file batch tasks)
await super().shutdown()
async def register_vector_db(self, vector_db: VectorDB) -> None:
# Persist vector DB metadata in the KV store

2
llama_stack/providers/remote/vector_io/qdrant/qdrant.py
View file

@ -191,6 +191,8 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
async def shutdown(self) -> None:
await self.client.close()
# Clean up mixin resources (file batch tasks)
await super().shutdown()
async def register_vector_db(
self,

2
llama_stack/providers/remote/vector_io/weaviate/weaviate.py
View file

@ -347,6 +347,8 @@ class WeaviateVectorIOAdapter(
async def shutdown(self) -> None:
for client in self.client_cache.values():
client.close()
# Clean up mixin resources (file batch tasks)
await super().shutdown()
async def register_vector_db(
self,