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 2025/07/18 13:11:36 UTC
[PR] Query skeleton for window-pessimistic task instance fetching [airflow]
Asquator opened a new pull request, #53492:
URL: https://github.com/apache/airflow/pull/53492
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! 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 committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2252064905
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_add_ti_max_active_tis_per_dagrun.py:
##########
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add callback_state to deadline.
+
+Revision ID: 2f49f2dae90c
+Revises: f56f68b9e02f
+Create Date: 2025-07-28 16:39:01.181132
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2f49f2dae90c"
+down_revision = "f56f68b9e02f"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+ """Add callback_state to deadline."""
+ with op.batch_alter_table("task_instance", schema=None) as batch_op:
+ batch_op.add_column(sa.Column("max_active_tis_per_dag", sa.Integer, nullable=True))
Review Comment:
but it seems pretty uncontroversial that it doesn't need to be on task instance table given that, right...?
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_add_ti_max_active_tis_per_dagrun.py:
##########
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add callback_state to deadline.
+
+Revision ID: 2f49f2dae90c
+Revises: f56f68b9e02f
+Create Date: 2025-07-28 16:39:01.181132
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2f49f2dae90c"
+down_revision = "f56f68b9e02f"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+ """Add callback_state to deadline."""
+ with op.batch_alter_table("task_instance", schema=None) as batch_op:
+ batch_op.add_column(sa.Column("max_active_tis_per_dag", sa.Integer, nullable=True))
Review Comment:
wouldn't this info always be the same for a given dag run? why not store it on the dag run table. you could also argue for it being on the dag but maybe you want to preserve ability for things to change per dag run.
but it seems pretty uncontroversial that it doesn't need to be on task instance table given that, right...?
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3168943144
> @dstandish
>
> > This query is not a starvation panacea, I don't think.
> > For example what if there are two conditions interplaying.
> > Suppose there's only one pool slot available in pool A. Tasks X and Y have pool A. Task X sorts first. So task Y would be filtered out by the pool constraint.
>
> If task x sorts first it means its at a higher priority and needs to be run first, which is expected behaviour, I may have not fully understood what the problem here, and I would appreciate it if you could clarify further
>
> > But suppose further that for task X, total_tis_per_dagrun_count is beyond dag max active tasks. So nothing will be scheduled for pool A even though Y should be scheduled.
>
@Nataneljpwd
> That is a good point, however, I don't think this case can happen as the query won't allow more than dag Max active tasks to be run for a specific day, correct me if I'm wrong and if you think I didn't undestand. In this case task x is filtered out and task y is checked
The fear is that row_number() will be less than #pool slots, so Y won't be even looked at. I think we whether must nest or find another way to cope with the situation.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3176223969
Hi in your calculation of `tis_per_dag_count` you are partitioning by dag and task but i think maybe it's supposed to be just 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] Query skeleton for window-pessimistic task instance fetching [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3089475085
Solves:
https://github.com/apache/airflow/issues/45636
TODO:
- Add `DagModel.max_active_tasks_per_dag`, `TI.max_active_tis_per_dagrun` to DB and replace the dummy values
- Make the query robust to missing values (coalesce, additional filters)
- Figure out which metrics we should emit, because some of them (`starved_X`) were removed as irrelevant. If we still want to emit them, it should be outside of the critical section
- Make the task fetching strategy experimental and injectable
- Cosmetics - lint, pre-commit hooks, consider writing the query as a whole instead of the dynamic building process
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250071293
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -1023,15 +898,15 @@ def _execute(self) -> int | None:
def _update_dag_run_state_for_paused_dags(self, session: Session = NEW_SESSION) -> None:
try:
paused_runs = session.scalars(
- select(DagRun)
+ select(DR)
Review Comment:
I'd propose to keep it as it was because (1) better readable and (2) smaller change set :-D
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250099050
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +377,171 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class LimitWindowDescriptor:
+ running_now_join: CTE
+ running_now_join_predicates: Collection[str]
+ limit_column: Column
+ window: expression.ColumnElement
+ limit_join_model: Base | None = None
+
+ def running_tasks_group(group_fields: Collection[Column], states: Collection[TaskInstanceState]=EXECUTION_STATES) -> CTE:
Review Comment:
This is still a draft PR, nothing has yet been stabilized, we do manual tests first and then go with the automated tests after we fix all the bugs, as we first want to measure performance to know if it is even worth the performance hit
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250139465
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -1023,15 +898,15 @@ def _execute(self) -> int | None:
def _update_dag_run_state_for_paused_dags(self, session: Session = NEW_SESSION) -> None:
try:
paused_runs = session.scalars(
- select(DagRun)
+ select(DR)
Review Comment:
I agree, reverted
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3155838953
> I printed out the updated query and posted it here
>
> https://gist.github.com/dstandish/a748988614f9f718f6b96b1330c04eec
>
> It's pretty ugly. I would guess performance is not gonna be good, and suspect there is room for optimization.
>
> Maybe try writing the query first in sql optimally, then figure out how to implement it in sqlalchemy.
>
> Among other things, it seems you nest many window functions over and over, but you likely can include most of them in the same select, i would bet.
>
> Lots of stuff in there that i think needs to be improved upon. Hope that is helpful.
The query looks awful because SQLalchemy insists on listing all the existing fields when doing `select(TI)`.
If we combine those, the query will be much shorter:
```sql
WITH
anon_5 AS
(SELECT
task_instance.dag_id AS dag_id,
task_instance.run_id AS run_id,
count(:count_1) AS now_running
FROM task_instance
WHERE
task_instance.state IN (__[postcompile_state_4])
GROUP BY task_instance.dag_id, task_instance.run_id
),
anon_6 AS
(SELECT
task_instance.dag_id AS dag_id,
task_instance.task_id AS task_id,
count(:count_2) AS now_running
FROM task_instance
WHERE
task_instance.state IN (__[postcompile_state_5])
GROUP BY task_instance.dag_id, task_instance.task_id
),
anon_7 AS
(SELECT
task_instance.dag_id AS dag_id,
task_instance.run_id AS run_id,
task_instance.task_id AS task_id,
count(:count_3) AS now_running
FROM task_instance
WHERE
task_instance.state IN (__[postcompile_state_6])
GROUP BY task_instance.dag_id, task_instance.run_id, task_instance.task_id
),
anon_8 AS
(SELECT
task_instance.pool AS pool,
count(:count_4) AS now_running
FROM task_instance
WHERE
task_instance.state IN (__[postcompile_state_7])
GROUP BY task_instance.pool
)
SELECT
task_instance.*
dag_run_1.*
FROM
task_instance
JOIN (SELECT
task_instance.*
sum(task_instance.pool_slots)
OVER (PARTITION BY task_instance.pool ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index) AS pool_slots_taken
FROM
task_instance
JOIN (SELECT
task_instance.*
row_number()
OVER (PARTITION BY task_instance.dag_id, task_instance.run_id, task_instance.task_id ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index) AS mapped_tis_per_dagrun_count
FROM
task_instance
JOIN (SELECT
task_instance.*
row_number()
OVER (PARTITION BY task_instance.dag_id, task_instance.task_id ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index) AS tis_per_dag_count
FROM
task_instance
JOIN (SELECT
task_instance.*
row_number()
OVER (PARTITION BY task_instance.dag_id, task_instance.run_id ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index) AS total_tis_per_dagrun_count
FROM
task_instance
JOIN dag_run ON dag_run.dag_id = task_instance.dag_id
AND dag_run.run_id = task_instance.run_id
JOIN dag ON task_instance.dag_id = dag.dag_id
WHERE
dag_run.state = :state_2
AND NOT dag.is_paused
AND task_instance.state = :state_3
AND dag.bundle_name IS NOT NULL
) AS anon_4 ON task_instance.id = anon_4.id
LEFT OUTER JOIN anon_5
ON task_instance.dag_id = anon_5.dag_id AND task_instance.run_id = anon_5.run_id
JOIN dag_run ON task_instance.run_id = dag_run.run_id
JOIN dag ON task_instance.dag_id = dag.dag_id
WHERE
coalesce(anon_4.total_tis_per_dagrun_count, 0) + coalesce(anon_5.now_running, 0)
<= coalesce(dag.max_active_tasks, :coalesce_1)
) AS anon_3 ON task_instance.id = anon_3.id
LEFT OUTER JOIN anon_6
ON task_instance.dag_id = anon_6.dag_id AND task_instance.task_id = anon_6.task_id
JOIN dag_run ON task_instance.run_id = dag_run.run_id
WHERE
coalesce(anon_3.tis_per_dag_count, 0) + coalesce(anon_6.now_running, 0)
<= coalesce(task_instance.max_active_tis_per_dag, :coalesce_2)
) AS anon_2 ON task_instance.id = anon_2.id
LEFT OUTER JOIN anon_7 ON task_instance.dag_id = anon_7.dag_id AND task_instance.run_id = anon_7.run_id
AND task_instance.task_id = anon_7.task_id
JOIN dag_run ON task_instance.run_id = dag_run.run_id
WHERE
coalesce(anon_2.mapped_tis_per_dagrun_count, 0) + coalesce(anon_7.now_running, 0)
<= coalesce(task_instance.max_active_tis_per_dagrun, :coalesce_3)
) AS anon_1 ON task_instance.id = anon_1.id
LEFT OUTER JOIN anon_8 ON task_instance.pool = anon_8.pool
JOIN dag_run ON task_instance.run_id = dag_run.run_id
JOIN slot_pool ON task_instance.pool = slot_pool.pool
JOIN dag_run AS dag_run_1 ON dag_run_1.dag_id = task_instance.dag_id AND dag_run_1.run_id = task_instance.run_id
WHERE
coalesce(anon_1.pool_slots_taken, 0) + coalesce(anon_8.now_running, 0) <= coalesce(slot_pool.slots, :coalesce_4)
LIMIT :param_1```
Regarding placing all windows in the same select - it won't work due to how filtering is done in the `where` clause. To correctly check all limits (an AND predicate), we have to nest them.
I already referred to this issue and some performance considerations in https://github.com/apache/airflow/issues/45636.
Anyway, I'm starting to work on the performance tests which will show us how slower this query will be.
--
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] Pessimistic task instance fetching with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2238893449
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
+
+ def running_tasks_group(*group_fields):
+ return (
+ select(TI, func.count("*").label("now_running"))
Review Comment:
The query won't be faster as we're counting rows. The SQL engine is smart enough to optimize these.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2257869034
##########
airflow-core/docs/migrations-ref.rst:
##########
@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| Revision ID | Revises ID | Airflow Version | Description |
+=========================+==================+===================+==============================================================+
-| ``f56f68b9e02f`` (head) | ``09fa89ba1710`` | ``3.1.0`` | Add callback_state to deadline. |
+| ``2f49f2dae90c`` (head) | ``f56f68b9e02f`` | ``3.1.0`` | Add callback_state to deadline. |
Review Comment:
Fixed by rebasing the revisions from scratch
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3161548920
I also, well, I tried to take a look at this question of whether any of the window functions can be collapsed into the same select.
So I took the two innermost window functions in your query.
This is what it looks like when I pull out only the two innermost selects. Here it is with no changes except naming of the aliases.
```sql
WITH
cte_running AS
(SELECT
task_instance.dag_id AS dag_id,
task_instance.run_id AS run_id,
count(:count_1) AS running_this_run
FROM task_instance
WHERE
task_instance.state IN ('RUNNING')
GROUP BY task_instance.dag_id, task_instance.run_id
)
SELECT
task_instance.*,
row_number() OVER (
PARTITION BY task_instance.dag_id, task_instance.task_id
ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index
) AS tis_per_dag_count
FROM
task_instance
JOIN (SELECT
task_instance.*,
row_number() OVER (
PARTITION BY task_instance.dag_id, task_instance.run_id
ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index
) AS total_tis_per_dagrun_count
FROM
task_instance
JOIN dag_run ON dag_run.dag_id = task_instance.dag_id
AND dag_run.run_id = task_instance.run_id
JOIN dag ON task_instance.dag_id = dag.dag_id
WHERE
dag_run.state = :state_2
AND NOT dag.is_paused
AND task_instance.state = :state_3
AND dag.bundle_name IS NOT NULL
) AS anon_4 ON task_instance.id = anon_4.id
LEFT JOIN cte_running
ON task_instance.dag_id = cte_running.dag_id AND task_instance.run_id = cte_running.run_id
JOIN dag_run ON task_instance.run_id = dag_run.run_id
JOIN dag ON task_instance.dag_id = dag.dag_id
WHERE
coalesce(anon_4.total_tis_per_dagrun_count, 0) + coalesce(cte_running.running_this_run, 0) <= coalesce(dag.max_active_tasks, 10)
```
So then I explored, what would happen if I collapsed them into one. Is it actually not possible?
So this is what it looks like when they are collapsed into one:
```sql
WITH
cte_running AS
(SELECT
task_instance.dag_id AS dag_id,
task_instance.run_id AS run_id,
count(:count_1) AS running_this_run
FROM task_instance
WHERE
task_instance.state IN ('RUNNING')
GROUP BY task_instance.dag_id, task_instance.run_id
)
select *
from
(SELECT
dag.max_active_tasks,
task_instance.*,
row_number() OVER (
PARTITION BY task_instance.dag_id, task_instance.run_id
ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index
) AS total_tis_per_dagrun_count,
row_number() OVER (
PARTITION BY task_instance.dag_id, task_instance.task_id
ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index
) AS tis_per_dag_count
FROM
task_instance
JOIN dag_run ON dag_run.dag_id = task_instance.dag_id
AND dag_run.run_id = task_instance.run_id
JOIN dag ON task_instance.dag_id = dag.dag_id
LEFT JOIN cte_running
ON task_instance.dag_id = cte_running.dag_id AND task_instance.run_id = cte_running.run_id
WHERE
dag_run.state = :state_2
AND NOT dag.is_paused
AND task_instance.state = :state_3
AND dag.bundle_name IS NOT NULL
) base
where
coalesce(base.total_tis_per_dagrun_count, 0) + coalesce(cte_running.running_this_run, 0)
<= coalesce(base.max_active_tasks, 10)
```
But then I noticed something. The two window functions are identical. So, it does seem a bit doubtful that we actually need both of them. Couldn't the outer queries just reference `total_tis_per_dagrun_count` instead of `tis_per_dag_count`, given that they are the same thing?
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3167965635
This query is a million times better, thank you. The other one was just like, not acceptable. It was too messy to be reasonably understood. This one, it's fine. You can understand what it is doing.
And that is important. Because now we can look at it and understand it and think about what it is doing, and evaluate whether we like it.
And now that we can do that, one observation.
This query is not a starvation panacea, I don't think.
For example what if there are two conditions interplaying.
Suppose there's only one pool slot available in pool A. Tasks X and Y have pool A. Task X sorts first. So task Y would be filtered out by the pool constraint.
But suppose further that for task X, total_tis_per_dagrun_count is beyond dag max active tasks. So nothing will be scheduled for pool A even though Y should be scheduled.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3168304983
> > as of now we want to get performance data and we also want to get it ready to be as an experimental feature in airflow 3.1 preferrably
>
> I would just say, we don't need to necessarily rush this out. Just wondering, have you considered patching your airflow and running it in your shop to see what the performance is? This is a thing that folks do sometimes. Can help make the case. Not saying we can't just add it but just asking the question.
@dstandish that is part of the plan and a great idea indeed, first I want to test it locally and see that it gives decent performance, then we will try running it on the shop with a cluster policy that sets all operators to empty operator, and testing at a higher scale (maybe enable backfill) to really test it under very high load with multiple instances of the DB on the same server and multiple airflow deployments just like @ashb proposed, though first I will do it locally with a smaller scale of 100k or more schedulable tasks on a decently powerful linux machine, and collect metrics.
I will publish the test on a different branch in case someone wants to test it as well.
Testing it locally first to tune (or understand how to tune) Max tis.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250099600
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_add_ti_max_active_tis_per_dagrun.py:
##########
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add callback_state to deadline.
+
+Revision ID: 2f49f2dae90c
+Revises: f56f68b9e02f
+Create Date: 2025-07-28 16:39:01.181132
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2f49f2dae90c"
+down_revision = "f56f68b9e02f"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+ """Add callback_state to deadline."""
+ with op.batch_alter_table("task_instance", schema=None) as batch_op:
+ batch_op.add_column(sa.Column("max_active_tis_per_dag", sa.Integer, nullable=True))
Review Comment:
It does not add much data, a million of 8 byte integers is still only 8MB which is a very small amount and that is assuming no encoding or compression at all, these fields are mandatory as we want to fix starvation bugs in this pr
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250135977
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +377,171 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class LimitWindowDescriptor:
Review Comment:
> I'd propose to make this rather a top level but private class and not a class in the class.
I think the dynamic query building process should be taken outside the scheduler, as it's just not wise to build the query every iteration from scratch (even if the interpreter can optimize it, which I doubt). The query should rather be a constant imported from elsewhere as it's a pretty static object (after the injection of `max_tis`)
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2252062326
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_add_ti_max_active_tis_per_dagrun.py:
##########
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add callback_state to deadline.
+
+Revision ID: 2f49f2dae90c
+Revises: f56f68b9e02f
+Create Date: 2025-07-28 16:39:01.181132
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2f49f2dae90c"
+down_revision = "f56f68b9e02f"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+ """Add callback_state to deadline."""
+ with op.batch_alter_table("task_instance", schema=None) as batch_op:
+ batch_op.add_column(sa.Column("max_active_tis_per_dag", sa.Integer, nullable=True))
Review Comment:
wouldn't this info always be the same for a given dag run? why not store it on the dag run table. you could also argue for it being on the dag but maybe you want to preserve ability for things to change per dag run.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3165220892
Optimized the query:
- No more nested windows, all of them are computed in a flat manner
- No more passing task instance.* fields across subqueries
- Fixed flawed joins
The runnable query version for max_tis=128:
```sql
WITH dag_run_active_tasks AS (
SELECT
task_instance.dag_id AS dag_id,
task_instance.run_id AS run_id,
count(*) AS now_running
FROM task_instance
WHERE task_instance.state IN ('queued', 'running')
GROUP BY task_instance.dag_id, task_instance.run_id
),
active_tis_across_dag_runs AS (
SELECT
task_instance.dag_id AS dag_id,
task_instance.task_id AS task_id,
count(*) AS now_running
FROM task_instance
WHERE task_instance.state IN ('queued', 'running')
GROUP BY task_instance.dag_id, task_instance.task_id
),
active_tis_in_one_task AS (
SELECT
task_instance.dag_id AS dag_id,
task_instance.run_id AS run_id,
task_instance.task_id AS task_id,
count(*) AS now_running
FROM task_instance
WHERE task_instance.state IN ('queued', 'running')
GROUP BY task_instance.dag_id, task_instance.run_id, task_instance.task_id
),
pool_active_tasks AS (
SELECT
task_instance.pool AS pool,
count(*) AS now_running
FROM task_instance
WHERE task_instance.state IN ('queued', 'running', 'deferred')
GROUP BY task_instance.pool
)
SELECT
task_instance.*,
dag_run_1.*
FROM
task_instance
JOIN (
SELECT
task_instance.id AS id,
row_number() OVER (
PARTITION BY task_instance.dag_id, task_instance.run_id
ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index
) AS total_tis_per_dagrun_count,
row_number() OVER (
PARTITION BY task_instance.dag_id, task_instance.task_id
ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index
) AS tis_per_dag_count,
row_number() OVER (
PARTITION BY task_instance.dag_id, task_instance.run_id, task_instance.task_id
ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index
) AS mapped_tis_per_dagrun_count,
sum(task_instance.pool_slots) OVER (
PARTITION BY task_instance.pool
ORDER BY -task_instance.priority_weight, dag_run.logical_date, task_instance.map_index
) AS pool_slots_taken_sum
FROM
task_instance
JOIN dag_run ON dag_run.dag_id = task_instance.dag_id AND dag_run.run_id = task_instance.run_id
JOIN dag ON task_instance.dag_id = dag.dag_id
WHERE
dag_run.state = 'running'
AND NOT dag.is_paused
AND task_instance.state = 'scheduled'
AND dag.bundle_name IS NOT NULL
) AS windowed_tis ON task_instance.id = windowed_tis.id
JOIN dag ON task_instance.dag_id = dag.dag_id
JOIN slot_pool ON task_instance.pool = slot_pool.pool
LEFT OUTER JOIN dag_run_active_tasks ON task_instance.dag_id = dag_run_active_tasks.dag_id AND task_instance.run_id = dag_run_active_tasks.run_id
LEFT OUTER JOIN active_tis_across_dag_runs ON task_instance.dag_id = active_tis_across_dag_runs.dag_id AND task_instance.task_id = active_tis_across_dag_runs.task_id
LEFT OUTER JOIN active_tis_in_one_task ON task_instance.dag_id = active_tis_in_one_task.dag_id AND task_instance.run_id = active_tis_in_one_task.run_id AND task_instance.task_id = active_tis_in_one_task.task_id
LEFT OUTER JOIN pool_active_tasks ON task_instance.pool = pool_active_tasks.pool
JOIN dag_run AS dag_run_1 ON dag_run_1.dag_id = task_instance.dag_id AND dag_run_1.run_id = task_instance.run_id
WHERE
coalesce(windowed_tis.total_tis_per_dagrun_count, 0) + coalesce(dag_run_active_tasks.now_running, 0) <= coalesce(dag.max_active_tasks, 128)
AND coalesce(windowed_tis.tis_per_dag_count, 0) + coalesce(active_tis_across_dag_runs.now_running, 0) <= coalesce(task_instance.max_active_tis_per_dag, 128)
AND coalesce(windowed_tis.mapped_tis_per_dagrun_count, 0) + coalesce(active_tis_in_one_task.now_running, 0) <= coalesce(task_instance.max_active_tis_per_dagrun, 128)
AND coalesce(windowed_tis.pool_slots_taken_sum, 0) + coalesce(pool_active_tasks.now_running, 0) <= coalesce(slot_pool.slots, 128)
LIMIT 128;
```
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3162669328
> You seem to pretty quickly dismiss the notion that the query is .... pretty gnarly. You don't think it can be simplified? Or maybe you just want to get some data on performance first.
Yes, we know the query is a monsterous one, as of now we want to get performance data and we also want to get it ready to be as an experimental feature in airflow 3.1 preferrably, we are working on thinking how to optimize the query, temp tables and materialized views may help with the performance, as of now from a very small test suite, the query takes between around half a second and a few seconds on different datasets, more data will be provided when further testing is done.
> It's also, you know, I'm not sure that piling so much into one query is really required to solve starvation problems. But it is not a positive for intelligibility, and can make it harder to optimize. The nesting in particular I think is pretty tough for intelligibility. Maybe it's necessary but I'd be surprised.
Me and @Asquator have talked about it in #45636 and we do need the 4 windows to solve all possible starvations (other than executor slots), if you find any flaw in there we would be very happy to hear your suggetsions.
>
> I did notice you are missing a join condition here:
>
> ```
> JOIN dag_run ON task_instance.run_id = dag_run.run_id
> ```
>
> This would result in join explosion if you have same run id for diff dags (and since airlfow's run id is derived from a date, it happens commonly).
Hmm, don't we have this join? If we didnt have it, wouldnt it cause a Cartesian join? @Asquator worked on removing all Cartesian joins, maybe the query posted above has some parts missing of the actual query.
> You may be able to avoid some joins to dag run if you are just after TI. Some of our relationships are eager loading and result in joins we don't need (but it can be disabled with query options).
>
> After taking a little closer look at the query....
>
> I was struck by how often we go back and do a row number over the task instance table. Is that really necessary? I would have expected that maybe once we grab the pool of candidate TIs in the innermost select, that we'd be able to just successively filter that down to the final set, even if we do have to involve some nested window functions. Over and over it's select *, row_number() from TI, DR, inner_query. Is there no way to stop going back to TI and DR after the first result set?
We will check if we can de-nest the windows, each window has a different order by (with some exceptions) and maybe reduce computations, but we do need to do some aggregation for each window in order for the query to compile, and we do need the window for the group by and partition by, so row number is used.
> It also can be a nice thing to do (for readability and debugging) to provide an alias or label for your subselects and ctes (e.g. to avoid anon_6 etc.
That will be done, we just want to get some testing done and then make it beautiful.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250072086
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_add_ti_max_active_tis_per_dagrun.py:
##########
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add callback_state to deadline.
+
+Revision ID: 2f49f2dae90c
+Revises: f56f68b9e02f
+Create Date: 2025-07-28 16:39:01.181132
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2f49f2dae90c"
+down_revision = "f56f68b9e02f"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+ """Add callback_state to deadline."""
+ with op.batch_alter_table("task_instance", schema=None) as batch_op:
+ batch_op.add_column(sa.Column("max_active_tis_per_dag", sa.Integer, nullable=True))
Review Comment:
Skimming over this PR... I am always very reluctant to change the task_instance table because this is probably the largest table in the DB schema in regards of volume. Mostly there are millions of DB records in. Even if this is just adding a few bytes per tuple... is there an option to make the same function w/o DB change? This adds a huge migration burden and increases the DB complexity.
During the migration as well the fields are not initialized. How does migrated data harm the scheduling?
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2252629525
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -554,6 +559,14 @@ class TaskInstance(Base, LoggingMixin):
innerjoin=True,
viewonly=True,
)
+ pool_model: Pool = relationship(
+ "Pool",
+ primaryjoin="TaskInstance.pool == Pool.pool",
+ foreign_keys=pool,
+ uselist=False,
+ innerjoin=True,
Review Comment:
Mhm, yeah I remember you get warnigns frmo Dag Parser anyway if the mentioned pool is not existing. If this makes ORM query better then it might be OK. We might leave a not that it is not made with a FK constraint by intend (because most probably we should not add a constraint)
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Ariel2400 (via GitHub)" <gi...@apache.org>.
Ariel2400 commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2257886015
##########
airflow-core/src/airflow/task/pessimistic_task_selector.py:
##########
@@ -0,0 +1,216 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from collections.abc import Collection
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from sqlalchemy import Column, and_, func, select, text
+from sqlalchemy.orm import Query, selectinload
+
+from airflow.models import DagRun, TaskInstance
+from airflow.models.dag import DagModel
+from airflow.models.pool import Pool
+from airflow.task.task_selector_strategy import TaskSelectorStrategy
+from airflow.ti_deps.dependencies_states import EXECUTION_STATES
+from airflow.utils.state import DagRunState, TaskInstanceState
+
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Query
+ from sqlalchemy.sql import expression
+ from sqlalchemy.sql.selectable import CTE, Select
+
+ from airflow.models.base import Base
+
+
+@dataclass
+class LimitWindowDescriptor:
+ """
+ Describes a limit window.
+
+ Args:
+ running_now_join (Subquery): the subquery on which we join for the window to get the additional parallelism
+ limits data.
+ running_now_join_predicates (Collection[str]): on what columns we do the join for for the
+ running_now_join, also used in the group by.
+ limit_column (Column): The column which decides how many rows should be chosen per window for
+ concurency limits.
+ window (expression.ColumnElement): the column by which we window.
+ limit_join_model (Base | None): the model on which we join to get a concurency limit.
+ """
+
+ running_now_join: CTE
+ running_now_join_predicates: Collection[str]
+ limit_column: Column
+ window: expression.ColumnElement
+ limit_join_model: Base | None = None
+
+
+TI = TaskInstance
+DR = DagRun
+DM = DagModel
+
+
+class PessimisticTaskSelector(TaskSelectorStrategy):
+ """
+ Pessimisticly query task instances ready for scheduling.
+
+ Works by delegating almost all the work from python do sql.
Review Comment:
I aasume you meant:
```suggestion
Works by delegating almost all the work from python to sql.
```
?
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3291350849
> Some starvation cases remain with this kind of query.
> We'll try another strategy featuring stored SQL procedures.
Just for clarification, we have solved all starvation cases, however they are solved between query runs rather than in a single query run.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator closed pull request #53492: Pessimistic task instance scheduling with window functions
URL: https://github.com/apache/airflow/pull/53492
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250151423
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_add_ti_max_active_tis_per_dagrun.py:
##########
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add callback_state to deadline.
+
+Revision ID: 2f49f2dae90c
+Revises: f56f68b9e02f
+Create Date: 2025-07-28 16:39:01.181132
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2f49f2dae90c"
+down_revision = "f56f68b9e02f"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+ """Add callback_state to deadline."""
+ with op.batch_alter_table("task_instance", schema=None) as batch_op:
+ batch_op.add_column(sa.Column("max_active_tis_per_dag", sa.Integer, nullable=True))
Review Comment:
> During the migration as well the fields are not initialized. How does migrated data harm the scheduling?
Do you mean existing task instances in one of the execution states that have concurrency limits? I can see the bug where the new scheduler will ignore the limits. We should manually reparse all the dags and add the fields to DB, it seems.
> Even if this is just adding a few bytes per tuple... is there an option to make the same function w/o DB change?
We need the fields in the DB to use them in a query. Regarding the table size, I see no other means to reduce it because we're in a system that heavily relies on SQL as its SSOT, and as long as we want features to be added, we must readily accept table extensions. task_instance is a flat table with a lot of params, some of them are duplicated across many TIs. Maybe it will be desirable to normalize it and extract the notorious task_template table or just a task_limits table in the long term.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250153533
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -554,6 +559,14 @@ class TaskInstance(Base, LoggingMixin):
innerjoin=True,
viewonly=True,
)
+ pool_model: Pool = relationship(
+ "Pool",
+ primaryjoin="TaskInstance.pool == Pool.pool",
+ foreign_keys=pool,
+ uselist=False,
+ innerjoin=True,
Review Comment:
Another direction is having consistency checks performed outside the scheduler by an "SRE checker" or something alike (can it be a scheduled DAG? :D )
From another perspective, it would be great to implement option 2 with an optimistic where clause we have now, but some persistence across iterations, such that inconsistent tasks are "punished" and not getting into the batch again.
--
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] Pessimistic task instance fetching with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2238908454
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
Review Comment:
I agree
--
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] Pessimistic task instance fetching with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2239169269
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
+
+ def running_tasks_group(*group_fields):
Review Comment:
The method builds a CTE to count concurrent tasks in different contexts, grouped by the passed fields.
Added the hinting
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3170695625
Posted the last checked query, returned nesting because it's needed to solve starvation:
https://gist.github.com/Asquator/156a3fcd99f6822cf477e7a6ccfdfff3
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2261333734
##########
airflow-core/src/airflow/task/pessimistic_task_selector.py:
##########
@@ -0,0 +1,204 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from collections.abc import Collection
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from sqlalchemy import Column, and_, func, select, text
+from sqlalchemy.orm import Query, selectinload
+
+from airflow.models import DagRun, TaskInstance
+from airflow.models.dag import DagModel
+from airflow.models.pool import Pool
+from airflow.task.task_selector_strategy import TaskSelectorStrategy
+from airflow.ti_deps.dependencies_states import EXECUTION_STATES
+from airflow.utils.state import DagRunState, TaskInstanceState
+
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Query
+ from sqlalchemy.sql import expression
+ from sqlalchemy.sql.selectable import CTE, Select
+
+ from airflow.models.base import Base
+
+
+@dataclass
+class LimitWindowDescriptor:
+ """
+ Describes a limit window.
+
+ Args:
+ running_now_join (Subquery): the subquery on which we join for the window to get the additional parallelism
+ limits data.
+ running_now_join_predicates (Collection[str]): on what columns we do the join for for the
+ running_now_join, also used in the group by.
+ limit_column (Column): The column which decides how many rows should be chosen per window for
+ concurency limits.
+ window (expression.ColumnElement): the column by which we window.
+ limit_join_model (Base | None): the model on which we join to get a concurency limit.
+ """
+
+ running_now_join: CTE
+ running_now_join_predicates: Collection[str]
+ limit_column: Column
+ window: expression.ColumnElement
+ limit_join_model: Base | None = None
+
+
+TI = TaskInstance
+DR = DagRun
+DM = DagModel
+
+
+class PessimisticTaskSelector(TaskSelectorStrategy):
+ """
+ Pessimisticly query task instances ready for scheduling.
+
+ Works by delegating almost all the work from python do sql.
+ Uses a few nested window functions to query only ready tasks.
+ """
+
+ def __init__(self) -> None:
+ self.priority_order = [-TaskInstance.priority_weight, DagRun.logical_date, TaskInstance.map_index]
+
+ def get_query(self, **additional_params) -> Query:
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
+ max_tis = additional_params["max_tis"]
+
+ base_query = (
+ select(TI.id)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ )
+
+ def running_tasks_group(
+ name: str, group_fields: Collection[Column], states: Collection[TaskInstanceState] = EXECUTION_STATES
+ ) -> CTE:
+ return (
+ select(*group_fields, func.count("*").label("now_running"))
+ .where(TI.state.in_(states))
+ .group_by(*group_fields)
+ .cte(name)
+ )
+
+ running_total_tis_per_dagrun = running_tasks_group("dag_run_active_tasks", [TI.dag_id, TI.run_id])
+ running_tis_per_dag = running_tasks_group("active_tis_across_dag_runs", [TI.dag_id, TI.task_id])
+ running_total_tis_per_task_run = running_tasks_group("active_tis_in_one_task", [TI.dag_id, TI.run_id, TI.task_id])
+ running_tis_per_pool = running_tasks_group("pool_active_tasks", [TI.pool], [*EXECUTION_STATES, TaskInstanceState.DEFERRED])
+
+ total_tis_per_dagrun_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.run_id), order_by=priority_order)
+ .label("total_tis_per_dagrun_count")
+ )
+ tis_per_dag_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.task_id), order_by=priority_order)
+ .label("tis_per_dag_count")
+ )
+ mapped_tis_per_task_run_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.run_id, TI.task_id), order_by=priority_order)
+ .label("mapped_tis_per_dagrun_count")
+ )
+ pool_slots_taken = (
+ func.sum(TI.pool_slots)
+ .over(partition_by=(TI.pool), order_by=priority_order)
+ .label("pool_slots_taken_sum")
+ )
+
+ limits = [
+ LimitWindowDescriptor(
+ running_total_tis_per_dagrun,
+ ["dag_id", "run_id"],
+ DagModel.max_active_tasks,
+ total_tis_per_dagrun_count,
+ TI.dag_model
+ ),
+ LimitWindowDescriptor(
+ running_tis_per_dag,
+ ["dag_id", "task_id"],
+ TI.max_active_tis_per_dag,
+ tis_per_dag_count
+ ),
+ LimitWindowDescriptor(
+ running_total_tis_per_task_run,
+ ["dag_id", "run_id", "task_id"],
+ TI.max_active_tis_per_dagrun,
+ mapped_tis_per_task_run_count,
+ ),
+ LimitWindowDescriptor(
+ running_tis_per_pool,
+ ["pool"],
+ Pool.slots,
+ pool_slots_taken,
+ TI.pool_model
+ ),
+ ]
+
+ base_query = base_query.add_columns(
+ total_tis_per_dagrun_count,
+ tis_per_dag_count,
+ mapped_tis_per_task_run_count,
+ pool_slots_taken
+ ).subquery().alias("windowed_tis")
+
+ query = (
+ select(TI)
+ .join(base_query, TI.id == base_query.c.id)
+ )
+
+ for limit in limits:
+ if limit.limit_join_model is not None:
+ query = query.join(limit.limit_join_model)
+ for limit in limits:
+ query = query.outerjoin(
+ limit.running_now_join,
+ and_(
+ *(
+ getattr(TI, predicate) == getattr(limit.running_now_join.c, predicate)
+ for predicate in limit.running_now_join_predicates
+ )
+ ),
+ )
+ query = query.where(
+ and_(
+ func.coalesce(getattr(base_query.c, limit.window.name), text("0"))
+ + func.coalesce(limit.running_now_join.c.now_running, text("0"))
+ <= func.coalesce(limit.limit_column, max_tis)
+ )
+ )
+
+ query = query.options(selectinload(TI.dag_model))
+ query = query.limit(max_tis)
+ return query
Review Comment:
Have you run any performance benchmarks with this sort of query? I have to say right now I am _very_ worried about merging a change that makes _such_ a monster query just on maintainability ground, even before I've seen what it does to the DB server. I'd especially like to see what happens if you run 3 or 10 deployments against a Single DB server (in different logical databases, i.e. `CREATE DATABASE airflow_n`);
Also how does mysql cope with this query?
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250072263
##########
airflow-core/docs/migrations-ref.rst:
##########
@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| Revision ID | Revises ID | Airflow Version | Description |
+=========================+==================+===================+==============================================================+
-| ``f56f68b9e02f`` (head) | ``09fa89ba1710`` | ``3.1.0`` | Add callback_state to deadline. |
+| ``2f49f2dae90c`` (head) | ``f56f68b9e02f`` | ``3.1.0`` | Add callback_state to deadline. |
Review Comment:
The comment seems to be "wrong". This is not what you are changing in the DB scheme.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250146754
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -554,6 +559,14 @@ class TaskInstance(Base, LoggingMixin):
innerjoin=True,
viewonly=True,
)
+ pool_model: Pool = relationship(
+ "Pool",
+ primaryjoin="TaskInstance.pool == Pool.pool",
+ foreign_keys=pool,
+ uselist=False,
+ innerjoin=True,
Review Comment:
> At the moment there is no constraint on the relation between pool and TI. Would it mean the scheduler will not "see" the TIs when a pool is referenced that an admin has deleted? Do we need to add a constraint to pool table as well?
The only reason I added the relationship is easier join syntax in the query. I'm not sure about the correct alchemy construct for this case, we should probably remove the `foreign_keys` argument?
Regarding "ghost" pools - it's actually a deeper concern and I'm glad you raised it. If a task references a pool that does not exist, you usually get a corresponding message, but only when the task is "accidentally" pulled by the scheduler. Let's discuss the possible solutions:
1. Don't show any message if a task references a pool that does not exist by not pulling them at all, means non-indicative task banning (bad, current PR state)
2. Do an outer join, and for tasks with no corresponding pool emit the message - makes the scheduler optimistic as to inconsistency handling, with many inconsistent tasks we get more slots "flooded" with pool-ghosted tasks (bad, current main state)
3. Show the message on dag parsing (why would you want to create a DAG with such tasks?) (better, can it go wrong?)
4. Handle pools in a smarter way than it's done now - don't allow deleting pools when there are tasks associated with it, but allow to turn it off (mainly for disaster handling). It's difficult to implement because there is no table for tasks. (best, complex and timely to implement)
We have different options and we have to choose one to merge this change ASAP, avoid breaking the current behavior and leave space for introducing new changes in subsequent PRs. I'll leave this open
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250098141
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +377,171 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class LimitWindowDescriptor:
Review Comment:
This will change in upcoming commits, I just need to find some time and fix a different issue
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3169114882
> @Nataneljpwd
> > That is a good point, however, I don't think this case can happen as the query won't allow more than dag Max active tasks to be run for a specific day, correct me if I'm wrong and if you think I didn't undestand. In this case task x is filtered out and task y is checked
>
> The fear is that row_number() will be less than #pool slots, so Y won't be even looked at. I think we whether must nest or find another way to cope with the situation.
@Asquator but don't we sum over pool slots? Rather than row number?
Maybe I'm wrong, if so you are more than welcome to correct me.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3168264230
@dstandish
> This query is not a starvation panacea, I don't think.
>
> For example what if there are two conditions interplaying.
>
> Suppose there's only one pool slot available in pool A. Tasks X and Y have pool A. Task X sorts first. So task Y would be filtered out by the pool constraint.
If task x sorts first it means its at a higher priority and needs to be run first, which is expected behaviour, I may have not fully understood what the problem here, and I would appreciate it if you could clarify further
> But suppose further that for task X, total_tis_per_dagrun_count is beyond dag max active tasks. So nothing will be scheduled for pool A even though Y should be scheduled.
That is a good point, however, I don't think this case can happen as the query won't allow more than dag Max active tasks to be run for a specific day, correct me if I'm wrong and if you think I didn't undestand.
In this case task x is filtered out and task y is checked, the query checks all tasks and returns when it has Max tis tasks (or no tasks left) so I don't see this becoming a problem.
The query works in such a way that it checks all tasks that are scheduled and does not drop any, until either the limit is satisfied or no data is left.
--
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] Pessimistic task instance fetching with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2238900449
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
+
+ def running_tasks_group(*group_fields):
Review Comment:
Please add another suggestion to import CTE and I will commit them in batch
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
+
+ def running_tasks_group(*group_fields):
Review Comment:
Please add another suggestion to import `CTE` and I will commit them in batch
--
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] Pessimistic task instance fetching with window functions [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2237286013
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
+
+ def running_tasks_group(*group_fields):
Review Comment:
```suggestion
def running_tasks_group(*group_fields: list[Column]) -> CTE:
```
what exactly does this method do? does it just group tasks by the groups given? what is the EXECUTION_STATES?
that would make it a little nicer to read
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
+
+ def running_tasks_group(*group_fields):
+ return (
+ select(TI, func.count("*").label("now_running"))
+ .where(TI.state.in_(EXECUTION_STATES))
+ .group_by(*group_fields)
+ .cte()
+ )
- query = (
+ def add_window_limit(query: Select, limit: Limit):
Review Comment:
```suggestion
def add_window_limit(query: Select, limit: Limit) -> Select:
```
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
Review Comment:
why does a Limit hold a column?
why does a limit hold a subquery as well? maybe the name should be something more along the lines of `WindowDescriptor` or something along those lines?
a limit should not hold columns or column elements, it would also help if you added a docstring on this dataclass, and give the variables more general names, so that it would be simpler to understand
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
+
+ def running_tasks_group(*group_fields):
+ return (
+ select(TI, func.count("*").label("now_running"))
Review Comment:
```suggestion
select(TI, func.count("id").label("now_running"))
```
to enable projection in sql, will make the query slightly faster
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
Review Comment:
```suggestion
class LimitWindowDescriptor:
join_on: Subquery
choose_up_to: Column
window_by: expression.ColumnElement
```
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
+
+ def running_tasks_group(*group_fields):
+ return (
+ select(TI, func.count("*").label("now_running"))
+ .where(TI.state.in_(EXECUTION_STATES))
+ .group_by(*group_fields)
+ .cte()
+ )
- query = (
+ def add_window_limit(query: Select, limit: Limit):
+ inner_query = (
+ query.add_columns(limit.window)
+ .join(limit.running_now_join, TI.id == limit.running_now_join.c.id)
+ .order_by(*priority_order)
+ .subquery()
+ )
+ return (
select(TI)
- .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
- .join(TI.dag_run)
- .where(DR.state == DagRunState.RUNNING)
- .join(TI.dag_model)
- .where(~DM.is_paused)
- .where(TI.state == TaskInstanceState.SCHEDULED)
- .where(DM.bundle_name.is_not(None))
- .options(selectinload(TI.dag_model))
- .order_by(-TI.priority_weight, DR.logical_date, TI.map_index)
+ .join(inner_query, TI.id == inner_query.c.id)
+ .where(
+ getattr(inner_query.c, limit.window.name) + limit.running_now_join.c.now_running
+ < limit.max_units
+ )
)
- if starved_pools:
- query = query.where(TI.pool.not_in(starved_pools))
+ running_total_tis_per_dagrun = running_tasks_group(TI.dag_id, TI.run_id)
+ running_tis_per_dag = running_tasks_group(TI.dag_id, TI.task_id)
+ running_total_tis_per_task_run = running_tasks_group(TI.dag_id, TI.run_id, TI.task_id)
+ running_tis_per_pool = running_tasks_group(TI.pool)
- if starved_dags:
- query = query.where(TI.dag_id.not_in(starved_dags))
+ total_tis_per_dagrun_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.run_id), order_by=priority_order)
+ .label("total_tis_per_dagrun_count")
+ )
+ tis_per_dag_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.task_id), order_by=priority_order)
+ .label("tis_per_dag_count")
+ )
+ mapped_tis_per_task_run_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.run_id, TI.task_id), order_by=priority_order)
+ .label("mapped_tis_per_dagrun_count")
+ )
+ pool_slots_taken = (
+ func.sum(TI.pool_slots)
+ .over(partition_by=(TI.pool), order_by=priority_order)
+ .label("pool_slots_taken")
+ )
- if starved_tasks:
- query = query.where(tuple_(TI.dag_id, TI.task_id).not_in(starved_tasks))
+ limits = [
+ Limit(running_total_tis_per_dagrun, DagModel.max_active_tasks, total_tis_per_dagrun_count),
+ Limit(
+ running_tis_per_dag, TaskInstance.max_active_tis_per_dag, tis_per_dag_count
+ ), # TODO: Add to DB model: TaskInstance.max_active_tis_per_dag, DUMMY: DagModel.max_active_tasks
+ Limit(
+ running_total_tis_per_task_run, TI.max_tries, mapped_tis_per_task_run_count
+ ), # TODO: Add to DB model: TI.max_active_tis_per_dagrun, DUMMY: TI.max_tries
+ Limit(running_tis_per_pool, Pool.slots, pool_slots_taken),
+ ]
- if starved_tasks_task_dagrun_concurrency:
- query = query.where(
- tuple_(TI.dag_id, TI.run_id, TI.task_id).not_in(starved_tasks_task_dagrun_concurrency)
- )
+ for limit in limits:
+ query = add_window_limit(query, limit)
- query = query.limit(max_tis)
+ query = query.limit(max_tis)
- timer = Stats.timer("scheduler.critical_section_query_duration")
- timer.start()
+ timer = Stats.timer("scheduler.critical_section_query_duration")
+ timer.start()
- try:
- query = with_row_locks(query, of=TI, session=session, skip_locked=True)
- task_instances_to_examine: list[TI] = session.scalars(query).all()
+ try:
+ query = with_row_locks(query, of=TI, session=session, skip_locked=True)
+ task_instances_to_examine: list[TI] = session.scalars(query).all()
- timer.stop(send=True)
- except OperationalError as e:
- timer.stop(send=False)
- raise e
+ timer.stop(send=True)
+ except OperationalError as e:
+ timer.stop(send=False)
+ raise e
- # TODO[HA]: This was wrong before anyway, as it only looked at a sub-set of dags, not everything.
- # Stats.gauge('scheduler.tasks.pending', len(task_instances_to_examine))
+ # TODO[HA]: This was wrong before anyway, as it only looked at a sub-set of dags, not everything.
+ # Stats.gauge('scheduler.tasks.pending', len(task_instances_to_examine))
- if not task_instances_to_examine:
- self.log.debug("No tasks to consider for execution.")
- break
+ if not task_instances_to_examine:
+ self.log.debug("No tasks to consider for execution.")
+ return []
- # Put one task instance on each line
- task_instance_str = "\n".join(f"\t{x!r}" for x in task_instances_to_examine)
- self.log.info("%s tasks up for execution:\n%s", len(task_instances_to_examine), task_instance_str)
+ # Put one task instance on each line
+ task_instance_str = "\n".join(f"\t{x!r}" for x in task_instances_to_examine)
+ self.log.info("%s tasks up for execution:\n%s", len(task_instances_to_examine), task_instance_str)
- executor_slots_available: dict[ExecutorName, int] = {}
- # First get a mapping of executor names to slots they have available
- for executor in self.job.executors:
+ executor_slots_available: dict[ExecutorName, int] = {}
+ # First get a mapping of executor names to slots they have available
+ for executor in self.job.executors:
+ if TYPE_CHECKING:
+ # All executors should have a name if they are initted from the executor_loader.
+ # But we need to check for None to make mypy happy.
+ assert executor.name
+ executor_slots_available[executor.name] = executor.slots_available
+
+ for task_instance in task_instances_to_examine:
+ pool_name = task_instance.pool
+
+ pool_stats = pools.get(pool_name)
+ if not pool_stats:
+ self.log.warning("Tasks using non-existent pool '%s' will not be scheduled", pool_name)
+ starved_pools.add(pool_name)
+ continue
+
+ if executor_obj := self._try_to_load_executor(task_instance.executor):
if TYPE_CHECKING:
# All executors should have a name if they are initted from the executor_loader.
# But we need to check for None to make mypy happy.
- assert executor.name
- executor_slots_available[executor.name] = executor.slots_available
-
- for task_instance in task_instances_to_examine:
- pool_name = task_instance.pool
-
- pool_stats = pools.get(pool_name)
- if not pool_stats:
- self.log.warning("Tasks using non-existent pool '%s' will not be scheduled", pool_name)
- starved_pools.add(pool_name)
- continue
-
- # Make sure to emit metrics if pool has no starving tasks
- pool_num_starving_tasks.setdefault(pool_name, 0)
-
- pool_total = pool_stats["total"]
- open_slots = pool_stats["open"]
-
- if open_slots <= 0:
- self.log.info(
- "Not scheduling since there are %s open slots in pool %s", open_slots, pool_name
- )
- # Can't schedule any more since there are no more open slots.
- pool_num_starving_tasks[pool_name] += 1
- num_starving_tasks_total += 1
- starved_pools.add(pool_name)
- continue
-
- if task_instance.pool_slots > pool_total:
- self.log.warning(
- "Not executing %s. Requested pool slots (%s) are greater than "
- "total pool slots: '%s' for pool: %s.",
- task_instance,
- task_instance.pool_slots,
- pool_total,
- pool_name,
- )
-
- pool_num_starving_tasks[pool_name] += 1
- num_starving_tasks_total += 1
- starved_tasks.add((task_instance.dag_id, task_instance.task_id))
- continue
-
- if task_instance.pool_slots > open_slots:
- self.log.info(
- "Not executing %s since it requires %s slots "
- "but there are %s open slots in the pool %s.",
- task_instance,
- task_instance.pool_slots,
- open_slots,
- pool_name,
+ assert executor_obj.name
+ if executor_slots_available[executor_obj.name] <= 0:
+ self.log.debug(
Review Comment:
I think here it should be a warning rather than a debug, as if we start getting a lot of those, we would like to know about it and know how to adjust, rather than just write to debug.
```suggestion
if executor_slots_available[executor_obj.name] <= 0:
self.log.warn(
"Not scheduling %s since its executor %s does not currently have any more "
"available slots"
)
```
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -1023,15 +898,15 @@ def _execute(self) -> int | None:
def _update_dag_run_state_for_paused_dags(self, session: Session = NEW_SESSION) -> None:
try:
paused_runs = session.scalars(
- select(DagRun)
+ select(DR)
Review Comment:
why not just keep it `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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3153062016
I printed out the updated query and posted it here
https://gist.github.com/dstandish/a748988614f9f718f6b96b1330c04eec
It's pretty ugly. I would guess performance is not gonna be good, and suspect there is room for optimization.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3156707658
That is _one hell of a query_.
I'm frankly amazed that gives any performance at all.
Dare I mention the fact that we still support MySQL? How does it behave with this?
--
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] Pessimistic task instance fetching with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2238903981
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
Review Comment:
Renamed to `WindowLimitDescriptor`
--
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] Pessimistic task instance selection with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3144513095
Looks like the query works, basic sanity checks pass and tasks are getting scheduled with all the limits correctly checked.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250073230
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -554,6 +559,14 @@ class TaskInstance(Base, LoggingMixin):
innerjoin=True,
viewonly=True,
)
+ pool_model: Pool = relationship(
+ "Pool",
+ primaryjoin="TaskInstance.pool == Pool.pool",
+ foreign_keys=pool,
+ uselist=False,
+ innerjoin=True,
Review Comment:
(Which then would mean the other way around you can not delete a pool until you have archived all old records which maybe reference a pool...)
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250100311
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -554,6 +559,14 @@ class TaskInstance(Base, LoggingMixin):
innerjoin=True,
viewonly=True,
)
+ pool_model: Pool = relationship(
+ "Pool",
+ primaryjoin="TaskInstance.pool == Pool.pool",
+ foreign_keys=pool,
+ uselist=False,
+ innerjoin=True,
Review Comment:
There is no direct relation, however what about pool slots? And slots taken by a TI? You do need the relation for calculations that allow to avoid starvation of pools
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3168927238
@dstandish
> This query is not a starvation panacea, I don't think.
>
> For example what if there are two conditions interplaying.
>
> Suppose there's only one pool slot available in pool A. Tasks X and Y have pool A. Task X sorts first. So task Y would be filtered out by the pool constraint.
>
> But suppose further that for task X, total_tis_per_dagrun_count is beyond dag max active tasks. So nothing will be scheduled for pool A even though Y should be scheduled.
I think this is a very likely situation, and I think this was the reason I decided to nest the windows in the beginning. I'll have to think it over. If we nest the windows, the problem doesn't exist anymore since every window gets the filtered TI.ids.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3162666876
I think @dstandish has a point here. In the beginning we planned to make four flat windows, but it was decided to nest them, and now I think for no reason at all. As the window values (row_numbers, sums) are assigned after partitioning is done, consistency and scheduling guarantees should be achieved. I will fix it ASAP to measure performance again, since now it's pretty slow, and the plan looks intimidating.
Regarding the two windows that look similar - I doubt they can be replaced by one window due to different _partition by_ columns.
Thanks for looking at the join conditions. I'll try to fix them too.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2259212542
##########
airflow-core/src/airflow/task/pessimistic_task_selector.py:
##########
@@ -0,0 +1,216 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from collections.abc import Collection
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from sqlalchemy import Column, and_, func, select, text
+from sqlalchemy.orm import Query, selectinload
+
+from airflow.models import DagRun, TaskInstance
+from airflow.models.dag import DagModel
+from airflow.models.pool import Pool
+from airflow.task.task_selector_strategy import TaskSelectorStrategy
+from airflow.ti_deps.dependencies_states import EXECUTION_STATES
+from airflow.utils.state import DagRunState, TaskInstanceState
+
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Query
+ from sqlalchemy.sql import expression
+ from sqlalchemy.sql.selectable import CTE, Select
+
+ from airflow.models.base import Base
+
+
+@dataclass
+class LimitWindowDescriptor:
+ """
+ Describes a limit window.
+
+ Args:
+ running_now_join (Subquery): the subquery on which we join for the window to get the additional parallelism
+ limits data.
+ running_now_join_predicates (Collection[str]): on what columns we do the join for for the
+ running_now_join, also used in the group by.
+ limit_column (Column): The column which decides how many rows should be chosen per window for
+ concurency limits.
+ window (expression.ColumnElement): the column by which we window.
+ limit_join_model (Base | None): the model on which we join to get a concurency limit.
+ """
+
+ running_now_join: CTE
+ running_now_join_predicates: Collection[str]
+ limit_column: Column
+ window: expression.ColumnElement
+ limit_join_model: Base | None = None
+
+
+TI = TaskInstance
+DR = DagRun
+DM = DagModel
+
+
+class PessimisticTaskSelector(TaskSelectorStrategy):
+ """
+ Pessimisticly query task instances ready for scheduling.
+
+ Works by delegating almost all the work from python do sql.
Review Comment:
Yes, now it looks as if a drunk man wrote these docstrings.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3167978103
> as of now we want to get performance data and we also want to get it ready to be as an experimental feature in airflow 3.1 preferrably
I would just say, we don't need to necessarily rush this out. Just wondering, have you considered patching your airflow and running it in your shop to see what the performance is? This is a thing that folks do sometimes. Can help make the case. Not saying we can't just add it but just asking the question.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250074104
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +377,171 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class LimitWindowDescriptor:
Review Comment:
I'd propose to make this rather a top level but private class and not a class in the class.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250151423
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_add_ti_max_active_tis_per_dagrun.py:
##########
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add callback_state to deadline.
+
+Revision ID: 2f49f2dae90c
+Revises: f56f68b9e02f
+Create Date: 2025-07-28 16:39:01.181132
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2f49f2dae90c"
+down_revision = "f56f68b9e02f"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+ """Add callback_state to deadline."""
+ with op.batch_alter_table("task_instance", schema=None) as batch_op:
+ batch_op.add_column(sa.Column("max_active_tis_per_dag", sa.Integer, nullable=True))
Review Comment:
> During the migration as well the fields are not initialized. How does migrated data harm the scheduling?
Do you mean existing task instances in one of the execution states that have a concurrency limits? I can see the bug where the new scheduler will ignore the limits. We should manually reparse all the dags and add the fields to DB, it seems.
> Even if this is just adding a few bytes per tuple... is there an option to make the same function w/o DB change?
We need the fields in the DB to use them in a query. Regarding the table size, I see no other means to reduce it because we're in a system that heavily relies on SQL as its SSOT, and as long as we want features to be added, we must readily accept table extensions. task_instance is a flat table with a lot of params, some of them are duplicated across many TIs. Maybe it will be desirable to normalize it and extract the notorious task_template table or just a task_limits table in the long term.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250146754
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -554,6 +559,14 @@ class TaskInstance(Base, LoggingMixin):
innerjoin=True,
viewonly=True,
)
+ pool_model: Pool = relationship(
+ "Pool",
+ primaryjoin="TaskInstance.pool == Pool.pool",
+ foreign_keys=pool,
+ uselist=False,
+ innerjoin=True,
Review Comment:
> At the moment there is no constraint on the relation between pool and TI. Would it mean the scheduler will not "see" the TIs when a pool is referenced that an admin has deleted? Do we need to add a constraint to pool table as well?
The only reason I added the relationship is easier join syntax in the query. I'm not sure about the correct alchemy construct for this case, we should probably remove the `foreign_keys` argument?
Regarding "ghost" pools - it's actually a deeper concern and I'm glad you raised it. If a task references a pool that does not exist, you usually get a corresponding message, but only when the task is "accidentally" pulled by the scheduler. Let's discuss the possible solutions:
1. Don't show any message if a task references a pool that does not exist by not pulling them at all, means non-indicative task banning (bad, current PR state)
2. . Do an outer join, and for tasks with no corresponding pool emit the message - makes the scheduler optimistic as to inconsistency handling, with many inconsistent tasks we get more slots "flooded" with pool-ghosted tasks (bad, current main state)
3. Show the message on dag parsing (why would you want to create a DAG with such tasks?) (better, can it go wrong?)
4. Handle pools in a smarter way than it's done now - don't allow deleting pools when there are tasks associated with it, but allow to turn it off (mainly for disaster handling). It's difficult to implement because there is no table for tasks.
We have different options and we have to choose one to merge this change ASAP, avoid breaking the current behavior and leave space for introducing new changes in subsequent PRs. I'll leave this open
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3169162959
> > @Nataneljpwd
> > > That is a good point, however, I don't think this case can happen as the query won't allow more than dag Max active tasks to be run for a specific day, correct me if I'm wrong and if you think I didn't undestand. In this case task x is filtered out and task y is checked
> >
> >
> > The fear is that row_number() will be less than #pool slots, so Y won't be even looked at. I think we whether must nest or find another way to cope with the situation.
>
> @Asquator but don't we sum over pool slots? Rather than row number?
>
> Maybe I'm wrong, if so you are more than welcome to correct me.
We do sum over pool slots, but it doesn't change nothing because it's a fixed, static, cumulative sum. If a task is dropped due to another limit, the pool slots are gone.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3275916179
Some starvation cases remain with this kind of query.
We'll try another strategy featuring stored SQL procedures.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2252064905
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_add_ti_max_active_tis_per_dagrun.py:
##########
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add callback_state to deadline.
+
+Revision ID: 2f49f2dae90c
+Revises: f56f68b9e02f
+Create Date: 2025-07-28 16:39:01.181132
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2f49f2dae90c"
+down_revision = "f56f68b9e02f"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+ """Add callback_state to deadline."""
+ with op.batch_alter_table("task_instance", schema=None) as batch_op:
+ batch_op.add_column(sa.Column("max_active_tis_per_dag", sa.Integer, nullable=True))
Review Comment:
but it seems pretty uncontroversial that it doesn't need to be on task instance table given that, right...?
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Ariel2400 (via GitHub)" <gi...@apache.org>.
Ariel2400 commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2257911624
##########
airflow-core/src/airflow/task/pessimistic_task_selector.py:
##########
@@ -0,0 +1,216 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from collections.abc import Collection
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from sqlalchemy import Column, and_, func, select, text
+from sqlalchemy.orm import Query, selectinload
+
+from airflow.models import DagRun, TaskInstance
+from airflow.models.dag import DagModel
+from airflow.models.pool import Pool
+from airflow.task.task_selector_strategy import TaskSelectorStrategy
+from airflow.ti_deps.dependencies_states import EXECUTION_STATES
+from airflow.utils.state import DagRunState, TaskInstanceState
+
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Query
+ from sqlalchemy.sql import expression
+ from sqlalchemy.sql.selectable import CTE, Select
+
+ from airflow.models.base import Base
+
+
+@dataclass
+class LimitWindowDescriptor:
+ """
+ Describes a limit window.
+
+ Args:
+ running_now_join (Subquery): the subquery on which we join for the window to get the additional parallelism
+ limits data.
+ running_now_join_predicates (Collection[str]): on what columns we do the join for for the
+ running_now_join, also used in the group by.
+ limit_column (Column): The column which decides how many rows should be chosen per window for
+ concurency limits.
+ window (expression.ColumnElement): the column by which we window.
+ limit_join_model (Base | None): the model on which we join to get a concurency limit.
+ """
+
+ running_now_join: CTE
+ running_now_join_predicates: Collection[str]
+ limit_column: Column
+ window: expression.ColumnElement
+ limit_join_model: Base | None = None
+
+
+TI = TaskInstance
+DR = DagRun
+DM = DagModel
+
+
+class PessimisticTaskSelector(TaskSelectorStrategy):
+ """
+ Pessimisticly query task instances ready for scheduling.
+
+ Works by delegating almost all the work from python do sql.
Review Comment:
also, please run spellchecker on the rest of the docstrings
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2252151781
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_add_ti_max_active_tis_per_dagrun.py:
##########
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add callback_state to deadline.
+
+Revision ID: 2f49f2dae90c
+Revises: f56f68b9e02f
+Create Date: 2025-07-28 16:39:01.181132
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2f49f2dae90c"
+down_revision = "f56f68b9e02f"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+ """Add callback_state to deadline."""
+ with op.batch_alter_table("task_instance", schema=None) as batch_op:
+ batch_op.add_column(sa.Column("max_active_tis_per_dag", sa.Integer, nullable=True))
Review Comment:
`max_active_tis_per_dag` and `max_active_tis_per_dagrun` are a per-task parameters. They can't be stored in the dag_run table because different tasks may have different values for these fields. The overhead is clear, and it's even more evident with mapped tasks, where you can have thousands of tasks duplicated with same values. I agree it's bad, and I'd like to hear from Airflow experts who specialize in DB performance. Another discussion will be the need to normalize the `task_instance` table and maybe derive the `task` table. It's a totally different change that deserves an AIP of its own, because it means tasks become standalone objects. I think I can see the good and the very good sides of it, but it's a completely different topic.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3161258751
You seem to pretty quickly dismiss the notion that the query is .... pretty gnarly. You don't think it can be simplified? Or maybe you just want to get some data on performance first.
It's also, you know, I'm not sure that piling so much into one query is really required to solve starvation problems. But it is not a positive for intelligibility, and can make it harder to optimize. The nesting in particular I think is pretty tough for intelligibility. Maybe it's necessary but I'd be surprised.
I did notice you are missing a join condition here:
```
JOIN dag_run ON task_instance.run_id = dag_run.run_id
```
This would result in join explosion if you have same run id for diff dags (and since airlfow's run id is derived from a date, it happens commonly).
You may be able to avoid some joins to dag run if you are just after TI. Some of our relationships are eager loading and result in joins we don't need (but it can be disabled with query options).
After taking a little closer look at the query....
I was struck by how often we go back and do a row number over the task instance table. Is that really necessary? I would have expected that maybe once we grab the pool of candidate TIs in the innermost select, that we'd be able to just successively filter that down to the final set, even if we do have to involve some nested window functions. Over and over it's select *, row_number() from TI, DR, inner_query. Is there no way to stop going back to TI and DR after the first result set?
It also can be a nice thing to do (for readability and debugging) to provide an alias or label for your subselects and ctes (e.g. to avoid anon_6 etc.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2261384655
##########
airflow-core/src/airflow/task/pessimistic_task_selector.py:
##########
@@ -0,0 +1,204 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from collections.abc import Collection
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from sqlalchemy import Column, and_, func, select, text
+from sqlalchemy.orm import Query, selectinload
+
+from airflow.models import DagRun, TaskInstance
+from airflow.models.dag import DagModel
+from airflow.models.pool import Pool
+from airflow.task.task_selector_strategy import TaskSelectorStrategy
+from airflow.ti_deps.dependencies_states import EXECUTION_STATES
+from airflow.utils.state import DagRunState, TaskInstanceState
+
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Query
+ from sqlalchemy.sql import expression
+ from sqlalchemy.sql.selectable import CTE, Select
+
+ from airflow.models.base import Base
+
+
+@dataclass
+class LimitWindowDescriptor:
+ """
+ Describes a limit window.
+
+ Args:
+ running_now_join (Subquery): the subquery on which we join for the window to get the additional parallelism
+ limits data.
+ running_now_join_predicates (Collection[str]): on what columns we do the join for for the
+ running_now_join, also used in the group by.
+ limit_column (Column): The column which decides how many rows should be chosen per window for
+ concurency limits.
+ window (expression.ColumnElement): the column by which we window.
+ limit_join_model (Base | None): the model on which we join to get a concurency limit.
+ """
+
+ running_now_join: CTE
+ running_now_join_predicates: Collection[str]
+ limit_column: Column
+ window: expression.ColumnElement
+ limit_join_model: Base | None = None
+
+
+TI = TaskInstance
+DR = DagRun
+DM = DagModel
+
+
+class PessimisticTaskSelector(TaskSelectorStrategy):
+ """
+ Pessimisticly query task instances ready for scheduling.
+
+ Works by delegating almost all the work from python do sql.
+ Uses a few nested window functions to query only ready tasks.
+ """
+
+ def __init__(self) -> None:
+ self.priority_order = [-TaskInstance.priority_weight, DagRun.logical_date, TaskInstance.map_index]
+
+ def get_query(self, **additional_params) -> Query:
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
+ max_tis = additional_params["max_tis"]
+
+ base_query = (
+ select(TI.id)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ )
+
+ def running_tasks_group(
+ name: str, group_fields: Collection[Column], states: Collection[TaskInstanceState] = EXECUTION_STATES
+ ) -> CTE:
+ return (
+ select(*group_fields, func.count("*").label("now_running"))
+ .where(TI.state.in_(states))
+ .group_by(*group_fields)
+ .cte(name)
+ )
+
+ running_total_tis_per_dagrun = running_tasks_group("dag_run_active_tasks", [TI.dag_id, TI.run_id])
+ running_tis_per_dag = running_tasks_group("active_tis_across_dag_runs", [TI.dag_id, TI.task_id])
+ running_total_tis_per_task_run = running_tasks_group("active_tis_in_one_task", [TI.dag_id, TI.run_id, TI.task_id])
+ running_tis_per_pool = running_tasks_group("pool_active_tasks", [TI.pool], [*EXECUTION_STATES, TaskInstanceState.DEFERRED])
+
+ total_tis_per_dagrun_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.run_id), order_by=priority_order)
+ .label("total_tis_per_dagrun_count")
+ )
+ tis_per_dag_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.task_id), order_by=priority_order)
+ .label("tis_per_dag_count")
+ )
+ mapped_tis_per_task_run_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.run_id, TI.task_id), order_by=priority_order)
+ .label("mapped_tis_per_dagrun_count")
+ )
+ pool_slots_taken = (
+ func.sum(TI.pool_slots)
+ .over(partition_by=(TI.pool), order_by=priority_order)
+ .label("pool_slots_taken_sum")
+ )
+
+ limits = [
+ LimitWindowDescriptor(
+ running_total_tis_per_dagrun,
+ ["dag_id", "run_id"],
+ DagModel.max_active_tasks,
+ total_tis_per_dagrun_count,
+ TI.dag_model
+ ),
+ LimitWindowDescriptor(
+ running_tis_per_dag,
+ ["dag_id", "task_id"],
+ TI.max_active_tis_per_dag,
+ tis_per_dag_count
+ ),
+ LimitWindowDescriptor(
+ running_total_tis_per_task_run,
+ ["dag_id", "run_id", "task_id"],
+ TI.max_active_tis_per_dagrun,
+ mapped_tis_per_task_run_count,
+ ),
+ LimitWindowDescriptor(
+ running_tis_per_pool,
+ ["pool"],
+ Pool.slots,
+ pool_slots_taken,
+ TI.pool_model
+ ),
+ ]
+
+ base_query = base_query.add_columns(
+ total_tis_per_dagrun_count,
+ tis_per_dag_count,
+ mapped_tis_per_task_run_count,
+ pool_slots_taken
+ ).subquery().alias("windowed_tis")
+
+ query = (
+ select(TI)
+ .join(base_query, TI.id == base_query.c.id)
+ )
+
+ for limit in limits:
+ if limit.limit_join_model is not None:
+ query = query.join(limit.limit_join_model)
+ for limit in limits:
+ query = query.outerjoin(
+ limit.running_now_join,
+ and_(
+ *(
+ getattr(TI, predicate) == getattr(limit.running_now_join.c, predicate)
+ for predicate in limit.running_now_join_predicates
+ )
+ ),
+ )
+ query = query.where(
+ and_(
+ func.coalesce(getattr(base_query.c, limit.window.name), text("0"))
+ + func.coalesce(limit.running_now_join.c.now_running, text("0"))
+ <= func.coalesce(limit.limit_column, max_tis)
+ )
+ )
+
+ query = query.options(selectinload(TI.dag_model))
+ query = query.limit(max_tis)
+ return query
Review Comment:
I share the same worries
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Nataneljpwd (via GitHub)" <gi...@apache.org>.
Nataneljpwd commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3176810402
> Hi in your calculation of `tis_per_dag_count` you are partitioning by dag and task but i think maybe it's supposed to be just dag?
>
> er... oh boy... what's the difference between `max_active_tis_per_dagrun` and `max_active_tasks`
>
> ok i guess `max_active_tis_per_dagrun` is per task across all dag runs.... that's a pain -- we should maybe deprecate that
@dstandish
We know it is a little confusing, sometimes we still mix them up, there are discussions in the PRs of the features, Max active tasks is the maximum active tasks across a DAG across all Dagruns (globally, there is also a per dag setting) and Max active tis per dag is for mapped tasks, a better name is max active (mapped) tis per dag per dagrun, @Asquator even asked the PR author somewhere but I couldn't find it.
--
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] Pessimistic task instance fetching with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2238938830
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
+
+ def running_tasks_group(*group_fields):
+ return (
+ select(TI, func.count("*").label("now_running"))
+ .where(TI.state.in_(EXECUTION_STATES))
+ .group_by(*group_fields)
+ .cte()
+ )
- query = (
+ def add_window_limit(query: Select, limit: Limit):
+ inner_query = (
+ query.add_columns(limit.window)
+ .join(limit.running_now_join, TI.id == limit.running_now_join.c.id)
+ .order_by(*priority_order)
+ .subquery()
+ )
+ return (
select(TI)
- .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
- .join(TI.dag_run)
- .where(DR.state == DagRunState.RUNNING)
- .join(TI.dag_model)
- .where(~DM.is_paused)
- .where(TI.state == TaskInstanceState.SCHEDULED)
- .where(DM.bundle_name.is_not(None))
- .options(selectinload(TI.dag_model))
- .order_by(-TI.priority_weight, DR.logical_date, TI.map_index)
+ .join(inner_query, TI.id == inner_query.c.id)
+ .where(
+ getattr(inner_query.c, limit.window.name) + limit.running_now_join.c.now_running
+ < limit.max_units
+ )
)
- if starved_pools:
- query = query.where(TI.pool.not_in(starved_pools))
+ running_total_tis_per_dagrun = running_tasks_group(TI.dag_id, TI.run_id)
+ running_tis_per_dag = running_tasks_group(TI.dag_id, TI.task_id)
+ running_total_tis_per_task_run = running_tasks_group(TI.dag_id, TI.run_id, TI.task_id)
+ running_tis_per_pool = running_tasks_group(TI.pool)
- if starved_dags:
- query = query.where(TI.dag_id.not_in(starved_dags))
+ total_tis_per_dagrun_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.run_id), order_by=priority_order)
+ .label("total_tis_per_dagrun_count")
+ )
+ tis_per_dag_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.task_id), order_by=priority_order)
+ .label("tis_per_dag_count")
+ )
+ mapped_tis_per_task_run_count = (
+ func.row_number()
+ .over(partition_by=(TI.dag_id, TI.run_id, TI.task_id), order_by=priority_order)
+ .label("mapped_tis_per_dagrun_count")
+ )
+ pool_slots_taken = (
+ func.sum(TI.pool_slots)
+ .over(partition_by=(TI.pool), order_by=priority_order)
+ .label("pool_slots_taken")
+ )
- if starved_tasks:
- query = query.where(tuple_(TI.dag_id, TI.task_id).not_in(starved_tasks))
+ limits = [
+ Limit(running_total_tis_per_dagrun, DagModel.max_active_tasks, total_tis_per_dagrun_count),
+ Limit(
+ running_tis_per_dag, TaskInstance.max_active_tis_per_dag, tis_per_dag_count
+ ), # TODO: Add to DB model: TaskInstance.max_active_tis_per_dag, DUMMY: DagModel.max_active_tasks
+ Limit(
+ running_total_tis_per_task_run, TI.max_tries, mapped_tis_per_task_run_count
+ ), # TODO: Add to DB model: TI.max_active_tis_per_dagrun, DUMMY: TI.max_tries
+ Limit(running_tis_per_pool, Pool.slots, pool_slots_taken),
+ ]
- if starved_tasks_task_dagrun_concurrency:
- query = query.where(
- tuple_(TI.dag_id, TI.run_id, TI.task_id).not_in(starved_tasks_task_dagrun_concurrency)
- )
+ for limit in limits:
+ query = add_window_limit(query, limit)
- query = query.limit(max_tis)
+ query = query.limit(max_tis)
- timer = Stats.timer("scheduler.critical_section_query_duration")
- timer.start()
+ timer = Stats.timer("scheduler.critical_section_query_duration")
+ timer.start()
- try:
- query = with_row_locks(query, of=TI, session=session, skip_locked=True)
- task_instances_to_examine: list[TI] = session.scalars(query).all()
+ try:
+ query = with_row_locks(query, of=TI, session=session, skip_locked=True)
+ task_instances_to_examine: list[TI] = session.scalars(query).all()
- timer.stop(send=True)
- except OperationalError as e:
- timer.stop(send=False)
- raise e
+ timer.stop(send=True)
+ except OperationalError as e:
+ timer.stop(send=False)
+ raise e
- # TODO[HA]: This was wrong before anyway, as it only looked at a sub-set of dags, not everything.
- # Stats.gauge('scheduler.tasks.pending', len(task_instances_to_examine))
+ # TODO[HA]: This was wrong before anyway, as it only looked at a sub-set of dags, not everything.
+ # Stats.gauge('scheduler.tasks.pending', len(task_instances_to_examine))
- if not task_instances_to_examine:
- self.log.debug("No tasks to consider for execution.")
- break
+ if not task_instances_to_examine:
+ self.log.debug("No tasks to consider for execution.")
+ return []
- # Put one task instance on each line
- task_instance_str = "\n".join(f"\t{x!r}" for x in task_instances_to_examine)
- self.log.info("%s tasks up for execution:\n%s", len(task_instances_to_examine), task_instance_str)
+ # Put one task instance on each line
+ task_instance_str = "\n".join(f"\t{x!r}" for x in task_instances_to_examine)
+ self.log.info("%s tasks up for execution:\n%s", len(task_instances_to_examine), task_instance_str)
- executor_slots_available: dict[ExecutorName, int] = {}
- # First get a mapping of executor names to slots they have available
- for executor in self.job.executors:
+ executor_slots_available: dict[ExecutorName, int] = {}
+ # First get a mapping of executor names to slots they have available
+ for executor in self.job.executors:
+ if TYPE_CHECKING:
+ # All executors should have a name if they are initted from the executor_loader.
+ # But we need to check for None to make mypy happy.
+ assert executor.name
+ executor_slots_available[executor.name] = executor.slots_available
+
+ for task_instance in task_instances_to_examine:
+ pool_name = task_instance.pool
+
+ pool_stats = pools.get(pool_name)
+ if not pool_stats:
+ self.log.warning("Tasks using non-existent pool '%s' will not be scheduled", pool_name)
+ starved_pools.add(pool_name)
+ continue
+
+ if executor_obj := self._try_to_load_executor(task_instance.executor):
if TYPE_CHECKING:
# All executors should have a name if they are initted from the executor_loader.
# But we need to check for None to make mypy happy.
- assert executor.name
- executor_slots_available[executor.name] = executor.slots_available
-
- for task_instance in task_instances_to_examine:
- pool_name = task_instance.pool
-
- pool_stats = pools.get(pool_name)
- if not pool_stats:
- self.log.warning("Tasks using non-existent pool '%s' will not be scheduled", pool_name)
- starved_pools.add(pool_name)
- continue
-
- # Make sure to emit metrics if pool has no starving tasks
- pool_num_starving_tasks.setdefault(pool_name, 0)
-
- pool_total = pool_stats["total"]
- open_slots = pool_stats["open"]
-
- if open_slots <= 0:
- self.log.info(
- "Not scheduling since there are %s open slots in pool %s", open_slots, pool_name
- )
- # Can't schedule any more since there are no more open slots.
- pool_num_starving_tasks[pool_name] += 1
- num_starving_tasks_total += 1
- starved_pools.add(pool_name)
- continue
-
- if task_instance.pool_slots > pool_total:
- self.log.warning(
- "Not executing %s. Requested pool slots (%s) are greater than "
- "total pool slots: '%s' for pool: %s.",
- task_instance,
- task_instance.pool_slots,
- pool_total,
- pool_name,
- )
-
- pool_num_starving_tasks[pool_name] += 1
- num_starving_tasks_total += 1
- starved_tasks.add((task_instance.dag_id, task_instance.task_id))
- continue
-
- if task_instance.pool_slots > open_slots:
- self.log.info(
- "Not executing %s since it requires %s slots "
- "but there are %s open slots in the pool %s.",
- task_instance,
- task_instance.pool_slots,
- open_slots,
- pool_name,
+ assert executor_obj.name
+ if executor_slots_available[executor_obj.name] <= 0:
+ self.log.debug(
Review Comment:
Not in the scope of this 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] Pessimistic task instance fetching with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2238957160
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -1023,15 +898,15 @@ def _execute(self) -> int | None:
def _update_dag_run_state_for_paused_dags(self, session: Session = NEW_SESSION) -> None:
try:
paused_runs = session.scalars(
- select(DagRun)
+ select(DR)
Review Comment:
I wanted it to be consistent across usages
--
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] Pessimistic task instance fetching with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2238913248
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
Review Comment:
Done
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
+
+ def running_tasks_group(*group_fields):
Review Comment:
Please add another suggestion to import `CTE` and I will commit them in batch
--
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] Pessimistic task instance fetching with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2238903981
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +375,165 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ .options(selectinload(TI.dag_model))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class Limit:
+ running_now_join: Subquery
+ max_units: Column
+ window: expression.ColumnElement
Review Comment:
Renamed to `WindowLimitDescriptor`
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250074710
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +377,171 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class LimitWindowDescriptor:
+ running_now_join: CTE
+ running_now_join_predicates: Collection[str]
+ limit_column: Column
+ window: expression.ColumnElement
+ limit_join_model: Base | None = None
+
+ def running_tasks_group(group_fields: Collection[Column], states: Collection[TaskInstanceState]=EXECUTION_STATES) -> CTE:
Review Comment:
Can you make these inner functions also rather private class functions - then they are also testable via pytest.
Alongside saying also: Please add appropriate pytests.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "jscheffl (via GitHub)" <gi...@apache.org>.
jscheffl commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250073077
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -554,6 +559,14 @@ class TaskInstance(Base, LoggingMixin):
innerjoin=True,
viewonly=True,
)
+ pool_model: Pool = relationship(
+ "Pool",
+ primaryjoin="TaskInstance.pool == Pool.pool",
+ foreign_keys=pool,
+ uselist=False,
+ innerjoin=True,
Review Comment:
At the moment there is no constraint on the relation between pool and TI. Would it mean the scheduler will not "see" the TIs when a pool is referenced that an admin has deleted? Do we need to add a constraint to pool table as well?
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250151423
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_add_ti_max_active_tis_per_dagrun.py:
##########
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add callback_state to deadline.
+
+Revision ID: 2f49f2dae90c
+Revises: f56f68b9e02f
+Create Date: 2025-07-28 16:39:01.181132
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2f49f2dae90c"
+down_revision = "f56f68b9e02f"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+ """Add callback_state to deadline."""
+ with op.batch_alter_table("task_instance", schema=None) as batch_op:
+ batch_op.add_column(sa.Column("max_active_tis_per_dag", sa.Integer, nullable=True))
Review Comment:
> During the migration as well the fields are not initialized. How does migrated data harm the scheduling?
Do you mean existing task instances in one of the execution states that have a concurrency limits? I can see the bug where the new scheduler will ignore the limits. We should manually reparse all the dags and add the fields to DB, it seems.
> Even if this is just adding a few bytes per tuple... is there an option to make the same function w/o DB change?
We need the fields in the DB to use them in a query. Regarding the table size, I see no other means to reduce it because we're in a system that heavily relies on SQL as its SSOT, and as long as we want features to be added, we must readily accept table extensions. task_instance is a flat table with a lot of params, some of them are duplicated across many TIs. Maybe it will be desirable to normalize it and extract the notorious task_template table or just a task_limits table in the long term.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250147441
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -373,292 +377,171 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) -
# If the pools are full, there is no point doing anything!
# If _somehow_ the pool is overfull, don't let the limit go negative - it breaks SQL
pool_slots_free = sum(max(0, pool["open"]) for pool in pools.values())
+ starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
+ starved_tasks: set[tuple[str, str]] = set()
if pool_slots_free == 0:
self.log.debug("All pools are full!")
return []
- max_tis = min(max_tis, pool_slots_free)
-
- starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0}
-
- # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks.
- concurrency_map = ConcurrencyMap()
- concurrency_map.load(session=session)
+ priority_order = [-TI.priority_weight, DR.logical_date, TI.map_index]
- # Number of tasks that cannot be scheduled because of no open slot in pool
- num_starving_tasks_total = 0
-
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
- starved_tasks: set[tuple[str, str]] = set()
- starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] = set()
-
- pool_num_starving_tasks: dict[str, int] = Counter()
+ query = (
+ select(TI)
+ .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
+ .join(TI.dag_run)
+ .where(DR.state == DagRunState.RUNNING)
+ .join(TI.dag_model)
+ .where(~DM.is_paused)
+ .where(TI.state == TaskInstanceState.SCHEDULED)
+ .where(DM.bundle_name.is_not(None))
+ )
- for loop_count in itertools.count(start=1):
- num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
- num_starved_tasks = len(starved_tasks)
- num_starved_tasks_task_dagrun_concurrency = len(starved_tasks_task_dagrun_concurrency)
+ @dataclass
+ class LimitWindowDescriptor:
+ running_now_join: CTE
+ running_now_join_predicates: Collection[str]
+ limit_column: Column
+ window: expression.ColumnElement
+ limit_join_model: Base | None = None
+
+ def running_tasks_group(group_fields: Collection[Column], states: Collection[TaskInstanceState]=EXECUTION_STATES) -> CTE:
Review Comment:
These auxiliary constructs will be taken out of the scheduler loop as soon as strategy injection is implemented for task selection.
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on code in PR #53492:
URL: https://github.com/apache/airflow/pull/53492#discussion_r2250146754
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -554,6 +559,14 @@ class TaskInstance(Base, LoggingMixin):
innerjoin=True,
viewonly=True,
)
+ pool_model: Pool = relationship(
+ "Pool",
+ primaryjoin="TaskInstance.pool == Pool.pool",
+ foreign_keys=pool,
+ uselist=False,
+ innerjoin=True,
Review Comment:
> At the moment there is no constraint on the relation between pool and TI. Would it mean the scheduler will not "see" the TIs when a pool is referenced that an admin has deleted? Do we need to add a constraint to pool table as well?
The only reason I added the relationship is easier join syntax in the query. I'm not sure about the correct alchemy construct for this case, we should probably remove the `foreign_keys` argument?
Regarding "ghost" pools - it's actually a deeper concern and I'm glad you raised it. If a task references a pool that does not exist, you usually get a corresponding message, but only when the task is "accidentally" pulled by the scheduler. Let's discuss the possible solutions:
1. Don't show any message if a task references a pool that does not exist by not pulling them at all, means non-indicative task banning (bad, current PR state)
2. Do an outer join, and for tasks with no corresponding pool emit the message - makes the scheduler optimistic as to inconsistency handling, with many inconsistent tasks we get more slots "flooded" with pool-ghosted tasks (bad, current main state)
3. Show the message on dag parsing (why would you want to create a DAG with such tasks?) (better, can it go wrong?)
4. Handle pools in a smarter way than it's done now - don't allow deleting pools when there are tasks associated with it, but allow to turn it off (mainly for disaster handling). It's difficult to implement because there is no table for tasks.
We have different options and we have to choose one to merge this change ASAP, avoid breaking the current behavior and leave space for introducing new changes in subsequent PRs. I'll leave this open
--
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] Pessimistic task instance scheduling with window functions [airflow]
Posted by "Asquator (via GitHub)" <gi...@apache.org>.
Asquator commented on PR #53492:
URL: https://github.com/apache/airflow/pull/53492#issuecomment-3293688742
Problems with this approach:
- It may be slow. Especially in MySQL where this query might take 15x more time than the optimistic one.
- Windows function don't really fit the use case, as we have different orthogonal limits that have to be enforced. Big thanks to @dstandish about paying attention and pushing us to find more flaws.
- Some starvation cases still persist as in certain cases we're at the mercy of a query planner. Consider the following example:
>Pool A: 16/32
Pool B: 16:32
>Task A: 32 mapped Pool A
Task B: 32 mapped Pool B
>Task A sorts first: 1-32
Task B sorts second: 33-64
>Task A gets 16 scheduled TIs
Task B: 16 potential TIs are starved
So we don't schedule 32 tasks alghough we technically could.
Take it further:
>Pools A,B,C,D 8/32
Tasks A,B,C,D accordingly, 32 each, sort in this order
>A: 1-32
B: 33-64
>C: 65-96
D: 97-128
>Dagrun max_active_tasks: 32
Ideally we want to schedule 32 tasks (8 of A, 8 of B, 8 of C, 8 of D)
>In reality we schedule 8 tasks of A and starve the other ready tasks
Some may not consider it starvation, but I'm personally not ready to accept this approach knowing it's both slower and doesn't completely solve the issue.
For any future developments I have to say that performance isn't completely terrible. We often got only 3x-6x slowdown.
The thing that bothers me the most is maintainability and unsolved cases. IMO WFs just don't fit for task selection with multidimensional limits. I'd rather try an exhaustive scan approach, see: https://github.com/apache/airflow/pull/55537
--
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