Skip to content

Commit 9d68492

Browse files
misc: update MLEngine system tests (#32881)
1 parent b733667 commit 9d68492

File tree

3 files changed

+22
-19
lines changed

3 files changed

+22
-19
lines changed

β€Žairflow/providers/google/cloud/operators/mlengine.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1002,6 +1002,9 @@ class MLEngineStartTrainingJobOperator(GoogleCloudBaseOperator):
10021002
For more information on how to use this operator, take a look at the guide:
10031003
:ref:`howto/operator:MLEngineStartTrainingJobOperator`
10041004
1005+
For more information about used parameters, check:
1006+
https://cloud.google.com/sdk/gcloud/reference/ml-engine/jobs/submit/training
1007+
10051008
:param job_id: A unique templated id for the submitted Google MLEngine
10061009
training job. (templated)
10071010
:param region: The Google Compute Engine region to run the MLEngine training

β€Žtests/system/providers/google/cloud/ml_engine/example_mlengine.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -49,14 +49,14 @@
4949

5050
DAG_ID = "example_gcp_mlengine"
5151
PREDICT_FILE_NAME = "predict.json"
52-
MODEL_NAME = f"example_ml_model_{ENV_ID}"
53-
BUCKET_NAME = f"example_ml_bucket_{ENV_ID}"
52+
MODEL_NAME = f"model_{DAG_ID}_{ENV_ID}".replace("_", "-")
53+
BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}".replace("_", "-")
5454
BUCKET_PATH = f"gs://{BUCKET_NAME}"
5555
JOB_DIR = f"{BUCKET_PATH}/job-dir"
5656
SAVED_MODEL_PATH = f"{JOB_DIR}/"
5757
PREDICTION_INPUT = f"{BUCKET_PATH}/{PREDICT_FILE_NAME}"
5858
PREDICTION_OUTPUT = f"{BUCKET_PATH}/prediction_output/"
59-
TRAINER_URI = "gs://system-tests-resources/example_gcp_mlengine/trainer-0.2.tar.gz"
59+
TRAINER_URI = "gs://airflow-system-tests-resources/ml-engine/trainer-0.2.tar.gz"
6060
TRAINER_PY_MODULE = "trainer.task"
6161
SUMMARY_TMP = f"{BUCKET_PATH}/tmp/"
6262
SUMMARY_STAGING = f"{BUCKET_PATH}/staging/"
@@ -104,8 +104,8 @@ def write_predict_file(path_to_file: str):
104104
project_id=PROJECT_ID,
105105
region="us-central1",
106106
job_id="training-job-{{ ts_nodash }}-{{ params.model_name }}",
107-
runtime_version="2.1",
108-
python_version="3.8",
107+
runtime_version="1.15",
108+
python_version="3.7",
109109
job_dir=JOB_DIR,
110110
package_uris=[TRAINER_URI],
111111
training_python_module=TRAINER_PY_MODULE,
@@ -148,10 +148,10 @@ def write_predict_file(path_to_file: str):
148148
"name": "v1",
149149
"description": "First-version",
150150
"deployment_uri": JOB_DIR,
151-
"runtime_version": "2.1",
151+
"runtime_version": "1.15",
152152
"machineType": "mls1-c1-m2",
153153
"framework": "TENSORFLOW",
154-
"pythonVersion": "3.8",
154+
"pythonVersion": "3.7",
155155
},
156156
)
157157
# [END howto_operator_gcp_mlengine_create_version1]
@@ -165,10 +165,10 @@ def write_predict_file(path_to_file: str):
165165
"name": "v2",
166166
"description": "Second version",
167167
"deployment_uri": JOB_DIR,
168-
"runtime_version": "2.1",
168+
"runtime_version": "1.15",
169169
"machineType": "mls1-c1-m2",
170170
"framework": "TENSORFLOW",
171-
"pythonVersion": "3.8",
171+
"pythonVersion": "3.7",
172172
},
173173
)
174174
# [END howto_operator_gcp_mlengine_create_version2]

β€Žtests/system/providers/google/cloud/ml_engine/example_mlengine_async.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,14 +49,14 @@
4949

5050
DAG_ID = "async_example_gcp_mlengine"
5151
PREDICT_FILE_NAME = "async_predict.json"
52-
MODEL_NAME = f"example_async_ml_model_{ENV_ID}"
53-
BUCKET_NAME = f"example_async_ml_bucket_{ENV_ID}"
52+
MODEL_NAME = f"model_{DAG_ID}_{ENV_ID}".replace("-", "_")
53+
BUCKET_NAME = f"bucket_{DAG_ID}_{ENV_ID}".replace("_", "-")
5454
BUCKET_PATH = f"gs://{BUCKET_NAME}"
5555
JOB_DIR = f"{BUCKET_PATH}/job-dir"
5656
SAVED_MODEL_PATH = f"{JOB_DIR}/"
5757
PREDICTION_INPUT = f"{BUCKET_PATH}/{PREDICT_FILE_NAME}"
5858
PREDICTION_OUTPUT = f"{BUCKET_PATH}/prediction_output/"
59-
TRAINER_URI = "gs://system-tests-resources/example_gcp_mlengine/async-trainer-0.2.tar.gz"
59+
TRAINER_URI = "gs://airflow-system-tests-resources/ml-engine/async-trainer-0.2.tar.gz"
6060
TRAINER_PY_MODULE = "trainer.task"
6161
SUMMARY_TMP = f"{BUCKET_PATH}/tmp/"
6262
SUMMARY_STAGING = f"{BUCKET_PATH}/staging/"
@@ -74,7 +74,7 @@ def generate_model_predict_input_data() -> list[int]:
7474
schedule="@once",
7575
start_date=datetime(2021, 1, 1),
7676
catchup=False,
77-
tags=["example", "ml_engine"],
77+
tags=["example", "ml_engine", "deferrable"],
7878
params={"model_name": MODEL_NAME},
7979
) as dag:
8080
create_bucket = GCSCreateBucketOperator(
@@ -104,8 +104,8 @@ def write_predict_file(path_to_file: str):
104104
project_id=PROJECT_ID,
105105
region="us-central1",
106106
job_id="async_training-job-{{ ts_nodash }}-{{ params.model_name }}",
107-
runtime_version="2.1",
108-
python_version="3.8",
107+
runtime_version="1.15",
108+
python_version="3.7",
109109
job_dir=JOB_DIR,
110110
package_uris=[TRAINER_URI],
111111
training_python_module=TRAINER_PY_MODULE,
@@ -149,10 +149,10 @@ def write_predict_file(path_to_file: str):
149149
"name": "v1",
150150
"description": "First-version",
151151
"deployment_uri": JOB_DIR,
152-
"runtime_version": "2.1",
152+
"runtime_version": "1.15",
153153
"machineType": "mls1-c1-m2",
154154
"framework": "TENSORFLOW",
155-
"pythonVersion": "3.8",
155+
"pythonVersion": "3.7",
156156
},
157157
)
158158
# [END howto_operator_gcp_mlengine_create_version1]
@@ -166,10 +166,10 @@ def write_predict_file(path_to_file: str):
166166
"name": "v2",
167167
"description": "Second version",
168168
"deployment_uri": JOB_DIR,
169-
"runtime_version": "2.1",
169+
"runtime_version": "1.15",
170170
"machineType": "mls1-c1-m2",
171171
"framework": "TENSORFLOW",
172-
"pythonVersion": "3.8",
172+
"pythonVersion": "3.7",
173173
},
174174
)
175175
# [END howto_operator_gcp_mlengine_create_version2]

0 commit comments

Comments
 (0)