From fbfc9bf9eb84ea5dc0b257ab67ad2ae06b4d9992 Mon Sep 17 00:00:00 2001 From: KSouleymane Date: Mon, 2 Nov 2026 12:29:53 -0500 Subject: [PATCH] feat(worker): support async def activities Registering an `async def` activity used to fail at runtime: the activity executor called it without awaiting, and the data converter then tried to serialize the coroutine ("Failed to serialize object of type 'coroutine'"). Coroutine activities now run on the worker's event loop instead of the thread pool, so an I/O-bound activity no longer holds a worker thread while it waits. Their concurrency stays bounded by `maximum_concurrent_activity_work_items`. Synchronous activities keep their existing path unchanged. - `_ActivityExecutor.execute_async` awaits awaitable results; input deserialization and output serialization are shared with `execute`. - The worker dispatches coroutine activities to `_execute_activity_async`, which keeps tracing, failure details and the activity execution hooks, and offloads the blocking calls (payload store, `CompleteActivityTask`) to a thread. - `wrap_with_release` keeps coroutine handlers awaitable; it previously wrapped every handler in a synchronous function, which would have sent coroutine handlers to the thread pool unawaited. - `call_activity()` gains an overload so an async activity's task is typed with the awaited result rather than a coroutine. - A synchronous activity that returns an awaitable now fails with a `TypeError` asking for `async def`, instead of a serialization error. Tests cover the executor, the activity hooks, handler selection, the activity span on both paths, end-to-end runs on the in-memory backend (including activities that must all be in flight at once to pass a barrier, with a single pool thread) and the static result types. --- CHANGELOG.md | 7 + docs/features.md | 16 ++ durabletask-azuremanaged/CHANGELOG.md | 5 + durabletask/task.py | 13 +- durabletask/worker.py | 242 ++++++++++++++---- examples/async_activities.py | 59 +++++ tests/durabletask/test_activity_executor.py | 80 ++++++ tests/durabletask/test_orchestration_e2e.py | 91 +++++++ tests/durabletask/test_tracing.py | 49 ++++ .../durabletask/test_worker_activity_hooks.py | 53 ++++ tests/durabletask/typing/activity_results.py | 24 ++ 11 files changed, 593 insertions(+), 46 deletions(-) create mode 100644 examples/async_activities.py create mode 100644 tests/durabletask/typing/activity_results.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 864adff1..28597e8b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,11 @@ ADDED - Added `OrchestrationQuery.task_hub_names` to filter orchestration instances by one or more task hubs. +- Activities can now be `async def` functions. An async activity no longer +occupies a worker thread while it waits on I/O; the number running at once is +bounded by `ConcurrencyOptions.maximum_concurrent_activity_work_items`, and +`call_activity()` types its task with the awaited result. See +[activities](docs/features.md#activities). CHANGED @@ -20,6 +25,8 @@ Python versions use the standard-library module. - Context managers now preserve subclass types in type checkers, including Azure Managed clients and workers. - Entity ID copy and pickle hooks now preserve subclass types in type checkers. +- A synchronous activity that returns an awaitable now fails with a `TypeError` +asking for `async def`, instead of a JSON serialization error. - Multiple task completion callback failures now raise the built-in `ExceptionGroup`, preserving each exception and its traceback in memory and the first callback error as the explicit cause and serialized inner failure. diff --git a/docs/features.md b/docs/features.md index e3b10b2c..fdecc52e 100644 --- a/docs/features.md +++ b/docs/features.md @@ -76,6 +76,22 @@ guarantees, meaning that they will be executed at least once but may be executed the event of a transient failure. Activity functions are where the real "work" of any orchestration is done. +Activities can also be `async def` functions. The worker awaits them on its event loop instead of +running them in its thread pool, so I/O-bound activities (HTTP or database calls, async SDK clients) +don't each hold a thread while they wait. Their concurrency is bounded by +`ConcurrencyOptions.maximum_concurrent_activity_work_items`. + +```python +async def fetch_profile(ctx: task.ActivityContext, user_id: str) -> dict: + async with httpx.AsyncClient() as http: + response = await http.get(f"https://example.com/users/{user_id}") + return response.json() +``` + +Async activities share the worker's event loop: avoid blocking calls in them (synchronous I/O or +long CPU-bound work), which would delay every other async activity and the dispatch of new work +items. Keep such work in ordinary `def` activities. Orchestrators remain generator functions. + ### Durable timers Orchestrations can schedule durable timers using the `create_timer` API. These timers are durable, diff --git a/durabletask-azuremanaged/CHANGELOG.md b/durabletask-azuremanaged/CHANGELOG.md index f2f59c94..bd95b66d 100644 --- a/durabletask-azuremanaged/CHANGELOG.md +++ b/durabletask-azuremanaged/CHANGELOG.md @@ -7,6 +7,11 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## Unreleased +ADDED + +- With the corresponding core `durabletask` SDK change, `DurableTaskSchedulerWorker` +runs `async def` activities without occupying a worker thread while they wait on I/O. + CHANGED - Dropped Python 3.10 support. Python 3.11 or later is now required. diff --git a/durabletask/task.py b/durabletask/task.py index 17cc7f3d..fb6dc17c 100644 --- a/durabletask/task.py +++ b/durabletask/task.py @@ -7,7 +7,7 @@ import logging import math from abc import ABC, abstractmethod -from collections.abc import Callable, Generator, Sequence +from collections.abc import Awaitable, Callable, Generator, Sequence from dataclasses import dataclass from datetime import UTC, datetime, timedelta from typing import Any, Generic, TypeAlias, TypeVar, cast, overload @@ -140,6 +140,14 @@ def call_activity(self, activity: Activity[TInput, TOutput] | str, *, return_type: type[T]) -> CompletableTask[T]: ... + @overload + def call_activity(self, activity: AsyncActivity[TInput, TOutput], *, + input: TInput | None = ..., + retry_policy: RetryPolicy | None = ..., + tags: dict[str, str] | None = ..., + return_type: None = ...) -> CompletableTask[TOutput]: + ... + @overload def call_activity(self, activity: Activity[TInput, TOutput] | str, *, input: TInput | None = ..., @@ -948,6 +956,9 @@ def task_id(self) -> int: # Activities are simple functions that can be scheduled by orchestrators Activity: TypeAlias = Callable[[ActivityContext, TInput], TOutput] +# Async activities are awaited by the worker; their tasks complete with the awaited result +AsyncActivity: TypeAlias = Callable[[ActivityContext, TInput], Awaitable[TOutput]] + Entity: TypeAlias = Callable[[EntityContext, TInput], TOutput] | type[DurableEntity] diff --git a/durabletask/worker.py b/durabletask/worker.py index 1d4b59f8..e44cd601 100644 --- a/durabletask/worker.py +++ b/durabletask/worker.py @@ -10,6 +10,7 @@ import time from collections.abc import Callable, Generator, Sequence from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta from threading import Event, Lock, Thread @@ -830,6 +831,18 @@ def wrap_with_release( handler: Callable[..., Any], release: Callable[[], None], ) -> Callable[..., Any]: + if inspect.iscoroutinefunction(handler): + # Keep coroutine handlers recognizable as such, so the worker + # manager awaits them on its event loop rather than calling + # them from the thread pool. + async def wrapped_async(*args: Any, **kwargs: Any) -> Any: + try: + return await handler(*args, **kwargs) + finally: + release() + + return wrapped_async + def wrapped(*args: Any, **kwargs: Any) -> Any: try: return handler(*args, **kwargs) @@ -1025,7 +1038,7 @@ def stream_reader(response_stream: Any = response_stream) -> None: elif work_item.HasField("activityRequest"): submit_work_item( self._async_worker_manager.submit_activity, - self._execute_activity, + self._get_activity_handler(work_item.activityRequest), self._cancel_activity, work_item.activityRequest, stub, @@ -1399,54 +1412,130 @@ def _execute_activity( payload_helpers.deexternalize_payloads(req, self._payload_store) try: executor = _ActivityExecutor(self._registry, self._logger, self._data_converter) - with tracing.suppress_span_emission(not self._emit_trace_spans): - with tracing.start_span( - tracing.create_span_name("activity", req.name), - trace_context=req.parentTraceContext, - kind=tracing.SpanKind.SERVER, - attributes={ - tracing.ATTR_TASK_TYPE: "activity", - tracing.ATTR_TASK_INSTANCE_ID: instance_id, - tracing.ATTR_TASK_NAME: req.name, - tracing.ATTR_TASK_TASK_ID: str(req.taskId), - }, - ) as span: - try: - result = executor.execute( - instance_id, req.name, req.taskId, req.input.value - ) - except Exception as ex: - tracing.set_span_error(span, ex) - raise - res = pb.ActivityResponse( - instanceId=instance_id, - taskId=req.taskId, - result=ph.get_string_value(result), - completionToken=completionToken, - ) + with self._start_activity_span(req) as span: + try: + result = executor.execute( + instance_id, req.name, req.taskId, req.input.value + ) + except Exception as ex: + tracing.set_span_error(span, ex) + raise + res = self._new_activity_response(req, completionToken, result) except Exception as ex: - res = pb.ActivityResponse( - instanceId=instance_id, - taskId=req.taskId, - failureDetails=ph.new_failure_details( - ex, self._exception_properties_provider, self._logger), - completionToken=completionToken, - ) + res = self._new_activity_failure_response(req, completionToken, ex) + + self._complete_activity_task(req, stub, res) + finally: + self._on_activity_execution_completed(req) + + async def _execute_activity_async( + self, + req: pb.ActivityRequest, + stub: Any, + completionToken: Any, + ) -> None: + """Executes an ``async def`` activity on the worker's event loop. + Mirrors :meth:`_execute_activity`, but awaits the activity instead of + occupying a thread-pool worker. The blocking calls around it (payload + store access and the gRPC completion) are offloaded to a thread so + they don't stall the event loop. + """ + instance_id = req.orchestrationInstance.instanceId + self._on_activity_execution_started(req) + try: + # De-externalize any large-payload tokens in the incoming request + if self._payload_store is not None: + await asyncio.to_thread( + payload_helpers.deexternalize_payloads, req, self._payload_store + ) try: - # Externalize any large payloads in the response - if self._payload_store is not None: - payload_helpers.externalize_payloads( - res, self._payload_store, instance_id=instance_id, - ) - stub.CompleteActivityTask(res) + executor = _ActivityExecutor(self._registry, self._logger, self._data_converter) + with self._start_activity_span(req) as span: + try: + result = await executor.execute_async( + instance_id, req.name, req.taskId, req.input.value + ) + except Exception as ex: + tracing.set_span_error(span, ex) + raise + res = self._new_activity_response(req, completionToken, result) except Exception as ex: - self._logger.exception( - f"Failed to deliver activity response for '{req.name}#{req.taskId}' of orchestration ID '{instance_id}' to sidecar: {ex}" - ) + res = self._new_activity_failure_response(req, completionToken, ex) + + await asyncio.to_thread(self._complete_activity_task, req, stub, res) finally: self._on_activity_execution_completed(req) + def _get_activity_handler(self, req: pb.ActivityRequest) -> Callable[..., Any]: + """Returns the handler that executes the activity named by the request.""" + fn = self._registry.get_activity(req.name) + if fn is not None and _is_async_callable(fn): + return self._execute_activity_async + return self._execute_activity + + @contextmanager + def _start_activity_span(self, req: pb.ActivityRequest): + with tracing.suppress_span_emission(not self._emit_trace_spans): + with tracing.start_span( + tracing.create_span_name("activity", req.name), + trace_context=req.parentTraceContext, + kind=tracing.SpanKind.SERVER, + attributes={ + tracing.ATTR_TASK_TYPE: "activity", + tracing.ATTR_TASK_INSTANCE_ID: req.orchestrationInstance.instanceId, + tracing.ATTR_TASK_NAME: req.name, + tracing.ATTR_TASK_TASK_ID: str(req.taskId), + }, + ) as span: + yield span + + def _new_activity_response( + self, + req: pb.ActivityRequest, + completionToken: Any, + result: str | None, + ) -> pb.ActivityResponse: + return pb.ActivityResponse( + instanceId=req.orchestrationInstance.instanceId, + taskId=req.taskId, + result=ph.get_string_value(result), + completionToken=completionToken, + ) + + def _new_activity_failure_response( + self, + req: pb.ActivityRequest, + completionToken: Any, + ex: Exception, + ) -> pb.ActivityResponse: + return pb.ActivityResponse( + instanceId=req.orchestrationInstance.instanceId, + taskId=req.taskId, + failureDetails=ph.new_failure_details( + ex, self._exception_properties_provider, self._logger), + completionToken=completionToken, + ) + + def _complete_activity_task( + self, + req: pb.ActivityRequest, + stub: Any, + res: pb.ActivityResponse, + ) -> None: + instance_id = req.orchestrationInstance.instanceId + try: + # Externalize any large payloads in the response + if self._payload_store is not None: + payload_helpers.externalize_payloads( + res, self._payload_store, instance_id=instance_id, + ) + stub.CompleteActivityTask(res) + except Exception as ex: + self._logger.exception( + f"Failed to deliver activity response for '{req.name}#{req.taskId}' of orchestration ID '{instance_id}' to sidecar: {ex}" + ) + def _cancel_activity( self, req: pb.ActivityRequest, @@ -1878,6 +1967,18 @@ def call_activity( ) -> task.CompletableTask[T]: ... + @overload + def call_activity( + self, + activity: task.AsyncActivity[TInput, TOutput], + *, + input: TInput | None = ..., + retry_policy: task.RetryPolicy | None = ..., + tags: dict[str, str] | None = ..., + return_type: None = ..., + ) -> task.CompletableTask[TOutput]: + ... + @overload def call_activity( self, @@ -3401,6 +3502,13 @@ def compare_versions(self, source_version: str | None, default_version: str | No return (source_version > default_version) - (source_version < default_version) +def _is_async_callable(fn: Callable[..., Any]) -> bool: + """Returns True for ``async def`` functions and objects with an ``async def __call__``.""" + return inspect.iscoroutinefunction(fn) or inspect.iscoroutinefunction( + getattr(fn, "__call__", None) + ) + + class _ActivityExecutor: def __init__(self, registry: _Registry, logger: logging.Logger, data_converter: DataConverter): @@ -3416,6 +3524,45 @@ def execute( encoded_input: str | None, ) -> str | None: """Executes an activity function and returns the serialized result, if any.""" + fn, ctx, activity_input = self._prepare(orchestration_id, name, task_id, encoded_input) + + # Execute the activity function + activity_output = fn(ctx, activity_input) + if inspect.isawaitable(activity_output): + if inspect.iscoroutine(activity_output): + activity_output.close() + raise TypeError( + f"Activity '{name}' returned an awaitable from a synchronous call. " + "Declare the activity with 'async def' so the worker awaits it." + ) + + return self._serialize_output(orchestration_id, name, task_id, activity_output) + + async def execute_async( + self, + orchestration_id: str, + name: str, + task_id: int, + encoded_input: str | None, + ) -> str | None: + """Executes an activity function, awaiting its result if it is awaitable, + and returns the serialized result, if any.""" + fn, ctx, activity_input = self._prepare(orchestration_id, name, task_id, encoded_input) + + # Execute the activity function + activity_output = fn(ctx, activity_input) + if inspect.isawaitable(activity_output): + activity_output = await activity_output + + return self._serialize_output(orchestration_id, name, task_id, activity_output) + + def _prepare( + self, + orchestration_id: str, + name: str, + task_id: int, + encoded_input: str | None, + ) -> tuple[task.Activity[Any, Any], task.ActivityContext, Any]: self._logger.debug( f"{orchestration_id}/{task_id}: Executing activity '{name}'..." ) @@ -3428,10 +3575,15 @@ def execute( input_type = type_discovery.activity_input_type(fn, self._data_converter) if encoded_input else None activity_input = self._data_converter.deserialize(encoded_input, input_type) ctx = task.ActivityContext(orchestration_id, task_id) + return fn, ctx, activity_input - # Execute the activity function - activity_output = fn(ctx, activity_input) - + def _serialize_output( + self, + orchestration_id: str, + name: str, + task_id: int, + activity_output: Any, + ) -> str | None: encoded_output = self._data_converter.serialize(activity_output) chars = len(encoded_output) if encoded_output else 0 self._logger.debug( diff --git a/examples/async_activities.py b/examples/async_activities.py new file mode 100644 index 00000000..7df88a76 --- /dev/null +++ b/examples/async_activities.py @@ -0,0 +1,59 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""End-to-end sample that demonstrates how to write activities as `async def` +functions and fan them out from an orchestrator.""" +import asyncio +import os +import random +from collections.abc import Generator +from typing import Any + +from azure.identity import DefaultAzureCredential + +from durabletask import client, task +from durabletask.azuremanaged.client import DurableTaskSchedulerClient +from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker + + +async def fetch_temperature(ctx: task.ActivityContext, city: str) -> float: + """Async activity function that simulates a call to a weather API. + + The worker awaits async activities on its event loop, so many of them can + wait on I/O at the same time without each holding a worker thread.""" + await asyncio.sleep(random.uniform(0.5, 1.5)) + return round(random.uniform(-5, 30), 1) + + +def forecast(ctx: task.OrchestrationContext, cities: list[str]) -> Generator[task.Task[Any], Any, dict[str, float]]: + """Orchestrator function that fetches the temperature of each city in parallel""" + tasks: list[task.Task[float]] = [ctx.call_activity(fetch_temperature, input=city) for city in cities] + temperatures: list[float] = yield task.when_all(tasks) + return dict(zip(cities, temperatures)) + + +# Use environment variables if provided, otherwise use default emulator values +taskhub_name = os.getenv("TASKHUB", "default") +endpoint = os.getenv("ENDPOINT", "http://localhost:8080") + +print(f"Using taskhub: {taskhub_name}") +print(f"Using endpoint: {endpoint}") + +# Set credential to None for emulator, or DefaultAzureCredential for Azure +secure_channel = endpoint.startswith("https://") +credential = DefaultAzureCredential() if secure_channel else None +with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=secure_channel, + taskhub=taskhub_name, token_credential=credential) as w: + w.add_orchestrator(forecast) + w.add_activity(fetch_temperature) + w.start() + + # Construct the client and run the orchestration + c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=secure_channel, + taskhub=taskhub_name, token_credential=credential) + instance_id = c.schedule_new_orchestration(forecast, input=["Tokyo", "Seattle", "London", "Dakar"]) + state = c.wait_for_orchestration_completion(instance_id, timeout=60) + if state and state.runtime_status == client.OrchestrationStatus.COMPLETED: + print(f'Orchestration completed! Result: {state.serialized_output}') + elif state: + print(f'Orchestration failed: {state.failure_details}') diff --git a/tests/durabletask/test_activity_executor.py b/tests/durabletask/test_activity_executor.py index ff775adc..c11d05fc 100644 --- a/tests/durabletask/test_activity_executor.py +++ b/tests/durabletask/test_activity_executor.py @@ -1,10 +1,15 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import asyncio +import functools import json import logging +import warnings from typing import Any +import pytest + from durabletask import task, worker from durabletask.serialization import JsonDataConverter @@ -51,6 +56,81 @@ def test_activity(ctx: task.ActivityContext, _): assert "Bogus" in str(caught_exception) +def test_async_activity_inputs(): + """Validates that async activity functions are awaited and get their inputs""" + async def test_activity(ctx: task.ActivityContext, test_input: Any): + await asyncio.sleep(0) + return test_input, ctx.orchestration_id, ctx.task_id + + activity_input = "Hello, 世界!" + executor, name = _get_activity_executor(test_activity) + result = asyncio.run( + executor.execute_async(TEST_INSTANCE_ID, name, TEST_TASK_ID, json.dumps(activity_input)) + ) + assert result is not None + + result_input, result_orchestration_id, result_task_id = json.loads(result) + assert activity_input == result_input + assert TEST_INSTANCE_ID == result_orchestration_id + assert TEST_TASK_ID == result_task_id + + +def test_async_activity_exception_propagates(): + async def test_activity(ctx: task.ActivityContext, _): + await asyncio.sleep(0) + raise ValueError("boom") + + executor, name = _get_activity_executor(test_activity) + with pytest.raises(ValueError, match="boom"): + asyncio.run(executor.execute_async(TEST_INSTANCE_ID, name, TEST_TASK_ID, None)) + + +def test_execute_async_runs_sync_activities(): + def test_activity(ctx: task.ActivityContext, test_input: int): + return test_input + 1 + + executor, name = _get_activity_executor(test_activity) + result = asyncio.run(executor.execute_async(TEST_INSTANCE_ID, name, TEST_TASK_ID, json.dumps(1))) + assert result == json.dumps(2) + + +def test_sync_execute_rejects_awaitable_result(): + """A sync call that returns a coroutine fails with a clear error instead of a serialization error""" + async def test_activity(ctx: task.ActivityContext, _): + return "never" + + def wrapper(ctx: task.ActivityContext, test_input: Any): + return test_activity(ctx, test_input) + + executor, name = _get_activity_executor(wrapper) + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) # the coroutine is closed, not leaked + with pytest.raises(TypeError, match="async def"): + executor.execute(TEST_INSTANCE_ID, name, TEST_TASK_ID, None) + + +def test_is_async_callable(): + async def async_fn(ctx, _): + pass + + def sync_fn(ctx, _): + pass + + class AsyncCallable: + async def __call__(self, ctx, _): + pass + + class SyncCallable: + def __call__(self, ctx, _): + pass + + assert worker._is_async_callable(async_fn) + assert worker._is_async_callable(functools.partial(async_fn)) + assert worker._is_async_callable(AsyncCallable()) + assert not worker._is_async_callable(sync_fn) + assert not worker._is_async_callable(SyncCallable()) + + def _get_activity_executor(fn: task.Activity) -> tuple[worker._ActivityExecutor, str]: registry = worker._Registry() name = registry.add_activity(fn) diff --git a/tests/durabletask/test_orchestration_e2e.py b/tests/durabletask/test_orchestration_e2e.py index 1f898fdd..251a5b71 100644 --- a/tests/durabletask/test_orchestration_e2e.py +++ b/tests/durabletask/test_orchestration_e2e.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import asyncio import json import threading import time @@ -185,6 +186,96 @@ def orchestrator(ctx: task.OrchestrationContext, input: int): assert compensation_counter == 2 +def test_async_activity_sequence(): + + async def plus_one(_: task.ActivityContext, input: int) -> int: + await asyncio.sleep(0) + return input + 1 + + def sequence(ctx: task.OrchestrationContext, start_val: int): + numbers = [start_val] + current = start_val + for _ in range(10): + current = yield ctx.call_activity(plus_one, input=current) + numbers.append(current) + return numbers + + with worker.TaskHubGrpcWorker(host_address=HOST) as w: + w.add_orchestrator(sequence) + w.add_activity(plus_one) + w.start() + + with client.TaskHubGrpcClient(host_address=HOST) as task_hub_client: + id = task_hub_client.schedule_new_orchestration(sequence, input=1) + state = task_hub_client.wait_for_orchestration_completion( + id, timeout=30) + + assert state is not None + assert state.runtime_status == client.OrchestrationStatus.COMPLETED + assert state.failure_details is None + assert state.serialized_output == json.dumps([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) + + +def test_async_activity_error_handling(): + + async def throw(_: task.ActivityContext, input: int) -> int: + await asyncio.sleep(0) + raise RuntimeError("Kah-BOOOOM!!!") + + def orchestrator(ctx: task.OrchestrationContext, input: int): + try: + yield ctx.call_activity(throw, input=input) + except task.TaskFailedError as e: + return e.details.message + + with worker.TaskHubGrpcWorker(host_address=HOST) as w: + w.add_orchestrator(orchestrator) + w.add_activity(throw) + w.start() + + with client.TaskHubGrpcClient(host_address=HOST) as task_hub_client: + id = task_hub_client.schedule_new_orchestration(orchestrator, input=1) + state = task_hub_client.wait_for_orchestration_completion(id, timeout=30) + + assert state is not None + assert state.runtime_status == client.OrchestrationStatus.COMPLETED + assert state.serialized_output == json.dumps("Kah-BOOOOM!!!") + + +def test_async_activities_do_not_occupy_thread_pool(): + """Async activities run concurrently on the worker's event loop, even with a single pool thread.""" + activity_count = 5 + # Every activity must be in flight at once to pass the barrier: run one after the other on + # the single pool thread, the first would wait for the others until its timeout. + all_in_flight = asyncio.Barrier(activity_count) + activity_threads: set[int] = set() + + async def wait_for_the_others(_: task.ActivityContext, _input) -> None: + activity_threads.add(threading.get_ident()) + await asyncio.wait_for(all_in_flight.wait(), timeout=10) + + def fan_out(ctx: task.OrchestrationContext, _): + yield task.when_all([ctx.call_activity(wait_for_the_others) for _ in range(activity_count)]) + + options = worker.ConcurrencyOptions( + maximum_concurrent_activity_work_items=activity_count, + maximum_thread_pool_workers=1, + ) + with worker.TaskHubGrpcWorker(host_address=HOST, concurrency_options=options) as w: + w.add_orchestrator(fan_out) + w.add_activity(wait_for_the_others) + w.start() + + with client.TaskHubGrpcClient(host_address=HOST) as task_hub_client: + id = task_hub_client.schedule_new_orchestration(fan_out) + state = task_hub_client.wait_for_orchestration_completion(id, timeout=30) + + assert state is not None + assert state.runtime_status == client.OrchestrationStatus.COMPLETED + assert state.failure_details is None + assert len(activity_threads) == 1 + + def test_get_orchestration_history(): def plus_one(_: task.ActivityContext, input: int) -> int: diff --git a/tests/durabletask/test_tracing.py b/tests/durabletask/test_tracing.py index dfdeae25..001c9173 100644 --- a/tests/durabletask/test_tracing.py +++ b/tests/durabletask/test_tracing.py @@ -3,6 +3,7 @@ """Tests for distributed tracing utilities and integration.""" +import asyncio import json import logging from datetime import datetime, timezone @@ -1553,6 +1554,54 @@ def orchestrator(ctx: task.OrchestrationContext, _): # --------------------------------------------------------------------------- +def _sync_echo(ctx: task.ActivityContext, value: str) -> str: + return value + + +async def _async_echo(ctx: task.ActivityContext, value: str) -> str: + return value + + +def _run_sync_activity(w: worker.TaskHubGrpcWorker, req: pb.ActivityRequest, stub: MagicMock) -> None: + w._execute_activity(req, stub, "token1") + + +def _run_async_activity(w: worker.TaskHubGrpcWorker, req: pb.ActivityRequest, stub: MagicMock) -> None: + asyncio.run(w._execute_activity_async(req, stub, "token1")) + + +class TestActivityServerSpan: + """Tests for the activity server span, on the sync and async execution paths.""" + + @pytest.mark.parametrize("activity, run", [ + (_sync_echo, _run_sync_activity), + (_async_echo, _run_async_activity), + ]) + @pytest.mark.parametrize("emit_trace_spans, span_count", [(True, 1), (False, 0)]) + def test_activity_span_follows_emit_trace_spans( + self, otel_setup, activity, run, emit_trace_spans, span_count): + w = worker.TaskHubGrpcWorker( + host_address="localhost:4001", + emit_trace_spans=emit_trace_spans, + ) + name = w.add_activity(activity) + stub = MagicMock() + req = pb.ActivityRequest( + name=name, + taskId=1, + orchestrationInstance=pb.OrchestrationInstance(instanceId=TEST_INSTANCE_ID), + input=wrappers_pb2.StringValue(value=json.dumps("hello")), + parentTraceContext=_make_parent_trace_ctx(), + ) + + run(w, req, stub) + + spans = otel_setup.get_finished_spans() + assert [span.name for span in spans] == [tracing.create_span_name("activity", name)] * span_count + response = stub.CompleteActivityTask.call_args.args[0] + assert response.result.value == json.dumps("hello") + + class TestParseTraceparent: """Tests for tracing._parse_traceparent.""" diff --git a/tests/durabletask/test_worker_activity_hooks.py b/tests/durabletask/test_worker_activity_hooks.py index 5e59e731..19d87c17 100644 --- a/tests/durabletask/test_worker_activity_hooks.py +++ b/tests/durabletask/test_worker_activity_hooks.py @@ -1,6 +1,8 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import asyncio + from google.protobuf import wrappers_pb2 import durabletask.internal.orchestrator_service_pb2 as pb @@ -59,6 +61,57 @@ def fails(_ctx, _value): assert stub.completed[0].failureDetails.errorType == "builtins.ValueError" +def test_activity_execution_hooks_run_around_successful_async_activity() -> None: + async def echo(_ctx, value): + await asyncio.sleep(0) + return value + + worker = HookedWorker() + worker.add_activity(echo) + stub = RecordingStub() + request = _activity_request("echo", "hello") + + asyncio.run(worker._execute_activity_async(request, stub, b"token")) + + assert worker.events == [("started", "echo"), ("completed", "echo")] + assert len(stub.completed) == 1 + assert stub.completed[0].result.value == "\"hello\"" + + +def test_activity_execution_completed_hook_runs_after_failed_async_activity() -> None: + async def fails(_ctx, _value): + await asyncio.sleep(0) + raise ValueError("boom") + + worker = HookedWorker() + worker.add_activity(fails) + stub = RecordingStub() + request = _activity_request("fails", "hello") + + asyncio.run(worker._execute_activity_async(request, stub, b"token")) + + assert worker.events == [("started", "fails"), ("completed", "fails")] + assert len(stub.completed) == 1 + assert stub.completed[0].failureDetails.errorType == "builtins.ValueError" + + +def test_activity_handler_matches_activity_kind() -> None: + def sync_activity(_ctx, value): + return value + + async def async_activity(_ctx, value): + return value + + worker = HookedWorker() + worker.add_activity(sync_activity) + worker.add_activity(async_activity) + + assert worker._get_activity_handler(_activity_request("sync_activity", "x")) == worker._execute_activity + assert worker._get_activity_handler(_activity_request("async_activity", "x")) == worker._execute_activity_async + # Unknown activities keep the sync path, which reports ActivityNotRegisteredError. + assert worker._get_activity_handler(_activity_request("missing", "x")) == worker._execute_activity + + def _activity_request(name: str, value: str) -> pb.ActivityRequest: return pb.ActivityRequest( orchestrationInstance=pb.OrchestrationInstance(instanceId="instance-1"), diff --git a/tests/durabletask/typing/activity_results.py b/tests/durabletask/typing/activity_results.py new file mode 100644 index 00000000..c690abf8 --- /dev/null +++ b/tests/durabletask/typing/activity_results.py @@ -0,0 +1,24 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Static regression checks for the result types of scheduled activities.""" + +from typing import Any, assert_type + +from durabletask import task + + +def sync_activity(ctx: task.ActivityContext, value: int) -> float: + return float(value) + + +async def async_activity(ctx: task.ActivityContext, value: int) -> float: + return float(value) + + +def check_activity_result_types(ctx: task.OrchestrationContext) -> None: + assert_type(ctx.call_activity(sync_activity, input=1), task.CompletableTask[float]) + # An async activity's task completes with the awaited result, not a coroutine. + assert_type(ctx.call_activity(async_activity, input=1), task.CompletableTask[float]) + assert_type(ctx.call_activity(async_activity, input=1, return_type=int), task.CompletableTask[int]) + assert_type(ctx.call_activity("async_activity", input=1), task.CompletableTask[Any])