Simplify A/B-Test orchestration

.buildkite/common.py

@@ -180,16 +180,7 @@ def random_str(k: int):
 
 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 -",
-        "git branch -D $$branch_name",
-    ]
+    return [f"./tools/devtool -y build --rev {revision} --release"]
 
 
 def shared_build():

.buildkite/pipeline_perf.py

@@ -98,7 +98,9 @@
     pytest_opts = ""
     if REVISION_A:
         devtool_opts += " --ab"
-        pytest_opts = f"{ab_opts} run {REVISION_A} {REVISION_B} --test {test_path}"
+        pytest_opts = (
+            f"{ab_opts} run build/{REVISION_A}/ build/{REVISION_B} --test {test_path}"
+        )
     else:
         # Passing `-m ''` below instructs pytest to collect tests regardless of
         # their markers (e.g. it will collect both tests marked as nonci, and

tests/README.md

@@ -169,16 +169,21 @@ function to [`.buildkite/pipeline_perf.py`](../.buildkite/pipeline_perf.py). To
 manually run an A/B-Test, use
 
 ```sh
-tools/devtool -y test --ab [optional arguments to ab_test.py] run <revision A> <revision B> --test <test specification>
+tools/devtool -y test --ab [optional arguments to ab_test.py] run <dir A> <dir B> --test <test specification>
 ```
 
-Here, _revision A_ and _revision B_ can be arbitrary git objects, such as commit
-SHAs, branches or tags. For example, to compare boottime of microVMs between
-Firecracker binaries compiled from the `main` branch and the `HEAD` of your
-current branch, run
+Here, _dir A_ and _dir B_ are directories containing firecracker and jailer
+binaries whose performance characteristics you wish to compare. You can use
+`./tools/devtool build --rev <revision> --release` to compile binaries from an
+arbitrary git object (commit SHAs, branches, tags etc.). This will create
+sub-directories in `build` containing the binaries. For example, to compare
+boottime of microVMs between Firecracker binaries compiled from the `main`
+branch and the `HEAD` of your current branch, run
 
 ```sh
-tools/devtool -y test --ab run main HEAD --test integration_tests/performance/test_boottime.py::test_boottime
+tools/devtool -y build --rev main --release
+tools/devtool -y build --rev HEAD --release
+tools/devtool -y test --ab -- run build/main build/HEAD --test integration_tests/performance/test_boottime.py::test_boottime
 ```
 
 #### How to Write an A/B-Compatible Test and Common Pitfalls

tests/conftest.py

@@ -99,19 +99,7 @@ def record_props(request, record_property):
 
     # Extract attributes from the docstrings
     function_docstring = inspect.getdoc(request.function)
-    description = []
-    attributes = {}
-    for line in function_docstring.split("\n"):
-        # extract tags like @type, @issue, etc
-        match = re.match(r"\s*@(?P<attr>\w+):\s*(?P<value>\w+)", line)
-        if match:
-            attr, value = match["attr"], match["value"]
-            attributes[attr] = value
-        else:
-            description.append(line)
-    for attr_name, attr_value in attributes.items():
-        record_property(attr_name, attr_value)
-    record_property("description", "".join(description))
+    record_property("description", function_docstring)
 
 
 def pytest_runtest_logreport(report):

tests/framework/ab_test.py

@@ -21,19 +21,20 @@
 of both invocations is the same, the test passes (with us being alerted to this situtation via a special pipeline that
 does not block PRs). If not, it fails, preventing PRs from introducing new vulnerable dependencies.
 """
-import contextlib
 import os
 import statistics
 from pathlib import Path
+from tempfile import TemporaryDirectory
 from typing import Callable, List, Optional, TypeVar
 
 import scipy
 
 from framework import utils
+from framework.defs import FC_WORKSPACE_DIR
 from framework.microvm import Microvm
 from framework.utils import CommandReturn
 from framework.with_filelock import with_filelock
-from host_tools.cargo_build import get_binary, get_firecracker_binaries
+from host_tools.cargo_build import DEFAULT_TARGET_DIR, get_firecracker_binaries
 
 # Locally, this will always compare against main, even if we try to merge into, say, a feature branch.
 # We might want to do a more sophisticated way to determine a "parent" branch here.
@@ -82,19 +83,41 @@ def git_ab_test(
              (alternatively, your comparator can perform any required assertions and not return anything).
     """
 
-    dir_a = git_clone(Path("../build") / a_revision, a_revision)
-    result_a = test_runner(dir_a, True)
+    with TemporaryDirectory() as tmp_dir:
+        dir_a = git_clone(Path(tmp_dir) / a_revision, a_revision)
+        result_a = test_runner(dir_a, True)
 
-    if b_revision:
-        dir_b = git_clone(Path("../build") / b_revision, b_revision)
-    else:
-        # By default, pytest execution happens inside the `tests` subdirectory. Pass the repository root, as
-        # documented.
-        dir_b = Path.cwd().parent
-    result_b = test_runner(dir_b, False)
+        if b_revision:
+            dir_b = git_clone(Path(tmp_dir) / b_revision, b_revision)
+        else:
+            # By default, pytest execution happens inside the `tests` subdirectory. Pass the repository root, as
+            # documented.
+            dir_b = Path.cwd().parent
+        result_b = test_runner(dir_b, False)
 
-    comparison = comparator(result_a, result_b)
-    return result_a, result_b, comparison
+        comparison = comparator(result_a, result_b)
+        return result_a, result_b, comparison
+
+
+DEFAULT_A_DIRECTORY = FC_WORKSPACE_DIR / "build" / "main"
+DEFAULT_B_DIRECTORY = FC_WORKSPACE_DIR / "build" / "cargo_target" / DEFAULT_TARGET_DIR
+
+
+def binary_ab_test(
+    test_runner: Callable[[Path, bool], T],
+    comparator: Callable[[T, T], U] = default_comparator,
+    *,
+    a_directory: Path = DEFAULT_A_DIRECTORY,
+    b_directory: Path = DEFAULT_B_DIRECTORY,
+):
+    """
+    Similar to `git_ab_test`, but instead of locally checking out different revisions, it operates on
+    directories containing firecracker/jailer binaries
+    """
+    result_a = test_runner(a_directory, True)
+    result_b = test_runner(b_directory, False)
+
+    return result_a, result_b, comparator(result_a, result_b)
 
 
 def is_pr() -> bool:
@@ -153,7 +176,7 @@ def set_did_not_grow_comparator(
     )
 
 
-def git_ab_test_guest_command(
+def precompiled_ab_test_guest_command(
     microvm_factory: Callable[[Path, Path], Microvm],
     command: str,
     *,
@@ -163,30 +186,29 @@ def git_ab_test_guest_command(
 ):
     """The same as git_ab_test_command, but via SSH. The closure argument should setup a microvm using the passed
     paths to firecracker and jailer binaries."""
+    b_directory = (
+        DEFAULT_B_DIRECTORY
+        if b_revision is None
+        else FC_WORKSPACE_DIR / "build" / b_revision
+    )
 
-    @with_filelock
-    def build_firecracker(workspace_dir):
-        utils.check_output("./tools/release.sh --profile release", cwd=workspace_dir)
-
-    def test_runner(workspace_dir, _is_a: bool):
-        firecracker = get_binary("firecracker", workspace_dir=workspace_dir)
-        if not firecracker.exists():
-            build_firecracker(workspace_dir)
-        bin_dir = firecracker.parent.resolve()
-        firecracker, jailer = bin_dir / "firecracker", bin_dir / "jailer"
-        microvm = microvm_factory(firecracker, jailer)
+    def test_runner(bin_dir, _is_a: bool):
+        microvm = microvm_factory(bin_dir / "firecracker", bin_dir / "jailer")
         return microvm.ssh.run(command)
 
-    (_, old_out, old_err), (_, new_out, new_err), the_same = git_ab_test(
-        test_runner, comparator, a_revision=a_revision, b_revision=b_revision
+    (_, old_out, old_err), (_, new_out, new_err), the_same = binary_ab_test(
+        test_runner,
+        comparator,
+        a_directory=FC_WORKSPACE_DIR / "build" / a_revision,
+        b_directory=b_directory,
     )
 
     assert (
         the_same
     ), f"The output of running command `{command}` changed:\nOld:\nstdout:\n{old_out}\nstderr\n{old_err}\n\nNew:\nstdout:\n{new_out}\nstderr:\n{new_err}"
 
 
-def git_ab_test_guest_command_if_pr(
+def precompiled_ab_test_guest_command_if_pr(
     microvm_factory: Callable[[Path, Path], Microvm],
     command: str,
     *,
@@ -195,7 +217,9 @@ def git_ab_test_guest_command_if_pr(
 ):
     """The same as git_ab_test_command_if_pr, but via SSH"""
     if is_pr():
-        git_ab_test_guest_command(microvm_factory, command, comparator=comparator)
+        precompiled_ab_test_guest_command(
+            microvm_factory, command, comparator=comparator
+        )
         return None
 
     microvm = microvm_factory(*get_firecracker_binaries())
@@ -249,17 +273,3 @@ def git_clone(clone_path, commitish):
         )
         utils.check_output(f"git branch -D {branch_name}")
     return clone_path
-
-
-# Once we upgrade to python 3.11, this will be in contextlib:
-# https://docs.python.org/3/library/contextlib.html#contextlib.chdir
-@contextlib.contextmanager
-def chdir(to):
-    """Context manager that temporarily `chdir`s to the specified path"""
-    cur = os.getcwd()
-
-    try:
-        os.chdir(to)
-        yield
-    finally:
-        os.chdir(cur)

tests/framework/properties.py

@@ -48,11 +48,11 @@ def get_host_os(kv: str = None):
         kv = platform.release()
     parts = kv.split("-")
     if len(parts) < 2:
-        return None
+        return kv
     misc = parts[1].split(".")
     if len(misc) > 2 and misc[2] in {"amzn2", "amzn2023"}:
         return misc[2]
-    return None
+    return kv
 
 
 class GlobalProps:

tests/framework/utils.py

@@ -234,17 +234,14 @@ def _cpus(cls):
         #  - cgroupsv1: /cpuset.cpus
         #  - cgroupsv2: /cpuset.cpus.effective
         # For more details, see https://docs.kernel.org/admin-guide/cgroup-v2.html#cpuset-interface-files
-        cpulist = None
         for path in [
             Path("/sys/fs/cgroup/cpuset/cpuset.cpus"),
             Path("/sys/fs/cgroup/cpuset.cpus.effective"),
         ]:
             if path.exists():
-                cpulist = path.read_text("ascii").strip()
-                break
-        else:
-            raise RuntimeError("Could not find cgroups cpuset")
-        return ListFormatParser(cpulist).parse()
+                return ListFormatParser(path.read_text("ascii").strip()).parse()
+
+        raise RuntimeError("Could not find cgroups cpuset")
 
 
 class ListFormatParser:

tests/integration_tests/performance/test_benchmarks.py

@@ -1,6 +1,7 @@
 # Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 """Optional benchmarks-do-not-regress test"""
+import contextlib
 import json
 import logging
 import platform
@@ -10,7 +11,7 @@
 import pytest
 
 from framework import utils
-from framework.ab_test import chdir, git_ab_test
+from framework.ab_test import git_ab_test
 from host_tools.cargo_build import cargo
 
 LOGGER = logging.getLogger(__name__)
@@ -32,7 +33,7 @@ def run_criterion(firecracker_checkout: Path, is_a: bool) -> Path:
     """
     baseline_name = "a_baseline" if is_a else "b_baseline"
 
-    with chdir(firecracker_checkout):
+    with contextlib.chdir(firecracker_checkout):
         # Passing --message-format json to cargo tells it to print its log in a json format. At the end, instead of the
         # usual "placed executable <...> at <...>" we'll get a json object with an 'executable' key, from which we
         # extract the path to the compiled benchmark binary.

tests/integration_tests/security/test_jail.py

@@ -64,19 +64,14 @@ def test_empty_jailer_id(uvm_plain):
         exec_file=test_microvm.fc_binary_path,
     )
 
-    # pylint: disable=W0703
-    try:
+    # If the exception is not thrown, it means that Firecracker was
+    # started successfully, hence there's a bug in the code due to which
+    # we can set an empty ID.
+    with pytest.raises(
+        ChildProcessError,
+        match=r"Jailer error: Invalid instance ID: Invalid len \(0\);  the length must be between 1 and 64",
+    ):
         test_microvm.spawn()
-        # If the exception is not thrown, it means that Firecracker was
-        # started successfully, hence there's a bug in the code due to which
-        # we can set an empty ID.
-        assert False
-    except Exception as err:
-        expected_err = (
-            "Jailer error: Invalid instance ID: Invalid len (0);"
-            "  the length must be between 1 and 64"
-        )
-        assert expected_err in str(err)
 
 
 def test_exec_file_not_exist(uvm_plain, tmp_path):