apache-airflow-providers-google 10.10.0rc1__py3-none-any.whl → 10.10.1__py3-none-any.whl

airflow/providers/google/cloud/links/vertex_ai.py

@@ -49,6 +49,10 @@ VERTEX_AI_ENDPOINT_LINK = (
     VERTEX_AI_BASE_LINK + "/locations/{region}/endpoints/{endpoint_id}?project={project_id}"
 )
 VERTEX_AI_ENDPOINT_LIST_LINK = VERTEX_AI_BASE_LINK + "/endpoints?project={project_id}"
+VERTEX_AI_PIPELINE_JOB_LINK = (
+    VERTEX_AI_BASE_LINK + "/locations/{region}/pipelines/runs/{pipeline_id}?project={project_id}"
+)
+VERTEX_AI_PIPELINE_JOB_LIST_LINK = VERTEX_AI_BASE_LINK + "/pipelines/runs?project={project_id}"
 
 
 class VertexAIModelLink(BaseGoogleLink):
@@ -319,3 +323,48 @@ class VertexAIEndpointListLink(BaseGoogleLink):
                 "project_id": task_instance.project_id,
             },
         )
+
+
+class VertexAIPipelineJobLink(BaseGoogleLink):
+    """Helper class for constructing Vertex AI PipelineJob link."""
+
+    name = "Pipeline Job"
+    key = "pipeline_job_conf"
+    format_str = VERTEX_AI_PIPELINE_JOB_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+        pipeline_id: str,
+    ):
+        task_instance.xcom_push(
+            context=context,
+            key=VertexAIPipelineJobLink.key,
+            value={
+                "pipeline_id": pipeline_id,
+                "region": task_instance.region,
+                "project_id": task_instance.project_id,
+            },
+        )
+
+
+class VertexAIPipelineJobListLink(BaseGoogleLink):
+    """Helper class for constructing Vertex AI PipelineJobList link."""
+
+    name = "Pipeline Job List"
+    key = "pipeline_job_list_conf"
+    format_str = VERTEX_AI_PIPELINE_JOB_LIST_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance,
+    ):
+        task_instance.xcom_push(
+            context=context,
+            key=VertexAIPipelineJobListLink.key,
+            value={
+                "project_id": task_instance.project_id,
+            },
+        )

airflow/providers/google/cloud/operators/dataproc.py

@@ -1790,6 +1790,7 @@ class DataprocInstantiateWorkflowTemplateOperator(GoogleCloudBaseOperator):
         account from the list granting this role to the originating account (templated).
     :param deferrable: Run operator in the deferrable mode.
     :param polling_interval_seconds: Time (seconds) to wait between calls to check the run status.
+    :param cancel_on_kill: Flag which indicates whether cancel the workflow, when on_kill is called
     """
 
     template_fields: Sequence[str] = ("template_id", "impersonation_chain", "request_id", "parameters")
@@ -1812,6 +1813,7 @@ class DataprocInstantiateWorkflowTemplateOperator(GoogleCloudBaseOperator):
         impersonation_chain: str | Sequence[str] | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         polling_interval_seconds: int = 10,
+        cancel_on_kill: bool = True,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -1830,6 +1832,8 @@ class DataprocInstantiateWorkflowTemplateOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.deferrable = deferrable
         self.polling_interval_seconds = polling_interval_seconds
+        self.cancel_on_kill = cancel_on_kill
+        self.operation_name: str | None = None
 
     def execute(self, context: Context):
         hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
@@ -1845,24 +1849,26 @@ class DataprocInstantiateWorkflowTemplateOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        self.workflow_id = operation.operation.name.split("/")[-1]
+        operation_name = operation.operation.name
+        self.operation_name = operation_name
+        workflow_id = operation_name.split("/")[-1]
         project_id = self.project_id or hook.project_id
         if project_id:
             DataprocWorkflowLink.persist(
                 context=context,
                 operator=self,
-                workflow_id=self.workflow_id,
+                workflow_id=workflow_id,
                 region=self.region,
                 project_id=project_id,
             )
-        self.log.info("Template instantiated. Workflow Id : %s", self.workflow_id)
+        self.log.info("Template instantiated. Workflow Id : %s", workflow_id)
         if not self.deferrable:
             hook.wait_for_operation(timeout=self.timeout, result_retry=self.retry, operation=operation)
-            self.log.info("Workflow %s completed successfully", self.workflow_id)
+            self.log.info("Workflow %s completed successfully", workflow_id)
         else:
             self.defer(
                 trigger=DataprocWorkflowTrigger(
-                    name=operation.operation.name,
+                    name=operation_name,
                     project_id=self.project_id,
                     region=self.region,
                     gcp_conn_id=self.gcp_conn_id,
@@ -1884,6 +1890,11 @@ class DataprocInstantiateWorkflowTemplateOperator(GoogleCloudBaseOperator):
 
         self.log.info("Workflow %s completed successfully", event["operation_name"])
 
+    def on_kill(self) -> None:
+        if self.cancel_on_kill and self.operation_name:
+            hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
+            hook.get_operations_client(region=self.region).cancel_operation(name=self.operation_name)
+
 
 class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator):
     """Instantiate a WorkflowTemplate Inline on Google Cloud Dataproc.
@@ -1926,6 +1937,7 @@ class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator)
         account from the list granting this role to the originating account (templated).
     :param deferrable: Run operator in the deferrable mode.
     :param polling_interval_seconds: Time (seconds) to wait between calls to check the run status.
+    :param cancel_on_kill: Flag which indicates whether cancel the workflow, when on_kill is called
     """
 
     template_fields: Sequence[str] = ("template", "impersonation_chain")
@@ -1946,6 +1958,7 @@ class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator)
         impersonation_chain: str | Sequence[str] | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         polling_interval_seconds: int = 10,
+        cancel_on_kill: bool = True,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -1963,6 +1976,8 @@ class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator)
         self.impersonation_chain = impersonation_chain
         self.deferrable = deferrable
         self.polling_interval_seconds = polling_interval_seconds
+        self.cancel_on_kill = cancel_on_kill
+        self.operation_name: str | None = None
 
     def execute(self, context: Context):
         self.log.info("Instantiating Inline Template")
@@ -1977,23 +1992,25 @@ class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator)
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        self.workflow_id = operation.operation.name.split("/")[-1]
+        operation_name = operation.operation.name
+        self.operation_name = operation_name
+        workflow_id = operation_name.split("/")[-1]
         if project_id:
             DataprocWorkflowLink.persist(
                 context=context,
                 operator=self,
-                workflow_id=self.workflow_id,
+                workflow_id=workflow_id,
                 region=self.region,
                 project_id=project_id,
             )
         if not self.deferrable:
-            self.log.info("Template instantiated. Workflow Id : %s", self.workflow_id)
+            self.log.info("Template instantiated. Workflow Id : %s", workflow_id)
             operation.result()
-            self.log.info("Workflow %s completed successfully", self.workflow_id)
+            self.log.info("Workflow %s completed successfully", workflow_id)
         else:
             self.defer(
                 trigger=DataprocWorkflowTrigger(
-                    name=operation.operation.name,
+                    name=operation_name,
                     project_id=self.project_id or hook.project_id,
                     region=self.region,
                     gcp_conn_id=self.gcp_conn_id,
@@ -2015,6 +2032,11 @@ class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator)
 
         self.log.info("Workflow %s completed successfully", event["operation_name"])
 
+    def on_kill(self) -> None:
+        if self.cancel_on_kill and self.operation_name:
+            hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
+            hook.get_operations_client(region=self.region).cancel_operation(name=self.operation_name)
+
 
 class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
     """Submit a job to a cluster.

airflow/providers/google/cloud/operators/gcs.py

@@ -919,7 +919,7 @@ class GCSSynchronizeBucketsOperator(GoogleCloudBaseOperator):
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
-        :ref:`howto/operator:GCSSynchronizeBuckets`
+        :ref:`howto/operator:GCSSynchronizeBucketsOperator`
 
     :param source_bucket: The name of the bucket containing the source objects.
     :param destination_bucket: The name of the bucket containing the destination objects.

airflow/providers/google/cloud/operators/mlengine.py

@@ -82,6 +82,10 @@ class MLEngineStartBatchPredictionJobOperator(GoogleCloudBaseOperator):
     """
     Start a Google Cloud ML Engine prediction job.
 
+    This operator is deprecated. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.batch_prediction.CreateBatchPredictionJobOperator`
+    instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:MLEngineStartBatchPredictionJobOperator`
@@ -210,6 +214,14 @@ class MLEngineStartBatchPredictionJobOperator(GoogleCloudBaseOperator):
         self._labels = labels
         self._impersonation_chain = impersonation_chain
 
+        warnings.warn(
+            "This operator is deprecated. All the functionality of legacy "
+            "MLEngine and new features are available on the Vertex AI platform. "
+            "Please use `CreateBatchPredictionJobOperator`",
+            AirflowProviderDeprecationWarning,
+            stacklevel=3,
+        )
+
         if not self._project_id:
             raise AirflowException("Google Cloud project id is required.")
         if not self._job_id:
@@ -364,6 +376,9 @@ class MLEngineCreateModelOperator(GoogleCloudBaseOperator):
     """
     Creates a new model.
 
+    This operator is deprecated. Please use appropriate VertexAI operator from
+    :class:`airflow.providers.google.cloud.operators.vertex_ai` instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:MLEngineCreateModelOperator`
@@ -407,6 +422,14 @@ class MLEngineCreateModelOperator(GoogleCloudBaseOperator):
         self._gcp_conn_id = gcp_conn_id
         self._impersonation_chain = impersonation_chain
 
+        warnings.warn(
+            "This operator is deprecated. All the functionality of legacy "
+            "MLEngine and new features are available on the Vertex AI platform. "
+            "Please use appropriate VertexAI operator.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=3,
+        )
+
     def execute(self, context: Context):
         hook = MLEngineHook(
             gcp_conn_id=self._gcp_conn_id,
@@ -429,6 +452,9 @@ class MLEngineGetModelOperator(GoogleCloudBaseOperator):
     """
     Gets a particular model.
 
+    This operator is deprecated. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.GetModelOperator` instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:MLEngineGetModelOperator`
@@ -472,6 +498,14 @@ class MLEngineGetModelOperator(GoogleCloudBaseOperator):
         self._gcp_conn_id = gcp_conn_id
         self._impersonation_chain = impersonation_chain
 
+        warnings.warn(
+            "This operator is deprecated. All the functionality of legacy "
+            "MLEngine and new features are available on the Vertex AI platform. "
+            "Please use `GetModelOperator`",
+            AirflowProviderDeprecationWarning,
+            stacklevel=3,
+        )
+
     def execute(self, context: Context):
         hook = MLEngineHook(
             gcp_conn_id=self._gcp_conn_id,
@@ -493,6 +527,10 @@ class MLEngineDeleteModelOperator(GoogleCloudBaseOperator):
     """
     Deletes a model.
 
+    This operator is deprecated. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelOperator` instead.
+
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:MLEngineDeleteModelOperator`
@@ -541,6 +579,14 @@ class MLEngineDeleteModelOperator(GoogleCloudBaseOperator):
         self._gcp_conn_id = gcp_conn_id
         self._impersonation_chain = impersonation_chain
 
+        warnings.warn(
+            "This operator is deprecated. All the functionality of legacy "
+            "MLEngine and new features are available on the Vertex AI platform. "
+            "Please use `DeleteModelOperator`",
+            AirflowProviderDeprecationWarning,
+            stacklevel=3,
+        )
+
     def execute(self, context: Context):
         hook = MLEngineHook(
             gcp_conn_id=self._gcp_conn_id,
@@ -682,6 +728,8 @@ class MLEngineCreateVersionOperator(GoogleCloudBaseOperator):
     """
     Creates a new version in the model.
 
+    This operator is deprecated. Please use parent_model parameter of VertexAI operators instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:MLEngineCreateVersionOperator`
@@ -731,6 +779,14 @@ class MLEngineCreateVersionOperator(GoogleCloudBaseOperator):
         self._impersonation_chain = impersonation_chain
         self._validate_inputs()
 
+        warnings.warn(
+            "This operator is deprecated. All the functionality of legacy "
+            "MLEngine and new features are available on the Vertex AI platform. "
+            "Please use parent_model parameter for VertexAI operators instead.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=3,
+        )
+
     def _validate_inputs(self):
         if not self._model_name:
             raise AirflowException("The model_name parameter could not be empty.")
@@ -763,6 +819,10 @@ class MLEngineSetDefaultVersionOperator(GoogleCloudBaseOperator):
     """
     Sets a version in the model.
 
+    This operator is deprecated. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.SetDefaultVersionOnModelOperator`
+    instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:MLEngineSetDefaultVersionOperator`
@@ -812,6 +872,14 @@ class MLEngineSetDefaultVersionOperator(GoogleCloudBaseOperator):
         self._impersonation_chain = impersonation_chain
         self._validate_inputs()
 
+        warnings.warn(
+            "This operator is deprecated. All the functionality of legacy "
+            "MLEngine and new features are available on the Vertex AI platform. "
+            "Please use `SetDefaultVersionOnModelOperator` instead.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=3,
+        )
+
     def _validate_inputs(self):
         if not self._model_name:
             raise AirflowException("The model_name parameter could not be empty.")
@@ -844,6 +912,10 @@ class MLEngineListVersionsOperator(GoogleCloudBaseOperator):
     """
     Lists all available versions of the model.
 
+    This operator is deprecated. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.ListModelVersionsOperator`
+    instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:MLEngineListVersionsOperator`
@@ -889,6 +961,14 @@ class MLEngineListVersionsOperator(GoogleCloudBaseOperator):
         self._impersonation_chain = impersonation_chain
         self._validate_inputs()
 
+        warnings.warn(
+            "This operator is deprecated. All the functionality of legacy "
+            "MLEngine and new features are available on the Vertex AI platform. "
+            "Please use `ListModelVersionsOperator` instead.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=3,
+        )
+
     def _validate_inputs(self):
         if not self._model_name:
             raise AirflowException("The model_name parameter could not be empty.")
@@ -918,6 +998,10 @@ class MLEngineDeleteVersionOperator(GoogleCloudBaseOperator):
     """
     Deletes the version from the model.
 
+    This operator is deprecated. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.model_service.DeleteModelVersionOperator`
+    instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:MLEngineDeleteVersionOperator`
@@ -967,6 +1051,14 @@ class MLEngineDeleteVersionOperator(GoogleCloudBaseOperator):
         self._impersonation_chain = impersonation_chain
         self._validate_inputs()
 
+        warnings.warn(
+            "This operator is deprecated. All the functionality of legacy "
+            "MLEngine and new features are available on the Vertex AI platform. "
+            "Please use `DeleteModelVersionOperator` instead.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=3,
+        )
+
     def _validate_inputs(self):
         if not self._model_name:
             raise AirflowException("The model_name parameter could not be empty.")
@@ -998,6 +1090,10 @@ class MLEngineStartTrainingJobOperator(GoogleCloudBaseOperator):
     """
     Operator for launching a MLEngine training job.
 
+    This operator is deprecated. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.custom_job.CreateCustomPythonPackageTrainingJobOperator`
+    instead.
+
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:MLEngineStartTrainingJobOperator`
@@ -1124,6 +1220,14 @@ class MLEngineStartTrainingJobOperator(GoogleCloudBaseOperator):
         self.deferrable = deferrable
         self.cancel_on_kill = cancel_on_kill
 
+        warnings.warn(
+            "This operator is deprecated. All the functionality of legacy "
+            "MLEngine and new features are available on the Vertex AI platform. "
+            "Please use `CreateCustomPythonPackageTrainingJobOperator` instead.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=3,
+        )
+
         custom = self._scale_tier is not None and self._scale_tier.upper() == "CUSTOM"
         custom_image = (
             custom
@@ -1328,6 +1432,10 @@ class MLEngineTrainingCancelJobOperator(GoogleCloudBaseOperator):
     """
     Operator for cleaning up failed MLEngine training job.
 
+    This operator is deprecated. Please use
+    :class:`airflow.providers.google.cloud.operators.vertex_ai.custom_job.CancelCustomTrainingJobOperator`
+    instead.
+
     :param job_id: A unique templated id for the submitted Google MLEngine
         training job. (templated)
     :param project_id: The Google Cloud project name within which MLEngine training job should run.
@@ -1366,6 +1474,14 @@ class MLEngineTrainingCancelJobOperator(GoogleCloudBaseOperator):
         self._gcp_conn_id = gcp_conn_id
         self._impersonation_chain = impersonation_chain
 
+        warnings.warn(
+            "This operator is deprecated. All the functionality of legacy "
+            "MLEngine and new features are available on the Vertex AI platform. "
+            "Please use `CancelCustomTrainingJobOperator` instead.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=3,
+        )
+
         if not self._project_id:
             raise AirflowException("Google Cloud project id is required.")
 

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

@@ -15,7 +15,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
 """This module contains Google Vertex AI operators."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence
@@ -50,6 +52,10 @@ class AutoMLTrainingJobBaseOperator(GoogleCloudBaseOperator):
         region: str,
         display_name: str,
         labels: dict[str, str] | None = None,
+        parent_model: str | None = None,
+        is_default_version: bool | None = None,
+        model_version_aliases: list[str] | None = None,
+        model_version_description: str | None = None,
         training_encryption_spec_key_name: str | None = None,
         model_encryption_spec_key_name: str | None = None,
         # RUN
@@ -67,6 +73,10 @@ class AutoMLTrainingJobBaseOperator(GoogleCloudBaseOperator):
         self.region = region
         self.display_name = display_name
         self.labels = labels
+        self.parent_model = parent_model
+        self.is_default_version = is_default_version
+        self.model_version_aliases = model_version_aliases
+        self.model_version_description = model_version_description
         self.training_encryption_spec_key_name = training_encryption_spec_key_name
         self.model_encryption_spec_key_name = model_encryption_spec_key_name
         # START Run param
@@ -90,6 +100,7 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator):
     """Create AutoML Forecasting Training job."""
 
     template_fields = (
+        "parent_model",
         "dataset_id",
         "region",
         "impersonation_chain",
@@ -158,11 +169,16 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        self.parent_model = self.parent_model.rpartition("@")[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,
             display_name=self.display_name,
             dataset=datasets.TimeSeriesDataset(dataset_name=self.dataset_id),
+            parent_model=self.parent_model,
+            is_default_version=self.is_default_version,
+            model_version_aliases=self.model_version_aliases,
+            model_version_description=self.model_version_description,
             target_column=self.target_column,
             time_column=self.time_column,
             time_series_identifier_column=self.time_series_identifier_column,
@@ -202,6 +218,7 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
+            self.xcom_push(context, key="model_id", value=model_id)
             VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
         else:
             result = model  # type: ignore
@@ -214,6 +231,7 @@ class CreateAutoMLImageTrainingJobOperator(AutoMLTrainingJobBaseOperator):
     """Create Auto ML Image Training job."""
 
     template_fields = (
+        "parent_model",
         "dataset_id",
         "region",
         "impersonation_chain",
@@ -254,11 +272,16 @@ class CreateAutoMLImageTrainingJobOperator(AutoMLTrainingJobBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        self.parent_model = self.parent_model.rpartition("@")[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,
             display_name=self.display_name,
             dataset=datasets.ImageDataset(dataset_name=self.dataset_id),
+            parent_model=self.parent_model,
+            is_default_version=self.is_default_version,
+            model_version_aliases=self.model_version_aliases,
+            model_version_description=self.model_version_description,
             prediction_type=self.prediction_type,
             multi_label=self.multi_label,
             model_type=self.model_type,
@@ -282,6 +305,7 @@ class CreateAutoMLImageTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
+            self.xcom_push(context, key="model_id", value=model_id)
             VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
         else:
             result = model  # type: ignore
@@ -294,6 +318,7 @@ class CreateAutoMLTabularTrainingJobOperator(AutoMLTrainingJobBaseOperator):
     """Create Auto ML Tabular Training job."""
 
     template_fields = (
+        "parent_model",
         "dataset_id",
         "region",
         "impersonation_chain",
@@ -351,6 +376,7 @@ class CreateAutoMLTabularTrainingJobOperator(AutoMLTrainingJobBaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
         credentials, _ = self.hook.get_credentials_and_project_id()
+        self.parent_model = self.parent_model.rpartition("@")[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,
@@ -360,6 +386,10 @@ class CreateAutoMLTabularTrainingJobOperator(AutoMLTrainingJobBaseOperator):
                 project=self.project_id,
                 credentials=credentials,
             ),
+            parent_model=self.parent_model,
+            is_default_version=self.is_default_version,
+            model_version_aliases=self.model_version_aliases,
+            model_version_description=self.model_version_description,
             target_column=self.target_column,
             optimization_prediction_type=self.optimization_prediction_type,
             optimization_objective=self.optimization_objective,
@@ -393,6 +423,7 @@ class CreateAutoMLTabularTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
+            self.xcom_push(context, key="model_id", value=model_id)
             VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
         else:
             result = model  # type: ignore
@@ -405,6 +436,7 @@ class CreateAutoMLTextTrainingJobOperator(AutoMLTrainingJobBaseOperator):
     """Create Auto ML Text Training job."""
 
     template_fields = [
+        "parent_model",
         "dataset_id",
         "region",
         "impersonation_chain",
@@ -439,6 +471,7 @@ class CreateAutoMLTextTrainingJobOperator(AutoMLTrainingJobBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        self.parent_model = self.parent_model.rpartition("@")[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,
@@ -459,11 +492,16 @@ class CreateAutoMLTextTrainingJobOperator(AutoMLTrainingJobBaseOperator):
             model_display_name=self.model_display_name,
             model_labels=self.model_labels,
             sync=self.sync,
+            parent_model=self.parent_model,
+            is_default_version=self.is_default_version,
+            model_version_aliases=self.model_version_aliases,
+            model_version_description=self.model_version_description,
         )
 
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
+            self.xcom_push(context, key="model_id", value=model_id)
             VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
         else:
             result = model  # type: ignore
@@ -476,6 +514,7 @@ class CreateAutoMLVideoTrainingJobOperator(AutoMLTrainingJobBaseOperator):
     """Create Auto ML Video Training job."""
 
     template_fields = (
+        "parent_model",
         "dataset_id",
         "region",
         "impersonation_chain",
@@ -504,6 +543,7 @@ class CreateAutoMLVideoTrainingJobOperator(AutoMLTrainingJobBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        self.parent_model = self.parent_model.rpartition("@")[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,
@@ -521,11 +561,16 @@ class CreateAutoMLVideoTrainingJobOperator(AutoMLTrainingJobBaseOperator):
             model_display_name=self.model_display_name,
             model_labels=self.model_labels,
             sync=self.sync,
+            parent_model=self.parent_model,
+            is_default_version=self.is_default_version,
+            model_version_aliases=self.model_version_aliases,
+            model_version_description=self.model_version_description,
         )
 
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
+            self.xcom_push(context, key="model_id", value=model_id)
             VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
         else:
             result = model  # type: ignore

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

@@ -15,15 +15,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module contains Google Vertex AI operators.
 
-.. spelling:word-list::
+"""This module contains Google Vertex AI operators."""
 
-    jsonl
-    codepoints
-    aiplatform
-    gapic
-"""
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence
@@ -54,7 +48,7 @@ class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator):
     :param region: Required. The ID of the Google Cloud region that the service belongs to.
     :param batch_prediction_job: Required. The BatchPredictionJob to create.
     :param job_display_name: Required. The user-defined name of the BatchPredictionJob. The name can be
-        up to 128 characters long and can be consist of any UTF-8 characters.
+        up to 128 characters long and can consist of any UTF-8 characters.
     :param model_name: Required. A fully-qualified model resource name or model ID.
     :param instances_format: Required. The format in which instances are provided. Must be one of the
         formats listed in `Model.supported_input_storage_formats`. Default is "jsonl" when using