Fix validating parent_model parameter in UploadModelOperator (#43473)

VladaZakharova · Ulada Zakharava · web-flow · commit b344cc1474de · 2024-10-29T20:21:56.000+02:00
Co-authored-by: Ulada Zakharava &lt;vlada_zakharava@epam.com&gt;
diff --git a/providers/src/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py b/providers/src/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py
@@ -219,7 +219,7 @@ def upload_model(
         :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
         :param region: Required. The ID of the Google Cloud region that the service belongs to.
         :param model:  Required. The Model to create.
-        :param parent_model: The name of the parent model to create a new version under.
+        :param parent_model: The ID of the parent model to create a new version under.
         :param retry: Designation of what errors, if any, should be retried.
         :param timeout: The timeout for this request.
         :param metadata: Strings which should be sent along with the request as metadata.
@@ -233,7 +233,7 @@ def upload_model(
         }
 
         if parent_model:
-            request["parent_model"] = parent_model
+            request["parent_model"] = client.model_path(project_id, region, parent_model)
 
         result = client.upload_model(
             request=request,
diff --git a/providers/src/airflow/providers/google/cloud/operators/vertex_ai/model_service.py b/providers/src/airflow/providers/google/cloud/operators/vertex_ai/model_service.py
@@ -361,8 +361,9 @@ class UploadModelOperator(GoogleCloudBaseOperator):
 
     :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
     :param region: Required. The ID of the Google Cloud region that the service belongs to.
-    :param model: Required. The Model to create.
-    :param parent_model: The name of the parent model to create a new version under.
+    :param model: Required. The Model to create. Creating model with the name that already
+        exists leads to creating new version of existing model.
+    :param parent_model: The ID of the parent model to create a new version under.
     :param retry: Designation of what errors, if any, should be retried.
     :param timeout: The timeout for this request.
     :param metadata: Strings which should be sent along with the request as metadata.
@@ -377,7 +378,7 @@ class UploadModelOperator(GoogleCloudBaseOperator):
         account from the list granting this role to the originating account (templated).
     """
 
-    template_fields = ("region", "project_id", "model", "impersonation_chain")
+    template_fields = ("region", "project_id", "model", "parent_model", "impersonation_chain")
     operator_extra_links = (VertexAIModelLink(),)
 
     def __init__(
diff --git a/providers/tests/google/cloud/hooks/vertex_ai/test_model_service.py b/providers/tests/google/cloud/hooks/vertex_ai/test_model_service.py
@@ -37,7 +37,7 @@
 TEST_REGION: str = "test-region"
 TEST_PROJECT_ID: str = "test-project-id"
 TEST_MODEL = None
-TEST_PARENT_MODEL = "test-parent-model"
+TEST_PARENT_MODEL = "projects/test-project-id/locations/test-region/models/test-parent-model"
 TEST_MODEL_NAME: str = "test_model_name"
 TEST_OUTPUT_CONFIG: dict = {}
 
@@ -148,7 +148,7 @@ def test_upload_model_with_parent_model(self, mock_client) -> None:
             request=dict(
                 parent=mock_client.return_value.common_location_path.return_value,
                 model=TEST_MODEL,
-                parent_model=TEST_PARENT_MODEL,
+                parent_model=mock_client.return_value.model_path.return_value,
             ),
             metadata=(),
             retry=DEFAULT,
@@ -352,7 +352,7 @@ def test_upload_model_with_parent_model(self, mock_client) -> None:
             request=dict(
                 parent=mock_client.return_value.common_location_path.return_value,
                 model=TEST_MODEL,
-                parent_model=TEST_PARENT_MODEL,
+                parent_model=mock_client.return_value.model_path.return_value,
             ),
             metadata=(),
             retry=DEFAULT,
diff --git a/providers/tests/system/google/cloud/vertex_ai/example_vertex_ai_model_service.py b/providers/tests/system/google/cloud/vertex_ai/example_vertex_ai_model_service.py
@@ -113,6 +113,19 @@
         "health_route": "",
     },
 }
+MODEL_OBJ_V2 = {
+    "display_name": f"model-{ENV_ID}-v2",
+    "artifact_uri": "{{ti.xcom_pull('custom_task')['artifactUri']}}",
+    "container_spec": {
+        "image_uri": MODEL_SERVING_CONTAINER_URI,
+        "command": [],
+        "args": [],
+        "env": [],
+        "ports": [],
+        "predict_route": "",
+        "health_route": "",
+    },
+}
 
 
 with DAG(
@@ -229,13 +242,14 @@
         project_id=PROJECT_ID,
         model=MODEL_OBJ,
     )
+    upload_model_v1 = upload_model.output["model_id"]
     # [END how_to_cloud_vertex_ai_upload_model_operator]
     upload_model_with_parent_model = UploadModelOperator(
         task_id="upload_model_with_parent_model",
         region=REGION,
         project_id=PROJECT_ID,
-        model=MODEL_OBJ,
-        parent_model=MODEL_DISPLAY_NAME,
+        model=MODEL_OBJ_V2,
+        parent_model=upload_model_v1,
     )
 
     # [START how_to_cloud_vertex_ai_export_model_operator]
