Allow concurrent workers in one shared runtime (#54)

Author: lukasbindreiter

CHANGELOG.md

@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - `tilebox-datasets`: Allow `iter_datapoints()` to handle empty query results.
+- `tilebox-workflows`: Make built-in worker state safe for concurrent task execution in a shared Python runtime and
+  document the concurrency contract for custom runner contexts, caches, and shared task state.
 
 ## [0.60.0] - 2026-08-25
 

tilebox-workflows/README.md

@@ -71,6 +71,21 @@ runner = client.runner(tasks=[MyFirstTask])
 runner.run_all()
 ```
 
+## Concurrent worker execution
+
+A worker runtime can execute multiple tasks concurrently in one Python process. Each execution receives a newly
+deserialized task instance and its own `ExecutionContext`, including task-local subtask and progress state. The
+`RunnerContext`, configured `JobCache`, and any class or module state are process-level resources shared by those
+executions.
+
+Custom runner contexts, caches, and shared task state must therefore support concurrent access from multiple threads.
+Asynchronous task executions may also run on different event loops. Configure and register these resources before the
+worker starts; do not mutate runner configuration while tasks are executing. Compound cache operations are not atomic
+unless the cache implementation explicitly provides that guarantee.
+
+Concurrency in one runtime avoids repeated process initialization and allows overlapping I/O or native code that
+releases Python's GIL. CPU-bound Python code still needs multiple runtime processes for parallel execution.
+
 ## Documentation
 
 Check out the [Tilebox Workflows documentation](https://docs.tilebox.com/workflows/introduction) for more information.

tilebox-workflows/tests/runner/test_worker_concurrency.py

@@ -0,0 +1,175 @@
+import asyncio
+import logging
+import socket
+import threading
+from datetime import datetime, timedelta, timezone
+from typing import ClassVar
+from unittest.mock import MagicMock, patch
+from uuid import UUID, uuid4
+
+import grpc
+import pytest
+from google.protobuf.empty_pb2 import Empty
+
+from tilebox.datasets.uuid import must_uuid_to_uuid_message
+from tilebox.workflows import ExecutionContext, Runner, Task
+from tilebox.workflows.cache import InMemoryCache, JobCache
+from tilebox.workflows.data import ExecutionStats, Job, JobState, RunnerContext, TaskState
+from tilebox.workflows.data import Task as TaskData
+from tilebox.workflows.observability.tracing import NoopWorkflowTracer
+from tilebox.workflows.runner.executor import LazyStorageLocations
+from tilebox.workflows.runner.worker_server import serve_runner
+from tilebox.workflows.task import TaskMeta
+from tilebox.workflows.workflows.v1 import core_pb2, worker_pb2, worker_pb2_grpc
+
+
+def test_worker_executes_tasks_concurrently_with_isolated_execution_state(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    class SharedRunnerContext(RunnerContext):
+        instances: ClassVar[list["SharedRunnerContext"]] = []
+
+        def __init__(self, tracer: NoopWorkflowTracer) -> None:
+            super().__init__(tracer)
+            self.instances.append(self)
+
+    class ConcurrentTask(Task):
+        label: str
+
+        barrier: ClassVar[threading.Barrier] = threading.Barrier(2)
+        observations: ClassVar[list[tuple[str, int, int, int, int]]] = []
+        observations_lock = threading.Lock()
+
+        async def execute(self, context: ExecutionContext) -> None:
+            await asyncio.sleep(0)
+            with self.observations_lock:
+                self.observations.append(
+                    (self.label, id(self), id(context), id(context.runner_context), id(asyncio.get_running_loop()))
+                )
+
+            cache: JobCache = context.job_cache  # ty: ignore[unresolved-attribute]
+            cache[self.label] = self.label.encode()
+            context.logger.info("Concurrent task executing", label=self.label)
+            context.progress(self.label).add(1)
+            self.barrier.wait(timeout=5)
+            context.progress(self.label).done(1)
+
+    cache = InMemoryCache()
+    runner = Runner(tasks=[ConcurrentTask], cache=cache, context=SharedRunnerContext)
+    fake_client = MagicMock()
+    fake_client._tracer = NoopWorkflowTracer()
+    fake_client._task_logger = logging.getLogger("tilebox.workflows.tests.shared-worker")
+    caplog.set_level(logging.INFO, logger=fake_client._task_logger.name)
+
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as free_socket:
+        free_socket.bind(("127.0.0.1", 0))
+        address = f"127.0.0.1:{free_socket.getsockname()[1]}"
+
+    server_thread = threading.Thread(target=serve_runner, args=(runner, address), daemon=True)
+
+    with patch("tilebox.workflows.runner.worker_service.Client", return_value=fake_client):
+        server_thread.start()
+        channel = grpc.insecure_channel(address)
+        grpc.channel_ready_future(channel).result(timeout=5)
+        stub = worker_pb2_grpc.WorkerServiceStub(channel)
+        stub.InitializeWorker(
+            worker_pb2.InitializeRunnerRequest(runner_id=must_uuid_to_uuid_message(uuid4())),
+            timeout=5,
+        )
+
+        job = _job()
+        tasks = [_task_message(ConcurrentTask(label), job) for label in ("first", "second")]
+        responses = [stub.ExecuteTask.future(task, timeout=5) for task in tasks]
+
+        try:
+            results = [response.result() for response in responses]
+        finally:
+            stub.ShutdownWorker(Empty(), timeout=5)
+            channel.close()
+            server_thread.join(timeout=10)
+
+    assert not server_thread.is_alive()
+    assert len(SharedRunnerContext.instances) == 1
+    assert all(result.HasField("computed_task") for result in results)
+    assert [result.computed_task.progress_updates[0].label for result in results] == ["first", "second"]
+    assert all(result.computed_task.progress_updates[0].total == 1 for result in results)
+    assert all(result.computed_task.progress_updates[0].done == 1 for result in results)
+
+    assert {observation[0] for observation in ConcurrentTask.observations} == {"first", "second"}
+    assert len({observation[1] for observation in ConcurrentTask.observations}) == 2
+    assert len({observation[2] for observation in ConcurrentTask.observations}) == 2
+    assert {observation[3] for observation in ConcurrentTask.observations} == {id(SharedRunnerContext.instances[0])}
+    assert len({observation[4] for observation in ConcurrentTask.observations}) == 2
+    assert sorted(cache.group(str(job.id)).items()) == [("first", b"first"), ("second", b"second")]
+
+    log_attributes = [
+        record.tilebox_structured_log_attributes  # ty: ignore[unresolved-attribute]
+        for record in caplog.records
+        if record.message == "Concurrent task executing"
+    ]
+    assert {attributes["label"] for attributes in log_attributes} == {"first", "second"}
+    assert {attributes["task_id"] for attributes in log_attributes} == {str(UUID(bytes=task.id.uuid)) for task in tasks}
+
+
+def test_lazy_storage_locations_are_loaded_once_during_concurrent_access() -> None:
+    storage_location = MagicMock()
+    storage_location.id = UUID(int=1)
+    storage_location._with_runner_context.return_value = storage_location
+
+    load_started = threading.Event()
+    release_load = threading.Event()
+
+    def storage_locations() -> list[MagicMock]:
+        load_started.set()
+        release_load.wait(timeout=5)
+        return [storage_location]
+
+    client = MagicMock()
+    client.automations.return_value.storage_locations.side_effect = storage_locations
+    locations = LazyStorageLocations(client, RunnerContext())
+
+    first = threading.Thread(target=len, args=(locations,))
+    first.start()
+    assert load_started.wait(timeout=5)
+
+    second_access_started = threading.Event()
+
+    def read_locations() -> None:
+        second_access_started.set()
+        len(locations)
+
+    second = threading.Thread(target=read_locations)
+    second.start()
+    assert second_access_started.wait(timeout=5)
+    release_load.set()
+    first.join(timeout=5)
+    second.join(timeout=5)
+
+    assert not first.is_alive()
+    assert not second.is_alive()
+    client.automations.return_value.storage_locations.assert_called_once_with()
+    assert list(locations) == [storage_location.id]
+
+
+def _job() -> Job:
+    return Job(
+        id=uuid4(),
+        name="concurrent worker test",
+        trace_parent="00-0123456789abcdef0123456789abcdef-0123456789abcdef-01",
+        state=JobState.RUNNING,
+        submitted_at=datetime.now(tz=timezone.utc),
+        progress=[],
+        execution_stats=ExecutionStats(None, None, timedelta(), timedelta(), 0, 1, {}),
+    )
+
+
+def _task_message(task: Task, job: Job) -> core_pb2.Task:
+    identifier = TaskMeta.for_task(task).identifier
+    return TaskData(
+        id=uuid4(),
+        identifier=identifier,
+        state=TaskState.RUNNING,
+        input=task._serialize(),
+        display=type(task).__name__,
+        job=job,
+    ).to_message()

tilebox-workflows/tilebox/workflows/cache.py

@@ -5,6 +5,7 @@
 from io import BytesIO
 from pathlib import Path
 from pathlib import PurePosixPath as ObjectPath
+from threading import RLock
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
@@ -17,6 +18,13 @@
 
 
 class JobCache(ABC):
+    """Cache shared by tasks belonging to the same job.
+
+    Task executions may access a cache concurrently, including from different threads in a shared worker runtime.
+    Implementations must therefore make individual cache operations and :meth:`group` thread-safe. Sequences of
+    operations, such as checking for a key before setting it, are not atomic.
+    """
+
     @abstractmethod
     def __contains__(self, key: str) -> bool: ...
     @abstractmethod
@@ -121,46 +129,50 @@ def group(self, key: str) -> "ObstoreCache":
 
 
 class InMemoryCache(JobCache):
-    def __init__(self) -> None:
+    def __init__(self, *, _lock: Any | None = None) -> None:
         """A simple in-memory cache implementation.
 
         Useful for testing and development. Provides no persistence, and
         no way of sharing data between multiple task runners.
         """
         self.cache: dict[str, bytes | InMemoryCache] = {}
+        self._lock = _lock or RLock()
 
     def __contains__(self, key: str) -> bool:
-        return key in self.cache
+        with self._lock:
+            return key in self.cache
 
     def __setitem__(self, key: str, value: bytes) -> None:
-        parent_group, key = self._resolve_slashes(key, create_missing=False)
-        parent_group.cache[key] = value
+        with self._lock:
+            parent_group, key = self._resolve_slashes(key, create_missing=False)
+            parent_group.cache[key] = value
 
     def __getitem__(self, key: str) -> bytes:
-        parent_group, key = self._resolve_slashes(key, create_missing=False)
-        item = parent_group.cache[key]
-        if not isinstance(item, bytes):
-            # item is a directory
-            raise KeyError(f"{key} is not cached!")
-        return item
+        with self._lock:
+            parent_group, key = self._resolve_slashes(key, create_missing=False)
+            item = parent_group.cache[key]
+            if not isinstance(item, bytes):
+                # item is a directory
+                raise KeyError(f"{key} is not cached!")
+            return item
 
     def __iter__(self) -> Iterator[str]:
-        for k, v in self.cache.items():
-            if isinstance(v, bytes):
-                yield k
+        with self._lock:
+            return iter([key for key, value in self.cache.items() if isinstance(value, bytes)])
 
     def group(self, key: str) -> "InMemoryCache":
-        parent_group, key = self._resolve_slashes(key, create_missing=True)
-        try:
-            group = parent_group.cache[key]
-        except KeyError:
-            group = InMemoryCache()
-            parent_group.cache[key] = group
+        with self._lock:
+            parent_group, key = self._resolve_slashes(key, create_missing=True)
+            try:
+                group = parent_group.cache[key]
+            except KeyError:
+                group = InMemoryCache(_lock=self._lock)
+                parent_group.cache[key] = group
 
-        if not isinstance(group, InMemoryCache):
-            # if key is a file, we return an empty group
-            return InMemoryCache()
-        return group
+            if not isinstance(group, InMemoryCache):
+                # if key is a file, we return an empty group
+                return InMemoryCache(_lock=self._lock)
+            return group
 
     def _resolve_slashes(self, key: str, create_missing: bool = False) -> tuple["InMemoryCache", str]:
         """Resolve slashes in a given cache key, by converting them into nested groups.
@@ -187,7 +199,7 @@ def _resolve_slashes(self, key: str, create_missing: bool = False) -> tuple["InM
             except KeyError:
                 if create_missing:
                     # create a new group for this key if it doesn't exist
-                    sub_group = InMemoryCache()
+                    sub_group = InMemoryCache(_lock=self._lock)
                     group.cache[part] = sub_group
                 else:
                     raise KeyError(f"{part} is not cached!") from None

tilebox-workflows/tilebox/workflows/data.py

@@ -1136,6 +1136,12 @@ def to_message(self) -> automation_pb.AutomationPrototype:
 
 
 class RunnerContext:
+    """Process-level context shared by task executions in a runner.
+
+    A runner creates one context instance during initialization. Worker runtimes may access that instance concurrently
+    from multiple threads, so subclasses must synchronize mutable state and use clients that support concurrent access.
+    """
+
     def __init__(
         self,
         tracer: WorkflowTracer | None = None,

tilebox-workflows/tilebox/workflows/runner/executor.py

@@ -8,6 +8,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import AbstractContextManager, contextmanager
 from contextvars import copy_context
+from threading import RLock
 from typing import TYPE_CHECKING
 from uuid import UUID
 from warnings import warn
@@ -241,35 +242,42 @@ def __init__(self, client: Client, runner_context: RunnerContext) -> None:
         self._runner_context = runner_context
         self._locations: dict[UUID, StorageLocation] = {}
         self._loaded = False
+        self._lock = RLock()
 
     def _load(self) -> None:
-        if self._loaded:
-            return
-        self._locations = {
-            location.id: location._with_runner_context(self._runner_context)  # noqa: SLF001
-            for location in self._client.automations().storage_locations()
-        }
-        self._loaded = True
+        with self._lock:
+            if self._loaded:
+                return
+            self._locations = {
+                location.id: location._with_runner_context(self._runner_context)  # noqa: SLF001
+                for location in self._client.automations().storage_locations()
+            }
+            self._loaded = True
 
     def __getitem__(self, key: UUID) -> StorageLocation:
-        self._load()
-        return self._locations[key]
+        with self._lock:
+            self._load()
+            return self._locations[key]
 
     def __setitem__(self, key: UUID, value: StorageLocation) -> None:
-        self._load()
-        self._locations[key] = value
+        with self._lock:
+            self._load()
+            self._locations[key] = value
 
     def __delitem__(self, key: UUID) -> None:
-        self._load()
-        del self._locations[key]
+        with self._lock:
+            self._load()
+            del self._locations[key]
 
     def __iter__(self) -> Iterator[UUID]:
-        self._load()
-        return iter(self._locations)
+        with self._lock:
+            self._load()
+            return iter(tuple(self._locations))
 
     def __len__(self) -> int:
-        self._load()
-        return len(self._locations)
+        with self._lock:
+            self._load()
+            return len(self._locations)
 
 
 def _finalize_mutable_progress_trackers(