From 3492c9c8d1c9c561e42f36988b2a463e94729c8c Mon Sep 17 00:00:00 2001
From: Omar Abdelwahab <omara@fb.com>
Date: Mon, 15 Sep 2025 16:47:35 -0700
Subject: [PATCH] Added a bug fix when registering new models

---
 llama_stack/core/store/registry.py | 51 ++++++++++++++++++++++--------
 1 file changed, 38 insertions(+), 13 deletions(-)

diff --git a/llama_stack/core/store/registry.py b/llama_stack/core/store/registry.py
index 5f4abe9aa..ddb863349 100644
--- a/llama_stack/core/store/registry.py
+++ b/llama_stack/core/store/registry.py
@@ -14,7 +14,10 @@ from llama_stack.core.datatypes import RoutableObjectWithProvider
 from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl
-from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig
+from llama_stack.providers.utils.kvstore.config import (
+    KVStoreConfig,
+    SqliteKVStoreConfig,
+)
 
 logger = get_logger(__name__, category="core::registry")
 
@@ -28,7 +31,9 @@ class DistributionRegistry(Protocol):
 
     def get_cached(self, identifier: str) -> RoutableObjectWithProvider | None: ...
 
-    async def update(self, obj: RoutableObjectWithProvider) -> RoutableObjectWithProvider: ...
+    async def update(
+        self, obj: RoutableObjectWithProvider
+    ) -> RoutableObjectWithProvider: ...
 
     async def register(self, obj: RoutableObjectWithProvider) -> bool: ...
 
@@ -54,7 +59,9 @@ def _parse_registry_values(values: list[str]) -> list[RoutableObjectWithProvider
             obj = pydantic.TypeAdapter(RoutableObjectWithProvider).validate_json(value)
             all_objects.append(obj)
         except pydantic.ValidationError as e:
-            logger.error(f"Error parsing registry value, raw value: {value}. Error: {e}")
+            logger.error(
+                f"Error parsing registry value, raw value: {value}. Error: {e}"
+            )
             continue
 
     return all_objects
@@ -67,7 +74,9 @@ class DiskDistributionRegistry(DistributionRegistry):
     async def initialize(self) -> None:
         pass
 
-    def get_cached(self, type: str, identifier: str) -> RoutableObjectWithProvider | None:
+    def get_cached(
+        self, type: str, identifier: str
+    ) -> RoutableObjectWithProvider | None:
         # Disk registry does not have a cache
         raise NotImplementedError("Disk registry does not have a cache")
 
@@ -76,15 +85,23 @@ class DiskDistributionRegistry(DistributionRegistry):
         values = await self.kvstore.values_in_range(start_key, end_key)
         return _parse_registry_values(values)
 
-    async def get(self, type: str, identifier: str) -> RoutableObjectWithProvider | None:
-        json_str = await self.kvstore.get(KEY_FORMAT.format(type=type, identifier=identifier))
+    async def get(
+        self, type: str, identifier: str
+    ) -> RoutableObjectWithProvider | None:
+        json_str = await self.kvstore.get(
+            KEY_FORMAT.format(type=type, identifier=identifier)
+        )
         if not json_str:
             return None
 
         try:
-            return pydantic.TypeAdapter(RoutableObjectWithProvider).validate_json(json_str)
+            return pydantic.TypeAdapter(RoutableObjectWithProvider).validate_json(
+                json_str
+            )
         except pydantic.ValidationError as e:
-            logger.error(f"Error parsing registry value for {type}:{identifier}, raw value: {json_str}. Error: {e}")
+            logger.error(
+                f"Error parsing registry value for {type}:{identifier}, raw value: {json_str}. Error: {e}"
+            )
             return None
 
     async def update(self, obj: RoutableObjectWithProvider) -> None:
@@ -96,9 +113,11 @@ class DiskDistributionRegistry(DistributionRegistry):
 
     async def register(self, obj: RoutableObjectWithProvider) -> bool:
         existing_obj = await self.get(obj.type, obj.identifier)
-        # dont register if the object's providerid already exists
+        # warn if the object's providerid already exists but proceed with registration
         if existing_obj and existing_obj.provider_id == obj.provider_id:
-            return False
+            logger.warning(
+                f"Object {obj.type}:{obj.identifier} with provider_id {obj.provider_id} already exists, overwriting"
+            )
 
         await self.kvstore.set(
             KEY_FORMAT.format(type=obj.type, identifier=obj.identifier),
@@ -147,7 +166,9 @@ class CachedDiskDistributionRegistry(DiskDistributionRegistry):
     async def initialize(self) -> None:
         await self._ensure_initialized()
 
-    def get_cached(self, type: str, identifier: str) -> RoutableObjectWithProvider | None:
+    def get_cached(
+        self, type: str, identifier: str
+    ) -> RoutableObjectWithProvider | None:
         return self.cache.get((type, identifier), None)
 
     async def get_all(self) -> list[RoutableObjectWithProvider]:
@@ -155,7 +176,9 @@ class CachedDiskDistributionRegistry(DiskDistributionRegistry):
         async with self._locked_cache() as cache:
             return list(cache.values())
 
-    async def get(self, type: str, identifier: str) -> RoutableObjectWithProvider | None:
+    async def get(
+        self, type: str, identifier: str
+    ) -> RoutableObjectWithProvider | None:
         await self._ensure_initialized()
         cache_key = (type, identifier)
 
@@ -197,7 +220,9 @@ async def create_dist_registry(
         dist_kvstore = await kvstore_impl(metadata_store)
     else:
         dist_kvstore = await kvstore_impl(
-            SqliteKVStoreConfig(db_path=(DISTRIBS_BASE_DIR / image_name / "kvstore.db").as_posix())
+            SqliteKVStoreConfig(
+                db_path=(DISTRIBS_BASE_DIR / image_name / "kvstore.db").as_posix()
+            )
         )
     dist_registry = CachedDiskDistributionRegistry(dist_kvstore)
     await dist_registry.initialize()