feat: allow task group to be overridden in create_tasks/taskgraph_dec…

src/taskgraph/create.py

@@ -21,7 +21,14 @@
 testing = False
 
 
-def create_tasks(graph_config, taskgraph, label_to_taskid, params, decision_task_id):
+def create_tasks(
+    graph_config,
+    taskgraph,
+    label_to_taskid,
+    params,
+    decision_task_id,
+    task_group_id=None,
+):
     taskid_to_label = {t: l for l, t in label_to_taskid.items()}
 
     # when running as an actual decision task, we use the decision task's
@@ -43,7 +50,7 @@ def create_tasks(graph_config, taskgraph, label_to_taskid, params, decision_task
         if not any(t in taskgraph.tasks for t in task_def.get("dependencies", [])):
             task_def.setdefault("dependencies", []).append(decision_task_id)
 
-        task_def["taskGroupId"] = decision_task_id
+        task_def["taskGroupId"] = task_group_id or decision_task_id
         task_def["schedulerId"] = scheduler_id
 
     # If `testing` is True, then run without parallelization

src/taskgraph/decision.py

@@ -63,7 +63,7 @@ def full_task_graph_to_runnable_tasks(full_task_json):
     return runnable_tasks
 
 
-def taskgraph_decision(options, parameters=None):
+def taskgraph_decision(options, parameters=None, task_group_id=None):
     """
     Run the decision task.  This function implements `mach taskgraph decision`,
     and is responsible for
@@ -152,6 +152,7 @@ def taskgraph_decision(options, parameters=None):
         tgg.label_to_taskid,
         tgg.parameters,
         decision_task_id=decision_task_id,
+        task_group_id=task_group_id,
     )
 
 

test/test_create.py

@@ -58,6 +58,31 @@ def test_create_tasks(self):
                     continue
                 self.assertIn(depid, self.created_tasks)
 
+    def test_create_tasks_explicit_group_id(self):
+        tasks = {
+            "tid-a": Task(
+                kind="test", label="a", attributes={}, task={"payload": "hello world"}
+            ),
+            "tid-b": Task(
+                kind="test", label="b", attributes={}, task={"payload": "hello world"}
+            ),
+        }
+        label_to_taskid = {"a": "tid-a", "b": "tid-b"}
+        graph = Graph(nodes={"tid-a", "tid-b"}, edges={("tid-a", "tid-b", "edge")})
+        taskgraph = TaskGraph(tasks, graph)
+
+        create.create_tasks(
+            GRAPH_CONFIG,
+            taskgraph,
+            label_to_taskid,
+            {"level": "4"},
+            decision_task_id="decisiontask",
+            task_group_id="task_group",
+        )
+
+        for _, task in self.created_tasks.items():
+            self.assertEqual(task["taskGroupId"], "task_group")
+
     def test_create_task_without_dependencies(self):
         "a task with no dependencies depends on the decision task"
         tasks = {