Dataset aware scheduling - is there a way to reset DAG?

[funes79]

In a simple use case - where DAG is allowed to run in max 1 instance, and when there is a daily processing i.e. DAG should be executed once a day - I would like to know if there is a way to reset DAG's last dataset event trigger - i.e. manually define that your last dataset triggered date is midnight.

What is the problem

Two waiting DAGs - WAIT1 and WAIT2 are waiting for something and then setting DATASET1 and DATASET2 as outlets.

MyDAG should be triggered after both updates, but in the current day. This works when you set up everything at once, but got out of sync if you re-execute one of the WAIT dags.

For example:

Day1: 
WAIT1 - starts at midnight and waits, and sets ready DATASET1 at 7:00 AM. (MyDAG = 1/2)
WAIT2 - starts at midnight and waits and sets ready DATASET2 at 15:00 AM. (MyDAG= 2/2)
MyDAG triggers at 15:00.  This is ok. (MyDAG=0/2)

Now somebody repeats WAIT2 manually at 16:00. (MyDAG=1/2)

Day2:
WAIT1 starts at midnight again and waits and sets ready let say DATASET1 at 6:30 AM. (MyDAG = 2/2)
!! and here is a problem: now we see the MyDAG triggered because, since
 the **last dataset triggered run** both datasets has been updated.

WAIT2 -  sets DATASET2 as ready let say 14:30.

And as you can see the MyDAG in day two starts at 6:30 instead of 14:30, because it is out of sync.

I tried to manually ruyn MyDAG - hopin that it will reset the Datasets time, so basically it will show 0/2 datasets ready, but no, does not help.

Possible workaround

So the only option is to rerun at the same time ALL WAIT dags and then let MyDAG to trigger again. But this is costly (as there can be many downstream dags) - and it can be out of sync if somebody manually runs WAIT2.

[funes79]

And my propose solution with "RESET" function:

Day1: 
WAIT1 - starts at midnight and waits, and sets ready DATASET1 at 7:00 AM. (MyDAG = 1/2)
WAIT2 - starts at midnight and waits and sets ready DATASET2 at 15:00 AM. (MyDAG= 2/2)
MyDAG triggers at 15:00.  This is ok. (MyDAG=0/2)

Now somebody repeats WAIT2 manually at 16:00. (MyDAG=1/2)

Day2:
=> Immediately after midnight:
=> We call MyDAG.reset_last_dataset_trigger_date("Day2-Month-2024 00:00:00");
=>  so this way we get back the DAG into (MyDAG=0/2)
WAIT1 starts at midnight again and waits and sets ready let say DATASET1 at 6:30 AM. (MyDAG=1/2)
WAIT2 -  sets DATASET2 as ready let say 14:30 (MyDAG=2/2)
And now the MyDAG is triggered. Correctly.

[funes79]

Came up with this workaround, so whenever the dataset scheduled DAG gets " out of sync" I reset the datasets state by running this helper DAG. Essentially I am clearing out the DAG's dataset queue.

def reset_dataset_queues(**kwargs):
    params = kwargs.get("params", {})
    print("Resetting for ", params)
    from airflow.utils.db import create_session
    from sqlalchemy import delete, select
    from airflow.models.dataset import DatasetDagRunQueue

    dag_id = params.get("dag_id")
    with create_session() as session:
        session.execute(delete(DatasetDagRunQueue).where(DatasetDagRunQueue.target_dag_id == dag_id))
    print(session.execute(select(DatasetDagRunQueue).where(DatasetDagRunQueue.target_dag_id == dag_id)).all())
    return "OK"



params = {"dag_id": Param(type=["string", "null"], enum=["DagA", "DagB"], default=None)}
with DAG(
    "reset_datasets",
    start_date=pendulum.now(),
    schedule_interval=None,
    catchup=False,
    doc_md=doc_md_DAG,
    params=params,
) as dag:
    pythonOperator = PythonOperator(
        task_id="task",
        python_callable=reset_dataset_queues,
    )

[cmarteepants]

Glad that you were able to resolve it this way, but agree that this is just a workaround and not a solution to a real problem. There should be a way to "reset" things without manipulating the Airflow DB directly.

[Amin-Siddique]

Nice solution, i agree it would be really nice to have a generic solution as our team is also facing the same issue causing unbalance. For a workaround, we tried creating a custom datasetsensor where the consumer dag checks whether it's execution date matches with tha of producer dag.But again, if producer dags are scheduled multiples times in a day, one if one dataset update fails, our sensor will not be able to catch that.

[SuaveCharlie]

@Amin-Siddique that's exactly what I've done too; task at the start of any Dataset triggered dags that inspects when the dataset was updated and whether it's within an expected tolerance.
I'm surprised there isn't more noise about this key flaw of Datasets!