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])