Repository: meta-llama/llama-stack
Pull request #1859: feat: NVIDIA allow non-llama model registration
Author: login: raspawar
Labels: name: CLA Signed
description: This label is managed by the Meta Open Source bot.
Created: 2025-04-02T10:05:41Z | Merged: 2025-04-25T00:13:33Z

# What does this PR do?
Adds custom model registration functionality to NVIDIAInferenceAdapter which let's the inference happen on:
- post-training model
- non-llama models in API Catalogue(behind https://integrate.api.nvidia.com and endpoints compatible with AyncOpenAI)

## Example Usage:
```python
from llama_stack.apis.models import Model, ModelType
from llama_stack.distribution.library_client import LlamaStackAsLibraryClient
client = LlamaStackAsLibraryClient("nvidia")
_ = client.initialize()

client.models.register(
        model_id=model_name,
        model_type=ModelType.llm,
        provider_id="nvidia"
)

response = client.inference.chat_completion(
    model_id=model_name,
    messages=[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"Write a limerick about the wonders of GPU computing."}],
)
```

## Test Plan
```bash
pytest tests/unit/providers/nvidia/test_supervised_fine_tuning.py 
========================================================== test session starts ===========================================================
platform linux -- Python 3.10.0, pytest-8.3.5, pluggy-1.5.0
rootdir: /home/ubuntu/llama-stack
configfile: pyproject.toml
plugins: anyio-4.9.0
collected 6 items                                                                                                                        

tests/unit/providers/nvidia/test_supervised_fine_tuning.py ......                                                                  [100%]

============================================================ warnings summary ============================================================
../miniconda/envs/nvidia-1/lib/python3.10/site-packages/pydantic/fields.py:1076
  /home/ubuntu/miniconda/envs/nvidia-1/lib/python3.10/site-packages/pydantic/fields.py:1076: PydanticDeprecatedSince20: Using extra keyword arguments on `Field` is deprecated and will be removed. Use `json_schema_extra` instead. (Extra keys: 'contentEncoding'). Deprecated in Pydantic V2.0 to be removed in V3.0. See Pydantic V2 Migration Guide at https://errors.pydantic.dev/2.11/migration/
    warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
====================================================== 6 passed, 1 warning in 1.51s ======================================================
```

[//]: # (## Documentation)
Updated Readme.md

cc: @dglogo, @sumitb, @mattf



Diff:
diff --git a/docs/source/distributions/self_hosted_distro/nvidia.md b/docs/source/distributions/self_hosted_distro/nvidia.md
index 0922cb5127..c6c92f5d21 100644
--- a/docs/source/distributions/self_hosted_distro/nvidia.md
+++ b/docs/source/distributions/self_hosted_distro/nvidia.md
@@ -22,9 +22,8 @@ The `llamastack/distribution-nvidia` distribution consists of the following prov
 The following environment variables can be configured:
 
 - `NVIDIA_API_KEY`: NVIDIA API Key (default: ``)
-- `NVIDIA_USER_ID`: NVIDIA User ID (default: `llama-stack-user`)
+- `NVIDIA_APPEND_API_VERSION`: Whether to append the API version to the base_url (default: `True`)
 - `NVIDIA_DATASET_NAMESPACE`: NVIDIA Dataset Namespace (default: `default`)
-- `NVIDIA_ACCESS_POLICIES`: NVIDIA Access Policies (default: `{}`)
 - `NVIDIA_PROJECT_ID`: NVIDIA Project ID (default: `test-project`)
 - `NVIDIA_CUSTOMIZER_URL`: NVIDIA Customizer URL (default: `https://customizer.api.nvidia.com`)
 - `NVIDIA_OUTPUT_MODEL_DIR`: NVIDIA Output Model Directory (default: `test-example-model@v1`)
diff --git a/llama_stack/providers/remote/inference/nvidia/config.py b/llama_stack/providers/remote/inference/nvidia/config.py
index abd34b498c..8f80408d45 100644
--- a/llama_stack/providers/remote/inference/nvidia/config.py
+++ b/llama_stack/providers/remote/inference/nvidia/config.py
@@ -47,10 +47,15 @@ class NVIDIAConfig(BaseModel):
         default=60,
         description="Timeout for the HTTP requests",
     )
+    append_api_version: bool = Field(
+        default_factory=lambda: os.getenv("NVIDIA_APPEND_API_VERSION", "True").lower() != "false",
+        description="When set to false, the API version will not be appended to the base_url. By default, it is true.",
+    )
 
     @classmethod
     def sample_run_config(cls, **kwargs) -> Dict[str, Any]:
         return {
             "url": "${env.NVIDIA_BASE_URL:https://integrate.api.nvidia.com}",
             "api_key": "${env.NVIDIA_API_KEY:}",
+            "append_api_version": "${env.NVIDIA_APPEND_API_VERSION:True}",
         }
diff --git a/llama_stack/providers/remote/inference/nvidia/nvidia.py b/llama_stack/providers/remote/inference/nvidia/nvidia.py
index c91b4d7688..4a62ad6cb3 100644
--- a/llama_stack/providers/remote/inference/nvidia/nvidia.py
+++ b/llama_stack/providers/remote/inference/nvidia/nvidia.py
@@ -33,7 +33,6 @@
     TextTruncation,
     ToolChoice,
     ToolConfig,
-    ToolDefinition,
 )
 from llama_stack.apis.inference.inference import (
     OpenAIChatCompletion,
@@ -42,7 +41,11 @@
     OpenAIMessageParam,
     OpenAIResponseFormatParam,
 )
-from llama_stack.models.llama.datatypes import ToolPromptFormat
+from llama_stack.apis.models import Model, ModelType
+from llama_stack.models.llama.datatypes import ToolDefinition, ToolPromptFormat
+from llama_stack.providers.utils.inference import (
+    ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR,
+)
 from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
 )
@@ -120,10 +123,10 @@ def _get_client_for_base_url(base_url: str) -> AsyncOpenAI:
             "meta/llama-3.2-90b-vision-instruct": "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct",
         }
 
-        base_url = f"{self._config.url}/v1"
+        base_url = f"{self._config.url}/v1" if self._config.append_api_version else self._config.url
+
         if _is_nvidia_hosted(self._config) and provider_model_id in special_model_urls:
             base_url = special_model_urls[provider_model_id]
-
         return _get_client_for_base_url(base_url)
 
     async def _get_provider_model_id(self, model_id: str) -> str:
@@ -387,3 +390,44 @@ async def openai_chat_completion(
             return await self._get_client(provider_model_id).chat.completions.create(**params)
         except APIConnectionError as e:
             raise ConnectionError(f"Failed to connect to NVIDIA NIM at {self._config.url}: {e}") from e
+
+    async def register_model(self, model: Model) -> Model:
+        """
+        Allow non-llama model registration.
+
+        Non-llama model registration: API Catalogue models, post-training models, etc.
+            client = LlamaStackAsLibraryClient("nvidia")
+            client.models.register(
+                    model_id="mistralai/mixtral-8x7b-instruct-v0.1",
+                    model_type=ModelType.llm,
+                    provider_id="nvidia",
+                    provider_model_id="mistralai/mixtral-8x7b-instruct-v0.1"
+            )
+
+            NOTE: Only supports models endpoints compatible with AsyncOpenAI base_url format.
+        """
+        if model.model_type == ModelType.embedding:
+            # embedding models are always registered by their provider model id and does not need to be mapped to a llama model
+            provider_resource_id = model.provider_resource_id
+        else:
+            provider_resource_id = self.get_provider_model_id(model.provider_resource_id)
+
+        if provider_resource_id:
+            model.provider_resource_id = provider_resource_id
+        else:
+            llama_model = model.metadata.get("llama_model")
+            existing_llama_model = self.get_llama_model(model.provider_resource_id)
+            if existing_llama_model:
+                if existing_llama_model != llama_model:
+                    raise ValueError(
+                        f"Provider model id '{model.provider_resource_id}' is already registered to a different llama model: '{existing_llama_model}'"
+                    )
+            else:
+                # not llama model
+                if llama_model in ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR:
+                    self.provider_id_to_llama_model_map[model.provider_resource_id] = (
+                        ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR[llama_model]
+                    )
+                else:
+                    self.alias_to_provider_id_map[model.provider_model_id] = model.provider_model_id
+        return model
diff --git a/llama_stack/providers/remote/post_training/nvidia/README.md b/llama_stack/providers/remote/post_training/nvidia/README.md
index 230587d66b..3ef538d293 100644
--- a/llama_stack/providers/remote/post_training/nvidia/README.md
+++ b/llama_stack/providers/remote/post_training/nvidia/README.md
@@ -36,7 +36,6 @@ import os
 
 os.environ["NVIDIA_API_KEY"] = "your-api-key"
 os.environ["NVIDIA_CUSTOMIZER_URL"] = "http://nemo.test"
-os.environ["NVIDIA_USER_ID"] = "llama-stack-user"
 os.environ["NVIDIA_DATASET_NAMESPACE"] = "default"
 os.environ["NVIDIA_PROJECT_ID"] = "test-project"
 os.environ["NVIDIA_OUTPUT_MODEL_DIR"] = "test-example-model@v1"
@@ -125,6 +124,21 @@ client.post_training.job.cancel(job_uuid="your-job-id")
 
 ### Inference with the fine-tuned model
 
+#### 1. Register the model
+
+```python
+from llama_stack.apis.models import Model, ModelType
+
+client.models.register(
+    model_id="test-example-model@v1",
+    provider_id="nvidia",
+    provider_model_id="test-example-model@v1",
+    model_type=ModelType.llm,
+)
+```
+
+#### 2. Inference with the fine-tuned model
+
 ```python
 response = client.inference.completion(
     content="Complete the sentence using one word: Roses are red, violets are ",
diff --git a/llama_stack/templates/nvidia/nvidia.py b/llama_stack/templates/nvidia/nvidia.py
index a0cefba52f..46a5ad2c42 100644
--- a/llama_stack/templates/nvidia/nvidia.py
+++ b/llama_stack/templates/nvidia/nvidia.py
@@ -90,19 +90,15 @@ def get_distribution_template() -> DistributionTemplate:
                 "",
                 "NVIDIA API Key",
             ),
-            ## Nemo Customizer related variables
-            "NVIDIA_USER_ID": (
-                "llama-stack-user",
-                "NVIDIA User ID",
+            "NVIDIA_APPEND_API_VERSION": (
+                "True",
+                "Whether to append the API version to the base_url",
             ),
+            ## Nemo Customizer related variables
             "NVIDIA_DATASET_NAMESPACE": (
                 "default",
                 "NVIDIA Dataset Namespace",
             ),
-            "NVIDIA_ACCESS_POLICIES": (
-                "{}",
-                "NVIDIA Access Policies",
-            ),
             "NVIDIA_PROJECT_ID": (
                 "test-project",
                 "NVIDIA Project ID",
diff --git a/llama_stack/templates/nvidia/run-with-safety.yaml b/llama_stack/templates/nvidia/run-with-safety.yaml
index 658d9377eb..2f7a62ce04 100644
--- a/llama_stack/templates/nvidia/run-with-safety.yaml
+++ b/llama_stack/templates/nvidia/run-with-safety.yaml
@@ -18,6 +18,7 @@ providers:
     config:
       url: ${env.NVIDIA_BASE_URL:https://integrate.api.nvidia.com}
       api_key: ${env.NVIDIA_API_KEY:}
+      append_api_version: ${env.NVIDIA_APPEND_API_VERSION:True}
   - provider_id: nvidia
     provider_type: remote::nvidia
     config:
diff --git a/llama_stack/templates/nvidia/run.yaml b/llama_stack/templates/nvidia/run.yaml
index ff548d82ec..bc5dec8e4f 100644
--- a/llama_stack/templates/nvidia/run.yaml
+++ b/llama_stack/templates/nvidia/run.yaml
@@ -18,6 +18,7 @@ providers:
     config:
       url: ${env.NVIDIA_BASE_URL:https://integrate.api.nvidia.com}
       api_key: ${env.NVIDIA_API_KEY:}
+      append_api_version: ${env.NVIDIA_APPEND_API_VERSION:True}
   vector_io:
   - provider_id: faiss
     provider_type: inline::faiss
diff --git a/tests/unit/providers/nvidia/test_supervised_fine_tuning.py b/tests/unit/providers/nvidia/test_supervised_fine_tuning.py
index 43e0ac11cc..09f67e4e68 100644
--- a/tests/unit/providers/nvidia/test_supervised_fine_tuning.py
+++ b/tests/unit/providers/nvidia/test_supervised_fine_tuning.py
@@ -17,6 +17,8 @@
     TrainingConfigOptimizerConfig,
 )
 
+from llama_stack.apis.models import Model, ModelType
+from llama_stack.providers.remote.inference.nvidia.nvidia import NVIDIAConfig, NVIDIAInferenceAdapter
 from llama_stack.providers.remote.post_training.nvidia.post_training import (
     ListNvidiaPostTrainingJobs,
     NvidiaPostTrainingAdapter,
@@ -40,8 +42,22 @@ def setUp(self):
         )
         self.mock_make_request = self.make_request_patcher.start()
 
+        # Mock the inference client
+        inference_config = NVIDIAConfig(base_url=os.environ["NVIDIA_BASE_URL"], api_key=None)
+        self.inference_adapter = NVIDIAInferenceAdapter(inference_config)
+
+        self.mock_client = unittest.mock.MagicMock()
+        self.mock_client.chat.completions.create = unittest.mock.AsyncMock()
+        self.inference_mock_make_request = self.mock_client.chat.completions.create
+        self.inference_make_request_patcher = patch(
+            "llama_stack.providers.remote.inference.nvidia.nvidia.NVIDIAInferenceAdapter._get_client",
+            return_value=self.mock_client,
+        )
+        self.inference_make_request_patcher.start()
+
     def tearDown(self):
         self.make_request_patcher.stop()
+        self.inference_make_request_patcher.stop()
 
     @pytest.fixture(autouse=True)
     def inject_fixtures(self, run_async):
@@ -303,6 +319,31 @@ def test_cancel_training_job(self):
             expected_params={"job_id": job_id},
         )
 
+    def test_inference_register_model(self):
+        model_id = "default/job-1234"
+        model_type = ModelType.llm
+        model = Model(
+            identifier=model_id,
+            provider_id="nvidia",
+            provider_model_id=model_id,
+            provider_resource_id=model_id,
+            model_type=model_type,
+        )
+        result = self.run_async(self.inference_adapter.register_model(model))
+        assert result == model
+        assert len(self.inference_adapter.alias_to_provider_id_map) > 1
+        assert self.inference_adapter.get_provider_model_id(model.provider_model_id) == model_id
+
+        with patch.object(self.inference_adapter, "chat_completion") as mock_chat_completion:
+            self.run_async(
+                self.inference_adapter.chat_completion(
+                    model_id=model_id,
+                    messages=[{"role": "user", "content": "Hello, model"}],
+                )
+            )
+
+            mock_chat_completion.assert_called()
+
 
 if __name__ == "__main__":
     unittest.main()


Reviews:
user: login: mattf
body: thanks for adding this. few comments inline for you.
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-02T15:29:34Z
user: login: raspawar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T17:39:04Z
user: login: raspawar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T17:44:24Z
user: login: mattf
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T13:16:31Z
user: login: raghotham
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-12T00:16:13Z
user: login: raspawar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-14T11:12:32Z
user: login: mattf
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T13:29:54Z
user: login: raspawar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-14T13:35:39Z
user: login: mattf
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T11:23:30Z
user: login: mattf
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-16T10:49:20Z
user: login: raspawar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-16T11:57:35Z
user: login: mattf
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-16T13:45:35Z
user: login: ashwinb
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-17T18:40:20Z
user: login: shethaadit
state: APPROVED
author_association: NONE
submitted_at: 2025-04-21T00:49:00Z

Review comments:
diff_hunk: @@ -114,10 +118,13 @@ def _get_client_for_base_url(base_url: str) -> AsyncOpenAI:
             "meta/llama-3.2-90b-vision-instruct": "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct",
         }
 
-        base_url = f"{self._config.url}/v1"
-        if _is_nvidia_hosted(self._config) and provider_model_id in special_model_urls:
-            base_url = special_model_urls[provider_model_id]
-
+        # add /v1 in case of hosted models
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: mattf
body: this is a behavior change.

current behavior: always add /v1
new behavior: add /v1 for hosted and don't for non-hosted

the behavior should be consistent between hosted and non-hosted. for instance, a user should not need to know they're talking to https://integrate.api.nv.c and therefore don't need to supply the /v1 or since they're talking to http://localhost that they do need to provide the /v1.

is there an issue w/ /v1 and customizer?
created_at: 2025-04-02T15:15:15Z
author_association: COLLABORATOR
diff_hunk: @@ -265,3 +272,44 @@ async def chat_completion(
         else:
             # we pass n=1 to get only one completion
             return convert_openai_chat_completion_choice(response.choices[0])
+
+    async def register_model(self, model: Model) -> Model:
+        """
+        Allow non-llama model registration.
+
+        Non-llama model registration: API Catalogue models, post-training models, etc.
+            client = LlamaStackAsLibraryClient("nvidia")
+            client.models.register(
+                    model_id="mistralai/mixtral-8x7b-instruct-v0.1",
+                    model_type=ModelType.llm,
+                    provider_id="nvidia",
+                    provider_model_id="mistralai/mixtral-8x7b-instruct-v0.1"
+            )
+
+            NOTE: Only supports models endpoints compatible with AsyncOpenAI base_url format.
+        """
+        if model.model_type == ModelType.embedding:
+            # embedding models are always registered by their provider model id and does not need to be mapped to a llama model
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: mattf
body: this should be handled within the provider model id function
created_at: 2025-04-02T15:18:44Z
author_association: COLLABORATOR
diff_hunk: @@ -265,3 +272,44 @@ async def chat_completion(
         else:
             # we pass n=1 to get only one completion
             return convert_openai_chat_completion_choice(response.choices[0])
+
+    async def register_model(self, model: Model) -> Model:
+        """
+        Allow non-llama model registration.
+
+        Non-llama model registration: API Catalogue models, post-training models, etc.
+            client = LlamaStackAsLibraryClient("nvidia")
+            client.models.register(
+                    model_id="mistralai/mixtral-8x7b-instruct-v0.1",
+                    model_type=ModelType.llm,
+                    provider_id="nvidia",
+                    provider_model_id="mistralai/mixtral-8x7b-instruct-v0.1"
+            )
+
+            NOTE: Only supports models endpoints compatible with AsyncOpenAI base_url format.
+        """
+        if model.model_type == ModelType.embedding:
+            # embedding models are always registered by their provider model id and does not need to be mapped to a llama model
+            provider_resource_id = model.provider_resource_id
+        else:
+            provider_resource_id = self.get_provider_model_id(model.provider_resource_id)
+
+        if provider_resource_id:
+            model.provider_resource_id = provider_resource_id
+        else:
+            llama_model = model.metadata.get("llama_model")
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: mattf
body: i believe model  is a https://github.com/meta-llama/llama-stack/blob/main/llama_stack/apis/models/models.py#L31, which does not have a metadata

https://github.com/meta-llama/llama-stack/blob/main/llama_stack/models/llama/datatypes.py#L346 has a metadata, and confusingly the same class name.

suggestion: trust the input config. it should fail at inference if it's incorrect.
created_at: 2025-04-02T15:29:06Z
author_association: COLLABORATOR
diff_hunk: @@ -114,10 +118,13 @@ def _get_client_for_base_url(base_url: str) -> AsyncOpenAI:
             "meta/llama-3.2-90b-vision-instruct": "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct",
         }
 
-        base_url = f"{self._config.url}/v1"
-        if _is_nvidia_hosted(self._config) and provider_model_id in special_model_urls:
-            base_url = special_model_urls[provider_model_id]
-
+        # add /v1 in case of hosted models
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: raspawar
body: Adding /v1 may produce errors, specially when user is specifying `NVIDIA_BASE_URL`
Can we remove the /v1 entirely? And add it in default base_url, what do we miss in that case?
As some models endpoints on API catalogue follow /chat/completion.
created_at: 2025-04-02T17:39:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -265,3 +272,44 @@ async def chat_completion(
         else:
             # we pass n=1 to get only one completion
             return convert_openai_chat_completion_choice(response.choices[0])
+
+    async def register_model(self, model: Model) -> Model:
+        """
+        Allow non-llama model registration.
+
+        Non-llama model registration: API Catalogue models, post-training models, etc.
+            client = LlamaStackAsLibraryClient("nvidia")
+            client.models.register(
+                    model_id="mistralai/mixtral-8x7b-instruct-v0.1",
+                    model_type=ModelType.llm,
+                    provider_id="nvidia",
+                    provider_model_id="mistralai/mixtral-8x7b-instruct-v0.1"
+            )
+
+            NOTE: Only supports models endpoints compatible with AsyncOpenAI base_url format.
+        """
+        if model.model_type == ModelType.embedding:
+            # embedding models are always registered by their provider model id and does not need to be mapped to a llama model
+            provider_resource_id = model.provider_resource_id
+        else:
+            provider_resource_id = self.get_provider_model_id(model.provider_resource_id)
+
+        if provider_resource_id:
+            model.provider_resource_id = provider_resource_id
+        else:
+            llama_model = model.metadata.get("llama_model")
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: raspawar
body: This is base register_model() logic: https://github.com/meta-llama/llama-stack/blob/main/llama_stack/providers/utils/inference/model_registry.py#L76

I modified only parts to allow non-llama models. 
created_at: 2025-04-02T17:44:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -114,10 +118,13 @@ def _get_client_for_base_url(base_url: str) -> AsyncOpenAI:
             "meta/llama-3.2-90b-vision-instruct": "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct",
         }
 
-        base_url = f"{self._config.url}/v1"
-        if _is_nvidia_hosted(self._config) and provider_model_id in special_model_urls:
-            base_url = special_model_urls[provider_model_id]
-
+        # add /v1 in case of hosted models
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: mattf
body: i'm advocating for a consistent user experience. it can be either w/ or w/o /v1.

the nvdia nim (local) requires .../v1/chat/completions /v1/completions /v1/embedding, nvidia nim (hosted) mostly matches that (excpetion is models not on integrate.api.nvidia.com).

users may deploy a proxy, resulting in .../a/b/c/v1/{chat/completions,completions,embedding}, or even w/o the v1 component

how does customizer expose the models?
created_at: 2025-04-11T13:16:31Z
author_association: COLLABORATOR
diff_hunk: @@ -114,10 +118,13 @@ def _get_client_for_base_url(base_url: str) -> AsyncOpenAI:
             "meta/llama-3.2-90b-vision-instruct": "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct",
         }
 
-        base_url = f"{self._config.url}/v1"
-        if _is_nvidia_hosted(self._config) and provider_model_id in special_model_urls:
-            base_url = special_model_urls[provider_model_id]
-
+        # add /v1 in case of hosted models
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: raghotham
body: @raspawar can you please respond to this?
created_at: 2025-04-12T00:16:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -114,10 +118,13 @@ def _get_client_for_base_url(base_url: str) -> AsyncOpenAI:
             "meta/llama-3.2-90b-vision-instruct": "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct",
         }
 
-        base_url = f"{self._config.url}/v1"
-        if _is_nvidia_hosted(self._config) and provider_model_id in special_model_urls:
-            base_url = special_model_urls[provider_model_id]
-
+        # add /v1 in case of hosted models
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: raspawar
body: Custmomizer exposes as /v1/completion which works perfectly fine with the given code.
However, for many nvdev endpoinds this was breaking as they have /chat/completion endoints.
created_at: 2025-04-14T11:12:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -114,10 +118,13 @@ def _get_client_for_base_url(base_url: str) -> AsyncOpenAI:
             "meta/llama-3.2-90b-vision-instruct": "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct",
         }
 
-        base_url = f"{self._config.url}/v1"
-        if _is_nvidia_hosted(self._config) and provider_model_id in special_model_urls:
-            base_url = special_model_urls[provider_model_id]
-
+        # add /v1 in case of hosted models
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: mattf
body: i recall nvdev being attached to the model name, not the path.

what are the paths that need to be supported?
created_at: 2025-04-14T13:29:54Z
author_association: COLLABORATOR
diff_hunk: @@ -114,10 +118,13 @@ def _get_client_for_base_url(base_url: str) -> AsyncOpenAI:
             "meta/llama-3.2-90b-vision-instruct": "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct",
         }
 
-        base_url = f"{self._config.url}/v1"
-        if _is_nvidia_hosted(self._config) and provider_model_id in special_model_urls:
-            base_url = special_model_urls[provider_model_id]
-
+        # add /v1 in case of hosted models
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: raspawar
body: The few model endpoints have format: - v1/*/{model_name}/chat/completions
created_at: 2025-04-14T13:35:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -114,10 +118,13 @@ def _get_client_for_base_url(base_url: str) -> AsyncOpenAI:
             "meta/llama-3.2-90b-vision-instruct": "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct",
         }
 
-        base_url = f"{self._config.url}/v1"
-        if _is_nvidia_hosted(self._config) and provider_model_id in special_model_urls:
-            base_url = special_model_urls[provider_model_id]
-
+        # add /v1 in case of hosted models
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: mattf
body: @raspawar let's go with the env override discussed on discord. the /v1 is appended by default, a user can set NVIDIA_APPEND_API_VERSION=False to override.
created_at: 2025-04-15T11:17:19Z
author_association: COLLABORATOR
diff_hunk: @@ -265,3 +272,44 @@ async def chat_completion(
         else:
             # we pass n=1 to get only one completion
             return convert_openai_chat_completion_choice(response.choices[0])
+
+    async def register_model(self, model: Model) -> Model:
+        """
+        Allow non-llama model registration.
+
+        Non-llama model registration: API Catalogue models, post-training models, etc.
+            client = LlamaStackAsLibraryClient("nvidia")
+            client.models.register(
+                    model_id="mistralai/mixtral-8x7b-instruct-v0.1",
+                    model_type=ModelType.llm,
+                    provider_id="nvidia",
+                    provider_model_id="mistralai/mixtral-8x7b-instruct-v0.1"
+            )
+
+            NOTE: Only supports models endpoints compatible with AsyncOpenAI base_url format.
+        """
+        if model.model_type == ModelType.embedding:
+            # embedding models are always registered by their provider model id and does not need to be mapped to a llama model
+            provider_resource_id = model.provider_resource_id
+        else:
+            provider_resource_id = self.get_provider_model_id(model.provider_resource_id)
+
+        if provider_resource_id:
+            model.provider_resource_id = provider_resource_id
+        else:
+            llama_model = model.metadata.get("llama_model")
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: mattf
body: fyi, the Model class in https://github.com/meta-llama/llama-stack/blob/main/llama_stack/models/llama/datatypes.py has been removed.
created_at: 2025-04-15T11:23:17Z
author_association: COLLABORATOR
diff_hunk: @@ -47,10 +47,15 @@ class NVIDIAConfig(BaseModel):
         default=60,
         description="Timeout for the HTTP requests",
     )
+    append_api_version: bool = Field(
+        default=lambda: os.getenv("NVIDIA_APPEND_API_VERSION", True),
path: llama_stack/providers/remote/inference/nvidia/config.py
user: login: mattf
body: is the only way to disable api appending to set `NVIDIA_APPEND_API_VERSION=""`?

if so, the result of getenv should be coerced to a bool.
created_at: 2025-04-16T10:46:35Z
author_association: COLLABORATOR
diff_hunk: @@ -379,3 +382,44 @@ async def openai_chat_completion(
             return await self._get_client(provider_model_id).chat.completions.create(**params)
         except APIConnectionError as e:
             raise ConnectionError(f"Failed to connect to NVIDIA NIM at {self._config.url}: {e}") from e
+
+    async def register_model(self, model: Model) -> Model:
path: llama_stack/providers/remote/inference/nvidia/nvidia.py
user: login: mattf
body: the bug here should be fixed in ModelRegistryHelper, we can clean it up in a follow up pr
created_at: 2025-04-16T10:48:50Z
author_association: COLLABORATOR
diff_hunk: @@ -47,10 +47,15 @@ class NVIDIAConfig(BaseModel):
         default=60,
         description="Timeout for the HTTP requests",
     )
+    append_api_version: bool = Field(
+        default=lambda: os.getenv("NVIDIA_APPEND_API_VERSION", True),
path: llama_stack/providers/remote/inference/nvidia/config.py
user: login: raspawar
body: fixed in latest commit:
sharing example usage here

```python
>>> import os
>>> os.environ["NVIDIA_APPEND_API_VERSION"]="false"
>>> os.environ["NVIDIA_BASE_URL"]="http://abc.com/chat/completion"
>>> from llama_stack.providers.remote.inference.nvidia.nvidia import NVIDIAConfig, NVIDIAInferenceAdapter
>>> inference_config = NVIDIAConfig(base_url=os.environ["NVIDIA_BASE_URL"], api_key=None)
>>> inference_adapter = NVIDIAInferenceAdapter(inference_config)
>>> client = inference_adapter._get_client("meta/llama-3.1-8b-instruct")
>>> client.base_url
URL('http://abc.com/chat/completion/')
>>> os.environ["NVIDIA_APPEND_API_VERSION"]="true"
>>> from llama_stack.providers.remote.inference.nvidia.nvidia import NVIDIAConfig, NVIDIAInferenceAdapter
>>> inference_config = NVIDIAConfig(base_url=os.environ["NVIDIA_BASE_URL"], api_key=None)
>>> inference_adapter = NVIDIAInferenceAdapter(inference_config)
>>> client = inference_adapter._get_client("meta/llama-3.1-8b-instruct")
>>> client.base_url
URL('http://abc.com/chat/completion/v1/')
```

Need to add unit test cases for inference adapater and this too, will do in follow-up PR.
created_at: 2025-04-16T11:57:35Z
author_association: CONTRIBUTOR

PR comments:
user: login: ashwinb
created_at: 2025-04-14T20:03:31Z
body: Going to let @mattf approve this one first before I take a look. 
author_association: CONTRIBUTOR
user: login: ashwinb
created_at: 2025-04-14T21:04:01Z
body: Could you resolve the conflicts? I will take a look then and get this in.
author_association: CONTRIBUTOR