From 4d0d9f6fd0451df7f22341551e5035cb5d047a43 Mon Sep 17 00:00:00 2001
From: Matthieu Darbois <mayeut@users.noreply.github.com>
Date: Fri, 10 Jan 2025 16:24:02 +0100
Subject: [PATCH] ci: warm-up docker images (#2099)

* ci: warm-up docker images

* don't cache when running emulation tests

* skip graalpy cache

* add comment about the single worker case

* enforce docker warm-up succeeds in CI

* skip warm-up on architectures where it does not do any good
---
 pyproject.toml   |  1 +
 test/conftest.py | 81 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 3b0d72d1f..59ae18c79 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -84,6 +84,7 @@ docs = [
 ]
 test = [
     "build",
+    "filelock",
     "jinja2",
     "pytest-timeout",
     "pytest-xdist",
diff --git a/test/conftest.py b/test/conftest.py
index a42a3ce9a..dd884f464 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -5,7 +5,10 @@
 from typing import Generator
 
 import pytest
+from filelock import FileLock
 
+from cibuildwheel.architecture import Architecture
+from cibuildwheel.options import CommandLineArguments, Options
 from cibuildwheel.util import detect_ci_provider, find_uv
 
 from .utils import EMULATED_ARCHS, platform
@@ -28,6 +31,84 @@ def pytest_addoption(parser: pytest.Parser) -> None:
     )
 
 
+def docker_warmup(request: pytest.FixtureRequest) -> None:
+    machine = request.config.getoption("--run-emulation", default=None)
+    if machine is None:
+        archs = {arch.value for arch in Architecture.auto_archs("linux")}
+    elif machine == "all":
+        archs = set(EMULATED_ARCHS)
+    else:
+        archs = {machine}
+
+    # Only include architectures where there are missing pre-installed interpreters
+    archs &= {"x86_64", "i686", "aarch64"}
+    if not archs:
+        return
+
+    options = Options(
+        platform="linux",
+        command_line_arguments=CommandLineArguments.defaults(),
+        env={},
+        defaults=True,
+    )
+    build_options = options.build_options(None)
+    assert build_options.manylinux_images is not None
+    assert build_options.musllinux_images is not None
+    images = [build_options.manylinux_images[arch] for arch in archs] + [
+        build_options.musllinux_images[arch] for arch in archs
+    ]
+    # exclude GraalPy as it's not a target for cibuildwheel
+    command = (
+        "manylinux-interpreters ensure $(manylinux-interpreters list 2>/dev/null | grep -v graalpy) &&"
+        "cpython3.13 -m pip download -d /tmp setuptools wheel pytest"
+    )
+    for image in images:
+        container_id = subprocess.run(
+            ["docker", "create", image, "bash", "-c", command],
+            text=True,
+            check=True,
+            stdout=subprocess.PIPE,
+        ).stdout.strip()
+        try:
+            subprocess.run(["docker", "start", container_id], check=True, stdout=subprocess.DEVNULL)
+            exit_code = subprocess.run(
+                ["docker", "wait", container_id], text=True, check=True, stdout=subprocess.PIPE
+            ).stdout.strip()
+            assert exit_code == "0"
+            subprocess.run(
+                ["docker", "commit", container_id, image], check=True, stdout=subprocess.DEVNULL
+            )
+        finally:
+            subprocess.run(["docker", "rm", container_id], check=True, stdout=subprocess.DEVNULL)
+
+
+@pytest.fixture(scope="session", autouse=True)
+def docker_warmup_fixture(
+    request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory, worker_id: str
+) -> None:
+    # if we're in CI testing linux, let's warm-up docker images
+    if detect_ci_provider() is None or platform != "linux":
+        return None
+    if request.config.getoption("--run-emulation", default=None) is not None:
+        # emulation tests only run one test in CI, caching the image only slows down the test
+        return None
+
+    if worker_id == "master":
+        # not executing with multiple workers
+        # it might be unsafe to write to tmp_path_factory.getbasetemp().parent
+        return docker_warmup(request)
+
+    # get the temp directory shared by all workers
+    root_tmp_dir = tmp_path_factory.getbasetemp().parent
+
+    fn = root_tmp_dir / "warmup.done"
+    with FileLock(str(fn) + ".lock"):
+        if not fn.is_file():
+            docker_warmup(request)
+            fn.write_text("done")
+    return None
+
+
 @pytest.fixture(params=["pip", "build"])
 def build_frontend_env_nouv(request: pytest.FixtureRequest) -> dict[str, str]:
     frontend = request.param