4949
5050DAG_ID = "async_example_gcp_mlengine"
5151PREDICT_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 ( "_" , "-" )
5454BUCKET_PATH = f"gs://{ BUCKET_NAME } "
5555JOB_DIR = f"{ BUCKET_PATH } /job-dir"
5656SAVED_MODEL_PATH = f"{ JOB_DIR } /"
5757PREDICTION_INPUT = f"{ BUCKET_PATH } /{ PREDICT_FILE_NAME } "
5858PREDICTION_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"
6060TRAINER_PY_MODULE = "trainer.task"
6161SUMMARY_TMP = f"{ BUCKET_PATH } /tmp/"
6262SUMMARY_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