diff --git a/src/taskgraph/create.py b/src/taskgraph/create.py index ac026e10..e0faa3dc 100644 --- a/src/taskgraph/create.py +++ b/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 diff --git a/src/taskgraph/decision.py b/src/taskgraph/decision.py index 5398bd70..5ec59e87 100644 --- a/src/taskgraph/decision.py +++ b/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, ) diff --git a/test/test_create.py b/test/test_create.py index 1b93c9b1..eb646bac 100644 --- a/test/test_create.py +++ b/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 = {