query_available_models() -> list[str] -> check_model_availability(model) -> bool

Author: mattf

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

@@ -89,9 +89,14 @@ def __init__(self, config: NVIDIAConfig) -> None:
 
         self._config = config
 
-    async def query_available_models(self) -> list[str]:
-        """Query available models from the NVIDIA API."""
-        return [model.id async for model in self._get_client().models.list()]
+    async def check_model_availability(self, model: str) -> bool:
+        """Check if a specific model is available from the NVIDIA API."""
+        try:
+            await self._get_client().models.retrieve(model)
+            return True
+        except Exception:
+            # If we can't retrieve the model, it's not available
+            return False
 
     @lru_cache  # noqa: B019
     def _get_client(self, provider_model_id: str | None = None) -> AsyncOpenAI:

tests/unit/providers/nvidia/test_supervised_fine_tuning.py

@@ -344,8 +344,8 @@ def test_inference_register_model(self):
         )
 
         # simulate a NIM where default/job-1234 is an available model
-        with patch.object(self.inference_adapter, "query_available_models", new_callable=AsyncMock) as mock_query:
-            mock_query.return_value = [model_id]
+        with patch.object(self.inference_adapter, "check_model_availability", new_callable=AsyncMock) as mock_check:
+            mock_check.return_value = True
             result = self.run_async(self.inference_adapter.register_model(model))
         assert result == model
         assert len(self.inference_adapter.alias_to_provider_id_map) > 1