| # 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 unittest import mock |
| from unittest.mock import MagicMock, call |
| |
| import pytest |
| |
| # For no Pydantic environment, we need to skip the tests |
| pytest.importorskip("google.cloud.aiplatform_v1") |
| |
| from google.api_core.gapic_v1.method import DEFAULT |
| from google.api_core.retry import Retry |
| |
| from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, TaskDeferred |
| from airflow.providers.google.cloud.operators.vertex_ai.auto_ml import ( |
| CreateAutoMLForecastingTrainingJobOperator, |
| CreateAutoMLImageTrainingJobOperator, |
| CreateAutoMLTabularTrainingJobOperator, |
| CreateAutoMLTextTrainingJobOperator, |
| CreateAutoMLVideoTrainingJobOperator, |
| DeleteAutoMLTrainingJobOperator, |
| ListAutoMLTrainingJobOperator, |
| ) |
| from airflow.providers.google.cloud.operators.vertex_ai.batch_prediction_job import ( |
| CreateBatchPredictionJobOperator, |
| DeleteBatchPredictionJobOperator, |
| ListBatchPredictionJobsOperator, |
| ) |
| from airflow.providers.google.cloud.operators.vertex_ai.custom_job import ( |
| CreateCustomContainerTrainingJobOperator, |
| CreateCustomPythonPackageTrainingJobOperator, |
| CreateCustomTrainingJobOperator, |
| DeleteCustomTrainingJobOperator, |
| ListCustomTrainingJobOperator, |
| ) |
| from airflow.providers.google.cloud.operators.vertex_ai.dataset import ( |
| CreateDatasetOperator, |
| DeleteDatasetOperator, |
| ExportDataOperator, |
| ImportDataOperator, |
| ListDatasetsOperator, |
| UpdateDatasetOperator, |
| ) |
| from airflow.providers.google.cloud.operators.vertex_ai.endpoint_service import ( |
| CreateEndpointOperator, |
| DeleteEndpointOperator, |
| DeployModelOperator, |
| ListEndpointsOperator, |
| UndeployModelOperator, |
| ) |
| from airflow.providers.google.cloud.operators.vertex_ai.hyperparameter_tuning_job import ( |
| CreateHyperparameterTuningJobOperator, |
| DeleteHyperparameterTuningJobOperator, |
| GetHyperparameterTuningJobOperator, |
| ListHyperparameterTuningJobOperator, |
| ) |
| from airflow.providers.google.cloud.operators.vertex_ai.model_service import ( |
| AddVersionAliasesOnModelOperator, |
| DeleteModelOperator, |
| DeleteModelVersionOperator, |
| DeleteVersionAliasesOnModelOperator, |
| ExportModelOperator, |
| GetModelOperator, |
| ListModelsOperator, |
| ListModelVersionsOperator, |
| SetDefaultVersionOnModelOperator, |
| UploadModelOperator, |
| ) |
| from airflow.providers.google.cloud.operators.vertex_ai.pipeline_job import ( |
| DeletePipelineJobOperator, |
| GetPipelineJobOperator, |
| ListPipelineJobOperator, |
| RunPipelineJobOperator, |
| ) |
| from airflow.providers.google.cloud.triggers.vertex_ai import ( |
| CustomContainerTrainingJobTrigger, |
| CustomPythonPackageTrainingJobTrigger, |
| CustomTrainingJobTrigger, |
| RunPipelineJobTrigger, |
| ) |
| from airflow.utils import timezone |
| |
| VERTEX_AI_PATH = "airflow.providers.google.cloud.operators.vertex_ai.{}" |
| VERTEX_AI_LINKS_PATH = "airflow.providers.google.cloud.links.vertex_ai.{}" |
| TIMEOUT = 120 |
| RETRY = mock.MagicMock(Retry) |
| METADATA = [("key", "value")] |
| |
| TASK_ID = "test_task_id" |
| GCP_PROJECT = "test-project" |
| GCP_LOCATION = "test-location" |
| GCP_CONN_ID = "test-conn" |
| IMPERSONATION_CHAIN = ["ACCOUNT_1", "ACCOUNT_2", "ACCOUNT_3"] |
| STAGING_BUCKET = "gs://test-vertex-ai-bucket" |
| DISPLAY_NAME = "display_name_1" # Create random display name |
| DISPLAY_NAME_2 = "display_nmae_2" |
| ARGS = ["--tfds", "tf_flowers:3.*.*"] |
| CONTAINER_URI = "gcr.io/cloud-aiplatform/training/tf-cpu.2-2:latest" |
| REPLICA_COUNT = 1 |
| MACHINE_TYPE = "n1-standard-4" |
| ACCELERATOR_TYPE = "ACCELERATOR_TYPE_UNSPECIFIED" |
| ACCELERATOR_COUNT = 0 |
| TRAINING_FRACTION_SPLIT = 0.7 |
| TEST_FRACTION_SPLIT = 0.15 |
| VALIDATION_FRACTION_SPLIT = 0.15 |
| COMMAND_2 = ["echo", "Hello World"] |
| |
| TEST_API_ENDPOINT: str = "test-api-endpoint" |
| TEST_PIPELINE_JOB: str = "test-pipeline-job" |
| TEST_TRAINING_PIPELINE: str = "test-training-pipeline" |
| TEST_PIPELINE_JOB_ID: str = "test-pipeline-job-id" |
| |
| PYTHON_PACKAGE = "/files/trainer-0.1.tar.gz" |
| PYTHON_PACKAGE_CMDARGS = "test-python-cmd" |
| PYTHON_PACKAGE_GCS_URI = "gs://test-vertex-ai-bucket/trainer-0.1.tar.gz" |
| PYTHON_MODULE_NAME = "trainer.task" |
| |
| TRAINING_PIPELINE_ID = "test-training-pipeline-id" |
| CUSTOM_JOB_ID = "test-custom-job-id" |
| |
| TEST_DATASET = { |
| "display_name": "test-dataset-name", |
| "metadata_schema_uri": "gs://google-cloud-aiplatform/schema/dataset/metadata/image_1.0.0.yaml", |
| "metadata": "test-image-dataset", |
| } |
| TEST_DATASET_ID = "test-dataset-id" |
| TEST_PARENT_MODEL = "test-parent-model" |
| TEST_EXPORT_CONFIG = { |
| "annotationsFilter": "test-filter", |
| "gcs_destination": {"output_uri_prefix": "airflow-system-tests-data"}, |
| } |
| TEST_IMPORT_CONFIG = [ |
| { |
| "data_item_labels": { |
| "test-labels-name": "test-labels-value", |
| }, |
| "import_schema_uri": "test-shema-uri", |
| "gcs_source": {"uris": ["test-string"]}, |
| }, |
| {}, |
| ] |
| TEST_UPDATE_MASK = "test-update-mask" |
| |
| TEST_TRAINING_TARGET_COLUMN = "target" |
| TEST_TRAINING_TIME_COLUMN = "time" |
| TEST_TRAINING_TIME_SERIES_IDENTIFIER_COLUMN = "time_series_identifier" |
| TEST_TRAINING_UNAVAILABLE_AT_FORECAST_COLUMNS: list[str] = [] |
| TEST_TRAINING_AVAILABLE_AT_FORECAST_COLUMNS: list[str] = [] |
| TEST_TRAINING_FORECAST_HORIZON = 10 |
| TEST_TRAINING_DATA_GRANULARITY_UNIT = "day" |
| TEST_TRAINING_DATA_GRANULARITY_COUNT = 1 |
| |
| TEST_MODEL_ID = "test_model_id" |
| TEST_MODEL_NAME = f"projects/{GCP_PROJECT}/locations/{GCP_LOCATION}/models/test_model_id" |
| TEST_MODEL_OBJ: dict = {} |
| TEST_JOB_DISPLAY_NAME = "temp_create_batch_prediction_job_test" |
| TEST_BATCH_PREDICTION_JOB_ID = "test_batch_prediction_job_id" |
| |
| TEST_ENDPOINT = { |
| "display_name": "endpoint_test", |
| } |
| TEST_ENDPOINT_ID = "1234567890" |
| TEST_DEPLOYED_MODEL = { |
| # format: 'projects/{project}/locations/{location}/models/{model}' |
| "model": f"projects/{GCP_PROJECT}/locations/{GCP_LOCATION}/models/test_model_id", |
| "display_name": "temp_endpoint_test", |
| "dedicated_resources": { |
| "min_replica_count": 1, |
| "max_replica_count": 1, |
| }, |
| } |
| TEST_DEPLOYED_MODEL_ID = "test_deployed_model_id" |
| TEST_HYPERPARAMETER_TUNING_JOB_ID = "test_hyperparameter_tuning_job_id" |
| |
| TEST_OUTPUT_CONFIG = { |
| "artifact_destination": {"output_uri_prefix": "gcs_destination_output_uri_prefix"}, |
| "export_format_id": "tf-saved-model", |
| } |
| |
| TEST_CREATE_REQUEST_TIMEOUT = 100.5 |
| TEST_BATCH_SIZE = 4000 |
| TEST_VERSION_ALIASES = ["new-alias"] |
| |
| TEST_TEMPLATE_PATH = "test_template_path" |
| |
| SYNC_DEPRECATION_WARNING = "The 'sync' parameter is deprecated and will be removed after 01.10.2024." |
| |
| |
| class TestVertexAICreateCustomContainerTrainingJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.Dataset")) |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CustomJobHook")) |
| def test_execute(self, mock_hook, mock_dataset): |
| mock_hook.return_value.create_custom_container_training_job.return_value = ( |
| None, |
| "training_id", |
| "custom_job_id", |
| ) |
| op = CreateCustomContainerTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| args=ARGS, |
| container_uri=CONTAINER_URI, |
| model_serving_container_image_uri=CONTAINER_URI, |
| command=COMMAND_2, |
| model_display_name=DISPLAY_NAME_2, |
| replica_count=REPLICA_COUNT, |
| machine_type=MACHINE_TYPE, |
| accelerator_type=ACCELERATOR_TYPE, |
| accelerator_count=ACCELERATOR_COUNT, |
| training_fraction_split=TRAINING_FRACTION_SPLIT, |
| validation_fraction_split=VALIDATION_FRACTION_SPLIT, |
| test_fraction_split=TEST_FRACTION_SPLIT, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| dataset_id=TEST_DATASET_ID, |
| parent_model=TEST_PARENT_MODEL, |
| ) |
| with pytest.warns(AirflowProviderDeprecationWarning, match=SYNC_DEPRECATION_WARNING): |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_dataset.assert_called_once_with(name=TEST_DATASET_ID) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.create_custom_container_training_job.assert_called_once_with( |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| args=ARGS, |
| container_uri=CONTAINER_URI, |
| model_serving_container_image_uri=CONTAINER_URI, |
| command=COMMAND_2, |
| dataset=mock_dataset.return_value, |
| model_display_name=DISPLAY_NAME_2, |
| replica_count=REPLICA_COUNT, |
| machine_type=MACHINE_TYPE, |
| accelerator_type=ACCELERATOR_TYPE, |
| accelerator_count=ACCELERATOR_COUNT, |
| training_fraction_split=TRAINING_FRACTION_SPLIT, |
| validation_fraction_split=VALIDATION_FRACTION_SPLIT, |
| test_fraction_split=TEST_FRACTION_SPLIT, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| parent_model=TEST_PARENT_MODEL, |
| model_serving_container_predict_route=None, |
| model_serving_container_health_route=None, |
| model_serving_container_command=None, |
| model_serving_container_args=None, |
| model_serving_container_environment_variables=None, |
| model_serving_container_ports=None, |
| model_description=None, |
| model_instance_schema_uri=None, |
| model_parameters_schema_uri=None, |
| model_prediction_schema_uri=None, |
| labels=None, |
| training_encryption_spec_key_name=None, |
| model_encryption_spec_key_name=None, |
| # RUN |
| annotation_schema_uri=None, |
| model_labels=None, |
| base_output_dir=None, |
| service_account=None, |
| network=None, |
| bigquery_destination=None, |
| environment_variables=None, |
| boot_disk_type="pd-ssd", |
| boot_disk_size_gb=100, |
| training_filter_split=None, |
| validation_filter_split=None, |
| test_filter_split=None, |
| predefined_split_column_name=None, |
| timestamp_split_column_name=None, |
| tensorboard=None, |
| sync=True, |
| is_default_version=None, |
| model_version_aliases=None, |
| model_version_description=None, |
| ) |
| |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CreateCustomContainerTrainingJobOperator.hook")) |
| def test_execute_enters_deferred_state(self, mock_hook): |
| task = CreateCustomContainerTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| args=ARGS, |
| container_uri=CONTAINER_URI, |
| model_serving_container_image_uri=CONTAINER_URI, |
| command=COMMAND_2, |
| model_display_name=DISPLAY_NAME_2, |
| replica_count=REPLICA_COUNT, |
| machine_type=MACHINE_TYPE, |
| accelerator_type=ACCELERATOR_TYPE, |
| accelerator_count=ACCELERATOR_COUNT, |
| training_fraction_split=TRAINING_FRACTION_SPLIT, |
| validation_fraction_split=VALIDATION_FRACTION_SPLIT, |
| test_fraction_split=TEST_FRACTION_SPLIT, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| deferrable=True, |
| ) |
| mock_hook.return_value.exists.return_value = False |
| with pytest.raises(TaskDeferred) as exc: |
| with pytest.warns(AirflowProviderDeprecationWarning, match=SYNC_DEPRECATION_WARNING): |
| task.execute(context={"ti": mock.MagicMock()}) |
| assert isinstance( |
| exc.value.trigger, CustomContainerTrainingJobTrigger |
| ), "Trigger is not a CustomContainerTrainingJobTrigger" |
| |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CreateCustomContainerTrainingJobOperator.xcom_push")) |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CreateCustomContainerTrainingJobOperator.hook")) |
| def test_execute_complete_success(self, mock_hook, mock_xcom_push): |
| task = CreateCustomContainerTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| args=ARGS, |
| container_uri=CONTAINER_URI, |
| model_serving_container_image_uri=CONTAINER_URI, |
| command=COMMAND_2, |
| model_display_name=DISPLAY_NAME_2, |
| replica_count=REPLICA_COUNT, |
| machine_type=MACHINE_TYPE, |
| accelerator_type=ACCELERATOR_TYPE, |
| accelerator_count=ACCELERATOR_COUNT, |
| training_fraction_split=TRAINING_FRACTION_SPLIT, |
| validation_fraction_split=VALIDATION_FRACTION_SPLIT, |
| test_fraction_split=TEST_FRACTION_SPLIT, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| deferrable=True, |
| ) |
| expected_result = {} |
| mock_hook.return_value.exists.return_value = False |
| mock_xcom_push.return_value = None |
| actual_result = task.execute_complete( |
| context=None, event={"status": "success", "message": "", "job": {}} |
| ) |
| assert actual_result == expected_result |
| |
| def test_execute_complete_error_status_raises_exception(self): |
| task = CreateCustomContainerTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| args=ARGS, |
| container_uri=CONTAINER_URI, |
| model_serving_container_image_uri=CONTAINER_URI, |
| command=COMMAND_2, |
| model_display_name=DISPLAY_NAME_2, |
| replica_count=REPLICA_COUNT, |
| machine_type=MACHINE_TYPE, |
| accelerator_type=ACCELERATOR_TYPE, |
| accelerator_count=ACCELERATOR_COUNT, |
| training_fraction_split=TRAINING_FRACTION_SPLIT, |
| validation_fraction_split=VALIDATION_FRACTION_SPLIT, |
| test_fraction_split=TEST_FRACTION_SPLIT, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| deferrable=True, |
| ) |
| with pytest.raises(AirflowException): |
| task.execute_complete(context=None, event={"status": "error", "message": "test message"}) |
| |
| |
| class TestVertexAICreateCustomPythonPackageTrainingJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.Dataset")) |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CustomJobHook")) |
| def test_execute(self, mock_hook, mock_dataset): |
| mock_hook.return_value.create_custom_python_package_training_job.return_value = ( |
| None, |
| "training_id", |
| "custom_job_id", |
| ) |
| op = CreateCustomPythonPackageTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| python_package_gcs_uri=PYTHON_PACKAGE_GCS_URI, |
| python_module_name=PYTHON_MODULE_NAME, |
| container_uri=CONTAINER_URI, |
| args=ARGS, |
| model_serving_container_image_uri=CONTAINER_URI, |
| model_display_name=DISPLAY_NAME_2, |
| replica_count=REPLICA_COUNT, |
| machine_type=MACHINE_TYPE, |
| accelerator_type=ACCELERATOR_TYPE, |
| accelerator_count=ACCELERATOR_COUNT, |
| training_fraction_split=TRAINING_FRACTION_SPLIT, |
| validation_fraction_split=VALIDATION_FRACTION_SPLIT, |
| test_fraction_split=TEST_FRACTION_SPLIT, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| dataset_id=TEST_DATASET_ID, |
| parent_model=TEST_PARENT_MODEL, |
| ) |
| with pytest.warns(AirflowProviderDeprecationWarning, match=SYNC_DEPRECATION_WARNING): |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_dataset.assert_called_once_with(name=TEST_DATASET_ID) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.create_custom_python_package_training_job.assert_called_once_with( |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| args=ARGS, |
| container_uri=CONTAINER_URI, |
| model_serving_container_image_uri=CONTAINER_URI, |
| python_package_gcs_uri=PYTHON_PACKAGE_GCS_URI, |
| python_module_name=PYTHON_MODULE_NAME, |
| dataset=mock_dataset.return_value, |
| model_display_name=DISPLAY_NAME_2, |
| replica_count=REPLICA_COUNT, |
| machine_type=MACHINE_TYPE, |
| accelerator_type=ACCELERATOR_TYPE, |
| accelerator_count=ACCELERATOR_COUNT, |
| training_fraction_split=TRAINING_FRACTION_SPLIT, |
| validation_fraction_split=VALIDATION_FRACTION_SPLIT, |
| test_fraction_split=TEST_FRACTION_SPLIT, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| parent_model=TEST_PARENT_MODEL, |
| is_default_version=None, |
| model_version_aliases=None, |
| model_version_description=None, |
| model_serving_container_predict_route=None, |
| model_serving_container_health_route=None, |
| model_serving_container_command=None, |
| model_serving_container_args=None, |
| model_serving_container_environment_variables=None, |
| model_serving_container_ports=None, |
| model_description=None, |
| model_instance_schema_uri=None, |
| model_parameters_schema_uri=None, |
| model_prediction_schema_uri=None, |
| labels=None, |
| training_encryption_spec_key_name=None, |
| model_encryption_spec_key_name=None, |
| # RUN |
| annotation_schema_uri=None, |
| model_labels=None, |
| base_output_dir=None, |
| service_account=None, |
| network=None, |
| bigquery_destination=None, |
| environment_variables=None, |
| boot_disk_type="pd-ssd", |
| boot_disk_size_gb=100, |
| training_filter_split=None, |
| validation_filter_split=None, |
| test_filter_split=None, |
| predefined_split_column_name=None, |
| timestamp_split_column_name=None, |
| tensorboard=None, |
| sync=True, |
| ) |
| |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CreateCustomPythonPackageTrainingJobOperator.hook")) |
| def test_execute_enters_deferred_state(self, mock_hook): |
| task = CreateCustomPythonPackageTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| python_package_gcs_uri=PYTHON_PACKAGE_GCS_URI, |
| python_module_name=PYTHON_MODULE_NAME, |
| container_uri=CONTAINER_URI, |
| args=ARGS, |
| model_serving_container_image_uri=CONTAINER_URI, |
| model_display_name=DISPLAY_NAME_2, |
| replica_count=REPLICA_COUNT, |
| machine_type=MACHINE_TYPE, |
| accelerator_type=ACCELERATOR_TYPE, |
| accelerator_count=ACCELERATOR_COUNT, |
| training_fraction_split=TRAINING_FRACTION_SPLIT, |
| validation_fraction_split=VALIDATION_FRACTION_SPLIT, |
| test_fraction_split=TEST_FRACTION_SPLIT, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| deferrable=True, |
| ) |
| mock_hook.return_value.exists.return_value = False |
| with pytest.raises(TaskDeferred) as exc: |
| with pytest.warns(AirflowProviderDeprecationWarning, match=SYNC_DEPRECATION_WARNING): |
| task.execute(context={"ti": mock.MagicMock()}) |
| assert isinstance( |
| exc.value.trigger, CustomPythonPackageTrainingJobTrigger |
| ), "Trigger is not a CustomPythonPackageTrainingJobTrigger" |
| |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CreateCustomPythonPackageTrainingJobOperator.xcom_push")) |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CreateCustomPythonPackageTrainingJobOperator.hook")) |
| def test_execute_complete_success(self, mock_hook, mock_xcom_push): |
| task = CreateCustomPythonPackageTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| python_package_gcs_uri=PYTHON_PACKAGE_GCS_URI, |
| python_module_name=PYTHON_MODULE_NAME, |
| container_uri=CONTAINER_URI, |
| args=ARGS, |
| model_serving_container_image_uri=CONTAINER_URI, |
| model_display_name=DISPLAY_NAME_2, |
| replica_count=REPLICA_COUNT, |
| machine_type=MACHINE_TYPE, |
| accelerator_type=ACCELERATOR_TYPE, |
| accelerator_count=ACCELERATOR_COUNT, |
| training_fraction_split=TRAINING_FRACTION_SPLIT, |
| validation_fraction_split=VALIDATION_FRACTION_SPLIT, |
| test_fraction_split=TEST_FRACTION_SPLIT, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| deferrable=True, |
| ) |
| expected_result = {} |
| mock_hook.return_value.exists.return_value = False |
| mock_xcom_push.return_value = None |
| actual_result = task.execute_complete( |
| context=None, event={"status": "success", "message": "", "job": {}} |
| ) |
| assert actual_result == expected_result |
| |
| def test_execute_complete_error_status_raises_exception(self): |
| task = CreateCustomPythonPackageTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| python_package_gcs_uri=PYTHON_PACKAGE_GCS_URI, |
| python_module_name=PYTHON_MODULE_NAME, |
| container_uri=CONTAINER_URI, |
| args=ARGS, |
| model_serving_container_image_uri=CONTAINER_URI, |
| model_display_name=DISPLAY_NAME_2, |
| replica_count=REPLICA_COUNT, |
| machine_type=MACHINE_TYPE, |
| accelerator_type=ACCELERATOR_TYPE, |
| accelerator_count=ACCELERATOR_COUNT, |
| training_fraction_split=TRAINING_FRACTION_SPLIT, |
| validation_fraction_split=VALIDATION_FRACTION_SPLIT, |
| test_fraction_split=TEST_FRACTION_SPLIT, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| deferrable=True, |
| ) |
| with pytest.raises(AirflowException): |
| task.execute_complete(context=None, event={"status": "error", "message": "test message"}) |
| |
| |
| class TestVertexAICreateCustomTrainingJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.Dataset")) |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CustomJobHook")) |
| def test_execute(self, mock_hook, mock_dataset): |
| mock_hook.return_value.create_custom_training_job.return_value = ( |
| None, |
| "training_id", |
| "custom_job_id", |
| ) |
| op = CreateCustomTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| script_path=PYTHON_PACKAGE, |
| args=PYTHON_PACKAGE_CMDARGS, |
| container_uri=CONTAINER_URI, |
| model_serving_container_image_uri=CONTAINER_URI, |
| requirements=[], |
| replica_count=1, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| dataset_id=TEST_DATASET_ID, |
| parent_model=TEST_PARENT_MODEL, |
| ) |
| with pytest.warns(AirflowProviderDeprecationWarning, match=SYNC_DEPRECATION_WARNING): |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_dataset.assert_called_once_with(name=TEST_DATASET_ID) |
| mock_hook.return_value.create_custom_training_job.assert_called_once_with( |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| args=PYTHON_PACKAGE_CMDARGS, |
| container_uri=CONTAINER_URI, |
| model_serving_container_image_uri=CONTAINER_URI, |
| script_path=PYTHON_PACKAGE, |
| requirements=[], |
| dataset=mock_dataset.return_value, |
| model_display_name=None, |
| replica_count=REPLICA_COUNT, |
| machine_type=MACHINE_TYPE, |
| accelerator_type=ACCELERATOR_TYPE, |
| accelerator_count=ACCELERATOR_COUNT, |
| training_fraction_split=None, |
| validation_fraction_split=None, |
| test_fraction_split=None, |
| parent_model=TEST_PARENT_MODEL, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model_serving_container_predict_route=None, |
| model_serving_container_health_route=None, |
| model_serving_container_command=None, |
| model_serving_container_args=None, |
| model_serving_container_environment_variables=None, |
| model_serving_container_ports=None, |
| model_description=None, |
| model_instance_schema_uri=None, |
| model_parameters_schema_uri=None, |
| model_prediction_schema_uri=None, |
| labels=None, |
| training_encryption_spec_key_name=None, |
| model_encryption_spec_key_name=None, |
| # RUN |
| annotation_schema_uri=None, |
| model_labels=None, |
| base_output_dir=None, |
| service_account=None, |
| network=None, |
| bigquery_destination=None, |
| environment_variables=None, |
| boot_disk_type="pd-ssd", |
| boot_disk_size_gb=100, |
| training_filter_split=None, |
| validation_filter_split=None, |
| test_filter_split=None, |
| predefined_split_column_name=None, |
| timestamp_split_column_name=None, |
| tensorboard=None, |
| sync=True, |
| is_default_version=None, |
| model_version_aliases=None, |
| model_version_description=None, |
| ) |
| |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CreateCustomTrainingJobOperator.hook")) |
| def test_execute_enters_deferred_state(self, mock_hook): |
| task = CreateCustomTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| script_path=PYTHON_PACKAGE, |
| args=PYTHON_PACKAGE_CMDARGS, |
| container_uri=CONTAINER_URI, |
| model_serving_container_image_uri=CONTAINER_URI, |
| requirements=[], |
| replica_count=1, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| deferrable=True, |
| ) |
| mock_hook.return_value.exists.return_value = False |
| with pytest.raises(TaskDeferred) as exc: |
| with pytest.warns(AirflowProviderDeprecationWarning, match=SYNC_DEPRECATION_WARNING): |
| task.execute(context={"ti": mock.MagicMock()}) |
| assert isinstance( |
| exc.value.trigger, CustomTrainingJobTrigger |
| ), "Trigger is not a CustomTrainingJobTrigger" |
| |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CreateCustomTrainingJobOperator.xcom_push")) |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CreateCustomTrainingJobOperator.hook")) |
| def test_execute_complete_success(self, mock_hook, mock_xcom_push): |
| task = CreateCustomTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| script_path=PYTHON_PACKAGE, |
| args=PYTHON_PACKAGE_CMDARGS, |
| container_uri=CONTAINER_URI, |
| model_serving_container_image_uri=CONTAINER_URI, |
| requirements=[], |
| replica_count=1, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| deferrable=True, |
| ) |
| expected_result = {} |
| mock_hook.return_value.exists.return_value = False |
| mock_xcom_push.return_value = None |
| actual_result = task.execute_complete( |
| context=None, event={"status": "success", "message": "", "job": {}} |
| ) |
| assert actual_result == expected_result |
| |
| def test_execute_complete_error_status_raises_exception(self): |
| task = CreateCustomTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| script_path=PYTHON_PACKAGE, |
| args=PYTHON_PACKAGE_CMDARGS, |
| container_uri=CONTAINER_URI, |
| model_serving_container_image_uri=CONTAINER_URI, |
| requirements=[], |
| replica_count=1, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| deferrable=True, |
| ) |
| with pytest.raises(AirflowException): |
| task.execute_complete(context=None, event={"status": "error", "message": "test message"}) |
| |
| |
| class TestVertexAIDeleteCustomTrainingJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CustomJobHook")) |
| def test_execute(self, mock_hook): |
| op = DeleteCustomTrainingJobOperator( |
| task_id=TASK_ID, |
| training_pipeline_id=TRAINING_PIPELINE_ID, |
| custom_job_id=CUSTOM_JOB_ID, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.delete_training_pipeline.assert_called_once_with( |
| training_pipeline=TRAINING_PIPELINE_ID, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| mock_hook.return_value.delete_custom_job.assert_called_once_with( |
| custom_job=CUSTOM_JOB_ID, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| @pytest.mark.db_test |
| def test_templating(self, create_task_instance_of_operator): |
| ti = create_task_instance_of_operator( |
| DeleteCustomTrainingJobOperator, |
| # Templated fields |
| training_pipeline_id="{{ 'training-pipeline-id' }}", |
| custom_job_id="{{ 'custom_job_id' }}", |
| region="{{ 'region' }}", |
| project_id="{{ 'project_id' }}", |
| impersonation_chain="{{ 'impersonation-chain' }}", |
| # Other parameters |
| dag_id="test_template_body_templating_dag", |
| task_id="test_template_body_templating_task", |
| execution_date=timezone.datetime(2024, 2, 1, tzinfo=timezone.utc), |
| ) |
| ti.render_templates() |
| task: DeleteCustomTrainingJobOperator = ti.task |
| assert task.training_pipeline_id == "training-pipeline-id" |
| assert task.custom_job_id == "custom_job_id" |
| assert task.region == "region" |
| assert task.project_id == "project_id" |
| assert task.impersonation_chain == "impersonation-chain" |
| |
| # Deprecated aliases |
| with pytest.warns(AirflowProviderDeprecationWarning): |
| assert task.training_pipeline == "training-pipeline-id" |
| with pytest.warns(AirflowProviderDeprecationWarning): |
| assert task.custom_job == "custom_job_id" |
| |
| |
| class TestVertexAIListCustomTrainingJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("custom_job.CustomJobHook")) |
| def test_execute(self, mock_hook): |
| page_token = "page_token" |
| page_size = 42 |
| filter = "filter" |
| read_mask = "read_mask" |
| |
| op = ListCustomTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| read_mask=read_mask, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.list_training_pipelines.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| read_mask=read_mask, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAICreateDatasetOperator: |
| @mock.patch(VERTEX_AI_PATH.format("dataset.Dataset.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("dataset.DatasetHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = CreateDatasetOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| dataset=TEST_DATASET, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.create_dataset.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| dataset=TEST_DATASET, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIDeleteDatasetOperator: |
| @mock.patch(VERTEX_AI_PATH.format("dataset.Dataset.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("dataset.DatasetHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = DeleteDatasetOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| dataset_id=TEST_DATASET_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.delete_dataset.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| dataset=TEST_DATASET_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIExportDataOperator: |
| @mock.patch(VERTEX_AI_PATH.format("dataset.Dataset.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("dataset.DatasetHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = ExportDataOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| dataset_id=TEST_DATASET_ID, |
| export_config=TEST_EXPORT_CONFIG, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.export_data.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| dataset=TEST_DATASET_ID, |
| export_config=TEST_EXPORT_CONFIG, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIImportDataOperator: |
| @mock.patch(VERTEX_AI_PATH.format("dataset.Dataset.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("dataset.DatasetHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = ImportDataOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| dataset_id=TEST_DATASET_ID, |
| import_configs=TEST_IMPORT_CONFIG, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.import_data.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| dataset=TEST_DATASET_ID, |
| import_configs=TEST_IMPORT_CONFIG, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIListDatasetsOperator: |
| @mock.patch(VERTEX_AI_PATH.format("dataset.Dataset.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("dataset.DatasetHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| page_token = "page_token" |
| page_size = 42 |
| filter = "filter" |
| read_mask = "read_mask" |
| order_by = "order_by" |
| |
| op = ListDatasetsOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| filter=filter, |
| page_size=page_size, |
| page_token=page_token, |
| read_mask=read_mask, |
| order_by=order_by, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.list_datasets.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| filter=filter, |
| page_size=page_size, |
| page_token=page_token, |
| read_mask=read_mask, |
| order_by=order_by, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIUpdateDatasetOperator: |
| @mock.patch(VERTEX_AI_PATH.format("dataset.Dataset.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("dataset.DatasetHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = UpdateDatasetOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| dataset_id=TEST_DATASET_ID, |
| dataset=TEST_DATASET, |
| update_mask=TEST_UPDATE_MASK, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.update_dataset.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| dataset_id=TEST_DATASET_ID, |
| dataset=TEST_DATASET, |
| update_mask=TEST_UPDATE_MASK, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAICreateAutoMLForecastingTrainingJobOperator: |
| @mock.patch("google.cloud.aiplatform.datasets.TimeSeriesDataset") |
| @mock.patch(VERTEX_AI_PATH.format("auto_ml.AutoMLHook")) |
| def test_execute(self, mock_hook, mock_dataset): |
| mock_hook.return_value.create_auto_ml_forecasting_training_job.return_value = (None, "training_id") |
| op = CreateAutoMLForecastingTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| display_name=DISPLAY_NAME, |
| dataset_id=TEST_DATASET_ID, |
| target_column=TEST_TRAINING_TARGET_COLUMN, |
| time_column=TEST_TRAINING_TIME_COLUMN, |
| time_series_identifier_column=TEST_TRAINING_TIME_SERIES_IDENTIFIER_COLUMN, |
| unavailable_at_forecast_columns=TEST_TRAINING_UNAVAILABLE_AT_FORECAST_COLUMNS, |
| available_at_forecast_columns=TEST_TRAINING_AVAILABLE_AT_FORECAST_COLUMNS, |
| forecast_horizon=TEST_TRAINING_FORECAST_HORIZON, |
| data_granularity_unit=TEST_TRAINING_DATA_GRANULARITY_UNIT, |
| data_granularity_count=TEST_TRAINING_DATA_GRANULARITY_COUNT, |
| sync=True, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| parent_model=TEST_PARENT_MODEL, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_dataset.assert_called_once_with(dataset_name=TEST_DATASET_ID) |
| mock_hook.return_value.create_auto_ml_forecasting_training_job.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| display_name=DISPLAY_NAME, |
| dataset=mock_dataset.return_value, |
| target_column=TEST_TRAINING_TARGET_COLUMN, |
| time_column=TEST_TRAINING_TIME_COLUMN, |
| time_series_identifier_column=TEST_TRAINING_TIME_SERIES_IDENTIFIER_COLUMN, |
| unavailable_at_forecast_columns=TEST_TRAINING_UNAVAILABLE_AT_FORECAST_COLUMNS, |
| available_at_forecast_columns=TEST_TRAINING_AVAILABLE_AT_FORECAST_COLUMNS, |
| forecast_horizon=TEST_TRAINING_FORECAST_HORIZON, |
| data_granularity_unit=TEST_TRAINING_DATA_GRANULARITY_UNIT, |
| data_granularity_count=TEST_TRAINING_DATA_GRANULARITY_COUNT, |
| parent_model=TEST_PARENT_MODEL, |
| optimization_objective=None, |
| column_specs=None, |
| column_transformations=None, |
| labels=None, |
| training_encryption_spec_key_name=None, |
| model_encryption_spec_key_name=None, |
| training_fraction_split=None, |
| validation_fraction_split=None, |
| test_fraction_split=None, |
| predefined_split_column_name=None, |
| weight_column=None, |
| time_series_attribute_columns=None, |
| context_window=None, |
| export_evaluated_data_items=False, |
| export_evaluated_data_items_bigquery_destination_uri=None, |
| export_evaluated_data_items_override_destination=False, |
| quantiles=None, |
| validation_options=None, |
| budget_milli_node_hours=1000, |
| model_display_name=None, |
| model_labels=None, |
| sync=True, |
| is_default_version=None, |
| model_version_aliases=None, |
| model_version_description=None, |
| ) |
| |
| |
| class TestVertexAICreateAutoMLImageTrainingJobOperator: |
| @mock.patch("google.cloud.aiplatform.datasets.ImageDataset") |
| @mock.patch(VERTEX_AI_PATH.format("auto_ml.AutoMLHook")) |
| def test_execute(self, mock_hook, mock_dataset): |
| mock_hook.return_value.create_auto_ml_image_training_job.return_value = (None, "training_id") |
| op = CreateAutoMLImageTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| display_name=DISPLAY_NAME, |
| dataset_id=TEST_DATASET_ID, |
| prediction_type="classification", |
| multi_label=False, |
| model_type="CLOUD", |
| sync=True, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| parent_model=TEST_PARENT_MODEL, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_dataset.assert_called_once_with(dataset_name=TEST_DATASET_ID) |
| mock_hook.return_value.create_auto_ml_image_training_job.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| display_name=DISPLAY_NAME, |
| dataset=mock_dataset.return_value, |
| prediction_type="classification", |
| parent_model=TEST_PARENT_MODEL, |
| multi_label=False, |
| model_type="CLOUD", |
| base_model=None, |
| labels=None, |
| training_encryption_spec_key_name=None, |
| model_encryption_spec_key_name=None, |
| training_fraction_split=None, |
| validation_fraction_split=None, |
| test_fraction_split=None, |
| training_filter_split=None, |
| validation_filter_split=None, |
| test_filter_split=None, |
| budget_milli_node_hours=None, |
| model_display_name=None, |
| model_labels=None, |
| disable_early_stopping=False, |
| sync=True, |
| is_default_version=None, |
| model_version_aliases=None, |
| model_version_description=None, |
| ) |
| |
| |
| class TestVertexAICreateAutoMLTabularTrainingJobOperator: |
| @mock.patch("google.cloud.aiplatform.datasets.TabularDataset") |
| @mock.patch(VERTEX_AI_PATH.format("auto_ml.AutoMLHook")) |
| def test_execute(self, mock_hook, mock_dataset): |
| mock_hook.return_value = MagicMock( |
| **{ |
| "create_auto_ml_tabular_training_job.return_value": (None, "training_id"), |
| "get_credentials_and_project_id.return_value": ("creds", "project_id"), |
| } |
| ) |
| op = CreateAutoMLTabularTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| display_name=DISPLAY_NAME, |
| dataset_id=TEST_DATASET_ID, |
| target_column=None, |
| optimization_prediction_type=None, |
| sync=True, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| parent_model=TEST_PARENT_MODEL, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_dataset.assert_called_once_with( |
| dataset_name=TEST_DATASET_ID, project=GCP_PROJECT, credentials="creds" |
| ) |
| mock_hook.return_value.create_auto_ml_tabular_training_job.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| display_name=DISPLAY_NAME, |
| dataset=mock_dataset.return_value, |
| parent_model=TEST_PARENT_MODEL, |
| target_column=None, |
| optimization_prediction_type=None, |
| optimization_objective=None, |
| column_specs=None, |
| column_transformations=None, |
| optimization_objective_recall_value=None, |
| optimization_objective_precision_value=None, |
| labels=None, |
| training_encryption_spec_key_name=None, |
| model_encryption_spec_key_name=None, |
| training_fraction_split=None, |
| validation_fraction_split=None, |
| test_fraction_split=None, |
| predefined_split_column_name=None, |
| timestamp_split_column_name=None, |
| weight_column=None, |
| budget_milli_node_hours=1000, |
| model_display_name=None, |
| model_labels=None, |
| disable_early_stopping=False, |
| export_evaluated_data_items=False, |
| export_evaluated_data_items_bigquery_destination_uri=None, |
| export_evaluated_data_items_override_destination=False, |
| sync=True, |
| is_default_version=None, |
| model_version_aliases=None, |
| model_version_description=None, |
| ) |
| |
| |
| class TestVertexAICreateAutoMLTextTrainingJobOperator: |
| @mock.patch("google.cloud.aiplatform.datasets.TextDataset") |
| @mock.patch(VERTEX_AI_PATH.format("auto_ml.AutoMLHook")) |
| def test_execute(self, mock_hook, mock_dataset): |
| mock_hook.return_value.create_auto_ml_text_training_job.return_value = (None, "training_id") |
| op = CreateAutoMLTextTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| display_name=DISPLAY_NAME, |
| dataset_id=TEST_DATASET_ID, |
| prediction_type=None, |
| multi_label=False, |
| sentiment_max=10, |
| sync=True, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| parent_model=TEST_PARENT_MODEL, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_dataset.assert_called_once_with(dataset_name=TEST_DATASET_ID) |
| mock_hook.return_value.create_auto_ml_text_training_job.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| display_name=DISPLAY_NAME, |
| dataset=mock_dataset.return_value, |
| parent_model=TEST_PARENT_MODEL, |
| prediction_type=None, |
| multi_label=False, |
| sentiment_max=10, |
| labels=None, |
| training_encryption_spec_key_name=None, |
| model_encryption_spec_key_name=None, |
| training_fraction_split=None, |
| validation_fraction_split=None, |
| test_fraction_split=None, |
| training_filter_split=None, |
| validation_filter_split=None, |
| test_filter_split=None, |
| model_display_name=None, |
| model_labels=None, |
| sync=True, |
| is_default_version=None, |
| model_version_aliases=None, |
| model_version_description=None, |
| ) |
| |
| |
| class TestVertexAICreateAutoMLVideoTrainingJobOperator: |
| @mock.patch("google.cloud.aiplatform.datasets.VideoDataset") |
| @mock.patch(VERTEX_AI_PATH.format("auto_ml.AutoMLHook")) |
| def test_execute(self, mock_hook, mock_dataset): |
| mock_hook.return_value.create_auto_ml_video_training_job.return_value = (None, "training_id") |
| op = CreateAutoMLVideoTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| display_name=DISPLAY_NAME, |
| dataset_id=TEST_DATASET_ID, |
| prediction_type="classification", |
| model_type="CLOUD", |
| sync=True, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| parent_model=TEST_PARENT_MODEL, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_dataset.assert_called_once_with(dataset_name=TEST_DATASET_ID) |
| mock_hook.return_value.create_auto_ml_video_training_job.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| display_name=DISPLAY_NAME, |
| dataset=mock_dataset.return_value, |
| parent_model=TEST_PARENT_MODEL, |
| prediction_type="classification", |
| model_type="CLOUD", |
| labels=None, |
| training_encryption_spec_key_name=None, |
| model_encryption_spec_key_name=None, |
| training_fraction_split=None, |
| test_fraction_split=None, |
| training_filter_split=None, |
| test_filter_split=None, |
| model_display_name=None, |
| model_labels=None, |
| sync=True, |
| is_default_version=None, |
| model_version_aliases=None, |
| model_version_description=None, |
| ) |
| |
| |
| class TestVertexAIDeleteAutoMLTrainingJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("auto_ml.AutoMLHook")) |
| def test_execute(self, mock_hook): |
| op = DeleteAutoMLTrainingJobOperator( |
| task_id=TASK_ID, |
| training_pipeline_id=TRAINING_PIPELINE_ID, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.delete_training_pipeline.assert_called_once_with( |
| training_pipeline=TRAINING_PIPELINE_ID, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| @pytest.mark.db_test |
| def test_templating(self, create_task_instance_of_operator): |
| ti = create_task_instance_of_operator( |
| DeleteAutoMLTrainingJobOperator, |
| # Templated fields |
| training_pipeline_id="{{ 'training-pipeline-id' }}", |
| region="{{ 'region' }}", |
| project_id="{{ 'project-id' }}", |
| impersonation_chain="{{ 'impersonation-chain' }}", |
| # Other parameters |
| dag_id="test_template_body_templating_dag", |
| task_id="test_template_body_templating_task", |
| execution_date=timezone.datetime(2024, 2, 1, tzinfo=timezone.utc), |
| ) |
| ti.render_templates() |
| task: DeleteAutoMLTrainingJobOperator = ti.task |
| assert task.training_pipeline_id == "training-pipeline-id" |
| assert task.region == "region" |
| assert task.project_id == "project-id" |
| assert task.impersonation_chain == "impersonation-chain" |
| |
| with pytest.warns(AirflowProviderDeprecationWarning): |
| assert task.training_pipeline == "training-pipeline-id" |
| |
| |
| class TestVertexAIListAutoMLTrainingJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("auto_ml.AutoMLHook")) |
| def test_execute(self, mock_hook): |
| page_token = "page_token" |
| page_size = 42 |
| filter = "filter" |
| read_mask = "read_mask" |
| |
| op = ListAutoMLTrainingJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| read_mask=read_mask, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.list_training_pipelines.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| read_mask=read_mask, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAICreateBatchPredictionJobOperator: |
| @mock.patch(VERTEX_AI_LINKS_PATH.format("VertexAIBatchPredictionJobLink.persist")) |
| @mock.patch(VERTEX_AI_PATH.format("batch_prediction_job.BatchPredictionJobHook")) |
| def test_execute(self, mock_hook, mock_link_persist): |
| mock_job = mock_hook.return_value.submit_batch_prediction_job.return_value |
| mock_job.name = TEST_BATCH_PREDICTION_JOB_ID |
| |
| op = CreateBatchPredictionJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| job_display_name=TEST_JOB_DISPLAY_NAME, |
| model_name=TEST_MODEL_NAME, |
| instances_format="jsonl", |
| predictions_format="jsonl", |
| create_request_timeout=TEST_CREATE_REQUEST_TIMEOUT, |
| batch_size=TEST_BATCH_SIZE, |
| ) |
| context = {"ti": mock.MagicMock()} |
| op.execute(context=context) |
| |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.submit_batch_prediction_job.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| job_display_name=TEST_JOB_DISPLAY_NAME, |
| model_name=TEST_MODEL_NAME, |
| instances_format="jsonl", |
| predictions_format="jsonl", |
| gcs_source=None, |
| bigquery_source=None, |
| gcs_destination_prefix=None, |
| bigquery_destination_prefix=None, |
| model_parameters=None, |
| machine_type=None, |
| accelerator_type=None, |
| accelerator_count=None, |
| starting_replica_count=None, |
| max_replica_count=None, |
| generate_explanation=False, |
| explanation_metadata=None, |
| explanation_parameters=None, |
| labels=None, |
| encryption_spec_key_name=None, |
| create_request_timeout=TEST_CREATE_REQUEST_TIMEOUT, |
| batch_size=TEST_BATCH_SIZE, |
| ) |
| mock_job.wait_for_completion.assert_called_once() |
| mock_job.to_dict.assert_called_once() |
| mock_link_persist.assert_called_once_with( |
| context=context, task_instance=op, batch_prediction_job_id=TEST_BATCH_PREDICTION_JOB_ID |
| ) |
| |
| @mock.patch(VERTEX_AI_LINKS_PATH.format("VertexAIBatchPredictionJobLink.persist")) |
| @mock.patch(VERTEX_AI_PATH.format("batch_prediction_job.BatchPredictionJobHook")) |
| def test_execute_deferrable(self, mock_hook, mock_link_persist): |
| mock_job = mock_hook.return_value.submit_batch_prediction_job.return_value |
| mock_job.name = TEST_BATCH_PREDICTION_JOB_ID |
| op = CreateBatchPredictionJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| job_display_name=TEST_JOB_DISPLAY_NAME, |
| model_name=TEST_MODEL_NAME, |
| instances_format="jsonl", |
| predictions_format="jsonl", |
| create_request_timeout=TEST_CREATE_REQUEST_TIMEOUT, |
| batch_size=TEST_BATCH_SIZE, |
| deferrable=True, |
| ) |
| context = {"ti": mock.MagicMock()} |
| with pytest.raises(TaskDeferred) as exception_info: |
| op.execute(context=context) |
| |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.submit_batch_prediction_job.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| job_display_name=TEST_JOB_DISPLAY_NAME, |
| model_name=TEST_MODEL_NAME, |
| instances_format="jsonl", |
| predictions_format="jsonl", |
| gcs_source=None, |
| bigquery_source=None, |
| gcs_destination_prefix=None, |
| bigquery_destination_prefix=None, |
| model_parameters=None, |
| machine_type=None, |
| accelerator_type=None, |
| accelerator_count=None, |
| starting_replica_count=None, |
| max_replica_count=None, |
| generate_explanation=False, |
| explanation_metadata=None, |
| explanation_parameters=None, |
| labels=None, |
| encryption_spec_key_name=None, |
| create_request_timeout=TEST_CREATE_REQUEST_TIMEOUT, |
| batch_size=TEST_BATCH_SIZE, |
| ) |
| mock_job.wait_for_completion.assert_not_called() |
| mock_job.to_dict.assert_not_called() |
| mock_link_persist.assert_called_once_with( |
| batch_prediction_job_id=TEST_BATCH_PREDICTION_JOB_ID, |
| context=context, |
| task_instance=op, |
| ) |
| assert hasattr(exception_info.value, "trigger") |
| assert exception_info.value.trigger.conn_id == GCP_CONN_ID |
| assert exception_info.value.trigger.project_id == GCP_PROJECT |
| assert exception_info.value.trigger.location == GCP_LOCATION |
| assert exception_info.value.trigger.job_id == TEST_BATCH_PREDICTION_JOB_ID |
| assert exception_info.value.trigger.poll_interval == 10 |
| assert exception_info.value.trigger.impersonation_chain == IMPERSONATION_CHAIN |
| |
| @mock.patch(VERTEX_AI_PATH.format("batch_prediction_job.CreateBatchPredictionJobOperator.xcom_push")) |
| @mock.patch(VERTEX_AI_PATH.format("batch_prediction_job.BatchPredictionJobHook")) |
| def test_execute_complete(self, mock_hook, mock_xcom_push): |
| context = mock.MagicMock() |
| mock_job = {"name": TEST_JOB_DISPLAY_NAME} |
| event = { |
| "status": "success", |
| "job": mock_job, |
| } |
| mock_hook.return_value.extract_batch_prediction_job_id.return_value = TEST_BATCH_PREDICTION_JOB_ID |
| |
| op = CreateBatchPredictionJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| job_display_name=TEST_JOB_DISPLAY_NAME, |
| model_name=TEST_MODEL_NAME, |
| instances_format="jsonl", |
| predictions_format="jsonl", |
| create_request_timeout=TEST_CREATE_REQUEST_TIMEOUT, |
| batch_size=TEST_BATCH_SIZE, |
| ) |
| execute_complete_result = op.execute_complete(context=context, event=event) |
| |
| mock_hook.return_value.extract_batch_prediction_job_id.assert_called_once_with(mock_job) |
| mock_xcom_push.assert_has_calls( |
| [ |
| call(context, key="batch_prediction_job_id", value=TEST_BATCH_PREDICTION_JOB_ID), |
| call( |
| context, |
| key="training_conf", |
| value={ |
| "training_conf_id": TEST_BATCH_PREDICTION_JOB_ID, |
| "region": GCP_LOCATION, |
| "project_id": GCP_PROJECT, |
| }, |
| ), |
| ] |
| ) |
| assert execute_complete_result == mock_job |
| |
| |
| class TestVertexAIDeleteBatchPredictionJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("batch_prediction_job.BatchPredictionJobHook")) |
| def test_execute(self, mock_hook): |
| op = DeleteBatchPredictionJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| batch_prediction_job_id=TEST_BATCH_PREDICTION_JOB_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.delete_batch_prediction_job.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| batch_prediction_job=TEST_BATCH_PREDICTION_JOB_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIListBatchPredictionJobsOperator: |
| @mock.patch(VERTEX_AI_PATH.format("batch_prediction_job.BatchPredictionJobHook")) |
| def test_execute(self, mock_hook): |
| page_token = "page_token" |
| page_size = 42 |
| filter = "filter" |
| read_mask = "read_mask" |
| |
| op = ListBatchPredictionJobsOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| read_mask=read_mask, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.list_batch_prediction_jobs.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| read_mask=read_mask, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAICreateEndpointOperator: |
| @mock.patch(VERTEX_AI_PATH.format("endpoint_service.Endpoint.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("endpoint_service.EndpointServiceHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = CreateEndpointOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| endpoint=TEST_ENDPOINT, |
| endpoint_id=TEST_ENDPOINT_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.create_endpoint.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| endpoint=TEST_ENDPOINT, |
| endpoint_id=TEST_ENDPOINT_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIDeleteEndpointOperator: |
| @mock.patch(VERTEX_AI_PATH.format("endpoint_service.EndpointServiceHook")) |
| def test_execute(self, mock_hook): |
| op = DeleteEndpointOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| endpoint_id=TEST_ENDPOINT_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.delete_endpoint.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| endpoint=TEST_ENDPOINT_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIDeployModelOperator: |
| @mock.patch(VERTEX_AI_PATH.format("endpoint_service.endpoint_service.DeployModelResponse.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("endpoint_service.EndpointServiceHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = DeployModelOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| endpoint_id=TEST_ENDPOINT_ID, |
| deployed_model=TEST_DEPLOYED_MODEL, |
| traffic_split=None, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.deploy_model.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| endpoint=TEST_ENDPOINT_ID, |
| deployed_model=TEST_DEPLOYED_MODEL, |
| traffic_split=None, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIListEndpointsOperator: |
| @mock.patch(VERTEX_AI_PATH.format("endpoint_service.EndpointServiceHook")) |
| def test_execute(self, mock_hook): |
| page_token = "page_token" |
| page_size = 42 |
| filter = "filter" |
| read_mask = "read_mask" |
| order_by = "order_by" |
| |
| op = ListEndpointsOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| read_mask=read_mask, |
| order_by=order_by, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.list_endpoints.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| read_mask=read_mask, |
| order_by=order_by, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIUndeployModelOperator: |
| @mock.patch(VERTEX_AI_PATH.format("endpoint_service.EndpointServiceHook")) |
| def test_execute(self, mock_hook): |
| op = UndeployModelOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| endpoint_id=TEST_ENDPOINT_ID, |
| deployed_model_id=TEST_DEPLOYED_MODEL_ID, |
| traffic_split=None, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.undeploy_model.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| endpoint=TEST_ENDPOINT_ID, |
| deployed_model_id=TEST_DEPLOYED_MODEL_ID, |
| traffic_split=None, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAICreateHyperparameterTuningJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("hyperparameter_tuning_job.types.HyperparameterTuningJob.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("hyperparameter_tuning_job.HyperparameterTuningJobHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = CreateHyperparameterTuningJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| worker_pool_specs=[], |
| sync=False, |
| parameter_spec={}, |
| metric_spec={}, |
| max_trial_count=15, |
| parallel_trial_count=3, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.create_hyperparameter_tuning_job.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| display_name=DISPLAY_NAME, |
| metric_spec={}, |
| parameter_spec={}, |
| max_trial_count=15, |
| parallel_trial_count=3, |
| worker_pool_specs=[], |
| base_output_dir=None, |
| custom_job_labels=None, |
| custom_job_encryption_spec_key_name=None, |
| staging_bucket=STAGING_BUCKET, |
| max_failed_trial_count=0, |
| search_algorithm=None, |
| measurement_selection="best", |
| hyperparameter_tuning_job_labels=None, |
| hyperparameter_tuning_job_encryption_spec_key_name=None, |
| service_account=None, |
| network=None, |
| timeout=None, |
| restart_job_on_worker_restart=False, |
| enable_web_access=False, |
| tensorboard=None, |
| sync=False, |
| wait_job_completed=False, |
| ) |
| |
| @mock.patch( |
| VERTEX_AI_PATH.format("hyperparameter_tuning_job.CreateHyperparameterTuningJobOperator.defer") |
| ) |
| @mock.patch(VERTEX_AI_PATH.format("hyperparameter_tuning_job.HyperparameterTuningJobHook")) |
| def test_deferrable(self, mock_hook, mock_defer): |
| op = CreateHyperparameterTuningJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| worker_pool_specs=[], |
| sync=False, |
| parameter_spec={}, |
| metric_spec={}, |
| max_trial_count=15, |
| parallel_trial_count=3, |
| deferrable=True, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_defer.assert_called_once() |
| |
| @pytest.mark.db_test |
| def test_deferrable_sync_error(self): |
| op = CreateHyperparameterTuningJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| worker_pool_specs=[], |
| sync=True, |
| parameter_spec={}, |
| metric_spec={}, |
| max_trial_count=15, |
| parallel_trial_count=3, |
| deferrable=True, |
| ) |
| with pytest.raises(AirflowException): |
| op.execute(context={"ti": mock.MagicMock()}) |
| |
| @mock.patch(VERTEX_AI_PATH.format("hyperparameter_tuning_job.HyperparameterTuningJobHook")) |
| def test_execute_complete(self, mock_hook): |
| test_job_id = "test_job_id" |
| test_job = {"name": f"test/{test_job_id}"} |
| event = { |
| "status": "success", |
| "message": "test message", |
| "job": test_job, |
| } |
| mock_hook.return_value.extract_hyperparameter_tuning_job_id.return_value = test_job_id |
| mock_context = mock.MagicMock() |
| |
| op = CreateHyperparameterTuningJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| worker_pool_specs=[], |
| sync=False, |
| parameter_spec={}, |
| metric_spec={}, |
| max_trial_count=15, |
| parallel_trial_count=3, |
| ) |
| |
| result = op.execute_complete(context=mock_context, event=event) |
| |
| assert result == test_job |
| |
| def test_execute_complete_error(self): |
| event = { |
| "status": "error", |
| "message": "test error message", |
| } |
| |
| op = CreateHyperparameterTuningJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| staging_bucket=STAGING_BUCKET, |
| display_name=DISPLAY_NAME, |
| worker_pool_specs=[], |
| sync=False, |
| parameter_spec={}, |
| metric_spec={}, |
| max_trial_count=15, |
| parallel_trial_count=3, |
| ) |
| |
| with pytest.raises(AirflowException): |
| op.execute_complete(context=mock.MagicMock(), event=event) |
| |
| |
| class TestVertexAIGetHyperparameterTuningJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("hyperparameter_tuning_job.types.HyperparameterTuningJob.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("hyperparameter_tuning_job.HyperparameterTuningJobHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = GetHyperparameterTuningJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| hyperparameter_tuning_job_id=TEST_HYPERPARAMETER_TUNING_JOB_ID, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.get_hyperparameter_tuning_job.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| hyperparameter_tuning_job=TEST_HYPERPARAMETER_TUNING_JOB_ID, |
| retry=DEFAULT, |
| timeout=None, |
| metadata=(), |
| ) |
| |
| |
| class TestVertexAIDeleteHyperparameterTuningJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("hyperparameter_tuning_job.HyperparameterTuningJobHook")) |
| def test_execute(self, mock_hook): |
| op = DeleteHyperparameterTuningJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| hyperparameter_tuning_job_id=TEST_HYPERPARAMETER_TUNING_JOB_ID, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.delete_hyperparameter_tuning_job.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| hyperparameter_tuning_job=TEST_HYPERPARAMETER_TUNING_JOB_ID, |
| retry=DEFAULT, |
| timeout=None, |
| metadata=(), |
| ) |
| |
| |
| class TestVertexAIListHyperparameterTuningJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("hyperparameter_tuning_job.HyperparameterTuningJobHook")) |
| def test_execute(self, mock_hook): |
| page_token = "page_token" |
| page_size = 42 |
| filter = "filter" |
| read_mask = "read_mask" |
| |
| op = ListHyperparameterTuningJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| read_mask=read_mask, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.list_hyperparameter_tuning_jobs.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| read_mask=read_mask, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIExportModelOperator: |
| @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) |
| def test_execute(self, mock_hook): |
| op = ExportModelOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model_id=TEST_MODEL_ID, |
| output_config=TEST_OUTPUT_CONFIG, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.export_model.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model=TEST_MODEL_ID, |
| output_config=TEST_OUTPUT_CONFIG, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIDeleteModelOperator: |
| @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) |
| def test_execute(self, mock_hook): |
| op = DeleteModelOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model_id=TEST_MODEL_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.delete_model.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model=TEST_MODEL_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIListModelsOperator: |
| @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| page_token = "page_token" |
| page_size = 42 |
| filter = "filter" |
| read_mask = "read_mask" |
| order_by = "order_by" |
| |
| op = ListModelsOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| filter=filter, |
| page_size=page_size, |
| page_token=page_token, |
| read_mask=read_mask, |
| order_by=order_by, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.list_models.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| filter=filter, |
| page_size=page_size, |
| page_token=page_token, |
| read_mask=read_mask, |
| order_by=order_by, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIUploadModelOperator: |
| @mock.patch(VERTEX_AI_PATH.format("model_service.model_service.UploadModelResponse.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = UploadModelOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model=TEST_MODEL_OBJ, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.upload_model.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model=TEST_MODEL_OBJ, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIGetModelsOperator: |
| @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = GetModelOperator( |
| task_id=TASK_ID, |
| model_id=TEST_MODEL_NAME, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.get_model.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model_id=TEST_MODEL_NAME, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIListModelVersionsOperator: |
| @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = ListModelVersionsOperator( |
| task_id=TASK_ID, |
| model_id=TEST_MODEL_NAME, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.list_model_versions.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model_id=TEST_MODEL_NAME, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAISetDefaultVersionOnModelOperator: |
| @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = SetDefaultVersionOnModelOperator( |
| task_id=TASK_ID, |
| model_id=TEST_MODEL_NAME, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.set_version_as_default.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model_id=TEST_MODEL_NAME, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIAddVersionAliasesOnModelOperator: |
| @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = AddVersionAliasesOnModelOperator( |
| task_id=TASK_ID, |
| model_id=TEST_MODEL_NAME, |
| version_aliases=TEST_VERSION_ALIASES, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.add_version_aliases.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model_id=TEST_MODEL_NAME, |
| version_aliases=TEST_VERSION_ALIASES, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIDeleteVersionAliasesOnModelOperator: |
| @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = DeleteVersionAliasesOnModelOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model_id=TEST_MODEL_ID, |
| version_aliases=TEST_VERSION_ALIASES, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.delete_version_aliases.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model_id=TEST_MODEL_ID, |
| version_aliases=TEST_VERSION_ALIASES, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIDeleteModelVersionOperator: |
| @mock.patch(VERTEX_AI_PATH.format("model_service.Model.to_dict")) |
| @mock.patch(VERTEX_AI_PATH.format("model_service.ModelServiceHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = DeleteModelVersionOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model_id=TEST_MODEL_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.delete_model_version.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| model_id=TEST_MODEL_ID, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| |
| |
| class TestVertexAIRunPipelineJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("pipeline_job.PipelineJobHook")) |
| @mock.patch("google.cloud.aiplatform_v1.types.PipelineJob.to_dict") |
| def test_execute(self, to_dict_mock, mock_hook): |
| op = RunPipelineJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| display_name=DISPLAY_NAME, |
| template_path=TEST_TEMPLATE_PATH, |
| job_id=TEST_PIPELINE_JOB_ID, |
| pipeline_root="", |
| parameter_values={}, |
| input_artifacts={}, |
| enable_caching=False, |
| encryption_spec_key_name="", |
| labels={}, |
| failure_policy="", |
| service_account="", |
| network="", |
| create_request_timeout=None, |
| experiment=None, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.submit_pipeline_job.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| display_name=DISPLAY_NAME, |
| template_path=TEST_TEMPLATE_PATH, |
| job_id=TEST_PIPELINE_JOB_ID, |
| pipeline_root="", |
| parameter_values={}, |
| input_artifacts={}, |
| enable_caching=False, |
| encryption_spec_key_name="", |
| labels={}, |
| failure_policy="", |
| service_account="", |
| network="", |
| create_request_timeout=None, |
| experiment=None, |
| ) |
| |
| @mock.patch(VERTEX_AI_PATH.format("pipeline_job.PipelineJobHook")) |
| def test_execute_enters_deferred_state(self, mock_hook): |
| task = RunPipelineJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| display_name=DISPLAY_NAME, |
| template_path=TEST_TEMPLATE_PATH, |
| job_id=TEST_PIPELINE_JOB_ID, |
| deferrable=True, |
| ) |
| mock_hook.return_value.exists.return_value = False |
| with pytest.raises(TaskDeferred) as exc: |
| task.execute(context={"ti": mock.MagicMock()}) |
| assert isinstance(exc.value.trigger, RunPipelineJobTrigger), "Trigger is not a RunPipelineJobTrigger" |
| |
| @mock.patch(VERTEX_AI_PATH.format("pipeline_job.RunPipelineJobOperator.xcom_push")) |
| @mock.patch(VERTEX_AI_PATH.format("pipeline_job.PipelineJobHook")) |
| def test_execute_complete_success(self, mock_hook, mock_xcom_push): |
| task = RunPipelineJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| display_name=DISPLAY_NAME, |
| template_path=TEST_TEMPLATE_PATH, |
| job_id=TEST_PIPELINE_JOB_ID, |
| deferrable=True, |
| ) |
| expected_pipeline_job = expected_result = { |
| "name": f"projects/{GCP_PROJECT}/locations/{GCP_LOCATION}/pipelineJobs/{TEST_PIPELINE_JOB_ID}", |
| } |
| mock_hook.return_value.exists.return_value = False |
| mock_xcom_push.return_value = None |
| actual_result = task.execute_complete( |
| context=None, event={"status": "success", "message": "", "job": expected_pipeline_job} |
| ) |
| assert actual_result == expected_result |
| |
| def test_execute_complete_error_status_raises_exception(self): |
| task = RunPipelineJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| display_name=DISPLAY_NAME, |
| template_path=TEST_TEMPLATE_PATH, |
| job_id=TEST_PIPELINE_JOB_ID, |
| deferrable=True, |
| ) |
| with pytest.raises(AirflowException): |
| task.execute_complete( |
| context=None, event={"status": "error", "message": "test message", "job": None} |
| ) |
| |
| |
| class TestVertexAIGetPipelineJobOperator: |
| @mock.patch("google.cloud.aiplatform_v1.types.PipelineJob.to_dict") |
| @mock.patch(VERTEX_AI_PATH.format("pipeline_job.PipelineJobHook")) |
| def test_execute(self, mock_hook, to_dict_mock): |
| op = GetPipelineJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| pipeline_job_id=TEST_PIPELINE_JOB_ID, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.get_pipeline_job.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| pipeline_job_id=TEST_PIPELINE_JOB_ID, |
| retry=DEFAULT, |
| timeout=None, |
| metadata=(), |
| ) |
| |
| |
| class TestVertexAIDeletePipelineJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("pipeline_job.PipelineJobHook")) |
| def test_execute(self, mock_hook): |
| op = DeletePipelineJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| pipeline_job_id=TEST_PIPELINE_JOB_ID, |
| ) |
| op.execute(context={}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.delete_pipeline_job.assert_called_once_with( |
| project_id=GCP_PROJECT, |
| region=GCP_LOCATION, |
| pipeline_job_id=TEST_PIPELINE_JOB_ID, |
| retry=DEFAULT, |
| timeout=None, |
| metadata=(), |
| ) |
| |
| |
| class TestVertexAIListPipelineJobOperator: |
| @mock.patch(VERTEX_AI_PATH.format("pipeline_job.PipelineJobHook")) |
| def test_execute(self, mock_hook): |
| page_token = "page_token" |
| page_size = 42 |
| filter = "filter" |
| order_by = "order_by" |
| |
| op = ListPipelineJobOperator( |
| task_id=TASK_ID, |
| gcp_conn_id=GCP_CONN_ID, |
| impersonation_chain=IMPERSONATION_CHAIN, |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| order_by=order_by, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
| op.execute(context={"ti": mock.MagicMock()}) |
| mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) |
| mock_hook.return_value.list_pipeline_jobs.assert_called_once_with( |
| region=GCP_LOCATION, |
| project_id=GCP_PROJECT, |
| page_size=page_size, |
| page_token=page_token, |
| filter=filter, |
| order_by=order_by, |
| retry=RETRY, |
| timeout=TIMEOUT, |
| metadata=METADATA, |
| ) |
