From 49e3d183e5bca2c6e475649eda4b6d1a42323dae Mon Sep 17 00:00:00 2001 From: "maciej.dembiczak" Date: Fri, 19 May 2023 09:49:54 +0200 Subject: [PATCH 1/3] assign filter output to variable --- airflow/models/taskinstance.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index bb7175a4a330e..2e416ce2eccf5 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2551,7 +2551,9 @@ def get_num_running_task_instances(self, session: Session, same_dagrun=False) -> TaskInstance.state == State.RUNNING, ) if same_dagrun: - num_running_task_instances_query.filter(TaskInstance.run_id == self.run_id) + num_running_task_instances_query = num_running_task_instances_query.filter( + TaskInstance.run_id == self.run_id + ) return num_running_task_instances_query.scalar() def init_run_context(self, raw: bool = False) -> None: From 2b1e74f6acb91a9337c188cbd1f6aa9df66bec11 Mon Sep 17 00:00:00 2001 From: "maciej.dembiczak" Date: Tue, 30 May 2023 14:28:29 +0200 Subject: [PATCH 2/3] add tests for max_active_tis_per_dagrun --- tests/models/test_taskinstance.py | 86 +++++++++++++++++++++---------- 1 file changed, 58 insertions(+), 28 deletions(-) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index 0dcfb632142de..9f62fa9833ca3 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -82,7 +82,7 @@ from airflow.utils.db import merge_conn from airflow.utils.module_loading import qualname from airflow.utils.session import create_session, provide_session -from airflow.utils.state import State, TaskInstanceState +from airflow.utils.state import State, TaskInstanceState, DagRunState from airflow.utils.task_group import TaskGroup from airflow.utils.types import DagRunType from airflow.utils.xcom import XCOM_RETURN_KEY @@ -1770,6 +1770,59 @@ def test_get_num_running_task_instances(self, create_task_instance): assert 1 == ti2.get_num_running_task_instances(session=session) assert 1 == ti3.get_num_running_task_instances(session=session) + def test_get_num_running_task_instances_per_dagrun(self, create_task_instance, dag_maker): + session = settings.Session() + + @task + def dummy_mapped_task(x: int): + return x + + with dag_maker(dag_id="test_dag") as dag: + dummy_mapped_task.expand(x=[1, 2]) + dagrun1 = dag.create_dagrun(execution_date=timezone.utcnow(), + state=DagRunState.RUNNING, + run_id="run_id_1", + session=session) + + ti1 = dagrun1.task_instances[0] + ti2 = dagrun1.task_instances[1] + + with dag_maker(dag_id="test_dag") as dag: + dummy_mapped_task.expand(x=[3]) + dagrun2 = dag.create_dagrun(execution_date=timezone.utcnow(), + state=DagRunState.RUNNING, + run_id="run_id_2", + session=session) + + ti3 = dagrun2.task_instances[0] + + with dag_maker(dag_id="test_dag") as dag: + dummy_mapped_task.expand(x=[4]) + dagrun3 = dag.create_dagrun(execution_date=timezone.utcnow(), + state=DagRunState.RUNNING, + run_id="run_id_3", + session=session) + + ti4 = dagrun3.task_instances[0] + + assert ti1 in session + assert ti2 in session + assert ti3 in session + assert ti4 in session + + ti1.state = State.RUNNING + ti2.state = State.QUEUED + ti3.state = State.RUNNING + ti4.state = State.RUNNING + + session.commit() + + assert 1 == ti1.get_num_running_task_instances(session=session, same_dagrun=True) + assert 1 == ti2.get_num_running_task_instances(session=session, same_dagrun=True) + assert 3 == ti3.get_num_running_task_instances(session=session) + assert 1 == ti4.get_num_running_task_instances(session=session, same_dagrun=True) + + def test_log_url(self, create_task_instance): ti = create_task_instance(dag_id="dag", task_id="op", execution_date=timezone.datetime(2018, 1, 1)) @@ -1903,7 +1956,6 @@ def test_email_alert_with_filenotfound_config(self, mock_send_email, dag_maker): @patch("airflow.models.taskinstance.send_email") def test_failure_mapped_taskflow(self, mock_send_email, dag_maker, session, task_id): with dag_maker(session=session) as dag: - @dag.task(email="to") def test_email_alert(x): raise RuntimeError("Fail please") @@ -2005,10 +2057,10 @@ def test_outlet_datasets(self, create_task_instance): assert session.query(DatasetDagRunQueue.target_dag_id).filter_by( dataset_id=event.dataset.id ).order_by(DatasetDagRunQueue.target_dag_id).all() == [ - ("dataset_consumes_1",), - ("dataset_consumes_1_and_2",), - ("dataset_consumes_1_never_scheduled",), - ] + ("dataset_consumes_1",), + ("dataset_consumes_1_and_2",), + ("dataset_consumes_1_never_scheduled",), + ] # check that one event record created for dataset1 and this TI assert session.query(DatasetModel.uri).join(DatasetEvent.dataset).filter( @@ -2113,7 +2165,6 @@ def test_changing_of_dataset_when_ddrq_is_already_populated(self, dag_maker, ses from airflow import Dataset with dag_maker(schedule=None, serialized=True) as dag1: - @task(outlets=Dataset("test/1")) def test_task1(): print(1) @@ -2124,7 +2175,6 @@ def test_task1(): test_task1 = dag1.get_task("test_task1") with dag_maker(dag_id="testdag", schedule=[Dataset("test/1")], serialized=True): - @task def test_task2(): print(1) @@ -2135,7 +2185,6 @@ def test_task2(): ti.run() # Change the dataset. with dag_maker(dag_id="testdag", schedule=[Dataset("test2/1")], serialized=True): - @task def test_task2(): print(1) @@ -3265,7 +3314,6 @@ def setup_class(self): def test_not_recorded_if_leaf(self, dag_maker, xcom_value): """Return value should not be recorded if there are no downstreams.""" with dag_maker(dag_id="test_not_recorded_for_unused") as dag: - @dag.task() def push_something(): return xcom_value @@ -3281,7 +3329,6 @@ def push_something(): def test_not_recorded_if_not_used(self, dag_maker, xcom_value): """Return value should not be recorded if no downstreams are mapped.""" with dag_maker(dag_id="test_not_recorded_for_unused") as dag: - @dag.task() def push_something(): return xcom_value @@ -3302,7 +3349,6 @@ def completely_different(): def test_not_recorded_if_irrelevant(self, dag_maker, xcom_1, xcom_4): """Return value should only be recorded if a mapped downstream uses the it.""" with dag_maker(dag_id="test_not_recorded_for_unused") as dag: - @dag.task() def push_1(): return xcom_1 @@ -3355,7 +3401,6 @@ def tg(arg): def test_expand_error_if_unmappable_type(self, dag_maker, return_value, exception_type, error_message): """If an unmappable return value is used for expand(), fail the task that pushed the XCom.""" with dag_maker(dag_id="test_expand_error_if_unmappable_type") as dag: - @dag.task() def push_something(): return return_value @@ -3390,7 +3435,6 @@ def test_expand_kwargs_error_if_unmappable_type( ): """If an unmappable return value is used for expand_kwargs(), fail the task that pushed the XCom.""" with dag_maker(dag_id="test_expand_kwargs_error_if_unmappable_type") as dag: - @dag.task() def push(): return return_value @@ -3421,7 +3465,6 @@ def test_task_group_expand_error_if_unmappable_type( ): """If an unmappable return value is used , fail the task that pushed the XCom.""" with dag_maker(dag_id="test_task_group_expand_error_if_unmappable_type") as dag: - @dag.task() def push(): return return_value @@ -3456,7 +3499,6 @@ def test_task_group_expand_kwargs_error_if_unmappable_type( ): """If an unmappable return value is used, fail the task that pushed the XCom.""" with dag_maker(dag_id="test_task_group_expand_kwargs_error_if_unmappable_type") as dag: - @dag.task() def push(): return return_value @@ -3568,7 +3610,6 @@ def pull(arg1, arg2): def test_error_if_upstream_does_not_push(self, dag_maker): """Fail the upstream task if it fails to push the XCom used for task mapping.""" with dag_maker(dag_id="test_not_recorded_for_unused") as dag: - @dag.task(do_xcom_push=False) def push_something(): return [1, 2] @@ -3591,7 +3632,6 @@ def pull_something(value): def test_error_if_unmappable_length(self, dag_maker): """If an unmappable return value is used to map, fail the task that pushed the XCom.""" with dag_maker(dag_id="test_not_recorded_for_unused") as dag: - @dag.task() def push_something(): return [1, 2] @@ -3620,7 +3660,6 @@ def pull_something(value): def test_written_task_map(self, dag_maker, xcom_value, expected_length, expected_keys): """Return value should be recorded in TaskMap if it's used by a downstream to map.""" with dag_maker(dag_id="test_written_task_map") as dag: - @dag.task() def push_something(): return xcom_value @@ -3656,7 +3695,6 @@ def test_map_literal(self, literal, expected_outputs, dag_maker, session): outputs = [] with dag_maker(dag_id="literal", session=session) as dag: - @dag.task def show(value): outputs.append(value) @@ -3689,7 +3727,6 @@ def test_map_xcom(self, upstream_return, expected_outputs, dag_maker, session): outputs = [] with dag_maker(dag_id="xcom", session=session) as dag: - @dag.task def emit(): return upstream_return @@ -3760,7 +3797,6 @@ def test_map_product_same(self, dag_maker, session): outputs = [] with dag_maker(dag_id="product_same", session=session) as dag: - @dag.task def emit_numbers(): return [1, 2] @@ -3793,7 +3829,6 @@ def test_map_literal_cross_product(self, dag_maker, session): outputs = [] with dag_maker(dag_id="product_same_types", session=session) as dag: - @dag.task def show(a, b): outputs.append((a, b)) @@ -4028,7 +4063,6 @@ class SimpleBashOperator(BashOperator): template_fields = () with dag_maker(dag_id="product_same_types", session=session) as dag: - @dag.task def get_extra_env(): return [{"foo": "bar"}, {"foo": "biz"}] @@ -4048,7 +4082,6 @@ def test_mapped_task_does_not_error_in_mini_scheduler_if_upstreams_are_not_done( not marked as `upstream_failed' """ with dag_maker() as dag: - @dag.task def second_task(): return [0, 1, 2] @@ -4090,7 +4123,6 @@ def test_empty_operator_is_not_considered_in_mini_scheduler(dag_maker, caplog, s submit them directly to worker. """ with dag_maker() as dag: - @dag.task def first_task(): print(2) @@ -4130,7 +4162,6 @@ def second_task(): def test_mapped_task_expands_in_mini_scheduler_if_upstreams_are_done(dag_maker, caplog, session): """Test that mini scheduler expands mapped task""" with dag_maker() as dag: - @dag.task def second_task(): return [0, 1, 2] @@ -4170,7 +4201,6 @@ def last_task(): def test_mini_scheduler_not_skip_mapped_downstream_until_all_upstreams_finish(dag_maker, session): with dag_maker(session=session): - @task def generate() -> list[list[int]]: return [] From a9c4573943708e213fc0b71e718d3de9ef0a3cfa Mon Sep 17 00:00:00 2001 From: "maciej.dembiczak" Date: Tue, 30 May 2023 15:33:27 +0200 Subject: [PATCH 3/3] fix tests and code style --- tests/models/test_taskinstance.py | 115 +++++++++++++++++------------- 1 file changed, 67 insertions(+), 48 deletions(-) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index 9f62fa9833ca3..3ae696a27c142 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -82,7 +82,7 @@ from airflow.utils.db import merge_conn from airflow.utils.module_loading import qualname from airflow.utils.session import create_session, provide_session -from airflow.utils.state import State, TaskInstanceState, DagRunState +from airflow.utils.state import DagRunState, State, TaskInstanceState from airflow.utils.task_group import TaskGroup from airflow.utils.types import DagRunType from airflow.utils.xcom import XCOM_RETURN_KEY @@ -1773,55 +1773,51 @@ def test_get_num_running_task_instances(self, create_task_instance): def test_get_num_running_task_instances_per_dagrun(self, create_task_instance, dag_maker): session = settings.Session() - @task - def dummy_mapped_task(x: int): - return x - - with dag_maker(dag_id="test_dag") as dag: - dummy_mapped_task.expand(x=[1, 2]) - dagrun1 = dag.create_dagrun(execution_date=timezone.utcnow(), - state=DagRunState.RUNNING, - run_id="run_id_1", - session=session) - - ti1 = dagrun1.task_instances[0] - ti2 = dagrun1.task_instances[1] - - with dag_maker(dag_id="test_dag") as dag: - dummy_mapped_task.expand(x=[3]) - dagrun2 = dag.create_dagrun(execution_date=timezone.utcnow(), - state=DagRunState.RUNNING, - run_id="run_id_2", - session=session) - - ti3 = dagrun2.task_instances[0] + with dag_maker(dag_id="test_dag"): + MockOperator.partial(task_id="task_1").expand_kwargs([{"a": 1, "b": 2}, {"a": 3, "b": 4}]) + MockOperator.partial(task_id="task_2").expand_kwargs([{"a": 1, "b": 2}]) + MockOperator.partial(task_id="task_3").expand_kwargs([{"a": 1, "b": 2}]) - with dag_maker(dag_id="test_dag") as dag: - dummy_mapped_task.expand(x=[4]) - dagrun3 = dag.create_dagrun(execution_date=timezone.utcnow(), - state=DagRunState.RUNNING, - run_id="run_id_3", - session=session) - - ti4 = dagrun3.task_instances[0] - - assert ti1 in session - assert ti2 in session - assert ti3 in session - assert ti4 in session + dr1 = dag_maker.create_dagrun( + execution_date=timezone.utcnow(), state=DagRunState.RUNNING, run_id="run_id_1", session=session + ) + tis1 = {(ti.task_id, ti.map_index): ti for ti in dr1.task_instances} + print(f"tis1: {tis1}") - ti1.state = State.RUNNING - ti2.state = State.QUEUED - ti3.state = State.RUNNING - ti4.state = State.RUNNING + dr2 = dag_maker.create_dagrun( + execution_date=timezone.utcnow(), state=DagRunState.RUNNING, run_id="run_id_2", session=session + ) + tis2 = {(ti.task_id, ti.map_index): ti for ti in dr2.task_instances} + + assert tis1[("task_1", 0)] in session + assert tis1[("task_1", 1)] in session + assert tis1[("task_2", 0)] in session + assert tis1[("task_3", 0)] in session + assert tis2[("task_1", 0)] in session + assert tis2[("task_1", 1)] in session + assert tis2[("task_2", 0)] in session + assert tis2[("task_3", 0)] in session + + tis1[("task_1", 0)].state = State.RUNNING + tis1[("task_1", 1)].state = State.QUEUED + tis1[("task_2", 0)].state = State.RUNNING + tis1[("task_3", 0)].state = State.RUNNING + tis2[("task_1", 0)].state = State.RUNNING + tis2[("task_1", 1)].state = State.QUEUED + tis2[("task_2", 0)].state = State.RUNNING + tis2[("task_3", 0)].state = State.RUNNING session.commit() - assert 1 == ti1.get_num_running_task_instances(session=session, same_dagrun=True) - assert 1 == ti2.get_num_running_task_instances(session=session, same_dagrun=True) - assert 3 == ti3.get_num_running_task_instances(session=session) - assert 1 == ti4.get_num_running_task_instances(session=session, same_dagrun=True) + assert 1 == tis1[("task_1", 0)].get_num_running_task_instances(session=session, same_dagrun=True) + assert 1 == tis1[("task_1", 1)].get_num_running_task_instances(session=session, same_dagrun=True) + assert 2 == tis1[("task_2", 0)].get_num_running_task_instances(session=session) + assert 1 == tis1[("task_3", 0)].get_num_running_task_instances(session=session, same_dagrun=True) + assert 1 == tis2[("task_1", 0)].get_num_running_task_instances(session=session, same_dagrun=True) + assert 1 == tis2[("task_1", 1)].get_num_running_task_instances(session=session, same_dagrun=True) + assert 2 == tis2[("task_2", 0)].get_num_running_task_instances(session=session) + assert 1 == tis2[("task_3", 0)].get_num_running_task_instances(session=session, same_dagrun=True) def test_log_url(self, create_task_instance): ti = create_task_instance(dag_id="dag", task_id="op", execution_date=timezone.datetime(2018, 1, 1)) @@ -1956,6 +1952,7 @@ def test_email_alert_with_filenotfound_config(self, mock_send_email, dag_maker): @patch("airflow.models.taskinstance.send_email") def test_failure_mapped_taskflow(self, mock_send_email, dag_maker, session, task_id): with dag_maker(session=session) as dag: + @dag.task(email="to") def test_email_alert(x): raise RuntimeError("Fail please") @@ -2057,10 +2054,10 @@ def test_outlet_datasets(self, create_task_instance): assert session.query(DatasetDagRunQueue.target_dag_id).filter_by( dataset_id=event.dataset.id ).order_by(DatasetDagRunQueue.target_dag_id).all() == [ - ("dataset_consumes_1",), - ("dataset_consumes_1_and_2",), - ("dataset_consumes_1_never_scheduled",), - ] + ("dataset_consumes_1",), + ("dataset_consumes_1_and_2",), + ("dataset_consumes_1_never_scheduled",), + ] # check that one event record created for dataset1 and this TI assert session.query(DatasetModel.uri).join(DatasetEvent.dataset).filter( @@ -2165,6 +2162,7 @@ def test_changing_of_dataset_when_ddrq_is_already_populated(self, dag_maker, ses from airflow import Dataset with dag_maker(schedule=None, serialized=True) as dag1: + @task(outlets=Dataset("test/1")) def test_task1(): print(1) @@ -2175,6 +2173,7 @@ def test_task1(): test_task1 = dag1.get_task("test_task1") with dag_maker(dag_id="testdag", schedule=[Dataset("test/1")], serialized=True): + @task def test_task2(): print(1) @@ -2185,6 +2184,7 @@ def test_task2(): ti.run() # Change the dataset. with dag_maker(dag_id="testdag", schedule=[Dataset("test2/1")], serialized=True): + @task def test_task2(): print(1) @@ -3314,6 +3314,7 @@ def setup_class(self): def test_not_recorded_if_leaf(self, dag_maker, xcom_value): """Return value should not be recorded if there are no downstreams.""" with dag_maker(dag_id="test_not_recorded_for_unused") as dag: + @dag.task() def push_something(): return xcom_value @@ -3329,6 +3330,7 @@ def push_something(): def test_not_recorded_if_not_used(self, dag_maker, xcom_value): """Return value should not be recorded if no downstreams are mapped.""" with dag_maker(dag_id="test_not_recorded_for_unused") as dag: + @dag.task() def push_something(): return xcom_value @@ -3349,6 +3351,7 @@ def completely_different(): def test_not_recorded_if_irrelevant(self, dag_maker, xcom_1, xcom_4): """Return value should only be recorded if a mapped downstream uses the it.""" with dag_maker(dag_id="test_not_recorded_for_unused") as dag: + @dag.task() def push_1(): return xcom_1 @@ -3401,6 +3404,7 @@ def tg(arg): def test_expand_error_if_unmappable_type(self, dag_maker, return_value, exception_type, error_message): """If an unmappable return value is used for expand(), fail the task that pushed the XCom.""" with dag_maker(dag_id="test_expand_error_if_unmappable_type") as dag: + @dag.task() def push_something(): return return_value @@ -3435,6 +3439,7 @@ def test_expand_kwargs_error_if_unmappable_type( ): """If an unmappable return value is used for expand_kwargs(), fail the task that pushed the XCom.""" with dag_maker(dag_id="test_expand_kwargs_error_if_unmappable_type") as dag: + @dag.task() def push(): return return_value @@ -3465,6 +3470,7 @@ def test_task_group_expand_error_if_unmappable_type( ): """If an unmappable return value is used , fail the task that pushed the XCom.""" with dag_maker(dag_id="test_task_group_expand_error_if_unmappable_type") as dag: + @dag.task() def push(): return return_value @@ -3499,6 +3505,7 @@ def test_task_group_expand_kwargs_error_if_unmappable_type( ): """If an unmappable return value is used, fail the task that pushed the XCom.""" with dag_maker(dag_id="test_task_group_expand_kwargs_error_if_unmappable_type") as dag: + @dag.task() def push(): return return_value @@ -3610,6 +3617,7 @@ def pull(arg1, arg2): def test_error_if_upstream_does_not_push(self, dag_maker): """Fail the upstream task if it fails to push the XCom used for task mapping.""" with dag_maker(dag_id="test_not_recorded_for_unused") as dag: + @dag.task(do_xcom_push=False) def push_something(): return [1, 2] @@ -3632,6 +3640,7 @@ def pull_something(value): def test_error_if_unmappable_length(self, dag_maker): """If an unmappable return value is used to map, fail the task that pushed the XCom.""" with dag_maker(dag_id="test_not_recorded_for_unused") as dag: + @dag.task() def push_something(): return [1, 2] @@ -3660,6 +3669,7 @@ def pull_something(value): def test_written_task_map(self, dag_maker, xcom_value, expected_length, expected_keys): """Return value should be recorded in TaskMap if it's used by a downstream to map.""" with dag_maker(dag_id="test_written_task_map") as dag: + @dag.task() def push_something(): return xcom_value @@ -3695,6 +3705,7 @@ def test_map_literal(self, literal, expected_outputs, dag_maker, session): outputs = [] with dag_maker(dag_id="literal", session=session) as dag: + @dag.task def show(value): outputs.append(value) @@ -3727,6 +3738,7 @@ def test_map_xcom(self, upstream_return, expected_outputs, dag_maker, session): outputs = [] with dag_maker(dag_id="xcom", session=session) as dag: + @dag.task def emit(): return upstream_return @@ -3797,6 +3809,7 @@ def test_map_product_same(self, dag_maker, session): outputs = [] with dag_maker(dag_id="product_same", session=session) as dag: + @dag.task def emit_numbers(): return [1, 2] @@ -3829,6 +3842,7 @@ def test_map_literal_cross_product(self, dag_maker, session): outputs = [] with dag_maker(dag_id="product_same_types", session=session) as dag: + @dag.task def show(a, b): outputs.append((a, b)) @@ -4063,6 +4077,7 @@ class SimpleBashOperator(BashOperator): template_fields = () with dag_maker(dag_id="product_same_types", session=session) as dag: + @dag.task def get_extra_env(): return [{"foo": "bar"}, {"foo": "biz"}] @@ -4082,6 +4097,7 @@ def test_mapped_task_does_not_error_in_mini_scheduler_if_upstreams_are_not_done( not marked as `upstream_failed' """ with dag_maker() as dag: + @dag.task def second_task(): return [0, 1, 2] @@ -4123,6 +4139,7 @@ def test_empty_operator_is_not_considered_in_mini_scheduler(dag_maker, caplog, s submit them directly to worker. """ with dag_maker() as dag: + @dag.task def first_task(): print(2) @@ -4162,6 +4179,7 @@ def second_task(): def test_mapped_task_expands_in_mini_scheduler_if_upstreams_are_done(dag_maker, caplog, session): """Test that mini scheduler expands mapped task""" with dag_maker() as dag: + @dag.task def second_task(): return [0, 1, 2] @@ -4201,6 +4219,7 @@ def last_task(): def test_mini_scheduler_not_skip_mapped_downstream_until_all_upstreams_finish(dag_maker, session): with dag_maker(session=session): + @task def generate() -> list[list[int]]: return []