How to pull XCOM value from on_failure_callback - airflow

When a task fails, is it possible to pull a XCOM value that has been previously set in an other task during the on_failure_callback execution?
To be more specific, exemple:
dag: task1 >> task2
task1 runs successfully and set key="test" value=123 in Xcom
task2 fails
on_failure_callback is called
Is this possible to retrieve the value of key test in the on_failure_callback ?
I tried like this, but it seems it didn't find any value:
# Daf configuration
...
"on_failure_callback": deploy_failure,
...
# In task1
kwargs["ti"].xcom_push(key="test", value=123)
# on_failure_callback method
def deploy_failure(context):
print("/!\ Deploy failure callback triggered...")
test_value = context.get("ti").xcom_pull(key="test")
print(test_value)
test_value is None
I am sure the Xcom value is set because I can see it on the Airflow backend.
Any idea?

I guess there is some issue with provide_context in the failure_callback. You can work around this by accessing directly XCom class:
from airflow.models import XCom
def deploy_failure(context):
print("/!\ Deploy failure callback triggered...")
test_value = XCom.get_one(execution_date = context.get('execution_date'), key='test')
print("ALERT: {0}".format(test_value))

Related

Airflow: Get the status of the prior run for a task

I'm working with Airflow 2.1.4 and looking to find the status of the prior task run (Task Run, not Task Instance and not Dag Run).
I.e., DAG MorningWorkflow runs a 9:00am, and task ConditionalTask is in that dag. There is some precondition logic that will throw an AirflowSkipException in a number of situations (including timeframe of day and other context-specific information to reduce the likelihood of collisions with independent processes)
If ConditionalTask fails, we can fix the issue, clear the failed run, and re-run it without running the entire DAG. However, the skip logic reruns and will often now skip it, even though the original conditions were non-skipping.
So, I want to update the precondition logic to never skip if this taskinstance ran previously and failed. I can determine if the taskinstance ran previously using TaskInstance.try_number orTaskInstance.prev_attempted_tries, but this doesn't tell me whether it actually tried to run originally or if it skipped (i.e., if we clear the entire DagRun to rerun the whole workflow, we would want it to still skip).
An alternative would be to determine whether the first attempted run was skipped or not.
#Kevin Crouse
In order to answer your question, we can take advantage of from airflow.models import DagRun
To provide you with a complete, answer I have created two functions to assist you in resolving similar quandaries in the future.
How to return the overall state/success of a specific dag_id passed as a function arg?
def get_last_dag_run_status(dag_id):
""" Returns the status of the last dag run for the given dag_id
1. Utilise the find method of DagRun class
2. Step 1 returns a list, so we sort it by the last execution date
3. I have returned 2 examples for you to see a) the state, b) the last execution date, you can explore this further by just returning last_dag_run[0]
Args:
dag_id (str): The dag_id to check
Returns:
List - The status of the last dag run for the given dag_id
List - The last execution date of the dag run for the given dag_id
"""
last_dag_run = DagRun.find(dag_id=dag_id)
last_dag_run.sort(key=lambda x: x.execution_date, reverse=True)
return [last_dag_run[0].state, last_dag_run[0].execution_date]
How to return the status of a specific task_id, within a specific dag_id?
def get_task_status(dag_id, task_id):
""" Returns the status of the last dag run for the given dag_id
1. The code is very similar to the above function, I use it as the foundation for many similar problems/solutions
2. The key difference is that in the return statement, we can directly access the .get_task_instance passing our desired task_id and its state
Args:
dag_id (str): The dag_id to check
task_id (str): The task_id to check
Returns:
List - The status of the last dag run for the given dag_id
"""
last_dag_run = DagRun.find(dag_id=dag_id)
last_dag_run.sort(key=lambda x: x.execution_date, reverse=True)
return last_dag_run[0].get_task_instance(task_id).state
I hope this helps you in your journey to resolve your issues.
For posterity, here is a complete dummy Dag to demonstrate the 2 functions working.
from airflow import DAG
from airflow.operators.dummy import DummyOperator
from airflow.operators.python import PythonOperator
from airflow.models import DagRun
from datetime import datetime
def get_last_dag_run_status(dag_id):
""" Returns the status of the last dag run for the given dag_id
Args:
dag_id (str): The dag_id to check
Returns:
List - The status of the last dag run for the given dag_id
List - The last execution date of the dag run for the given dag_id
"""
last_dag_run = DagRun.find(dag_id=dag_id)
last_dag_run.sort(key=lambda x: x.execution_date, reverse=True)
return [last_dag_run[0].state, last_dag_run[0].execution_date]
def get_task_status(dag_id, task_id):
""" Returns the status of the last dag run for the given dag_id
Args:
dag_id (str): The dag_id to check
task_id (str): The task_id to check
Returns:
List - The status of the last dag run for the given dag_id
"""
last_dag_run = DagRun.find(dag_id=dag_id)
last_dag_run.sort(key=lambda x: x.execution_date, reverse=True)
return last_dag_run[0].get_task_instance(task_id).state
with DAG(
'stack_overflow_ans_1',
tags = ['SO'],
start_date = datetime(2022, 1, 1),
schedule_interval = None,
catchup = False,
is_paused_upon_creation = False
) as dag:
t1 = DummyOperator(
task_id = 'start'
)
t2 = PythonOperator(
task_id = 'get_last_dag_run_status',
python_callable = get_last_dag_run_status,
op_args = ['YOUR_DAG_NAME'],
do_xcom_push = False
)
t3 = PythonOperator(
task_id = 'get_task_status',
python_callable = get_task_status,
op_args = ['YOUR_DAG_NAME', 'YOUR_DAG_TASK_WITHIN_THE_DAG'],
do_xcom_push = False
)
t4 = DummyOperator(
task_id = 'end'
)
t1 >> t2 >> t3 >> t4

Airflow: Get status of a task from an external DAG

I want to get the status of a task from an external DAG. I have the same tasks running in 2 different DAGs based on some conditions. So, I want to check the status of this task in DAG2 from DAG1. If the task status is 'running' in DAG2, then I will skip this task in DAG1.
I tried using:
dag_runs = DagRun.find(dag_id=dag_id,execution_date=exec_dt)
for dag_run in dag_runs:
dag_run.state
I couldn't figure out if we can get task status using DagRun.
If I use TaskDependencySensor, the DAG will have to wait until it finds the allowed_states of the task.
Is there a way to get the current status of a task in another DAG?
I used below code to get the status of a task from another DAG:
from airflow.api.common.experimental.get_task_instance import get_task_instance
def get_dag_state(execution_date, **kwargs):
ti = get_task_instance('dag_id', 'task_id', execution_date)
task_status = ti.current_state()
return task_status
dag_status = BranchPythonOperator(
task_id='dag_status',
python_callable=get_dag_state,
dag=dag
)
More details can be found here

Airflow Exception not being thrown when encountering None/Falsy values

I am trying to pass data between a PythonOperator, _etl_lasic to another PythonOperator, _download_s3_data, which works fine but I want to throw an exception when the value passed is None which should mark the task as a failure.
import airflow
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.exceptions import AirflowFailException
def _etl_lasic(**context):
path_s3 = None
context["task_instance"].xcom_push(
key="path_s3",
value=path_s3,
)
def _download_s3_data(templates_dict, **context):
path_s3 = templates_dict["path_s3"]
if not path_s3:
raise AirflowFailException("Path to S3 was not passed!")
else:
print(f"Path to S3: {path_s3}")
with DAG(
dag_id="02_lasic_retraining_without_etl",
start_date=airflow.utils.dates.days_ago(3),
schedule_interval="#once",
) as dag:
etl_lasic = PythonOperator(
task_id="etl_lasic",
python_callable=_etl_lasic,
)
download_s3_data = PythonOperator(
task_id="download_s3_data",
python_callable=_download_s3_data,
templates_dict={
"path_s3": "{{task_instance.xcom_pull(task_ids='etl_lasic',key='path_s3')}}"
},
)
etl_lasic >> download_s3_data
Logs:
[2021-08-17 04:04:41,128] {logging_mixin.py:103} INFO - Path to S3: None
[2021-08-17 04:04:41,128] {python.py:118} INFO - Done. Returned value was: None
[2021-08-17 04:04:41,143] {taskinstance.py:1135} INFO - Marking task as SUCCESS. dag_id=02_lasic_retraining_without_etl, task_id=download_s3_data, execution_date=20210817T040439, start_date=20210817T040440, end_date=20210817T040441
[2021-08-17 04:04:41,189] {taskinstance.py:1195} INFO - 0 downstream tasks scheduled from follow-on schedule check
[2021-08-17 04:04:41,212] {local_task_job.py:118} INFO - Task exited with return code 0
Jinja-templated values are rendered as strings by default. In your case, even though you push an XCom value of None, when the value is pulled via "{{task_instance.xcom_pull(task_ids='etl_lasic',key='path_s3')}}" the value is actually rendered as "None" which doesn't throw an exception based on the current logic.
There are two options that will solve this:
Instead of setting path_s3 to None in the "_etl_lasic" function, set it to an empty string.
If you are using Airflow 2.1+, there is a parameter, render_template_as_native_obj, that can be set at the DAG level which will render Jinja-templated values as native Python types (list, dict, etc.). Setting that parameter to True will do the trick without changing how path_s3 is set in the function. A conceptual example is documented here.

How to setup Nagios Alerts for Apache Airflow Dags

Is it possible to setup Nagios alerts for airflow dags?
In case the dag is failed, I need to alert the respective groups.
You can add an "on_failure_callback" to any task which will call an arbitrary failure handling function. In that function you can then send an error call to Nagios.
For example:
dag = DAG(dag_id="failure_handling",
schedule_interval='#daily')
def handle_failure(context):
# first get useful fields to send to nagios/elsewhere
dag_id = context['dag'].dag_id
ds = context['ds']
task_id = context['ti'].task_id
# instead of printing these out - you can send these to somewhere else
logging.info("dag_id={}, ds={}, task_id={}".format(dag_id, ds, task_id))
def task_that_fails(**kwargs):
raise Exception("failing test")
task_to_fail = PythonOperator(
task_id='python_task_to_fail',
python_callable=task_that_fails,
provide_context=True,
on_failure_callback=handle_failure,
dag=dag)
If you run a test on this:
airflow test failure_handling task_to_fail 2018-08-10
You get the following in your log output:
INFO - dag_id=failure_handling, ds=2018-08-10, task_id=task_to_fail

Is it possible to have a pipeline in Airflow that does not tie to any schedule?

I need to have pipeline that will be executed either manually or programmatically, is possible with Airflow? Looks like right now each workflow MUST be tied to a schedule.
Just set the schedule_interval to None when you create the DAG:
dag = DAG('workflow_name',
template_searchpath='path',
schedule_interval=None,
default_args=default_args)
From the Airflow Manual:
Each DAG may or may not have a schedule, which informs how DAG Runs
are created. schedule_interval is defined as a DAG arguments, and
receives preferably a cron expression as a str, or a
datetime.timedelta object.
The manual then goes on to list some cron 'presets' one of which is None.
Yes, this can be achieved by passing None to schedule_interval in default_args.
Check this documation on DAG run.
For example:
default_args = {
'owner': 'airflow',
'depends_on_past': False,
'start_date': datetime(2015, 12, 1),
'email': ['airflow#example.com'],
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
'retry_delay': timedelta(minutes=5),
'schedule_interval': None, # Check this line
}
In Airflow, every DAG is required to have a start date and schedule interval*, for example hourly:
import datetime
dag = DAG(
dag_id='my_dag',
schedule_interval=datetime.timedelta(hours=1),
start_date=datetime(2018, 5, 23),
)
(Without a schedule how would it know when to run?)
Alternatively to a cron schedule, you can set the schedule to #once to only run once.
*One exception: You can omit the schedule for externally triggered DAGs because Airflow will not schedule them itself.
However, that said, if you omit the schedule, then you need to trigger the DAG externally somehow. If you want to be able to call a DAG programmatically, for instance, as a result of a separate condition occurring in another DAG, you can do that with the TriggerDagRunOperator. You might also hear this idea called externally triggered DAGs.
Here's a usage example from the Airflow Example DAGs:
File 1 - example_trigger_controller_dag.py:
"""This example illustrates the use of the TriggerDagRunOperator. There are 2
entities at work in this scenario:
1. The Controller DAG - the DAG that conditionally executes the trigger
2. The Target DAG - DAG being triggered (in example_trigger_target_dag.py)
This example illustrates the following features :
1. A TriggerDagRunOperator that takes:
a. A python callable that decides whether or not to trigger the Target DAG
b. An optional params dict passed to the python callable to help in
evaluating whether or not to trigger the Target DAG
c. The id (name) of the Target DAG
d. The python callable can add contextual info to the DagRun created by
way of adding a Pickleable payload (e.g. dictionary of primitives). This
state is then made available to the TargetDag
2. A Target DAG : c.f. example_trigger_target_dag.py
"""
from airflow import DAG
from airflow.operators.dagrun_operator import TriggerDagRunOperator
from datetime import datetime
import pprint
pp = pprint.PrettyPrinter(indent=4)
def conditionally_trigger(context, dag_run_obj):
"""This function decides whether or not to Trigger the remote DAG"""
c_p = context['params']['condition_param']
print("Controller DAG : conditionally_trigger = {}".format(c_p))
if context['params']['condition_param']:
dag_run_obj.payload = {'message': context['params']['message']}
pp.pprint(dag_run_obj.payload)
return dag_run_obj
# Define the DAG
dag = DAG(dag_id='example_trigger_controller_dag',
default_args={"owner": "airflow",
"start_date": datetime.utcnow()},
schedule_interval='#once')
# Define the single task in this controller example DAG
trigger = TriggerDagRunOperator(task_id='test_trigger_dagrun',
trigger_dag_id="example_trigger_target_dag",
python_callable=conditionally_trigger,
params={'condition_param': True,
'message': 'Hello World'},
dag=dag)
File 2 - example_trigger_target_dag.py:
from airflow.operators.bash_operator import BashOperator
from airflow.operators.python_operator import PythonOperator
from airflow.models import DAG
from datetime import datetime
import pprint
pp = pprint.PrettyPrinter(indent=4)
# This example illustrates the use of the TriggerDagRunOperator. There are 2
# entities at work in this scenario:
# 1. The Controller DAG - the DAG that conditionally executes the trigger
# (in example_trigger_controller.py)
# 2. The Target DAG - DAG being triggered
#
# This example illustrates the following features :
# 1. A TriggerDagRunOperator that takes:
# a. A python callable that decides whether or not to trigger the Target DAG
# b. An optional params dict passed to the python callable to help in
# evaluating whether or not to trigger the Target DAG
# c. The id (name) of the Target DAG
# d. The python callable can add contextual info to the DagRun created by
# way of adding a Pickleable payload (e.g. dictionary of primitives). This
# state is then made available to the TargetDag
# 2. A Target DAG : c.f. example_trigger_target_dag.py
args = {
'start_date': datetime.utcnow(),
'owner': 'airflow',
}
dag = DAG(
dag_id='example_trigger_target_dag',
default_args=args,
schedule_interval=None)
def run_this_func(ds, **kwargs):
print("Remotely received value of {} for key=message".
format(kwargs['dag_run'].conf['message']))
run_this = PythonOperator(
task_id='run_this',
provide_context=True,
python_callable=run_this_func,
dag=dag)
# You can also access the DagRun object in templates
bash_task = BashOperator(
task_id="bash_task",
bash_command='echo "Here is the message: '
'{{ dag_run.conf["message"] if dag_run else "" }}" ',
dag=dag)

Resources