Add additional type checks

Author: lukasbindreiter

tilebox-workflows/tests/test_task.py

@@ -428,6 +428,7 @@ def test_merge_future_tasks_to_submissions() -> None:
     tasks_3 = context.submit_subtasks([TaskA(6, "six"), TaskB(8.12)], cluster="other")
 
     submissions = merge_future_tasks_to_submissions(tasks_1 + tasks_2 + tasks_3, fallback_cluster="test")
+    assert submissions is not None
     assert len(submissions.task_groups) == 1
     group = submissions.task_groups[0]
     assert group.dependencies_on_other_groups == []
@@ -465,6 +466,7 @@ def test_merge_future_tasks_to_submissions_dependencies() -> None:
     # tasks_1, tasks_2 and tasks_4 should not be merged, because they have different dependants
     # tasks_3 and tasks_5 should not be merged, because they have different dependencies
 
+    assert submissions is not None
     assert len(submissions.task_groups) == 5
     assert submissions.task_groups[0].dependencies_on_other_groups == []
     assert submissions.task_groups[0].inputs == [serialize_task(TaskA(2, "two")), serialize_task(TaskA(3, "three"))]
@@ -485,6 +487,7 @@ def test_merge_future_tasks_to_submissions_many_tasks() -> None:
     tasks_2 = context.submit_subtasks([TaskB(i / 3) for i in range(n)], depends_on=tasks_1)
 
     submissions = merge_future_tasks_to_submissions(tasks_1 + tasks_2, fallback_cluster="test")
+    assert submissions is not None
     assert len(submissions.task_groups) == 2
     assert submissions.task_groups[0].dependencies_on_other_groups == []
     assert submissions.task_groups[0].identifier_pointers == [0] * n
@@ -500,4 +503,26 @@ def test_merge_future_tasks_to_submissions_many_non_mergeable_dependency_groups(
         context.submit_subtasks([TaskB(i / 3)], depends_on=task_1)
 
     submissions = merge_future_tasks_to_submissions(context._sub_tasks, fallback_cluster="test")
+    assert submissions is not None
     assert len(submissions.task_groups) == 2 * n
+
+
+def test_merge_future_tasks_two_separate_branches() -> None:
+    context = RunnerExecutionContext(None, None, job_cache=InMemoryCache())  # type: ignore[arg-type]
+    task_a = context.submit_subtasks([TaskA(0, "Task 0")])
+    # left branch
+    task_b_left = context.submit_subtasks([TaskB(0.0)], depends_on=task_a)
+    context.submit_subtasks([TaskB(1.0)], depends_on=task_b_left)
+
+    # right branch
+    task_b_right = context.submit_subtasks([TaskB(2.0)], depends_on=task_a)
+    context.submit_subtasks([TaskB(3.0)], depends_on=task_b_right)
+
+    submissions = merge_future_tasks_to_submissions(context._sub_tasks, fallback_cluster="test")
+    assert submissions is not None
+    assert len(submissions.task_groups) == 5
+    assert submissions.task_groups[0].dependencies_on_other_groups == []
+    assert submissions.task_groups[1].dependencies_on_other_groups == [0]
+    assert submissions.task_groups[2].dependencies_on_other_groups == [1]
+    assert submissions.task_groups[3].dependencies_on_other_groups == [0]
+    assert submissions.task_groups[4].dependencies_on_other_groups == [3]

tilebox-workflows/tilebox/workflows/jobs/client.py

@@ -78,12 +78,13 @@ def submit(
             )
 
         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, merge_future_tasks_to_submissions(task_submissions, default_cluster)
-            )
+            return self._service.submit(job_name, trace_parent, submissions_merged)
 
     def retry(self, job_or_id: JobIDLike) -> int:
         """Retry a job.

tilebox-workflows/tilebox/workflows/runner/task_runner.py

@@ -526,12 +526,20 @@ def __init__(self, runner: TaskRunner, task: Task, job_cache: JobCache) -> None:
     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] = []
-        for dep in depends_on or []:
+
+        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):
@@ -553,7 +561,7 @@ def submit_subtasks(
         tasks: Sequence[TaskInstance],
         cluster: str | None = None,
         max_retries: int = 0,
-        depends_on: list[FutureTask] | None = None,
+        depends_on: FutureTask | list[FutureTask] | None = None,
     ) -> list[FutureTask]:
         return [
             self.submit_subtask(task, cluster=cluster, max_retries=max_retries, depends_on=depends_on) for task in tasks

tilebox-workflows/tilebox/workflows/task.py

@@ -265,7 +265,10 @@ def append_if_unique(self, value: _T) -> int:
         return index
 
 
-def merge_future_tasks_to_submissions(future_tasks: list[FutureTask], fallback_cluster: str) -> TaskSubmissions:
+def merge_future_tasks_to_submissions(future_tasks: list[FutureTask], fallback_cluster: str) -> TaskSubmissions | None:
+    if len(future_tasks) == 0:
+        return None
+
     dependants = defaultdict(set)
     for task in future_tasks:
         for dep in task.depends_on:
@@ -292,23 +295,25 @@ def merge_future_tasks_to_submissions(future_tasks: list[FutureTask], fallback_c
         group_index = group_keys.append_if_unique(group_key)
         if group_index == len(groups):  # it was a new unique group
             groups.append(TaskSubmissionGroup(dependencies_on_other_groups=task.depends_on))
-
-        group = groups[group_index]
         task_index_to_group[task.index] = group_index
 
-        group.inputs.append(task.input())
-        group.identifier_pointers.append(identifiers.append_if_unique(task.identifier()))
-        group.cluster_slug_pointers.append(cluster_slugs.append_if_unique(task.cluster or fallback_cluster))
-        group.display_pointers.append(displays.append_if_unique(task.display()))
-        group.max_retries_values.append(task.max_retries)
-
     for i in range(len(groups)):
         group = groups[i]
         group.dependencies_on_other_groups = list(
             # convert the task dependencies to group dependencies, deduplicate and sort them
             {task_index_to_group[dep] for dep in group.dependencies_on_other_groups}
         )
 
+    for task in future_tasks:
+        group_index = task_index_to_group[task.index]
+        group = groups[group_index]
+
+        group.inputs.append(task.input())
+        group.identifier_pointers.append(identifiers.append_if_unique(task.identifier()))
+        group.cluster_slug_pointers.append(cluster_slugs.append_if_unique(task.cluster or fallback_cluster))
+        group.display_pointers.append(displays.append_if_unique(task.display()))
+        group.max_retries_values.append(task.max_retries)
+
     return TaskSubmissions(
         task_groups=groups,
         cluster_slug_lookup=cluster_slugs.values,