Coverage for src / invariant / scheduler.py: 67.82%

1"""Invocation schedulers for async execution.""" 

2 

3import asyncio 

4from collections.abc import Callable 

5from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor 

6from dataclasses import dataclass 

7from typing import Any, Protocol 

8 

9from invariant.invocation import invoke_op, invoke_op_async 

10from invariant.registry import OpRegistry, import_implementation_ref 

11from invariant.store.codec import deserialize, serialize 

12from invariant.traits import OpTrait 

13 

14 

15@dataclass(frozen=True) 

16class InvocationRequest: 

17 """A scheduler-facing operation invocation request.""" 

18 

19 op_name: str 

20 op: Callable[..., Any] 

21 manifest: dict[str, Any] 

22 traits: frozenset[str] 

23 implementation_ref: str | None = None 

24 cache_key: tuple[str, str] | None = None 

25 

26 

27class InvocationScheduler(Protocol): 

28 """Protocol implemented by local and remote invocation schedulers.""" 

29 

30 async def invoke(self, request: InvocationRequest) -> Any: 

31 """Invoke an operation and return its artifact.""" 

32 ... 

33 

34 

35class InlineScheduler: 

36 """Invoke operations directly on the event loop thread.""" 

37 

38 def __init__(self, max_concurrency: int = 16): 

39 if ( 

40 isinstance(max_concurrency, bool) 

41 or not isinstance(max_concurrency, int) 

42 or max_concurrency < 1 

43 ): 

44 raise ValueError("max_concurrency must be a positive integer") 

45 self._slots = asyncio.Semaphore(max_concurrency) 

46 

47 async def invoke(self, request: InvocationRequest) -> Any: 

48 """Invoke cheap synchronous or asynchronous operations with bounded concurrency.""" 

49 async with self._slots: 

50 return await invoke_op_async(request.op, request.op_name, request.manifest) 

51 

52 

53class ThreadPoolScheduler: 

54 """Invoke operations in a thread pool.""" 

55 

56 def __init__( 

57 self, 

58 max_workers: int | None = None, 

59 executor: ThreadPoolExecutor | None = None, 

60 ) -> None: 

61 if max_workers is not None and executor is not None: 

62 raise ValueError("max_workers cannot be set when executor is provided") 

63 self._executor = executor or ThreadPoolExecutor(max_workers=max_workers) 

64 self._owns_executor = executor is None 

65 

66 async def invoke(self, request: InvocationRequest) -> Any: 

67 """Invoke an operation in the thread pool.""" 

68 loop = asyncio.get_running_loop() 

69 return await loop.run_in_executor( 

70 self._executor, 

71 invoke_op, 

72 request.op, 

73 request.op_name, 

74 request.manifest, 

75 ) 

76 

77 async def aclose(self) -> None: 

78 """Shut down the owned thread pool.""" 

79 if self._owns_executor: 

80 await asyncio.to_thread(self._executor.shutdown, wait=True) 

81 

82 

83class ProcessPoolScheduler: 

84 """Invoke worker-resolvable operations in a process pool.""" 

85 

86 def __init__( 

87 self, 

88 max_workers: int | None = None, 

89 executor: ProcessPoolExecutor | None = None, 

90 ) -> None: 

91 if max_workers is not None and executor is not None: 

92 raise ValueError("max_workers cannot be set when executor is provided") 

93 self._executor = executor or ProcessPoolExecutor(max_workers=max_workers) 

94 self._owns_executor = executor is None 

95 

96 async def invoke(self, request: InvocationRequest) -> Any: 

97 """Invoke an operation through an Invariant codec process boundary.""" 

98 if not request.implementation_ref: 

99 raise ValueError( 

100 f"Op '{request.op_name}' cannot run in a process because it has " 

101 "no worker-resolvable implementation_ref" 

102 ) 

103 

104 manifest_payload = serialize(request.manifest) 

105 loop = asyncio.get_running_loop() 

106 artifact_payload = await loop.run_in_executor( 

107 self._executor, 

108 _process_worker_invoke, 

109 request.op_name, 

110 request.implementation_ref, 

111 manifest_payload, 

112 ) 

113 return deserialize(artifact_payload) 

114 

115 async def aclose(self) -> None: 

116 """Shut down the owned process pool.""" 

117 if self._owns_executor: 

118 await asyncio.to_thread(self._executor.shutdown, wait=True) 

119 

120 

121class RoutingScheduler: 

122 """Route invocations to local schedulers according to traits.""" 

123 

124 def __init__( 

125 self, 

126 *, 

127 inline_scheduler: InvocationScheduler | None = None, 

128 thread_scheduler: InvocationScheduler | None = None, 

129 process_scheduler: InvocationScheduler | None = None, 

130 ) -> None: 

131 self.inline_scheduler = inline_scheduler or InlineScheduler() 

132 self.thread_scheduler = thread_scheduler 

133 self.process_scheduler = process_scheduler 

134 

135 async def invoke(self, request: InvocationRequest) -> Any: 

136 """Route an invocation to the first configured matching scheduler.""" 

137 if ( 

138 OpTrait.PROCESS_SAFE.value in request.traits 

139 and self.process_scheduler is not None 

140 ): 

141 return await self.process_scheduler.invoke(request) 

142 

143 if self.thread_scheduler is not None and ( 

144 OpTrait.BLOCKING.value in request.traits 

145 or OpTrait.IO_BOUND.value in request.traits 

146 ): 

147 return await self.thread_scheduler.invoke(request) 

148 

149 return await self.inline_scheduler.invoke(request) 

150 

151 async def aclose(self) -> None: 

152 """Close child schedulers that expose ``aclose``.""" 

153 for scheduler in ( 

154 self.process_scheduler, 

155 self.thread_scheduler, 

156 self.inline_scheduler, 

157 ): 

158 close = getattr(scheduler, "aclose", None) 

159 if close is not None: 

160 await close() 

161 

162 

163def _process_worker_invoke( 

164 op_name: str, 

165 implementation_ref: str, 

166 manifest_payload: bytes, 

167) -> bytes: 

168 """Process worker entrypoint. 

169 

170 The parent sends only simple strings and Invariant codec bytes. The worker 

171 resolves the exact callable locally, invokes it, and returns codec bytes. 

172 """ 

173 registry = OpRegistry() 

174 registry.clear() 

175 registry.auto_discover() 

176 

177 if registry.has(op_name): 

178 binding = registry.get_binding(op_name) 

179 if binding.implementation_ref != implementation_ref: 

180 raise ValueError( 

181 f"Worker discovered op '{op_name}' as " 

182 f"{binding.implementation_ref!r}, but request requires " 

183 f"{implementation_ref!r}" 

184 ) 

185 op = binding.op 

186 else: 

187 op = import_implementation_ref(implementation_ref) 

188 registry.register( 

189 op_name, 

190 op, 

191 implementation_ref=implementation_ref, 

192 ) 

193 

194 manifest = deserialize(manifest_payload) 

195 artifact = invoke_op(op, op_name, manifest) 

196 return serialize(artifact) 

197 

198 

199__all__ = [ 

200 "InlineScheduler", 

201 "InvocationRequest", 

202 "InvocationScheduler", 

203 "ProcessPoolScheduler", 

204 "RoutingScheduler", 

205 "ThreadPoolScheduler", 

206]