# launch sagemaker hyperparameter job and wait until it completes
tune_model_task = SageMakerTuningOperator(task_id='model_tuning',
dag=dag,
config=tuner_config,
aws_conn_id='airflow-sagemaker',
wait_for_completion=True,
check_interval=30)
# launch sagemaker batch transform job and wait until it completes
batch_transform_task = SageMakerTransformOperator(
task_id='predicting',
dag=dag,
config=transform_config,
aws_conn_id='airflow-sagemaker',
wait_for_completion=True,
check_interval=30,
trigger_rule=TriggerRule.ONE_SUCCESS)
cleanup_task = DummyOperator(task_id='cleaning_up', dag=dag)
# set the dependencies between tasks
init.set_downstream(preprocess_task)
preprocess_task.set_downstream(prepare_task)
prepare_task.set_downstream(branching)
branching.set_downstream(tune_model_task)
branching.set_downstream(train_model_task)
tune_model_task.set_downstream(batch_transform_task)
train_model_task.set_downstream(batch_transform_task)
batch_transform_task.set_downstream(cleanup_task)
'region': region,
'bucket': bucket
})
# launch sagemaker batch transform job and wait until it completes
batch_transform_task = SageMakerTransformOperator(
task_id='batch_predicting',
dag=dag,
config=transform_config,
aws_conn_id='airflow-sagemaker',
wait_for_completion=True,
check_interval=30)
# Cleanup task, deletes ALL SageMaker endpoints and model artifacts
# Uncomment below clean_up_task to clean up sagemaker endpoint resources and model artifacts
# clean_up_task = PythonOperator(
# task_id='clean_up',
# dag=dag,
# python_callable=clean_up.clean_up,
# op_kwargs={'region': region, "bucket": bucket}
# )
init.set_downstream(sm_proc_job_task)
sm_proc_job_task.set_downstream(train_model_task)
train_model_task.set_downstream(inference_pipeline_task)
inference_pipeline_task.set_downstream(batch_transform_task)
# Uncomment line below to disable clean up task
# batch_transform_task.set_downstream(clean_up_task)
# check_interval=30
# )
# # launch sagemaker batch transform job and wait until it completes
# batch_transform_task = SageMakerTransformOperator(
# task_id='predicting',
# dag=dag,
# config=transform_config,
# aws_conn_id='airflow-sagemaker',
# wait_for_completion=True,
# check_interval=30,
# trigger_rule=TriggerRule.ONE_SUCCESS
# )
cleanup_task = DummyOperator(
task_id='cleaning_up',
dag=dag)
# set the dependencies between tasks
# init.set_downstream(preprocess_task)
# preprocess_task.set_downstream(prepare_task)
# prepare_task.set_downstream(branching)
# branching.set_downstream(tune_model_task)
# branching.set_downstream(train_model_task)
# tune_model_task.set_downstream(batch_transform_task)
# train_model_task.set_downstream(batch_transform_task)
# batch_transform_task.set_downstream(cleanup_task)
init.set_downstream(train_model_task)
train_model_task.set_downstream(cleanup_task)