From 8be00c3cbdba937110c22c70d157fa87d06766ba Mon Sep 17 00:00:00 2001 From: mday-io Date: Tue, 13 Oct 2026 19:18:22 -0400 Subject: [PATCH] Fix: Propagate terminal worker failures and stop queued DAG work Signed-off-by: mday-io --- sqlmesh/utils/concurrency.py | 31 ++++- tests/utils/test_concurrency.py | 215 ++++++++++++++++++++++++++++++++ 2 files changed, 240 insertions(+), 6 deletions(-) diff --git a/sqlmesh/utils/concurrency.py b/sqlmesh/utils/concurrency.py index c5f78645f6..1a9a5e2134 100644 --- a/sqlmesh/utils/concurrency.py +++ b/sqlmesh/utils/concurrency.py @@ -65,26 +65,40 @@ def run(self) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]: return self._node_errors, self._skipped_nodes def _process_node(self, node: H, executor: Executor) -> None: + with self._unprocessed_nodes_lock: + if self._finished_future.done(): + return + try: self.fn(node) with self._unprocessed_nodes_lock: + if self._finished_future.done(): + return self._unprocessed_nodes_num -= 1 self._submit_next_nodes(executor, node) except Exception as ex: error = NodeExecutionFailedError(node) error.__cause__ = ex - if self.raise_on_error: - self._finished_future.set_exception(error) - return - with self._unprocessed_nodes_lock: + if self._finished_future.done(): + return + if self.raise_on_error: + self._finished_future.set_exception(error) + return self._unprocessed_nodes_num -= 1 self._node_errors.append(error) self._skip_next_nodes(node) + except BaseException as ex: + with self._unprocessed_nodes_lock: + if not self._finished_future.done(): + self._finished_future.set_exception(ex) def _submit_next_nodes(self, executor: Executor, processed_node: t.Optional[H] = None) -> None: + if self._finished_future.done(): + return + if not self._unprocessed_nodes_num: self._finished_future.set_result(None) return @@ -101,6 +115,9 @@ def _submit_next_nodes(self, executor: Executor, processed_node: t.Optional[H] = executor.submit(self._process_node, submitted_node, executor) def _skip_next_nodes(self, parent: H) -> None: + if self._finished_future.done(): + return + if not self._unprocessed_nodes_num: self._finished_future.set_result(None) return @@ -267,9 +284,11 @@ def concurrent_apply_to_values( def _process_value(value: A, index: int) -> None: try: - futures[index].set_result(fn(value)) - except Exception as ex: + result = fn(value) + except BaseException as ex: futures[index].set_exception(ex) + else: + futures[index].set_result(result) with ThreadPoolExecutor(max_workers=tasks_num) as pool: for index, value in enumerate(values): diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py index 5e1e4326f7..89f9789725 100644 --- a/tests/utils/test_concurrency.py +++ b/tests/utils/test_concurrency.py @@ -1,3 +1,9 @@ +import os +import subprocess +import sys +import textwrap +from pathlib import Path + import pytest from pytest_mock.plugin import MockerFixture @@ -147,3 +153,212 @@ def test_concurrent_apply(tasks_num: int): values = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] results = concurrent_apply_to_values(values, lambda x: x * 2, tasks_num) assert results == [x * 2 for x in values] + + +def _run_bounded_python(program: str) -> None: + repo_root = Path(__file__).resolve().parents[2] + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + path for path in (str(repo_root), env.get("PYTHONPATH")) if path + ) + result = subprocess.run( + [sys.executable, "-c", textwrap.dedent(program)], + cwd=repo_root, + env=env, + capture_output=True, + text=True, + timeout=8, + ) + assert result.returncode == 0, f"{result.stdout}\n{result.stderr}" + + +def test_concurrent_apply_to_values_propagates_base_exception(): + _run_bounded_python( + """ + from sqlmesh.utils.concurrency import concurrent_apply_to_values + + error = SystemExit("cancel value execution") + + def process(value): + if value == 0: + raise error + return value + + try: + concurrent_apply_to_values([0, 1], process, 2) + except SystemExit as raised: + assert raised is error + else: + raise AssertionError("the value worker cancellation was not propagated") + """ + ) + + +@pytest.mark.parametrize("raise_on_error", [True, False]) +def test_concurrent_dag_propagates_base_exception_drains_workers_and_can_be_reused( + raise_on_error: bool, +): + _run_bounded_python( + f""" + from threading import Event, Thread + + from sqlmesh.utils.concurrency import ConcurrentDAGExecutor + from sqlmesh.utils.dag import DAG + + dag = DAG() + dag.add("cancel") + dag.add("sibling") + dag.add("dependent", ["sibling"]) + error = KeyboardInterrupt("cancel execution") + release_sibling = Event() + sibling_started = Event() + called = [] + + def execute(node): + if node == "cancel": + assert sibling_started.wait(timeout=3) + raise error + if node == "sibling": + sibling_started.set() + assert release_sibling.wait(timeout=6) + called.append(node) + + executor = ConcurrentDAGExecutor(dag, execute, 2, {raise_on_error!r}) + run_errors = [] + + def run(): + try: + executor.run() + except BaseException as ex: + run_errors.append(ex) + + runner = Thread(target=run) + runner.start() + try: + assert sibling_started.wait(timeout=3) + try: + executor._finished_future.result(timeout=3) + except KeyboardInterrupt as raised: + assert raised is error + else: + raise AssertionError("the worker cancellation was not propagated") + assert runner.is_alive(), "run returned before the active worker drained" + finally: + release_sibling.set() + + runner.join(timeout=3) + assert not runner.is_alive() + assert run_errors == [error] + assert called == ["sibling"] + + second_error = SystemExit("cancel reused executor") + + def cancel_again(node): + raise second_error + + executor.fn = cancel_again + try: + executor.run() + except SystemExit as raised: + assert raised is second_error + else: + raise AssertionError("the reused executor did not propagate cancellation") + """ + ) + + +def test_concurrent_dag_fail_fast_stops_queued_work_while_aggregate_mode_continues(): + _run_bounded_python( + """ + from threading import Event, Thread + + from sqlmesh.utils.concurrency import ConcurrentDAGExecutor + from sqlmesh.utils.dag import DAG + + def make_dag(): + dag = DAG() + dag.add("fail") + dag.add("running") + dag.add("queued") + dag.add("failed_child", ["fail"]) + dag.add("running_child", ["running"]) + dag.add("queued_child", ["queued"]) + return dag + + failure = RuntimeError("fail node") + release_running = Event() + running_started = Event() + queued_started = Event() + called = [] + + def execute(node): + if node == "fail": + assert running_started.wait(timeout=3) + raise failure + if node == "running": + running_started.set() + assert release_running.wait(timeout=6) + if node == "queued": + queued_started.set() + called.append(node) + + executor = ConcurrentDAGExecutor(make_dag(), execute, 2, True) + run_errors = [] + + def run_fail_fast(): + try: + executor.run() + except BaseException as ex: + run_errors.append(ex) + + runner = Thread(target=run_fail_fast) + runner.start() + try: + assert running_started.wait(timeout=3) + try: + executor._finished_future.result(timeout=3) + except Exception as raised: + assert raised.node == "fail" + assert raised.__cause__ is failure + else: + raise AssertionError("fail-fast execution did not raise") + assert runner.is_alive(), "run returned before the active worker drained" + assert not queued_started.is_set(), "queued work started after fail-fast" + finally: + release_running.set() + + runner.join(timeout=3) + assert not runner.is_alive() + assert len(run_errors) == 1 + assert run_errors[0].node == "fail" + assert run_errors[0].__cause__ is failure + assert called == ["running"] + + release_running = Event() + running_started = Event() + queued_started = Event() + called = [] + executor = ConcurrentDAGExecutor(make_dag(), execute, 2, False) + run_results = [] + + def run_aggregate(): + run_results.append(executor.run()) + + runner = Thread(target=run_aggregate) + runner.start() + try: + assert running_started.wait(timeout=3) + assert queued_started.wait(timeout=3), "aggregate mode stopped independent work" + finally: + release_running.set() + + runner.join(timeout=3) + assert not runner.is_alive() + errors, skipped = run_results[0] + assert len(errors) == 1 + assert errors[0].node == "fail" + assert errors[0].__cause__ is failure + assert skipped == ["failed_child"] + assert set(called) == {"running", "queued", "running_child", "queued_child"} + """ + )