Source code
Revision control
Copy as Markdown
Other Tools
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
import copy
from textwrap import dedent
from voluptuous import ALLOW_EXTRA, Optional, Required
from taskgraph.transforms.base import TransformSequence
from taskgraph.util.schema import Schema
from taskgraph.util.templates import substitute
CHUNK_SCHEMA = Schema(
{
# Optional, so it can be used for a subset of tasks in a kind
Optional(
"chunk",
description=dedent(
"""
`chunk` can be used to split one task into `total-chunks`
tasks, substituting `this_chunk` and `total_chunks` into any
fields in `substitution-fields`.
""".lstrip()
),
): {
Required(
"total-chunks",
description=dedent(
"""
The total number of chunks to split the task into.
""".lstrip()
),
): int,
Optional(
"substitution-fields",
description=dedent(
"""
A list of fields that need to have `{this_chunk}` and/or
`{total_chunks}` replaced in them.
""".lstrip()
),
): [str],
}
},
extra=ALLOW_EXTRA,
)
transforms = TransformSequence()
transforms.add_validate(CHUNK_SCHEMA)
@transforms.add
def chunk_tasks(config, tasks):
for task in tasks:
chunk_config = task.pop("chunk", None)
if not chunk_config:
yield task
continue
total_chunks = chunk_config["total-chunks"]
for this_chunk in range(1, total_chunks + 1):
subtask = copy.deepcopy(task)
subs = {
"this_chunk": this_chunk,
"total_chunks": total_chunks,
}
subtask.setdefault("attributes", {})
subtask["attributes"].update(subs)
for field in chunk_config["substitution-fields"]:
container, subfield = subtask, field
while "." in subfield:
f, subfield = subfield.split(".", 1)
container = container[f]
subcontainer = copy.deepcopy(container[subfield])
subfield = substitute(subfield, **subs)
container[subfield] = substitute(subcontainer, **subs)
yield subtask