From dbaccc93964a40d7e2c7905cc6f78ac823573943 Mon Sep 17 00:00:00 2001 From: Alex Wang Date: Fri, 11 Sep 2026 00:13:14 +0000 Subject: [PATCH] feat(insight): export records asynchronously --- .../README.md | 10 +- .../_export_scheduler.py | 162 ++++++++++++++++++ .../plugin.py | 21 +-- .../tests/test_export_scheduler.py | 159 +++++++++++++++++ .../tests/test_plugin.py | 26 +-- 5 files changed, 345 insertions(+), 33 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py create mode 100644 packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py diff --git a/packages/aws-durable-execution-sdk-python-insight/README.md b/packages/aws-durable-execution-sdk-python-insight/README.md index ed47c053..be1027df 100644 --- a/packages/aws-durable-execution-sdk-python-insight/README.md +++ b/packages/aws-durable-execution-sdk-python-insight/README.md @@ -55,10 +55,12 @@ and `top-level` vs `full-tree` operation detail all mirror the JS plugin. Behavior is validated cross-SDK by the `insight` conformance suite (`aws-durable-execution-conformance-tests-insight`). -> **Note (`on-change` emission).** In `on-change` mode, exporter calls currently -> run synchronously on the SDK checkpoint path, so a slow exporter can delay -> workflow progress. Asynchronous scheduling/coalescing is deferred and tracked -> in [issue #687](https://github.com/aws/aws-durable-execution-sdk-python/issues/687). +> **Note (asynchronous export).** Export rendering, truncation, `export()`, and +> `flush()` run on one lazy background worker per plugin. Checkpoint hooks only +> replace the latest pending snapshot and wake the worker. Consecutive +> `on-change` snapshots may coalesce while an export is in flight. An invocation +> that emits a record drains the latest snapshot and flushes exporters before it +> returns; invocations that emit nothing do not start or flush the worker. ## Requirements diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py new file mode 100644 index 00000000..ef22ea54 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/_export_scheduler.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 +"""Latest-pending asynchronous export scheduling for Workflow Insight.""" + +from __future__ import annotations + +import logging +import threading +from typing import Any + +from aws_durable_execution_sdk_python_insight.truncation import truncate_record +from aws_durable_execution_sdk_python_insight.types import InsightExporter + + +_logger = logging.getLogger("aws_durable_execution_sdk_python_insight") + + +class _ExportScheduler: + """Run all exporters on one lazy worker with one latest pending record.""" + + def __init__(self, exporters: list[InsightExporter]) -> None: + self._exporters = exporters + self._condition = threading.Condition(threading.Lock()) + self._pending: dict[str, Any] | None = None + self._flush_requested = False + self._flush_event: threading.Event | None = None + self._worker: threading.Thread | None = None + self._disabled = False + + def schedule(self, record: dict[str, Any]) -> None: + """Replace the pending snapshot and return without running exporters.""" + displaced: dict[str, Any] | None = None + failed_pending: dict[str, Any] | None = None + start_error: Exception | None = None + with self._condition: + if self._disabled: + return + displaced = self._pending + self._pending = record + failed_pending, start_error = self._ensure_worker_locked() + self._condition.notify() + # Releasing either record may run custom finalizers, so do it unlocked. + del displaced, failed_pending + if start_error is not None: + _logger.warning( + "workflow-insight: could not start export worker; disabling " + "asynchronous export: %s", + start_error, + ) + + def drain(self) -> None: + """Wait until the latest pending record is exported and exporters flush.""" + failed_pending: dict[str, Any] | None = None + start_error: Exception | None = None + with self._condition: + if self._disabled: + return + if not self._flush_requested: + self._flush_requested = True + self._flush_event = threading.Event() + flush_event = self._flush_event + assert flush_event is not None + failed_pending, start_error = self._ensure_worker_locked() + started = not self._disabled + self._condition.notify() + del failed_pending + if start_error is not None: + _logger.warning( + "workflow-insight: could not start export worker; disabling " + "asynchronous export: %s", + start_error, + ) + if started: + flush_event.wait() + + def _ensure_worker_locked( + self, + ) -> tuple[dict[str, Any] | None, Exception | None]: + if self._worker is not None and self._worker.is_alive(): + return None, None + worker = threading.Thread( + target=self._run, + name=f"workflow-insight-export-{id(self)}", + daemon=True, + ) + self._worker = worker + try: + worker.start() + except Exception as exc: # noqa: BLE001 - instrumentation must not escape hooks + self._disabled = True + self._worker = None + failed_pending = self._pending + self._pending = None + failed_event = self._flush_event + self._flush_event = None + self._flush_requested = False + if failed_event is not None: + failed_event.set() + return failed_pending, exc + return None, None + + def _run(self) -> None: + while True: + record: dict[str, Any] | None = None + flush_event: threading.Event | None = None + with self._condition: + while self._pending is None and not self._flush_requested: + self._condition.wait() + if self._pending is not None: + record = self._pending + self._pending = None + else: + flush_event = self._flush_event + self._flush_event = None + self._flush_requested = False + + if record is not None: + self._export(record) + continue + + self._flush() + if flush_event is not None: + flush_event.set() + with self._condition: + if self._pending is None and not self._flush_requested: + self._worker = None + return + + def _export(self, record: dict[str, Any]) -> None: + for exporter in self._exporters: + try: + shaped = truncate_record( + record, exporter.max_record_size_bytes, exporter.render + ) + exporter.export(shaped) + except Exception as exc: # noqa: BLE001 - one exporter must not break others + _logger.warning( + "workflow-insight: exporter %s failed: %s", + type(exporter).__name__, + exc, + ) + + def _flush(self) -> None: + for exporter in self._exporters: + try: + exporter.flush() + except Exception as exc: # noqa: BLE001 - one exporter must not break others + _logger.warning( + "workflow-insight: exporter %s flush failed: %s", + type(exporter).__name__, + exc, + ) + + # Test helpers. + def _worker_alive(self) -> bool: + with self._condition: + return self._worker is not None and self._worker.is_alive() + + def _pending_count(self) -> int: + with self._condition: + return int(self._pending is not None) diff --git a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py index 13f16a06..f19796a0 100644 --- a/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/src/aws_durable_execution_sdk_python_insight/plugin.py @@ -33,7 +33,6 @@ import datetime import json import math -import sys import threading from typing import Any, Callable @@ -47,10 +46,10 @@ OperationType, ) +from aws_durable_execution_sdk_python_insight._export_scheduler import _ExportScheduler from aws_durable_execution_sdk_python_insight.exporters.lambda_log_exporter import ( LambdaLogExporter, ) -from aws_durable_execution_sdk_python_insight.truncation import truncate_record from aws_durable_execution_sdk_python_insight.types import ( ContentConfig, EmitMode, @@ -205,6 +204,7 @@ def __init__(self, config: WorkflowInsightConfig) -> None: self._exporters: list[InsightExporter] = ( list(config.exporters) if config.exporters else [LambdaLogExporter()] ) + self._scheduler = _ExportScheduler(self._exporters) self._state: dict[str, _ExecutionState] = {} self._lock = threading.Lock() @@ -322,6 +322,7 @@ def on_invocation_end(self, info: InvocationEndInfo) -> None: output_raw=info.execution_result if is_terminal else None, error=info.error if is_terminal else None, ) + self._scheduler.drain() # Clear state after EVERY invocation end, including PENDING/RETRY. The # next invocation rebuilds it from InvocationStartInfo.operations, so a @@ -435,21 +436,7 @@ def _emit( record["error"] = {"name": error.type, "message": error.message} record["operations"] = self._build_operations(operations) - for exporter in self._exporters: - try: - shaped = truncate_record( - record, exporter.max_record_size_bytes, exporter.render - ) - exporter.export(shaped) - except Exception as exc: # noqa: BLE001 - one exporter must not break others / the execution - # NOTE (parity gap, same as JS Promise.allSettled): exporter - # failures are swallowed so instrumentation never breaks the - # execution. A silently broken exporter is indistinguishable - # from success; we at least log to stderr. - print( - f"[workflow-insight] exporter {type(exporter).__name__} failed: {exc}", - file=sys.stderr, - ) # noqa: T201 + self._scheduler.schedule(record) def workflow_insight(config: WorkflowInsightConfig) -> WorkflowInsightPlugin: diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py new file mode 100644 index 00000000..34879059 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_export_scheduler.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: 2026-present Amazon.com, Inc. or its affiliates. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import threading +import time +from typing import Any + +from aws_durable_execution_sdk_python_insight._export_scheduler import ( + _ExportScheduler, +) + + +def _record(value: str) -> dict[str, Any]: + return {"status": "RUNNING", "value": value, "operations": []} + + +def _wait_until(predicate, timeout: float = 5.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.005) + return predicate() + + +class CaptureExporter: + max_record_size_bytes: int | None = None + + def __init__(self) -> None: + self.calls: list[tuple[str, str | None]] = [] + + def render(self, record: dict[str, Any]) -> dict[str, Any]: + return record + + def export(self, record: dict[str, Any]) -> None: + self.calls.append(("export", record["value"])) + + def flush(self) -> None: + self.calls.append(("flush", None)) + + +class BlockingExporter(CaptureExporter): + def __init__(self) -> None: + super().__init__() + self.started = threading.Event() + self.release = threading.Event() + + def export(self, record: dict[str, Any]) -> None: + self.started.set() + self.release.wait(5.0) + super().export(record) + + +class FailingExporter(CaptureExporter): + def export(self, record: dict[str, Any]) -> None: + raise RuntimeError("export failed") + + def flush(self) -> None: + raise RuntimeError("flush failed") + + +def test_latest_pending_coalesces_without_blocking_schedule() -> None: + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(_record("first")) + assert exporter.started.wait(5.0) + + start = time.monotonic() + scheduler.schedule(_record("middle")) + scheduler.schedule(_record("latest")) + assert time.monotonic() - start < 0.5 + assert scheduler._pending_count() == 1 + + exporter.release.set() + scheduler.drain() + assert exporter.calls == [ + ("export", "first"), + ("export", "latest"), + ("flush", None), + ] + assert _wait_until(lambda: not scheduler._worker_alive()) + + +def test_exporter_failure_does_not_block_other_exporters() -> None: + failing = FailingExporter() + capture = CaptureExporter() + scheduler = _ExportScheduler([failing, capture]) + + scheduler.schedule(_record("terminal")) + + scheduler.drain() + assert capture.calls == [("export", "terminal"), ("flush", None)] + + +def test_drain_flushes_after_export() -> None: + capture = CaptureExporter() + scheduler = _ExportScheduler([capture]) + + scheduler.schedule(_record("terminal")) + + scheduler.drain() + assert capture.calls == [("export", "terminal"), ("flush", None)] + + +def test_worker_start_failure_never_escapes_hook(monkeypatch) -> None: + def fail_start(self) -> None: # noqa: ARG001 + raise RuntimeError("cannot start") + + monkeypatch.setattr(threading.Thread, "start", fail_start) + scheduler = _ExportScheduler([CaptureExporter()]) + + scheduler.schedule(_record("dropped")) + scheduler.drain() + + assert scheduler._pending_count() == 0 + + +def test_superseded_record_finalizes_after_lane_unlock() -> None: + scheduler = _ExportScheduler([BlockingExporter()]) + exporter = scheduler._exporters[0] + assert isinstance(exporter, BlockingExporter) + scheduler.schedule(_record("inflight")) + assert exporter.started.wait(5.0) + finalized = threading.Event() + + class ReentrantValue: + def __del__(self) -> None: + scheduler.schedule(_record("from-finalizer")) + finalized.set() + + pending = _record("superseded") + pending["payload"] = ReentrantValue() + scheduler.schedule(pending) + del pending + + scheduler.schedule(_record("replacement")) + + assert finalized.wait(5.0) + exporter.release.set() + scheduler.drain() + + +def test_drain_waits_for_blocked_exporter() -> None: + exporter = BlockingExporter() + scheduler = _ExportScheduler([exporter]) + scheduler.schedule(_record("terminal")) + assert exporter.started.wait(5.0) + drain_thread = threading.Thread(target=scheduler.drain) + + drain_thread.start() + assert _wait_until(drain_thread.is_alive) + exporter.release.set() + drain_thread.join(5.0) + + assert not drain_thread.is_alive() + assert exporter.calls == [("export", "terminal"), ("flush", None)] diff --git a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py index 485bcaab..9d7638eb 100644 --- a/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py +++ b/packages/aws-durable-execution-sdk-python-insight/tests/test_plugin.py @@ -51,6 +51,7 @@ def __init__(self, max_record_size_bytes: int | None = None, render=None) -> Non self.max_record_size_bytes = max_record_size_bytes self._render = render or (lambda r: r) self.records: list[dict[str, Any]] = [] + self.flush_count = 0 def render(self, record: dict[str, Any]) -> Any: return self._render(record) @@ -59,7 +60,7 @@ def export(self, record: dict[str, Any]) -> None: self.records.append(record) def flush(self) -> None: - return None + self.flush_count += 1 def _step( @@ -186,6 +187,8 @@ def test_on_failure_success_emits_nothing(): ) _run(plugin, ops=[_step("greet")], status=InvocationStatus.SUCCEEDED) assert exporter.records == [] + assert exporter.flush_count == 0 + assert not plugin._scheduler._worker_alive() def test_sampling_zero_emits_nothing(): @@ -365,10 +368,10 @@ def test_cold_resume_reports_prior_terminal_ops_with_fresh_plugin(): assert rec["durationMs"] is not None and rec["durationMs"] >= 0 -# -- on-change emits an updated record per change (comment 2) ---------------- +# -- on-change schedules progress and delivers terminal state ---------------- -def test_on_change_emits_running_on_each_change(): +def test_on_change_schedules_running_and_delivers_terminal(): exporter = CaptureExporter() plugin = workflow_insight( WorkflowInsightConfig(exporters=[exporter], emit_mode="on-change") @@ -376,26 +379,25 @@ def test_on_change_emits_running_on_each_change(): op1 = _step("s1", op_id="1") op2 = _step("s2", op_id="2") - plugin.on_invocation_start(_start(operations={})) # RUNNING #1 (start) + plugin.on_invocation_start(_start(operations={})) plugin.on_operation_change( OperationChangeInfo( execution_arn=ARN, updated_operations=_ops(op1), operations=_ops(op1) ) - ) # RUNNING #2 + ) plugin.on_operation_change( OperationChangeInfo( execution_arn=ARN, updated_operations=_ops(op2), operations=_ops(op1, op2) ) - ) # RUNNING #3 - plugin.on_invocation_end(_end(operations=_ops(op1, op2))) # SUCCEEDED #4 + ) + plugin.on_invocation_end(_end(operations=_ops(op1, op2))) - statuses = [r["status"] for r in exporter.records] - assert statuses == ["RUNNING", "RUNNING", "RUNNING", "SUCCEEDED"] - # The record emitted after the 2nd change already carries both operations. - assert [op["name"] for op in exporter.records[2]["operations"]] == ["s1", "s2"] + statuses = [record["status"] for record in exporter.records] + assert statuses + assert statuses[-1] == "SUCCEEDED" + assert set(statuses[:-1]) <= {"RUNNING"} final = exporter.records[-1] assert [op["name"] for op in final["operations"]] == ["s1", "s2"] - # No duplicate operation entries within a record (no end/change double-count). ids = [op["id"] for op in final["operations"]] assert len(ids) == len(set(ids))