Coverage for src / invariant / async_executor.py: 94.59%
148 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-10-25 08:17 +0000
1"""Async executor: scheduler-driven alternative to the safe sync Executor."""
3import asyncio
4from collections.abc import Iterable
5from typing import Any
7from invariant.cacheable import is_cacheable
8from invariant.executor import Executor
9from invariant.graph import Graph
10from invariant.hashing import hash_manifest
11from invariant.node import Node, SubGraphNode, SwitchNode
12from invariant.scheduler import InlineScheduler, InvocationRequest, InvocationScheduler
15class AsyncExecutor(Executor):
16 """Async runtime engine for executing DAGs.
18 The async executor preserves the synchronous executor's graph and cache
19 semantics, while delegating operation placement to an InvocationScheduler.
20 """
22 def __init__(
23 self,
24 registry,
25 store,
26 scheduler: InvocationScheduler | None = None,
27 *,
28 cache: bool = True,
29 ) -> None:
30 """Initialize AsyncExecutor."""
31 super().__init__(registry, store, cache=cache)
32 self.scheduler = scheduler or InlineScheduler()
33 self._singleflight: dict[tuple[str, str], asyncio.Task[Any]] = {}
34 self._singleflight_lock = asyncio.Lock()
35 self._waiters = {}
37 async def __aenter__(self) -> "AsyncExecutor":
38 """Enter an async context manager."""
39 return self
41 async def __aexit__(self, *exc_info: object) -> None:
42 """Close scheduler resources on context exit."""
43 await self.aclose()
45 async def aclose(self) -> None:
46 """Close scheduler resources when supported."""
47 close = getattr(self.scheduler, "aclose", None)
48 if close is not None:
49 await close()
51 async def execute(
52 self,
53 graph: Graph,
54 outputs: Iterable[str],
55 context: dict[str, Any] | None = None,
56 ) -> dict[str, Any]:
57 """Execute requested graph outputs and return their artifacts."""
58 output_ids = self._normalize_outputs(outputs, graph)
59 artifacts_by_node, _ = await self._execute_requested_outputs_async(
60 graph,
61 output_ids,
62 context=context,
63 )
64 return {output: artifacts_by_node[output] for output in output_ids}
66 async def _execute_requested_outputs_async(
67 self,
68 graph: Graph,
69 outputs: tuple[str, ...],
70 context: dict[str, Any] | None = None,
71 uncacheable_context_keys: set[str] | None = None,
72 ) -> tuple[dict[str, Any], set[str]]:
73 """Demand-execute active paths for requested graph outputs."""
74 context = context or {}
75 uncacheable_nodes = set(uncacheable_context_keys or set())
76 artifacts_by_node: dict[str, Any] = {}
77 node_tasks: dict[str, asyncio.Task[Any]] = {}
78 waiting_for: dict[str, set[str]] = {}
80 def check_wait_cycle(waiter: str, target: str) -> None:
81 stack = [target]
82 seen: set[str] = set()
83 while stack:
84 current = stack.pop()
85 if current == waiter:
86 raise ValueError(
87 "Graph contains cycles on active path: "
88 f"{waiter} -> {target} -> {waiter}"
89 )
90 if current in seen:
91 continue
92 seen.add(current)
93 stack.extend(waiting_for.get(current, ()))
95 async def await_existing_task(
96 node_id: str,
97 task: asyncio.Task[Any],
98 waiter: str | None,
99 ) -> Any:
100 if task.done() or waiter is None:
101 return await task
102 check_wait_cycle(waiter, node_id)
103 waiting_for.setdefault(waiter, set()).add(node_id)
104 try:
105 return await task
106 finally:
107 waiters = waiting_for.get(waiter)
108 if waiters is not None:
109 waiters.discard(node_id)
110 if not waiters:
111 waiting_for.pop(waiter, None)
113 async def resolve_artifact(
114 node_id: str,
115 path: tuple[str, ...] = (),
116 waiter: str | None = None,
117 ) -> Any:
118 if node_id in artifacts_by_node:
119 return artifacts_by_node[node_id]
121 if node_id in path:
122 cycle = " -> ".join([*path, node_id])
123 raise ValueError(f"Graph contains cycles on active path: {cycle}")
125 if node_id in graph:
126 existing = node_tasks.get(node_id)
127 if existing is not None:
128 return await await_existing_task(node_id, existing, waiter)
130 task = asyncio.create_task(run_graph_node(node_id, (*path, node_id)))
131 node_tasks[node_id] = task
132 return await await_existing_task(node_id, task, waiter)
134 if node_id in context:
135 value = context[node_id]
136 if not is_cacheable(value):
137 raise ValueError(
138 f"Context value for '{node_id}' is not cacheable, "
139 f"got {type(value)}"
140 )
141 artifacts_by_node[node_id] = value
142 return value
144 raise ValueError(
145 f"Node depends on '{node_id}' but it doesn't exist in graph "
146 f"or context. Available: graph={sorted(graph)}, "
147 f"context={sorted(context)}"
148 )
150 async def run_graph_node(node_id: str, path: tuple[str, ...]) -> Any:
151 node = graph[node_id]
152 await asyncio.gather(
153 *(
154 resolve_artifact(dep_id, path, waiter=node_id)
155 for dep_id in node.deps
156 )
157 )
159 depends_on_uncacheable = any(
160 dep_id in uncacheable_nodes for dep_id in node.deps
161 )
163 if isinstance(node, SwitchNode):
164 self._validate_switch_targets(node, node_id, graph)
165 target_id = self._select_switch_target(
166 node,
167 node_id,
168 artifacts_by_node,
169 )
170 if target_id not in graph:
171 raise ValueError(
172 f"SwitchNode '{node_id}' targets '{target_id}' "
173 "which doesn't exist in graph"
174 )
175 await resolve_artifact(target_id, path, waiter=node_id)
176 artifacts_by_node[node_id] = artifacts_by_node[target_id]
177 if depends_on_uncacheable or target_id in uncacheable_nodes:
178 uncacheable_nodes.add(node_id)
179 elif isinstance(node, SubGraphNode):
180 manifest = self._build_manifest(
181 node,
182 node_id,
183 graph,
184 artifacts_by_node,
185 )
186 inner_uncacheable_context_keys = (
187 set(manifest.keys()) if depends_on_uncacheable else set()
188 )
189 (
190 inner_results,
191 inner_uncacheable_nodes,
192 ) = await self._execute_requested_outputs_async(
193 node.graph,
194 (node.output,),
195 context=manifest,
196 uncacheable_context_keys=inner_uncacheable_context_keys,
197 )
198 artifacts_by_node[node_id] = inner_results[node.output]
199 if depends_on_uncacheable or node.output in inner_uncacheable_nodes:
200 uncacheable_nodes.add(node_id)
201 else:
202 manifest = self._build_manifest(
203 node,
204 node_id,
205 graph,
206 artifacts_by_node,
207 )
208 artifact = await self._execute_node_async(
209 node,
210 node_id,
211 manifest,
212 depends_on_uncacheable,
213 uncacheable_nodes,
214 )
215 artifacts_by_node[node_id] = artifact
217 return artifacts_by_node[node_id]
219 try:
220 await asyncio.gather(*(resolve_artifact(output) for output in outputs))
221 except BaseException:
222 for task in tuple(node_tasks.values()):
223 if not task.done():
224 task.cancel()
225 await asyncio.gather(*node_tasks.values(), return_exceptions=True)
226 raise
228 return artifacts_by_node, uncacheable_nodes
230 async def _execute_node_async(
231 self,
232 node: Node,
233 node_id: str,
234 manifest: dict[str, Any],
235 depends_on_uncacheable: bool,
236 uncacheable_nodes: set[str],
237 ) -> Any:
238 """Execute one Node using scheduler-driven cache semantics."""
239 if not self.registry.has(node.op_name):
240 raise ValueError(
241 f"Node '{node_id}' references unregistered op '{node.op_name}'"
242 )
244 binding = self.registry.get_binding(node.op_name)
245 should_cache = self.cache and node.cache and not depends_on_uncacheable
246 if not should_cache:
247 artifact = await self.scheduler.invoke(
248 InvocationRequest(
249 op_name=node.op_name,
250 op=binding.op,
251 manifest=manifest,
252 traits=binding.traits,
253 implementation_ref=binding.implementation_ref,
254 )
255 )
256 uncacheable_nodes.add(node_id)
257 return artifact
259 digest = hash_manifest(manifest)
260 key = (node.op_name, digest)
262 async with self._singleflight_lock:
263 task = self._singleflight.get(key)
264 if task is None:
265 if self.store.exists(node.op_name, digest):
266 return self.store.get(node.op_name, digest)
267 task = asyncio.create_task(
268 self._invoke_and_store(binding, manifest, key)
269 )
270 self._singleflight[key] = task
271 self._waiters[key] = self._waiters.get(key, 0) + 1
273 try:
274 # One cancelled graph must not cancel another graph's shared work.
275 return await asyncio.shield(task)
276 finally:
277 abandoned = False
278 async with self._singleflight_lock:
279 self._waiters[key] -= 1
280 if not self._waiters[key]:
281 self._waiters.pop(key)
282 self._singleflight.pop(key, None)
283 if not task.done():
284 task.cancel()
285 abandoned = True
286 if abandoned:
287 await asyncio.gather(task, return_exceptions=True)
289 async def _invoke_and_store(self, binding, manifest, key: tuple[str, str]) -> Any:
290 """Invoke an op through the scheduler and store the resulting artifact."""
291 op_name, digest = key
292 artifact = await self.scheduler.invoke(
293 InvocationRequest(
294 op_name=op_name,
295 op=binding.op,
296 manifest=manifest,
297 traits=binding.traits,
298 implementation_ref=binding.implementation_ref,
299 cache_key=key,
300 )
301 )
302 self.store.put(op_name, digest, artifact)
303 return artifact
306__all__ = ["AsyncExecutor"]