From 0cba76fef981b7af440c09c86185ceb9e8ed2503 Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:19:35 -0700 Subject: [PATCH] refactor: move workflow task duration logging to core --- temporalio/bridge/proto/common/__init__.py | 2 + temporalio/bridge/proto/common/common_pb2.py | 20 +- temporalio/bridge/proto/common/common_pb2.pyi | 53 ++++++ .../workflow_completion_pb2.py | 12 +- .../workflow_completion_pb2.pyi | 36 +++- temporalio/bridge/sdk-core | 2 +- temporalio/worker/_workflow.py | 105 +++-------- tests/worker/test_extstore.py | 174 ++++++++---------- 8 files changed, 218 insertions(+), 186 deletions(-) diff --git a/temporalio/bridge/proto/common/__init__.py b/temporalio/bridge/proto/common/__init__.py index 5622fffb8..a8506090d 100644 --- a/temporalio/bridge/proto/common/__init__.py +++ b/temporalio/bridge/proto/common/__init__.py @@ -1,10 +1,12 @@ from .common_pb2 import ( + ExternalStorageMetrics, NamespacedWorkflowExecution, VersioningIntent, WorkerDeploymentVersion, ) __all__ = [ + "ExternalStorageMetrics", "NamespacedWorkflowExecution", "VersioningIntent", "WorkerDeploymentVersion", diff --git a/temporalio/bridge/proto/common/common_pb2.py b/temporalio/bridge/proto/common/common_pb2.py index c56456fce..481cf216d 100644 --- a/temporalio/bridge/proto/common/common_pb2.py +++ b/temporalio/bridge/proto/common/common_pb2.py @@ -18,7 +18,7 @@ from google.protobuf import duration_pb2 as google_dot_protobuf_dot_duration__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n%temporal/sdk/core/common/common.proto\x12\x0e\x63oresdk.common\x1a\x1egoogle/protobuf/duration.proto"U\n\x1bNamespacedWorkflowExecution\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x13\n\x0bworkflow_id\x18\x02 \x01(\t\x12\x0e\n\x06run_id\x18\x03 \x01(\t"D\n\x17WorkerDeploymentVersion\x12\x17\n\x0f\x64\x65ployment_name\x18\x01 \x01(\t\x12\x10\n\x08\x62uild_id\x18\x02 \x01(\t*@\n\x10VersioningIntent\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x0e\n\nCOMPATIBLE\x10\x01\x12\x0b\n\x07\x44\x45\x46\x41ULT\x10\x02\x42,\xea\x02)Temporalio::Internal::Bridge::Api::Commonb\x06proto3' + b'\n%temporal/sdk/core/common/common.proto\x12\x0e\x63oresdk.common\x1a\x1egoogle/protobuf/duration.proto"U\n\x1bNamespacedWorkflowExecution\x12\x11\n\tnamespace\x18\x01 \x01(\t\x12\x13\n\x0bworkflow_id\x18\x02 \x01(\t\x12\x0e\n\x06run_id\x18\x03 \x01(\t"D\n\x17WorkerDeploymentVersion\x12\x17\n\x0f\x64\x65ployment_name\x18\x01 \x01(\t\x12\x10\n\x08\x62uild_id\x18\x02 \x01(\t"\x92\x01\n\x16\x45xternalStorageMetrics\x12\x15\n\rpayload_count\x18\x01 \x01(\x04\x12\x18\n\x10total_size_bytes\x18\x02 \x01(\x04\x12\x31\n\x0etotal_duration\x18\x03 \x01(\x0b\x32\x19.google.protobuf.Duration\x12\x14\n\x0c\x64river_names\x18\x04 \x03(\t*@\n\x10VersioningIntent\x12\x0f\n\x0bUNSPECIFIED\x10\x00\x12\x0e\n\nCOMPATIBLE\x10\x01\x12\x0b\n\x07\x44\x45\x46\x41ULT\x10\x02\x42,\xea\x02)Temporalio::Internal::Bridge::Api::Commonb\x06proto3' ) _VERSIONINGINTENT = DESCRIPTOR.enum_types_by_name["VersioningIntent"] @@ -32,6 +32,7 @@ "NamespacedWorkflowExecution" ] _WORKERDEPLOYMENTVERSION = DESCRIPTOR.message_types_by_name["WorkerDeploymentVersion"] +_EXTERNALSTORAGEMETRICS = DESCRIPTOR.message_types_by_name["ExternalStorageMetrics"] NamespacedWorkflowExecution = _reflection.GeneratedProtocolMessageType( "NamespacedWorkflowExecution", (_message.Message,), @@ -54,15 +55,28 @@ ) _sym_db.RegisterMessage(WorkerDeploymentVersion) +ExternalStorageMetrics = _reflection.GeneratedProtocolMessageType( + "ExternalStorageMetrics", + (_message.Message,), + { + "DESCRIPTOR": _EXTERNALSTORAGEMETRICS, + "__module__": "temporal.sdk.core.common.common_pb2", + # @@protoc_insertion_point(class_scope:coresdk.common.ExternalStorageMetrics) + }, +) +_sym_db.RegisterMessage(ExternalStorageMetrics) + if _descriptor._USE_C_DESCRIPTORS == False: DESCRIPTOR._options = None DESCRIPTOR._serialized_options = ( b"\352\002)Temporalio::Internal::Bridge::Api::Common" ) - _VERSIONINGINTENT._serialized_start = 246 - _VERSIONINGINTENT._serialized_end = 310 + _VERSIONINGINTENT._serialized_start = 395 + _VERSIONINGINTENT._serialized_end = 459 _NAMESPACEDWORKFLOWEXECUTION._serialized_start = 89 _NAMESPACEDWORKFLOWEXECUTION._serialized_end = 174 _WORKERDEPLOYMENTVERSION._serialized_start = 176 _WORKERDEPLOYMENTVERSION._serialized_end = 244 + _EXTERNALSTORAGEMETRICS._serialized_start = 247 + _EXTERNALSTORAGEMETRICS._serialized_end = 393 # @@protoc_insertion_point(module_scope) diff --git a/temporalio/bridge/proto/common/common_pb2.pyi b/temporalio/bridge/proto/common/common_pb2.pyi index 739a129e1..8862fa036 100644 --- a/temporalio/bridge/proto/common/common_pb2.pyi +++ b/temporalio/bridge/proto/common/common_pb2.pyi @@ -4,10 +4,13 @@ isort:skip_file """ import builtins +import collections.abc import sys import typing import google.protobuf.descriptor +import google.protobuf.duration_pb2 +import google.protobuf.internal.containers import google.protobuf.internal.enum_type_wrapper import google.protobuf.message @@ -121,3 +124,53 @@ class WorkerDeploymentVersion(google.protobuf.message.Message): ) -> None: ... global___WorkerDeploymentVersion = WorkerDeploymentVersion + +class ExternalStorageMetrics(google.protobuf.message.Message): + """Metrics for a set of external payload storage operations (all uploads and downloads) + performed while processing a task, so core can emit unified logging and metrics. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PAYLOAD_COUNT_FIELD_NUMBER: builtins.int + TOTAL_SIZE_BYTES_FIELD_NUMBER: builtins.int + TOTAL_DURATION_FIELD_NUMBER: builtins.int + DRIVER_NAMES_FIELD_NUMBER: builtins.int + payload_count: builtins.int + """Number of payloads stored or retrieved externally.""" + total_size_bytes: builtins.int + """Total size in bytes of the externally stored or retrieved payloads.""" + @property + def total_duration(self) -> google.protobuf.duration_pb2.Duration: + """Wall-clock time spent on the external storage operations.""" + @property + def driver_names( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """Names of the drivers that participated in the operations.""" + def __init__( + self, + *, + payload_count: builtins.int = ..., + total_size_bytes: builtins.int = ..., + total_duration: google.protobuf.duration_pb2.Duration | None = ..., + driver_names: collections.abc.Iterable[builtins.str] | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["total_duration", b"total_duration"] + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "driver_names", + b"driver_names", + "payload_count", + b"payload_count", + "total_duration", + b"total_duration", + "total_size_bytes", + b"total_size_bytes", + ], + ) -> None: ... + +global___ExternalStorageMetrics = ExternalStorageMetrics diff --git a/temporalio/bridge/proto/workflow_completion/workflow_completion_pb2.py b/temporalio/bridge/proto/workflow_completion/workflow_completion_pb2.py index ce26b220d..057b301e4 100644 --- a/temporalio/bridge/proto/workflow_completion/workflow_completion_pb2.py +++ b/temporalio/bridge/proto/workflow_completion/workflow_completion_pb2.py @@ -31,7 +31,7 @@ ) DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n?temporal/sdk/core/workflow_completion/workflow_completion.proto\x12\x1b\x63oresdk.workflow_completion\x1a%temporal/api/failure/v1/message.proto\x1a(temporal/api/enums/v1/failed_cause.proto\x1a$temporal/api/enums/v1/workflow.proto\x1a%temporal/sdk/core/common/common.proto\x1a;temporal/sdk/core/workflow_commands/workflow_commands.proto"\xac\x01\n\x1cWorkflowActivationCompletion\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12:\n\nsuccessful\x18\x02 \x01(\x0b\x32$.coresdk.workflow_completion.SuccessH\x00\x12\x36\n\x06\x66\x61iled\x18\x03 \x01(\x0b\x32$.coresdk.workflow_completion.FailureH\x00\x42\x08\n\x06status"\xac\x01\n\x07Success\x12<\n\x08\x63ommands\x18\x01 \x03(\x0b\x32*.coresdk.workflow_commands.WorkflowCommand\x12\x1b\n\x13used_internal_flags\x18\x06 \x03(\r\x12\x46\n\x13versioning_behavior\x18\x07 \x01(\x0e\x32).temporal.api.enums.v1.VersioningBehavior"\x81\x01\n\x07\x46\x61ilure\x12\x31\n\x07\x66\x61ilure\x18\x01 \x01(\x0b\x32 .temporal.api.failure.v1.Failure\x12\x43\n\x0b\x66orce_cause\x18\x02 \x01(\x0e\x32..temporal.api.enums.v1.WorkflowTaskFailedCauseB8\xea\x02\x35Temporalio::Internal::Bridge::Api::WorkflowCompletionb\x06proto3' + b'\n?temporal/sdk/core/workflow_completion/workflow_completion.proto\x12\x1b\x63oresdk.workflow_completion\x1a%temporal/api/failure/v1/message.proto\x1a(temporal/api/enums/v1/failed_cause.proto\x1a$temporal/api/enums/v1/workflow.proto\x1a%temporal/sdk/core/common/common.proto\x1a;temporal/sdk/core/workflow_commands/workflow_commands.proto"\xbe\x02\n\x1cWorkflowActivationCompletion\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12:\n\nsuccessful\x18\x02 \x01(\x0b\x32$.coresdk.workflow_completion.SuccessH\x00\x12\x36\n\x06\x66\x61iled\x18\x03 \x01(\x0b\x32$.coresdk.workflow_completion.FailureH\x00\x12H\n\x18payload_download_metrics\x18\x04 \x01(\x0b\x32&.coresdk.common.ExternalStorageMetrics\x12\x46\n\x16payload_upload_metrics\x18\x05 \x01(\x0b\x32&.coresdk.common.ExternalStorageMetricsB\x08\n\x06status"\xac\x01\n\x07Success\x12<\n\x08\x63ommands\x18\x01 \x03(\x0b\x32*.coresdk.workflow_commands.WorkflowCommand\x12\x1b\n\x13used_internal_flags\x18\x06 \x03(\r\x12\x46\n\x13versioning_behavior\x18\x07 \x01(\x0e\x32).temporal.api.enums.v1.VersioningBehavior"\x81\x01\n\x07\x46\x61ilure\x12\x31\n\x07\x66\x61ilure\x18\x01 \x01(\x0b\x32 .temporal.api.failure.v1.Failure\x12\x43\n\x0b\x66orce_cause\x18\x02 \x01(\x0e\x32..temporal.api.enums.v1.WorkflowTaskFailedCauseB8\xea\x02\x35Temporalio::Internal::Bridge::Api::WorkflowCompletionb\x06proto3' ) @@ -79,9 +79,9 @@ b"\352\0025Temporalio::Internal::Bridge::Api::WorkflowCompletion" ) _WORKFLOWACTIVATIONCOMPLETION._serialized_start = 316 - _WORKFLOWACTIVATIONCOMPLETION._serialized_end = 488 - _SUCCESS._serialized_start = 491 - _SUCCESS._serialized_end = 663 - _FAILURE._serialized_start = 666 - _FAILURE._serialized_end = 795 + _WORKFLOWACTIVATIONCOMPLETION._serialized_end = 634 + _SUCCESS._serialized_start = 637 + _SUCCESS._serialized_end = 809 + _FAILURE._serialized_start = 812 + _FAILURE._serialized_end = 941 # @@protoc_insertion_point(module_scope) diff --git a/temporalio/bridge/proto/workflow_completion/workflow_completion_pb2.pyi b/temporalio/bridge/proto/workflow_completion/workflow_completion_pb2.pyi index 5b438f360..8e12736aa 100644 --- a/temporalio/bridge/proto/workflow_completion/workflow_completion_pb2.pyi +++ b/temporalio/bridge/proto/workflow_completion/workflow_completion_pb2.pyi @@ -14,6 +14,7 @@ import google.protobuf.message import temporalio.api.enums.v1.failed_cause_pb2 import temporalio.api.enums.v1.workflow_pb2 import temporalio.api.failure.v1.message_pb2 +import temporalio.bridge.proto.common.common_pb2 import temporalio.bridge.proto.workflow_commands.workflow_commands_pb2 if sys.version_info >= (3, 8): @@ -31,23 +32,52 @@ class WorkflowActivationCompletion(google.protobuf.message.Message): RUN_ID_FIELD_NUMBER: builtins.int SUCCESSFUL_FIELD_NUMBER: builtins.int FAILED_FIELD_NUMBER: builtins.int + PAYLOAD_DOWNLOAD_METRICS_FIELD_NUMBER: builtins.int + PAYLOAD_UPLOAD_METRICS_FIELD_NUMBER: builtins.int run_id: builtins.str """The run id from the workflow activation you are completing""" @property def successful(self) -> global___Success: ... @property def failed(self) -> global___Failure: ... + @property + def payload_download_metrics( + self, + ) -> temporalio.bridge.proto.common.common_pb2.ExternalStorageMetrics: + """Metrics for external payload storage downloads (retrievals) performed while processing + this activation. Only set when external storage retrieved payloads. + """ + @property + def payload_upload_metrics( + self, + ) -> temporalio.bridge.proto.common.common_pb2.ExternalStorageMetrics: + """Metrics for external payload storage uploads (stores) performed while processing this + activation. Only set when external storage stored payloads. + """ def __init__( self, *, run_id: builtins.str = ..., successful: global___Success | None = ..., failed: global___Failure | None = ..., + payload_download_metrics: temporalio.bridge.proto.common.common_pb2.ExternalStorageMetrics + | None = ..., + payload_upload_metrics: temporalio.bridge.proto.common.common_pb2.ExternalStorageMetrics + | None = ..., ) -> None: ... def HasField( self, field_name: typing_extensions.Literal[ - "failed", b"failed", "status", b"status", "successful", b"successful" + "failed", + b"failed", + "payload_download_metrics", + b"payload_download_metrics", + "payload_upload_metrics", + b"payload_upload_metrics", + "status", + b"status", + "successful", + b"successful", ], ) -> builtins.bool: ... def ClearField( @@ -55,6 +85,10 @@ class WorkflowActivationCompletion(google.protobuf.message.Message): field_name: typing_extensions.Literal[ "failed", b"failed", + "payload_download_metrics", + b"payload_download_metrics", + "payload_upload_metrics", + b"payload_upload_metrics", "run_id", b"run_id", "status", diff --git a/temporalio/bridge/sdk-core b/temporalio/bridge/sdk-core index d2769368d..ce69d10f0 160000 --- a/temporalio/bridge/sdk-core +++ b/temporalio/bridge/sdk-core @@ -1 +1 @@ -Subproject commit d2769368df9077a311537431ff4594c9c14db4e7 +Subproject commit ce69d10f0e80ec154264c3a7ed395af1e18aa796 diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index c031b5653..1b217b4a5 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -9,13 +9,13 @@ import os import sys import threading -import time from collections.abc import Awaitable, Callable, MutableMapping, Sequence from dataclasses import dataclass -from datetime import timedelta, timezone +from datetime import timezone from types import TracebackType import temporalio.api.common.v1 +import temporalio.bridge.proto.common import temporalio.bridge.proto.workflow_activation import temporalio.bridge.proto.workflow_completion import temporalio.bridge.runtime @@ -64,6 +64,17 @@ _DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY: int = 3 +def _set_external_storage_metrics( + target: temporalio.bridge.proto.common.ExternalStorageMetrics, + metrics: temporalio.converter._extstore.StorageOperationMetrics, +) -> None: + """Populate a proto ``ExternalStorageMetrics`` from measured storage metrics.""" + target.payload_count = metrics.payload_count + target.total_size_bytes = metrics.total_size + target.total_duration.FromTimedelta(metrics.total_duration) + target.driver_names.extend(sorted(metrics.driver_names)) + + class _WorkflowWorker: # type:ignore[reportUnusedClass] def __init__( self, @@ -325,7 +336,6 @@ async def _handle_activation( completion.successful.SetInParent() workflow = None data_converter = self._data_converter - task_start_time = time.monotonic() download_metrics = temporalio.converter._extstore.StorageOperationMetrics() try: if LOG_PROTOS: @@ -500,6 +510,17 @@ async def _handle_activation( completion.failed.Clear() completion.failed.failure.message = f"Failed encoding completion: {err}" + # Reported on the completion so core can include them in its workflow-task duration + # log; core measures the duration itself. + if download_metrics.payload_count > 0: + _set_external_storage_metrics( + completion.payload_download_metrics, download_metrics + ) + if upload_metrics.payload_count > 0: + _set_external_storage_metrics( + completion.payload_upload_metrics, upload_metrics + ) + # Send off completion if LOG_PROTOS: logger.debug("Sending workflow completion:\n%s", completion) @@ -511,84 +532,6 @@ async def _handle_activation( "Failed completing activation on workflow with run ID %s", act.run_id ) - # Log workflow task duration with external storage metrics - self._log_workflow_task_duration( - act, workflow, task_start_time, download_metrics, upload_metrics - ) - - def _log_workflow_task_duration( - self, - act: temporalio.bridge.proto.workflow_activation.WorkflowActivation, - workflow: _RunningWorkflow | None, - task_start_time: float, - download_metrics: temporalio.converter._extstore.StorageOperationMetrics, - upload_metrics: temporalio.converter._extstore.StorageOperationMetrics, - ) -> None: - task_duration = timedelta(seconds=time.monotonic() - task_start_time) - - def _fmt_duration(td: timedelta) -> str: - secs = td.total_seconds() - if secs >= 1: - return f"{secs:.3f}s" - return f"{secs * 1000:.3f}ms" - - completed_event_id = act.history_length + 1 - _info = workflow.get_info() if workflow is not None else None - attempt = _info.attempt if _info is not None else "unknown" - log_id = f"{act.run_id}:{completed_event_id}:{attempt}" - msg_details, extra = temporalio.workflow._build_log_context( - _info._logger_details() if _info is not None else None, - full_workflow_info=_info, - ) - msg_details["event_id"] = completed_event_id - msg_details["workflow_task_duration"] = _fmt_duration(task_duration) - msg_details["workflow_history_size"] = act.history_size_bytes - extra["event_id"] = completed_event_id - extra["workflow_task_duration"] = task_duration - extra["workflow_history_size"] = act.history_size_bytes - if download_metrics.payload_count > 0: - msg_details["payload_download_count"] = download_metrics.payload_count - msg_details["payload_download_size"] = download_metrics.total_size - msg_details["payload_download_duration"] = _fmt_duration( - download_metrics.total_duration - ) - msg_details["payload_download_drivers"] = sorted( - download_metrics.driver_names - ) - extra["payload_download_count"] = download_metrics.payload_count - extra["payload_download_size"] = download_metrics.total_size - extra["payload_download_duration"] = download_metrics.total_duration - extra["payload_download_drivers"] = sorted(download_metrics.driver_names) - if upload_metrics.payload_count > 0: - msg_details["payload_upload_count"] = upload_metrics.payload_count - msg_details["payload_upload_size"] = upload_metrics.total_size - msg_details["payload_upload_duration"] = _fmt_duration( - upload_metrics.total_duration - ) - msg_details["payload_upload_drivers"] = sorted(upload_metrics.driver_names) - extra["payload_upload_count"] = upload_metrics.payload_count - extra["payload_upload_size"] = upload_metrics.total_size - extra["payload_upload_duration"] = upload_metrics.total_duration - extra["payload_upload_drivers"] = sorted(upload_metrics.driver_names) - if task_duration.total_seconds() > 10: - logger.warning( - f"[TMPRL1104] {log_id} Workflow task exceeded 10 seconds (%s)", - msg_details, - extra=extra, - ) - elif task_duration.total_seconds() > 5: - logger.info( - f"[TMPRL1104] {log_id} Workflow task exceeded 5 seconds (%s)", - msg_details, - extra=extra, - ) - else: - logger.debug( - f"[TMPRL1104] {log_id} Workflow task duration information (%s)", - msg_details, - extra=extra, - ) - async def _handle_cache_eviction( self, act: temporalio.bridge.proto.workflow_activation.WorkflowActivation, diff --git a/tests/worker/test_extstore.py b/tests/worker/test_extstore.py index 998a2bd27..50041b749 100644 --- a/tests/worker/test_extstore.py +++ b/tests/worker/test_extstore.py @@ -1,8 +1,7 @@ +import contextlib import dataclasses -import logging -import re import uuid -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from dataclasses import dataclass from datetime import timedelta from unittest import mock @@ -11,10 +10,10 @@ import temporalio import temporalio.bridge.client +import temporalio.bridge.proto.workflow_completion import temporalio.bridge.worker import temporalio.client import temporalio.converter -import temporalio.worker._workflow from temporalio import activity, workflow from temporalio.api.common.v1 import Payload from temporalio.client import Client, WorkflowFailureError, WorkflowHandle @@ -31,7 +30,7 @@ from temporalio.exceptions import ActivityError, ApplicationError from temporalio.testing._workflow import WorkflowEnvironment from temporalio.worker import Replayer -from tests.helpers import LogCapturer, assert_task_fail_eventually, new_worker +from tests.helpers import assert_task_fail_eventually, new_worker from tests.test_extstore import InMemoryTestDriver @@ -616,19 +615,32 @@ async def test_worker_storage_drivers_empty_without_external_storage( # TMPRL1104 workflow task duration logging # --------------------------------------------------------------------------- -_workflow_logger = logging.getLogger(temporalio.worker._workflow.__name__) +# The duration log itself is emitted (and tested) in sdk-core. The Python worker's part is +# attaching the external-storage metrics to the completion, so these tests capture the +# completion and assert on its fields directly rather than on core's asynchronously +# forwarded log, which would be nondeterministic to observe here. -def _tmprl1104_records(capturer: LogCapturer) -> list[logging.LogRecord]: - """Return all TMPRL1104 log records from the capturer.""" - return capturer.find_all(lambda r: r.getMessage().startswith("[TMPRL1104]")) +@contextlib.contextmanager +def _capture_completions() -> Iterator[ + list[temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion] +]: + """Capture every WorkflowActivationCompletion the worker hands to core.""" + completions: list[ + temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion + ] = [] + original = temporalio.bridge.worker.Worker.complete_workflow_activation + async def capturing(self, completion): # type: ignore[no-untyped-def] + completions.append(completion) + return await original(self, completion) -# Accept any duration-bucket wording: a loaded host can push a trivial task past 5s. -_TMPRL1104_DURATION_MESSAGE = re.compile( - r"\[TMPRL1104\] [^:]+:\d+:\d+ Workflow task " - r"(?:duration information|exceeded \d+ seconds) \(" -) + with mock.patch.object( + temporalio.bridge.worker.Worker, + "complete_workflow_activation", + capturing, + ): + yield completions async def _expected_payload_size( @@ -639,44 +651,33 @@ async def _expected_payload_size( return payloads[0].ByteSize() -@workflow.defn -class SimpleWorkflow: - """Minimal workflow for testing logging without external storage.""" - - @workflow.run - async def run(self) -> str: - return "done" - - async def test_tmprl1104_no_extstore(env: WorkflowEnvironment) -> None: - """Without external storage, TMPRL1104 logs contain duration but no - download/upload metrics.""" - with LogCapturer().logs_captured(_workflow_logger, level=logging.DEBUG) as capturer: - async with new_worker(env.client, SimpleWorkflow) as worker: + """Without external storage configured, completions carry no storage metrics.""" + with _capture_completions() as completions: + async with new_worker( + env.client, ExtStoreWorkflow, activities=[ext_store_activity] + ) as worker: await env.client.execute_workflow( - SimpleWorkflow.run, + ExtStoreWorkflow.run, + ExtStoreWorkflowInput( + input_data="small", + activity_input_size=10, + activity_output_size=10, + output_size=10, + ), id=f"workflow-{uuid.uuid4()}", task_queue=worker.task_queue, ) - records = _tmprl1104_records(capturer) - assert len(records) == 1 - record = records[0] - assert _TMPRL1104_DURATION_MESSAGE.match(record.getMessage()) - assert hasattr(record, "workflow_task_duration") - assert hasattr(record, "event_id") - # No external storage — download/upload fields must be absent - assert not hasattr(record, "payload_download_count") - assert not hasattr(record, "payload_download_size") - assert not hasattr(record, "payload_download_duration") - assert not hasattr(record, "payload_upload_count") - assert not hasattr(record, "payload_upload_size") - assert not hasattr(record, "payload_upload_duration") + assert completions, "expected the worker to complete at least one activation" + for c in completions: + assert not c.HasField("payload_download_metrics") + assert not c.HasField("payload_upload_metrics") async def test_tmprl1104_with_extstore_download(env: WorkflowEnvironment) -> None: - """When external storage decodes payloads, TMPRL1104 logs include download - metrics on the activation that retrieves them.""" + """When external storage retrieves payloads, the completion for the WFT that + retrieved them carries download metrics.""" driver = InMemoryTestDriver() data_converter = dataclasses.replace( temporalio.converter.default(), @@ -699,7 +700,7 @@ async def test_tmprl1104_with_extstore_download(env: WorkflowEnvironment) -> Non ) expected_input_size = await _expected_payload_size(data_converter, wf_input) - with LogCapturer().logs_captured(_workflow_logger, level=logging.DEBUG) as capturer: + with _capture_completions() as completions: async with new_worker( client, ExtStoreWorkflow, activities=[ext_store_activity] ) as worker: @@ -710,25 +711,19 @@ async def test_tmprl1104_with_extstore_download(env: WorkflowEnvironment) -> Non task_queue=worker.task_queue, ) - records = _tmprl1104_records(capturer) - assert len(records) == 2 - - # WFT 1: retrieves the externalized workflow input - assert _TMPRL1104_DURATION_MESSAGE.match(records[0].getMessage()) - assert getattr(records[0], "payload_download_count") == 1 - assert getattr(records[0], "payload_download_size") == expected_input_size - assert getattr(records[0], "payload_download_duration") > timedelta(0) - assert not hasattr(records[0], "payload_upload_count") - - # WFT 2: activity result is small — no external storage - assert _TMPRL1104_DURATION_MESSAGE.match(records[1].getMessage()) - assert not hasattr(records[1], "payload_download_count") - assert not hasattr(records[1], "payload_upload_count") + downloads = [c for c in completions if c.HasField("payload_download_metrics")] + assert len(downloads) == 1 + m = downloads[0].payload_download_metrics + assert m.payload_count == 1 + assert m.total_size_bytes == expected_input_size + assert m.total_duration.ToTimedelta() > timedelta(0) + assert list(m.driver_names) == [driver.name()] + assert not any(c.HasField("payload_upload_metrics") for c in completions) async def test_tmprl1104_with_extstore_upload(env: WorkflowEnvironment) -> None: - """When external storage encodes payloads, TMPRL1104 logs include upload - metrics on the WFT that produces them.""" + """When external storage stores payloads, the completion for the WFT that + produced them carries upload metrics.""" driver = InMemoryTestDriver() data_converter = dataclasses.replace( temporalio.converter.default(), @@ -746,7 +741,7 @@ async def test_tmprl1104_with_extstore_upload(env: WorkflowEnvironment) -> None: wf_output = "wo" * 1024 # 2048 bytes → stored externally expected_output_size = await _expected_payload_size(data_converter, wf_output) - with LogCapturer().logs_captured(_workflow_logger, level=logging.DEBUG) as capturer: + with _capture_completions() as completions: async with new_worker( client, ExtStoreWorkflow, activities=[ext_store_activity] ) as worker: @@ -762,27 +757,21 @@ async def test_tmprl1104_with_extstore_upload(env: WorkflowEnvironment) -> None: task_queue=worker.task_queue, ) - records = _tmprl1104_records(capturer) - assert len(records) == 2 - - # WFT 1: small input — no external storage - assert _TMPRL1104_DURATION_MESSAGE.match(records[0].getMessage()) - assert not hasattr(records[0], "payload_download_count") - assert not hasattr(records[0], "payload_upload_count") - - # WFT 2: workflow returns large result → uploaded - assert _TMPRL1104_DURATION_MESSAGE.match(records[1].getMessage()) - assert not hasattr(records[1], "payload_download_count") - assert getattr(records[1], "payload_upload_count") == 1 - assert getattr(records[1], "payload_upload_size") == expected_output_size - assert getattr(records[1], "payload_upload_duration") > timedelta(0) + uploads = [c for c in completions if c.HasField("payload_upload_metrics")] + assert len(uploads) == 1 + m = uploads[0].payload_upload_metrics + assert m.payload_count == 1 + assert m.total_size_bytes == expected_output_size + assert m.total_duration.ToTimedelta() > timedelta(0) + assert list(m.driver_names) == [driver.name()] + assert not any(c.HasField("payload_download_metrics") for c in completions) async def test_tmprl1104_with_extstore_download_and_upload( env: WorkflowEnvironment, ) -> None: - """When both download and upload happen across WFTs, TMPRL1104 logs include - both sets of metrics.""" + """When both download and upload happen across WFTs, the respective completions + carry the matching metrics.""" driver = InMemoryTestDriver() data_converter = dataclasses.replace( temporalio.converter.default(), @@ -807,7 +796,7 @@ async def test_tmprl1104_with_extstore_download_and_upload( wf_output = "wo" * 1024 expected_output_size = await _expected_payload_size(data_converter, wf_output) - with LogCapturer().logs_captured(_workflow_logger, level=logging.DEBUG) as capturer: + with _capture_completions() as completions: async with new_worker( client, ExtStoreWorkflow, activities=[ext_store_activity] ) as worker: @@ -818,22 +807,19 @@ async def test_tmprl1104_with_extstore_download_and_upload( task_queue=worker.task_queue, ) - records = _tmprl1104_records(capturer) - assert len(records) == 2 - - # WFT 1: retrieves externalized workflow input - assert _TMPRL1104_DURATION_MESSAGE.match(records[0].getMessage()) - assert getattr(records[0], "payload_download_count") == 1 - assert getattr(records[0], "payload_download_size") == expected_input_size - assert getattr(records[0], "payload_download_duration") > timedelta(0) - assert not hasattr(records[0], "payload_upload_count") - - # WFT 2: uploads externalized workflow result - assert _TMPRL1104_DURATION_MESSAGE.match(records[1].getMessage()) - assert not hasattr(records[1], "payload_download_count") - assert getattr(records[1], "payload_upload_count") == 1 - assert getattr(records[1], "payload_upload_size") == expected_output_size - assert getattr(records[1], "payload_upload_duration") > timedelta(0) + downloads = [c for c in completions if c.HasField("payload_download_metrics")] + assert len(downloads) == 1 + dm = downloads[0].payload_download_metrics + assert dm.payload_count == 1 + assert dm.total_size_bytes == expected_input_size + assert dm.total_duration.ToTimedelta() > timedelta(0) + + uploads = [c for c in completions if c.HasField("payload_upload_metrics")] + assert len(uploads) == 1 + um = uploads[0].payload_upload_metrics + assert um.payload_count == 1 + assert um.total_size_bytes == expected_output_size + assert um.total_duration.ToTimedelta() > timedelta(0) # ---------------------------------------------------------------------------