diff --git a/CHANGELOG.md b/CHANGELOG.md index 6d0767e..beb10bd 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/tilebox-workflows/README.md b/tilebox-workflows/README.md index 7f2e491..a1a3dd2 100644 --- a/tilebox-workflows/README.md +++ b/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. diff --git a/tilebox-workflows/tests/runner/test_worker_concurrency.py b/tilebox-workflows/tests/runner/test_worker_concurrency.py new file mode 100644 index 0000000..df0a356 --- /dev/null +++ b/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() diff --git a/tilebox-workflows/tilebox/workflows/cache.py b/tilebox-workflows/tilebox/workflows/cache.py index 50c64fc..d59a399 100644 --- a/tilebox-workflows/tilebox/workflows/cache.py +++ b/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 diff --git a/tilebox-workflows/tilebox/workflows/data.py b/tilebox-workflows/tilebox/workflows/data.py index 450fb30..5b52c6f 100644 --- a/tilebox-workflows/tilebox/workflows/data.py +++ b/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, diff --git a/tilebox-workflows/tilebox/workflows/runner/executor.py b/tilebox-workflows/tilebox/workflows/runner/executor.py index 74b415d..211174a 100644 --- a/tilebox-workflows/tilebox/workflows/runner/executor.py +++ b/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( diff --git a/tilebox-workflows/tilebox/workflows/runner/runner.py b/tilebox-workflows/tilebox/workflows/runner/runner.py index 1d6b179..579d5e7 100644 --- a/tilebox-workflows/tilebox/workflows/runner/runner.py +++ b/tilebox-workflows/tilebox/workflows/runner/runner.py @@ -13,6 +13,12 @@ class Runner: + """Registry and process-level resources for executing workflow tasks. + + Register tasks and configure the runner before starting it. A worker runtime can execute several registered tasks + concurrently; task classes and configured context and cache implementations must follow their concurrency contracts. + """ + def __init__( self, *, diff --git a/tilebox-workflows/tilebox/workflows/runner/worker_service.py b/tilebox-workflows/tilebox/workflows/runner/worker_service.py index eabc8ce..b5cad6b 100644 --- a/tilebox-workflows/tilebox/workflows/runner/worker_service.py +++ b/tilebox-workflows/tilebox/workflows/runner/worker_service.py @@ -72,7 +72,8 @@ def ExecuteTask( # noqa: N802 ) -> worker_pb2.ExecuteTaskResponse: logger.debug("ExecuteTask RPC called") task = Task.from_message(request) - if self._executor is None: + executor = self._executor + if executor is None: failed_task = FailedTask.from_task_error( task, RuntimeError("Worker is not initialized"), @@ -82,7 +83,7 @@ def ExecuteTask( # noqa: N802 logger.debug(f"ExecuteTask RPC returning failed task for uninitialized worker, task_id={task.id}") return worker_pb2.ExecuteTaskResponse(failed_task=failed_task.to_message()) - result = self._executor.execute_task(task) + result = executor.execute_task(task) if isinstance(result, ComputedTask): logger.debug(f"ExecuteTask RPC returning computed task, task_id={task.id}") return worker_pb2.ExecuteTaskResponse(computed_task=result.to_message()) diff --git a/tilebox-workflows/tilebox/workflows/task.py b/tilebox-workflows/tilebox/workflows/task.py index 3113c80..86c8b7b 100644 --- a/tilebox-workflows/tilebox/workflows/task.py +++ b/tilebox-workflows/tilebox/workflows/task.py @@ -108,6 +108,9 @@ 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. + A fresh task instance and execution context are created for every execution. Worker runtimes may execute + multiple tasks concurrently on different threads, and asynchronous tasks may use different event loops. + Mutable class or module state, runner context state, and custom caches must be safe for concurrent access. Args: context: The execution context for the task. It provides access to an API for submitting new tasks as part