feat(starter)!: simplify starter distro; litellm model registry changes (#2916)

llama_stack/distribution/routing_tables/models.py

@@ -25,7 +25,8 @@ class ModelsRoutingTable(CommonRoutingTableImpl, Models):
     async def refresh(self) -> None:
         for provider_id, provider in self.impls_by_provider_id.items():
             refresh = await provider.should_refresh_models()
-            if not (refresh or provider_id in self.listed_providers):
+            refresh = refresh or provider_id not in self.listed_providers
+            if not refresh:
                 continue
 
             try:
@@ -138,6 +139,9 @@ class ModelsRoutingTable(CommonRoutingTableImpl, Models):
                 # avoid overwriting a non-provider-registered model entry
                 continue
 
+            if model.identifier == model.provider_resource_id:
+                model.identifier = f"{provider_id}/{model.provider_resource_id}"
+
             logger.debug(f"registering model {model.identifier} ({model.provider_resource_id})")
             await self.register_object(
                 ModelWithOwner(

llama_stack/distribution/server/server.py

@@ -611,11 +611,8 @@ def extract_path_params(route: str) -> list[str]:
 
 def remove_disabled_providers(obj):
     if isinstance(obj, dict):
-        if (
-            obj.get("provider_id") == "__disabled__"
-            or obj.get("shield_id") == "__disabled__"
-            or obj.get("provider_model_id") == "__disabled__"
-        ):
+        keys = ["provider_id", "shield_id", "provider_model_id", "model_id"]
+        if any(k in obj and obj[k] in ("__disabled__", "", None) for k in keys):
             return None
         return {k: v for k, v in ((k, remove_disabled_providers(v)) for k, v in obj.items()) if v is not None}
     elif isinstance(obj, list):

llama_stack/distribution/stack.py

@@ -105,23 +105,10 @@ async def register_resources(run_config: StackRunConfig, impls: dict[Api, Any]):
         method = getattr(impls[api], register_method)
         for obj in objects:
             logger.debug(f"registering {rsrc.capitalize()} {obj} for provider {obj.provider_id}")
-            # Do not register models on disabled providers
-            if hasattr(obj, "provider_id") and obj.provider_id is not None and obj.provider_id == "__disabled__":
-                logger.debug(f"Skipping {rsrc.capitalize()} registration for disabled provider.")
-                continue
-            # In complex templates, like our starter template, we may have dynamic model ids
-            # given by environment variables. This allows those environment variables to have
-            # a default value of __disabled__ to skip registration of the model if not set.
-            if (
-                hasattr(obj, "provider_model_id")
-                and obj.provider_model_id is not None
-                and "__disabled__" in obj.provider_model_id
-            ):
-                logger.debug(f"Skipping {rsrc.capitalize()} registration for disabled model.")
-                continue
 
-            if hasattr(obj, "shield_id") and obj.shield_id is not None and obj.shield_id == "__disabled__":
-                logger.debug(f"Skipping {rsrc.capitalize()} registration for disabled shield.")
+            # Do not register models on disabled providers
+            if hasattr(obj, "provider_id") and (not obj.provider_id or obj.provider_id == "__disabled__"):
+                logger.debug(f"Skipping {rsrc.capitalize()} registration for disabled provider.")
                 continue
 
             # we want to maintain the type information in arguments to method.
@@ -331,8 +318,10 @@ async def construct_stack(
 
     await register_resources(run_config, impls)
 
+    await refresh_registry_once(impls)
+
     global REGISTRY_REFRESH_TASK
-    REGISTRY_REFRESH_TASK = asyncio.create_task(refresh_registry(impls))
+    REGISTRY_REFRESH_TASK = asyncio.create_task(refresh_registry_task(impls))
 
     def cb(task):
         import traceback
@@ -368,11 +357,17 @@ async def shutdown_stack(impls: dict[Api, Any]):
         REGISTRY_REFRESH_TASK.cancel()
 
 
-async def refresh_registry(impls: dict[Api, Any]):
+async def refresh_registry_once(impls: dict[Api, Any]):
+    logger.info("refreshing registry")
     routing_tables = [v for v in impls.values() if isinstance(v, CommonRoutingTableImpl)]
+    for routing_table in routing_tables:
+        await routing_table.refresh()
+
+
+async def refresh_registry_task(impls: dict[Api, Any]):
+    logger.info("starting registry refresh task")
     while True:
-        for routing_table in routing_tables:
-            await routing_table.refresh()
+        await refresh_registry_once(impls)
 
         await asyncio.sleep(REGISTRY_REFRESH_INTERVAL_SECONDS)