diff --git a/generated/known_airflow_exceptions.txt b/generated/known_airflow_exceptions.txt index 202df42602ecc..a4073767d523d 100644 --- a/generated/known_airflow_exceptions.txt +++ b/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 diff --git a/providers/google/provider.yaml b/providers/google/provider.yaml index 80624bb5ffb67..f31d6b0bf7b26 100644 --- a/providers/google/provider.yaml +++ b/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 diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py index ea446d3f5e2e0..ad949662e89eb 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +++ b/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) diff --git a/providers/google/src/airflow/providers/google/cloud/links/vertex_ai.py b/providers/google/src/airflow/providers/google/cloud/links/vertex_ai.py index d749dda1e1848..32d122f9a7248 100644 --- a/providers/google/src/airflow/providers/google/cloud/links/vertex_ai.py +++ b/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 diff --git a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py index ec88395d9751e..953fe8441eaab 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py @@ -26,12 +26,14 @@ 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, @@ -39,6 +41,7 @@ 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) diff --git a/providers/google/src/airflow/providers/google/cloud/triggers/vertex_ai.py b/providers/google/src/airflow/providers/google/cloud/triggers/vertex_ai.py index 6e7ecff02fc60..1ba0819ac4d1c 100644 --- a/providers/google/src/airflow/providers/google/cloud/triggers/vertex_ai.py +++ b/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), + } + ) diff --git a/providers/google/src/airflow/providers/google/get_provider_info.py b/providers/google/src/airflow/providers/google/get_provider_info.py index e2bfa861e248c..7dcaaa450b95c 100644 --- a/providers/google/src/airflow/providers/google/get_provider_info.py +++ b/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", diff --git a/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_create_custom_job.py b/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_create_custom_job.py new file mode 100644 index 0000000000000..85fc8d975424f --- /dev/null +++ b/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_create_custom_job.py @@ -0,0 +1,134 @@ +# +# 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. + + +"""Example Airflow DAG for Google Vertex AI service testing create Custom Jobs operator.""" + +from __future__ import annotations + +import os +from datetime import datetime + +from airflow.models.dag import DAG +from airflow.providers.google.cloud.operators.vertex_ai.custom_job import ( + CreateCustomJobOperator, +) + +ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default") +PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default") +REGION = "us-central1" +DAG_ID = "vertex_ai_create_custom_job" + +REPLICA_COUNT = 1 +MACHINE_TYPE = "n1-standard-4" +ACCELERATOR_TYPE = "ACCELERATOR_TYPE_UNSPECIFIED" +ACCELERATOR_COUNT = 0 +IMAGE_URI = "us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-16.py310:latest" + +test_python_code = ( + "import sys; " + "print('=== VERTEX AI RAW CUSTOM_JOB RUNNING SUCCESSFULLY ==='); " + "import math; " + "print(f'Sanity check calculation (pi): {math.pi}'); " + "print('=== TEST COMPLETED CLEANLY ==='); " + "sys.exit(0);" +) + + +with DAG( + DAG_ID, + schedule="@once", + start_date=datetime(2021, 1, 1), + catchup=False, + tags=["example", "vertex_ai", "custom_job"], +) as dag: + create_custom_job = CreateCustomJobOperator( + task_id="create_custom_job", + region=REGION, + project_id=PROJECT_ID, + custom_job={ + "display_name": f"{DAG_ID}_{ENV_ID}", + "labels": { + "vertex_pipelines": "", + "airflow_dag_id": DAG_ID, + "airflow_task_id": DAG_ID, + }, + "job_spec": { + "scheduling": {"disable_retries": True}, + "worker_pool_specs": [ + { + "machine_spec": { + "machine_type": MACHINE_TYPE, + "accelerator_type": ACCELERATOR_TYPE, + "accelerator_count": ACCELERATOR_COUNT, + }, + "replica_count": REPLICA_COUNT, + "container_spec": { + "image_uri": IMAGE_URI, + "command": ["python3", "-c", test_python_code], + }, + } + ], + }, + }, + ) + + create_custom_job_def = CreateCustomJobOperator( + task_id="create_custom_job_def", + region=REGION, + project_id=PROJECT_ID, + custom_job={ + "display_name": f"{DAG_ID}_{ENV_ID}_def", + "labels": { + "vertex_pipelines": "", + "airflow_dag_id": DAG_ID, + "airflow_task_id": DAG_ID, + }, + "job_spec": { + "scheduling": {"disable_retries": True}, + "worker_pool_specs": [ + { + "machine_spec": { + "machine_type": MACHINE_TYPE, + "accelerator_type": ACCELERATOR_TYPE, + "accelerator_count": ACCELERATOR_COUNT, + }, + "replica_count": REPLICA_COUNT, + "container_spec": { + "image_uri": IMAGE_URI, + "command": ["python3", "-c", test_python_code], + }, + } + ], + }, + }, + deferrable=True, + ) + + # ### Everything below this line is not part of example ### + # ### Just for system tests purpose ### + from tests_common.test_utils.watcher import watcher + + # This test needs watcher in order to properly mark success/failure + # when "tearDown" task with trigger rule is part of the DAG + list(dag.tasks) >> watcher() + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst) +test_run = get_test_run(dag) diff --git a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py index 3859946731727..fb949a3ec73eb 100644 --- a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py +++ b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_custom_job.py @@ -206,6 +206,79 @@ def test_list_training_pipelines(self, mock_client) -> None: ) mock_client.return_value.common_location_path.assert_called_once_with(TEST_PROJECT_ID, TEST_REGION) + @pytest.mark.parametrize( + "job_state_value", + [ + JobState.JOB_STATE_SUCCEEDED, + ], + ) + @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_custom_job")) + def test_wait_for_custom_job( + self, + mock_get_custom_job, + job_state_value, + test_custom_job_name, + ): + expected_obj = types.CustomJob( + state=job_state_value, + name=test_custom_job_name, + ) + mock_get_custom_job.return_value = expected_obj + actual_obj = self.hook.wait_for_custom_job( + project_id=TEST_PROJECT_ID, + region=TEST_REGION, + custom_job_id=TEST_PIPELINE_JOB_ID, + ) + assert actual_obj == expected_obj + + @pytest.mark.parametrize( + ("job_state_value", "error_message"), + [ + ( + JobState.JOB_STATE_FAILED, + "Error message from VertexAI", + ), + ( + JobState.JOB_STATE_CANCELLED, + "The CustomJob has been cancelled.", + ), + ( + JobState.JOB_STATE_PAUSED, + "The CustomJob has been stopped, and can be resumed.", + ), + ( + JobState.JOB_STATE_EXPIRED, + "The CustomJob has expired.", + ), + ( + JobState.JOB_STATE_PARTIALLY_SUCCEEDED, + "Error message from VertexAI", + ), + ], + ) + @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_custom_job")) + def test_wait_for_custom_job_failed_states( + self, + mock_get_custom_job, + job_state_value, + error_message, + test_custom_job_name, + ): + expected_obj = types.CustomJob( + state=job_state_value, + name=test_custom_job_name, + error={ + "message": "Error message from VertexAI", + }, + ) + mock_get_custom_job.return_value = expected_obj + with pytest.raises(RuntimeError, match=error_message): + self.hook.wait_for_custom_job( + project_id=TEST_PROJECT_ID, + region=TEST_REGION, + custom_job_id=TEST_PIPELINE_JOB_ID, + ) + class TestCustomJobWithoutDefaultProjectIdHook: def setup_method(self): @@ -314,6 +387,79 @@ def test_list_training_pipelines(self, mock_client) -> None: ) mock_client.return_value.common_location_path.assert_called_once_with(TEST_PROJECT_ID, TEST_REGION) + @pytest.mark.parametrize( + "job_state_value", + [ + JobState.JOB_STATE_SUCCEEDED, + ], + ) + @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_custom_job")) + def test_wait_for_custom_job( + self, + mock_get_custom_job, + job_state_value, + test_custom_job_name, + ): + expected_obj = types.CustomJob( + state=job_state_value, + name=test_custom_job_name, + ) + mock_get_custom_job.return_value = expected_obj + actual_obj = self.hook.wait_for_custom_job( + project_id=TEST_PROJECT_ID, + region=TEST_REGION, + custom_job_id=TEST_PIPELINE_JOB_ID, + ) + assert actual_obj == expected_obj + + @pytest.mark.parametrize( + ("job_state_value", "error_message"), + [ + ( + JobState.JOB_STATE_FAILED, + "Error message from VertexAI", + ), + ( + JobState.JOB_STATE_CANCELLED, + "The CustomJob has been cancelled.", + ), + ( + JobState.JOB_STATE_PAUSED, + "The CustomJob has been stopped, and can be resumed.", + ), + ( + JobState.JOB_STATE_EXPIRED, + "The CustomJob has expired.", + ), + ( + JobState.JOB_STATE_PARTIALLY_SUCCEEDED, + "Error message from VertexAI", + ), + ], + ) + @mock.patch(CUSTOM_JOB_STRING.format("CustomJobHook.get_custom_job")) + def test_wait_for_custom_job_failed_states( + self, + mock_get_custom_job, + job_state_value, + error_message, + test_custom_job_name, + ): + expected_obj = types.CustomJob( + state=job_state_value, + name=test_custom_job_name, + error={ + "message": "Error message from VertexAI", + }, + ) + mock_get_custom_job.return_value = expected_obj + with pytest.raises(RuntimeError, match=error_message): + self.hook.wait_for_custom_job( + project_id=TEST_PROJECT_ID, + region=TEST_REGION, + custom_job_id=TEST_PIPELINE_JOB_ID, + ) + class TestCustomJobAsyncHook: @pytest.mark.asyncio @@ -398,9 +544,6 @@ async def test_wait_for_training_pipeline_returns_pipeline_if_in_complete_state( @pytest.mark.parametrize( "job_state_value", [ - JobState.JOB_STATE_CANCELLED, - JobState.JOB_STATE_FAILED, - JobState.JOB_STATE_PAUSED, JobState.JOB_STATE_SUCCEEDED, ], ) @@ -427,6 +570,59 @@ async def test_wait_for_custom_job_returns_job_if_in_complete_state( mock_get_job_service_client.assert_awaited_once_with(region=TEST_REGION) assert actual_obj == expected_obj + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("job_state_value", "error_message"), + [ + ( + JobState.JOB_STATE_FAILED, + "Error message from VertexAI", + ), + ( + JobState.JOB_STATE_CANCELLED, + "The CustomJob has been cancelled.", + ), + ( + JobState.JOB_STATE_PAUSED, + "The CustomJob has been stopped, and can be resumed.", + ), + ( + JobState.JOB_STATE_EXPIRED, + "The CustomJob has expired.", + ), + ( + JobState.JOB_STATE_PARTIALLY_SUCCEEDED, + "Error message from VertexAI", + ), + ], + ) + @mock.patch(CUSTOM_JOB_STRING.format("CustomJobAsyncHook.get_custom_job")) + @mock.patch(CUSTOM_JOB_STRING.format("CustomJobAsyncHook.get_job_service_client")) + async def test_wait_for_custom_job_failed_states( + self, + mock_get_job_service_client, + mock_get_custom_job, + job_state_value, + error_message, + test_async_hook, + test_custom_job_name, + ): + expected_obj = types.CustomJob( + state=job_state_value, + name=test_custom_job_name, + error={ + "message": "Error message from VertexAI", + }, + ) + mock_get_custom_job.return_value = expected_obj + with pytest.raises(RuntimeError, match=error_message): + await test_async_hook.wait_for_custom_job( + project_id=TEST_PROJECT_ID, + location=TEST_REGION, + job_id=TEST_PIPELINE_JOB_ID, + ) + mock_get_job_service_client.assert_awaited_once_with(region=TEST_REGION) + @pytest.mark.asyncio @pytest.mark.parametrize( "pipeline_state_value", diff --git a/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py b/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py index 762648f619229..7fee751c47dac 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py +++ b/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py @@ -45,6 +45,7 @@ ) from airflow.providers.google.cloud.operators.vertex_ai.custom_job import ( CreateCustomContainerTrainingJobOperator, + CreateCustomJobOperator, CreateCustomPythonPackageTrainingJobOperator, CreateCustomTrainingJobOperator, DeleteCustomTrainingJobOperator, @@ -233,6 +234,27 @@ TEST_CLUSTER_NAME: str = "test-cluster-name" TEST_CLUSTER_ID: str = "test-cluster-id" +TEST_CUSTOM_JOB = { + "display_name": DISPLAY_NAME, + "job_spec": { + "scheduling": {"disable_retries": True}, + "worker_pool_specs": [ + { + "machine_spec": { + "machine_type": MACHINE_TYPE, + "accelerator_type": ACCELERATOR_TYPE, + "accelerator_count": ACCELERATOR_COUNT, + }, + "replica_count": REPLICA_COUNT, + "container_spec": { + "image_uri": "test_image_uri", + "command": ["python3", "-c", "test_python_code"], + }, + } + ], + }, +} + class TestVertexAICreateCustomContainerTrainingJobOperator: @mock.patch(VERTEX_AI_PATH.format("custom_job.Dataset")) @@ -3172,3 +3194,30 @@ def test_execute(self, mock_hook): project_id=GCP_PROJECT, cluster_id=TEST_CLUSTER_ID, ) + + +class TestVertexAICreateCustomJobOperator: + @mock.patch(VERTEX_AI_PATH.format("custom_job.CustomJob.to_dict")) + @mock.patch(VERTEX_AI_PATH.format("custom_job.CustomJobHook")) + def test_execute(self, mock_hook, to_dict_mock): + op = CreateCustomJobOperator( + task_id=TASK_ID, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + region=GCP_LOCATION, + project_id=GCP_PROJECT, + custom_job=TEST_CUSTOM_JOB, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) + op.execute(context={"ti": mock.MagicMock(), "task": mock.MagicMock()}) + mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) + mock_hook.return_value.create_custom_job.assert_called_once_with( + region=GCP_LOCATION, + project_id=GCP_PROJECT, + custom_job=TEST_CUSTOM_JOB, + retry=RETRY, + timeout=TIMEOUT, + metadata=METADATA, + ) diff --git a/providers/google/tests/unit/google/cloud/triggers/test_vertex_ai.py b/providers/google/tests/unit/google/cloud/triggers/test_vertex_ai.py index 1fb59d6ea6b4c..b34c9b875e47f 100644 --- a/providers/google/tests/unit/google/cloud/triggers/test_vertex_ai.py +++ b/providers/google/tests/unit/google/cloud/triggers/test_vertex_ai.py @@ -42,6 +42,7 @@ CreateBatchPredictionJobTrigger, CreateHyperparameterTuningJobTrigger, CustomContainerTrainingJobTrigger, + CustomJobTrigger, CustomPythonPackageTrainingJobTrigger, CustomTrainingJobTrigger, RunPipelineJobTrigger, @@ -113,6 +114,18 @@ def custom_training_job_trigger(): ) +@pytest.fixture +def custom_job_trigger(): + return CustomJobTrigger( + gcp_conn_id=TEST_CONN_ID, + project_id=TEST_PROJECT_ID, + location=TEST_LOCATION, + custom_job_id=TEST_HPT_JOB_ID, + poll_interval=TEST_POLL_INTERVAL, + impersonation_chain=TEST_IMPERSONATION_CHAIN, + ) + + @pytest_asyncio.fixture async def custom_job_async_hook(): return CustomJobAsyncHook( @@ -998,3 +1011,20 @@ async def test_wait_training_pipeline( pipeline_id=custom_python_package_training_job_trigger.job_id, poll_interval=custom_python_package_training_job_trigger.poll_interval, ) + + +class TestCustomJobTrigger: + def test_serialize(self, custom_job_trigger): + actual_data = custom_job_trigger.serialize() + expected_data = ( + "airflow.providers.google.cloud.triggers.vertex_ai.CustomJobTrigger", + { + "project_id": TEST_PROJECT_ID, + "location": TEST_LOCATION, + "custom_job_id": TEST_HPT_JOB_ID, + "gcp_conn_id": TEST_CONN_ID, + "impersonation_chain": TEST_IMPERSONATION_CHAIN, + "poll_interval": TEST_POLL_INTERVAL, + }, + ) + assert actual_data == expected_data