ci: revamp passing options to Buildkite steps

Until now, we could only pass top-level parameters to a Buildkite step.

After this change we can pass nested parameters to the buildkite
pipeline. For example:

    pipeline_pr.py --step-param agents/queue=q --step-param env/VAR=7

    "agents": {
      "os": "al2",
      "queue": "q"
    },
    "env": {"VAR": "7"},

It supersedes the previous `--step-env` workaround to pass environment
variables.

Signed-off-by: Pablo Barbáchano <[email protected]>

Author: pb8o

.buildkite/common.py

@@ -23,22 +23,42 @@
 ]
 
 
+def overlay_dict(base: dict, update: dict):
+    """Overlay a dict over a base one"""
+    base = base.copy()
+    for key, val in update.items():
+        if key in base and isinstance(val, dict):
+            base[key] = overlay_dict(base.get(key, {}), val)
+        else:
+            base[key] = val
+    return base
+
+
 def field_fmt(field, args):
     """If `field` is a string, interpolate variables in `args`"""
     if not isinstance(field, str):
         return field
     return field.format(**args)
 
 
-def group(label, command, instances, platforms, agent_tags=None, **kwargs):
+def dict_fmt(dict_tmpl, args):
+    """Apply field_fmt over a hole dict"""
+    res = {}
+    for key, val in dict_tmpl.items():
+        if isinstance(val, dict):
+            res[key] = dict_fmt(val, args)
+        else:
+            res[key] = field_fmt(val, args)
+    return res
+
+
+def group(label, command, instances, platforms, **kwargs):
     """
     Generate a group step with specified parameters, for each instance+kernel
     combination
 
     https://buildkite.com/docs/pipelines/group-step
     """
-    if agent_tags is None:
-        agent_tags = []
     # Use the 1st character of the group name (should be an emoji)
     label1 = label[0]
     steps = []
@@ -48,16 +68,14 @@ def group(label, command, instances, platforms, agent_tags=None, **kwargs):
     for instance in instances:
         for os, kv in platforms:
             # fill any templated variables
-            args = {"os": os, "kv": kv, "instance": instance}
-            step_commands = [cmd.format(**args) for cmd in commands]
-            step_kwargs = {key: field_fmt(val, args) for key, val in kwargs.items()}
-            agents = [f"instance={instance}", f"kv={kv}", f"os={os}"] + agent_tags
+            args = {"instance": instance, "os": os, "kv": kv}
             step = {
-                "command": step_commands,
+                "command": [cmd.format(**args) for cmd in commands],
                 "label": f"{label1} {instance} {os} {kv}",
-                "agents": agents,
-                **step_kwargs,
+                "agents": args,
             }
+            step_kwargs = dict_fmt(kwargs, args)
+            step = overlay_dict(step_kwargs, step)
             steps.append(step)
 
     return {"group": label, "steps": steps}
@@ -68,6 +86,31 @@ def pipeline_to_json(pipeline):
     return json.dumps(pipeline, indent=4, sort_keys=True, ensure_ascii=False)
 
 
+class DictAction(argparse.Action):
+    """An argparse action that can receive a nested dictionary
+
+    Examples:
+
+        --step-param a/b/c=3
+        {"a": {"b": {"c": 3}}}
+    """
+
+    def __init__(self, option_strings, dest, nargs=None, **kwargs):
+        if nargs is not None:
+            raise ValueError("nargs not allowed")
+        super().__init__(option_strings, dest, **kwargs)
+
+    def __call__(self, parser, namespace, value, option_string=None):
+        res = getattr(namespace, self.dest, {})
+        key_str, val = value.split("=", maxsplit=1)
+        keys = key_str.split("/")
+        update = {keys[-1]: val}
+        for key in list(reversed(keys))[1:]:
+            update = {key: update}
+        res = overlay_dict(res, update)
+        setattr(namespace, self.dest, res)
+
+
 COMMON_PARSER = argparse.ArgumentParser()
 COMMON_PARSER.add_argument(
     "--instances",
@@ -88,16 +131,7 @@ def pipeline_to_json(pipeline):
     metavar="PARAM=VALUE",
     help="parameters to add to each step",
     required=False,
-    action="append",
-    default=[],
-    type=lambda arg: tuple(arg.split("=", maxsplit=1)),
-)
-COMMON_PARSER.add_argument(
-    "--step-env",
-    metavar="KEY=VALUE",
-    help="environment to use in each step",
-    required=False,
-    action="append",
-    default=[],
-    type=lambda arg: tuple(arg.split("=", maxsplit=1)),
+    action=DictAction,
+    default={},
+    type=str,
 )

.buildkite/pipeline_cross.py

@@ -30,7 +30,7 @@ def restore_step(label, src_instance, src_kv, dst_instance, dst_os, dst_kv):
         ],
         "label": label,
         "timeout": 30,
-        "agents": [f"instance={dst_instance}", f"kv={dst_kv}", f"os={dst_os}"],
+        "agents": {"instance": dst_instance, "kv": dst_kv, "os": dst_os},
     }
 
 

.buildkite/pipeline_perf.py

@@ -4,7 +4,7 @@
 
 """Generate Buildkite performance pipelines dynamically"""
 
-from common import COMMON_PARSER, group, pipeline_to_json
+from common import COMMON_PARSER, group, overlay_dict, pipeline_to_json
 
 perf_test = {
     "block": {
@@ -48,7 +48,6 @@ def build_group(test):
     return group(
         label=test.pop("label"),
         command=f"./tools/devtool -y test {devtool_opts} -- -m nonci --reruns {retries} {test_path}",
-        agent_tags=["ag=1"],
         artifacts=["./test_results/*"],
         instances=test.pop("instances"),
         platforms=test.pop("platforms"),
@@ -72,10 +71,10 @@ def build_group(test):
 for test_data in tests:
     test_data.setdefault("platforms", args.platforms)
     test_data.setdefault("instances", args.instances)
-    test_data["env"] = dict(args.step_env)
+    test_data.setdefault("agents", {"ag": 1})
     test_data["retries"] = args.retries
     test_data["timeout_in_minutes"] *= args.retries + 1
-    test_data.update(args.step_param)
+    test_data = overlay_dict(test_data, args.step_param)
     test_data["retry"] = {
         "automatic": [
             # Agent was lost, retry one time

.buildkite/pipeline_pr.py

@@ -7,7 +7,7 @@
 import subprocess
 from pathlib import Path
 
-from common import COMMON_PARSER, group, pipeline_to_json
+from common import COMMON_PARSER, group, overlay_dict, pipeline_to_json
 
 # Buildkite default job priority is 0. Setting this to 1 prioritizes PRs over
 # scheduled jobs and other batch jobs.
@@ -36,10 +36,9 @@ def get_changed_files(branch):
     # buildkite step parameters
     "priority": DEFAULT_PRIORITY,
     "timeout_in_minutes": 45,
-    "env": dict(args.step_env),
     "artifacts": ["./test_results/**/*"],
 }
-defaults.update(args.step_param)
+defaults = overlay_dict(defaults, args.step_param)
 
 build_grp = group(
     "📦 Build",
@@ -65,12 +64,14 @@ def get_changed_files(branch):
     **defaults,
 )
 
-defaults_for_performance = defaults.copy()
-defaults_for_performance.update(
-    # We specify higher priority so the ag=1 jobs get picked up before the ag=n
-    # jobs in ag=1 agents
-    priority=DEFAULT_PRIORITY + 1,
-    agent_tags=["ag=1"],
+defaults_for_performance = overlay_dict(
+    defaults,
+    {
+        # We specify higher priority so the ag=1 jobs get picked up before the ag=n
+        # jobs in ag=1 agents
+        "priority": DEFAULT_PRIORITY + 1,
+        "agents": {"ag": 1},
+    },
 )
 
 performance_grp = group(
@@ -79,13 +80,14 @@ def get_changed_files(branch):
     **defaults_for_performance,
 )
 
-defaults_for_kani = defaults.copy()
-defaults_for_kani.update(
-    # Kani runs fastest on m6i.metal
-    instances=["m6i.metal"],
-    platforms=[("al2", "linux_5.10")],
-    timeout_in_minutes=300,
-    agent_tags=["ag=1"],
+defaults_for_kani = overlay_dict(
+    defaults_for_performance,
+    {
+        # Kani runs fastest on m6i.metal
+        "instances": ["m6i.metal"],
+        "platforms": [("al2", "linux_5.10")],
+        "timeout_in_minutes": 300,
+    },
 )
 
 kani_grp = group(

.buildkite/pipeline_pr_no_block.py

@@ -4,7 +4,7 @@
 
 """Generate Buildkite pipelines dynamically"""
 
-from common import COMMON_PARSER, group, pipeline_to_json
+from common import COMMON_PARSER, group, overlay_dict, pipeline_to_json
 
 # Buildkite default job priority is 0. Setting this to 1 prioritizes PRs over
 # scheduled jobs and other batch jobs.
@@ -17,12 +17,12 @@
     "platforms": args.platforms,
     # buildkite step parameters
     "timeout_in_minutes": 45,
-    "env": dict(args.step_env),
     # some non-blocking tests are performance, so make sure they get ag=1 instances
     "priority": DEFAULT_PRIORITY + 1,
-    "agent_tags": ["ag=1"],
+    "agents": {"ag": 1},
 }
-defaults.update(args.step_param)
+defaults = overlay_dict(defaults, args.step_param)
+
 
 optional_grp = group(
     "❓ Optional",