diff --git a/src/taskgraph/util/parameterization.py b/src/taskgraph/util/parameterization.py index 48e8f8966..ab9df23e3 100644 --- a/src/taskgraph/util/parameterization.py +++ b/src/taskgraph/util/parameterization.py @@ -15,20 +15,18 @@ def _recurse(val, param_fns): def recurse(val): - if isinstance(val, list): - return [recurse(v) for v in val] - elif isinstance(val, dict): + if isinstance(val, dict): if len(val) == 1: - for param_key, param_fn in param_fns.items(): - if set(val.keys()) == {param_key}: - if isinstance(val[param_key], dict): - # handle `{"task-reference": {"": "bar"}}` - return { - param_fn(key): recurse(v) - for key, v in val[param_key].items() - } - return param_fn(val[param_key]) + ((key, value),) = val.items() + param_fn = param_fns.get(key) + if param_fn is not None: + if isinstance(value, dict): + # handle `{"task-reference": {"": "bar"}}` + return {param_fn(k): recurse(v) for k, v in value.items()} + return param_fn(value) return {k: recurse(v) for k, v in val.items()} + elif isinstance(val, list): + return [recurse(v) for v in val] else: return val diff --git a/test/test_graph_perf.py b/test/test_graph_perf.py index 83d400f1b..3559f2d7b 100644 --- a/test/test_graph_perf.py +++ b/test/test_graph_perf.py @@ -12,6 +12,11 @@ from taskgraph.task import Task from taskgraph.taskgraph import TaskGraph from taskgraph.transforms.base import TransformSequence +from taskgraph.util.parameterization import ( + resolve_task_references, + resolve_timestamps, +) +from taskgraph.util.time import current_json_time # --------------------------------------------------------------------------- # Graph builders – each returns (tasks_dict, Graph, TaskGraph) for 1000 nodes @@ -179,6 +184,48 @@ def test_taskgraph_to_json(geometry): assert len(data) == N +# --------------------------------------------------------------------------- +# Benchmarks – task definition parameterization +# --------------------------------------------------------------------------- + + +def _make_task_def(i): + return { + "created": {"relative-datestamp": "0 seconds"}, + "deadline": {"relative-datestamp": "1 day"}, + "expires": {"relative-datestamp": "28 days"}, + "metadata": {"name": f"task-{i}", "description": "d", "owner": "o"}, + "routes": [f"index.domain.v2.project.task-{i}"], + "payload": { + "command": ["run-task", "--", "bash", "-c", f"echo {i}"], + "env": { + "BUILD": {"task-reference": ""}, + "ARTIFACT": {"artifact-reference": ""}, + **{f"VAR{j}": f"value-{j}" for j in range(20)}, + }, + "artifacts": [{"name": f"public/a{j}", "path": "/x"} for j in range(5)], + }, + } + + +TASK_DEFS = [_make_task_def(i) for i in range(N)] + + +@pytest.mark.benchmark +def test_resolve_task_references(): + for task_def in TASK_DEFS: + resolve_task_references( + "label", task_def, "task-id", "decision-id", {"build": "build-id"} + ) + + +@pytest.mark.benchmark +def test_resolve_timestamps(): + now = current_json_time(datetime_format=True) + for task_def in TASK_DEFS: + resolve_timestamps(now, task_def) + + # --------------------------------------------------------------------------- # Benchmarks – TransformSequence with a simple transform # ---------------------------------------------------------------------------