firecracker

1
# Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2
# SPDX-License-Identifier: Apache-2.0
3

4
"""
5
Common helpers to create Buildkite pipelines
6
"""
7

8
import argparse
9
import json
10
import subprocess
11
from pathlib import Path
12

13
DEFAULT_INSTANCES = [
14
    "c5n.metal",  # Intel Skylake
15
    "m5n.metal",  # Intel Cascade Lake
16
    "m6i.metal",  # Intel Icelake
17
    "m6a.metal",  # AMD Milan
18
    "m6g.metal",  # Graviton2
19
    "m7g.metal",  # Graviton3
20
]
21

22
DEFAULT_PLATFORMS = [
23
    ("al2", "linux_4.14"),
24
    ("al2", "linux_5.10"),
25
    ("al2023", "linux_6.1"),
26
]
27

28

29
def overlay_dict(base: dict, update: dict):
30
    """Overlay a dict over a base one"""
31
    base = base.copy()
32
    for key, val in update.items():
33
        if key in base and isinstance(val, dict):
34
            base[key] = overlay_dict(base.get(key, {}), val)
35
        else:
36
            base[key] = val
37
    return base
38

39

40
def field_fmt(field, args):
41
    """If `field` is a string, interpolate variables in `args`"""
42
    if not isinstance(field, str):
43
        return field
44
    return field.format(**args)
45

46

47
def dict_fmt(dict_tmpl, args):
48
    """Apply field_fmt over a hole dict"""
49
    res = {}
50
    for key, val in dict_tmpl.items():
51
        if isinstance(val, dict):
52
            res[key] = dict_fmt(val, args)
53
        else:
54
            res[key] = field_fmt(val, args)
55
    return res
56

57

58
def group(label, command, instances, platforms, **kwargs):
59
    """
60
    Generate a group step with specified parameters, for each instance+kernel
61
    combination
62

63
    https://buildkite.com/docs/pipelines/group-step
64
    """
65
    # Use the 1st character of the group name (should be an emoji)
66
    label1 = label[0]
67
    steps = []
68
    commands = command
69
    if isinstance(command, str):
70
        commands = [command]
71
    for instance in instances:
72
        for os, kv in platforms:
73
            # fill any templated variables
74
            args = {"instance": instance, "os": os, "kv": kv}
75
            step = {
76
                "command": [cmd.format(**args) for cmd in commands],
77
                "label": f"{label1} {instance} {os} {kv}",
78
                "agents": args,
79
            }
80
            step_kwargs = dict_fmt(kwargs, args)
81
            step = overlay_dict(step_kwargs, step)
82
            steps.append(step)
83

84
    return {"group": label, "steps": steps}
85

86

87
def pipeline_to_json(pipeline):
88
    """Serialize a pipeline dictionary to JSON"""
89
    return json.dumps(pipeline, indent=4, sort_keys=True, ensure_ascii=False)
90

91

92
def get_changed_files(branch):
93
    """
94
    Get all files changed since `branch`
95
    """
96
    stdout = subprocess.check_output(["git", "diff", "--name-only", branch])
97
    return [Path(line) for line in stdout.decode().splitlines()]
98

99

100
def run_all_tests(changed_files):
101
    """
102
    Check if we should run all tests, based on the files that have been changed
103
    """
104

105
    # run the whole test suite if either of:
106
    # - any file changed that is not documentation nor GitHub action config file
107
    # - no files changed
108
    return not changed_files or any(
109
        x.suffix != ".md" and not (x.parts[0] == ".github" and x.suffix == ".yml")
110
        for x in changed_files
111
    )
112

113

114
def devtool_test(devtool_opts=None, pytest_opts=None, binary_dir=None):
115
    """Generate a `devtool test` command"""
116
    cmds = []
117
    parts = ["./tools/devtool -y test"]
118
    if devtool_opts:
119
        parts.append(devtool_opts)
120
    parts.append("--")
121
    if binary_dir is not None:
122
        cmds.append(f'buildkite-agent artifact download "{binary_dir}/$(uname -m)/*" .')
123
        cmds.append(f"chmod -v a+x {binary_dir}/**/*")
124
        parts.append(f"--binary-dir=../{binary_dir}/$(uname -m)")
125
    if pytest_opts:
126
        parts.append(pytest_opts)
127
    cmds.append(" ".join(parts))
128
    return cmds
129

130

131
class DictAction(argparse.Action):
132
    """An argparse action that can receive a nested dictionary
133

134
    Examples:
135

136
        --step-param a/b/c=3
137
        {"a": {"b": {"c": 3}}}
138
    """
139

140
    def __init__(self, option_strings, dest, nargs=None, **kwargs):
141
        if nargs is not None:
142
            raise ValueError("nargs not allowed")
143
        super().__init__(option_strings, dest, **kwargs)
144

145
    def __call__(self, parser, namespace, value, option_string=None):
146
        res = getattr(namespace, self.dest, {})
147
        key_str, val = value.split("=", maxsplit=1)
148
        keys = key_str.split("/")
149
        update = {keys[-1]: val}
150
        for key in list(reversed(keys))[1:]:
151
            update = {key: update}
152
        res = overlay_dict(res, update)
153
        setattr(namespace, self.dest, res)
154

155

156
COMMON_PARSER = argparse.ArgumentParser()
157
COMMON_PARSER.add_argument(
158
    "--instances",
159
    required=False,
160
    nargs="+",
161
    default=DEFAULT_INSTANCES,
162
)
163
COMMON_PARSER.add_argument(
164
    "--platforms",
165
    metavar="OS-KV",
166
    required=False,
167
    nargs="+",
168
    default=DEFAULT_PLATFORMS,
169
    type=lambda arg: tuple(arg.split("-", maxsplit=1)),
170
)
171
COMMON_PARSER.add_argument(
172
    "--step-param",
173
    metavar="PARAM=VALUE",
174
    help="parameters to add to each step",
175
    required=False,
176
    action=DictAction,
177
    default={},
178
    type=str,
179
)
180
COMMON_PARSER.add_argument(
181
    "--binary-dir",
182
    help="Use the Firecracker binaries from this path",
183
    required=False,
184
    default=None,
185
    type=str,
186
)
187