feat: Add shutdown functionality to LlamaStackAsLibraryClient and AsyncLlamaStackAsLibraryClient (ogx-ai#4642)

# What does this PR do?
Adds shutdown functionality to LlamaStackAsLibraryClient and
AsyncLlamaStackAsLibraryClient. See the change in
docs/docs/distributions/importing_as_library.mdx for details.

Closes ogx-ai#4641

## Test Plan
The test script from ogx-ai#4641, modified to shut down
AsyncLlamaStackAsLibraryClient, exits.
```
"""
Minimal test to reproduce the race condition in _load_sentence_transformer_model
via the Llama Stack API.
"""

import asyncio
import tempfile
from io import BytesIO

from llama_stack.core.library_client import AsyncLlamaStackAsLibraryClient

CONFIG = """
version: 2
image_name: test
apis: [files, vector_io, inference]
providers:
  inference:
  - provider_id: st
    provider_type: inline::sentence-transformers
  files:
  - provider_id: fs
    provider_type: inline::localfs
    config:
      storage_dir: {tmpdir}/files
      metadata_store: {{table_name: files, backend: sql_default}}
  vector_io:
  - provider_id: faiss
    provider_type: inline::faiss
    config:
      persistence: {{namespace: faiss, backend: kv_default}}
storage:
  backends:
    kv_default: {{type: kv_sqlite, db_path: {tmpdir}/kv.db}}
    sql_default: {{type: sql_sqlite, db_path: {tmpdir}/sql.db}}
  stores:
    metadata: {{namespace: registry, backend: kv_default}}
    inference: {{table_name: inference, backend: sql_default}}
    conversations: {{table_name: conversations, backend: sql_default}}
    prompts: {{namespace: prompts, backend: kv_default}}
metadata_store: {{type: sqlite, db_path: {tmpdir}/meta.db}}
vector_stores:
  default_provider_id: faiss
  default_embedding_model: {{provider_id: st, model_id: all-MiniLM-L6-v2}}
  file_batch_params: {{max_concurrent_files_per_batch: 5}}
registered_resources:
  models:
  - model_id: all-MiniLM-L6-v2
    provider_id: st
    provider_model_id: all-MiniLM-L6-v2
    model_type: embedding
    metadata: {{embedding_dimension: 384}}
"""

async def main():
    with tempfile.TemporaryDirectory() as tmpdir:
        cfg_path = f"{tmpdir}/config.yaml"
        with open(cfg_path, "w") as f:
            f.write(CONFIG.format(tmpdir=tmpdir))

        async with AsyncLlamaStackAsLibraryClient(cfg_path) as client:
            # Upload 5 files
            file_ids = []
            for i in range(5):
                with BytesIO(f"Test content {i}".encode() * 50) as buf:
                    buf.name = f"test_{i}.txt"
                    file_ids.append((await client.files.create(file=buf, purpose="assistants")).id)

            # Create vector store and batch - triggers concurrent embedding loads
            vs = await client.vector_stores.create(name="test")
            batch = await client.vector_stores.file_batches.create(
                vector_store_id=vs.id, file_ids=file_ids
            )

            # Wait for completion
            while batch.status == "in_progress":
                await asyncio.sleep(0.5)
                batch = await client.vector_stores.file_batches.retrieve(
                    vector_store_id=vs.id, batch_id=batch.id
                )

            print(f"Result: {batch.file_counts.completed} succeeded, {batch.file_counts.failed} failed")
            if batch.file_counts.failed:
                print("*** RACE CONDITION DETECTED ***")


if __name__ == "__main__":
    import faulthandler
    import signal
    faulthandler.register(signal.SIGUSR1)

    asyncio.run(main())
```

Author: syedriko

docs/docs/distributions/importing_as_library.mdx

@@ -38,3 +38,67 @@ If you've created a [custom distribution](./building_distro), you can also use t
 ```python
 client = LlamaStackAsLibraryClient(config_path)
 ```
+
+## Resource Management
+
+When you're done using the client, you should properly release resources such as database connections. There are two ways to do this:
+
+### Using Context Managers (Recommended)
+
+The easiest and most Pythonic way is to use the client as a context manager, which automatically handles cleanup:
+
+```python
+# Synchronous client
+from llama_stack.core.library_client import LlamaStackAsLibraryClient
+
+with LlamaStackAsLibraryClient("starter") as client:
+    response = client.models.list()
+# Client is automatically shut down here
+```
+
+For the async client:
+
+```python
+from llama_stack.core.library_client import AsyncLlamaStackAsLibraryClient
+
+async with AsyncLlamaStackAsLibraryClient("starter") as client:
+    response = await client.models.list()
+# Client is automatically shut down here
+```
+
+### Using Explicit shutdown()
+
+Alternatively, you can manually call `shutdown()` when you're done:
+
+```python
+# Synchronous client
+client = LlamaStackAsLibraryClient("starter")
+try:
+    # ... use the client ...
+    response = client.models.list()
+finally:
+    client.shutdown()
+```
+
+For the async client:
+
+```python
+from llama_stack.core.library_client import AsyncLlamaStackAsLibraryClient
+
+client = AsyncLlamaStackAsLibraryClient("starter")
+await client.initialize()
+try:
+    # ... use the client ...
+    response = await client.models.list()
+finally:
+    await client.shutdown()
+```
+
+The `shutdown()` method:
+- Closes all database connections (SQLite, PostgreSQL, etc.)
+- Releases any held resources
+- Can be called multiple times safely (idempotent)
+
+:::tip
+If you don't call `shutdown()` or use a context manager, your program may hang on exit while waiting for background threads to complete, especially when using SQLite-based storage backends.
+:::

src/llama_stack/core/library_client.py

@@ -161,6 +161,45 @@ def initialize(self):
         """
         pass
 
+    def shutdown(self) -> None:
+        """Shutdown the client and release all resources.
+
+        This method should be called when you're done using the client to properly
+        close database connections and release other resources. Failure to call this
+        method may result in the program hanging on exit while waiting for background
+        threads to complete.
+
+        This method is idempotent and can be called multiple times safely.
+
+        Example:
+            client = LlamaStackAsLibraryClient("starter")
+            # ... use the client ...
+            client.shutdown()
+        """
+        loop = self.loop
+        asyncio.set_event_loop(loop)
+        try:
+            loop.run_until_complete(self.async_client.shutdown())
+        finally:
+            loop.close()
+            asyncio.set_event_loop(None)
+
+    def __enter__(self) -> "LlamaStackAsLibraryClient":
+        """Enter the context manager.
+
+        The client is already initialized in __init__, so this just returns self.
+
+        Example:
+            with LlamaStackAsLibraryClient("starter") as client:
+                response = client.models.list()
+            # Client is automatically shut down here
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the context manager and shut down the client."""
+        self.shutdown()
+
     def request(self, *args, **kwargs):
         loop = self.loop
         asyncio.set_event_loop(loop)
@@ -224,6 +263,7 @@ def __init__(
         self.custom_provider_registry = custom_provider_registry
         self.provider_data = provider_data
         self.route_impls: RouteImpls | None = None  # Initialize to None to prevent AttributeError
+        self.stack: Stack | None = None
 
     def _remove_root_logger_handlers(self):
         """
@@ -246,9 +286,9 @@ async def initialize(self) -> bool:
         try:
             self.route_impls = None
 
-            stack = Stack(self.config, self.custom_provider_registry)
-            await stack.initialize()
-            self.impls = stack.impls
+            self.stack = Stack(self.config, self.custom_provider_registry)
+            await self.stack.initialize()
+            self.impls = self.stack.impls
         except ModuleNotFoundError as _e:
             cprint(_e.msg, color="red", file=sys.stderr)
             cprint(
@@ -283,6 +323,43 @@ async def initialize(self) -> bool:
         self.route_impls = initialize_route_impls(self.impls)
         return True
 
+    async def shutdown(self) -> None:
+        """Shutdown the client and release all resources.
+
+        This method should be called when you're done using the client to properly
+        close database connections and release other resources. Failure to call this
+        method may result in the program hanging on exit while waiting for background
+        threads to complete.
+
+        This method is idempotent and can be called multiple times safely.
+
+        Example:
+            client = AsyncLlamaStackAsLibraryClient("starter")
+            await client.initialize()
+            # ... use the client ...
+            await client.shutdown()
+        """
+        if self.stack:
+            await self.stack.shutdown()
+            self.stack = None
+
+    async def __aenter__(self) -> "AsyncLlamaStackAsLibraryClient":
+        """Enter the async context manager.
+
+        Initializes the client and returns it.
+
+        Example:
+            async with AsyncLlamaStackAsLibraryClient("starter") as client:
+                response = await client.models.list()
+            # Client is automatically shut down here
+        """
+        await self.initialize()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the async context manager and shut down the client."""
+        await self.shutdown()
+
     async def request(
         self,
         cast_to: Any,

src/llama_stack/core/stack.py

@@ -654,7 +654,7 @@ def cb(task):
     async def shutdown(self):
         for impl in self.impls.values():
             impl_name = impl.__class__.__name__
-            logger.info(f"Shutting down {impl_name}")
+            logger.debug(f"Shutting down {impl_name}")
             try:
                 if hasattr(impl, "shutdown"):
                     await asyncio.wait_for(impl.shutdown(), timeout=5)
@@ -676,6 +676,20 @@ async def shutdown(self):
         if REGISTRY_REFRESH_TASK:
             REGISTRY_REFRESH_TASK.cancel()
 
+        # Shutdown storage backends
+        from llama_stack.core.storage.kvstore.kvstore import shutdown_kvstore_backends
+        from llama_stack.core.storage.sqlstore.sqlstore import shutdown_sqlstore_backends
+
+        try:
+            await shutdown_kvstore_backends()
+        except Exception as e:
+            logger.exception(f"Failed to shutdown KV store backends: {e}")
+
+        try:
+            await shutdown_sqlstore_backends()
+        except Exception as e:
+            logger.exception(f"Failed to shutdown SQL store backends: {e}")
+
 
 async def refresh_registry_once(impls: dict[Api, Any]):
     logger.debug("refreshing registry")

src/llama_stack/core/storage/kvstore/kvstore.py

@@ -62,6 +62,9 @@ async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
     async def delete(self, key: str) -> None:
         del self._store[key]
 
+    async def shutdown(self) -> None:
+        self._store.clear()
+
 
 _KVSTORE_BACKENDS: dict[str, KVStoreConfig] = {}
 _KVSTORE_INSTANCES: dict[tuple[str, str], KVStore] = {}
@@ -126,3 +129,11 @@ async def kvstore_impl(reference: KVStoreReference) -> KVStore:
         await impl.initialize()
         _KVSTORE_INSTANCES[cache_key] = impl
         return impl
+
+
+async def shutdown_kvstore_backends() -> None:
+    """Shutdown all cached KV store instances."""
+    global _KVSTORE_INSTANCES
+    for instance in _KVSTORE_INSTANCES.values():
+        await instance.shutdown()
+    _KVSTORE_INSTANCES.clear()

src/llama_stack/core/storage/kvstore/mongodb/mongodb.py

@@ -83,3 +83,8 @@ async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
         async for doc in cursor:
             result.append(doc["key"])
         return result
+
+    async def shutdown(self) -> None:
+        if self.conn:
+            await self.conn.close()
+            self.conn = None

src/llama_stack/core/storage/kvstore/postgres/postgres.py

@@ -123,3 +123,11 @@ async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
             (start_key, end_key),
         )
         return [row[0] for row in cursor.fetchall()]
+
+    async def shutdown(self) -> None:
+        if self._cursor:
+            self._cursor.close()
+            self._cursor = None
+        if self._conn:
+            self._conn.close()
+            self._conn = None

src/llama_stack/core/storage/kvstore/redis/redis.py

@@ -99,3 +99,8 @@ async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
             if cursor == 0:
                 break
         return result
+
+    async def shutdown(self) -> None:
+        if self._redis:
+            await self._redis.close()
+            self._redis = None

src/llama_stack/core/storage/sqlstore/sqlalchemy_sqlstore.py

@@ -107,6 +107,14 @@ def set_sqlite_pragma(dbapi_conn, connection_record):
 
         return engine
 
+    async def shutdown(self) -> None:
+        """Dispose the session maker's engine and close all connections."""
+        # The async_session holds a reference to the engine created in __init__
+        if self.async_session:
+            engine = self.async_session.kw.get("bind")
+            if engine:
+                await engine.dispose()
+
     async def create_table(
         self,
         table: str,

src/llama_stack/core/storage/sqlstore/sqlstore.py

@@ -85,3 +85,11 @@ def register_sqlstore_backends(backends: dict[str, StorageBackendConfig]) -> Non
     _SQLSTORE_LOCKS.clear()
     for name, cfg in backends.items():
         _SQLSTORE_BACKENDS[name] = cfg
+
+
+async def shutdown_sqlstore_backends() -> None:
+    """Shutdown all cached SQL store instances."""
+    global _SQLSTORE_INSTANCES
+    for instance in _SQLSTORE_INSTANCES.values():
+        await instance.shutdown()
+    _SQLSTORE_INSTANCES.clear()

src/llama_stack_api/internal/kvstore.py

@@ -22,5 +22,7 @@ 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 shutdown(self) -> None: ...
+
 
 __all__ = ["KVStore"]