[mypyc] Cache built librt in tests (#20626)

JukkaL · web-flow · commit 6da2283cfcd6 · 2026-01-21T15:29:28.000Z
Store built `librt` under `build/librt-cache`, so that it can be reused
to speed up mypyc tests when there are no changes under `mypyc/lib-rt`.
Also invalidate the cache if the build environment changes.

This also ensures that each test run only compiles `librt` at most twice
(with and without experimental features), instead of once per test case
that requires `librt`.

Use `filelock` to support parallel tests.

Building `librt` can take 6+ seconds, so the savings can be quite
significant. We are also planning to add several additional `librt`
submodules, which will further slow down the build, and this will result
in a larger number of tests that require `librt`. The main benefit is
that adding more tests or submodules doesn't significantly slow down
tests.
diff --git a/mypyc/test/librt_cache.py b/mypyc/test/librt_cache.py
@@ -0,0 +1,233 @@
+"""Build and cache librt for use in tests.
+
+This module provides a way to build librt extension modules once and cache
+them across test runs, and across different test cases in a single run. The
+cache is invalidated when source files or details of the build environment change.
+
+Note: Tests must run in a subprocess to use the cached librt, since importing
+this module also triggers the import of the regular installed librt.
+
+Usage:
+    from mypyc.test.librt_cache import get_librt_path, run_with_librt
+
+    # Get path to built librt (builds if needed)
+    path = get_librt_path()
+
+    # Run a test file in subprocess with built librt
+    result = run_with_librt("test_librt.py")
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import shutil
+import subprocess
+import sys
+import sysconfig
+from typing import Any
+
+import filelock
+
+from mypyc.build import LIBRT_MODULES, get_cflags, include_dir
+from mypyc.common import RUNTIME_C_FILES
+
+
+def _librt_build_hash(experimental: bool) -> str:
+    """Compute hash for librt build, including sources and build environment."""
+    # Import lazily to ensure mypyc.build has ensured that distutils is correctly set up
+    from distutils import ccompiler
+
+    h = hashlib.sha256()
+    # Include experimental flag
+    h.update(b"exp" if experimental else b"noexp")
+    # Include full Python version string (includes git hash for dev builds)
+    h.update(sys.version.encode())
+    # Include debug build status (gettotalrefcount only exists in debug builds)
+    is_debug = hasattr(sys, "gettotalrefcount")
+    h.update(b"debug" if is_debug else b"release")
+    # Include free-threading status (Python 3.13+)
+    is_free_threaded = bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
+    h.update(b"freethreaded" if is_free_threaded else b"gil")
+    # Include compiler type (e.g., "unix" or "msvc")
+    compiler: Any = ccompiler.new_compiler()
+    h.update(compiler.compiler_type.encode())
+    # Include environment variables that affect C compilation
+    for var in ("CC", "CXX", "CFLAGS", "CPPFLAGS", "LDFLAGS"):
+        val = os.environ.get(var, "")
+        h.update(f"{var}={val}".encode())
+    # Hash runtime files
+    for name in RUNTIME_C_FILES:
+        path = os.path.join(include_dir(), name)
+        h.update(name.encode() + b"|")
+        with open(path, "rb") as f:
+            h.update(f.read())
+    # Hash librt module files
+    for mod, files, extra, includes in LIBRT_MODULES:
+        for fname in files + extra:
+            path = os.path.join(include_dir(), fname)
+            h.update(fname.encode() + b"|")
+            with open(path, "rb") as f:
+                h.update(f.read())
+    return h.hexdigest()[:16]
+
+
+def _generate_setup_py(build_dir: str, experimental: bool) -> str:
+    """Generate setup.py content for building librt directly.
+
+    We inline LIBRT_MODULES/RUNTIME_C_FILES/include_dir/cflags values to avoid
+    importing mypyc.build, which recursively imports lots of things.
+    """
+    lib_rt_dir = include_dir()
+
+    # Get compiler flags using the shared helper (with -O0 for faster builds)
+    cflags = get_cflags(opt_level="0", experimental_features=experimental)
+
+    # Serialize values to inline in generated setup.py
+    librt_modules_repr = repr(
+        [(m.module, m.c_files, m.other_files, m.include_dirs) for m in LIBRT_MODULES]
+    )
+    runtime_files_repr = repr(RUNTIME_C_FILES)
+    cflags_repr = repr(cflags)
+
+    return f"""\
+import os
+from setuptools import setup, Extension
+import build_setup  # noqa: F401  # Monkey-patches compiler for per-file SIMD flags
+
+build_dir = {build_dir!r}
+lib_rt_dir = {lib_rt_dir!r}
+
+RUNTIME_C_FILES = {runtime_files_repr}
+LIBRT_MODULES = {librt_modules_repr}
+CFLAGS = {cflags_repr}
+
+def write_file(path, contents):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "wb") as f:
+        f.write(contents)
+
+# Copy runtime C files
+for name in RUNTIME_C_FILES:
+    src = os.path.join(lib_rt_dir, name)
+    dst = os.path.join(build_dir, name)
+    with open(src, "rb") as f:
+        write_file(dst, f.read())
+
+# Build extensions for each librt module
+extensions = []
+for mod, file_names, extra_files, includes in LIBRT_MODULES:
+    # Copy source files
+    for fname in file_names + extra_files:
+        src = os.path.join(lib_rt_dir, fname)
+        dst = os.path.join(build_dir, fname)
+        with open(src, "rb") as f:
+            write_file(dst, f.read())
+
+    extensions.append(Extension(
+        mod,
+        sources=[os.path.join(build_dir, f) for f in file_names + RUNTIME_C_FILES],
+        include_dirs=[lib_rt_dir] + [os.path.join(lib_rt_dir, d) for d in includes],
+        extra_compile_args=CFLAGS,
+    ))
+
+setup(name='librt_cached', ext_modules=extensions)
+"""
+
+
+def get_librt_path(experimental: bool = True) -> str:
+    """Get path to librt built from the repository, building and caching if necessary.
+
+    Uses build/librt-cache/ under the repo root (gitignored). The cache is
+    keyed by a hash of sources and build environment, so it auto-invalidates
+    when relevant factors change.
+
+    Safe to call from multiple parallel pytest workers - uses file locking.
+
+    Args:
+        experimental: Whether to enable experimental features.
+
+    Returns:
+        Path to directory containing built librt modules.
+    """
+    # Use build/librt-cache/ under the repo root (gitignored)
+    repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    cache_root = os.path.join(repo_root, "build", "librt-cache")
+    build_hash = _librt_build_hash(experimental)
+    build_dir = os.path.join(cache_root, f"librt-{build_hash}")
+    lock_file = os.path.join(cache_root, f"librt-{build_hash}.lock")
+    marker = os.path.join(build_dir, ".complete")
+
+    os.makedirs(cache_root, exist_ok=True)
+
+    with filelock.FileLock(lock_file, timeout=300):  # 5 min timeout
+        if os.path.exists(marker):
+            return build_dir
+
+        # Clean up any partial build
+        if os.path.exists(build_dir):
+            shutil.rmtree(build_dir)
+
+        os.makedirs(build_dir)
+
+        # Create librt package directory for --inplace to copy .so files into
+        librt_pkg = os.path.join(build_dir, "librt")
+        os.makedirs(librt_pkg)
+        with open(os.path.join(librt_pkg, "__init__.py"), "w") as f:
+            pass
+
+        # Copy build_setup.py for per-file SIMD compiler flags
+        build_setup_src = os.path.join(
+            os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "build_setup.py"
+        )
+        build_setup_dst = os.path.join(build_dir, "build_setup.py")
+        shutil.copy(build_setup_src, build_setup_dst)
+
+        # Write setup.py
+        setup_py = os.path.join(build_dir, "setup.py")
+        with open(setup_py, "w") as f:
+            f.write(_generate_setup_py(build_dir, experimental))
+
+        # Build (parallel builds don't work well because multiple extensions
+        # share the same runtime C files, causing race conditions)
+        result = subprocess.run(
+            [sys.executable, setup_py, "build_ext", "--inplace"],
+            cwd=build_dir,
+            capture_output=True,
+            text=True,
+        )
+        if result.returncode != 0:
+            raise RuntimeError(f"librt build failed:\n{result.stdout}\n{result.stderr}")
+
+        # Mark complete
+        with open(marker, "w") as f:
+            f.write("ok")
+
+    return build_dir
+
+
+def run_with_librt(
+    file_path: str, experimental: bool = True, check: bool = True
+) -> subprocess.CompletedProcess[str]:
+    """Run a Python file in a subprocess with built librt available.
+
+    This runs the file in a fresh Python process where the built librt
+    is at the front of sys.path, avoiding conflicts with any system librt.
+
+    Args:
+        file_path: Path to Python file to execute.
+        experimental: Whether to use experimental features.
+        check: If True, raise CalledProcessError on non-zero exit.
+
+    Returns:
+        CompletedProcess with stdout, stderr, and returncode.
+    """
+    librt_path = get_librt_path(experimental)
+    # Prepend librt path to PYTHONPATH
+    env = os.environ.copy()
+    existing = env.get("PYTHONPATH", "")
+    env["PYTHONPATH"] = librt_path + (os.pathsep + existing if existing else "")
+
+    return subprocess.run(
+        [sys.executable, file_path], capture_output=True, text=True, check=check, env=env
+    )
diff --git a/mypyc/test/test_run.py b/mypyc/test/test_run.py
@@ -26,6 +26,7 @@
 from mypyc.errors import Errors
 from mypyc.options import CompilerOptions
 from mypyc.test.config import test_data_prefix
+from mypyc.test.librt_cache import get_librt_path
 from mypyc.test.test_serialization import check_serialization_roundtrip
 from mypyc.test.testutil import (
     ICODE_GEN_BUILTINS,
@@ -289,6 +290,7 @@ def run_case_step(self, testcase: DataDrivenTestCase, incremental_step: int) ->
 
         setup_file = os.path.abspath(os.path.join(WORKDIR, "setup.py"))
         # We pass the C file information to the build script via setup.py unfortunately
+        # Note: install_librt is always False since we use cached librt from librt_cache
         with open(setup_file, "w", encoding="utf-8") as f:
             f.write(
                 setup_format.format(
@@ -297,16 +299,15 @@ def run_case_step(self, testcase: DataDrivenTestCase, incremental_step: int) ->
                     (cfiles, deps),
                     self.multi_file,
                     opt_level,
-                    librt,
+                    False,  # install_librt - use cached version instead
                     experimental_features,
                 )
             )
 
         if librt:
-            # This hack forces Python to prefer the local "installation".
-            os.makedirs("librt", exist_ok=True)
-            with open(os.path.join("librt", "__init__.py"), "a"):
-                pass
+            # Use cached pre-built librt instead of rebuilding for each test
+            cached_librt = get_librt_path(experimental_features)
+            shutil.copytree(os.path.join(cached_librt, "librt"), "librt")
 
         if not run_setup(setup_file, ["build_ext", "--inplace"]):
             if testcase.config.getoption("--mypyc-showc"):
