diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a1a92a95..cf1b7bafe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,10 @@ to include examples, links to docs, or any other relevant information. ### Added +- Added the experimental `Worker` `patch_activation_callback` option, allowing workers + to decide whether a first non-replay `workflow.patched` call should activate a patch + during rolling deployments. + ### Changed ### Deprecated diff --git a/temporalio/worker/__init__.py b/temporalio/worker/__init__.py index 55966b35d..4f6efe68c 100644 --- a/temporalio/worker/__init__.py +++ b/temporalio/worker/__init__.py @@ -58,6 +58,7 @@ WorkerDeploymentConfig, ) from ._workflow_instance import ( + PatchActivationInput, UnsandboxedWorkflowRunner, WorkflowInstance, WorkflowInstanceDetails, @@ -77,6 +78,7 @@ "PollerBehavior", "PollerBehaviorSimpleMaximum", "PollerBehaviorAutoscaling", + "PatchActivationInput", # Interceptor base classes "Interceptor", "ActivityInboundInterceptor", diff --git a/temporalio/worker/_replayer.py b/temporalio/worker/_replayer.py index 508d5f708..a9fc11b49 100644 --- a/temporalio/worker/_replayer.py +++ b/temporalio/worker/_replayer.py @@ -255,6 +255,7 @@ def on_eviction_hook( workflow_failure_exception_types=self._config.get( "workflow_failure_exception_types", [] ), + patch_activation_callback=None, debug_mode=self._config.get("debug_mode", False), metric_meter=runtime.metric_meter, on_eviction_hook=on_eviction_hook, diff --git a/temporalio/worker/_worker.py b/temporalio/worker/_worker.py index 2ad1d42c6..ff5ffb8f0 100644 --- a/temporalio/worker/_worker.py +++ b/temporalio/worker/_worker.py @@ -40,7 +40,11 @@ _DEFAULT_WORKFLOW_TASK_EXTERNAL_STORAGE_CONCURRENCY, _WorkflowWorker, ) -from ._workflow_instance import UnsandboxedWorkflowRunner, WorkflowRunner +from ._workflow_instance import ( + PatchActivationInput, + UnsandboxedWorkflowRunner, + WorkflowRunner, +) from .workflow_sandbox import SandboxedWorkflowRunner logger = logging.getLogger(__name__) @@ -135,6 +139,7 @@ def __init__( use_worker_versioning: bool = False, disable_safe_workflow_eviction: bool = False, deployment_config: WorkerDeploymentConfig | None = None, + patch_activation_callback: Callable[[PatchActivationInput], bool] | None = None, workflow_task_poller_behavior: PollerBehavior = PollerBehaviorSimpleMaximum( maximum=5 ), @@ -307,6 +312,11 @@ def __init__( deployment_config: Deployment config for the worker. Exclusive with ``build_id`` and ``use_worker_versioning``. WARNING: This is an experimental feature and may change in the future. + patch_activation_callback: Experimental callback to decide whether the first + non-replay call to :py:func:`workflow.patched` + for a patch ID should activate that patch. The callback receives a + :py:class:`PatchActivationInput` and must return ``True`` to activate + the patch or ``False`` to leave it inactive. workflow_task_poller_behavior: Specify the behavior of workflow task polling. Defaults to a 5-poller maximum. activity_task_poller_behavior: Specify the behavior of activity task polling. @@ -368,6 +378,7 @@ def __init__( use_worker_versioning=use_worker_versioning, disable_safe_workflow_eviction=disable_safe_workflow_eviction, deployment_config=deployment_config, + patch_activation_callback=patch_activation_callback, workflow_task_poller_behavior=workflow_task_poller_behavior, activity_task_poller_behavior=activity_task_poller_behavior, nexus_task_poller_behavior=nexus_task_poller_behavior, @@ -528,6 +539,7 @@ def check_activity(activity: str): workflow_failure_exception_types=config[ "workflow_failure_exception_types" ], # type: ignore[reportTypedDictNotRequiredAccess] + patch_activation_callback=config.get("patch_activation_callback"), debug_mode=config["debug_mode"], # type: ignore[reportTypedDictNotRequiredAccess] disable_eager_activity_execution=config[ "disable_eager_activity_execution" @@ -982,6 +994,7 @@ class WorkerConfig(TypedDict, total=False): use_worker_versioning: bool disable_safe_workflow_eviction: bool deployment_config: WorkerDeploymentConfig | None + patch_activation_callback: Callable[[PatchActivationInput], bool] | None workflow_task_poller_behavior: PollerBehavior activity_task_poller_behavior: PollerBehavior nexus_task_poller_behavior: PollerBehavior diff --git a/temporalio/worker/_workflow.py b/temporalio/worker/_workflow.py index 8e6ba2726..953e07a5a 100644 --- a/temporalio/worker/_workflow.py +++ b/temporalio/worker/_workflow.py @@ -42,6 +42,7 @@ WorkflowInterceptorClassInput, ) from ._workflow_instance import ( + PatchActivationInput, WorkflowInstance, WorkflowInstanceDetails, WorkflowRunner, @@ -78,6 +79,7 @@ def __init__( data_converter: temporalio.converter.DataConverter, interceptors: Sequence[Interceptor], workflow_failure_exception_types: Sequence[type[BaseException]], + patch_activation_callback: Callable[[PatchActivationInput], bool] | None, debug_mode: bool, disable_eager_activity_execution: bool, metric_meter: temporalio.common.MetricMeter, @@ -145,6 +147,7 @@ def __init__( ) self._workflow_failure_exception_types = workflow_failure_exception_types + self._patch_activation_callback = patch_activation_callback self._running_workflows: dict[str, _RunningWorkflow] = {} self._disable_eager_activity_execution = disable_eager_activity_execution self._on_eviction_hook = on_eviction_hook @@ -798,6 +801,7 @@ def _create_workflow_instance( extern_functions=self._extern_functions, disable_eager_activity_execution=self._disable_eager_activity_execution, worker_level_failure_exception_types=self._workflow_failure_exception_types, + patch_activation_callback=self._patch_activation_callback, last_completion_result=init.last_completion_result, last_failure=last_failure, ) diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index 74edc66b7..3083c527a 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -131,6 +131,17 @@ def set_worker_level_failure_exception_types( pass +@dataclass(frozen=True) +class PatchActivationInput: + """Input for the experimental worker patch activation callback.""" + + workflow_info: temporalio.workflow.Info + """Information about the workflow execution calling ``patched``.""" + + patch_id: str + """Patch ID passed to ``patched``.""" + + @dataclass(frozen=True) class WorkflowInstanceDetails: """Immutable details for creating a workflow instance.""" @@ -144,6 +155,7 @@ class WorkflowInstanceDetails: extern_functions: Mapping[str, Callable] disable_eager_activity_execution: bool worker_level_failure_exception_types: Sequence[type[BaseException]] + patch_activation_callback: Callable[[PatchActivationInput], bool] | None last_completion_result: temporalio.api.common.v1.Payloads last_failure: Failure | None @@ -264,6 +276,7 @@ def __init__(self, det: WorkflowInstanceDetails) -> None: self._worker_level_failure_exception_types = ( det.worker_level_failure_exception_types ) + self._patch_activation_callback = det.patch_activation_callback self._primary_task: asyncio.Task[None] | None = None self._time_ns = 0 self._cancel_reason: str | None = None @@ -1363,7 +1376,17 @@ def workflow_patch(self, id: str, *, deprecated: bool) -> bool: if use_patch is not None: return use_patch - use_patch = not self._is_replaying or id in self._patches_notified + if deprecated or self._is_replaying or id in self._patches_notified: + use_patch = not self._is_replaying or id in self._patches_notified + elif self._patch_activation_callback is not None: + with self._as_read_only(in_query_or_validator=False): + use_patch = self._patch_activation_callback( + PatchActivationInput(workflow_info=self._info, patch_id=id) + ) + if type(use_patch) is not bool: + raise TypeError("Patch activation callback must return true or false") + else: + use_patch = True self._patches_memoized[id] = use_patch if use_patch: command = self._add_command() diff --git a/temporalio/worker/workflow_sandbox/_runner.py b/temporalio/worker/workflow_sandbox/_runner.py index b11c9b8c4..17f473d64 100644 --- a/temporalio/worker/workflow_sandbox/_runner.py +++ b/temporalio/worker/workflow_sandbox/_runner.py @@ -89,6 +89,7 @@ def prepare_workflow(self, defn: temporalio.workflow._Definition) -> None: extern_functions={}, disable_eager_activity_execution=False, worker_level_failure_exception_types=self._worker_level_failure_exception_types, + patch_activation_callback=None, last_completion_result=Payloads(), last_failure=Failure(), ), diff --git a/tests/worker/test_workflow.py b/tests/worker/test_workflow.py index d20077cf5..fb345cb5e 100644 --- a/tests/worker/test_workflow.py +++ b/tests/worker/test_workflow.py @@ -3192,6 +3192,328 @@ async def waiting_signal() -> bool: ] == await post_patch_handle.result() +@workflow.defn +class PatchActivationWorkflow: + def __init__(self) -> None: + self._ready = False + self._released = False + + @workflow.run + async def run(self, patch_id: str, wait_for_release: bool = False) -> list[bool]: + first = workflow.patched(patch_id) + self._ready = True + if wait_for_release: + await workflow.wait_condition(lambda: self._released) + return [first, workflow.patched(patch_id)] + + @workflow.query + def ready(self) -> bool: + return self._ready + + @workflow.signal + def release(self) -> None: + self._released = True + + +@workflow.defn +class PatchActivationDeprecateWorkflow: + @workflow.run + async def run(self, patch_id: str) -> bool: + workflow.deprecate_patch(patch_id) + return workflow.patched(patch_id) + + +@workflow.defn(name="PatchActivationRolloutWorkflow") +class PatchActivationRolloutWorkflow: + def __init__(self) -> None: + self._ready = False + self._released = False + + @workflow.run + async def run(self) -> str: + workflow.patched("rollout-patch") + self._ready = True + await workflow.wait_condition(lambda: self._released) + return "new" if workflow.patched("rollout-patch") else "old" + + @workflow.query + def ready(self) -> bool: + return self._ready + + @workflow.signal + def release(self) -> None: + self._released = True + + +@workflow.defn(name="PatchActivationRolloutWorkflow") +class PatchActivationOldRolloutWorkflow: + def __init__(self) -> None: + self._ready = False + self._released = False + + @workflow.run + async def run(self) -> str: + self._ready = True + await workflow.wait_condition(lambda: self._released) + return "old" + + @workflow.query + def ready(self) -> bool: + return self._ready + + @workflow.signal + def release(self) -> None: + self._released = True + + +async def patch_marker_count(handle: WorkflowHandle) -> int: + count = 0 + async for event in handle.fetch_history_events(): + if event.event_type is EventType.EVENT_TYPE_MARKER_RECORDED: + count += 1 + return count + + +def recording_patch_activation_callback( + calls: list[temporalio.worker.PatchActivationInput], decision: bool +) -> typing.Callable[[temporalio.worker.PatchActivationInput], bool]: + def callback(input: temporalio.worker.PatchActivationInput) -> bool: + calls.append(input) + return decision + + return callback + + +async def test_workflow_patch_activation_callback(client: Client): + workflow_id = f"workflow-{uuid.uuid4()}" + calls: list[temporalio.worker.PatchActivationInput] = [] + async with new_worker( + client, + PatchActivationWorkflow, + patch_activation_callback=recording_patch_activation_callback(calls, True), + ) as worker: + result = await client.execute_workflow( + PatchActivationWorkflow.run, + args=["my-patch", False], + id=workflow_id, + task_queue=worker.task_queue, + ) + + assert result == [True, True] + assert len(calls) == 1 + assert calls[0].workflow_info.workflow_id == workflow_id + assert calls[0].patch_id == "my-patch" + + +async def test_workflow_patch_activation_callback_can_decline(client: Client): + calls: list[temporalio.worker.PatchActivationInput] = [] + async with new_worker( + client, + PatchActivationWorkflow, + patch_activation_callback=recording_patch_activation_callback(calls, False), + ) as worker: + handle = await client.start_workflow( + PatchActivationWorkflow.run, + args=["my-patch", False], + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == [False, False] + + assert len(calls) == 1 + assert await patch_marker_count(handle) == 0 + + +async def test_workflow_patch_activation_default_activates(client: Client): + async with new_worker(client, PatchActivationWorkflow) as worker: + handle = await client.start_workflow( + PatchActivationWorkflow.run, + args=["my-patch", False], + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await handle.result() == [True, True] + + assert await patch_marker_count(handle) == 1 + + +async def test_workflow_patch_activation_callback_not_recalled_on_replay( + client: Client, +): + calls: list[temporalio.worker.PatchActivationInput] = [] + async with new_worker( + client, + PatchActivationWorkflow, + max_cached_workflows=0, + patch_activation_callback=recording_patch_activation_callback(calls, False), + ) as worker: + handle = await client.start_workflow( + PatchActivationWorkflow.run, + args=["my-patch", True], + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await assert_eq_eventually( + True, lambda: handle.query(PatchActivationWorkflow.ready) + ) + await handle.signal(PatchActivationWorkflow.release) + assert await handle.result() == [False, False] + + assert len(calls) == 1 + + +async def test_workflow_patch_activation_callback_bypassed_for_deprecate( + client: Client, +): + def unexpected_callback(_: temporalio.worker.PatchActivationInput) -> bool: + raise AssertionError("Patch activation callback should not be called") + + async with new_worker( + client, + PatchActivationDeprecateWorkflow, + patch_activation_callback=unexpected_callback, + ) as worker: + result = await client.execute_workflow( + PatchActivationDeprecateWorkflow.run, + "my-patch", + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + + assert result is True + + +async def test_workflow_patch_activation_callback_must_return_bool(client: Client): + def invalid_callback(_: temporalio.worker.PatchActivationInput) -> bool: + return cast(bool, 1) + + async with new_worker( + client, + PatchActivationWorkflow, + patch_activation_callback=invalid_callback, + ) as worker: + handle = await client.start_workflow( + PatchActivationWorkflow.run, + args=["my-patch", False], + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await assert_task_fail_eventually( + handle, + message_contains="Patch activation callback must return true or false", + ) + + +async def test_workflow_patch_activation_callback_is_read_only(client: Client): + def make_command(_: temporalio.worker.PatchActivationInput) -> bool: + workflow.upsert_memo({"foo": "bar"}) + return True + + def schedule_task(_: temporalio.worker.PatchActivationInput) -> bool: + asyncio.get_running_loop().call_soon(lambda: None) + return True + + def wait_condition(_: temporalio.worker.PatchActivationInput) -> bool: + coroutine = workflow.wait_condition(lambda: True) + try: + coroutine.send(None) + finally: + coroutine.close() + return True + + def use_random(_: temporalio.worker.PatchActivationInput) -> bool: + workflow.random().random() + return True + + callbacks = [ + (make_command, "action attempted: add command"), + (schedule_task, "action attempted: schedule task"), + (wait_condition, "action attempted: wait condition"), + (use_random, "action attempted: random"), + ] + for callback, message in callbacks: + async with new_worker( + client, + PatchActivationWorkflow, + patch_activation_callback=callback, + ) as worker: + handle = await client.start_workflow( + PatchActivationWorkflow.run, + args=["my-patch", False], + id=f"workflow-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + await assert_task_fail_eventually(handle, message_contains=message) + + +async def test_workflow_declined_patch_rolls_out_to_old_worker(client: Client): + task_queue = f"tq-{uuid.uuid4()}" + calls: list[temporalio.worker.PatchActivationInput] = [] + async with new_worker( + client, + PatchActivationRolloutWorkflow, + task_queue=task_queue, + max_cached_workflows=0, + patch_activation_callback=recording_patch_activation_callback(calls, False), + ): + handle = await client.start_workflow( + PatchActivationRolloutWorkflow.run, + id=f"workflow-{uuid.uuid4()}", + task_queue=task_queue, + ) + await assert_eq_eventually( + True, lambda: handle.query(PatchActivationRolloutWorkflow.ready) + ) + + assert len(calls) == 1 + async with new_worker( + client, + PatchActivationOldRolloutWorkflow, + task_queue=task_queue, + max_cached_workflows=0, + ): + await handle.signal("release") + assert await handle.result() == "old" + + +async def test_workflow_activated_patch_ignores_declining_worker(client: Client): + task_queue = f"tq-{uuid.uuid4()}" + activated_calls: list[temporalio.worker.PatchActivationInput] = [] + declining_calls: list[temporalio.worker.PatchActivationInput] = [] + async with new_worker( + client, + PatchActivationRolloutWorkflow, + task_queue=task_queue, + max_cached_workflows=0, + patch_activation_callback=recording_patch_activation_callback( + activated_calls, True + ), + ): + handle = await client.start_workflow( + PatchActivationRolloutWorkflow.run, + id=f"workflow-{uuid.uuid4()}", + task_queue=task_queue, + ) + await assert_eq_eventually( + True, lambda: handle.query(PatchActivationRolloutWorkflow.ready) + ) + + assert len(activated_calls) == 1 + async with new_worker( + client, + PatchActivationRolloutWorkflow, + task_queue=task_queue, + max_cached_workflows=0, + patch_activation_callback=recording_patch_activation_callback( + declining_calls, False + ), + ): + await handle.signal(PatchActivationRolloutWorkflow.release) + assert await handle.result() == "new" + + assert not declining_calls + + @workflow.defn class UUIDWorkflow: def __init__(self) -> None: