diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml index cfadcb16fd28a..837fedfac2b77 100644 --- a/airflow/config_templates/config.yml +++ b/airflow/config_templates/config.yml @@ -2418,15 +2418,6 @@ scheduler: type: integer default: "20" see_also: ":ref:`scheduler:ha:tunables`" - schedule_after_task_execution: - description: | - Should the Task supervisor process perform a "mini scheduler" to attempt to schedule more tasks of the - same DAG. Leaving this on will mean tasks in the same DAG execute quicker, but might starve out other - dags in some circumstances - example: ~ - version_added: 2.0.0 - type: boolean - default: "True" parsing_pre_import_modules: description: | The scheduler reads dag files to extract the airflow modules that are going to be used, diff --git a/airflow/jobs/local_task_job_runner.py b/airflow/jobs/local_task_job_runner.py index 599493ea58c4f..5d1f15fe8f8a5 100644 --- a/airflow/jobs/local_task_job_runner.py +++ b/airflow/jobs/local_task_job_runner.py @@ -253,7 +253,7 @@ def handle_task_exit(self, return_code: int) -> None: self.terminating = True self._log_return_code_metric(return_code) - if is_deferral := return_code == TaskReturnCode.DEFERRED.value: + if return_code == TaskReturnCode.DEFERRED.value: self.log.info("Task exited with return code %s (task deferral)", return_code) _set_task_deferred_context_var() else: @@ -262,10 +262,6 @@ def handle_task_exit(self, return_code: int) -> None: message += ". For more information, see https://airflow.apache.org/docs/apache-airflow/stable/troubleshooting.html#LocalTaskJob-killed" self.log.info(message) - if not (self.task_instance.test_mode or is_deferral): - if conf.getboolean("scheduler", "schedule_after_task_execution", fallback=True): - self.task_instance.schedule_downstream_tasks(max_tis_per_query=self.job.max_tis_per_query) - def on_kill(self): self.task_runner.terminate() self.task_runner.on_finish() diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index c525a40a14ab5..1ca90e7f1c4df 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -3639,100 +3639,6 @@ def ti_selector_condition(cls, vals: Collection[str | tuple[str, int]]) -> Colum return filters[0] return or_(*filters) - @classmethod - @provide_session - def _schedule_downstream_tasks( - cls, - ti: TaskInstance | TaskInstancePydantic, - session: Session = NEW_SESSION, - max_tis_per_query: int | None = None, - ): - from sqlalchemy.exc import OperationalError - - from airflow.models.dagrun import DagRun - - try: - # Re-select the row with a lock - dag_run = with_row_locks( - session.query(DagRun).filter_by( - dag_id=ti.dag_id, - run_id=ti.run_id, - ), - session=session, - skip_locked=True, - ).one_or_none() - - if not dag_run: - cls.logger().debug("Skip locked rows, rollback") - session.rollback() - return - - task = ti.task - if TYPE_CHECKING: - assert task - assert task.dag - - # Previously, this section used task.dag.partial_subset to retrieve a partial DAG. - # However, this approach is unsafe as it can result in incomplete or incorrect task execution, - # leading to potential bad cases. As a result, the operation has been removed. - # For more details, refer to the discussion in PR #[https://github.com/apache/airflow/pull/42582]. - dag_run.dag = task.dag - info = dag_run.task_instance_scheduling_decisions(session) - - skippable_task_ids = { - task_id for task_id in task.dag.task_ids if task_id not in task.downstream_task_ids - } - - schedulable_tis = [ - ti - for ti in info.schedulable_tis - if ti.task_id not in skippable_task_ids - and not ( - ti.task.inherits_from_empty_operator - and not ti.task.on_execute_callback - and not ti.task.on_success_callback - and not ti.task.outlets - ) - ] - for schedulable_ti in schedulable_tis: - if getattr(schedulable_ti, "task", None) is None: - schedulable_ti.task = task.dag.get_task(schedulable_ti.task_id) - - num = dag_run.schedule_tis(schedulable_tis, session=session, max_tis_per_query=max_tis_per_query) - cls.logger().info("%d downstream tasks scheduled from follow-on schedule check", num) - - session.flush() - - except OperationalError as e: - # Any kind of DB error here is _non fatal_ as this block is just an optimisation. - cls.logger().warning( - "Skipping mini scheduling run due to exception: %s", - e.statement, - exc_info=True, - ) - session.rollback() - - @provide_session - def schedule_downstream_tasks(self, session: Session = NEW_SESSION, max_tis_per_query: int | None = None): - """ - Schedule downstream tasks of this task instance. - - :meta: private - """ - try: - return TaskInstance._schedule_downstream_tasks( - ti=self, session=session, max_tis_per_query=max_tis_per_query - ) - except Exception: - self.log.exception( - "Error scheduling downstream tasks. Skipping it as this is entirely optional optimisation. " - "There might be various reasons for it, please take a look at the stack trace to figure " - "out if the root cause can be diagnosed and fixed. See the issue " - "https://github.com/apache/airflow/issues/39717 for details and an example problem. If you " - "would like to get help in solving root cause, open discussion with all details with your " - "managed service support or in Airflow repository." - ) - def get_relevant_upstream_map_indexes( self, upstream: Operator, diff --git a/airflow/serialization/pydantic/taskinstance.py b/airflow/serialization/pydantic/taskinstance.py index d5573922b839d..b3b60383ea6c1 100644 --- a/airflow/serialization/pydantic/taskinstance.py +++ b/airflow/serialization/pydantic/taskinstance.py @@ -467,16 +467,6 @@ def check_and_change_state_before_execution( session=session, ) - def schedule_downstream_tasks(self, session: Session | None = None, max_tis_per_query: int | None = None): - """ - Schedule downstream tasks of this task instance. - - :meta: private - """ - # we should not schedule downstream tasks with Pydantic model because it will not be able to - # get the DAG object (we do not serialize it currently). - return - def command_as_list( self, mark_success: bool = False, diff --git a/docs/apache-airflow/administration-and-deployment/scheduler.rst b/docs/apache-airflow/administration-and-deployment/scheduler.rst index e269740819901..476e2e662790e 100644 --- a/docs/apache-airflow/administration-and-deployment/scheduler.rst +++ b/docs/apache-airflow/administration-and-deployment/scheduler.rst @@ -382,8 +382,3 @@ However, you can also look at other non-performance-related scheduler configurat in the loop. i.e. if it scheduled something then it will start the next loop iteration straight away. This parameter is badly named (historical reasons) and it will be renamed in the future with deprecation of the current name. - -- :ref:`config:scheduler__schedule_after_task_execution` - Should the Task supervisor process perform a "mini scheduler" to attempt to schedule more tasks of - the same DAG. Leaving this on will mean tasks in the same DAG execute quicker, - but might starve out other DAGs in some circumstances. diff --git a/providers/src/airflow/providers/edge/cli/edge_command.py b/providers/src/airflow/providers/edge/cli/edge_command.py index 487b3adde70cd..69f12acc3fad2 100644 --- a/providers/src/airflow/providers/edge/cli/edge_command.py +++ b/providers/src/airflow/providers/edge/cli/edge_command.py @@ -77,8 +77,6 @@ def force_use_internal_api_on_edge_worker(): os.environ["AIRFLOW_ENABLE_AIP_44"] = "True" os.environ["AIRFLOW__CORE__INTERNAL_API_URL"] = api_url InternalApiConfig.set_use_internal_api("edge-worker") - # Disable mini-scheduler post task execution and leave next task schedule to core scheduler - os.environ["AIRFLOW__SCHEDULER__SCHEDULE_AFTER_TASK_EXECUTION"] = "False" force_use_internal_api_on_edge_worker() diff --git a/providers/tests/cncf/kubernetes/decorators/test_kubernetes.py b/providers/tests/cncf/kubernetes/decorators/test_kubernetes.py index 9545e9adbf3d7..8f412c81b7f3c 100644 --- a/providers/tests/cncf/kubernetes/decorators/test_kubernetes.py +++ b/providers/tests/cncf/kubernetes/decorators/test_kubernetes.py @@ -215,28 +215,3 @@ def f(): assert len(dag.task_group.children) == 1 teardown_task = dag.task_group.children["f"] assert teardown_task.is_teardown - - -# Database isolation mode does not support mini-scheduler -@pytest.mark.skip_if_database_isolation_mode -def test_kubernetes_with_mini_scheduler( - dag_maker, session, mock_create_pod: mock.Mock, mock_hook: mock.Mock -) -> None: - with dag_maker(session=session): - - @task.kubernetes( - image="python:3.10-slim-buster", - in_cluster=False, - cluster_context="default", - config_file="/tmp/fake_file", - ) - def f(arg1, arg2, kwarg1=None, kwarg2=None): - return {"key1": "value1", "key2": "value2"} - - f1 = f.override(task_id="my_task_id", do_xcom_push=True)("arg1", "arg2", kwarg1="kwarg1") - f.override(task_id="my_task_id2", do_xcom_push=False)("arg1", "arg2", kwarg1=f1) - - dr = dag_maker.create_dagrun() - (ti, _) = dr.task_instances - # check that mini-scheduler works - ti.schedule_downstream_tasks() diff --git a/tests/jobs/test_local_task_job.py b/tests/jobs/test_local_task_job.py index 7ee037478833a..936f773e55fc8 100644 --- a/tests/jobs/test_local_task_job.py +++ b/tests/jobs/test_local_task_job.py @@ -37,11 +37,9 @@ from airflow.executors.sequential_executor import SequentialExecutor from airflow.jobs.job import Job, run_job from airflow.jobs.local_task_job_runner import SIGSEGV_MESSAGE, LocalTaskJobRunner -from airflow.jobs.scheduler_job_runner import SchedulerJobRunner from airflow.listeners.listener import get_listener_manager from airflow.models.dag import DAG from airflow.models.dagbag import DagBag -from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskinstance import TaskInstance from airflow.operators.empty import EmptyOperator from airflow.providers.standard.operators.python import PythonOperator @@ -845,162 +843,6 @@ def send_signal(ti, signal_sent, sig): lines = f.readlines() assert len(lines) == 0 - @pytest.mark.skip_if_database_isolation_mode # Does not work in db isolation mode - @pytest.mark.parametrize( - "conf, init_state, first_run_state, second_run_state, task_ids_to_run, error_message", - [ - ( - {("scheduler", "schedule_after_task_execution"): "True"}, - {"A": State.QUEUED, "B": State.NONE, "C": State.NONE}, - {"A": State.SUCCESS, "B": State.SCHEDULED, "C": State.NONE}, - {"A": State.SUCCESS, "B": State.SUCCESS, "C": State.SCHEDULED}, - ["A", "B"], - "A -> B -> C, with fast-follow ON when A runs, B should be QUEUED. Same for B and C.", - ), - ( - {("scheduler", "schedule_after_task_execution"): "False"}, - {"A": State.QUEUED, "B": State.NONE, "C": State.NONE}, - {"A": State.SUCCESS, "B": State.NONE, "C": State.NONE}, - None, - ["A", "B"], - "A -> B -> C, with fast-follow OFF, when A runs, B shouldn't be QUEUED.", - ), - ( - {("scheduler", "schedule_after_task_execution"): "True"}, - {"D": State.QUEUED, "E": State.NONE, "F": State.NONE, "G": State.NONE}, - {"D": State.SUCCESS, "E": State.NONE, "F": State.NONE, "G": State.NONE}, - None, - ["D", "E"], - "G -> F -> E & D -> E, when D runs but F isn't QUEUED yet, E shouldn't be QUEUED.", - ), - ( - {("scheduler", "schedule_after_task_execution"): "True"}, - {"H": State.QUEUED, "I": State.FAILED, "J": State.NONE}, - {"H": State.SUCCESS, "I": State.FAILED, "J": State.UPSTREAM_FAILED}, - None, - ["H", "I"], - "H -> J & I -> J, when H is QUEUED but I has FAILED, J is marked UPSTREAM_FAILED.", - ), - ], - ) - def test_fast_follow( - self, - conf, - init_state, - first_run_state, - second_run_state, - task_ids_to_run, - error_message, - get_test_dag, - ): - with conf_vars(conf): - dag = get_test_dag( - "test_dagrun_fast_follow", - ) - - scheduler_job = Job() - scheduler_job_runner = SchedulerJobRunner(job=scheduler_job, subdir=os.devnull) - scheduler_job_runner.dagbag.bag_dag(dag) - triggered_by_kwargs = {"triggered_by": DagRunTriggeredByType.TEST} if AIRFLOW_V_3_0_PLUS else {} - - dag_run = dag.create_dagrun( - run_id="test_dagrun_fast_follow", state=State.RUNNING, **triggered_by_kwargs - ) - - ti_by_task_id = {} - with create_session() as session: - for task_id in init_state: - ti = TaskInstance(dag.get_task(task_id), run_id=dag_run.run_id) - ti.refresh_from_db() - ti.state = init_state[task_id] - session.merge(ti) - ti_by_task_id[task_id] = ti - - ti = TaskInstance(task=dag.get_task(task_ids_to_run[0]), run_id=dag_run.run_id) - ti.refresh_from_db() - job1 = Job(executor=SequentialExecutor(), dag_id=ti.dag_id) - job_runner = LocalTaskJobRunner(job=job1, task_instance=ti, ignore_ti_state=True) - job1.task_runner = StandardTaskRunner(job_runner) - - run_job(job=job1, execute_callable=job_runner._execute) - self.validate_ti_states(dag_run, first_run_state, error_message) - if second_run_state: - ti = TaskInstance(task=dag.get_task(task_ids_to_run[1]), run_id=dag_run.run_id) - ti.refresh_from_db() - job2 = Job(dag_id=ti.dag_id, executor=SequentialExecutor()) - job_runner = LocalTaskJobRunner(job=job2, task_instance=ti, ignore_ti_state=True) - job2.task_runner = StandardTaskRunner(job_runner) - run_job(job2, execute_callable=job_runner._execute) - self.validate_ti_states(dag_run, second_run_state, error_message) - if scheduler_job_runner.processor_agent: - scheduler_job_runner.processor_agent.end() - - @pytest.mark.skip_if_database_isolation_mode # Does not work in db isolation mode - @conf_vars({("scheduler", "schedule_after_task_execution"): "True"}) - def test_mini_scheduler_works_with_wait_for_upstream(self, caplog, get_test_dag): - dag = get_test_dag("test_dagrun_fast_follow") - data_interval = dag.infer_automated_data_interval(DEFAULT_LOGICAL_DATE) - dag.catchup = False - SerializedDagModel.write_dag(dag) - triggered_by_kwargs = {"triggered_by": DagRunTriggeredByType.TEST} if AIRFLOW_V_3_0_PLUS else {} - - dr = dag.create_dagrun( - run_id="test_1", - state=State.RUNNING, - execution_date=DEFAULT_DATE, - data_interval=data_interval, - **triggered_by_kwargs, - ) - dr2 = dag.create_dagrun( - run_id="test_2", - state=State.RUNNING, - execution_date=DEFAULT_DATE + datetime.timedelta(hours=1), - data_interval=data_interval, - **triggered_by_kwargs, - ) - task_k = dag.get_task("K") - task_l = dag.get_task("L") - with create_session() as session: - ti_k = dr.get_task_instance(task_k.task_id, session=session) - ti_k.refresh_from_task(task_k) - ti_k.state = State.SUCCESS - - ti_b = dr.get_task_instance(task_l.task_id, session=session) - ti_b.refresh_from_task(task_l) - ti_b.state = State.SUCCESS - - ti2_k = dr2.get_task_instance(task_k.task_id, session=session) - ti2_k.refresh_from_task(task_k) - ti2_k.state = State.NONE - - ti2_l = dr2.get_task_instance(task_l.task_id, session=session) - ti2_l.refresh_from_task(task_l) - ti2_l.state = State.NONE - - session.merge(ti_k) - session.merge(ti_b) - - session.merge(ti2_k) - session.merge(ti2_l) - - job1 = Job( - executor=SequentialExecutor(), - dag_id=ti2_k.dag_id, - ) - job_runner = LocalTaskJobRunner(job=job1, task_instance=ti2_k, ignore_ti_state=True) - job1.task_runner = StandardTaskRunner(job_runner) - run_job(job=job1, execute_callable=job_runner._execute) - - ti2_k.refresh_from_db() - ti2_l.refresh_from_db() - assert ti2_k.state == State.SUCCESS - assert ti2_l.state == State.NONE - - failed_deps = list(ti2_l.get_failed_dep_statuses()) - assert len(failed_deps) == 1 - assert failed_deps[0].dep_name == "Previous Dagrun State" - assert not failed_deps[0].passed - def test_process_sigsegv_error_message(self, caplog, dag_maker): """Test that shows error if process failed with segmentation fault.""" caplog.set_level(logging.CRITICAL, logger="local_task_job.py") diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index 8a1df0594e4e9..3a99ac392269f 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -19,7 +19,6 @@ import contextlib import datetime -import logging import operator import os import pathlib @@ -78,7 +77,7 @@ from airflow.notifications.basenotifier import BaseNotifier from airflow.operators.empty import EmptyOperator from airflow.providers.standard.operators.bash import BashOperator -from airflow.providers.standard.operators.python import BranchPythonOperator, PythonOperator +from airflow.providers.standard.operators.python import PythonOperator from airflow.providers.standard.sensors.python import PythonSensor from airflow.sensors.base import BaseSensorOperator from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG @@ -5121,252 +5120,6 @@ def get_extra_env(): assert "get_extra_env" in echo_task.upstream_task_ids -@pytest.mark.skip_if_database_isolation_mode # Does not work in db isolation mode -def test_mapped_task_does_not_error_in_mini_scheduler_if_upstreams_are_not_done(dag_maker, caplog, session): - """ - This tests that when scheduling child tasks of a task and there's a mapped downstream task, - if the mapped downstream task has upstreams that are not yet done, the mapped downstream task is - not marked as `upstream_failed' - """ - with dag_maker() as dag: - - @dag.task - def second_task(): - return [0, 1, 2] - - @dag.task - def first_task(): - print(2) - - @dag.task - def middle_task(id): - return id - - middle = middle_task.expand(id=second_task()) - - @dag.task - def last_task(): - print(3) - - [first_task(), middle] >> last_task() - - dag_run = dag_maker.create_dagrun() - first_ti = dag_run.get_task_instance(task_id="first_task") - second_ti = dag_run.get_task_instance(task_id="second_task") - first_ti.state = State.SUCCESS - second_ti.state = State.RUNNING - session.merge(first_ti) - session.merge(second_ti) - session.commit() - first_ti.schedule_downstream_tasks(session=session) - middle_ti = dag_run.get_task_instance(task_id="middle_task") - assert middle_ti.state != State.UPSTREAM_FAILED - assert "0 downstream tasks scheduled from follow-on schedule" in caplog.text - - -@pytest.mark.skip_if_database_isolation_mode # Does not work in db isolation mode -def test_empty_operator_is_not_considered_in_mini_scheduler(dag_maker, caplog, session): - """ - This tests verify that operators with inherits_from_empty_operator are not considered by mini scheduler. - Such operators should not run on workers thus the mini scheduler optimization should skip them and not - submit them directly to worker. - """ - with dag_maker() as dag: - - @dag.task - def first_task(): - print(2) - - @dag.task - def second_task(): - print(2) - - third_task = EmptyOperator(task_id="third_task") - forth_task = EmptyOperator(task_id="forth_task", on_success_callback=lambda x: print("hi")) - - first_task() >> [second_task(), third_task, forth_task] - dag_run = dag_maker.create_dagrun() - first_ti = dag_run.get_task_instance(task_id="first_task") - second_ti = dag_run.get_task_instance(task_id="second_task") - third_ti = dag_run.get_task_instance(task_id="third_task") - forth_ti = dag_run.get_task_instance(task_id="forth_task") - first_ti.state = State.SUCCESS - second_ti.state = State.NONE - third_ti.state = State.NONE - forth_ti.state = State.NONE - session.merge(first_ti) - session.merge(second_ti) - session.merge(third_ti) - session.merge(forth_ti) - session.commit() - first_ti.schedule_downstream_tasks(session=session) - second_task = dag_run.get_task_instance(task_id="second_task") - third_task = dag_run.get_task_instance(task_id="third_task") - forth_task = dag_run.get_task_instance(task_id="forth_task") - assert second_task.state == State.SCHEDULED - assert third_task.state == State.NONE - assert forth_task.state == State.SCHEDULED - assert "2 downstream tasks scheduled from follow-on schedule" in caplog.text - - -@pytest.mark.skip_if_database_isolation_mode # Does not work in db isolation mode -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] - - @dag.task - def first_task(): - print(2) - - @dag.task - def middle_task(id): - return id - - middle = middle_task.expand(id=second_task()) - - @dag.task - def last_task(): - print(3) - - [first_task(), middle] >> last_task() - - dr = dag_maker.create_dagrun() - - first_ti = dr.get_task_instance(task_id="first_task") - first_ti.state = State.SUCCESS - session.merge(first_ti) - session.commit() - second_task = dag.get_task("second_task") - second_ti = dr.get_task_instance(task_id="second_task") - second_ti.refresh_from_task(second_task) - second_ti.run() - second_ti.schedule_downstream_tasks(session=session) - for i in range(3): - middle_ti = dr.get_task_instance(task_id="middle_task", map_index=i) - assert middle_ti.state == State.SCHEDULED - assert "3 downstream tasks scheduled from follow-on schedule" in caplog.text - - -@pytest.mark.skip_if_database_isolation_mode -def test_one_success_task_in_mini_scheduler_if_upstreams_are_done(dag_maker, caplog, session): - """Test that mini scheduler with one_success task""" - with dag_maker() as dag: - branch = BranchPythonOperator(task_id="branch", python_callable=lambda: "task_run") - task_run = BashOperator(task_id="task_run", bash_command="echo 0") - task_skip = BashOperator(task_id="task_skip", bash_command="echo 0") - task_1 = BashOperator(task_id="task_1", bash_command="echo 0") - task_one_success = BashOperator( - task_id="task_one_success", bash_command="echo 0", trigger_rule="one_success" - ) - task_2 = BashOperator(task_id="task_2", bash_command="echo 0") - - task_1 >> task_2 - branch >> task_skip - branch >> task_run - task_run >> task_one_success - task_skip >> task_one_success - task_one_success >> task_2 - task_skip >> task_2 - - dr = dag_maker.create_dagrun() - - branch = dr.get_task_instance(task_id="branch") - task_1 = dr.get_task_instance(task_id="task_1") - task_skip = dr.get_task_instance(task_id="task_skip") - branch.state = State.SUCCESS - task_1.state = State.SUCCESS - task_skip.state = State.SKIPPED - session.merge(branch) - session.merge(task_1) - session.merge(task_skip) - session.commit() - task_1.refresh_from_task(dag.get_task("task_1")) - task_1.schedule_downstream_tasks(session=session) - - branch = dr.get_task_instance(task_id="branch") - task_run = dr.get_task_instance(task_id="task_run") - task_skip = dr.get_task_instance(task_id="task_skip") - task_1 = dr.get_task_instance(task_id="task_1") - task_one_success = dr.get_task_instance(task_id="task_one_success") - task_2 = dr.get_task_instance(task_id="task_2") - assert branch.state == State.SUCCESS - assert task_run.state == State.NONE - assert task_skip.state == State.SKIPPED - assert task_1.state == State.SUCCESS - # task_one_success should not be scheduled - assert task_one_success.state == State.NONE - assert task_2.state == State.SKIPPED - assert "0 downstream tasks scheduled from follow-on schedule" in caplog.text - - task_run = dr.get_task_instance(task_id="task_run") - task_run.state = State.SUCCESS - session.merge(task_run) - session.commit() - task_run.refresh_from_task(dag.get_task("task_run")) - task_run.schedule_downstream_tasks(session=session) - - branch = dr.get_task_instance(task_id="branch") - task_run = dr.get_task_instance(task_id="task_run") - task_skip = dr.get_task_instance(task_id="task_skip") - task_1 = dr.get_task_instance(task_id="task_1") - task_one_success = dr.get_task_instance(task_id="task_one_success") - task_2 = dr.get_task_instance(task_id="task_2") - assert branch.state == State.SUCCESS - assert task_run.state == State.SUCCESS - assert task_skip.state == State.SKIPPED - assert task_1.state == State.SUCCESS - # task_one_success should not be scheduled - assert task_one_success.state == State.SCHEDULED - assert task_2.state == State.SKIPPED - assert "1 downstream tasks scheduled from follow-on schedule" in caplog.text - - -@pytest.mark.skip_if_database_isolation_mode # Does not work in db isolation mode -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 [] - - @task - def a_sum(numbers: list[int]) -> int: - return sum(numbers) - - @task - def b_double(summed: int) -> int: - return summed * 2 - - @task - def c_gather(result) -> None: - pass - - static = EmptyOperator(task_id="static") - - summed = a_sum.expand(numbers=generate()) - doubled = b_double.expand(summed=summed) - static >> c_gather(doubled) - - dr: DagRun = dag_maker.create_dagrun() - tis = {(ti.task_id, ti.map_index): ti for ti in dr.task_instances} - - static_ti = tis[("static", -1)] - static_ti.run(session=session) - static_ti.schedule_downstream_tasks(session=session) - # No tasks should be skipped yet! - assert not dr.get_task_instances([TaskInstanceState.SKIPPED], session=session) - - generate_ti = tis[("generate", -1)] - generate_ti.run(session=session) - generate_ti.schedule_downstream_tasks(session=session) - # Now downstreams can be skipped. - assert dr.get_task_instances([TaskInstanceState.SKIPPED], session=session) - - def test_taskinstance_with_note(create_task_instance, session): ti: TaskInstance = create_task_instance(session=session) ti.note = "ti with note" @@ -5399,18 +5152,3 @@ def test__refresh_from_db_should_not_increment_try_number(dag_maker, session): assert ti.try_number == 1 # stays 1 ti.refresh_from_db() assert ti.try_number == 1 # stays 1 - - -@mock.patch("airflow.models.taskinstance.TaskInstance._schedule_downstream_tasks") -def test_swallow_mini_scheduler_exceptions(_schedule_downstream_mock, create_task_instance, caplog): - _schedule_downstream_mock.side_effect = Exception("To be swallowed") - caplog.set_level(logging.ERROR) - ti = create_task_instance( - dag_id="dag_for_testing_swallowing_exception", - task_id="task_for_testing_swallowing_exception", - run_type=DagRunType.SCHEDULED, - execution_date=DEFAULT_DATE, - ) - ti.schedule_downstream_tasks() - assert "Error scheduling downstream tasks." in caplog.text - assert "To be swallowed" in caplog.text