test: Enable to run CI with custom CPU templates for testing / debugging purpose

.buildkite/common.py

@@ -176,6 +176,22 @@ def __call__(self, parser, namespace, value, option_string=None):
     default=None,
     type=str,
 )
+COMMON_PARSER.add_argument(
+    "--additional-prepend",
+    help="Commands to be prepended additionally",
+    required=False,
+    nargs="+",
+    default=[],
+    type=str,
+)
+COMMON_PARSER.add_argument(
+    "-k",
+    "--keywords",
+    help="Keywords to filter pytest tests to run",
+    required=False,
+    default=None,
+    type=str,
+)
 
 
 def random_str(k: int):
@@ -288,6 +304,7 @@ def _adapt_group(self, group):
                     f"chmod -v a+x {self.binary_dir}/**/*",
                 ]
             )
+        prepend.extend(self.args.additional_prepend)
 
         for step in group["steps"]:
             step["command"] = prepend + step["command"]
@@ -351,6 +368,8 @@ def devtool_test(self, devtool_opts=None, pytest_opts=None):
         parts.append("--")
         if self.binary_dir is not None:
             parts.append(f"--binary-dir=../{self.binary_dir}/$(uname -m)")
+        if self.args.keywords is not None:
+            parts.append(f"-k {self.args.keywords}")
         if pytest_opts:
             parts.append(pytest_opts)
         cmds.append(" ".join(parts))

tests/README.md

@@ -306,10 +306,14 @@ that are pre-initialized with specific guest kernels and rootfs:
   24.04 squashfs as rootfs,
 - `uvm_plain` yields a Firecracker process pre-initialized with a 5.10 kernel
   and the same Ubuntu 24.04 squashfs.
+- `uvm_any` yields started microvms, parametrized by all supported kernels, all
+  CPU templates (static, custom and none), and either booted or restored from a
+  snapshot.
+- `uvm_any_booted` works the same as `uvm_any`, but only for booted VMs.
 
-Generally, tests should use the former if you are testing some interaction
-between the guest and Firecracker, while the latter should be used if
-Firecracker functionality unrelated to the guest is being tested.
+Generally, tests should use `uvm_plain_any` if you are testing some interaction
+between the guest and Firecracker, and `uvm_plain` should be used if Firecracker
+functionality unrelated to the guest is being tested.
 
 ### Markers
 
@@ -642,3 +646,49 @@ sudo env PYTHONPATH=tests HOME=$HOME ~/.local/bin/ipython3 -i tools/sandbox.py -
 > \[!WARNING\]
 >
 > **Notice this runs as root!**
+
+## How to run python tests with custom CPU templates
+
+By placing custom CPU templates under `tests/data/custom_cpu_templates/`
+directory, you can run the CI with them for testing / debugging purposes.
+
+Using the pytest keyword filtering option `-k`, you can run only python tests
+that use a specific CPU template. For example:
+
+```sh
+tools/devtool -y test -- integration_tests/functional -k unique_template_name
+```
+
+You can also do it from buildkite using the scripts under `.buildkite/`
+directory. The easiest way is to commit custom CPU template JSON files in
+question to your forked repo. Note that you should specify platforms on which
+the custom CPU templates are expected to work. For example:
+
+```yaml
+steps:
+- label: "Run test with custom CPU templates"
+  command: |
+    .buildkite/pipeline_pr.py \
+      --instances m6g.metal m7g.metal \
+      -k unique_template_name \
+      | buildkite-agent pipeline upload
+```
+
+Even without making any commit, you can inject the custom CPU template at
+runtime via `--additional-prepend` option of the buildkite step generation
+scripts. For example:
+
+```yaml
+steps:
+- label: "Run test with custom CPU templates"
+  command: |
+    .buildkite/pipeline_pr.py \
+      --instances m6g.metal m7g.metal \
+      -k unique_template_name \
+      --additional-prepend 'echo "{\"kvm_capabilities\": [\"170\", \"171\", \"172\"], \"vcpu_features\": [{\"index\": 0, \"bitmap\": \"0b111xxxx\"}]}" > tests/data/custom_cpu_templates/unique_template_name.json \
+      | buildkite-agent pipeline upload
+```
+
+In case that a CPU template written directly looks ugly or too lengthy, an
+alternative way is to download or copy it from somewhere at runtime also via the
+prepended command, although it is almost same as committing to your forked repo.

tests/conftest.py

@@ -35,7 +35,7 @@
 
 import host_tools.cargo_build as build_tools
 from framework import defs, utils
-from framework.artifacts import kernel_params, rootfs_params
+from framework.artifacts import disks, kernel_params
 from framework.microvm import MicroVMFactory
 from framework.properties import global_props
 from framework.utils_cpu_templates import (
@@ -292,13 +292,6 @@ def microvm_factory(request, record_property, results_dir):
     uvm_factory.kill()
 
 
-@pytest.fixture(params=static_cpu_templates_params())
-def cpu_template(request, record_property):
-    """Return all static CPU templates supported by the vendor."""
-    record_property("static_cpu_template", request.param)
-    return request.param
-
-
 @pytest.fixture(params=custom_cpu_templates_params())
 def custom_cpu_template(request, record_property):
     """Return all dummy custom CPU templates supported by the vendor."""
@@ -361,13 +354,6 @@ def guest_kernel_fxt(request, record_property):
     return kernel
 
 
-def rootfs_fxt(request, record_property):
-    """Return all supported rootfs."""
-    fs = request.param
-    record_property("rootfs", fs.name)
-    return fs
-
-
 # Fixtures for all guest kernels, and specific versions
 guest_kernel = pytest.fixture(guest_kernel_fxt, params=kernel_params("vmlinux-*"))
 guest_kernel_acpi = pytest.fixture(
@@ -387,9 +373,17 @@ def rootfs_fxt(request, record_property):
     params=kernel_params("vmlinux-6.1*"),
 )
 
-# Fixtures for all Ubuntu rootfs, and specific versions
-rootfs = pytest.fixture(rootfs_fxt, params=rootfs_params("ubuntu-24*.squashfs"))
-rootfs_rw = pytest.fixture(rootfs_fxt, params=rootfs_params("*.ext4"))
+
+@pytest.fixture
+def rootfs():
+    """Return an Ubuntu 24.04 read-only rootfs"""
+    return disks("ubuntu-24.04.squashfs")[0]
+
+
+@pytest.fixture
+def rootfs_rw():
+    """Return an Ubuntu 24.04 ext4 rootfs"""
+    return disks("ubuntu-24.04.ext4")[0]
 
 
 @pytest.fixture
@@ -459,3 +453,80 @@ def uvm_with_initrd(
     uvm = microvm_factory.build(guest_kernel_linux_5_10)
     uvm.initrd_file = fs
     yield uvm
+
+
+@pytest.fixture
+def vcpu_count():
+    """Return default vcpu_count. Use indirect parametrization to override."""
+    return 2
+
+
+@pytest.fixture
+def mem_size_mib():
+    """Return memory size. Use indirect parametrization to override."""
+    return 256
+
+
+def uvm_booted(
+    microvm_factory, guest_kernel, rootfs, cpu_template, vcpu_count=2, mem_size_mib=256
+):
+    """Return a booted uvm"""
+    uvm = microvm_factory.build(guest_kernel, rootfs)
+    uvm.spawn()
+    uvm.basic_config(vcpu_count=vcpu_count, mem_size_mib=mem_size_mib)
+    uvm.set_cpu_template(cpu_template)
+    uvm.add_net_iface()
+    uvm.start()
+    return uvm
+
+
+def uvm_restored(microvm_factory, guest_kernel, rootfs, cpu_template, **kwargs):
+    """Return a restored uvm"""
+    uvm = uvm_booted(microvm_factory, guest_kernel, rootfs, cpu_template, **kwargs)
+    snapshot = uvm.snapshot_full()
+    uvm.kill()
+    uvm2 = microvm_factory.build_from_snapshot(snapshot)
+    uvm2.cpu_template_name = uvm.cpu_template_name
+    return uvm2
+
+
+@pytest.fixture(params=[uvm_booted, uvm_restored])
+def uvm_ctor(request):
+    """Fixture to return uvms with different constructors"""
+    return request.param
+
+
+@pytest.fixture
+def uvm_any(
+    microvm_factory,
+    uvm_ctor,
+    guest_kernel,
+    rootfs,
+    cpu_template_any,
+    vcpu_count,
+    mem_size_mib,
+):
+    """Return booted and restored uvms"""
+    return uvm_ctor(
+        microvm_factory,
+        guest_kernel,
+        rootfs,
+        cpu_template_any,
+        vcpu_count=vcpu_count,
+        mem_size_mib=mem_size_mib,
+    )
+
+
+@pytest.fixture
+def uvm_any_booted(
+    microvm_factory, guest_kernel, rootfs, cpu_template_any, vcpu_count, mem_size_mib
+):
+    """Return booted uvms"""
+    return uvm_booted(
+        microvm_factory,
+        guest_kernel,
+        rootfs,
+        cpu_template_any,
+        vcpu_count=vcpu_count,
+        mem_size_mib=mem_size_mib,
+    )

tests/framework/ab_test.py

@@ -31,10 +31,9 @@
 
 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 DEFAULT_TARGET_DIR, get_firecracker_binaries
+from host_tools.cargo_build import DEFAULT_TARGET_DIR
 
 # 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.
@@ -176,56 +175,6 @@ def set_did_not_grow_comparator(
     )
 
 
-def precompiled_ab_test_guest_command(
-    microvm_factory: Callable[[Path, Path], Microvm],
-    command: str,
-    *,
-    comparator: Callable[[CommandReturn, CommandReturn], bool] = default_comparator,
-    a_revision: str = DEFAULT_A_REVISION,
-    b_revision: Optional[str] = None,
-):
-    """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
-    )
-
-    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 = 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 precompiled_ab_test_guest_command_if_pr(
-    microvm_factory: Callable[[Path, Path], Microvm],
-    command: str,
-    *,
-    comparator=default_comparator,
-    check_in_nonpr=True,
-):
-    """The same as git_ab_test_command_if_pr, but via SSH"""
-    if is_pr():
-        precompiled_ab_test_guest_command(
-            microvm_factory, command, comparator=comparator
-        )
-        return None
-
-    microvm = microvm_factory(*get_firecracker_binaries())
-    return microvm.ssh.run(command, check=check_in_nonpr)
-
-
 def check_regression(
     a_samples: List[float], b_samples: List[float], *, n_resamples: int = 9999
 ):

tests/framework/artifacts.py

@@ -44,9 +44,9 @@ def kernels(glob, artifact_dir: Path = ARTIFACT_DIR) -> Iterator:
                 break
 
 
-def disks(glob) -> Iterator:
+def disks(glob) -> list:
     """Return supported rootfs"""
-    yield from sorted(ARTIFACT_DIR.glob(glob))
+    return sorted(ARTIFACT_DIR.glob(glob))
 
 
 def kernel_params(
@@ -57,12 +57,6 @@ def kernel_params(
         yield pytest.param(kernel, id=kernel.name)
 
 
-def rootfs_params(glob="ubuntu-*.squashfs") -> Iterator:
-    """Return supported rootfs as pytest parameters"""
-    for rootfs in disks(glob=glob):
-        yield pytest.param(rootfs, id=rootfs.name)
-
-
 @dataclass(frozen=True, repr=True)
 class FirecrackerArtifact:
     """Utility class for Firecracker binary artifacts."""