Remove parent_model version suffix if it is passed to Vertex AI opera…

…tors (#39640)

airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py

@@ -176,6 +176,7 @@ def execute(self, context: Context):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        self.parent_model = self.parent_model.split("@")[0] if self.parent_model else None
         model, training_id = self.hook.create_auto_ml_forecasting_training_job(
             project_id=self.project_id,
             region=self.region,
@@ -283,6 +284,7 @@ def execute(self, context: Context):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        self.parent_model = self.parent_model.split("@")[0] if self.parent_model else None
         model, training_id = self.hook.create_auto_ml_image_training_job(
             project_id=self.project_id,
             region=self.region,
@@ -391,6 +393,7 @@ def execute(self, context: Context):
             impersonation_chain=self.impersonation_chain,
         )
         credentials, _ = self.hook.get_credentials_and_project_id()
+        self.parent_model = self.parent_model.split("@")[0] if self.parent_model else None
         model, training_id = self.hook.create_auto_ml_tabular_training_job(
             project_id=self.project_id,
             region=self.region,
@@ -485,6 +488,7 @@ def execute(self, context: Context):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        self.parent_model = self.parent_model.split("@")[0] if self.parent_model else None
         model, training_id = self.hook.create_auto_ml_text_training_job(
             project_id=self.project_id,
             region=self.region,
@@ -561,6 +565,7 @@ def execute(self, context: Context):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        self.parent_model = self.parent_model.split("@")[0] if self.parent_model else None
         model, training_id = self.hook.create_auto_ml_video_training_job(
             project_id=self.project_id,
             region=self.region,

airflow/providers/google/cloud/operators/vertex_ai/custom_job.py

@@ -493,6 +493,8 @@ def __init__(
     def execute(self, context: Context):
         super().execute(context)
 
+        self.parent_model = self.parent_model.split("@")[0] if self.parent_model else None
+
         if self.deferrable:
             self.invoke_defer(context=context)
 
@@ -966,6 +968,8 @@ def __init__(
     def execute(self, context: Context):
         super().execute(context)
 
+        self.parent_model = self.parent_model.split("@")[0] if self.parent_model else None
+
         if self.deferrable:
             self.invoke_defer(context=context)
 
@@ -1446,6 +1450,8 @@ def __init__(
     def execute(self, context: Context):
         super().execute(context)
 
+        self.parent_model = self.parent_model.split("@")[0] if self.parent_model else None
+
         if self.deferrable:
             self.invoke_defer(context=context)
 

docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst

@@ -187,8 +187,8 @@ The same operation can be performed in the deferrable mode:
 .. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job.py
     :language: python
     :dedent: 4
-    :start-after: [START how_to_cloud_vertex_ai_create_custom_training_job_v2_operator_deferrable]
-    :end-before: [END how_to_cloud_vertex_ai_create_custom_training_job_v2_operator_deferrable]
+    :start-after: [START how_to_cloud_vertex_ai_create_custom_training_job_v2_deferrable_operator]
+    :end-before: [END how_to_cloud_vertex_ai_create_custom_training_job_v2_deferrable_operator]
 
 
 You can get a list of Training Jobs using