diff --git a/sqlmesh/utils/concurrency.py b/sqlmesh/utils/concurrency.py index c5f78645f6..a4271f8072 100644 --- a/sqlmesh/utils/concurrency.py +++ b/sqlmesh/utils/concurrency.py @@ -1,6 +1,7 @@ import typing as t from concurrent.futures import Executor, Future, ThreadPoolExecutor from threading import Lock +from contextvars import Context from sqlmesh.core.snapshot import SnapshotId, SnapshotInfoLike from sqlmesh.utils.dag import DAG @@ -30,6 +31,9 @@ class ConcurrentDAGExecutor(t.Generic[H]): raise_on_error: If set to True raises an exception on a first encountered error, otherwises returns a tuple which contains a list of failed nodes and a list of skipped nodes. + context_factory: Optional factory called at each run admission. Each node runs + in a separate copy of that context, including serial execution. By default + no context is captured or copied. """ def __init__( @@ -38,11 +42,14 @@ def __init__( fn: t.Callable[[H], None], tasks_num: int, raise_on_error: bool, + context_factory: t.Optional[t.Callable[[], Context]] = None, ): self.dag = dag self.fn = fn self.tasks_num = tasks_num self.raise_on_error = raise_on_error + self.context_factory = context_factory + self._context: t.Optional[Context] = None self._init_state() @@ -55,36 +62,54 @@ def run(self) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]: Returns: A pair which contains a list of node errors and a list of skipped nodes. """ + self._context = self.context_factory() if self.context_factory else None if self._finished_future.done(): self._init_state() - with ThreadPoolExecutor(max_workers=self.tasks_num) as pool: - with self._unprocessed_nodes_lock: - self._submit_next_nodes(pool) - self._finished_future.result() - return self._node_errors, self._skipped_nodes + try: + with ThreadPoolExecutor(max_workers=self.tasks_num) as pool: + with self._unprocessed_nodes_lock: + self._submit_next_nodes(pool) + self._finished_future.result() + return self._node_errors, self._skipped_nodes + finally: + self._context = None 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 @@ -98,9 +123,17 @@ def _submit_next_nodes(self, executor: Executor, processed_node: t.Optional[H] = for submitted_node in submitted_nodes: self._unprocessed_nodes.pop(submitted_node) - executor.submit(self._process_node, submitted_node, executor) + if self._context is None: + executor.submit(self._process_node, submitted_node, executor) + else: + executor.submit( + self._context.copy().run, 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 @@ -139,6 +172,7 @@ def concurrent_apply_to_snapshots( tasks_num: int, reverse_order: bool = False, raise_on_error: bool = True, + context_factory: t.Optional[t.Callable[[], Context]] = None, ) -> t.Tuple[t.List[NodeExecutionFailedError[SnapshotId]], t.List[SnapshotId]]: """Applies a function to the given collection of snapshots concurrently while preserving the topological order between snapshots. @@ -151,6 +185,9 @@ def concurrent_apply_to_snapshots( raise_on_error: If set to True raises an exception on a first encountered error, otherwises returns a tuple which contains a list of failed nodes and a list of skipped nodes. + context_factory: Optional factory called at each run admission. Each node runs + in a separate copy of that context, including serial execution. By default + no context is captured or copied. Raises: NodeExecutionFailedError if `raise_on_error` is set to True and execution fails for any snapshot. @@ -172,6 +209,7 @@ def concurrent_apply_to_snapshots( lambda s_id: fn(snapshots_by_id[s_id]), tasks_num, raise_on_error=raise_on_error, + context_factory=context_factory, ) @@ -180,6 +218,7 @@ def concurrent_apply_to_dag( fn: t.Callable[[H], None], tasks_num: int, raise_on_error: bool = True, + context_factory: t.Optional[t.Callable[[], Context]] = None, ) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]: """Applies a function to the given DAG concurrently while preserving the topological order between snapshots. @@ -191,6 +230,9 @@ def concurrent_apply_to_dag( raise_on_error: If set to True raises an exception on a first encountered error, otherwises returns a tuple which contains a list of failed nodes and a list of skipped nodes. + context_factory: Optional factory called at each run admission. Each node runs + in a separate copy of that context, including serial execution. By default + no context is captured or copied. Raises: NodeExecutionFailedError if `raise_on_error` is set to True and execution fails for any snapshot. @@ -202,13 +244,14 @@ def concurrent_apply_to_dag( raise ConfigError(f"Invalid number of concurrent tasks {tasks_num}") if tasks_num == 1: - return sequential_apply_to_dag(dag, fn, raise_on_error) + return sequential_apply_to_dag(dag, fn, raise_on_error, context_factory) return ConcurrentDAGExecutor( dag, fn, tasks_num, raise_on_error, + context_factory, ).run() @@ -216,7 +259,9 @@ def sequential_apply_to_dag( dag: DAG[H], fn: t.Callable[[H], None], raise_on_error: bool = True, + context_factory: t.Optional[t.Callable[[], Context]] = None, ) -> t.Tuple[t.List[NodeExecutionFailedError[H]], t.List[H]]: + context = context_factory() if context_factory else None dependencies = dag.graph node_errors: t.List[NodeExecutionFailedError[H]] = [] @@ -231,7 +276,10 @@ def sequential_apply_to_dag( continue try: - fn(node) + if context is None: + fn(node) + else: + context.copy().run(fn, node) except Exception as ex: error = NodeExecutionFailedError(node) error.__cause__ = ex @@ -249,6 +297,7 @@ def concurrent_apply_to_values( values: t.Sequence[A], fn: t.Callable[[A], R], tasks_num: int, + context_factory: t.Optional[t.Callable[[], Context]] = None, ) -> t.List[R]: """Applies a function to the given collection of values concurrently. @@ -256,23 +305,35 @@ def concurrent_apply_to_values( values: Target values. fn: The function that will be applied concurrently to each value. tasks_num: The number of concurrent tasks. + context_factory: Optional factory called once when values are admitted. Each value runs + in a separate copy of that context, including serial execution. By default no + context is captured or copied. Returns: A list of results. """ + context = context_factory() if context_factory else None + if tasks_num == 1: - return [fn(value) for value in values] + if context is None: + return [fn(value) for value in values] + return [context.copy().run(fn, value) for value in values] futures: t.List[Future] = [Future() for _ in 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): - pool.submit(_process_value, value, index) + if context is None: + pool.submit(_process_value, value, index) + else: + pool.submit(context.copy().run, _process_value, value, index) return [f.result() for f in futures] diff --git a/tests/utils/test_concurrency.py b/tests/utils/test_concurrency.py index 5e1e4326f7..083897fe56 100644 --- a/tests/utils/test_concurrency.py +++ b/tests/utils/test_concurrency.py @@ -1,3 +1,10 @@ +import os +import subprocess +import sys +import textwrap +from contextvars import ContextVar, copy_context +from pathlib import Path + import pytest from pytest_mock.plugin import MockerFixture @@ -147,3 +154,261 @@ 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"} + """ + ) + + +@pytest.mark.parametrize("tasks_num", [1, 3]) +def test_concurrent_apply_to_values_copies_admission_context(tasks_num: int): + value = ContextVar("admission", default="worker") + token = value.set("admitted") + + def process(item: int) -> str: + assert value.get() == "admitted" + value.set(str(item)) + return value.get() + + try: + assert concurrent_apply_to_values( + [1, 2, 3], process, tasks_num, context_factory=copy_context + ) == ["1", "2", "3"] + assert value.get() == "admitted" + default_results = concurrent_apply_to_values([1, 2, 3], lambda _: value.get(), tasks_num) + expected = "admitted" if tasks_num == 1 else "worker" + assert default_results == [expected, expected, expected] + finally: + value.reset(token) + + +@pytest.mark.parametrize("tasks_num", [1, 2]) +def test_context_admission_and_sibling_isolation(tasks_num: int): + from sqlmesh.utils.concurrency import ConcurrentDAGExecutor, concurrent_apply_to_dag + from sqlmesh.utils.dag import DAG + + consumer = ContextVar("admission", default="construction") + dag: DAG[str] = DAG() + dag.add("second", {"first"}) + seen = [] + + def execute(node): + seen.append(consumer.get()) + consumer.set(node) + + executor = ConcurrentDAGExecutor(dag, execute, tasks_num, False, copy_context) + token = consumer.set("first admission") + try: + executor.run() + consumer.set("second admission") + executor.run() + concurrent_apply_to_dag(dag, execute, tasks_num, context_factory=copy_context) + assert consumer.get() == "second admission" + finally: + consumer.reset(token) + assert seen == ["first admission"] * 2 + ["second admission"] * 4