viettelcloud-aiplatform 0.3.0__py3-none-any.whl

viettelcloud/aiplatform/hub/api/model_registry_client.py

@@ -0,0 +1,561 @@
+# Copyright 2025 The Kubeflow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterator, Mapping
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    from model_registry.types import (
+        ModelArtifact,
+        ModelVersion,
+        RegisteredModel,
+        SupportedTypes,
+    )
+
+
+class ModelRegistryClient:
+    """Client for Model Registry operations.
+
+    Supports two modes of operation:
+
+    1. **Direct mode**: Connect directly to a Model Registry server.
+       Requires the model-registry package to be installed:
+       ``pip install 'viettelcloud-aiplatform[hub]'``
+
+    2. **Proxy mode**: Connect through cmp-backend which handles authentication
+       and routes to the correct Model Registry instance based on project.
+       Does not require the model-registry package.
+
+    Examples:
+        Direct mode::
+
+            client = ModelRegistryClient(
+                base_url="https://registry.example.com",
+                user_token="...",
+            )
+
+        Proxy mode::
+
+            client = ModelRegistryClient(
+                cmp_backend_url="https://api.viettelcloud.vn",
+                pat_token="vtp_xxx...",
+                project_id="my-ml-project",
+                region="HN",  # or "HCM", etc.
+            )
+
+        Proxy mode with environment variables::
+
+            # Set these environment variables:
+            # VIETTELCLOUD_CMP_URL=https://api.viettelcloud.vn
+            # VIETTELCLOUD_PAT_TOKEN=vtp_xxx...
+            # VIETTELCLOUD_PROJECT_ID=my-ml-project
+            # VIETTELCLOUD_REGION=HN
+
+            client = ModelRegistryClient()  # Uses env vars
+    """
+
+    def __init__(
+        self,
+        base_url: str | None = None,
+        port: int | None = None,
+        *,
+        author: str | None = None,
+        is_secure: bool | None = None,
+        user_token: str | None = None,
+        custom_ca: str | None = None,
+        # Proxy mode parameters
+        cmp_backend_url: str | None = None,
+        pat_token: str | None = None,
+        project_id: str | None = None,
+        region: str | None = None,
+    ):
+        """Initialize the ModelRegistryClient.
+
+        Two modes of operation:
+
+        1. **Direct mode** (existing): Connect directly to MR server
+
+           - Requires: ``base_url``
+           - Optional: ``port``, ``author``, ``is_secure``, ``user_token``, ``custom_ca``
+
+        2. **Proxy mode** (new): Connect through cmp-backend
+
+           - Requires: ``cmp_backend_url``, ``pat_token``, ``project_id``
+           - cmp-backend handles authentication and routing to correct MR instance
+           - Supports environment variables for configuration
+
+        Args:
+            base_url: Base URL of the model registry server including scheme.
+                     Examples: "https://registry.example.com", "http://localhost"
+                     Required for direct mode.
+
+        Keyword Args:
+            port: Server port (direct mode). If not provided, inferred from scheme:
+                 - https:// defaults to 443
+                 - http:// defaults to 8080
+            author: Name of the author.
+            is_secure: Whether to use TLS (direct mode). Inferred from base_url if not provided.
+            user_token: The PEM-encoded user token (direct mode).
+            custom_ca: Path to PEM-encoded root certificates (direct mode).
+            cmp_backend_url: URL of cmp-backend (proxy mode).
+                            Can also be set via ``VIETTELCLOUD_CMP_URL`` env var.
+            pat_token: Personal Access Token (proxy mode).
+                      Can also be set via ``VIETTELCLOUD_PAT_TOKEN`` env var.
+            project_id: Project UUID or slug (proxy mode).
+                       Can also be set via ``VIETTELCLOUD_PROJECT_ID`` env var.
+            region: Region UUID or name (proxy mode). E.g., "HN", "HCM".
+                   Can also be set via ``VIETTELCLOUD_REGION`` env var.
+
+        Raises:
+            ValueError: If neither direct nor proxy mode parameters are provided.
+            ImportError: If model-registry is not installed (direct mode only).
+        """
+        # Support environment variables for proxy mode
+        cmp_backend_url = cmp_backend_url or os.environ.get("VIETTELCLOUD_CMP_URL")
+        pat_token = pat_token or os.environ.get("VIETTELCLOUD_PAT_TOKEN")
+        project_id = project_id or os.environ.get("VIETTELCLOUD_PROJECT_ID")
+        region = region or os.environ.get("VIETTELCLOUD_REGION")
+
+        self._author = author
+
+        # Determine mode
+        if cmp_backend_url and pat_token and project_id and region:
+            # Proxy mode
+            self._mode: Literal["direct", "proxy"] = "proxy"
+            self._registry = None
+
+            from ._proxy_client import ProxyHTTPClient
+
+            self._proxy_client = ProxyHTTPClient(
+                cmp_backend_url=cmp_backend_url,
+                pat_token=pat_token,
+                project_id=project_id,
+                region=region,
+            )
+
+        elif base_url:
+            # Direct mode
+            self._mode = "direct"
+            self._proxy_client = None
+
+            try:
+                from model_registry import ModelRegistry
+            except ImportError as e:
+                raise ImportError(
+                    "model-registry is not installed. Install it with:\n\n"
+                    "  pip install 'viettelcloud-aiplatform[hub]'\n"
+                ) from e
+
+            is_http = base_url.startswith("http://")
+            if is_secure is None:
+                is_secure = not is_http
+            if port is None:
+                port = 8080 if is_http else 443
+
+            self._registry = ModelRegistry(
+                server_address=base_url,
+                port=port,
+                author=author,  # type: ignore[arg-type]
+                is_secure=is_secure,
+                user_token=user_token,
+                custom_ca=custom_ca,
+            )
+
+        else:
+            raise ValueError(
+                "Must provide either:\n"
+                "  - base_url (direct mode), or\n"
+                "  - cmp_backend_url + pat_token + project_id + region (proxy mode)\n\n"
+                "Proxy mode also supports environment variables:\n"
+                "  VIETTELCLOUD_CMP_URL, VIETTELCLOUD_PAT_TOKEN, VIETTELCLOUD_PROJECT_ID, VIETTELCLOUD_REGION"
+            )
+
+    @property
+    def mode(self) -> Literal["direct", "proxy"]:
+        """Return the client mode ('direct' or 'proxy')."""
+        return self._mode
+
+    # =========================================================================
+    # Helper methods for proxy mode
+    # =========================================================================
+
+    def _dict_to_registered_model(self, data: dict) -> RegisteredModel:
+        """Convert proxy API response dict to RegisteredModel."""
+        from model_registry.types import RegisteredModel
+
+        return RegisteredModel(
+            name=data.get("name", ""),
+            owner=data.get("owner"),
+            description=data.get("description"),
+            external_id=data.get("externalId"),
+            id=data.get("id"),
+            create_time_since_epoch=data.get("createTimeSinceEpoch"),
+            last_update_time_since_epoch=data.get("lastUpdateTimeSinceEpoch"),
+            custom_properties=data.get("customProperties"),
+        )
+
+    def _dict_to_model_version(self, data: dict) -> ModelVersion:
+        """Convert proxy API response dict to ModelVersion."""
+        from model_registry.types import ModelVersion
+
+        return ModelVersion(
+            name=data.get("name", ""),
+            author=data.get("author"),
+            description=data.get("description"),
+            external_id=data.get("externalId"),
+            id=data.get("id"),
+            create_time_since_epoch=data.get("createTimeSinceEpoch"),
+            last_update_time_since_epoch=data.get("lastUpdateTimeSinceEpoch"),
+            custom_properties=data.get("customProperties"),
+            registered_model_id=data.get("registeredModelId"),
+        )
+
+    def _dict_to_model_artifact(self, data: dict) -> ModelArtifact:
+        """Convert proxy API response dict to ModelArtifact."""
+        from model_registry.types import ModelArtifact
+
+        return ModelArtifact(
+            name=data.get("name", ""),
+            uri=data.get("uri", ""),
+            description=data.get("description"),
+            external_id=data.get("externalId"),
+            id=data.get("id"),
+            create_time_since_epoch=data.get("createTimeSinceEpoch"),
+            last_update_time_since_epoch=data.get("lastUpdateTimeSinceEpoch"),
+            custom_properties=data.get("customProperties"),
+            model_format_name=data.get("modelFormatName"),
+            model_format_version=data.get("modelFormatVersion"),
+            storage_key=data.get("storageKey"),
+            storage_path=data.get("storagePath"),
+            service_account_name=data.get("serviceAccountName"),
+        )
+
+    # =========================================================================
+    # Public API
+    # =========================================================================
+
+    def register_model(
+        self,
+        name: str,
+        uri: str,
+        *,
+        version: str,
+        model_format_name: str | None = None,
+        model_format_version: str | None = None,
+        author: str | None = None,
+        owner: str | None = None,
+        version_description: str | None = None,
+        metadata: Mapping[str, SupportedTypes] | None = None,
+    ) -> RegisteredModel:
+        """Register a model.
+
+        This registers a model in the model registry. The model is not downloaded,
+        and has to be stored prior to registration.
+
+        Args:
+            name: Name of the model.
+            uri: URI of the model.
+
+        Keyword Args:
+            version: Version of the model. Has to be unique.
+            model_format_name: Name of the model format (e.g., "pytorch", "tensorflow", "onnx").
+            model_format_version: Version of the model format (e.g., "2.0", "1.15").
+            author: Author of the model. Defaults to the client author.
+            owner: Owner of the model. Defaults to the client author.
+            version_description: Description of the model version.
+            metadata: Additional version metadata.
+
+        Returns:
+            Registered model.
+        """
+        if self._mode == "proxy":
+            # Create model
+            model_data = self._proxy_client.create_registered_model(
+                name=name,
+                owner=owner or self._author,
+                custom_properties=dict(metadata) if metadata else None,
+            )
+            # Create version
+            self._proxy_client.create_model_version(
+                model_name=name,
+                version_name=version,
+                description=version_description,
+                author=author or self._author,
+            )
+            # Create artifact
+            self._proxy_client.create_model_artifact(
+                name=f"{name}-{version}",
+                uri=uri,
+                model_format_name=model_format_name,
+                model_format_version=model_format_version,
+            )
+            return self._dict_to_registered_model(model_data)
+        else:
+            return self._registry.register_model(
+                name=name,
+                uri=uri,
+                model_format_name=model_format_name,  # type: ignore[arg-type]
+                model_format_version=model_format_version,  # type: ignore[arg-type]
+                version=version,
+                author=author,
+                owner=owner,
+                description=version_description,
+                metadata=metadata,
+            )
+
+    def update_model(self, model: RegisteredModel) -> RegisteredModel:
+        """Update a registered model.
+
+        Args:
+            model: The registered model to update. Must have an ID.
+
+        Returns:
+            Updated registered model.
+
+        Raises:
+            TypeError: If model is not a RegisteredModel instance.
+            model_registry.exceptions.StoreError: If model does not have an ID.
+        """
+        from model_registry.types import RegisteredModel
+
+        if not isinstance(model, RegisteredModel):
+            raise TypeError(f"Expected RegisteredModel, got {type(model).__name__}. ")
+
+        if self._mode == "proxy":
+            if not model.id:
+                raise ValueError("Model must have an ID to update")
+            data = self._proxy_client.update_registered_model(
+                model_id=model.id,
+                description=model.description,
+                custom_properties=model.custom_properties,
+            )
+            return self._dict_to_registered_model(data)
+        else:
+            return self._registry.update(model)
+
+    def update_model_version(self, model_version: ModelVersion) -> ModelVersion:
+        """Update a model version.
+
+        Args:
+            model_version: The model version to update. Must have an ID.
+
+        Returns:
+            Updated model version.
+
+        Raises:
+            TypeError: If model_version is not a ModelVersion instance.
+            model_registry.exceptions.StoreError: If model version does not have an ID.
+        """
+        from model_registry.types import ModelVersion
+
+        if not isinstance(model_version, ModelVersion):
+            raise TypeError(f"Expected ModelVersion, got {type(model_version).__name__}. ")
+
+        if self._mode == "proxy":
+            # Proxy mode requires model name and version name
+            # This is a limitation - need to fetch model first or change API
+            raise NotImplementedError(
+                "update_model_version in proxy mode requires model name. "
+                "Use the proxy client directly for this operation."
+            )
+        else:
+            return self._registry.update(model_version)
+
+    def update_model_artifact(self, model_artifact: ModelArtifact) -> ModelArtifact:
+        """Update a model artifact.
+
+        Args:
+            model_artifact: The model artifact to update. Must have an ID.
+
+        Returns:
+            Updated model artifact.
+
+        Raises:
+            TypeError: If model_artifact is not a ModelArtifact instance.
+            model_registry.exceptions.StoreError: If model artifact does not have an ID.
+        """
+        from model_registry.types import ModelArtifact
+
+        if not isinstance(model_artifact, ModelArtifact):
+            raise TypeError(f"Expected ModelArtifact, got {type(model_artifact).__name__}. ")
+
+        if self._mode == "proxy":
+            if not model_artifact.id:
+                raise ValueError("ModelArtifact must have an ID to update")
+            data = self._proxy_client.update_model_artifact(
+                artifact_id=model_artifact.id,
+                description=model_artifact.description,
+                custom_properties=model_artifact.custom_properties,
+            )
+            return self._dict_to_model_artifact(data)
+        else:
+            return self._registry.update(model_artifact)
+
+    def get_model(self, name: str) -> RegisteredModel:
+        """Get a registered model.
+
+        Args:
+            name: Name of the model.
+
+        Returns:
+            Registered model.
+
+        Raises:
+            ValueError: If the model does not exist.
+        """
+        if self._mode == "proxy":
+            try:
+                data = self._proxy_client.get_registered_model(name)
+                return self._dict_to_registered_model(data)
+            except Exception as e:
+                if "404" in str(e) or "not found" in str(e).lower():
+                    raise ValueError(f"Model {name!r} not found") from e
+                raise
+        else:
+            model = self._registry.get_registered_model(name)
+            if model is None:
+                raise ValueError(f"Model {name!r} not found")
+            return model
+
+    def get_model_version(self, name: str, version: str) -> ModelVersion:
+        """Get a model version.
+
+        Args:
+            name: Name of the model.
+            version: Version of the model.
+
+        Returns:
+            Model version.
+
+        Raises:
+            model_registry.exceptions.StoreError: If the model does not exist.
+            ValueError: If the version does not exist.
+        """
+        if self._mode == "proxy":
+            try:
+                data = self._proxy_client.get_model_version(name, version)
+                return self._dict_to_model_version(data)
+            except Exception as e:
+                if "404" in str(e) or "not found" in str(e).lower():
+                    raise ValueError(f"Model version {version!r} not found for model {name!r}") from e
+                raise
+        else:
+            model_version = self._registry.get_model_version(name, version)
+            if model_version is None:
+                raise ValueError(f"Model version {version!r} not found for model {name!r}")
+            return model_version
+
+    def get_model_artifact(self, name: str, version: str) -> ModelArtifact:
+        """Get a model artifact.
+
+        Args:
+            name: Name of the model.
+            version: Version of the model.
+
+        Returns:
+            Model artifact.
+
+        Raises:
+            model_registry.exceptions.StoreError: If either the model or the version don't exist.
+            ValueError: If the artifact does not exist.
+        """
+        if self._mode == "proxy":
+            # In proxy mode, we need artifact ID, not model name/version
+            # This is a limitation of the current API design
+            raise NotImplementedError(
+                "get_model_artifact by name/version not supported in proxy mode. "
+                "Use get_model_artifact_by_id() instead."
+            )
+        else:
+            artifact = self._registry.get_model_artifact(name, version)
+            if artifact is None:
+                raise ValueError(f"Model artifact not found for model {name!r} version {version!r}")
+            return artifact
+
+    def get_model_artifact_by_id(self, artifact_id: str) -> ModelArtifact:
+        """Get a model artifact by ID.
+
+        This method is available in proxy mode.
+
+        Args:
+            artifact_id: ID of the artifact.
+
+        Returns:
+            Model artifact.
+        """
+        if self._mode == "proxy":
+            data = self._proxy_client.get_model_artifact(artifact_id)
+            return self._dict_to_model_artifact(data)
+        else:
+            # Direct mode doesn't have this method in the library
+            raise NotImplementedError(
+                "get_model_artifact_by_id not available in direct mode. "
+                "Use get_model_artifact(name, version) instead."
+            )
+
+    def list_models(self) -> Iterator[RegisteredModel]:
+        """Get an iterator for registered models.
+
+        Yields:
+            Registered models.
+        """
+        if self._mode == "proxy":
+            next_page_token = None
+            while True:
+                result = self._proxy_client.list_registered_models(
+                    page_size=100,
+                    next_page_token=next_page_token,
+                )
+                items = result.get("items", [])
+                for item in items:
+                    yield self._dict_to_registered_model(item)
+
+                next_page_token = result.get("nextPageToken")
+                if not next_page_token:
+                    break
+        else:
+            yield from self._registry.get_registered_models()
+
+    def list_model_versions(self, name: str) -> Iterator[ModelVersion]:
+        """Get an iterator for model versions.
+
+        Args:
+            name: Name of the model.
+
+        Yields:
+            Model versions.
+
+        Raises:
+            model_registry.exceptions.StoreError: If the model does not exist.
+        """
+        if self._mode == "proxy":
+            next_page_token = None
+            while True:
+                result = self._proxy_client.list_model_versions(
+                    model_name=name,
+                    page_size=100,
+                    next_page_token=next_page_token,
+                )
+                items = result.get("items", [])
+                for item in items:
+                    yield self._dict_to_model_version(item)
+
+                next_page_token = result.get("nextPageToken")
+                if not next_page_token:
+                    break
+        else:
+            yield from self._registry.get_model_versions(name)