default_args = {
   'owner': 'airflow',
    'depends_on_past': False,
    'start_date': days_ago(2),
    'email': ['airflow@example.com'],
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

training_dag = DAG(
    'my_training_dag',
    default_args=default_args,
    description='A simple tutorial DAG',
    schedule_interval=timedelta(days=1),
)

input_param = {'input_data':'/tmp/1.parquet'}

with training_dag as dag:
   # First data preparation task
    prep_data = PythonOperator(
        task_id="prep_data",
        python_callable=prep_data.prep_data,
        op_kwargs=training_config,
        provide_context=True
    )
    
    #Second model training task
    training_model = PythonOperator(
        task_id="prep_data",
        python_callable=tasks.prep_data,
        op_kwargs=input_param,
        provide_context=True
    )

    [prep_data, training_model]