From 1b37d5343f991cfe827328f3a3af609e88ef15ff Mon Sep 17 00:00:00 2001 From: Lukas Bindreiter Date: Wed, 12 Aug 2026 16:13:48 +0200 Subject: [PATCH 1/3] Allow async task execution --- CHANGELOG.md | 5 ++ tilebox-workflows/tests/runner/test_runner.py | 24 +++++- tilebox-workflows/tests/test_task.py | 73 +++++++++++++++++-- .../tilebox/workflows/runner/executor.py | 12 ++- .../tilebox/workflows/runner/task_runner.py | 11 +++ tilebox-workflows/tilebox/workflows/task.py | 32 +++++++- 6 files changed, 141 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 84e4dd3..89fd873 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- `tilebox-workflows`: Added support for asynchronous task `execute()` methods, allowing tasks to await async APIs + directly. + ## [0.58.0] - 2026-07-31 ### Added diff --git a/tilebox-workflows/tests/runner/test_runner.py b/tilebox-workflows/tests/runner/test_runner.py index 780bfb4..45a6559 100644 --- a/tilebox-workflows/tests/runner/test_runner.py +++ b/tilebox-workflows/tests/runner/test_runner.py @@ -1,3 +1,4 @@ +import asyncio import os import re import subprocess @@ -42,7 +43,8 @@ def bytes_to_int(b: bytes) -> int: class FibonacciTask(Task): n: int - def execute(self, context: ExecutionContext) -> None: + async def execute(self, context: ExecutionContext) -> None: + await asyncio.sleep(0) cache: JobCache = context.job_cache # ty: ignore[unresolved-attribute] key = f"fib_{self.n}" if f"fib_{self.n}" in cache: @@ -91,7 +93,8 @@ def test_runner_with_fibonacci_workflow() -> None: class FlakyTask(Task): - def execute(self, context: ExecutionContext) -> None: + async def execute(self, context: ExecutionContext) -> None: + await asyncio.sleep(0) cache: JobCache = context.job_cache # ty: ignore[unresolved-attribute] if "succeed" in cache: return # finally succeed @@ -219,8 +222,8 @@ def execute(self, context: ExecutionContext) -> None: pass -def test_runner_disallow_duplicate_task_identifiers() -> None: - runner = TaskRunner( +def _mock_task_runner() -> TaskRunner: + return TaskRunner( MagicMock(), "dummy-cluster", InMemoryCache(), @@ -231,6 +234,19 @@ def test_runner_disallow_duplicate_task_identifiers() -> None: MagicMock(), ) + +@pytest.mark.asyncio +async def test_runner_must_be_called_from_synchronous_code() -> None: + runner = _mock_task_runner() + + for run in (runner.run_all, runner.run_forever): + with pytest.raises(RuntimeError, match="must be called from synchronous code"): + run() + + +def test_runner_disallow_duplicate_task_identifiers() -> None: + runner = _mock_task_runner() + runner.register(FlakyTask) with pytest.raises( ValueError, diff --git a/tilebox-workflows/tests/test_task.py b/tilebox-workflows/tests/test_task.py index 7b8bc42..d138abb 100644 --- a/tilebox-workflows/tests/test_task.py +++ b/tilebox-workflows/tests/test_task.py @@ -1,12 +1,13 @@ -import __future__ - import json +from collections.abc import Awaitable from dataclasses import dataclass +from datetime import datetime, timezone from typing import Annotated import pytest from tests.proto.test_pb2 import SampleArgs +from tilebox.types import CRS, GeographicArea, GridSpec, PixelWindow, SpatialResolution, TimeInterval from tilebox.workflows.cache import InMemoryCache from tilebox.workflows.data import TaskIdentifier from tilebox.workflows.runner.task_runner import ExecutionContext as RunnerExecutionContext @@ -48,18 +49,47 @@ def execute(self, context: ExecutionContext) -> None: assert TaskMeta.for_task(SimpleTask).executable is True -def _compile_task_with_postponed_return_annotation(return_annotation: str) -> type: +def test_task_validation_execute_awaitable_return_type() -> None: + class AwaitableTask(Task): + def execute(self, context: ExecutionContext) -> Awaitable[None]: + _ = context + + async def execute_async() -> None: + pass + + return execute_async() + + assert TaskMeta.for_task(AwaitableTask).executable is True + + +def test_task_validation_execute_awaitable_with_value_return_type() -> None: + with pytest.raises(TypeError, match="to not have a return value"): + + class InvalidAwaitableTask(Task): + def execute(self, context: ExecutionContext) -> Awaitable[int]: + _ = context + + async def execute_async() -> int: + return 1 + + return execute_async() + + +def _compile_task_with_postponed_return_annotation( + return_annotation: str, context_annotation: str = "ExecutionContext" +) -> type: source = f""" +from __future__ import annotations + class PostponedAnnotationsTask(Task): - def execute(self, context: ExecutionContext) -> {return_annotation}: + def execute(self, context: {context_annotation}) -> {return_annotation}: pass """ - namespace: dict[str, type] = {"Task": Task, "ExecutionContext": ExecutionContext} + namespace: dict[str, type] = {"Task": Task, "ExecutionContext": ExecutionContext, "Awaitable": Awaitable} code = compile( source, filename="", mode="exec", - flags=__future__.annotations.compiler_flag, dont_inherit=True, ) exec(code, namespace) # noqa: S102 @@ -72,6 +102,19 @@ def test_task_validation_execute_none_return_type_with_postponed_annotations() - assert TaskMeta.for_task(task_class).executable is True +def test_task_validation_execute_awaitable_return_type_with_postponed_annotations() -> None: + task_class = _compile_task_with_postponed_return_annotation("Awaitable[None]") + + assert TaskMeta.for_task(task_class).executable is True + + +@pytest.mark.parametrize("return_annotation", ["None", "Awaitable[None]"]) +def test_task_validation_postponed_return_annotation_fallback(return_annotation: str) -> None: + task_class = _compile_task_with_postponed_return_annotation(return_annotation, "MissingContext") + + assert TaskMeta.for_task(task_class).executable is True + + def test_task_validation_execute_invalid_return_type_with_postponed_annotations() -> None: with pytest.raises(TypeError, match="to not have a return value"): _compile_task_with_postponed_return_annotation("int") @@ -253,6 +296,24 @@ def test_serialize_deserialize_task_nested_json() -> None: assert deserialize_task(ExampleTaskWithNestedJson, serialize_task(task)) == task +class ExampleTaskWithSharedTypes(Task): + area: GeographicArea + time: TimeInterval + grid: GridSpec + window: PixelWindow + + +def test_serialize_deserialize_task_shared_types() -> None: + task = ExampleTaskWithSharedTypes( + area=GeographicArea.from_bounds(16.1, 48.0, 16.7, 48.4), + time=TimeInterval(datetime(2026, 1, 1, tzinfo=timezone.utc), datetime(2026, 2, 1, tzinfo=timezone.utc)), + grid=GridSpec(CRS("EPSG:3857"), SpatialResolution.square(10, unit="metre")), + window=PixelWindow(0, 0, 256, 256), + ) + + assert deserialize_task(ExampleTaskWithSharedTypes, serialize_task(task)) == task + + class ExampleTaskWithNestedProtobuf(Task): x: str nested: SampleArgs diff --git a/tilebox-workflows/tilebox/workflows/runner/executor.py b/tilebox-workflows/tilebox/workflows/runner/executor.py index 72cab8b..d84a6a9 100644 --- a/tilebox-workflows/tilebox/workflows/runner/executor.py +++ b/tilebox-workflows/tilebox/workflows/runner/executor.py @@ -1,9 +1,11 @@ from __future__ import annotations +import asyncio +import inspect import json import logging from base64 import b64encode -from collections.abc import Callable, Iterator, MutableMapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, MutableMapping, Sequence from contextlib import AbstractContextManager, contextmanager from typing import TYPE_CHECKING from uuid import UUID @@ -287,7 +289,13 @@ def _set_task_input_span_attribute(span: object, task_input: bytes | None) -> No def _execute(task: TaskInstance, context: ExecutionContext) -> None: - return task.execute(context) + result = task.execute(context) + if inspect.isawaitable(result): + asyncio.run(_await_execute(result)) + + +async def _await_execute(result: Awaitable[None]) -> None: + await result @contextmanager diff --git a/tilebox-workflows/tilebox/workflows/runner/task_runner.py b/tilebox-workflows/tilebox/workflows/runner/task_runner.py index 6f5894c..04dbcd8 100644 --- a/tilebox-workflows/tilebox/workflows/runner/task_runner.py +++ b/tilebox-workflows/tilebox/workflows/runner/task_runner.py @@ -1,3 +1,4 @@ +import asyncio import random import signal import threading @@ -448,6 +449,7 @@ def run_forever(self) -> None: Run the task runner forever. This will poll for new tasks and execute them as they come in. If no tasks are available, it will sleep for a short time and then try again. """ + _ensure_no_running_event_loop() with _GracefulShutdown(_SHUTDOWN_GRACE_PERIOD, self._polling_runner) as shutdown_context: self._polling_runner.run_forever(shutdown_context) @@ -455,5 +457,14 @@ def run_all(self) -> None: """ Run the task runner and execute all tasks, until there are no more tasks available. """ + _ensure_no_running_event_loop() with _GracefulShutdown(_SHUTDOWN_GRACE_PERIOD, self._polling_runner) as shutdown_context: self._polling_runner.run_all(shutdown_context) + + +def _ensure_no_running_event_loop() -> None: + try: + asyncio.get_running_loop() + except RuntimeError: + return + raise RuntimeError("TaskRunner.run_all() and TaskRunner.run_forever() must be called from synchronous code.") diff --git a/tilebox-workflows/tilebox/workflows/task.py b/tilebox-workflows/tilebox/workflows/task.py index 3a7e0b7..dd7be4f 100644 --- a/tilebox-workflows/tilebox/workflows/task.py +++ b/tilebox-workflows/tilebox/workflows/task.py @@ -4,8 +4,10 @@ from abc import ABC, ABCMeta, abstractmethod from base64 import b64decode, b64encode from collections import defaultdict -from collections.abc import Sequence +from collections.abc import Awaitable, Sequence +from contextlib import suppress from dataclasses import dataclass, fields, is_dataclass +from datetime import datetime from types import NoneType, UnionType from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast, get_args, get_origin @@ -104,7 +106,7 @@ class Task(metaclass=_ABCTaskify): This class is a dataclass. The task is automatically assigned an identifier based on the class name. """ - def execute(self, context: "ExecutionContext") -> None: + def execute(self, context: "ExecutionContext") -> Awaitable[None] | None: """The entry point for the execution of the task. It is called when the task is executed and is responsible for performing the task's operation. @@ -141,13 +143,31 @@ def _validate_execute_method( f"but got {class_name}.execute{signature}!" ) - # `from __future__ import annotations` stores `-> None` as the string "None". - if signature.return_annotation not in (None, "None", inspect.Signature.empty): + return_annotation = signature.return_annotation + with suppress(NameError, TypeError): + return_annotation = typing.get_type_hints(execute).get("return", return_annotation) + + if not _is_valid_execute_return_annotation(return_annotation): raise TypeError(f"Expected {class_name}.execute{signature} to not have a return value!") return True +def _is_valid_execute_return_annotation(annotation: Any) -> bool: + if annotation in (None, NoneType, "None", "Awaitable[None]", inspect.Signature.empty): + return True + + origin = get_origin(annotation) + if origin in (typing.Union, UnionType): + return all(_is_valid_execute_return_annotation(member) for member in get_args(annotation)) + + if not isinstance(origin, type) or not issubclass(origin, Awaitable): + return False + + result_types = get_args(annotation) + return bool(result_types) and result_types[-1] in (None, NoneType) + + @dataclass class TaskMeta: identifier: TaskIdentifier @@ -458,6 +478,8 @@ def _serialize_as_dict(task: Task) -> dict[str, Any]: def _serialize_value(value: Any, base64_encode_protobuf: bool) -> Any: # noqa: PLR0911 + if isinstance(value, datetime): + return value.isoformat() if isinstance(value, list): return [_serialize_value(v, base64_encode_protobuf) for v in value] if isinstance(value, tuple): @@ -515,6 +537,8 @@ def _deserialize_value(field_type: type, value: Any) -> Any: # noqa: PLR0911 return None field_type = _get_deserialization_field_type(field_type) + if field_type is datetime and isinstance(value, str): + return datetime.fromisoformat(value) if hasattr(field_type, "FromString"): return field_type.FromString(b64decode(value)) # ty: ignore[call-non-callable] if is_dataclass(field_type) and isinstance(value, dict): From 5299ab682a38d1c580d6f63ce60a86103f732a0c Mon Sep 17 00:00:00 2001 From: Lukas Bindreiter Date: Wed, 12 Aug 2026 16:22:01 +0200 Subject: [PATCH 2/3] Fix lint issues --- .../protobuf_conversion/protobuf_xarray.py | 27 +++++++++---------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py b/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py index 517512e..640bb5a 100644 --- a/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py +++ b/tilebox-datasets/tilebox/datasets/protobuf_conversion/protobuf_xarray.py @@ -229,13 +229,11 @@ def resize(self, buffer_size: int) -> None: if self._data.shape == (): # the first time we actually allocate a buffer self._data = np.full((buffer_size, self._type.value_dim), self._type.fill_value, dtype=self._type.dtype) elif buffer_size > len(self._data): - # resize the data buffer to the new capacity, by just padding it with zeros at the end - missing = buffer_size - len(self._data) - self._data = np.pad( - self._data, - ((0, missing), (0, 0)), - constant_values=self._type.fill_value, - ) + current_size = len(self._data) + data = np.empty((buffer_size, self._type.value_dim), dtype=self._type.dtype) + data[:current_size] = self._data + data[current_size:] = self._type.fill_value + self._data = data class _ArrayFieldConverter(_FieldConverter): @@ -306,14 +304,13 @@ def _resize(self) -> None: self._data = np.full( (self._capacity, self._array_dim, self._type.value_dim), self._type.fill_value, dtype=self._type.dtype ) - else: # resize the data buffer to the new capacity, by just padding it with zeros at the end - missing_capacity = self._capacity - self._data.shape[0] - missing_array_dim = self._array_dim - self._data.shape[1] - self._data = np.pad( - self._data, - ((0, missing_capacity), (0, missing_array_dim), (0, 0)), - constant_values=self._type.fill_value, - ) + else: + current_capacity, current_array_dim = self._data.shape[:2] + data = np.empty((self._capacity, self._array_dim, self._type.value_dim), dtype=self._type.dtype) + data[:current_capacity, :current_array_dim] = self._data + data[current_capacity:] = self._type.fill_value + data[:current_capacity, current_array_dim:] = self._type.fill_value + self._data = data class _EnumFieldConverter(_SimpleFieldConverter): From d3c803dfeed6a1b0a74052f775732f91d9054d2f Mon Sep 17 00:00:00 2001 From: Lukas Bindreiter Date: Wed, 12 Aug 2026 16:38:20 +0200 Subject: [PATCH 3/3] Fix failing test --- tilebox-workflows/tests/test_task.py | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/tilebox-workflows/tests/test_task.py b/tilebox-workflows/tests/test_task.py index d138abb..2dbe7ca 100644 --- a/tilebox-workflows/tests/test_task.py +++ b/tilebox-workflows/tests/test_task.py @@ -1,13 +1,11 @@ import json from collections.abc import Awaitable from dataclasses import dataclass -from datetime import datetime, timezone from typing import Annotated import pytest from tests.proto.test_pb2 import SampleArgs -from tilebox.types import CRS, GeographicArea, GridSpec, PixelWindow, SpatialResolution, TimeInterval from tilebox.workflows.cache import InMemoryCache from tilebox.workflows.data import TaskIdentifier from tilebox.workflows.runner.task_runner import ExecutionContext as RunnerExecutionContext @@ -296,24 +294,6 @@ def test_serialize_deserialize_task_nested_json() -> None: assert deserialize_task(ExampleTaskWithNestedJson, serialize_task(task)) == task -class ExampleTaskWithSharedTypes(Task): - area: GeographicArea - time: TimeInterval - grid: GridSpec - window: PixelWindow - - -def test_serialize_deserialize_task_shared_types() -> None: - task = ExampleTaskWithSharedTypes( - area=GeographicArea.from_bounds(16.1, 48.0, 16.7, 48.4), - time=TimeInterval(datetime(2026, 1, 1, tzinfo=timezone.utc), datetime(2026, 2, 1, tzinfo=timezone.utc)), - grid=GridSpec(CRS("EPSG:3857"), SpatialResolution.square(10, unit="metre")), - window=PixelWindow(0, 0, 256, 256), - ) - - assert deserialize_task(ExampleTaskWithSharedTypes, serialize_task(task)) == task - - class ExampleTaskWithNestedProtobuf(Task): x: str nested: SampleArgs