Skip to content

Pass Dag runs directly to Deadline references - #70714

Open
imrichardwu wants to merge 19 commits into
apache:mainfrom
imrichardwu:fix/deadline-reference-context
Open

Pass Dag runs directly to Deadline references#70714
imrichardwu wants to merge 19 commits into
apache:mainfrom
imrichardwu:fix/deadline-reference-context

Conversation

@imrichardwu

Copy link
Copy Markdown
Contributor

Preserve the Dag run evaluation context for custom references that accept additional keyword arguments, so they continue to evaluate correctly.

Pass the in-memory Dag run to Deadline references so built-in references avoid redundant database reads and custom references can evaluate the current state.

Deprecate required_kwargs while preserving it for existing custom references. Use the shared evaluator when recalculating queued-at deadlines.

closes: #70710


Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

  • Read the 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) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@imrichardwu
imrichardwu marked this pull request as draft July 30, 2026 05:29
@imrichardwu
imrichardwuforce-pushed the fix/deadline-reference-context branch 2 times, most recently from e9c48fd to 3b99b93CompareJuly 30, 2026 05:33
@imrichardwu
imrichardwu marked this pull request as ready for review July 30, 2026 05:33
@imrichardwu

Copy link
Copy Markdown
ContributorAuthor

@ferruzzi can I get a review?

Preserve the Dag run evaluation context for custom references that accept
additional keyword arguments, so they continue to evaluate correctly.
@imrichardwu
imrichardwuforce-pushed the fix/deadline-reference-context branch from 3b99b93 to db15212CompareJuly 31, 2026 01:39
Comment threadairflow-core/docs/howto/deadline-alerts.rst Outdated
Comment threadairflow-core/src/airflow/models/deadline.py Outdated
Comment threadairflow-core/src/airflow/models/deadline.py
Comment threadairflow-core/src/airflow/models/deadline.py Outdated
imrichardwuand others added 6 commits August 6, 2026 20:55
Co-authored-by: D. Ferruzzi <ferruzzi@amazon.com>
Co-authored-by: D. Ferruzzi <ferruzzi@amazon.com>
Co-authored-by: D. Ferruzzi <ferruzzi@amazon.com>
 The custom deadline reference examples returned the dagrun attribute
directly, which read as if the attribute was the required return value
rather than an input available to the author's business logic. Type the
evaluation dagrun with the shared DagRunProtocol so authors get a
concrete contract for what is available, and show the attribute feeding
into business logic in the examples.
A custom deadline reference whose _evaluate_with declares **kwargs is
expected to receive the full evaluation context, matching the standard
Python "accept anything" convention rather than having the context
filtered away. Align the test with that behaviour.
Comment threadtask-sdk/src/airflow/sdk/definitions/deadline.py Outdated
Comment threadairflow-core/src/airflow/serialization/definitions/deadline.py Outdated
Comment threadairflow-core/docs/howto/deadline-alerts.rst
Clarify the supported Dag run interface for deadline references and document the migration away from required_kwargs before Airflow 4.0.
@github-actions

github-actionsBot commented Aug 13, 2026

Copy link
Copy Markdown
Contributor

uv.lock on main just moved via #71520 ("Restore the evaluation extra of google-cloud-aiplatform"), commit b69253b and this PR currently conflicts.

Quickest fix:

git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-lease

Automated nudge — ignore if you're not ready to rebase. This comment is updated in place on future uv.lock bumps.

Updated the DeadlineDagRunProtocol class to eliminate its inheritance from DagRunProtocol, simplifying the interface for deadline references. The class now directly defines the necessary attributes for evaluating deadlines without relying on the ORM DagRun structure.
Comment threadairflow-core/src/airflow/models/deadline.py
Comment threadairflow-core/src/airflow/models/taskinstance.py Outdated
Comment threadairflow-core/tests/unit/models/test_taskinstance.py Outdated
Comment threadairflow-core/tests/unit/models/test_taskinstance.py Outdated
@ferruzzi

Copy link
Copy Markdown
Contributor

Heads up on sequencing: this PR now overlaps two others in _recalculate_dagrun_queued_at_deadlines.

This PR fixes the same TypeError incidentally, and I want to be explicit that that does not make #70370 redundant. This is a refactor with a significant.rst newsfragment and a deprecation, so it is not backport material; #70370 is the only route the fix has to 3.3.x, where it currently affects every queued-at deadline. Please do not let this PR be used as a reason to close that one.

The cheapest merge order looks like it will be to merge #70370 first, then rebase here. #70370 adds decode_deadline_alert_model() and resolve_deadline_alert_interval() to serialization/decoders.py, which is exactly what this PR hand-rolls at taskinstance.py:255-260. Rebasing onto it lets you delete those lines.

@ferruzzi

Copy link
Copy Markdown
Contributor

Adding a fourth PR to the picture since it breaks something here concretely.

#71802 adds a core-side SerializedVariableInterval, makes decode_deadline_alert return that instead of VariableInterval, and deletes resolve() from the SDK entirely.

That will break test_clear_task_instances_recalculates_variable_interval_deadline, which patches airflow.sdk.definitions.deadline.Variable.get. Worth knowing now rather than after a rebase.

Also worth noting, #71850's diff still has required_kwargs: set[str] = set() on
SerializedBaseDeadlineReference, which this PR removes, and it adds an evaluation_timing class attribute immediately below it. So be careful when rebasing, whichever lands first.

…evaluations
This update introduces warning logs in the SerializedReferenceModels for cases where a DagRun's logical date or queued at time is missing. Additionally, the test suite has been enhanced to verify that these warnings are correctly logged when the respective attributes are None.

@ferruzziferruzzi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Final nit, sorry.

Comment threadairflow-core/tests/unit/models/test_deadline.py Outdated
@ferruzziferruzzi added this to the Airflow 3.3.2 milestone Aug 28, 2026
@ferruzziferruzzi removed this from the Airflow 3.3.2 milestone Aug 29, 2026
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Deadline reference evaluation should receive the DagRun instead of dag_id/run_id

3 participants

@imrichardwu@ferruzzi@Prab-27