Implement deferrable mode for CreateHyperparameterTuningJobOperator (#…

…36594)

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

@@ -15,22 +15,32 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module contains a Google Cloud Vertex AI hook."""
+"""
+This module contains a Google Cloud Vertex AI hook.
+
+.. spelling:word-list::
+
+    JobServiceAsyncClient
+"""
 from __future__ import annotations
 
+import asyncio
+from functools import lru_cache
 from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import CustomJob, HyperparameterTuningJob, gapic, hyperparameter_tuning
-from google.cloud.aiplatform_v1 import JobServiceClient, types
+from google.cloud.aiplatform_v1 import JobServiceAsyncClient, JobServiceClient, JobState, types
 
 from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry
+    from google.api_core.retry_async import AsyncRetry
     from google.cloud.aiplatform_v1.services.job_service.pagers import ListHyperparameterTuningJobsPager
 
 
@@ -172,6 +182,7 @@ def create_hyperparameter_tuning_job(
         tensorboard: str | None = None,
         sync: bool = True,
         # END: run param
+        wait_job_completed: bool = True,
     ) -> HyperparameterTuningJob:
         """
         Create a HyperparameterTuningJob.
@@ -256,6 +267,7 @@ def create_hyperparameter_tuning_job(
             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 wait_job_completed: Whether to wait for the job completed.
         """
         custom_job = self.get_custom_job_object(
             project=project_id,
@@ -292,7 +304,11 @@ def create_hyperparameter_tuning_job(
             tensorboard=tensorboard,
             sync=sync,
         )
-        self._hyperparameter_tuning_job.wait()
+
+        if wait_job_completed:
+            self._hyperparameter_tuning_job.wait()
+        else:
+            self._hyperparameter_tuning_job._wait_for_resource_creation()
         return self._hyperparameter_tuning_job
 
     @GoogleBaseHook.fallback_to_default_project_id
@@ -413,3 +429,104 @@ def delete_hyperparameter_tuning_job(
             metadata=metadata,
         )
         return result
+
+
+class HyperparameterTuningJobAsyncHook(GoogleBaseHook):
+    """Async hook for Google Cloud Vertex AI Hyperparameter Tuning Job APIs."""
+
+    def __init__(
+        self,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ):
+        super().__init__(
+            gcp_conn_id=gcp_conn_id,
+            impersonation_chain=impersonation_chain,
+            **kwargs,
+        )
+
+    @lru_cache
+    def get_job_service_client(self, region: str | None = None) -> JobServiceAsyncClient:
+        """
+        Retrieves Vertex AI async client.
+
+        :return: Google Cloud Vertex AI client object.
+        """
+        endpoint = f"{region}-aiplatform.googleapis.com:443" if region and region != "global" else None
+        return JobServiceAsyncClient(
+            credentials=self.get_credentials(),
+            client_info=CLIENT_INFO,
+            client_options=ClientOptions(api_endpoint=endpoint),
+        )
+
+    async def get_hyperparameter_tuning_job(
+        self,
+        project_id: str,
+        location: str,
+        job_id: str,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> types.HyperparameterTuningJob:
+        """
+        Retrieves a hyperparameter tuning job.
+
+        :param project_id: Required. The ID of the Google Cloud project that the job belongs to.
+        :param location: Required. The ID of the Google Cloud region that the job belongs to.
+        :param job_id: Required. The hyperparameter tuning job id.
+        :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.
+        """
+        client: JobServiceAsyncClient = self.get_job_service_client(region=location)
+        job_name = client.hyperparameter_tuning_job_path(project_id, location, job_id)
+
+        result = await client.get_hyperparameter_tuning_job(
+            request={
+                "name": job_name,
+            },
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
+        return result
+
+    async def wait_hyperparameter_tuning_job(
+        self,
+        project_id: str,
+        location: str,
+        job_id: str,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+        poll_interval: int = 10,
+    ) -> types.HyperparameterTuningJob:
+        statuses_complete = {
+            JobState.JOB_STATE_CANCELLED,
+            JobState.JOB_STATE_FAILED,
+            JobState.JOB_STATE_PAUSED,
+            JobState.JOB_STATE_SUCCEEDED,
+        }
+        while True:
+            try:
+                self.log.info("Requesting hyperparameter tuning job with id %s", job_id)
+                job: types.HyperparameterTuningJob = await self.get_hyperparameter_tuning_job(
+                    project_id=project_id,
+                    location=location,
+                    job_id=job_id,
+                    retry=retry,
+                    timeout=timeout,
+                    metadata=metadata,
+                )
+            except Exception as ex:
+                self.log.exception("Exception occurred while requesting job %s", job_id)
+                raise AirflowException(ex)
+
+            self.log.info("Status of the hyperparameter tuning job %s is %s", job.name, job.state.name)
+            if job.state in statuses_complete:
+                return job
+
+            self.log.info("Sleeping for %s seconds.", poll_interval)
+            await asyncio.sleep(poll_interval)

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

@@ -20,12 +20,14 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+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 airflow.configuration import conf
+from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job import (
     HyperparameterTuningJobHook,
 )
@@ -34,6 +36,7 @@
     VertexAITrainingLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.cloud.triggers.vertex_ai import CreateHyperparameterTuningJobTrigger
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
@@ -124,7 +127,7 @@ 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
+    :param sync: 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
@@ -135,6 +138,9 @@ 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 poll_interval: Interval size which defines how often job status is checked in deferrable mode.
     """
 
     template_fields = [
@@ -177,6 +183,8 @@ def __init__(
         # END: run param
         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)
@@ -209,8 +217,17 @@ def __init__(
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
         self.hook: HyperparameterTuningJobHook | None = None
+        self.deferrable = deferrable
+        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`."
+            )
+
         self.log.info("Creating Hyperparameter Tuning job")
         self.hook = HyperparameterTuningJobHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -243,12 +260,26 @@ def execute(self, context: Context):
             enable_web_access=self.enable_web_access,
             tensorboard=self.tensorboard,
             sync=self.sync,
+            wait_job_completed=not self.deferrable,
         )
 
         hyperparameter_tuning_job = result.to_dict()
         hyperparameter_tuning_job_id = self.hook.extract_hyperparameter_tuning_job_id(
             hyperparameter_tuning_job
         )
+        if self.deferrable:
+            self.defer(
+                trigger=CreateHyperparameterTuningJobTrigger(
+                    conn_id=self.gcp_conn_id,
+                    project_id=self.project_id,
+                    location=self.region,
+                    job_id=hyperparameter_tuning_job_id,
+                    poll_interval=self.poll_interval,
+                    impersonation_chain=self.impersonation_chain,
+                ),
+                method_name="execute_complete",
+            )
+
         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)
@@ -262,6 +293,32 @@ def on_kill(self) -> None:
         if self.hook:
             self.hook.cancel_hyperparameter_tuning_job()
 
+    def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any]:
+        if event and event["status"] == "error":
+            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"]
+
 
 class GetHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
     """

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

@@ -0,0 +1,99 @@
+# 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.
+from __future__ import annotations
+
+from typing import Any, AsyncIterator, Sequence
+
+from google.cloud.aiplatform_v1 import HyperparameterTuningJob, JobState
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job import (
+    HyperparameterTuningJobAsyncHook,
+)
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class CreateHyperparameterTuningJobTrigger(BaseTrigger):
+    """CreateHyperparameterTuningJobTrigger run on the trigger worker to perform create operation."""
+
+    statuses_success = {
+        JobState.JOB_STATE_PAUSED,
+        JobState.JOB_STATE_SUCCEEDED,
+    }
+
+    def __init__(
+        self,
+        conn_id: str,
+        project_id: str,
+        location: str,
+        job_id: str,
+        poll_interval: int,
+        impersonation_chain: str | Sequence[str] | None = None,
+    ):
+        super().__init__()
+        self.conn_id = conn_id
+        self.project_id = project_id
+        self.location = location
+        self.job_id = job_id
+        self.poll_interval = poll_interval
+        self.impersonation_chain = impersonation_chain
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        return (
+            "airflow.providers.google.cloud.triggers.vertex_ai.CreateHyperparameterTuningJobTrigger",
+            {
+                "conn_id": self.conn_id,
+                "project_id": self.project_id,
+                "location": self.location,
+                "job_id": self.job_id,
+                "poll_interval": self.poll_interval,
+                "impersonation_chain": self.impersonation_chain,
+            },
+        )
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:
+        hook = self._get_async_hook()
+        try:
+            job = await hook.wait_hyperparameter_tuning_job(
+                project_id=self.project_id,
+                location=self.location,
+                job_id=self.job_id,
+                poll_interval=self.poll_interval,
+            )
+        except AirflowException as ex:
+            yield TriggerEvent(
+                {
+                    "status": "error",
+                    "message": str(ex),
+                }
+            )
+            return
+
+        status = "success" if job.state in self.statuses_success else "error"
+        message = f"Hyperparameter tuning job {job.name} completed with status {job.state.name}"
+        yield TriggerEvent(
+            {
+                "status": status,
+                "message": message,
+                "job": HyperparameterTuningJob.to_dict(job),
+            }
+        )
+
+    def _get_async_hook(self) -> HyperparameterTuningJobAsyncHook:
+        return HyperparameterTuningJobAsyncHook(
+            gcp_conn_id=self.conn_id, impersonation_chain=self.impersonation_chain
+        )