misc: update MLEngine system tests (#32881)

apache · Jul 27, 2023 · 9d68492 · 9d68492
1 parent b733667
commit 9d68492
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 19 deletions.
diff --git 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
@@ -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
@@ -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]