From fa100c77fd815c8a5a887a342c05a19bc59969ab Mon Sep 17 00:00:00 2001 From: Raghotham Murthy Date: Tue, 7 Oct 2025 08:29:06 -0700 Subject: [PATCH] close connections --- llama_stack/providers/inline/agents/meta_reference/agents.py | 2 +- llama_stack/providers/inline/batches/reference/batches.py | 1 + llama_stack/providers/inline/eval/meta_reference/eval.py | 3 ++- llama_stack/providers/utils/kvstore/api.py | 4 ++++ llama_stack/providers/utils/kvstore/kvstore.py | 4 ++++ tests/unit/fixtures.py | 1 + .../vector_io/test_vector_io_openai_vector_stores.py | 3 --- 7 files changed, 13 insertions(+), 5 deletions(-) diff --git a/llama_stack/providers/inline/agents/meta_reference/agents.py b/llama_stack/providers/inline/agents/meta_reference/agents.py index 27d3a94cc..0c37b05bc 100644 --- a/llama_stack/providers/inline/agents/meta_reference/agents.py +++ b/llama_stack/providers/inline/agents/meta_reference/agents.py @@ -314,7 +314,7 @@ class MetaReferenceAgentsImpl(Agents): return paginate_records(session_dicts, start_index, limit) async def shutdown(self) -> None: - pass + await self.persistence_store.close() # OpenAI responses async def get_openai_response( diff --git a/llama_stack/providers/inline/batches/reference/batches.py b/llama_stack/providers/inline/batches/reference/batches.py index 39f45d7d1..501b17dc6 100644 --- a/llama_stack/providers/inline/batches/reference/batches.py +++ b/llama_stack/providers/inline/batches/reference/batches.py @@ -129,6 +129,7 @@ class ReferenceBatchesImpl(Batches): # don't cancel tasks - just let them stop naturally on shutdown # cancelling would mark batches as "cancelled" in the database logger.info(f"Shutdown initiated with {len(self._processing_tasks)} active batch processing tasks") + await self.kvstore.close() # TODO (SECURITY): this currently works w/ configured api keys, not with x-llamastack-provider-data or with user policy restrictions async def create_batch( diff --git a/llama_stack/providers/inline/eval/meta_reference/eval.py b/llama_stack/providers/inline/eval/meta_reference/eval.py index 0dfe23dca..ff91a8da6 100644 --- a/llama_stack/providers/inline/eval/meta_reference/eval.py +++ b/llama_stack/providers/inline/eval/meta_reference/eval.py @@ -64,7 +64,8 @@ class MetaReferenceEvalImpl( benchmark = Benchmark.model_validate_json(benchmark) self.benchmarks[benchmark.identifier] = benchmark - async def shutdown(self) -> None: ... + async def shutdown(self) -> None: + await self.kvstore.close() async def register_benchmark(self, task_def: Benchmark) -> None: # Store in kvstore diff --git a/llama_stack/providers/utils/kvstore/api.py b/llama_stack/providers/utils/kvstore/api.py index d17dc66e1..06bce2891 100644 --- a/llama_stack/providers/utils/kvstore/api.py +++ b/llama_stack/providers/utils/kvstore/api.py @@ -19,3 +19,7 @@ class KVStore(Protocol): async def values_in_range(self, start_key: str, end_key: str) -> list[str]: ... async def keys_in_range(self, start_key: str, end_key: str) -> list[str]: ... + + async def close(self) -> None: + """Close any persistent connections. Optional method for cleanup.""" + ... diff --git a/llama_stack/providers/utils/kvstore/kvstore.py b/llama_stack/providers/utils/kvstore/kvstore.py index 426523d8e..475d5df96 100644 --- a/llama_stack/providers/utils/kvstore/kvstore.py +++ b/llama_stack/providers/utils/kvstore/kvstore.py @@ -43,6 +43,10 @@ class InmemoryKVStoreImpl(KVStore): async def delete(self, key: str) -> None: del self._store[key] + async def close(self) -> None: + """No-op for in-memory store.""" + pass + async def kvstore_impl(config: KVStoreConfig) -> KVStore: if config.type == KVStoreType.redis.value: diff --git a/tests/unit/fixtures.py b/tests/unit/fixtures.py index 443a1d371..ff86f564f 100644 --- a/tests/unit/fixtures.py +++ b/tests/unit/fixtures.py @@ -18,6 +18,7 @@ async def sqlite_kvstore(tmp_path): kvstore = SqliteKVStoreImpl(kvstore_config) await kvstore.initialize() yield kvstore + await kvstore.close() @pytest.fixture(scope="function") diff --git a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py index ed0934224..436f32e27 100644 --- a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py +++ b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py @@ -46,12 +46,9 @@ async def test_initialize_index(vector_index): async def test_add_chunks_query_vector(vector_index, sample_chunks, sample_embeddings): - vector_index.delete() - vector_index.initialize() await vector_index.add_chunks(sample_chunks, sample_embeddings) resp = await vector_index.query_vector(sample_embeddings[0], k=1, score_threshold=-1) assert resp.chunks[0].content == sample_chunks[0].content - vector_index.delete() async def test_chunk_id_conflict(vector_index, sample_chunks, embedding_dimension):