Coverage for src/invariant/executor.py: 91.18%

1"""Executor: The runtime engine for executing DAGs."""

2

3 import inspect

4 from typing import TYPE_CHECKING, Any

5

6 from invariant.cacheable import is_cacheable

7 from invariant.expressions import resolve_params

8 from invariant.graph import Graph, GraphResolver

9 from invariant.hashing import hash_manifest

10 from invariant.node import Node, SubGraphNode

11

12 if TYPE_CHECKING:

13 from invariant.registry import OpRegistry

14 from invariant.store.base import ArtifactStore

15

16

17 class Executor:

18 """Runtime engine for executing DAGs.

19

20 Manages the two-phase execution:

21 - Phase 1: Context Resolution (Graph -> Manifest)

22 - Phase 2: Action Execution (Manifest -> Artifact)

23 """

24

25 def __init__(

26 self,

27 registry: "OpRegistry",

28 store: "ArtifactStore",

29 resolver: "GraphResolver | None" = None,

30 ) -> None:

31 """Initialize Executor.

32

33 Args:

34 registry: OpRegistry for looking up operations.

35 store: ArtifactStore for caching artifacts.

36 resolver: Optional GraphResolver. If None, creates one with registry.

37 """

38 self.registry = registry

39 self.store = store

40 self.resolver = resolver or GraphResolver(registry)

41

42 def execute(

43 self, graph: Graph, context: dict[str, Any] | None = None

44 ) -> dict[str, Any]:

45 """Execute a graph and return artifacts for each node.

46

47 Args:

48 graph: Dictionary mapping node IDs to Node or SubGraphNode objects.

49 context: Optional dictionary of external dependencies (values not in graph).

50 These are injected as artifacts available to any node that declares

51 them in deps.

52

53 Returns:

54 Dictionary mapping node IDs to their resulting artifacts.

55

56 Raises:

57 ValueError: If graph validation fails or execution errors occur.

58 """

59 # Validate and sort graph (pass context for validation)

60 context = context or {}

61 sorted_nodes = self.resolver.resolve(graph, context_keys=set(context.keys()))

62

63 # Track artifacts by node ID

64 artifacts_by_node: dict[str, Any] = {}

65

66 # Inject context values into artifacts_by_node before execution

67 # This makes external dependencies available to any node that declares them in deps

68 for key, value in context.items():

69 # Context values must be cacheable

70 if not is_cacheable(value):

71 raise ValueError(

72 f"Context value for '{key}' is not cacheable, got {type(value)}"

73 )

74 # Store native types as-is (no wrapping)

75 artifacts_by_node[key] = value

76

77 # Execute nodes in topological order

78 for node_id in sorted_nodes:

79 node = graph[node_id]

80

81 if isinstance(node, SubGraphNode):

82 # SubGraphNode: run internal graph with resolved params as context

83 manifest = self._build_manifest(node, node_id, graph, artifacts_by_node)

84 inner_results = self.execute(node.graph, context=manifest)

85 if node.output not in inner_results:

86 raise ValueError(

87 f"SubGraphNode '{node_id}' output '{node.output}' not in "

88 f"internal results. Keys: {list(inner_results.keys())}."

89 )

90 artifacts_by_node[node_id] = inner_results[node.output]

91 else:

92 # Node: Phase 1 build manifest, Phase 2 cache lookup or execute op

93 manifest = self._build_manifest(node, node_id, graph, artifacts_by_node)

94 if not node.cache:

95 # Ephemeral node: always execute, never cache

96 op = self.registry.get(node.op_name)

97 artifact = self._invoke_op(op, node.op_name, manifest)

98 else:

99 digest = hash_manifest(manifest)

100 if self.store.exists(node.op_name, digest):

101 artifact = self.store.get(node.op_name, digest)

102 else:

103 op = self.registry.get(node.op_name)

104 artifact = self._invoke_op(op, node.op_name, manifest)

105 self.store.put(node.op_name, digest, artifact)

106 artifacts_by_node[node_id] = artifact

107

108 return artifacts_by_node

109

110 def _build_manifest(

111 self,

112 node: Node | SubGraphNode,

113 node_id: str,

114 graph: Graph,

115 artifacts_by_node: dict[str, Any],

116 ) -> dict[str, Any]:

117 """Build the input manifest for a node (Phase 1).

118

119 The manifest is built entirely from resolved params. Dependencies are NOT

120 injected into the manifest directly - they are only available for ref()/cel()

121 resolution within params.

122

123 Args:

124 node: The node to build manifest for.

125 node_id: The ID of the node.

126 graph: The full graph (for reference).

127 artifacts_by_node: Already computed artifacts for upstream nodes.

128

129 Returns:

130 The manifest dictionary mapping parameter names to resolved values.

131 """

132 # Collect dependency artifacts for ref()/cel() resolution

133 dependencies: dict[str, Any] = {}

134 for dep_id in node.deps:

135 if dep_id not in artifacts_by_node:

136 raise ValueError(

137 f"Node '{node_id}' depends on '{dep_id}' but artifact not found. "

138 f"This should not happen if graph is topologically sorted or "

139 f"if '{dep_id}' is provided in context."

140 )

141 dependencies[dep_id] = artifacts_by_node[dep_id]

142

143 # Manifest = resolved params only. No dependency injection.

144 # ref() and cel() markers in params are resolved using dependencies.

145 return resolve_params(node.params, dependencies)

146

147 def _invoke_op(self, op: Any, op_name: str, manifest: dict[str, Any]) -> Any:

148 """Invoke an operation with kwargs dispatch and return validation.

149

150 Args:

151 op: The callable operation to invoke.

152 op_name: The name of the operation (for error messages).

153 manifest: The manifest dictionary mapping parameter names to values.

154

155 Returns:

156 The operation result (native type or ICacheable domain type).

157

158 Raises:

159 ValueError: If required parameters are missing.

160 TypeError: If return value is not cacheable.

161 """

162 # Inspect function signature to map manifest keys to function parameters

163 sig = inspect.signature(op)

164 kwargs: dict[str, Any] = {}

165

166 # Map manifest keys to function parameters by name

167 for name, param in sig.parameters.items():

168 if name in manifest:

169 value = manifest[name]

170 kwargs[name] = value

171 elif param.default is not inspect.Parameter.empty:

172 # Parameter has a default value, skip it

173 pass

174 elif param.kind == inspect.Parameter.VAR_KEYWORD:

175 # Function accepts **kwargs, will handle below

176 pass

177 else:

178 # Required parameter missing

179 raise ValueError(f"Op '{op_name}': missing required parameter '{name}'")

180

181 # If function has **kwargs, pass remaining manifest keys

182 has_var_kwargs = any(

183 p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()

184 )

185 if has_var_kwargs:

186 for key, val in manifest.items():

187 if key not in kwargs:

188 kwargs[key] = val

189

190 # Invoke the operation

191 result = op(**kwargs)

192

193 # Validate return value is cacheable

194 if not is_cacheable(result):

195 raise TypeError(

196 f"Op '{op_name}' returned {type(result).__name__}, "

197 f"which is not a cacheable type"

198 )

199

200 # Return as-is (no wrapping needed)

201 return result