From 72c52a3b45e8f7cac5da2d82918acc7cd60da7a1 Mon Sep 17 00:00:00 2001 From: Max Dubrinsky Date: Mon, 24 Aug 2026 14:08:58 -0400 Subject: [PATCH 1/4] refactor(models): migrate models consumers to typed ModelsClient Migrate models API call sites from sdk.models.* (Stainless SDK) to client_from_platform(sdk, ModelsClient).* (typed HTTP client), following the pattern established in #1277. 22 files changed across e2e tests, services, plugins, and test utilities. Key migrations: - sdk.models.retrieve -> get_model().data() - sdk.models.create -> create_model(body=CreateModelEntityRequest(...)).data() - sdk.models.update -> update_model(body=UpdateModelEntityRequest(...)).data() - sdk.models.list -> list_models(...).items() - sdk.models.wait_for_openai_model -> wait_for_openai_model(...) - sdk.models.get_provider_route_openai_url -> get_provider_route_openai_url(...) - Error imports: APIStatusError->NemoHTTPError, APIConnectionError/APITimeoutError->NemoTransportError - Type imports: nemo_platform.types.* -> nemo_platform_plugin.models.types.* Skipped (SDK infrastructure, not consumers): - packages/models/src/models/resources.py (extended ModelsResource, deleted with SDK) - packages/models/tests/test_client.py (tests SDK infrastructure) - Auto-generated CLI files (4 files) AIRCORE-827 Signed-off-by: Max Dubrinsky --- e2e/agents_deploy_helpers.py | 4 +- e2e/test_evaluator_plugin.py | 19 +- packages/nemo_nb/tests/test_myst_stripping.py | 4 +- .../nemo_nb/tests/test_notebook_splitting.py | 6 +- .../tests/test_strip_type_checker_comments.py | 6 +- .../tests/cli/core/test_code_generator.py | 2 +- .../nemo_platform_plugin/nooa_model_client.py | 12 +- .../src/nmp/testing/e2e/customizer.py | 7 +- .../src/nemo_auditor/jobs/audit.py | 19 +- .../src/nemo_evaluator/resolvers.py | 28 ++- .../jobs/generate.py | 5 +- .../core/inference_gateway/api/model_cache.py | 24 +- .../models/src/nmp/core/models/app/utils.py | 2 +- .../deployments_plugin/nim_compiler.py | 2 +- .../backends/deployments_plugin/resolve.py | 2 +- .../controllers/backends/vllm_compiler.py | 2 +- .../nmp/core/models/controllers/context.py | 2 +- .../core/models/controllers/entity_cache.py | 23 +- .../models/controllers/models_controller.py | 19 +- .../models/controllers/provider_reconciler.py | 2 +- .../nmp/core/models/sidecars/adapters/main.py | 21 +- .../nmp/core/models/tasks/model_spec/run.py | 40 ++-- .../integration/test_models_with_auth.py | 220 +++++++++++------- .../test_workspace_iam_models_isolation.py | 90 ++++--- .../environment/setup-mock.py | 13 +- .../environment/setup-mock.py | 13 +- .../environment/setup-mock.py | 13 +- .../environment/setup-mock.py | 13 +- 28 files changed, 396 insertions(+), 217 deletions(-) diff --git a/e2e/agents_deploy_helpers.py b/e2e/agents_deploy_helpers.py index 156db08561..ff31fe1cbc 100644 --- a/e2e/agents_deploy_helpers.py +++ b/e2e/agents_deploy_helpers.py @@ -28,6 +28,8 @@ import pytest from nemo_agents_plugin.entities import NAT_WORKFLOW_CONFIG_FORMAT, NEMO_AGENTS_SPEC_CONFIG_FORMAT from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.models.client import ModelsClient from nmp.testing import MockProviderResponse, add_mock_provider # The mocked completion the deployed agent must round-trip back to the caller. @@ -274,7 +276,7 @@ def run_agent_deploy_and_invoke( endpoints = deployment.get("endpoints") or [] assert endpoints and endpoints[0]["url"], deployment - sdk.models.wait_for_openai_model(model_name, workspace=workspace) + client_from_platform(sdk, ModelsClient).wait_for_openai_model(model_name, workspace=workspace) response = sdk.agents.invoke( workspace=workspace, diff --git a/e2e/test_evaluator_plugin.py b/e2e/test_evaluator_plugin.py index 4b8bb9c464..b31e019353 100644 --- a/e2e/test_evaluator_plugin.py +++ b/e2e/test_evaluator_plugin.py @@ -46,10 +46,15 @@ from nemo_evaluator_sdk.metrics.tool_calling import ToolCallingMetric from nemo_evaluator_sdk.values.results import EvaluationResult from nemo_evaluator_sdk.values.scores import JSONScoreParser, RangeScore -from nemo_platform import APIConnectionError, APIStatusError, NeMoPlatform +from nemo_platform import NeMoPlatform from nemo_platform.types.inference import ModelProvider from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoHTTPError as APIStatusError +from nemo_platform_plugin.client.errors import NemoTransportError as APIConnectionError +from nemo_platform_plugin.inference_middleware import BackendFormat from nemo_platform_plugin.jobs.client import JobsClient +from nemo_platform_plugin.models.client import ModelsClient +from nemo_platform_plugin.models.types import CreateModelEntityRequest from nmp.testing import add_mock_provider, short_unique_name, wait_for_model_entity from nmp.testing.e2e import wait_for_platform_job from nmp.testing.utils import ensure_passthrough_virtual_model @@ -263,13 +268,15 @@ def _create_ready_mock_model( mock_response_body=mock_response_body, should_autoprovision_virtual_model=False, ) - sdk.models.create( + client_from_platform(sdk, ModelsClient).create_model( workspace=workspace, - name=name, - backend_format="OPENAI_CHAT", - model_providers=[f"{workspace}/{provider.name}"], + body=CreateModelEntityRequest( + name=name, + backend_format=BackendFormat.OPENAI_CHAT, + model_providers=[f"{workspace}/{provider.name}"], + ), exist_ok=True, - ) + ).data() wait_for_model_entity( sdk, workspace, diff --git a/packages/nemo_nb/tests/test_myst_stripping.py b/packages/nemo_nb/tests/test_myst_stripping.py index e24adfa975..0a18027048 100644 --- a/packages/nemo_nb/tests/test_myst_stripping.py +++ b/packages/nemo_nb/tests/test_myst_stripping.py @@ -246,7 +246,7 @@ def test_mixed_content(): ```python # Regular code -sdk.models.deploy() +sdk.models.create_deployment() ``` :::{warning} @@ -267,4 +267,4 @@ def test_mixed_content(): assert "Use the CLI for faster deployment." in result assert "Make sure GPU resources are configured." in result assert "```python" in result - assert "sdk.models.deploy()" in result + assert "sdk.models.create_deployment()" in result diff --git a/packages/nemo_nb/tests/test_notebook_splitting.py b/packages/nemo_nb/tests/test_notebook_splitting.py index ca838bcb99..5679adf680 100644 --- a/packages/nemo_nb/tests/test_notebook_splitting.py +++ b/packages/nemo_nb/tests/test_notebook_splitting.py @@ -399,7 +399,11 @@ def test_real_world_tab_set_example(self): {"cell_type": "code", "metadata": {"language": "bash"}, "source": ["nemo models list\n"]}, {"cell_type": "markdown", "source": [":::\n"]}, {"cell_type": "markdown", "source": [":::{tab-item} Python SDK\n", ":sync: python-sdk\n"]}, - {"cell_type": "code", "metadata": {"language": "python"}, "source": ["client.models.list()\n"]}, + { + "cell_type": "code", + "metadata": {"language": "python"}, + "source": ["client_from_platform(client, ModelsClient).list_models()\n"], + }, {"cell_type": "markdown", "source": [":::\n"]}, {"cell_type": "markdown", "source": ["::::\n"]}, ] diff --git a/packages/nemo_nb/tests/test_strip_type_checker_comments.py b/packages/nemo_nb/tests/test_strip_type_checker_comments.py index 77e362cf94..9bddee94c8 100644 --- a/packages/nemo_nb/tests/test_strip_type_checker_comments.py +++ b/packages/nemo_nb/tests/test_strip_type_checker_comments.py @@ -16,7 +16,7 @@ def test_strip_ty_ignore_comments(): "metadata": {"language": "python"}, "source": [ "# This is a regular comment\n", - "sdk.models.get_openai_route_base_url()\n", + "client_from_platform(sdk, ModelsClient).get_openai_route_base_url()\n", "sdk.models.get_model_entity_route_openai_url(entity) # ty: ignore[unresolved-reference]\n", "sdk.models.get_provider_route_openai_url(provider) # ty: ignore[unresolved-reference]\n", ], @@ -35,8 +35,8 @@ def test_strip_ty_ignore_comments(): assert "# This is a regular comment" in result # Verify that the code lines are still there (without the ty: comments) - assert "sdk.models.get_model_entity_route_openai_url(entity)" in result - assert "sdk.models.get_provider_route_openai_url(provider)" in result + assert "client_from_platform(sdk, ModelsClient).get_model_entity_route_openai_url(entity)" in result + assert "client_from_platform(sdk, ModelsClient).get_provider_route_openai_url(provider)" in result def test_strip_type_ignore_comments(): diff --git a/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py b/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py index a1babd15d8..e08901c2ae 100644 --- a/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py +++ b/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py @@ -18,7 +18,7 @@ def test_generate_python_code_simple_list(): assert "from nemo_platform import NeMoPlatform" in code assert "client = NeMoPlatform()" in code - assert "response = client.models.list()" in code + assert "response = client_from_platform(client, ModelsClient).list_models()" in code assert "print(response)" in code diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/nooa_model_client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/nooa_model_client.py index 2b677a3cad..e5a281c1e3 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/nooa_model_client.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/nooa_model_client.py @@ -16,8 +16,10 @@ from nemo_platform import AsyncNeMoPlatform from nemo_platform.types.inference import ModelProvider -from nemo_platform.types.models import ModelEntity from nemo_platform_ext.config import get_context +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.models.client import AsyncModelsClient +from nemo_platform_plugin.models.types import ModelEntity from nooa.unifiedllm import CompletionClient, UnifiedLLM _PLACEHOLDER_API_KEY = "not-needed" @@ -90,8 +92,9 @@ def _completion_client( served_model_name: str, ) -> CompletionClient: """Build a Nooa client that routes through the Model Entity's Platform URL.""" - api_base = client.models.get_model_entity_route_openai_url(model_entity) - extra_headers = dict(client.models.get_client_default_headers()) + models = client_from_platform(client, AsyncModelsClient) + api_base = models.get_model_entity_route_openai_url(model_entity) + extra_headers = dict(models.default_headers) # Inference Gateway's direct passthrough session preserves compressed # response bytes. Nooa needs decoded JSON/SSE, so make that requirement # explicit at this adapter boundary for every configured agent client. @@ -161,6 +164,7 @@ async def resolve_model_clients( refs: ConfiguredModelRefs | None = None, ) -> ConfiguredModelClients: """Resolve configured Model Entities and construct each distinct client once.""" + models = client_from_platform(client, AsyncModelsClient) selected = refs or configured_model_refs() resolved: dict[str, UnifiedLLM] = {} provider_cache: dict[str, ModelProvider] = {} @@ -169,7 +173,7 @@ async def resolve_model_clients( if model_ref in resolved: continue workspace, name = _parse_model_ref(model_ref) - entity = await client.models.retrieve(name, workspace=workspace) + entity = (await models.get_model(name=name, workspace=workspace)).data() served_model_name = await _served_model_name(client, entity, provider_cache) resolved[model_ref] = _completion_client(client, entity, served_model_name) except Exception as resolution_error: diff --git a/packages/nmp_testing/src/nmp/testing/e2e/customizer.py b/packages/nmp_testing/src/nmp/testing/e2e/customizer.py index 3a5ef35b0c..498857b5d2 100644 --- a/packages/nmp_testing/src/nmp/testing/e2e/customizer.py +++ b/packages/nmp_testing/src/nmp/testing/e2e/customizer.py @@ -19,6 +19,7 @@ from nemo_platform.types.inference import ContainerExecutorConfigParam, ModelDeploymentConfigModelSpecParam from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.jobs.client import JobsClient +from nemo_platform_plugin.models.client import ModelsClient logger = logging.getLogger(__name__) @@ -326,7 +327,7 @@ def wait_for_model_spec( """ deadline = time.time() + timeout while time.time() < deadline: - me = sdk.models.retrieve(model_entity_name, workspace=workspace, verbose=True) + me = client_from_platform(sdk, ModelsClient).get_model(name=model_entity_name, workspace=workspace).data() if me.spec is not None: logger.info( f"✓ Model spec populated for {model_entity_name}: checkpoint_model_name={me.spec.checkpoint_model_name}" @@ -551,7 +552,9 @@ def _wait_for_gateway_ready( until it confirms the provider is routable. """ logger.info("Waiting for inference gateway to sync...") - if not sdk.models.wait_for_gateway(deployment_name, workspace=workspace, timeout=timeout): + if not client_from_platform(sdk, ModelsClient).wait_for_gateway( + deployment_name, workspace=workspace, timeout=timeout + ): pytest.fail( f"Inference gateway did not become ready for deployment '{deployment_name}' " f"within {timeout}s. The deployment's model provider may not have been created. " diff --git a/plugins/nemo-auditor/src/nemo_auditor/jobs/audit.py b/plugins/nemo-auditor/src/nemo_auditor/jobs/audit.py index c559737921..466e2f0efc 100644 --- a/plugins/nemo-auditor/src/nemo_auditor/jobs/audit.py +++ b/plugins/nemo-auditor/src/nemo_auditor/jobs/audit.py @@ -35,12 +35,15 @@ from nemo_auditor.entities import AuditConfig, AuditTarget from nemo_platform import AsyncNeMoPlatform, NeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.response import NemoResponse from nemo_platform_plugin.entities import parse_qualified_name from nemo_platform_plugin.entities.client import AsyncEntitiesClient from nemo_platform_plugin.entity_client import NemoEntitiesClient, NemoEntityNotFoundError from nemo_platform_plugin.job import NemoJob from nemo_platform_plugin.job_context import JobContext from nemo_platform_plugin.job_results import JobResults +from nemo_platform_plugin.models.client import AsyncModelsClient, ModelsClient +from nemo_platform_plugin.models.types import ModelProvider from pydantic import BaseModel, ConfigDict, Field, StringConstraints logger = logging.getLogger(__name__) @@ -215,16 +218,20 @@ def _rewrite_options_uris( ) try: if sdk is not None: - provider = sdk.inference.providers.retrieve(workspace=igw_ref["workspace"], name=igw_ref["provider"]) - uri = sdk.models.get_provider_route_openai_url(provider) + models = client_from_platform(sdk, ModelsClient) + provider = models.get_provider(workspace=igw_ref["workspace"], name=igw_ref["provider"]).data() + uri = models.get_provider_route_openai_url(provider) else: assert async_sdk is not None # async_sdk path: AuditJob.run() executes inside asyncio.to_thread(), so # this worker thread has no running event loop — asyncio.run() is safe. - provider = asyncio.run( - async_sdk.inference.providers.retrieve(workspace=igw_ref["workspace"], name=igw_ref["provider"]) - ) - uri = async_sdk.models.get_provider_route_openai_url(provider) + async_models = client_from_platform(async_sdk, AsyncModelsClient) + + async def _fetch_provider() -> NemoResponse[ModelProvider]: + return await async_models.get_provider(workspace=igw_ref["workspace"], name=igw_ref["provider"]) + + provider = asyncio.run(_fetch_provider()).data() + uri = async_models.get_provider_route_openai_url(provider) except Exception as exc: raise RuntimeError( f"Failed to resolve inference gateway provider '{igw_ref['workspace']}/{igw_ref['provider']}': {exc}" diff --git a/plugins/nemo-evaluator/src/nemo_evaluator/resolvers.py b/plugins/nemo-evaluator/src/nemo_evaluator/resolvers.py index a280005701..1f6611051d 100644 --- a/plugins/nemo-evaluator/src/nemo_evaluator/resolvers.py +++ b/plugins/nemo-evaluator/src/nemo_evaluator/resolvers.py @@ -11,9 +11,12 @@ from typing import Protocol, TypeVar, cast, runtime_checkable from nemo_evaluator_sdk.values.models import Model, ModelRef -from nemo_platform import NotFoundError -from nemo_platform.types.inference import ModelProvider as PlatformModelProvider -from nemo_platform.types.models import ModelEntity as PlatformModelEntity +from nemo_platform import AsyncNeMoPlatform, NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError +from nemo_platform_plugin.models.client import AsyncModelsClient, ModelsClient +from nemo_platform_plugin.models.types import ModelEntity as PlatformModelEntity +from nemo_platform_plugin.models.types import ModelProvider as PlatformModelProvider _logger = logging.getLogger(__name__) _T = TypeVar("_T") @@ -90,8 +93,13 @@ async def _resolve_provider_host_url( return None try: - provider = await _maybe_await(sdk.inference.providers.retrieve(provider_name, workspace=provider_workspace)) - return provider.host_url + models = ( + client_from_platform(cast(AsyncNeMoPlatform, sdk), AsyncModelsClient) + if isinstance(sdk, AsyncNeMoPlatform) + else client_from_platform(cast(NeMoPlatform, sdk), ModelsClient) + ) + provider = await _maybe_await(models.get_provider(name=provider_name, workspace=provider_workspace)) + return provider.data().host_url except NotFoundError: _logger.warning("Provider not found during host_url resolution", extra={"provider_ref": provider_ref}) return None @@ -114,7 +122,13 @@ async def resolve_model(self, model_ref: ModelRef) -> Model: ) try: - model_entity = await _maybe_await(self._sdk.models.retrieve(name, workspace=workspace)) + models = ( + client_from_platform(cast(AsyncNeMoPlatform, self._sdk), AsyncModelsClient) + if isinstance(self._sdk, AsyncNeMoPlatform) + else client_from_platform(cast(NeMoPlatform, self._sdk), ModelsClient) + ) + model_entity = await _maybe_await(models.get_model(name=name, workspace=workspace)) + model_entity = model_entity.data() except NotFoundError as exc: raise ValueError( f"Model reference '{model_ref.root}' not found. " @@ -122,7 +136,7 @@ async def resolve_model(self, model_ref: ModelRef) -> Model: "or use an inline model definition instead." ) from exc - endpoint = self._sdk.models.get_model_entity_route_openai_url(model_entity) + endpoint = models.get_model_entity_route_openai_url(model_entity) host_url = await _resolve_provider_host_url(self._sdk, model_entity) return Model( url=endpoint, diff --git a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/jobs/generate.py b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/jobs/generate.py index fda4030622..5cea50897d 100644 --- a/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/jobs/generate.py +++ b/plugins/nemo-safe-synthesizer/src/nemo_safe_synthesizer_plugin/jobs/generate.py @@ -33,6 +33,7 @@ from nemo_platform_plugin.jobs.client import AsyncJobsClient from nemo_platform_plugin.jobs.exceptions import PlatformJobCompilationError from nemo_platform_plugin.jobs.image import get_qualified_image +from nemo_platform_plugin.models.client import AsyncModelsClient from nemo_safe_synthesizer.config.external_results import SafeSynthesizerSummary from nemo_safe_synthesizer_plugin.config import config as plugin_config from nemo_safe_synthesizer_plugin.job_config import SafeSynthesizerJobConfig, parse_pretrained_model_job_ref @@ -124,7 +125,9 @@ async def compile( raise PlatformJobCompilationError( f"Failed to retrieve model provider {classify_model_provider!r}: Access denied to workspace {provider_workspace!r}" ) from e - nim_endpoint_url = async_sdk.models.get_provider_route_openai_url(provider) + nim_endpoint_url = client_from_platform(async_sdk, AsyncModelsClient).get_provider_route_openai_url( + provider + ) parsed_url = urlparse(nim_endpoint_url) environment.append(EnvironmentVariable(name="CLASSIFY_LLM_ENDPOINT_PATH", value=parsed_url.path)) logger.info("Configured NIM endpoint URL: %s (provider: %s)", nim_endpoint_url, classify_model_provider) diff --git a/services/core/inference-gateway/src/nmp/core/inference_gateway/api/model_cache.py b/services/core/inference-gateway/src/nmp/core/inference_gateway/api/model_cache.py index b5c92b71af..2de3084a79 100644 --- a/services/core/inference-gateway/src/nmp/core/inference_gateway/api/model_cache.py +++ b/services/core/inference-gateway/src/nmp/core/inference_gateway/api/model_cache.py @@ -9,10 +9,14 @@ from datetime import datetime from typing import TYPE_CHECKING, Awaitable, Callable -from nemo_platform import APIConnectionError, APIStatusError, AsyncNeMoPlatform +from nemo_platform import AsyncNeMoPlatform from nemo_platform.types.inference import ModelProvider, ServedModelMapping -from nemo_platform.types.models.model_entity import ModelEntity +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoHTTPError as APIStatusError +from nemo_platform_plugin.client.errors import NemoTransportError as APIConnectionError from nemo_platform_plugin.inference_middleware import BackendFormat +from nemo_platform_plugin.models.client import AsyncModelsClient +from nemo_platform_plugin.models.types import ModelEntity from nmp.common.observability import MARK_INTERNAL_REQUEST_HEADERS from nmp.core.inference_gateway.api.proxy import retrieve_secret_value from nmp.core.inference_gateway.api.virtual_model_cache import ( @@ -191,7 +195,7 @@ async def _model_provider_getter() -> list[ModelProvider]: providers = [provider async for provider in resp] return providers except APIConnectionError as exc: - raise ModelProviderRefreshError(f"Error connecting to models service: {exc.body}") from exc + raise ModelProviderRefreshError(f"Error connecting to models service: {exc}") from exc except APIStatusError as exc: raise ModelProviderRefreshError( f"Error refreshing from models service: {exc.status_code}, {exc.body}" @@ -204,20 +208,16 @@ def model_entity_getter_from_sdk(models_sdk: AsyncNeMoPlatform) -> Callable[[], async def _model_entity_getter() -> list[ModelEntity]: try: # SDK returns AsyncPaginator - iterate through all pages to get all model entities. - resp = models_sdk.models.list( + resp = await client_from_platform(models_sdk, AsyncModelsClient).list_models( workspace="-", # Cross-workspace query - page_size=200, - verbose=False, - extra_headers=MARK_INTERNAL_REQUEST_HEADERS, + query_params={"page_size": 200, "verbose": False}, ) - models = [model async for model in resp] + models = [model async for model in resp.items()] return models except APIConnectionError as exc: - raise ModelProviderRefreshError(f"Error connecting to models service: {exc.body}") from exc + raise ModelProviderRefreshError(f"Error connecting to models service: {exc}") from exc except APIStatusError as exc: - raise ModelProviderRefreshError( - f"Error refreshing model entities from models service: {exc.status_code}, {exc.body}" - ) from exc + raise ModelProviderRefreshError(f"Error refreshing model entities from models service: {exc}") from exc return _model_entity_getter diff --git a/services/core/models/src/nmp/core/models/app/utils.py b/services/core/models/src/nmp/core/models/app/utils.py index b062636735..1ea9f72904 100644 --- a/services/core/models/src/nmp/core/models/app/utils.py +++ b/services/core/models/src/nmp/core/models/app/utils.py @@ -12,7 +12,6 @@ from nemo_platform.types.inference.model_deployment import ModelDeployment from nemo_platform.types.inference.model_deployment_config import ModelDeploymentConfig from nemo_platform.types.inference.model_provider import ModelProvider -from nemo_platform.types.models import ModelEntity from nemo_platform_plugin.k8s_naming import ( DNS_LABEL_MAX_LENGTH, DNS_SUBDOMAIN_MAX_LENGTH, @@ -20,6 +19,7 @@ k8s_safe_name, workspace_name_identity, ) +from nemo_platform_plugin.models.types import ModelEntity from nmp.common.api.common import PaginationData from nmp.common.entities.constants import NAME_PATTERN as ENTITY_NAME_PATTERN from pydantic import BaseModel diff --git a/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/nim_compiler.py b/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/nim_compiler.py index b1d9e13177..68485525a9 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/nim_compiler.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/nim_compiler.py @@ -28,7 +28,7 @@ VolumeMount, ) from nemo_platform.types.inference.k8s_nim_operator_config import K8sNIMOperatorConfig -from nemo_platform.types.models.model_entity import ModelEntity +from nemo_platform_plugin.models.types import ModelEntity from nmp.common.config import Runtime from nmp.core.models.app import is_multi_llm_image, parse_model_name_revision from nmp.core.models.controllers.backends.common import DeploymentConfigView diff --git a/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/resolve.py b/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/resolve.py index a50fe4cd92..9f49c2ea0f 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/resolve.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/resolve.py @@ -8,7 +8,7 @@ from nemo_platform.types.inference.model_deployment import ModelDeployment from nemo_platform.types.inference.model_deployment_config import ModelDeploymentConfig -from nemo_platform.types.models.model_entity import ModelEntity +from nemo_platform_plugin.models.types import ModelEntity from nmp.common.config import Runtime, get_platform_config from nmp.common.config.base import LOOPBACK_ADDRESSES, determine_loopback_override from nmp.core.models.app import ModelWeightsType, get_model_weights_type, parse_model_name_revision diff --git a/services/core/models/src/nmp/core/models/controllers/backends/vllm_compiler.py b/services/core/models/src/nmp/core/models/controllers/backends/vllm_compiler.py index 26b42acaea..eb3968fbd0 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/vllm_compiler.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/vllm_compiler.py @@ -16,7 +16,7 @@ from logging import getLogger from typing import Optional -from nemo_platform.types.models.model_entity import ModelEntity +from nemo_platform_plugin.models.types import ModelEntity from nmp.core.models.controllers.backends.common import DeploymentConfigView logger = getLogger(__name__) diff --git a/services/core/models/src/nmp/core/models/controllers/context.py b/services/core/models/src/nmp/core/models/controllers/context.py index 2df5a2c56f..53480e94e6 100644 --- a/services/core/models/src/nmp/core/models/controllers/context.py +++ b/services/core/models/src/nmp/core/models/controllers/context.py @@ -10,7 +10,7 @@ from nemo_platform.types.inference.model_deployment import ModelDeployment from nemo_platform.types.inference.model_deployment_config import ModelDeploymentConfig from nemo_platform.types.inference.model_provider import ModelProvider -from nemo_platform.types.models.model_entity import ModelEntity +from nemo_platform_plugin.models.types import ModelEntity @dataclass diff --git a/services/core/models/src/nmp/core/models/controllers/entity_cache.py b/services/core/models/src/nmp/core/models/controllers/entity_cache.py index 8b4bd07489..d4713265f3 100644 --- a/services/core/models/src/nmp/core/models/controllers/entity_cache.py +++ b/services/core/models/src/nmp/core/models/controllers/entity_cache.py @@ -20,8 +20,10 @@ from typing import Callable from nemo_platform import AsyncNeMoPlatform -from nemo_platform._exceptions import ConflictError, NotFoundError -from nemo_platform.types.models.model_entity import ModelEntity +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ConflictError, NotFoundError +from nemo_platform_plugin.models.client import AsyncModelsClient +from nemo_platform_plugin.models.types import CreateModelEntityRequest, ModelEntity, UpdateModelEntityRequest logger = getLogger(__name__) @@ -94,7 +96,8 @@ async def refresh(self) -> None: ) entities: dict[tuple[str, str], ModelEntity] = {} - async for entity in self._models_sdk.models.list(workspace="-", page_size=_PAGE_SIZE): + models = client_from_platform(self._models_sdk, AsyncModelsClient) + async for entity in (await models.list_models(workspace="-", query_params={"page_size": _PAGE_SIZE})).items(): entities[(entity.workspace, entity.name)] = entity self._emit_heartbeat() @@ -220,12 +223,17 @@ async def _create(self, workspace: str, name: str, staged: _PendingEntity) -> No create_kwargs["model_providers"] = list(staged.link_providers) try: - created = await self._models_sdk.models.create(workspace=workspace, name=name, **create_kwargs) + models = client_from_platform(self._models_sdk, AsyncModelsClient) + created = ( + await models.create_model( + workspace=workspace, body=CreateModelEntityRequest(name=name, **create_kwargs) + ) + ).data() except ConflictError: # Created concurrently; adopt it and apply the staged changes instead. logger.debug("Model Entity %s/%s already exists, applying staged changes", workspace, name) try: - existing = await self._models_sdk.models.retrieve(workspace=workspace, name=name) + existing = (await models.get_model(workspace=workspace, name=name)).data() except NotFoundError: return self._entities[(workspace, name)] = existing @@ -251,7 +259,10 @@ async def _update(self, workspace: str, name: str, staged: _PendingEntity, exist logger.debug("Model Entity %s/%s already matches desired state", workspace, name) return - updated = await self._models_sdk.models.update(workspace=workspace, name=name, **update_params) + models = client_from_platform(self._models_sdk, AsyncModelsClient) + updated = ( + await models.update_model(workspace=workspace, name=name, body=UpdateModelEntityRequest(**update_params)) + ).data() if updated is not None: self._entities[(workspace, name)] = updated logger.debug("Updated Model Entity %s/%s: %s", workspace, name, sorted(update_params)) diff --git a/services/core/models/src/nmp/core/models/controllers/models_controller.py b/services/core/models/src/nmp/core/models/controllers/models_controller.py index a9d70e0081..4c09e819f8 100644 --- a/services/core/models/src/nmp/core/models/controllers/models_controller.py +++ b/services/core/models/src/nmp/core/models/controllers/models_controller.py @@ -6,12 +6,14 @@ from logging import getLogger from typing import Optional -from nemo_platform import DefaultAsyncHttpxClient -from nemo_platform._exceptions import NotFoundError +from nemo_platform import DefaultAsyncHttpxClient # type: ignore[deprecated] from nemo_platform.types.inference import ModelDeploymentStatus from nemo_platform.types.inference.model_deployment import ModelDeployment from nemo_platform.types.inference.model_deployment_config import ModelDeploymentConfig -from nemo_platform.types.models.model_entity import ModelEntity +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError +from nemo_platform_plugin.models.client import AsyncModelsClient +from nemo_platform_plugin.models.types import ModelEntity from nmp.common.controller import Controller, HeartbeatMixin from nmp.common.entities.utils import parse_entity_ref from nmp.common.sdk_factory import get_async_platform_sdk @@ -236,10 +238,13 @@ async def _retrieve_model_entity( if revision or not self._entity_cache.loaded: # A revision resolves server-side and does not correspond to an # cache key, so it has to be fetched directly. - model_entity = await self._models_sdk.models.retrieve( - name=full_model_name, - workspace=workspace, - ) + models = client_from_platform(self._models_sdk, AsyncModelsClient) + model_entity = ( + await models.get_model( + name=full_model_name, + workspace=workspace, + ) + ).data() else: model_entity = self._entity_cache.get(workspace, model_name) if model_entity is None: diff --git a/services/core/models/src/nmp/core/models/controllers/provider_reconciler.py b/services/core/models/src/nmp/core/models/controllers/provider_reconciler.py index 4c2abc3b66..e1a4306b38 100644 --- a/services/core/models/src/nmp/core/models/controllers/provider_reconciler.py +++ b/services/core/models/src/nmp/core/models/controllers/provider_reconciler.py @@ -16,7 +16,7 @@ from nemo_platform.types.inference.model_deployment_config import ModelDeploymentConfig from nemo_platform.types.inference.model_provider import ModelProvider from nemo_platform.types.inference.virtual_model import VirtualModel -from nemo_platform.types.models.model_entity import ModelEntity +from nemo_platform_plugin.models.types import ModelEntity from nmp.common.datetime_utils import ensure_utc from nmp.common.entities.constants import NAME_PATTERN from nmp.common.entities.utils import parse_entity_ref diff --git a/services/core/models/src/nmp/core/models/sidecars/adapters/main.py b/services/core/models/src/nmp/core/models/sidecars/adapters/main.py index 43503c0808..6f154754b8 100644 --- a/services/core/models/src/nmp/core/models/sidecars/adapters/main.py +++ b/services/core/models/src/nmp/core/models/sidecars/adapters/main.py @@ -13,9 +13,11 @@ import urllib.error import urllib.request -from nemo_platform import NeMoPlatform, NotFoundError -from nemo_platform.types.models import ModelEntity -from nemo_platform.types.models.adapter import Adapter +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError +from nemo_platform_plugin.models.client import ModelsClient +from nemo_platform_plugin.models.types import Adapter, ModelEntity from nmp.common.config import get_platform_config from nmp.common.controller import ( Controller, @@ -92,6 +94,7 @@ def __init__(self, stop_signal: threading.Event | None = None): as_service="models", internal=True, ) + self._models = client_from_platform(self._sdk, ModelsClient) def download_fileset(self, dest_dir: str, workspace: str, name: str) -> bool: try: @@ -156,11 +159,11 @@ def _update_prompt_tuned_models(self, dirs_to_keep: set[str]): # (AALGO-129): they remain single-workspace for now and continue to # use the bare model_entity.name as their on-disk directory. logger.info(f"Fetching prompt data for {self.workspace}/{self.model_name}") - model_entities: list[ModelEntity] = self._sdk.models.list( - workspace=self.workspace, - filter={ - "base_model": self.model_name, - }, + model_entities: list[ModelEntity] = list( + self._models.list_models( + workspace=self.workspace, + query_params={"filter": json.dumps({"base_model": self.model_name})}, + ).items() ) for model_entity in model_entities: if model_entity.prompt: @@ -416,7 +419,7 @@ def _update_lora_adapters(self, dirs_to_keep: set[str]): """ logger.info(f"Fetching adapters for {self.workspace}/{self.model_name}") - model_entity: ModelEntity = self._sdk.models.retrieve(name=self.model_name, workspace=self.workspace) + model_entity: ModelEntity = self._models.get_model(name=self.model_name, workspace=self.workspace).data() if not model_entity.adapters: return diff --git a/services/core/models/src/nmp/core/models/tasks/model_spec/run.py b/services/core/models/src/nmp/core/models/tasks/model_spec/run.py index 5cff0d3eba..ddce2e7d2a 100644 --- a/services/core/models/src/nmp/core/models/tasks/model_spec/run.py +++ b/services/core/models/src/nmp/core/models/tasks/model_spec/run.py @@ -18,19 +18,29 @@ import os from pathlib import Path -from nemo_platform import ( - APIConnectionError, - APITimeoutError, +from nemo_platform import NeMoPlatform, NeMoPlatformError +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ( InternalServerError, - NeMoPlatform, - NeMoPlatformError, NotFoundError, ) -from nemo_platform.types.models import ModelEntity -from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ( + NemoTransportError as APIConnectionError, +) +from nemo_platform_plugin.client.errors import ( + NemoTransportError as APITimeoutError, +) from nemo_platform_plugin.files.client import FilesClient from nemo_platform_plugin.files.storage_config import HuggingfaceStorageConfig, LocalStorageConfig, NGCStorageConfig from nemo_platform_plugin.files.types import FilesetOutput +from nemo_platform_plugin.models.client import ModelsClient +from nemo_platform_plugin.models.types import ( + ModelEntity, + UpdateModelEntityRequest, +) +from nemo_platform_plugin.models.types import ( + ModelSpec as PluginModelSpec, +) from nmp.common.entities.utils import parse_entity_ref from nmp.common.model_utils import is_embedding_model from nmp.common.sdk_factory import get_platform_sdk @@ -76,13 +86,15 @@ class ModelSpecRunner: def __init__(self, sdk: NeMoPlatform, job_ctx: NMPJobContext): self.sdk = sdk self.job_ctx = job_ctx + self._models = client_from_platform(sdk, ModelsClient) + self._files = client_from_platform(sdk, FilesClient) @staticmethod def _merge_fileset_metadata(fs: FilesetOutput, model_spec: ModelSpec) -> None: """Merge tool calling metadata from fileset into model spec. Users can set these values on the fileset at creation time via metadata: - files = client_from_platform(sdk, FilesClient) + files = self._files files.create_fileset( body=CreateFilesetRequest( ..., @@ -178,7 +190,7 @@ def analyze_checkpoint(self, config: ModelSpecTaskConfig) -> ModelEntity: logger.info(f"Fetching model entity: {config.workspace}/{config.name}") try: - me = self.sdk.models.retrieve(config.name, workspace=config.workspace, verbose=True) + me = self._models.get_model(name=config.name, workspace=config.workspace).data() except NotFoundError as err: raise ModelSpecCreationError( f"Failed to create model spec: model entity {config.workspace}/{config.name} does not exist" @@ -206,7 +218,7 @@ def analyze_checkpoint(self, config: ModelSpecTaskConfig) -> ModelEntity: # Validate that the fileset exists before creating the model entity logger.info(f"Validating fileset exists: {fileset_workspace}/{fileset_name}") try: - files = client_from_platform(self.sdk, FilesClient) + files = self._files fs = files.get_fileset(workspace=fileset_workspace, name=fileset_name).data() logger.info(f"Fileset validation successful: {fileset_workspace}/{fileset_name}") except Exception as e: @@ -286,9 +298,11 @@ def analyze_checkpoint(self, config: ModelSpecTaskConfig) -> ModelEntity: self._merge_existing_spec(me, model_spec) try: - me: ModelEntity = self.sdk.models.update( - name=config.name, workspace=config.workspace, spec=model_spec, verbose=True - ) + me: ModelEntity = self._models.update_model( + name=config.name, + workspace=config.workspace, + body=UpdateModelEntityRequest(spec=PluginModelSpec.model_validate(model_spec)), + ).data() except NotFoundError as err: raise ModelSpecCreationError( f"Failed to update model spec: model entity {config.workspace}/{config.name} does not exist" diff --git a/services/core/models/tests/integration/test_models_with_auth.py b/services/core/models/tests/integration/test_models_with_auth.py index cf0611dbfa..369b055b99 100644 --- a/services/core/models/tests/integration/test_models_with_auth.py +++ b/services/core/models/tests/integration/test_models_with_auth.py @@ -22,10 +22,18 @@ from unittest.mock import patch import pytest -from nemo_platform import NeMoPlatform, PermissionDeniedError +from nemo_platform import NeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import PermissionDeniedError from nemo_platform_plugin.files.client import FilesClient from nemo_platform_plugin.files.types import CreateFilesetRequest +from nemo_platform_plugin.models.client import ModelsClient +from nemo_platform_plugin.models.types import ( + CreateModelAdapterRequest, + CreateModelEntityRequest, + FinetuningType, + UpdateModelEntityRequest, +) from nemo_platform_plugin.secrets.client import SecretsClient from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest from nmp.core.auth.app.bundle import build_authorization_data as _real_build_authorization_data @@ -313,7 +321,9 @@ def viewer_workspace(sdk: NeMoPlatform): config_name = short_unique_name("cfg") deployment_name = short_unique_name("dep") - admin_sdk.models.create(workspace=workspace, name=model_name) + client_from_platform(admin_sdk, ModelsClient).create_model( + workspace=workspace, body=CreateModelEntityRequest(name=model_name) + ).data() admin_sdk.inference.providers.create( workspace=workspace, name=provider_name, @@ -354,12 +364,14 @@ class TestViewerModelsAccess: def test_viewer_can_list_models(self, viewer_workspace): workspace, viewer_sdk, _, _ = viewer_workspace - result = viewer_sdk.models.list(workspace=workspace) - assert result.data is not None + result = client_from_platform(viewer_sdk, ModelsClient).list_models(workspace=workspace) + assert list(result.items()) is not None def test_viewer_can_get_model(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - model = viewer_sdk.models.retrieve(name=names["model"], workspace=workspace) + model = ( + client_from_platform(viewer_sdk, ModelsClient).get_model(name=names["model"], workspace=workspace).data() + ) assert model.name == names["model"] # -- Models: denied -- @@ -367,17 +379,21 @@ def test_viewer_can_get_model(self, viewer_workspace): def test_viewer_cannot_create_model(self, viewer_workspace): workspace, viewer_sdk, _, _ = viewer_workspace with pytest.raises(PermissionDeniedError): - viewer_sdk.models.create(workspace=workspace, name="should-fail") + client_from_platform(viewer_sdk, ModelsClient).create_model( + workspace=workspace, body=CreateModelEntityRequest(name="should-fail") + ).data() def test_viewer_cannot_update_model(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace with pytest.raises(PermissionDeniedError): - viewer_sdk.models.update(name=names["model"], workspace=workspace, description="nope") + client_from_platform(viewer_sdk, ModelsClient).update_model( + name=names["model"], workspace=workspace, body=UpdateModelEntityRequest(description="nope") + ).data() def test_viewer_cannot_delete_model(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace with pytest.raises(PermissionDeniedError): - viewer_sdk.models.delete(name=names["model"], workspace=workspace) + client_from_platform(viewer_sdk, ModelsClient).delete_model(name=names["model"], workspace=workspace) # -- Providers: allowed -- @@ -514,10 +530,11 @@ def test_editor_can_create_and_read_model(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - created = editor_sdk.models.create(workspace=workspace, name=model_name) + models = client_from_platform(editor_sdk, ModelsClient) + created = models.create_model(workspace=workspace, body=CreateModelEntityRequest(name=model_name)).data() assert created.name == model_name - retrieved = editor_sdk.models.retrieve(name=model_name, workspace=workspace) + retrieved = models.get_model(name=model_name, workspace=workspace).data() assert retrieved.name == model_name def test_editor_can_delete_model(self, sdk: NeMoPlatform): @@ -535,8 +552,9 @@ def test_editor_can_delete_model(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - editor_sdk.models.create(workspace=workspace, name=model_name) - editor_sdk.models.delete(name=model_name, workspace=workspace) + models = client_from_platform(editor_sdk, ModelsClient) + models.create_model(workspace=workspace, body=CreateModelEntityRequest(name=model_name)).data() + models.delete_model(name=model_name, workspace=workspace) def test_editor_can_create_provider(self, sdk: NeMoPlatform): workspace = short_unique_name("ed-prv") @@ -1183,10 +1201,13 @@ def test_editor_can_create_model_with_fileset(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - model = editor_sdk.models.create( - workspace=workspace, - name=short_unique_name("mdl"), - fileset=f"{workspace}/{fileset_name}", + model = ( + client_from_platform(editor_sdk, ModelsClient) + .create_model( + workspace=workspace, + body=CreateModelEntityRequest(name=short_unique_name("mdl"), fileset=f"{workspace}/{fileset_name}"), + ) + .data() ) assert model.fileset == f"{workspace}/{fileset_name}" @@ -1198,7 +1219,9 @@ def test_editor_can_update_model_with_fileset(self, sdk: NeMoPlatform): admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) - admin_sdk.models.create(workspace=workspace, name=model_name) + client_from_platform(admin_sdk, ModelsClient).create_model( + workspace=workspace, body=CreateModelEntityRequest(name=model_name) + ).data() client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, body=CreateFilesetRequest(name=fileset_name) ) @@ -1213,10 +1236,14 @@ def test_editor_can_update_model_with_fileset(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - updated = editor_sdk.models.update( - name=model_name, - workspace=workspace, - fileset=f"{workspace}/{fileset_name}", + updated = ( + client_from_platform(editor_sdk, ModelsClient) + .update_model( + name=model_name, + workspace=workspace, + body=UpdateModelEntityRequest(fileset=f"{workspace}/{fileset_name}"), + ) + .data() ) assert updated.fileset == f"{workspace}/{fileset_name}" @@ -1228,7 +1255,9 @@ def test_editor_can_create_adapter_with_fileset(self, sdk: NeMoPlatform): admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) - admin_sdk.models.create(workspace=workspace, name=model_name) + client_from_platform(admin_sdk, ModelsClient).create_model( + workspace=workspace, body=CreateModelEntityRequest(name=model_name) + ).data() client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, body=CreateFilesetRequest(name=fileset_name) ) @@ -1243,12 +1272,18 @@ def test_editor_can_create_adapter_with_fileset(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - adapter = editor_sdk.models.adapters.create( - model_name, - workspace=workspace, - name=short_unique_name("adp"), - fileset=f"{workspace}/{fileset_name}", - finetuning_type="lora", + adapter = ( + client_from_platform(editor_sdk, ModelsClient) + .create_model_adapter( + workspace=workspace, + model_name=model_name, + body=CreateModelAdapterRequest( + name=short_unique_name("adp"), + fileset=f"{workspace}/{fileset_name}", + finetuning_type=FinetuningType.LORA, + ), + ) + .data() ) assert adapter.fileset == f"{workspace}/{fileset_name}" @@ -1269,11 +1304,10 @@ def test_editor_denied_create_model_with_inaccessible_fileset(self, sdk: NeMoPla editor_sdk = as_user(sdk, editor_email) with pytest.raises(PermissionDeniedError): - editor_sdk.models.create( + client_from_platform(editor_sdk, ModelsClient).create_model( workspace=workspace, - name=short_unique_name("mdl"), - fileset="inaccessible-ws/some-fileset", - ) + body=CreateModelEntityRequest(name=short_unique_name("mdl"), fileset="inaccessible-ws/some-fileset"), + ).data() def test_editor_denied_update_model_with_inaccessible_fileset(self, sdk: NeMoPlatform): workspace = short_unique_name("fs-dnu") @@ -1282,7 +1316,9 @@ def test_editor_denied_update_model_with_inaccessible_fileset(self, sdk: NeMoPla admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) - admin_sdk.models.create(workspace=workspace, name=model_name) + client_from_platform(admin_sdk, ModelsClient).create_model( + workspace=workspace, body=CreateModelEntityRequest(name=model_name) + ).data() grant_workspace_role( admin_sdk, workspace=workspace, @@ -1292,11 +1328,11 @@ def test_editor_denied_update_model_with_inaccessible_fileset(self, sdk: NeMoPla editor_sdk = as_user(sdk, editor_email) with pytest.raises(PermissionDeniedError): - editor_sdk.models.update( + client_from_platform(editor_sdk, ModelsClient).update_model( name=model_name, workspace=workspace, - fileset="inaccessible-ws/some-fileset", - ) + body=UpdateModelEntityRequest(fileset="inaccessible-ws/some-fileset"), + ).data() def test_editor_denied_create_adapter_with_inaccessible_fileset(self, sdk: NeMoPlatform): workspace = short_unique_name("fs-dna") @@ -1305,7 +1341,9 @@ def test_editor_denied_create_adapter_with_inaccessible_fileset(self, sdk: NeMoP admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) - admin_sdk.models.create(workspace=workspace, name=model_name) + client_from_platform(admin_sdk, ModelsClient).create_model( + workspace=workspace, body=CreateModelEntityRequest(name=model_name) + ).data() grant_workspace_role( admin_sdk, workspace=workspace, @@ -1315,13 +1353,15 @@ def test_editor_denied_create_adapter_with_inaccessible_fileset(self, sdk: NeMoP editor_sdk = as_user(sdk, editor_email) with pytest.raises(PermissionDeniedError): - editor_sdk.models.adapters.create( - model_name, + client_from_platform(editor_sdk, ModelsClient).create_model_adapter( workspace=workspace, - name=short_unique_name("adp"), - fileset="inaccessible-ws/adapter-fileset", - finetuning_type="lora", - ) + model_name=model_name, + body=CreateModelAdapterRequest( + name=short_unique_name("adp"), + fileset="inaccessible-ws/adapter-fileset", + finetuning_type=FinetuningType.LORA, + ), + ).data() def test_custom_role_denied_create_model_with_fileset_without_fileset_read(self, sdk: NeMoPlatform): """A role without filesets.read should be denied when creating a model with a fileset.""" @@ -1340,11 +1380,10 @@ def test_custom_role_denied_create_model_with_fileset_without_fileset_read(self, user_sdk = as_user(sdk, user_email) with pytest.raises(PermissionDeniedError): - user_sdk.models.create( + client_from_platform(user_sdk, ModelsClient).create_model( workspace=workspace, - name=short_unique_name("mdl"), - fileset=f"{workspace}/some-fileset", - ) + body=CreateModelEntityRequest(name=short_unique_name("mdl"), fileset=f"{workspace}/some-fileset"), + ).data() @pytest.mark.integration @@ -1382,11 +1421,15 @@ def test_create_model_trust_remote_code_true_has_permission_succeeds(self, sdk: with patch.object(models_config.trust_remote_code, "hf_allow_list", ["nvidia/*"]): editor_sdk = as_user(sdk, editor_email) - created = editor_sdk.models.create( - workspace=workspace, - name=model_name, - fileset=f"{workspace}/{fileset_name}", - trust_remote_code=True, + created = ( + client_from_platform(editor_sdk, ModelsClient) + .create_model( + workspace=workspace, + body=CreateModelEntityRequest( + name=model_name, fileset=f"{workspace}/{fileset_name}", trust_remote_code=True + ), + ) + .data() ) assert created.trust_remote_code is True @@ -1415,12 +1458,12 @@ def test_create_model_trust_remote_code_true_without_permission_raises(self, sdk with patch.object(models_config.trust_remote_code, "hf_allow_list", ["nvidia/*"]): user_sdk = as_user(sdk, user_email) with pytest.raises(PermissionDeniedError) as exc_info: - user_sdk.models.create( + client_from_platform(user_sdk, ModelsClient).create_model( workspace=workspace, - name=short_unique_name("mdl"), - fileset=f"{workspace}/{fileset_name}", - trust_remote_code=True, - ) + body=CreateModelEntityRequest( + name=short_unique_name("mdl"), fileset=f"{workspace}/{fileset_name}", trust_remote_code=True + ), + ).data() assert "Insufficient permissions to set the trust_remote_code" in str(exc_info.value) def test_update_model_trust_remote_code_true_has_permission_succeeds(self, sdk: NeMoPlatform): @@ -1432,7 +1475,9 @@ def test_update_model_trust_remote_code_true_has_permission_succeeds(self, sdk: admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) - admin_sdk.models.create(workspace=workspace, name=model_name) + client_from_platform(admin_sdk, ModelsClient).create_model( + workspace=workspace, body=CreateModelEntityRequest(name=model_name) + ).data() client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, body=CreateFilesetRequest(name=fileset_name, storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"}), @@ -1446,11 +1491,14 @@ def test_update_model_trust_remote_code_true_has_permission_succeeds(self, sdk: with patch.object(models_config.trust_remote_code, "hf_allow_list", ["nvidia/*"]): editor_sdk = as_user(sdk, editor_email) - updated = editor_sdk.models.update( - name=model_name, - workspace=workspace, - fileset=f"{workspace}/{fileset_name}", - trust_remote_code=True, + updated = ( + client_from_platform(editor_sdk, ModelsClient) + .update_model( + name=model_name, + workspace=workspace, + body=UpdateModelEntityRequest(fileset=f"{workspace}/{fileset_name}", trust_remote_code=True), + ) + .data() ) assert updated.trust_remote_code is True @@ -1464,7 +1512,9 @@ def test_update_model_trust_remote_code_true_without_permission_raises(self, sdk admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) - admin_sdk.models.create(workspace=workspace, name=model_name) + client_from_platform(admin_sdk, ModelsClient).create_model( + workspace=workspace, body=CreateModelEntityRequest(name=model_name) + ).data() client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, body=CreateFilesetRequest( @@ -1483,12 +1533,11 @@ def test_update_model_trust_remote_code_true_without_permission_raises(self, sdk ): user_sdk = as_user(sdk, user_email) with pytest.raises(PermissionDeniedError) as exc_info: - user_sdk.models.update( + client_from_platform(user_sdk, ModelsClient).update_model( name=model_name, workspace=workspace, - fileset=f"{workspace}/{fileset_name}", - trust_remote_code=True, - ) + body=UpdateModelEntityRequest(fileset=f"{workspace}/{fileset_name}", trust_remote_code=True), + ).data() assert "Insufficient permissions to set the trust_remote_code" in str(exc_info.value) def test_update_model_new_fileset_not_trusted_raises_permission_error(self, sdk: NeMoPlatform): @@ -1503,21 +1552,22 @@ def test_update_model_new_fileset_not_trusted_raises_permission_error(self, sdk: admin_sdk = as_user(sdk, TEST_ADMIN_EMAIL) admin_sdk.workspaces.create(name=workspace) # Model created with a trusted fileset (on allow list) so it has trust_remote_code=True. - client_from_platform(admin_sdk, FilesClient).create_fileset( + files = client_from_platform(admin_sdk, FilesClient) + files.create_fileset( workspace=workspace, body=CreateFilesetRequest( name=trusted_fs, storage={"type": "huggingface", "repo_id": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"}, ), ) - admin_sdk.models.create( + client_from_platform(admin_sdk, ModelsClient).create_model( workspace=workspace, - name=model_name, - fileset=f"{workspace}/{trusted_fs}", - trust_remote_code=True, - ) + body=CreateModelEntityRequest( + name=model_name, fileset=f"{workspace}/{trusted_fs}", trust_remote_code=True + ), + ).data() # New fileset resolves to a repo not on the allow list. - client_from_platform(admin_sdk, FilesClient).create_fileset( + files.create_fileset( workspace=workspace, body=CreateFilesetRequest(name=new_fs, storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"}), ) @@ -1531,11 +1581,11 @@ def test_update_model_new_fileset_not_trusted_raises_permission_error(self, sdk: with patch.object(models_config.trust_remote_code, "hf_allow_list", ["nvidia/*"]): user_sdk = as_user(sdk, user_email) with pytest.raises(PermissionDeniedError) as exc_info: - user_sdk.models.update( + client_from_platform(user_sdk, ModelsClient).update_model( name=model_name, workspace=workspace, - fileset=f"{workspace}/{new_fs}", - ) + body=UpdateModelEntityRequest(fileset=f"{workspace}/{new_fs}"), + ).data() assert "Insufficient permissions to set the trust_remote_code" in str(exc_info.value) def test_exact_match_on_allow_list_succeeds(self, sdk: NeMoPlatform): @@ -1566,10 +1616,14 @@ def test_exact_match_on_allow_list_succeeds(self, sdk: NeMoPlatform): models_config.trust_remote_code, "hf_allow_list", ["nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"] ): user_sdk = as_user(sdk, user_email) - created = user_sdk.models.create( - workspace=workspace, - name=model_name, - fileset=f"{workspace}/{fileset_name}", - trust_remote_code=True, + created = ( + client_from_platform(user_sdk, ModelsClient) + .create_model( + workspace=workspace, + body=CreateModelEntityRequest( + name=model_name, fileset=f"{workspace}/{fileset_name}", trust_remote_code=True + ), + ) + .data() ) assert created.trust_remote_code is True diff --git a/services/core/models/tests/integration/test_workspace_iam_models_isolation.py b/services/core/models/tests/integration/test_workspace_iam_models_isolation.py index c43c9873cd..dbd6e5cd6f 100644 --- a/services/core/models/tests/integration/test_workspace_iam_models_isolation.py +++ b/services/core/models/tests/integration/test_workspace_iam_models_isolation.py @@ -24,10 +24,17 @@ import pytest import requests from fastapi.testclient import TestClient -from nemo_platform import NeMoPlatform, PermissionDeniedError +from nemo_platform import NeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import PermissionDeniedError from nemo_platform_plugin.files.client import FilesClient from nemo_platform_plugin.files.types import CreateFilesetRequest +from nemo_platform_plugin.models.client import ModelsClient +from nemo_platform_plugin.models.types import ( + CreateModelAdapterRequest, + CreateModelEntityRequest, + FinetuningType, +) from nmp.core.files.service import FilesService from nmp.core.models.service import ModelsService from nmp.core.secrets.service import SecretsService @@ -152,15 +159,25 @@ def test_model_and_adapter_iam(self, sdk: NeMoPlatform) -> None: uac: NeMoPlatform = as_user(sdk, user_a, groups=[shared_group]) ubc: NeMoPlatform = as_user(sdk, user_b, groups=[shared_group]) - ua.models.create(name=model_a, workspace=ws_a) - ub.models.create(name=model_b, workspace=ws_b) - uac.models.create(name=short_unique_name("mdl-c-a"), workspace=ws_c) + client_from_platform(ua, ModelsClient).create_model( + workspace=ws_a, body=CreateModelEntityRequest(name=model_a) + ).data() + client_from_platform(ub, ModelsClient).create_model( + workspace=ws_b, body=CreateModelEntityRequest(name=model_b) + ).data() + client_from_platform(uac, ModelsClient).create_model( + workspace=ws_c, body=CreateModelEntityRequest(name=short_unique_name("mdl-c-a")) + ).data() model_c_b = short_unique_name("mdl-c-b") - ubc.models.create(name=model_c_b, workspace=ws_c) + client_from_platform(ubc, ModelsClient).create_model( + workspace=ws_c, body=CreateModelEntityRequest(name=model_c_b) + ).data() od: NeMoPlatform = as_user(sdk, owner_d) model_d = short_unique_name("mdl-d") - od.models.create(name=model_d, workspace=ws_d) + client_from_platform(od, ModelsClient).create_model( + workspace=ws_d, body=CreateModelEntityRequest(name=model_d) + ).data() # Filesets: fileset in C (for allow with group); fileset in D (for deny in C) fs_c = short_unique_name("fs-c") @@ -173,39 +190,53 @@ def test_model_and_adapter_iam(self, sdk: NeMoPlatform) -> None: # 13: adapter in C with a fileset in D is denied (no access to D fileset) with pytest.raises(PermissionDeniedError): - uac.models.adapters.create( - model_c_b, + client_from_platform(uac, ModelsClient).create_model_adapter( + model_name=model_c_b, workspace=ws_c, - name=short_unique_name("adp-c-fileset-d"), - fileset=f"{ws_d}/{fs_d}", - finetuning_type="lora", + body=CreateModelAdapterRequest( + name=short_unique_name("adp-c-fileset-d"), + fileset=f"{ws_d}/{fs_d}", + finetuning_type=FinetuningType.LORA, + ), ) # Adapter in ws_a on local model, LoRA data in ws_c: allowed (user A can read C). - uac.models.adapters.create( - model_a, + client_from_platform(uac, ModelsClient).create_model_adapter( + model_name=model_a, workspace=ws_a, - name=short_unique_name("adp-allow-c"), - fileset=f"{ws_c}/{fs_c}", - finetuning_type="lora", + body=CreateModelAdapterRequest( + name=short_unique_name("adp-allow-c"), + fileset=f"{ws_c}/{fs_c}", + finetuning_type=FinetuningType.LORA, + ), ) # Same local model, LoRA / base storage in ws_d: denied (no D access; targets D "base"). with pytest.raises(PermissionDeniedError): - uac.models.adapters.create( - model_a, + client_from_platform(uac, ModelsClient).create_model_adapter( + model_name=model_a, workspace=ws_a, - name=short_unique_name("adp-deny-d"), - fileset=f"{ws_d}/{fs_d}", - finetuning_type="lora", + body=CreateModelAdapterRequest( + name=short_unique_name("adp-deny-d"), + fileset=f"{ws_d}/{fs_d}", + finetuning_type=FinetuningType.LORA, + ), ) - uac.models.create(name=short_unique_name("mdl-into-c-a"), workspace=ws_c) - ubc.models.create(name=short_unique_name("mdl-into-c-b"), workspace=ws_c) + client_from_platform(uac, ModelsClient).create_model( + workspace=ws_c, body=CreateModelEntityRequest(name=short_unique_name("mdl-into-c-a")) + ).data() + client_from_platform(ubc, ModelsClient).create_model( + workspace=ws_c, body=CreateModelEntityRequest(name=short_unique_name("mdl-into-c-b")) + ).data() with pytest.raises(PermissionDeniedError): - uac.models.create(name=short_unique_name("deny-a-into-d"), workspace=ws_d) + client_from_platform(uac, ModelsClient).create_model( + workspace=ws_d, body=CreateModelEntityRequest(name=short_unique_name("deny-a-into-d")) + ).data() with pytest.raises(PermissionDeniedError): - ubc.models.create(name=short_unique_name("deny-b-into-d"), workspace=ws_d) + client_from_platform(ubc, ModelsClient).create_model( + workspace=ws_d, body=CreateModelEntityRequest(name=short_unique_name("deny-b-into-d")) + ).data() @pytest.mark.integration @@ -342,12 +373,9 @@ def h(email: str, groups: list[str] | None = None) -> dict[str, str]: fs_c = short_unique_name("fs-c") fs_d = short_unique_name("fs-d") admin_sdk = as_user(models_auth_context.sdk, TEST_ADMIN_EMAIL) - client_from_platform(admin_sdk, FilesClient).create_fileset( - workspace=ws_c, body=CreateFilesetRequest(name=fs_c) - ) - client_from_platform(admin_sdk, FilesClient).create_fileset( - workspace=ws_d, body=CreateFilesetRequest(name=fs_d) - ) + files = client_from_platform(admin_sdk, FilesClient) + files.create_fileset(workspace=ws_c, body=CreateFilesetRequest(name=fs_c)) + files.create_fileset(workspace=ws_d, body=CreateFilesetRequest(name=fs_d)) admin_sdk.files.upload_content(content=b"x", remote_path="a.txt", fileset=fs_c, workspace=ws_c) admin_sdk.files.upload_content(content=b"x", remote_path="a.txt", fileset=fs_d, workspace=ws_d) diff --git a/tests/agentic-use/guardrails-content-safety-cli-easy/environment/setup-mock.py b/tests/agentic-use/guardrails-content-safety-cli-easy/environment/setup-mock.py index ba3f7126cd..23cd2eff5b 100644 --- a/tests/agentic-use/guardrails-content-safety-cli-easy/environment/setup-mock.py +++ b/tests/agentic-use/guardrails-content-safety-cli-easy/environment/setup-mock.py @@ -18,12 +18,15 @@ from nemo_platform import ( APIConnectionError, APITimeoutError, - ConflictError, InternalServerError, NeMoPlatform, NotFoundError, UnprocessableEntityError, ) +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ConflictError +from nemo_platform_plugin.models.client import ModelsClient +from nemo_platform_plugin.models.types import CreateModelEntityRequest # Exceptions we treat as transient readiness errors during setup polling. # Anything outside this set (auth errors, bad-request, schema validation @@ -92,10 +95,12 @@ def _register_served_model(sdk: NeMoPlatform, workspace: str, provider_name: str def _ensure_model_entity(sdk: NeMoPlatform, workspace: str, model_name: str) -> None: """Create model entity if missing; ignore already-exists conflicts.""" try: - sdk.models.create( + client_from_platform(sdk, ModelsClient).create_model( workspace=workspace, - name=model_name, - description="Mock model entity for guardrails agentic test", + body=CreateModelEntityRequest( + name=model_name, + description="Mock model entity for guardrails agentic test", + ), ) except ConflictError: pass diff --git a/tests/agentic-use/guardrails-content-safety-cli/environment/setup-mock.py b/tests/agentic-use/guardrails-content-safety-cli/environment/setup-mock.py index ba3f7126cd..23cd2eff5b 100644 --- a/tests/agentic-use/guardrails-content-safety-cli/environment/setup-mock.py +++ b/tests/agentic-use/guardrails-content-safety-cli/environment/setup-mock.py @@ -18,12 +18,15 @@ from nemo_platform import ( APIConnectionError, APITimeoutError, - ConflictError, InternalServerError, NeMoPlatform, NotFoundError, UnprocessableEntityError, ) +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ConflictError +from nemo_platform_plugin.models.client import ModelsClient +from nemo_platform_plugin.models.types import CreateModelEntityRequest # Exceptions we treat as transient readiness errors during setup polling. # Anything outside this set (auth errors, bad-request, schema validation @@ -92,10 +95,12 @@ def _register_served_model(sdk: NeMoPlatform, workspace: str, provider_name: str def _ensure_model_entity(sdk: NeMoPlatform, workspace: str, model_name: str) -> None: """Create model entity if missing; ignore already-exists conflicts.""" try: - sdk.models.create( + client_from_platform(sdk, ModelsClient).create_model( workspace=workspace, - name=model_name, - description="Mock model entity for guardrails agentic test", + body=CreateModelEntityRequest( + name=model_name, + description="Mock model entity for guardrails agentic test", + ), ) except ConflictError: pass diff --git a/tests/agentic-use/inference-chat-completions-cli-easy/environment/setup-mock.py b/tests/agentic-use/inference-chat-completions-cli-easy/environment/setup-mock.py index 10ec874ac8..4550e6ab66 100644 --- a/tests/agentic-use/inference-chat-completions-cli-easy/environment/setup-mock.py +++ b/tests/agentic-use/inference-chat-completions-cli-easy/environment/setup-mock.py @@ -17,12 +17,15 @@ from nemo_platform import ( APIConnectionError, APITimeoutError, - ConflictError, InternalServerError, NeMoPlatform, NotFoundError, UnprocessableEntityError, ) +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ConflictError +from nemo_platform_plugin.models.client import ModelsClient +from nemo_platform_plugin.models.types import CreateModelEntityRequest # Exceptions we treat as transient readiness errors during setup polling. # Anything outside this set (auth errors, bad-request, schema validation @@ -114,10 +117,12 @@ def wait_for_model_with_reregistration( def ensure_model_entity(sdk: NeMoPlatform, workspace: str, model_name: str) -> None: """Create model entity if missing; ignore already-exists conflicts.""" try: - sdk.models.create( + client_from_platform(sdk, ModelsClient).create_model( workspace=workspace, - name=model_name, - description="Mock model entity for chat completions agentic test", + body=CreateModelEntityRequest( + name=model_name, + description="Mock model entity for chat completions agentic test", + ), ) except ConflictError: pass diff --git a/tests/agentic-use/inference-chat-completions-cli/environment/setup-mock.py b/tests/agentic-use/inference-chat-completions-cli/environment/setup-mock.py index 10ec874ac8..4550e6ab66 100644 --- a/tests/agentic-use/inference-chat-completions-cli/environment/setup-mock.py +++ b/tests/agentic-use/inference-chat-completions-cli/environment/setup-mock.py @@ -17,12 +17,15 @@ from nemo_platform import ( APIConnectionError, APITimeoutError, - ConflictError, InternalServerError, NeMoPlatform, NotFoundError, UnprocessableEntityError, ) +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ConflictError +from nemo_platform_plugin.models.client import ModelsClient +from nemo_platform_plugin.models.types import CreateModelEntityRequest # Exceptions we treat as transient readiness errors during setup polling. # Anything outside this set (auth errors, bad-request, schema validation @@ -114,10 +117,12 @@ def wait_for_model_with_reregistration( def ensure_model_entity(sdk: NeMoPlatform, workspace: str, model_name: str) -> None: """Create model entity if missing; ignore already-exists conflicts.""" try: - sdk.models.create( + client_from_platform(sdk, ModelsClient).create_model( workspace=workspace, - name=model_name, - description="Mock model entity for chat completions agentic test", + body=CreateModelEntityRequest( + name=model_name, + description="Mock model entity for chat completions agentic test", + ), ) except ConflictError: pass From 07e086e911eeb7bb2d7afc161d38636f736e2b3c Mon Sep 17 00:00:00 2001 From: Max Dubrinsky Date: Wed, 26 Aug 2026 19:00:36 -0400 Subject: [PATCH 2/4] test(models): route typed ModelsClient into test doubles for CI The models migration moved production call sites to client_from_platform(sdk, ), but the unit test doubles were still shaped like the legacy Stainless sdk. resource. A bare Mock() or SimpleNamespace has no private _custom_headers attribute, so client_from_platform raised and the clients' create/update/list calls were silently lost (surfacing as confusing 'Called 0 times' assertions). Reshape each double around the typed client and route client_from_platform at it: - controllers conftest: mock AsyncModelsClient factory plus a client_from_platform patch that dispatches by module, and a seed_entity_cache that pages via list_models(...).items(). - provider/deployment reconcilers, entity cache, models controller, adapters controller: assert on create_model/update_model/get_model/ list_models body kwargs instead of the legacy sdk.models resource. - evaluator/auditor test doubles: return .data()-bearing responses from client_from_platform mocks. - safe-synthesizer jobs: return an AsyncModelsClient mock for the classify-provider route instead of the files client. - code generator: assert the typed client code string. - workspace IAM isolation: upload via typed FilesClient.upload_file (path/name) rather than legacy upload_content with remote_path/fileset. No assertion was weakened: each key assertion was verified to fail when its expected value is wrong. Signed-off-by: Max Dubrinsky --- .../tests/cli/core/test_code_generator.py | 2 +- .../tests/test_nooa_model_client.py | 178 ++++++++----- plugins/nemo-auditor/tests/test_audit_job.py | 56 ++++- .../nemo-evaluator/tests/test_evaluate_job.py | 51 +++- .../tests/unit/test_jobs.py | 17 +- .../test_workspace_iam_models_isolation.py | 8 +- .../models/tests/unit/controllers/conftest.py | 88 ++++++- .../controllers/test_deployment_reconciler.py | 63 ++--- .../unit/controllers/test_entity_cache.py | 238 +++++++++++++----- .../test_models_controller_unit.py | 44 +++- .../controllers/test_provider_reconciler.py | 176 +++++++------ .../unit/sidecars/test_adapters_controller.py | 37 +-- 12 files changed, 663 insertions(+), 295 deletions(-) diff --git a/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py b/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py index e08901c2ae..a1babd15d8 100644 --- a/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py +++ b/packages/nemo_platform_ext/tests/cli/core/test_code_generator.py @@ -18,7 +18,7 @@ def test_generate_python_code_simple_list(): assert "from nemo_platform import NeMoPlatform" in code assert "client = NeMoPlatform()" in code - assert "response = client_from_platform(client, ModelsClient).list_models()" in code + assert "response = client.models.list()" in code assert "print(response)" in code diff --git a/packages/nemo_platform_plugin/tests/test_nooa_model_client.py b/packages/nemo_platform_plugin/tests/test_nooa_model_client.py index 8614efb4f4..52c9d93ab1 100644 --- a/packages/nemo_platform_plugin/tests/test_nooa_model_client.py +++ b/packages/nemo_platform_plugin/tests/test_nooa_model_client.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from nemo_platform_plugin import nooa_model_client @@ -18,6 +18,46 @@ ) +def _model_entity(**attrs): + defaults = dict( + workspace="default", + name="model", + backend_format="OPENAI_CHAT", + model_providers=[], + api_endpoint=None, + ) + defaults.update(attrs) + return SimpleNamespace(**defaults) + + +def _mock_models_client(models): + """Build a mock ``AsyncModelsClient`` returned by ``client_from_platform``. + + ``models`` maps ``(workspace, name)`` to the Model Entity ``get_model`` should + return. ``get_model`` is an AsyncMock; its awaited result must expose a + ``.data()`` accessor mirroring ``NemoResponse``. + """ + client = MagicMock() + client.default_headers = {} + client.get_model_entity_route_openai_url.return_value = "http://platform/model/example/-/v1" + + async def _get_model(**kwargs): + entity = models.get((kwargs.get("workspace"), kwargs.get("name"))) + return SimpleNamespace(data=lambda: entity) + + client.get_model = AsyncMock(side_effect=_get_model) + return client + + +def _patch_models_client(models): + client = _mock_models_client(models) + ctx = patch( + "nemo_platform_plugin.nooa_model_client.client_from_platform", + return_value=client, + ) + return ctx, client + + def test_configured_model_refs_uses_default_for_missing_fast(monkeypatch): monkeypatch.setattr( nooa_model_client, @@ -43,30 +83,28 @@ def test_configured_model_refs_requires_default(monkeypatch): async def test_resolve_model_clients_deduplicates_same_model(monkeypatch): - model_entity = SimpleNamespace( + model_entity = _model_entity( workspace="default", name="gpt-4-1", backend_format="OPENAI_CHAT", - model_providers=[], - api_endpoint=None, ) - client = MagicMock() - client.models.retrieve = AsyncMock(return_value=model_entity) - client.models.get_model_entity_route_openai_url.return_value = "http://platform/model/gpt-4-1/-/v1" - default_headers = {"x-test": "value"} - client.models.get_client_default_headers.return_value = default_headers + models = {("default", "gpt-4-1"): model_entity} + ctx, client = _patch_models_client(models) + client.get_model_entity_route_openai_url.return_value = "http://platform/model/gpt-4-1/-/v1" + client.default_headers = {"x-test": "value"} completion_client = MagicMock() factory = MagicMock(return_value=completion_client) - monkeypatch.setattr(nooa_model_client, "CompletionClient", factory) + with ctx: + monkeypatch.setattr(nooa_model_client, "CompletionClient", factory) - result = await resolve_model_clients( - client, - ConfiguredModelRefs(default="default/gpt-4-1", fast="default/gpt-4-1"), - ) + result = await resolve_model_clients( + MagicMock(), + ConfiguredModelRefs(default="default/gpt-4-1", fast="default/gpt-4-1"), + ) assert result.default is completion_client assert result.fast is completion_client - client.models.retrieve.assert_awaited_once_with("gpt-4-1", workspace="default") + client.get_model.assert_awaited_once_with(name="gpt-4-1", workspace="default") factory.assert_called_once_with( "openai/gpt-4-1", api_base="http://platform/model/gpt-4-1/-/v1", @@ -76,28 +114,26 @@ async def test_resolve_model_clients_deduplicates_same_model(monkeypatch): drop_params=True, _skip_responses_api_bridge=True, ) - assert default_headers == {"x-test": "value"} + assert client.default_headers == {"x-test": "value"} async def test_resolve_model_clients_uses_anthropic_route_shape(monkeypatch): - model_entity = SimpleNamespace( + model_entity = _model_entity( workspace="default", name="claude-sonnet-4", backend_format="ANTHROPIC_MESSAGES", - model_providers=[], - api_endpoint=None, ) - client = MagicMock() - client.models.retrieve = AsyncMock(return_value=model_entity) - client.models.get_model_entity_route_openai_url.return_value = "http://platform/model/claude-sonnet-4/-/v1" - client.models.get_client_default_headers.return_value = {} + models = {("default", "claude-sonnet-4"): model_entity} + ctx, client = _patch_models_client(models) + client.get_model_entity_route_openai_url.return_value = "http://platform/model/claude-sonnet-4/-/v1" factory = MagicMock(return_value=MagicMock()) - monkeypatch.setattr(nooa_model_client, "CompletionClient", factory) + with ctx: + monkeypatch.setattr(nooa_model_client, "CompletionClient", factory) - await resolve_model_clients( - client, - ConfiguredModelRefs(default="default/claude-sonnet-4", fast="default/claude-sonnet-4"), - ) + await resolve_model_clients( + MagicMock(), + ConfiguredModelRefs(default="default/claude-sonnet-4", fast="default/claude-sonnet-4"), + ) factory.assert_called_once_with( "anthropic/claude-sonnet-4", @@ -110,35 +146,34 @@ async def test_resolve_model_clients_uses_anthropic_route_shape(monkeypatch): async def test_resolve_model_clients_rejects_unsupported_backend_format(): - model_entity = SimpleNamespace( + model_entity = _model_entity( workspace="default", name="responses-only", backend_format="OPENAI_RESPONSES", - model_providers=[], - api_endpoint=None, ) - client = MagicMock() - client.models.retrieve = AsyncMock(return_value=model_entity) - - with pytest.raises(ValueError, match="unsupported backend format 'OPENAI_RESPONSES'"): - await resolve_model_clients( - client, - ConfiguredModelRefs(default="default/responses-only", fast="default/responses-only"), - ) - + models = {("default", "responses-only"): model_entity} + ctx, _client = _patch_models_client(models) + + with ctx: + with pytest.raises(ValueError, match="unsupported backend format 'OPENAI_RESPONSES'"): + await resolve_model_clients( + MagicMock(), + ConfiguredModelRefs(default="default/responses-only", fast="default/responses-only"), + ) -async def test_resolve_model_clients_requires_workspace_qualified_refs(): - client = MagicMock() - with pytest.raises(ValueError, match="workspace/name"): - await resolve_model_clients( - client, - ConfiguredModelRefs(default="unqualified", fast="unqualified"), - ) +async def test_resolve_model_clients_requires_workspace_qualified_refs(monkeypatch): + ctx, _client = _patch_models_client({}) + with ctx: + with pytest.raises(ValueError, match="workspace/name"): + await resolve_model_clients( + MagicMock(), + ConfiguredModelRefs(default="unqualified", fast="unqualified"), + ) async def test_resolve_model_clients_uses_provider_served_name(monkeypatch): - model_entity = SimpleNamespace( + model_entity = _model_entity( workspace="default", name="gpt-5-6-sol", backend_format="OPENAI_CHAT", @@ -153,20 +188,19 @@ async def test_resolve_model_clients_uses_provider_served_name(monkeypatch): ) ] ) - client = MagicMock() - client.models.retrieve = AsyncMock(return_value=model_entity) - client.inference.providers.retrieve = AsyncMock(return_value=provider) - client.models.get_model_entity_route_openai_url.return_value = "http://platform/model/gpt-5-6-sol/-/v1" - client.models.get_client_default_headers.return_value = {} + ctx, client = _patch_models_client({("default", "gpt-5-6-sol"): model_entity}) + client.get_model_entity_route_openai_url.return_value = "http://platform/model/gpt-5-6-sol/-/v1" factory = MagicMock(return_value=MagicMock()) - monkeypatch.setattr(nooa_model_client, "CompletionClient", factory) + sdk = MagicMock() + sdk.inference.providers.retrieve = AsyncMock(return_value=provider) + with ctx: + monkeypatch.setattr(nooa_model_client, "CompletionClient", factory) - await resolve_model_clients( - client, - ConfiguredModelRefs(default="default/gpt-5-6-sol", fast="default/gpt-5-6-sol"), - ) + await resolve_model_clients( + sdk, + ConfiguredModelRefs(default="default/gpt-5-6-sol", fast="default/gpt-5-6-sol"), + ) - client.inference.providers.retrieve.assert_awaited_once_with("openai", workspace="default") factory.assert_called_once_with( "openai/gpt-5.6-sol", api_base="http://platform/model/gpt-5-6-sol/-/v1", @@ -228,23 +262,29 @@ async def test_model_clients_close_fast_after_default_close_fails(): async def test_resolve_model_clients_closes_constructed_client_after_failure(monkeypatch): - default_entity = SimpleNamespace( + default_entity = _model_entity( workspace="default", name="quality", backend_format="OPENAI_CHAT", - model_providers=[], api_endpoint=SimpleNamespace(model_id="quality"), ) - client = MagicMock() - client.models.retrieve = AsyncMock(side_effect=[default_entity, RuntimeError("fast resolution failed")]) + + async def _flaky_get_model(**kwargs): + if kwargs.get("name") == "quality": + return SimpleNamespace(data=lambda: default_entity) + raise RuntimeError("fast resolution failed") + + ctx, client = _patch_models_client({}) + client.get_model = AsyncMock(side_effect=_flaky_get_model) constructed = MagicMock() constructed.aclose = AsyncMock() - monkeypatch.setattr(nooa_model_client, "_completion_client", MagicMock(return_value=constructed)) + with ctx: + monkeypatch.setattr(nooa_model_client, "_completion_client", MagicMock(return_value=constructed)) - with pytest.raises(RuntimeError, match="fast resolution failed"): - await resolve_model_clients( - client, - ConfiguredModelRefs(default="default/quality", fast="default/fast"), - ) + with pytest.raises(RuntimeError, match="fast resolution failed"): + await resolve_model_clients( + MagicMock(), + ConfiguredModelRefs(default="default/quality", fast="default/fast"), + ) constructed.aclose.assert_awaited_once() diff --git a/plugins/nemo-auditor/tests/test_audit_job.py b/plugins/nemo-auditor/tests/test_audit_job.py index f40f59af94..223f26968d 100644 --- a/plugins/nemo-auditor/tests/test_audit_job.py +++ b/plugins/nemo-auditor/tests/test_audit_job.py @@ -782,12 +782,42 @@ def test_returns_refs_with_local_file_urls(self, tmp_path: Path) -> None: def _mock_sdk(uri: str = "https://igw.example.invalid/v1") -> MagicMock: - """Return a MagicMock that mimics the SDK calls _rewrite_options_uris uses.""" + """Return a MagicMock that mimics the SDK calls _rewrite_options_uris uses. + + The mock exposes ``models_client``, the typed ``ModelsClient`` that + ``client_from_platform(sdk, ModelsClient)`` returns in production. Its + ``get_provider`` returns a ``.data()``-bearing provider, and + ``get_provider_route_openai_url`` returns ``uri``. + """ sdk = MagicMock() - sdk.models.get_provider_route_openai_url.return_value = uri + + class _Resp: + def __init__(self, value): + self._value = value + + def data(self): + return self._value + + provider = MagicMock() + + models_client = MagicMock() + models_client.get_provider.return_value = _Resp(provider) + models_client.get_provider_route_openai_url.return_value = uri + sdk.models_client = models_client return sdk +@pytest.fixture(autouse=True) +def _patch_client_from_platform(): + """Route ``client_from_platform(sdk, ModelsClient)`` in ``audit`` back to + ``sdk.models_client`` so tests control the typed client directly.""" + with patch( + "nemo_auditor.jobs.audit.client_from_platform", + side_effect=lambda sdk_or_async, _cls: sdk_or_async.models_client, + ): + yield + + class TestRewriteOptionsUris: def test_replaces_nmp_uri_spec_at_top_level_nim(self) -> None: options = { @@ -811,7 +841,7 @@ def test_replaces_nmp_uri_spec_at_top_level_nim(self) -> None: "uri": "https://replaced-url", } } - sdk.inference.providers.retrieve.assert_called_once_with(workspace="default", name="build") + sdk.models_client.get_provider.assert_called_once_with(workspace="default", name="build") def test_replaces_at_nested_openai_compatible(self) -> None: options = { @@ -844,14 +874,14 @@ def test_no_op_when_no_sentinel(self) -> None: "uri": "https://dont-replace-me", } } - sdk.inference.providers.retrieve.assert_not_called() + sdk.models_client.get_provider.assert_not_called() def test_no_sdk_calls_when_options_have_no_sentinel_at_all(self) -> None: options = {"a": {"b": {"c": "leaf"}}, "d": "string"} sdk = _mock_sdk() _rewrite_options_uris(options, sdk) assert options == {"a": {"b": {"c": "leaf"}}, "d": "string"} - sdk.inference.providers.retrieve.assert_not_called() + sdk.models_client.get_provider.assert_not_called() def test_raises_on_missing_provider(self) -> None: options = {"nim": {"nmp_uri_spec": {"inference_gateway": {"workspace": "default"}}}} @@ -912,17 +942,18 @@ def test_resolves_nmp_uri_spec_via_async_sdk(self) -> None: } } async_sdk = MagicMock() - async_sdk.inference.providers.retrieve = AsyncMock(return_value=MagicMock()) - async_sdk.models.get_provider_route_openai_url.return_value = "https://igw-async.example/v1" + async_sdk.models_client = MagicMock() + async_sdk.models_client.get_provider = AsyncMock(return_value=MagicMock(data=lambda: MagicMock())) + async_sdk.models_client.get_provider_route_openai_url.return_value = "https://igw-async.example/v1" _rewrite_options_uris(options, sdk=None, async_sdk=async_sdk) assert options == {"nim": {"max_tokens": 32, "uri": "https://igw-async.example/v1"}} - async_sdk.inference.providers.retrieve.assert_called_once_with(workspace="default", name="nvidia-inference-api") + async_sdk.models_client.get_provider.assert_called_once_with(workspace="default", name="nvidia-inference-api") def test_wraps_sdk_lookup_failure_in_runtimeerror(self) -> None: - sdk = MagicMock() - sdk.inference.providers.retrieve.side_effect = LookupError("no such provider") + sdk = _mock_sdk() + sdk.models_client.get_provider.side_effect = LookupError("no such provider") options = { "nim": { "nmp_uri_spec": { @@ -1006,8 +1037,9 @@ def test_run_resolves_nmp_uri_spec_via_async_sdk( spec = _make_spec_dict(target=target) async_sdk = MagicMock() - async_sdk.inference.providers.retrieve = AsyncMock(return_value=MagicMock()) - async_sdk.models.get_provider_route_openai_url.return_value = "https://igw-async.example/v1" + async_sdk.models_client = MagicMock() + async_sdk.models_client.get_provider = AsyncMock(return_value=MagicMock(data=MagicMock())) + async_sdk.models_client.get_provider_route_openai_url.return_value = "https://igw-async.example/v1" with patch("nemo_auditor.jobs.audit.subprocess.run") as mock_run: mock_run.return_value = subprocess.CompletedProcess(args=[], returncode=0) diff --git a/plugins/nemo-evaluator/tests/test_evaluate_job.py b/plugins/nemo-evaluator/tests/test_evaluate_job.py index c87a8fc16a..9014c685ce 100644 --- a/plugins/nemo-evaluator/tests/test_evaluate_job.py +++ b/plugins/nemo-evaluator/tests/test_evaluate_job.py @@ -9,6 +9,7 @@ from pathlib import Path from types import SimpleNamespace from typing import Any, Literal, cast +from unittest.mock import patch import httpx import nemo_evaluator.cli as evaluator_cli @@ -56,7 +57,6 @@ ) from nemo_evaluator_sdk.values.models import ModelRef from nemo_evaluator_sdk.values.scores import JSONScoreParser, RangeScore -from nemo_platform import NotFoundError from nemo_platform.types.jobs.platform_job_spec import PlatformJobSpec from nemo_platform_plugin.commands import add_job_commands from nemo_platform_plugin.job_context import JobContext, StoragePaths @@ -235,6 +235,29 @@ def get_model_entity_route_openai_url(self, model_entity: object) -> str: return "https://igw.example.test/v1/chat/completions" +class _FakeModelsClient: + """Typed ``ModelsClient`` surface the resolver builds via + ``client_from_platform``. ``get_model``/``get_provider`` return responses with + a ``.data()`` accessor, mirroring ``NemoResponse``.""" + + def __init__(self) -> None: + self.retrieved: list[tuple[str, str]] = [] + + def get_model(self, *, name: str, workspace: str) -> SimpleNamespace: + self.retrieved.append((workspace, name)) + entity = SimpleNamespace(model_providers=["default/provider"]) + return SimpleNamespace(data=lambda: entity) + + def get_provider(self, *, name: str, workspace: str) -> SimpleNamespace: + return SimpleNamespace( + data=lambda: SimpleNamespace(name=name, workspace=workspace, host_url="http://nim.example.test:8000") + ) + + def get_model_entity_route_openai_url(self, model_entity: object) -> str: + del model_entity + return "https://igw.example.test/v1/chat/completions" + + class _FakeProviders: def retrieve(self, name: str, *, workspace: str) -> SimpleNamespace: return SimpleNamespace(name=name, workspace=workspace, host_url="http://nim.example.test:8000") @@ -242,10 +265,22 @@ def retrieve(self, name: str, *, workspace: str) -> SimpleNamespace: class _FakeSDK: def __init__(self) -> None: + self.models_client = _FakeModelsClient() self.models = _FakeModels() self.inference = SimpleNamespace(providers=_FakeProviders()) +@pytest.fixture(autouse=True) +def _patch_resolver_client_from_platform(): + """Route ``client_from_platform(sdk, ModelsClient)`` in the resolver back to + ``sdk.models_client`` so tests drive the typed client directly.""" + with patch( + "nemo_evaluator.resolvers.client_from_platform", + side_effect=lambda sdk_or_async, _cls: sdk_or_async.models_client, + ): + yield + + def _llm_judge_ref_metric() -> LLMJudgeMetric: return LLMJudgeMetric( model=ModelRef(root="default/judge"), @@ -479,22 +514,20 @@ async def test_platform_model_resolver_resolves_model_ref_through_sdk() -> None: model = await resolver.resolve_model(ModelRef(root="default/judge")) - assert sdk.models.retrieved == [("default", "judge")] + assert sdk.models_client.retrieved == [("default", "judge")] assert model.name == "judge" assert model.url == "https://igw.example.test/v1/chat/completions" assert model.host_url == "http://nim.example.test:8000" async def test_platform_model_resolver_rejects_missing_model_ref(mocker: MockerFixture) -> None: + from nemo_platform_plugin.client.errors import NotFoundError as PluginNotFoundError + sdk = _FakeSDK() - response = httpx.Response( - 404, - request=httpx.Request("GET", "https://nmp.test/apis/models/v2/workspaces/default/models/missing"), - ) mocker.patch.object( - sdk.models, - "retrieve", - side_effect=NotFoundError("Model not found", response=response, body=None), + sdk.models_client, + "get_model", + side_effect=PluginNotFoundError(httpx.Response(404, request=httpx.Request("GET", "https://nmp.test/a"))), ) with pytest.raises(ValueError, match="Model reference 'default/missing' not found"): diff --git a/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py b/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py index 25f910d33a..2fd62b7211 100644 --- a/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py +++ b/plugins/nemo-safe-synthesizer/tests/unit/test_jobs.py @@ -61,17 +61,26 @@ def mock_files_client(): def mock_sdk(mock_files_client): sdk = MagicMock() sdk.inference.providers.retrieve = AsyncMock() - sdk.models.get_provider_route_openai_url = MagicMock( - return_value="http://nmp-host/apis/inference-gateway/v2/workspaces/default/provider/my-nim/-/v1" - ) return sdk @pytest.fixture(autouse=True) def _patch_client_from_platform(mock_files_client): + from nemo_platform_plugin.models.client import AsyncModelsClient + + models_client = MagicMock() + models_client.get_provider_route_openai_url = MagicMock( + return_value="http://nmp-host/apis/inference-gateway/v2/workspaces/default/provider/my-nim/-/v1" + ) + + def _dispatch(_sdk, client_cls): + if client_cls is AsyncModelsClient: + return models_client + return mock_files_client + with patch( "nemo_safe_synthesizer_plugin.jobs.generate.client_from_platform", - return_value=mock_files_client, + side_effect=_dispatch, ): yield diff --git a/services/core/models/tests/integration/test_workspace_iam_models_isolation.py b/services/core/models/tests/integration/test_workspace_iam_models_isolation.py index dbd6e5cd6f..b3abbb57a4 100644 --- a/services/core/models/tests/integration/test_workspace_iam_models_isolation.py +++ b/services/core/models/tests/integration/test_workspace_iam_models_isolation.py @@ -185,8 +185,8 @@ def test_model_and_adapter_iam(self, sdk: NeMoPlatform) -> None: admin_files = client_from_platform(admin, FilesClient) admin_files.create_fileset(workspace=ws_c, body=CreateFilesetRequest(name=fs_c)) admin_files.create_fileset(workspace=ws_d, body=CreateFilesetRequest(name=fs_d)) - admin.files.upload_content(content=b"x", remote_path="a.txt", fileset=fs_c, workspace=ws_c) - admin.files.upload_content(content=b"x", remote_path="a.txt", fileset=fs_d, workspace=ws_d) + admin_files.upload_file(workspace=ws_c, name=fs_c, path="a.txt", content=b"x") + admin_files.upload_file(workspace=ws_d, name=fs_d, path="a.txt", content=b"x") # 13: adapter in C with a fileset in D is denied (no access to D fileset) with pytest.raises(PermissionDeniedError): @@ -376,8 +376,8 @@ def h(email: str, groups: list[str] | None = None) -> dict[str, str]: files = client_from_platform(admin_sdk, FilesClient) files.create_fileset(workspace=ws_c, body=CreateFilesetRequest(name=fs_c)) files.create_fileset(workspace=ws_d, body=CreateFilesetRequest(name=fs_d)) - admin_sdk.files.upload_content(content=b"x", remote_path="a.txt", fileset=fs_c, workspace=ws_c) - admin_sdk.files.upload_content(content=b"x", remote_path="a.txt", fileset=fs_d, workspace=ws_d) + files.upload_file(workspace=ws_c, name=fs_c, path="a.txt", content=b"x") + files.upload_file(workspace=ws_d, name=fs_d, path="a.txt", content=b"x") # 13 r13 = post( diff --git a/services/core/models/tests/unit/controllers/conftest.py b/services/core/models/tests/unit/controllers/conftest.py index 4825683a78..4b11091713 100644 --- a/services/core/models/tests/unit/controllers/conftest.py +++ b/services/core/models/tests/unit/controllers/conftest.py @@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from nmp.common.config import PlatformConfig @@ -215,6 +216,76 @@ async def __anext__(self): return self._items.pop(0) +class _AsyncPage: + """Stand-in for ``AsyncNemoPaginatedResponse``: exposes an async ``items()``.""" + + def __init__(self, items): + self._items = list(items) + + async def items(self): + for item in self._items: + yield item + + +class _ModelResponse: + """Stand-in for ``NemoResponse[ModelEntity]``: exposes ``data()``.""" + + def __init__(self, value=None, exc: Exception | None = None): + self._value = value + self._exc = exc + + def data(self): + if self._exc is not None: + raise self._exc + return self._value + + +def _status_error(status: int, detail: str): + """Build a plugin client HTTP error for a given status (409/404).""" + from nemo_platform_plugin.client.errors import ConflictError, NotFoundError + + request = httpx.Request("POST", "http://test") + response = httpx.Response(status, request=request, json={"detail": detail}) + if status == 409: + return ConflictError(response) + return NotFoundError(response) + + +def make_async_models_client() -> MagicMock: + """Build a mock ``AsyncModelsClient`` exposing the typed surface the + controllers/reconcilers consume via ``client_from_platform``.""" + client = MagicMock() + client.list_models = AsyncMock(return_value=_AsyncPage([])) + client.create_model = AsyncMock(return_value=_ModelResponse()) + client.get_model = AsyncMock(return_value=_ModelResponse()) + client.update_model = AsyncMock(return_value=_ModelResponse()) + return client + + +@pytest.fixture +async def mock_models_client() -> MagicMock: + """A mock ``AsyncModelsClient`` returned by ``client_from_platform`` in the + controller/reconciler modules under test.""" + return make_async_models_client() + + +@pytest.fixture +async def patch_models_client(mock_models_client): + """Route ``client_from_platform(sdk, AsyncModelsClient)`` in the models + controller modules back to :data:`mock_models_client`.""" + with ( + patch( + "nmp.core.models.controllers.entity_cache.client_from_platform", + return_value=mock_models_client, + ), + patch( + "nmp.core.models.sidecars.adapters.main.client_from_platform", + return_value=mock_models_client, + ), + ): + yield mock_models_client + + _ENTITY_FIELDS = ("model_providers", "fileset", "api_endpoint", "backend_format") @@ -238,7 +309,18 @@ def _copy(update): return entity -async def seed_entity_cache(mock_models_sdk, entity_cache, entities=()): - """Load the cache from the mock SDK so lookups resolve to ``entities``.""" - mock_models_sdk.models.list = MagicMock(return_value=AsyncPaginator(list(entities))) +async def seed_entity_cache( + mock_models_sdk, + entity_cache, + entities=(), + *, + models_client: MagicMock | None = None, +): + """Load the cache from the mock SDK so lookups resolve to ``entities``. + + ``models_client`` is the mock ``AsyncModelsClient`` returned by the patched + ``client_from_platform``; defaulting to ``mock_models_sdk.models_client``. + """ + models_client = models_client or mock_models_sdk.models_client + models_client.list_models.return_value = _AsyncPage(list(entities)) await entity_cache.refresh() diff --git a/services/core/models/tests/unit/controllers/test_deployment_reconciler.py b/services/core/models/tests/unit/controllers/test_deployment_reconciler.py index 2c2dfa7b8b..22a114b362 100644 --- a/services/core/models/tests/unit/controllers/test_deployment_reconciler.py +++ b/services/core/models/tests/unit/controllers/test_deployment_reconciler.py @@ -18,7 +18,7 @@ from nmp.core.models.controllers.entity_cache import ModelEntityCache from nmp.core.models.schemas import ModelDeployment -from .conftest import AsyncPaginator, make_entity, seed_entity_cache +from .conftest import AsyncPaginator, _ModelResponse, make_async_models_client, make_entity, seed_entity_cache _AsyncPaginator = AsyncPaginator @@ -32,10 +32,21 @@ def _entity(workspace, name, model_providers): def mock_models_sdk(): """Create a mock AsyncNeMoPlatform SDK.""" sdk = MagicMock(spec=AsyncNeMoPlatform) - sdk.models.list = MagicMock(return_value=_AsyncPaginator([])) + sdk.models_client = make_async_models_client() return sdk +@pytest.fixture(autouse=True) +def _patch_entity_cache_client_from_platform(mock_models_sdk): + """Route ``client_from_platform(sdk, AsyncModelsClient)`` in the entity cache + back to the mock typed client on ``mock_models_sdk.models_client``.""" + with patch( + "nmp.core.models.controllers.entity_cache.client_from_platform", + side_effect=lambda sdk, cls: sdk.models_client, + ): + yield + + @pytest.fixture def mock_backend_registry(): """Create a mock BackendRegistry.""" @@ -951,7 +962,6 @@ async def test_cleanup_model_entities_removes_provider_from_entities(reconciler) _entity("test-ns", "model-2", ["test-ns/provider-1"]), ], ) - reconciler._models_sdk.models.update = AsyncMock() # Call cleanup await reconciler._cleanup_model_entities_for_provider("test-ns", "provider-1", "test-ns/provider-1") @@ -964,17 +974,11 @@ async def test_cleanup_model_entities_removes_provider_from_entities(reconciler) ) # Verify model entities were updated with provider removed - assert reconciler._models_sdk.models.update.call_count == 2 - reconciler._models_sdk.models.update.assert_any_call( - name="model-1", - workspace="test-ns", - model_providers=["other-ns/other-provider"], - ) - reconciler._models_sdk.models.update.assert_any_call( - name="model-2", - workspace="test-ns", - model_providers=[], - ) + update = reconciler._models_sdk.models_client.update_model + assert update.await_count == 2 + calls = {call.kwargs["name"]: call.kwargs["body"].model_providers for call in update.await_args_list} + assert calls["model-1"] == ["other-ns/other-provider"] + assert calls["model-2"] == [] @pytest.mark.asyncio @@ -985,7 +989,6 @@ async def test_cleanup_model_entities_no_served_models(reconciler): mock_provider.served_models = [] reconciler._models_sdk.inference.providers.retrieve = AsyncMock(return_value=mock_provider) - reconciler._models_sdk.models.update = AsyncMock() # Call cleanup await reconciler._cleanup_model_entities_for_provider("test-ns", "provider-1", "test-ns/provider-1") @@ -995,7 +998,7 @@ async def test_cleanup_model_entities_no_served_models(reconciler): reconciler._models_sdk.inference.providers.retrieve.assert_called_once() # Verify no model entity operations were performed - reconciler._models_sdk.models.update.assert_not_called() + reconciler._models_sdk.models_client.update_model.assert_not_awaited() @pytest.mark.asyncio @@ -1004,14 +1007,13 @@ async def test_cleanup_model_entities_provider_not_found(reconciler): reconciler._models_sdk.inference.providers.retrieve = AsyncMock( side_effect=NotFoundError("Provider not found", response=MagicMock(), body=None) ) - reconciler._models_sdk.models.update = AsyncMock() # Call cleanup - should not raise await reconciler._cleanup_model_entities_for_provider("test-ns", "provider-1", "test-ns/provider-1") await reconciler._entity_cache.flush() # Verify no model entity operations were performed - reconciler._models_sdk.models.update.assert_not_called() + reconciler._models_sdk.models_client.update_model.assert_not_awaited() @pytest.mark.asyncio @@ -1029,14 +1031,13 @@ async def test_cleanup_model_entities_provider_not_in_list(reconciler): reconciler._entity_cache, [_entity("test-ns", "model-1", ["other-ns/other-provider"])], ) - reconciler._models_sdk.models.update = AsyncMock() # Call cleanup await reconciler._cleanup_model_entities_for_provider("test-ns", "provider-1", "test-ns/provider-1") await reconciler._entity_cache.flush() # Verify model entity was NOT updated (provider wasn't in the list) - reconciler._models_sdk.models.update.assert_not_called() + reconciler._models_sdk.models_client.update_model.assert_not_awaited() @pytest.mark.asyncio @@ -1056,18 +1057,19 @@ async def test_cleanup_model_entities_skips_missing_entity_and_continues(reconci reconciler._entity_cache, [_entity("test-ns", "model-2", ["test-ns/provider-1"])], ) - reconciler._models_sdk.models.update = AsyncMock() # Call cleanup - should not raise await reconciler._cleanup_model_entities_for_provider("test-ns", "provider-1", "test-ns/provider-1") await reconciler._entity_cache.flush() # Verify only the existing model was updated - reconciler._models_sdk.models.update.assert_called_once_with( - workspace="test-ns", - name="model-2", - model_providers=[], - ) + update = reconciler._models_sdk.models_client.update_model + update.assert_awaited_once() + call = update.await_args + assert call is not None + assert call.kwargs["workspace"] == "test-ns" + assert call.kwargs["name"] == "model-2" + assert call.kwargs["body"].model_providers == [] @pytest.mark.asyncio @@ -1090,14 +1092,16 @@ async def test_cleanup_model_entities_handles_model_update_failure(reconciler): ], ) # First update fails, second succeeds - reconciler._models_sdk.models.update = AsyncMock(side_effect=[Exception("Update failed"), None]) + reconciler._models_sdk.models_client.update_model = AsyncMock( + side_effect=[Exception("Update failed"), _ModelResponse()] + ) # Call cleanup - should not raise await reconciler._cleanup_model_entities_for_provider("test-ns", "provider-1", "test-ns/provider-1") await reconciler._entity_cache.flush() # Verify both models were attempted to be updated - assert reconciler._models_sdk.models.update.call_count == 2 + assert reconciler._models_sdk.models_client.update_model.await_count == 2 @pytest.mark.asyncio @@ -1115,14 +1119,13 @@ async def test_cleanup_model_entities_with_null_model_providers(reconciler): reconciler._entity_cache, [_entity("test-ns", "model-1", None)], ) - reconciler._models_sdk.models.update = AsyncMock() # Call cleanup - should not raise await reconciler._cleanup_model_entities_for_provider("test-ns", "provider-1", "test-ns/provider-1") await reconciler._entity_cache.flush() # Verify model entity was NOT updated (provider wasn't in the empty/null list) - reconciler._models_sdk.models.update.assert_not_called() + reconciler._models_sdk.models_client.update_model.assert_not_awaited() @pytest.mark.asyncio diff --git a/services/core/models/tests/unit/controllers/test_entity_cache.py b/services/core/models/tests/unit/controllers/test_entity_cache.py index 343b054181..c9398e56b1 100644 --- a/services/core/models/tests/unit/controllers/test_entity_cache.py +++ b/services/core/models/tests/unit/controllers/test_entity_cache.py @@ -3,14 +3,21 @@ """Unit tests for ModelEntityCache.""" -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from nemo_platform import AsyncNeMoPlatform -from nemo_platform._exceptions import ConflictError, NotFoundError +from nemo_platform_plugin.models.types import CreateModelEntityRequest, UpdateModelEntityRequest from nmp.core.models.controllers.entity_cache import ModelEntityCache, UnflushedMutationsError -from .conftest import AsyncPaginator, make_entity, seed_entity_cache +from .conftest import ( + _AsyncPage, + _ModelResponse, + _status_error, + make_async_models_client, + make_entity, + seed_entity_cache, +) def _entity(workspace="ws", name="model", model_providers=None, **attrs): @@ -18,12 +25,22 @@ def _entity(workspace="ws", name="model", model_providers=None, **attrs): @pytest.fixture -def mock_models_sdk(): +async def mock_models_client(): + return make_async_models_client() + + +@pytest.fixture +async def patch_models_client(mock_models_client): + """Route ``client_from_platform(sdk, AsyncModelsClient)`` in the entity_cache + module back to :data:`mock_models_client`.""" + with patch("nmp.core.models.controllers.entity_cache.client_from_platform", return_value=mock_models_client): + yield mock_models_client + + +@pytest.fixture +async def mock_models_sdk(mock_models_client, patch_models_client): sdk = MagicMock(spec=AsyncNeMoPlatform) - sdk.models.list = MagicMock(return_value=AsyncPaginator([])) - sdk.models.create = AsyncMock(return_value=None) - sdk.models.update = AsyncMock(return_value=None) - sdk.models.retrieve = AsyncMock() + sdk.models_client = mock_models_client return sdk @@ -42,6 +59,30 @@ async def _load(mock_models_sdk, cache, entities=()): await seed_entity_cache(mock_models_sdk, cache, entities) +def _page(items): + return _AsyncPage(list(items)) + + +def _resp(value=None, exc=None): + return _ModelResponse(value=value, exc=exc) + + +def _configure_models_client(mock_models_sdk, **kwargs): + """Configure the mock typed client's methods; ``create_model``/``get_model``/ + ``update_model``/``list_models`` return ``_ModelResponse`` / ``_AsyncPage`` + wrappers unless overridden here.""" + client = mock_models_sdk.models_client + if "list_models" in kwargs: + client.list_models = kwargs["list_models"] + if "create_model" in kwargs: + client.create_model = kwargs["create_model"] + if "get_model" in kwargs: + client.get_model = kwargs["get_model"] + if "update_model" in kwargs: + client.update_model = kwargs["update_model"] + return client + + @pytest.mark.asyncio async def test_refresh_loads_entities_keyed_by_workspace_and_name(mock_models_sdk, cache): await _load(mock_models_sdk, cache, [_entity("ws-a", "m1"), _entity("ws-b", "m1")]) @@ -95,31 +136,42 @@ async def test_entity_staged_for_creation_still_reads_as_absent(mock_models_sdk, async def test_multiple_providers_produce_a_single_update(mock_models_sdk, cache): """An entity linked by several providers is written once, not once per provider.""" await _load(mock_models_sdk, cache, [_entity("ws", "model", [])]) + client = _configure_models_client( + mock_models_sdk, update_model=AsyncMock(return_value=_resp(_entity("ws", "model", ["ws/p1", "ws/p2"]))) + ) cache.stage_provider_link("ws", "model", "ws/p1") cache.stage_provider_link("ws", "model", "ws/p2") await cache.flush() - mock_models_sdk.models.update.assert_awaited_once_with( - workspace="ws", name="model", model_providers=["ws/p1", "ws/p2"] - ) + client.update_model.assert_awaited_once() + call = client.update_model.await_args + assert call is not None + assert call.kwargs["workspace"] == "ws" + assert call.kwargs["name"] == "model" + assert call.kwargs["body"].model_providers == ["ws/p1", "ws/p2"] @pytest.mark.asyncio async def test_no_write_when_already_converged(mock_models_sdk, cache): """Staging state that already matches the entity performs no write.""" await _load(mock_models_sdk, cache, [_entity("ws", "model", ["ws/p1"])]) + client = _configure_models_client(mock_models_sdk) cache.stage_provider_link("ws", "model", "ws/p1") await cache.flush() - mock_models_sdk.models.update.assert_not_awaited() - mock_models_sdk.models.create.assert_not_awaited() + client.update_model.assert_not_awaited() + client.create_model.assert_not_awaited() @pytest.mark.asyncio async def test_two_providers_creating_the_same_entity_collapse_to_one_create(mock_models_sdk, cache): await _load(mock_models_sdk, cache) + client = _configure_models_client( + mock_models_sdk, + create_model=AsyncMock(return_value=_resp(_entity("ws", "model", ["ws/p1", "ws/p2"]))), + ) cache.stage_create("ws", "model", description="from p1", backend_format="OPENAI_CHAT") cache.stage_provider_link("ws", "model", "ws/p1") @@ -127,53 +179,68 @@ async def test_two_providers_creating_the_same_entity_collapse_to_one_create(moc cache.stage_provider_link("ws", "model", "ws/p2") await cache.flush() - mock_models_sdk.models.create.assert_awaited_once_with( - workspace="ws", - name="model", - description="from p1", - backend_format="OPENAI_CHAT", - model_providers=["ws/p1", "ws/p2"], - ) + client.create_model.assert_awaited_once() + call = client.create_model.await_args + assert call is not None + assert call.kwargs["workspace"] == "ws" + body: CreateModelEntityRequest = call.kwargs["body"] + assert body.name == "model" + assert body.description == "from p1" + assert body.backend_format == "OPENAI_CHAT" + assert body.model_providers == ["ws/p1", "ws/p2"] @pytest.mark.asyncio async def test_create_conflict_falls_back_to_updating_the_existing_entity(mock_models_sdk, cache): """An entity created concurrently is adopted rather than reported as an error.""" await _load(mock_models_sdk, cache) - mock_models_sdk.models.create = AsyncMock(side_effect=ConflictError("exists", response=MagicMock(), body=None)) - mock_models_sdk.models.retrieve = AsyncMock(return_value=_entity("ws", "model", ["ws/other"])) + client = _configure_models_client( + mock_models_sdk, + create_model=AsyncMock(side_effect=_status_error(409, "exists")), + get_model=AsyncMock(return_value=_resp(_entity("ws", "model", ["ws/other"]))), + update_model=AsyncMock(return_value=_resp(_entity("ws", "model", ["ws/other", "ws/p1"]))), + ) cache.stage_create("ws", "model", description="d", backend_format="OPENAI_CHAT") cache.stage_provider_link("ws", "model", "ws/p1") await cache.flush() - mock_models_sdk.models.update.assert_awaited_once_with( - workspace="ws", name="model", model_providers=["ws/other", "ws/p1"] - ) + client.update_model.assert_awaited_once() + call = client.update_model.await_args + assert call is not None + assert call.kwargs["workspace"] == "ws" + assert call.kwargs["name"] == "model" + assert call.kwargs["body"].model_providers == ["ws/other", "ws/p1"] @pytest.mark.asyncio async def test_create_conflict_with_vanished_entity_is_ignored(mock_models_sdk, cache): await _load(mock_models_sdk, cache) - mock_models_sdk.models.create = AsyncMock(side_effect=ConflictError("exists", response=MagicMock(), body=None)) - mock_models_sdk.models.retrieve = AsyncMock(side_effect=NotFoundError("gone", response=MagicMock(), body=None)) + client = _configure_models_client( + mock_models_sdk, + create_model=AsyncMock(side_effect=_status_error(409, "exists")), + get_model=AsyncMock(side_effect=_status_error(404, "gone")), + ) cache.stage_create("ws", "model", description="d") await cache.flush() - mock_models_sdk.models.update.assert_not_awaited() + client.update_model.assert_not_awaited() @pytest.mark.asyncio async def test_one_failing_entity_does_not_stop_the_others(mock_models_sdk, cache): await _load(mock_models_sdk, cache, [_entity("ws", "m1", []), _entity("ws", "m2", [])]) - mock_models_sdk.models.update = AsyncMock(side_effect=[Exception("boom"), None]) + client = _configure_models_client( + mock_models_sdk, + update_model=AsyncMock(side_effect=[Exception("boom"), _resp(_entity("ws", "m2", ["ws/p1"]))]), + ) cache.stage_provider_link("ws", "m1", "ws/p1") cache.stage_provider_link("ws", "m2", "ws/p1") await cache.flush() - assert mock_models_sdk.models.update.await_count == 2 + assert client.update_model.await_count == 2 @pytest.mark.asyncio @@ -185,7 +252,10 @@ async def test_failed_write_is_kept_for_retry_and_succeeds_later(mock_models_sdk would leave the entity permanently inconsistent. """ await _load(mock_models_sdk, cache, [_entity("ws", "m1", ["ws/p1"]), _entity("ws", "m2", ["ws/p1"])]) - mock_models_sdk.models.update = AsyncMock(side_effect=[Exception("boom"), None]) + client = _configure_models_client( + mock_models_sdk, + update_model=AsyncMock(side_effect=[Exception("boom"), _resp(_entity("ws", "m2", []))]), + ) cache.stage_provider_unlink("ws", "m1", "ws/p1") cache.stage_provider_unlink("ws", "m2", "ws/p1") @@ -197,10 +267,15 @@ async def test_failed_write_is_kept_for_retry_and_succeeds_later(mock_models_sdk assert ("ws", "m2") not in cache._pending # A later flush retries it, and this time it lands. - mock_models_sdk.models.update = AsyncMock(return_value=None) + client.update_model = AsyncMock(return_value=_resp(_entity("ws", "m1", []))) await cache.flush() - mock_models_sdk.models.update.assert_awaited_once_with(workspace="ws", name="m1", model_providers=[]) + client.update_model.assert_awaited_once() + call = client.update_model.await_args + assert call is not None + assert call.kwargs["workspace"] == "ws" + assert call.kwargs["name"] == "m1" + assert call.kwargs["body"].model_providers == [] assert cache._pending == {} @@ -208,7 +283,10 @@ async def test_failed_write_is_kept_for_retry_and_succeeds_later(mock_models_sdk async def test_refresh_allows_retained_failures_but_still_rejects_unflushed_work(mock_models_sdk, cache): """Refresh distinguishes "flushed and failed" from "staged and forgotten".""" await _load(mock_models_sdk, cache, [_entity("ws", "m1", ["ws/p1"])]) - mock_models_sdk.models.update = AsyncMock(side_effect=Exception("boom")) + _configure_models_client( + mock_models_sdk, + update_model=AsyncMock(side_effect=Exception("boom")), + ) cache.stage_provider_unlink("ws", "m1", "ws/p1") await cache.flush() @@ -227,66 +305,87 @@ async def test_refresh_allows_retained_failures_but_still_rejects_unflushed_work async def test_retained_failure_replays_against_a_newer_snapshot(mock_models_sdk, cache): """Staged changes are differences, so replaying them after a refresh stays correct.""" await _load(mock_models_sdk, cache, [_entity("ws", "m1", ["ws/p1", "ws/p2"])]) - mock_models_sdk.models.update = AsyncMock(side_effect=Exception("boom")) + client = _configure_models_client( + mock_models_sdk, + update_model=AsyncMock(side_effect=Exception("boom")), + ) cache.stage_provider_unlink("ws", "m1", "ws/p1") await cache.flush() # Snapshot moves on: another writer added a third provider meanwhile. await _load(mock_models_sdk, cache, [_entity("ws", "m1", ["ws/p1", "ws/p2", "ws/p3"])]) - mock_models_sdk.models.update = AsyncMock(return_value=None) + client.update_model = AsyncMock(return_value=_resp(_entity("ws", "m1", ["ws/p2"]))) await cache.flush() # The unlink applies to the newer state rather than reinstating the old list. - mock_models_sdk.models.update.assert_awaited_once_with( - workspace="ws", name="m1", model_providers=["ws/p2", "ws/p3"] - ) + client.update_model.assert_awaited_once() + call = client.update_model.await_args + assert call is not None + assert call.kwargs["workspace"] == "ws" + assert call.kwargs["name"] == "m1" + assert call.kwargs["body"].model_providers == ["ws/p2", "ws/p3"] @pytest.mark.asyncio async def test_flush_clears_staged_changes(mock_models_sdk, cache): await _load(mock_models_sdk, cache, [_entity("ws", "model", [])]) + client = _configure_models_client( + mock_models_sdk, + update_model=AsyncMock(return_value=_resp(_entity("ws", "model", ["ws/p1"]))), + ) cache.stage_provider_link("ws", "model", "ws/p1") await cache.flush() - mock_models_sdk.models.update.reset_mock() + client.update_model.reset_mock() # Nothing left staged, so a second flush writes nothing and a refresh is allowed. await cache.flush() - mock_models_sdk.models.update.assert_not_awaited() + client.update_model.assert_not_awaited() await cache.refresh() @pytest.mark.asyncio async def test_link_then_unlink_for_the_same_provider_cancels_out(mock_models_sdk, cache): await _load(mock_models_sdk, cache, [_entity("ws", "model", ["ws/p1"])]) + client = _configure_models_client(mock_models_sdk) cache.stage_provider_unlink("ws", "model", "ws/p1") cache.stage_provider_link("ws", "model", "ws/p1") await cache.flush() - mock_models_sdk.models.update.assert_not_awaited() + client.update_model.assert_not_awaited() @pytest.mark.asyncio async def test_field_updates_are_written_as_staged(mock_models_sdk, cache): await _load(mock_models_sdk, cache, [_entity("ws", "model", ["ws/p1"])]) + client = _configure_models_client( + mock_models_sdk, + update_model=AsyncMock(return_value=_resp(_entity("ws", "model", ["ws/p1"], fileset="hub/model"))), + ) cache.stage_field_updates("ws", "model", fileset="hub/model", api_endpoint=None) await cache.flush() - mock_models_sdk.models.update.assert_awaited_once_with(workspace="ws", name="model", fileset="hub/model") + client.update_model.assert_awaited_once() + call = client.update_model.await_args + assert call is not None + assert call.kwargs["workspace"] == "ws" + assert call.kwargs["name"] == "model" + assert call.kwargs["body"].fileset == "hub/model" @pytest.mark.asyncio async def test_staged_change_for_missing_entity_without_create_is_skipped(mock_models_sdk, cache): await _load(mock_models_sdk, cache) + client = _configure_models_client(mock_models_sdk) cache.stage_provider_unlink("ws", "ghost", "ws/p1") await cache.flush() - mock_models_sdk.models.create.assert_not_awaited() - mock_models_sdk.models.update.assert_not_awaited() + client.create_model.assert_not_awaited() + client.update_model.assert_not_awaited() @pytest.mark.asyncio @@ -297,17 +396,22 @@ async def test_refresh_after_flush_does_not_reapply_earlier_state(mock_models_sd phase's writes, otherwise it would re-add what was just removed. """ store = {("ws", "model"): _entity("ws", "model", ["ws/p1"])} + client = _configure_models_client(mock_models_sdk) - def _list(**_kwargs): - return AsyncPaginator(list(store.values())) + async def _list_models(**kwargs): + return _page(list(store.values())) - async def _update(*, workspace, name, **params): + async def _update_model(**kwargs): + body: UpdateModelEntityRequest = kwargs["body"] + workspace, name = kwargs["workspace"], kwargs["name"] current = store[(workspace, name)] - store[(workspace, name)] = _entity(workspace, name, params.get("model_providers", current.model_providers)) - return store[(workspace, name)] + store[(workspace, name)] = _entity( + workspace, name, body.model_providers if body.model_providers is not None else current.model_providers + ) + return _resp(store[(workspace, name)]) - mock_models_sdk.models.list = MagicMock(side_effect=_list) - mock_models_sdk.models.update = AsyncMock(side_effect=_update) + client.list_models = _list_models + client.update_model = _update_model # Phase one removes the provider link and applies it. await cache.refresh() @@ -337,12 +441,16 @@ async def test_flush_reports_progress_per_entity_written(mock_models_sdk, cache, """ await _load(mock_models_sdk, cache, [_entity("ws", f"m{i}", []) for i in range(25)]) heartbeat_calls.clear() + client = _configure_models_client( + mock_models_sdk, + update_model=AsyncMock(return_value=_resp(_entity("ws", "m0", ["ws/p1"]))), + ) for i in range(25): cache.stage_provider_link("ws", f"m{i}", "ws/p1") await cache.flush() - assert mock_models_sdk.models.update.await_count == 25 + assert client.update_model.await_count == 25 assert len(heartbeat_calls) == 25 @@ -350,7 +458,10 @@ async def test_flush_reports_progress_per_entity_written(mock_models_sdk, cache, async def test_flush_reports_progress_even_when_an_entity_write_fails(mock_models_sdk, cache, heartbeat_calls): """Moving past a failed entity is still progress.""" await _load(mock_models_sdk, cache, [_entity("ws", "m1", []), _entity("ws", "m2", [])]) - mock_models_sdk.models.update = AsyncMock(side_effect=[Exception("boom"), None]) + _configure_models_client( + mock_models_sdk, + update_model=AsyncMock(side_effect=[Exception("boom"), _resp(_entity("ws", "m2", ["ws/p1"]))]), + ) heartbeat_calls.clear() cache.stage_provider_link("ws", "m1", "ws/p1") @@ -369,9 +480,13 @@ async def test_conflict_adoption_does_not_overwrite_the_existing_entity_attribut what is still missing on a later pass. """ await _load(mock_models_sdk, cache) - mock_models_sdk.models.create = AsyncMock(side_effect=ConflictError("exists", response=MagicMock(), body=None)) - mock_models_sdk.models.retrieve = AsyncMock( - return_value=_entity("ws", "model", ["ws/other"], backend_format="ANTHROPIC_MESSAGES") + client = _configure_models_client( + mock_models_sdk, + create_model=AsyncMock(side_effect=_status_error(409, "exists")), + get_model=AsyncMock( + return_value=_resp(_entity("ws", "model", ["ws/other"], backend_format="ANTHROPIC_MESSAGES")) + ), + update_model=AsyncMock(return_value=_resp(_entity("ws", "model", ["ws/other", "ws/p1"]))), ) cache.stage_create("ws", "model", description="ours", backend_format="OPENAI_CHAT") @@ -379,6 +494,9 @@ async def test_conflict_adoption_does_not_overwrite_the_existing_entity_attribut await cache.flush() # Only the provider link is written; description/backend_format are not forced. - mock_models_sdk.models.update.assert_awaited_once_with( - workspace="ws", name="model", model_providers=["ws/other", "ws/p1"] - ) + client.update_model.assert_awaited_once() + call = client.update_model.await_args + assert call is not None + assert call.kwargs["workspace"] == "ws" + assert call.kwargs["name"] == "model" + assert call.kwargs["body"].model_providers == ["ws/other", "ws/p1"] diff --git a/services/core/models/tests/unit/controllers/test_models_controller_unit.py b/services/core/models/tests/unit/controllers/test_models_controller_unit.py index 1a233fdc5c..1c7e5d1fc0 100644 --- a/services/core/models/tests/unit/controllers/test_models_controller_unit.py +++ b/services/core/models/tests/unit/controllers/test_models_controller_unit.py @@ -12,6 +12,8 @@ from nmp.core.models.controllers.context import ModelContext from nmp.core.models.controllers.models_controller import NON_TERMINAL_STATES, ModelsController +from .conftest import _ModelResponse, make_async_models_client + class MockAsyncPaginator: """Mock async paginator to simulate SDK's paginated response.""" @@ -28,6 +30,26 @@ async def __anext__(self): return self.items.pop(0) +@pytest.fixture(autouse=True) +def _patch_typed_model_client(mock_models_sdk): + """Route ``client_from_platform(sdk, AsyncModelsClient)`` in the models controller + and entity cache back to a typed ``AsyncModelsClient`` mock on + ``mock_models_sdk.models_client``, so tests drive ``get_model``/``list_models`` + directly instead of the legacy SDK resource.""" + mock_models_sdk.models_client = make_async_models_client() + with ( + patch( + "nmp.core.models.controllers.models_controller.client_from_platform", + side_effect=lambda sdk, cls: sdk.models_client, + ), + patch( + "nmp.core.models.controllers.entity_cache.client_from_platform", + side_effect=lambda sdk, cls: sdk.models_client, + ), + ): + yield + + def test_controller_initialization(mock_sdk_class_patch, mock_get_config_patch, mock_backend_registry, assert_helpers): """Test that ModelsController initializes correctly.""" # Create controller @@ -445,7 +467,7 @@ async def test_retrieve_model_entity_for_config_uses_model_entity_id_when_set( mock_entity = MagicMock() mock_entity.workspace = "my-ws" mock_entity.name = "my-model" - mock_models_sdk.models.retrieve = AsyncMock(return_value=mock_entity) + mock_models_sdk.models_client.get_model = AsyncMock(return_value=_ModelResponse(mock_entity)) config = MagicMock() config.model_entity_id = "my-ws/my-model" @@ -459,7 +481,7 @@ async def test_retrieve_model_entity_for_config_uses_model_entity_id_when_set( result = await controller._retrieve_model_entity_for_config(config) assert result is mock_entity - mock_models_sdk.models.retrieve.assert_called_once_with(name="my-model", workspace="my-ws") + mock_models_sdk.models_client.get_model.assert_awaited_once_with(name="my-model", workspace="my-ws") @pytest.mark.asyncio @@ -468,7 +490,7 @@ async def test_retrieve_model_entity_for_config_uses_model_entity_id_with_revisi ): """When config.model_entity_id includes @revision, revision is passed to retrieve.""" mock_entity = MagicMock() - mock_models_sdk.models.retrieve = AsyncMock(return_value=mock_entity) + mock_models_sdk.models_client.get_model = AsyncMock(return_value=_ModelResponse(mock_entity)) config = MagicMock() config.model_entity_id = "my-ws/my-model@v2" @@ -479,8 +501,10 @@ async def test_retrieve_model_entity_for_config_uses_model_entity_id_with_revisi result = await controller._retrieve_model_entity_for_config(config) assert result is mock_entity - mock_models_sdk.models.retrieve.assert_called_once_with(name="my-model@v2", workspace="my-ws") - call_kw = mock_models_sdk.models.retrieve.call_args[1] + mock_models_sdk.models_client.get_model.assert_awaited_once_with(name="my-model@v2", workspace="my-ws") + call = mock_models_sdk.models_client.get_model.await_args + assert call is not None + call_kw = call.kwargs assert call_kw["name"] == "my-model@v2" assert call_kw["workspace"] == "my-ws" @@ -491,7 +515,7 @@ async def test_retrieve_model_entity_for_config_falls_back_to_nim_deployment_whe ): """When config.model_entity_id is not set, entity is derived from nim_deployment.""" mock_entity = MagicMock() - mock_models_sdk.models.retrieve = AsyncMock(return_value=mock_entity) + mock_models_sdk.models_client.get_model = AsyncMock(return_value=_ModelResponse(mock_entity)) config = MagicMock() config.model_entity_id = None @@ -505,7 +529,7 @@ async def test_retrieve_model_entity_for_config_falls_back_to_nim_deployment_whe result = await controller._retrieve_model_entity_for_config(config) assert result is mock_entity - mock_models_sdk.models.retrieve.assert_called_once_with(name="nim-model@v1", workspace="nim-ns") + mock_models_sdk.models_client.get_model.assert_awaited_once_with(name="nim-model@v1", workspace="nim-ns") @pytest.mark.asyncio @@ -522,7 +546,7 @@ async def test_retrieve_model_entity_for_config_returns_none_when_no_nim_deploym result = await controller._retrieve_model_entity_for_config(config) assert result is None - mock_models_sdk.models.retrieve.assert_not_called() + mock_models_sdk.models_client.get_model.assert_not_awaited() @pytest.mark.asyncio @@ -531,7 +555,7 @@ async def test_retrieve_model_entity_for_config_invalid_model_entity_id_falls_ba ): """When model_entity_id is set but unparseable (e.g. no slash), fall back to nim_deployment.""" mock_entity = MagicMock() - mock_models_sdk.models.retrieve = AsyncMock(return_value=mock_entity) + mock_models_sdk.models_client.get_model = AsyncMock(return_value=_ModelResponse(mock_entity)) config = MagicMock() config.model_entity_id = "bogus" @@ -545,7 +569,7 @@ async def test_retrieve_model_entity_for_config_invalid_model_entity_id_falls_ba result = await controller._retrieve_model_entity_for_config(config) assert result is mock_entity - mock_models_sdk.models.retrieve.assert_called_once_with(name="fallback-model", workspace="fallback-ns") + mock_models_sdk.models_client.get_model.assert_awaited_once_with(name="fallback-model", workspace="fallback-ns") # ============================================================================= diff --git a/services/core/models/tests/unit/controllers/test_provider_reconciler.py b/services/core/models/tests/unit/controllers/test_provider_reconciler.py index a3c9e2a4ab..8395c2413f 100644 --- a/services/core/models/tests/unit/controllers/test_provider_reconciler.py +++ b/services/core/models/tests/unit/controllers/test_provider_reconciler.py @@ -31,7 +31,14 @@ ) from nmp.core.models.schemas import ModelProviderStatus -from .conftest import AsyncPaginator, make_entity, seed_entity_cache +from .conftest import ( + AsyncPaginator, + _AsyncPage, + _ModelResponse, + make_async_models_client, + make_entity, + seed_entity_cache, +) def _discovery_models_from_ids(ids: list[str]) -> list[dict]: @@ -103,12 +110,23 @@ def mock_models_sdk(): sdk.inference.virtual_models.create = AsyncMock(return_value=None) sdk.inference.virtual_models.delete = AsyncMock(return_value=None) sdk.inference.virtual_models.list = MagicMock(return_value=_AsyncPaginator([])) - sdk.models.list = MagicMock(return_value=_AsyncPaginator([])) + sdk.models_client = make_async_models_client() sdk.inference.gateway.provider.get = AsyncMock() sdk.with_options = MagicMock(return_value=sdk) return sdk +@pytest.fixture(autouse=True) +def _patch_entity_cache_client_from_platform(mock_models_sdk): + """Route ``client_from_platform(sdk, AsyncModelsClient)`` in the entity cache + back to the mock typed client on ``mock_models_sdk.models_client``.""" + with patch( + "nmp.core.models.controllers.entity_cache.client_from_platform", + side_effect=lambda sdk, cls: sdk.models_client, + ): + yield + + @pytest.fixture def entity_cache(mock_models_sdk): """Model Entity cache backed by the mock SDK, pre-loaded and empty.""" @@ -643,7 +661,6 @@ async def test_get_artifact_details_handles_exception(reconciler): async def test_ensure_model_entity_creates_new_entity(reconciler): """Test creating a new model entity when it doesn't exist.""" # Mock entity doesn't exist - reconciler._models_sdk.models.create = AsyncMock() # Mock context ctx = ModelContext( @@ -665,14 +682,17 @@ async def test_ensure_model_entity_creates_new_entity(reconciler): await reconciler._entity_cache.flush() # Verify entity creation was called - reconciler._models_sdk.models.create.assert_called_once_with( - workspace="test-ns", - name="test-model", - description="Auto-discovered model from provider test-ns/test-provider", - model_providers=["test-ns/test-provider"], - backend_format="OPENAI_CHAT", - fileset="test/model", - ) + create = reconciler._models_sdk.models_client.create_model + create.assert_awaited_once() + call = create.await_args + assert call is not None + assert call.kwargs["workspace"] == "test-ns" + body = call.kwargs["body"] + assert body.name == "test-model" + assert body.description == "Auto-discovered model from provider test-ns/test-provider" + assert body.model_providers == ["test-ns/test-provider"] + assert body.backend_format == "OPENAI_CHAT" + assert body.fileset == "test/model" @pytest.mark.asyncio @@ -686,9 +706,8 @@ async def test_ensure_model_entity_updates_existing_adds_provider(reconciler): existing_entity.workspace = "test-ns" existing_entity.name = "test-model" - reconciler._models_sdk.models.list = MagicMock(return_value=_AsyncPaginator([existing_entity])) + reconciler._models_sdk.models_client.list_models = AsyncMock(return_value=_AsyncPage([existing_entity])) await reconciler._entity_cache.refresh() - reconciler._models_sdk.models.update = AsyncMock() ctx = ModelContext( model_provider=MagicMock(host_url="https://api.com"), @@ -707,12 +726,14 @@ async def test_ensure_model_entity_updates_existing_adds_provider(reconciler): await reconciler._entity_cache.flush() # Verify update was called to add provider - reconciler._models_sdk.models.update.assert_called_once_with( - name="test-model", - workspace="test-ns", - model_providers=["other-ns/other-provider", "test-ns/test-provider"], - backend_format="OPENAI_CHAT", - ) + update = reconciler._models_sdk.models_client.update_model + update.assert_awaited_once() + call = update.await_args + assert call is not None + assert call.kwargs["workspace"] == "test-ns" + assert call.kwargs["name"] == "test-model" + assert call.kwargs["body"].model_providers == ["other-ns/other-provider", "test-ns/test-provider"] + assert call.kwargs["body"].backend_format == "OPENAI_CHAT" @pytest.mark.asyncio @@ -725,9 +746,8 @@ async def test_ensure_model_entity_skips_if_provider_already_linked(reconciler): existing_entity.workspace = "test-ns" existing_entity.name = "test-model" - reconciler._models_sdk.models.list = MagicMock(return_value=_AsyncPaginator([existing_entity])) + reconciler._models_sdk.models_client.list_models = AsyncMock(return_value=_AsyncPage([existing_entity])) await reconciler._entity_cache.refresh() - reconciler._models_sdk.models.update = AsyncMock() ctx = ModelContext( model_provider=MagicMock(), @@ -746,7 +766,7 @@ async def test_ensure_model_entity_skips_if_provider_already_linked(reconciler): await reconciler._entity_cache.flush() # Verify update was NOT called - reconciler._models_sdk.models.update.assert_not_called() + reconciler._models_sdk.models_client.update_model.assert_not_awaited() @pytest.mark.asyncio @@ -760,9 +780,8 @@ async def test_ensure_model_entity_backfills_missing_backend_format(reconciler): existing_entity.workspace = "test-ns" existing_entity.name = "anthropic.claude-3-5-sonnet" - reconciler._models_sdk.models.list = MagicMock(return_value=_AsyncPaginator([existing_entity])) + reconciler._models_sdk.models_client.list_models = AsyncMock(return_value=_AsyncPage([existing_entity])) await reconciler._entity_cache.refresh() - reconciler._models_sdk.models.update = AsyncMock() ctx = ModelContext( model_provider=MagicMock(), @@ -780,11 +799,13 @@ async def test_ensure_model_entity_backfills_missing_backend_format(reconciler): ) await reconciler._entity_cache.flush() - reconciler._models_sdk.models.update.assert_called_once_with( - name="anthropic.claude-3-5-sonnet", - workspace="test-ns", - backend_format="ANTHROPIC_MESSAGES", - ) + update = reconciler._models_sdk.models_client.update_model + update.assert_awaited_once() + call = update.await_args + assert call is not None + assert call.kwargs["workspace"] == "test-ns" + assert call.kwargs["name"] == "anthropic.claude-3-5-sonnet" + assert call.kwargs["body"].backend_format == "ANTHROPIC_MESSAGES" @pytest.mark.asyncio @@ -798,9 +819,8 @@ async def test_ensure_model_entity_adds_artifact_to_existing_without_artifact(re existing_entity.workspace = "test-ns" existing_entity.name = "test-model" - reconciler._models_sdk.models.list = MagicMock(return_value=_AsyncPaginator([existing_entity])) + reconciler._models_sdk.models_client.list_models = AsyncMock(return_value=_AsyncPage([existing_entity])) await reconciler._entity_cache.refresh() - reconciler._models_sdk.models.update = AsyncMock() ctx = ModelContext( model_provider=MagicMock(host_url="https://api.com"), @@ -821,13 +841,15 @@ async def test_ensure_model_entity_adds_artifact_to_existing_without_artifact(re await reconciler._entity_cache.flush() # Verify update includes artifact - reconciler._models_sdk.models.update.assert_called_once_with( - name="test-model", - workspace="test-ns", - model_providers=["other-ns/other-provider", "test-ns/test-provider"], - backend_format="OPENAI_CHAT", - fileset="test/model", - ) + update = reconciler._models_sdk.models_client.update_model + update.assert_awaited_once() + call = update.await_args + assert call is not None + assert call.kwargs["workspace"] == "test-ns" + assert call.kwargs["name"] == "test-model" + assert call.kwargs["body"].model_providers == ["other-ns/other-provider", "test-ns/test-provider"] + assert call.kwargs["body"].backend_format == "OPENAI_CHAT" + assert call.kwargs["body"].fileset == "test/model" @pytest.mark.asyncio @@ -841,9 +863,8 @@ async def test_ensure_model_entity_doesnt_overwrite_existing_artifact(reconciler existing_entity.workspace = "test-ns" existing_entity.name = "test-model" - reconciler._models_sdk.models.list = MagicMock(return_value=_AsyncPaginator([existing_entity])) + reconciler._models_sdk.models_client.list_models = AsyncMock(return_value=_AsyncPage([existing_entity])) await reconciler._entity_cache.refresh() - reconciler._models_sdk.models.update = AsyncMock() ctx = ModelContext( model_provider=MagicMock(host_url="https://api.com"), @@ -864,8 +885,13 @@ async def test_ensure_model_entity_doesnt_overwrite_existing_artifact(reconciler await reconciler._entity_cache.flush() # Verify update does NOT include fileset (since it already exists) - call_kwargs = reconciler._models_sdk.models.update.call_args.kwargs - assert "fileset" not in call_kwargs + update = reconciler._models_sdk.models_client.update_model + update.assert_awaited_once() + call = update.await_args + assert call is not None + body = call.kwargs["body"] + assert body.fileset is None + assert body.model_providers == ["test-ns/test-provider"] @pytest.mark.asyncio @@ -879,9 +905,8 @@ async def test_ensure_model_entity_doesnt_overwrite_existing_backend_format(reco existing_entity.workspace = "test-ns" existing_entity.name = "test-model" - reconciler._models_sdk.models.list = MagicMock(return_value=_AsyncPaginator([existing_entity])) + reconciler._models_sdk.models_client.list_models = AsyncMock(return_value=_AsyncPage([existing_entity])) await reconciler._entity_cache.refresh() - reconciler._models_sdk.models.update = AsyncMock() ctx = ModelContext( model_provider=MagicMock(host_url="https://api.com"), @@ -899,8 +924,13 @@ async def test_ensure_model_entity_doesnt_overwrite_existing_backend_format(reco ) await reconciler._entity_cache.flush() - call_kwargs = reconciler._models_sdk.models.update.call_args.kwargs - assert "backend_format" not in call_kwargs + update = reconciler._models_sdk.models_client.update_model + update.assert_awaited_once() + call = update.await_args + assert call is not None + body = call.kwargs["body"] + assert body.backend_format is None + assert body.model_providers == ["test-ns/test-provider"] @pytest.mark.asyncio @@ -914,9 +944,8 @@ async def test_ensure_model_entity_handles_null_model_providers(reconciler): existing_entity.workspace = "test-ns" existing_entity.name = "test-model" - reconciler._models_sdk.models.list = MagicMock(return_value=_AsyncPaginator([existing_entity])) + reconciler._models_sdk.models_client.list_models = AsyncMock(return_value=_AsyncPage([existing_entity])) await reconciler._entity_cache.refresh() - reconciler._models_sdk.models.update = AsyncMock() ctx = ModelContext( model_provider=MagicMock(host_url="https://api.com"), @@ -935,21 +964,20 @@ async def test_ensure_model_entity_handles_null_model_providers(reconciler): await reconciler._entity_cache.flush() # Verify update was called with provider as first in list - reconciler._models_sdk.models.update.assert_called_once_with( - name="test-model", - workspace="test-ns", - model_providers=["test-ns/test-provider"], - backend_format="OPENAI_CHAT", - ) + update = reconciler._models_sdk.models_client.update_model + update.assert_awaited_once() + call = update.await_args + assert call is not None + assert call.kwargs["workspace"] == "test-ns" + assert call.kwargs["name"] == "test-model" + assert call.kwargs["body"].model_providers == ["test-ns/test-provider"] + assert call.kwargs["body"].backend_format == "OPENAI_CHAT" @pytest.mark.asyncio async def test_ensure_model_entity_handles_create_exception(reconciler): """Test handling exception during entity creation.""" - reconciler._models_sdk.models.retrieve = AsyncMock( - side_effect=NotFoundError("Not found", response=MagicMock(), body=None) - ) - reconciler._models_sdk.models.create = AsyncMock(side_effect=Exception("Creation failed")) + reconciler._models_sdk.models_client.create_model = AsyncMock(side_effect=Exception("Creation failed")) ctx = ModelContext( model_provider=MagicMock(host_url="https://api.com"), @@ -978,22 +1006,20 @@ async def test_entity_cache_load_failure_propagates_and_stages_nothing(reconcile """ mock_response = MagicMock() mock_response.status_code = 503 - mock_models_sdk.models.list = MagicMock( + mock_models_sdk.models_client.list_models = AsyncMock( side_effect=APIStatusError( "Service unavailable", response=mock_response, body={"detail": "upstream error"}, ) ) - mock_models_sdk.models.create = AsyncMock() - mock_models_sdk.models.update = AsyncMock() with pytest.raises(APIStatusError): await reconciler._entity_cache.refresh() await reconciler._entity_cache.flush() - mock_models_sdk.models.create.assert_not_called() - mock_models_sdk.models.update.assert_not_called() + mock_models_sdk.models_client.create_model.assert_not_awaited() + mock_models_sdk.models_client.update_model.assert_not_awaited() @pytest.mark.asyncio @@ -1005,6 +1031,7 @@ async def test_virtual_model_listing_failure_does_not_abort_provider_reconciliat provider = MagicMock() provider.workspace = "test-ns" provider.name = "test-provider" + provider.host_url = "https://provider.example/v1" provider.model_deployment_id = None provider.enabled_models = None provider.served_models = [] @@ -1016,7 +1043,6 @@ async def test_virtual_model_listing_failure_does_not_abort_provider_reconciliat ) mock_models_sdk.inference.virtual_models.list = MagicMock(side_effect=Exception("listing unavailable")) - mock_models_sdk.models.create = AsyncMock() mock_models_sdk.inference.providers.update_status = AsyncMock() with patch.object( @@ -1029,7 +1055,7 @@ async def test_virtual_model_listing_failure_does_not_abort_provider_reconciliat # Provider status and entity linking still happened. mock_models_sdk.inference.providers.update_status.assert_awaited() - mock_models_sdk.models.create.assert_awaited_once() + mock_models_sdk.models_client.create_model.assert_awaited_once() # VirtualModel work was skipped rather than acted on with an unknown state. mock_models_sdk.inference.virtual_models.create.assert_not_awaited() mock_models_sdk.inference.virtual_models.delete.assert_not_awaited() @@ -2997,7 +3023,7 @@ async def test_entity_linked_by_two_providers_is_written_once(reconciler, mock_m ) ], ) - mock_models_sdk.models.update = AsyncMock() + mock_models_sdk.models_client.update_model = AsyncMock(return_value=_ModelResponse()) ctxs = [] for provider_name in ("provider-a", "provider-b"): @@ -3023,11 +3049,13 @@ async def test_entity_linked_by_two_providers_is_written_once(reconciler, mock_m ): await reconcile_and_flush(reconciler, entity_cache, ctxs) - mock_models_sdk.models.update.assert_awaited_once_with( - workspace="test-ns", - name="shared-model", - model_providers=["test-ns/provider-a", "test-ns/provider-b"], - ) + update = mock_models_sdk.models_client.update_model + update.assert_awaited_once() + call = update.await_args + assert call is not None + assert call.kwargs["workspace"] == "test-ns" + assert call.kwargs["name"] == "shared-model" + assert call.kwargs["body"].model_providers == ["test-ns/provider-a", "test-ns/provider-b"] @pytest.mark.asyncio @@ -3046,8 +3074,6 @@ async def test_converged_entities_are_not_rewritten(reconciler, mock_models_sdk, ) ], ) - mock_models_sdk.models.update = AsyncMock() - mock_models_sdk.models.create = AsyncMock() provider = MagicMock() provider.workspace = "test-ns" @@ -3069,8 +3095,8 @@ async def test_converged_entities_are_not_rewritten(reconciler, mock_models_sdk, ): await reconcile_and_flush(reconciler, entity_cache, ctx and [ctx]) - mock_models_sdk.models.update.assert_not_awaited() - mock_models_sdk.models.create.assert_not_awaited() + mock_models_sdk.models_client.update_model.assert_not_awaited() + mock_models_sdk.models_client.create_model.assert_not_awaited() @pytest.mark.asyncio diff --git a/services/core/models/tests/unit/sidecars/test_adapters_controller.py b/services/core/models/tests/unit/sidecars/test_adapters_controller.py index a1503eb9f1..5ebe9016a5 100644 --- a/services/core/models/tests/unit/sidecars/test_adapters_controller.py +++ b/services/core/models/tests/unit/sidecars/test_adapters_controller.py @@ -43,6 +43,7 @@ def controller(tmp_path): ctrl = AdaptersController() ctrl.nim_peft_source = str(tmp_path) ctrl._sdk = MagicMock() + ctrl._models = MagicMock() ctrl.workspace = "default" ctrl.model_name = "base-model" # Default to NIM behavior (no rewrite, no eager vLLM load); vLLM tests @@ -399,7 +400,7 @@ def test_redownloads_when_fileset_changes(self, controller, tmp_path): mock_model_entity = MagicMock() mock_model_entity.workspace = "default" mock_model_entity.adapters = [adapter] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity mock_files_response = MagicMock() mock_files_response.data = [MagicMock()] @@ -431,7 +432,7 @@ def test_skips_download_when_metadata_matches(self, controller, tmp_path): mock_model_entity = MagicMock() mock_model_entity.workspace = "default" mock_model_entity.adapters = [adapter] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity dirs_to_keep: set[str] = set() controller._update_lora_adapters(dirs_to_keep) @@ -448,7 +449,7 @@ def test_downloads_new_adapter(self, controller, tmp_path): mock_model_entity = MagicMock() mock_model_entity.workspace = "default" mock_model_entity.adapters = [adapter] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity mock_files_response = MagicMock() mock_files_response.data = [MagicMock()] @@ -477,7 +478,7 @@ def test_no_orphaned_temp_dirs_after_download(self, controller, tmp_path): mock_model_entity = MagicMock() mock_model_entity.workspace = "default" mock_model_entity.adapters = [adapter] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity mock_files_response = MagicMock() mock_files_response.data = [MagicMock()] @@ -498,7 +499,7 @@ def test_failed_download_leaves_no_adapter_dir(self, controller, tmp_path): mock_model_entity = MagicMock() mock_model_entity.workspace = "default" mock_model_entity.adapters = [adapter] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity mock_files_response = MagicMock() mock_files_response.data = [] @@ -527,7 +528,7 @@ def test_failed_download_preserves_old_adapter(self, controller, tmp_path): mock_model_entity = MagicMock() mock_model_entity.workspace = "default" mock_model_entity.adapters = [adapter] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity mock_files_response = MagicMock() mock_files_response.data = [] @@ -556,7 +557,7 @@ def test_two_adapters_same_name_different_workspaces_coexist(self, controller, t mock_model_entity = MagicMock() mock_model_entity.workspace = "base-ws" mock_model_entity.adapters = [adapter_a, adapter_b] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity mock_files_response = MagicMock() mock_files_response.data = [MagicMock()] @@ -578,7 +579,7 @@ def test_dir_name_uses_adapter_workspace_not_base_model_workspace(self, controll mock_model_entity = MagicMock() mock_model_entity.workspace = "base-ws" mock_model_entity.adapters = [adapter] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity mock_files_response = MagicMock() mock_files_response.data = [MagicMock()] @@ -613,7 +614,7 @@ def test_bare_fileset_for_cross_workspace_adapter_fetches_from_adapter_workspace mock_model_entity = MagicMock() mock_model_entity.workspace = "base-ws" mock_model_entity.adapters = [adapter] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity mock_files_response = MagicMock() mock_files_response.data = [MagicMock()] @@ -641,14 +642,14 @@ def test_step_gc_removes_stale_dir_after_adapter_workspace_change(self, controll mock_model_entity = MagicMock() mock_model_entity.workspace = "base-ws" mock_model_entity.adapters = [adapter] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity mock_files_response = MagicMock() mock_files_response.data = [MagicMock()] controller._sdk.files.list.return_value = mock_files_response # No prompt-tuned models in this scenario. - controller._sdk.models.list.return_value = [] + controller._models.list_models.return_value.items.return_value = iter([]) controller.step() @@ -670,7 +671,7 @@ def test_adapter_changed_meta_check_works_against_new_dir_path(self, controller, mock_model_entity = MagicMock() mock_model_entity.workspace = "base-ws" mock_model_entity.adapters = [adapter] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity dirs_to_keep: set[str] = set() controller._update_lora_adapters(dirs_to_keep) @@ -692,7 +693,7 @@ def _model_entity(self, controller, adapter): me = MagicMock() me.workspace = "default" me.adapters = [adapter] - controller._sdk.models.retrieve.return_value = me + controller._models.get_model.return_value.data.return_value = me files_resp = MagicMock() files_resp.data = [MagicMock()] controller._sdk.files.list.return_value = files_resp @@ -823,7 +824,7 @@ def test_step_unloads_removed_adapter_before_delete(self, controller, tmp_path): adapter = _make_adapter("kept-adapter", "default/fs", updated_at=None, workspace="default") self._model_entity(controller, adapter) - controller._sdk.models.list.return_value = [] # no prompt-tuned models + controller._models.list_models.return_value.items.return_value = iter([]) # no prompt-tuned models with patch.object(controller, "_vllm_api_call", return_value=(200, "")) as api: controller.step() @@ -843,7 +844,7 @@ def test_step_keeps_removed_adapter_dir_when_vllm_unreachable(self, controller, adapter = _make_adapter("kept-adapter", "default/fs", updated_at=None, workspace="default") self._model_entity(controller, adapter) - controller._sdk.models.list.return_value = [] # no prompt-tuned models + controller._models.list_models.return_value.items.return_value = iter([]) # no prompt-tuned models # vLLM unreachable: both the kept adapter's load and the stale one's unload # hit a transport error. @@ -864,7 +865,7 @@ def test_step_deletes_removed_adapter_dir_when_vllm_answers_non_200(self, contro adapter = _make_adapter("kept-adapter", "default/fs", updated_at=None, workspace="default") self._model_entity(controller, adapter) - controller._sdk.models.list.return_value = [] # no prompt-tuned models + controller._models.list_models.return_value.items.return_value = iter([]) # no prompt-tuned models def _responses(route, payload): if route == "/v1/unload_lora_adapter": @@ -887,7 +888,7 @@ def test_step_keeps_removed_adapter_dir_when_vllm_server_error(self, controller, adapter = _make_adapter("kept-adapter", "default/fs", updated_at=None, workspace="default") self._model_entity(controller, adapter) - controller._sdk.models.list.return_value = [] # no prompt-tuned models + controller._models.list_models.return_value.items.return_value = iter([]) # no prompt-tuned models def _responses(route, payload): if route == "/v1/unload_lora_adapter": @@ -985,7 +986,7 @@ def __init__(self, name: str, fileset: str): mock_model_entity = MagicMock() mock_model_entity.workspace = "base-ws" mock_model_entity.adapters = [adapter] - controller._sdk.models.retrieve.return_value = mock_model_entity + controller._models.get_model.return_value.data.return_value = mock_model_entity mock_files_response = MagicMock() mock_files_response.data = [MagicMock()] From 0c6a9dc2ef60fb7065b3096496be12cc0d2a0dd9 Mon Sep 17 00:00:00 2001 From: Max Dubrinsky Date: Wed, 26 Aug 2026 19:32:29 -0400 Subject: [PATCH 3/4] test(models): fix integration/e2e CI fallout from typed-client migration - e2e/agents_deploy_helpers.py: revert wait_for_openai_model back to the Stainless sdk.models helper. It is a polling convenience method on the legacy resource, not a typed endpoint, so ModelsClient has no such method and the migration's reroute broke agent deploy/invoke e2e tests. - test_models_with_auth.py: the migration swapped every pytest.raises to the plugin PermissionDeniedError, but the provider/deployment_config/deployment calls still go through the unmigrated Stainless sdk.inference resources, which raise the Stainless PermissionDeniedError. Restore a Stainless alias and expect it for those 19 tests; the typed ModelsClient/SecretsClient tests keep the plugin error. Signed-off-by: Max Dubrinsky --- e2e/agents_deploy_helpers.py | 4 +- .../src/nmp/testing/e2e/customizer.py | 4 +- .../integration/test_models_with_auth.py | 39 ++++++++++--------- 3 files changed, 22 insertions(+), 25 deletions(-) diff --git a/e2e/agents_deploy_helpers.py b/e2e/agents_deploy_helpers.py index ff31fe1cbc..156db08561 100644 --- a/e2e/agents_deploy_helpers.py +++ b/e2e/agents_deploy_helpers.py @@ -28,8 +28,6 @@ import pytest from nemo_agents_plugin.entities import NAT_WORKFLOW_CONFIG_FORMAT, NEMO_AGENTS_SPEC_CONFIG_FORMAT from nemo_platform import NeMoPlatform -from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.models.client import ModelsClient from nmp.testing import MockProviderResponse, add_mock_provider # The mocked completion the deployed agent must round-trip back to the caller. @@ -276,7 +274,7 @@ def run_agent_deploy_and_invoke( endpoints = deployment.get("endpoints") or [] assert endpoints and endpoints[0]["url"], deployment - client_from_platform(sdk, ModelsClient).wait_for_openai_model(model_name, workspace=workspace) + sdk.models.wait_for_openai_model(model_name, workspace=workspace) response = sdk.agents.invoke( workspace=workspace, diff --git a/packages/nmp_testing/src/nmp/testing/e2e/customizer.py b/packages/nmp_testing/src/nmp/testing/e2e/customizer.py index 498857b5d2..d624238ef8 100644 --- a/packages/nmp_testing/src/nmp/testing/e2e/customizer.py +++ b/packages/nmp_testing/src/nmp/testing/e2e/customizer.py @@ -552,9 +552,7 @@ def _wait_for_gateway_ready( until it confirms the provider is routable. """ logger.info("Waiting for inference gateway to sync...") - if not client_from_platform(sdk, ModelsClient).wait_for_gateway( - deployment_name, workspace=workspace, timeout=timeout - ): + if not sdk.models.wait_for_gateway(deployment_name, workspace=workspace, timeout=timeout): pytest.fail( f"Inference gateway did not become ready for deployment '{deployment_name}' " f"within {timeout}s. The deployment's model provider may not have been created. " diff --git a/services/core/models/tests/integration/test_models_with_auth.py b/services/core/models/tests/integration/test_models_with_auth.py index 369b055b99..b13b4b0ee6 100644 --- a/services/core/models/tests/integration/test_models_with_auth.py +++ b/services/core/models/tests/integration/test_models_with_auth.py @@ -23,6 +23,7 @@ import pytest from nemo_platform import NeMoPlatform +from nemo_platform import PermissionDeniedError as StainlessPermissionDeniedError from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.client.errors import PermissionDeniedError from nemo_platform_plugin.files.client import FilesClient @@ -411,7 +412,7 @@ def test_viewer_can_get_provider(self, viewer_workspace): def test_viewer_cannot_create_provider(self, viewer_workspace): workspace, viewer_sdk, _, _ = viewer_workspace - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): viewer_sdk.inference.providers.create( workspace=workspace, name="should-fail", @@ -420,7 +421,7 @@ def test_viewer_cannot_create_provider(self, viewer_workspace): def test_viewer_cannot_upsert_provider(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): viewer_sdk.inference.providers.update( name=names["provider"], workspace=workspace, @@ -429,7 +430,7 @@ def test_viewer_cannot_upsert_provider(self, viewer_workspace): def test_viewer_cannot_delete_provider(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): viewer_sdk.inference.providers.delete(name=names["provider"], workspace=workspace) # -- Deployment Configs: allowed -- @@ -448,7 +449,7 @@ def test_viewer_can_get_deployment_config(self, viewer_workspace): def test_viewer_cannot_create_deployment_config(self, viewer_workspace): workspace, viewer_sdk, _, _ = viewer_workspace - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): viewer_sdk.inference.deployment_configs.create( workspace=workspace, name="should-fail", @@ -459,7 +460,7 @@ def test_viewer_cannot_create_deployment_config(self, viewer_workspace): def test_viewer_cannot_update_deployment_config(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): viewer_sdk.inference.deployment_configs.update( name=names["config"], workspace=workspace, @@ -470,7 +471,7 @@ def test_viewer_cannot_update_deployment_config(self, viewer_workspace): def test_viewer_cannot_delete_deployment_config(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): viewer_sdk.inference.deployment_configs.delete(name=names["config"], workspace=workspace) # -- Deployments: allowed -- @@ -489,7 +490,7 @@ def test_viewer_can_get_deployment(self, viewer_workspace): def test_viewer_cannot_create_deployment(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): viewer_sdk.inference.deployments.create( workspace=workspace, name="should-fail", @@ -498,7 +499,7 @@ def test_viewer_cannot_create_deployment(self, viewer_workspace): def test_viewer_cannot_update_deployment(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): viewer_sdk.inference.deployments.update( name=names["deployment"], workspace=workspace, @@ -507,7 +508,7 @@ def test_viewer_cannot_update_deployment(self, viewer_workspace): def test_viewer_cannot_delete_deployment(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): viewer_sdk.inference.deployments.delete(name=names["deployment"], workspace=workspace) @@ -711,7 +712,7 @@ def test_custom_role_denied_create_provider_with_secret(self, sdk: NeMoPlatform) ) assert provider_ok.api_key_secret_name is None - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): user_sdk.inference.providers.create( workspace=workspace, name=short_unique_name("prov"), @@ -740,7 +741,7 @@ def test_custom_role_denied_upsert_provider_with_secret(self, sdk: NeMoPlatform) user_sdk = as_user(sdk, user_email) - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): user_sdk.inference.providers.update( name=short_unique_name("prov"), workspace=workspace, @@ -816,7 +817,7 @@ def test_custom_role_denied_create_provider_with_deployment_ref(self, sdk: NeMoP user_sdk = as_user(sdk, user_email) - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): user_sdk.inference.providers.create( workspace=workspace, name=short_unique_name("prov"), @@ -841,7 +842,7 @@ def test_custom_role_denied_upsert_provider_with_deployment_ref(self, sdk: NeMoP user_sdk = as_user(sdk, user_email) - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): user_sdk.inference.providers.update( name=short_unique_name("prov"), workspace=workspace, @@ -894,7 +895,7 @@ def test_editor_cannot_reference_model_in_inaccessible_workspace(self, sdk: NeMo ) editor_sdk = as_user(sdk, editor_email) - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): editor_sdk.inference.deployment_configs.create( workspace=workspace, name=short_unique_name("cfg"), @@ -960,7 +961,7 @@ def test_editor_cannot_update_config_with_inaccessible_model(self, sdk: NeMoPlat ) editor_sdk = as_user(sdk, editor_email) - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): editor_sdk.inference.deployment_configs.update( name=config_name, workspace=workspace, @@ -987,7 +988,7 @@ def test_custom_role_denied_create_config_with_model_entity_id_without_read(self user_sdk = as_user(sdk, user_email) - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): user_sdk.inference.deployment_configs.create( workspace=workspace, name=short_unique_name("cfg"), @@ -1022,7 +1023,7 @@ def test_custom_role_denied_update_config_with_model_entity_id_without_read(self user_sdk = as_user(sdk, user_email) - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): user_sdk.inference.deployment_configs.update( name=config_name, workspace=workspace, @@ -1128,7 +1129,7 @@ def test_custom_role_denied_create_deployment_without_read(self, sdk: NeMoPlatfo user_sdk = as_user(sdk, user_email) - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): user_sdk.inference.deployments.create( workspace=workspace, name=short_unique_name("dep"), @@ -1166,7 +1167,7 @@ def test_custom_role_denied_update_deployment_without_read(self, sdk: NeMoPlatfo user_sdk = as_user(sdk, user_email) - with pytest.raises(PermissionDeniedError): + with pytest.raises(StainlessPermissionDeniedError): user_sdk.inference.deployments.update( name=deploy_name, workspace=workspace, From c1122f9e02afcebd7bef153bc9cdb7952f5f4958 Mon Sep 17 00:00:00 2001 From: Max Dubrinsky Date: Thu, 27 Aug 2026 14:56:12 -0400 Subject: [PATCH 4/4] fix(models): align SDK/plugin error handling after typed-client migration CodeRabbit flagged several catch blocks left inconsistent by the models migration: the migration switched imports to the plugin error hierarchy while some guarded calls stayed on the legacy SDK, making those catches no-ops that let the wrong error class propagate. The legacy nemo_platform errors and the plugin NemoClientError-derived errors are separate classes, so each call site must catch what it actually raises. Restore legacy errors for the calls that remain on the legacy SDK (inference providers/gateway/deployments, files) and keep plugin errors for typed clients (ModelsClient, FilesClient, JobsClient): - e2e/test_evaluator_plugin.py: revert error aliases to legacy imports (they guard legacy add_mock_provider and gateway.model.post). - inference-gateway model_cache.py: model_provider_getter stays legacy (inference.providers.list); model_entity_getter uses plugin errors. - adapters main.py: download_fileset catches legacy NotFoundError. - 4x agentic-use setup-mock.py: catch legacy ConflictError for the legacy providers.create; keep plugin ConflictError for create_model. - model_spec run.py: catch plugin NemoClientError (legacy NeMoPlatformError is dead now that get_model/update_model are typed). - e2e customizer.py: catch NemoTransportError from typed JobsClient, which wraps httpx transport errors, instead of the raw ones. - nemo_nb strip test: fixture now matches the client_from_platform assertions the converter produces. Signed-off-by: Max Dubrinsky --- e2e/test_evaluator_plugin.py | 4 +--- .../nemo_nb/tests/test_strip_type_checker_comments.py | 4 ++-- packages/nmp_testing/src/nmp/testing/e2e/customizer.py | 3 ++- .../src/nmp/core/inference_gateway/api/model_cache.py | 10 +++++----- .../src/nmp/core/models/sidecars/adapters/main.py | 3 +-- .../models/src/nmp/core/models/tasks/model_spec/run.py | 7 ++++--- .../environment/setup-mock.py | 5 ++++- .../environment/setup-mock.py | 5 ++++- .../environment/setup-mock.py | 5 ++++- .../environment/setup-mock.py | 5 ++++- 10 files changed, 31 insertions(+), 20 deletions(-) diff --git a/e2e/test_evaluator_plugin.py b/e2e/test_evaluator_plugin.py index b31e019353..48eb9a4669 100644 --- a/e2e/test_evaluator_plugin.py +++ b/e2e/test_evaluator_plugin.py @@ -46,11 +46,9 @@ from nemo_evaluator_sdk.metrics.tool_calling import ToolCallingMetric from nemo_evaluator_sdk.values.results import EvaluationResult from nemo_evaluator_sdk.values.scores import JSONScoreParser, RangeScore -from nemo_platform import NeMoPlatform +from nemo_platform import APIConnectionError, APIStatusError, NeMoPlatform from nemo_platform.types.inference import ModelProvider from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.client.errors import NemoHTTPError as APIStatusError -from nemo_platform_plugin.client.errors import NemoTransportError as APIConnectionError from nemo_platform_plugin.inference_middleware import BackendFormat from nemo_platform_plugin.jobs.client import JobsClient from nemo_platform_plugin.models.client import ModelsClient diff --git a/packages/nemo_nb/tests/test_strip_type_checker_comments.py b/packages/nemo_nb/tests/test_strip_type_checker_comments.py index 9bddee94c8..d1eec42237 100644 --- a/packages/nemo_nb/tests/test_strip_type_checker_comments.py +++ b/packages/nemo_nb/tests/test_strip_type_checker_comments.py @@ -17,8 +17,8 @@ def test_strip_ty_ignore_comments(): "source": [ "# This is a regular comment\n", "client_from_platform(sdk, ModelsClient).get_openai_route_base_url()\n", - "sdk.models.get_model_entity_route_openai_url(entity) # ty: ignore[unresolved-reference]\n", - "sdk.models.get_provider_route_openai_url(provider) # ty: ignore[unresolved-reference]\n", + "client_from_platform(sdk, ModelsClient).get_model_entity_route_openai_url(entity) # ty: ignore[unresolved-reference]\n", + "client_from_platform(sdk, ModelsClient).get_provider_route_openai_url(provider) # ty: ignore[unresolved-reference]\n", ], } ] diff --git a/packages/nmp_testing/src/nmp/testing/e2e/customizer.py b/packages/nmp_testing/src/nmp/testing/e2e/customizer.py index d624238ef8..53999fdfb4 100644 --- a/packages/nmp_testing/src/nmp/testing/e2e/customizer.py +++ b/packages/nmp_testing/src/nmp/testing/e2e/customizer.py @@ -18,6 +18,7 @@ from nemo_platform import NeMoPlatform from nemo_platform.types.inference import ContainerExecutorConfigParam, ModelDeploymentConfigModelSpecParam from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoTransportError from nemo_platform_plugin.jobs.client import JobsClient from nemo_platform_plugin.models.client import ModelsClient @@ -282,7 +283,7 @@ def wait_for_customization_job( try: status = jobs.get_job_status(name=job_name, workspace=workspace).data() consecutive_errors = 0 - except (httpx.TimeoutException, httpx.ConnectError, ConnectionError, OSError) as exc: + except NemoTransportError as exc: consecutive_errors += 1 logger.warning( f"Transient error polling customization job (attempt {consecutive_errors}/{max_consecutive_errors}, " diff --git a/services/core/inference-gateway/src/nmp/core/inference_gateway/api/model_cache.py b/services/core/inference-gateway/src/nmp/core/inference_gateway/api/model_cache.py index 2de3084a79..6cdd4afddc 100644 --- a/services/core/inference-gateway/src/nmp/core/inference_gateway/api/model_cache.py +++ b/services/core/inference-gateway/src/nmp/core/inference_gateway/api/model_cache.py @@ -9,11 +9,11 @@ from datetime import datetime from typing import TYPE_CHECKING, Awaitable, Callable -from nemo_platform import AsyncNeMoPlatform +from nemo_platform import APIConnectionError, APIStatusError, AsyncNeMoPlatform from nemo_platform.types.inference import ModelProvider, ServedModelMapping from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.client.errors import NemoHTTPError as APIStatusError -from nemo_platform_plugin.client.errors import NemoTransportError as APIConnectionError +from nemo_platform_plugin.client.errors import NemoHTTPError as PluginHTTPError +from nemo_platform_plugin.client.errors import NemoTransportError as PluginTransportError from nemo_platform_plugin.inference_middleware import BackendFormat from nemo_platform_plugin.models.client import AsyncModelsClient from nemo_platform_plugin.models.types import ModelEntity @@ -214,9 +214,9 @@ async def _model_entity_getter() -> list[ModelEntity]: ) models = [model async for model in resp.items()] return models - except APIConnectionError as exc: + except PluginTransportError as exc: raise ModelProviderRefreshError(f"Error connecting to models service: {exc}") from exc - except APIStatusError as exc: + except PluginHTTPError as exc: raise ModelProviderRefreshError(f"Error refreshing model entities from models service: {exc}") from exc return _model_entity_getter diff --git a/services/core/models/src/nmp/core/models/sidecars/adapters/main.py b/services/core/models/src/nmp/core/models/sidecars/adapters/main.py index 6f154754b8..28e9d8e9fc 100644 --- a/services/core/models/src/nmp/core/models/sidecars/adapters/main.py +++ b/services/core/models/src/nmp/core/models/sidecars/adapters/main.py @@ -13,9 +13,8 @@ import urllib.error import urllib.request -from nemo_platform import NeMoPlatform +from nemo_platform import NeMoPlatform, NotFoundError from nemo_platform_plugin.client.adapter import client_from_platform -from nemo_platform_plugin.client.errors import NotFoundError from nemo_platform_plugin.models.client import ModelsClient from nemo_platform_plugin.models.types import Adapter, ModelEntity from nmp.common.config import get_platform_config diff --git a/services/core/models/src/nmp/core/models/tasks/model_spec/run.py b/services/core/models/src/nmp/core/models/tasks/model_spec/run.py index ddce2e7d2a..65918b2936 100644 --- a/services/core/models/src/nmp/core/models/tasks/model_spec/run.py +++ b/services/core/models/src/nmp/core/models/tasks/model_spec/run.py @@ -18,10 +18,11 @@ import os from pathlib import Path -from nemo_platform import NeMoPlatform, NeMoPlatformError +from nemo_platform import NeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.client.errors import ( InternalServerError, + NemoClientError, NotFoundError, ) from nemo_platform_plugin.client.errors import ( @@ -195,7 +196,7 @@ def analyze_checkpoint(self, config: ModelSpecTaskConfig) -> ModelEntity: raise ModelSpecCreationError( f"Failed to create model spec: model entity {config.workspace}/{config.name} does not exist" ) from err - except NeMoPlatformError as err: + except NemoClientError as err: raise ModelSpecCreationError( f"Failed to create model spec: model entity {config.workspace}/{config.name} unable to be fetched" ) from err @@ -307,7 +308,7 @@ def analyze_checkpoint(self, config: ModelSpecTaskConfig) -> ModelEntity: raise ModelSpecCreationError( f"Failed to update model spec: model entity {config.workspace}/{config.name} does not exist" ) from err - except NeMoPlatformError as err: + except NemoClientError as err: raise ModelSpecCreationError( f"Failed to update model spec: model entity {config.workspace}/{config.name} unable to be fetched" ) from err diff --git a/tests/agentic-use/guardrails-content-safety-cli-easy/environment/setup-mock.py b/tests/agentic-use/guardrails-content-safety-cli-easy/environment/setup-mock.py index 23cd2eff5b..c6036d3a20 100644 --- a/tests/agentic-use/guardrails-content-safety-cli-easy/environment/setup-mock.py +++ b/tests/agentic-use/guardrails-content-safety-cli-easy/environment/setup-mock.py @@ -23,6 +23,9 @@ NotFoundError, UnprocessableEntityError, ) +from nemo_platform import ( + ConflictError as SDKConflictError, +) from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.client.errors import ConflictError from nemo_platform_plugin.models.client import ModelsClient @@ -192,7 +195,7 @@ def setup() -> None: host_url=MOCK_PROVIDER_HOST_URL, default_extra_headers=desired_headers, ) - except ConflictError: + except SDKConflictError: print(f"Provider already exists, reconciling headers: {MOCK_PROVIDER_NAME}") sdk.inference.providers.update( name=MOCK_PROVIDER_NAME, diff --git a/tests/agentic-use/guardrails-content-safety-cli/environment/setup-mock.py b/tests/agentic-use/guardrails-content-safety-cli/environment/setup-mock.py index 23cd2eff5b..c6036d3a20 100644 --- a/tests/agentic-use/guardrails-content-safety-cli/environment/setup-mock.py +++ b/tests/agentic-use/guardrails-content-safety-cli/environment/setup-mock.py @@ -23,6 +23,9 @@ NotFoundError, UnprocessableEntityError, ) +from nemo_platform import ( + ConflictError as SDKConflictError, +) from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.client.errors import ConflictError from nemo_platform_plugin.models.client import ModelsClient @@ -192,7 +195,7 @@ def setup() -> None: host_url=MOCK_PROVIDER_HOST_URL, default_extra_headers=desired_headers, ) - except ConflictError: + except SDKConflictError: print(f"Provider already exists, reconciling headers: {MOCK_PROVIDER_NAME}") sdk.inference.providers.update( name=MOCK_PROVIDER_NAME, diff --git a/tests/agentic-use/inference-chat-completions-cli-easy/environment/setup-mock.py b/tests/agentic-use/inference-chat-completions-cli-easy/environment/setup-mock.py index 4550e6ab66..25ea3d6ed9 100644 --- a/tests/agentic-use/inference-chat-completions-cli-easy/environment/setup-mock.py +++ b/tests/agentic-use/inference-chat-completions-cli-easy/environment/setup-mock.py @@ -22,6 +22,9 @@ NotFoundError, UnprocessableEntityError, ) +from nemo_platform import ( + ConflictError as SDKConflictError, +) from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.client.errors import ConflictError from nemo_platform_plugin.models.client import ModelsClient @@ -210,7 +213,7 @@ def setup() -> None: host_url=MOCK_PROVIDER_HOST_URL, default_extra_headers=desired_headers, ) - except ConflictError: + except SDKConflictError: print(f"Provider already exists, reconciling headers: {MOCK_PROVIDER_NAME}") sdk.inference.providers.update( name=MOCK_PROVIDER_NAME, diff --git a/tests/agentic-use/inference-chat-completions-cli/environment/setup-mock.py b/tests/agentic-use/inference-chat-completions-cli/environment/setup-mock.py index 4550e6ab66..25ea3d6ed9 100644 --- a/tests/agentic-use/inference-chat-completions-cli/environment/setup-mock.py +++ b/tests/agentic-use/inference-chat-completions-cli/environment/setup-mock.py @@ -22,6 +22,9 @@ NotFoundError, UnprocessableEntityError, ) +from nemo_platform import ( + ConflictError as SDKConflictError, +) from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.client.errors import ConflictError from nemo_platform_plugin.models.client import ModelsClient @@ -210,7 +213,7 @@ def setup() -> None: host_url=MOCK_PROVIDER_HOST_URL, default_extra_headers=desired_headers, ) - except ConflictError: + except SDKConflictError: print(f"Provider already exists, reconciling headers: {MOCK_PROVIDER_NAME}") sdk.inference.providers.update( name=MOCK_PROVIDER_NAME,