Worker execution runtime

Author: lukasbindreiter

tilebox-datasets/tilebox/datasets/datasets/v1/collections_pb2_grpc.py

@@ -6,7 +6,7 @@
 from tilebox.datasets.datasets.v1 import core_pb2 as datasets_dot_v1_dot_core__pb2
 
 
-class CollectionServiceStub(object):
+class CollectionServiceStub:
     """CollectionService is the service definition for the Tilebox datasets service, which provides access to datasets
     """
 
@@ -38,7 +38,7 @@ def __init__(self, channel):
                 _registered_method=True)
 
 
-class CollectionServiceServicer(object):
+class CollectionServiceServicer:
     """CollectionService is the service definition for the Tilebox datasets service, which provides access to datasets
     """
 
@@ -97,7 +97,7 @@ def add_CollectionServiceServicer_to_server(servicer, server):
 
 
  # This class is part of an EXPERIMENTAL API.
-class CollectionService(object):
+class CollectionService:
     """CollectionService is the service definition for the Tilebox datasets service, which provides access to datasets
     """
 

tilebox-datasets/tilebox/datasets/datasets/v1/data_access_pb2_grpc.py

@@ -6,7 +6,7 @@
 from tilebox.datasets.datasets.v1 import data_access_pb2 as datasets_dot_v1_dot_data__access__pb2
 
 
-class DataAccessServiceStub(object):
+class DataAccessServiceStub:
     """DataAccessService provides data access and querying capabilities for Tilebox datasets.
     """
 
@@ -28,7 +28,7 @@ def __init__(self, channel):
                 _registered_method=True)
 
 
-class DataAccessServiceServicer(object):
+class DataAccessServiceServicer:
     """DataAccessService provides data access and querying capabilities for Tilebox datasets.
     """
 
@@ -67,7 +67,7 @@ def add_DataAccessServiceServicer_to_server(servicer, server):
 
 
  # This class is part of an EXPERIMENTAL API.
-class DataAccessService(object):
+class DataAccessService:
     """DataAccessService provides data access and querying capabilities for Tilebox datasets.
     """
 

tilebox-datasets/tilebox/datasets/datasets/v1/data_ingestion_pb2_grpc.py

@@ -5,7 +5,7 @@
 from tilebox.datasets.datasets.v1 import data_ingestion_pb2 as datasets_dot_v1_dot_data__ingestion__pb2
 
 
-class DataIngestionServiceStub(object):
+class DataIngestionServiceStub:
     """DataIngestionService provides data ingestion and deletion capabilities for Tilebox datasets.
     """
 
@@ -27,7 +27,7 @@ def __init__(self, channel):
                 _registered_method=True)
 
 
-class DataIngestionServiceServicer(object):
+class DataIngestionServiceServicer:
     """DataIngestionService provides data ingestion and deletion capabilities for Tilebox datasets.
     """
 
@@ -64,7 +64,7 @@ def add_DataIngestionServiceServicer_to_server(servicer, server):
 
 
  # This class is part of an EXPERIMENTAL API.
-class DataIngestionService(object):
+class DataIngestionService:
     """DataIngestionService provides data ingestion and deletion capabilities for Tilebox datasets.
     """
 

tilebox-datasets/tilebox/datasets/datasets/v1/datasets_pb2_grpc.py

@@ -6,7 +6,7 @@
 from tilebox.datasets.datasets.v1 import datasets_pb2 as datasets_dot_v1_dot_datasets__pb2
 
 
-class DatasetServiceStub(object):
+class DatasetServiceStub:
     """DatasetsService is the CRUD service for Tilebox datasets.
     """
 
@@ -43,7 +43,7 @@ def __init__(self, channel):
                 _registered_method=True)
 
 
-class DatasetServiceServicer(object):
+class DatasetServiceServicer:
     """DatasetsService is the CRUD service for Tilebox datasets.
     """
 
@@ -113,7 +113,7 @@ def add_DatasetServiceServicer_to_server(servicer, server):
 
 
  # This class is part of an EXPERIMENTAL API.
-class DatasetService(object):
+class DatasetService:
     """DatasetsService is the CRUD service for Tilebox datasets.
     """
 

tilebox-workflows/tilebox/workflows/__init__.py

@@ -5,9 +5,10 @@
 
 from tilebox.workflows.client import Client
 from tilebox.workflows.data import Job
+from tilebox.workflows.runner.runner import Runner
 from tilebox.workflows.task import ExecutionContext, Task
 
-__all__ = ["Client", "ExecutionContext", "Job", "Task"]
+__all__ = ["Client", "ExecutionContext", "Job", "Runner", "Task"]
 
 
 def _init_logging(level: str = "INFO") -> None:

tilebox-workflows/tilebox/workflows/client.py

@@ -4,7 +4,6 @@
 from uuid import uuid4
 
 from _tilebox.grpc.channel import open_channel, parse_channel_info
-from tilebox.datasets.sync.client import Client as DatasetsClient
 from tilebox.workflows.automations.client import AutomationClient, AutomationService
 from tilebox.workflows.cache import JobCache, NoCache
 from tilebox.workflows.clusters.client import ClusterClient, ClusterSlugLike, to_cluster_slug
@@ -22,8 +21,10 @@
     _create_tilebox_logger_provider,
 )
 from tilebox.workflows.observability.tracing import WorkflowTracer
+from tilebox.workflows.runner.runner import Runner
 from tilebox.workflows.runner.task_runner import TaskRunner, _LeaseRenewer
 from tilebox.workflows.runner.task_service import TaskService
+from tilebox.workflows.task import Task
 
 
 class Client:
@@ -107,9 +108,10 @@ def jobs(self) -> JobClient:
     def runner(
         self,
         cluster: ClusterSlugLike | None = None,
-        tasks: list[type] | None = None,
+        tasks: list[type[Task]] | None = None,
         cache: JobCache | None = None,
         context: type[RunnerContext] | None = None,
+        runner: Runner | None = None,
     ) -> TaskRunner:
         """Initialize a task runner.
 
@@ -118,12 +120,17 @@ def runner(
             tasks: A list of task the runner is able to execute.
             cache: The cache to share between tasks.
             context: The type of the runner context to use. Defaults to RunnerContext.
+            runner: A runner definition containing tasks, cache and context configuration.
 
         Returns:
             A task runner.
         """
+        if runner is not None and (tasks is not None or cache is not None or context is not None):
+            raise ValueError("Pass either runner or tasks/cache/context, not both.")
+
+        runner_definition = runner or Runner(tasks=tasks, cache=cache, context=context)
         if cache is None:
-            cache = NoCache()  # a no-op cache that will raise an error if it's used
+            cache = runner_definition.cache or NoCache()  # a no-op cache that will raise an error if it's used
 
         found_cluster = self.clusters().find(to_cluster_slug(cluster or ""))
 
@@ -134,14 +141,13 @@ def runner(
             # lets refactor this to a lazy loading mechanism in the future
             storage_locations = []
 
-        runner_context_type = context or RunnerContext
+        runner_context_type = runner_definition.context or RunnerContext
         runner_context = runner_context_type(
             self._tracer,
-            datasets_client=DatasetsClient(**self._auth),  # ty: ignore[invalid-argument-type]
             storage_locations=storage_locations,
         )
 
-        runner = TaskRunner(
+        task_runner = TaskRunner(
             TaskService(self._channel),
             found_cluster.slug,
             cache,
@@ -152,11 +158,10 @@ def runner(
             runner_logger=StructuredLogger(self._runner_logger, {}),
         )
 
-        if tasks is not None:
-            for task in tasks:
-                runner.register(task)
+        for task in runner_definition.tasks_by_identifier.values():
+            task_runner.register(task)
 
-        return runner
+        return task_runner
 
     def clusters(self) -> ClusterClient:
         """

tilebox-workflows/tilebox/workflows/data.py

@@ -36,7 +36,6 @@
 except ModuleNotFoundError:
     from typing import Any as S3Client
 
-from tilebox.datasets.sync.client import Client as DatasetsClient
 from tilebox.workflows.observability.tracing import NoopWorkflowTracer, WorkflowTracer
 from tilebox.workflows.workflows.v1 import automation_pb2 as automation_pb
 from tilebox.workflows.workflows.v1 import core_pb2, job_pb2, task_pb2, workflows_pb2
@@ -332,10 +331,26 @@ def to_message(self) -> workflows_pb2.Cluster:
         return workflows_pb2.Cluster(slug=self.slug, display_name=self.display_name, deletable=self.deletable)
 
 
+@dataclass(order=True, frozen=True)
+class Workflow:
+    slug: str
+    name: str
+    description: str
+
+    @classmethod
+    def from_message(cls, workflow: workflows_pb2.Workflow) -> "Workflow":
+        """Convert a Workflow protobuf message to a Workflow object."""
+        return cls(slug=workflow.slug, name=workflow.name, description=workflow.description)
+
+    def to_message(self) -> workflows_pb2.Workflow:
+        """Convert a Workflow object to a Workflow protobuf message."""
+        return workflows_pb2.Workflow(slug=self.slug, name=self.name, description=self.description)
+
+
 @dataclass
 class NextTaskToRun:
     cluster_slug: str
-    identifiers: dict[TaskIdentifier, type]
+    identifiers: dict[TaskIdentifier, type[Any]]
 
     # from message not needed, as we never return this from the server
 
@@ -474,6 +489,51 @@ def to_message(self) -> task_pb2.ComputedTask:
         )
 
 
+@dataclass
+class FailedTask:
+    task_id: UUID
+    display: str | None
+    was_workflow_error: bool
+    progress_updates: list[ProgressIndicator]
+
+    @classmethod
+    def from_message(cls, failed_task: task_pb2.TaskFailedRequest) -> "FailedTask":
+        """Convert a TaskFailedRequest protobuf message to a FailedTask object."""
+        return cls(
+            task_id=uuid_message_to_uuid(failed_task.task_id),
+            display=failed_task.display,
+            was_workflow_error=failed_task.was_workflow_error,
+            progress_updates=[ProgressIndicator.from_message(progress) for progress in failed_task.progress_updates],
+        )
+
+    @classmethod
+    def from_task_error(
+        cls,
+        task: Task,
+        error: Exception,
+        was_workflow_error: bool,
+        progress_updates: list[ProgressIndicator],
+    ) -> "FailedTask":
+        # job output is limited to 1KB, so truncate the error message if necessary
+        error_message = repr(error)[: (1024 - len(task.display or "None") - 1)]
+        display = f"{task.display}" if error_message == "" else f"{task.display}\n{error_message}"
+        return cls(
+            task_id=task.id,
+            display=display,
+            was_workflow_error=was_workflow_error,
+            progress_updates=progress_updates,
+        )
+
+    def to_message(self) -> task_pb2.TaskFailedRequest:
+        """Convert a FailedTask object to a TaskFailedRequest protobuf message."""
+        return task_pb2.TaskFailedRequest(
+            task_id=uuid_to_uuid_message(self.task_id),
+            display=self.display,
+            was_workflow_error=self.was_workflow_error,
+            progress_updates=[progress.to_message() for progress in self.progress_updates],
+        )
+
+
 def _parse_version(version: str) -> tuple[int, int]:
     """
     Parse the major and minor version from a string in the format "vMajor.Minor" and returns them as tuple of ints.
@@ -924,13 +984,11 @@ class RunnerContext:
     def __init__(
         self,
         tracer: WorkflowTracer | None = None,
-        datasets_client: DatasetsClient | None = None,
         storage_locations: list[StorageLocation] | None = None,
     ) -> None:
         if tracer is None:
             tracer = NoopWorkflowTracer()
         self.tracer = tracer
-        self.datasets_client = datasets_client
         self.storage_locations = {
             sl.id: sl._with_runner_context(self)  # noqa: SLF001
             for sl in storage_locations or []

tilebox-workflows/tilebox/workflows/interceptors.py

@@ -0,0 +1,4 @@
+from tilebox.workflows.runner.runner import Runner
+from tilebox.workflows.runner.task_runner import TaskRunner
+
+__all__ = ["Runner", "TaskRunner"]