You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2023/09/14 15:34:30 UTC
[airflow] branch main updated: Update Vertex AI system tests (#34364)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 013c95bfe6 Update Vertex AI system tests (#34364)
013c95bfe6 is described below
commit 013c95bfe6d2b4de44633ab1a0c326cd0969fcf3
Author: Maksim <ma...@google.com>
AuthorDate: Thu Sep 14 17:34:19 2023 +0200
Update Vertex AI system tests (#34364)
* Update Vertex AI system tests
* Update image uri for "Hyperparameter" test
---
.../google/cloud/example_dags/example_vertex_ai.py | 754 ---------------------
.../example_vertex_ai_batch_prediction_job.py | 47 +-
.../example_vertex_ai_custom_container.py | 52 +-
.../vertex_ai/example_vertex_ai_custom_job.py | 56 +-
.../example_vertex_ai_custom_job_python_package.py | 55 +-
.../cloud/vertex_ai/example_vertex_ai_dataset.py | 59 +-
.../cloud/vertex_ai/example_vertex_ai_endpoint.py | 50 +-
.../example_vertex_ai_hyperparameter_tuning_job.py | 12 +-
.../vertex_ai/example_vertex_ai_model_service.py | 57 +-
.../resources/California-housing-ai-model.zip | Bin 366644 -> 0 bytes
.../California-housing-custom-container.zip | Bin 366644 -> 0 bytes
.../resources/California-housing-custom-job.zip | Bin 366644 -> 0 bytes
.../California-housing-python-package.zip | Bin 366644 -> 0 bytes
.../google/cloud/vertex_ai/resources/__init__.py | 16 -
.../cloud/vertex_ai/resources/all-datasets.zip | Bin 84305 -> 0 bytes
.../vertex_ai/resources/forecast-dataset.csv.zip | Bin 11538 -> 0 bytes
.../vertex_ai/resources/image-dataset.csv.zip | Bin 16626 -> 0 bytes
.../vertex_ai/resources/tabular-dataset.csv.zip | Bin 9403 -> 0 bytes
.../cloud/vertex_ai/resources/text-dataset.csv.zip | Bin 41210 -> 0 bytes
.../vertex_ai/resources/video-dataset.csv.zip | Bin 5616 -> 0 bytes
20 files changed, 169 insertions(+), 989 deletions(-)
diff --git a/airflow/providers/google/cloud/example_dags/example_vertex_ai.py b/airflow/providers/google/cloud/example_dags/example_vertex_ai.py
deleted file mode 100644
index 4265ac31ce..0000000000
--- a/airflow/providers/google/cloud/example_dags/example_vertex_ai.py
+++ /dev/null
@@ -1,754 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""
-Example Airflow DAG that demonstrates operators for the Google Vertex AI service in the Google
-Cloud Platform.
-
-This DAG relies on the following OS environment variables:
-
-* GCP_VERTEX_AI_BUCKET - Google Cloud Storage bucket where the model will be saved
- after training process was finished.
-* CUSTOM_CONTAINER_URI - path to container with model.
-* PYTHON_PACKAGE_GSC_URI - path to test model in archive.
-* LOCAL_TRAINING_SCRIPT_PATH - path to local training script.
-* DATASET_ID - ID of dataset which will be used in training process.
-* MODEL_ID - ID of model which will be used in predict process.
-* MODEL_ARTIFACT_URI - The artifact_uri should be the path to a GCS directory containing saved model
- artifacts.
-"""
-from __future__ import annotations
-
-import os
-from datetime import datetime
-from uuid import uuid4
-
-from google.cloud import aiplatform
-from google.protobuf.struct_pb2 import Value
-
-from airflow import models
-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,
- GetDatasetOperator,
- 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 (
- DeleteModelOperator,
- ExportModelOperator,
- ListModelsOperator,
- UploadModelOperator,
-)
-
-# mypy ignore arg types (for templated fields)
-# type: ignore[arg-type]
-
-
-PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "an-id")
-REGION = os.environ.get("GCP_LOCATION", "us-central1")
-BUCKET = os.environ.get("GCP_VERTEX_AI_BUCKET", "vertex-ai-system-tests")
-
-STAGING_BUCKET = f"gs://{BUCKET}"
-DISPLAY_NAME = str(uuid4()) # Create random display name
-CONTAINER_URI = "gcr.io/cloud-aiplatform/training/tf-cpu.2-2:latest"
-CUSTOM_CONTAINER_URI = os.environ.get("CUSTOM_CONTAINER_URI", "path_to_container_with_model")
-MODEL_SERVING_CONTAINER_URI = "gcr.io/cloud-aiplatform/prediction/tf2-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
-
-PYTHON_PACKAGE_GCS_URI = os.environ.get("PYTHON_PACKAGE_GSC_URI", "path_to_test_model_in_arch")
-PYTHON_MODULE_NAME = "aiplatform_custom_trainer_script.task"
-
-LOCAL_TRAINING_SCRIPT_PATH = os.environ.get("LOCAL_TRAINING_SCRIPT_PATH", "path_to_training_script")
-
-TRAINING_PIPELINE_ID = "test-training-pipeline-id"
-CUSTOM_JOB_ID = "test-custom-job-id"
-
-IMAGE_DATASET = {
- "display_name": str(uuid4()),
- "metadata_schema_uri": "gs://google-cloud-aiplatform/schema/dataset/metadata/image_1.0.0.yaml",
- "metadata": Value(string_value="test-image-dataset"),
-}
-TABULAR_DATASET = {
- "display_name": str(uuid4()),
- "metadata_schema_uri": "gs://google-cloud-aiplatform/schema/dataset/metadata/tabular_1.0.0.yaml",
- "metadata": Value(string_value="test-tabular-dataset"),
-}
-TEXT_DATASET = {
- "display_name": str(uuid4()),
- "metadata_schema_uri": "gs://google-cloud-aiplatform/schema/dataset/metadata/text_1.0.0.yaml",
- "metadata": Value(string_value="test-text-dataset"),
-}
-VIDEO_DATASET = {
- "display_name": str(uuid4()),
- "metadata_schema_uri": "gs://google-cloud-aiplatform/schema/dataset/metadata/video_1.0.0.yaml",
- "metadata": Value(string_value="test-video-dataset"),
-}
-TIME_SERIES_DATASET = {
- "display_name": str(uuid4()),
- "metadata_schema_uri": "gs://google-cloud-aiplatform/schema/dataset/metadata/time_series_1.0.0.yaml",
- "metadata": Value(string_value="test-video-dataset"),
-}
-DATASET_ID = os.environ.get("DATASET_ID", "test-dataset-id")
-TEST_EXPORT_CONFIG = {"gcs_destination": {"output_uri_prefix": "gs://test-vertex-ai-bucket/exports"}}
-TEST_IMPORT_CONFIG = [
- {
- "data_item_labels": {
- "test-labels-name": "test-labels-value",
- },
- "import_schema_uri": (
- "gs://google-cloud-aiplatform/schema/dataset/ioformat/image_bounding_box_io_format_1.0.0.yaml"
- ),
- "gcs_source": {
- "uris": ["gs://ucaip-test-us-central1/dataset/salads_oid_ml_use_public_unassigned.jsonl"]
- },
- },
-]
-DATASET_TO_UPDATE = {"display_name": "test-name"}
-TEST_UPDATE_MASK = {"paths": ["displayName"]}
-
-TEST_TIME_COLUMN = "date"
-TEST_TIME_SERIES_IDENTIFIER_COLUMN = "store_name"
-TEST_TARGET_COLUMN = "sale_dollars"
-
-COLUMN_SPECS = {
- TEST_TIME_COLUMN: "timestamp",
- TEST_TARGET_COLUMN: "numeric",
- "city": "categorical",
- "zip_code": "categorical",
- "county": "categorical",
-}
-
-COLUMN_TRANSFORMATIONS = [
- {"categorical": {"column_name": "Type"}},
- {"numeric": {"column_name": "Age"}},
- {"categorical": {"column_name": "Breed1"}},
- {"categorical": {"column_name": "Color1"}},
- {"categorical": {"column_name": "Color2"}},
- {"categorical": {"column_name": "MaturitySize"}},
- {"categorical": {"column_name": "FurLength"}},
- {"categorical": {"column_name": "Vaccinated"}},
- {"categorical": {"column_name": "Sterilized"}},
- {"categorical": {"column_name": "Health"}},
- {"numeric": {"column_name": "Fee"}},
- {"numeric": {"column_name": "PhotoAmt"}},
-]
-
-MODEL_ID = os.environ.get("MODEL_ID", "test-model-id")
-MODEL_ARTIFACT_URI = os.environ.get("MODEL_ARTIFACT_URI", "path_to_folder_with_model_artifacts")
-MODEL_NAME = f"projects/{PROJECT_ID}/locations/{REGION}/models/{MODEL_ID}"
-JOB_DISPLAY_NAME = f"temp_create_batch_prediction_job_test_{uuid4()}"
-BIGQUERY_SOURCE = f"bq://{PROJECT_ID}.test_iowa_liquor_sales_forecasting_us.2021_sales_predict"
-GCS_DESTINATION_PREFIX = "gs://test-vertex-ai-bucket-us/output"
-MODEL_PARAMETERS: dict | None = {}
-
-ENDPOINT_CONF = {
- "display_name": f"endpoint_test_{uuid4()}",
-}
-DEPLOYED_MODEL = {
- # format: 'projects/{project}/locations/{location}/models/{model}'
- "model": f"projects/{PROJECT_ID}/locations/{REGION}/models/{MODEL_ID}",
- "display_name": f"temp_endpoint_test_{uuid4()}",
- "dedicated_resources": {
- "machine_spec": {
- "machine_type": "n1-standard-2",
- "accelerator_type": aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_K80,
- "accelerator_count": 1,
- },
- "min_replica_count": 1,
- "max_replica_count": 1,
- },
-}
-
-MODEL_OUTPUT_CONFIG = {
- "artifact_destination": {
- "output_uri_prefix": STAGING_BUCKET,
- },
- "export_format_id": "custom-trained",
-}
-MODEL_OBJ = {
- "display_name": f"model-{str(uuid4())}",
- "artifact_uri": MODEL_ARTIFACT_URI,
- "container_spec": {
- "image_uri": MODEL_SERVING_CONTAINER_URI,
- "command": [],
- "args": [],
- "env": [],
- "ports": [],
- "predict_route": "",
- "health_route": "",
- },
-}
-
-with models.DAG(
- "example_gcp_vertex_ai_custom_jobs",
- start_date=datetime(2021, 1, 1),
- catchup=False,
-) as custom_jobs_dag:
- # [START how_to_cloud_vertex_ai_create_custom_container_training_job_operator]
- create_custom_container_training_job = CreateCustomContainerTrainingJobOperator(
- task_id="custom_container_task",
- staging_bucket=STAGING_BUCKET,
- display_name=f"train-housing-container-{DISPLAY_NAME}",
- container_uri=CUSTOM_CONTAINER_URI,
- model_serving_container_image_uri=MODEL_SERVING_CONTAINER_URI,
- # run params
- dataset_id=DATASET_ID,
- command=["python3", "task.py"],
- model_display_name=f"container-housing-model-{DISPLAY_NAME}",
- 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=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_create_custom_container_training_job_operator]
-
- # [START how_to_cloud_vertex_ai_create_custom_python_package_training_job_operator]
- create_custom_python_package_training_job = CreateCustomPythonPackageTrainingJobOperator(
- task_id="python_package_task",
- staging_bucket=STAGING_BUCKET,
- display_name=f"train-housing-py-package-{DISPLAY_NAME}",
- python_package_gcs_uri=PYTHON_PACKAGE_GCS_URI,
- python_module_name=PYTHON_MODULE_NAME,
- container_uri=CONTAINER_URI,
- model_serving_container_image_uri=MODEL_SERVING_CONTAINER_URI,
- # run params
- dataset_id=DATASET_ID,
- model_display_name=f"py-package-housing-model-{DISPLAY_NAME}",
- 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=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_create_custom_python_package_training_job_operator]
-
- # [START how_to_cloud_vertex_ai_create_custom_training_job_operator]
- create_custom_training_job = CreateCustomTrainingJobOperator(
- task_id="custom_task",
- staging_bucket=STAGING_BUCKET,
- display_name=f"train-housing-custom-{DISPLAY_NAME}",
- script_path=LOCAL_TRAINING_SCRIPT_PATH,
- container_uri=CONTAINER_URI,
- requirements=["gcsfs==0.7.1"],
- model_serving_container_image_uri=MODEL_SERVING_CONTAINER_URI,
- # run params
- dataset_id=DATASET_ID,
- replica_count=1,
- model_display_name=f"custom-housing-model-{DISPLAY_NAME}",
- sync=False,
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_create_custom_training_job_operator]
-
- # [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=TRAINING_PIPELINE_ID,
- custom_job_id=CUSTOM_JOB_ID,
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_delete_custom_training_job_operator]
-
- # [START how_to_cloud_vertex_ai_list_custom_training_job_operator]
- list_custom_training_job = ListCustomTrainingJobOperator(
- task_id="list_custom_training_job",
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_list_custom_training_job_operator]
-
-with models.DAG(
- "example_gcp_vertex_ai_dataset",
- start_date=datetime(2021, 1, 1),
- catchup=False,
-) as dataset_dag:
- # [START how_to_cloud_vertex_ai_create_dataset_operator]
- create_image_dataset_job = CreateDatasetOperator(
- task_id="image_dataset",
- dataset=IMAGE_DATASET,
- region=REGION,
- project_id=PROJECT_ID,
- )
- create_tabular_dataset_job = CreateDatasetOperator(
- task_id="tabular_dataset",
- dataset=TABULAR_DATASET,
- region=REGION,
- project_id=PROJECT_ID,
- )
- create_text_dataset_job = CreateDatasetOperator(
- task_id="text_dataset",
- dataset=TEXT_DATASET,
- region=REGION,
- project_id=PROJECT_ID,
- )
- create_video_dataset_job = CreateDatasetOperator(
- task_id="video_dataset",
- dataset=VIDEO_DATASET,
- region=REGION,
- project_id=PROJECT_ID,
- )
- create_time_series_dataset_job = CreateDatasetOperator(
- task_id="time_series_dataset",
- dataset=TIME_SERIES_DATASET,
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_create_dataset_operator]
-
- # [START how_to_cloud_vertex_ai_delete_dataset_operator]
- delete_dataset_job = DeleteDatasetOperator(
- task_id="delete_dataset",
- dataset_id=create_text_dataset_job.output["dataset_id"],
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_delete_dataset_operator]
-
- # [START how_to_cloud_vertex_ai_get_dataset_operator]
- get_dataset = GetDatasetOperator(
- task_id="get_dataset",
- project_id=PROJECT_ID,
- region=REGION,
- dataset_id=create_tabular_dataset_job.output["dataset_id"],
- )
- # [END how_to_cloud_vertex_ai_get_dataset_operator]
-
- # [START how_to_cloud_vertex_ai_export_data_operator]
- export_data_job = ExportDataOperator(
- task_id="export_data",
- dataset_id=create_image_dataset_job.output["dataset_id"],
- region=REGION,
- project_id=PROJECT_ID,
- export_config=TEST_EXPORT_CONFIG,
- )
- # [END how_to_cloud_vertex_ai_export_data_operator]
-
- # [START how_to_cloud_vertex_ai_import_data_operator]
- import_data_job = ImportDataOperator(
- task_id="import_data",
- dataset_id=create_image_dataset_job.output["dataset_id"],
- region=REGION,
- project_id=PROJECT_ID,
- import_configs=TEST_IMPORT_CONFIG,
- )
- # [END how_to_cloud_vertex_ai_import_data_operator]
-
- # [START how_to_cloud_vertex_ai_list_dataset_operator]
- list_dataset_job = ListDatasetsOperator(
- task_id="list_dataset",
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_list_dataset_operator]
-
- # [START how_to_cloud_vertex_ai_update_dataset_operator]
- update_dataset_job = UpdateDatasetOperator(
- task_id="update_dataset",
- project_id=PROJECT_ID,
- region=REGION,
- dataset_id=create_video_dataset_job.output["dataset_id"],
- dataset=DATASET_TO_UPDATE,
- update_mask=TEST_UPDATE_MASK,
- )
- # [END how_to_cloud_vertex_ai_update_dataset_operator]
-
- create_time_series_dataset_job
- create_text_dataset_job >> delete_dataset_job
- create_tabular_dataset_job >> get_dataset
- create_image_dataset_job >> import_data_job >> export_data_job
- create_video_dataset_job >> update_dataset_job
- list_dataset_job
-
-with models.DAG(
- "example_gcp_vertex_ai_auto_ml",
- start_date=datetime(2021, 1, 1),
- catchup=False,
-) as auto_ml_dag:
- # [START how_to_cloud_vertex_ai_create_auto_ml_forecasting_training_job_operator]
- create_auto_ml_forecasting_training_job = CreateAutoMLForecastingTrainingJobOperator(
- task_id="auto_ml_forecasting_task",
- display_name=f"auto-ml-forecasting-{DISPLAY_NAME}",
- optimization_objective="minimize-rmse",
- column_specs=COLUMN_SPECS,
- # run params
- dataset_id=DATASET_ID,
- target_column=TEST_TARGET_COLUMN,
- time_column=TEST_TIME_COLUMN,
- time_series_identifier_column=TEST_TIME_SERIES_IDENTIFIER_COLUMN,
- available_at_forecast_columns=[TEST_TIME_COLUMN],
- unavailable_at_forecast_columns=[TEST_TARGET_COLUMN],
- time_series_attribute_columns=["city", "zip_code", "county"],
- forecast_horizon=30,
- context_window=30,
- data_granularity_unit="day",
- data_granularity_count=1,
- weight_column=None,
- budget_milli_node_hours=1000,
- model_display_name=f"auto-ml-forecasting-model-{DISPLAY_NAME}",
- predefined_split_column_name=None,
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_create_auto_ml_forecasting_training_job_operator]
-
- # [START how_to_cloud_vertex_ai_create_auto_ml_image_training_job_operator]
- create_auto_ml_image_training_job = CreateAutoMLImageTrainingJobOperator(
- task_id="auto_ml_image_task",
- display_name=f"auto-ml-image-{DISPLAY_NAME}",
- dataset_id=DATASET_ID,
- prediction_type="classification",
- multi_label=False,
- model_type="CLOUD",
- training_fraction_split=0.6,
- validation_fraction_split=0.2,
- test_fraction_split=0.2,
- budget_milli_node_hours=8000,
- model_display_name=f"auto-ml-image-model-{DISPLAY_NAME}",
- disable_early_stopping=False,
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_create_auto_ml_image_training_job_operator]
-
- # [START how_to_cloud_vertex_ai_create_auto_ml_tabular_training_job_operator]
- create_auto_ml_tabular_training_job = CreateAutoMLTabularTrainingJobOperator(
- task_id="auto_ml_tabular_task",
- display_name=f"auto-ml-tabular-{DISPLAY_NAME}",
- optimization_prediction_type="classification",
- column_transformations=COLUMN_TRANSFORMATIONS,
- dataset_id=DATASET_ID,
- target_column="Adopted",
- training_fraction_split=0.8,
- validation_fraction_split=0.1,
- test_fraction_split=0.1,
- model_display_name="adopted-prediction-model",
- disable_early_stopping=False,
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_create_auto_ml_tabular_training_job_operator]
-
- # [START how_to_cloud_vertex_ai_create_auto_ml_text_training_job_operator]
- create_auto_ml_text_training_job = CreateAutoMLTextTrainingJobOperator(
- task_id="auto_ml_text_task",
- display_name=f"auto-ml-text-{DISPLAY_NAME}",
- prediction_type="classification",
- multi_label=False,
- dataset_id=DATASET_ID,
- model_display_name=f"auto-ml-text-model-{DISPLAY_NAME}",
- training_fraction_split=0.7,
- validation_fraction_split=0.2,
- test_fraction_split=0.1,
- sync=True,
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_create_auto_ml_text_training_job_operator]
-
- # [START how_to_cloud_vertex_ai_create_auto_ml_video_training_job_operator]
- create_auto_ml_video_training_job = CreateAutoMLVideoTrainingJobOperator(
- task_id="auto_ml_video_task",
- display_name=f"auto-ml-video-{DISPLAY_NAME}",
- prediction_type="classification",
- model_type="CLOUD",
- dataset_id=DATASET_ID,
- model_display_name=f"auto-ml-video-model-{DISPLAY_NAME}",
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_create_auto_ml_video_training_job_operator]
-
- # [START how_to_cloud_vertex_ai_delete_auto_ml_training_job_operator]
- delete_auto_ml_training_job = DeleteAutoMLTrainingJobOperator(
- task_id="delete_auto_ml_training_job",
- training_pipeline_id=TRAINING_PIPELINE_ID,
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_delete_auto_ml_training_job_operator]
-
- # [START how_to_cloud_vertex_ai_list_auto_ml_training_job_operator]
- list_auto_ml_training_job = ListAutoMLTrainingJobOperator(
- task_id="list_auto_ml_training_job",
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_list_auto_ml_training_job_operator]
-
-with models.DAG(
- "example_gcp_vertex_ai_batch_prediction_job",
- start_date=datetime(2021, 1, 1),
- catchup=False,
-) as batch_prediction_job_dag:
- # [START how_to_cloud_vertex_ai_create_batch_prediction_job_operator]
- create_batch_prediction_job = CreateBatchPredictionJobOperator(
- task_id="create_batch_prediction_job",
- job_display_name=JOB_DISPLAY_NAME,
- model_name=MODEL_NAME,
- predictions_format="csv",
- bigquery_source=BIGQUERY_SOURCE,
- gcs_destination_prefix=GCS_DESTINATION_PREFIX,
- model_parameters=MODEL_PARAMETERS,
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_create_batch_prediction_job_operator]
-
- # [START how_to_cloud_vertex_ai_list_batch_prediction_job_operator]
- list_batch_prediction_job = ListBatchPredictionJobsOperator(
- task_id="list_batch_prediction_jobs",
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_list_batch_prediction_job_operator]
-
- # [START how_to_cloud_vertex_ai_delete_batch_prediction_job_operator]
- delete_batch_prediction_job = DeleteBatchPredictionJobOperator(
- task_id="delete_batch_prediction_job",
- batch_prediction_job_id=create_batch_prediction_job.output["batch_prediction_job_id"],
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_delete_batch_prediction_job_operator]
-
- create_batch_prediction_job >> delete_batch_prediction_job
- list_batch_prediction_job
-
-with models.DAG(
- "example_gcp_vertex_ai_endpoint",
- start_date=datetime(2021, 1, 1),
- catchup=False,
-) as endpoint_dag:
- # [START how_to_cloud_vertex_ai_create_endpoint_operator]
- create_endpoint = CreateEndpointOperator(
- task_id="create_endpoint",
- endpoint=ENDPOINT_CONF,
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_create_endpoint_operator]
-
- # [START how_to_cloud_vertex_ai_delete_endpoint_operator]
- delete_endpoint = DeleteEndpointOperator(
- task_id="delete_endpoint",
- endpoint_id=create_endpoint.output["endpoint_id"],
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_delete_endpoint_operator]
-
- # [START how_to_cloud_vertex_ai_list_endpoints_operator]
- list_endpoints = ListEndpointsOperator(
- task_id="list_endpoints",
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_list_endpoints_operator]
-
- # [START how_to_cloud_vertex_ai_deploy_model_operator]
- deploy_model = DeployModelOperator(
- task_id="deploy_model",
- endpoint_id=create_endpoint.output["endpoint_id"],
- deployed_model=DEPLOYED_MODEL,
- traffic_split={"0": 100},
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_deploy_model_operator]
-
- # [START how_to_cloud_vertex_ai_undeploy_model_operator]
- undeploy_model = UndeployModelOperator(
- task_id="undeploy_model",
- endpoint_id=create_endpoint.output["endpoint_id"],
- deployed_model_id=deploy_model.output["deployed_model_id"],
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_undeploy_model_operator]
-
- create_endpoint >> deploy_model >> undeploy_model >> delete_endpoint
- list_endpoints
-
-with models.DAG(
- "example_gcp_vertex_ai_hyperparameter_tuning_job",
- start_date=datetime(2021, 1, 1),
- catchup=False,
-) as hyperparameter_tuning_job_dag:
- # [START how_to_cloud_vertex_ai_create_hyperparameter_tuning_job_operator]
- create_hyperparameter_tuning_job = CreateHyperparameterTuningJobOperator(
- task_id="create_hyperparameter_tuning_job",
- staging_bucket=STAGING_BUCKET,
- display_name=f"horses-humans-hyptertune-{DISPLAY_NAME}",
- worker_pool_specs=[
- {
- "machine_spec": {
- "machine_type": MACHINE_TYPE,
- "accelerator_type": ACCELERATOR_TYPE,
- "accelerator_count": ACCELERATOR_COUNT,
- },
- "replica_count": REPLICA_COUNT,
- "container_spec": {
- "image_uri": f"gcr.io/{PROJECT_ID}/horse-human:hypertune",
- },
- }
- ],
- sync=False,
- region=REGION,
- project_id=PROJECT_ID,
- parameter_spec={
- "learning_rate": aiplatform.hyperparameter_tuning.DoubleParameterSpec(
- min=0.01, max=1, scale="log"
- ),
- "momentum": aiplatform.hyperparameter_tuning.DoubleParameterSpec(min=0, max=1, scale="linear"),
- "num_neurons": aiplatform.hyperparameter_tuning.DiscreteParameterSpec(
- values=[64, 128, 512], scale="linear"
- ),
- },
- metric_spec={
- "accuracy": "maximize",
- },
- max_trial_count=15,
- parallel_trial_count=3,
- )
- # [END how_to_cloud_vertex_ai_create_hyperparameter_tuning_job_operator]
-
- # [START how_to_cloud_vertex_ai_get_hyperparameter_tuning_job_operator]
- get_hyperparameter_tuning_job = GetHyperparameterTuningJobOperator(
- task_id="get_hyperparameter_tuning_job",
- project_id=PROJECT_ID,
- region=REGION,
- hyperparameter_tuning_job_id=create_hyperparameter_tuning_job.output["hyperparameter_tuning_job_id"],
- )
- # [END how_to_cloud_vertex_ai_get_hyperparameter_tuning_job_operator]
-
- # [START how_to_cloud_vertex_ai_delete_hyperparameter_tuning_job_operator]
- delete_hyperparameter_tuning_job = DeleteHyperparameterTuningJobOperator(
- task_id="delete_hyperparameter_tuning_job",
- project_id=PROJECT_ID,
- region=REGION,
- hyperparameter_tuning_job_id=create_hyperparameter_tuning_job.output["hyperparameter_tuning_job_id"],
- )
- # [END how_to_cloud_vertex_ai_delete_hyperparameter_tuning_job_operator]
-
- # [START how_to_cloud_vertex_ai_list_hyperparameter_tuning_job_operator]
- list_hyperparameter_tuning_job = ListHyperparameterTuningJobOperator(
- task_id="list_hyperparameter_tuning_job",
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_list_hyperparameter_tuning_job_operator]
-
- create_hyperparameter_tuning_job >> get_hyperparameter_tuning_job >> delete_hyperparameter_tuning_job
- list_hyperparameter_tuning_job
-
-with models.DAG(
- "example_gcp_vertex_ai_model_service",
- start_date=datetime(2021, 1, 1),
- catchup=False,
-) as model_service_dag:
- # [START how_to_cloud_vertex_ai_upload_model_operator]
- upload_model = UploadModelOperator(
- task_id="upload_model",
- region=REGION,
- project_id=PROJECT_ID,
- model=MODEL_OBJ,
- )
- # [END how_to_cloud_vertex_ai_upload_model_operator]
-
- # [START how_to_cloud_vertex_ai_export_model_operator]
- export_model = ExportModelOperator(
- task_id="export_model",
- project_id=PROJECT_ID,
- region=REGION,
- model_id=upload_model.output["model_id"],
- output_config=MODEL_OUTPUT_CONFIG,
- )
- # [END how_to_cloud_vertex_ai_export_model_operator]
-
- # [START how_to_cloud_vertex_ai_delete_model_operator]
- delete_model = DeleteModelOperator(
- task_id="delete_model",
- project_id=PROJECT_ID,
- region=REGION,
- model_id=upload_model.output["model_id"],
- )
- # [END how_to_cloud_vertex_ai_delete_model_operator]
-
- # [START how_to_cloud_vertex_ai_list_models_operator]
- list_models = ListModelsOperator(
- task_id="list_models",
- region=REGION,
- project_id=PROJECT_ID,
- )
- # [END how_to_cloud_vertex_ai_list_models_operator]
-
- upload_model >> export_model >> delete_model
- list_models
diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_batch_prediction_job.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_batch_prediction_job.py
index 41f6ff7854..5d696f3457 100644
--- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_batch_prediction_job.py
+++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_batch_prediction_job.py
@@ -26,15 +26,17 @@ from __future__ import annotations
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,
@@ -48,22 +50,20 @@ from airflow.providers.google.cloud.operators.vertex_ai.dataset import (
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}",
@@ -109,18 +109,15 @@ with models.DAG(
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,
@@ -186,7 +183,8 @@ with models.DAG(
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,
@@ -204,16 +202,10 @@ with models.DAG(
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
@@ -224,7 +216,6 @@ with models.DAG(
>> delete_auto_ml_forecasting_training_job
>> delete_forecast_dataset
>> delete_bucket
- >> clear_folder
)
diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_container.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_container.py
index 94e2b9fffc..612e3e4f9a 100644
--- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_container.py
+++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_container.py
@@ -26,15 +26,17 @@ from __future__ import annotations
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,
@@ -43,22 +45,19 @@ from airflow.providers.google.cloud.operators.vertex_ai.dataset import (
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):
@@ -97,17 +96,16 @@ with models.DAG(
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),
@@ -141,8 +139,10 @@ with models.DAG(
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,
@@ -160,16 +160,11 @@ with models.DAG(
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
@@ -177,7 +172,6 @@ with models.DAG(
>> delete_custom_training_job
>> delete_tabular_dataset
>> delete_bucket
- >> clear_folder
)
diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job.py
index 9617d67387..ad611eb67a 100644
--- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job.py
+++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job.py
@@ -26,15 +26,17 @@ from __future__ import annotations
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,
@@ -43,22 +45,20 @@ from airflow.providers.google.cloud.operators.vertex_ai.dataset import (
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):
@@ -76,7 +76,7 @@ CONTAINER_URI = "gcr.io/cloud-aiplatform/training/tf-cpu.2-2:latest"
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(
@@ -92,16 +92,23 @@ with models.DAG(
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),
@@ -132,8 +139,8 @@ with models.DAG(
# [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,
@@ -152,16 +159,12 @@ with models.DAG(
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
@@ -169,7 +172,6 @@ with models.DAG(
>> delete_custom_training_job
>> delete_tabular_dataset
>> delete_bucket
- >> clear_folder
)
diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job_python_package.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job_python_package.py
index ff42dfbc42..99a1a734eb 100644
--- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job_python_package.py
+++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_custom_job_python_package.py
@@ -26,15 +26,17 @@ from __future__ import annotations
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 (
CreateCustomPythonPackageTrainingJobOperator,
DeleteCustomTrainingJobOperator,
@@ -43,27 +45,19 @@ from airflow.providers.google.cloud.operators.vertex_ai.dataset import (
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"
PACKAGE_DISPLAY_NAME = f"train-housing-py-package-{ENV_ID}"
MODEL_DISPLAY_NAME = f"py-package-housing-model-{ENV_ID}"
-CUSTOM_PYTHON_GCS_BUCKET_NAME = f"bucket_python_{DAG_ID}_{ENV_ID}"
+RESOURCE_DATA_BUCKET = "airflow-system-tests-resources"
+CUSTOM_PYTHON_GCS_BUCKET_NAME = f"bucket_python_{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-python-package.zip")
-CSV_FILE_LOCAL_PATH = "/custom-job-python/california_housing_train.csv"
-TAR_FILE_LOCAL_PATH = "/custom-job-python/custom_trainer_script-0.1.tar"
-FILES_TO_UPLOAD = [
- CSV_FILE_LOCAL_PATH,
- TAR_FILE_LOCAL_PATH,
-]
def TABULAR_DATASET(bucket_name):
@@ -104,16 +98,16 @@ with models.DAG(
storage_class="REGIONAL",
location=REGION,
)
- unzip_file = BashOperator(
- task_id="unzip_csv_data_file",
- bash_command=f"mkdir -p /custom-job-python && unzip {CSV_ZIP_FILE_LOCAL_PATH} -d /custom-job-python/",
- )
- upload_files = LocalFilesystemToGCSOperator(
- task_id="upload_file_to_bucket",
- src=FILES_TO_UPLOAD,
- dst="vertex-ai/",
- bucket=CUSTOM_PYTHON_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_PYTHON_GCS_BUCKET_NAME,
+ destination_object="vertex-ai",
+ recursive=True,
)
+
create_tabular_dataset = CreateDatasetOperator(
task_id="tabular_dataset",
dataset=TABULAR_DATASET(CUSTOM_PYTHON_GCS_BUCKET_NAME),
@@ -148,8 +142,9 @@ with models.DAG(
delete_custom_training_job = DeleteCustomTrainingJobOperator(
task_id="delete_custom_training_job",
- training_pipeline_id=create_custom_python_package_training_job.output["training_id"],
- custom_job_id=create_custom_python_package_training_job.output["custom_job_id"],
+ training_pipeline_id="{{ task_instance.xcom_pull(task_ids='python_package_task', "
+ "key='training_id') }}",
+ custom_job_id="{{ task_instance.xcom_pull(task_ids='python_package_task', key='custom_job_id') }}",
region=REGION,
project_id=PROJECT_ID,
trigger_rule=TriggerRule.ALL_DONE,
@@ -167,16 +162,11 @@ with models.DAG(
bucket_name=CUSTOM_PYTHON_GCS_BUCKET_NAME,
trigger_rule=TriggerRule.ALL_DONE,
)
- clear_folder = BashOperator(
- task_id="clear_folder",
- bash_command="rm -r /custom-job-python/*",
- )
(
# TEST SETUP
create_bucket
- >> unzip_file
- >> upload_files
+ >> move_data_files
>> create_tabular_dataset
# TEST BODY
>> create_custom_python_package_training_job
@@ -184,7 +174,6 @@ with models.DAG(
>> delete_custom_training_job
>> delete_tabular_dataset
>> delete_bucket
- >> clear_folder
)
diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_dataset.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_dataset.py
index 71992673a8..153788835b 100644
--- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_dataset.py
+++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_dataset.py
@@ -26,15 +26,17 @@ from __future__ import annotations
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.dataset import (
CreateDatasetOperator,
DeleteDatasetOperator,
@@ -44,26 +46,15 @@ from airflow.providers.google.cloud.operators.vertex_ai.dataset import (
ListDatasetsOperator,
UpdateDatasetOperator,
)
-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_dataset_operations"
+DAG_ID = "example_vertex_ai_dataset_operations"
REGION = "us-central1"
-DATA_SAMPLE_GCS_BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}"
-
-RESOURCES_PATH = Path(__file__).parent / "resources"
-ALL_DATASETS_ZIP_CSV_FILE_LOCAL_PATH = str(RESOURCES_PATH / "all-datasets.zip")
-
-CSV_FILES_LOCAL_PATH = [
- "/all-datasets/forecast-dataset.csv",
- "/all-datasets/image-dataset.csv",
- "/all-datasets/tabular-dataset.csv",
- "/all-datasets/text-dataset.csv",
- "/all-datasets/video-dataset.csv",
-]
+RESOURCE_DATA_BUCKET = "airflow-system-tests-resources"
+DATA_SAMPLE_GCS_BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}".replace("_", "-")
TIME_SERIES_DATASET = {
"display_name": f"time-series-dataset-{ENV_ID}",
@@ -113,11 +104,7 @@ TEST_IMPORT_CONFIG = [
"import_schema_uri": (
"gs://google-cloud-aiplatform/schema/dataset/ioformat/image_bounding_box_io_format_1.0.0.yaml"
),
- "gcs_source": {
- "uris": [
- "gs://system-tests-resources/vertex-ai/dataset/salads_oid_ml_use_public_unassigned.jsonl"
- ]
- },
+ "gcs_source": {"uris": ["gs://cloud-samples-data/vision/salads.csv"]},
},
]
DATASET_TO_UPDATE = {"display_name": "test-name"}
@@ -137,15 +124,14 @@ with models.DAG(
storage_class="REGIONAL",
location=REGION,
)
- unzip_file = BashOperator(
- task_id="unzip_csv_data_file",
- bash_command=f"unzip {ALL_DATASETS_ZIP_CSV_FILE_LOCAL_PATH} -d /all-datasets/",
- )
- upload_files = LocalFilesystemToGCSOperator(
- task_id="upload_file_to_bucket",
- src=CSV_FILES_LOCAL_PATH,
- dst="vertex-ai/",
- bucket=DATA_SAMPLE_GCS_BUCKET_NAME,
+
+ move_datasets_files = GCSSynchronizeBucketsOperator(
+ task_id="move_datasets_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,
)
# [START how_to_cloud_vertex_ai_create_dataset_operator]
@@ -276,16 +262,10 @@ with models.DAG(
trigger_rule=TriggerRule.ALL_DONE,
)
- clear_folder = BashOperator(
- task_id="clear_folder",
- bash_command="rm -r /all-datasets/*",
- )
-
(
# TEST SETUP
create_bucket
- >> unzip_file
- >> upload_files
+ >> move_datasets_files
# TEST BODY
>> [
create_time_series_dataset_job >> delete_time_series_dataset_job,
@@ -297,7 +277,6 @@ with models.DAG(
]
# TEST TEARDOWN
>> delete_bucket
- >> clear_folder
)
diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_endpoint.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_endpoint.py
index 5d40393fd0..8be2357dd8 100644
--- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_endpoint.py
+++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_endpoint.py
@@ -26,14 +26,16 @@ from __future__ import annotations
import os
from datetime import datetime
-from pathlib import Path
from google.cloud.aiplatform import schema
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 (
CreateAutoMLImageTrainingJobOperator,
DeleteAutoMLTrainingJobOperator,
@@ -50,21 +52,17 @@ from airflow.providers.google.cloud.operators.vertex_ai.endpoint_service import
ListEndpointsOperator,
UndeployModelOperator,
)
-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_endpoint_service_operations"
+DAG_ID = "example_vertex_ai_endpoint_service_operations"
REGION = "us-central1"
IMAGE_DISPLAY_NAME = f"auto-ml-image-{ENV_ID}"
MODEL_DISPLAY_NAME = f"auto-ml-image-model-{ENV_ID}"
-DATA_SAMPLE_GCS_BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}"
-DATA_SAMPLE_GCS_OBJECT_NAME = "vertex-ai/image-dataset.csv"
-RESOURCES_PATH = Path(__file__).parent / "resources"
-IMAGE_ZIP_CSV_FILE_LOCAL_PATH = str(RESOURCES_PATH / "image-dataset.csv.zip")
-IMAGE_CSV_FILE_LOCAL_PATH = "/endpoint/image-dataset.csv"
+RESOURCE_DATA_BUCKET = "airflow-system-tests-resources"
+DATA_SAMPLE_GCS_BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}".replace("_", "-")
IMAGE_DATASET = {
"display_name": f"image-dataset-{ENV_ID}",
@@ -97,16 +95,16 @@ with models.DAG(
storage_class="REGIONAL",
location=REGION,
)
- unzip_file = BashOperator(
- task_id="unzip_csv_data_file",
- bash_command=f"unzip {IMAGE_ZIP_CSV_FILE_LOCAL_PATH} -d /endpoint/",
- )
- upload_files = LocalFilesystemToGCSOperator(
- task_id="upload_file_to_bucket",
- src=IMAGE_CSV_FILE_LOCAL_PATH,
- dst=DATA_SAMPLE_GCS_OBJECT_NAME,
- bucket=DATA_SAMPLE_GCS_BUCKET_NAME,
+
+ 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,
)
+
create_image_dataset = CreateDatasetOperator(
task_id="image_dataset",
dataset=IMAGE_DATASET,
@@ -198,7 +196,8 @@ with models.DAG(
delete_auto_ml_image_training_job = DeleteAutoMLTrainingJobOperator(
task_id="delete_auto_ml_training_job",
- training_pipeline_id=create_auto_ml_image_training_job.output["training_id"],
+ training_pipeline_id="{{ task_instance.xcom_pull(task_ids='auto_ml_image_task', "
+ "key='training_id') }}",
region=REGION,
project_id=PROJECT_ID,
trigger_rule=TriggerRule.ALL_DONE,
@@ -216,19 +215,13 @@ with models.DAG(
bucket_name=DATA_SAMPLE_GCS_BUCKET_NAME,
trigger_rule=TriggerRule.ALL_DONE,
)
- clear_folder = BashOperator(
- task_id="clear_folder",
- bash_command="rm -r /endpoint/*",
- )
(
# TEST SETUP
[
- create_bucket,
+ create_bucket >> move_dataset_file,
create_image_dataset,
]
- >> unzip_file
- >> upload_files
>> import_image_dataset
>> create_auto_ml_image_training_job
# TEST BODY
@@ -241,7 +234,6 @@ with models.DAG(
>> delete_auto_ml_image_training_job
>> delete_image_dataset
>> delete_bucket
- >> clear_folder
)
diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_hyperparameter_tuning_job.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_hyperparameter_tuning_job.py
index 1976a026ec..6239f263ee 100644
--- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_hyperparameter_tuning_job.py
+++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_hyperparameter_tuning_job.py
@@ -39,9 +39,9 @@ from airflow.providers.google.cloud.operators.vertex_ai.hyperparameter_tuning_jo
)
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_hyperparameter_tuning_job_operations"
+DAG_ID = "example_vertex_ai_hyperparameter_tuning_job_operations"
REGION = "us-central1"
DISPLAY_NAME = f"hyperparameter-tuning-job-{ENV_ID}"
@@ -60,7 +60,7 @@ WORKER_POOL_SPECS = [
},
"replica_count": REPLICA_COUNT,
"container_spec": {
- "image_uri": f"gcr.io/{PROJECT_ID}/horse-human:hypertune",
+ "image_uri": "us-docker.pkg.dev/composer-256318/horse-human/horse-human-image:latest",
},
}
]
@@ -111,7 +111,8 @@ with models.DAG(
task_id="get_hyperparameter_tuning_job",
project_id=PROJECT_ID,
region=REGION,
- hyperparameter_tuning_job_id=create_hyperparameter_tuning_job.output["hyperparameter_tuning_job_id"],
+ hyperparameter_tuning_job_id="{{ task_instance.xcom_pull("
+ "task_ids='create_hyperparameter_tuning_job', key='hyperparameter_tuning_job_id') }}",
)
# [END how_to_cloud_vertex_ai_get_hyperparameter_tuning_job_operator]
@@ -120,7 +121,8 @@ with models.DAG(
task_id="delete_hyperparameter_tuning_job",
project_id=PROJECT_ID,
region=REGION,
- hyperparameter_tuning_job_id=create_hyperparameter_tuning_job.output["hyperparameter_tuning_job_id"],
+ hyperparameter_tuning_job_id="{{ task_instance.xcom_pull("
+ "task_ids='create_hyperparameter_tuning_job', key='hyperparameter_tuning_job_id') }}",
trigger_rule=TriggerRule.ALL_DONE,
)
# [END how_to_cloud_vertex_ai_delete_hyperparameter_tuning_job_operator]
diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py
index fc690d36b7..98004acac2 100644
--- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py
+++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_model_service.py
@@ -26,15 +26,17 @@ from __future__ import annotations
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,
@@ -49,23 +51,21 @@ from airflow.providers.google.cloud.operators.vertex_ai.model_service import (
ListModelsOperator,
UploadModelOperator,
)
-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")
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", "default")
PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default")
-DAG_ID = "vertex_ai_model_service_operations"
+DAG_ID = "example_vertex_ai_model_service_operations"
REGION = "us-central1"
TRAIN_DISPLAY_NAME = f"train-housing-custom-{ENV_ID}"
MODEL_DISPLAY_NAME = f"custom-housing-model-{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("_", "-")
STAGING_BUCKET = f"gs://{DATA_SAMPLE_GCS_BUCKET_NAME}"
DATA_SAMPLE_GCS_OBJECT_NAME = "vertex-ai/california_housing_train.csv"
-CSV_FILE_LOCAL_PATH = "/model_service/california_housing_train.csv"
-RESOURCES_PATH = Path(__file__).parent / "resources"
-CSV_ZIP_FILE_LOCAL_PATH = str(RESOURCES_PATH / "California-housing-ai-model.zip")
TABULAR_DATASET = {
"display_name": f"tabular-dataset-{ENV_ID}",
@@ -82,7 +82,7 @@ TABULAR_DATASET = {
CONTAINER_URI = "gcr.io/cloud-aiplatform/training/tf-cpu.2-2:latest"
-LOCAL_TRAINING_SCRIPT_PATH = "/model_service/california_housing_training_script.py"
+LOCAL_TRAINING_SCRIPT_PATH = "california_housing_training_script.py"
MODEL_OUTPUT_CONFIG = {
"artifact_destination": {
@@ -120,16 +120,23 @@ with models.DAG(
storage_class="REGIONAL",
location=REGION,
)
- unzip_file = BashOperator(
- task_id="unzip_csv_data_file",
- bash_command=f"mkdir -p /model_service && unzip {CSV_ZIP_FILE_LOCAL_PATH} -d /model_service/",
+
+ move_data_files = GCSSynchronizeBucketsOperator(
+ task_id="move_files_to_bucket",
+ source_bucket=RESOURCE_DATA_BUCKET,
+ source_object="vertex-ai/california-housing-data",
+ destination_bucket=DATA_SAMPLE_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=DATA_SAMPLE_GCS_BUCKET_NAME,
+ filename=LOCAL_TRAINING_SCRIPT_PATH,
)
+
create_tabular_dataset = CreateDatasetOperator(
task_id="tabular_dataset",
dataset=TABULAR_DATASET,
@@ -194,8 +201,8 @@ with models.DAG(
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,
@@ -215,16 +222,11 @@ with models.DAG(
trigger_rule=TriggerRule.ALL_DONE,
)
- clear_folder = BashOperator(
- task_id="clear_folder",
- bash_command="rm -r /model_service/*",
- )
-
(
# TEST SETUP
create_bucket
- >> unzip_file
- >> upload_files
+ >> move_data_files
+ >> download_training_script_file
>> create_tabular_dataset
>> create_custom_training_job
# TEST BODY
@@ -236,7 +238,6 @@ with models.DAG(
>> delete_custom_training_job
>> delete_tabular_dataset
>> delete_bucket
- >> clear_folder
)
diff --git a/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-ai-model.zip b/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-ai-model.zip
deleted file mode 100644
index 1ac6fc83a4..0000000000
Binary files a/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-ai-model.zip and /dev/null differ
diff --git a/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-custom-container.zip b/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-custom-container.zip
deleted file mode 100644
index 1ac6fc83a4..0000000000
Binary files a/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-custom-container.zip and /dev/null differ
diff --git a/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-custom-job.zip b/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-custom-job.zip
deleted file mode 100644
index 1ac6fc83a4..0000000000
Binary files a/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-custom-job.zip and /dev/null differ
diff --git a/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-python-package.zip b/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-python-package.zip
deleted file mode 100644
index 1ac6fc83a4..0000000000
Binary files a/tests/system/providers/google/cloud/vertex_ai/resources/California-housing-python-package.zip and /dev/null differ
diff --git a/tests/system/providers/google/cloud/vertex_ai/resources/__init__.py b/tests/system/providers/google/cloud/vertex_ai/resources/__init__.py
deleted file mode 100644
index 13a83393a9..0000000000
--- a/tests/system/providers/google/cloud/vertex_ai/resources/__init__.py
+++ /dev/null
@@ -1,16 +0,0 @@
-# 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.
diff --git a/tests/system/providers/google/cloud/vertex_ai/resources/all-datasets.zip b/tests/system/providers/google/cloud/vertex_ai/resources/all-datasets.zip
deleted file mode 100644
index 00c7a27029..0000000000
Binary files a/tests/system/providers/google/cloud/vertex_ai/resources/all-datasets.zip and /dev/null differ
diff --git a/tests/system/providers/google/cloud/vertex_ai/resources/forecast-dataset.csv.zip b/tests/system/providers/google/cloud/vertex_ai/resources/forecast-dataset.csv.zip
deleted file mode 100644
index fbb59761a9..0000000000
Binary files a/tests/system/providers/google/cloud/vertex_ai/resources/forecast-dataset.csv.zip and /dev/null differ
diff --git a/tests/system/providers/google/cloud/vertex_ai/resources/image-dataset.csv.zip b/tests/system/providers/google/cloud/vertex_ai/resources/image-dataset.csv.zip
deleted file mode 100644
index 5dc45deeb7..0000000000
Binary files a/tests/system/providers/google/cloud/vertex_ai/resources/image-dataset.csv.zip and /dev/null differ
diff --git a/tests/system/providers/google/cloud/vertex_ai/resources/tabular-dataset.csv.zip b/tests/system/providers/google/cloud/vertex_ai/resources/tabular-dataset.csv.zip
deleted file mode 100644
index 33f50119ab..0000000000
Binary files a/tests/system/providers/google/cloud/vertex_ai/resources/tabular-dataset.csv.zip and /dev/null differ
diff --git a/tests/system/providers/google/cloud/vertex_ai/resources/text-dataset.csv.zip b/tests/system/providers/google/cloud/vertex_ai/resources/text-dataset.csv.zip
deleted file mode 100644
index 2a38cbb6e7..0000000000
Binary files a/tests/system/providers/google/cloud/vertex_ai/resources/text-dataset.csv.zip and /dev/null differ
diff --git a/tests/system/providers/google/cloud/vertex_ai/resources/video-dataset.csv.zip b/tests/system/providers/google/cloud/vertex_ai/resources/video-dataset.csv.zip
deleted file mode 100644
index 658ce398d9..0000000000
Binary files a/tests/system/providers/google/cloud/vertex_ai/resources/video-dataset.csv.zip and /dev/null differ