From 9d68492f875464f505afef2ecd81a28d8e4922b8 Mon Sep 17 00:00:00 2001 From: VladaZakharova <80038284+VladaZakharova@users.noreply.github.com> Date: Thu, 27 Jul 2023 14:17:16 +0200 Subject: [PATCH] misc: update MLEngine system tests (#32881) --- .../google/cloud/operators/mlengine.py | 3 +++ .../cloud/ml_engine/example_mlengine.py | 18 ++++++++--------- .../cloud/ml_engine/example_mlengine_async.py | 20 +++++++++---------- 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/airflow/providers/google/cloud/operators/mlengine.py b/airflow/providers/google/cloud/operators/mlengine.py index 554167662e17e..b391041372273 100644 --- a/airflow/providers/google/cloud/operators/mlengine.py +++ b/airflow/providers/google/cloud/operators/mlengine.py @@ -1002,6 +1002,9 @@ class MLEngineStartTrainingJobOperator(GoogleCloudBaseOperator): For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:MLEngineStartTrainingJobOperator` + For more information about used parameters, check: + https://cloud.google.com/sdk/gcloud/reference/ml-engine/jobs/submit/training + :param job_id: A unique templated id for the submitted Google MLEngine training job. (templated) :param region: The Google Compute Engine region to run the MLEngine training diff --git a/tests/system/providers/google/cloud/ml_engine/example_mlengine.py b/tests/system/providers/google/cloud/ml_engine/example_mlengine.py index 6edd147bbf3a0..51bf14fb5e88e 100644 --- a/tests/system/providers/google/cloud/ml_engine/example_mlengine.py +++ b/tests/system/providers/google/cloud/ml_engine/example_mlengine.py @@ -49,14 +49,14 @@ DAG_ID = "example_gcp_mlengine" PREDICT_FILE_NAME = "predict.json" -MODEL_NAME = f"example_ml_model_{ENV_ID}" -BUCKET_NAME = f"example_ml_bucket_{ENV_ID}" +MODEL_NAME = f"model_{DAG_ID}_{ENV_ID}".replace("_", "-") +BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}".replace("_", "-") BUCKET_PATH = f"gs://{BUCKET_NAME}" JOB_DIR = f"{BUCKET_PATH}/job-dir" SAVED_MODEL_PATH = f"{JOB_DIR}/" PREDICTION_INPUT = f"{BUCKET_PATH}/{PREDICT_FILE_NAME}" PREDICTION_OUTPUT = f"{BUCKET_PATH}/prediction_output/" -TRAINER_URI = "gs://system-tests-resources/example_gcp_mlengine/trainer-0.2.tar.gz" +TRAINER_URI = "gs://airflow-system-tests-resources/ml-engine/trainer-0.2.tar.gz" TRAINER_PY_MODULE = "trainer.task" SUMMARY_TMP = f"{BUCKET_PATH}/tmp/" SUMMARY_STAGING = f"{BUCKET_PATH}/staging/" @@ -104,8 +104,8 @@ def write_predict_file(path_to_file: str): project_id=PROJECT_ID, region="us-central1", job_id="training-job-{{ ts_nodash }}-{{ params.model_name }}", - runtime_version="2.1", - python_version="3.8", + runtime_version="1.15", + python_version="3.7", job_dir=JOB_DIR, package_uris=[TRAINER_URI], training_python_module=TRAINER_PY_MODULE, @@ -148,10 +148,10 @@ def write_predict_file(path_to_file: str): "name": "v1", "description": "First-version", "deployment_uri": JOB_DIR, - "runtime_version": "2.1", + "runtime_version": "1.15", "machineType": "mls1-c1-m2", "framework": "TENSORFLOW", - "pythonVersion": "3.8", + "pythonVersion": "3.7", }, ) # [END howto_operator_gcp_mlengine_create_version1] @@ -165,10 +165,10 @@ def write_predict_file(path_to_file: str): "name": "v2", "description": "Second version", "deployment_uri": JOB_DIR, - "runtime_version": "2.1", + "runtime_version": "1.15", "machineType": "mls1-c1-m2", "framework": "TENSORFLOW", - "pythonVersion": "3.8", + "pythonVersion": "3.7", }, ) # [END howto_operator_gcp_mlengine_create_version2] diff --git a/tests/system/providers/google/cloud/ml_engine/example_mlengine_async.py b/tests/system/providers/google/cloud/ml_engine/example_mlengine_async.py index 873b574016013..7c2caef846cfc 100644 --- a/tests/system/providers/google/cloud/ml_engine/example_mlengine_async.py +++ b/tests/system/providers/google/cloud/ml_engine/example_mlengine_async.py @@ -49,14 +49,14 @@ DAG_ID = "async_example_gcp_mlengine" PREDICT_FILE_NAME = "async_predict.json" -MODEL_NAME = f"example_async_ml_model_{ENV_ID}" -BUCKET_NAME = f"example_async_ml_bucket_{ENV_ID}" +MODEL_NAME = f"model_{DAG_ID}_{ENV_ID}".replace("-", "_") +BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}".replace("_", "-") BUCKET_PATH = f"gs://{BUCKET_NAME}" JOB_DIR = f"{BUCKET_PATH}/job-dir" SAVED_MODEL_PATH = f"{JOB_DIR}/" PREDICTION_INPUT = f"{BUCKET_PATH}/{PREDICT_FILE_NAME}" PREDICTION_OUTPUT = f"{BUCKET_PATH}/prediction_output/" -TRAINER_URI = "gs://system-tests-resources/example_gcp_mlengine/async-trainer-0.2.tar.gz" +TRAINER_URI = "gs://airflow-system-tests-resources/ml-engine/async-trainer-0.2.tar.gz" TRAINER_PY_MODULE = "trainer.task" SUMMARY_TMP = f"{BUCKET_PATH}/tmp/" SUMMARY_STAGING = f"{BUCKET_PATH}/staging/" @@ -74,7 +74,7 @@ def generate_model_predict_input_data() -> list[int]: schedule="@once", start_date=datetime(2021, 1, 1), catchup=False, - tags=["example", "ml_engine"], + tags=["example", "ml_engine", "deferrable"], params={"model_name": MODEL_NAME}, ) as dag: create_bucket = GCSCreateBucketOperator( @@ -104,8 +104,8 @@ def write_predict_file(path_to_file: str): project_id=PROJECT_ID, region="us-central1", job_id="async_training-job-{{ ts_nodash }}-{{ params.model_name }}", - runtime_version="2.1", - python_version="3.8", + runtime_version="1.15", + python_version="3.7", job_dir=JOB_DIR, package_uris=[TRAINER_URI], training_python_module=TRAINER_PY_MODULE, @@ -149,10 +149,10 @@ def write_predict_file(path_to_file: str): "name": "v1", "description": "First-version", "deployment_uri": JOB_DIR, - "runtime_version": "2.1", + "runtime_version": "1.15", "machineType": "mls1-c1-m2", "framework": "TENSORFLOW", - "pythonVersion": "3.8", + "pythonVersion": "3.7", }, ) # [END howto_operator_gcp_mlengine_create_version1] @@ -166,10 +166,10 @@ def write_predict_file(path_to_file: str): "name": "v2", "description": "Second version", "deployment_uri": JOB_DIR, - "runtime_version": "2.1", + "runtime_version": "1.15", "machineType": "mls1-c1-m2", "framework": "TENSORFLOW", - "pythonVersion": "3.8", + "pythonVersion": "3.7", }, ) # [END howto_operator_gcp_mlengine_create_version2]