test_benchmarks.py

# Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
"""Optional benchmarks-do-not-regress test"""
importcontextlib
importjson
importlogging
importplatform
importshutil
frompathlibimportPath

importpytest

fromframeworkimportutils
fromframework.ab_testimportgit_ab_test
fromhost_tools.cargo_buildimportcargo

LOGGER=logging.getLogger(__name__)


@pytest.mark.no_block_pr
@pytest.mark.timeout(900)
deftest_no_regression_relative_to_target_branch():
"""
 Run the microbenchmarks in this repository, comparing results from pull
 request target branch against what's achieved on HEAD
 """
git_ab_test(run_criterion, compare_results)


defrun_criterion(firecracker_checkout: Path, is_a: bool) ->Path:
"""
 Executes all benchmarks by running "cargo bench --no-run", finding the executables, and running them pinned to some CPU
 """
baseline_name="a_baseline"ifis_aelse"b_baseline"

withcontextlib.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.
_, stdout, _=cargo(
"bench",
f"--all --quiet --target {platform.machine()}-unknown-linux-musl --message-format json --no-run",
 )

executables= []
forlineinstdout.split("\n"):
ifline:
msg=json.loads(line)
executable=msg.get("executable")
ifexecutable:
executables.append(executable)

forexecutableinexecutables:
utils.check_output(
f"CARGO_TARGET_DIR=build/cargo_target taskset -c 1 {executable} --bench --save-baseline {baseline_name}"
 )

returnfirecracker_checkout/"build"/"cargo_target"/"criterion"


defcompare_results(location_a_baselines: Path, location_b_baselines: Path):
"""Compares the two recorded criterion baselines for regressions, assuming that "A" is the baseline from main"""
forbenchmarkinlocation_b_baselines.glob("*"):
data=json.loads(
 (benchmark/"b_baseline"/"estimates.json").read_text("utf-8")
 )

average_ns=data["mean"]["point_estimate"]

LOGGER.info("%s mean: %iµs", benchmark.name, average_ns/1000)

# Assumption: location_b_baseline = cargo_target of current working directory. So just copy the a_baselines here
# to do the comparison
forbenchmarkinlocation_a_baselines.glob("*"):
shutil.copytree(
benchmark/"a_baseline",
location_b_baselines/benchmark.name/"a_baseline",
 )

_, stdout, _=cargo(
"bench",
f"--all --target {platform.machine()}-unknown-linux-musl",
"--baseline a_baseline --load-baseline b_baseline",
 )

regressions_only="\n\n".join(
result
forresultinstdout.split("\n\n")
if"Performance has regressed."inresult
 )

# If this string is anywhere in stdout, then at least one of our benchmarks
# is now performing worse with the PR changes.
assertnotregressions_only, "\n"+regressions_only