apache-airflow-providers-google 10.16.0rc1__py3-none-any.whl → 10.17.0__py3-none-any.whl

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

@@ -0,0 +1,306 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# 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 Generative AI operators."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import GenerativeModelHook
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class PromptLanguageModelOperator(GoogleCloudBaseOperator):
+    """
+    Uses the Vertex AI PaLM API to generate natural language text.
+
+    :param project_id: Required. The ID of the Google Cloud project that the
+        service belongs to.
+    :param location: Required. The ID of the Google Cloud location that the
+        service belongs to.
+    :param prompt: Required. Inputs or queries that a user or a program gives
+        to the Vertex AI PaLM API, in order to elicit a specific response.
+    :param pretrained_model: By default uses the pre-trained model `text-bison`,
+        optimized for performing natural language tasks such as classification,
+        summarization, extraction, content creation, and ideation.
+    :param temperature: Temperature controls the degree of randomness in token
+        selection. Defaults to 0.0.
+    :param max_output_tokens: Token limit determines the maximum amount of text
+        output. Defaults to 256.
+    :param top_p: Tokens are selected from most probable to least until the sum
+        of their probabilities equals the top_p value. Defaults to 0.8.
+    :param top_k: A top_k of 1 means the selected token is the most probable
+        among all tokens. Defaults to 0.4.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        prompt: str,
+        pretrained_model: str = "text-bison",
+        temperature: float = 0.0,
+        max_output_tokens: int = 256,
+        top_p: float = 0.8,
+        top_k: int = 40,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.prompt = prompt
+        self.pretrained_model = pretrained_model
+        self.temperature = temperature
+        self.max_output_tokens = max_output_tokens
+        self.top_p = top_p
+        self.top_k = top_k
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        self.hook = GenerativeModelHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        self.log.info("Submitting prompt")
+        response = self.hook.prompt_language_model(
+            project_id=self.project_id,
+            location=self.location,
+            prompt=self.prompt,
+            pretrained_model=self.pretrained_model,
+            temperature=self.temperature,
+            max_output_tokens=self.max_output_tokens,
+            top_p=self.top_p,
+            top_k=self.top_k,
+        )
+
+        self.log.info("Model response: %s", response)
+        self.xcom_push(context, key="prompt_response", value=response)
+
+        return response
+
+
+class GenerateTextEmbeddingsOperator(GoogleCloudBaseOperator):
+    """
+    Uses the Vertex AI PaLM API to generate natural language text.
+
+    :param project_id: Required. The ID of the Google Cloud project that the
+        service belongs to.
+    :param location: Required. The ID of the Google Cloud location that the
+        service belongs to.
+    :param prompt: Required. Inputs or queries that a user or a program gives
+        to the Vertex AI PaLM API, in order to elicit a specific response.
+    :param pretrained_model: By default uses the pre-trained model `textembedding-gecko`,
+        optimized for performing text embeddings.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        prompt: str,
+        pretrained_model: str = "textembedding-gecko",
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.prompt = prompt
+        self.pretrained_model = pretrained_model
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        self.hook = GenerativeModelHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        self.log.info("Generating text embeddings")
+        response = self.hook.generate_text_embeddings(
+            project_id=self.project_id,
+            location=self.location,
+            prompt=self.prompt,
+            pretrained_model=self.pretrained_model,
+        )
+
+        self.log.info("Model response: %s", response)
+        self.xcom_push(context, key="prompt_response", value=response)
+
+        return response
+
+
+class PromptMultimodalModelOperator(GoogleCloudBaseOperator):
+    """
+    Use the Vertex AI Gemini Pro foundation model to generate natural language text.
+
+    :param project_id: Required. The ID of the Google Cloud project that the
+        service belongs to.
+    :param location: Required. The ID of the Google Cloud location that the
+        service belongs to.
+    :param prompt: Required. Inputs or queries that a user or a program gives
+        to the Multi-modal model, in order to elicit a specific response.
+    :param pretrained_model: By default uses the pre-trained model `gemini-pro`,
+        supporting prompts with text-only input, including natural language
+        tasks, multi-turn text and code chat, and code generation. It can
+        output text and code.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        prompt: str,
+        pretrained_model: str = "gemini-pro",
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.prompt = prompt
+        self.pretrained_model = pretrained_model
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        self.hook = GenerativeModelHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        response = self.hook.prompt_multimodal_model(
+            project_id=self.project_id,
+            location=self.location,
+            prompt=self.prompt,
+            pretrained_model=self.pretrained_model,
+        )
+
+        self.log.info("Model response: %s", response)
+        self.xcom_push(context, key="prompt_response", value=response)
+
+        return response
+
+
+class PromptMultimodalModelWithMediaOperator(GoogleCloudBaseOperator):
+    """
+    Use the Vertex AI Gemini Pro foundation model to generate natural language text.
+
+    :param project_id: Required. The ID of the Google Cloud project that the
+        service belongs to.
+    :param location: Required. The ID of the Google Cloud location that the
+        service belongs to.
+    :param prompt: Required. Inputs or queries that a user or a program gives
+        to the Multi-modal model, in order to elicit a specific response.
+    :param pretrained_model: By default uses the pre-trained model `gemini-pro-vision`,
+        supporting prompts with text-only input, including natural language
+        tasks, multi-turn text and code chat, and code generation. It can
+        output text and code.
+    :param media_gcs_path: A GCS path to a media file such as an image or a video.
+        Can be passed to the multi-modal model as part of the prompt. Used with vision models.
+    :param mime_type: Validates the media type presented by the file in the media_gcs_path.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        prompt: str,
+        media_gcs_path: str,
+        mime_type: str,
+        pretrained_model: str = "gemini-pro-vision",
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.prompt = prompt
+        self.pretrained_model = pretrained_model
+        self.media_gcs_path = media_gcs_path
+        self.mime_type = mime_type
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context):
+        self.hook = GenerativeModelHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        response = self.hook.prompt_multimodal_model_with_media(
+            project_id=self.project_id,
+            location=self.location,
+            prompt=self.prompt,
+            pretrained_model=self.pretrained_model,
+            media_gcs_path=self.media_gcs_path,
+            mime_type=self.mime_type,
+        )
+
+        self.log.info("Model response: %s", response)
+        self.xcom_push(context, key="prompt_response", value=response)
+
+        return response

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

@@ -20,14 +20,15 @@
 
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING, Any, Sequence
 
 from google.api_core.exceptions import NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform_v1.types import HyperparameterTuningJob
+from google.cloud.aiplatform_v1 import types
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job import (
     HyperparameterTuningJobHook,
 )
@@ -40,7 +41,7 @@ from airflow.providers.google.cloud.triggers.vertex_ai import CreateHyperparamet
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
-    from google.cloud.aiplatform import gapic, hyperparameter_tuning
+    from google.cloud.aiplatform import HyperparameterTuningJob, gapic, hyperparameter_tuning
 
     from airflow.utils.context import Context
 
@@ -127,8 +128,8 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
         `service_account` is required with provided `tensorboard`. For more information on configuring
         your service account please visit:
         https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
-    :param sync: Whether to execute this method synchronously. If False, this method will unblock, and it
-        will be executed in a concurrent Future.
+    :param sync: (Deprecated) Whether to execute this method synchronously. If False, this method will
+        unblock, and it will be executed in a concurrent Future.
     :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
     :param impersonation_chain: Optional service account to impersonate using short-term
         credentials, or chained list of accounts required to get the access_token
@@ -138,8 +139,7 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
         If set as a sequence, the identities from the list must grant
         Service Account Token Creator IAM role to the directly preceding identity, with first
         account from the list granting this role to the originating account (templated).
-    :param deferrable: Run operator in the deferrable mode. Note that it requires calling the operator
-        with `sync=False` parameter.
+    :param deferrable: Run operator in the deferrable mode.
     :param poll_interval: Interval size which defines how often job status is checked in deferrable mode.
     """
 
@@ -221,19 +221,18 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
         self.poll_interval = poll_interval
 
     def execute(self, context: Context):
-        if self.deferrable and self.sync:
-            raise AirflowException(
-                "Deferrable mode can be used only with sync=False option. "
-                "If you are willing to run the operator in deferrable mode, please, set sync=False. "
-                "Otherwise, disable deferrable mode `deferrable=False`."
-            )
+        warnings.warn(
+            "The 'sync' parameter is deprecated and will be removed after 01.09.2024.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=2,
+        )
 
         self.log.info("Creating Hyperparameter Tuning job")
         self.hook = HyperparameterTuningJobHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
-        result = self.hook.create_hyperparameter_tuning_job(
+        hyperparameter_tuning_job: HyperparameterTuningJob = self.hook.create_hyperparameter_tuning_job(
             project_id=self.project_id,
             region=self.region,
             display_name=self.display_name,
@@ -259,14 +258,19 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
             restart_job_on_worker_restart=self.restart_job_on_worker_restart,
             enable_web_access=self.enable_web_access,
             tensorboard=self.tensorboard,
-            sync=self.sync,
-            wait_job_completed=not self.deferrable,
+            sync=False,
+            wait_job_completed=False,
         )
 
-        hyperparameter_tuning_job = result.to_dict()
-        hyperparameter_tuning_job_id = self.hook.extract_hyperparameter_tuning_job_id(
-            hyperparameter_tuning_job
+        hyperparameter_tuning_job.wait_for_resource_creation()
+        hyperparameter_tuning_job_id = hyperparameter_tuning_job.name
+        self.log.info("Hyperparameter Tuning job was created. Job id: %s", hyperparameter_tuning_job_id)
+
+        self.xcom_push(context, key="hyperparameter_tuning_job_id", value=hyperparameter_tuning_job_id)
+        VertexAITrainingLink.persist(
+            context=context, task_instance=self, training_id=hyperparameter_tuning_job_id
         )
+
         if self.deferrable:
             self.defer(
                 trigger=CreateHyperparameterTuningJobTrigger(
@@ -279,14 +283,10 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
                 ),
                 method_name="execute_complete",
             )
+            return
 
-        self.log.info("Hyperparameter Tuning job was created. Job id: %s", hyperparameter_tuning_job_id)
-
-        self.xcom_push(context, key="hyperparameter_tuning_job_id", value=hyperparameter_tuning_job_id)
-        VertexAITrainingLink.persist(
-            context=context, task_instance=self, training_id=hyperparameter_tuning_job_id
-        )
-        return hyperparameter_tuning_job
+        hyperparameter_tuning_job.wait_for_completion()
+        return hyperparameter_tuning_job.to_dict()
 
     def on_kill(self) -> None:
         """Act as a callback called when the operator is killed; cancel any running job."""
@@ -298,26 +298,7 @@ class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
             raise AirflowException(event["message"])
         job: dict[str, Any] = event["job"]
         self.log.info("Hyperparameter tuning job %s created and completed successfully.", job["name"])
-        hook = HyperparameterTuningJobHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
-        job_id = hook.extract_hyperparameter_tuning_job_id(job)
-        self.xcom_push(
-            context,
-            key="hyperparameter_tuning_job_id",
-            value=job_id,
-        )
-        self.xcom_push(
-            context,
-            key="training_conf",
-            value={
-                "training_conf_id": job_id,
-                "region": self.region,
-                "project_id": self.project_id,
-            },
-        )
-        return event["job"]
+        return job
 
 
 class GetHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
@@ -387,7 +368,7 @@ class GetHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
                 context=context, task_instance=self, training_id=self.hyperparameter_tuning_job_id
             )
             self.log.info("Hyperparameter tuning job was gotten.")
-            return HyperparameterTuningJob.to_dict(result)
+            return types.HyperparameterTuningJob.to_dict(result)
         except NotFound:
             self.log.info(
                 "The Hyperparameter tuning job %s does not exist.", self.hyperparameter_tuning_job_id
@@ -532,4 +513,4 @@ class ListHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         VertexAIHyperparameterTuningJobListLink.persist(context=context, task_instance=self)
-        return [HyperparameterTuningJob.to_dict(result) for result in results]
+        return [types.HyperparameterTuningJob.to_dict(result) for result in results]

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

@@ -16,23 +16,29 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Vertex AI operators."""
+
 from __future__ import annotations
 
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, Sequence
 
 from google.api_core.exceptions import NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform_v1.types import PipelineJob
+from google.cloud.aiplatform_v1 import types
 
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.pipeline_job import PipelineJobHook
 from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIPipelineJobLink,
     VertexAIPipelineJobListLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.cloud.triggers.vertex_ai import RunPipelineJobTrigger
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
+    from google.cloud.aiplatform import PipelineJob
     from google.cloud.aiplatform.metadata import experiment_resources
 
     from airflow.utils.context import Context
@@ -40,7 +46,7 @@ if TYPE_CHECKING:
 
 class RunPipelineJobOperator(GoogleCloudBaseOperator):
     """
-    Run Pipeline job.
+    Create and run a Pipeline job.
 
     :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.
@@ -82,9 +88,9 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
         Private services access must already be configured for the network. If left unspecified, the
         network set in aiplatform.init will be used. Otherwise, the job is not peered with any network.
     :param create_request_timeout: Optional. The timeout for the create request in seconds.
-    :param experiment: Optional. The Vertex AI experiment name or instance to associate to this
-        PipelineJob. Metrics produced by the PipelineJob as system.Metric Artifacts will be associated as
-        metrics to the current Experiment Run. Pipeline parameters will be associated as parameters to
+    :param experiment: Optional. The Vertex AI experiment name or instance to associate to this PipelineJob.
+        Metrics produced by the PipelineJob as system.Metric Artifacts will be associated as metrics
+        to the current Experiment Run. Pipeline parameters will be associated as parameters to
         the current Experiment Run.
     :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
     :param impersonation_chain: Optional service account to impersonate using short-term
@@ -95,6 +101,10 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
         If set as a sequence, the identities from the list must grant
         Service Account Token Creator IAM role to the directly preceding identity, with first
         account from the list granting this role to the originating account (templated).
+    :param deferrable: If True, run the task in the deferrable mode.
+        Note that it requires calling the operator with `sync=False` parameter.
+    :param poll_interval: Time (seconds) to wait between two consecutive calls to check the job.
+        The default is 300 seconds.
     """
 
     template_fields = [
@@ -126,6 +136,8 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
         experiment: str | experiment_resources.Experiment | None = None,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        poll_interval: int = 5 * 60,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -147,15 +159,12 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
         self.experiment = experiment
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
-        self.hook: PipelineJobHook | None = None
+        self.deferrable = deferrable
+        self.poll_interval = poll_interval
 
     def execute(self, context: Context):
         self.log.info("Running Pipeline job")
-        self.hook = PipelineJobHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
-        result = self.hook.run_pipeline_job(
+        pipeline_job_obj: PipelineJob = self.hook.submit_pipeline_job(
             project_id=self.project_id,
             region=self.region,
             display_name=self.display_name,
@@ -173,20 +182,46 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
             create_request_timeout=self.create_request_timeout,
             experiment=self.experiment,
         )
-
-        pipeline_job = result.to_dict()
-        pipeline_job_id = self.hook.extract_pipeline_job_id(pipeline_job)
+        pipeline_job_id = pipeline_job_obj.job_id
         self.log.info("Pipeline job was created. Job id: %s", pipeline_job_id)
-
         self.xcom_push(context, key="pipeline_job_id", value=pipeline_job_id)
         VertexAIPipelineJobLink.persist(context=context, task_instance=self, pipeline_id=pipeline_job_id)
+
+        if self.deferrable:
+            pipeline_job_obj.wait_for_resource_creation()
+            self.defer(
+                trigger=RunPipelineJobTrigger(
+                    conn_id=self.gcp_conn_id,
+                    project_id=self.project_id,
+                    location=pipeline_job_obj.location,
+                    job_id=pipeline_job_id,
+                    poll_interval=self.poll_interval,
+                    impersonation_chain=self.impersonation_chain,
+                ),
+                method_name="execute_complete",
+            )
+
+        pipeline_job_obj.wait()
+        pipeline_job = pipeline_job_obj.to_dict()
         return pipeline_job
 
+    def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any]:
+        if event["status"] == "error":
+            raise AirflowException(event["message"])
+        return event["job"]
+
     def on_kill(self) -> None:
         """Act as a callback called when the operator is killed; cancel any running job."""
         if self.hook:
             self.hook.cancel_pipeline_job()
 
+    @cached_property
+    def hook(self) -> PipelineJobHook:
+        return PipelineJobHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
 
 class GetPipelineJobOperator(GoogleCloudBaseOperator):
     """
@@ -261,7 +296,7 @@ class GetPipelineJobOperator(GoogleCloudBaseOperator):
                 context=context, task_instance=self, pipeline_id=self.pipeline_job_id
             )
             self.log.info("Pipeline job was gotten.")
-            return PipelineJob.to_dict(result)
+            return types.PipelineJob.to_dict(result)
         except NotFound:
             self.log.info("The Pipeline job %s does not exist.", self.pipeline_job_id)
 
@@ -389,7 +424,7 @@ class ListPipelineJobOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         VertexAIPipelineJobListLink.persist(context=context, task_instance=self)
-        return [PipelineJob.to_dict(result) for result in results]
+        return [types.PipelineJob.to_dict(result) for result in results]
 
 
 class DeletePipelineJobOperator(GoogleCloudBaseOperator):

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud Vision operators."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Vision operator."""
+
 from __future__ import annotations
 
 from copy import deepcopy

airflow/providers/google/cloud/secrets/secret_manager.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Objects relating to sourcing connections from Google Cloud Secrets Manager."""
+
 from __future__ import annotations
 
 import logging

airflow/providers/google/cloud/sensors/bigquery.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google BigQuery sensors."""
+
 from __future__ import annotations
 
 import warnings

airflow/providers/google/cloud/sensors/bigquery_dts.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google BigQuery Data Transfer Service sensor."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

airflow/providers/google/cloud/sensors/bigtable.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud Bigtable sensor."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence