Async task.execute() methods

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

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

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,

tilebox-workflows/tests/test_task.py

@@ -1,6 +1,5 @@
-import __future__
-
 import json
+from collections.abc import Awaitable
 from dataclasses import dataclass
 from typing import Annotated
 
@@ -48,18 +47,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="<postponed-annotations-test>",
  mode="exec",
- flags=__future__.annotations.compiler_flag,
  dont_inherit=True,
  )
  exec(code, namespace) # noqa: S102
@@ -72,6 +100,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")

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

tilebox-workflows/tilebox/workflows/runner/task_runner.py

@@ -1,3 +1,4 @@
+import asyncio
 import random
 import signal
 import threading
@@ -448,12 +449,22 @@ 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)
 
  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.")

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