dheerajturaga opened a new pull request, #68198:
URL: https://github.com/apache/airflow/pull/68198
Cluster policies could mutate a TaskInstance before scheduling but had no
supported way to read the DagRun it belongs to -- so routing a task to a queue
based on the run's configuration was impossible. Reaching the DagRun via
TaskInstance.get_dagrun() opens a nested committing session that trips the
scheduler's prohibit_commit guard and crashes the scheduler.
Add an optional dag_run argument to the task_instance_mutation_hook policy
hook and thread the in-scope DagRun through every scheduler-side call site (no
extra database access). A hook can now route on dag_run.conf, for example:
def task_instance_mutation_hook(task_instance, dag_run=None):
if dag_run is not None and (dag_run.conf or {}).get("route") ==
"high":
task_instance.queue = "high_priority_queue"
Backward compatibility is preserved through the local-settings plugin shim.
Pluggy passes a hookimpl only the parameters it declares without a default, so
this commit also makes make_plugin_from_local_settings register a function
as-is when its parameters are a subset of the hookspec's, and shim it
(forwarding the call positionally) when it renames a parameter or gives a
hookspec parameter a default value. As a result every signature keeps working:
the legacy single-argument hook, an explicit (task_instance, dag_run) hook, and
the ergonomic (task_instance, dag_run=None) hook all behave correctly; a hook
declaring an unknown parameter is still rejected.
The dag_run may be None during early task-instance construction (when the
hook is re-applied from refresh_from_task before the instance is bound to a
run), and the hook must not open a new session; both are documented on the
hookspec and in the cluster-policies guide, which also corrects the prior claim
that the hook runs on the worker.
---
##### 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)
ClaudeCode Opus 4.8
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
