<reponame>KoiDev13/airflow_lab<filename>airflow/dags/xcom_dag.py
from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator, BranchPythonOperator #Looking for a way to choose one task or another according to condition
# from airflow.operators.subdag import SubDagOperator
from airflow.utils.task_group import TaskGroup #Create multiple tasks in a group
from airflow.operators.dummy import DummyOperator #For dummy test

from random import uniform
from datetime import datetime

default_args = {
    'start_date': datetime(2020, 1, 1)
}

def _training_model(ti): 
    accuracy = uniform(0.1, 10.0) #Uniform method from random
    print(f'model\'s accuracy: {accuracy}')
    ti.xcom_push(key='model_accuracy', value=accuracy) #truyền tham số xcom vào kết quả return

def _choose_best_model(ti):
    print('choose best model')
    accuracies=ti.xcom_pull(key='model_accuracy', task_ids=[
        'processing_tasks.training_model_a',
        'processing_tasks.training_model_b',
        'processing_tasks.training_model_c'
    ])
    for accuracy in accuracies:
        if accuracy > 5:
            return 'accurate'
    return 'inaccurate'
    print(accuracies)

with DAG('xcom_dag', schedule_interval='@daily', default_args=default_args, catchup=False) as dag:

    downloading_data = BashOperator(
        task_id='downloading_data',
        bash_command='sleep 3',
        do_xcom_push=False
    )

    with TaskGroup('processing_tasks') as processing_tasks:
        training_model_a = PythonOperator(
            task_id='training_model_a',
            python_callable=_training_model
        )

        training_model_b = PythonOperator(
            task_id='training_model_b',
            python_callable=_training_model
        )

        training_model_c = PythonOperator(
            task_id='training_model_c',
            python_callable=_training_model #automatically push to xcom
        )

    choose_model = BranchPythonOperator(
        task_id='choose_best_model',
        python_callable=_choose_best_model
    )

    accurate = DummyOperator(
        task_id = 'accurate',
    )
    
    inaccurate = DummyOperator(
        task_id = 'inaccurate',
    )

    storing = DummyOperator(
        task_id='storing',
        trigger_rule = 'none_failed_or_skipped'
    )

    downloading_data >> processing_tasks >> choose_model
    choose_model >> [accurate, inaccurate] >> storing
from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator, BranchPythonOperator #Looking for a way to choose one task or another according to condition
# from airflow.operators.subdag import SubDagOperator
from airflow.utils.task_group import TaskGroup #Create multiple tasks in a group
from airflow.operators.dummy import DummyOperator #For dummy test

from random import uniform
from datetime import datetime

default_args = {
    'start_date': datetime(2020, 1, 1)
}

def _training_model(ti): 
    accuracy = uniform(0.1, 10.0) #Uniform method from random
    print(f'model\'s accuracy: {accuracy}')
    ti.xcom_push(key='model_accuracy', value=accuracy) #truyền tham số xcom vào kết quả return

def _choose_best_model(ti):
    print('choose best model')
    accuracies=ti.xcom_pull(key='model_accuracy', task_ids=[
        'processing_tasks.training_model_a',
        'processing_tasks.training_model_b',
        'processing_tasks.training_model_c'
    ])
    for accuracy in accuracies:
        if accuracy > 5:
            return 'accurate'
    return 'inaccurate'
    print(accuracies)

with DAG('xcom_dag', schedule_interval='@daily', default_args=default_args, catchup=False) as dag:

    downloading_data = BashOperator(
        task_id='downloading_data',
        bash_command='sleep 3',
        do_xcom_push=False
    )

    with TaskGroup('processing_tasks') as processing_tasks:
        training_model_a = PythonOperator(
            task_id='training_model_a',
            python_callable=_training_model
        )

        training_model_b = PythonOperator(
            task_id='training_model_b',
            python_callable=_training_model
        )

        training_model_c = PythonOperator(
            task_id='training_model_c',
            python_callable=_training_model #automatically push to xcom
        )

    choose_model = BranchPythonOperator(
        task_id='choose_best_model',
        python_callable=_choose_best_model
    )

    accurate = DummyOperator(
        task_id = 'accurate',
    )
    
    inaccurate = DummyOperator(
        task_id = 'inaccurate',
    )

    storing = DummyOperator(
        task_id='storing',
        trigger_rule = 'none_failed_or_skipped'
    )

    downloading_data >> processing_tasks >> choose_model
    choose_model >> [accurate, inaccurate] >> storing
#Looking for a way to choose one task or another according to condition # from airflow.operators.subdag import SubDagOperator #Create multiple tasks in a group #For dummy test #Uniform method from random #truyền tham số xcom vào kết quả return #automatically push to xcom