Merge branch 'main' into add-nvidia-inference-adapter

llama_stack/providers/remote/inference/fireworks/config.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
@@ -20,3 +20,10 @@ class FireworksImplConfig(BaseModel):
         default=None,
         description="The Fireworks.ai API Key",
     )
+
+    @classmethod
+    def sample_run_config(cls) -> Dict[str, Any]:
+        return {
+            "url": "https://api.fireworks.ai/inference",
+            "api_key": "${env.FIREWORKS_API_KEY}",
+        }

llama_stack/providers/remote/inference/fireworks/fireworks.py

@@ -35,7 +35,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
 from .config import FireworksImplConfig
 
 
-model_aliases = [
+MODEL_ALIASES = [
     build_model_alias(
         "fireworks/llama-v3p1-8b-instruct",
         CoreModelId.llama3_1_8b_instruct.value,
@@ -79,7 +79,7 @@ class FireworksInferenceAdapter(
     ModelRegistryHelper, Inference, NeedsRequestProviderData
 ):
     def __init__(self, config: FireworksImplConfig) -> None:
-        ModelRegistryHelper.__init__(self, model_aliases)
+        ModelRegistryHelper.__init__(self, MODEL_ALIASES)
         self.config = config
         self.formatter = ChatFormat(Tokenizer.get_instance())
 

llama_stack/providers/remote/inference/nvidia/_nvidia.py

@@ -30,7 +30,7 @@ from llama_stack.apis.inference import (
     ResponseFormat,
 )
 from llama_stack.providers.utils.inference.model_registry import (
-    build_model_alias,
+    build_model_alias_with_just_provider_model_id,
     ModelRegistryHelper,
 )
 
@@ -43,39 +43,39 @@ from ._openai_utils import (
 from ._utils import check_health
 
 _MODEL_ALIASES = [
-    build_model_alias(
+    build_model_alias_with_just_provider_model_id(
         "meta/llama3-8b-instruct",
         CoreModelId.llama3_8b_instruct.value,
     ),
-    build_model_alias(
+    build_model_alias_with_just_provider_model_id(
         "meta/llama3-70b-instruct",
         CoreModelId.llama3_70b_instruct.value,
     ),
-    build_model_alias(
+    build_model_alias_with_just_provider_model_id(
         "meta/llama-3.1-8b-instruct",
         CoreModelId.llama3_1_8b_instruct.value,
     ),
-    build_model_alias(
+    build_model_alias_with_just_provider_model_id(
         "meta/llama-3.1-70b-instruct",
         CoreModelId.llama3_1_70b_instruct.value,
     ),
-    build_model_alias(
+    build_model_alias_with_just_provider_model_id(
         "meta/llama-3.1-405b-instruct",
         CoreModelId.llama3_1_405b_instruct.value,
     ),
-    build_model_alias(
+    build_model_alias_with_just_provider_model_id(
         "meta/llama-3.2-1b-instruct",
         CoreModelId.llama3_2_1b_instruct.value,
     ),
-    build_model_alias(
+    build_model_alias_with_just_provider_model_id(
         "meta/llama-3.2-3b-instruct",
         CoreModelId.llama3_2_3b_instruct.value,
     ),
-    build_model_alias(
+    build_model_alias_with_just_provider_model_id(
         "meta/llama-3.2-11b-vision-instruct",
         CoreModelId.llama3_2_11b_vision_instruct.value,
     ),
-    build_model_alias(
+    build_model_alias_with_just_provider_model_id(
         "meta/llama-3.2-90b-vision-instruct",
         CoreModelId.llama3_2_90b_vision_instruct.value,
     ),

llama_stack/providers/remote/inference/ollama/__init__.py

@@ -4,14 +4,10 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_stack.distribution.datatypes import RemoteProviderConfig
+from .config import OllamaImplConfig
 
 
-class OllamaImplConfig(RemoteProviderConfig):
-    port: int = 11434
-
-
-async def get_adapter_impl(config: RemoteProviderConfig, _deps):
+async def get_adapter_impl(config: OllamaImplConfig, _deps):
     from .ollama import OllamaInferenceAdapter
 
     impl = OllamaInferenceAdapter(config.url)

llama_stack/providers/remote/inference/ollama/config.py

@@ -0,0 +1,22 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any, Dict
+
+from pydantic import BaseModel
+
+
+DEFAULT_OLLAMA_URL = "http://localhost:11434"
+
+
+class OllamaImplConfig(BaseModel):
+    url: str = DEFAULT_OLLAMA_URL
+
+    @classmethod
+    def sample_run_config(
+        cls, url: str = "${env.OLLAMA_URL:http://localhost:11434}", **kwargs
+    ) -> Dict[str, Any]:
+        return {"url": url}

llama_stack/providers/remote/inference/ollama/ollama.py

@@ -16,6 +16,7 @@ from ollama import AsyncClient
 
 from llama_stack.providers.utils.inference.model_registry import (
     build_model_alias,
+    build_model_alias_with_just_provider_model_id,
     ModelRegistryHelper,
 )
 
@@ -44,10 +45,18 @@ model_aliases = [
         "llama3.1:8b-instruct-fp16",
         CoreModelId.llama3_1_8b_instruct.value,
     ),
+    build_model_alias_with_just_provider_model_id(
+        "llama3.1:8b",
+        CoreModelId.llama3_1_8b_instruct.value,
+    ),
     build_model_alias(
         "llama3.1:70b-instruct-fp16",
         CoreModelId.llama3_1_70b_instruct.value,
     ),
+    build_model_alias_with_just_provider_model_id(
+        "llama3.1:70b",
+        CoreModelId.llama3_1_70b_instruct.value,
+    ),
     build_model_alias(
         "llama3.2:1b-instruct-fp16",
         CoreModelId.llama3_2_1b_instruct.value,
@@ -56,6 +65,24 @@ model_aliases = [
         "llama3.2:3b-instruct-fp16",
         CoreModelId.llama3_2_3b_instruct.value,
     ),
+    build_model_alias_with_just_provider_model_id(
+        "llama3.2:1b",
+        CoreModelId.llama3_2_1b_instruct.value,
+    ),
+    build_model_alias_with_just_provider_model_id(
+        "llama3.2:3b",
+        CoreModelId.llama3_2_3b_instruct.value,
+    ),
+    build_model_alias(
+        "llama3.2-vision:11b-instruct-fp16",
+        CoreModelId.llama3_2_11b_vision_instruct.value,
+    ),
+    build_model_alias_with_just_provider_model_id(
+        "llama3.2-vision",
+        CoreModelId.llama3_2_11b_vision_instruct.value,
+    ),
+    # The Llama Guard models don't have their full fp16 versions
+    # so we are going to alias their default version to the canonical SKU
     build_model_alias(
         "llama-guard3:8b",
         CoreModelId.llama_guard_3_8b.value,
@@ -64,10 +91,6 @@ model_aliases = [
         "llama-guard3:1b",
         CoreModelId.llama_guard_3_1b.value,
     ),
-    build_model_alias(
-        "x/llama3.2-vision:11b-instruct-fp16",
-        CoreModelId.llama3_2_11b_vision_instruct.value,
-    ),
 ]
 
 
@@ -82,7 +105,7 @@ class OllamaInferenceAdapter(Inference, ModelsProtocolPrivate):
         return AsyncClient(host=self.url)
 
     async def initialize(self) -> None:
-        print("Initializing Ollama, checking connectivity to server...")
+        print(f"checking connectivity to Ollama at `{self.url}`...")
         try:
             await self.client.ps()
         except httpx.ConnectError as e:

llama_stack/providers/remote/inference/tgi/config.py

@@ -12,19 +12,20 @@ from pydantic import BaseModel, Field
 
 @json_schema_type
 class TGIImplConfig(BaseModel):
-    host: str = "localhost"
-    port: int = 8080
-    protocol: str = "http"
-
-    @property
-    def url(self) -> str:
-        return f"{self.protocol}://{self.host}:{self.port}"
-
+    url: str = Field(
+        description="The URL for the TGI serving endpoint",
+    )
     api_token: Optional[str] = Field(
         default=None,
         description="A bearer token if your TGI endpoint is protected.",
     )
 
+    @classmethod
+    def sample_run_config(cls, url: str = "${env.TGI_URL}", **kwargs):
+        return {
+            "url": url,
+        }
+
 
 @json_schema_type
 class InferenceEndpointImplConfig(BaseModel):

llama_stack/providers/remote/inference/together/config.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
@@ -20,3 +20,10 @@ class TogetherImplConfig(BaseModel):
         default=None,
         description="The Together AI API Key",
     )
+
+    @classmethod
+    def sample_run_config(cls, **kwargs) -> Dict[str, Any]:
+        return {
+            "url": "https://api.together.xyz/v1",
+            "api_key": "${env.TOGETHER_API_KEY}",
+        }

llama_stack/providers/remote/inference/together/together.py

@@ -38,7 +38,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
 from .config import TogetherImplConfig
 
 
-model_aliases = [
+MODEL_ALIASES = [
     build_model_alias(
         "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
         CoreModelId.llama3_1_8b_instruct.value,
@@ -78,7 +78,7 @@ class TogetherInferenceAdapter(
     ModelRegistryHelper, Inference, NeedsRequestProviderData
 ):
     def __init__(self, config: TogetherImplConfig) -> None:
-        ModelRegistryHelper.__init__(self, model_aliases)
+        ModelRegistryHelper.__init__(self, MODEL_ALIASES)
         self.config = config
         self.formatter = ChatFormat(Tokenizer.get_instance())
 

llama_stack/providers/remote/inference/vllm/config.py

@@ -24,3 +24,15 @@ class VLLMInferenceAdapterConfig(BaseModel):
         default="fake",
         description="The API token",
     )
+
+    @classmethod
+    def sample_run_config(
+        cls,
+        url: str = "${env.VLLM_URL}",
+        **kwargs,
+    ):
+        return {
+            "url": url,
+            "max_tokens": "${env.VLLM_MAX_TOKENS:4096}",
+            "api_token": "${env.VLLM_API_TOKEN:fake}",
+        }