111 lines
3.1 KiB
Python
111 lines
3.1 KiB
Python
import os
|
|
import sys
|
|
from datetime import datetime
|
|
from datetime import timedelta
|
|
from pprint import pprint
|
|
|
|
from airflow import DAG
|
|
|
|
from airflow.operators.dummy_operator import DummyOperator
|
|
from airflow.operators.python_operator import PythonOperator
|
|
from airflow.contrib.operators.ecs_operator import ECSOperator
|
|
|
|
DAG_NAME = 'Test_Dag'
|
|
|
|
default_args = {
|
|
'owner': 'CM',
|
|
'start_date': datetime(2019, 6, 8),
|
|
'email': ['xyz@amazon.com'],
|
|
'email_on_failure': False,
|
|
'email_on_retry': False,
|
|
'retries': 3,
|
|
'retry_delay': timedelta(minutes=1)
|
|
}
|
|
|
|
|
|
def get_ecs_operator_args(taskDefinitionName, taskContainerName, entryFile, param):
|
|
return dict(
|
|
launch_type="FARGATE",
|
|
# The name of your task as defined in ECS
|
|
task_definition=taskDefinitionName,
|
|
# The name of your ECS cluster
|
|
cluster=os.environ['CLUSTER'],
|
|
network_configuration={
|
|
'awsvpcConfiguration': {
|
|
'securityGroups': [os.environ['SECURITY_GROUP']],
|
|
'subnets': os.environ['SUBNETS'].split(","),
|
|
'assignPublicIp': "DISABLED"
|
|
}
|
|
},
|
|
overrides={
|
|
'containerOverrides': [
|
|
{
|
|
'name': taskContainerName,
|
|
'command': ["python", entryFile, param]
|
|
}
|
|
]
|
|
}
|
|
)
|
|
|
|
oddTaskConfig = {
|
|
'taskDefinitionName': 'FarFlowCombinedTask',
|
|
'taskContainerName': 'MultiTaskContainer',
|
|
'entryFile': 'odd_numbers.py',
|
|
'param': '10'
|
|
}
|
|
evenTaskConfig = {
|
|
'taskDefinitionName': 'FarFlowCombinedTask',
|
|
'taskContainerName': 'MultiTaskContainer',
|
|
'entryFile': 'even_numbers.py',
|
|
'param': '10'
|
|
}
|
|
numbersTaskConfig = {
|
|
'taskDefinitionName': 'FarFlowNumbersTask',
|
|
'taskContainerName': 'NumbersContainer',
|
|
'entryFile': 'numbers.py',
|
|
'param': '10'
|
|
}
|
|
|
|
oddTask_args = get_ecs_operator_args(**oddTaskConfig)
|
|
evenTask_args = get_ecs_operator_args(**evenTaskConfig)
|
|
numbersTask_args = get_ecs_operator_args(**numbersTaskConfig)
|
|
|
|
dag = DAG( DAG_NAME,
|
|
schedule_interval=None,
|
|
default_args=default_args)
|
|
|
|
start_process = DummyOperator(task_id="start_process", dag=dag)
|
|
|
|
# Following tasks will get triggered from worker and runs on OnDemand Fargate Task
|
|
odd_task = ECSOperator(task_id="odd_task", **oddTask_args, dag=dag)
|
|
even_task = ECSOperator(task_id="even_task", **evenTask_args, dag=dag)
|
|
numbers_task = ECSOperator(task_id="numbers_task", **numbersTask_args, dag=dag)
|
|
|
|
|
|
# [START howto_operator_python]
|
|
# Pulled from : https://github.com/apache/airflow/blob/master/airflow/example_dags/example_python_operator.py#L40
|
|
def print_context(ds, **kwargs):
|
|
"""Print the Airflow context and ds variable from the context."""
|
|
pprint(kwargs)
|
|
print(ds)
|
|
return 'Whatever you return gets printed in the logs'
|
|
|
|
|
|
task_config = {
|
|
"key1": "value1",
|
|
"key2": "value2",
|
|
"key3": "value3",
|
|
"key4": "value4"
|
|
}
|
|
|
|
on_worker_task = PythonOperator(
|
|
task_id='runs_on_worker',
|
|
python_callable=print_context,
|
|
dag=dag,
|
|
op_args=[task_config]
|
|
)
|
|
# [END howto_operator_python]
|
|
|
|
|
|
start_process >> [odd_task, even_task] >> numbers_task >> on_worker_task
|