Feat: Add optional per-task context propagation

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,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]] = []
@@ -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,30 +297,43 @@ 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]
 
     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]