ephraimbuddy commented on code in PR #61769:
URL: https://github.com/apache/airflow/pull/61769#discussion_r3166531903
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -1924,22 +1926,53 @@ def xcom_pull(
@provide_session
def get_num_running_task_instances(self, session: Session, same_dagrun:
bool = False) -> int:
- """Return Number of running TIs from the DB."""
- # .count() is inefficient
- num_running_task_instances_query = (
+ """Count running TIs from the DB."""
+ warnings.warn(
+ "This function is deprecated and will be removed in Airflow.",
+ RemovedInAirflow4Warning,
+ stacklevel=2,
+ )
+ return self._get_num_task_instances_of_state(
+ [TaskInstanceState.RUNNING],
+ same_dagrun=same_dagrun,
+ session=session,
+ )
+
+ def get_num_active_task_instances(self, *, same_dagrun: bool = False,
session: Session) -> int:
+ """
+ Count active TIs for this task from the DB.
+
+ Counts task instances in running, queued, or deferred state.
Review Comment:
```suggestion
Counts task instances in running or deferred state.
```
##########
airflow-core/src/airflow/ti_deps/deps/task_concurrency_dep.py:
##########
@@ -36,13 +36,13 @@ def _get_dep_statuses(self, ti, session, dep_context):
if (
ti.task.max_active_tis_per_dag is not None
- and ti.get_num_running_task_instances(session) >=
ti.task.max_active_tis_per_dag
+ and ti.get_num_active_task_instances(session) >=
ti.task.max_active_tis_per_dag
Review Comment:
```suggestion
and ti.get_num_active_task_instances(session=session) >=
ti.task.max_active_tis_per_dag
```
keyword-only method?
##########
airflow-core/src/airflow/ti_deps/deps/task_concurrency_dep.py:
##########
@@ -36,13 +36,13 @@ def _get_dep_statuses(self, ti, session, dep_context):
if (
ti.task.max_active_tis_per_dag is not None
- and ti.get_num_running_task_instances(session) >=
ti.task.max_active_tis_per_dag
+ and ti.get_num_active_task_instances(session) >=
ti.task.max_active_tis_per_dag
):
yield self._failing_status(reason="The max task concurrency has
been reached.")
return
if (
ti.task.max_active_tis_per_dagrun is not None
- and ti.get_num_running_task_instances(session, same_dagrun=True)
+ and ti.get_num_active_task_instances(session, same_dagrun=True)
Review Comment:
```suggestion
and ti.get_num_active_task_instances(session=session,
same_dagrun=True)
```
##########
airflow-core/tests/unit/models/test_taskinstance.py:
##########
@@ -1580,11 +1580,14 @@ def test_get_num_running_task_instances(self,
dag_maker, create_task_instance):
assert ti3 in session
session.commit()
- assert ti1.get_num_running_task_instances(session=session) == 1
- assert ti2.get_num_running_task_instances(session=session) == 1
- assert ti3.get_num_running_task_instances(session=session) == 1
+ # get_num_active_task_instances now counts RUNNING + QUEUED + DEFERRED.
+ # ti1 (RUNNING) and ti2 (QUEUED) share the same dag_id/task_id, so both
Review Comment:
I think `queued` is now out of 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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
