feat: add plugin interface (#371)

Author: zhongkechen

.github/hooks/pre-commit

@@ -0,0 +1,9 @@
+#!/bin/sh
+
+if hatch fmt --check; then
+    echo "Hatch fmt check passed!"
+else
+    hatch fmt
+    echo "Error: hatch fmt modified your files. Please re-stage and commit again."
+    exit 1
+fi

packages/aws-durable-execution-sdk-python-examples/examples-catalog.json

@@ -602,6 +602,17 @@
         "ExecutionTimeout": 300
       },
       "path": "./src/parallel/parallel_with_named_branches.py"
+    },
+    {
+      "name": "Plugin",
+      "description": "Test plugin",
+      "handler": "execution_with_plugin.handler",
+      "integration": true,
+      "durableConfig": {
+        "RetentionPeriodInDays": 7,
+        "ExecutionTimeout": 300
+      },
+      "path": "./src/plugin/execution_with_plugin.py"
     }
   ]
 }

packages/aws-durable-execution-sdk-python-examples/src/plugin/execution_with_plugin.py

@@ -0,0 +1,63 @@
+"""Demonstrates handler execution without any durable operations."""
+
+import logging
+from typing import Any
+
+from aws_durable_execution_sdk_python import StepContext
+from aws_durable_execution_sdk_python.context import (
+    DurableContext,
+    durable_step,
+    durable_with_child_context,
+)
+from aws_durable_execution_sdk_python.execution import durable_execution
+from aws_durable_execution_sdk_python.plugin import (
+    DurableInstrumentationPlugin,
+)
+
+
+class MyPlugin(DurableInstrumentationPlugin):
+    logger = logging.getLogger("MyPlugin")
+
+    def on_operation_start(self, info):
+        self.logger.info(f"Operation started: {info}")
+
+    def on_operation_end(self, info):
+        self.logger.info(f"Operation ended: {info}")
+
+    def on_invocation_start(self, info):
+        self.logger.info(f"Invocation started: {info}")
+
+    def on_invocation_end(self, info):
+        self.logger.info(f"Invocation ended: {info}")
+
+    def on_user_function_start(self, info) -> None:
+        self.logger.info(f"User function started: {info}")
+
+    def on_user_function_end(self, info) -> None:
+        self.logger.info(f"User function ended: {info}")
+
+
+@durable_step
+def add_numbers(_step_context: StepContext, a: int, b: int) -> int:
+    return a + b
+
+
+@durable_with_child_context
+def add_numbers_in_child(child_context: DurableContext, a: int, b: int):
+    result: int = child_context.step(
+        add_numbers(a, b),
+        name="add-a-and-b",
+    )
+    return result
+
+
+@durable_execution(plugins=[MyPlugin()])
+def handler(_event: Any, context: DurableContext) -> int:
+    result: int = context.run_in_child_context(
+        add_numbers_in_child(6, 4),
+        name="add-6-and-4",
+    )
+    return context.step(
+        add_numbers(result, 2),
+        name="add-result-to-2",
+    )

packages/aws-durable-execution-sdk-python-examples/template.yaml

@@ -977,6 +977,24 @@
           "ExecutionTimeout": 300
         }
       }
+    },
+    "ExecutionWithPlugin": {
+      "Type": "AWS::Serverless::Function",
+      "Properties": {
+        "CodeUri": "build/",
+        "Handler": "execution_with_plugin.handler",
+        "Description": "Test plugin",
+        "Role": {
+          "Fn::GetAtt": [
+            "DurableFunctionRole",
+            "Arn"
+          ]
+        },
+        "DurableConfig": {
+          "RetentionPeriodInDays": 7,
+          "ExecutionTimeout": 300
+        }
+      }
     }
   }
 }

packages/aws-durable-execution-sdk-python-examples/test/plugin/test_plugin.py

@@ -0,0 +1,24 @@
+"""Tests for step example."""
+
+import pytest
+from aws_durable_execution_sdk_python.execution import InvocationStatus
+
+from src.plugin import execution_with_plugin
+from test.conftest import deserialize_operation_payload
+
+
+@pytest.mark.example
+@pytest.mark.durable_execution(
+    handler=execution_with_plugin.handler,
+    lambda_function_name="Plugin",
+)
+def test_plugin(durable_runner):
+    """Test basic step example."""
+    with durable_runner:
+        result = durable_runner.run(input="{}", timeout=10)
+
+    assert result.status is InvocationStatus.SUCCEEDED
+    assert deserialize_operation_payload(result.result) == 12
+
+    step_result = result.get_step("add-result-to-2")
+    assert deserialize_operation_payload(step_result.result) == 12

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/concurrency/executor.py

@@ -30,7 +30,7 @@
     TimedSuspendExecution,
 )
 from aws_durable_execution_sdk_python.identifier import OperationIdentifier
-from aws_durable_execution_sdk_python.lambda_service import ErrorObject
+from aws_durable_execution_sdk_python.lambda_service import ErrorObject, OperationType
 from aws_durable_execution_sdk_python.operation.child import child_handler
 
 
@@ -428,9 +428,10 @@ def _execute_item_in_child_context(
         # For FLAT `child_handler` skips checkpoints, so not used.
         # Construct it unconditionally to keep the call simple.
         operation_identifier = OperationIdentifier(
-            operation_id,
-            executor_context._parent_id,  # noqa: SLF001
-            name,
+            operation_id=operation_id,
+            sub_type=self.sub_type_iteration,
+            parent_id=executor_context._parent_id,  # noqa: SLF001
+            name=name,
         )
 
         def run_in_child_handler() -> ResultType:

packages/aws-durable-execution-sdk-python/src/aws_durable_execution_sdk_python/context.py

@@ -23,7 +23,10 @@
     ValidationError,
 )
 from aws_durable_execution_sdk_python.identifier import OperationIdentifier
-from aws_durable_execution_sdk_python.lambda_service import OperationSubType
+from aws_durable_execution_sdk_python.lambda_service import (
+    OperationSubType,
+    OperationType,
+)
 from aws_durable_execution_sdk_python.logger import Logger, LogInfo
 from aws_durable_execution_sdk_python.operation.callback import (
     CallbackOperationExecutor,
@@ -443,6 +446,7 @@ def create_callback(
             state=self.state,
             operation_identifier=OperationIdentifier(
                 operation_id=operation_id,
+                sub_type=OperationSubType.CALLBACK,
                 parent_id=self._parent_id,
                 name=name,
             ),
@@ -485,6 +489,7 @@ def invoke(
             state=self.state,
             operation_identifier=OperationIdentifier(
                 operation_id=operation_id,
+                sub_type=OperationSubType.CHAINED_INVOKE,
                 parent_id=self._parent_id,
                 name=name,
             ),
@@ -507,6 +512,7 @@ def map(
         operation_id = self._create_step_id()
         operation_identifier = OperationIdentifier(
             operation_id=operation_id,
+            sub_type=OperationSubType.MAP,
             parent_id=self._parent_id,
             name=map_name,
         )
@@ -553,7 +559,10 @@ def parallel(
         operation_id = self._create_step_id()
         parallel_context = self.create_child_context(operation_id=operation_id)
         operation_identifier = OperationIdentifier(
-            operation_id=operation_id, parent_id=self._parent_id, name=name
+            operation_id=operation_id,
+            sub_type=OperationSubType.PARALLEL,
+            parent_id=self._parent_id,
+            name=name,
         )
 
         def parallel_in_child_context() -> BatchResult[T]:
@@ -606,6 +615,11 @@ def run_in_child_context(
         step_name: str | None = self._resolve_step_name(name, func)
         # _create_step_id() is thread-safe. rest of method is safe, since using local copy of parent id
         operation_id = self._create_step_id()
+        sub_type = (
+            config.sub_type
+            if config and config.sub_type
+            else OperationSubType.RUN_IN_CHILD_CONTEXT
+        )
 
         is_virtual: bool = config.is_virtual if config else False
 
@@ -621,6 +635,7 @@ def callable_with_child_context():
             state=self.state,
             operation_identifier=OperationIdentifier(
                 operation_id=operation_id,
+                sub_type=sub_type,
                 parent_id=self._parent_id,
                 name=step_name,
             ),
@@ -646,6 +661,7 @@ def step(
             state=self.state,
             operation_identifier=OperationIdentifier(
                 operation_id=operation_id,
+                sub_type=OperationSubType.STEP,
                 parent_id=self._parent_id,
                 name=step_name,
             ),
@@ -673,6 +689,7 @@ def wait(self, duration: Duration, name: str | None = None) -> None:
             state=self.state,
             operation_identifier=OperationIdentifier(
                 operation_id=operation_id,
+                sub_type=OperationSubType.WAIT,
                 parent_id=self._parent_id,
                 name=name,
             ),
@@ -728,6 +745,7 @@ def wait_for_condition(
                 state=self.state,
                 operation_identifier=OperationIdentifier(
                     operation_id=operation_id,
+                    sub_type=OperationSubType.WAIT_FOR_CONDITION,
                     parent_id=self._parent_id,
                     name=name,
                 ),