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

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

@@ -19,7 +19,9 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+import warnings
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Sequence
 
 from deprecated import deprecated
 from google.api_core.exceptions import NotFound
@@ -28,7 +30,8 @@ from google.cloud.aiplatform.models import Model
 from google.cloud.aiplatform_v1.types.dataset import Dataset
 from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
 
-from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.vertex_ai.custom_job import CustomJobHook
 from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIModelLink,
@@ -36,9 +39,19 @@ from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAITrainingPipelinesLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.cloud.triggers.vertex_ai import (
+    CustomContainerTrainingJobTrigger,
+    CustomPythonPackageTrainingJobTrigger,
+    CustomTrainingJobTrigger,
+)
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
+    from google.cloud.aiplatform import (
+        CustomContainerTrainingJob,
+        CustomPythonPackageTrainingJob,
+        CustomTrainingJob,
+    )
 
     from airflow.utils.context import Context
 
@@ -160,6 +173,13 @@ class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    def execute(self, context: Context) -> None:
+        warnings.warn(
+            "The 'sync' parameter is deprecated and will be removed after 01.10.2024.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=2,
+        )
+
 
 class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
     """Create Custom Container Training job.
@@ -421,9 +441,6 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
             ``projects/{project}/locations/{location}/tensorboards/{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 the AI Platform job synchronously. If False, this method
-            will be executed in concurrent Future and any downstream object will
-            be immediately returned and synced when the Future has completed.
     :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
@@ -433,6 +450,9 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
         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.
+    :param poll_interval: Time (seconds) to wait between two consecutive calls to check the job.
+        The default is 60 seconds.
     """
 
     template_fields = (
@@ -442,7 +462,10 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
         "dataset_id",
         "impersonation_chain",
     )
-    operator_extra_links = (VertexAIModelLink(), VertexAITrainingLink())
+    operator_extra_links = (
+        VertexAIModelLink(),
+        VertexAITrainingLink(),
+    )
 
     def __init__(
         self,
@@ -452,6 +475,8 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
         parent_model: str | None = None,
         impersonation_chain: str | Sequence[str] | None = None,
         dataset_id: str | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        poll_interval: int = 60,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -462,12 +487,15 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
             **kwargs,
         )
         self.command = command
+        self.deferrable = deferrable
+        self.poll_interval = poll_interval
 
     def execute(self, context: Context):
-        self.hook = CustomJobHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
+        super().execute(context)
+
+        if self.deferrable:
+            self.invoke_defer(context=context)
+
         model, training_id, custom_job_id = self.hook.create_custom_container_training_job(
             project_id=self.project_id,
             region=self.region,
@@ -539,6 +567,94 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
         if self.hook:
             self.hook.cancel_job()
 
+    def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any] | None:
+        if event["status"] == "error":
+            raise AirflowException(event["message"])
+        result = event["job"]
+        model_id = self.hook.extract_model_id_from_training_pipeline(result)
+        custom_job_id = self.hook.extract_custom_job_id_from_training_pipeline(result)
+        self.xcom_push(context, key="model_id", value=model_id)
+        VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+        # push custom_job_id to xcom so it could be pulled by other tasks
+        self.xcom_push(context, key="custom_job_id", value=custom_job_id)
+        return result
+
+    def invoke_defer(self, context: Context) -> None:
+        custom_container_training_job_obj: CustomContainerTrainingJob = self.hook.submit_custom_container_training_job(
+            project_id=self.project_id,
+            region=self.region,
+            display_name=self.display_name,
+            command=self.command,
+            container_uri=self.container_uri,
+            model_serving_container_image_uri=self.model_serving_container_image_uri,
+            model_serving_container_predict_route=self.model_serving_container_predict_route,
+            model_serving_container_health_route=self.model_serving_container_health_route,
+            model_serving_container_command=self.model_serving_container_command,
+            model_serving_container_args=self.model_serving_container_args,
+            model_serving_container_environment_variables=self.model_serving_container_environment_variables,
+            model_serving_container_ports=self.model_serving_container_ports,
+            model_description=self.model_description,
+            model_instance_schema_uri=self.model_instance_schema_uri,
+            model_parameters_schema_uri=self.model_parameters_schema_uri,
+            model_prediction_schema_uri=self.model_prediction_schema_uri,
+            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,
+            labels=self.labels,
+            training_encryption_spec_key_name=self.training_encryption_spec_key_name,
+            model_encryption_spec_key_name=self.model_encryption_spec_key_name,
+            staging_bucket=self.staging_bucket,
+            # RUN
+            dataset=Dataset(name=self.dataset_id) if self.dataset_id else None,
+            annotation_schema_uri=self.annotation_schema_uri,
+            model_display_name=self.model_display_name,
+            model_labels=self.model_labels,
+            base_output_dir=self.base_output_dir,
+            service_account=self.service_account,
+            network=self.network,
+            bigquery_destination=self.bigquery_destination,
+            args=self.args,
+            environment_variables=self.environment_variables,
+            replica_count=self.replica_count,
+            machine_type=self.machine_type,
+            accelerator_type=self.accelerator_type,
+            accelerator_count=self.accelerator_count,
+            boot_disk_type=self.boot_disk_type,
+            boot_disk_size_gb=self.boot_disk_size_gb,
+            training_fraction_split=self.training_fraction_split,
+            validation_fraction_split=self.validation_fraction_split,
+            test_fraction_split=self.test_fraction_split,
+            training_filter_split=self.training_filter_split,
+            validation_filter_split=self.validation_filter_split,
+            test_filter_split=self.test_filter_split,
+            predefined_split_column_name=self.predefined_split_column_name,
+            timestamp_split_column_name=self.timestamp_split_column_name,
+            tensorboard=self.tensorboard,
+        )
+        custom_container_training_job_obj.wait_for_resource_creation()
+        training_pipeline_id: str = custom_container_training_job_obj.name
+        self.xcom_push(context, key="training_id", value=training_pipeline_id)
+        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_pipeline_id)
+        self.defer(
+            trigger=CustomContainerTrainingJobTrigger(
+                conn_id=self.gcp_conn_id,
+                project_id=self.project_id,
+                location=self.region,
+                job_id=training_pipeline_id,
+                poll_interval=self.poll_interval,
+                impersonation_chain=self.impersonation_chain,
+            ),
+            method_name="execute_complete",
+        )
+
+    @cached_property
+    def hook(self) -> CustomJobHook:
+        return CustomJobHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
 
 class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator):
     """Create Custom Python Package Training job.
@@ -800,9 +916,6 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
             ``projects/{project}/locations/{location}/tensorboards/{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 the AI Platform job synchronously. If False, this method
-            will be executed in concurrent Future and any downstream object will
-            be immediately returned and synced when the Future has completed.
     :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
@@ -812,6 +925,9 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
         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.
+    :param poll_interval: Time (seconds) to wait between two consecutive calls to check the job.
+        The default is 60 seconds.
     """
 
     template_fields = (
@@ -831,6 +947,8 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
         parent_model: str | None = None,
         impersonation_chain: str | Sequence[str] | None = None,
         dataset_id: str | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        poll_interval: int = 60,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -842,12 +960,15 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
         )
         self.python_package_gcs_uri = python_package_gcs_uri
         self.python_module_name = python_module_name
+        self.deferrable = deferrable
+        self.poll_interval = poll_interval
 
     def execute(self, context: Context):
-        self.hook = CustomJobHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
+        super().execute(context)
+
+        if self.deferrable:
+            self.invoke_defer(context=context)
+
         model, training_id, custom_job_id = self.hook.create_custom_python_package_training_job(
             project_id=self.project_id,
             region=self.region,
@@ -920,9 +1041,98 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
         if self.hook:
             self.hook.cancel_job()
 
+    def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any] | None:
+        if event["status"] == "error":
+            raise AirflowException(event["message"])
+        result = event["job"]
+        model_id = self.hook.extract_model_id_from_training_pipeline(result)
+        custom_job_id = self.hook.extract_custom_job_id_from_training_pipeline(result)
+        self.xcom_push(context, key="model_id", value=model_id)
+        VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+        # push custom_job_id to xcom so it could be pulled by other tasks
+        self.xcom_push(context, key="custom_job_id", value=custom_job_id)
+        return result
+
+    def invoke_defer(self, context: Context) -> None:
+        custom_python_training_job_obj: CustomPythonPackageTrainingJob = self.hook.submit_custom_python_package_training_job(
+            project_id=self.project_id,
+            region=self.region,
+            display_name=self.display_name,
+            python_package_gcs_uri=self.python_package_gcs_uri,
+            python_module_name=self.python_module_name,
+            container_uri=self.container_uri,
+            model_serving_container_image_uri=self.model_serving_container_image_uri,
+            model_serving_container_predict_route=self.model_serving_container_predict_route,
+            model_serving_container_health_route=self.model_serving_container_health_route,
+            model_serving_container_command=self.model_serving_container_command,
+            model_serving_container_args=self.model_serving_container_args,
+            model_serving_container_environment_variables=self.model_serving_container_environment_variables,
+            model_serving_container_ports=self.model_serving_container_ports,
+            model_description=self.model_description,
+            model_instance_schema_uri=self.model_instance_schema_uri,
+            model_parameters_schema_uri=self.model_parameters_schema_uri,
+            model_prediction_schema_uri=self.model_prediction_schema_uri,
+            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,
+            labels=self.labels,
+            training_encryption_spec_key_name=self.training_encryption_spec_key_name,
+            model_encryption_spec_key_name=self.model_encryption_spec_key_name,
+            staging_bucket=self.staging_bucket,
+            # RUN
+            dataset=Dataset(name=self.dataset_id) if self.dataset_id else None,
+            annotation_schema_uri=self.annotation_schema_uri,
+            model_display_name=self.model_display_name,
+            model_labels=self.model_labels,
+            base_output_dir=self.base_output_dir,
+            service_account=self.service_account,
+            network=self.network,
+            bigquery_destination=self.bigquery_destination,
+            args=self.args,
+            environment_variables=self.environment_variables,
+            replica_count=self.replica_count,
+            machine_type=self.machine_type,
+            accelerator_type=self.accelerator_type,
+            accelerator_count=self.accelerator_count,
+            boot_disk_type=self.boot_disk_type,
+            boot_disk_size_gb=self.boot_disk_size_gb,
+            training_fraction_split=self.training_fraction_split,
+            validation_fraction_split=self.validation_fraction_split,
+            test_fraction_split=self.test_fraction_split,
+            training_filter_split=self.training_filter_split,
+            validation_filter_split=self.validation_filter_split,
+            test_filter_split=self.test_filter_split,
+            predefined_split_column_name=self.predefined_split_column_name,
+            timestamp_split_column_name=self.timestamp_split_column_name,
+            tensorboard=self.tensorboard,
+        )
+        custom_python_training_job_obj.wait_for_resource_creation()
+        training_pipeline_id: str = custom_python_training_job_obj.name
+        self.xcom_push(context, key="training_id", value=training_pipeline_id)
+        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_pipeline_id)
+        self.defer(
+            trigger=CustomPythonPackageTrainingJobTrigger(
+                conn_id=self.gcp_conn_id,
+                project_id=self.project_id,
+                location=self.region,
+                job_id=training_pipeline_id,
+                poll_interval=self.poll_interval,
+                impersonation_chain=self.impersonation_chain,
+            ),
+            method_name="execute_complete",
+        )
+
+    @cached_property
+    def hook(self) -> CustomJobHook:
+        return CustomJobHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
 
 class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
-    """Create Custom Training job.
+    """Create a Custom Training Job pipeline.
 
     :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.
@@ -1181,9 +1391,6 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
             ``projects/{project}/locations/{location}/tensorboards/{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 the AI Platform job synchronously. If False, this method
-            will be executed in concurrent Future and any downstream object will
-            be immediately returned and synced when the Future has completed.
     :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
@@ -1193,6 +1400,9 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
         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.
+    :param poll_interval: Time (seconds) to wait between two consecutive calls to check the job.
+        The default is 60 seconds.
     """
 
     template_fields = (
@@ -1203,7 +1413,10 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
         "dataset_id",
         "impersonation_chain",
     )
-    operator_extra_links = (VertexAIModelLink(), VertexAITrainingLink())
+    operator_extra_links = (
+        VertexAIModelLink(),
+        VertexAITrainingLink(),
+    )
 
     def __init__(
         self,
@@ -1214,6 +1427,8 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
         parent_model: str | None = None,
         impersonation_chain: str | Sequence[str] | None = None,
         dataset_id: str | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        poll_interval: int = 60,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -1225,12 +1440,15 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
         )
         self.requirements = requirements
         self.script_path = script_path
+        self.deferrable = deferrable
+        self.poll_interval = poll_interval
 
     def execute(self, context: Context):
-        self.hook = CustomJobHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
+        super().execute(context)
+
+        if self.deferrable:
+            self.invoke_defer(context=context)
+
         model, training_id, custom_job_id = self.hook.create_custom_training_job(
             project_id=self.project_id,
             region=self.region,
@@ -1303,6 +1521,95 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
         if self.hook:
             self.hook.cancel_job()
 
+    def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any] | None:
+        if event["status"] == "error":
+            raise AirflowException(event["message"])
+        result = event["job"]
+        model_id = self.hook.extract_model_id_from_training_pipeline(result)
+        custom_job_id = self.hook.extract_custom_job_id_from_training_pipeline(result)
+        self.xcom_push(context, key="model_id", value=model_id)
+        VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+        # push custom_job_id to xcom so it could be pulled by other tasks
+        self.xcom_push(context, key="custom_job_id", value=custom_job_id)
+        return result
+
+    def invoke_defer(self, context: Context) -> None:
+        custom_training_job_obj: CustomTrainingJob = self.hook.submit_custom_training_job(
+            project_id=self.project_id,
+            region=self.region,
+            display_name=self.display_name,
+            script_path=self.script_path,
+            container_uri=self.container_uri,
+            requirements=self.requirements,
+            model_serving_container_image_uri=self.model_serving_container_image_uri,
+            model_serving_container_predict_route=self.model_serving_container_predict_route,
+            model_serving_container_health_route=self.model_serving_container_health_route,
+            model_serving_container_command=self.model_serving_container_command,
+            model_serving_container_args=self.model_serving_container_args,
+            model_serving_container_environment_variables=self.model_serving_container_environment_variables,
+            model_serving_container_ports=self.model_serving_container_ports,
+            model_description=self.model_description,
+            model_instance_schema_uri=self.model_instance_schema_uri,
+            model_parameters_schema_uri=self.model_parameters_schema_uri,
+            model_prediction_schema_uri=self.model_prediction_schema_uri,
+            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,
+            labels=self.labels,
+            training_encryption_spec_key_name=self.training_encryption_spec_key_name,
+            model_encryption_spec_key_name=self.model_encryption_spec_key_name,
+            staging_bucket=self.staging_bucket,
+            # RUN
+            dataset=Dataset(name=self.dataset_id) if self.dataset_id else None,
+            annotation_schema_uri=self.annotation_schema_uri,
+            model_display_name=self.model_display_name,
+            model_labels=self.model_labels,
+            base_output_dir=self.base_output_dir,
+            service_account=self.service_account,
+            network=self.network,
+            bigquery_destination=self.bigquery_destination,
+            args=self.args,
+            environment_variables=self.environment_variables,
+            replica_count=self.replica_count,
+            machine_type=self.machine_type,
+            accelerator_type=self.accelerator_type,
+            accelerator_count=self.accelerator_count,
+            boot_disk_type=self.boot_disk_type,
+            boot_disk_size_gb=self.boot_disk_size_gb,
+            training_fraction_split=self.training_fraction_split,
+            validation_fraction_split=self.validation_fraction_split,
+            test_fraction_split=self.test_fraction_split,
+            training_filter_split=self.training_filter_split,
+            validation_filter_split=self.validation_filter_split,
+            test_filter_split=self.test_filter_split,
+            predefined_split_column_name=self.predefined_split_column_name,
+            timestamp_split_column_name=self.timestamp_split_column_name,
+            tensorboard=self.tensorboard,
+        )
+        custom_training_job_obj.wait_for_resource_creation()
+        training_pipeline_id: str = custom_training_job_obj.name
+        self.xcom_push(context, key="training_id", value=training_pipeline_id)
+        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_pipeline_id)
+        self.defer(
+            trigger=CustomTrainingJobTrigger(
+                conn_id=self.gcp_conn_id,
+                project_id=self.project_id,
+                location=self.region,
+                job_id=training_pipeline_id,
+                poll_interval=self.poll_interval,
+                impersonation_chain=self.impersonation_chain,
+            ),
+            method_name="execute_complete",
+        )
+
+    @cached_property
+    def hook(self) -> CustomJobHook:
+        return CustomJobHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
 
 class DeleteCustomTrainingJobOperator(GoogleCloudBaseOperator):
     """

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

@@ -33,11 +33,11 @@ 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.
+        service belongs to (templated).
     :param location: Required. The ID of the Google Cloud location that the
-        service belongs to.
+        service belongs to (templated).
     :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.
+        to the Vertex AI PaLM API, in order to elicit a specific response (templated).
     :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.
@@ -60,6 +60,8 @@ class PromptLanguageModelOperator(GoogleCloudBaseOperator):
         account from the list granting this role to the originating account (templated).
     """
 
+    template_fields = ("location", "project_id", "impersonation_chain", "prompt")
+
     def __init__(
         self,
         *,
@@ -116,11 +118,11 @@ 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.
+        service belongs to (templated).
     :param location: Required. The ID of the Google Cloud location that the
-        service belongs to.
+        service belongs to (templated).
     :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.
+        to the Vertex AI PaLM API, in order to elicit a specific response (templated).
     :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.
@@ -134,6 +136,8 @@ class GenerateTextEmbeddingsOperator(GoogleCloudBaseOperator):
         account from the list granting this role to the originating account (templated).
     """
 
+    template_fields = ("location", "project_id", "impersonation_chain", "prompt")
+
     def __init__(
         self,
         *,
@@ -178,11 +182,11 @@ 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.
+        service belongs to (templated).
     :param location: Required. The ID of the Google Cloud location that the
-        service belongs to.
+        service belongs to (templated).
     :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.
+        to the Multi-modal model, in order to elicit a specific response (templated).
     :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
@@ -198,6 +202,8 @@ class PromptMultimodalModelOperator(GoogleCloudBaseOperator):
         account from the list granting this role to the originating account (templated).
     """
 
+    template_fields = ("location", "project_id", "impersonation_chain", "prompt")
+
     def __init__(
         self,
         *,
@@ -240,11 +246,11 @@ 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.
+        service belongs to (templated).
     :param location: Required. The ID of the Google Cloud location that the
-        service belongs to.
+        service belongs to (templated).
     :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.
+        to the Multi-modal model, in order to elicit a specific response (templated).
     :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
@@ -263,6 +269,8 @@ class PromptMultimodalModelWithMediaOperator(GoogleCloudBaseOperator):
         account from the list granting this role to the originating account (templated).
     """
 
+    template_fields = ("location", "project_id", "impersonation_chain", "prompt")
+
     def __init__(
         self,
         *,

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

@@ -102,7 +102,6 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
         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.
     """