forked from mirrors/gecko-dev
82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
# 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
|
|
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
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
|