Dynamic nvidia nim sut factory

src/modelgauge/suts/nvidia_nim_api_client.py

@@ -58,28 +58,6 @@ def _translate_request(self, messages, options: ModelOptions) -> NIMOpenAIChatRe
         )
 
 
-SUTS.register(
-    NvidiaNIMApiClient,
-    "nvidia-llama-3.1-nemotron-70b-instruct",
-    "nvidia/llama-3.1-nemotron-70b-instruct",
-    InjectSecret(NvidiaNIMApiKey),
-)
-
-
-SUTS.register(
-    NvidiaNIMApiClient,
-    "nvidia-nemotron-4-340b-instruct",
-    "nvidia/nemotron-4-340b-instruct",
-    InjectSecret(NvidiaNIMApiKey),
-)
-
-SUTS.register(
-    NvidiaNIMApiClient,
-    "nvidia-mistral-nemo-minitron-8b-8k-instruct",
-    "nvidia/mistral-nemo-minitron-8b-8k-instruct",
-    InjectSecret(NvidiaNIMApiKey),
-)
-
 SUTS.register(
     NvidiaNIMApiClient,
     "nvidia-nemotron-mini-4b-instruct",

src/modelgauge/suts/nvidia_nim_sut_factory.py

@@ -0,0 +1,36 @@
+from openai import OpenAI
+
+from modelgauge.dynamic_sut_factory import DynamicDriverSUTFactory, ModelNotSupportedError
+from modelgauge.secret_values import InjectSecret
+from modelgauge.sut_definition import SUTDefinition
+from modelgauge.suts.nvidia_nim_api_client import BASE_URL, NvidiaNIMApiKey, NvidiaNIMApiClient
+
+
+class NvidiaNIMSUTFactory(DynamicDriverSUTFactory):
+    DRIVER_NAME = "nvidia-nim"
+
+    def __init__(self, raw_secrets):
+        super().__init__(raw_secrets)
+        self._client = None
+
+    @property
+    def client(self) -> OpenAI:
+        if self._client is None:
+            self._client = OpenAI(api_key=self.injected_secrets()[0].value, base_url=BASE_URL)
+        return self._client
+
+    def get_secrets(self) -> list[InjectSecret]:
+        return [InjectSecret(NvidiaNIMApiKey)]
+
+    def _model_exists(self, model_name: str) -> bool:
+        try:
+            self.client.models.retrieve(model_name)  # type: ignore
+            return True
+        except:
+            return False
+
+    def make_sut(self, sut_definition: SUTDefinition) -> NvidiaNIMApiClient:
+        model_name = sut_definition.external_model_name()
+        if not self._model_exists(model_name):
+            raise ModelNotSupportedError(f"Model {model_name} not found or not available on nvidia NIM.")
+        return NvidiaNIMApiClient(sut_definition.uid, model_name, *self.injected_secrets())

tests/modelgauge_tests/sut_tests/test_nvidia_nim_sut_factory.py

@@ -0,0 +1,32 @@
+import pytest
+from unittest.mock import MagicMock
+
+from modelgauge.dynamic_sut_factory import ModelNotSupportedError
+from modelgauge.sut_definition import SUTDefinition
+from modelgauge.suts.nvidia_nim_api_client import NvidiaNIMApiClient
+from modelgauge.suts.nvidia_nim_sut_factory import NvidiaNIMSUTFactory
+
+
+@pytest.fixture
+def factory():
+    return NvidiaNIMSUTFactory({"nvidia-nim-api": {"api_key": "value"}})
+
+
+def test_make_sut(factory):
+    factory._client = MagicMock()
+    factory._client.models.retrieve.return_value = "model exists"
+
+    sut_definition = SUTDefinition(model="bar", maker="foo", driver="nvidia-nim")
+    sut = factory.make_sut(sut_definition)
+
+    assert isinstance(sut, NvidiaNIMApiClient)
+    assert sut.uid == "foo/bar:nvidia-nim"
+    assert sut.model == "foo/bar"
+
+
+def test_make_sut_bad_model(factory):
+    sut_definition = SUTDefinition(model="bogus", maker="fake", driver="nvidia-nim")
+    factory._client = MagicMock()
+    factory._client.models.retrieve.side_effect = Exception()
+    with pytest.raises(ModelNotSupportedError):
+        factory.make_sut(sut_definition)