Batch Python workflow payload encoding

Author: durable-workflow-ops

src/durable_workflow/_avro.py

@@ -30,6 +30,7 @@
 import base64
 import io
 import json
+from functools import lru_cache
 from typing import Any
 
 from .errors import AvroNotInstalledError
@@ -44,6 +45,7 @@
 _PREFIX_TYPED_SCHEMA = b"\x01"
 
 
+@lru_cache(maxsize=1)
 def _load_avro_schema() -> Any:
     try:
         import avro.schema

src/durable_workflow/serializer.py

@@ -20,9 +20,9 @@
 
 import json
 import logging
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, TypeGuard, cast
 
 from . import _avro
 
@@ -87,6 +87,8 @@ def to_log_context(self) -> dict[str, str]:
 
 
 DEFAULT_PAYLOAD_SIZE_WARNING = PayloadSizeWarningConfig()
+PayloadWarningContext = PayloadSizeWarningContext | Mapping[str, Any] | None
+PayloadWarningContexts = PayloadWarningContext | Sequence[PayloadWarningContext]
 
 
 def encode(
@@ -119,6 +121,32 @@ def encode(
     return blob
 
 
+def encode_many(
+    values: Sequence[Any],
+    codec: str = AVRO_CODEC,
+    *,
+    size_warning: PayloadSizeWarningConfig | None = DEFAULT_PAYLOAD_SIZE_WARNING,
+    warning_context: PayloadWarningContexts = None,
+) -> list[str]:
+    """Encode several payload blobs through one codec hook.
+
+    The default implementation intentionally preserves the single-value
+    encoder semantics and warning behavior. Codecs that can safely batch or
+    parallelize work can specialize behind this boundary without changing
+    call sites.
+    """
+    contexts = _warning_contexts_for_values(values, warning_context)
+    return [
+        encode(
+            value,
+            codec=codec,
+            size_warning=size_warning,
+            warning_context=contexts[index],
+        )
+        for index, value in enumerate(values)
+    ]
+
+
 def envelope(
     value: Any,
     codec: str = AVRO_CODEC,
@@ -138,6 +166,25 @@ def envelope(
     }
 
 
+def envelope_many(
+    values: Sequence[Any],
+    codec: str = AVRO_CODEC,
+    *,
+    size_warning: PayloadSizeWarningConfig | None = DEFAULT_PAYLOAD_SIZE_WARNING,
+    warning_context: PayloadWarningContexts = None,
+) -> list[dict[str, str]]:
+    """Wrap several values in ``{codec, blob}`` payload envelopes."""
+    return [
+        {"codec": codec, "blob": blob}
+        for blob in encode_many(
+            values,
+            codec=codec,
+            size_warning=size_warning,
+            warning_context=warning_context,
+        )
+    ]
+
+
 def warn_if_json_payload_near_limit(
     value: Any,
     *,
@@ -205,6 +252,27 @@ def _normalize_warning_context(
     return normalized
 
 
+def _warning_contexts_for_values(
+    values: Sequence[Any],
+    context: PayloadWarningContexts,
+) -> list[PayloadWarningContext]:
+    if not _is_context_sequence(context):
+        single_context = cast(PayloadWarningContext, context)
+        return [single_context] * len(values)
+    if len(context) != len(values):
+        raise ValueError("payload warning context count must match value count")
+    return list(context)
+
+
+def _is_context_sequence(
+    context: PayloadWarningContexts,
+) -> TypeGuard[Sequence[PayloadWarningContext]]:
+    return isinstance(context, Sequence) and not isinstance(
+        context,
+        (str, bytes, bytearray, PayloadSizeWarningContext, Mapping),
+    )
+
+
 def decode_envelope(value: Any, codec: str | None = None) -> Any:
     """Decode a value that may be a ``{codec, blob}`` envelope or a raw blob.
 

src/durable_workflow/worker.py

@@ -52,7 +52,7 @@
     WORKER_TASKS,
     MetricsRecorder,
 )
-from .workflow import apply_update, query_state, replay
+from .workflow import apply_update, commands_to_server_commands, query_state, replay
 
 log = logging.getLogger("durable_workflow.worker")
 
@@ -515,18 +515,16 @@ async def _run_workflow_task_core(self, task: dict[str, Any]) -> list[dict[str,
                 log.warning("failed to report replay failure: %s", fe)
             return None
 
-        commands = [
-            c.to_server_command(
-                self.task_queue,
-                payload_codec=command_codec,
-                size_warning=self._payload_size_warning_config(),
-                warning_context=self._workflow_payload_warning_context(
-                    task,
-                    kind="workflow_command",
-                ),
-            )
-            for c in outcome.commands
-        ]
+        commands = commands_to_server_commands(
+            outcome.commands,
+            self.task_queue,
+            payload_codec=command_codec,
+            size_warning=self._payload_size_warning_config(),
+            warning_context=self._workflow_payload_warning_context(
+                task,
+                kind="workflow_command",
+            ),
+        )
         log.info(
             "completing workflow task %s with %d command(s): %s",
             task_id,

src/durable_workflow/workflow.py

@@ -24,7 +24,7 @@
 import math
 import random
 import uuid
-from collections.abc import Callable, Iterable, Mapping
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from typing import Any
@@ -685,6 +685,177 @@ def to_server_command(
 )
 
 
+def commands_to_server_commands(
+    commands: Sequence[Command],
+    task_queue: str,
+    *,
+    payload_codec: str = serializer.AVRO_CODEC,
+    size_warning: serializer.PayloadSizeWarningConfig | None = serializer.DEFAULT_PAYLOAD_SIZE_WARNING,
+    warning_context: PayloadWarningContext = None,
+) -> list[dict[str, Any]]:
+    """Convert workflow commands to the server wire shape with batched payload encoding."""
+    server_commands: list[dict[str, Any]] = []
+    envelope_jobs: list[tuple[int, str, Any, dict[str, str]]] = []
+    encode_jobs: list[tuple[int, str, Any, dict[str, str]]] = []
+
+    for command in commands:
+        if isinstance(command, ScheduleActivity):
+            queue = command.queue or task_queue
+            server_command: dict[str, Any] = {
+                "type": "schedule_activity",
+                "activity_type": command.activity_type,
+                "queue": queue,
+            }
+            envelope_jobs.append((
+                len(server_commands),
+                "arguments",
+                command.arguments,
+                _payload_warning_context(
+                    warning_context,
+                    kind="activity_input",
+                    task_queue=queue,
+                    activity_name=command.activity_type,
+                ),
+            ))
+            if command.retry_policy is not None:
+                server_command["retry_policy"] = (
+                    command.retry_policy.to_dict()
+                    if isinstance(command.retry_policy, ActivityRetryPolicy)
+                    else dict(command.retry_policy)
+                )
+            if command.start_to_close_timeout is not None:
+                server_command["start_to_close_timeout"] = command.start_to_close_timeout
+            if command.schedule_to_start_timeout is not None:
+                server_command["schedule_to_start_timeout"] = command.schedule_to_start_timeout
+            if command.schedule_to_close_timeout is not None:
+                server_command["schedule_to_close_timeout"] = command.schedule_to_close_timeout
+            if command.heartbeat_timeout is not None:
+                server_command["heartbeat_timeout"] = command.heartbeat_timeout
+            server_commands.append(server_command)
+            continue
+
+        if isinstance(command, CompleteWorkflow):
+            server_commands.append({"type": "complete_workflow"})
+            envelope_jobs.append((
+                len(server_commands) - 1,
+                "result",
+                command.result,
+                _payload_warning_context(
+                    warning_context,
+                    kind="workflow_result",
+                    task_queue=task_queue,
+                ),
+            ))
+            continue
+
+        if isinstance(command, CompleteUpdate):
+            server_commands.append({"type": "complete_update", "update_id": command.update_id})
+            envelope_jobs.append((
+                len(server_commands) - 1,
+                "result",
+                command.result,
+                _payload_warning_context(
+                    warning_context,
+                    kind="update_result",
+                    task_queue=task_queue,
+                ),
+            ))
+            continue
+
+        if isinstance(command, ContinueAsNew):
+            queue = command.task_queue or task_queue
+            server_command = {"type": "continue_as_new", "queue": queue}
+            if command.workflow_type is not None:
+                server_command["workflow_type"] = command.workflow_type
+            server_commands.append(server_command)
+            envelope_jobs.append((
+                len(server_commands) - 1,
+                "arguments",
+                command.arguments,
+                _payload_warning_context(
+                    warning_context,
+                    kind="continue_as_new_input",
+                    task_queue=queue,
+                ),
+            ))
+            continue
+
+        if isinstance(command, RecordSideEffect):
+            server_commands.append({"type": "record_side_effect"})
+            encode_jobs.append((
+                len(server_commands) - 1,
+                "result",
+                command.result,
+                _payload_warning_context(
+                    warning_context,
+                    kind="side_effect_result",
+                    task_queue=task_queue,
+                ),
+            ))
+            continue
+
+        if isinstance(command, StartChildWorkflow):
+            queue = command.task_queue or task_queue
+            server_command = {
+                "type": "start_child_workflow",
+                "workflow_type": command.workflow_type,
+                "queue": queue,
+            }
+            envelope_jobs.append((
+                len(server_commands),
+                "arguments",
+                command.arguments,
+                _payload_warning_context(
+                    warning_context,
+                    kind="child_workflow_input",
+                    task_queue=queue,
+                ),
+            ))
+            if command.parent_close_policy is not None:
+                server_command["parent_close_policy"] = command.parent_close_policy
+            if command.retry_policy is not None:
+                server_command["retry_policy"] = (
+                    command.retry_policy.to_dict()
+                    if isinstance(command.retry_policy, ActivityRetryPolicy)
+                    else dict(command.retry_policy)
+                )
+            if command.execution_timeout_seconds is not None:
+                server_command["execution_timeout_seconds"] = command.execution_timeout_seconds
+            if command.run_timeout_seconds is not None:
+                server_command["run_timeout_seconds"] = command.run_timeout_seconds
+            server_commands.append(server_command)
+            continue
+
+        server_commands.append(command.to_server_command(
+            task_queue,
+            payload_codec=payload_codec,
+            size_warning=size_warning,
+            warning_context=warning_context,
+        ))
+
+    if envelope_jobs:
+        envelopes = serializer.envelope_many(
+            [value for _, _, value, _ in envelope_jobs],
+            codec=payload_codec,
+            size_warning=size_warning,
+            warning_context=[context for _, _, _, context in envelope_jobs],
+        )
+        for (index, key, _, _), envelope_value in zip(envelope_jobs, envelopes, strict=True):
+            server_commands[index][key] = envelope_value
+
+    if encode_jobs:
+        blobs = serializer.encode_many(
+            [value for _, _, value, _ in encode_jobs],
+            codec=payload_codec,
+            size_warning=size_warning,
+            warning_context=[context for _, _, _, context in encode_jobs],
+        )
+        for (index, key, _, _), blob in zip(encode_jobs, blobs, strict=True):
+            server_commands[index][key] = blob
+
+    return server_commands
+
+
 # ── Context passed to the workflow's run() ───────────────────────────
 
 _REPLAY_LOGGER = logging.getLogger("durable_workflow.workflow.replay")

tests/test_serializer.py

@@ -123,6 +123,49 @@ def test_none_value(self) -> None:
         assert serializer.decode(env["blob"], codec="avro") is None
 
 
+class TestBatchEncoding:
+    def test_encode_many_preserves_order(self) -> None:
+        blobs = serializer.encode_many([["a"], ["b"]], codec="json")
+        assert blobs == ['["a"]', '["b"]']
+
+    def test_envelope_many_wraps_each_value(self) -> None:
+        envelopes = serializer.envelope_many([["a"], ["b"]], codec="json")
+        assert envelopes == [
+            {"codec": "json", "blob": '["a"]'},
+            {"codec": "json", "blob": '["b"]'},
+        ]
+
+    def test_encode_many_accepts_per_payload_warning_context(
+        self, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        config = serializer.PayloadSizeWarningConfig(limit_bytes=10, threshold_percent=50)
+        contexts = [
+            serializer.PayloadSizeWarningContext(kind="signal", signal_name="one"),
+            serializer.PayloadSizeWarningContext(kind="signal", signal_name="two"),
+        ]
+
+        with caplog.at_level(logging.WARNING, logger="durable_workflow.serializer"):
+            serializer.encode_many(
+                ["abcdef", "ghijkl"],
+                codec="json",
+                size_warning=config,
+                warning_context=contexts,
+            )
+
+        assert [record.durable_workflow_payload["signal_name"] for record in caplog.records] == [
+            "one",
+            "two",
+        ]
+
+    def test_encode_many_rejects_context_count_mismatch(self) -> None:
+        with pytest.raises(ValueError, match="context count"):
+            serializer.encode_many(
+                ["a", "b"],
+                codec="json",
+                warning_context=[serializer.PayloadSizeWarningContext(kind="payload")],
+            )
+
+
 class TestPayloadSizeWarning:
     def test_encode_warns_with_structured_context(self, caplog: pytest.LogCaptureFixture) -> None:
         config = serializer.PayloadSizeWarningConfig(limit_bytes=10, threshold_percent=50)