Coverage for async_durable_execution/_operation/map.py: 100%
72 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-30 23:43 +0000
1"""Implementation for Durable Map operation."""
3from __future__ import annotations
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)
20from .parallel import (
21 _BATCH_RESULT_SERDES,
22 BatchResult,
23 CompletionConfig,
24 NestingType,
25 _validate_max_concurrency,
26)
27from .parallel import parallel_handler
28from .._core import (
29 DurableContext,
30 ExecutionState,
31 OperationIdentifier,
32 OperationSubType,
33 SerDes,
34 bind_current_context,
35 durable_callable,
36 get_current_context,
37 get_durable_context,
38)
39from ..extension import get_extension_context
41if TYPE_CHECKING:
42 from .child import SummaryGenerator
44logger = logging.getLogger(__name__)
46# Input item type
47T = TypeVar("T")
48# Result type
49R = TypeVar("R")
50U = TypeVar("U")
53def _run_in_child_context(
54 func: Callable[[], Awaitable[T]],
55 *,
56 sub_type: OperationSubType,
57 name: str | None = None,
58 serdes: SerDes | None = None,
59 summary_generator: SummaryGenerator | None = None,
60 is_virtual: bool = False,
61) -> asyncio.Task[T]:
62 """Run an SDK-owned child operation through the stable operation SPI."""
63 return (
64 get_extension_context()
65 ._reserve_sdk_operation(name) # noqa: SLF001
66 ._run_in_child_context( # noqa: SLF001
67 func,
68 sub_type=sub_type,
69 serdes=serdes,
70 summary_generator=summary_generator,
71 is_virtual=is_virtual,
72 )
73 )
76@dataclass(frozen=True)
77class BatchedInput(Generic[T, U]):
78 """Wrapper passed to batched map handlers."""
80 batch_input: T
81 items: list[U]
84@dataclass(frozen=True)
85class MapItemContext(DurableContext, Generic[T]):
86 """Context exposed while a map item function is executing."""
88 index: int = 0
89 items: Sequence[T] = field(default_factory=tuple)
92def get_map_item_context() -> MapItemContext[Any]:
93 """Return the active `MapItemContext`."""
94 current_context = get_current_context()
95 if not isinstance(current_context, MapItemContext):
96 msg = (
97 "get_map_item_context() can only be used while a map item "
98 "function is executing."
99 )
100 raise RuntimeError(msg)
101 return current_context
104def _bind_map_item_to_branch(
105 items: Sequence[T],
106 index: int,
107 func: Callable[[T], Awaitable[R]],
108) -> Callable[[], Awaitable[R]]:
109 async def run_branch() -> R:
110 logger.debug("🗺️ Processing map item: %s", index)
111 item = items[index]
112 child_context = get_durable_context()
113 map_item_context = MapItemContext(
114 execution_state=child_context.execution_state,
115 operation_identifier=child_context.operation_identifier,
116 step_id_prefix=child_context.step_id_prefix,
117 replaying=child_context.is_replaying(),
118 index=index,
119 items=items,
120 )
121 with bind_current_context(map_item_context):
122 result: R = await func(item)
123 logger.debug("✅ Processed map item: %s", index)
124 return result
126 return run_branch
129def _create_map_branches(
130 items: Sequence[T],
131 func: Callable[[T], Awaitable[R]],
132) -> list[Callable[[], Awaitable[R]]]:
133 return [
134 _bind_map_item_to_branch(items=items, index=index, func=func)
135 for index in range(len(items))
136 ]
139def _create_map_branch_namer(
140 items: Sequence[T],
141 item_namer: Callable[[T, int], str] | None,
142) -> Callable[[int], str] | None:
143 if item_namer is None:
144 return None
146 def name_branch(index: int) -> str:
147 return item_namer(items[index], index)
149 return name_branch
152class MapSummaryGenerator:
153 """Default summary generator for oversized `BatchResult` map payloads."""
155 def __call__(self, result: BatchResult) -> str:
156 fields = {
157 "totalCount": result.total_count,
158 "successCount": result.success_count,
159 "failureCount": result.failure_count,
160 "completionReason": result.completion_reason.value,
161 "status": result.status.value,
162 "type": "MapResult",
163 }
164 return json.dumps(fields)
167@durable_callable
168async def map_handler(
169 items: Sequence[T],
170 func: Callable[[T], Awaitable[R]],
171 execution_state: ExecutionState,
172 map_context: DurableContext,
173 operation_identifier: OperationIdentifier,
174 *,
175 max_concurrency: int | None = None,
176 completion_config: CompletionConfig | None = None,
177 serdes: SerDes | None = None,
178 item_serdes: SerDes | None = None,
179 summary_generator: SummaryGenerator | None = MapSummaryGenerator(),
180 nesting_type: NestingType = NestingType.NESTED,
181 item_namer: Callable[[T, int], str] | None = None,
182) -> BatchResult[R]:
183 """Execute a callable for each item through the parallel handler."""
184 handler = parallel_handler(
185 callables=_create_map_branches(items, func),
186 max_concurrency=max_concurrency,
187 completion_config=completion_config or CompletionConfig(),
188 serdes=serdes,
189 summary_generator=summary_generator,
190 item_serdes=item_serdes,
191 nesting_type=nesting_type,
192 execution_state=execution_state,
193 parallel_context=map_context,
194 operation_identifier=operation_identifier,
195 top_level_sub_type=OperationSubType.MAP,
196 iteration_sub_type=OperationSubType.MAP_ITERATION,
197 name_prefix="map-item-",
198 branch_namer=_create_map_branch_namer(items, item_namer),
199 )
201 return await handler()
204def map(
205 func: Callable[[U | BatchedInput[Any, U]], Awaitable[T]],
206 items: Iterable[U],
207 *,
208 name: str | None = None,
209 max_concurrency: int | None = None,
210 completion_config: CompletionConfig | None = None,
211 serdes: SerDes | None = None,
212 item_serdes: SerDes | None = None,
213 summary_generator: SummaryGenerator | None = MapSummaryGenerator(),
214 nesting_type: NestingType = NestingType.NESTED,
215 item_namer: Callable[[U, int], str] | None = None,
216) -> asyncio.Task[BatchResult[T]]:
217 """Start a durable map operation over a collection of items.
219 `map()` creates one durable child context per item and calls `func` with that
220 item. The item function must be async and may contain durable operations such
221 as `step()` or `wait()`.
223 The returned object is an `asyncio.Task`; awaiting it yields a `BatchResult`.
224 Calling `map()` without immediately awaiting it schedules the durable
225 operation in the background, consistent with other operation helpers.
227 By default, `map()` uses `CompletionConfig()` with no explicit success
228 threshold or failure tolerance: all-successful completion produces
229 `CompletionReason.ALL_COMPLETED`, while any observed failure completes the
230 operation as failed. Pass `completion_config` to use threshold-based or
231 custom completion.
233 Args:
234 func: Async callable that processes each item. It receives the original
235 item value and returns that item's result.
236 items: Items to process.
237 name: Optional durable operation name.
238 max_concurrency: Optional limit for in-flight items. A suspended item
239 retains its slot until it reaches a terminal state.
240 completion_config: Optional completion policy. Use
241 `CompletionConfig.thresholds()`, `first_successful()`,
242 `all_completed()`, `all_successful()`, or `custom()`.
243 serdes: Optional serializer for the final `BatchResult`.
244 item_serdes: Optional serializer for each item result.
245 summary_generator: Optional callable used to summarize oversized
246 checkpoint payloads.
247 nesting_type: Whether map iterations use nested or flat operation
248 identifiers.
249 item_namer: Optional callable for naming map item iterations.
251 Returns:
252 An `asyncio.Task` that resolves to a `BatchResult` containing one
253 `BatchItem` per input item.
255 Raises:
256 RuntimeError: If called outside a durable context.
257 """
258 _validate_max_concurrency(max_concurrency)
259 get_durable_context()
260 items_sequence = list(items)
261 map_name = name if name is not None else getattr(func, "__name__", None)
263 async def run_map_handler() -> BatchResult[T]:
264 map_context = get_durable_context()
265 operation_id = map_context.step_id_prefix
266 if operation_id is None:
267 msg = "map operation id is not available in the current context"
268 raise RuntimeError(msg)
269 operation_identifier = OperationIdentifier(
270 operation_id=operation_id,
271 sub_type=OperationSubType.MAP,
272 parent_id=map_context.parent_id,
273 name=map_name,
274 )
276 handler = map_handler(
277 items=items_sequence,
278 func=func,
279 execution_state=map_context.execution_state,
280 map_context=map_context,
281 operation_identifier=operation_identifier,
282 max_concurrency=max_concurrency,
283 completion_config=completion_config,
284 serdes=serdes,
285 item_serdes=item_serdes,
286 summary_generator=summary_generator,
287 nesting_type=nesting_type,
288 item_namer=item_namer,
289 )
290 return await handler()
292 return _run_in_child_context(
293 run_map_handler,
294 sub_type=OperationSubType.MAP,
295 name=map_name,
296 serdes=serdes if serdes is not None else _BATCH_RESULT_SERDES,
297 )