Merge branch 'firecracker-microvm:main' into rustify-snapshot-module

Author: beamandala

.buildkite/common.py

@@ -8,6 +8,8 @@
 import argparse
 import json
 import os
+import random
+import string
 import subprocess
 from pathlib import Path
 
@@ -21,7 +23,6 @@
 ]
 
 DEFAULT_PLATFORMS = [
-    ("al2", "linux_4.14"),
     ("al2", "linux_5.10"),
     ("al2023", "linux_6.1"),
 ]
@@ -70,12 +71,12 @@ def group(label, command, instances, platforms, **kwargs):
     if isinstance(command, str):
         commands = [command]
     for instance in instances:
-        for os, kv in platforms:
+        for os_, kv in platforms:
             # fill any templated variables
-            args = {"instance": instance, "os": os, "kv": kv}
+            args = {"instance": instance, "os": os_, "kv": kv}
             step = {
                 "command": [cmd.format(**args) for cmd in commands],
-                "label": f"{label1} {instance} {os} {kv}",
+                "label": f"{label1} {instance} {os_} {kv}",
                 "agents": args,
             }
             step_kwargs = dict_fmt(kwargs, args)
@@ -85,11 +86,6 @@ def group(label, command, instances, platforms, **kwargs):
     return {"group": label, "steps": steps}
 
 
-def pipeline_to_json(pipeline):
-    """Serialize a pipeline dictionary to JSON"""
-    return json.dumps(pipeline, indent=4, sort_keys=True, ensure_ascii=False)
-
-
 def get_changed_files():
     """
     Get all files changed since `branch`
@@ -119,23 +115,6 @@ def run_all_tests(changed_files):
     )
 
 
-def devtool_test(devtool_opts=None, pytest_opts=None, binary_dir=None):
-    """Generate a `devtool test` command"""
-    cmds = []
-    parts = ["./tools/devtool -y test"]
-    if devtool_opts:
-        parts.append(devtool_opts)
-    parts.append("--")
-    if binary_dir is not None:
-        cmds.append(f'buildkite-agent artifact download "{binary_dir}/$(uname -m)/*" .')
-        cmds.append(f"chmod -v a+x {binary_dir}/**/*")
-        parts.append(f"--binary-dir=../{binary_dir}/$(uname -m)")
-    if pytest_opts:
-        parts.append(pytest_opts)
-    cmds.append(" ".join(parts))
-    return cmds
-
-
 class DictAction(argparse.Action):
     """An argparse action that can receive a nested dictionary
 
@@ -192,3 +171,169 @@ def __call__(self, parser, namespace, value, option_string=None):
     default=None,
     type=str,
 )
+
+
+def random_str(k: int):
+    """Generate a random string of hex characters."""
+    return "".join(random.choices(string.hexdigits, k=k))
+
+
+def ab_revision_build(revision):
+    """Generate steps for building an A/B-test revision"""
+    # Copied from framework/ab_test. Double dollar signs needed for Buildkite (otherwise it will try to interpolate itself)
+    return [
+        f"commitish={revision}",
+        f"if ! git cat-file -t $$commitish; then commitish=origin/{revision}; fi",
+        "branch_name=tmp-$$commitish",
+        "git branch $$branch_name $$commitish",
+        f"git clone -b $$branch_name . build/{revision}",
+        f"cd build/{revision} && ./tools/devtool -y build --release && cd -",
+    ]
+
+
+def shared_build():
+    """Helper function to make it simple to share a compilation artifacts for a
+    whole Buildkite build
+    """
+
+    # We need to support 3 scenarios here:
+    # 1. We are running in the nightly pipeline - only compile the HEAD of main.
+    # 2. We are running in a PR pipeline - compile HEAD of main as revision A and HEAD of PR branch as revision B.
+    # 3. We are running in an A/B-test pipeline - compile what is passed via REVISION_{A,B} environment variables.
+    rev_a = os.environ.get("REVISION_A")
+    if rev_a is not None:
+        rev_b = os.environ.get("REVISION_B")
+        assert rev_b is not None, "REVISION_B environment variable not set"
+        build_cmds = ab_revision_build(rev_a) + ab_revision_build(rev_b)
+    elif os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false":
+        build_cmds = ab_revision_build(
+            os.environ.get("BUILDKITE_PULL_REQUEST_BASE_BRANCH", "main")
+        ) + ["./tools/devtool -y build --release"]
+    else:
+        build_cmds = ["./tools/devtool -y build --release"]
+    binary_dir = f"build_$(uname -m)_{random_str(k=8)}.tar.gz"
+    build_cmds += [
+        "du -sh build/*",
+        f"tar czf {binary_dir} build",
+        f"buildkite-agent artifact upload {binary_dir}",
+    ]
+    return build_cmds, binary_dir
+
+
+class BKPipeline:
+    """
+    Buildkite Pipeline class abstraction
+
+    Helper class to easily construct pipelines.
+    """
+
+    parser = COMMON_PARSER
+
+    def __init__(self, initial_steps=None, **kwargs):
+        self.steps = []
+        self.args = args = self.parser.parse_args()
+        # Retry one time if agent was lost. This can happen if we terminate the
+        # instance or the agent gets disconnected for whatever reason
+        retry = {
+            "automatic": [{"exit_status": -1, "limit": 1}],
+        }
+        retry = overlay_dict(retry, kwargs.pop("retry", {}))
+        # Calculate step defaults with parameters and kwargs
+        per_instance = {
+            "instances": args.instances,
+            "platforms": args.platforms,
+            "artifact_paths": ["./test_results/**/*"],
+            "retry": retry,
+            **kwargs,
+        }
+        self.per_instance = overlay_dict(per_instance, args.step_param)
+        self.per_arch = self.per_instance.copy()
+        self.per_arch["instances"] = ["m6i.metal", "m7g.metal"]
+        self.per_arch["platforms"] = [("al2", "linux_5.10")]
+        self.binary_dir = args.binary_dir
+        # Build sharing
+        build_cmds, self.shared_build = shared_build()
+        step_build = group("🏗️ Build", build_cmds, **self.per_arch)
+        self.steps += [step_build, "wait"]
+
+        # If we run initial_steps before the "wait" step above, then a failure of the initial steps
+        # would result in the build not progressing past the "wait" step (as buildkite only proceeds past a wait step
+        # if everything before it passed). Thus put the initial steps after the "wait" step, but set `"depends_on": null`
+        # to start running them immediately (e.g. without waiting for the "wait" step to unblock).
+        #
+        # See also https://buildkite.com/docs/pipelines/dependencies#explicit-dependencies-in-uploaded-steps
+        if initial_steps:
+            for step in initial_steps:
+                step["depends_on"] = None
+
+            self.steps += initial_steps
+
+    def add_step(self, step, decorate=True):
+        """
+        Add a step to the pipeline.
+
+        https://buildkite.com/docs/pipelines/step-reference
+
+        :param step: a Buildkite step
+        :param decorate: inject needed commands for sharing builds
+        """
+        if decorate and isinstance(step, dict):
+            step = self._adapt_group(step)
+        self.steps.append(step)
+        return step
+
+    def _adapt_group(self, group):
+        """"""
+        prepend = [
+            f'buildkite-agent artifact download "{self.shared_build}" .',
+            f"tar xzf {self.shared_build}",
+        ]
+        if self.binary_dir is not None:
+            prepend.extend(
+                [
+                    f'buildkite-agent artifact download "{self.binary_dir}/$(uname -m)/*" .',
+                    f"chmod -v a+x {self.binary_dir}/**/*",
+                ]
+            )
+
+        for step in group["steps"]:
+            step["command"] = prepend + step["command"]
+        return group
+
+    def build_group(self, *args, **kwargs):
+        """
+        Build a group, parametrizing over the selected instances/platforms.
+
+        https://buildkite.com/docs/pipelines/group-step
+        """
+        combined = overlay_dict(self.per_instance, kwargs)
+        return self.add_step(group(*args, **combined))
+
+    def build_group_per_arch(self, *args, **kwargs):
+        """
+        Build a group, parametrizing over the architectures only.
+        """
+        combined = overlay_dict(self.per_arch, kwargs)
+        return self.add_step(group(*args, **combined))
+
+    def to_dict(self):
+        """Render the pipeline as a dictionary."""
+        return {"steps": self.steps}
+
+    def to_json(self):
+        """Serialize the pipeline to JSON"""
+        return json.dumps(self.to_dict(), indent=4, sort_keys=True, ensure_ascii=False)
+
+    def devtool_test(self, devtool_opts=None, pytest_opts=None):
+        """Generate a `devtool test` command"""
+        cmds = []
+        parts = ["./tools/devtool -y test", "--no-build"]
+        if devtool_opts:
+            parts.append(devtool_opts)
+        parts.append("--")
+        if self.binary_dir is not None:
+            parts.append(f"--binary-dir=../{self.binary_dir}/$(uname -m)")
+        if pytest_opts:
+            parts.append(pytest_opts)
+        cmds.append(" ".join(parts))
+        return cmds

.buildkite/pipeline_cpu_template.py

@@ -4,10 +4,9 @@
 
 """Generate Buildkite CPU Template pipelines dynamically"""
 
-import argparse
 from enum import Enum
 
-from common import DEFAULT_INSTANCES, DEFAULT_PLATFORMS, group, pipeline_to_json
+from common import DEFAULT_INSTANCES, DEFAULT_PLATFORMS, BKPipeline, group
 
 
 class BkStep(str, Enum):
@@ -17,23 +16,31 @@ class BkStep(str, Enum):
 
     LABEL = "label"
     TIMEOUT = "timeout"
-    COMMAND = "commands"
+    COMMAND = "command"
     ARTIFACTS = "artifact_paths"
 
 
 cpu_template_test = {
     "rdmsr": {
         BkStep.COMMAND: [
-            "tools/devtool -y test -- -s -ra -m nonci -n4 --log-cli-level=INFO integration_tests/functional/test_cpu_features.py -k 'test_cpu_rdmsr' "
+            "tools/devtool -y test --no-build -- -s -ra -m nonci -n4 --log-cli-level=INFO integration_tests/functional/test_cpu_features.py -k 'test_cpu_rdmsr' "
         ],
         BkStep.LABEL: "📖 rdmsr",
         "instances": ["c5n.metal", "m5n.metal", "m6a.metal", "m6i.metal"],
         "platforms": DEFAULT_PLATFORMS,
     },
+    "fingerprint": {
+        BkStep.COMMAND: [
+            "tools/devtool -y test --no-build -- -m no_block_pr integration_tests/functional/test_cpu_template_helper.py -k test_guest_cpu_config_change",
+        ],
+        BkStep.LABEL: "🖐️ fingerprint",
+        "instances": DEFAULT_INSTANCES,
+        "platforms": DEFAULT_PLATFORMS,
+    },
     "cpuid_wrmsr": {
         "snapshot": {
             BkStep.COMMAND: [
-                "tools/devtool -y test -- -s -ra -m nonci -n4 --log-cli-level=INFO integration_tests/functional/test_cpu_features.py -k 'test_cpu_wrmsr_snapshot or test_cpu_cpuid_snapshot'",
+                "tools/devtool -y test --no-build -- -s -ra -m nonci -n4 --log-cli-level=INFO integration_tests/functional/test_cpu_features.py -k 'test_cpu_wrmsr_snapshot or test_cpu_cpuid_snapshot'",
                 "mkdir -pv tests/snapshot_artifacts_upload/{instance}_{os}_{kv}",
                 "sudo mv tests/snapshot_artifacts/* tests/snapshot_artifacts_upload/{instance}_{os}_{kv}",
             ],
@@ -45,7 +52,7 @@ class BkStep(str, Enum):
             BkStep.COMMAND: [
                 "buildkite-agent artifact download tests/snapshot_artifacts_upload/{instance}_{os}_{kv}/**/* .",
                 "mv tests/snapshot_artifacts_upload/{instance}_{os}_{kv} tests/snapshot_artifacts",
-                "tools/devtool -y test -- -s -ra -m nonci -n4 --log-cli-level=INFO integration_tests/functional/test_cpu_features.py -k 'test_cpu_wrmsr_restore or test_cpu_cpuid_restore'",
+                "tools/devtool -y test --no-build -- -s -ra -m nonci -n4 --log-cli-level=INFO integration_tests/functional/test_cpu_features.py -k 'test_cpu_wrmsr_restore or test_cpu_cpuid_restore'",
             ],
             BkStep.LABEL: "📸 load snapshot artifacts created on {instance} {snapshot_os} {snapshot_kv} to {restore_instance} {restore_os} {restore_kv}",
             BkStep.TIMEOUT: 30,
@@ -57,33 +64,9 @@ class BkStep(str, Enum):
         },
         "instances": ["c5n.metal", "m5n.metal", "m6i.metal", "m6a.metal"],
     },
-    "fingerprint": {
-        BkStep.COMMAND: [
-            "tools/devtool -y test -- -m no_block_pr integration_tests/functional/test_cpu_template_helper.py -k test_guest_cpu_config_change",
-        ],
-        BkStep.LABEL: "🖐️ fingerprint",
-        "instances": DEFAULT_INSTANCES,
-        "platforms": DEFAULT_PLATFORMS,
-    },
 }
 
 
-def group_single(tests):
-    """
-    Generate a group step with specified parameters for each instance
-    and kernel combination
-    https://buildkite.com/docs/pipelines/group-step
-    """
-    group_step = group(
-        label=tests[BkStep.LABEL],
-        command=tests[BkStep.COMMAND],
-        instances=tests["instances"],
-        platforms=tests["platforms"],
-        artifacts=["./test_results/**/*"],
-    )
-    return [group_step]
-
-
 def group_snapshot_restore(test_step):
     """
     Generate a group step with specified parameters for each instance
@@ -146,30 +129,20 @@ def group_snapshot_restore(test_step):
     return groups
 
 
-def main():
-    """
-    Generate group template required to trigger pipelines for
-    the requested CPU template test.
-    """
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
+if __name__ == "__main__":
+    BKPipeline.parser.add_argument(
         "--test",
-        required=True,
         choices=list(cpu_template_test),
         help="CPU template test",
+        action="append",
     )
-    test_args = parser.parse_args()
-
-    if test_args.test == "rdmsr":
-        test_group = group_single(cpu_template_test[test_args.test])
-    elif test_args.test == "cpuid_wrmsr":
-        test_group = group_snapshot_restore(cpu_template_test[test_args.test])
-    elif test_args.test == "fingerprint":
-        test_group = group_single(cpu_template_test[test_args.test])
-
-    pipeline = {"steps": test_group}
-    print(pipeline_to_json(pipeline))
-
-
-if __name__ == "__main__":
-    main()
+    pipeline = BKPipeline()
+    for test in pipeline.args.test or list(cpu_template_test):
+        if test == "cpuid_wrmsr":
+            groups = group_snapshot_restore(cpu_template_test[test])
+            for grp in groups:
+                pipeline.add_step(grp)
+        else:
+            test_data = cpu_template_test[test]
+            pipeline.build_group(**test_data, artifacts=["./test_results/**/*"])
+    print(pipeline.to_json())