feat: Allow setting Vertex Model Garden source model name during model upload

KCFindstr · copybara-github · commit 3af9cc875649 · 2024-11-18T16:01:00.000-08:00
PiperOrigin-RevId: 697786055
diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py
@@ -4663,6 +4663,7 @@ def upload(
         serving_container_health_probe_exec: Optional[Sequence[str]] = None,
         serving_container_health_probe_period_seconds: Optional[int] = None,
         serving_container_health_probe_timeout_seconds: Optional[int] = None,
+        model_garden_source_model_name: Optional[str] = None,
     ) -> "Model":
         """Uploads a model and returns a Model representing the uploaded Model
         resource.
@@ -4875,6 +4876,10 @@ def upload(
             serving_container_health_probe_timeout_seconds (int):
                 Optional. Number of seconds after which the health probe times
                 out. Defaults to 1 second. Minimum value is 1.
+            model_garden_source_model_name:
+                Optional. The model garden source model resource name if the
+                model is from Vertex Model Garden.
+
 
         Returns:
             model (aiplatform.Model):
@@ -5003,6 +5008,14 @@ def upload(
             version_aliases=version_aliases, is_default_version=is_default_version
         )
 
+        base_model_source = None
+        if model_garden_source_model_name:
+            base_model_source = gca_model_compat.Model.BaseModelSource(
+                model_garden_source=gca_model_compat.ModelGardenSource(
+                    public_model_name=model_garden_source_model_name
+                )
+            )
+
         managed_model = gca_model_compat.Model(
             display_name=display_name,
             description=description,
@@ -5012,6 +5025,7 @@ def upload(
             predict_schemata=model_predict_schemata,
             labels=labels,
             encryption_spec=encryption_spec,
+            base_model_source=base_model_source,
         )
 
         if artifact_uri and not artifact_uri.startswith("gs://"):
diff --git a/tests/unit/aiplatform/test_models.py b/tests/unit/aiplatform/test_models.py
@@ -174,6 +174,7 @@
 
 _TEST_CREDENTIALS = mock.Mock(spec=auth_credentials.AnonymousCredentials())
 _TEST_SERVICE_ACCOUNT = "vinnys@my-project.iam.gserviceaccount.com"
+_TEST_MODEL_GARDEN_SOURCE_MODEL_NAME = "publishers/meta/models/llama3_1"
 
 
 _TEST_EXPLANATION_METADATA = explain.ExplanationMetadata(
@@ -1900,6 +1901,53 @@ def test_upload_uploads_and_gets_model_with_custom_location(
             name=test_model_resource_name, retry=base._DEFAULT_RETRY
         )
 
+    @pytest.mark.parametrize("sync", [True, False])
+    def test_upload_with_model_garden_source(
+        self, upload_model_mock, get_model_mock, sync
+    ):
+
+        my_model = models.Model.upload(
+            display_name=_TEST_MODEL_NAME,
+            serving_container_image_uri=_TEST_SERVING_CONTAINER_IMAGE,
+            serving_container_predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE,
+            serving_container_health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE,
+            sync=sync,
+            upload_request_timeout=None,
+            model_garden_source_model_name=_TEST_MODEL_GARDEN_SOURCE_MODEL_NAME,
+        )
+
+        if not sync:
+            my_model.wait()
+
+        container_spec = gca_model.ModelContainerSpec(
+            image_uri=_TEST_SERVING_CONTAINER_IMAGE,
+            predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE,
+            health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE,
+        )
+
+        managed_model = gca_model.Model(
+            display_name=_TEST_MODEL_NAME,
+            container_spec=container_spec,
+            version_aliases=["default"],
+            base_model_source=gca_model.Model.BaseModelSource(
+                model_garden_source=gca_model.ModelGardenSource(
+                    public_model_name=_TEST_MODEL_GARDEN_SOURCE_MODEL_NAME
+                )
+            ),
+        )
+
+        upload_model_mock.assert_called_once_with(
+            request=gca_model_service.UploadModelRequest(
+                parent=initializer.global_config.common_location_path(),
+                model=managed_model,
+            ),
+            timeout=None,
+        )
+
+        get_model_mock.assert_called_once_with(
+            name=_TEST_MODEL_RESOURCE_NAME, retry=base._DEFAULT_RETRY
+        )
+
     @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock")
     @pytest.mark.parametrize("sync", [True, False])
     def test_deploy(self, deploy_model_mock, sync):
