Feat: Add optional per-task context propagation to concurrency helpers

Signed-off-by: mday-io <mdaytn@gmail.com>

Author: mday-io

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,21 +244,24 @@ 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()
 
 
 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,19 +297,27 @@ 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.
 
     Args:
         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]

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