diff --git a/sqlmesh/utils/concurrency.py b/sqlmesh/utils/concurrency.py index 1a9a5e2134..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,14 +62,18 @@ 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: @@ -112,7 +123,12 @@ 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(): @@ -156,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. @@ -168,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. @@ -189,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, ) @@ -197,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. @@ -208,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. @@ -219,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() @@ -233,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]] = [] @@ -248,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 @@ -266,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. @@ -273,12 +305,19 @@ 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] @@ -292,6 +331,9 @@ def _process_value(value: A, index: int) -> None: 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 89f9789725..083897fe56 100644 --- a/tests/utils/test_concurrency.py +++ b/tests/utils/test_concurrency.py @@ -2,6 +2,7 @@ import subprocess import sys import textwrap +from contextvars import ContextVar, copy_context from pathlib import Path import pytest @@ -362,3 +363,52 @@ def run_aggregate(): 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