chore: add _PublisherModel class in preview module

PiperOrigin-RevId: 528939747

Author: jaycee-li

google/cloud/aiplatform/compat/__init__.py

@@ -36,6 +36,7 @@
     services.featurestore_service_client = services.featurestore_service_client_v1beta1
     services.job_service_client = services.job_service_client_v1beta1
     services.model_service_client = services.model_service_client_v1beta1
+    services.model_garden_service_client = services.model_garden_service_client_v1beta1
     services.pipeline_service_client = services.pipeline_service_client_v1beta1
     services.prediction_service_client = services.prediction_service_client_v1beta1
     services.specialist_pool_service_client = (
@@ -103,6 +104,7 @@
     types.model_deployment_monitoring_job = (
         types.model_deployment_monitoring_job_v1beta1
     )
+    types.model_garden_service = types.model_garden_service_v1beta1
     types.model_monitoring = types.model_monitoring_v1beta1
     types.model_service = types.model_service_v1beta1
     types.operation = types.operation_v1beta1
@@ -111,6 +113,7 @@
     types.pipeline_service = types.pipeline_service_v1beta1
     types.pipeline_state = types.pipeline_state_v1beta1
     types.prediction_service = types.prediction_service_v1beta1
+    types.publisher_model = types.publisher_model_v1beta1
     types.specialist_pool = types.specialist_pool_v1beta1
     types.specialist_pool_service = types.specialist_pool_service_v1beta1
     types.study = types.study_v1beta1

google/cloud/aiplatform/compat/services/__init__.py

@@ -45,6 +45,9 @@
 from google.cloud.aiplatform_v1beta1.services.metadata_service import (
     client as metadata_service_client_v1beta1,
 )
+from google.cloud.aiplatform_v1beta1.services.model_garden_service import (
+    client as model_garden_service_client_v1beta1,
+)
 from google.cloud.aiplatform_v1beta1.services.model_service import (
     client as model_service_client_v1beta1,
 )
@@ -133,6 +136,7 @@
     index_endpoint_service_client_v1beta1,
     job_service_client_v1beta1,
     match_service_client_v1beta1,
+    model_garden_service_client_v1beta1,
     model_service_client_v1beta1,
     pipeline_service_client_v1beta1,
     prediction_service_client_v1beta1,

google/cloud/aiplatform/compat/types/__init__.py

@@ -65,6 +65,7 @@
     model_evaluation as model_evaluation_v1beta1,
     model_evaluation_slice as model_evaluation_slice_v1beta1,
     model_deployment_monitoring_job as model_deployment_monitoring_job_v1beta1,
+    model_garden_service as model_garden_service_v1beta1,
     model_service as model_service_v1beta1,
     model_monitoring as model_monitoring_v1beta1,
     operation as operation_v1beta1,
@@ -73,6 +74,7 @@
     pipeline_service as pipeline_service_v1beta1,
     pipeline_state as pipeline_state_v1beta1,
     prediction_service as prediction_service_v1beta1,
+    publisher_model as publisher_model_v1beta1,
     specialist_pool as specialist_pool_v1beta1,
     specialist_pool_service as specialist_pool_service_v1beta1,
     study as study_v1beta1,
@@ -204,7 +206,7 @@
     model_service_v1,
     model_monitoring_v1,
     operation_v1,
-    pipeline_failure_policy_v1beta1,
+    pipeline_failure_policy_v1,
     pipeline_job_v1,
     pipeline_service_v1,
     pipeline_state_v1,
@@ -219,6 +221,8 @@
     tensorboard_time_series_v1,
     training_pipeline_v1,
     types_v1,
+    study_v1,
+    vizier_service_v1,
     # v1beta1
     accelerator_type_v1beta1,
     annotation_v1beta1,
@@ -269,6 +273,7 @@
     model_evaluation_v1beta1,
     model_evaluation_slice_v1beta1,
     model_deployment_monitoring_job_v1beta1,
+    model_garden_service_v1beta1,
     model_service_v1beta1,
     model_monitoring_v1beta1,
     operation_v1beta1,
@@ -277,8 +282,10 @@
     pipeline_service_v1beta1,
     pipeline_state_v1beta1,
     prediction_service_v1beta1,
+    publisher_model_v1beta1,
     specialist_pool_v1beta1,
     specialist_pool_service_v1beta1,
+    study_v1beta1,
     tensorboard_v1beta1,
     tensorboard_data_v1beta1,
     tensorboard_experiment_v1beta1,
@@ -287,4 +294,5 @@
     tensorboard_time_series_v1beta1,
     training_pipeline_v1beta1,
     types_v1beta1,
+    vizier_service_v1beta1,
 )

google/cloud/aiplatform/preview/_publisher_model.py

@@ -0,0 +1,79 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2023 Google LLC
+#
+# 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.
+#
+
+import re
+from typing import Optional
+
+from google.auth import credentials as auth_credentials
+from google.cloud.aiplatform import base
+from google.cloud.aiplatform import utils
+
+
+class _PublisherModel(base.VertexAiResourceNoun):
+    """Publisher Model Resource for Vertex AI."""
+
+    client_class = utils.ModelGardenClientWithOverride
+
+    _resource_noun = "publisher_model"
+    _getter_method = "get_publisher_model"
+    _delete_method = None
+    _parse_resource_name_method = "parse_publisher_model_path"
+    _format_resource_name_method = "publisher_model_path"
+
+    def __init__(
+        self,
+        resource_name: str,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ):
+        """Retrieves an existing PublisherModel resource given a resource name or model garden id.
+
+        Args:
+            resource_name (str):
+                Required. A fully-qualified PublisherModel resource name or
+                model garden id. Format:
+                `publishers/{publisher}/models/{publisher_model}` or
+                `{publisher}/{publisher_model}`.
+            project (str):
+                Optional. Project to retrieve the resource from. If not set,
+                project set in aiplatform.init will be used.
+            location (str):
+                Optional. Location to retrieve the resource from. If not set,
+                location set in aiplatform.init will be used.
+            credentials (auth_credentials.Credentials):
+                Optional. Custom credentials to use to retrieve the resource.
+                Overrides credentials set in aiplatform.init.
+        """
+
+        super().__init__(project=project, location=location, credentials=credentials)
+
+        if self._parse_resource_name(resource_name):
+            full_resource_name = resource_name
+        else:
+            m = re.match(r"^(?P<publisher>.+?)/(?P<model>.+?)$", resource_name)
+            if m:
+                full_resource_name = self._format_resource_name(**m.groupdict())
+            else:
+                raise ValueError(
+                    f"`{resource_name}` is not a valid PublisherModel resource "
+                    "name or model garden id."
+                )
+
+        self._gca_resource = getattr(self.api_client, self._getter_method)(
+            name=full_resource_name, retry=base._DEFAULT_RETRY
+        )

google/cloud/aiplatform/utils/__init__.py

@@ -51,6 +51,7 @@
     prediction_service_client_v1beta1,
     tensorboard_service_client_v1beta1,
     vizier_service_client_v1beta1,
+    model_garden_service_client_v1beta1,
 )
 from google.cloud.aiplatform.compat.services import (
     dataset_service_client_v1,
@@ -633,6 +634,14 @@ class VizierClientWithOverride(ClientWithOverride):
     )
 
 
+class ModelGardenClientWithOverride(ClientWithOverride):
+    _is_temporary = True
+    _default_version = compat.V1BETA1
+    _version_map = (
+        (compat.V1BETA1, model_garden_service_client_v1beta1.ModelGardenServiceClient),
+    )
+
+
 VertexAiServiceClientWithOverride = TypeVar(
     "VertexAiServiceClientWithOverride",
     DatasetClientWithOverride,
@@ -647,6 +656,7 @@ class VizierClientWithOverride(ClientWithOverride):
     MetadataClientWithOverride,
     TensorboardClientWithOverride,
     VizierClientWithOverride,
+    ModelGardenClientWithOverride,
 )
 
 

tests/unit/aiplatform/test_publisher_model.py

@@ -0,0 +1,93 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2023 Google LLC
+#
+# 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.
+#
+
+import pytest
+
+from unittest import mock
+from importlib import reload
+
+from google.cloud import aiplatform
+from google.cloud.aiplatform import base
+from google.cloud.aiplatform import initializer
+from google.cloud.aiplatform.preview import _publisher_model
+
+from google.cloud.aiplatform.compat.services import (
+    model_garden_service_client_v1beta1,
+)
+
+
+_TEST_PROJECT = "test-project"
+_TEST_LOCATION = "us-central1"
+
+_TEST_RESOURCE_NAME = "publishers/google/models/chat-bison@001"
+_TEST_MODEL_GARDEN_ID = "google/chat-bison@001"
+_TEST_INVALID_RESOURCE_NAME = "google.chat-bison@001"
+
+
+@pytest.fixture
+def mock_get_publisher_model():
+    with mock.patch.object(
+        model_garden_service_client_v1beta1.ModelGardenServiceClient,
+        "get_publisher_model",
+    ) as mock_get_publisher_model:
+        yield mock_get_publisher_model
+
+
+@pytest.mark.usefixtures("google_auth_mock")
+class TestPublisherModel:
+    def setup_method(self):
+        reload(initializer)
+        reload(aiplatform)
+
+    def teardown_method(self):
+        initializer.global_pool.shutdown(wait=True)
+
+    def test_init_publisher_model_with_resource_name(self, mock_get_publisher_model):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+        _ = _publisher_model._PublisherModel(_TEST_RESOURCE_NAME)
+        mock_get_publisher_model.assert_called_once_with(
+            name=_TEST_RESOURCE_NAME, retry=base._DEFAULT_RETRY
+        )
+
+    def test_init_publisher_model_with_model_garden_id(self, mock_get_publisher_model):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+        _ = _publisher_model._PublisherModel(_TEST_MODEL_GARDEN_ID)
+        mock_get_publisher_model.assert_called_once_with(
+            name=_TEST_RESOURCE_NAME, retry=base._DEFAULT_RETRY
+        )
+
+    def test_init_publisher_model_with_invalid_resource_name(
+        self, mock_get_publisher_model
+    ):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+        with pytest.raises(
+            ValueError,
+            match=(
+                f"`{_TEST_INVALID_RESOURCE_NAME}` is not a valid PublisherModel "
+                "resource name or model garden id."
+            ),
+        ):
+            _ = _publisher_model._PublisherModel(_TEST_INVALID_RESOURCE_NAME)