Coverage for async-durable-execution/src/async_durable_execution/primitive/child.py: 100%
187 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 16:54 +0000
1"""Implementation for run_in_child_context."""
3from __future__ import annotations
5import asyncio
6import functools
7import hashlib
8import logging
9from contextlib import contextmanager
10from dataclasses import dataclass
11from typing import TYPE_CHECKING, TypeVar, cast
13from .base import (
14 OperationExecutor,
15 OperationContext,
16)
17from ..context import bind_current_context, get_current_context
18from ..exceptions import (
19 CallableRuntimeError,
20 InvocationError,
21)
22from ..models import (
23 ContextOptions,
24 ErrorObject,
25 Operation,
26 OperationIdentifier,
27 OperationStatus,
28 OperationSubType,
29 OperationUpdate,
30)
31from ..serdes import deserialize, serialize
32from ..task import create_eager_task
34if TYPE_CHECKING:
35 from collections.abc import Awaitable, Callable
37 from ..composite.parallel import SummaryGenerator
38 from ..serdes import SerDes
39 from ..state import ExecutionState
41logger = logging.getLogger(__name__)
43T = TypeVar("T")
45# Checkpoint size limit in bytes (256KB)
46CHECKPOINT_SIZE_LIMIT = 256 * 1024
49class OrphanedChildException(BaseException):
50 """Raised when a child operation checkpoints after its parent context completed.
52 This inherits from BaseException so user code does not accidentally catch it
53 with broad exception handlers like ``except Exception``.
54 """
56 def __init__(self, message: str, operation_id: str):
57 super().__init__(message)
58 self.operation_id = operation_id
61class OperationIdGenerator:
62 """Generate deterministic operation ids within a durable execution scope."""
64 def __init__(self, prefix: str | None) -> None:
65 self._prefix = prefix
66 self._counter = 0
68 def increment(self) -> int:
69 self._counter += 1
70 return self._counter
72 def get_current(self) -> int:
73 return self._counter
75 def _create_step_id_for_logical_step(self, step: int) -> str:
76 """
77 Generate a step_id based on the given logical step.
78 This allows us to recover operation ids or even look
79 forward without changing the internal state of this context.
80 """
81 prefix: str | None = self._prefix
82 step_id: str = f"{prefix}-{step}" if prefix else str(step)
83 return hashlib.blake2b(step_id.encode()).hexdigest()[:64]
85 def create_step_id(self) -> str:
86 """Generate a step id, incrementing in order of invocation.
88 This method is an internal implementation detail. Do not rely the exact format of
89 the id generated by this method. It is subject to change without notice.
90 """
91 new_counter: int = self.increment()
92 return self._create_step_id_for_logical_step(new_counter)
95class ChildOperationExecutor(OperationExecutor[T]):
96 """Executor for child context operations."""
98 def __init__(
99 self,
100 func: Callable[[], Awaitable[T]],
101 state: ExecutionState,
102 operation_identifier: OperationIdentifier,
103 *,
104 serdes: SerDes | None = None,
105 summary_generator: SummaryGenerator | None = None,
106 is_virtual: bool = False,
107 ):
108 """Initialize the child operation executor.
110 Args:
111 func: The child context function to execute
112 state: The execution state
113 operation_identifier: The operation identifier
114 serdes: Optional serializer for the child context result.
115 summary_generator: Optional summary generator for large child results.
116 is_virtual: Whether this child context should skip lifecycle checkpoints.
117 """
118 super().__init__(state=state, operation_identifier=operation_identifier)
119 self.func = func
120 self.serdes = serdes
121 self.summary_generator = summary_generator
122 self.is_virtual = is_virtual
124 async def start(self) -> T:
125 """Start a new child context operation."""
126 if not self.is_virtual:
127 start_operation: OperationUpdate = OperationUpdate.create_context_start(
128 identifier=self.operation_identifier,
129 sub_type=self.operation_identifier.sub_type,
130 )
131 await self.create_checkpoint(start_operation, is_sync=False)
133 return await self.execute(None)
135 async def replay(self, operation: Operation) -> T:
136 """Replay an existing child context operation from its checkpoint."""
137 if (
138 operation.status is OperationStatus.SUCCEEDED
139 and not self._is_replay_children(operation)
140 ):
141 logger.debug(
142 "Child context already completed, skipping execution for id: %s, name: %s",
143 self.operation_identifier.operation_id,
144 self.operation_name,
145 )
146 result_payload = self._get_result(operation)
147 if result_payload is None:
148 return cast("T", None)
150 result: T = await deserialize(
151 serdes=self.serdes,
152 data=result_payload,
153 operation_id=self.operation_id,
154 durable_execution_arn=self.durable_execution_arn,
155 )
156 return result
158 if operation.status is OperationStatus.SUCCEEDED and self._is_replay_children(
159 operation
160 ):
161 return await self.execute(operation)
163 if operation.status is OperationStatus.FAILED:
164 self._raise_callable_error(operation)
166 return await self.execute(operation)
168 async def execute(self, operation: Operation | None) -> T:
169 """Execute child context function with error handling and large payload support.
171 Args:
172 operation: The checkpointed operation state, if any
174 Returns:
175 The result of executing the child context function
177 Raises:
178 SuspendExecution: Re-raised without checkpointing
179 InvocationError: Re-raised after checkpointing FAIL
180 CallableRuntimeError: Raised for other exceptions after checkpointing FAIL
181 """
182 logger.debug(
183 "▶️ Executing child context for id: %s, name: %s",
184 self.operation_identifier.operation_id,
185 self.operation_identifier.name,
186 )
187 try:
188 replaying_children = self._is_replay_children(operation)
189 raw_result: T = await self.func()
191 if self.is_virtual:
192 logger.debug(
193 "Virtual context: Exiting child context without creating another checkpoint. id: %s, name: %s",
194 self.operation_identifier.operation_id,
195 self.operation_identifier.name,
196 )
197 return raw_result
199 # If in replay_children mode, return without checkpointing
200 if replaying_children:
201 logger.debug(
202 "ReplayChildren mode: Executed child context again on replay due to large payload. Exiting child context without creating another checkpoint. id: %s, name: %s",
203 self.operation_identifier.operation_id,
204 self.operation_identifier.name,
205 )
206 return raw_result
208 # Serialize result
209 serialized_result: str = await serialize(
210 serdes=self.serdes,
211 value=raw_result,
212 operation_id=self.operation_id,
213 durable_execution_arn=self.durable_execution_arn,
214 )
216 # Check payload size and use ReplayChildren mode if needed
217 # Summary Generator Logic:
218 # When the serialized result exceeds 256KB, we use ReplayChildren mode to avoid
219 # checkpointing large payloads. Instead, we checkpoint a compact summary and mark
220 # the operation for replay. This matches the TypeScript implementation behavior.
221 #
222 # See TypeScript reference:
223 # - aws-durable-execution-sdk-js/src/handlers/run-in-child-context-handler/run-in-child-context-handler.ts (lines ~200-220)
224 #
225 # The summary generator creates a JSON summary with metadata (type, counts, status)
226 # instead of the full BatchResult. During replay, the child context is re-executed
227 # to reconstruct the full result rather than deserializing from the checkpoint.
228 replay_children: bool = False
229 if len(serialized_result) > CHECKPOINT_SIZE_LIMIT:
230 logger.debug(
231 "Large payload detected, using ReplayChildren mode: id: %s, name: %s, payload_size: %d, limit: %d",
232 self.operation_identifier.operation_id,
233 self.operation_identifier.name,
234 len(serialized_result),
235 CHECKPOINT_SIZE_LIMIT,
236 )
237 replay_children = True
238 # Use summary generator if provided, otherwise use empty string (matches TypeScript)
239 serialized_result = (
240 self.summary_generator(raw_result) if self.summary_generator else ""
241 )
243 # Checkpoint SUCCEED
244 success_operation: OperationUpdate = OperationUpdate.create_context_succeed(
245 identifier=self.operation_identifier,
246 payload=serialized_result,
247 sub_type=self.operation_identifier.sub_type,
248 context_options=ContextOptions(replay_children=replay_children),
249 )
250 # Checkpoint child context SUCCEED with blocking (is_sync=True, default).
251 # Must ensure the child context result is persisted before returning to the parent.
252 # This guarantees the result is durable and child operations won't be re-executed on replay
253 # (unless replay_children=True for large payloads).
254 await self.create_checkpoint(success_operation)
256 logger.debug(
257 "✅ Successfully completed child context for id: %s, name: %s",
258 self.operation_identifier.operation_id,
259 self.operation_identifier.name,
260 )
261 if replay_children:
262 return raw_result
264 return await deserialize( # noqa: TRY300
265 serdes=self.serdes,
266 data=serialized_result,
267 operation_id=self.operation_id,
268 durable_execution_arn=self.durable_execution_arn,
269 )
270 except Exception as e:
271 error_object = ErrorObject.from_exception(e)
272 # Virtual deliberately does not write checkpoints, but exception still propagates below
273 if not self.is_virtual:
274 fail_operation: OperationUpdate = OperationUpdate.create_context_fail(
275 identifier=self.operation_identifier,
276 error=error_object,
277 sub_type=self.operation_identifier.sub_type,
278 )
279 # Checkpoint child context FAIL with blocking (is_sync=True, default).
280 # Must ensure the failure state is persisted before raising the exception.
281 # This guarantees the error is durable and child operations won't be re-executed on replay.
282 await self.create_checkpoint(fail_operation)
284 # InvocationError and its derivatives can be retried.
285 # When we encounter an invocation error (in all of its forms), we
286 # bubble that error upwards (with the checkpoint in place for
287 # non-virtual) such that we reach the execution handler at the
288 # very top, which will then make the backend retry.
289 if isinstance(e, InvocationError):
290 raise
291 raise CallableRuntimeError.from_error_object(error_object) from e
293 @staticmethod
294 def _is_replay_children(operation: Operation | None) -> bool:
295 if operation is None or operation.context_details is None:
296 return False
297 return operation.context_details.replay_children
299 @staticmethod
300 def _get_result(operation: Operation) -> str | None:
301 if operation.context_details is None:
302 return None
303 return operation.context_details.result
305 @staticmethod
306 def _raise_callable_error(operation: Operation) -> None:
307 error = operation.context_details.error if operation.context_details else None
308 if error is None:
309 msg = "Unknown error. No ErrorObject exists on the Checkpoint Operation."
310 raise CallableRuntimeError(
311 message=msg,
312 error_type=None,
313 data=None,
314 stack_trace=None,
315 )
317 raise CallableRuntimeError.from_error_object(error)
320@dataclass(frozen=True)
321class DurableContext(OperationContext):
322 """Runtime context available to a durable handler or child context."""
324 step_id_prefix: str | None = None
325 replaying: bool = False
327 @functools.cached_property
328 def step_counter(self):
329 return OperationIdGenerator(self.operation_id_generator_prefix)
331 @property
332 def is_virtual(self) -> bool:
333 return self.operation_identifier.parent_id != self.operation_id_generator_prefix
335 @property
336 def operation_id_generator_prefix(self):
337 return (
338 self.step_id_prefix
339 if self.step_id_prefix is not None
340 else self.operation_identifier.parent_id
341 )
343 def create_child_context(
344 self, operation_id: str, *, is_virtual: bool = False
345 ) -> DurableContext:
346 """Create a child context for the given operation."""
347 child_parent_id: str | None = self.parent_id if is_virtual else operation_id
348 logger.debug(
349 "Creating child context for operation %s (is_virtual=%s)",
350 operation_id,
351 is_virtual,
352 )
353 return DurableContext(
354 execution_state=self.execution_state,
355 operation_identifier=OperationIdentifier(
356 operation_id=None,
357 sub_type=OperationSubType.EXECUTION,
358 parent_id=child_parent_id,
359 ),
360 step_id_prefix=operation_id,
361 replaying=self.is_replaying(),
362 )
364 def is_replaying(self) -> bool:
365 """Return True while this context is replaying prior operations."""
366 return self.replaying
368 def _set_replay_status_new(self) -> None:
369 object.__setattr__(self, "replaying", False)
371 def _peek_next_operation_id(self) -> str:
372 return self.step_counter._create_step_id_for_logical_step( # noqa: SLF001
373 self.step_counter.get_current() + 1
374 )
376 def _next_operation_result(self) -> Operation | None:
377 return self.execution_state.operations.get(self._peek_next_operation_id())
379 def _next_operation_exists(self) -> bool:
380 return self._next_operation_result() is not None
382 def _next_operation_is_terminal_checkpoint(self) -> bool:
383 operation = self._next_operation_result()
384 if operation is None:
385 return False
386 return operation.status in {
387 OperationStatus.SUCCEEDED,
388 OperationStatus.FAILED,
389 OperationStatus.CANCELLED,
390 OperationStatus.STOPPED,
391 OperationStatus.TIMED_OUT,
392 }
394 @contextmanager
395 def _replay_aware(self, *, executes_user_code: bool = False):
396 """Update this context's replay status around one durable operation."""
397 was_replaying = self.is_replaying()
398 next_exists = was_replaying and self._next_operation_exists()
399 next_terminal = was_replaying and self._next_operation_is_terminal_checkpoint()
400 flip_after = (
401 was_replaying
402 and not executes_user_code
403 and next_exists
404 and not next_terminal
405 )
407 if was_replaying and (
408 not next_exists or (executes_user_code and not next_terminal)
409 ):
410 self._set_replay_status_new()
412 try:
413 yield
414 finally:
415 if flip_after:
416 self._set_replay_status_new()
417 elif self.is_replaying() and not self._next_operation_exists():
418 self._set_replay_status_new()
421def run_in_child_context(
422 func: Callable[[], Awaitable[T]],
423 *,
424 name: str | None = None,
425 serdes: SerDes | None = None,
426 summary_generator: SummaryGenerator | None = None,
427 is_virtual: bool = False,
428) -> asyncio.Task[T]:
429 """Execute a durable sub-workflow inside its own child context.
431 Args:
432 func: The child context function to execute.
433 name: Optional durable operation name.
434 serdes: Optional serializer for the child context result.
435 summary_generator: Optional summary generator for large child results.
436 is_virtual: Whether the child context should skip lifecycle checkpoints.
437 """
439 step_name = name if name is not None else getattr(func, "__name__", None)
440 return _create_child_context_task(
441 func,
442 sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT,
443 name=step_name,
444 serdes=serdes,
445 summary_generator=summary_generator,
446 is_virtual=is_virtual,
447 operation_name="run_in_child_context",
448 )
451def _create_child_context_task(
452 func: Callable[[], Awaitable[T]],
453 *,
454 sub_type: OperationSubType,
455 name: str | None = None,
456 serdes: SerDes | None = None,
457 summary_generator: SummaryGenerator | None = None,
458 is_virtual: bool = False,
459 operation_name: str | None = None,
460) -> asyncio.Task[T]:
461 context = get_durable_context(operation_name or "run_in_child_context")
463 operation_id = context.step_counter.create_step_id()
464 child_context = context.create_child_context(
465 operation_id=operation_id,
466 is_virtual=is_virtual,
467 )
468 operation_identifier = OperationIdentifier(
469 operation_id=operation_id,
470 sub_type=sub_type,
471 parent_id=context.parent_id,
472 name=name,
473 )
475 async def run_child_context() -> T:
476 with context._replay_aware():
477 return await _run_child_context(
478 func,
479 context=context,
480 child_context=child_context,
481 operation_identifier=operation_identifier,
482 serdes=serdes,
483 summary_generator=summary_generator,
484 is_virtual=is_virtual,
485 )
487 return create_eager_task(run_child_context)
490async def _run_in_child_context(
491 func: Callable[[], Awaitable[T]],
492 *,
493 sub_type: OperationSubType,
494 name: str | None = None,
495 serdes: SerDes | None = None,
496 summary_generator: SummaryGenerator | None = None,
497 is_virtual: bool = False,
498) -> T:
499 """Execute a durable sub-workflow with an explicit operation subtype."""
500 context = get_durable_context("run_in_child_context")
501 with context._replay_aware():
502 operation_id = context.step_counter.create_step_id()
504 child_context = context.create_child_context(
505 operation_id=operation_id,
506 is_virtual=is_virtual,
507 )
508 operation_identifier = OperationIdentifier(
509 operation_id=operation_id,
510 sub_type=sub_type,
511 parent_id=context.parent_id,
512 name=name,
513 )
515 return await _run_child_context(
516 func,
517 context=context,
518 child_context=child_context,
519 operation_identifier=operation_identifier,
520 serdes=serdes,
521 summary_generator=summary_generator,
522 is_virtual=is_virtual,
523 )
526async def _run_child_context(
527 func: Callable[[], Awaitable[T]],
528 *,
529 context: DurableContext,
530 child_context: DurableContext,
531 operation_identifier: OperationIdentifier,
532 serdes: SerDes | None = None,
533 summary_generator: SummaryGenerator | None = None,
534 is_virtual: bool = False,
535) -> T:
536 async def callable_with_child_context():
537 with bind_current_context(child_context):
538 return await func()
540 executor: ChildOperationExecutor[T] = ChildOperationExecutor(
541 callable_with_child_context,
542 context.execution_state,
543 operation_identifier,
544 serdes=serdes,
545 summary_generator=summary_generator,
546 is_virtual=is_virtual,
547 )
548 return await executor.process()
551def get_durable_context(operation_name: str | None = None) -> DurableContext:
552 current_context = get_current_context()
553 if (
554 current_context is None
555 or not hasattr(current_context, "execution_state")
556 or not hasattr(current_context, "operation_identifier")
557 or not hasattr(current_context, "step_counter")
558 or not hasattr(current_context, "create_child_context")
559 ):
560 msg = (
561 f"{operation_name or 'Durable operations'} can only be used while a durable function or child "
562 "context is executing."
563 )
564 raise RuntimeError(msg)
565 return cast(DurableContext, current_context)