tilebox-workflows 0.44.0__py3-none-any.whl → 0.46.0__py3-none-any.whl

tilebox/workflows/automations/client.py

@@ -7,9 +7,9 @@ from tilebox.workflows.clusters.client import ClusterSlugLike, to_cluster_slug
 from tilebox.workflows.data import (
     AutomationPrototype,
     CronTrigger,
+    SingleTaskSubmission,
     StorageEventTrigger,
     StorageLocation,
-    TaskSubmission,
 )
 from tilebox.workflows.task import _task_meta
 
@@ -76,7 +76,7 @@ class AutomationClient:
         automation = AutomationPrototype(
             id=UUID(int=0),
             name=name,
-            prototype=TaskSubmission(
+            prototype=SingleTaskSubmission(
                 cluster_slug=to_cluster_slug(cluster or ""),
                 identifier=_task_meta(task).identifier,
                 input=task._serialize_args(),  # noqa: SLF001
@@ -118,7 +118,7 @@ class AutomationClient:
         automation = AutomationPrototype(
             id=UUID(int=0),
             name=name,
-            prototype=TaskSubmission(
+            prototype=SingleTaskSubmission(
                 cluster_slug=to_cluster_slug(cluster or ""),
                 identifier=_task_meta(task).identifier,
                 input=task._serialize_args(),  # noqa: SLF001

tilebox/workflows/data.py

@@ -6,7 +6,7 @@ from datetime import datetime, timedelta
 from enum import Enum
 from functools import lru_cache
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 from uuid import UUID
 
 import boto3
@@ -23,7 +23,7 @@ from tilebox.datasets.query.time_interval import (
     timedelta_to_duration,
     timestamp_to_datetime,
 )
-from tilebox.datasets.uuid import uuid_message_to_optional_uuid, uuid_message_to_uuid, uuid_to_uuid_message
+from tilebox.datasets.uuid import must_uuid_to_uuid_message, uuid_message_to_uuid, uuid_to_uuid_message
 
 try:
     # let's not make this a hard dependency, but if it's installed we can use its types
@@ -162,7 +162,7 @@ class Task:
         return core_pb2.Task(
             id=uuid_to_uuid_message(self.id),
             identifier=self.identifier.to_message(),
-            state=f"TASK_STATE_{self.state.name}",
+            state=cast(core_pb2.TaskState, self.state.value),
             input=self.input,
             display=self.display,
             job=self.job.to_message() if self.job else None,
@@ -189,13 +189,66 @@ class Idling:
 
 class JobState(Enum):
     UNSPECIFIED = 0
-    QUEUED = 1
-    STARTED = 2
-    COMPLETED = 3
+    SUBMITTED = 1
+    RUNNING = 2
+    STARTED = 3
+    COMPLETED = 4
+    FAILED = 5
+    CANCELED = 6
 
 
 _JOB_STATES = {state.value: state for state in JobState}
 
+# JobState.QUEUED is deprecated and has been renamed to SUBMITTED, but we keep it around for backwards compatibility
+JobState.QUEUED = JobState.SUBMITTED  # type: ignore[assignment]
+
+
+@dataclass(order=True, frozen=True)
+class ExecutionStats:
+    first_task_started_at: datetime | None
+    last_task_stopped_at: datetime | None
+    compute_time: timedelta
+    elapsed_time: timedelta
+    parallelism: float
+    total_tasks: int
+    tasks_by_state: dict[TaskState, int]
+
+    @classmethod
+    def from_message(cls, execution_stats: core_pb2.ExecutionStats) -> "ExecutionStats":
+        """Convert a ExecutionStats protobuf message to a ExecutionStats object."""
+        return cls(
+            first_task_started_at=timestamp_to_datetime(execution_stats.first_task_started_at)
+            if execution_stats.HasField("first_task_started_at")
+            else None,
+            last_task_stopped_at=timestamp_to_datetime(execution_stats.last_task_stopped_at)
+            if execution_stats.HasField("last_task_stopped_at")
+            else None,
+            compute_time=duration_to_timedelta(execution_stats.compute_time),
+            elapsed_time=duration_to_timedelta(execution_stats.elapsed_time),
+            parallelism=execution_stats.parallelism,
+            total_tasks=execution_stats.total_tasks,
+            tasks_by_state={_TASK_STATES[state.state]: state.count for state in execution_stats.tasks_by_state},
+        )
+
+    def to_message(self) -> core_pb2.ExecutionStats:
+        """Convert a ExecutionStats object to a ExecutionStats protobuf message."""
+        return core_pb2.ExecutionStats(
+            first_task_started_at=datetime_to_timestamp(self.first_task_started_at)
+            if self.first_task_started_at
+            else None,
+            last_task_stopped_at=datetime_to_timestamp(self.last_task_stopped_at)
+            if self.last_task_stopped_at
+            else None,
+            compute_time=timedelta_to_duration(self.compute_time),
+            elapsed_time=timedelta_to_duration(self.elapsed_time),
+            parallelism=self.parallelism,
+            total_tasks=self.total_tasks,
+            tasks_by_state=[
+                core_pb2.TaskStateCount(state=cast(core_pb2.TaskState, state.value), count=count)
+                for state, count in self.tasks_by_state.items()
+            ],
+        )
+
 
 @dataclass(order=True, frozen=True)
 class Job:
@@ -204,9 +257,8 @@ class Job:
     trace_parent: str
     state: JobState
     submitted_at: datetime
-    started_at: datetime | None
-    canceled: bool
     progress: list[ProgressIndicator]
+    execution_stats: ExecutionStats
 
     @classmethod
     def from_message(
@@ -219,9 +271,8 @@ class Job:
             trace_parent=job.trace_parent,
             state=_JOB_STATES[job.state],
             submitted_at=timestamp_to_datetime(job.submitted_at),
-            started_at=timestamp_to_datetime(job.started_at) if job.HasField("started_at") else None,
-            canceled=job.canceled,
             progress=[ProgressIndicator.from_message(progress) for progress in job.progress],
+            execution_stats=ExecutionStats.from_message(job.execution_stats),
             **extra_kwargs,
         )
 
@@ -231,13 +282,33 @@ class Job:
             id=uuid_to_uuid_message(self.id),
             name=self.name,
             trace_parent=self.trace_parent,
-            state=f"JOB_STATE_{self.state.name}",
+            state=cast(core_pb2.JobState, self.state.value),
             submitted_at=datetime_to_timestamp(self.submitted_at),
-            started_at=datetime_to_timestamp(self.started_at) if self.started_at else None,
-            canceled=self.canceled,
             progress=[progress.to_message() for progress in self.progress],
+            execution_stats=self.execution_stats.to_message(),
         )
 
+    @property
+    def canceled(self) -> bool:
+        warnings.warn(
+            "The canceled property on a job has been deprecated, and will be removed in a future version. "
+            "Use job.state == JobState.CANCELED, or job.state == JobState.FAILED instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        # the deprecated canceled property was also true for failed jobs, so we keep that behavior
+        return self.state in (JobState.CANCELED, JobState.FAILED)
+
+    @property
+    def started_at(self) -> datetime | None:
+        warnings.warn(
+            "The started_at property on a job has been deprecated, use `job.execution_stats.first_task_started_at` "
+            "instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.execution_stats.first_task_started_at
+
 
 @dataclass(order=True)
 class Cluster:
@@ -271,7 +342,69 @@ class NextTaskToRun:
 
 
 @dataclass
-class TaskSubmission:
+class TaskSubmissionGroup:
+    dependencies_on_other_groups: list[int]
+    inputs: list[bytes] = field(default_factory=list)
+    identifier_pointers: list[int] = field(default_factory=list)
+    cluster_slug_pointers: list[int] = field(default_factory=list)
+    display_pointers: list[int] = field(default_factory=list)
+    max_retries_values: list[int] = field(default_factory=list)
+
+    @classmethod
+    def from_message(cls, group: core_pb2.TaskSubmissionGroup) -> "TaskSubmissionGroup":
+        """Convert a TaskSubmissionGroup protobuf message to a TaskSubmissionGroup object."""
+        return cls(
+            dependencies_on_other_groups=list(group.dependencies_on_other_groups),
+            inputs=list(group.inputs),
+            identifier_pointers=list(group.identifier_pointers),
+            cluster_slug_pointers=list(group.cluster_slug_pointers),
+            display_pointers=list(group.display_pointers),
+            max_retries_values=list(group.max_retries_values),
+        )
+
+    def to_message(self) -> core_pb2.TaskSubmissionGroup:
+        """Convert a TaskSubmissionGroup object to a TaskSubmissionGroup protobuf message."""
+        return core_pb2.TaskSubmissionGroup(
+            dependencies_on_other_groups=self.dependencies_on_other_groups,
+            inputs=self.inputs,
+            identifier_pointers=self.identifier_pointers,
+            cluster_slug_pointers=self.cluster_slug_pointers,
+            display_pointers=self.display_pointers,
+            max_retries_values=self.max_retries_values,
+        )
+
+
+@dataclass
+class TaskSubmissions:
+    task_groups: list[TaskSubmissionGroup]
+    cluster_slug_lookup: list[str]
+    identifier_lookup: list[TaskIdentifier]
+    display_lookup: list[str]
+
+    @classmethod
+    def from_message(cls, sub_task: core_pb2.TaskSubmissions) -> "TaskSubmissions":
+        """Convert a TaskSubmission protobuf message to a TaskSubmission object."""
+        return cls(
+            task_groups=[TaskSubmissionGroup.from_message(group) for group in sub_task.task_groups],
+            cluster_slug_lookup=list(sub_task.cluster_slug_lookup),
+            identifier_lookup=[TaskIdentifier.from_message(identifier) for identifier in sub_task.identifier_lookup],
+            display_lookup=list(sub_task.display_lookup),
+        )
+
+    def to_message(self) -> core_pb2.TaskSubmissions:
+        """Convert a TaskSubmissions object to a TaskSubmissions protobuf message."""
+        return core_pb2.TaskSubmissions(
+            task_groups=[group.to_message() for group in self.task_groups],
+            cluster_slug_lookup=self.cluster_slug_lookup,
+            identifier_lookup=[identifier.to_message() for identifier in self.identifier_lookup],
+            display_lookup=self.display_lookup,
+        )
+
+
+@dataclass
+class SingleTaskSubmission:
+    """A submission of a single task. Used for automations."""
+
     cluster_slug: str
     identifier: TaskIdentifier
     input: bytes
@@ -280,8 +413,8 @@ class TaskSubmission:
     max_retries: int = 0
 
     @classmethod
-    def from_message(cls, sub_task: core_pb2.TaskSubmission) -> "TaskSubmission":
-        """Convert a TaskSubmission protobuf message to a TaskSubmission object."""
+    def from_message(cls, sub_task: core_pb2.SingleTaskSubmission) -> "SingleTaskSubmission":
+        """Convert a TaskSubmission protobuf message to a SingleTaskSubmission object."""
         return cls(
             cluster_slug=sub_task.cluster_slug,
             identifier=TaskIdentifier.from_message(sub_task.identifier),
@@ -291,9 +424,9 @@ class TaskSubmission:
             max_retries=sub_task.max_retries,
         )
 
-    def to_message(self) -> core_pb2.TaskSubmission:
-        """Convert a TaskSubmission object to a TaskSubmission protobuf message."""
-        return core_pb2.TaskSubmission(
+    def to_message(self) -> core_pb2.SingleTaskSubmission:
+        """Convert a SingleTaskSubmission object to a TaskSubmission protobuf message."""
+        return core_pb2.SingleTaskSubmission(
             cluster_slug=self.cluster_slug,
             identifier=self.identifier.to_message(),
             input=self.input,
@@ -307,7 +440,7 @@ class TaskSubmission:
 class ComputedTask:
     id: UUID
     display: str | None
-    sub_tasks: list[TaskSubmission]
+    sub_tasks: TaskSubmissions | None
     progress_updates: list[ProgressIndicator]
 
     @classmethod
@@ -316,7 +449,9 @@ class ComputedTask:
         return cls(
             id=uuid_message_to_uuid(computed_task.id),
             display=computed_task.display,
-            sub_tasks=[TaskSubmission.from_message(sub_task) for sub_task in computed_task.sub_tasks],
+            sub_tasks=TaskSubmissions.from_message(computed_task.sub_tasks)
+            if computed_task.HasField("sub_tasks")
+            else None,
             progress_updates=[ProgressIndicator.from_message(progress) for progress in computed_task.progress_updates],
         )
 
@@ -325,7 +460,7 @@ class ComputedTask:
         return task_pb2.ComputedTask(
             id=uuid_to_uuid_message(self.id),
             display=self.display,
-            sub_tasks=[sub_task.to_message() for sub_task in self.sub_tasks],
+            sub_tasks=self.sub_tasks.to_message() if self.sub_tasks else None,
             progress_updates=[progress.to_message() for progress in self.progress_updates],
         )
 
@@ -505,7 +640,7 @@ class TriggeredCronEvent:
 class AutomationPrototype:
     id: UUID
     name: str
-    prototype: TaskSubmission
+    prototype: SingleTaskSubmission
     storage_event_triggers: list[StorageEventTrigger]
     cron_triggers: list[CronTrigger]
 
@@ -515,7 +650,7 @@ class AutomationPrototype:
         return cls(
             id=uuid_message_to_uuid(task.id),
             name=task.name,
-            prototype=TaskSubmission.from_message(task.prototype),
+            prototype=SingleTaskSubmission.from_message(task.prototype),
             storage_event_triggers=[
                 StorageEventTrigger.from_message(trigger) for trigger in task.storage_event_triggers
             ],
@@ -595,21 +730,31 @@ class QueryJobsResponse:
 
 @dataclass(frozen=True)
 class QueryFilters:
-    time_interval: TimeInterval | None = None
-    id_interval: IDInterval | None = None
-    automation_id: UUID | None = None
+    time_interval: TimeInterval | None
+    id_interval: IDInterval | None
+    automation_ids: list[UUID]
+    job_states: list[JobState]
+    name: str | None
 
     @classmethod
     def from_message(cls, filters: job_pb2.QueryFilters) -> "QueryFilters":
         return cls(
-            time_interval=TimeInterval.from_message(filters.time_interval),
-            id_interval=IDInterval.from_message(filters.id_interval),
-            automation_id=uuid_message_to_optional_uuid(filters.automation_id),
+            time_interval=TimeInterval.from_message(filters.time_interval)
+            if filters.HasField("time_interval")
+            else None,
+            id_interval=IDInterval.from_message(filters.id_interval) if filters.HasField("id_interval") else None,
+            automation_ids=[uuid_message_to_uuid(uuid) for uuid in filters.automation_ids],
+            job_states=[_JOB_STATES[state] for state in filters.states],
+            name=filters.name or None,
         )
 
     def to_message(self) -> job_pb2.QueryFilters:
         return job_pb2.QueryFilters(
             time_interval=self.time_interval.to_message() if self.time_interval else None,
             id_interval=self.id_interval.to_message() if self.id_interval else None,
-            automation_id=uuid_to_uuid_message(self.automation_id) if self.automation_id else None,
+            automation_ids=[must_uuid_to_uuid_message(uuid) for uuid in self.automation_ids]
+            if self.automation_ids
+            else None,
+            states=[cast(core_pb2.JobState, state.value) for state in self.job_states] if self.job_states else None,
+            name=self.name or None,  # empty string becomes None
         )

tilebox/workflows/formatting/job.py

@@ -35,9 +35,7 @@ class JobWidget:
             self.widgets.append(HTML(_render_job_details_html(self.job)))
             self.widgets.append(HTML(_render_job_progress(self.job, False)))
             self.widgets.extend(
-                _progress_indicator_bar(
-                    progress.label or self.job.name, progress.done, progress.total, self.job.canceled
-                )
+                _progress_indicator_bar(progress.label or self.job.name, progress.done, progress.total, self.job.state)
                 for progress in self.job.progress
             )
             self.layout = VBox(self.widgets)
@@ -63,14 +61,17 @@ class JobWidget:
             if last_progress is None:  # first time, don't add the refresh time
                 self.widgets[1] = HTML(_render_job_progress(progress, False))
                 updated = True
-            elif progress.state != last_progress.state or progress.started_at != last_progress.started_at:
+            elif (
+                progress.state != last_progress.state
+                or progress.execution_stats.first_task_started_at != last_progress.execution_stats.first_task_started_at
+            ):
                 self.widgets[1] = HTML(_render_job_progress(progress, True))
                 updated = True
 
             if last_progress is None or progress.progress != last_progress.progress:
                 self.widgets[2:] = [
                     _progress_indicator_bar(
-                        progress.label or self.job.name, progress.done, progress.total, self.job.canceled
+                        progress.label or self.job.name, progress.done, progress.total, self.job.state
                     )
                     for progress in progress.progress
                 ]
@@ -193,7 +194,7 @@ body.vscode-dark {
   padding: 2px 10px;
 }
 
-.tbx-job-state-queued {
+.tbx-job-state-submitted {
   background-color: #f1f5f9;
   color: #0f172a;
 }
@@ -203,6 +204,11 @@ body.vscode-dark {
   color: #f8fafc;
 }
 
+.tbx-job-state-started {
+  background-color: #fd9b11;
+  color: #f8fafc;
+}
+
 .tbx-job-state-completed {
   background-color: #21c45d;
   color: #f8fafc;
@@ -213,6 +219,11 @@ body.vscode-dark {
   color: #f8fafc;
 }
 
+.tbx-job-state-canceled {
+  background-color: #94a2b3;
+  color: #f8fafc;
+}
+
 .tbx-job-progress a {
   text-decoration: underline;
 }
@@ -285,20 +296,20 @@ def _render_job_progress(job: Job, include_refresh_time: bool) -> str:
         refresh = f" <span class='tbx-detail-value-muted'>(refreshed at {current_time.strftime('%H:%M:%S')})</span> {_info_icon}"
 
     state_name = job.state.name
-    if job.state == JobState.STARTED:
-        state_name = "RUNNING" if not job.canceled else "FAILED"
 
     no_progress = ""
     if not job.progress:
         no_progress = "<span class='tbx-detail-value-muted'>No user defined progress indicators. <a href='https://docs.tilebox.com/workflows/progress' target='_blank'>Learn more</a></span>"
 
+    started_at = job.execution_stats.first_task_started_at
+
     """Render a job's progress as HTML, needs to be called after render_job_details_html since that injects the necessary CSS."""
     return f"""
 <div class="tbx-wrap">
     <div class="tbx-obj-type">Progress{refresh}</div>
     <div class="tbx-job-progress">
         <div><span class="tbx-detail-key tbx-detail-mono">state:</span> <span class="tbx-job-state tbx-job-state-{state_name.lower()}">{state_name}</span><div>
-        <div><span class="tbx-detail-key tbx-detail-mono">started_at:</span> {_render_datetime(job.started_at) if job.started_at else "<span class='tbx-detail-value-muted tbx-detail-mono'>None</span>"}<div>
+        <div><span class="tbx-detail-key tbx-detail-mono">started_at:</span> {_render_datetime(started_at) if started_at else "<span class='tbx-detail-value-muted tbx-detail-mono'>None</span>"}<div>
         <div><span class="tbx-detail-key tbx-detail-mono">progress:</span> {no_progress}</div>
     </div>
 </div>
@@ -312,9 +323,11 @@ _BAR_COLORS = {
 }
 
 
-def _progress_indicator_bar(label: str, done: int, total: int, job_cancelled: bool) -> HBox:
+def _progress_indicator_bar(label: str, done: int, total: int, state: JobState) -> HBox:
     percentage = done / total if total > 0 else 0 if done <= total else 1
-    non_completed_color = _BAR_COLORS["running"] if not job_cancelled else _BAR_COLORS["failed"]
+    non_completed_color = (
+        _BAR_COLORS["failed"] if state in (JobState.FAILED, JobState.CANCELED) else _BAR_COLORS["running"]
+    )
     progress = IntProgress(
         min=0,
         max=total,

tilebox/workflows/jobs/client.py

@@ -9,12 +9,13 @@ from tilebox.datasets.query.time_interval import TimeInterval, TimeIntervalLike
 from tilebox.workflows.clusters.client import ClusterSlugLike, to_cluster_slug
 from tilebox.workflows.data import (
     Job,
+    JobState,
     QueryFilters,
     QueryJobsResponse,
 )
 from tilebox.workflows.jobs.service import JobService
 from tilebox.workflows.observability.tracing import WorkflowTracer, get_trace_parent_of_current_span
-from tilebox.workflows.task import FutureTask
+from tilebox.workflows.task import FutureTask, merge_future_tasks_to_submissions
 from tilebox.workflows.task import Task as TaskInstance
 
 try:
@@ -64,8 +65,9 @@ class JobClient:
         """
         tasks = root_task_or_tasks if isinstance(root_task_or_tasks, list) else [root_task_or_tasks]
 
+        default_cluster = ""
         if isinstance(cluster, ClusterSlugLike | None):
-            slugs = [to_cluster_slug(cluster or "")] * len(tasks)
+            slugs = [to_cluster_slug(cluster or default_cluster)] * len(tasks)
         else:
             slugs = [to_cluster_slug(c) for c in cluster]
 
@@ -75,13 +77,14 @@ class JobClient:
                 f"or exactly one cluster per task. But got {len(tasks)} tasks and {len(slugs)} clusters."
             )
 
-        task_submissions = [
-            FutureTask(i, task, [], slugs[i], max_retries).to_submission() for i, task in enumerate(tasks)
-        ]
+        task_submissions = [FutureTask(i, task, [], slugs[i], max_retries) for i, task in enumerate(tasks)]
+        submissions_merged = merge_future_tasks_to_submissions(task_submissions, default_cluster)
+        if submissions_merged is None:
+            raise ValueError("At least one task must be submitted.")
 
         with self._tracer.start_as_current_span(f"job/{job_name}"):
             trace_parent = get_trace_parent_of_current_span()
-            return self._service.submit(job_name, trace_parent, task_submissions)
+            return self._service.submit(job_name, trace_parent, submissions_merged)
 
     def retry(self, job_or_id: JobIDLike) -> int:
         """Retry a job.
@@ -154,7 +157,13 @@ class JobClient:
         """
         return self._service.visualize(_to_uuid(job), direction, layout, sketchy)
 
-    def query(self, temporal_extent: TimeIntervalLike | IDIntervalLike, automation_id: UUID | None = None) -> list[Job]:
+    def query(
+        self,
+        temporal_extent: TimeIntervalLike | IDIntervalLike,
+        automation_ids: UUID | list[UUID] | None = None,
+        job_states: JobState | list[JobState] | None = None,
+        name: str | None = None,
+    ) -> list[Job]:
         """List jobs in the given temporal extent.
 
         Args:
@@ -170,11 +179,14 @@ class JobClient:
                 - tuple of two UUIDs: [start, end) -> Construct an IDInterval with the given start and end id
                 - tuple of two strings: [start, end) -> Construct an IDInterval with the given start and end id
                     parsed from the strings
-            automation_id: The automation id to filter jobs by. If specified, only jobs created by the given automation
-                are returned.
-
+            automation_ids: An automation id or list of automation ids to filter jobs by.
+                If specified, only jobs created by one of the selected automations are returned.
+            job_states: A job state or list of job states to filter jobs by. If specified, only jobs in one of the
+                selected states are returned.
+            name: A name to filter jobs by. If specified, only jobs with a matching name are returned. The match is
+                case-insensitive and uses a fuzzy matching scheme.
         Returns:
-            A list of jobs.
+            A list of jobs matching the given filters.
         """
         time_interval: TimeInterval | None = None
         id_interval: IDInterval | None = None
@@ -202,7 +214,21 @@ class JobClient:
                     end_inclusive=dataset_time_interval.end_inclusive,
                 )
 
-        filters = QueryFilters(time_interval=time_interval, id_interval=id_interval, automation_id=automation_id)
+        automation_ids = automation_ids or []
+        if not isinstance(automation_ids, list):
+            automation_ids = [automation_ids]
+
+        job_states = job_states or []
+        if not isinstance(job_states, list):
+            job_states = [job_states]
+
+        filters = QueryFilters(
+            time_interval=time_interval,
+            id_interval=id_interval,
+            automation_ids=automation_ids,
+            job_states=job_states,
+            name=name,
+        )
 
         def request(page: PaginationProtocol) -> QueryJobsResponse:
             query_page = Pagination(page.limit, page.starting_after)

tilebox/workflows/jobs/service.py

@@ -8,7 +8,7 @@ from tilebox.workflows.data import (
     Job,
     QueryFilters,
     QueryJobsResponse,
-    TaskSubmission,
+    TaskSubmissions,
     uuid_to_uuid_message,
 )
 from tilebox.workflows.formatting.job import JobWidget, RichDisplayJob
@@ -39,9 +39,9 @@ class JobService:
         """
         self.service = with_pythonic_errors(JobServiceStub(channel))
 
-    def submit(self, job_name: str, trace_parent: str, tasks: list[TaskSubmission]) -> Job:
+    def submit(self, job_name: str, trace_parent: str, tasks: TaskSubmissions) -> Job:
         request = SubmitJobRequest(
-            tasks=[task.to_message() for task in tasks],
+            tasks=tasks.to_message(),
             job_name=job_name,
             trace_parent=trace_parent,
         )

tilebox/workflows/runner/task_runner.py

@@ -34,7 +34,13 @@ from tilebox.workflows.observability.logging import get_logger
 from tilebox.workflows.observability.tracing import WorkflowTracer
 from tilebox.workflows.runner.task_service import TaskService
 from tilebox.workflows.task import ExecutionContext as ExecutionContextBase
-from tilebox.workflows.task import FutureTask, ProgressUpdate, RunnerContext, TaskMeta
+from tilebox.workflows.task import (
+    FutureTask,
+    ProgressUpdate,
+    RunnerContext,
+    TaskMeta,
+    merge_future_tasks_to_submissions,
+)
 from tilebox.workflows.task import Task as TaskInstance
 
 # The time we give a task to finish it's execution when a runner shutdown is requested before we forcefully stop it
@@ -489,10 +495,12 @@ class TaskRunner:
                 computed_task = ComputedTask(
                     id=task.id,
                     display=task.display,
-                    sub_tasks=[
-                        task.to_submission(self.tasks_to_run.cluster_slug)
-                        for task in context._sub_tasks  # noqa: SLF001
-                    ],
+                    sub_tasks=merge_future_tasks_to_submissions(
+                        context._sub_tasks,  # noqa: SLF001
+                        # if not otherwise specified, we use the cluster of the runner for all subtasks, which is also
+                        # the cluster of the parent task
+                        self.tasks_to_run.cluster_slug,
+                    ),
                     progress_updates=_finalize_mutable_progress_trackers(context._progress_indicators),  # noqa: SLF001
                 )
 
@@ -518,15 +526,30 @@ class ExecutionContext(ExecutionContextBase):
     def submit_subtask(
         self,
         task: TaskInstance,
-        depends_on: list[FutureTask] | None = None,
+        depends_on: FutureTask | list[FutureTask] | None = None,
         cluster: str | None = None,
         max_retries: int = 0,
     ) -> FutureTask:
+        dependencies: list[int] = []
+
+        if depends_on is None:
+            depends_on = []
+        elif isinstance(depends_on, FutureTask):
+            depends_on = [depends_on]
+        elif not isinstance(depends_on, list):
+            raise TypeError(f"Invalid dependency. Expected FutureTask or list[FutureTask], got {type(depends_on)}")
+
+        for dep in depends_on:
+            if not isinstance(dep, FutureTask):
+                raise TypeError(f"Invalid dependency. Expected FutureTask, got {type(dep)}")
+            if dep.index >= len(self._sub_tasks):
+                raise ValueError(f"Dependent task {dep.index} does not exist")
+            dependencies.append(dep.index)
         subtask = FutureTask(
             index=len(self._sub_tasks),
             task=task,
             # cyclic dependencies are not allowed, they are detected by the server and will result in an error
-            depends_on=[d.index for d in depends_on] if depends_on is not None else [],
+            depends_on=dependencies,
             cluster=cluster,
             max_retries=max_retries,
         )
@@ -534,9 +557,15 @@ class ExecutionContext(ExecutionContextBase):
         return subtask
 
     def submit_subtasks(
-        self, tasks: Sequence[TaskInstance], cluster: str | None = None, max_retries: int = 0
+        self,
+        tasks: Sequence[TaskInstance],
+        cluster: str | None = None,
+        max_retries: int = 0,
+        depends_on: FutureTask | list[FutureTask] | None = None,
     ) -> list[FutureTask]:
-        return [self.submit_subtask(task, cluster=cluster, max_retries=max_retries) for task in tasks]
+        return [
+            self.submit_subtask(task, cluster=cluster, max_retries=max_retries, depends_on=depends_on) for task in tasks
+        ]
 
     def submit_batch(
         self, tasks: Sequence[TaskInstance], cluster: str | None = None, max_retries: int = 0