You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "Asquator (via GitHub)" <gi...@apache.org> on 2026/01/31 00:33:16 UTC
[PR] Always provide a relevant task instance in Dag callback [airflow]
Asquator opened a new pull request, #61274:
URL: https://github.com/apache/airflow/pull/61274
DRAFT
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2770913506
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2283,7 +2283,10 @@ def _schedule_dag_run(
select(TI)
.where(TI.dag_id == dag_run.dag_id)
.where(TI.run_id == dag_run.run_id)
- .where(TI.state.in_(State.unfinished))
+ .where(TI.state.in_(State.unfinished) | (TI.state.is_(None)))
Review Comment:
From what I tried, the `in` clause doesn't catch the None value. It's in accordance with SQL semantics for NULL values. See also: [https://stackoverflow.com/questions/6362112/in-clause-with-null-or-is-null](https://stackoverflow.com/questions/6362112/in-clause-with-null-or-is-null)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2777429309
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
Review Comment:
It's inside the `dagrun` class, so I decided to keep it shorter
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
Review Comment:
It's inside the `dagrun` module, so I decided to keep it shorter
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-4019233852
Newsfragment validation should pass now.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3977634143
So, after some consideration, I changed the logic for deadlocked tasks. Instead of just passing the latest finished task, we're now passing the next task that couldn't run due to the deadlock. This way the callback has more idea of an actual task that gone wrong.
To achieve this, I turned the fields of `TISchedulingDecision` from `list`s to `set`s, to allow for more indexing. Let me know if it's better to do it locally rather than changing on the class level.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3977664346
Something got wrong, I did not intend to request these reviews.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3980702300
Another pair of exes might be good, I am only 90% in this topic. But would be really good tomerge prior 3.1.8 cut.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Ajay9704 (via GitHub)" <gi...@apache.org>.
Ajay9704 commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2838308369
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2311,7 +2311,10 @@ def _schedule_dag_run(
select(TI)
.where(TI.dag_id == dag_run.dag_id)
.where(TI.run_id == dag_run.run_id)
- .where(TI.state.in_(State.unfinished))
+ .where(TI.state.in_(State.unfinished) | (TI.state.is_(None)))
+ ).all()
+ last_unfinished_ti = (
+ max(unfinished_task_instances, key=lambda ti: ti.start_date, default=None)
Review Comment:
last_unfinished_ti = max(
unfinished_task_instances, key=lambda ti: ti.start_date, default=None
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1215,22 +1215,36 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="task_failure")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(dag=cast("SDKDAG", dag), success=False, reason="task_failure")
- elif dag.has_on_failure_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=True,
- msg="task_failure",
+ if dag.has_on_failure_callback:
+ ti_causing_failure = (
+ max(
+ (ti for ti in tis_for_dagrun_state
+ if ti.state == TaskInstanceState.FAILED and ti.end_date is not None),
+ key=lambda ti: ti.end_date,
+ default=None
+ )
Review Comment:
ti_causing_failure = max(
(
ti
for ti in tis_for_dagrun_state
if ti.state == TaskInstanceState.FAILED and ti.end_date is not None
),
key=lambda ti: ti.end_date,
default=None,
\> fix as per logs
"-" ti_causing_failure = (
"-" max(
"-" (ti for ti in tis_for_dagrun_state
"-" if ti.state == TaskInstanceState.FAILED and ti.end_date is not None),
"-" key=lambda ti: ti.end_date,
"-" default=None
"-" )
"+" ti_causing_failure = max(
"+" (
"+" ti
"+" for ti in tis_for_dagrun_state
"+" if ti.state == TaskInstanceState.FAILED and ti.end_date is not None
"+" ),
"+" key=lambda ti: ti.end_date,
"+" default=None,
)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2836572530
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1245,22 +1252,31 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.SUCCESS)
self.notify_dagrun_state_changed(msg="success")
- if execute_callbacks and dag.has_on_success_callback:
- self.handle_dag_callback(dag=cast("SDKDAG", dag), success=True, reason="success")
- elif dag.has_on_success_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=False,
- msg="success",
+ if dag.has_on_success_callback:
+ last_ti_to_run: TI | None = (
+ max(tis_for_dagrun_state, key=lambda ti: ti.end_date, default=None)
)
+ if execute_callbacks:
Review Comment:
Same as https://github.com/apache/airflow/pull/61274#discussion_r2778072290
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1280,26 +1296,31 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="all_tasks_deadlocked")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(
- dag=cast("SDKDAG", dag),
- success=False,
- reason="all_tasks_deadlocked",
- )
- elif dag.has_on_failure_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=True,
- msg="all_tasks_deadlocked",
+ if dag.has_on_failure_callback:
+ last_finished_ti: TI | None = (
+ max(info.finished_tis, key=lambda ti: ti.end_date, default=None)
)
+ if execute_callbacks:
Review Comment:
Same as https://github.com/apache/airflow/pull/61274#discussion_r2778072290
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3959438684
> I think you need to set up the pre-commit hooks locally and then run the command:
>
> prek run --all-files
>
> This should fix the static errors you’re receiving.
I think it is also possible to run
prek ruff-format --all-files
Or
prek airflow-core --all-files
To not run on everything
In general, you can list prek hooks using
prek list
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2777429148
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
+
+ return min(failed_on_paths, key=lambda ti: ti.end_date, default=None)
Review Comment:
Switched to `start_date` now
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2748629648
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2283,7 +2283,10 @@ def _schedule_dag_run(
select(TI)
.where(TI.dag_id == dag_run.dag_id)
.where(TI.run_id == dag_run.run_id)
- .where(TI.state.in_(State.unfinished))
+ .where(TI.state.in_(State.unfinished) | (TI.state.is_(None)))
Review Comment:
I thought that tasks with `null` state should also be always marked as `skipped`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2776788765
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1287,6 +1290,9 @@ def recalculate(self) -> _UnfinishedStates:
reason="all_tasks_deadlocked",
)
elif dag.has_on_failure_callback:
+ last_finished_ti: TI | None = (
+ max(info.finished_tis, key=lambda ti: ti.end_date, default=None)
Review Comment:
We do this check once for a failed DAG. Wdym?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2836572530
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1245,22 +1252,31 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.SUCCESS)
self.notify_dagrun_state_changed(msg="success")
- if execute_callbacks and dag.has_on_success_callback:
- self.handle_dag_callback(dag=cast("SDKDAG", dag), success=True, reason="success")
- elif dag.has_on_success_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=False,
- msg="success",
+ if dag.has_on_success_callback:
+ last_ti_to_run: TI | None = (
+ max(tis_for_dagrun_state, key=lambda ti: ti.end_date, default=None)
)
+ if execute_callbacks:
Review Comment:
https://github.com/apache/airflow/pull/61274#discussion_r2778072290
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2867804484
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -433,9 +433,16 @@ def on_success_callable(context):
}
dag_run = self.create_dag_run(dag=dag, task_states=initial_task_states, session=session)
- with mock.patch.object(dag_run, "handle_dag_callback") as handle_dag_callback:
+ with mock.patch.object(dag_run, "execute_dag_callbacks") as execute_dag_callbacks:
_, callback = dag_run.update_state()
- assert handle_dag_callback.mock_calls == [mock.call(dag=dag, success=True, reason="success")]
+ assert execute_dag_callbacks.mock_calls == [
Review Comment:
Added a test for deadlocked tasks.
As for timeouts, it's tested in `test_scheduler_job.py`. Let me know if I missed something.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2892647195
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1287,25 +1280,24 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="all_tasks_deadlocked")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(
- dag=cast("SDKDAG", dag),
+ if dag.has_on_failure_callback:
+ unfinished_non_schedulable = (ti for ti in unfinished.tis if ti not in set(schedulable_tis))
+ finished_task_ids = {ti.task_id for ti in finished_tis}
+ blocking_ti = next(
+ iter(
Review Comment:
This one is pending, rest lgtm
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2895849655
##########
airflow-core/docs/administration-and-deployment/logging-monitoring/callbacks.rst:
##########
@@ -53,23 +53,47 @@ Callback Types
There are six types of events that can trigger a callback:
-=========================================== ================================================================
-Name Description
-=========================================== ================================================================
-``on_success_callback`` Invoked when the :ref:`Dag succeeds <dag-run:dag-run-status>` or :ref:`task succeeds <concepts:task-instances>`.
- Available at the Dag or task level.
-``on_failure_callback`` Invoked when the task :ref:`fails <concepts:task-instances>`.
- Available at the Dag or task level.
-``on_retry_callback`` Invoked when the task is :ref:`up for retry <concepts:task-instances>`.
- Available only at the task level.
-``on_execute_callback`` Invoked right before the task begins executing.
- Available only at the task level.
-``on_skipped_callback`` Invoked when the task is :ref:`running <concepts:task-instances>` and AirflowSkipException raised.
- Explicitly it is NOT called if a task is not started to be executed because of a preceding branching
- decision in the Dag or a trigger rule which causes execution to skip so that the task execution
- is never scheduled.
- Available only at the task level.
-=========================================== ================================================================
+=========================================== ======================================================================= =================
+Name Description Availability
+=========================================== ======================================================================= =================
+``on_success_callback`` Invoked when the :ref:`Dag succeeds <dag-run:dag-run-status>` Dag or Task
+ or :ref:`task succeeds <concepts:task-instances>`.
+``on_failure_callback`` Invoked when the :ref:`Dag fails <dag-run:dag-run-status>` Dag or Task
+ or task :ref:`fails <concepts:task-instances>`.
+``on_retry_callback`` Invoked when the task is :ref:`up for retry <concepts:task-instances>`. Task
+``on_execute_callback`` Invoked right before the task begins executing. Task
+``on_skipped_callback`` Invoked when the task is :ref:`running <concepts:task-instances>` Task
+ and AirflowSkipException raised. Explicitly it is NOT called if a task
+ is not started to be executed because of a preceding branching
+ decision in the Dag or a trigger rule which causes execution
+ to skip so that the task execution is never scheduled.
+=========================================== ======================================================================= =================
+
+
+Context Mapping
+---------------
+
+A context mapping that contains runtime information about a task instance is passed to every callback.
+Full list of variables available in ``context`` are in :doc:`docs <../../templates-ref>` and `code <https://github.com/apache/airflow/blob/main/task-sdk/src/airflow/sdk/definitions/context.py>`_.
+
+
+Dag Callbacks
+^^^^^^^^^^^^^
+
+As the context mapping describes execution of a task instance, contexts passed to Dag callbacks will also contain task instance variables,
+and the task selected depends on the state of a Dag:
+
+#. On regular failure, the latest failed task is selected.
+#. On Dag run timeout, the latest started but not finished task is passed.
+#. If tasks are deadlocked, a task that should have run next but couldn't is passed.
+#. On success, the latest succeeded task is passed.
+
+It's not recommended to rely on task instance variables in Dag callbacks except for human analysis, as they reflect only partial information about the Dag's state.
+For example, a timeout may be caused by a number of stalling tasks, but only one will eventually be selected for context.
+
+.. note::
+ Before Airflow 3.1.9, the rules above did not apply and the task instance passed to Dag callback was not related to Dag state, rather being selected as the latest task in the Dag
Review Comment:
```suggestion
Before Airflow 3.2.0, the rules above did not apply and the task instance passed to Dag callback was not related to Dag state, rather being selected as the latest task in the Dag
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-4011793632
Since technically, this is a chance in behavior, let's put it in 3.2, can you add a newsfragment too please @Asquator : https://github.com/apache/airflow/tree/main/airflow-core/newsfragments
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-4018623428
static checks are failing
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2771489777
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2282,7 +2282,10 @@ def _schedule_dag_run(
select(TI)
.where(TI.dag_id == dag_run.dag_id)
.where(TI.run_id == dag_run.run_id)
- .where(TI.state.in_(State.unfinished))
+ .where(TI.state.in_(State.unfinished) | (TI.state.is_(None)))
+ ).all()
+ last_unfinished_ti = (
Review Comment:
Well, `start_date` is exactly the field I'm using here. On timeout, we assume there are unfinished tasks, which means they have no `end_date` yet.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2777430246
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
Review Comment:
Well, it's a list of failed leaf tasks whose failure indicates the the dagrun failed. I think `last_failed_tasks` may be confusing, or maybe I do not understand the rationale.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-4019651459
Awesome work, congrats on your first merged pull request! You are invited to check our [Issue Tracker](https://github.com/apache/airflow/issues) for additional contributions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-4010762133
Slightly confused about the milestone
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2867804484
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -433,9 +433,16 @@ def on_success_callable(context):
}
dag_run = self.create_dag_run(dag=dag, task_states=initial_task_states, session=session)
- with mock.patch.object(dag_run, "handle_dag_callback") as handle_dag_callback:
+ with mock.patch.object(dag_run, "execute_dag_callbacks") as execute_dag_callbacks:
_, callback = dag_run.update_state()
- assert handle_dag_callback.mock_calls == [mock.call(dag=dag, success=True, reason="success")]
+ assert execute_dag_callbacks.mock_calls == [
Review Comment:
Added a test for deadlocked tasks.
As for timeouts, it's checked in `test_scheduler_job.py`. Let me know if I missed something.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3977842490
It's not just from the merge. I just used a shortcut to squash commits with `reset --soft` and then just pushed all the changes. As main advanced since then, it looked like I reverted all those changes...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2894919963
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1416,27 +1408,40 @@ def notify_dagrun_state_changed(self, msg: str):
# we can't get all the state changes on SchedulerJob,
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
- @provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
- tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
- if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
+ def produce_dag_callback(
+ self,
+ dag: SerializedDAG,
+ success: bool = True,
+ relevant_ti: TI | None = None,
+ reason: str = "success",
+ execute=False,
Review Comment:
Fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3977759416
> But yeah. What works better is not to merge but rebase your changes - if you merge, your change contains all the changed files from already merged changes, when you rebase, you have clean commits from your changes only.
>
> We strongly recommend using rebase.
Interesting, I always merge and it never does that, which is what the 'update branch' button does, I think the squash was wrong, as it looks like all rebased changes were squashed into 1 commit
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2836546068
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1212,22 +1212,29 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="task_failure")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(dag=cast("SDKDAG", dag), success=False, reason="task_failure")
- elif dag.has_on_failure_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=True,
- msg="task_failure",
- )
+ if dag.has_on_failure_callback:
+ ti_causing_failure = self.get_first_ti_causing_failure(dag=dag, session=session)
+ if execute_callbacks:
Review Comment:
Looks way better and simpler
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2879755565
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1416,27 +1408,40 @@ def notify_dagrun_state_changed(self, msg: str):
# we can't get all the state changes on SchedulerJob,
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
- @provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
- tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
- if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
+ def produce_dag_callback(
+ self,
+ dag: SerializedDAG,
+ success: bool = True,
+ relevant_ti: TI | None = None,
+ reason: str = "success",
+ execute=False,
Review Comment:
```suggestion
execute: bool = False,
```
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1287,25 +1280,24 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="all_tasks_deadlocked")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(
- dag=cast("SDKDAG", dag),
+ if dag.has_on_failure_callback:
+ unfinished_non_schedulable = (ti for ti in unfinished.tis if ti not in set(schedulable_tis))
+ finished_task_ids = {ti.task_id for ti in finished_tis}
+ blocking_ti = next(
+ iter(
Review Comment:
`iter()` around a generator expression is redundant — generators are already iterators. `next(genexpr, None)` works directly.
Also, in the deadlock branch (`unfinished.should_schedule and not are_runnable_tasks`), `schedulable_tis` is guaranteed to be empty (since `are_runnable_tasks` is `False`), making the `set(schedulable_tis)` filter on line 1284 a no-op. You could simplify `unfinished_non_schedulable` to just `unfinished.tis`. If you want to keep it as defensive code, at least pre-compute the set outside the generator to avoid recreating it per element:
```python
schedulable_set = set(schedulable_tis)
unfinished_non_schedulable = (ti for ti in unfinished.tis if ti not in schedulable_set)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2807965193
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
Review Comment:
This logic was removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2771502039
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
Review Comment:
It's a shortcut to avoid the logic below.
In the end, I use the `min` function that doesn't operate on empty collections. the `default` argument catches that, and yet there's no reason to scan the DAG if for some reason we have no failed tasks. I think I should add a warning there, because this function shouldn't be called on non-failed DAGs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2778062045
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
Review Comment:
> We iterate over the leaves of the DAG that define its failure, and go up until we find the first failed task. First we collect all the failed upstream dependencies of the leaves and then do the sorting. `failure_path_task_ids` is the set of all upstream dependencies.
This can be simplified, whatvyou are doing here is filtering for upstream failed, what if we just select failed tasks?
A downstream task will only be marked as failed of the correct trigger rule is set and the direct upstream failed, this whole loop may be removed in my opinion
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2778016236
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2283,7 +2283,10 @@ def _schedule_dag_run(
select(TI)
.where(TI.dag_id == dag_run.dag_id)
.where(TI.run_id == dag_run.run_id)
- .where(TI.state.in_(State.unfinished))
+ .where(TI.state.in_(State.unfinished) | (TI.state.is_(None)))
Review Comment:
Seems good, checked and verified it myself
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3849587044
I want to get some feedback on the behavior, whether it's the desired one or if anything can be made better.
I will update the docs to explain the behavior of passed TIs to dag callbacks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2842597816
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1248,22 +1262,35 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.SUCCESS)
self.notify_dagrun_state_changed(msg="success")
- if execute_callbacks and dag.has_on_success_callback:
- self.handle_dag_callback(dag=cast("SDKDAG", dag), success=True, reason="success")
- elif dag.has_on_success_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=False,
- msg="success",
+ if dag.has_on_success_callback:
+ last_succeeded_ti: TI | None = (
+ max(
+ (ti for ti in tis_for_dagrun_state if ti.state == TaskInstanceState.SUCCESS),
+ key=lambda ti: ti.end_date,
+ default=None
+ )
)
+ if execute_callbacks:
+ self.handle_dag_callback(
+ dag=cast("SDKDAG", dag),
+ success=True,
+ relevant_ti=last_succeeded_ti,
+ reason="success"
+ )
+ else:
+ callback = DagCallbackRequest(
+ filepath=self.dag_model.relative_fileloc,
+ dag_id=self.dag_id,
+ run_id=self.run_id,
+ bundle_name=self.dag_model.bundle_name,
+ bundle_version=self.bundle_version,
+ context_from_server=DagRunContext(
+ dag_run=self,
+ last_ti=last_succeeded_ti,
+ ),
+ is_failure_callback=False,
+ msg="success",
+ )
Review Comment:
Introduced `produce_dag_callback` as a template method.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1283,26 +1310,31 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="all_tasks_deadlocked")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(
- dag=cast("SDKDAG", dag),
- success=False,
- reason="all_tasks_deadlocked",
- )
- elif dag.has_on_failure_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=True,
- msg="all_tasks_deadlocked",
+ if dag.has_on_failure_callback:
+ last_finished_ti: TI | None = (
+ max(info.finished_tis, key=lambda ti: ti.end_date, default=None)
)
+ if execute_callbacks:
+ self.handle_dag_callback(
+ dag=cast("SDKDAG", dag),
+ success=False,
+ relevant_ti=last_finished_ti,
+ reason="all_tasks_deadlocked",
+ )
+ else:
+ callback = DagCallbackRequest(
+ filepath=self.dag_model.relative_fileloc,
+ dag_id=self.dag_id,
+ run_id=self.run_id,
+ bundle_name=self.dag_model.bundle_name,
+ bundle_version=self.bundle_version,
+ context_from_server=DagRunContext(
+ dag_run=self,
+ last_ti=last_finished_ti
+ ),
+ is_failure_callback=True,
+ msg="all_tasks_deadlocked",
+ )
Review Comment:
Introduced `produce_dag_callback` as a template method.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3977738401
> Something got wrong, I did not intend to request these reviews. Just reset soft and pushed to squash many merge commits.
No. Nothing wrong. CODEOWNERS decide who gets asked for review when certain files change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3961729312
> > I think you need to set up the pre-commit hooks locally and then run the command:
> > prek run --all-files
Yes, I just encountered a problem where some files in main fail some of the mypy checks, probably related to https://github.com/apache/airflow/issues/62272
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3946874218
Renamed `handle_dag_callback` to `execute_dag_callbacks`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2776788765
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1287,6 +1290,9 @@ def recalculate(self) -> _UnfinishedStates:
reason="all_tasks_deadlocked",
)
elif dag.has_on_failure_callback:
+ last_finished_ti: TI | None = (
+ max(info.finished_tis, key=lambda ti: ti.end_date, default=None)
Review Comment:
We do this check once for a finished DAG. Wdym?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2778054981
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
Review Comment:
Btw, isn't the None issue also persisting in here?
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
Review Comment:
Btw, isn't the None issue also persisting here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2778072973
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1245,22 +1252,31 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.SUCCESS)
self.notify_dagrun_state_changed(msg="success")
- if execute_callbacks and dag.has_on_success_callback:
- self.handle_dag_callback(dag=cast("SDKDAG", dag), success=True, reason="success")
- elif dag.has_on_success_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=False,
- msg="success",
+ if dag.has_on_success_callback:
+ last_ti_to_run: TI | None = (
+ max(tis_for_dagrun_state, key=lambda ti: ti.end_date, default=None)
)
+ if execute_callbacks:
Review Comment:
Same as the comment above
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1280,26 +1296,31 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="all_tasks_deadlocked")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(
- dag=cast("SDKDAG", dag),
- success=False,
- reason="all_tasks_deadlocked",
- )
- elif dag.has_on_failure_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=True,
- msg="all_tasks_deadlocked",
+ if dag.has_on_failure_callback:
+ last_finished_ti: TI | None = (
+ max(info.finished_tis, key=lambda ti: ti.end_date, default=None)
)
+ if execute_callbacks:
Review Comment:
Same here
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1212,22 +1212,29 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="task_failure")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(dag=cast("SDKDAG", dag), success=False, reason="task_failure")
- elif dag.has_on_failure_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=True,
- msg="task_failure",
- )
+ if dag.has_on_failure_callback:
+ ti_causing_failure = self.get_first_ti_causing_failure(dag=dag, session=session)
+ if execute_callbacks:
Review Comment:
Why isn't this checked in the upper if statement? It can be less nested while remaining readable, the else just becomes an else if for the upper if
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2775447876
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
+
+ return min(failed_on_paths, key=lambda ti: ti.end_date, default=None)
Review Comment:
> Why shouldn't we return multiple tasks if a few of them were running concurrently?
Unfortunately, the current `Context` has to provide at least one TI under the `"ti"` key, even on DAG failure callbacks. I don't know whether it's good or bad, probably bad, but I decided to adhere to this rule for back comp. I think the context user can still access other tasks in the DAG, can't he?
> And why filter by end_date and not start_date?
A good point. Logically, I wanted to return the task that ended last and could cause the DAG to fail. On the second thought, maybe it will be wiser to filter by `start_date` since many tasks will be failed manually by the DagRun.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
+
+ return min(failed_on_paths, key=lambda ti: ti.end_date, default=None)
Review Comment:
> Why shouldn't we return multiple tasks if a few of them were running concurrently?
Unfortunately, the current `Context` has to provide at least one TI under the `"ti"` key, even on DAG failure callbacks. I don't know whether it's good or bad, probably bad, but I decided to adhere to this rule for back comp. I think the context user can still access other tasks in the DAG, can't he?
> And why filter by end_date and not start_date?
A good point. Logically, I wanted to return the task that ended last and could cause the DAG to fail. On the second thought, maybe it will be wiser to filter by `start_date` since many tasks will be failed manually by the DagRun.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2777422835
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
Review Comment:
We iterate over the leaves of the DAG that define its failure, and go up until we find the first failed task. First we collect all the failed upstream dependencies of the leaves and then do the sorting. `failure_path_task_ids` is the set of all upstream dependencies.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2807816813
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1212,22 +1212,29 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="task_failure")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(dag=cast("SDKDAG", dag), success=False, reason="task_failure")
- elif dag.has_on_failure_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=True,
- msg="task_failure",
- )
+ if dag.has_on_failure_callback:
+ ti_causing_failure = self.get_first_ti_causing_failure(dag=dag, session=session)
+ if execute_callbacks:
Review Comment:
It's what was being done before the change. After the change I included another statement that is needed in both cases, whether `execute_callbacks` is true or not. Look at it another time, maybe it will be clearer now?
Or can you suggest a change so I understand what you mean?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3977741622
But yeah. What works better is not to merge but rebase your changes - if you merge, your change contains all the changed files from already merged changes, when you rebase, you have clean commits from your changes only.
We strongly recommend using rebase.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-4019100213
The test is failing because my newsfragment has multiple lines, and it can't...
Do I change the type to `significant`, or avoid the elaborate explanations by including just one general sentence?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2767500941
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
Review Comment:
Why is the "if" necessary?
Maybe we can just return an empty array at the end if the check yields no results?
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
+
+ return min(failed_on_paths, key=lambda ti: ti.end_date, default=None)
Review Comment:
Why shouldn't we return multiple tasks if a few of them were running concurrently?
And why filter by end_date and not start_date?
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1287,6 +1290,9 @@ def recalculate(self) -> _UnfinishedStates:
reason="all_tasks_deadlocked",
)
elif dag.has_on_failure_callback:
+ last_finished_ti: TI | None = (
+ max(info.finished_tis, key=lambda ti: ti.end_date, default=None)
Review Comment:
Why do we make this check over and over again?
Can't we do it once?
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2282,7 +2282,10 @@ def _schedule_dag_run(
select(TI)
.where(TI.dag_id == dag_run.dag_id)
.where(TI.run_id == dag_run.run_id)
- .where(TI.state.in_(State.unfinished))
+ .where(TI.state.in_(State.unfinished) | (TI.state.is_(None)))
+ ).all()
+ last_unfinished_ti = (
Review Comment:
What if it was a mapped task? Or what if it was concurrently running tasks, how do you decide?
As the end date is not the only factor, the start date may be a better option here, or maybe even return a few tasks if they were running concurrently, and check by the dependencies of the task rather than solely rely on end date
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
Review Comment:
This part confuses me a little, why do we get ALL task instances (`tis`), just to get all of the previous task instances from the failed tasks (`failure_path_task_ids`), just to get the failed tasks instances again?
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
Review Comment:
So leaf here is meant as if the leaf task for a given dagrun?
If so maybe a better name could be `last_failed_tasks`
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
Review Comment:
Minor nit but aybe the name of the method can be changed to `get_first_ti_causing_dagrun_failure`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2775447876
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
+
+ return min(failed_on_paths, key=lambda ti: ti.end_date, default=None)
Review Comment:
> Why shouldn't we return multiple tasks if a few of them were running concurrently?
Unfortunately, the current `Context` has to provide at least one TI under the `"ti"` key, even on DAG failure callbacks. I don't know whether it's good or bad, probably bad, but I decided to adhere to this rule for back comp. I think the context user can still access other tasks in the DAG, can't he?
> And why filter by end_date and not start_date?
A good point. Logically, I wanted to return the task that ended last and could cause the DAG to fail. On the second thought, maybe it will be wiser to filter by `start_date` since many tasks will be failed manually after the DAG fails.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2778016771
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2282,7 +2282,10 @@ def _schedule_dag_run(
select(TI)
.where(TI.dag_id == dag_run.dag_id)
.where(TI.run_id == dag_run.run_id)
- .where(TI.state.in_(State.unfinished))
+ .where(TI.state.in_(State.unfinished) | (TI.state.is_(None)))
+ ).all()
+ last_unfinished_ti = (
Review Comment:
MIght have looked at a different line before commenting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2778054981
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
Review Comment:
Btw, isn't the None issue also persisting here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2778054981
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
Review Comment:
Btw, isn't the None issue also persisting here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2778055310
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
Review Comment:
Btw, isn't the None issue also persisting here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3954157476
Going to fix lint errors soon
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal merged PR #61274:
URL: https://github.com/apache/airflow/pull/61274
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2842677923
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2311,7 +2311,10 @@ def _schedule_dag_run(
select(TI)
.where(TI.dag_id == dag_run.dag_id)
.where(TI.run_id == dag_run.run_id)
- .where(TI.state.in_(State.unfinished))
+ .where(TI.state.in_(State.unfinished) | (TI.state.is_(None)))
+ ).all()
+ last_unfinished_ti = (
+ max(unfinished_task_instances, key=lambda ti: ti.start_date, default=None)
Review Comment:
Applied, thank you
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1215,22 +1215,36 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="task_failure")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(dag=cast("SDKDAG", dag), success=False, reason="task_failure")
- elif dag.has_on_failure_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=True,
- msg="task_failure",
+ if dag.has_on_failure_callback:
+ ti_causing_failure = (
+ max(
+ (ti for ti in tis_for_dagrun_state
+ if ti.state == TaskInstanceState.FAILED and ti.end_date is not None),
+ key=lambda ti: ti.end_date,
+ default=None
+ )
Review Comment:
Applied, thank you
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-3992368525
Would happily get some more review of the logic before the merge, just to make sure it works as expected, as it's four different cases we handle.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2877487167
##########
airflow-core/docs/administration-and-deployment/logging-monitoring/callbacks.rst:
##########
@@ -53,23 +53,47 @@ Callback Types
There are six types of events that can trigger a callback:
-=========================================== ================================================================
-Name Description
-=========================================== ================================================================
-``on_success_callback`` Invoked when the :ref:`Dag succeeds <dag-run:dag-run-status>` or :ref:`task succeeds <concepts:task-instances>`.
- Available at the Dag or task level.
-``on_failure_callback`` Invoked when the task :ref:`fails <concepts:task-instances>`.
- Available at the Dag or task level.
-``on_retry_callback`` Invoked when the task is :ref:`up for retry <concepts:task-instances>`.
- Available only at the task level.
-``on_execute_callback`` Invoked right before the task begins executing.
- Available only at the task level.
-``on_skipped_callback`` Invoked when the task is :ref:`running <concepts:task-instances>` and AirflowSkipException raised.
- Explicitly it is NOT called if a task is not started to be executed because of a preceding branching
- decision in the Dag or a trigger rule which causes execution to skip so that the task execution
- is never scheduled.
- Available only at the task level.
-=========================================== ================================================================
+=========================================== ======================================================================= =================
+Name Description Availability
+=========================================== ======================================================================= =================
+``on_success_callback`` Invoked when the :ref:`Dag succeeds <dag-run:dag-run-status>` Dag or Task
+ or :ref:`task succeeds <concepts:task-instances>`.
+``on_failure_callback`` Invoked when the :ref:`Dag fails <dag-run:dag-run-status>` Dag or Task
+ or task :ref:`fails <concepts:task-instances>`.
+``on_retry_callback`` Invoked when the task is :ref:`up for retry <concepts:task-instances>`. Task
+``on_execute_callback`` Invoked right before the task begins executing. Task
+``on_skipped_callback`` Invoked when the task is :ref:`running <concepts:task-instances>` Task
+ and AirflowSkipException raised. Explicitly it is NOT called if a task
+ is not started to be executed because of a preceding branching
+ decision in the Dag or a trigger rule which causes execution
+ to skip so that the task execution is never scheduled.
+=========================================== ======================================================================= =================
+
+
+Context Mapping
+---------------
+
+A context mapping that contains runtime information about a task instance is passed to every callback.
+Full list of variables available in ``context`` are in :doc:`docs <../../templates-ref>` and `code <https://github.com/apache/airflow/blob/main/task-sdk/src/airflow/sdk/definitions/context.py>`_.
+
+
+Dag Callbacks
+^^^^^^^^^^^^^
+
+As the context mapping describes execution of a task instance, contexts passed to Dag callbacks will also contain task instance variables,
+and the task selected depends on the state of a Dag:
+
+#. On regular failure, the latest failed task is selected.
+#. On Dag run timeout, the latest started but not finished task is passed.
+#. If tasks are deadlocked, a task that should have run next but couldn't is passed.
+#. On success, the latest succeeded task is passed.
+
+It's not recommended to rely on task instance variables in Dag callbacks except for human analysis, as they reflect only partial information about the Dag's state.
+For example, a timeout may be caused by a number of stalling tasks, but only one will eventually be selected for context.
+
+.. note::
+ Before Airflow 3.1.9, the rules above did not apply and the task instance passed to Dag callback was not related to Dag state, rather being selected as the latest task in the Dag
Review Comment:
cc @vatsrahul1001 if this PR is merged before cut of 3.1.8 we'll need to change this entry just before mrege
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2807960289
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
Review Comment:
The thing I wanted to do is to return the latest failed task that actually caused the DAG run to fail.
The simplest way is just returning the latest failed task in the DAG run. I switched to this logic now, but we have to acknowledge some edge cases where a task is marked as failed just after the DAG run fails and this task is sent to the callback.
The previous logic wouldn't catch the case too, so I removed it. Do you think we should consider this edge case?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2894878772
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1287,25 +1280,24 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="all_tasks_deadlocked")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(
- dag=cast("SDKDAG", dag),
+ if dag.has_on_failure_callback:
+ unfinished_non_schedulable = (ti for ti in unfinished.tis if ti not in set(schedulable_tis))
+ finished_task_ids = {ti.task_id for ti in finished_tis}
+ blocking_ti = next(
+ iter(
Review Comment:
Fixed both, thank you for noticing!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on PR #61274:
URL: https://github.com/apache/airflow/pull/61274#issuecomment-4010991796
> Slightly confused about the milestone
I'll set it for 3.1.9 but this is subject to confirmation from the release manager of 3.1.9
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2836552762
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1248,22 +1262,35 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.SUCCESS)
self.notify_dagrun_state_changed(msg="success")
- if execute_callbacks and dag.has_on_success_callback:
- self.handle_dag_callback(dag=cast("SDKDAG", dag), success=True, reason="success")
- elif dag.has_on_success_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=False,
- msg="success",
+ if dag.has_on_success_callback:
+ last_succeeded_ti: TI | None = (
+ max(
+ (ti for ti in tis_for_dagrun_state if ti.state == TaskInstanceState.SUCCESS),
+ key=lambda ti: ti.end_date,
+ default=None
+ )
)
+ if execute_callbacks:
+ self.handle_dag_callback(
+ dag=cast("SDKDAG", dag),
+ success=True,
+ relevant_ti=last_succeeded_ti,
+ reason="success"
+ )
+ else:
+ callback = DagCallbackRequest(
+ filepath=self.dag_model.relative_fileloc,
+ dag_id=self.dag_id,
+ run_id=self.run_id,
+ bundle_name=self.dag_model.bundle_name,
+ bundle_version=self.bundle_version,
+ context_from_server=DagRunContext(
+ dag_run=self,
+ last_ti=last_succeeded_ti,
+ ),
+ is_failure_callback=False,
+ msg="success",
+ )
Review Comment:
This looks exactly like what is above, maybe it can be extracted to a template method or something similar
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1283,26 +1310,31 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.FAILED)
self.notify_dagrun_state_changed(msg="all_tasks_deadlocked")
- if execute_callbacks and dag.has_on_failure_callback:
- self.handle_dag_callback(
- dag=cast("SDKDAG", dag),
- success=False,
- reason="all_tasks_deadlocked",
- )
- elif dag.has_on_failure_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=True,
- msg="all_tasks_deadlocked",
+ if dag.has_on_failure_callback:
+ last_finished_ti: TI | None = (
+ max(info.finished_tis, key=lambda ti: ti.end_date, default=None)
)
+ if execute_callbacks:
+ self.handle_dag_callback(
+ dag=cast("SDKDAG", dag),
+ success=False,
+ relevant_ti=last_finished_ti,
+ reason="all_tasks_deadlocked",
+ )
+ else:
+ callback = DagCallbackRequest(
+ filepath=self.dag_model.relative_fileloc,
+ dag_id=self.dag_id,
+ run_id=self.run_id,
+ bundle_name=self.dag_model.bundle_name,
+ bundle_version=self.bundle_version,
+ context_from_server=DagRunContext(
+ dag_run=self,
+ last_ti=last_finished_ti
+ ),
+ is_failure_callback=True,
+ msg="all_tasks_deadlocked",
+ )
Review Comment:
Again, looks exactly the same multiple times and probably can be extracted
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2836694097
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1248,22 +1262,35 @@ def recalculate(self) -> _UnfinishedStates:
self.set_state(DagRunState.SUCCESS)
self.notify_dagrun_state_changed(msg="success")
- if execute_callbacks and dag.has_on_success_callback:
- self.handle_dag_callback(dag=cast("SDKDAG", dag), success=True, reason="success")
- elif dag.has_on_success_callback:
- callback = DagCallbackRequest(
- filepath=self.dag_model.relative_fileloc,
- dag_id=self.dag_id,
- run_id=self.run_id,
- bundle_name=self.dag_model.bundle_name,
- bundle_version=self.bundle_version,
- context_from_server=DagRunContext(
- dag_run=self,
- last_ti=self.get_last_ti(dag=dag, session=session),
- ),
- is_failure_callback=False,
- msg="success",
+ if dag.has_on_success_callback:
+ last_succeeded_ti: TI | None = (
+ max(
+ (ti for ti in tis_for_dagrun_state if ti.state == TaskInstanceState.SUCCESS),
+ key=lambda ti: ti.end_date,
+ default=None
+ )
)
+ if execute_callbacks:
+ self.handle_dag_callback(
+ dag=cast("SDKDAG", dag),
+ success=True,
+ relevant_ti=last_succeeded_ti,
+ reason="success"
+ )
+ else:
+ callback = DagCallbackRequest(
+ filepath=self.dag_model.relative_fileloc,
+ dag_id=self.dag_id,
+ run_id=self.run_id,
+ bundle_name=self.dag_model.bundle_name,
+ bundle_version=self.bundle_version,
+ context_from_server=DagRunContext(
+ dag_run=self,
+ last_ti=last_succeeded_ti,
+ ),
+ is_failure_callback=False,
+ msg="success",
+ )
Review Comment:
I think you're right. I'm going to introduce a template method for the shared logic, thanks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2748629648
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2283,7 +2283,10 @@ def _schedule_dag_run(
select(TI)
.where(TI.dag_id == dag_run.dag_id)
.where(TI.run_id == dag_run.run_id)
- .where(TI.state.in_(State.unfinished))
+ .where(TI.state.in_(State.unfinished) | (TI.state.is_(None)))
Review Comment:
I thought that tasks with `null` state should also be always marked as `skipped` (the `in_` clause won't fetch them because of SQL semantics).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2767482328
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2283,7 +2283,10 @@ def _schedule_dag_run(
select(TI)
.where(TI.dag_id == dag_run.dag_id)
.where(TI.run_id == dag_run.run_id)
- .where(TI.state.in_(State.unfinished))
+ .where(TI.state.in_(State.unfinished) | (TI.state.is_(None)))
Review Comment:
None is in unfinished states, so there is no need to change anything [here](https://airflow.apache.org/docs/apache-airflow/stable/_modules/airflow/utils/state.html#State.unfinished)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2777430246
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
Review Comment:
Well, it's a list of failed leaf tasks. I think `last_failed_tasks` may be confusing, or maybe I do not understand the rationale.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Always provide a relevant TI context in Dag callback [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #61274:
URL: https://github.com/apache/airflow/pull/61274#discussion_r2778068043
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1410,24 +1416,44 @@ def notify_dagrun_state_changed(self, msg: str):
# or LocalTaskJob, so we don't want to "falsely advertise" we notify about that
@provide_session
- def get_last_ti(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
- """Get Last TI from the dagrun to build and pass Execution context object from server to then run callbacks."""
+ def get_first_ti_causing_failure(self, dag: SerializedDAG, session: Session = NEW_SESSION) -> TI | None:
+ """
+ Get the first task instance that would cause a leaf task to fail the run.
+ """
+
tis = self.get_task_instances(session=session)
- # tis from a dagrun may not be a part of dag.partial_subset,
- # since dag.partial_subset is a subset of the dag.
- # This ensures that we will only use the accessible TI
- # context for the callback.
+
+ failed_leaf_tis = [
+ ti for ti in self._tis_for_dagrun_state(dag=dag, tis=tis)
+ if ti.state in State.failed_states
+ ]
+
+ if not failed_leaf_tis:
+ return None
+
if dag.partial:
- tis = [ti for ti in tis if not ti.state == State.NONE]
- # filter out removed tasks
- tis = natsorted(
- (ti for ti in tis if ti.state != TaskInstanceState.REMOVED),
- key=lambda ti: ti.task_id,
- )
- if not tis:
- return None
- ti = tis[-1] # get last TaskInstance of DagRun
- return ti
+ tis = [
+ ti for ti in tis if not ti.state in (
+ State.NONE, TaskInstanceState.REMOVED
+ )
+ ]
+
+ # Collect all task IDs on failure paths
+ failure_path_task_ids = set()
+ for failed_leaf in failed_leaf_tis:
+ leaf_task = dag.get_task(failed_leaf.task_id)
+ upstream_ids = leaf_task.get_flat_relative_ids(upstream=True)
+ failure_path_task_ids.update(upstream_ids)
+ failure_path_task_ids.add(failed_leaf.task_id)
+
+ # Find failed tasks on possible failure paths
+ failed_on_paths = [
+ ti for ti in tis
+ if ti.task_id in failure_path_task_ids and ti.state == State.FAILED
+ ]
Review Comment:
My point is, one of the last failed tasks is the one that you need, and you do not need to get all task instances of a dag for that
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org