diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/inference_gateway/client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/inference_gateway/client.py new file mode 100644 index 0000000000..f8bff69315 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/inference_gateway/client.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed clients for narrow Inference Gateway provider proxy calls.""" + +from __future__ import annotations + +import json + +import httpx +from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient +from nemo_platform_plugin.client.method import method +from nemo_platform_plugin.inference_gateway import endpoints + + +def _decode_provider_proxy_body(content: bytes) -> object: + """Decode a provider proxy body. + + Provider proxy responses are intentionally dynamic because the gateway + forwards upstream provider payloads. Decode JSON when possible and return + text otherwise so callers can validate the small shape they need. + """ + text = content.decode("utf-8", errors="replace") + try: + return json.loads(text) + except json.JSONDecodeError: + return text + + +class _InferenceGatewayProviderMethods: + get_provider_models_raw = method(endpoints.get_provider_models_raw) + + +class InferenceGatewayProviderClient(_InferenceGatewayProviderMethods, NemoClient): + """Sync client for Inference Gateway provider proxy reads.""" + + def get_provider_models( + self, + *, + workspace: str | None = None, + name: str, + timeout: float | httpx.Timeout | None = None, + ) -> object: + """Return the decoded ``GET /v1/models`` payload for a provider.""" + client = self.with_options(timeout=timeout) if timeout is not None else self + response = client.get_provider_models_raw(workspace=workspace, name=name) + return _decode_provider_proxy_body(response.read()) + + +class AsyncInferenceGatewayProviderClient(_InferenceGatewayProviderMethods, AsyncNemoClient): + """Async client for Inference Gateway provider proxy reads.""" + + async def get_provider_models( + self, + *, + workspace: str | None = None, + name: str, + timeout: float | httpx.Timeout | None = None, + ) -> object: + """Return the decoded ``GET /v1/models`` payload for a provider.""" + client = self.with_options(timeout=timeout) if timeout is not None else self + response = await client.get_provider_models_raw(workspace=workspace, name=name) + return _decode_provider_proxy_body(await response.read()) diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/inference_gateway/endpoints.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/inference_gateway/endpoints.py new file mode 100644 index 0000000000..d3af3a2ad1 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/inference_gateway/endpoints.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed endpoint definitions for narrow Inference Gateway provider calls.""" + +from __future__ import annotations + +from abc import abstractmethod + +from nemo_platform_plugin.client.endpoint import get +from nemo_platform_plugin.client.types import BinaryContent + +_PROVIDER = "/apis/inference-gateway/v2/workspaces/{workspace}/provider/{name}/-" + + +@get(_PROVIDER + "/v1/models") +@abstractmethod +def get_provider_models_raw(*, workspace: str | None = None, name: str) -> BinaryContent: ... diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/models/client.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/models/client.py index 236be3c17c..c567b31f60 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/models/client.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/models/client.py @@ -37,7 +37,11 @@ from datetime import datetime from typing import Protocol -from models import ( +from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient +from nemo_platform_plugin.client.errors import NotFoundError +from nemo_platform_plugin.client.method import method +from nemo_platform_plugin.models import endpoints +from nemo_platform_plugin.models.refs import ( ResolvedModelReference, first_provider_ref, model_entity_route_openai_url, @@ -45,10 +49,6 @@ resolved_model_reference, warn_provider_host_url_resolution_failure, ) -from nemo_platform_plugin.client.client import AsyncNemoClient, NemoClient -from nemo_platform_plugin.client.errors import NotFoundError -from nemo_platform_plugin.client.method import method -from nemo_platform_plugin.models import endpoints from nemo_platform_plugin.models.types import ModelDeployment, ModelEntity _INFERENCE_GATEWAY_PREFIX = "/apis/inference-gateway/v2/workspaces" diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/models/refs.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/models/refs.py new file mode 100644 index 0000000000..1c260dbcf3 --- /dev/null +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/models/refs.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pure Models route-reference helpers. + +These helpers mirror the convenience functions historically exported from the +Stainless-backed ``models`` package. They live in the plugin client package so +typed clients can build route references without importing generated resources. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass + +_logger = logging.getLogger(__name__) + + +@dataclass(frozen=True, slots=True) +class ResolvedModelReference: + """Inference route details for a workspace-qualified model reference.""" + + url: str + name: str + host_url: str | None + + +def parse_workspace_name_ref(ref: str, *, label: str, expected_format: str = "workspace/name") -> tuple[str, str]: + """Parse a strict workspace-qualified reference.""" + workspace, separator, name = ref.partition("/") + if separator != "/" or not workspace or not name or "/" in name: + raise ValueError(f"{label} must be in format '{expected_format}'") + return workspace, name + + +def first_provider_ref(model_providers: list[str] | None) -> tuple[str, str, str] | None: + """Return the first valid ``(ref, workspace, name)`` provider reference, if present.""" + if not model_providers: + return None + + provider_ref = model_providers[0] + try: + provider_workspace, provider_name = parse_workspace_name_ref(provider_ref, label="Provider reference") + except ValueError: + _logger.warning("Invalid provider reference format", extra={"provider_ref": provider_ref}) + return None + return provider_ref, provider_workspace, provider_name + + +def model_entity_route_openai_url(*, base_url: str, workspace: str, name: str) -> str: + """OpenAI SDK-compatible URL for a model-entity proxy route.""" + return f"{base_url.rstrip('/')}/apis/inference-gateway/v2/workspaces/{workspace}/model/{name}/-/v1" + + +def resolved_model_reference( + *, + base_url: str, + name: str, + route_workspace: str, + route_model_name: str, + host_url: str | None, +) -> ResolvedModelReference: + """Build route details for a resolved model entity.""" + return ResolvedModelReference( + url=model_entity_route_openai_url(base_url=base_url, workspace=route_workspace, name=route_model_name), + name=name, + host_url=host_url, + ) + + +def warn_provider_host_url_resolution_failure( + provider_ref: str, + exc: Exception, + *, + not_found_error_type: type[Exception], +) -> None: + """Log a provider host-url lookup failure with the expected severity.""" + if isinstance(exc, not_found_error_type): + _logger.warning("Provider not found during host_url resolution", extra={"provider_ref": provider_ref}) + return + _logger.warning("Failed to resolve provider host_url", extra={"provider_ref": provider_ref}, exc_info=True) 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 22fdc04065..02090b0256 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 @@ -14,11 +14,11 @@ from contextvars import ContextVar from dataclasses import dataclass -from models import parse_workspace_name_ref from nemo_platform import AsyncNeMoPlatform 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.refs import parse_workspace_name_ref from nemo_platform_plugin.models.types import ModelEntity, ModelProvider from nooa.unifiedllm import CompletionClient, UnifiedLLM diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/virtual_models/endpoints.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/virtual_models/endpoints.py index 5972c6814b..eda8a30735 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/virtual_models/endpoints.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/virtual_models/endpoints.py @@ -11,6 +11,7 @@ from nemo_platform_plugin.client.types import Paginated from nemo_platform_plugin.virtual_models.types import ( CreateVirtualModelRequest, + DeleteVirtualModelQueryParams, ListVirtualModelsQueryParams, UpdateVirtualModelRequest, VirtualModel, @@ -45,4 +46,6 @@ def update_virtual_model( @delete(_VIRTUAL_MODELS + "/{name}") @abstractmethod -def delete_virtual_model(*, workspace: str | None = None, name: str) -> None: ... +def delete_virtual_model( + *, workspace: str | None = None, name: str, query_params: DeleteVirtualModelQueryParams | None = None +) -> None: ... diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/virtual_models/types.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/virtual_models/types.py index f8587b5dab..31ad7a0e87 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/virtual_models/types.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/virtual_models/types.py @@ -27,6 +27,7 @@ __all__ = [ "CreateVirtualModelRequest", + "DeleteVirtualModelQueryParams", "ListVirtualModelsQueryParams", "MiddlewareCall", "UpdateVirtualModelRequest", @@ -117,3 +118,9 @@ class ListVirtualModelsQueryParams(TypedDict, total=False): sort: NotRequired[str] filter: NotRequired[str] exclude_autoprovisioned: NotRequired[bool] + + +class DeleteVirtualModelQueryParams(TypedDict, total=False): + """Query parameters accepted by the VirtualModel delete operation.""" + + expected_db_version: NotRequired[int] diff --git a/packages/nemo_platform_plugin/tests/inference_gateway/test_client.py b/packages/nemo_platform_plugin/tests/inference_gateway/test_client.py new file mode 100644 index 0000000000..40f55e9fac --- /dev/null +++ b/packages/nemo_platform_plugin/tests/inference_gateway/test_client.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Inference Gateway provider client tests.""" + +from __future__ import annotations + +import httpx +from nemo_platform_plugin.inference_gateway.client import ( + AsyncInferenceGatewayProviderClient, + InferenceGatewayProviderClient, +) + +BASE = "http://test:8000" + + +def test_get_provider_models_decodes_json_and_uses_provider_route() -> None: + seen: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request) + return httpx.Response(200, request=request, json={"object": "list", "data": [{"id": "model-a"}]}) + + client = InferenceGatewayProviderClient( + base_url=BASE, + workspace="default", + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + + result = client.get_provider_models(workspace="team-a", name="provider-a") + + assert result == {"object": "list", "data": [{"id": "model-a"}]} + assert seen[0].url.path == "/apis/inference-gateway/v2/workspaces/team-a/provider/provider-a/-/v1/models" + + +async def test_async_get_provider_models_returns_text_for_non_json_body() -> None: + async def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, request=request, content=b"not-json") + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client = AsyncInferenceGatewayProviderClient( + base_url=BASE, + workspace="default", + http_client=http_client, + ) + + result = await client.get_provider_models(name="provider-a") + + assert result == "not-json" 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 1ea9f72904..51b271e4ab 100644 --- a/services/core/models/src/nmp/core/models/app/utils.py +++ b/services/core/models/src/nmp/core/models/app/utils.py @@ -9,9 +9,6 @@ from logging import getLogger from typing import Generic, List, Optional, TypeVar -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_plugin.k8s_naming import ( DNS_LABEL_MAX_LENGTH, DNS_SUBDOMAIN_MAX_LENGTH, @@ -19,7 +16,7 @@ k8s_safe_name, workspace_name_identity, ) -from nemo_platform_plugin.models.types import ModelEntity +from nemo_platform_plugin.models.types import ModelDeployment, ModelDeploymentConfig, ModelEntity, ModelProvider from nmp.common.api.common import PaginationData from nmp.common.entities.constants import NAME_PATTERN as ENTITY_NAME_PATTERN from pydantic import BaseModel @@ -108,7 +105,7 @@ def parse_model_name_revision( parsed_name = name_without_revision # Parse namespace prefix only if explicit model_namespace was NOT provided - if not model_namespace and "/" in parsed_name: + if not model_namespace and parsed_name is not None and "/" in parsed_name: # Split on first / to extract namespace parts = parsed_name.split("/", 1) parsed_namespace = parts[0] @@ -169,12 +166,10 @@ def get_model_weights_type( if model_entity and model_entity.fileset: return ModelWeightsType.FILES_SERVICE - # Guard the nested groups: a partial/legacy config may omit executor_config or - # model_spec, and we must not raise AttributeError while resolving weights. - executor_cfg = getattr(model_deployment_config, "executor_config", None) - model_spec_cfg = getattr(model_deployment_config, "model_spec", None) - image_name = getattr(executor_cfg, "image_name", None) - model_name = getattr(model_spec_cfg, "model_name", None) + executor_cfg = model_deployment_config.executor_config if model_deployment_config else None + model_spec_cfg = model_deployment_config.model_spec if model_deployment_config else None + image_name = executor_cfg.image_name if executor_cfg else None + model_name = model_spec_cfg.model_name if model_spec_cfg else None # If the model is a multi-LLM, we have already ruled out HF weights, so we download from Files service if is_multi_llm_image(image_name) and model_name: diff --git a/services/core/models/src/nmp/core/models/controllers/backends/backends.py b/services/core/models/src/nmp/core/models/controllers/backends/backends.py index 44d4dc5fe3..62817ed2fb 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/backends.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/backends.py @@ -7,7 +7,7 @@ from typing import Any, Dict from nemo_platform import AsyncNeMoPlatform -from nemo_platform.types.inference import ModelDeploymentStatus +from nemo_platform_plugin.models.types import ModelDeploymentStatus from nmp.core.models.controllers.context import ModelContext from pydantic import BaseModel diff --git a/services/core/models/src/nmp/core/models/controllers/backends/common.py b/services/core/models/src/nmp/core/models/controllers/backends/common.py index fd9a903604..497d8d8b61 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/common.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/common.py @@ -7,9 +7,7 @@ from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Protocol -from nemo_platform.types.inference.k8s_nim_operator_config import K8sNIMOperatorConfig -from nemo_platform.types.inference.model_deployment import ModelDeployment -from nemo_platform.types.shared.tool_call_config import ToolCallConfig +from nemo_platform_plugin.models.types import K8sNIMOperatorConfig, ModelDeployment, ToolCallConfig LOG_TAIL_LINES = 80 LOG_MAX_CHARS = 2048 diff --git a/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/backend.py b/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/backend.py index 3fef6a916d..13210da443 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/backend.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/backend.py @@ -8,11 +8,11 @@ from nemo_deployments_plugin.entities import Deployment, DeploymentConfig, Prerequisite, Volume from nemo_platform import AsyncNeMoPlatform -from nemo_platform.types.inference.model_deployment import ModelDeployment from nemo_platform_plugin.auth import AuthContext as DeploymentAuthContext from nemo_platform_plugin.client.adapter import client_from_platform from nemo_platform_plugin.entities.client import AsyncEntitiesClient from nemo_platform_plugin.entity_client import NemoEntitiesClient, NemoEntityConflictError, NemoEntityNotFoundError +from nemo_platform_plugin.models.types import ModelDeployment, ModelDeploymentStatus from nemo_platform_plugin.sdk_provider import get_async_platform_sdk from nmp.common.config import Runtime from nmp.core.models.app.constants import MODEL_MANAGED_BY_LABEL, MODEL_MANAGED_BY_MODELS_CONTROLLER @@ -81,18 +81,19 @@ async def create_model_deployment(self, ctx: ModelContext) -> DeploymentStatusUp resolved = resolve_plugin_deployment(ctx, self._huggingface_model_puller) if resolved.runtime == Runtime.NONE: return DeploymentStatusUpdate( - status="UNKNOWN", status_message="Deployments plugin is unavailable for runtime none." + status=ModelDeploymentStatus.UNKNOWN, + status_message="Deployments plugin is unavailable for runtime none.", ) teardown = await self.delete_model_deployment(resolved.deployment.workspace, resolved.deployment.name) - if teardown.status == "DELETING": + if teardown.status == ModelDeploymentStatus.DELETING: return DeploymentStatusUpdate( - status="PENDING", + status=ModelDeploymentStatus.PENDING, status_message="Waiting for prior deployments-plugin substrate teardown before recreate.", ) executor = executor_for_runtime(self._cfg, resolved.runtime) if executor is None: return DeploymentStatusUpdate( - status="ERROR", + status=ModelDeploymentStatus.ERROR, status_message=( "No deployments-plugin executor configured for the current runtime. " "Set docker_executor, k8s_executor, or default_executor under " @@ -142,11 +143,14 @@ async def create_model_deployment(self, ctx: ModelContext) -> DeploymentStatusUp except Exception as exc: await self._rollback_create(ctx) return DeploymentStatusUpdate( - status="ERROR", + status=ModelDeploymentStatus.ERROR, status_message=f"Unable to create deployments-plugin entities: {exc}", error_details={"error": str(exc)}, ) - return DeploymentStatusUpdate(status="PENDING", status_message="Created deployments-plugin entities.") + return DeploymentStatusUpdate( + status=ModelDeploymentStatus.PENDING, + status_message="Created deployments-plugin entities.", + ) async def _rollback_create(self, ctx: ModelContext) -> None: """Best-effort controlled teardown after a partial create failure.""" @@ -171,12 +175,20 @@ async def get_model_deployment_status(self, ctx: ModelContext) -> DeploymentStat ``pending_timeout_seconds`` when the deployment remains PENDING too long. """ if ctx.model_deployment is None: - return DeploymentStatusUpdate(status="UNKNOWN", status_message="Model deployment unavailable.") + return DeploymentStatusUpdate( + status=ModelDeploymentStatus.UNKNOWN, + status_message="Model deployment unavailable.", + ) names = entity_names(ctx.model_deployment.name) server = await self._get_optional(Deployment, ctx.model_deployment.workspace, names.server) puller = await self._get_optional(Deployment, ctx.model_deployment.workspace, names.puller) volume = await self._get_optional(Volume, ctx.model_deployment.workspace, names.volume) - result = aggregate_status(volume, puller, server, previously_ready=ctx.model_deployment.status == "READY") + result = aggregate_status( + volume, + puller, + server, + previously_ready=ctx.model_deployment.status == ModelDeploymentStatus.READY, + ) elapsed = deployment_elapsed_seconds(ctx.model_deployment) return apply_pending_timeout( result, @@ -187,7 +199,10 @@ async def get_model_deployment_status(self, ctx: ModelContext) -> DeploymentStat async def update_model_deployment(self, ctx: ModelContext) -> DeploymentStatusUpdate: del ctx - return DeploymentStatusUpdate(status="ERROR", status_message="Update via recreate not yet supported.") + return DeploymentStatusUpdate( + status=ModelDeploymentStatus.ERROR, + status_message="Update via recreate not yet supported.", + ) async def delete_model_deployment( self, @@ -201,7 +216,8 @@ async def delete_model_deployment( for deployment_name, config_name in ((names.server, names.server), (names.puller, names.puller)): if not await self._complete_deployment_delete(workspace, deployment_name, config_name): result = DeploymentStatusUpdate( - status="DELETING", status_message="Waiting for plugin deployment teardown." + status=ModelDeploymentStatus.DELETING, + status_message="Waiting for plugin deployment teardown.", ) return apply_deleting_timeout( result, @@ -214,7 +230,10 @@ async def delete_model_deployment( await self._entity_client().delete(Volume, name=volume_name, workspace=workspace) except NemoEntityNotFoundError: pass - return DeploymentStatusUpdate(status="DELETED", status_message="Deleted deployments-plugin entities.") + return DeploymentStatusUpdate( + status=ModelDeploymentStatus.DELETED, + status_message="Deleted deployments-plugin entities.", + ) async def _complete_deployment_delete(self, workspace: str, deployment_name: str, config_name: str) -> bool: """Initiate plugin deployment stop and return True once config can be removed.""" 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 b2b76625ec..5ac03a67e0 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 @@ -27,8 +27,7 @@ Toleration, VolumeMount, ) -from nemo_platform.types.inference.k8s_nim_operator_config import K8sNIMOperatorConfig -from nemo_platform_plugin.models.types import ModelEntity +from nemo_platform_plugin.models.types import K8sNIMOperatorConfig, 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 @@ -251,12 +250,10 @@ def _image(name: str, tag: str) -> str: return name if "@" in name or name.endswith(f":{tag}") else f"{name}:{tag}" -def _k8s_config_dict(k8s_config: K8sNIMOperatorConfig | dict[str, Any] | Any) -> dict[str, Any]: - if hasattr(k8s_config, "model_dump"): +def _k8s_config_dict(k8s_config: K8sNIMOperatorConfig | dict[str, Any]) -> dict[str, Any]: + if isinstance(k8s_config, K8sNIMOperatorConfig): return k8s_config.model_dump(exclude_none=True) - if isinstance(k8s_config, dict): - return {key: value for key, value in k8s_config.items() if value is not None} - return {} + return {key: value for key, value in k8s_config.items() if value is not None} def _tolerations_from_config(raw: list[dict[str, Any]]) -> list[Toleration]: @@ -268,16 +265,19 @@ def _tolerations_from_config(raw: list[dict[str, Any]]) -> list[Toleration]: def _affinity_from_node_selector(node_selector: dict[str, str]) -> Affinity: - return Affinity( - node_affinity={ - "requiredDuringSchedulingIgnoredDuringExecution": { - "nodeSelectorTerms": [ - { - "matchExpressions": [ - {"key": key, "operator": "In", "values": [value]} for key, value in node_selector.items() - ] - } - ] + return Affinity.model_validate( + { + "nodeAffinity": { + "requiredDuringSchedulingIgnoredDuringExecution": { + "nodeSelectorTerms": [ + { + "matchExpressions": [ + {"key": key, "operator": "In", "values": [value]} + for key, value in node_selector.items() + ] + } + ] + } } } ) @@ -323,7 +323,13 @@ def pod_security_context_for_engine( group_id = view.run_as_group if view.run_as_group is not None else config.default_group_id if user_id is None and group_id is None: return None - return PodSecurityContext(run_as_user=user_id, run_as_group=group_id, fs_group=group_id) + return PodSecurityContext.model_validate( + { + "runAsUser": user_id, + "runAsGroup": group_id, + "fsGroup": group_id, + } + ) def _default_tolerations(config: DeploymentsPluginConfig) -> list[Toleration]: @@ -408,8 +414,10 @@ def startup_probe_failure_threshold(view: DeploymentConfigView, *, period_second def apply_container_resources(container: Container, resources: dict[str, Any]) -> None: """Apply k8s resource requirements to a plugin container.""" - requests = resources.get("requests") if isinstance(resources.get("requests"), dict) else {} - limits = resources.get("limits") if isinstance(resources.get("limits"), dict) else {} + raw_requests = resources.get("requests") + raw_limits = resources.get("limits") + requests = raw_requests if isinstance(raw_requests, dict) else {} + limits = raw_limits if isinstance(raw_limits, dict) else {} if not requests and not limits: return existing = container.resources 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 9f49c2ea0f..26c5e8a90f 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 @@ -6,9 +6,7 @@ from dataclasses import dataclass from urllib.parse import SplitResult, urljoin, urlsplit -from nemo_platform.types.inference.model_deployment import ModelDeployment -from nemo_platform.types.inference.model_deployment_config import ModelDeploymentConfig -from nemo_platform_plugin.models.types import ModelEntity +from nemo_platform_plugin.models.types import ModelDeployment, ModelDeploymentConfig, 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/deployments_plugin/status.py b/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/status.py index 4a3fbce966..aed560d42c 100644 --- a/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/status.py +++ b/services/core/models/src/nmp/core/models/controllers/backends/deployments_plugin/status.py @@ -7,25 +7,25 @@ from nemo_deployments_plugin.entities import Deployment, Volume from nemo_deployments_plugin.types import Endpoint -from nemo_platform.types.inference import ModelDeploymentStatus +from nemo_platform_plugin.models.types import ModelDeploymentStatus from nmp.core.models.controllers.backends.backends import DeploymentStatusUpdate from nmp.core.models.controllers.backends.common import format_duration _STATUS_MAP: dict[str, ModelDeploymentStatus] = { - "PENDING": "PENDING", - "STARTING": "PENDING", - "READY": "READY", - "FAILED": "ERROR", - "LOST": "LOST", - "UNKNOWN": "UNKNOWN", - "DELETING": "DELETING", - "SUCCEEDED": "PENDING", + "PENDING": ModelDeploymentStatus.PENDING, + "STARTING": ModelDeploymentStatus.PENDING, + "READY": ModelDeploymentStatus.READY, + "FAILED": ModelDeploymentStatus.ERROR, + "LOST": ModelDeploymentStatus.LOST, + "UNKNOWN": ModelDeploymentStatus.UNKNOWN, + "DELETING": ModelDeploymentStatus.DELETING, + "SUCCEEDED": ModelDeploymentStatus.PENDING, } def map_status(status: str) -> ModelDeploymentStatus: """Map a deployments-plugin status to a ModelDeployment status.""" - return _STATUS_MAP.get(status, "UNKNOWN") + return _STATUS_MAP.get(status, ModelDeploymentStatus.UNKNOWN) def project_host_url(endpoints: Iterable[Endpoint]) -> str | None: @@ -57,7 +57,7 @@ def _substrate_issue( ) -> DeploymentStatusUpdate | None: if entity is None or entity.status not in _ATTENTION_SUBSTRATE_STATUSES: return None - status: ModelDeploymentStatus = "UNKNOWN" if entity.status == "UNKNOWN" else "ERROR" + status = ModelDeploymentStatus.UNKNOWN if entity.status == "UNKNOWN" else ModelDeploymentStatus.ERROR return DeploymentStatusUpdate( status=status, status_message=entity.status_message or f"{label} is {entity.status}.", @@ -80,11 +80,11 @@ def aggregate_status( status=status, status_message=server.status_message or f"Server deployment is {server.status}.", error_details={"substrate": substrate}, - host_url=project_host_url(server.endpoints) if status == "READY" else None, + host_url=project_host_url(server.endpoints) if status == ModelDeploymentStatus.READY else None, ) if previously_ready: return DeploymentStatusUpdate( - status="LOST", + status=ModelDeploymentStatus.LOST, status_message="Serving deployment is missing after reporting READY.", error_details={"substrate": substrate}, ) @@ -94,7 +94,7 @@ def aggregate_status( if issue is not None: return issue return DeploymentStatusUpdate( - status="PENDING", + status=ModelDeploymentStatus.PENDING, status_message="Waiting for deployments-plugin substrate resources.", error_details={"substrate": substrate}, ) @@ -121,7 +121,7 @@ def build_pending_timeout_error( if substrate is not None: error_details["substrate"] = substrate return DeploymentStatusUpdate( - status="ERROR", + status=ModelDeploymentStatus.ERROR, status_message=status_msg, error_details=error_details, ) @@ -135,7 +135,7 @@ def apply_pending_timeout( deployment_name: str, ) -> DeploymentStatusUpdate: """Escalate a PENDING projection to ERROR once the deployment ages out.""" - if result.status != "PENDING" or elapsed_seconds < timeout_seconds: + if result.status != ModelDeploymentStatus.PENDING or elapsed_seconds < timeout_seconds: return result substrate = result.error_details.get("substrate") if result.error_details else None return build_pending_timeout_error( @@ -158,7 +158,7 @@ def build_deleting_timeout_error( f"deployments-plugin substrate teardown (timeout: {format_duration(timeout_seconds)})." ) return DeploymentStatusUpdate( - status="ERROR", + status=ModelDeploymentStatus.ERROR, status_message=status_msg, error_details={ "reason": "deleting_timeout", @@ -177,7 +177,7 @@ def apply_deleting_timeout( deployment_name: str, ) -> DeploymentStatusUpdate: """Escalate a DELETING delete result to ERROR once teardown ages out.""" - if result.status != "DELETING" or timeout_seconds <= 0 or elapsed_seconds < timeout_seconds: + if result.status != ModelDeploymentStatus.DELETING or timeout_seconds <= 0 or elapsed_seconds < timeout_seconds: return result return build_deleting_timeout_error( deployment_name=deployment_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 53480e94e6..67f09c6b84 100644 --- a/services/core/models/src/nmp/core/models/controllers/context.py +++ b/services/core/models/src/nmp/core/models/controllers/context.py @@ -6,11 +6,13 @@ from dataclasses import dataclass from typing import Optional -from nemo_platform.types.inference import ServedModelMapping -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_plugin.models.types import ModelEntity +from nemo_platform_plugin.models.types import ( + ModelDeployment, + ModelDeploymentConfig, + ModelEntity, + ModelProvider, + ServedModelMapping, +) @dataclass diff --git a/services/core/models/src/nmp/core/models/controllers/deployment_reconciler.py b/services/core/models/src/nmp/core/models/controllers/deployment_reconciler.py index a916362d6e..03f0747ec9 100644 --- a/services/core/models/src/nmp/core/models/controllers/deployment_reconciler.py +++ b/services/core/models/src/nmp/core/models/controllers/deployment_reconciler.py @@ -11,9 +11,18 @@ from typing import Awaitable, Callable, Optional from nemo_platform import AsyncNeMoPlatform -from nemo_platform._exceptions import ConflictError, NotFoundError -from nemo_platform.types.inference.model_deployment import ModelDeployment -from nemo_platform.types.inference.model_provider import ModelProvider +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 ( + CreateModelProviderRequest, + ModelDeployment, + ModelDeploymentStatus, + ModelProvider, + ModelProviderStatus, + UpdateModelDeploymentStatusRequest, + UpsertModelProviderRequest, +) from nmp.common.entities.utils import parse_entity_ref from nmp.core.models.config import ControllerConfig from nmp.core.models.controllers.backends.backends import DeploymentStatusUpdate, ServiceBackend @@ -159,6 +168,7 @@ def __init__( distinguishable from a stalled one """ self._models_sdk = models_sdk + self._models_client = client_from_platform(models_sdk, AsyncModelsClient) self._backend_registry = backend_registry self._controller_config = controller_config self._entity_cache = entity_cache @@ -178,6 +188,28 @@ def get_service_backend(self) -> ServiceBackend: """ return self._backend_registry.get_backend() + async def _update_deployment_status( + self, + deployment: ModelDeployment, + *, + status: ModelDeploymentStatus, + status_message: str = "", + model_provider_id: str | None = None, + ) -> ModelDeployment: + """Write a controller-owned deployment status update via the typed Models client.""" + return ( + await self._models_client.update_deployment_status( + name=deployment.name, + workspace=deployment.workspace, + body=UpdateModelDeploymentStatusRequest( + status=status, + status_message=status_message, + model_provider_id=model_provider_id, + ), + query_params={"version": str(deployment.entity_version)}, + ) + ).data() + async def reconcile_deployments(self, deployment_contexts: list[ModelContext]) -> None: """Process deployments and reconcile their state with backends. @@ -189,12 +221,16 @@ async def reconcile_deployments(self, deployment_contexts: list[ModelContext]) - """ for ctx in deployment_contexts: deployment = ctx.model_deployment + if deployment is None: + logger.warning("Skipping deployment reconciliation for context with no model_deployment") + self._emit_heartbeat() + continue model_deployment_id = f"{deployment.workspace}/{deployment.name}" try: backend = self.get_service_backend() match deployment.status: - case "CREATED": + case ModelDeploymentStatus.CREATED: # Lambda needed to bind ctx (the reconcile context bundles # the deployment, config, and model entity). await self._reconcile_individual_deployment( @@ -203,22 +239,22 @@ async def reconcile_deployments(self, deployment_contexts: list[ModelContext]) - "create", existing_provider=ctx.model_provider, ) - case "PENDING" | "READY" | "UNKNOWN": + case ModelDeploymentStatus.PENDING | ModelDeploymentStatus.READY | ModelDeploymentStatus.UNKNOWN: # Check status and handle drift/backend issues. The ctx # carries the config + entity so backends that advance # creation in the status path (k8s vLLM) can compile the # serving objects. status_update = await backend.get_model_deployment_status(ctx) - if status_update.status == "LOST": + if status_update.status == ModelDeploymentStatus.LOST: # Drift detected - attempt recovery await self._handle_drift_recovery(deployment, ctx, backend) continue - elif status_update.status == "UNKNOWN": + elif status_update.status == ModelDeploymentStatus.UNKNOWN: # Backend communication failure - track attempts, eventually error out await self._handle_unknown_status(deployment, status_update) continue - elif status_update.status in ("READY", "ERROR"): + elif status_update.status in (ModelDeploymentStatus.READY, ModelDeploymentStatus.ERROR): # Clear recovery state - deployment is healthy or in terminal state self._drift_recovery_cache.remove(model_deployment_id) @@ -226,7 +262,7 @@ async def reconcile_deployments(self, deployment_contexts: list[ModelContext]) - # fetched above, so ``_reconcile_individual_deployment`` # won't invoke this callable; it's passed only for the # generic signature (bind ctx for type consistency). - action = "check status of" if deployment.status == "PENDING" else "monitor" + action = "check status of" if deployment.status == ModelDeploymentStatus.PENDING else "monitor" await self._reconcile_individual_deployment( deployment, lambda _dep, _ctx=ctx: backend.get_model_deployment_status(_ctx), @@ -234,7 +270,7 @@ async def reconcile_deployments(self, deployment_contexts: list[ModelContext]) - existing_provider=ctx.model_provider, status_update=status_update, ) - case "DELETING": + case ModelDeploymentStatus.DELETING: await self._reconcile_individual_deployment( deployment, lambda d: backend.delete_model_deployment( @@ -245,7 +281,7 @@ async def reconcile_deployments(self, deployment_contexts: list[ModelContext]) - "delete", existing_provider=ctx.model_provider, ) - case "DELETED": + case ModelDeploymentStatus.DELETED: # Check if deployment has been in DELETED state long enough to hard-delete await self._handle_deleted_deployment(deployment) except Exception as e: @@ -374,11 +410,9 @@ async def gc_error_deployments(self, error_deployments: list[ModelDeployment]) - if original_message: gc_message = f"{gc_message} Original error: {original_message}" - await self._models_sdk.inference.deployments.update_status( - name=deployment.name, - workspace=deployment.workspace, - status="DELETING", - version=deployment.entity_version, + await self._update_deployment_status( + deployment, + status=ModelDeploymentStatus.DELETING, status_message=gc_message, ) @@ -431,7 +465,7 @@ async def _reconcile_individual_deployment( ) if ( action_description == "monitor" - and status_update.status == "READY" + and status_update.status == ModelDeploymentStatus.READY and not (status_update.status_message or "").strip() ): # DEBUG level for routine READY monitoring with nothing to report. @@ -442,11 +476,9 @@ async def _reconcile_individual_deployment( model_provider_id = await self._reconcile_model_provider(deployment, status_update, existing_provider) - await self._models_sdk.inference.deployments.update_status( - name=deployment.name, - workspace=deployment.workspace, + await self._update_deployment_status( + deployment, status=status_update.status, - version=deployment.entity_version, status_message=status_update.status_message, model_provider_id=model_provider_id, ) @@ -457,11 +489,9 @@ async def _reconcile_individual_deployment( except Exception as e: logger.exception(f"Failed to {action_description} deployment {model_deployment_id}: {e}") try: - await self._models_sdk.inference.deployments.update_status( - name=deployment.name, - workspace=deployment.workspace, - status="ERROR", - version=deployment.entity_version, + await self._update_deployment_status( + deployment, + status=ModelDeploymentStatus.ERROR, status_message=f"Failed to {action_description} deployment {model_deployment_id}", ) except ConflictError as e: @@ -499,11 +529,9 @@ async def _handle_drift_recovery( attempts = cache.get_attempts(model_deployment_id) logger.error(f"Drift recovery failed for {model_deployment_id} after {attempts} attempts") try: - await self._models_sdk.inference.deployments.update_status( - name=deployment.name, - workspace=deployment.workspace, - status="ERROR", - version=deployment.entity_version, + await self._update_deployment_status( + deployment, + status=ModelDeploymentStatus.ERROR, status_message=( f"Drift recovery failed after {attempts} attempts. " f"Backend resources could not be recreated. Manual intervention required." @@ -538,11 +566,9 @@ async def _handle_drift_recovery( f"{status_update.status_message}" ) - await self._models_sdk.inference.deployments.update_status( - name=deployment.name, - workspace=deployment.workspace, + await self._update_deployment_status( + deployment, status=status_update.status, - version=deployment.entity_version, status_message=recovery_message, model_provider_id=None, # Provider will be recreated when READY ) @@ -558,11 +584,9 @@ async def _handle_drift_recovery( # Update status to PENDING with error info for visibility, but don't set ERROR # The next cycle will retry (respecting backoff) and can detect if recovery succeeded try: - await self._models_sdk.inference.deployments.update_status( - name=deployment.name, - workspace=deployment.workspace, - status="PENDING", - version=deployment.entity_version, + await self._update_deployment_status( + deployment, + status=ModelDeploymentStatus.PENDING, status_message=f"Recovery attempt {attempt_count}/{max_attempts} failed: {e}. Will retry.", ) except Exception: @@ -596,11 +620,9 @@ async def _handle_unknown_status( attempts = cache.get_attempts(model_deployment_id) logger.error(f"Backend communication failed for {model_deployment_id} after {attempts} attempts") try: - await self._models_sdk.inference.deployments.update_status( - name=deployment.name, - workspace=deployment.workspace, - status="ERROR", - version=deployment.entity_version, + await self._update_deployment_status( + deployment, + status=ModelDeploymentStatus.ERROR, status_message=( f"Unable to communicate with backend after {attempts} attempts. " f"Last error: {status_update.status_message}. Manual intervention required." @@ -628,11 +650,9 @@ async def _handle_unknown_status( ) try: - await self._models_sdk.inference.deployments.update_status( - name=deployment.name, - workspace=deployment.workspace, - status="UNKNOWN", - version=deployment.entity_version, + await self._update_deployment_status( + deployment, + status=ModelDeploymentStatus.UNKNOWN, status_message=( f"Unable to determine deployment status (attempt {attempt_count}/{max_attempts}). " f"{status_update.status_message}" @@ -663,9 +683,9 @@ async def _reconcile_model_provider( model_deployment_id = f"{deployment.workspace}/{deployment.name}" try: - if status_update.status == "READY": + if status_update.status == ModelDeploymentStatus.READY: return await self._ensure_model_provider(deployment, status_update.host_url, existing_provider) - elif status_update.status in ("DELETING", "DELETED"): + elif status_update.status in (ModelDeploymentStatus.DELETING, ModelDeploymentStatus.DELETED): await self._delete_model_provider(deployment) return None except Exception as e: @@ -704,23 +724,27 @@ async def _ensure_model_provider( provider_workspace, provider_name = _provider_ref.workspace, _provider_ref.name if not existing_provider: - existing_provider = await self._models_sdk.inference.providers.retrieve( - name=provider_name, - workspace=provider_workspace, - ) + existing_provider = ( + await self._models_client.get_provider( + name=provider_name, + workspace=provider_workspace, + ) + ).data() if existing_provider.host_url != host_url: logger.info( f"ModelProvider {deployment.model_provider_id} host_url changed from " f"{existing_provider.host_url} to {host_url}, updating provider" ) - await self._models_sdk.inference.providers.update( + await self._models_client.upsert_provider( name=provider_name, workspace=provider_workspace, - host_url=host_url, - description=existing_provider.description, - enabled_models=existing_provider.enabled_models, - status="READY", + body=UpsertModelProviderRequest( + host_url=host_url, + description=existing_provider.description, + enabled_models=existing_provider.enabled_models, + status=ModelProviderStatus.READY, + ), ) else: logger.debug( @@ -743,10 +767,7 @@ async def _ensure_model_provider( provider_workspace = deployment.workspace try: - await self._models_sdk.inference.providers.retrieve( - name=provider_name, - workspace=provider_workspace, - ) + await self._models_client.get_provider(name=provider_name, workspace=provider_workspace) unique_suffix = uuid.uuid4().hex[:8] provider_name = f"{deployment.name}_{unique_suffix}" logger.info( @@ -756,14 +777,16 @@ async def _ensure_model_provider( except NotFoundError: logger.debug(f"Creating ModelProvider {provider_workspace}/{provider_name} for deployment") - await self._models_sdk.inference.providers.create( + await self._models_client.create_provider( workspace=provider_workspace, - name=provider_name, - host_url=host_url, - description=f"Auto-created provider for deployment {deployment.name}", - project=deployment.project, - model_deployment_id=model_deployment_id, - status="READY", + body=CreateModelProviderRequest( + name=provider_name, + host_url=host_url, + description=f"Auto-created provider for deployment {deployment.name}", + project=deployment.project, + model_deployment_id=model_deployment_id, + status=ModelProviderStatus.READY, + ), ) model_provider_id = f"{provider_workspace}/{provider_name}" @@ -783,10 +806,7 @@ async def _cleanup_model_entities_for_provider( """ try: # Get the provider to see what models it was serving - provider = await self._models_sdk.inference.providers.retrieve( - name=provider_name, - workspace=provider_workspace, - ) + provider = (await self._models_client.get_provider(name=provider_name, workspace=provider_workspace)).data() if not provider.served_models: logger.debug(f"Provider {provider_id} has no served_models, no cleanup needed") @@ -868,7 +888,7 @@ async def _delete_model_provider(self, deployment: ModelDeployment) -> None: try: logger.info(f"Deleting ModelProvider {model_provider_id} for deployment {model_deployment_id}") - await self._models_sdk.inference.providers.delete( + await self._models_client.delete_provider( name=provider_name, workspace=provider_workspace, ) @@ -903,7 +923,7 @@ async def _handle_deleted_deployment(self, deployment: ModelDeployment) -> None: ) try: # Hard-delete this specific version by calling the delete API again on a DELETED deployment - await self._models_sdk.inference.deployments.versions.delete( + await self._models_client.delete_deployment_version( name=str(deployment.entity_version), # version number workspace=deployment.workspace, # workspace deployment=deployment.name, # deployment name 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 4c09e819f8..c433596136 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 @@ -2,18 +2,16 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio +import json import threading from logging import getLogger from typing import Optional -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 +import httpx 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 nemo_platform_plugin.models.types import ModelDeployment, ModelDeploymentConfig, ModelDeploymentStatus, 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 @@ -28,12 +26,15 @@ logger = getLogger(__name__) +_CONTROLLER_HTTP_TIMEOUT = httpx.Timeout(timeout=60, connect=5.0) +_CONTROLLER_HTTP_LIMITS = httpx.Limits(max_connections=100, max_keepalive_connections=20) + NON_TERMINAL_STATES: list[ModelDeploymentStatus] = [ - "CREATED", - "PENDING", - "READY", - "DELETING", - "DELETED", # Poll DELETED deployments to clean them up after grace period + ModelDeploymentStatus.CREATED, + ModelDeploymentStatus.PENDING, + ModelDeploymentStatus.READY, + ModelDeploymentStatus.DELETING, + ModelDeploymentStatus.DELETED, # Poll DELETED deployments to clean them up after grace period ] @@ -64,8 +65,13 @@ def __init__( self._models_sdk = get_async_platform_sdk( as_service="models", internal=True, - http_client=DefaultAsyncHttpxClient(), + http_client=httpx.AsyncClient( + timeout=_CONTROLLER_HTTP_TIMEOUT, + limits=_CONTROLLER_HTTP_LIMITS, + follow_redirects=True, + ), ) + self._models_client = client_from_platform(self._models_sdk, AsyncModelsClient) self._service_backends = backend_registry.list_backends() # Shared by both reconcilers; re-read at the start of each phase that @@ -126,7 +132,7 @@ def get_service_backend(self) -> ServiceBackend: return self._backend_registry.get_backend() async def _retrieve_deployment_config( - self, config_ref: str, config_version: str, deployment_workspace: str + self, config_ref: str, config_version: int | str, deployment_workspace: str ) -> ModelDeploymentConfig: """Retrieve the ModelDeploymentConfig from the API. @@ -143,11 +149,13 @@ async def _retrieve_deployment_config( workspace, name = ref.workspace, ref.name logger.debug(f"Fetching ModelDeploymentConfig {workspace}/{name}@{config_version}") - config = await self._models_sdk.inference.deployment_configs.versions.retrieve( - name=str(config_version), # version number - workspace=workspace, # workspace - config=name, # config name - ) + config = ( + await self._models_client.get_deployment_config_version( + name=str(config_version), + workspace=workspace, + config=name, + ) + ).data() return config except Exception as e: logger.error(f"Failed to fetch ModelDeploymentConfig {config_ref}@{config_version}: {e}") @@ -238,9 +246,8 @@ 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. - models = client_from_platform(self._models_sdk, AsyncModelsClient) model_entity = ( - await models.get_model( + await self._models_client.get_model( name=full_model_name, workspace=workspace, ) @@ -289,16 +296,18 @@ async def retrieve_non_terminal_deployments(self) -> list[ModelContext]: try: logger.debug(f"Querying ModelDeployments with status: {status} across all workspaces") # SDK returns AsyncPaginator - iterate through all pages - resp = self._models_sdk.inference.deployments.list( + resp = await self._models_client.list_deployments( workspace="-", # Cross-workspace query - filter={"status": status}, - all_versions=True, - page_size=1000, + query_params={ + "filter": json.dumps({"status": status.value}), + "all_versions": True, + "page_size": 1000, + }, ) logger.debug(f"Got paginator response for status {status}, iterating...") # Collect all deployments from paginator - deployments = [deployment async for deployment in resp] + deployments = [deployment async for deployment in resp.items()] logger.debug(f"Iteration complete for status {status}, got {len(deployments)} deployment(s)") if deployments: @@ -328,10 +337,12 @@ async def retrieve_non_terminal_deployments(self) -> list[ModelContext]: try: _prov_ref = parse_entity_ref(deployment.model_provider_id) provider_workspace, provider_name = _prov_ref.workspace, _prov_ref.name - provider = await self._models_sdk.inference.providers.retrieve( - name=provider_name, - workspace=provider_workspace, - ) + provider = ( + await self._models_client.get_provider( + name=provider_name, + workspace=provider_workspace, + ) + ).data() except Exception as e: logger.warning( f"Failed to fetch provider for deployment {deployment.workspace}/{deployment.name}: {e}" @@ -369,11 +380,11 @@ async def retrieve_model_providers(self) -> list[ModelContext] | None: provider_contexts: list[ModelContext] = [] try: - providers = self._models_sdk.inference.providers.list( + providers = await self._models_client.list_providers( workspace="-", # Cross-workspace query ) - async for provider in providers: + async for provider in providers.items(): deployment = None config = None entity = None @@ -383,10 +394,12 @@ async def retrieve_model_providers(self) -> list[ModelContext] | None: try: _depl_ref = parse_entity_ref(provider.model_deployment_id) deployment_workspace, deployment_name = _depl_ref.workspace, _depl_ref.name - deployment = await self._models_sdk.inference.deployments.retrieve( - deployment_name, - workspace=deployment_workspace, - ) + deployment = ( + await self._models_client.get_deployment( + name=deployment_name, + workspace=deployment_workspace, + ) + ).data() # Fetch config if deployment has config reference if deployment and deployment.config and deployment.config_version: @@ -429,13 +442,15 @@ async def retrieve_error_deployments(self) -> list[ModelDeployment]: since GC only needs the deployment itself and its timestamps. """ try: - resp = self._models_sdk.inference.deployments.list( + resp = await self._models_client.list_deployments( workspace="-", - filter={"status": "ERROR"}, - all_versions=True, - page_size=1000, + query_params={ + "filter": json.dumps({"status": ModelDeploymentStatus.ERROR.value}), + "all_versions": True, + "page_size": 1000, + }, ) - return [deployment async for deployment in resp] + return [deployment async for deployment in resp.items()] except Exception: logger.warning("Error querying ERROR deployments for GC", exc_info=True) return [] @@ -472,7 +487,9 @@ async def async_controller_step(self) -> None: await self._deployment_reconciler.reconcile_deployments(deployment_contexts) known_deployment_ids = { - f"{ctx.model_deployment.workspace}/{ctx.model_deployment.name}" for ctx in deployment_contexts + f"{ctx.model_deployment.workspace}/{ctx.model_deployment.name}" + for ctx in deployment_contexts + if ctx.model_deployment is not None } await self._deployment_reconciler.reconcile_orphans(known_deployment_ids) self.emit_heartbeat() 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 8c0826892e..67ab2aa4d1 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 @@ -5,19 +5,28 @@ import json import re +from collections.abc import Mapping from dataclasses import dataclass, field from datetime import datetime, timezone from logging import getLogger from typing import Callable, TypedDict from nemo_platform import AsyncNeMoPlatform -from nemo_platform._exceptions import APIStatusError, ConflictError, NotFoundError -from nemo_platform.types.inference import ServedModelMapping -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.inference.virtual_model import VirtualModel -from nemo_platform_plugin.models.types import ModelEntity +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ConflictError, NemoHTTPError, NotFoundError +from nemo_platform_plugin.inference_gateway.client import AsyncInferenceGatewayProviderClient +from nemo_platform_plugin.models.client import AsyncModelsClient +from nemo_platform_plugin.models.types import ( + ModelDeployment, + ModelDeploymentConfig, + ModelEntity, + ModelProvider, + ModelProviderStatus, + ServedModelMapping, + UpdateModelProviderStatusRequest, +) +from nemo_platform_plugin.virtual_models.client import AsyncVirtualModelsClient +from nemo_platform_plugin.virtual_models.types import CreateVirtualModelRequest, VirtualModel 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 @@ -30,7 +39,7 @@ from nmp.core.models.config import ControllerConfig from nmp.core.models.controllers.context import ModelContext from nmp.core.models.controllers.entity_cache import ModelEntityCache -from nmp.core.models.schemas import BackendFormat, ModelProviderStatus +from nmp.core.models.schemas import BackendFormat logger = getLogger(__name__) @@ -60,12 +69,12 @@ def _infer_backend_format(model_name: str) -> str: def _has_backend_format(model_entity: ModelEntity) -> bool: - value = getattr(model_entity, "backend_format", None) + value = model_entity.backend_format return isinstance(value, str) and bool(value) -def _get_virtual_model_db_version(virtual_model: object) -> int | None: - db_version = getattr(virtual_model, "db_version", None) +def _get_virtual_model_db_version(virtual_model: VirtualModel) -> int | None: + db_version = virtual_model.db_version if isinstance(db_version, bool): return None if isinstance(db_version, int): @@ -73,15 +82,8 @@ def _get_virtual_model_db_version(virtual_model: object) -> int | None: return None -def _get_datetime_attr(obj: object, attr_name: str) -> datetime | None: - value = getattr(obj, attr_name, None) - if isinstance(value, datetime): - return ensure_utc(value) - return None - - -def _get_virtual_model_observed_at(virtual_model: object) -> datetime | None: - return _get_datetime_attr(virtual_model, "updated_at") or _get_datetime_attr(virtual_model, "created_at") +def _get_virtual_model_observed_at(virtual_model: VirtualModel) -> datetime | None: + return ensure_utc(virtual_model.updated_at) or ensure_utc(virtual_model.created_at) # --------------------------------------------------------------------------- @@ -269,7 +271,7 @@ def _resolve_base_backend_model_id( ``model_spec.model_namespace`` / ``model_name`` / ``model_revision`` (same as config create when ``model_entity_id`` is inferred from the entity store). """ - if config and getattr(config, "model_entity_id", None): + if config and config.model_entity_id: model_workspace, model_name, _revision = parse_model_name_revision(model_name=config.model_entity_id) if model_workspace and model_name: return f"{model_workspace}/{model_name}" @@ -278,12 +280,12 @@ def _resolve_base_backend_model_id( if base_model_entity is not None: return f"{base_model_entity.workspace}/{base_model_entity.name}" - model_spec = getattr(config, "model_spec", None) if config else None + model_spec = config.model_spec if config else None if model_spec is not None: model_workspace, model_name, _revision = parse_model_name_revision( - model_namespace=getattr(model_spec, "model_namespace", None), - model_name=getattr(model_spec, "model_name", None), - model_revision=getattr(model_spec, "model_revision", None), + model_namespace=model_spec.model_namespace, + model_name=model_spec.model_name, + model_revision=model_spec.model_revision, ) if model_workspace and model_name: return f"{model_workspace}/{model_name}" @@ -324,9 +326,12 @@ def __init__( self._controller_config = controller_config self._entity_cache = entity_cache self._emit_heartbeat = emit_heartbeat - self._discovery_sdk = models_sdk.with_options( + self._models_client = client_from_platform(models_sdk, AsyncModelsClient) + self._virtual_models_client = client_from_platform(models_sdk, AsyncVirtualModelsClient) + discovery_sdk = models_sdk.with_options( max_retries=controller_config.provider_discovery_max_retries, ) + self._gateway_provider_client = client_from_platform(discovery_sdk, AsyncInferenceGatewayProviderClient) # ------------------------------------------------------------------------- # Public entry point @@ -400,9 +405,11 @@ async def _load_virtual_models(self) -> tuple[list[VirtualModel], set[tuple[str, """ vm_snapshot: list[VirtualModel] = [] try: - async for virtual_model in self._models_sdk.inference.virtual_models.list( - workspace="-", page_size=_VIRTUAL_MODEL_PAGE_SIZE - ): + response = await self._virtual_models_client.list_virtual_models( + workspace="-", + query_params={"page_size": _VIRTUAL_MODEL_PAGE_SIZE}, + ) + async for virtual_model in response.items(): vm_snapshot.append(virtual_model) self._emit_heartbeat() except Exception: @@ -454,13 +461,17 @@ async def _reconcile_single_provider( ) ctx.served_models = [] try: - ctx.model_provider = await self._models_sdk.inference.providers.update_status( - name=provider.name, - workspace=provider.workspace, - served_models=[], - status="READY", - status_message="Non-OpenAI compliant endpoint, model entity routing disabled", - ) + ctx.model_provider = ( + await self._models_client.update_provider_status( + name=provider.name, + workspace=provider.workspace, + body=UpdateModelProviderStatusRequest( + served_models=[], + status=ModelProviderStatus.READY, + status_message="Non-OpenAI compliant endpoint, model entity routing disabled", + ), + ) + ).data() except Exception as e: logger.error(f"Failed to update provider {provider_id} status: {e}") return @@ -507,12 +518,16 @@ async def _reconcile_single_provider( logger.debug(f"Provider {provider_id}: serving {len(served_models)} model(s)") try: - ctx.model_provider = await self._models_sdk.inference.providers.update_status( - name=provider.name, - workspace=provider.workspace, - served_models=served_models, - status="READY", - ) + ctx.model_provider = ( + await self._models_client.update_provider_status( + name=provider.name, + workspace=provider.workspace, + body=UpdateModelProviderStatusRequest( + served_models=served_models, + status=ModelProviderStatus.READY, + ), + ) + ).data() if served_models: logger.debug(f"Updated provider {provider_id} with {len(served_models)} served model(s)") else: @@ -575,12 +590,16 @@ def _is_past_lost_threshold(self, provider: ModelProvider, provider_id: str, now async def _mark_lost(self, ctx: ModelContext, provider: ModelProvider, provider_id: str) -> None: """Write LOST status for a provider that has permanently failed discovery.""" try: - ctx.model_provider = await self._models_sdk.inference.providers.update_status( - name=provider.name, - workspace=provider.workspace, - status="LOST", - status_message="Provider discovery permanently failed. Delete and recreate to retry.", - ) + ctx.model_provider = ( + await self._models_client.update_provider_status( + name=provider.name, + workspace=provider.workspace, + body=UpdateModelProviderStatusRequest( + status=ModelProviderStatus.LOST, + status_message="Provider discovery permanently failed. Delete and recreate to retry.", + ), + ) + ).data() except Exception: logger.exception( "Failed to transition provider to LOST", @@ -611,13 +630,15 @@ async def _on_transient_failure( updated_at = ensure_utc(provider.updated_at) if updated_at and (now - updated_at).total_seconds() > PROVIDER_ERROR_THRESHOLD_SECONDS: try: - await self._models_sdk.inference.providers.update_status( + await self._models_client.update_provider_status( name=provider.name, workspace=provider.workspace, - status="ERROR", - status_message=f"Provider discovery failed: {err.message}" - if err.message - else "Provider discovery failed: unable to reach GET /v1/models", + body=UpdateModelProviderStatusRequest( + status=ModelProviderStatus.ERROR, + status_message=f"Provider discovery failed: {err.message}" + if err.message + else "Provider discovery failed: unable to reach GET /v1/models", + ), ) logger.warning( "Provider escalated to ERROR after persistent discovery failures", @@ -639,13 +660,15 @@ async def _on_transient_failure( elif provider.status == ModelProviderStatus.ERROR: # Bump updated_at to pace the next retry try: - await self._models_sdk.inference.providers.update_status( + await self._models_client.update_provider_status( name=provider.name, workspace=provider.workspace, - status="ERROR", - status_message=f"Discovery retry failed: {err.message}" - if err.message - else "Discovery retry failed: still unable to reach GET /v1/models", + body=UpdateModelProviderStatusRequest( + status=ModelProviderStatus.ERROR, + status_message=f"Discovery retry failed: {err.message}" + if err.message + else "Discovery retry failed: still unable to reach GET /v1/models", + ), ) except Exception: logger.exception( @@ -688,8 +711,7 @@ async def _discover_models(self, provider: ModelProvider) -> DiscoveryResult: # This intentionally uses the controller's service principal to perform # infrastructure reconciliation. User-level secret access remains guarded # at provider create/upsert validation and by IGW when proxying requests. - models_response = await self._discovery_sdk.inference.gateway.provider.get( - "v1/models", + models_response = await self._gateway_provider_client.get_provider_models( workspace=provider.workspace, name=provider.name, timeout=self._controller_config.provider_discovery_timeout_seconds, @@ -704,39 +726,43 @@ async def _discover_models(self, provider: ModelProvider) -> DiscoveryResult: logger.warning(f"Non-OpenAI compliant response format from {provider_id}") return DiscoveryNonCompliant() - if not isinstance(models_response, dict) or "data" not in models_response: + if not isinstance(models_response, Mapping) or "data" not in models_response: logger.warning(f"Non-OpenAI compliant response format from {provider_id}") return DiscoveryNonCompliant() - discovered_models = models_response["data"] + discovered_models = models_response.get("data") if not isinstance(discovered_models, list): logger.warning(f"Non-OpenAI compliant data field from {provider_id}") return DiscoveryNonCompliant() - models = [] + models: list[DiscoveredModel] = [] for model in discovered_models: - if not isinstance(model, dict) or not isinstance(model.get("id"), str): + if not isinstance(model, Mapping): logger.warning(f"Skipping invalid model entry in {provider_id}: {model}") continue - models.append( - { - "id": model["id"], - "root": model.get("root"), - "parent": model.get("parent"), - } - ) + model_id = model.get("id") + if not isinstance(model_id, str): + logger.warning(f"Skipping invalid model entry in {provider_id}: {model}") + continue + root = model.get("root") + parent = model.get("parent") + discovered_model: DiscoveredModel = { + "id": model_id, + "root": root if isinstance(root, str) else None, + "parent": parent if isinstance(parent, str) else None, + } + models.append(discovered_model) return DiscoverySuccess(models) - except APIStatusError as e: + except NemoHTTPError as e: # 404 from the provider proxy is only returned when the provider is not in the gateway # cache yet (single code path in IGW). Preserve served_models. if e.status_code == 404: logger.warning(f"Provider {provider_id} not yet in gateway cache (404), preserving served_models") return DiscoveryTransientError("Provider not yet in gateway cache (404)") # IGW (FastAPI) returns 502 with body {"detail": "Backend returned 404: ..."} when backend has no /v1/models. - detail = str((e.body or {}).get("detail", "")) if isinstance(e.body, dict) else "" - if e.status_code == 502 and _GATEWAY_BACKEND_404_DETAIL in detail: + if e.status_code == 502 and _GATEWAY_BACKEND_404_DETAIL in e.detail: # Backend (NIM) returned 404 — no GET /v1/models or similar. Mark non-compliant. logger.info( f"Backend for {provider_id} returned 404 for GET /v1/models, disabling model entity routing" @@ -809,6 +835,7 @@ async def _ensure_external_entities( await self._ensure_model_entity_for_provider( model_workspace=provider.workspace, model_name=normalized, + provider=provider, provider_id=provider_id, ctx=ctx, ) @@ -981,7 +1008,12 @@ def _generate_deployment_served_model_mappings( # ------------------------------------------------------------------------- async def _ensure_model_entity_for_provider( - self, model_workspace: str, model_name: str, provider_id: str, ctx: ModelContext + self, + model_workspace: str, + model_name: str, + provider: ModelProvider, + provider_id: str, + ctx: ModelContext, ) -> None: """Ensure a Model Entity exists for an autodiscovered model and link it to the provider. @@ -1007,7 +1039,7 @@ async def _ensure_model_entity_for_provider( details = await self._build_artifact_details( model_name, provider_id, - ctx.model_provider, + provider, existing_model_entity, ctx.model_deployment, ctx.model_deployment_config, @@ -1039,7 +1071,7 @@ async def _ensure_model_entity_for_provider( # Only fill in what the entity is missing, so a value a user corrected is # never overwritten. ``backend_format`` treats None as missing so entities # registered before it existed get backfilled. - updates: dict = {} + updates: dict[str, object] = {} if fileset and not existing_model_entity.fileset: updates["fileset"] = fileset if details.api_endpoint and not existing_model_entity.api_endpoint: @@ -1063,7 +1095,7 @@ async def _ensure_passthrough_virtual_model( updated on creation so the same name is not attempted twice in one pass. A name held by a user-managed VirtualModel is left as it is. - This is idempotent: a :class:`~nemo_platform.ConflictError` (409) means + This is idempotent: a :class:`~nemo_platform_plugin.client.errors.ConflictError` (409) means the VirtualModel already exists and is silently ignored. Any other exception is logged as a warning and does not propagate — VirtualModel creation failures must not block provider reconciliation. @@ -1088,11 +1120,13 @@ async def _ensure_passthrough_virtual_model( return try: - await self._models_sdk.inference.virtual_models.create( + await self._virtual_models_client.create_virtual_model( workspace=workspace, - name=model_name, - default_model_entity=f"{workspace}/{model_name}", - autoprovisioned=True, + body=CreateVirtualModelRequest( + name=model_name, + default_model_entity=f"{workspace}/{model_name}", + autoprovisioned=True, + ), ) logger.info( "Auto-created passthrough VirtualModel %s/%s", @@ -1169,10 +1203,10 @@ async def _cleanup_orphaned_virtual_models( continue try: - await self._models_sdk.inference.virtual_models.delete( + await self._virtual_models_client.delete_virtual_model( name=virtual_model.name, workspace=virtual_model.workspace, - expected_db_version=expected_db_version, + query_params={"expected_db_version": expected_db_version}, ) logger.info( "Deleted orphaned autoprovisioned VirtualModel %s/%s", @@ -1237,7 +1271,7 @@ async def _build_artifact_details( logger.debug(f"Built api_endpoint for external provider: {provider.host_url}") elif weights_type == ModelWeightsType.HUGGINGFACE and config: - model_spec = getattr(config, "model_spec", None) + model_spec = config.model_spec if model_spec is None: logger.warning("Missing model_spec for HuggingFace weights; skipping fileset_url build") return details @@ -1259,7 +1293,7 @@ async def _build_artifact_details( elif weights_type == ModelWeightsType.FILES_SERVICE and config: # Files service models (including SFT) use hf:// prefix since Files service exposes # models via HuggingFace-compatible API - model_spec = getattr(config, "model_spec", None) + model_spec = config.model_spec if model_spec is None: logger.warning("Missing model_spec for Files service weights; skipping fileset_url build") return details diff --git a/services/core/models/tests/integration/conftest.py b/services/core/models/tests/integration/conftest.py index ccd6564c50..ab6594b77e 100644 --- a/services/core/models/tests/integration/conftest.py +++ b/services/core/models/tests/integration/conftest.py @@ -6,16 +6,31 @@ from __future__ import annotations from collections.abc import Callable -from typing import Any, Generator, Optional +from typing import Any, Generator, Optional, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest from nemo_deployments_plugin.config import ControllerConfig, DeploymentsConfig, ExecutorConfigEntry from nemo_deployments_plugin.controller import DeploymentsController from nemo_platform import AsyncNeMoPlatform, NeMoPlatform -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.models.client import ModelsClient +from nemo_platform_plugin.models.types import ( + ContainerExecutorConfig, + CreateModelDeploymentConfigRequest, + CreateModelDeploymentRequest, + CreateModelProviderRequest, + Engine, + ModelDeployment, + ModelDeploymentConfig, + ModelDeploymentConfigModelSpec, + ModelDeploymentStatus, + ModelEntity, + ModelProvider, + UpdateModelDeploymentConfigRequest, + UpdateModelDeploymentRequest, + UpsertModelProviderRequest, +) from nmp.common.config import Runtime from nmp.common.secrets.encryption import get_base64_encoded_random_bytes from nmp.core.files.app.backends.base import FileInfo @@ -106,6 +121,166 @@ def secrets_service_config() -> SecretsServiceConfig: ) +def models_client_from_sdk(sdk: NeMoPlatform) -> ModelsClient: + """Create a typed Models client sharing the test platform transport.""" + return client_from_platform(sdk, ModelsClient) + + +def create_provider( + client: ModelsClient, + *, + workspace: str = DEFAULT_WORKSPACE, + name: str, + host_url: str, + **kwargs: Any, +) -> ModelProvider: + return client.create_provider( + workspace=workspace, + body=CreateModelProviderRequest(name=name, host_url=host_url, **kwargs), + ).data() + + +def upsert_provider( + client: ModelsClient, + *, + workspace: str = DEFAULT_WORKSPACE, + name: str, + host_url: str, + **kwargs: Any, +) -> ModelProvider: + return client.upsert_provider( + workspace=workspace, + name=name, + body=UpsertModelProviderRequest(host_url=host_url, **kwargs), + ).data() + + +def get_provider(client: ModelsClient, *, workspace: str = DEFAULT_WORKSPACE, name: str) -> ModelProvider: + return client.get_provider(workspace=workspace, name=name).data() + + +def list_providers(client: ModelsClient, *, workspace: str = DEFAULT_WORKSPACE) -> list[ModelProvider]: + return list(client.list_providers(workspace=workspace).items()) + + +def delete_provider(client: ModelsClient, *, workspace: str = DEFAULT_WORKSPACE, name: str) -> None: + client.delete_provider(workspace=workspace, name=name).data() + + +def _engine(value: Engine | str) -> Engine: + return value if isinstance(value, Engine) else Engine(value) + + +def _model_spec( + value: ModelDeploymentConfigModelSpec | dict[str, Any] | None, +) -> ModelDeploymentConfigModelSpec: + if isinstance(value, ModelDeploymentConfigModelSpec): + return value + return ModelDeploymentConfigModelSpec.model_validate(value or {}) + + +def _executor_config( + value: ContainerExecutorConfig | dict[str, Any] | None, +) -> ContainerExecutorConfig: + if isinstance(value, ContainerExecutorConfig): + return value + return ContainerExecutorConfig.model_validate(value or {"gpu": 0}) + + +def create_deployment_config( + client: ModelsClient, + *, + workspace: str = DEFAULT_WORKSPACE, + name: str, + engine: Engine | str, + model_spec: ModelDeploymentConfigModelSpec | dict[str, Any] | None = None, + executor_config: ContainerExecutorConfig | dict[str, Any] | None = None, + model_entity_id: str | None = None, +) -> ModelDeploymentConfig: + body = CreateModelDeploymentConfigRequest( + name=name, + engine=_engine(engine), + model_spec=_model_spec(model_spec), + executor_config=_executor_config(executor_config), + model_entity_id=model_entity_id, + ) + return client.create_deployment_config(workspace=workspace, body=body).data() + + +def update_deployment_config( + client: ModelsClient, + *, + workspace: str = DEFAULT_WORKSPACE, + name: str, + engine: Engine | str, + model_spec: ModelDeploymentConfigModelSpec | dict[str, Any], + executor_config: ContainerExecutorConfig | dict[str, Any], + model_entity_id: str | None = None, +) -> ModelDeploymentConfig: + body = UpdateModelDeploymentConfigRequest( + engine=_engine(engine), + model_spec=_model_spec(model_spec), + executor_config=_executor_config(executor_config), + model_entity_id=model_entity_id, + ) + return client.update_deployment_config(workspace=workspace, name=name, body=body).data() + + +def get_deployment_config( + client: ModelsClient, *, workspace: str = DEFAULT_WORKSPACE, name: str +) -> ModelDeploymentConfig: + return client.get_deployment_config(workspace=workspace, name=name).data() + + +def list_deployment_configs(client: ModelsClient, *, workspace: str = DEFAULT_WORKSPACE) -> list[ModelDeploymentConfig]: + return list(client.list_deployment_configs(workspace=workspace).items()) + + +def delete_deployment_config(client: ModelsClient, *, workspace: str = DEFAULT_WORKSPACE, name: str) -> None: + client.delete_deployment_config(workspace=workspace, name=name).data() + + +def create_deployment( + client: ModelsClient, + *, + workspace: str = DEFAULT_WORKSPACE, + name: str, + config: str, + config_version: int | None = None, +) -> ModelDeployment: + return client.create_deployment( + workspace=workspace, + body=CreateModelDeploymentRequest(name=name, config=config, config_version=config_version), + ).data() + + +def update_deployment( + client: ModelsClient, + *, + workspace: str = DEFAULT_WORKSPACE, + name: str, + config: str, + config_version: int | None = None, +) -> ModelDeployment: + return client.update_deployment( + workspace=workspace, + name=name, + body=UpdateModelDeploymentRequest(config=config, config_version=config_version), + ).data() + + +def get_deployment(client: ModelsClient, *, workspace: str = DEFAULT_WORKSPACE, name: str) -> ModelDeployment: + return client.get_deployment(workspace=workspace, name=name).data() + + +def list_deployments(client: ModelsClient, *, workspace: str = DEFAULT_WORKSPACE) -> list[ModelDeployment]: + return list(client.list_deployments(workspace=workspace).items()) + + +def delete_deployment(client: ModelsClient, *, workspace: str = DEFAULT_WORKSPACE, name: str) -> None: + client.delete_deployment(workspace=workspace, name=name).data() + + # ============================================================================= # Mock Backend for Backend-Agnostic Tests # ============================================================================= @@ -133,7 +308,7 @@ def __init__( # Note: host_url is None by default to avoid port conflicts in parallel tests. # Tests that need a specific host_url should set it explicitly in status_responses. self.create_response = DeploymentStatusUpdate( - status="PENDING", + status=ModelDeploymentStatus.PENDING, status_message="Container created and starting", host_url=None, ) @@ -141,12 +316,12 @@ def __init__( # If a deployment name is not in this dict, falls back to default_status_response self.status_responses: dict[str, DeploymentStatusUpdate] = {} self.default_status_response = DeploymentStatusUpdate( - status="READY", + status=ModelDeploymentStatus.READY, status_message="Container is ready", host_url=None, ) self.delete_response = DeploymentStatusUpdate( - status="DELETED", + status=ModelDeploymentStatus.DELETED, status_message="Container deleted", ) @@ -160,12 +335,20 @@ def shutdown(self) -> None: async def create_model_deployment(self, ctx: ModelContext) -> DeploymentStatusUpdate: """Record call and return configured response.""" - self.create_calls.append((ctx.model_deployment, ctx.model_deployment_config, ctx.model_entity)) + deployment = ctx.model_deployment + config = ctx.model_deployment_config + assert deployment is not None + assert config is not None + self.create_calls.append((deployment, config, ctx.model_entity)) return self.create_response async def update_model_deployment(self, ctx: ModelContext) -> DeploymentStatusUpdate: """Record call and return configured response.""" - self.update_calls.append((ctx.model_deployment, ctx.model_deployment_config, ctx.model_entity)) + deployment = ctx.model_deployment + config = ctx.model_deployment_config + assert deployment is not None + assert config is not None + self.update_calls.append((deployment, config, ctx.model_entity)) return self.create_response # Update returns same as create async def get_model_deployment_status(self, ctx: ModelContext) -> DeploymentStatusUpdate: @@ -179,6 +362,7 @@ async def get_model_deployment_status(self, ctx: ModelContext) -> DeploymentStat otherwise falls back to default_status_response. """ deployment = ctx.model_deployment + assert deployment is not None self.status_calls.append(deployment) return self.status_responses.get(deployment.name, self.default_status_response) @@ -232,7 +416,7 @@ def mock_backend_registry(mock_backend: MockServiceBackend) -> BackendRegistry: @pytest.fixture def controller_with_mock_backend( test_clients: ClientContext, mock_backend_registry: BackendRegistry -) -> Generator[tuple[ModelsController, MockServiceBackend, NeMoPlatform], None, None]: +) -> Generator[tuple[ModelsController, MockServiceBackend, ModelsClient], None, None]: """Create a ModelsController wired to use the test SDK and mock backend. Note: The ProviderReconciler's autodiscovery is mocked to avoid issues when @@ -241,9 +425,9 @@ def controller_with_mock_backend( would also iterate over providers from other tests running in the same worker. Yields: - Tuple of (controller, mock_backend, sync_sdk) for testing + Tuple of (controller, mock_backend, models_client) for testing """ - mock_backend = mock_backend_registry.get_backend() + mock_backend = cast(MockServiceBackend, mock_backend_registry.get_backend()) # Create controller with mock backend registry # We need to patch the SDK factory and platform config (used in config and main modules) @@ -267,7 +451,7 @@ def controller_with_mock_backend( # which isn't available in models-only tests. controller._provider_reconciler.reconcile_model_providers = AsyncMock(return_value=None) - yield controller, mock_backend, test_clients.sdk + yield controller, mock_backend, models_client_from_sdk(test_clients.sdk) # Clean up controller resources (event loop, backend registry, etc.) controller.shutdown() @@ -467,7 +651,7 @@ def reconcile_stack(models_controller: ModelsController) -> None: yield ( models_controller, deployments_controller, - test_clients.sdk, + models_client_from_sdk(test_clients.sdk), mock_nim_image, docker_test_context, reconcile_stack, @@ -487,6 +671,6 @@ def reconcile_stack(models_controller: ModelsController) -> None: @pytest.hookimpl(tryfirst=True, hookwrapper=True) def pytest_runtest_makereport(item: pytest.Item, call: pytest.CallInfo[None]) -> Generator[None, None, None]: """Store test results on the item for fixture access.""" - outcome = yield + outcome: Any = yield rep = outcome.get_result() setattr(item, f"rep_{rep.when}", rep) diff --git a/services/core/models/tests/integration/test_deployments_plugin_lifecycle.py b/services/core/models/tests/integration/test_deployments_plugin_lifecycle.py index e0826fbe8e..e17fc5e3b3 100644 --- a/services/core/models/tests/integration/test_deployments_plugin_lifecycle.py +++ b/services/core/models/tests/integration/test_deployments_plugin_lifecycle.py @@ -10,12 +10,20 @@ import pytest from docker.errors import NotFound from nemo_deployments_plugin.backends.labels import container_name -from nemo_platform import NotFoundError +from nemo_platform_plugin.client.errors import NotFoundError from nmp.core.models.controllers.backends.deployments_plugin.naming import entity_names from tenacity import retry, stop_after_delay, wait_fixed import docker +from .conftest import ( + create_deployment, + create_deployment_config, + delete_deployment, + get_deployment, + get_provider, +) + try: docker.from_env().ping() _DOCKER_AVAILABLE = True @@ -36,7 +44,7 @@ def test_deployments_plugin_docker_lifecycle(controller_with_deployments_plugin, Tests: create → PENDING → READY → delete → cleanup Also verifies ModelProvider creation and deletion. """ - controller, _, sdk, mock_nim_image, ctx, reconcile = controller_with_deployments_plugin + controller, _, models_client, mock_nim_image, ctx, reconcile = controller_with_deployments_plugin test_uuid = uuid.uuid4().hex[:8] config_name = f"test-plugin-lifecycle-{test_uuid}" deployment_name = f"test-plugin-lifecycle-{test_uuid}" @@ -47,7 +55,8 @@ def test_deployments_plugin_docker_lifecycle(controller_with_deployments_plugin, ctx.register_container(server_container_name) image_name, image_tag = mock_nim_image.rsplit(":", 1) - sdk.inference.deployment_configs.create( + create_deployment_config( + models_client, name=config_name, workspace=workspace, engine="nim", @@ -58,7 +67,8 @@ def test_deployments_plugin_docker_lifecycle(controller_with_deployments_plugin, "image_tag": image_tag, }, ) - sdk.inference.deployments.create( + create_deployment( + models_client, name=deployment_name, workspace=workspace, config=config_name, @@ -83,7 +93,7 @@ def wait_for_container_running(): @retry(stop=stop_after_delay(45), wait=wait_fixed(0.2), reraise=True) def wait_for_deployment_ready(): reconcile(controller) - dep = sdk.inference.deployments.retrieve(deployment_name, workspace=workspace) + dep = get_deployment(models_client, name=deployment_name, workspace=workspace) assert dep.status == "READY", f"Deployment not READY: {dep.status} ({dep.status_message})" return dep @@ -92,14 +102,14 @@ def wait_for_deployment_ready(): provider_id = deployment.model_provider_id assert provider_id is not None, "ModelProvider should be created when deployment becomes READY" provider_workspace, provider_name = provider_id.split("/") - provider = sdk.inference.providers.retrieve(provider_name, workspace=provider_workspace) + provider = get_provider(models_client, name=provider_name, workspace=provider_workspace) assert provider.host_url is not None assert provider.status == "READY" reconcile(controller) - sdk.inference.providers.retrieve(provider_name, workspace=provider_workspace) + get_provider(models_client, name=provider_name, workspace=provider_workspace) - sdk.inference.deployments.delete(deployment_name, workspace=workspace) + delete_deployment(models_client, name=deployment_name, workspace=workspace) reconcile(controller) @retry(stop=stop_after_delay(30), wait=wait_fixed(0.2), reraise=True) @@ -113,7 +123,7 @@ def wait_for_delete_complete(): except NotFound: pass try: - sdk.inference.providers.retrieve(provider_name, workspace=provider_workspace) + get_provider(models_client, name=provider_name, workspace=provider_workspace) except NotFoundError: return raise AssertionError("ModelProvider still exists after deployment delete") @@ -121,4 +131,4 @@ def wait_for_delete_complete(): wait_for_delete_complete() with pytest.raises(NotFoundError): - sdk.inference.providers.retrieve(provider_name, workspace=provider_workspace) + get_provider(models_client, name=provider_name, workspace=provider_workspace) diff --git a/services/core/models/tests/integration/test_models.py b/services/core/models/tests/integration/test_models.py index 1ae166d95f..ad0a071f75 100644 --- a/services/core/models/tests/integration/test_models.py +++ b/services/core/models/tests/integration/test_models.py @@ -15,8 +15,8 @@ import uuid from unittest.mock import AsyncMock, patch -from nemo_platform import ConflictError from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ConflictError from nemo_platform_plugin.workspaces.client import WorkspacesClient from nemo_platform_plugin.workspaces.types import CreateWorkspaceRequest from nmp.core.models.config import ControllerConfig, ModelsConfig @@ -1529,13 +1529,16 @@ def test_backend_config_key_deployments_plugin_works_end_to_end(): backends={"deployments_plugin": DeploymentsPluginBackendConfigModel(enabled=True)}, ) ) + backend_configs: dict[str, DeploymentsPluginBackendConfigModel] = { + "deployments_plugin": config.controller.backends["deployments_plugin"], + } with patch( "nmp.core.models.controllers.backends.deployments_plugin.backend.NemoEntitiesClient", ): registry = BackendRegistry.from_config( nmp_sdk=AsyncMock(), - backend_configs=config.controller.backends, + backend_configs=backend_configs, huggingface_model_puller=config.huggingface_model_puller, ) diff --git a/services/core/models/tests/integration/test_models_auth_propagation.py b/services/core/models/tests/integration/test_models_auth_propagation.py index 03e7da0c2d..40df072061 100644 --- a/services/core/models/tests/integration/test_models_auth_propagation.py +++ b/services/core/models/tests/integration/test_models_auth_propagation.py @@ -18,10 +18,22 @@ import pytest from nemo_platform import NeMoPlatform -from nmp.core.models.schemas import ContainerExecutorConfig, ModelDeploymentConfigModelSpec, ModelType +from nemo_platform_plugin.models.client import ModelsClient from nmp.core.models.service import ModelsService from nmp.testing import as_user, create_test_client, short_unique_name, unique_email +from .conftest import ( + create_deployment, + create_deployment_config, + create_provider, + get_deployment, + get_provider, + list_deployments, + list_providers, + models_client_from_sdk, + upsert_provider, +) + @pytest.fixture(scope="module") def sdk() -> Generator[NeMoPlatform, None, None]: @@ -33,9 +45,9 @@ def sdk() -> Generator[NeMoPlatform, None, None]: yield sdk -def _as_service_principal(sdk: NeMoPlatform, service_name: str = "models-controller") -> NeMoPlatform: +def _as_service_principal(sdk: NeMoPlatform, service_name: str = "models-controller") -> ModelsClient: """Create an SDK client authenticated as a service principal.""" - return as_user(sdk, f"service:{service_name}") + return models_client_from_sdk(as_user(sdk, f"service:{service_name}")) def _create_deployment( @@ -44,24 +56,19 @@ def _create_deployment( prefix: str = "test", ): """Create a deployment config + deployment, returning the deployment name.""" + client = models_client_from_sdk(user_sdk) config_name = short_unique_name(f"{prefix}-config") deployment_name = short_unique_name(f"{prefix}-deploy") - user_sdk.inference.deployment_configs.create( + create_deployment_config( + client, workspace=workspace, name=config_name, engine="nim", - model_spec=ModelDeploymentConfigModelSpec( - model_type=ModelType.LLM, - model_namespace="nvidia", - model_name="test-model", - ), - executor_config=ContainerExecutorConfig( - image_name="nvcr.io/nvidia/nim/llm", - image_tag="latest", - gpu=0, - ), + model_spec={"model_type": "llm", "model_namespace": "nvidia", "model_name": "test-model"}, + executor_config={"image_name": "nvcr.io/nvidia/nim/llm", "image_tag": "latest", "gpu": 0}, ) - return user_sdk.inference.deployments.create( + return create_deployment( + client, workspace=workspace, name=deployment_name, config=config_name, @@ -77,17 +84,18 @@ def test_auth_context_sanitized_for_regular_user(self, sdk: NeMoPlatform): # Create response assert deployment.auth_context is None, "create: regular user should not see auth_context" + creator_client = models_client_from_sdk(creator_sdk) # Retrieve response - retrieved = creator_sdk.inference.deployments.retrieve( + retrieved = get_deployment( + creator_client, workspace="default", name=deployment.name, ) assert retrieved.auth_context is None, "retrieve: regular user should not see auth_context" # List response - result = creator_sdk.inference.deployments.list(workspace="default") - matching = [d for d in result.data if d.name == deployment.name] + matching = [d for d in list_deployments(creator_client, workspace="default") if d.name == deployment.name] assert len(matching) == 1 assert matching[0].auth_context is None, "list: regular user should not see auth_context" @@ -101,7 +109,8 @@ def test_auth_context_visible_to_service_principal(self, sdk: NeMoPlatform): service_sdk = _as_service_principal(sdk) # Retrieve response - retrieved = service_sdk.inference.deployments.retrieve( + retrieved = get_deployment( + service_sdk, workspace="default", name=deployment.name, ) @@ -111,8 +120,7 @@ def test_auth_context_visible_to_service_principal(self, sdk: NeMoPlatform): assert retrieved.auth_context.principal_groups == creator_groups # List response - result = service_sdk.inference.deployments.list(workspace="default") - matching = [d for d in result.data if d.name == deployment.name] + matching = [d for d in list_deployments(service_sdk, workspace="default") if d.name == deployment.name] assert len(matching) == 1 assert matching[0].auth_context is not None, "list: service principal should see auth_context" assert matching[0].auth_context.principal_id == creator_email @@ -127,7 +135,8 @@ def test_auth_context_persisted_across_users(self, sdk: NeMoPlatform): # Different regular user should not see auth_context other_user = as_user(sdk, unique_email("admin"), groups=["admins"]) - retrieved_by_user = other_user.inference.deployments.retrieve( + retrieved_by_user = get_deployment( + models_client_from_sdk(other_user), workspace="default", name=deployment.name, ) @@ -135,7 +144,8 @@ def test_auth_context_persisted_across_users(self, sdk: NeMoPlatform): # Service principal should see the original creator's auth_context service_sdk = _as_service_principal(sdk) - retrieved_by_service = service_sdk.inference.deployments.retrieve( + retrieved_by_service = get_deployment( + service_sdk, workspace="default", name=deployment.name, ) @@ -152,9 +162,11 @@ def test_auth_context_captured_at_creation(self, sdk: NeMoPlatform): provider_name = short_unique_name("auth-prov") creator_sdk = as_user(sdk, creator_email, groups=creator_groups) + creator_client = models_client_from_sdk(creator_sdk) # Regular user should not see auth_context in the create response - provider = creator_sdk.inference.providers.create( + provider = create_provider( + creator_client, workspace="default", name=provider_name, host_url="http://test.local:8000", @@ -163,7 +175,8 @@ def test_auth_context_captured_at_creation(self, sdk: NeMoPlatform): # Service principal should see it service_sdk = _as_service_principal(sdk) - retrieved = service_sdk.inference.providers.retrieve( + retrieved = get_provider( + service_sdk, workspace="default", name=provider_name, ) @@ -177,16 +190,17 @@ def test_auth_context_stripped_for_regular_user_on_list(self, sdk: NeMoPlatform) provider_name = short_unique_name("list-prov") creator_sdk = as_user(sdk, unique_email("creator"), groups=["team-gamma"]) + creator_client = models_client_from_sdk(creator_sdk) - creator_sdk.inference.providers.create( + create_provider( + creator_client, workspace="default", name=provider_name, host_url="http://test.local:8000", ) # List as regular user — auth_context should be stripped - result = creator_sdk.inference.providers.list(workspace="default") - matching = [p for p in result.data if p.name == provider_name] + matching = [p for p in list_providers(creator_client, workspace="default") if p.name == provider_name] assert len(matching) == 1 assert matching[0].auth_context is None, "Regular user should not see auth_context in list" @@ -197,9 +211,11 @@ def test_auth_context_on_upsert(self, sdk: NeMoPlatform): provider_name = short_unique_name("upsert-prov") creator_sdk = as_user(sdk, creator_email, groups=creator_groups) + creator_client = models_client_from_sdk(creator_sdk) # Upsert creates a new provider - provider = creator_sdk.inference.providers.update( + provider = upsert_provider( + creator_client, workspace="default", name=provider_name, host_url="http://upsert.local:8000", @@ -208,7 +224,8 @@ def test_auth_context_on_upsert(self, sdk: NeMoPlatform): # Service principal should see auth_context after create service_sdk = _as_service_principal(sdk) - retrieved = service_sdk.inference.providers.retrieve( + retrieved = get_provider( + service_sdk, workspace="default", name=provider_name, ) @@ -217,7 +234,8 @@ def test_auth_context_on_upsert(self, sdk: NeMoPlatform): assert retrieved.auth_context.principal_groups == creator_groups # Upsert updates the existing provider - updated = creator_sdk.inference.providers.update( + updated = upsert_provider( + creator_client, workspace="default", name=provider_name, host_url="http://upsert-updated.local:9000", @@ -225,7 +243,8 @@ def test_auth_context_on_upsert(self, sdk: NeMoPlatform): assert updated.auth_context is None, "Regular user should not see auth_context after update" # Service principal should still see auth_context after update - retrieved_after_update = service_sdk.inference.providers.retrieve( + retrieved_after_update = get_provider( + service_sdk, workspace="default", name=provider_name, ) diff --git a/services/core/models/tests/integration/test_models_controller.py b/services/core/models/tests/integration/test_models_controller.py index cfd05f8d58..904fc4f6e5 100644 --- a/services/core/models/tests/integration/test_models_controller.py +++ b/services/core/models/tests/integration/test_models_controller.py @@ -13,9 +13,18 @@ from unittest.mock import AsyncMock import pytest -from nemo_platform import NotFoundError +from nemo_platform_plugin.client.errors import NotFoundError +from nemo_platform_plugin.models.types import ModelDeploymentStatus from nmp.core.models.controllers.backends.backends import DeploymentStatusUpdate +from .conftest import ( + create_deployment, + create_deployment_config, + delete_deployment, + get_deployment, + get_provider, +) + # ============================================================================= # Backend-Agnostic Tests (Mock Backend) # ============================================================================= @@ -42,7 +51,7 @@ def test_controller_step_marks_healthy(controller_with_mock_backend): def test_controller_reconciles_created_deployment(controller_with_mock_backend): """Test that controller calls backend.create_model_deployment for CREATED deployments.""" - controller, mock_backend, sdk = controller_with_mock_backend + controller, mock_backend, models_client = controller_with_mock_backend test_uuid = uuid.uuid4().hex[:8] config_name = f"test-config-{test_uuid}" deployment_name = f"test-deployment-{test_uuid}" @@ -50,13 +59,14 @@ def test_controller_reconciles_created_deployment(controller_with_mock_backend): # Configure mock backend to keep this specific deployment in PENDING state # (otherwise the reconciler processes PENDING->READY in the same step) mock_backend.status_responses[deployment_name] = DeploymentStatusUpdate( - status="PENDING", + status=ModelDeploymentStatus.PENDING, status_message="Still starting", host_url="http://localhost:8500", ) # Create deployment config first - sdk.inference.deployment_configs.create( + create_deployment_config( + models_client, name=config_name, workspace="default", engine="nim", @@ -65,14 +75,15 @@ def test_controller_reconciles_created_deployment(controller_with_mock_backend): ) # Create deployment - starts in CREATED status - sdk.inference.deployments.create( + create_deployment( + models_client, name=deployment_name, workspace="default", config=config_name, ) # Verify deployment is in CREATED status - deployment = sdk.inference.deployments.retrieve(deployment_name, workspace="default") + deployment = get_deployment(models_client, name=deployment_name, workspace="default") assert deployment.status == "CREATED" # Run controller step - should call backend.create_model_deployment @@ -88,26 +99,26 @@ def test_controller_reconciles_created_deployment(controller_with_mock_backend): assert called_config.name == config_name # Verify deployment status was updated to PENDING - deployment = sdk.inference.deployments.retrieve(deployment_name, workspace="default") + deployment = get_deployment(models_client, name=deployment_name, workspace="default") assert deployment.status == "PENDING" def test_controller_polls_pending_deployment(controller_with_mock_backend): """Test that controller calls backend.get_model_deployment_status for PENDING deployments.""" - controller, mock_backend, sdk = controller_with_mock_backend + controller, mock_backend, models_client = controller_with_mock_backend test_uuid = uuid.uuid4().hex[:8] config_name = f"test-config-poll-{test_uuid}" deployment_name = f"test-deployment-poll-{test_uuid}" # Create config and deployment - sdk.inference.deployment_configs.create( - name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} + create_deployment_config( + models_client, name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} ) - sdk.inference.deployments.create(name=deployment_name, workspace="default", config=config_name) + create_deployment(models_client, name=deployment_name, workspace="default", config=config_name) # Configure status to PENDING so first step doesn't immediately go to READY mock_backend.status_responses[deployment_name] = DeploymentStatusUpdate( - status="PENDING", + status=ModelDeploymentStatus.PENDING, status_message="Still starting", host_url="http://localhost:8500", ) @@ -121,7 +132,7 @@ def test_controller_polls_pending_deployment(controller_with_mock_backend): # Configure status response to return READY for this specific deployment mock_backend.status_responses[deployment_name] = DeploymentStatusUpdate( - status="READY", + status=ModelDeploymentStatus.READY, status_message="Container ready", host_url="http://localhost:8500", ) @@ -134,26 +145,26 @@ def test_controller_polls_pending_deployment(controller_with_mock_backend): assert len(deployment_status_calls) == 1 # Verify deployment was updated to READY - deployment = sdk.inference.deployments.retrieve(deployment_name, workspace="default") + deployment = get_deployment(models_client, name=deployment_name, workspace="default") assert deployment.status == "READY" def test_controller_handles_backend_error(controller_with_mock_backend): """Test that controller handles backend errors gracefully.""" - controller, mock_backend, sdk = controller_with_mock_backend + controller, mock_backend, models_client = controller_with_mock_backend test_uuid = uuid.uuid4().hex[:8] config_name = f"test-config-err-{test_uuid}" deployment_name = f"test-deployment-err-{test_uuid}" # Create config and deployment - sdk.inference.deployment_configs.create( - name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} + create_deployment_config( + models_client, name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} ) - sdk.inference.deployments.create(name=deployment_name, workspace="default", config=config_name) + create_deployment(models_client, name=deployment_name, workspace="default", config=config_name) # Configure backend to return ERROR mock_backend.create_response = DeploymentStatusUpdate( - status="ERROR", + status=ModelDeploymentStatus.ERROR, status_message="Failed to create container", error_details={"error": "Image not found"}, ) @@ -162,39 +173,43 @@ def test_controller_handles_backend_error(controller_with_mock_backend): controller.step() # Verify deployment status was updated to ERROR - deployment = sdk.inference.deployments.retrieve(deployment_name, workspace="default") + deployment = get_deployment(models_client, name=deployment_name, workspace="default") assert deployment.status == "ERROR" assert "Failed to create container" in (deployment.status_message or "") def test_controller_deletes_when_deleting(controller_with_mock_backend): """Test that controller calls backend.delete_model_deployment for DELETING deployments.""" - controller, mock_backend, sdk = controller_with_mock_backend + controller, mock_backend, models_client = controller_with_mock_backend test_uuid = uuid.uuid4().hex[:8] config_name = f"test-config-del-{test_uuid}" deployment_name = f"test-deployment-del-{test_uuid}" # Create config and deployment - sdk.inference.deployment_configs.create( - name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} + create_deployment_config( + models_client, name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} ) - sdk.inference.deployments.create(name=deployment_name, workspace="default", config=config_name) + create_deployment(models_client, name=deployment_name, workspace="default", config=config_name) # Move to READY state - mock_backend.create_response = DeploymentStatusUpdate(status="PENDING", status_message="Starting") + mock_backend.create_response = DeploymentStatusUpdate( + status=ModelDeploymentStatus.PENDING, status_message="Starting" + ) controller.step() - mock_backend.status_responses[deployment_name] = DeploymentStatusUpdate(status="READY", status_message="Ready") + mock_backend.status_responses[deployment_name] = DeploymentStatusUpdate( + status=ModelDeploymentStatus.READY, status_message="Ready" + ) controller.step() # Delete the deployment (moves to DELETING) - sdk.inference.deployments.delete(deployment_name, workspace="default") + delete_deployment(models_client, name=deployment_name, workspace="default") # Clear call history mock_backend.delete_calls.clear() # Configure delete response mock_backend.delete_response = DeploymentStatusUpdate( - status="DELETED", + status=ModelDeploymentStatus.DELETED, status_message="Container deleted", ) @@ -208,34 +223,38 @@ def test_controller_deletes_when_deleting(controller_with_mock_backend): def test_controller_garbage_collects_deleted_deployment(controller_with_mock_backend): """Test that controller hard-deletes DELETED deployments after grace period expires.""" - controller, mock_backend, sdk = controller_with_mock_backend + controller, mock_backend, models_client = controller_with_mock_backend test_uuid = uuid.uuid4().hex[:8] config_name = f"test-config-gc-{test_uuid}" deployment_name = f"test-deployment-gc-{test_uuid}" # Create config and deployment - sdk.inference.deployment_configs.create( - name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} + create_deployment_config( + models_client, name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} ) - sdk.inference.deployments.create(name=deployment_name, workspace="default", config=config_name) + create_deployment(models_client, name=deployment_name, workspace="default", config=config_name) # Progress through lifecycle: CREATED → PENDING → READY → DELETING → DELETED - mock_backend.create_response = DeploymentStatusUpdate(status="PENDING", status_message="Starting") + mock_backend.create_response = DeploymentStatusUpdate( + status=ModelDeploymentStatus.PENDING, status_message="Starting" + ) controller.step() mock_backend.status_responses[deployment_name] = DeploymentStatusUpdate( - status="READY", status_message="Ready", host_url="http://localhost:8080" + status=ModelDeploymentStatus.READY, status_message="Ready", host_url="http://localhost:8080" ) controller.step() # Delete deployment (moves to DELETING) - sdk.inference.deployments.delete(deployment_name, workspace="default") + delete_deployment(models_client, name=deployment_name, workspace="default") - mock_backend.delete_response = DeploymentStatusUpdate(status="DELETED", status_message="Deleted") + mock_backend.delete_response = DeploymentStatusUpdate( + status=ModelDeploymentStatus.DELETED, status_message="Deleted" + ) controller.step() # DELETING → DELETED # Verify deployment is in DELETED state (soft-deleted, still exists) - deployment = sdk.inference.deployments.retrieve(deployment_name, workspace="default") + deployment = get_deployment(models_client, name=deployment_name, workspace="default") assert deployment.status == "DELETED" # Patch the controller's reconciler to have 0 second grace period @@ -246,7 +265,7 @@ def test_controller_garbage_collects_deleted_deployment(controller_with_mock_bac # Verify deployment is gone (hard-deleted) with pytest.raises(NotFoundError): - sdk.inference.deployments.retrieve(deployment_name, workspace="default") + get_deployment(models_client, name=deployment_name, workspace="default") def test_controller_orphan_cleanup_after_deleted(controller_with_mock_backend): @@ -256,33 +275,37 @@ def test_controller_orphan_cleanup_after_deleted(controller_with_mock_backend): → then simulate backend still reporting the deployment (orphan) → next step runs reconcile_orphans and calls delete_model_deployment(workspace, name) for the orphan. """ - controller, mock_backend, sdk = controller_with_mock_backend + controller, mock_backend, models_client = controller_with_mock_backend test_uuid = uuid.uuid4().hex[:8] config_name = f"test-config-orphan-{test_uuid}" deployment_name = f"test-deployment-orphan-{test_uuid}" # Create config and deployment - sdk.inference.deployment_configs.create( - name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} + create_deployment_config( + models_client, name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} ) - sdk.inference.deployments.create(name=deployment_name, workspace="default", config=config_name) + create_deployment(models_client, name=deployment_name, workspace="default", config=config_name) # CREATED → PENDING → READY - mock_backend.create_response = DeploymentStatusUpdate(status="PENDING", status_message="Starting") + mock_backend.create_response = DeploymentStatusUpdate( + status=ModelDeploymentStatus.PENDING, status_message="Starting" + ) controller.step() mock_backend.status_responses[deployment_name] = DeploymentStatusUpdate( - status="READY", status_message="Ready", host_url="http://localhost:8080" + status=ModelDeploymentStatus.READY, status_message="Ready", host_url="http://localhost:8080" ) controller.step() # Delete via API (moves to DELETING) - sdk.inference.deployments.delete(deployment_name, workspace="default") + delete_deployment(models_client, name=deployment_name, workspace="default") - mock_backend.delete_response = DeploymentStatusUpdate(status="DELETED", status_message="Deleted") + mock_backend.delete_response = DeploymentStatusUpdate( + status=ModelDeploymentStatus.DELETED, status_message="Deleted" + ) controller.step() # DELETING → DELETED - deployment = sdk.inference.deployments.retrieve(deployment_name, workspace="default") + deployment = get_deployment(models_client, name=deployment_name, workspace="default") assert deployment.status == "DELETED" # Hard-delete after grace period so deployment is no longer in API @@ -290,7 +313,7 @@ def test_controller_orphan_cleanup_after_deleted(controller_with_mock_backend): controller.step() with pytest.raises(NotFoundError): - sdk.inference.deployments.retrieve(deployment_name, workspace="default") + get_deployment(models_client, name=deployment_name, workspace="default") # Simulate backend still reporting this deployment (orphan) deployment_id = f"default/{deployment_name}" @@ -309,70 +332,73 @@ def test_controller_orphan_cleanup_after_deleted(controller_with_mock_backend): def test_controller_creates_model_provider_when_ready(controller_with_mock_backend): """Test that controller creates ModelProvider when deployment becomes READY.""" - controller, mock_backend, sdk = controller_with_mock_backend + controller, mock_backend, models_client = controller_with_mock_backend test_uuid = uuid.uuid4().hex[:8] config_name = f"test-config-prov-{test_uuid}" deployment_name = f"test-deployment-prov-{test_uuid}" # Create config and deployment - sdk.inference.deployment_configs.create( - name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} + create_deployment_config( + models_client, name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} ) - sdk.inference.deployments.create(name=deployment_name, workspace="default", config=config_name) + create_deployment(models_client, name=deployment_name, workspace="default", config=config_name) # Move to READY state with host_url - this should trigger provider creation mock_backend.create_response = DeploymentStatusUpdate( - status="READY", status_message="Ready", host_url="http://localhost:9000" + status=ModelDeploymentStatus.READY, status_message="Ready", host_url="http://localhost:9000" ) controller.step() # Verify deployment has model_provider_id set - deployment = sdk.inference.deployments.retrieve(deployment_name, workspace="default") + deployment = get_deployment(models_client, name=deployment_name, workspace="default") assert deployment.status == "READY" assert deployment.model_provider_id is not None # Verify provider was created with correct host_url and status provider_id = deployment.model_provider_id provider_workspace, provider_name = provider_id.split("/") - provider = sdk.inference.providers.retrieve(provider_name, workspace=provider_workspace) + provider = get_provider(models_client, name=provider_name, workspace=provider_workspace) assert provider.host_url == "http://localhost:9000" assert provider.status == "READY", "Provider should be READY when deployment is READY" def test_controller_deletes_model_provider_on_delete(controller_with_mock_backend): """Test that controller deletes ModelProvider when deployment is deleted.""" - controller, mock_backend, sdk = controller_with_mock_backend + controller, mock_backend, models_client = controller_with_mock_backend test_uuid = uuid.uuid4().hex[:8] config_name = f"test-config-delprov-{test_uuid}" deployment_name = f"test-deployment-delprov-{test_uuid}" # Create config and deployment - sdk.inference.deployment_configs.create( - name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} + create_deployment_config( + models_client, name=config_name, workspace="default", engine="nim", model_spec={}, executor_config={"gpu": 0} ) - sdk.inference.deployments.create(name=deployment_name, workspace="default", config=config_name) + create_deployment(models_client, name=deployment_name, workspace="default", config=config_name) # Move to READY state (creates provider) mock_backend.create_response = DeploymentStatusUpdate( - status="READY", status_message="Ready", host_url="http://localhost:9001" + status=ModelDeploymentStatus.READY, status_message="Ready", host_url="http://localhost:9001" ) controller.step() # Get provider info before deletion - deployment = sdk.inference.deployments.retrieve(deployment_name, workspace="default") + deployment = get_deployment(models_client, name=deployment_name, workspace="default") provider_id = deployment.model_provider_id + assert provider_id is not None provider_workspace, provider_name = provider_id.split("/") # Verify provider exists - sdk.inference.providers.retrieve(provider_name, workspace=provider_workspace) + get_provider(models_client, name=provider_name, workspace=provider_workspace) # Delete deployment (moves to DELETING) - sdk.inference.deployments.delete(deployment_name, workspace="default") + delete_deployment(models_client, name=deployment_name, workspace="default") # Configure delete response and run controller - mock_backend.delete_response = DeploymentStatusUpdate(status="DELETED", status_message="Deleted") + mock_backend.delete_response = DeploymentStatusUpdate( + status=ModelDeploymentStatus.DELETED, status_message="Deleted" + ) controller.step() # Verify provider was deleted with pytest.raises(NotFoundError): - sdk.inference.providers.retrieve(provider_name, workspace=provider_workspace) + get_provider(models_client, name=provider_name, workspace=provider_workspace) 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 e4deafa835..5810dbd6ec 100644 --- a/services/core/models/tests/integration/test_models_with_auth.py +++ b/services/core/models/tests/integration/test_models_with_auth.py @@ -23,10 +23,10 @@ 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 +from nemo_platform_plugin.files.storage_config import HuggingfaceStorageConfig from nemo_platform_plugin.files.types import CreateFilesetRequest from nemo_platform_plugin.models.client import ModelsClient from nemo_platform_plugin.models.types import ( @@ -54,6 +54,25 @@ ) from pydantic import SecretStr +from .conftest import ( + create_deployment, + create_deployment_config, + create_provider, + delete_deployment, + delete_deployment_config, + delete_provider, + get_deployment, + get_deployment_config, + get_provider, + list_deployment_configs, + list_deployments, + list_providers, + models_client_from_sdk, + update_deployment, + update_deployment_config, + upsert_provider, +) + async def _build_authorization_data_without_secrets(entities_client=None): """Wraps the real build_authorization_data to inject an EditorNoSecrets role. @@ -329,19 +348,22 @@ def viewer_workspace(sdk: NeMoPlatform): client_from_platform(admin_sdk, ModelsClient).create_model( workspace=workspace, body=CreateModelEntityRequest(name=model_name) ).data() - admin_sdk.inference.providers.create( + create_provider( + models_client_from_sdk(admin_sdk), workspace=workspace, name=provider_name, host_url="http://example.com", ) - admin_sdk.inference.deployment_configs.create( + create_deployment_config( + models_client_from_sdk(admin_sdk), workspace=workspace, name=config_name, engine="nim", model_spec={"model_name": "test"}, executor_config={"gpu": 1}, ) - admin_sdk.inference.deployments.create( + create_deployment( + models_client_from_sdk(admin_sdk), workspace=workspace, name=deployment_name, config=config_name, @@ -403,21 +425,22 @@ def test_viewer_cannot_delete_model(self, viewer_workspace): # -- Providers: allowed -- def test_viewer_can_list_providers(self, viewer_workspace): - workspace, viewer_sdk, _, _ = viewer_workspace - result = viewer_sdk.inference.providers.list(workspace=workspace) - assert result.data is not None + workspace, viewer_sdk, _, names = viewer_workspace + result = list_providers(models_client_from_sdk(viewer_sdk), workspace=workspace) + assert any(provider.name == names["provider"] for provider in result) def test_viewer_can_get_provider(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - provider = viewer_sdk.inference.providers.retrieve(name=names["provider"], workspace=workspace) + provider = get_provider(models_client_from_sdk(viewer_sdk), name=names["provider"], workspace=workspace) assert provider.name == names["provider"] # -- Providers: denied -- def test_viewer_cannot_create_provider(self, viewer_workspace): workspace, viewer_sdk, _, _ = viewer_workspace - with pytest.raises(StainlessPermissionDeniedError): - viewer_sdk.inference.providers.create( + with pytest.raises(PermissionDeniedError): + create_provider( + models_client_from_sdk(viewer_sdk), workspace=workspace, name="should-fail", host_url="http://example.com", @@ -425,8 +448,9 @@ 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(StainlessPermissionDeniedError): - viewer_sdk.inference.providers.update( + with pytest.raises(PermissionDeniedError): + upsert_provider( + models_client_from_sdk(viewer_sdk), name=names["provider"], workspace=workspace, host_url="http://updated.com", @@ -434,27 +458,28 @@ 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(StainlessPermissionDeniedError): - viewer_sdk.inference.providers.delete(name=names["provider"], workspace=workspace) + with pytest.raises(PermissionDeniedError): + delete_provider(models_client_from_sdk(viewer_sdk), name=names["provider"], workspace=workspace) # -- Deployment Configs: allowed -- def test_viewer_can_list_deployment_configs(self, viewer_workspace): - workspace, viewer_sdk, _, _ = viewer_workspace - result = viewer_sdk.inference.deployment_configs.list(workspace=workspace) - assert result.data is not None + workspace, viewer_sdk, _, names = viewer_workspace + result = list_deployment_configs(models_client_from_sdk(viewer_sdk), workspace=workspace) + assert any(config.name == names["config"] for config in result) def test_viewer_can_get_deployment_config(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - config = viewer_sdk.inference.deployment_configs.retrieve(name=names["config"], workspace=workspace) + config = get_deployment_config(models_client_from_sdk(viewer_sdk), name=names["config"], workspace=workspace) assert config.name == names["config"] # -- Deployment Configs: denied -- def test_viewer_cannot_create_deployment_config(self, viewer_workspace): workspace, viewer_sdk, _, _ = viewer_workspace - with pytest.raises(StainlessPermissionDeniedError): - viewer_sdk.inference.deployment_configs.create( + with pytest.raises(PermissionDeniedError): + create_deployment_config( + models_client_from_sdk(viewer_sdk), workspace=workspace, name="should-fail", engine="nim", @@ -464,8 +489,9 @@ 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(StainlessPermissionDeniedError): - viewer_sdk.inference.deployment_configs.update( + with pytest.raises(PermissionDeniedError): + update_deployment_config( + models_client_from_sdk(viewer_sdk), name=names["config"], workspace=workspace, engine="nim", @@ -475,27 +501,28 @@ 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(StainlessPermissionDeniedError): - viewer_sdk.inference.deployment_configs.delete(name=names["config"], workspace=workspace) + with pytest.raises(PermissionDeniedError): + delete_deployment_config(models_client_from_sdk(viewer_sdk), name=names["config"], workspace=workspace) # -- Deployments: allowed -- def test_viewer_can_list_deployments(self, viewer_workspace): - workspace, viewer_sdk, _, _ = viewer_workspace - result = viewer_sdk.inference.deployments.list(workspace=workspace) - assert result.data is not None + workspace, viewer_sdk, _, names = viewer_workspace + result = list_deployments(models_client_from_sdk(viewer_sdk), workspace=workspace) + assert any(deployment.name == names["deployment"] for deployment in result) def test_viewer_can_get_deployment(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - deployment = viewer_sdk.inference.deployments.retrieve(name=names["deployment"], workspace=workspace) + deployment = get_deployment(models_client_from_sdk(viewer_sdk), name=names["deployment"], workspace=workspace) assert deployment.name == names["deployment"] # -- Deployments: denied -- def test_viewer_cannot_create_deployment(self, viewer_workspace): workspace, viewer_sdk, _, names = viewer_workspace - with pytest.raises(StainlessPermissionDeniedError): - viewer_sdk.inference.deployments.create( + with pytest.raises(PermissionDeniedError): + create_deployment( + models_client_from_sdk(viewer_sdk), workspace=workspace, name="should-fail", config=names["config"], @@ -503,8 +530,9 @@ 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(StainlessPermissionDeniedError): - viewer_sdk.inference.deployments.update( + with pytest.raises(PermissionDeniedError): + update_deployment( + models_client_from_sdk(viewer_sdk), name=names["deployment"], workspace=workspace, config=names["config"], @@ -512,8 +540,8 @@ 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(StainlessPermissionDeniedError): - viewer_sdk.inference.deployments.delete(name=names["deployment"], workspace=workspace) + with pytest.raises(PermissionDeniedError): + delete_deployment(models_client_from_sdk(viewer_sdk), name=names["deployment"], workspace=workspace) @pytest.mark.integration @@ -582,7 +610,8 @@ def test_editor_can_create_provider(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - created = editor_sdk.inference.providers.create( + created = create_provider( + models_client_from_sdk(editor_sdk), workspace=workspace, name=provider_name, host_url="http://example.com", @@ -606,7 +635,8 @@ def test_editor_can_create_deployment_config(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - created = editor_sdk.inference.deployment_configs.create( + created = create_deployment_config( + models_client_from_sdk(editor_sdk), workspace=workspace, name=config_name, engine="nim", @@ -636,7 +666,8 @@ def test_editor_can_create_provider_without_secret(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - provider = editor_sdk.inference.providers.create( + provider = create_provider( + models_client_from_sdk(editor_sdk), workspace=workspace, name=short_unique_name("prov"), host_url="http://example.com", @@ -664,7 +695,8 @@ def test_editor_can_create_provider_with_secret(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - provider = editor_sdk.inference.providers.create( + provider = create_provider( + models_client_from_sdk(editor_sdk), workspace=workspace, name=short_unique_name("prov"), host_url="http://example.com", @@ -694,7 +726,8 @@ def test_editor_can_upsert_provider_with_secret(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - provider = editor_sdk.inference.providers.update( + provider = upsert_provider( + models_client_from_sdk(editor_sdk), name=provider_name, workspace=workspace, host_url="http://example.com", @@ -725,15 +758,17 @@ def test_custom_role_denied_create_provider_with_secret(self, sdk: NeMoPlatform) user_sdk = as_user(sdk, user_email) - provider_ok = user_sdk.inference.providers.create( + provider_ok = create_provider( + models_client_from_sdk(user_sdk), workspace=workspace, name=short_unique_name("prov"), host_url="http://example.com", ) assert provider_ok.api_key_secret_name is None - with pytest.raises(StainlessPermissionDeniedError): - user_sdk.inference.providers.create( + with pytest.raises(PermissionDeniedError): + create_provider( + models_client_from_sdk(user_sdk), workspace=workspace, name=short_unique_name("prov"), host_url="http://example.com", @@ -763,8 +798,9 @@ def test_custom_role_denied_upsert_provider_with_secret(self, sdk: NeMoPlatform) user_sdk = as_user(sdk, user_email) - with pytest.raises(StainlessPermissionDeniedError): - user_sdk.inference.providers.update( + with pytest.raises(PermissionDeniedError): + upsert_provider( + models_client_from_sdk(user_sdk), name=short_unique_name("prov"), workspace=workspace, host_url="http://example.com", @@ -793,7 +829,8 @@ def test_editor_can_create_provider_with_deployment_ref(self, sdk: NeMoPlatform) ) editor_sdk = as_user(sdk, editor_email) - provider = editor_sdk.inference.providers.create( + provider = create_provider( + models_client_from_sdk(editor_sdk), workspace=workspace, name=short_unique_name("prov"), host_url="http://example.com", @@ -818,7 +855,8 @@ def test_editor_can_upsert_provider_with_deployment_ref(self, sdk: NeMoPlatform) ) editor_sdk = as_user(sdk, editor_email) - provider = editor_sdk.inference.providers.update( + provider = upsert_provider( + models_client_from_sdk(editor_sdk), name=short_unique_name("prov"), workspace=workspace, host_url="http://example.com", @@ -845,8 +883,9 @@ def test_custom_role_denied_create_provider_with_deployment_ref(self, sdk: NeMoP user_sdk = as_user(sdk, user_email) - with pytest.raises(StainlessPermissionDeniedError): - user_sdk.inference.providers.create( + with pytest.raises(PermissionDeniedError): + create_provider( + models_client_from_sdk(user_sdk), workspace=workspace, name=short_unique_name("prov"), host_url="http://example.com", @@ -872,8 +911,9 @@ def test_custom_role_denied_upsert_provider_with_deployment_ref(self, sdk: NeMoP user_sdk = as_user(sdk, user_email) - with pytest.raises(StainlessPermissionDeniedError): - user_sdk.inference.providers.update( + with pytest.raises(PermissionDeniedError): + upsert_provider( + models_client_from_sdk(user_sdk), name=short_unique_name("prov"), workspace=workspace, host_url="http://example.com", @@ -902,7 +942,8 @@ def test_editor_can_create_config_with_model_entity_id(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - config = editor_sdk.inference.deployment_configs.create( + config = create_deployment_config( + models_client_from_sdk(editor_sdk), workspace=workspace, name=short_unique_name("cfg"), engine="nim", @@ -929,8 +970,9 @@ def test_editor_cannot_reference_model_in_inaccessible_workspace(self, sdk: NeMo ) editor_sdk = as_user(sdk, editor_email) - with pytest.raises(StainlessPermissionDeniedError): - editor_sdk.inference.deployment_configs.create( + with pytest.raises(PermissionDeniedError): + create_deployment_config( + models_client_from_sdk(editor_sdk), workspace=workspace, name=short_unique_name("cfg"), engine="nim", @@ -949,7 +991,8 @@ def test_editor_can_update_config_with_model_entity_id(self, sdk: NeMoPlatform): client_from_platform(admin_sdk, WorkspacesClient).create_workspace( body=CreateWorkspaceRequest(name=workspace) ).data() - admin_sdk.inference.deployment_configs.create( + create_deployment_config( + models_client_from_sdk(admin_sdk), workspace=workspace, name=config_name, engine="nim", @@ -964,7 +1007,8 @@ def test_editor_can_update_config_with_model_entity_id(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - updated = editor_sdk.inference.deployment_configs.update( + updated = update_deployment_config( + models_client_from_sdk(editor_sdk), name=config_name, workspace=workspace, engine="nim", @@ -984,7 +1028,8 @@ def test_editor_cannot_update_config_with_inaccessible_model(self, sdk: NeMoPlat client_from_platform(admin_sdk, WorkspacesClient).create_workspace( body=CreateWorkspaceRequest(name=workspace) ).data() - admin_sdk.inference.deployment_configs.create( + create_deployment_config( + models_client_from_sdk(admin_sdk), workspace=workspace, name=config_name, engine="nim", @@ -999,8 +1044,9 @@ def test_editor_cannot_update_config_with_inaccessible_model(self, sdk: NeMoPlat ) editor_sdk = as_user(sdk, editor_email) - with pytest.raises(StainlessPermissionDeniedError): - editor_sdk.inference.deployment_configs.update( + with pytest.raises(PermissionDeniedError): + update_deployment_config( + models_client_from_sdk(editor_sdk), name=config_name, workspace=workspace, engine="nim", @@ -1028,8 +1074,9 @@ def test_custom_role_denied_create_config_with_model_entity_id_without_read(self user_sdk = as_user(sdk, user_email) - with pytest.raises(StainlessPermissionDeniedError): - user_sdk.inference.deployment_configs.create( + with pytest.raises(PermissionDeniedError): + create_deployment_config( + models_client_from_sdk(user_sdk), workspace=workspace, name=short_unique_name("cfg"), engine="nim", @@ -1049,7 +1096,8 @@ def test_custom_role_denied_update_config_with_model_entity_id_without_read(self client_from_platform(admin_sdk, WorkspacesClient).create_workspace( body=CreateWorkspaceRequest(name=workspace) ).data() - admin_sdk.inference.deployment_configs.create( + create_deployment_config( + models_client_from_sdk(admin_sdk), workspace=workspace, name=config_name, engine="nim", @@ -1065,8 +1113,9 @@ def test_custom_role_denied_update_config_with_model_entity_id_without_read(self user_sdk = as_user(sdk, user_email) - with pytest.raises(StainlessPermissionDeniedError): - user_sdk.inference.deployment_configs.update( + with pytest.raises(PermissionDeniedError): + update_deployment_config( + models_client_from_sdk(user_sdk), name=config_name, workspace=workspace, engine="nim", @@ -1090,7 +1139,8 @@ def test_editor_can_create_deployment_with_config(self, sdk: NeMoPlatform): client_from_platform(admin_sdk, WorkspacesClient).create_workspace( body=CreateWorkspaceRequest(name=workspace) ).data() - admin_sdk.inference.deployment_configs.create( + create_deployment_config( + models_client_from_sdk(admin_sdk), workspace=workspace, name=config_name, engine="nim", @@ -1105,7 +1155,8 @@ def test_editor_can_create_deployment_with_config(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - deployment = editor_sdk.inference.deployments.create( + deployment = create_deployment( + models_client_from_sdk(editor_sdk), workspace=workspace, name=short_unique_name("dep"), config=config_name, @@ -1123,14 +1174,16 @@ def test_editor_can_update_deployment_with_config(self, sdk: NeMoPlatform): client_from_platform(admin_sdk, WorkspacesClient).create_workspace( body=CreateWorkspaceRequest(name=workspace) ).data() - admin_sdk.inference.deployment_configs.create( + create_deployment_config( + models_client_from_sdk(admin_sdk), workspace=workspace, name=config_name, engine="nim", model_spec={"model_name": "test-model"}, executor_config={"gpu": 1}, ) - admin_sdk.inference.deployments.create( + create_deployment( + models_client_from_sdk(admin_sdk), workspace=workspace, name=deploy_name, config=config_name, @@ -1143,7 +1196,8 @@ def test_editor_can_update_deployment_with_config(self, sdk: NeMoPlatform): ) editor_sdk = as_user(sdk, editor_email) - updated = editor_sdk.inference.deployments.update( + updated = update_deployment( + models_client_from_sdk(editor_sdk), name=deploy_name, workspace=workspace, config=config_name, @@ -1161,7 +1215,8 @@ def test_custom_role_denied_create_deployment_without_read(self, sdk: NeMoPlatfo client_from_platform(admin_sdk, WorkspacesClient).create_workspace( body=CreateWorkspaceRequest(name=workspace) ).data() - admin_sdk.inference.deployment_configs.create( + create_deployment_config( + models_client_from_sdk(admin_sdk), workspace=workspace, name=config_name, engine="nim", @@ -1177,8 +1232,9 @@ def test_custom_role_denied_create_deployment_without_read(self, sdk: NeMoPlatfo user_sdk = as_user(sdk, user_email) - with pytest.raises(StainlessPermissionDeniedError): - user_sdk.inference.deployments.create( + with pytest.raises(PermissionDeniedError): + create_deployment( + models_client_from_sdk(user_sdk), workspace=workspace, name=short_unique_name("dep"), config=config_name, @@ -1196,14 +1252,16 @@ def test_custom_role_denied_update_deployment_without_read(self, sdk: NeMoPlatfo client_from_platform(admin_sdk, WorkspacesClient).create_workspace( body=CreateWorkspaceRequest(name=workspace) ).data() - admin_sdk.inference.deployment_configs.create( + create_deployment_config( + models_client_from_sdk(admin_sdk), workspace=workspace, name=config_name, engine="nim", model_spec={"model_name": "test-model"}, executor_config={"gpu": 1}, ) - admin_sdk.inference.deployments.create( + create_deployment( + models_client_from_sdk(admin_sdk), workspace=workspace, name=deploy_name, config=config_name, @@ -1217,8 +1275,9 @@ def test_custom_role_denied_update_deployment_without_read(self, sdk: NeMoPlatfo user_sdk = as_user(sdk, user_email) - with pytest.raises(StainlessPermissionDeniedError): - user_sdk.inference.deployments.update( + with pytest.raises(PermissionDeniedError): + update_deployment( + models_client_from_sdk(user_sdk), name=deploy_name, workspace=workspace, config=config_name, @@ -1477,7 +1536,7 @@ def test_create_model_trust_remote_code_true_has_permission_succeeds(self, sdk: ).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"}), + body=CreateFilesetRequest(name=fileset_name, storage=HuggingfaceStorageConfig(repo_id="Qwen/Qwen3-0.6B")), ) grant_workspace_role( admin_sdk, @@ -1514,7 +1573,7 @@ def test_create_model_trust_remote_code_true_without_permission_raises(self, sdk client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, body=CreateFilesetRequest( - name=fileset_name, storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"} + name=fileset_name, storage=HuggingfaceStorageConfig(repo_id="Qwen/Qwen3-0.6B") ), ) grant_workspace_role( @@ -1551,7 +1610,7 @@ def test_update_model_trust_remote_code_true_has_permission_succeeds(self, sdk: ).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"}), + body=CreateFilesetRequest(name=fileset_name, storage=HuggingfaceStorageConfig(repo_id="Qwen/Qwen3-0.6B")), ) grant_workspace_role( admin_sdk, @@ -1591,7 +1650,7 @@ def test_update_model_trust_remote_code_true_without_permission_raises(self, sdk client_from_platform(admin_sdk, FilesClient).create_fileset( workspace=workspace, body=CreateFilesetRequest( - name=fileset_name, storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"} + name=fileset_name, storage=HuggingfaceStorageConfig(repo_id="Qwen/Qwen3-0.6B") ), ) grant_workspace_role( @@ -1633,7 +1692,7 @@ def test_update_model_new_fileset_not_trusted_raises_permission_error(self, sdk: workspace=workspace, body=CreateFilesetRequest( name=trusted_fs, - storage={"type": "huggingface", "repo_id": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"}, + storage=HuggingfaceStorageConfig(repo_id="nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"), ), ) client_from_platform(admin_sdk, ModelsClient).create_model( @@ -1645,7 +1704,7 @@ def test_update_model_new_fileset_not_trusted_raises_permission_error(self, sdk: # New fileset resolves to a repo not on the allow list. files.create_fileset( workspace=workspace, - body=CreateFilesetRequest(name=new_fs, storage={"type": "huggingface", "repo_id": "Qwen/Qwen3-0.6B"}), + body=CreateFilesetRequest(name=new_fs, storage=HuggingfaceStorageConfig(repo_id="Qwen/Qwen3-0.6B")), ) grant_workspace_role( admin_sdk, @@ -1680,7 +1739,7 @@ def test_exact_match_on_allow_list_succeeds(self, sdk: NeMoPlatform): workspace=workspace, body=CreateFilesetRequest( name=fileset_name, - storage={"type": "huggingface", "repo_id": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"}, + storage=HuggingfaceStorageConfig(repo_id="nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"), ), ) grant_workspace_role( diff --git a/services/core/models/tests/unit/common/test_utils.py b/services/core/models/tests/unit/common/test_utils.py index 1ef0dbf33e..8de0f08f08 100644 --- a/services/core/models/tests/unit/common/test_utils.py +++ b/services/core/models/tests/unit/common/test_utils.py @@ -8,13 +8,15 @@ from datetime import datetime import pytest -from nemo_platform.types.inference.container_executor_config import ContainerExecutorConfig -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_deployment_config_model_spec import ModelDeploymentConfigModelSpec -from nemo_platform.types.inference.model_provider import ModelProvider -from nemo_platform.types.models.model_entity import ModelEntity -from nemo_platform.types.shared import ModelSpec +from nemo_platform_plugin.models.types import ( + ContainerExecutorConfig, + ModelDeployment, + ModelDeploymentConfig, + ModelDeploymentConfigModelSpec, + ModelEntity, + ModelProvider, + ModelSpec, +) from nmp.core.models.app import normalize_model_entity_name from nmp.core.models.app.utils import ( ModelConfigParseError, diff --git a/services/core/models/tests/unit/controllers/backends/deployments_plugin/test_compiler.py b/services/core/models/tests/unit/controllers/backends/deployments_plugin/test_compiler.py index 097711b802..40e3e2e24f 100644 --- a/services/core/models/tests/unit/controllers/backends/deployments_plugin/test_compiler.py +++ b/services/core/models/tests/unit/controllers/backends/deployments_plugin/test_compiler.py @@ -8,8 +8,8 @@ import pytest import yaml from nemo_deployments_plugin.entities import SecretRef -from nemo_platform.types.inference.k8s_nim_operator_config import K8sNIMOperatorConfig from nemo_platform_plugin.auth import AuthContext +from nemo_platform_plugin.models.types import K8sNIMOperatorConfig from nmp.common.config import Runtime from nmp.core.models.app import ModelWeightsType from nmp.core.models.controllers.backends.common import DeploymentConfigView diff --git a/services/core/models/tests/unit/controllers/backends/deployments_plugin/test_nim_compiler.py b/services/core/models/tests/unit/controllers/backends/deployments_plugin/test_nim_compiler.py index 6aaeb3aad9..4d070beca8 100644 --- a/services/core/models/tests/unit/controllers/backends/deployments_plugin/test_nim_compiler.py +++ b/services/core/models/tests/unit/controllers/backends/deployments_plugin/test_nim_compiler.py @@ -13,7 +13,7 @@ Toleration, VolumeMount, ) -from nemo_platform.types.inference.k8s_nim_operator_config import K8sNIMOperatorConfig +from nemo_platform_plugin.models.types import K8sNIMOperatorConfig from nmp.common.config import Runtime from nmp.core.models.app import ModelWeightsType from nmp.core.models.controllers.backends.common import DeploymentConfigView diff --git a/services/core/models/tests/unit/controllers/conftest.py b/services/core/models/tests/unit/controllers/conftest.py index 4b11091713..60b299518c 100644 --- a/services/core/models/tests/unit/controllers/conftest.py +++ b/services/core/models/tests/unit/controllers/conftest.py @@ -3,10 +3,12 @@ """Test fixtures for Models Controller tests.""" +import inspect from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from nemo_platform_plugin.client.errors import ConflictError, NemoHTTPError, NotFoundError from nmp.common.config import PlatformConfig @@ -68,6 +70,13 @@ def mock_asyncio_run_patch(): """Patch event loop run_until_complete for controller step tests.""" # Create a mock event loop with run_until_complete method mock_loop = MagicMock() + + def _run_until_complete(awaitable): + if inspect.iscoroutine(awaitable): + awaitable.close() + return mock_loop.run_until_complete.return_value + + mock_loop.run_until_complete.side_effect = _run_until_complete with patch("nmp.core.models.controllers.models_controller.asyncio.new_event_loop", return_value=mock_loop): yield mock_loop.run_until_complete @@ -76,13 +85,6 @@ def mock_asyncio_run_patch(): def mock_models_sdk(): """Create a mock AsyncNeMoPlatform SDK for testing.""" mock_sdk = MagicMock() - - # Set up the nested structure for v2.inference.deployments - mock_sdk.v2 = MagicMock() - mock_sdk.v2.inference = MagicMock() - mock_sdk.v2.inference.deployments = MagicMock() - mock_sdk.v2.inference.deployments.list = MagicMock() - return mock_sdk @@ -168,11 +170,6 @@ def _assert_asyncio_run_called_once(mock_asyncio_run_patch): assert mock_asyncio_run_patch.call_count == 1 -def _assert_sdk_list_called_for_all_statuses(mock_models_sdk, non_terminal_states_count): - """Assert that SDK list method was called for each non-terminal status.""" - assert mock_models_sdk.inference.deployments.list.call_count == non_terminal_states_count - - def _assert_deployments_count(deployments, expected_count): """Assert the number of deployments returned.""" assert len(deployments) == expected_count @@ -186,7 +183,6 @@ class AssertHelpers: assert_controller_healthy = staticmethod(_assert_controller_healthy) assert_sdk_initialized_correctly = staticmethod(_assert_sdk_initialized_correctly) assert_asyncio_run_called_once = staticmethod(_assert_asyncio_run_called_once) - assert_sdk_list_called_for_all_statuses = staticmethod(_assert_sdk_list_called_for_all_statuses) assert_deployments_count = staticmethod(_assert_deployments_count) @@ -201,21 +197,6 @@ def test_something(assert_helpers): return AssertHelpers -class AsyncPaginator: - """Async iterator standing in for the SDK's paginated list() responses.""" - - def __init__(self, items): - self._items = list(items) - - def __aiter__(self): - return self - - async def __anext__(self): - if not self._items: - raise StopAsyncIteration - return self._items.pop(0) - - class _AsyncPage: """Stand-in for ``AsyncNemoPaginatedResponse``: exposes an async ``items()``.""" @@ -241,14 +222,14 @@ def data(self): 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 - + """Build a plugin client HTTP error for a given status.""" request = httpx.Request("POST", "http://test") response = httpx.Response(status, request=request, json={"detail": detail}) + if status == 404: + return NotFoundError(response) if status == 409: return ConflictError(response) - return NotFoundError(response) + return NemoHTTPError(response) def make_async_models_client() -> MagicMock: @@ -259,6 +240,17 @@ def make_async_models_client() -> MagicMock: client.create_model = AsyncMock(return_value=_ModelResponse()) client.get_model = AsyncMock(return_value=_ModelResponse()) client.update_model = AsyncMock(return_value=_ModelResponse()) + client.list_deployments = AsyncMock(return_value=_AsyncPage([])) + client.get_deployment = AsyncMock(return_value=_ModelResponse()) + client.get_deployment_config_version = AsyncMock(return_value=_ModelResponse()) + client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) + client.delete_deployment_version = AsyncMock(return_value=_ModelResponse()) + client.list_providers = AsyncMock(return_value=_AsyncPage([])) + client.get_provider = AsyncMock(return_value=_ModelResponse()) + client.create_provider = AsyncMock(return_value=_ModelResponse()) + client.upsert_provider = AsyncMock(return_value=_ModelResponse()) + client.update_provider_status = AsyncMock(return_value=_ModelResponse()) + client.delete_provider = AsyncMock(return_value=_ModelResponse()) return client 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 22a114b362..866b2830c3 100644 --- a/services/core/models/tests/unit/controllers/test_deployment_reconciler.py +++ b/services/core/models/tests/unit/controllers/test_deployment_reconciler.py @@ -5,11 +5,11 @@ import logging from datetime import datetime, timedelta, timezone +from enum import Enum from unittest.mock import AsyncMock, MagicMock, patch import pytest from nemo_platform import AsyncNeMoPlatform -from nemo_platform._exceptions import ConflictError, NotFoundError from nmp.core.models.config import ControllerConfig from nmp.core.models.controllers.backends.backends import DeploymentStatusUpdate from nmp.core.models.controllers.backends.registry import BackendRegistry @@ -18,9 +18,13 @@ from nmp.core.models.controllers.entity_cache import ModelEntityCache from nmp.core.models.schemas import ModelDeployment -from .conftest import AsyncPaginator, _ModelResponse, make_async_models_client, make_entity, seed_entity_cache - -_AsyncPaginator = AsyncPaginator +from .conftest import ( + _ModelResponse, + _status_error, + make_async_models_client, + make_entity, + seed_entity_cache, +) def _entity(workspace, name, model_providers): @@ -28,6 +32,35 @@ def _entity(workspace, name, model_providers): return make_entity(workspace, name, model_providers=model_providers) +def _enum_value(value: Enum | str) -> str: + return value.value if isinstance(value, Enum) else value + + +def _deployment_status_call(update_status: AsyncMock, index: int = -1) -> dict[str, str | None]: + call = update_status.call_args_list[index] + body = call.kwargs["body"] + query_params = call.kwargs["query_params"] + return { + "name": call.kwargs["name"], + "workspace": call.kwargs["workspace"], + "version": query_params.get("version"), + "status": _enum_value(body.status), + "status_message": body.status_message, + "model_provider_id": body.model_provider_id, + } + + +def _request_body_call(method: AsyncMock, index: int = -1) -> dict[str, object]: + call = method.call_args_list[index] + values: dict[str, object] = {} + if "workspace" in call.kwargs: + values["workspace"] = call.kwargs["workspace"] + if "name" in call.kwargs: + values["name"] = call.kwargs["name"] + values.update(call.kwargs["body"].model_dump(exclude_unset=True, mode="json")) + return values + + @pytest.fixture def mock_models_sdk(): """Create a mock AsyncNeMoPlatform SDK.""" @@ -40,9 +73,15 @@ def mock_models_sdk(): 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, + with ( + patch( + "nmp.core.models.controllers.entity_cache.client_from_platform", + side_effect=lambda sdk, cls: sdk.models_client, + ), + patch( + "nmp.core.models.controllers.deployment_reconciler.client_from_platform", + side_effect=lambda sdk, cls: sdk.models_client, + ), ): yield @@ -127,7 +166,7 @@ async def test_handle_created_deployment_success(reconciler, mock_backend_regist mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK update_status method - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Call the handler with the backend function await reconciler._reconcile_individual_deployment(deployment, mock_backend.create_model_deployment, "create") @@ -135,15 +174,14 @@ async def test_handle_created_deployment_success(reconciler, mock_backend_regist # Verify backend was called mock_backend.create_model_deployment.assert_called_once_with(deployment) - # Verify SDK update was called - reconciler._models_sdk.inference.deployments.update_status.assert_called_once_with( - name="test-deployment", - workspace="default", - status="PENDING", - version="v1", - status_message="Deployment created", - model_provider_id=None, # No provider created for PENDING status - ) + assert _deployment_status_call(reconciler._models_client.update_deployment_status) == { + "name": "test-deployment", + "workspace": "default", + "status": "PENDING", + "version": "v1", + "status_message": "Deployment created", + "model_provider_id": None, + } @pytest.mark.asyncio @@ -157,7 +195,7 @@ async def test_handle_created_deployment_backend_failure(reconciler, mock_backen mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK update_status method - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Call the handler - should not raise exception await reconciler._reconcile_individual_deployment(deployment, mock_backend.create_model_deployment, "create") @@ -166,8 +204,8 @@ async def test_handle_created_deployment_backend_failure(reconciler, mock_backen mock_backend.create_model_deployment.assert_called_once_with(deployment) # Verify SDK update was called with ERROR status - reconciler._models_sdk.inference.deployments.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + reconciler._models_client.update_deployment_status.assert_called_once() + call_kwargs = _deployment_status_call(reconciler._models_client.update_deployment_status) assert call_kwargs["name"] == "test-deployment" assert call_kwargs["workspace"] == "default" assert call_kwargs["version"] == "v1" @@ -180,7 +218,7 @@ async def test_reconcile_individual_deployment_monitor_ready_no_message_logs_deb """Routine monitor + READY with no status message should log at DEBUG, not INFO.""" deployment = make_deployment(status="READY") status_update = DeploymentStatusUpdate(status="READY", status_message="", host_url=None) - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) with caplog.at_level(logging.DEBUG, logger="nmp.core.models.controllers.deployment_reconciler"): await reconciler._reconcile_individual_deployment( @@ -202,7 +240,7 @@ async def test_reconcile_individual_deployment_monitor_ready_with_message_logs_i """Monitor + READY with a status message stays at INFO.""" deployment = make_deployment(status="READY") status_update = DeploymentStatusUpdate(status="READY", status_message="NIM loading", host_url=None) - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) with caplog.at_level(logging.INFO, logger="nmp.core.models.controllers.deployment_reconciler"): await reconciler._reconcile_individual_deployment( @@ -233,15 +271,13 @@ async def test_reconcile_individual_deployment_conflict_is_noop(reconciler, mock mock_backend_registry.get_backend.return_value = mock_backend # Main status update conflicts (deployment was marked DELETING server-side) - reconciler._models_sdk.inference.deployments.update_status = AsyncMock( - side_effect=ConflictError("Conflict", response=MagicMock(), body=None) - ) + reconciler._models_client.update_deployment_status = AsyncMock(side_effect=_status_error(409, "Conflict")) # Should not raise and should not attempt ERROR update await reconciler._reconcile_individual_deployment(deployment, mock_backend.create_model_deployment, "create") - reconciler._models_sdk.inference.deployments.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + reconciler._models_client.update_deployment_status.assert_called_once() + call_kwargs = _deployment_status_call(reconciler._models_client.update_deployment_status) assert call_kwargs["status"] == "PENDING" @@ -257,15 +293,13 @@ async def test_reconcile_individual_deployment_error_fallback_conflict_is_noop( mock_backend.create_model_deployment = AsyncMock(side_effect=Exception("Backend error")) mock_backend_registry.get_backend.return_value = mock_backend - reconciler._models_sdk.inference.deployments.update_status = AsyncMock( - side_effect=ConflictError("Conflict", response=MagicMock(), body=None) - ) + reconciler._models_client.update_deployment_status = AsyncMock(side_effect=_status_error(409, "Conflict")) # Should not raise if fallback ERROR update hits 409 conflict await reconciler._reconcile_individual_deployment(deployment, mock_backend.create_model_deployment, "create") - reconciler._models_sdk.inference.deployments.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + reconciler._models_client.update_deployment_status.assert_called_once() + call_kwargs = _deployment_status_call(reconciler._models_client.update_deployment_status) assert call_kwargs["status"] == "ERROR" @@ -281,11 +315,11 @@ async def test_reconcile_created_backend_error_persisted(reconciler, mock_backen ) ) mock_backend_registry.get_backend.return_value = mock_backend - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) await reconciler._reconcile_individual_deployment(deployment, mock_backend.create_model_deployment, "create") - call_kwargs = reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + call_kwargs = _deployment_status_call(reconciler._models_client.update_deployment_status) assert call_kwargs["status"] == "ERROR" assert call_kwargs["status_message"] == "Backend create failed for some reason" @@ -325,7 +359,7 @@ async def test_reconcile_deployments_with_created_status(reconciler, mock_backen mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Process deployments (now passing contexts with pre-fetched data) await reconciler.reconcile_deployments([created_context, pending_context]) @@ -339,7 +373,7 @@ async def test_reconcile_deployments_with_created_status(reconciler, mock_backen mock_backend.get_model_deployment_status.assert_called_once_with(pending_context) # Verify SDK update was called twice (once for each deployment) - assert reconciler._models_sdk.inference.deployments.update_status.call_count == 2 + assert reconciler._models_client.update_deployment_status.call_count == 2 # ============================================================================ @@ -351,10 +385,8 @@ async def test_reconcile_deployments_with_created_status(reconciler, mock_backen async def test_ensure_model_provider_creates_when_not_exists(reconciler, make_deployment): """Test that ensure_model_provider creates provider when it doesn't exist.""" # Mock provider doesn't exist (retrieve raises NotFoundError) - reconciler._models_sdk.inference.providers.retrieve = AsyncMock( - side_effect=NotFoundError("Not found", response=MagicMock(), body=None) - ) - reconciler._models_sdk.inference.providers.create = AsyncMock() + reconciler._models_client.get_provider = AsyncMock(side_effect=_status_error(404, "Not found")) + reconciler._models_client.create_provider = AsyncMock() deployment = make_deployment(workspace="test-ns", project="test-project") @@ -365,21 +397,20 @@ async def test_ensure_model_provider_creates_when_not_exists(reconciler, make_de assert model_provider_id == "test-ns/test-deployment" # Verify retrieve was called to check existence - reconciler._models_sdk.inference.providers.retrieve.assert_called_once_with( + reconciler._models_client.get_provider.assert_called_once_with( name="test-deployment", workspace="test-ns", ) - # Verify create was called with correct parameters including model_deployment_id and status - reconciler._models_sdk.inference.providers.create.assert_called_once_with( - workspace="test-ns", - name="test-deployment", - host_url="http://test-ns/test-deployment", - description="Auto-created provider for deployment test-deployment", - project="test-project", - model_deployment_id="test-ns/test-deployment", - status="READY", - ) + assert _request_body_call(reconciler._models_client.create_provider) == { + "workspace": "test-ns", + "name": "test-deployment", + "host_url": "http://test-ns/test-deployment", + "description": "Auto-created provider for deployment test-deployment", + "project": "test-project", + "model_deployment_id": "test-ns/test-deployment", + "status": "READY", + } @pytest.mark.asyncio @@ -393,8 +424,8 @@ async def test_ensure_model_provider_handles_name_collision(mock_uuid, reconcile # Mock provider exists (retrieve succeeds on first call - collision) mock_provider = MagicMock() - reconciler._models_sdk.inference.providers.retrieve = AsyncMock(return_value=mock_provider) - reconciler._models_sdk.inference.providers.create = AsyncMock() + reconciler._models_client.get_provider = AsyncMock(return_value=_ModelResponse(mock_provider)) + reconciler._models_client.create_provider = AsyncMock() deployment = make_deployment(workspace="test-ns", project="test-project") @@ -405,18 +436,17 @@ async def test_ensure_model_provider_handles_name_collision(mock_uuid, reconcile assert model_provider_id == "test-ns/test-deployment_abcdef12" # Verify retrieve was called to check existence - reconciler._models_sdk.inference.providers.retrieve.assert_called_once() + reconciler._models_client.get_provider.assert_called_once() - # Verify create was called with UUID-suffixed name and status - reconciler._models_sdk.inference.providers.create.assert_called_once_with( - workspace="test-ns", - name="test-deployment_abcdef12", - host_url="http://test-ns/test-deployment", - description="Auto-created provider for deployment test-deployment", - project="test-project", - model_deployment_id="test-ns/test-deployment", - status="READY", - ) + assert _request_body_call(reconciler._models_client.create_provider) == { + "workspace": "test-ns", + "name": "test-deployment_abcdef12", + "host_url": "http://test-ns/test-deployment", + "description": "Auto-created provider for deployment test-deployment", + "project": "test-project", + "model_deployment_id": "test-ns/test-deployment", + "status": "READY", + } @pytest.mark.asyncio @@ -427,9 +457,9 @@ async def test_ensure_model_provider_reuses_existing_when_already_set(reconciler mock_provider.host_url = "http://test-ns/test-deployment" mock_provider.description = "Existing provider" mock_provider.enabled_models = None - reconciler._models_sdk.inference.providers.retrieve = AsyncMock(return_value=mock_provider) - reconciler._models_sdk.inference.providers.create = AsyncMock() - reconciler._models_sdk.inference.providers.update = AsyncMock() + reconciler._models_client.get_provider = AsyncMock(return_value=_ModelResponse(mock_provider)) + reconciler._models_client.create_provider = AsyncMock() + reconciler._models_client.upsert_provider = AsyncMock() deployment = make_deployment( workspace="test-ns", project="test-project", model_provider_id="test-ns/existing-provider" @@ -442,16 +472,16 @@ async def test_ensure_model_provider_reuses_existing_when_already_set(reconciler assert model_provider_id == "test-ns/existing-provider" # Verify retrieve was called to check the existing provider exists - reconciler._models_sdk.inference.providers.retrieve.assert_called_once_with( + reconciler._models_client.get_provider.assert_called_once_with( name="existing-provider", workspace="test-ns", ) # Verify create was NOT called since we're reusing existing provider - reconciler._models_sdk.inference.providers.create.assert_not_called() + reconciler._models_client.create_provider.assert_not_called() # Verify update was NOT called since host_url matches - reconciler._models_sdk.inference.providers.update.assert_not_called() + reconciler._models_client.upsert_provider.assert_not_called() @pytest.mark.asyncio @@ -462,9 +492,9 @@ async def test_ensure_model_provider_updates_when_host_url_changes(reconciler, m mock_provider.host_url = "http://old-host/test-deployment" mock_provider.description = "Existing provider" mock_provider.enabled_models = ["model1", "model2"] - reconciler._models_sdk.inference.providers.retrieve = AsyncMock(return_value=mock_provider) - reconciler._models_sdk.inference.providers.create = AsyncMock() - reconciler._models_sdk.inference.providers.update = AsyncMock() + reconciler._models_client.get_provider = AsyncMock(return_value=_ModelResponse(mock_provider)) + reconciler._models_client.create_provider = AsyncMock() + reconciler._models_client.upsert_provider = AsyncMock() deployment = make_deployment( workspace="test-ns", project="test-project", model_provider_id="test-ns/existing-provider" @@ -477,33 +507,30 @@ async def test_ensure_model_provider_updates_when_host_url_changes(reconciler, m assert model_provider_id == "test-ns/existing-provider" # Verify retrieve was called to check the existing provider - reconciler._models_sdk.inference.providers.retrieve.assert_called_once_with( + reconciler._models_client.get_provider.assert_called_once_with( name="existing-provider", workspace="test-ns", ) - # Verify update was called with new host_url, existing metadata, and status - reconciler._models_sdk.inference.providers.update.assert_called_once_with( - name="existing-provider", - workspace="test-ns", - host_url=new_host_url, - description="Existing provider", - enabled_models=["model1", "model2"], - status="READY", - ) + assert _request_body_call(reconciler._models_client.upsert_provider) == { + "name": "existing-provider", + "workspace": "test-ns", + "host_url": new_host_url, + "description": "Existing provider", + "enabled_models": ["model1", "model2"], + "status": "READY", + } # Verify create was NOT called since we're updating existing provider - reconciler._models_sdk.inference.providers.create.assert_not_called() + reconciler._models_client.create_provider.assert_not_called() @pytest.mark.asyncio async def test_ensure_model_provider_creates_new_when_existing_not_found(reconciler, make_deployment): """Test that ensure_model_provider creates new provider when existing provider_id points to non-existent provider.""" # First retrieve (checking existing provider) fails, second retrieve (checking name collision) fails too - reconciler._models_sdk.inference.providers.retrieve = AsyncMock( - side_effect=NotFoundError("Not found", response=MagicMock(), body=None) - ) - reconciler._models_sdk.inference.providers.create = AsyncMock() + reconciler._models_client.get_provider = AsyncMock(side_effect=_status_error(404, "Not found")) + reconciler._models_client.create_provider = AsyncMock() deployment = make_deployment( workspace="test-ns", project="test-project", model_provider_id="test-ns/missing-provider" @@ -516,25 +543,24 @@ async def test_ensure_model_provider_creates_new_when_existing_not_found(reconci assert model_provider_id == "test-ns/test-deployment" # Verify retrieve was called twice (once for existing, once for name collision check) - assert reconciler._models_sdk.inference.providers.retrieve.call_count == 2 + assert reconciler._models_client.get_provider.call_count == 2 - # Verify create was called to create new provider with status - reconciler._models_sdk.inference.providers.create.assert_called_once_with( - workspace="test-ns", - name="test-deployment", - host_url="http://test-ns/test-deployment", - description="Auto-created provider for deployment test-deployment", - project="test-project", - model_deployment_id="test-ns/test-deployment", - status="READY", - ) + assert _request_body_call(reconciler._models_client.create_provider) == { + "workspace": "test-ns", + "name": "test-deployment", + "host_url": "http://test-ns/test-deployment", + "description": "Auto-created provider for deployment test-deployment", + "project": "test-project", + "model_deployment_id": "test-ns/test-deployment", + "status": "READY", + } @pytest.mark.asyncio async def test_delete_model_provider_deletes_when_exists(reconciler, make_deployment): """Test that delete_model_provider deletes provider when it exists.""" # Mock provider exists and delete succeeds - reconciler._models_sdk.inference.providers.delete = AsyncMock() + reconciler._models_client.delete_provider = AsyncMock() # Mock the cleanup method to track if it's called reconciler._cleanup_model_entities_for_provider = AsyncMock() @@ -549,7 +575,7 @@ async def test_delete_model_provider_deletes_when_exists(reconciler, make_deploy ) # Verify delete was called with correct parameters - reconciler._models_sdk.inference.providers.delete.assert_called_once_with( + reconciler._models_client.delete_provider.assert_called_once_with( name="test-deployment", workspace="test-ns", ) @@ -559,9 +585,7 @@ async def test_delete_model_provider_deletes_when_exists(reconciler, make_deploy async def test_delete_model_provider_handles_not_found(reconciler, make_deployment): """Test that delete_model_provider handles NotFoundError gracefully.""" # Mock provider doesn't exist (delete raises NotFoundError) - reconciler._models_sdk.inference.providers.delete = AsyncMock( - side_effect=NotFoundError("Not found", response=MagicMock(), body=None) - ) + reconciler._models_client.delete_provider = AsyncMock(side_effect=_status_error(404, "Not found")) # Mock the cleanup method to track if it's called reconciler._cleanup_model_entities_for_provider = AsyncMock() @@ -577,13 +601,13 @@ async def test_delete_model_provider_handles_not_found(reconciler, make_deployme ) # Verify delete was called - reconciler._models_sdk.inference.providers.delete.assert_called_once() + reconciler._models_client.delete_provider.assert_called_once() @pytest.mark.asyncio async def test_delete_model_provider_skips_when_no_provider_id(reconciler, make_deployment): """Test that delete_model_provider skips deletion when model_provider_id is not set.""" - reconciler._models_sdk.inference.providers.delete = AsyncMock() + reconciler._models_client.delete_provider = AsyncMock() # Mock the cleanup method to track if it's called reconciler._cleanup_model_entities_for_provider = AsyncMock() @@ -596,13 +620,13 @@ async def test_delete_model_provider_skips_when_no_provider_id(reconciler, make_ reconciler._cleanup_model_entities_for_provider.assert_not_called() # Verify delete was NOT called - reconciler._models_sdk.inference.providers.delete.assert_not_called() + reconciler._models_client.delete_provider.assert_not_called() @pytest.mark.asyncio async def test_delete_model_provider_handles_uuid_suffix(reconciler, make_deployment): """Test that delete_model_provider correctly parses provider ID with UUID suffix.""" - reconciler._models_sdk.inference.providers.delete = AsyncMock() + reconciler._models_client.delete_provider = AsyncMock() # Mock the cleanup method to track if it's called reconciler._cleanup_model_entities_for_provider = AsyncMock() @@ -617,7 +641,7 @@ async def test_delete_model_provider_handles_uuid_suffix(reconciler, make_deploy ) # Verify delete was called with UUID-suffixed name - reconciler._models_sdk.inference.providers.delete.assert_called_once_with( + reconciler._models_client.delete_provider.assert_called_once_with( name="test-deployment_abcdef12", workspace="test-ns", ) @@ -627,10 +651,8 @@ async def test_delete_model_provider_handles_uuid_suffix(reconciler, make_deploy async def test_reconcile_model_provider_creates_for_ready_status(reconciler, make_deployment): """Test that reconcile_model_provider creates provider when status is READY.""" # Mock provider doesn't exist - reconciler._models_sdk.inference.providers.retrieve = AsyncMock( - side_effect=NotFoundError("Not found", response=MagicMock(), body=None) - ) - reconciler._models_sdk.inference.providers.create = AsyncMock() + reconciler._models_client.get_provider = AsyncMock(side_effect=_status_error(404, "Not found")) + reconciler._models_client.create_provider = AsyncMock() deployment = make_deployment(workspace="test-ns", project="test-project") @@ -643,14 +665,14 @@ async def test_reconcile_model_provider_creates_for_ready_status(reconciler, mak assert model_provider_id == "test-ns/test-deployment" # Verify create was called - reconciler._models_sdk.inference.providers.create.assert_called_once() + reconciler._models_client.create_provider.assert_called_once() @pytest.mark.asyncio async def test_reconcile_model_provider_deletes_for_deleted_status(reconciler, make_deployment): """Test that reconcile_model_provider deletes provider when status is DELETED or DELETING.""" # Mock provider exists - reconciler._models_sdk.inference.providers.delete = AsyncMock() + reconciler._models_client.delete_provider = AsyncMock() deployment = make_deployment(workspace="test-ns", model_provider_id="test-ns/test-deployment") @@ -661,14 +683,14 @@ async def test_reconcile_model_provider_deletes_for_deleted_status(reconciler, m assert model_provider_id is None # Verify delete was called - reconciler._models_sdk.inference.providers.delete.assert_called_once() + reconciler._models_client.delete_provider.assert_called_once() @pytest.mark.asyncio async def test_reconcile_model_provider_deletes_for_deleting_status(reconciler, make_deployment): """Test that reconcile_model_provider deletes provider when status is DELETING.""" # Mock provider exists - reconciler._models_sdk.inference.providers.delete = AsyncMock() + reconciler._models_client.delete_provider = AsyncMock() deployment = make_deployment(workspace="test-ns", model_provider_id="test-ns/test-deployment") @@ -679,15 +701,15 @@ async def test_reconcile_model_provider_deletes_for_deleting_status(reconciler, assert model_provider_id is None # Verify delete was called - reconciler._models_sdk.inference.providers.delete.assert_called_once() + reconciler._models_client.delete_provider.assert_called_once() @pytest.mark.asyncio async def test_reconcile_model_provider_does_nothing_for_other_statuses(reconciler, make_deployment): """Test that reconcile_model_provider does nothing for statuses other than READY/DELETED/DELETING.""" - reconciler._models_sdk.inference.providers.retrieve = AsyncMock() - reconciler._models_sdk.inference.providers.create = AsyncMock() - reconciler._models_sdk.inference.providers.delete = AsyncMock() + reconciler._models_client.get_provider = AsyncMock() + reconciler._models_client.create_provider = AsyncMock() + reconciler._models_client.delete_provider = AsyncMock() deployment = make_deployment(workspace="test-ns") @@ -707,19 +729,17 @@ async def test_reconcile_model_provider_does_nothing_for_other_statuses(reconcil assert result is None # Verify no provider operations were called - reconciler._models_sdk.inference.providers.retrieve.assert_not_called() - reconciler._models_sdk.inference.providers.create.assert_not_called() - reconciler._models_sdk.inference.providers.delete.assert_not_called() + reconciler._models_client.get_provider.assert_not_called() + reconciler._models_client.create_provider.assert_not_called() + reconciler._models_client.delete_provider.assert_not_called() @pytest.mark.asyncio async def test_reconcile_model_provider_handles_errors_gracefully(reconciler, make_deployment): """Test that reconcile_model_provider handles errors without failing deployment update.""" # Mock provider creation fails with unexpected error - reconciler._models_sdk.inference.providers.retrieve = AsyncMock( - side_effect=NotFoundError("Not found", response=MagicMock(), body=None) - ) - reconciler._models_sdk.inference.providers.create = AsyncMock(side_effect=Exception("API Error")) + reconciler._models_client.get_provider = AsyncMock(side_effect=_status_error(404, "Not found")) + reconciler._models_client.create_provider = AsyncMock(side_effect=Exception("API Error")) deployment = make_deployment(workspace="test-ns", project="test-project") @@ -733,7 +753,7 @@ async def test_reconcile_model_provider_handles_errors_gracefully(reconciler, ma assert result is None # Verify create was attempted - reconciler._models_sdk.inference.providers.create.assert_called_once() + reconciler._models_client.create_provider.assert_called_once() @pytest.mark.asyncio @@ -767,12 +787,10 @@ async def test_full_deployment_lifecycle_with_provider_management(reconciler, mo mock_backend.delete_model_deployment = AsyncMock(return_value=deleted_status) # Mock provider operations - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() - reconciler._models_sdk.inference.providers.retrieve = AsyncMock( - side_effect=NotFoundError("Not found", response=MagicMock(), body=None) - ) - reconciler._models_sdk.inference.providers.create = AsyncMock() - reconciler._models_sdk.inference.providers.delete = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) + reconciler._models_client.get_provider = AsyncMock(side_effect=_status_error(404, "Not found")) + reconciler._models_client.create_provider = AsyncMock() + reconciler._models_client.delete_provider = AsyncMock() deployment = make_deployment(workspace="test-ns", status="CREATED", project="test-project") @@ -780,7 +798,7 @@ async def test_full_deployment_lifecycle_with_provider_management(reconciler, mo await reconciler._reconcile_individual_deployment(deployment, mock_backend.create_model_deployment, "create") # Provider should NOT be created for PENDING status - reconciler._models_sdk.inference.providers.create.assert_not_called() + reconciler._models_client.create_provider.assert_not_called() # Step 2: Deployment becomes READY deployment.status = "PENDING" @@ -788,20 +806,19 @@ async def test_full_deployment_lifecycle_with_provider_management(reconciler, mo deployment, mock_backend.get_model_deployment_status, "check status" ) - # Provider SHOULD be created for READY status with backend-provided host_url and status - reconciler._models_sdk.inference.providers.create.assert_called_once_with( - workspace="test-ns", - name="test-deployment", - host_url="http://test-ns/test-deployment", # From backend's status update - description="Auto-created provider for deployment test-deployment", - project="test-project", - model_deployment_id="test-ns/test-deployment", # New field linking to deployment - status="READY", - ) + assert _request_body_call(reconciler._models_client.create_provider) == { + "workspace": "test-ns", + "name": "test-deployment", + "host_url": "http://test-ns/test-deployment", + "description": "Auto-created provider for deployment test-deployment", + "project": "test-project", + "model_deployment_id": "test-ns/test-deployment", + "status": "READY", + } # Verify the status update for READY included model_provider_id - ready_call = reconciler._models_sdk.inference.deployments.update_status.call_args_list[1] - assert ready_call.kwargs["model_provider_id"] == "test-ns/test-deployment" + ready_call = _deployment_status_call(reconciler._models_client.update_deployment_status, 1) + assert ready_call["model_provider_id"] == "test-ns/test-deployment" # Step 3: Delete deployment (READY -> DELETED) # The deployment should now have the model_provider_id set from when it was READY @@ -812,13 +829,13 @@ async def test_full_deployment_lifecycle_with_provider_management(reconciler, mo ) # Provider SHOULD be deleted for DELETED status - reconciler._models_sdk.inference.providers.delete.assert_called_once_with( + reconciler._models_client.delete_provider.assert_called_once_with( name="test-deployment", workspace="test-ns", ) # Verify all deployment status updates were called - assert reconciler._models_sdk.inference.deployments.update_status.call_count == 3 + assert reconciler._models_client.update_deployment_status.call_count == 3 # ============================================================================ @@ -839,13 +856,13 @@ async def test_handle_deleted_deployment_cleanup_after_grace_period(reconciler, ) # Mock the SDK versions.delete method - reconciler._models_sdk.inference.deployments.versions.delete = AsyncMock() + reconciler._models_client.delete_deployment_version = AsyncMock() # Call handle_deleted_deployment await reconciler._handle_deleted_deployment(deployment) # Verify hard delete was called for the specific version - reconciler._models_sdk.inference.deployments.versions.delete.assert_called_once_with( + reconciler._models_client.delete_deployment_version.assert_called_once_with( name="1", workspace="test-workspace", deployment="test-deployment", @@ -865,13 +882,13 @@ async def test_handle_deleted_deployment_no_cleanup_within_grace_period(reconcil ) # Mock the SDK versions.delete method - reconciler._models_sdk.inference.deployments.versions.delete = AsyncMock() + reconciler._models_client.delete_deployment_version = AsyncMock() # Call handle_deleted_deployment await reconciler._handle_deleted_deployment(deployment) # Verify hard delete was NOT called - reconciler._models_sdk.inference.deployments.versions.delete.assert_not_called() + reconciler._models_client.delete_deployment_version.assert_not_called() @pytest.mark.asyncio @@ -890,13 +907,13 @@ async def test_handle_deleted_deployment_with_naive_datetime(reconciler, make_de ) # Mock the SDK versions.delete method - reconciler._models_sdk.inference.deployments.versions.delete = AsyncMock() + reconciler._models_client.delete_deployment_version = AsyncMock() # Call handle_deleted_deployment - should NOT raise TypeError await reconciler._handle_deleted_deployment(deployment) # Verify hard delete was called for the specific version (deployment is past grace period) - reconciler._models_sdk.inference.deployments.versions.delete.assert_called_once_with( + reconciler._models_client.delete_deployment_version.assert_called_once_with( name="1", workspace="test-workspace", deployment="test-deployment", @@ -925,13 +942,13 @@ async def test_reconcile_deployments_calls_handle_deleted(reconciler, make_deplo ) # Mock the SDK versions.delete method - reconciler._models_sdk.inference.deployments.versions.delete = AsyncMock() + reconciler._models_client.delete_deployment_version = AsyncMock() # Call reconcile_deployments with a list containing the DELETED deployment context await reconciler.reconcile_deployments([deleted_context]) # Verify hard delete was called for the specific version (since it's past grace period) - reconciler._models_sdk.inference.deployments.versions.delete.assert_called_once_with( + reconciler._models_client.delete_deployment_version.assert_called_once_with( name="1", workspace="test-workspace", deployment="deleted-deployment", @@ -953,7 +970,7 @@ async def test_cleanup_model_entities_removes_provider_from_entities(reconciler) MagicMock(model_entity_id="test-ns/model-2"), ] - reconciler._models_sdk.inference.providers.retrieve = AsyncMock(return_value=mock_provider) + reconciler._models_client.get_provider = AsyncMock(return_value=_ModelResponse(mock_provider)) await seed_entity_cache( reconciler._models_sdk, reconciler._entity_cache, @@ -968,7 +985,7 @@ async def test_cleanup_model_entities_removes_provider_from_entities(reconciler) await reconciler._entity_cache.flush() # Verify provider was retrieved - reconciler._models_sdk.inference.providers.retrieve.assert_called_once_with( + reconciler._models_client.get_provider.assert_called_once_with( name="provider-1", workspace="test-ns", ) @@ -988,14 +1005,14 @@ async def test_cleanup_model_entities_no_served_models(reconciler): mock_provider = MagicMock() mock_provider.served_models = [] - reconciler._models_sdk.inference.providers.retrieve = AsyncMock(return_value=mock_provider) + reconciler._models_client.get_provider = AsyncMock(return_value=_ModelResponse(mock_provider)) # Call cleanup await reconciler._cleanup_model_entities_for_provider("test-ns", "provider-1", "test-ns/provider-1") await reconciler._entity_cache.flush() # Verify provider was retrieved - reconciler._models_sdk.inference.providers.retrieve.assert_called_once() + reconciler._models_client.get_provider.assert_called_once() # Verify no model entity operations were performed reconciler._models_sdk.models_client.update_model.assert_not_awaited() @@ -1004,9 +1021,7 @@ async def test_cleanup_model_entities_no_served_models(reconciler): @pytest.mark.asyncio async def test_cleanup_model_entities_provider_not_found(reconciler): """Test that cleanup handles NotFoundError gracefully when provider doesn't exist.""" - reconciler._models_sdk.inference.providers.retrieve = AsyncMock( - side_effect=NotFoundError("Provider not found", response=MagicMock(), body=None) - ) + reconciler._models_client.get_provider = AsyncMock(side_effect=_status_error(404, "Provider not found")) # Call cleanup - should not raise await reconciler._cleanup_model_entities_for_provider("test-ns", "provider-1", "test-ns/provider-1") @@ -1025,7 +1040,7 @@ async def test_cleanup_model_entities_provider_not_in_list(reconciler): MagicMock(model_entity_id="test-ns/model-1"), ] - reconciler._models_sdk.inference.providers.retrieve = AsyncMock(return_value=mock_provider) + reconciler._models_client.get_provider = AsyncMock(return_value=_ModelResponse(mock_provider)) await seed_entity_cache( reconciler._models_sdk, reconciler._entity_cache, @@ -1050,7 +1065,7 @@ async def test_cleanup_model_entities_skips_missing_entity_and_continues(reconci MagicMock(model_entity_id="test-ns/model-2"), ] - reconciler._models_sdk.inference.providers.retrieve = AsyncMock(return_value=mock_provider) + reconciler._models_client.get_provider = AsyncMock(return_value=_ModelResponse(mock_provider)) # Only model-2 exists. await seed_entity_cache( reconciler._models_sdk, @@ -1082,7 +1097,7 @@ async def test_cleanup_model_entities_handles_model_update_failure(reconciler): MagicMock(model_entity_id="test-ns/model-2"), ] - reconciler._models_sdk.inference.providers.retrieve = AsyncMock(return_value=mock_provider) + reconciler._models_client.get_provider = AsyncMock(return_value=_ModelResponse(mock_provider)) await seed_entity_cache( reconciler._models_sdk, reconciler._entity_cache, @@ -1113,7 +1128,7 @@ async def test_cleanup_model_entities_with_null_model_providers(reconciler): MagicMock(model_entity_id="test-ns/model-1"), ] - reconciler._models_sdk.inference.providers.retrieve = AsyncMock(return_value=mock_provider) + reconciler._models_client.get_provider = AsyncMock(return_value=_ModelResponse(mock_provider)) await seed_entity_cache( reconciler._models_sdk, reconciler._entity_cache, @@ -1161,7 +1176,7 @@ async def test_lost_status_triggers_drift_recovery(reconciler, mock_backend_regi mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Process deployment await reconciler.reconcile_deployments([ctx]) @@ -1170,8 +1185,8 @@ async def test_lost_status_triggers_drift_recovery(reconciler, mock_backend_regi mock_backend.create_model_deployment.assert_called_once_with(ctx) # Verify status was updated to PENDING with recovery message - reconciler._models_sdk.inference.deployments.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + reconciler._models_client.update_deployment_status.assert_called_once() + call_kwargs = _deployment_status_call(reconciler._models_client.update_deployment_status) assert call_kwargs["status"] == "PENDING" assert "Recovering deployment" in call_kwargs["status_message"] assert "attempt 1/" in call_kwargs["status_message"] @@ -1205,11 +1220,9 @@ async def test_successful_status_clears_drift_state(reconciler, mock_backend_reg mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() - reconciler._models_sdk.inference.providers.retrieve = AsyncMock( - side_effect=NotFoundError("Not found", response=MagicMock(), body=None) - ) - reconciler._models_sdk.inference.providers.create = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) + reconciler._models_client.get_provider = AsyncMock(side_effect=_status_error(404, "Not found")) + reconciler._models_client.create_provider = AsyncMock() # Process deployment await reconciler.reconcile_deployments([ctx]) @@ -1250,7 +1263,7 @@ async def test_pending_status_preserves_drift_state(reconciler, mock_backend_reg mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Process deployment await reconciler.reconcile_deployments([ctx]) @@ -1295,7 +1308,7 @@ async def test_drift_recovery_max_retries_exceeded(reconciler, mock_backend_regi mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Process deployment await reconciler.reconcile_deployments([ctx]) @@ -1304,8 +1317,8 @@ async def test_drift_recovery_max_retries_exceeded(reconciler, mock_backend_regi mock_backend.create_model_deployment.assert_not_called() # Verify status was updated to ERROR - reconciler._models_sdk.inference.deployments.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + reconciler._models_client.update_deployment_status.assert_called_once() + call_kwargs = _deployment_status_call(reconciler._models_client.update_deployment_status) assert call_kwargs["status"] == "ERROR" assert "failed after 3 attempts" in call_kwargs["status_message"] @@ -1350,7 +1363,7 @@ async def test_drift_recovery_respects_backoff(reconciler, mock_backend_registry mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Process deployment await reconciler.reconcile_deployments([ctx]) @@ -1359,7 +1372,7 @@ async def test_drift_recovery_respects_backoff(reconciler, mock_backend_registry mock_backend.create_model_deployment.assert_not_called() # Verify status was NOT updated (skipped this cycle) - reconciler._models_sdk.inference.deployments.update_status.assert_not_called() + reconciler._models_client.update_deployment_status.assert_not_called() @pytest.mark.asyncio @@ -1407,7 +1420,7 @@ async def test_drift_recovery_proceeds_after_backoff(reconciler, mock_backend_re mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Process deployment await reconciler.reconcile_deployments([ctx]) @@ -1450,7 +1463,7 @@ async def test_drift_recovery_ready_deployment(reconciler, mock_backend_registry mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Process deployment await reconciler.reconcile_deployments([ctx]) @@ -1459,7 +1472,7 @@ async def test_drift_recovery_ready_deployment(reconciler, mock_backend_registry mock_backend.create_model_deployment.assert_called_once() # Verify status message indicates recovery - call_kwargs = reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + call_kwargs = _deployment_status_call(reconciler._models_client.update_deployment_status) assert "Recovering deployment" in call_kwargs["status_message"] @@ -1486,14 +1499,14 @@ async def test_unknown_status_triggers_handler_and_updates_status(reconciler, mo mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Process deployment await reconciler.reconcile_deployments([ctx]) # Verify status was updated to UNKNOWN with attempt info - reconciler._models_sdk.inference.deployments.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + reconciler._models_client.update_deployment_status.assert_called_once() + call_kwargs = _deployment_status_call(reconciler._models_client.update_deployment_status) assert call_kwargs["status"] == "UNKNOWN" assert "attempt 1/" in call_kwargs["status_message"] assert "Unable to determine deployment status" in call_kwargs["status_message"] @@ -1534,14 +1547,14 @@ async def test_unknown_status_max_retries_sets_error(reconciler, mock_backend_re mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Process deployment await reconciler.reconcile_deployments([ctx]) # Verify status was set to ERROR - reconciler._models_sdk.inference.deployments.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + reconciler._models_client.update_deployment_status.assert_called_once() + call_kwargs = _deployment_status_call(reconciler._models_client.update_deployment_status) assert call_kwargs["status"] == "ERROR" assert "Unable to communicate with backend after 3 attempts" in call_kwargs["status_message"] @@ -1581,13 +1594,13 @@ async def test_unknown_status_respects_backoff(reconciler, mock_backend_registry mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) # Process deployment await reconciler.reconcile_deployments([ctx]) # Verify status was NOT updated (in backoff period) - reconciler._models_sdk.inference.deployments.update_status.assert_not_called() + reconciler._models_client.update_deployment_status.assert_not_called() # Verify attempts was NOT incremented assert reconciler._drift_recovery_cache.get_attempts("default/test-deployment") == 1 @@ -1621,11 +1634,9 @@ async def test_unknown_status_clears_on_recovery(reconciler, mock_backend_regist mock_backend_registry.get_backend.return_value = mock_backend # Mock SDK - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() - reconciler._models_sdk.inference.providers.retrieve = AsyncMock( - side_effect=NotFoundError("Not found", response=MagicMock(), body=None) - ) - reconciler._models_sdk.inference.providers.create = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) + reconciler._models_client.get_provider = AsyncMock(side_effect=_status_error(404, "Not found")) + reconciler._models_client.create_provider = AsyncMock() # Process deployment await reconciler.reconcile_deployments([ctx]) @@ -1634,7 +1645,7 @@ async def test_unknown_status_clears_on_recovery(reconciler, mock_backend_regist assert "default/test-deployment" not in reconciler._drift_recovery_cache._states # Verify status was updated to READY - call_kwargs = reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + call_kwargs = _deployment_status_call(reconciler._models_client.update_deployment_status) assert call_kwargs["status"] == "READY" @@ -1741,7 +1752,7 @@ def gc_reconciler(mock_models_sdk, mock_backend_registry): return_value=DeploymentStatusUpdate(status="DELETED", status_message="") ) mock_backend_registry.get_backend.return_value = mock_backend - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) reconciler._delete_model_provider = AsyncMock() return reconciler @@ -1784,8 +1795,8 @@ async def test_gc_triggers_after_ttl(gc_reconciler, mock_backend_registry): mock_backend = mock_backend_registry.get_backend() mock_backend.delete_model_deployment.assert_called_once_with("default", "err-deploy") - gc_reconciler._models_sdk.inference.deployments.update_status.assert_called_once() - call_kw = gc_reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + gc_reconciler._models_client.update_deployment_status.assert_called_once() + call_kw = _deployment_status_call(gc_reconciler._models_client.update_deployment_status) assert call_kw["status"] == "DELETING" assert call_kw["name"] == "err-deploy" assert call_kw["workspace"] == "default" @@ -1813,7 +1824,7 @@ async def test_gc_skips_deployment_with_no_updated_at(gc_reconciler, mock_backen mock_backend = mock_backend_registry.get_backend() mock_backend.delete_model_deployment.assert_not_called() - gc_reconciler._models_sdk.inference.deployments.update_status.assert_not_called() + gc_reconciler._models_client.update_deployment_status.assert_not_called() @pytest.mark.asyncio @@ -1826,8 +1837,8 @@ async def test_gc_backend_delete_failure_still_transitions(gc_reconciler, mock_b await gc_reconciler.gc_error_deployments([dep]) - gc_reconciler._models_sdk.inference.deployments.update_status.assert_called_once() - call_kw = gc_reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + gc_reconciler._models_client.update_deployment_status.assert_called_once() + call_kw = _deployment_status_call(gc_reconciler._models_client.update_deployment_status) assert call_kw["status"] == "DELETING" @@ -1837,9 +1848,7 @@ async def test_gc_status_update_failure_does_not_block_others(gc_reconciler, moc dep1 = _make_error_deployment(name="dep-1") dep2 = _make_error_deployment(name="dep-2") - gc_reconciler._models_sdk.inference.deployments.update_status = AsyncMock( - side_effect=[Exception("version conflict"), None] - ) + gc_reconciler._models_client.update_deployment_status = AsyncMock(side_effect=[Exception("version conflict"), None]) await gc_reconciler.gc_error_deployments([dep1, dep2]) @@ -1863,7 +1872,7 @@ async def test_gc_mixed_ttl_only_expired_cleaned(gc_reconciler, mock_backend_reg mock_backend = mock_backend_registry.get_backend() mock_backend.delete_model_deployment.assert_called_once_with("default", "old-deploy") - gc_reconciler._models_sdk.inference.deployments.update_status.assert_called_once() + gc_reconciler._models_client.update_deployment_status.assert_called_once() @pytest.mark.asyncio @@ -1887,8 +1896,8 @@ async def test_gc_provider_cleanup_failure_is_non_fatal(gc_reconciler, mock_back mock_backend = mock_backend_registry.get_backend() mock_backend.delete_model_deployment.assert_called_once() - gc_reconciler._models_sdk.inference.deployments.update_status.assert_called_once() - call_kw = gc_reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + gc_reconciler._models_client.update_deployment_status.assert_called_once() + call_kw = _deployment_status_call(gc_reconciler._models_client.update_deployment_status) assert call_kw["status"] == "DELETING" assert "Provider cleanup failed" in call_kw["status_message"] @@ -1900,7 +1909,7 @@ async def test_gc_empty_list_no_ops(gc_reconciler, mock_backend_registry): mock_backend = mock_backend_registry.get_backend() mock_backend.delete_model_deployment.assert_not_called() - gc_reconciler._models_sdk.inference.deployments.update_status.assert_not_called() + gc_reconciler._models_client.update_deployment_status.assert_not_called() @pytest.mark.asyncio @@ -1921,7 +1930,7 @@ async def test_gc_custom_ttl_respected(mock_models_sdk, mock_backend_registry): return_value=DeploymentStatusUpdate(status="DELETED", status_message="") ) mock_backend_registry.get_backend.return_value = mock_backend - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) reconciler._delete_model_provider = AsyncMock() within_default_but_past_custom = _make_error_deployment( @@ -1940,7 +1949,7 @@ async def test_gc_status_message_includes_original_error(gc_reconciler, mock_bac await gc_reconciler.gc_error_deployments([dep]) - call_kw = gc_reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + call_kw = _deployment_status_call(gc_reconciler._models_client.update_deployment_status) assert "garbage collected" in call_kw["status_message"] assert "NIM health check timed out after 7200s" in call_kw["status_message"] assert "Original error:" in call_kw["status_message"] @@ -1953,7 +1962,7 @@ async def test_gc_status_message_without_original_error(gc_reconciler, mock_back await gc_reconciler.gc_error_deployments([dep]) - call_kw = gc_reconciler._models_sdk.inference.deployments.update_status.call_args.kwargs + call_kw = _deployment_status_call(gc_reconciler._models_client.update_deployment_status) assert "garbage collected" in call_kw["status_message"] assert "Original error:" not in call_kw["status_message"] @@ -1962,9 +1971,7 @@ async def test_gc_status_message_without_original_error(gc_reconciler, mock_back async def test_gc_not_found_on_status_update_handled(gc_reconciler, mock_backend_registry): """NotFoundError on status update (deployment deleted between query and GC) is handled.""" dep = _make_error_deployment() - gc_reconciler._models_sdk.inference.deployments.update_status = AsyncMock( - side_effect=NotFoundError("Not found", response=MagicMock(), body=None) - ) + gc_reconciler._models_client.update_deployment_status = AsyncMock(side_effect=_status_error(404, "Not found")) # Should not raise await gc_reconciler.gc_error_deployments([dep]) @@ -2010,7 +2017,7 @@ async def test_gc_ttl_boundary_parametrized(mock_models_sdk, mock_backend_regist return_value=DeploymentStatusUpdate(status="DELETED", status_message="") ) mock_backend_registry.get_backend.return_value = mock_backend - reconciler._models_sdk.inference.deployments.update_status = AsyncMock() + reconciler._models_client.update_deployment_status = AsyncMock(return_value=_ModelResponse()) reconciler._delete_model_provider = AsyncMock() dep = _make_error_deployment( @@ -2023,7 +2030,7 @@ async def test_gc_ttl_boundary_parametrized(mock_models_sdk, mock_backend_regist if should_gc: mock_backend.delete_model_deployment.assert_called_once() - reconciler._models_sdk.inference.deployments.update_status.assert_called_once() + reconciler._models_client.update_deployment_status.assert_called_once() else: mock_backend.delete_model_deployment.assert_not_called() - reconciler._models_sdk.inference.deployments.update_status.assert_not_called() + reconciler._models_client.update_deployment_status.assert_not_called() 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 1c7e5d1fc0..9abbc34009 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 @@ -4,7 +4,10 @@ """Unit tests for ModelsController.""" import asyncio +import inspect +import json import threading +from collections.abc import Mapping from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -12,22 +15,26 @@ 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 +from .conftest import _AsyncPage, _ModelResponse, make_async_models_client -class MockAsyncPaginator: - """Mock async paginator to simulate SDK's paginated response.""" +def _filter_status(kwargs: Mapping[str, object]) -> str | None: + query_params = kwargs.get("query_params") + if not isinstance(query_params, Mapping): + return None + raw_filter = query_params.get("filter") + if not isinstance(raw_filter, str): + return None + parsed = json.loads(raw_filter) + if not isinstance(parsed, dict): + return None + status = parsed.get("status") + return status if isinstance(status, str) else None - def __init__(self, items): - self.items = items - def __aiter__(self): - return self - - async def __anext__(self): - if not self.items: - raise StopAsyncIteration - return self.items.pop(0) +def _close_coro(awaitable: object) -> None: + if inspect.iscoroutine(awaitable): + awaitable.close() @pytest.fixture(autouse=True) @@ -37,14 +44,39 @@ def _patch_typed_model_client(mock_models_sdk): ``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() + mock_models_sdk.virtual_models_client = MagicMock() + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock(return_value=_AsyncPage([])) + mock_models_sdk.virtual_models_client.create_virtual_model = AsyncMock(return_value=_ModelResponse()) + mock_models_sdk.virtual_models_client.delete_virtual_model = AsyncMock(return_value=_ModelResponse()) + mock_models_sdk.gateway_provider_client = MagicMock() + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock() + + def _client_from_platform(sdk, cls): + match cls.__name__: + case "AsyncModelsClient": + return sdk.models_client + case "AsyncVirtualModelsClient": + return sdk.virtual_models_client + case "AsyncInferenceGatewayProviderClient": + return sdk.gateway_provider_client + raise AssertionError(f"Unexpected typed client class: {cls.__name__}") + with ( patch( "nmp.core.models.controllers.models_controller.client_from_platform", - side_effect=lambda sdk, cls: sdk.models_client, + side_effect=_client_from_platform, ), patch( "nmp.core.models.controllers.entity_cache.client_from_platform", - side_effect=lambda sdk, cls: sdk.models_client, + side_effect=_client_from_platform, + ), + patch( + "nmp.core.models.controllers.deployment_reconciler.client_from_platform", + side_effect=_client_from_platform, + ), + patch( + "nmp.core.models.controllers.provider_reconciler.client_from_platform", + side_effect=_client_from_platform, ), ): yield @@ -105,8 +137,13 @@ def test_step_with_exception( mock_sdk_class_patch, mock_get_config_patch, mock_asyncio_run_patch, mock_backend_registry, assert_helpers ): """Test step() when an exception occurs.""" + # Mock asyncio.run to raise exception - mock_asyncio_run_patch.side_effect = Exception("Test error") + def _raise_test_error(awaitable): + _close_coro(awaitable) + raise Exception("Test error") + + mock_asyncio_run_patch.side_effect = _raise_test_error controller = ModelsController(backend_registry=mock_backend_registry) @@ -126,10 +163,9 @@ async def test_get_non_terminal_deployments_calls_sdk( mock_get_config_patch, mock_models_sdk, mock_backend_registry, sample_deployment ): """Test that retrieve_non_terminal_deployments calls SDK with correct statuses.""" - # Setup SDK mock responses - SDK returns AsyncPaginator for each call - # Use MagicMock (not AsyncMock) because .list() returns an async iterator, not a coroutine - mock_models_sdk.inference.deployments.list = MagicMock( - side_effect=lambda **kwargs: MockAsyncPaginator([sample_deployment]) + # Setup typed-client mock responses. + mock_models_sdk.models_client.list_deployments = AsyncMock( + side_effect=lambda **kwargs: _AsyncPage([sample_deployment]) ) # Create controller and inject mock SDK @@ -140,7 +176,7 @@ async def test_get_non_terminal_deployments_calls_sdk( deployment_contexts = await controller.retrieve_non_terminal_deployments() # Verify SDK was called for each non-terminal status - assert mock_models_sdk.inference.deployments.list.call_count == len(NON_TERMINAL_STATES) + assert mock_models_sdk.models_client.list_deployments.call_count == len(NON_TERMINAL_STATES) # Verify we got ModelContext objects back assert len(deployment_contexts) > 0 @@ -157,14 +193,13 @@ async def test_get_non_terminal_deployments_handles_sdk_errors( # Setup SDK mock to raise exception on first call, succeed on others def side_effect(**kwargs): - filter_dict = kwargs.get("filter", {}) - status = filter_dict.get("status") + status = _filter_status(kwargs) if status == "CREATED": raise Exception("API Error") - return MockAsyncPaginator([]) + return _AsyncPage([]) # Use MagicMock (not AsyncMock) because .list() returns an async iterator, not a coroutine - mock_models_sdk.inference.deployments.list = MagicMock(side_effect=side_effect) + mock_models_sdk.models_client.list_deployments = AsyncMock(side_effect=side_effect) # Create controller and inject mock SDK with patch("nmp.core.models.controllers.models_controller.get_async_platform_sdk", return_value=mock_models_sdk): @@ -193,16 +228,15 @@ async def test_get_non_terminal_deployments_with_multiple_deployments( # Setup SDK mock responses - return different deployments for each status def list_side_effect(**kwargs): - filter_dict = kwargs.get("filter", {}) - status = filter_dict.get("status") + status = _filter_status(kwargs) if status == "CREATED": - return MockAsyncPaginator([sample_deployment]) + return _AsyncPage([sample_deployment]) elif status == "READY": - return MockAsyncPaginator([sample_deployment_ready]) - return MockAsyncPaginator([]) + return _AsyncPage([sample_deployment_ready]) + return _AsyncPage([]) # Use MagicMock (not AsyncMock) because .list() returns an async iterator, not a coroutine - mock_models_sdk.inference.deployments.list = MagicMock(side_effect=list_side_effect) + mock_models_sdk.models_client.list_deployments = AsyncMock(side_effect=list_side_effect) # Create controller and inject mock SDK with patch("nmp.core.models.controllers.models_controller.get_async_platform_sdk", return_value=mock_models_sdk): @@ -230,7 +264,7 @@ async def test_get_model_providers_calls_sdk(mock_get_config_patch, mock_models_ mock_provider.model_deployment_id = None # Use MagicMock (not AsyncMock) because .list() returns an async iterator, not a coroutine - mock_models_sdk.inference.providers.list = MagicMock(return_value=MockAsyncPaginator([mock_provider])) + mock_models_sdk.models_client.list_providers = AsyncMock(return_value=_AsyncPage([mock_provider])) # Create controller and inject mock SDK with patch("nmp.core.models.controllers.models_controller.get_async_platform_sdk", return_value=mock_models_sdk): @@ -240,7 +274,7 @@ async def test_get_model_providers_calls_sdk(mock_get_config_patch, mock_models_ provider_contexts = await controller.retrieve_model_providers() # Verify SDK was called - mock_models_sdk.inference.providers.list.assert_called_once() + mock_models_sdk.models_client.list_providers.assert_called_once() # Verify we got ModelContext objects back assert provider_contexts is not None @@ -264,15 +298,14 @@ async def test_async_controller_step_calls_reconcilers(mock_get_config_patch, mo # Mock SDK to return deployment only for CREATED status def list_deployments_side_effect(**kwargs): - filter_dict = kwargs.get("filter", {}) - status = filter_dict.get("status") + status = _filter_status(kwargs) if status == "CREATED": - return MockAsyncPaginator([mock_deployment]) - return MockAsyncPaginator([]) + return _AsyncPage([mock_deployment]) + return _AsyncPage([]) # Use MagicMock (not AsyncMock) because .list() returns an async iterator, not a coroutine - mock_models_sdk.inference.deployments.list = MagicMock(side_effect=list_deployments_side_effect) - mock_models_sdk.inference.providers.list = MagicMock(return_value=MockAsyncPaginator([mock_provider])) + mock_models_sdk.models_client.list_deployments = AsyncMock(side_effect=list_deployments_side_effect) + mock_models_sdk.models_client.list_providers = AsyncMock(return_value=_AsyncPage([mock_provider])) # Create controller and inject mock SDK with patch("nmp.core.models.controllers.models_controller.get_async_platform_sdk", return_value=mock_models_sdk): @@ -306,8 +339,8 @@ async def test_async_controller_step_runs_provider_reconciler_with_no_providers( mock_get_config_patch, mock_models_sdk, mock_backend_registry ): """The provider reconciler still runs with an empty list so VM orphan cleanup can execute.""" - mock_models_sdk.inference.deployments.list = MagicMock(return_value=MockAsyncPaginator([])) - mock_models_sdk.inference.providers.list = MagicMock(return_value=MockAsyncPaginator([])) + mock_models_sdk.models_client.list_deployments = AsyncMock(return_value=_AsyncPage([])) + mock_models_sdk.models_client.list_providers = AsyncMock(return_value=_AsyncPage([])) with patch("nmp.core.models.controllers.models_controller.get_async_platform_sdk", return_value=mock_models_sdk): controller = ModelsController(backend_registry=mock_backend_registry) @@ -327,8 +360,8 @@ async def test_async_controller_step_skips_provider_reconciler_when_provider_lis mock_get_config_patch, mock_models_sdk, mock_backend_registry ): """A provider list failure must not look like a successful empty list to cleanup.""" - mock_models_sdk.inference.deployments.list = MagicMock(return_value=MockAsyncPaginator([])) - mock_models_sdk.inference.providers.list = MagicMock(side_effect=RuntimeError("providers unavailable")) + mock_models_sdk.models_client.list_deployments = AsyncMock(return_value=_AsyncPage([])) + mock_models_sdk.models_client.list_providers = AsyncMock(side_effect=RuntimeError("providers unavailable")) with patch("nmp.core.models.controllers.models_controller.get_async_platform_sdk", return_value=mock_models_sdk): controller = ModelsController(backend_registry=mock_backend_registry) @@ -352,7 +385,12 @@ def test_step_handles_cancelled_error( mock_sdk_class_patch, mock_get_config_patch, mock_asyncio_run_patch, mock_backend_registry ): """Test that step() handles CancelledError gracefully without raising.""" - mock_asyncio_run_patch.side_effect = asyncio.CancelledError() + + def _raise_cancelled(awaitable): + _close_coro(awaitable) + raise asyncio.CancelledError() + + mock_asyncio_run_patch.side_effect = _raise_cancelled controller = ModelsController(backend_registry=mock_backend_registry) @@ -583,7 +621,7 @@ async def test_retrieve_error_deployments_calls_sdk(mock_get_config_patch, mock_ mock_deployment = MagicMock() mock_deployment.status = "ERROR" - mock_models_sdk.inference.deployments.list = MagicMock(return_value=MockAsyncPaginator([mock_deployment])) + mock_models_sdk.models_client.list_deployments = AsyncMock(return_value=_AsyncPage([mock_deployment])) with patch("nmp.core.models.controllers.models_controller.get_async_platform_sdk", return_value=mock_models_sdk): controller = ModelsController(backend_registry=mock_backend_registry) @@ -592,11 +630,13 @@ async def test_retrieve_error_deployments_calls_sdk(mock_get_config_patch, mock_ assert len(result) == 1 assert result[0] == mock_deployment - mock_models_sdk.inference.deployments.list.assert_called_once_with( + mock_models_sdk.models_client.list_deployments.assert_called_once_with( workspace="-", - filter={"status": "ERROR"}, - all_versions=True, - page_size=1000, + query_params={ + "filter": json.dumps({"status": "ERROR"}), + "all_versions": True, + "page_size": 1000, + }, ) @@ -605,7 +645,7 @@ async def test_retrieve_error_deployments_handles_sdk_error( mock_get_config_patch, mock_models_sdk, mock_backend_registry ): """Test that retrieve_error_deployments returns empty list on SDK error.""" - mock_models_sdk.inference.deployments.list = MagicMock(side_effect=Exception("API Error")) + mock_models_sdk.models_client.list_deployments = AsyncMock(side_effect=Exception("API Error")) with patch("nmp.core.models.controllers.models_controller.get_async_platform_sdk", return_value=mock_models_sdk): controller = ModelsController(backend_registry=mock_backend_registry) @@ -627,14 +667,13 @@ async def test_async_controller_step_calls_gc(mock_get_config_patch, mock_models def list_side_effect(**kwargs): nonlocal call_count - filter_dict = kwargs.get("filter", {}) - status = filter_dict.get("status") + status = _filter_status(kwargs) if status == "ERROR": - return MockAsyncPaginator([mock_error_deployment]) - return MockAsyncPaginator([]) + return _AsyncPage([mock_error_deployment]) + return _AsyncPage([]) - mock_models_sdk.inference.deployments.list = MagicMock(side_effect=list_side_effect) - mock_models_sdk.inference.providers.list = MagicMock(return_value=MockAsyncPaginator([mock_provider])) + mock_models_sdk.models_client.list_deployments = AsyncMock(side_effect=list_side_effect) + mock_models_sdk.models_client.list_providers = AsyncMock(return_value=_AsyncPage([mock_provider])) with patch("nmp.core.models.controllers.models_controller.get_async_platform_sdk", return_value=mock_models_sdk): controller = ModelsController(backend_registry=mock_backend_registry) @@ -660,8 +699,8 @@ async def test_async_controller_step_skips_gc_when_no_error_deployments( mock_provider = MagicMock() mock_provider.model_deployment_id = None - mock_models_sdk.inference.deployments.list = MagicMock(return_value=MockAsyncPaginator([])) - mock_models_sdk.inference.providers.list = MagicMock(return_value=MockAsyncPaginator([mock_provider])) + mock_models_sdk.models_client.list_deployments = AsyncMock(return_value=_AsyncPage([])) + mock_models_sdk.models_client.list_providers = AsyncMock(return_value=_AsyncPage([mock_provider])) with patch("nmp.core.models.controllers.models_controller.get_async_platform_sdk", return_value=mock_models_sdk): controller = ModelsController(backend_registry=mock_backend_registry) @@ -683,8 +722,8 @@ async def test_async_controller_step_stop_signal_skips_gc( """Test that GC is skipped when stop signal is set before GC runs.""" stop_signal = threading.Event() - mock_models_sdk.inference.deployments.list = MagicMock(return_value=MockAsyncPaginator([])) - mock_models_sdk.inference.providers.list = MagicMock(return_value=MockAsyncPaginator([])) + mock_models_sdk.models_client.list_deployments = AsyncMock(return_value=_AsyncPage([])) + mock_models_sdk.models_client.list_providers = AsyncMock(return_value=_AsyncPage([])) with patch("nmp.core.models.controllers.models_controller.get_async_platform_sdk", return_value=mock_models_sdk): controller = ModelsController(backend_registry=mock_backend_registry, stop_signal=stop_signal) 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 f345ae8400..476ad18eb6 100644 --- a/services/core/models/tests/unit/controllers/test_provider_reconciler.py +++ b/services/core/models/tests/unit/controllers/test_provider_reconciler.py @@ -6,13 +6,13 @@ import json import logging from datetime import datetime, timedelta, timezone +from enum import Enum from unittest.mock import AsyncMock, MagicMock, patch import pytest from nemo_platform import AsyncNeMoPlatform -from nemo_platform._exceptions import APIStatusError, ConflictError, NotFoundError -from nemo_platform.types.inference import ServedModelMapping -from nemo_platform.types.inference.model_provider import ModelProvider +from nemo_platform_plugin.client.errors import NemoHTTPError +from nemo_platform_plugin.models.types import ModelProvider, ModelProviderStatus, ServedModelMapping from nmp.core.models.config import ControllerConfig from nmp.core.models.controllers.context import ModelContext from nmp.core.models.controllers.entity_cache import ModelEntityCache @@ -30,12 +30,11 @@ _is_valid_served_model_entity_id, _resolve_base_backend_model_id, ) -from nmp.core.models.schemas import ModelProviderStatus from .conftest import ( - AsyncPaginator, _AsyncPage, _ModelResponse, + _status_error, make_async_models_client, make_entity, seed_entity_cache, @@ -47,7 +46,44 @@ def _discovery_models_from_ids(ids: list[str]) -> list[dict]: return [{"id": i, "root": None, "parent": None} for i in ids] -_AsyncPaginator = AsyncPaginator +def _enum_value(value: Enum | str) -> str: + return value.value if isinstance(value, Enum) else value + + +def _provider_status_call(update_status: AsyncMock, index: int = -1) -> dict[str, object]: + call = update_status.call_args_list[index] + body = call.kwargs["body"] + values: dict[str, object] = { + "name": call.kwargs["name"], + "workspace": call.kwargs["workspace"], + } + if "served_models" in body.model_fields_set: + values["served_models"] = body.served_models + if "status" in body.model_fields_set: + values["status"] = _enum_value(body.status) + if "status_message" in body.model_fields_set: + values["status_message"] = body.status_message + return values + + +def _request_body_call(method: AsyncMock, index: int = -1) -> dict[str, object]: + call = method.call_args_list[index] + values: dict[str, object] = {} + if "workspace" in call.kwargs: + values["workspace"] = call.kwargs["workspace"] + if "name" in call.kwargs: + values["name"] = call.kwargs["name"] + values.update(call.kwargs["body"].model_dump(exclude_unset=True, mode="json")) + return values + + +def _virtual_model_delete_call(delete_virtual_model: AsyncMock, index: int = -1) -> dict[str, object]: + call = delete_virtual_model.call_args_list[index] + return { + "name": call.kwargs["name"], + "workspace": call.kwargs["workspace"], + "expected_db_version": call.kwargs["query_params"]["expected_db_version"], + } def test_infer_backend_format(): @@ -80,7 +116,8 @@ def _make_discoverable_provider( def _configure_discovery_sdk(mock_models_sdk: MagicMock) -> MagicMock: """Wire mock_models_sdk.with_options to return a discovery-scoped SDK mock.""" discovery_sdk = MagicMock() - discovery_sdk.inference.gateway.provider.get = AsyncMock( + discovery_sdk.gateway_provider_client = MagicMock() + discovery_sdk.gateway_provider_client.get_provider_models = AsyncMock( return_value={"object": "list", "data": [{"id": "model-1"}]} ) mock_models_sdk.with_options = MagicMock(return_value=discovery_sdk) @@ -106,13 +143,13 @@ def controller_config(): def mock_models_sdk(): """Create a mock AsyncNeMoPlatform SDK.""" sdk = MagicMock(spec=AsyncNeMoPlatform) - # virtual_models.create must be an AsyncMock so tests that exercise the full - # reconcile path don't fail when _ensure_passthrough_virtual_model awaits it. - 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_client = make_async_models_client() - sdk.inference.gateway.provider.get = AsyncMock() + sdk.virtual_models_client = MagicMock() + sdk.virtual_models_client.list_virtual_models = AsyncMock(return_value=_AsyncPage([])) + sdk.virtual_models_client.create_virtual_model = AsyncMock(return_value=_ModelResponse()) + sdk.virtual_models_client.delete_virtual_model = AsyncMock(return_value=_ModelResponse()) + sdk.gateway_provider_client = MagicMock() + sdk.gateway_provider_client.get_provider_models = AsyncMock() sdk.with_options = MagicMock(return_value=sdk) return sdk @@ -121,9 +158,26 @@ def mock_models_sdk(): 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, + + def _client_from_platform(sdk, cls): + match cls.__name__: + case "AsyncModelsClient": + return sdk.models_client + case "AsyncVirtualModelsClient": + return sdk.virtual_models_client + case "AsyncInferenceGatewayProviderClient": + return sdk.gateway_provider_client + raise AssertionError(f"Unexpected typed client class: {cls.__name__}") + + with ( + patch( + "nmp.core.models.controllers.entity_cache.client_from_platform", + side_effect=_client_from_platform, + ), + patch( + "nmp.core.models.controllers.provider_reconciler.client_from_platform", + side_effect=_client_from_platform, + ), ): yield @@ -165,7 +219,7 @@ async def reconcile_and_flush(reconciler, entity_cache, provider_contexts): @pytest.mark.asyncio async def test_get_available_models_from_provider_success(reconciler, mock_models_sdk, controller_config): """Test successfully getting models from OpenAI-compliant provider.""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock( + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock( return_value={ "object": "list", "data": [ @@ -181,8 +235,7 @@ async def test_get_available_models_from_provider_success(reconciler, mock_model assert isinstance(result, DiscoverySuccess) assert result.model_ids == ["model-1", "model-2", "model-3"] - mock_models_sdk.inference.gateway.provider.get.assert_called_once_with( - "v1/models", + mock_models_sdk.gateway_provider_client.get_provider_models.assert_called_once_with( workspace="test-ns", name="test-provider", timeout=controller_config.provider_discovery_timeout_seconds, @@ -195,7 +248,7 @@ async def test_get_available_models_from_provider_success(reconciler, mock_model @pytest.mark.asyncio async def test_discover_models_passes_configured_timeout(mock_models_sdk): """Discovery should honor controller_config.provider_discovery_timeout_seconds.""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock( + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock( return_value={"object": "list", "data": [{"id": "model-1"}]} ) config = ControllerConfig(provider_discovery_timeout_seconds=240) @@ -208,8 +261,7 @@ async def test_discover_models_passes_configured_timeout(mock_models_sdk): await reconciler._discover_models(_make_discoverable_provider()) - mock_models_sdk.inference.gateway.provider.get.assert_called_once_with( - "v1/models", + mock_models_sdk.gateway_provider_client.get_provider_models.assert_called_once_with( workspace="test-ns", name="test-provider", timeout=240, @@ -223,7 +275,6 @@ async def test_discover_models_passes_configured_timeout(mock_models_sdk): ( 2, { - "path": "v1/models", "workspace": "test-ns", "name": "test-provider", }, @@ -248,10 +299,9 @@ async def test_discover_models_uses_discovery_sdk_with_configured_retries( mock_models_sdk.with_options.assert_called_once_with(max_retries=max_retries) if expect_get_call_kwargs is None: - discovery_sdk.inference.gateway.provider.get.assert_called_once() + discovery_sdk.gateway_provider_client.get_provider_models.assert_called_once() else: - discovery_sdk.inference.gateway.provider.get.assert_called_once_with( - expect_get_call_kwargs["path"], + discovery_sdk.gateway_provider_client.get_provider_models.assert_called_once_with( workspace=expect_get_call_kwargs["workspace"], name=expect_get_call_kwargs["name"], timeout=config.provider_discovery_timeout_seconds, @@ -262,11 +312,7 @@ async def test_discover_models_uses_discovery_sdk_with_configured_retries( "discovery_side_effect", [ pytest.param( - APIStatusError( - "Error code: 502 - {'detail': 'Backend networking error: Connection refused'}", - response=MagicMock(status_code=502), - body={"detail": "Backend networking error: Connection refused"}, - ), + _status_error(502, "Backend networking error: Connection refused"), id="http_502", ), pytest.param(Exception("Request timed out."), id="network_timeout"), @@ -277,7 +323,7 @@ async def test_discover_models_transient_errors_log_debug_not_warning( reconciler, mock_models_sdk, caplog, discovery_side_effect ): """Transient gateway and network failures during discovery must log at debug, not warning.""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock(side_effect=discovery_side_effect) + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock(side_effect=discovery_side_effect) with caplog.at_level(logging.DEBUG): result = await reconciler._discover_models(_make_discoverable_provider()) @@ -289,7 +335,7 @@ async def test_discover_models_transient_errors_log_debug_not_warning( @pytest.mark.asyncio async def test_get_available_models_from_provider_non_compliant_missing_data(reconciler, mock_models_sdk): """Test provider with non-OpenAI compliant response (missing 'data' field).""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock( + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock( return_value={"object": "list"} # Missing 'data' ) @@ -310,7 +356,7 @@ async def test_get_available_models_from_provider_parses_json_string_response(re """A valid JSON body served without an application/json Content-Type (e.g. some Ollama versions) arrives as a raw string from the SDK. Discovery should still parse it instead of treating it as non-compliant.""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock( + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock( return_value=json.dumps( { "object": "list", @@ -335,7 +381,7 @@ async def test_get_available_models_from_provider_parses_json_string_response(re @pytest.mark.asyncio async def test_get_available_models_from_provider_non_compliant_unparsable_string(reconciler, mock_models_sdk): """A non-JSON string response (genuinely non-compliant) is still rejected.""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock(return_value="not json") + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock(return_value="not json") model_provider = ModelProvider( name="test-provider", @@ -352,7 +398,7 @@ async def test_get_available_models_from_provider_non_compliant_unparsable_strin @pytest.mark.asyncio async def test_get_available_models_from_provider_non_compliant_wrong_type(reconciler, mock_models_sdk): """Test provider with non-dict response.""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock( + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock( return_value=["model-1", "model-2"] # Not a dict ) @@ -371,7 +417,9 @@ async def test_get_available_models_from_provider_non_compliant_wrong_type(recon @pytest.mark.asyncio async def test_get_available_models_from_provider_non_compliant_data_not_list(reconciler, mock_models_sdk): """Test provider with 'data' field that is not a list.""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock(return_value={"object": "list", "data": "not-a-list"}) + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock( + return_value={"object": "list", "data": "not-a-list"} + ) model_provider = ModelProvider( name="test-provider", @@ -388,7 +436,7 @@ async def test_get_available_models_from_provider_non_compliant_data_not_list(re @pytest.mark.asyncio async def test_get_available_models_from_provider_skips_invalid_entries(reconciler, mock_models_sdk): """Test provider response with some invalid model entries.""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock( + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock( return_value={ "object": "list", "data": [ @@ -419,7 +467,7 @@ async def test_get_available_models_from_provider_skips_invalid_entries(reconcil @pytest.mark.asyncio async def test_get_available_models_from_provider_handles_exception(reconciler, mock_models_sdk): """Test that exceptions from provider endpoint return DiscoveryTransientError.""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock(side_effect=Exception("Connection error")) + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock(side_effect=Exception("Connection error")) model_provider = ModelProvider( name="test-provider", @@ -436,14 +484,8 @@ async def test_get_available_models_from_provider_handles_exception(reconciler, @pytest.mark.asyncio async def test_query_available_models_gateway_404_provider_not_in_cache_is_transient(reconciler, mock_models_sdk): """Gateway 404 'Model provider not found' (cache miss) must be treated as transient.""" - mock_response = MagicMock() - mock_response.status_code = 404 - mock_models_sdk.inference.gateway.provider.get = AsyncMock( - side_effect=APIStatusError( - "Error code: 404 - {'detail': 'Model provider not found for test-ns/test-provider'}", - response=mock_response, - body={"detail": "Model provider not found for test-ns/test-provider"}, - ) + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock( + side_effect=_status_error(404, "Model provider not found for test-ns/test-provider") ) model_provider = ModelProvider( @@ -461,14 +503,8 @@ async def test_query_available_models_gateway_404_provider_not_in_cache_is_trans @pytest.mark.asyncio async def test_query_available_models_502_backend_404_is_non_compliant(reconciler, mock_models_sdk): """502 with 'Backend returned 404' means backend has no GET /v1/models — non-compliant.""" - mock_response = MagicMock() - mock_response.status_code = 502 - mock_models_sdk.inference.gateway.provider.get = AsyncMock( - side_effect=APIStatusError( - "Error code: 502 - {'detail': 'Backend returned 404: Not Found'}", - response=mock_response, - body={"detail": "Backend returned 404: Not Found"}, - ) + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock( + side_effect=_status_error(502, "Backend returned 404: Not Found") ) model_provider = ModelProvider( @@ -486,14 +522,8 @@ async def test_query_available_models_502_backend_404_is_non_compliant(reconcile @pytest.mark.asyncio async def test_query_available_models_502_other_detail_is_transient(reconciler, mock_models_sdk): """502 with detail other than 'Backend returned 404' is treated as transient.""" - mock_response = MagicMock() - mock_response.status_code = 502 - mock_models_sdk.inference.gateway.provider.get = AsyncMock( - side_effect=APIStatusError( - "Error code: 502 - {'detail': 'Backend networking error: Connection refused'}", - response=mock_response, - body={"detail": "Backend networking error: Connection refused"}, - ) + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock( + side_effect=_status_error(502, "Backend networking error: Connection refused") ) model_provider = ModelProvider( @@ -511,7 +541,7 @@ async def test_query_available_models_502_other_detail_is_transient(reconciler, @pytest.mark.asyncio async def test_get_available_models_from_provider_empty_list(reconciler, mock_models_sdk): """Test provider with no models.""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock(return_value={"object": "list", "data": []}) + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock(return_value={"object": "list", "data": []}) model_provider = ModelProvider( name="test-provider", @@ -721,6 +751,7 @@ async def test_ensure_model_entity_creates_new_entity(reconciler): await reconciler._ensure_model_entity_for_provider( model_workspace="test-ns", model_name="test-model", + provider=ctx.model_provider, provider_id="test-ns/test-provider", ctx=ctx, ) @@ -765,6 +796,7 @@ async def test_ensure_model_entity_updates_existing_adds_provider(reconciler): await reconciler._ensure_model_entity_for_provider( model_workspace="test-ns", model_name="test-model", + provider=ctx.model_provider, provider_id="test-ns/test-provider", ctx=ctx, ) @@ -805,6 +837,7 @@ async def test_ensure_model_entity_skips_if_provider_already_linked(reconciler): await reconciler._ensure_model_entity_for_provider( model_workspace="test-ns", model_name="test-model", + provider=ctx.model_provider, provider_id="test-ns/test-provider", ctx=ctx, ) @@ -839,6 +872,7 @@ async def test_ensure_model_entity_backfills_missing_backend_format(reconciler): await reconciler._ensure_model_entity_for_provider( model_workspace="test-ns", model_name="anthropic.claude-3-5-sonnet", + provider=ctx.model_provider, provider_id="test-ns/test-provider", ctx=ctx, ) @@ -880,6 +914,7 @@ async def test_ensure_model_entity_adds_artifact_to_existing_without_artifact(re await reconciler._ensure_model_entity_for_provider( model_workspace="test-ns", model_name="test-model", + provider=ctx.model_provider, provider_id="test-ns/test-provider", ctx=ctx, ) @@ -924,6 +959,7 @@ async def test_ensure_model_entity_doesnt_overwrite_existing_artifact(reconciler await reconciler._ensure_model_entity_for_provider( model_workspace="test-ns", model_name="test-model", + provider=ctx.model_provider, provider_id="test-ns/test-provider", ctx=ctx, ) @@ -964,6 +1000,7 @@ async def test_ensure_model_entity_doesnt_overwrite_existing_backend_format(reco await reconciler._ensure_model_entity_for_provider( model_workspace="test-ns", model_name="test-model", + provider=ctx.model_provider, provider_id="test-ns/test-provider", ctx=ctx, ) @@ -1003,6 +1040,7 @@ async def test_ensure_model_entity_handles_null_model_providers(reconciler): await reconciler._ensure_model_entity_for_provider( model_workspace="test-ns", model_name="test-model", + provider=ctx.model_provider, provider_id="test-ns/test-provider", ctx=ctx, ) @@ -1036,6 +1074,7 @@ async def test_ensure_model_entity_handles_create_exception(reconciler): await reconciler._ensure_model_entity_for_provider( model_workspace="test-ns", model_name="test-model", + provider=ctx.model_provider, provider_id="test-ns/test-provider", ctx=ctx, ) @@ -1049,17 +1088,9 @@ async def test_entity_cache_load_failure_propagates_and_stages_nothing(reconcile The controller loads the cache at the start of the phase, so this failure aborts the step before reconciliation runs; see the models controller tests for that. """ - mock_response = MagicMock() - mock_response.status_code = 503 - mock_models_sdk.models_client.list_models = AsyncMock( - side_effect=APIStatusError( - "Service unavailable", - response=mock_response, - body={"detail": "upstream error"}, - ) - ) + mock_models_sdk.models_client.list_models = AsyncMock(side_effect=_status_error(503, "upstream error")) - with pytest.raises(APIStatusError): + with pytest.raises(NemoHTTPError): await reconciler._entity_cache.refresh() await reconciler._entity_cache.flush() @@ -1087,8 +1118,8 @@ async def test_virtual_model_listing_failure_does_not_abort_provider_reconciliat model_entity=None, ) - mock_models_sdk.inference.virtual_models.list = MagicMock(side_effect=Exception("listing unavailable")) - mock_models_sdk.inference.providers.update_status = AsyncMock() + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock(side_effect=Exception("listing unavailable")) + mock_models_sdk.models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object( reconciler, @@ -1099,11 +1130,11 @@ async def test_virtual_model_listing_failure_does_not_abort_provider_reconciliat await entity_cache.flush() # Provider status and entity linking still happened. - mock_models_sdk.inference.providers.update_status.assert_awaited() + mock_models_sdk.models_client.update_provider_status.assert_awaited() 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() + mock_models_sdk.virtual_models_client.create_virtual_model.assert_not_awaited() + mock_models_sdk.virtual_models_client.delete_virtual_model.assert_not_awaited() # ============================================================================ @@ -1128,7 +1159,7 @@ async def test_update_model_providers_success(reconciler): model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object( reconciler, @@ -1145,8 +1176,8 @@ async def test_update_model_providers_success(reconciler): assert mock_ensure.call_count == 2 # Verify provider was updated with served models - reconciler._models_sdk.inference.providers.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + reconciler._models_client.update_provider_status.assert_called_once() + call_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) assert call_kwargs["name"] == "test-provider" assert call_kwargs["workspace"] == "test-ns" assert call_kwargs["status"] == "READY" @@ -1170,7 +1201,7 @@ async def test_update_model_providers_filters_by_enabled_models(reconciler): model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object( reconciler, @@ -1218,7 +1249,7 @@ async def test_ensure_external_entities_retries_after_transient_entity_failure(r model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object( reconciler, @@ -1258,7 +1289,7 @@ async def test_update_model_providers_removes_no_longer_served_models(reconciler model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) # Now only serving model-1 and model-2 (model-3 removed) with patch.object( @@ -1270,7 +1301,7 @@ async def test_update_model_providers_removes_no_longer_served_models(reconciler await reconciler.reconcile_model_providers([ctx]) # Verify only model-1 and model-2 are in final served_models - call_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + call_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) served_model_names = {m.served_model_name for m in call_kwargs["served_models"]} assert served_model_names == {"model-1", "model-2"} @@ -1290,7 +1321,7 @@ async def test_update_model_providers_handles_non_compliant_provider(reconciler) model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) # Provider returns DiscoveryNonCompliant (confirmed non-compliant) with patch.object(reconciler, "_discover_models", return_value=DiscoveryNonCompliant()): @@ -1301,8 +1332,8 @@ async def test_update_model_providers_handles_non_compliant_provider(reconciler) mock_ensure.assert_not_called() # Verify provider was updated with empty served_models and appropriate message - reconciler._models_sdk.inference.providers.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + reconciler._models_client.update_provider_status.assert_called_once() + call_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) assert call_kwargs["served_models"] == [] assert call_kwargs["status"] == "READY" assert "Non-OpenAI compliant" in call_kwargs["status_message"] @@ -1325,7 +1356,7 @@ async def test_update_model_providers_handles_update_exception(reconciler): model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock(side_effect=Exception("Update failed")) + reconciler._models_client.update_provider_status = AsyncMock(side_effect=Exception("Update failed")) with patch.object( reconciler, @@ -1354,7 +1385,7 @@ async def test_update_model_providers_normalizes_model_names(reconciler): model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) # Model with special characters that need normalization with patch.object( @@ -1370,7 +1401,7 @@ async def test_update_model_providers_normalizes_model_names(reconciler): assert "model-with-colons" in str(mock_ensure.call_args) # Normalized # Verify served_models keeps original name - call_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + call_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) served_models = call_kwargs["served_models"] assert len(served_models) == 1 assert served_models[0].served_model_name == "model:with:colons" # Original @@ -1394,7 +1425,7 @@ async def test_update_model_providers_strips_same_workspace_prefix_from_model_id model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) # Backend reports model id as workspace/name (e.g. NIM_SERVED_MODEL_NAME set to workspace/name) with patch.object( @@ -1411,7 +1442,7 @@ async def test_update_model_providers_strips_same_workspace_prefix_from_model_id assert call_kwargs["model_name"] == "qwen-2-5-1-5b" # served_models should have model_entity_id = workspace/name (no duplicate prefix in name) - update_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + update_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) served_models = update_kwargs["served_models"] assert len(served_models) == 1 assert served_models[0].model_entity_id == "test-ns/qwen-2-5-1-5b" @@ -1435,7 +1466,7 @@ async def test_update_model_providers_with_empty_discovery(reconciler): model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object( reconciler, @@ -1449,7 +1480,7 @@ async def test_update_model_providers_with_empty_discovery(reconciler): mock_ensure.assert_not_called() # Verify provider was updated with empty served_models - call_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + call_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) assert call_kwargs["served_models"] == [] assert call_kwargs["status"] == "READY" @@ -1474,7 +1505,7 @@ async def test_update_model_providers_multiple_providers(reconciler): ctx1 = ModelContext(model_provider=provider1) ctx2 = ModelContext(model_provider=provider2) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) async def get_models_side_effect(model_provider: ModelProvider): if model_provider.workspace == "ns1": @@ -1487,7 +1518,7 @@ async def get_models_side_effect(model_provider: ModelProvider): # Verify both providers were processed assert mock_get_models.call_count == 2 - assert reconciler._models_sdk.inference.providers.update_status.call_count == 2 + assert reconciler._models_client.update_provider_status.call_count == 2 @pytest.mark.asyncio @@ -1508,14 +1539,14 @@ async def test_reconcile_preserves_served_models_on_transient_error(reconciler): model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object(reconciler, "_discover_models", return_value=DiscoveryTransientError()): with patch.object(reconciler, "_ensure_model_entity_for_provider") as mock_ensure: await reconciler.reconcile_model_providers([ctx]) # Transient error must not trigger any status update — served_models are preserved implicitly - reconciler._models_sdk.inference.providers.update_status.assert_not_called() + reconciler._models_client.update_provider_status.assert_not_called() mock_ensure.assert_not_called() @@ -1540,7 +1571,7 @@ async def test_reconcile_preserves_served_models_when_deployment_base_id_unresol model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with ( patch.object(reconciler, "_discover_models", return_value=DiscoverySuccess([{"id": "test-ns/base"}])), @@ -1549,7 +1580,7 @@ async def test_reconcile_preserves_served_models_when_deployment_base_id_unresol ): await reconciler.reconcile_model_providers([ctx]) - reconciler._models_sdk.inference.providers.update_status.assert_not_called() + reconciler._models_client.update_provider_status.assert_not_called() mock_ensure.assert_not_called() # WARNING must surface the provider id so operators can correlate with # downstream "model not found" reports during a flaky prefetch tick. @@ -1577,7 +1608,7 @@ async def test_reconcile_clears_served_models_on_confirmed_non_compliant(reconci model_entity=None, ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object(reconciler, "_discover_models", return_value=DiscoveryNonCompliant()): with patch.object(reconciler, "_ensure_model_entity_for_provider") as mock_ensure: @@ -1585,8 +1616,8 @@ async def test_reconcile_clears_served_models_on_confirmed_non_compliant(reconci # Non-compliant must clear served_models mock_ensure.assert_not_called() - reconciler._models_sdk.inference.providers.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + reconciler._models_client.update_provider_status.assert_called_once() + call_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) assert call_kwargs["served_models"] == [] assert call_kwargs["status"] == "READY" assert "Non-OpenAI compliant" in call_kwargs["status_message"] @@ -1610,7 +1641,7 @@ async def test_reconcile_prunes_invalid_served_model_entity_ids_before_update_st ctx = ModelContext(model_provider=provider, model_deployment=None, model_deployment_config=None, model_entity=None) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) bad = ServedModelMapping(model_entity_id="ws/Bad.Name", served_model_name="Bad.Name") good = ServedModelMapping(model_entity_id="ws/model-a", served_model_name="model-a") @@ -1626,12 +1657,12 @@ async def test_reconcile_prunes_invalid_served_model_entity_ids_before_update_st ): await reconciler.reconcile_model_providers([ctx]) - reconciler._models_sdk.inference.providers.update_status.assert_called_once() - emitted = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs["served_models"] + reconciler._models_client.update_provider_status.assert_called_once() + emitted = _provider_status_call(reconciler._models_client.update_provider_status)["served_models"] assert [m.model_entity_id for m in emitted] == ["ws/model-a"] # Passthrough VirtualModel is attempted only for the surviving (non-LoRA) mapping. created_names = { - call.kwargs["name"] for call in reconciler._models_sdk.inference.virtual_models.create.call_args_list + call.kwargs["body"].name for call in reconciler._virtual_models_client.create_virtual_model.call_args_list } assert created_names == {"model-a"} @@ -1655,7 +1686,7 @@ async def test_reconcile_keeps_valid_lora_composite_through_gate(reconciler): model_provider=provider, model_deployment=None, model_deployment_config=config, model_entity=None ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object( reconciler, @@ -1669,12 +1700,12 @@ async def test_reconcile_keeps_valid_lora_composite_through_gate(reconciler): ): await reconciler.reconcile_model_providers([ctx]) - emitted = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs["served_models"] + emitted = _provider_status_call(reconciler._models_client.update_provider_status)["served_models"] eids = {m.model_entity_id for m in emitted} assert eids == {"ws/base", "ws/base&adapters/ws/lora-1"} # Only the base entity gets a passthrough VirtualModel; LoRA is skipped by design. created_names = { - call.kwargs["name"] for call in reconciler._models_sdk.inference.virtual_models.create.call_args_list + call.kwargs["body"].name for call in reconciler._virtual_models_client.create_virtual_model.call_args_list } assert created_names == {"base"} @@ -1727,7 +1758,7 @@ async def test_exception_in_one_provider_does_not_affect_others(reconciler): ctx_bad = ModelContext(model_provider=bad_provider) ctx_good = ModelContext(model_provider=good_provider) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) call_count = 0 @@ -1746,8 +1777,8 @@ async def query_side_effect(provider): # Both providers were attempted assert call_count == 2 # Good provider was still updated successfully - reconciler._models_sdk.inference.providers.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + reconciler._models_client.update_provider_status.assert_called_once() + call_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) assert call_kwargs["workspace"] == "ns-good" assert call_kwargs["status"] == "READY" @@ -1778,7 +1809,7 @@ async def test_created_provider_escalated_to_error_after_threshold(reconciler, _ provider = _make_provider(status=ModelProviderStatus.CREATED, created_at=stale, updated_at=stale) ctx = ModelContext(model_provider=provider) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object( reconciler, @@ -1787,8 +1818,8 @@ async def test_created_provider_escalated_to_error_after_threshold(reconciler, _ ): await reconciler.reconcile_model_providers([ctx]) - reconciler._models_sdk.inference.providers.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + reconciler._models_client.update_provider_status.assert_called_once() + call_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) assert call_kwargs["status"] == "ERROR" assert "connection refused" in call_kwargs["status_message"] @@ -1800,13 +1831,13 @@ async def test_created_provider_not_escalated_before_threshold(reconciler, _make provider = _make_provider(status=ModelProviderStatus.CREATED, created_at=recent, updated_at=recent) ctx = ModelContext(model_provider=provider) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object(reconciler, "_discover_models", return_value=DiscoveryTransientError()): await reconciler.reconcile_model_providers([ctx]) # Should NOT update status — still within grace period - reconciler._models_sdk.inference.providers.update_status.assert_not_called() + reconciler._models_client.update_provider_status.assert_not_called() @pytest.mark.asyncio @@ -1838,7 +1869,7 @@ async def test_error_provider_retried_after_cooldown(reconciler, _make_provider) ) ctx = ModelContext(model_provider=provider) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object( reconciler, @@ -1848,8 +1879,8 @@ async def test_error_provider_retried_after_cooldown(reconciler, _make_provider) await reconciler.reconcile_model_providers([ctx]) # Should update status to bump updated_at for next retry pacing - reconciler._models_sdk.inference.providers.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + reconciler._models_client.update_provider_status.assert_called_once() + call_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) assert call_kwargs["status"] == "ERROR" assert "still down" in call_kwargs["status_message"] @@ -1870,15 +1901,15 @@ async def test_error_provider_transitions_to_lost(reconciler, _make_provider): updated_at=datetime.now(timezone.utc), ) - reconciler._models_sdk.inference.providers.update_status = AsyncMock(return_value=updated_provider) + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse(updated_provider)) with patch.object(reconciler, "_discover_models") as mock_query: await reconciler.reconcile_model_providers([ctx]) # Should transition to LOST without attempting discovery mock_query.assert_not_called() - reconciler._models_sdk.inference.providers.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + reconciler._models_client.update_provider_status.assert_called_once() + call_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) assert call_kwargs["status"] == "LOST" assert "permanently failed" in call_kwargs["status_message"] assert ctx.model_provider is updated_provider @@ -1895,7 +1926,7 @@ async def test_error_provider_recovers_to_ready(reconciler, _make_provider): ) ctx = ModelContext(model_provider=provider) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object( reconciler, @@ -1905,8 +1936,8 @@ async def test_error_provider_recovers_to_ready(reconciler, _make_provider): with patch.object(reconciler, "_ensure_model_entity_for_provider"): await reconciler.reconcile_model_providers([ctx]) - reconciler._models_sdk.inference.providers.update_status.assert_called_once() - call_kwargs = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs + reconciler._models_client.update_provider_status.assert_called_once() + call_kwargs = _provider_status_call(reconciler._models_client.update_provider_status) assert call_kwargs["status"] == "READY" assert len(call_kwargs["served_models"]) == 1 @@ -1917,13 +1948,13 @@ async def test_lost_provider_skipped_entirely(reconciler, _make_provider): provider = _make_provider(status=ModelProviderStatus.LOST) ctx = ModelContext(model_provider=provider) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object(reconciler, "_discover_models") as mock_query: await reconciler.reconcile_model_providers([ctx]) mock_query.assert_not_called() - reconciler._models_sdk.inference.providers.update_status.assert_not_called() + reconciler._models_client.update_provider_status.assert_not_called() @pytest.mark.asyncio @@ -1943,19 +1974,19 @@ async def test_ready_provider_preserves_served_models_on_transient_error(reconci ) ctx = ModelContext(model_provider=provider) - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object(reconciler, "_discover_models", return_value=DiscoveryTransientError()): await reconciler.reconcile_model_providers([ctx]) # Should NOT update status — existing served_models preserved - reconciler._models_sdk.inference.providers.update_status.assert_not_called() + reconciler._models_client.update_provider_status.assert_not_called() @pytest.mark.asyncio async def test_discovery_transient_error_carries_message(reconciler, mock_models_sdk): """DiscoveryTransientError should carry the error message from the gateway.""" - mock_models_sdk.inference.gateway.provider.get = AsyncMock(side_effect=Exception("Connection refused")) + mock_models_sdk.gateway_provider_client.get_provider_models = AsyncMock(side_effect=Exception("Connection refused")) provider = ModelProvider( name="test-provider", @@ -1980,22 +2011,18 @@ async def test_ensure_passthrough_virtual_model_creates_when_not_exists(reconcil """Creates a passthrough VirtualModel with the correct arguments.""" await reconciler._ensure_passthrough_virtual_model("my-ws", "llama-3b", set()) - mock_models_sdk.inference.virtual_models.create.assert_awaited_once_with( - workspace="my-ws", - name="llama-3b", - default_model_entity="my-ws/llama-3b", - autoprovisioned=True, - ) + assert _request_body_call(mock_models_sdk.virtual_models_client.create_virtual_model) == { + "workspace": "my-ws", + "name": "llama-3b", + "default_model_entity": "my-ws/llama-3b", + "autoprovisioned": True, + } @pytest.mark.asyncio async def test_ensure_passthrough_virtual_model_ignores_conflict_error(reconciler, mock_models_sdk): """ConflictError (409) means the VirtualModel already exists — must not propagate.""" - mock_response = MagicMock() - mock_response.status_code = 409 - mock_models_sdk.inference.virtual_models.create = AsyncMock( - side_effect=ConflictError("Conflict", response=mock_response, body={}) - ) + mock_models_sdk.virtual_models_client.create_virtual_model = AsyncMock(side_effect=_status_error(409, "Conflict")) # Should not raise await reconciler._ensure_passthrough_virtual_model("my-ws", "llama-3b", set()) @@ -2004,7 +2031,7 @@ async def test_ensure_passthrough_virtual_model_ignores_conflict_error(reconcile @pytest.mark.asyncio async def test_ensure_passthrough_virtual_model_logs_warning_on_unexpected_error(reconciler, mock_models_sdk, caplog): """Unexpected exceptions are logged as warnings and must not propagate.""" - mock_models_sdk.inference.virtual_models.create = AsyncMock(side_effect=RuntimeError("network timeout")) + mock_models_sdk.virtual_models_client.create_virtual_model = AsyncMock(side_effect=RuntimeError("network timeout")) with caplog.at_level(logging.WARNING): # Should not raise @@ -2035,7 +2062,7 @@ async def test_reconcile_creates_passthrough_virtual_models_for_all_served_model model_entity=None, ) - mock_models_sdk.inference.providers.update_status = AsyncMock() + mock_models_sdk.models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object( reconciler, @@ -2045,13 +2072,15 @@ async def test_reconcile_creates_passthrough_virtual_models_for_all_served_model with patch.object(reconciler, "_ensure_model_entity_for_provider"): await reconciler.reconcile_model_providers([ctx]) - assert mock_models_sdk.inference.virtual_models.create.await_count == 2 - created_names = {call.kwargs["name"] for call in mock_models_sdk.inference.virtual_models.create.call_args_list} + assert mock_models_sdk.virtual_models_client.create_virtual_model.await_count == 2 + created_names = { + call.kwargs["body"].name for call in mock_models_sdk.virtual_models_client.create_virtual_model.call_args_list + } assert created_names == {"model-a", "model-b"} - for call in mock_models_sdk.inference.virtual_models.create.call_args_list: - assert call.kwargs["default_model_entity"] == f"test-ns/{call.kwargs['name']}" + for call in mock_models_sdk.virtual_models_client.create_virtual_model.call_args_list: + assert call.kwargs["body"].default_model_entity == f"test-ns/{call.kwargs['body'].name}" assert call.kwargs["workspace"] == "test-ns" - assert call.kwargs["autoprovisioned"] is True + assert call.kwargs["body"].autoprovisioned is True # ============================================================================ @@ -2113,8 +2142,8 @@ def _provider_context( @pytest.mark.asyncio async def test_reconcile_with_no_providers_deletes_orphaned_autoprovisioned_virtual_model(reconciler, mock_models_sdk): """When the last provider is gone, the final cleanup pass deletes its autoprovisioned VM.""" - mock_models_sdk.inference.virtual_models.list = MagicMock( - return_value=_AsyncPaginator( + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock( + return_value=_AsyncPage( [ _virtual_model( "model-a", @@ -2127,12 +2156,14 @@ async def test_reconcile_with_no_providers_deletes_orphaned_autoprovisioned_virt await reconciler.reconcile_model_providers([]) - mock_models_sdk.inference.virtual_models.list.assert_called_once_with(workspace="-", page_size=200) - mock_models_sdk.inference.virtual_models.delete.assert_awaited_once_with( - name="model-a", - workspace="ws", - expected_db_version=1, + mock_models_sdk.virtual_models_client.list_virtual_models.assert_called_once_with( + workspace="-", query_params={"page_size": 200} ) + assert _virtual_model_delete_call(mock_models_sdk.virtual_models_client.delete_virtual_model) == { + "name": "model-a", + "workspace": "ws", + "expected_db_version": 1, + } @pytest.mark.asyncio @@ -2143,8 +2174,8 @@ async def test_cleanup_keeps_autoprovisioned_virtual_model_served_by_remaining_p ServedModelMapping(model_entity_id="ws/model-a", served_model_name="model-a"), ] ) - mock_models_sdk.inference.virtual_models.list = MagicMock( - return_value=_AsyncPaginator( + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock( + return_value=_AsyncPage( [ _virtual_model( "model-a", @@ -2158,14 +2189,14 @@ async def test_cleanup_keeps_autoprovisioned_virtual_model_served_by_remaining_p vm_snapshot, _ = await reconciler._load_virtual_models() await reconciler._cleanup_orphaned_virtual_models([ctx], vm_snapshot) - mock_models_sdk.inference.virtual_models.delete.assert_not_awaited() + mock_models_sdk.virtual_models_client.delete_virtual_model.assert_not_awaited() @pytest.mark.asyncio async def test_cleanup_keeps_autoprovisioned_virtual_model_without_default_model_entity(reconciler, mock_models_sdk): """An adopted/customized autoprovisioned VM without a default route is not an orphan mismatch.""" - mock_models_sdk.inference.virtual_models.list = MagicMock( - return_value=_AsyncPaginator( + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock( + return_value=_AsyncPage( [ _virtual_model( "model-a", @@ -2179,7 +2210,7 @@ async def test_cleanup_keeps_autoprovisioned_virtual_model_without_default_model vm_snapshot, _ = await reconciler._load_virtual_models() await reconciler._cleanup_orphaned_virtual_models([], vm_snapshot) - mock_models_sdk.inference.virtual_models.delete.assert_not_awaited() + mock_models_sdk.virtual_models_client.delete_virtual_model.assert_not_awaited() @pytest.mark.asyncio @@ -2226,13 +2257,13 @@ async def test_cleanup_uses_virtual_model_updated_at_then_created_at_for_snapsho await reconciler._cleanup_orphaned_virtual_models([], [virtual_model], snapshot_taken_at=snapshot_taken_at) if should_delete: - mock_models_sdk.inference.virtual_models.delete.assert_awaited_once_with( - name="model-a", - workspace="ws", - expected_db_version=1, - ) + assert _virtual_model_delete_call(mock_models_sdk.virtual_models_client.delete_virtual_model) == { + "name": "model-a", + "workspace": "ws", + "expected_db_version": 1, + } else: - mock_models_sdk.inference.virtual_models.delete.assert_not_awaited() + mock_models_sdk.virtual_models_client.delete_virtual_model.assert_not_awaited() @pytest.mark.asyncio @@ -2244,8 +2275,8 @@ async def test_cleanup_lost_provider_does_not_protect_autoprovisioned_virtual_mo ServedModelMapping(model_entity_id="ws/model-a", served_model_name="model-a"), ], ) - mock_models_sdk.inference.virtual_models.list = MagicMock( - return_value=_AsyncPaginator( + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock( + return_value=_AsyncPage( [ _virtual_model( "model-a", @@ -2259,18 +2290,18 @@ async def test_cleanup_lost_provider_does_not_protect_autoprovisioned_virtual_mo vm_snapshot, _ = await reconciler._load_virtual_models() await reconciler._cleanup_orphaned_virtual_models([ctx], vm_snapshot) - mock_models_sdk.inference.virtual_models.delete.assert_awaited_once_with( - name="model-a", - workspace="ws", - expected_db_version=1, - ) + assert _virtual_model_delete_call(mock_models_sdk.virtual_models_client.delete_virtual_model) == { + "name": "model-a", + "workspace": "ws", + "expected_db_version": 1, + } @pytest.mark.asyncio async def test_cleanup_never_deletes_user_created_virtual_model(reconciler, mock_models_sdk): """Only autoprovisioned VirtualModels are eligible for orphan cleanup.""" - mock_models_sdk.inference.virtual_models.list = MagicMock( - return_value=_AsyncPaginator( + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock( + return_value=_AsyncPage( [ _virtual_model( "model-a", @@ -2284,14 +2315,14 @@ async def test_cleanup_never_deletes_user_created_virtual_model(reconciler, mock vm_snapshot, _ = await reconciler._load_virtual_models() await reconciler._cleanup_orphaned_virtual_models([], vm_snapshot) - mock_models_sdk.inference.virtual_models.delete.assert_not_awaited() + mock_models_sdk.virtual_models_client.delete_virtual_model.assert_not_awaited() @pytest.mark.asyncio async def test_cleanup_delete_failure_is_logged_and_non_fatal(reconciler, mock_models_sdk, caplog): """Delete failures are swallowed so the next reconcile cycle can retry.""" - mock_models_sdk.inference.virtual_models.list = MagicMock( - return_value=_AsyncPaginator( + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock( + return_value=_AsyncPage( [ _virtual_model( "model-a", @@ -2301,25 +2332,25 @@ async def test_cleanup_delete_failure_is_logged_and_non_fatal(reconciler, mock_m ] ) ) - mock_models_sdk.inference.virtual_models.delete = AsyncMock(side_effect=RuntimeError("delete failed")) + mock_models_sdk.virtual_models_client.delete_virtual_model = AsyncMock(side_effect=RuntimeError("delete failed")) with caplog.at_level(logging.WARNING): vm_snapshot, _ = await reconciler._load_virtual_models() await reconciler._cleanup_orphaned_virtual_models([], vm_snapshot) - mock_models_sdk.inference.virtual_models.delete.assert_awaited_once_with( - name="model-a", - workspace="ws", - expected_db_version=1, - ) + assert _virtual_model_delete_call(mock_models_sdk.virtual_models_client.delete_virtual_model) == { + "name": "model-a", + "workspace": "ws", + "expected_db_version": 1, + } assert any("Failed to delete orphaned autoprovisioned VirtualModel ws/model-a" in r.message for r in caplog.records) @pytest.mark.asyncio async def test_cleanup_skips_orphaned_virtual_model_without_db_version(reconciler, mock_models_sdk, caplog): """Cleanup must not fall back to an unconditional delete when the listed VM has no version.""" - mock_models_sdk.inference.virtual_models.list = MagicMock( - return_value=_AsyncPaginator( + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock( + return_value=_AsyncPage( [ _virtual_model( "model-a", @@ -2335,7 +2366,7 @@ async def test_cleanup_skips_orphaned_virtual_model_without_db_version(reconcile vm_snapshot, _ = await reconciler._load_virtual_models() await reconciler._cleanup_orphaned_virtual_models([], vm_snapshot) - mock_models_sdk.inference.virtual_models.delete.assert_not_awaited() + mock_models_sdk.virtual_models_client.delete_virtual_model.assert_not_awaited() assert any( "Skipping orphaned autoprovisioned VirtualModel ws/model-a because it has no database version" in r.message for r in caplog.records @@ -2367,7 +2398,7 @@ async def test_deployment_backed_never_calls_ensure_model_entity(reconciler, moc model_entity=None, ) - mock_models_sdk.inference.providers.update_status = AsyncMock() + mock_models_sdk.models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) discovered = [ {"id": "ws/base-entity", "root": "ws/base-entity", "parent": None}, @@ -2378,7 +2409,7 @@ async def test_deployment_backed_never_calls_ensure_model_entity(reconciler, moc await reconciler.reconcile_model_providers([ctx]) mock_ensure.assert_not_called() - call_kwargs = mock_models_sdk.inference.providers.update_status.call_args.kwargs + call_kwargs = _provider_status_call(mock_models_sdk.models_client.update_provider_status) assert len(call_kwargs["served_models"]) == 2 @@ -2428,7 +2459,7 @@ async def test_deployment_backed_links_base_entity_to_provider(reconciler, mock_ reconciler._models_sdk.models_client.list_models = AsyncMock(return_value=_AsyncPage([_base_entity([])])) await reconciler._entity_cache.refresh() - mock_models_sdk.inference.providers.update_status = AsyncMock() + mock_models_sdk.models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object(reconciler, "_discover_models", return_value=DiscoverySuccess(_DEPLOYMENT_DISCOVERED)): await reconciler.reconcile_model_providers([_deployment_backed_ctx()]) await reconciler._entity_cache.flush() @@ -2449,7 +2480,7 @@ async def test_deployment_backed_link_is_idempotent(reconciler, mock_models_sdk) reconciler._models_sdk.models_client.list_models = AsyncMock(return_value=_AsyncPage([entity])) await reconciler._entity_cache.refresh() - mock_models_sdk.inference.providers.update_status = AsyncMock() + mock_models_sdk.models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object(reconciler, "_discover_models", return_value=DiscoverySuccess(_DEPLOYMENT_DISCOVERED)): await reconciler.reconcile_model_providers([_deployment_backed_ctx()]) await reconciler._entity_cache.flush() @@ -2463,7 +2494,7 @@ async def test_deployment_backed_does_not_link_lora_composite_ids(reconciler, mo reconciler._models_sdk.models_client.list_models = AsyncMock(return_value=_AsyncPage([_base_entity([])])) await reconciler._entity_cache.refresh() - mock_models_sdk.inference.providers.update_status = AsyncMock() + mock_models_sdk.models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object(reconciler, "_discover_models", return_value=DiscoverySuccess(_DEPLOYMENT_DISCOVERED)): await reconciler.reconcile_model_providers([_deployment_backed_ctx()]) await reconciler._entity_cache.flush() @@ -2472,7 +2503,7 @@ async def test_deployment_backed_does_not_link_lora_composite_ids(reconciler, mo written = {c.kwargs["name"] for c in update.await_args_list} assert written == {"base-entity"} # The adapter is still routable via the composite served_models mapping. - served = mock_models_sdk.inference.providers.update_status.call_args.kwargs["served_models"] + served = _provider_status_call(mock_models_sdk.models_client.update_provider_status)["served_models"] assert any("&adapters/" in m.model_entity_id for m in served) @@ -2497,7 +2528,7 @@ async def test_reconcile_creates_virtual_models_for_previously_served_models(rec model_entity=None, ) - mock_models_sdk.inference.providers.update_status = AsyncMock() + mock_models_sdk.models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) # Discover old-model (already served) and new-model (new) with patch.object( @@ -2509,8 +2540,10 @@ async def test_reconcile_creates_virtual_models_for_previously_served_models(rec await reconciler.reconcile_model_providers([ctx]) # Both models (old + new) get a VirtualModel create attempt - assert mock_models_sdk.inference.virtual_models.create.await_count == 2 - created_names = {call.kwargs["name"] for call in mock_models_sdk.inference.virtual_models.create.call_args_list} + assert mock_models_sdk.virtual_models_client.create_virtual_model.await_count == 2 + created_names = { + call.kwargs["body"].name for call in mock_models_sdk.virtual_models_client.create_virtual_model.call_args_list + } assert created_names == {"old-model", "new-model"} @@ -2529,12 +2562,12 @@ async def test_reconcile_does_not_create_virtual_models_for_non_compliant_provid model_entity=None, ) - mock_models_sdk.inference.providers.update_status = AsyncMock() + mock_models_sdk.models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object(reconciler, "_discover_models", return_value=DiscoveryNonCompliant()): await reconciler.reconcile_model_providers([ctx]) - mock_models_sdk.inference.virtual_models.create.assert_not_awaited() + mock_models_sdk.virtual_models_client.create_virtual_model.assert_not_awaited() @pytest.mark.asyncio @@ -2555,8 +2588,8 @@ async def test_reconcile_creates_virtual_models_even_when_update_status_fails(re ) # update_status raises — VirtualModel creation must still run - mock_models_sdk.inference.providers.update_status = AsyncMock(side_effect=Exception("service unavailable")) - mock_models_sdk.inference.virtual_models.list = MagicMock(return_value=_AsyncPaginator([])) + mock_models_sdk.models_client.update_provider_status = AsyncMock(side_effect=Exception("service unavailable")) + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock(return_value=_AsyncPage([])) with patch.object( reconciler, @@ -2566,13 +2599,13 @@ async def test_reconcile_creates_virtual_models_even_when_update_status_fails(re with patch.object(reconciler, "_ensure_model_entity_for_provider"): await reconciler.reconcile_model_providers([ctx]) - mock_models_sdk.inference.virtual_models.create.assert_awaited_once_with( - workspace="test-ns", - name="model-x", - default_model_entity="test-ns/model-x", - autoprovisioned=True, - ) - mock_models_sdk.inference.virtual_models.delete.assert_not_awaited() + assert _request_body_call(mock_models_sdk.virtual_models_client.create_virtual_model) == { + "workspace": "test-ns", + "name": "model-x", + "default_model_entity": "test-ns/model-x", + "autoprovisioned": True, + } + mock_models_sdk.virtual_models_client.delete_virtual_model.assert_not_awaited() @pytest.mark.asyncio @@ -2605,12 +2638,12 @@ async def test_deployment_backed_served_models_base_lora_prompt_tuned(reconciler {"id": "qwen-lora-base-lora-e2e-dataset-5c30", "root": "/scratch/loras/...", "parent": "e2e-ws/qwen-lora-base"}, {"id": "qwen-lora-prompt-tuned", "root": "e2e-ws/qwen-lora-base", "parent": None}, ] - reconciler._models_sdk.inference.providers.update_status = AsyncMock() + reconciler._models_client.update_provider_status = AsyncMock(return_value=_ModelResponse()) with patch.object(reconciler, "_discover_models", return_value=DiscoverySuccess(discovered)): await reconciler.reconcile_model_providers([ctx]) - served = reconciler._models_sdk.inference.providers.update_status.call_args.kwargs["served_models"] + served = _provider_status_call(reconciler._models_client.update_provider_status)["served_models"] by_entity_id = {m.model_entity_id: m.served_model_name for m in served} assert by_entity_id["e2e-ws/qwen-lora-base"] == "e2e-ws/qwen-lora-base" lora_entity_id = "e2e-ws/qwen-lora-base&adapters/e2e-ws/qwen-lora-base-lora-e2e-dataset-5c30" @@ -2618,13 +2651,13 @@ async def test_deployment_backed_served_models_base_lora_prompt_tuned(reconciler assert by_entity_id["e2e-ws/qwen-lora-prompt-tuned"] == "qwen-lora-prompt-tuned" assert len(served) == 3 - vm_create_calls = reconciler._models_sdk.inference.virtual_models.create.call_args_list - created_names = {call.kwargs["name"] for call in vm_create_calls} + vm_create_calls = reconciler._virtual_models_client.create_virtual_model.call_args_list + created_names = {call.kwargs["body"].name for call in vm_create_calls} assert created_names == {"qwen-lora-base", "qwen-lora-prompt-tuned"} for call in vm_create_calls: assert call.kwargs["workspace"] == "e2e-ws" - assert call.kwargs["default_model_entity"] == f"e2e-ws/{call.kwargs['name']}" - assert call.kwargs["autoprovisioned"] is True + assert call.kwargs["body"].default_model_entity == f"e2e-ws/{call.kwargs['body'].name}" + assert call.kwargs["body"].autoprovisioned is True def test_handle_model_deployment_provider_base_only(reconciler): @@ -3027,7 +3060,7 @@ async def test_virtual_model_list_is_read_once_per_pass(reconciler, mock_models_ model_entity=None, ) - mock_models_sdk.inference.virtual_models.list = MagicMock(return_value=_AsyncPaginator([])) + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock(return_value=_AsyncPage([])) with patch.object( reconciler, "_discover_models", @@ -3035,19 +3068,21 @@ async def test_virtual_model_list_is_read_once_per_pass(reconciler, mock_models_ ): await reconciler.reconcile_model_providers([ctx]) - mock_models_sdk.inference.virtual_models.list.assert_called_once_with(workspace="-", page_size=200) + mock_models_sdk.virtual_models_client.list_virtual_models.assert_called_once_with( + workspace="-", query_params={"page_size": 200} + ) @pytest.mark.asyncio async def test_existing_virtual_model_is_not_recreated(reconciler, mock_models_sdk): """A VirtualModel already present is left alone instead of re-attempted.""" existing = _virtual_model("m1", workspace="test-ns", default_model_entity="test-ns/m1") - mock_models_sdk.inference.virtual_models.list = MagicMock(return_value=_AsyncPaginator([existing])) + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock(return_value=_AsyncPage([existing])) vm_snapshot, existing_vm_names = await reconciler._load_virtual_models() await reconciler._ensure_passthrough_virtual_model("test-ns", "m1", existing_vm_names) - mock_models_sdk.inference.virtual_models.create.assert_not_awaited() + mock_models_sdk.virtual_models_client.create_virtual_model.assert_not_awaited() assert vm_snapshot == [existing] @@ -3055,12 +3090,12 @@ async def test_existing_virtual_model_is_not_recreated(reconciler, mock_models_s async def test_user_managed_virtual_model_name_is_not_recreated(reconciler, mock_models_sdk): """A name held by a non-autoprovisioned VirtualModel is still treated as taken.""" manual = _virtual_model("m1", workspace="test-ns", default_model_entity="other/thing", autoprovisioned=False) - mock_models_sdk.inference.virtual_models.list = MagicMock(return_value=_AsyncPaginator([manual])) + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock(return_value=_AsyncPage([manual])) _, existing_vm_names = await reconciler._load_virtual_models() await reconciler._ensure_passthrough_virtual_model("test-ns", "m1", existing_vm_names) - mock_models_sdk.inference.virtual_models.create.assert_not_awaited() + mock_models_sdk.virtual_models_client.create_virtual_model.assert_not_awaited() @pytest.mark.asyncio @@ -3071,7 +3106,7 @@ async def test_virtual_model_created_in_pass_is_not_attempted_twice(reconciler, await reconciler._ensure_passthrough_virtual_model("test-ns", "m1", existing_vm_names) await reconciler._ensure_passthrough_virtual_model("test-ns", "m1", existing_vm_names) - mock_models_sdk.inference.virtual_models.create.assert_awaited_once() + mock_models_sdk.virtual_models_client.create_virtual_model.assert_awaited_once() @pytest.mark.asyncio @@ -3091,7 +3126,7 @@ async def test_virtual_model_created_in_pass_is_never_deleted_as_orphan(reconcil ) # Nothing exists up front, so the pass creates the VirtualModel itself. - mock_models_sdk.inference.virtual_models.list = MagicMock(return_value=_AsyncPaginator([])) + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock(return_value=_AsyncPage([])) with patch.object( reconciler, "_discover_models", @@ -3099,24 +3134,22 @@ async def test_virtual_model_created_in_pass_is_never_deleted_as_orphan(reconcil ): await reconciler.reconcile_model_providers([ctx]) - mock_models_sdk.inference.virtual_models.create.assert_awaited_once() - mock_models_sdk.inference.virtual_models.delete.assert_not_awaited() + mock_models_sdk.virtual_models_client.create_virtual_model.assert_awaited_once() + mock_models_sdk.virtual_models_client.delete_virtual_model.assert_not_awaited() @pytest.mark.asyncio async def test_cleanup_tolerates_virtual_model_deleted_concurrently(reconciler, mock_models_sdk): """A VirtualModel removed after the snapshot was taken is not an error.""" - mock_models_sdk.inference.virtual_models.list = MagicMock( - return_value=_AsyncPaginator([_virtual_model("model-a", default_model_entity="ws/model-a")]) - ) - mock_models_sdk.inference.virtual_models.delete = AsyncMock( - side_effect=NotFoundError("gone", response=MagicMock(), body=None) + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock( + return_value=_AsyncPage([_virtual_model("model-a", default_model_entity="ws/model-a")]) ) + mock_models_sdk.virtual_models_client.delete_virtual_model = AsyncMock(side_effect=_status_error(404, "gone")) vm_snapshot, _ = await reconciler._load_virtual_models() await reconciler._cleanup_orphaned_virtual_models([], vm_snapshot) - mock_models_sdk.inference.virtual_models.delete.assert_awaited_once() + mock_models_sdk.virtual_models_client.delete_virtual_model.assert_awaited_once() @pytest.mark.asyncio @@ -3134,15 +3167,15 @@ async def test_provider_skipped_before_discovery_keeps_served_models_unresolved( ctx.model_provider.created_at = datetime.now(timezone.utc) ctx.served_models = [] - mock_models_sdk.inference.virtual_models.list = MagicMock( - return_value=_AsyncPaginator([_virtual_model("model-a", default_model_entity="ws/model-a")]) + mock_models_sdk.virtual_models_client.list_virtual_models = AsyncMock( + return_value=_AsyncPage([_virtual_model("model-a", default_model_entity="ws/model-a")]) ) await reconciler.reconcile_model_providers([ctx]) # Still within the retry cooldown, so discovery never ran and nothing was deleted. assert ctx.served_models is None - mock_models_sdk.inference.virtual_models.delete.assert_not_awaited() + mock_models_sdk.virtual_models_client.delete_virtual_model.assert_not_awaited() @pytest.mark.asyncio