Add CreateCustomJobOperator for VertexAI service

generated/known_airflow_exceptions.txt

@@ -228,14 +228,13 @@ providers/google/src/airflow/providers/google/cloud/hooks/gdm.py::1
 providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py::2
 providers/google/src/airflow/providers/google/cloud/hooks/looker.py::8
 providers/google/src/airflow/providers/google/cloud/hooks/managed_kafka.py::1
-providers/google/src/airflow/providers/google/cloud/hooks/mlengine.py::2
 providers/google/src/airflow/providers/google/cloud/hooks/spanner.py::5
 providers/google/src/airflow/providers/google/cloud/hooks/stackdriver.py::2
 providers/google/src/airflow/providers/google/cloud/hooks/tasks.py::3
 providers/google/src/airflow/providers/google/cloud/hooks/translate.py::3
 providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py::3
 providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py::1
-providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py::8
+providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py::9
 providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py::2
 providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py::1
 providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py::2
@@ -409,7 +408,8 @@ providers/trino/src/airflow/providers/trino/hooks/trino.py::1
 providers/vespa/src/airflow/providers/vespa/operators/vespa_ingest.py::1
 providers/ydb/src/airflow/providers/ydb/hooks/ydb.py::1
 providers/ydb/src/airflow/providers/ydb/operators/ydb.py::1
-scripts/ci/prek/check_new_airflow_exception_usage.py::5
+scripts/ci/prek/check_new_airflow_exception_usage.py::4
+scripts/tests/ci/prek/test_check_new_airflow_exception_usage.py::9
 task-sdk/src/airflow/sdk/bases/sensor.py::1
 task-sdk/src/airflow/sdk/bases/skipmixin.py::3
 task-sdk/src/airflow/sdk/crypto.py::1

providers/google/provider.yaml

@@ -1416,6 +1416,7 @@ extra-links:
   - airflow.providers.google.cloud.links.vertex_ai.VertexAIPipelineJobListLink
   - airflow.providers.google.cloud.links.vertex_ai.VertexAIRayClusterLink
   - airflow.providers.google.cloud.links.vertex_ai.VertexAIRayClusterListLink
+  - airflow.providers.google.cloud.links.vertex_ai.VertexAICustomJobLink
   - airflow.providers.google.cloud.links.workflows.WorkflowsWorkflowDetailsLink
   - airflow.providers.google.cloud.links.workflows.WorkflowsListOfWorkflowsLink
   - airflow.providers.google.cloud.links.workflows.WorkflowsExecutionLink

providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py

@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 import asyncio
+import time
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
@@ -3093,17 +3094,63 @@ def list_custom_jobs(
         )
         return result
 
+    @GoogleBaseHook.fallback_to_default_project_id
+    def wait_for_custom_job(
+        self,
+        project_id: str,
+        region: str,
+        custom_job_id: str,
+        poll_interval: int = 10,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> CustomJob:
+        """
+        Wait until an VertexAI Custom job completes.
+
+        :param project_id: Required. The ID of the Google Cloud project.
+        :param region: Required. The ID of the Google Cloud location that the service belongs to.
+        :param custom_job_id: Required. The ID of the CustomJob to wait.
+        :param poll_interval: Time, in seconds, to wait between checks.
+        :param retry: Designation of what errors, if any, should be retried.
+        :param timeout: The timeout for this request.
+        :param metadata: Strings which should be sent along with the request as metadata.
+        """
+        CUSTOB_JOB_FAILED_STATES = {
+            JobState.JOB_STATE_FAILED: lambda: custom_job.error.message,  # type: ignore
+            JobState.JOB_STATE_CANCELLED: lambda: "The CustomJob has been cancelled.",
+            JobState.JOB_STATE_PAUSED: lambda: "The CustomJob has been stopped, and can be resumed.",
+            JobState.JOB_STATE_EXPIRED: lambda: "The CustomJob has expired.",
+            JobState.JOB_STATE_PARTIALLY_SUCCEEDED: lambda: custom_job.error.message,  # type: ignore
+        }
+
+        while True:
+            self.log.info("Waiting for custom job with id %s", custom_job_id)
+            try:
+                custom_job = self.get_custom_job(
+                    project_id=project_id,
+                    region=region,
+                    custom_job=custom_job_id,
+                    retry=retry,
+                    timeout=timeout,
+                    metadata=metadata,
+                )
+            except Exception as ex:
+                self.log.exception("Exception occurred while waiting job %s", custom_job_id)
+                raise AirflowException(ex)
+            self.log.info("Status of the custom job %s is %s", custom_job.name, custom_job.state.name)
+            if custom_job.state == JobState.JOB_STATE_SUCCEEDED:
+                return custom_job
+            if custom_job.state in CUSTOB_JOB_FAILED_STATES:
+                raise RuntimeError(CUSTOB_JOB_FAILED_STATES[custom_job.state]())
+            self.log.info("Sleeping for %s seconds.", poll_interval)
+            time.sleep(poll_interval)
+
 
 class CustomJobAsyncHook(GoogleBaseAsyncHook):
     """Async hook for Custom Job Service Client."""
 
     sync_hook_class = CustomJobHook
-    JOB_COMPLETE_STATES = {
-        JobState.JOB_STATE_CANCELLED,
-        JobState.JOB_STATE_FAILED,
-        JobState.JOB_STATE_PAUSED,
-        JobState.JOB_STATE_SUCCEEDED,
-    }
     PIPELINE_COMPLETE_STATES = (
         PipelineState.PIPELINE_STATE_CANCELLED,
         PipelineState.PIPELINE_STATE_FAILED,
@@ -3238,6 +3285,14 @@ async def wait_for_custom_job(
         poll_interval: int = 10,
     ) -> types.CustomJob:
         """Make async calls to Vertex AI to check the custom job state until it is complete."""
+        CUSTOB_JOB_FAILED_STATES = {
+            JobState.JOB_STATE_FAILED: lambda: job.error.message,
+            JobState.JOB_STATE_CANCELLED: lambda: "The CustomJob has been cancelled.",
+            JobState.JOB_STATE_PAUSED: lambda: "The CustomJob has been stopped, and can be resumed.",
+            JobState.JOB_STATE_EXPIRED: lambda: "The CustomJob has expired.",
+            JobState.JOB_STATE_PARTIALLY_SUCCEEDED: lambda: job.error.message,
+        }
+
         client = await self.get_job_service_client(region=location)
         while True:
             try:
@@ -3255,8 +3310,10 @@ async def wait_for_custom_job(
                 self.log.exception("Exception occurred while requesting job %s", job_id)
                 raise AirflowException(ex)
             self.log.info("Status of the custom job %s is %s", job.name, job.state.name)
-            if job.state in self.JOB_COMPLETE_STATES:
+            if job.state == JobState.JOB_STATE_SUCCEEDED:
                 return job
+            if job.state in CUSTOB_JOB_FAILED_STATES:
+                raise RuntimeError(CUSTOB_JOB_FAILED_STATES[job.state]())
             self.log.info("Sleeping for %s seconds.", poll_interval)
             await asyncio.sleep(poll_interval)
 

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

@@ -58,6 +58,9 @@
     VERTEX_AI_BASE_LINK + "/locations/{location}/ray-clusters/{cluster_id}?project={project_id}"
 )
 VERTEX_AI_RAY_CLUSTER_LIST_LINK = VERTEX_AI_BASE_LINK + "/ray?project={project_id}"
+VERTEX_AI_CUSTOM_JOB_LINK = (
+    VERTEX_AI_BASE_LINK + "/locations/{region}/training/{custom_job_id}/cpu?project={project_id}"
+)
 
 
 class VertexAIModelLink(BaseGoogleLink):
@@ -202,3 +205,11 @@ class VertexAIRayClusterListLink(BaseGoogleLink):
     name = "Ray Cluster List"
     key = "ray_cluster_list_conf"
     format_str = VERTEX_AI_RAY_CLUSTER_LIST_LINK
+
+
+class VertexAICustomJobLink(BaseGoogleLink):
+    """Helper class for constructing Vertex AI CustomJob link."""
+
+    name = "Vertex AI Custom Job"
+    key = "custom_job_conf"
+    format_str = VERTEX_AI_CUSTOM_JOB_LINK

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

@@ -26,19 +26,22 @@
 from google.api_core.exceptions import NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform.models import Model
+from google.cloud.aiplatform_v1.types.custom_job import CustomJob
 from google.cloud.aiplatform_v1.types.dataset import Dataset
 from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
 
 from airflow.providers.common.compat.sdk import AirflowException, conf
 from airflow.providers.google.cloud.hooks.vertex_ai.custom_job import CustomJobHook
 from airflow.providers.google.cloud.links.vertex_ai import (
+    VertexAICustomJobLink,
     VertexAIModelLink,
     VertexAITrainingLink,
     VertexAITrainingPipelinesLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.vertex_ai import (
     CustomContainerTrainingJobTrigger,
+    CustomJobTrigger,
     CustomPythonPackageTrainingJobTrigger,
     CustomTrainingJobTrigger,
 )
@@ -1793,3 +1796,123 @@ def execute(self, context: Context):
         )
         VertexAITrainingPipelinesLink.persist(context=context)
         return [TrainingPipeline.to_dict(result) for result in results]
+
+
+class CreateCustomJobOperator(GoogleCloudBaseOperator):
+    """
+    Create a CustomJob. A created CustomJob right away will be attempted to be run.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param region: Required. The ID of the Google Cloud region that the service belongs to.
+    :param custom_job:  Required. The CustomJob to create.
+    :param retry: Designation of what errors, if any, should be retried.
+    :param timeout: The timeout for this request.
+    :param metadata: Strings which should be sent along with the request as metadata.
+    """
+
+    template_fields = ("region", "project_id", "custom_job", "impersonation_chain")
+    operator_extra_links = (VertexAICustomJobLink(),)
+
+    def __init__(
+        self,
+        *,
+        region: str,
+        project_id: str,
+        custom_job: CustomJob | dict,
+        retry: Retry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        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 = 10,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.region = region
+        self.project_id = project_id
+        self.custom_job = custom_job
+        self.retry = retry
+        self.timeout = timeout
+        self.metadata = metadata
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+        self.deferrable = deferrable
+        self.poll_interval = poll_interval
+        self.custom_job_id: str | None = None
+
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
+    @cached_property
+    def hook(self) -> CustomJobHook:
+        return CustomJobHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+    def execute(self, context: Context):
+        self.log.info("Creating CustomJob")
+        custom_job_obj = self.hook.create_custom_job(
+            project_id=self.project_id,
+            region=self.region,
+            custom_job=self.custom_job,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+        self.custom_job_id = self.hook.extract_custom_job_id(custom_job_name=custom_job_obj.name)
+        self.log.info("Custom job was created. Job id: %s", self.custom_job_id)
+        context["ti"].xcom_push(key="custom_job_id", value=self.custom_job_id)
+        VertexAICustomJobLink.persist(context=context, custom_job_id=self.custom_job_id)
+
+        if self.deferrable:
+            self.defer(
+                trigger=CustomJobTrigger(
+                    gcp_conn_id=self.gcp_conn_id,
+                    project_id=self.project_id,
+                    location=self.region,
+                    custom_job_id=self.custom_job_id,
+                    poll_interval=self.poll_interval,
+                    impersonation_chain=self.impersonation_chain,
+                ),
+                method_name="execute_complete",
+            )
+
+        custom_job_obj = self.hook.wait_for_custom_job(
+            project_id=self.project_id,
+            region=self.region,
+            custom_job_id=self.custom_job_id,
+            poll_interval=self.poll_interval,
+            retry=self.retry,
+            timeout=self.timeout,
+            metadata=self.metadata,
+        )
+
+        self.log.info("Custom Job %s completed.", self.custom_job_id)
+        custom_job = CustomJob.to_dict(custom_job_obj)
+        return custom_job
+
+    def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any]:
+        if event["status"] == "error":
+            raise RuntimeError(event["message"])
+        self.log.info(event["message"])
+        return event["custom_job"]
+
+    def on_kill(self) -> None:
+        """Act as a callback called when the operator is killed; cancel any running job."""
+        if self.custom_job_id:
+            self.log.info("Cancelling CustomJob with JobID: %s", self.custom_job_id)
+            self.hook.cancel_custom_job(
+                project_id=self.project_id,
+                region=self.region,
+                custom_job=self.custom_job_id,
+                retry=self.retry,
+                timeout=self.timeout,
+                metadata=self.metadata,
+            )
+            self.log.info("Custom Job %s canceled.", self.custom_job_id)

providers/google/src/airflow/providers/google/cloud/triggers/vertex_ai.py

@@ -401,3 +401,70 @@ async def _wait_job(self) -> types.TrainingPipeline:
             poll_interval=self.poll_interval,
         )
         return pipeline
+
+
+class CustomJobTrigger(BaseTrigger):
+    """Trigger that waits until a Vertex AI Custom job completes."""
+
+    def __init__(
+        self,
+        project_id: str,
+        location: str,
+        custom_job_id: str,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        poll_interval: int = 10,
+    ):
+        super().__init__()
+        self.project_id = project_id
+        self.location = location
+        self.custom_job_id = custom_job_id
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+        self.poll_interval = poll_interval
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        return (
+            "airflow.providers.google.cloud.triggers.vertex_ai.CustomJobTrigger",
+            {
+                "project_id": self.project_id,
+                "location": self.location,
+                "custom_job_id": self.custom_job_id,
+                "gcp_conn_id": self.gcp_conn_id,
+                "impersonation_chain": self.impersonation_chain,
+                "poll_interval": self.poll_interval,
+            },
+        )
+
+    @cached_property
+    def async_hook(self) -> CustomJobAsyncHook:
+        return CustomJobAsyncHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:
+        try:
+            custom_job = await self.async_hook.wait_for_custom_job(
+                project_id=self.project_id,
+                location=self.location,
+                job_id=self.custom_job_id,
+                poll_interval=self.poll_interval,
+            )
+        except (AirflowException, RuntimeError) as ex:
+            yield TriggerEvent(
+                {
+                    "status": "error",
+                    "message": str(ex),
+                }
+            )
+            return
+
+        message = f"Custom Job {custom_job.name} completed with status {custom_job.state.name}"
+        yield TriggerEvent(
+            {
+                "status": "success",
+                "message": message,
+                "custom_job": types.custom_job.CustomJob.to_dict(custom_job),
+            }
+        )

providers/google/src/airflow/providers/google/get_provider_info.py

@@ -1635,6 +1635,7 @@ def get_provider_info():
             "airflow.providers.google.cloud.links.vertex_ai.VertexAIPipelineJobListLink",
             "airflow.providers.google.cloud.links.vertex_ai.VertexAIRayClusterLink",
             "airflow.providers.google.cloud.links.vertex_ai.VertexAIRayClusterListLink",
+            "airflow.providers.google.cloud.links.vertex_ai.VertexAICustomJobLink",
             "airflow.providers.google.cloud.links.workflows.WorkflowsWorkflowDetailsLink",
             "airflow.providers.google.cloud.links.workflows.WorkflowsListOfWorkflowsLink",
             "airflow.providers.google.cloud.links.workflows.WorkflowsExecutionLink",