diff --git a/.agents/setup b/.agents/setup index e238c02..da5ce44 100755 --- a/.agents/setup +++ b/.agents/setup @@ -1,6 +1,16 @@ #!/usr/bin/env bash set -euo pipefail +if ! git lfs version >/dev/null 2>&1; then + echo "Installing Git LFS" + sudo apt-get update + sudo apt-get install -y git-lfs +fi + +echo "Fetching Git LFS fixtures" +git lfs install --local +git lfs pull + if ! command -v uv >/dev/null 2>&1; then python -m pip install --user uv export PATH="$HOME/.local/bin:$PATH" diff --git a/tilebox-workflows/tests/runner/test_runner.py b/tilebox-workflows/tests/runner/test_runner.py index bc642fc..1ad57a0 100644 --- a/tilebox-workflows/tests/runner/test_runner.py +++ b/tilebox-workflows/tests/runner/test_runner.py @@ -5,6 +5,7 @@ import sys from pathlib import Path from unittest.mock import MagicMock, patch +from uuid import uuid4 import pytest @@ -12,11 +13,29 @@ from tilebox.workflows import ExecutionContext, Runner, Task from tilebox.workflows.cache import InMemoryCache, JobCache from tilebox.workflows.client import Client -from tilebox.workflows.data import JobState, ProgressIndicator, RunnerContext, TaskState +from tilebox.workflows.data import JobState, ProgressIndicator, RunnerContext, TaskIdentifier, TaskState +from tilebox.workflows.data import Task as TaskData from tilebox.workflows.observability.tracing import NoopWorkflowTracer +from tilebox.workflows.runner.executor import ExecutionContext as RunnerExecutionContext from tilebox.workflows.runner.task_runner import TaskRunner +def test_public_execution_context_attributes() -> None: + class RenameTask(Task): + def execute(self, context: ExecutionContext) -> None: + context.current_task.display = "Processing" + context.job_cache["result"] = b"processed" + + task = TaskData(uuid4(), TaskIdentifier("RenameTask", "v0.0"), display="Queued") + cache = InMemoryCache() + context = RunnerExecutionContext(MagicMock(), task, cache) + + RenameTask().execute(context) + + assert task.display == "Processing" + assert cache["result"] == b"processed" + + def test_task_authoring_imports_are_lazy() -> None: code = ( "import sys\n" @@ -45,7 +64,7 @@ class FibonacciTask(Task): async def execute(self, context: ExecutionContext) -> None: await asyncio.sleep(0) - cache: JobCache = context.job_cache # ty: ignore[unresolved-attribute] + cache: JobCache = context.job_cache key = f"fib_{self.n}" if f"fib_{self.n}" in cache: # If the result is already in the cache, we can skip the calculation @@ -66,7 +85,7 @@ class SumResultTask(Task): n: int def execute(self, context: ExecutionContext) -> None: - cache: JobCache = context.job_cache # ty: ignore[unresolved-attribute] + cache: JobCache = context.job_cache fib_n_1 = bytes_to_int(cache[f"fib_{self.n - 1}"]) fib_n_2 = bytes_to_int(cache[f"fib_{self.n - 2}"]) @@ -96,7 +115,7 @@ async def test_runner_with_fibonacci_workflow() -> None: class FlakyTask(Task): async def execute(self, context: ExecutionContext) -> None: await asyncio.sleep(0) - cache: JobCache = context.job_cache # ty: ignore[unresolved-attribute] + cache: JobCache = context.job_cache if "succeed" in cache: return # finally succeed @@ -288,14 +307,14 @@ def execute(self, context: ExecutionContext) -> None: class FailingTask(Task): def execute(self, context: ExecutionContext) -> None: - cache = context.job_cache # ty: ignore[unresolved-attribute] + cache = context.job_cache cache["failing_task"] = b"1" # to make sure it actually ran raise ValueError("This task always fails") class SucceedingTask(Task): def execute(self, context: ExecutionContext) -> None: - cache = context.job_cache # ty: ignore[unresolved-attribute] + cache = context.job_cache cache["succeeding_task"] = b"1" # to make sure it actually ran diff --git a/tilebox-workflows/tests/runner/test_worker_concurrency.py b/tilebox-workflows/tests/runner/test_worker_concurrency.py index df0a356..e510d9e 100644 --- a/tilebox-workflows/tests/runner/test_worker_concurrency.py +++ b/tilebox-workflows/tests/runner/test_worker_concurrency.py @@ -47,7 +47,7 @@ async def execute(self, context: ExecutionContext) -> None: (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: JobCache = context.job_cache cache[self.label] = self.label.encode() context.logger.info("Concurrent task executing", label=self.label) context.progress(self.label).add(1) diff --git a/tilebox-workflows/tilebox/workflows/runner/executor.py b/tilebox-workflows/tilebox/workflows/runner/executor.py index 211174a..627d65b 100644 --- a/tilebox-workflows/tilebox/workflows/runner/executor.py +++ b/tilebox-workflows/tilebox/workflows/runner/executor.py @@ -9,7 +9,7 @@ from contextlib import AbstractContextManager, contextmanager from contextvars import copy_context from threading import RLock -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from uuid import UUID from warnings import warn @@ -26,8 +26,14 @@ from tilebox.workflows.observability.logging import StructuredLogger from tilebox.workflows.observability.tracing import NoopWorkflowTracer, WorkflowTracer, start_job_span from tilebox.workflows.runner.runner import Runner +from tilebox.workflows.task import ( + CurrentTask, + FutureTask, + ProgressUpdate, + RunnerContext, + merge_future_tasks_to_submissions, +) from tilebox.workflows.task import ExecutionContext as ExecutionContextBase -from tilebox.workflows.task import FutureTask, ProgressUpdate, RunnerContext, merge_future_tasks_to_submissions from tilebox.workflows.task import Task as TaskInstance if TYPE_CHECKING: @@ -133,7 +139,8 @@ def execute_task( class ExecutionContext(ExecutionContextBase): def __init__(self, executor: TaskExecutor, task: Task, job_cache: JobCache) -> None: self._executor = executor - self.current_task = task + # Executing tasks have a job, unlike general Task data objects. + self.current_task = cast(CurrentTask, task) self.job_cache = job_cache self._sub_tasks: list[FutureTask] = [] self._progress_indicators: dict[str | None, ProgressUpdate] = {} diff --git a/tilebox-workflows/tilebox/workflows/task.py b/tilebox-workflows/tilebox/workflows/task.py index 86c8b7b..cc84e2f 100644 --- a/tilebox-workflows/tilebox/workflows/task.py +++ b/tilebox-workflows/tilebox/workflows/task.py @@ -6,7 +6,8 @@ from contextlib import suppress from dataclasses import dataclass, fields, is_dataclass from types import NoneType, UnionType -from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast, get_args, get_origin +from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast, get_args, get_origin +from uuid import UUID # from python 3.11 onwards this is available as typing.dataclass_transform: from typing_extensions import dataclass_transform @@ -15,6 +16,7 @@ from tilebox.workflows.data import RunnerContext, TaskIdentifier, TaskSubmissionGroup, TaskSubmissions if TYPE_CHECKING: + from tilebox.workflows.cache import JobCache from tilebox.workflows.observability.logging import StructuredLogger from tilebox.workflows.observability.tracing import WorkflowTracer else: @@ -375,9 +377,31 @@ def done(self, count: int) -> None: self._done += count +class CurrentJob(Protocol): + """Read-only job information available during task execution.""" + + @property + def id(self) -> UUID: ... + + @property + def name(self) -> str: ... + + +class CurrentTask(Protocol): + """Task information available during execution, with an editable display label.""" + + display: str | None + + @property + def job(self) -> CurrentJob: ... + + class ExecutionContext(ABC): """The execution context for a task.""" + current_task: CurrentTask + job_cache: "JobCache" + @abstractmethod def submit_subtask( self,