chore: add model garden and publisher model v1 support

PiperOrigin-RevId: 537955149

Author: sararob

google/cloud/aiplatform/compat/__init__.py

@@ -140,6 +140,7 @@
     )
     services.featurestore_service_client = services.featurestore_service_client_v1
     services.job_service_client = services.job_service_client_v1
+    services.model_garden_service_client = services.model_garden_service_client_v1
     services.model_service_client = services.model_service_client_v1
     services.pipeline_service_client = services.pipeline_service_client_v1
     services.prediction_service_client = services.prediction_service_client_v1
@@ -206,6 +207,7 @@
     types.pipeline_service = types.pipeline_service_v1
     types.pipeline_state = types.pipeline_state_v1
     types.prediction_service = types.prediction_service_v1
+    types.publisher_model = types.publisher_model_v1
     types.specialist_pool = types.specialist_pool_v1
     types.specialist_pool_service = types.specialist_pool_service_v1
     types.study = types.study_v1

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

@@ -94,6 +94,9 @@
 from google.cloud.aiplatform_v1.services.metadata_service import (
     client as metadata_service_client_v1,
 )
+from google.cloud.aiplatform_v1.services.model_garden_service import (
+    client as model_garden_service_client_v1,
+)
 from google.cloud.aiplatform_v1.services.model_service import (
     client as model_service_client_v1,
 )
@@ -123,6 +126,7 @@
     index_endpoint_service_client_v1,
     job_service_client_v1,
     metadata_service_client_v1,
+    model_garden_service_client_v1,
     model_service_client_v1,
     pipeline_service_client_v1,
     prediction_service_client_v1,

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

@@ -144,6 +144,7 @@
     pipeline_service as pipeline_service_v1,
     pipeline_state as pipeline_state_v1,
     prediction_service as prediction_service_v1,
+    publisher_model as publisher_model_v1,
     specialist_pool as specialist_pool_v1,
     specialist_pool_service as specialist_pool_service_v1,
     study as study_v1,
@@ -213,6 +214,7 @@
     pipeline_service_v1,
     pipeline_state_v1,
     prediction_service_v1,
+    publisher_model_v1,
     specialist_pool_v1,
     specialist_pool_service_v1,
     tensorboard_v1,

google/cloud/aiplatform/utils/__init__.py

@@ -63,6 +63,7 @@
     index_endpoint_service_client_v1,
     job_service_client_v1,
     metadata_service_client_v1,
+    model_garden_service_client_v1,
     model_service_client_v1,
     pipeline_service_client_v1,
     prediction_service_client_v1,
@@ -646,8 +647,9 @@ class VizierClientWithOverride(ClientWithOverride):
 
 class ModelGardenClientWithOverride(ClientWithOverride):
     _is_temporary = True
-    _default_version = compat.V1BETA1
+    _default_version = compat.DEFAULT_VERSION
     _version_map = (
+        (compat.V1, model_garden_service_client_v1.ModelGardenServiceClient),
         (compat.V1BETA1, model_garden_service_client_v1beta1.ModelGardenServiceClient),
     )
 

tests/unit/aiplatform/test_language_models.py

@@ -34,7 +34,7 @@
 import constants as test_constants
 
 from google.cloud.aiplatform.compat.services import (
-    model_garden_service_client_v1beta1,
+    model_garden_service_client,
     endpoint_service_client,
     model_service_client,
     pipeline_service_client,
@@ -46,9 +46,9 @@
     endpoint as gca_endpoint,
     pipeline_job as gca_pipeline_job,
     pipeline_state as gca_pipeline_state,
-    deployed_model_ref_v1beta1,
+    deployed_model_ref_v1,
 )
-from google.cloud.aiplatform_v1beta1.types import (
+from google.cloud.aiplatform_v1.types import (
     publisher_model as gca_publisher_model,
     model as gca_model,
 )
@@ -478,7 +478,7 @@ def get_model_with_tuned_version_label_mock():
             name=test_constants.ModelConstants._TEST_MODEL_RESOURCE_NAME,
             labels={"google-vertex-llm-tuning-base-model-id": "text-bison-001"},
             deployed_models=[
-                deployed_model_ref_v1beta1.DeployedModelRef(
+                deployed_model_ref_v1.DeployedModelRef(
                     endpoint=test_constants.EndpointConstants._TEST_ENDPOINT_NAME,
                     deployed_model_id=test_constants.ModelConstants._TEST_MODEL_RESOURCE_NAME,
                 )
@@ -525,7 +525,7 @@ def test_text_generation(self):
             location=_TEST_LOCATION,
         )
         with mock.patch.object(
-            target=model_garden_service_client_v1beta1.ModelGardenServiceClient,
+            target=model_garden_service_client.ModelGardenServiceClient,
             attribute="get_publisher_model",
             return_value=gca_publisher_model.PublisherModel(
                 _TEXT_BISON_PUBLISHER_MODEL_DICT
@@ -590,7 +590,7 @@ def test_tune_model(
             encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME,
         )
         with mock.patch.object(
-            target=model_garden_service_client_v1beta1.ModelGardenServiceClient,
+            target=model_garden_service_client.ModelGardenServiceClient,
             attribute="get_publisher_model",
             return_value=gca_publisher_model.PublisherModel(
                 _TEXT_BISON_PUBLISHER_MODEL_DICT
@@ -625,7 +625,7 @@ def test_get_tuned_model(
         )
 
         with mock.patch.object(
-            target=model_garden_service_client_v1beta1.ModelGardenServiceClient,
+            target=model_garden_service_client.ModelGardenServiceClient,
             attribute="get_publisher_model",
             return_value=gca_publisher_model.PublisherModel(
                 _TEXT_BISON_PUBLISHER_MODEL_DICT
@@ -662,7 +662,7 @@ def test_chat(self):
             location=_TEST_LOCATION,
         )
         with mock.patch.object(
-            target=model_garden_service_client_v1beta1.ModelGardenServiceClient,
+            target=model_garden_service_client.ModelGardenServiceClient,
             attribute="get_publisher_model",
             return_value=gca_publisher_model.PublisherModel(
                 _CHAT_BISON_PUBLISHER_MODEL_DICT
@@ -780,7 +780,7 @@ def test_code_chat(self):
             location=_TEST_LOCATION,
         )
         with mock.patch.object(
-            target=model_garden_service_client_v1beta1.ModelGardenServiceClient,
+            target=model_garden_service_client.ModelGardenServiceClient,
             attribute="get_publisher_model",
             return_value=gca_publisher_model.PublisherModel(
                 _CODECHAT_BISON_PUBLISHER_MODEL_DICT
@@ -876,7 +876,7 @@ def test_code_generation(self):
             location=_TEST_LOCATION,
         )
         with mock.patch.object(
-            target=model_garden_service_client_v1beta1.ModelGardenServiceClient,
+            target=model_garden_service_client.ModelGardenServiceClient,
             attribute="get_publisher_model",
             return_value=gca_publisher_model.PublisherModel(
                 _CODE_GENERATION_BISON_PUBLISHER_MODEL_DICT
@@ -942,7 +942,7 @@ def test_code_completion(self):
             location=_TEST_LOCATION,
         )
         with mock.patch.object(
-            target=model_garden_service_client_v1beta1.ModelGardenServiceClient,
+            target=model_garden_service_client.ModelGardenServiceClient,
             attribute="get_publisher_model",
             return_value=gca_publisher_model.PublisherModel(
                 _CODE_COMPLETION_BISON_PUBLISHER_MODEL_DICT
@@ -1008,7 +1008,7 @@ def test_text_embedding(self):
             location=_TEST_LOCATION,
         )
         with mock.patch.object(
-            target=model_garden_service_client_v1beta1.ModelGardenServiceClient,
+            target=model_garden_service_client.ModelGardenServiceClient,
             attribute="get_publisher_model",
             return_value=gca_publisher_model.PublisherModel(
                 _TEXT_EMBEDDING_GECKO_PUBLISHER_MODEL_DICT

tests/unit/aiplatform/test_model_garden_models.py

@@ -26,10 +26,10 @@
 import constants as test_constants
 
 from google.cloud.aiplatform.compat.services import (
-    model_garden_service_client_v1beta1,
+    model_garden_service_client_v1,
 )
 
-from google.cloud.aiplatform_v1beta1.types import (
+from google.cloud.aiplatform_v1.types import (
     publisher_model as gca_publisher_model,
 )
 
@@ -76,7 +76,7 @@ def test_init_model_garden_model_with_from_pretrained(self):
             location=test_constants.ProjectConstants._TEST_LOCATION,
         )
         with mock.patch.object(
-            target=model_garden_service_client_v1beta1.ModelGardenServiceClient,
+            target=model_garden_service_client_v1.ModelGardenServiceClient,
             attribute="get_publisher_model",
             return_value=gca_publisher_model.PublisherModel(
                 _TEXT_BISON_PUBLISHER_MODEL_DICT

tests/unit/aiplatform/test_publisher_model.py

@@ -26,7 +26,7 @@
 from google.cloud.aiplatform import _publisher_models
 
 from google.cloud.aiplatform.compat.services import (
-    model_garden_service_client_v1beta1,
+    model_garden_service_client_v1,
 )
 
 
@@ -41,7 +41,7 @@
 @pytest.fixture
 def mock_get_publisher_model():
     with mock.patch.object(
-        model_garden_service_client_v1beta1.ModelGardenServiceClient,
+        model_garden_service_client_v1.ModelGardenServiceClient,
         "get_publisher_model",
     ) as mock_get_publisher_model:
         yield mock_get_publisher_model