Coverage for async-durable-execution/src/async_durable_execution/composite/map.py: 100%

1"""Implementation for Durable Map operation.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6import json 

7import logging 

8from dataclasses import dataclass, field 

9from typing import ( 

10 TYPE_CHECKING, 

11 Generic, 

12 TypeVar, 

13 Sequence, 

14 Iterable, 

15 Callable, 

16 Any, 

17 Awaitable, 

18) 

19 

20from .parallel import ( 

21 BatchResult, 

22 CompletionConfig, 

23 NestingType, 

24) 

25from .parallel import parallel_handler 

26from ..context import bind_current_context 

27from ..execution import durable_callable 

28from ..models import OperationIdentifier, OperationSubType 

29from ..primitive.child import ( 

30 DurableContext, 

31 _create_child_context_task as _run_in_child_context, 

32 get_durable_context, 

33) 

34 

35if TYPE_CHECKING: 

36 from .parallel import SummaryGenerator 

37 from ..serdes import SerDes 

38 from ..state import ExecutionState 

39 

40logger = logging.getLogger(__name__) 

41 

42# Input item type 

43T = TypeVar("T") 

44# Result type 

45R = TypeVar("R") 

46U = TypeVar("U") 

47 

48 

49@dataclass(frozen=True) 

50class BatchedInput(Generic[T, U]): 

51 """Wrapper passed to batched map handlers.""" 

52 

53 batch_input: T 

54 items: list[U] 

55 

56 

57@dataclass(frozen=True) 

58class MapItemContext(DurableContext, Generic[T]): 

59 """Context exposed while a map item function is executing.""" 

60 

61 index: int = 0 

62 items: Sequence[T] = field(default_factory=tuple) 

63 

64 

65def _bind_map_item_to_branch( 

66 items: Sequence[T], 

67 index: int, 

68 func: Callable[[T], Awaitable[R]], 

69) -> Callable[[], Awaitable[R]]: 

70 async def run_branch() -> R: 

71 logger.debug("🗺️ Processing map item: %s", index) 

72 item = items[index] 

73 child_context = get_durable_context("map") 

74 map_item_context = MapItemContext( 

75 execution_state=child_context.execution_state, 

76 operation_identifier=child_context.operation_identifier, 

77 step_id_prefix=child_context.step_id_prefix, 

78 replaying=child_context.is_replaying(), 

79 index=index, 

80 items=items, 

81 ) 

82 with bind_current_context(map_item_context): 

83 result: R = await func(item) 

84 logger.debug("✅ Processed map item: %s", index) 

85 return result 

86 

87 return run_branch 

88 

89 

90def _create_map_branches( 

91 items: Sequence[T], 

92 func: Callable[[T], Awaitable[R]], 

93) -> list[Callable[[], Awaitable[R]]]: 

94 return [ 

95 _bind_map_item_to_branch(items=items, index=index, func=func) 

96 for index in range(len(items)) 

97 ] 

98 

99 

100def _create_map_branch_namer( 

101 items: Sequence[T], 

102 item_namer: Callable[[T, int], str] | None, 

103) -> Callable[[int], str] | None: 

104 if item_namer is None: 

105 return None 

106 

107 def name_branch(index: int) -> str: 

108 return item_namer(items[index], index) 

109 

110 return name_branch 

111 

112 

113class MapSummaryGenerator: 

114 """Default summary generator for oversized `BatchResult` map payloads.""" 

115 

116 def __call__(self, result: BatchResult) -> str: 

117 fields = { 

118 "totalCount": result.total_count, 

119 "successCount": result.success_count, 

120 "failureCount": result.failure_count, 

121 "completionReason": result.completion_reason.value, 

122 "status": result.status.value, 

123 "type": "MapResult", 

124 } 

125 return json.dumps(fields) 

126 

127 

128@durable_callable 

129async def map_handler( 

130 items: Sequence[T], 

131 func: Callable[[T], Awaitable[R]], 

132 execution_state: ExecutionState, 

133 map_context: DurableContext, 

134 operation_identifier: OperationIdentifier, 

135 *, 

136 max_concurrency: int | None = None, 

137 completion_config: CompletionConfig | None = None, 

138 serdes: SerDes | None = None, 

139 item_serdes: SerDes | None = None, 

140 summary_generator: SummaryGenerator | None = MapSummaryGenerator(), 

141 nesting_type: NestingType = NestingType.NESTED, 

142 item_namer: Callable[[T, int], str] | None = None, 

143): 

144 """Execute a callable for each item through the parallel handler.""" 

145 handler = parallel_handler( 

146 callables=_create_map_branches(items, func), 

147 max_concurrency=max_concurrency, 

148 completion_config=completion_config or CompletionConfig(), 

149 serdes=serdes, 

150 summary_generator=summary_generator, 

151 item_serdes=item_serdes, 

152 nesting_type=nesting_type, 

153 execution_state=execution_state, 

154 parallel_context=map_context, 

155 operation_identifier=operation_identifier, 

156 top_level_sub_type=OperationSubType.MAP, 

157 iteration_sub_type=OperationSubType.MAP_ITERATION, 

158 name_prefix="map-item-", 

159 branch_namer=_create_map_branch_namer(items, item_namer), 

160 ) 

161 

162 return await handler() 

163 

164 

165def map( 

166 func: Callable[[U | BatchedInput[Any, U]], Awaitable[T]], 

167 items: Iterable[U], 

168 *, 

169 name: str | None = None, 

170 max_concurrency: int | None = None, 

171 completion_config: CompletionConfig | None = None, 

172 serdes: SerDes | None = None, 

173 item_serdes: SerDes | None = None, 

174 summary_generator: SummaryGenerator | None = MapSummaryGenerator(), 

175 nesting_type: NestingType = NestingType.NESTED, 

176 item_namer: Callable[[U, int], str] | None = None, 

177) -> asyncio.Task[BatchResult[T]]: 

178 """Process a collection durably with optional concurrency controls. 

179 

180 Args: 

181 func: Async callable that processes each item. 

182 items: Items to process. 

183 name: Optional durable operation name. 

184 max_concurrency: Optional limit for concurrent item processing. 

185 completion_config: Optional completion criteria. 

186 serdes: Optional serializer for the map result. 

187 item_serdes: Optional serializer for individual map item results. 

188 summary_generator: Optional summary generator for large map results. 

189 nesting_type: Whether map iterations use nested or flat operation ids. 

190 item_namer: Optional callable for naming map item iterations. 

191 """ 

192 context = get_durable_context("map") 

193 items_sequence = list(items) 

194 map_name = name if name is not None else getattr(func, "__name__", None) 

195 

196 async def run_map_handler() -> BatchResult[T]: 

197 map_context = get_durable_context("map") 

198 operation_id = map_context.step_id_prefix 

199 if operation_id is None: 

200 msg = "map operation id is not available in the current context" 

201 raise RuntimeError(msg) 

202 operation_identifier = OperationIdentifier( 

203 operation_id=operation_id, 

204 sub_type=OperationSubType.MAP, 

205 parent_id=map_context.parent_id, 

206 name=map_name, 

207 ) 

208 

209 handler = map_handler( 

210 items=items_sequence, 

211 func=func, 

212 execution_state=context.execution_state, 

213 map_context=map_context, 

214 operation_identifier=operation_identifier, 

215 max_concurrency=max_concurrency, 

216 completion_config=completion_config, 

217 serdes=serdes, 

218 item_serdes=item_serdes, 

219 summary_generator=summary_generator, 

220 nesting_type=nesting_type, 

221 item_namer=item_namer, 

222 ) 

223 return await handler() 

224 

225 return _run_in_child_context( 

226 run_map_handler, 

227 sub_type=OperationSubType.MAP, 

228 name=map_name, 

229 serdes=serdes, 

230 operation_name="map", 

231 )