Skip to content

Commit

Permalink
Update Vertex AI system tests (#34364)
Browse files Browse the repository at this point in the history 
* Update Vertex AI system tests

* Update image uri for "Hyperparameter" test
  • Loading branch information
@MaksYermak
MaksYermak committed Sep 14, 2023
1 parent f93b046 commit 013c95b
Show file tree
Hide file tree
Showing 20 changed files with 169 additions and 989 deletions.
754 changes: 0 additions & 754 deletions airflow/providers/google/cloud/example_dags/example_vertex_ai.py
View file

This file was deleted.

47 changes: 19 additions & 28 deletions tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_batch_prediction_job.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,15 +26,17 @@

import os
from datetime import datetime
from pathlib import Path

from google.cloud.aiplatform import schema
from google.protobuf.json_format import ParseDict
from google.protobuf.struct_pb2 import Value

from airflow import models
from airflow.operators.bash import BashOperator
from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
from airflow.providers.google.cloud.operators.gcs import (
GCSCreateBucketOperator,
GCSDeleteBucketOperator,
GCSSynchronizeBucketsOperator,
)
from airflow.providers.google.cloud.operators.vertex_ai.auto_ml import (
CreateAutoMLForecastingTrainingJobOperator,
DeleteAutoMLTrainingJobOperator,
Expand All @@ -48,22 +50,20 @@
CreateDatasetOperator,
DeleteDatasetOperator,
)
from airflow.providers.google.cloud.transfers.local_to_gcs import LocalFilesystemToGCSOperator
from airflow.utils.trigger_rule import TriggerRule

ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default")
DAG_ID = "vertex_ai_batch_prediction_job_operations"
DAG_ID = "example_vertex_ai_batch_prediction_operations"
REGION = "us-central1"

FORECAST_DISPLAY_NAME = f"auto-ml-forecasting-{ENV_ID}"
MODEL_DISPLAY_NAME = f"auto-ml-forecasting-model-{ENV_ID}"

JOB_DISPLAY_NAME = f"batch_prediction_job_test_{ENV_ID}"
DATA_SAMPLE_GCS_BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}"
RESOURCE_DATA_BUCKET = "airflow-system-tests-resources"
DATA_SAMPLE_GCS_BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}".replace("_", "-")
DATA_SAMPLE_GCS_OBJECT_NAME = "vertex-ai/forecast-dataset.csv"
FORECAST_ZIP_CSV_FILE_LOCAL_PATH = str(Path(__file__).parent / "resources" / "forecast-dataset.csv.zip")
FORECAST_CSV_FILE_LOCAL_PATH = "/batch-prediction/forecast-dataset.csv"

FORECAST_DATASET = {
"display_name": f"forecast-dataset-{ENV_ID}",
Expand Down Expand Up @@ -109,18 +109,15 @@
location=REGION,
)

unzip_file = BashOperator(
task_id="unzip_csv_data_file",
bash_command=f"mkdir -p /batch-prediction && "
f"unzip {FORECAST_ZIP_CSV_FILE_LOCAL_PATH} -d /batch-prediction/",
move_dataset_file = GCSSynchronizeBucketsOperator(
task_id="move_dataset_to_bucket",
source_bucket=RESOURCE_DATA_BUCKET,
source_object="vertex-ai/datasets",
destination_bucket=DATA_SAMPLE_GCS_BUCKET_NAME,
destination_object="vertex-ai",
recursive=True,
)

upload_files = LocalFilesystemToGCSOperator(
task_id="upload_file_to_bucket",
src=FORECAST_CSV_FILE_LOCAL_PATH,
dst=DATA_SAMPLE_GCS_OBJECT_NAME,
bucket=DATA_SAMPLE_GCS_BUCKET_NAME,
)
create_forecast_dataset = CreateDatasetOperator(
task_id="forecast_dataset",
dataset=FORECAST_DATASET,
Expand Down Expand Up @@ -186,7 +183,8 @@

delete_auto_ml_forecasting_training_job = DeleteAutoMLTrainingJobOperator(
task_id="delete_auto_ml_forecasting_training_job",
training_pipeline_id=create_auto_ml_forecasting_training_job.output["training_id"],
training_pipeline_id="{{ task_instance.xcom_pull(task_ids='auto_ml_forecasting_task', "
"key='training_id') }}",
region=REGION,
project_id=PROJECT_ID,
trigger_rule=TriggerRule.ALL_DONE,
Expand All @@ -204,16 +202,10 @@
trigger_rule=TriggerRule.ALL_DONE,
)

clear_folder = BashOperator(
task_id="clear_folder",
bash_command="rm -r /batch-prediction/*",
)

(
# TEST SETUP
create_bucket
>> unzip_file
>> upload_files
>> move_dataset_file
>> create_forecast_dataset
>> create_auto_ml_forecasting_training_job
# TEST BODY
Expand All @@ -224,7 +216,6 @@
>> delete_auto_ml_forecasting_training_job
>> delete_forecast_dataset
>> delete_bucket
>> clear_folder
)


Expand Down
52 changes: 23 additions & 29 deletions tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_container.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,15 +26,17 @@

import os
from datetime import datetime
from pathlib import Path

from google.cloud.aiplatform import schema
from google.protobuf.json_format import ParseDict
from google.protobuf.struct_pb2 import Value

from airflow import models
from airflow.operators.bash import BashOperator
from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
from airflow.providers.google.cloud.operators.gcs import (
GCSCreateBucketOperator,
GCSDeleteBucketOperator,
GCSSynchronizeBucketsOperator,
)
from airflow.providers.google.cloud.operators.vertex_ai.custom_job import (
CreateCustomContainerTrainingJobOperator,
DeleteCustomTrainingJobOperator,
Expand All @@ -43,22 +45,19 @@
CreateDatasetOperator,
DeleteDatasetOperator,
)
from airflow.providers.google.cloud.transfers.local_to_gcs import LocalFilesystemToGCSOperator
from airflow.utils.trigger_rule import TriggerRule

ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default")
DAG_ID = "vertex_ai_custom_job_operations"
DAG_ID = "example_vertex_ai_custom_job_operations"
REGION = "us-central1"
CONTAINER_DISPLAY_NAME = f"train-housing-container-{ENV_ID}"
MODEL_DISPLAY_NAME = f"container-housing-model-{ENV_ID}"

CUSTOM_CONTAINER_GCS_BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}"
RESOURCE_DATA_BUCKET = "airflow-system-tests-resources"
CUSTOM_CONTAINER_GCS_BUCKET_NAME = f"bucket_cont_{DAG_ID}_{ENV_ID}".replace("_", "-")

DATA_SAMPLE_GCS_OBJECT_NAME = "vertex-ai/california_housing_train.csv"
CSV_FILE_LOCAL_PATH = "/custom-job-container/california_housing_train.csv"
RESOURCES_PATH = Path(__file__).parent / "resources"
CSV_ZIP_FILE_LOCAL_PATH = str(RESOURCES_PATH / "California-housing-custom-container.zip")


def TABULAR_DATASET(bucket_name):
Expand Down Expand Up @@ -97,17 +96,16 @@ def TABULAR_DATASET(bucket_name):
storage_class="REGIONAL",
location=REGION,
)
unzip_file = BashOperator(
task_id="unzip_csv_data_file",
bash_command=f"mkdir -p /custom-job-container/ && "
f"unzip {CSV_ZIP_FILE_LOCAL_PATH} -d /custom-job-container/",
)
upload_files = LocalFilesystemToGCSOperator(
task_id="upload_file_to_bucket",
src=CSV_FILE_LOCAL_PATH,
dst=DATA_SAMPLE_GCS_OBJECT_NAME,
bucket=CUSTOM_CONTAINER_GCS_BUCKET_NAME,

move_data_files = GCSSynchronizeBucketsOperator(
task_id="move_files_to_bucket",
source_bucket=RESOURCE_DATA_BUCKET,
source_object="vertex-ai/california-housing-data",
destination_bucket=CUSTOM_CONTAINER_GCS_BUCKET_NAME,
destination_object="vertex-ai",
recursive=True,
)

create_tabular_dataset = CreateDatasetOperator(
task_id="tabular_dataset",
dataset=TABULAR_DATASET(CUSTOM_CONTAINER_GCS_BUCKET_NAME),
Expand Down Expand Up @@ -141,8 +139,10 @@ def TABULAR_DATASET(bucket_name):

delete_custom_training_job = DeleteCustomTrainingJobOperator(
task_id="delete_custom_training_job",
training_pipeline_id=create_custom_container_training_job.output["training_id"],
custom_job_id=create_custom_container_training_job.output["custom_job_id"],
training_pipeline_id="{{ task_instance.xcom_pull(task_ids='custom_container_task', "
"key='training_id') }}",
custom_job_id="{{ task_instance.xcom_pull(task_ids='custom_container_task', "
"key='custom_job_id') }}",
region=REGION,
project_id=PROJECT_ID,
trigger_rule=TriggerRule.ALL_DONE,
Expand All @@ -160,24 +160,18 @@ def TABULAR_DATASET(bucket_name):
bucket_name=CUSTOM_CONTAINER_GCS_BUCKET_NAME,
trigger_rule=TriggerRule.ALL_DONE,
)
clear_folder = BashOperator(
task_id="clear_folder",
bash_command="rm -r /custom-job-container/*",
)

(
# TEST SETUP
create_bucket
>> unzip_file
>> upload_files
>> move_data_files
>> create_tabular_dataset
# TEST BODY
>> create_custom_container_training_job
# TEST TEARDOWN
>> delete_custom_training_job
>> delete_tabular_dataset
>> delete_bucket
>> clear_folder
)


Expand Down
56 changes: 29 additions & 27 deletions tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,15 +26,17 @@

import os
from datetime import datetime
from pathlib import Path

from google.cloud.aiplatform import schema
from google.protobuf.json_format import ParseDict
from google.protobuf.struct_pb2 import Value

from airflow import models
from airflow.operators.bash import BashOperator
from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
from airflow.providers.google.cloud.operators.gcs import (
GCSCreateBucketOperator,
GCSDeleteBucketOperator,
GCSSynchronizeBucketsOperator,
)
from airflow.providers.google.cloud.operators.vertex_ai.custom_job import (
CreateCustomTrainingJobOperator,
DeleteCustomTrainingJobOperator,
Expand All @@ -43,22 +45,20 @@
CreateDatasetOperator,
DeleteDatasetOperator,
)
from airflow.providers.google.cloud.transfers.local_to_gcs import LocalFilesystemToGCSOperator
from airflow.providers.google.cloud.transfers.gcs_to_local import GCSToLocalFilesystemOperator
from airflow.utils.trigger_rule import TriggerRule

ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
DAG_ID = "vertex_ai_custom_job_operations"
ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
DAG_ID = "example_vertex_ai_custom_job_operations"
PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default")
REGION = "us-central1"
CUSTOM_DISPLAY_NAME = f"train-housing-custom-{ENV_ID}"
MODEL_DISPLAY_NAME = f"custom-housing-model-{ENV_ID}"

CUSTOM_GCS_BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}"
RESOURCE_DATA_BUCKET = "airflow-system-tests-resources"
CUSTOM_GCS_BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}".replace("_", "-")

DATA_SAMPLE_GCS_OBJECT_NAME = "vertex-ai/california_housing_train.csv"
RESOURCES_PATH = Path(__file__).parent / "resources"
CSV_ZIP_FILE_LOCAL_PATH = str(RESOURCES_PATH / "California-housing-custom-job.zip")
CSV_FILE_LOCAL_PATH = "/custom-job/california_housing_train.csv"


def TABULAR_DATASET(bucket_name):
Expand All @@ -76,7 +76,7 @@ def TABULAR_DATASET(bucket_name):
MODEL_SERVING_CONTAINER_URI = "gcr.io/cloud-aiplatform/prediction/tf2-cpu.2-2:latest"
REPLICA_COUNT = 1

LOCAL_TRAINING_SCRIPT_PATH = "/custom-job/california_housing_training_script.py"
LOCAL_TRAINING_SCRIPT_PATH = "california_housing_training_script.py"


with models.DAG(
Expand All @@ -92,16 +92,23 @@ def TABULAR_DATASET(bucket_name):
storage_class="REGIONAL",
location=REGION,
)
unzip_file = BashOperator(
task_id="unzip_csv_data_file",
bash_command=f"mkdir -p /custom-job && unzip {CSV_ZIP_FILE_LOCAL_PATH} -d /custom-job/",

move_data_files = GCSSynchronizeBucketsOperator(
task_id="move_files_to_bucket",
source_bucket=RESOURCE_DATA_BUCKET,
source_object="vertex-ai/california-housing-data",
destination_bucket=CUSTOM_GCS_BUCKET_NAME,
destination_object="vertex-ai",
recursive=True,
)
upload_files = LocalFilesystemToGCSOperator(
task_id="upload_file_to_bucket",
src=CSV_FILE_LOCAL_PATH,
dst=DATA_SAMPLE_GCS_OBJECT_NAME,

download_training_script_file = GCSToLocalFilesystemOperator(
task_id="download_training_script_file",
object_name="vertex-ai/california_housing_training_script.py",
bucket=CUSTOM_GCS_BUCKET_NAME,
filename=LOCAL_TRAINING_SCRIPT_PATH,
)

create_tabular_dataset = CreateDatasetOperator(
task_id="tabular_dataset",
dataset=TABULAR_DATASET(CUSTOM_GCS_BUCKET_NAME),
Expand Down Expand Up @@ -132,8 +139,8 @@ def TABULAR_DATASET(bucket_name):
# [START how_to_cloud_vertex_ai_delete_custom_training_job_operator]
delete_custom_training_job = DeleteCustomTrainingJobOperator(
task_id="delete_custom_training_job",
training_pipeline_id=create_custom_training_job.output["training_id"],
custom_job_id=create_custom_training_job.output["custom_job_id"],
training_pipeline_id="{{ task_instance.xcom_pull(task_ids='custom_task', key='training_id') }}",
custom_job_id="{{ task_instance.xcom_pull(task_ids='custom_task', key='custom_job_id') }}",
region=REGION,
project_id=PROJECT_ID,
trigger_rule=TriggerRule.ALL_DONE,
Expand All @@ -152,24 +159,19 @@ def TABULAR_DATASET(bucket_name):
bucket_name=CUSTOM_GCS_BUCKET_NAME,
trigger_rule=TriggerRule.ALL_DONE,
)
clear_folder = BashOperator(
task_id="clear_folder",
bash_command="rm -r /custom-job/*",
)

(
# TEST SETUP
create_bucket
>> unzip_file
>> upload_files
>> move_data_files
>> download_training_script_file
>> create_tabular_dataset
# TEST BODY
>> create_custom_training_job
# TEST TEARDOWN
>> delete_custom_training_job
>> delete_tabular_dataset
>> delete_bucket
>> clear_folder
)


Expand Down

0 comments on commit 013c95b

Please sign in to comment.