From 7a592ff09a5af950f76289a446702b0fe4ff5361 Mon Sep 17 00:00:00 2001 From: Mac Lee Date: Sun, 26 Jul 2026 19:21:09 +0800 Subject: [PATCH] [ci]: cache FastVideo kernel builds in Modal (#1562) --- .github/workflows/infra-build-image.yml | 21 +- docker/Dockerfile | 35 +- docs/contributing/ci_architecture.md | 30 +- fastvideo-kernel/build.sh | 33 +- fastvideo/tests/modal/kernel_build_cache.py | 571 ++++++++++++++++++ fastvideo/tests/modal/kernel_cache_smoke.py | 253 ++++++++ fastvideo/tests/modal/pr_test.py | 45 +- fastvideo/tests/modal/ssim_test.py | 10 +- .../tests/modal/test_kernel_build_cache.py | 429 +++++++++++++ fastvideo/tests/modal/test_pr_test.py | 64 +- fastvideo/tests/modal/test_ssim_test.py | 33 + 11 files changed, 1470 insertions(+), 54 deletions(-) create mode 100644 fastvideo/tests/modal/kernel_build_cache.py create mode 100644 fastvideo/tests/modal/kernel_cache_smoke.py create mode 100644 fastvideo/tests/modal/test_kernel_build_cache.py diff --git a/.github/workflows/infra-build-image.yml b/.github/workflows/infra-build-image.yml index e2fb20f69..d2d477958 100644 --- a/.github/workflows/infra-build-image.yml +++ b/.github/workflows/infra-build-image.yml @@ -13,16 +13,23 @@ on: required: false default: false type: boolean - # Auto-rebuild the CUDA images when their Dockerfile changes on main. The CUDA - # matrix is the only lane that builds from docker/Dockerfile, so a path-scoped - # push trigger is a sufficient change detector on its own -- no separate - # detect-changes/paths-filter job is needed now that there is a single - # in-scope Dockerfile. Dreamverse (apps/dreamverse/docker/Dockerfile) and the - # rocm Dockerfile stay manual-dispatch only. + # Auto-rebuild the CUDA images when a repository-controlled image input + # changes on main. This includes the trusted SM89 kernel artifact's source, + # metadata/key helper, ABI dependency metadata, and build orchestration. + # Dreamverse (apps/dreamverse/docker/Dockerfile) and the ROCm Dockerfile stay + # manual-dispatch only. push: branches: [main] paths: + - '.dockerignore' + - '.github/workflows/_template-build-image.yml' + - '.github/workflows/infra-build-image.yml' + - '.gitmodules' - 'docker/Dockerfile' + - 'docker/uv-excludes' + - 'fastvideo-kernel/**' + - 'fastvideo/tests/modal/kernel_build_cache.py' + - 'pyproject.toml' permissions: @@ -50,7 +57,7 @@ jobs: # 2.8.3 comes from the architecture-specific prebuilt releases. build-cuda-images: # Runs on a manual dispatch when build_cuda_matrix is set, or automatically - # on a push that changed docker/Dockerfile (inputs are null on push). The + # on an in-scope main push (inputs are null on push). The # repository guard keeps fork syncs from auto-building; manual dispatch # still works in forks. if: ${{ (github.event_name == 'push' && github.repository == 'hao-ai-lab/FastVideo') || github.event.inputs.build_cuda_matrix == 'true' }} diff --git a/docker/Dockerfile b/docker/Dockerfile index 160f3af4b..e924beef1 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -186,15 +186,40 @@ RUN --mount=type=cache,target=/opt/uv/cache \ COPY . . -# Install FastVideo Unified Kernel exactly once. build.sh initializes only its -# CUTLASS/ThunderKittens submodules, then compiles for TORCH_CUDA_ARCH_LIST -# (default Hopper sm_90a) without requiring a live GPU. +# Build immutable FastVideo kernel wheels for the published image. The requested +# architecture remains installed for normal image users; amd64 images also carry +# an SM89 artifact so the predominant L40S Modal lanes can reuse it exactly. +ARG FASTVIDEO_KERNEL_PREBUILT_DIR=/opt/fastvideo-kernel-prebuilt RUN --mount=type=cache,target=/opt/uv/cache \ source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ + default_arch="${TORCH_CUDA_ARCH_LIST}" && \ + default_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/${default_arch}" && \ + export TORCH_CUDA_ARCH_LIST="${default_arch}" && \ cd fastvideo-kernel && \ - CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \ - TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} ./build.sh + CMAKE_ARGS= CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \ + ./build.sh --wheel-dir "${default_wheel_dir}" && \ + cd /FastVideo && \ + CMAKE_ARGS= python fastvideo/tests/modal/kernel_build_cache.py write-build-info \ + --wheel-dir "${default_wheel_dir}" \ + --output "${default_wheel_dir}/metadata.json" && \ + if [[ "${TARGETARCH:-amd64}" == "amd64" && "${default_arch}" != "8.9" ]]; then \ + export TORCH_CUDA_ARCH_LIST=8.9 && \ + l40s_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/8.9" && \ + cd /FastVideo/fastvideo-kernel && \ + CMAKE_ARGS= CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \ + ./build.sh --wheel-dir "${l40s_wheel_dir}" && \ + cd /FastVideo && \ + CMAKE_ARGS= python fastvideo/tests/modal/kernel_build_cache.py write-build-info \ + --wheel-dir "${l40s_wheel_dir}" \ + --output "${l40s_wheel_dir}/metadata.json"; \ + fi && \ + export TORCH_CUDA_ARCH_LIST="${default_arch}" && \ + default_wheel="$(find "${default_wheel_dir}" -maxdepth 1 -type f \ + \( -name 'fastvideo_kernel-*.whl' -o -name 'fastvideo-kernel-*.whl' \) \ + | sort | tail -n 1)" && \ + uv pip install "${default_wheel}" \ + --reinstall-package fastvideo-kernel --no-deps # Install FastVideo itself (editable) now that the source is present, and set up # shell configuration. Dependencies and the local kernel are already installed, diff --git a/docs/contributing/ci_architecture.md b/docs/contributing/ci_architecture.md index 6bbd77fa8..21e9abb45 100644 --- a/docs/contributing/ci_architecture.md +++ b/docs/contributing/ci_architecture.md @@ -269,16 +269,26 @@ The docs job: ### Docker Images `.github/workflows/infra-build-image.yml` supports manual `workflow_dispatch` -runs and automatically rebuilds the CUDA matrix when `docker/Dockerfile` -changes on `main` in the canonical repository. Manual runs let maintainers -choose which image families to build. The `fastvideo-dev` matrix builds Python -3.12 images for CUDA 12.6 and CUDA 13 on native `amd64` and `arm64` runners, -then publishes one multi-platform manifest per CUDA version. CUDA 12.6 owns the -`py3.12-latest` and global `latest` tags, as well as the explicit -`py3.12-cuda12.6.3-latest` alias. CUDA 13 is published under the explicit -`py3.12-cuda13.0.0-latest` tag. This publication policy does not change the -unparameterized `docker/Dockerfile` build defaults, which remain CUDA 13 and -`cu130`. +runs and automatically rebuilds the CUDA matrix when a repository-controlled +image input changes on `main` in the canonical repository. Those inputs include +the CUDA Dockerfile and reusable workflow, dependency metadata, Docker context +policy, `fastvideo-kernel/**`, and the kernel artifact metadata/key helper. +Manual runs let maintainers choose which image families to build. The +`fastvideo-dev` matrix builds Python 3.12 images for CUDA 12.6 and CUDA 13 on +native `amd64` and `arm64` runners, then publishes one multi-platform manifest +per CUDA version. CUDA 12.6 owns the `py3.12-latest` and global `latest` tags, as +well as the explicit `py3.12-cuda12.6.3-latest` alias. CUDA 13 is published under +the explicit `py3.12-cuda13.0.0-latest` tag. This publication policy does not +change the unparameterized `docker/Dockerfile` build defaults, which remain CUDA +13 and `cu130`. + +Published amd64 development images keep their configured Hopper kernel wheel +installed and also carry an immutable SM89 wheel under +`/opt/fastvideo-kernel-prebuilt`. Modal PR and SSIM jobs select the exact +source, ABI, and GPU-architecture match from that directory, so L40S jobs reuse +the trusted image artifact while kernel-changing PRs still build locally. Once +a kernel or artifact-key change reaches `main`, the image workflow republishes +the matching trusted artifact before later jobs consume the updated image tag. The optional Dreamverse matrix builds backend and UI images for CUDA 12.6 and CUDA 13 on `amd64`. Dreamverse remains `amd64`-only because its FA4 dependency diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh index 602000b26..c7ed847c3 100755 --- a/fastvideo-kernel/build.sh +++ b/fastvideo-kernel/build.sh @@ -4,6 +4,7 @@ set -ex # Simple build script wrapping uv/pip # Usage: # ./build.sh # local build (torch-based arch detection, TK only on SM90) +# ./build.sh --wheel-dir /tmp/wheels # build + install a reusable wheel # Environment overrides (if set, they win over auto-detection): # TORCH_CUDA_ARCH_LIST # CMAKE_ARGS (for FASTVIDEO_KERNEL_BUILD_TK / CMAKE_CUDA_ARCHITECTURES / GPU_BACKEND) @@ -67,12 +68,25 @@ fi # Install build dependencies uv pip install scikit-build-core cmake ninja -RELEASE=0 GPU_BACKEND=CUDA -for arg in "$@"; do - case "$arg" in +WHEEL_DIR="" +while [[ $# -gt 0 ]]; do + case "$1" in --rocm) GPU_BACKEND=ROCM + shift + ;; + --wheel-dir) + if [[ $# -lt 2 ]]; then + echo "ERROR: --wheel-dir requires a directory path" >&2 + exit 2 + fi + WHEEL_DIR="$2" + shift 2 + ;; + *) + echo "ERROR: unknown argument: $1" >&2 + exit 2 ;; esac done @@ -165,4 +179,15 @@ echo "CMAKE_ARGS: ${CMAKE_ARGS:-}" echo "GPU_BACKEND: ${GPU_BACKEND:-}" # Build and install # Use -v for verbose output -uv pip install . -v --no-build-isolation +if [ -n "${WHEEL_DIR}" ]; then + mkdir -p "${WHEEL_DIR}" + uv build --wheel -v --no-build-isolation --out-dir "${WHEEL_DIR}" . + wheel_path="$(find "${WHEEL_DIR}" -maxdepth 1 -type f \( -name 'fastvideo_kernel-*.whl' -o -name 'fastvideo-kernel-*.whl' \) | sort | tail -n 1)" + if [ -z "${wheel_path}" ]; then + echo "ERROR: no fastvideo-kernel wheel was built in ${WHEEL_DIR}" >&2 + exit 1 + fi + uv pip install "${wheel_path}" --reinstall-package fastvideo-kernel --no-deps +else + uv pip install . -v --no-build-isolation +fi diff --git a/fastvideo/tests/modal/kernel_build_cache.py b/fastvideo/tests/modal/kernel_build_cache.py new file mode 100644 index 000000000..936a7b5ad --- /dev/null +++ b/fastvideo/tests/modal/kernel_build_cache.py @@ -0,0 +1,571 @@ +"""Build and reuse FastVideo kernel wheels in Modal CI. + +This module stays standalone because Modal launchers execute it from a freshly +cloned checkout after dependency installation. Shared cache consumers are +read-only; only the explicit ``produce`` command may mutate cache state. +""" + +import argparse +import datetime +import hashlib +import json +import os +import platform +import re +import shlex +import shutil +import subprocess +import sys +import sysconfig +import tempfile +import zipfile +from email.parser import Parser +from pathlib import Path + + +CACHE_SCHEMA_VERSION = 3 +DEFAULT_PREBUILT_INFO_PATH = "/opt/fastvideo-kernel-prebuilt" +KERNEL_RELATIVE_DIR = "fastvideo-kernel" +DEFAULT_BUILD_INFO_OUTPUT = "/opt/fastvideo-kernel-prebuilt/default/metadata.json" +METADATA_FILE = "metadata.json" +EXPECTED_DISTRIBUTION = "fastvideo-kernel" + + +def _utc_now_isoformat() -> str: + return datetime.datetime.now(datetime.timezone.utc).isoformat() + + +def _log(message: str) -> None: + print(f"[fastvideo-kernel-cache] {message}", flush=True) + + +def _run(args: list[str], + *, + cwd: Path | None = None, + env: dict[str, str] | None = None, + capture: bool = False) -> str: + _log("$ " + " ".join(str(arg) for arg in args)) + result = subprocess.run( + args, + cwd=str(cwd) if cwd else None, + env=env, + check=True, + text=True, + stdout=subprocess.PIPE if capture else None, + stderr=subprocess.STDOUT if capture else None, + ) + if capture: + return result.stdout.strip() + return "" + + +def _run_optional(args: list[str], *, cwd: Path | None = None) -> str: + try: + return _run(args, cwd=cwd, capture=True) + except (FileNotFoundError, subprocess.CalledProcessError, OSError) as error: + return f"" + + +def _sha256_text(text: str) -> str: + return hashlib.sha256(text.encode("utf-8", errors="replace")).hexdigest() + + +def _sha256_file(path: Path) -> str: + hasher = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + hasher.update(chunk) + return hasher.hexdigest() + + +def _hash_directory(root: Path) -> str: + hasher = hashlib.sha256() + skip_dirs = {".git", ".mypy_cache", ".pytest_cache", "__pycache__", "build", "dist"} + for current_root, dirnames, filenames in os.walk(root): + dirnames[:] = sorted( + dirname for dirname in dirnames if dirname not in skip_dirs and not dirname.endswith(".egg-info")) + for filename in sorted(filenames): + path = Path(current_root) / filename + hasher.update(path.relative_to(root).as_posix().encode("utf-8")) + hasher.update(b"\0") + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + hasher.update(chunk) + hasher.update(b"\0") + return hasher.hexdigest() + + +def _kernel_source_hash(repo_root: Path) -> dict[str, str]: + kernel_root = repo_root / KERNEL_RELATIVE_DIR + try: + tree_hash = _run(["git", "rev-parse", f"HEAD:{KERNEL_RELATIVE_DIR}"], cwd=repo_root, capture=True) + diff = _run_optional(["git", "diff", "--binary", "--", KERNEL_RELATIVE_DIR], cwd=repo_root) + cached_diff = _run_optional(["git", "diff", "--cached", "--binary", "--", KERNEL_RELATIVE_DIR], + cwd=repo_root) + submodules = _run_optional( + [ + "git", + "submodule", + "status", + "--", + "fastvideo-kernel/include/cutlass", + "fastvideo-kernel/include/tk", + ], + cwd=repo_root, + ) + descriptor = { + "mode": "git", + "tree_hash": tree_hash, + "diff_hash": _sha256_text(diff) if diff else "", + "cached_diff_hash": _sha256_text(cached_diff) if cached_diff else "", + "submodules_hash": _sha256_text(submodules), + } + descriptor["source_hash"] = _sha256_text(json.dumps(descriptor, sort_keys=True)) + return descriptor + except (FileNotFoundError, subprocess.CalledProcessError): + descriptor = { + "mode": "directory", + "tree_hash": _hash_directory(kernel_root), + "diff_hash": "", + "cached_diff_hash": "", + "submodules_hash": "", + } + descriptor["source_hash"] = descriptor["tree_hash"] + return descriptor + + +def _read_kernel_version(repo_root: Path) -> str: + pyproject = repo_root / KERNEL_RELATIVE_DIR / "pyproject.toml" + for line in pyproject.read_text(encoding="utf-8").splitlines(): + stripped = line.strip() + if stripped.startswith("version"): + _, _, value = stripped.partition("=") + return value.strip().strip("\"'") + return "" + + +def _detect_arch_from_torch() -> str: + try: + import torch + + if not torch.cuda.is_available(): + return "" + major, minor = torch.cuda.get_device_capability(0) + if major == 9 and minor == 0: + return "9.0a" + if major == 12 and minor == 0: + return "12.0a" + return f"{major}.{minor}" + except Exception as error: # noqa: BLE001 - metadata should explain fallback + return f"" + + +def _torch_metadata() -> dict[str, object]: + try: + import torch + + compiled_with_cxx11_abi = getattr(torch, "compiled_with_cxx11_abi", None) + if callable(compiled_with_cxx11_abi): + cxx11_abi: object = bool(compiled_with_cxx11_abi()) + else: + cxx11_abi = getattr(getattr(torch, "_C", object()), "_GLIBCXX_USE_CXX11_ABI", "") + return { + "torch_version": str(torch.__version__), + "torch_cuda_version": str(torch.version.cuda), + "torch_file": str(getattr(torch, "__file__", "")), + "torch_config": str(torch.__config__.show()), + "cxx11_abi": cxx11_abi, + } + except Exception as error: # noqa: BLE001 - report in metadata + return { + "torch_version": f"", + "torch_cuda_version": "", + "torch_file": "", + "torch_config": "", + "cxx11_abi": "", + } + + +def _selected_command_metadata(environment_name: str, default_command: str) -> dict[str, str]: + raw = os.environ.get(environment_name, "").strip() or default_command + try: + command = shlex.split(raw) + except ValueError as error: + return { + "raw": raw, + "command": "", + "resolved_executable": "", + "version": f"", + } + if not command: + command = [default_command] + resolved_executable = shutil.which(command[0]) or command[0] + return { + "raw": raw, + "command": json.dumps(command), + "resolved_executable": resolved_executable, + "version": _run_optional([*command, "--version"]), + } + + +def _compiler_libc_metadata() -> dict[str, object]: + return { + "compiler": { + "cc": _selected_command_metadata("CC", "cc"), + "cxx": _selected_command_metadata("CXX", "c++"), + }, + "build_tools": { + "cmake_version": _run_optional(["cmake", "--version"]), + "ninja_version": _run_optional(["ninja", "--version"]), + "linker": _selected_command_metadata("LD", "ld"), + }, + "libc": { + "platform_libc": list(platform.libc_ver()), + "ldd_version": _run_optional(["ldd", "--version"]), + }, + } + + +def _build_metadata(repo_root: Path) -> dict[str, object]: + explicit_arch = os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip() + resolved_arch = explicit_arch or _detect_arch_from_torch() + cache_key_build = { + "gpu_backend": os.environ.get("GPU_BACKEND", "CUDA"), + "resolved_torch_cuda_arch_list": resolved_arch, + "cmake_args": os.environ.get("CMAKE_ARGS", ""), + "cflags": os.environ.get("CFLAGS", ""), + "cxxflags": os.environ.get("CXXFLAGS", ""), + "ldflags": os.environ.get("LDFLAGS", ""), + } + cache_key_metadata: dict[str, object] = { + "schema_version": CACHE_SCHEMA_VERSION, + "kernel_version": _read_kernel_version(repo_root), + "source": _kernel_source_hash(repo_root), + "python": { + "version": platform.python_version(), + "implementation": platform.python_implementation(), + "cache_tag": sys.implementation.cache_tag, + "platform": sysconfig.get_platform(), + "machine": platform.machine(), + }, + "torch": _torch_metadata(), + "cuda": { + "cuda_home": os.environ.get("CUDA_HOME", ""), + "nvcc": _selected_command_metadata("CUDACXX", "nvcc"), + }, + "abi": _compiler_libc_metadata(), + "build": cache_key_build, + } + metadata = { + **cache_key_metadata, + "build": { + **cache_key_build, + "torch_cuda_arch_list": explicit_arch, + }, + } + metadata["cache_key"] = _sha256_text(json.dumps(cache_key_metadata, sort_keys=True, separators=(",", ":"))) + return metadata + + +def _find_wheel(directory: Path) -> Path: + candidates = sorted( + path + for pattern in ("fastvideo_kernel-*.whl", "fastvideo-kernel-*.whl") + for path in directory.glob(pattern) + ) + if not candidates: + raise RuntimeError(f"No fastvideo-kernel wheel found in {directory}") + return candidates[-1] + + +def _normalize_distribution(name: str) -> str: + return re.sub(r"[-_.]+", "-", name).lower() + + +def _wheel_artifact(wheel: Path) -> dict[str, object]: + if not wheel.is_file() or wheel.is_symlink(): + raise RuntimeError(f"wheel is missing or is not a regular file: {wheel}") + try: + with zipfile.ZipFile(wheel) as archive: + bad_member = archive.testzip() + if bad_member is not None: + raise RuntimeError(f"wheel archive has a corrupt member: {bad_member}") + names = archive.namelist() + wheel_metadata = [name for name in names if name.endswith(".dist-info/WHEEL")] + package_metadata = [name for name in names if name.endswith(".dist-info/METADATA")] + if len(wheel_metadata) != 1 or len(package_metadata) != 1: + raise RuntimeError("wheel archive must contain one WHEEL and one METADATA file") + parsed_metadata = Parser().parsestr(archive.read(package_metadata[0]).decode("utf-8")) + except (OSError, UnicodeDecodeError, zipfile.BadZipFile) as error: + raise RuntimeError(f"invalid wheel archive: {error}") from error + + distribution = _normalize_distribution(parsed_metadata.get("Name", "")) + if distribution != EXPECTED_DISTRIBUTION: + raise RuntimeError(f"unexpected wheel distribution {distribution!r}") + return { + "wheel_name": wheel.name, + "size_bytes": wheel.stat().st_size, + "sha256": _sha256_file(wheel), + "distribution": distribution, + } + + +def _validate_wheel(wheel: Path, expected_artifact: object) -> tuple[bool, str]: + if not isinstance(expected_artifact, dict): + return False, "artifact metadata is missing" + try: + actual_artifact = _wheel_artifact(wheel) + except RuntimeError as error: + return False, str(error) + for field in ("wheel_name", "size_bytes", "sha256", "distribution"): + if actual_artifact[field] != expected_artifact.get(field): + return False, f"artifact {field} mismatch" + return True, "" + + +def _install_wheel(wheel: Path) -> None: + _log(f"installing wheel {wheel}") + _run([ + "uv", + "pip", + "install", + str(wheel), + "--reinstall-package", + "fastvideo-kernel", + "--no-deps", + ]) + + +def _load_metadata(path: Path) -> tuple[dict[str, object] | None, str]: + if path.is_symlink(): + return None, "metadata is a symlink" + try: + metadata = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + return None, f"metadata is unreadable: {error}" + if not isinstance(metadata, dict): + return None, "metadata is not an object" + return metadata, "" + + +def _validated_payload_wheel(payload: dict[str, object], wheel: Path, cache_key: str) -> tuple[Path | None, str]: + if payload.get("schema_version") != CACHE_SCHEMA_VERSION: + return None, "schema version mismatch" + if payload.get("cache_key") != cache_key: + return None, "cache key mismatch" + valid, reason = _validate_wheel(wheel, payload.get("artifact")) + if not valid: + return None, reason + return wheel, "" + + +def _cache_entry_wheel(cache_entry: Path, cache_key: str) -> tuple[Path | None, str]: + if cache_entry.is_symlink(): + return None, "cache entry is a symlink" + metadata, reason = _load_metadata(cache_entry / METADATA_FILE) + if metadata is None: + return None, reason + artifact = metadata.get("artifact") + wheel_name = artifact.get("wheel_name") if isinstance(artifact, dict) else None + if not isinstance(wheel_name, str): + return None, "artifact wheel name is missing" + wheel = cache_entry / wheel_name + return _validated_payload_wheel(metadata, wheel, cache_key) + + +def _prebuilt_metadata_paths(prebuilt_info_path: Path) -> list[Path]: + if prebuilt_info_path.is_dir() and not prebuilt_info_path.is_symlink(): + return sorted(prebuilt_info_path.glob(f"*/{METADATA_FILE}")) + if prebuilt_info_path.exists() or prebuilt_info_path.is_symlink(): + return [prebuilt_info_path] + return [] + + +def _try_install_prebuilt(metadata: dict[str, object], prebuilt_info_path: Path) -> bool: + cache_key = str(metadata["cache_key"]) + for metadata_path in _prebuilt_metadata_paths(prebuilt_info_path): + prebuilt, reason = _load_metadata(metadata_path) + if prebuilt is None: + _log(f"Docker-prebuilt kernel metadata is invalid at {metadata_path}: {reason}") + continue + wheel_path = Path(str(prebuilt.get("wheel_path", ""))) + wheel, reason = _validated_payload_wheel(prebuilt, wheel_path, cache_key) + if wheel is None: + _log(f"Docker-prebuilt kernel artifact rejected at {metadata_path}: {reason}") + continue + try: + _log(f"Docker-prebuilt cache hit for key {cache_key}: {wheel}") + _install_wheel(wheel) + except (FileNotFoundError, OSError, subprocess.CalledProcessError) as error: + _log(f"Docker-prebuilt kernel installation failed: {error}; falling back") + return False + return True + _log(f"Docker-prebuilt cache miss for key {cache_key}") + return False + + +def _store_cache_entry(cache_root: Path, metadata: dict[str, object], wheel: Path) -> Path: + cache_key = str(metadata["cache_key"]) + cache_entry = cache_root / cache_key + if cache_entry.exists() or cache_entry.is_symlink(): + cached_wheel, reason = _cache_entry_wheel(cache_entry, cache_key) + if cached_wheel is not None: + _log(f"cache entry {cache_entry} appeared after build; leaving existing entry in place") + return cache_entry + _log(f"evicting invalid cache entry {cache_entry}: {reason}") + if cache_entry.is_dir() and not cache_entry.is_symlink(): + shutil.rmtree(cache_entry) + else: + cache_entry.unlink() + + temp_entry = Path(tempfile.mkdtemp(prefix=f".{cache_key}.tmp-", dir=cache_root)) + try: + stored_wheel = temp_entry / wheel.name + shutil.copy2(wheel, stored_wheel) + stored_metadata = { + **metadata, + "created_at_utc": _utc_now_isoformat(), + "artifact": _wheel_artifact(stored_wheel), + } + (temp_entry / METADATA_FILE).write_text(json.dumps(stored_metadata, indent=2, sort_keys=True), + encoding="utf-8") + temp_entry.rename(cache_entry) + except OSError as error: + shutil.rmtree(temp_entry, ignore_errors=True) + cached_wheel, _ = _cache_entry_wheel(cache_entry, cache_key) + if cached_wheel is not None: + _log(f"cache entry {cache_entry} was created concurrently; leaving existing entry in place") + return cache_entry + raise RuntimeError(f"Failed to store kernel cache entry at {cache_entry}") from error + + cached_wheel, reason = _cache_entry_wheel(cache_entry, cache_key) + if cached_wheel is None: + shutil.rmtree(cache_entry, ignore_errors=True) + raise RuntimeError(f"Stored kernel cache entry {cache_entry} could not be validated: {reason}") + return cache_entry + + +def _build_wheel(repo_root: Path, metadata: dict[str, object]) -> Path: + with tempfile.TemporaryDirectory(prefix="fastvideo-kernel-wheel-") as temp_dir: + wheel_dir = Path(temp_dir) + env = os.environ.copy() + build_metadata = metadata["build"] + if not isinstance(build_metadata, dict): + raise RuntimeError("build metadata is invalid") + resolved_arch = str(build_metadata["resolved_torch_cuda_arch_list"]) + if not env.get("TORCH_CUDA_ARCH_LIST") and not resolved_arch.startswith(" None: + wheel = _build_wheel(repo_root, metadata) + try: + _wheel_artifact(wheel) + _install_wheel(wheel) + finally: + shutil.rmtree(wheel.parent, ignore_errors=True) + + +def install_cached_or_build(repo_root: Path, cache_root: Path | None, prebuilt_info_path: Path) -> None: + """Install from a validated artifact or build locally without writing shared state.""" + metadata = _build_metadata(repo_root) + cache_key = str(metadata["cache_key"]) + _log(f"resolved cache key {cache_key}") + + if _try_install_prebuilt(metadata, prebuilt_info_path): + return + + if cache_root is not None: + cache_entry = cache_root / cache_key + cached_wheel, reason = _cache_entry_wheel(cache_entry, cache_key) + if cached_wheel is not None: + try: + _log(f"cache hit: {cached_wheel}") + _install_wheel(cached_wheel) + return + except (FileNotFoundError, OSError, subprocess.CalledProcessError) as error: + _log(f"cached wheel installation failed: {error}; building locally") + else: + _log(f"cache miss: {cache_entry} ({reason})") + else: + _log("shared cache disabled; using Docker-prebuilt artifact or a job-local build") + + _build_and_install_local(repo_root, metadata) + + +def produce_cached_wheel(repo_root: Path, cache_root: Path) -> None: + """Build and store one validated cache entry from a trusted producer.""" + metadata = _build_metadata(repo_root) + cache_key = str(metadata["cache_key"]) + _log(f"resolved cache key {cache_key}") + cache_root.mkdir(parents=True, exist_ok=True) + cache_entry = cache_root / cache_key + cached_wheel, reason = _cache_entry_wheel(cache_entry, cache_key) + if cached_wheel is not None: + _log(f"cache hit: {cached_wheel}") + return + if cache_entry.exists() or cache_entry.is_symlink(): + _log(f"producer will repair invalid cache entry {cache_entry}: {reason}") + + _log(f"cache miss: {cache_entry}") + wheel = _build_wheel(repo_root, metadata) + try: + _wheel_artifact(wheel) + _install_wheel(wheel) + stored_entry = _store_cache_entry(cache_root, metadata, wheel) + _log(f"stored wheel cache entry: {stored_entry}") + finally: + shutil.rmtree(wheel.parent, ignore_errors=True) + + +def write_build_info(repo_root: Path, wheel_dir: Path, output: Path) -> None: + metadata = _build_metadata(repo_root) + wheel = _find_wheel(wheel_dir) + output.parent.mkdir(parents=True, exist_ok=True) + payload = { + **metadata, + "created_at_utc": _utc_now_isoformat(), + "artifact": _wheel_artifact(wheel), + "wheel_path": str(wheel), + } + output.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8") + _log(f"wrote Docker-prebuilt kernel metadata: {output}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("command", choices=("install", "produce", "write-build-info", "print-key")) + parser.add_argument("--repo-root", default=os.getcwd()) + parser.add_argument("--cache-root", default=os.environ.get("FASTVIDEO_KERNEL_CACHE_ROOT", "")) + parser.add_argument("--prebuilt-info-path", + default=os.environ.get("FASTVIDEO_KERNEL_PREBUILT_INFO", DEFAULT_PREBUILT_INFO_PATH)) + parser.add_argument("--wheel-dir", default="") + parser.add_argument("--output", default=DEFAULT_BUILD_INFO_OUTPUT) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + repo_root = Path(args.repo_root).resolve() + cache_root = Path(args.cache_root) if args.cache_root else None + if args.command == "install": + install_cached_or_build(repo_root, cache_root, Path(args.prebuilt_info_path)) + elif args.command == "produce": + if cache_root is None: + raise RuntimeError("--cache-root is required for produce") + produce_cached_wheel(repo_root, cache_root) + elif args.command == "write-build-info": + if not args.wheel_dir: + raise RuntimeError("--wheel-dir is required for write-build-info") + write_build_info(repo_root, Path(args.wheel_dir), Path(args.output)) + elif args.command == "print-key": + print(json.dumps(_build_metadata(repo_root), indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/tests/modal/kernel_cache_smoke.py b/fastvideo/tests/modal/kernel_cache_smoke.py new file mode 100644 index 000000000..c140caab3 --- /dev/null +++ b/fastvideo/tests/modal/kernel_cache_smoke.py @@ -0,0 +1,253 @@ +"""Cross-container smoke for the Modal fastvideo-kernel build cache. + +Example: + python -m modal run fastvideo/tests/modal/kernel_cache_smoke.py \ + --git-repo https://github.com/macthecadillac/FastVideo.git \ + --git-commit + +This is intentionally opt-in: it builds fastvideo-kernel on a real Modal GPU. +""" + +import os +import shlex +import shutil +import subprocess +import sys +import time +from typing import Any + +import modal + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +try: + from modal_image_utils import ( # noqa: E402 + resolve_image_ref, resolve_uv_torch_backend) +except ModuleNotFoundError: + # Remote Modal containers re-import this module but mount only the + # entrypoint file; the digest resolution already happened at local launch. + def resolve_image_ref(image_ref: str) -> str: + return image_ref + + def resolve_uv_torch_backend(image_tag: str) -> str | None: + return os.environ.get("UV_TORCH_BACKEND") + + +app = modal.App("fastvideo-kernel-cache-smoke") + +REPO_DIR = "/FastVideo" +KERNEL_CACHE_VOLUME_PATH = "/root/fastvideo-kernel-cache" +KERNEL_CACHE_VOLUME_NAME = "fastvideo-kernel-build-cache-smoke-v3" +IMAGE_VERSION = os.getenv("IMAGE_VERSION", "latest") +IMAGE_TAG = os.environ.get( + "FASTVIDEO_MODAL_IMAGE", + f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{IMAGE_VERSION}", +) +IMAGE_REF = resolve_image_ref(IMAGE_TAG) + +print(f"Using image: {IMAGE_REF}") +print(f"Using kernel cache volume: {KERNEL_CACHE_VOLUME_NAME}") + +kernel_cache_vol = modal.Volume.from_name(KERNEL_CACHE_VOLUME_NAME, create_if_missing=True) +uv_torch_backend_override = resolve_uv_torch_backend(IMAGE_TAG) + +image = ( + modal.Image.from_registry(IMAGE_REF, add_python="3.12") + .apt_install( + "cmake", + "pkg-config", + "build-essential", + "curl", + "git", + "libssl-dev", + "ffmpeg", + ) + .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable") + .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc") + .env({ + "PATH": "/root/.cargo/bin:$PATH", + **({"UV_TORCH_BACKEND": uv_torch_backend_override} if uv_torch_backend_override else {}), + }) +) + +PRODUCER_KWARGS = dict( + gpu="L40S:1", + image=image, + timeout=5400, + volumes={KERNEL_CACHE_VOLUME_PATH: kernel_cache_vol}, +) +CONSUMER_KWARGS = dict( + gpu="L40S:1", + image=image, + timeout=5400, + volumes={KERNEL_CACHE_VOLUME_PATH: kernel_cache_vol.with_mount_options(read_only=True)}, +) + + +def _run_local_git_command(args: list[str]) -> str: + result = subprocess.run( + ["git", *args], + check=True, + capture_output=True, + text=True, + ) + return result.stdout.strip() + + +def _normalize_git_repo_url(git_repo: str) -> str: + if git_repo.startswith("git@github.com:"): + return "https://github.com/" + git_repo[len("git@github.com:"):] + if git_repo.startswith("ssh://git@github.com/"): + return "https://github.com/" + git_repo[len("ssh://git@github.com/"):] + return git_repo + + +def _resolve_git_repo(git_repo: str) -> str: + if git_repo.strip(): + return _normalize_git_repo_url(git_repo.strip()) + env_repo = os.environ.get("BUILDKITE_REPO", "").strip() + if env_repo: + return _normalize_git_repo_url(env_repo) + discovered_repo = _run_local_git_command(["config", "--get", "remote.origin.url"]) + if discovered_repo: + return _normalize_git_repo_url(discovered_repo) + raise RuntimeError("Could not resolve git repo URL. Pass --git-repo or set BUILDKITE_REPO.") + + +def _resolve_git_commit(git_commit: str) -> str: + if git_commit.strip(): + return git_commit.strip() + env_commit = os.environ.get("BUILDKITE_COMMIT", "").strip() + if env_commit: + return env_commit + discovered_commit = _run_local_git_command(["rev-parse", "HEAD"]) + if discovered_commit: + return discovered_commit + raise RuntimeError("Could not resolve git commit. Pass --git-commit or set BUILDKITE_COMMIT.") + + +def _run(args: list[str], cwd: str | None = None) -> str: + print("$ " + " ".join(shlex.quote(arg) for arg in args), flush=True) + result = subprocess.run( + args, + cwd=cwd, + check=True, + stdout=subprocess.PIPE, + stderr=sys.stderr, + text=True, + ) + if result.stdout: + print(result.stdout, end="", flush=True) + return result.stdout.strip() + + +def _run_shell_capture(command: str, cwd: str) -> str: + print(f"$ {command}", flush=True) + result = subprocess.run( + ["/bin/bash", "-lc", command], + cwd=cwd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, + ) + if result.stdout: + print(result.stdout, end="", flush=True) + if result.returncode != 0: + raise RuntimeError(f"Command failed with exit code {result.returncode}: {command}") + return result.stdout + + +def _clone_checkout(git_repo: str, git_commit: str) -> str: + shutil.rmtree(REPO_DIR, ignore_errors=True) + _run(["git", "-c", "http.version=HTTP/1.1", "clone", git_repo, REPO_DIR], cwd="/") + _run(["git", "checkout", git_commit], cwd=REPO_DIR) + _run(["git", "submodule", "update", "--init", "--recursive"], cwd=REPO_DIR) + return _run(["git", "rev-parse", "HEAD"], cwd=REPO_DIR) + + +def _cache_command_output(command: str, cache_root: str) -> str: + # This smoke exercises the Modal volume cache specifically, so do not let + # a matching Docker-prebuilt wheel satisfy the miss/hit assertions. + disabled_prebuilt_info_path = f"{cache_root}/.docker-prebuilt-disabled.json" + return _run_shell_capture( + "source $HOME/.local/bin/env && " + "source /opt/venv/bin/activate && " + f"python fastvideo/tests/modal/kernel_build_cache.py {shlex.quote(command)} " + f"--cache-root {shlex.quote(cache_root)} " + f"--prebuilt-info-path {shlex.quote(disabled_prebuilt_info_path)}", + cwd=REPO_DIR, + ) + + +def _assert_in_output(output: str, needle: str) -> None: + if needle not in output: + raise RuntimeError(f"Expected output to contain {needle!r}") + + +def _assert_not_in_output(output: str, needle: str) -> None: + if needle in output: + raise RuntimeError(f"Expected output not to contain {needle!r}") + + +@app.function(**PRODUCER_KWARGS) +def build_and_store_cache(git_repo: str, git_commit: str, cache_root: str) -> dict[str, Any]: + checked_out_commit = _clone_checkout(git_repo, git_commit) + output = _cache_command_output("produce", cache_root) + _assert_in_output(output, "cache miss:") + _assert_in_output(output, "./build.sh --wheel-dir") + _assert_in_output(output, "stored wheel cache entry:") + print("Committing kernel cache volume after miss/store", flush=True) + kernel_cache_vol.commit() + return { + "checked_out_commit": checked_out_commit, + "cache_root": cache_root, + "cache_miss": True, + "stored": True, + } + + +@app.function(**CONSUMER_KWARGS) +def verify_cache_hit(git_repo: str, git_commit: str, cache_root: str) -> dict[str, Any]: + checked_out_commit = _clone_checkout(git_repo, git_commit) + output = _cache_command_output("install", cache_root) + _assert_in_output(output, "cache hit:") + _assert_not_in_output(output, "cache miss:") + _assert_not_in_output(output, "./build.sh --wheel-dir") + return { + "checked_out_commit": checked_out_commit, + "cache_root": cache_root, + "cache_hit": True, + } + + +@app.function(image=image, timeout=300, volumes={KERNEL_CACHE_VOLUME_PATH: kernel_cache_vol}) +def cleanup_cache(cache_root: str) -> None: + shutil.rmtree(cache_root, ignore_errors=True) + print(f"Removed smoke cache root {cache_root}", flush=True) + kernel_cache_vol.commit() + + +@app.local_entrypoint() +def main( + git_repo: str = "", + git_commit: str = "", + cache_namespace: str = "", + cleanup: bool = True, +): + resolved_git_repo = _resolve_git_repo(git_repo) + resolved_git_commit = _resolve_git_commit(git_commit) + namespace = cache_namespace.strip() or f"smoke-{resolved_git_commit[:12]}-{int(time.time())}" + cache_root = f"{KERNEL_CACHE_VOLUME_PATH}/{namespace}" + + print("Launching cross-container kernel cache smoke") + print(f"Repo: {resolved_git_repo}") + print(f"Commit: {resolved_git_commit}") + print(f"Cache root: {cache_root}") + + miss_result = build_and_store_cache.remote(resolved_git_repo, resolved_git_commit, cache_root) + hit_result = verify_cache_hit.remote(resolved_git_repo, resolved_git_commit, cache_root) + if cleanup: + cleanup_cache.remote(cache_root) + + print("Kernel cache smoke passed") + print({"miss": miss_result, "hit": hit_result, "cleanup": cleanup}) diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index b972458cf..3f334b584 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -209,6 +209,10 @@ def run_test_command(test_command: str, lane would then test stale kernels. Pass install_command="" for commands that manage their own installs. """ + import os + import subprocess + import sys + git_repo = os.environ.get("BUILDKITE_REPO", "") git_commit = os.environ.get("BUILDKITE_COMMIT", "") pr_number = os.environ.get("BUILDKITE_PULL_REQUEST") @@ -219,23 +223,32 @@ def run_test_command(test_command: str, print(f"PR number: {pr_number}") _checkout_repository(git_repo, git_commit, pr_number) - build_kernel_command = """ - cd fastvideo-kernel && - ./build.sh && - cd .. && - """ if build_kernel else "" + setup_steps = [ + "source $HOME/.local/bin/env", + "source /opt/venv/bin/activate", + "cd /FastVideo", + ] + if install_command: + setup_steps.append(install_command) + if build_kernel: + setup_steps.append("python fastvideo/tests/modal/kernel_build_cache.py install") + setup_command = " &&\n ".join(setup_steps) - install_clause = f"{install_command} &&" if install_command else "" + setup_result = subprocess.run(["/bin/bash", "-c", setup_command], + stdout=sys.stdout, + stderr=sys.stderr, + check=False) + if setup_result.returncode != 0: + raise RuntimeError( + f"Setup command failed with exit code {setup_result.returncode}") - command = f""" - source $HOME/.local/bin/env && - source /opt/venv/bin/activate && - cd /FastVideo && - {install_clause} - {build_kernel_command} - {test_command} - """ + command = " &&\n ".join([ + "source $HOME/.local/bin/env", + "source /opt/venv/bin/activate", + "cd /FastVideo", + test_command, + ]) result = subprocess.run(["/bin/bash", "-c", command], stdout=sys.stdout, stderr=sys.stderr, @@ -369,7 +382,9 @@ def run_unit_test(): "./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ " "./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/worker/ " "./fastvideo/tests/training/test_trackers.py " - "./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py ./fastvideo/tests/modal/test_pr_test.py " + "./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py " + "./fastvideo/tests/modal/test_kernel_build_cache.py ./fastvideo/tests/modal/test_pr_test.py " + "./fastvideo/tests/modal/test_ssim_test.py " "--ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py " "--ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs" ) diff --git a/fastvideo/tests/modal/ssim_test.py b/fastvideo/tests/modal/ssim_test.py index 1b904421c..91033dea2 100644 --- a/fastvideo/tests/modal/ssim_test.py +++ b/fastvideo/tests/modal/ssim_test.py @@ -507,9 +507,7 @@ def _prepare_ssim_workspace( rm -rf fastvideo/tests/ssim/reference_videos git_retry git submodule update --init --recursive uv pip install -e ".[test]" - cd fastvideo-kernel - ./build.sh - cd .. + python fastvideo/tests/modal/kernel_build_cache.py install uv pip install git+https://github.com/microsoft/MoGe.git # Stable Audio Open 1.0 inference deps (optional in basic install, # required by `StableAudioDenoisingStage`; consumed by @@ -527,9 +525,11 @@ def _prepare_ssim_workspace( "HF_API_KEY": hf_api_key, }, ) + if result.stdout: + print(result.stdout, end="" if result.stdout.endswith("\n") else "\n") + if result.stderr: + print(result.stderr, end="" if result.stderr.endswith("\n") else "\n", file=sys.stderr) if result.returncode != 0: - print(result.stdout) - print(result.stderr) raise RuntimeError(f"Workspace setup failed with exit code {result.returncode}") ssim_dir = os.path.join(repo_root, "fastvideo", "tests", "ssim") diff --git a/fastvideo/tests/modal/test_kernel_build_cache.py b/fastvideo/tests/modal/test_kernel_build_cache.py new file mode 100644 index 000000000..5047eb575 --- /dev/null +++ b/fastvideo/tests/modal/test_kernel_build_cache.py @@ -0,0 +1,429 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import errno +import importlib.util +import json +import subprocess +import zipfile +from pathlib import Path + +import pytest +import yaml + + +def _load_kernel_build_cache(): + module_path = Path(__file__).with_name("kernel_build_cache.py") + spec = importlib.util.spec_from_file_location("kernel_build_cache_for_test", module_path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +kernel_build_cache = _load_kernel_build_cache() +REPO_ROOT = Path(__file__).resolve().parents[3] +WHEEL_NAME = "fastvideo_kernel-0.3.2-py3-none-any.whl" +TRUSTED_KERNEL_IMAGE_INPUTS = { + ".dockerignore", + ".github/workflows/_template-build-image.yml", + ".github/workflows/infra-build-image.yml", + ".gitmodules", + "docker/Dockerfile", + "docker/uv-excludes", + "fastvideo-kernel/**", + "fastvideo/tests/modal/kernel_build_cache.py", + "pyproject.toml", +} + + +def _write_wheel(path: Path, *, distribution: str = "fastvideo-kernel", payload: str = "payload") -> Path: + dist_info = "fastvideo_kernel-0.3.2.dist-info" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("fastvideo_kernel/__init__.py", payload) + archive.writestr( + f"{dist_info}/WHEEL", + "Wheel-Version: 1.0\nGenerator: test\nRoot-Is-Purelib: true\nTag: py3-none-any\n", + ) + archive.writestr( + f"{dist_info}/METADATA", + f"Metadata-Version: 2.1\nName: {distribution}\nVersion: 0.3.2\n", + ) + return path + + +def _metadata(cache_key: str = "cache-key") -> dict[str, object]: + return { + "cache_key": cache_key, + "schema_version": kernel_build_cache.CACHE_SCHEMA_VERSION, + } + + +def _write_cache_entry(cache_root: Path, cache_key: str, *, payload: str = "payload") -> Path: + cache_entry = cache_root / cache_key + cache_entry.mkdir() + wheel = _write_wheel(cache_entry / WHEEL_NAME, payload=payload) + stored_metadata = { + **_metadata(cache_key), + "artifact": kernel_build_cache._wheel_artifact(wheel), + } + (cache_entry / kernel_build_cache.METADATA_FILE).write_text(json.dumps(stored_metadata), encoding="utf-8") + return cache_entry + + +def _write_prebuilt_artifact(prebuilt_root: Path, name: str, cache_key: str) -> Path: + artifact_dir = prebuilt_root / name + artifact_dir.mkdir(parents=True) + wheel = _write_wheel(artifact_dir / WHEEL_NAME) + payload = { + **_metadata(cache_key), + "artifact": kernel_build_cache._wheel_artifact(wheel), + "wheel_path": str(wheel), + } + (artifact_dir / kernel_build_cache.METADATA_FILE).write_text(json.dumps(payload), encoding="utf-8") + return wheel + + +def _patch_stable_metadata(monkeypatch) -> None: + for name in ( + "GPU_BACKEND", + "CMAKE_ARGS", + "CFLAGS", + "CXXFLAGS", + "LDFLAGS", + "TORCH_CUDA_ARCH_LIST", + "FASTVIDEO_CONTAINER_IMAGE_REF", + "CUDACXX", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(kernel_build_cache, "_kernel_source_hash", lambda repo_root: {"source_hash": "source"}) + monkeypatch.setattr(kernel_build_cache, "_read_kernel_version", lambda repo_root: "0.3.2") + monkeypatch.setattr( + kernel_build_cache, + "_torch_metadata", + lambda: { + "torch_version": "2.9.0", + "torch_cuda_version": "12.8", + "torch_file": "/opt/venv/lib/python3.12/site-packages/torch/__init__.py", + "torch_config": "USE_CUDA=ON", + "cxx11_abi": True, + }, + ) + monkeypatch.setattr( + kernel_build_cache, + "_compiler_libc_metadata", + lambda: { + "compiler": {"cc": {"path": "/usr/bin/gcc", "version": "gcc 11.4.0"}}, + "libc": {"ldd_version": "ldd 2.35"}, + }, + ) + monkeypatch.setattr( + kernel_build_cache, + "_selected_command_metadata", + lambda environment_name, default_command: { + "raw": default_command, + "command": json.dumps([default_command]), + "resolved_executable": f"/usr/bin/{default_command}", + "version": f"{default_command} 12.8", + }, + ) + + +def test_kernel_only_main_change_republishes_trusted_l40s_artifact() -> None: + workflow = yaml.load( + (REPO_ROOT / ".github/workflows/infra-build-image.yml").read_text(encoding="utf-8"), + Loader=yaml.BaseLoader, + ) + push = workflow["on"]["push"] + + assert push["branches"] == ["main"] + assert TRUSTED_KERNEL_IMAGE_INPUTS <= set(push["paths"]) + assert "github.event_name == 'push'" in workflow["jobs"]["build-cuda-images"]["if"] + + dockerfile = (REPO_ROOT / "docker/Dockerfile").read_text(encoding="utf-8") + assert 'l40s_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/8.9"' in dockerfile + assert 'export TORCH_CUDA_ARCH_LIST=8.9' in dockerfile + assert '--output "${l40s_wheel_dir}/metadata.json"' in dockerfile + + +def test_cache_key_uses_resolved_arch_not_raw_env(monkeypatch, tmp_path) -> None: + _patch_stable_metadata(monkeypatch) + monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "9.0a") + + monkeypatch.setenv("TORCH_CUDA_ARCH_LIST", "9.0a") + explicit_hopper = kernel_build_cache._build_metadata(tmp_path) + + monkeypatch.delenv("TORCH_CUDA_ARCH_LIST", raising=False) + detected_hopper = kernel_build_cache._build_metadata(tmp_path) + + monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "8.9") + detected_l40s = kernel_build_cache._build_metadata(tmp_path) + + assert explicit_hopper["cache_key"] == detected_hopper["cache_key"] + assert explicit_hopper["build"]["torch_cuda_arch_list"] == "9.0a" + assert detected_hopper["build"]["torch_cuda_arch_list"] == "" + assert detected_l40s["cache_key"] != detected_hopper["cache_key"] + + +@pytest.mark.parametrize("environment_name", ["CFLAGS", "CXXFLAGS", "LDFLAGS"]) +def test_cache_key_changes_with_build_flags(monkeypatch, tmp_path, environment_name) -> None: + _patch_stable_metadata(monkeypatch) + monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "9.0a") + baseline = kernel_build_cache._build_metadata(tmp_path) + + monkeypatch.setenv(environment_name, "-DFASTVIDEO_ABI_VARIANT=1") + changed = kernel_build_cache._build_metadata(tmp_path) + + assert baseline["cache_key"] != changed["cache_key"] + + +def test_cache_key_changes_with_compiler_or_torch_abi(monkeypatch, tmp_path) -> None: + _patch_stable_metadata(monkeypatch) + monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "9.0a") + baseline = kernel_build_cache._build_metadata(tmp_path) + + monkeypatch.setattr( + kernel_build_cache, + "_compiler_libc_metadata", + lambda: { + "compiler": {"cc": {"path": "/opt/gcc-12/bin/gcc", "version": "gcc 12.3.0"}}, + "libc": {"ldd_version": "ldd 2.35"}, + }, + ) + compiler_changed = kernel_build_cache._build_metadata(tmp_path) + + _patch_stable_metadata(monkeypatch) + torch_metadata = kernel_build_cache._torch_metadata() + monkeypatch.setattr(kernel_build_cache, "_torch_metadata", lambda: {**torch_metadata, "cxx11_abi": False}) + torch_abi_changed = kernel_build_cache._build_metadata(tmp_path) + + assert baseline["schema_version"] == 3 + assert baseline["cache_key"] != compiler_changed["cache_key"] + assert baseline["cache_key"] != torch_abi_changed["cache_key"] + + +def test_find_wheel_uses_fastvideo_kernel_wheel_patterns(tmp_path) -> None: + (tmp_path / "unrelated-0.1.0.whl").write_text("") + (tmp_path / "fastvideo_kernel-0.3.1-py3-none-any.whl").write_text("") + expected = tmp_path / WHEEL_NAME + expected.write_text("") + + assert kernel_build_cache._find_wheel(tmp_path) == expected + + +def test_wheel_artifact_validates_archive_digest_size_and_distribution(tmp_path) -> None: + wheel = _write_wheel(tmp_path / WHEEL_NAME) + artifact = kernel_build_cache._wheel_artifact(wheel) + + valid, reason = kernel_build_cache._validate_wheel(wheel, artifact) + + assert valid + assert reason == "" + assert artifact["distribution"] == "fastvideo-kernel" + assert artifact["size_bytes"] == wheel.stat().st_size + assert len(artifact["sha256"]) == 64 + + +@pytest.mark.parametrize("corruption", ["digest", "size", "archive", "distribution"]) +def test_wheel_validation_rejects_corrupt_or_unexpected_artifacts(tmp_path, corruption) -> None: + wheel = _write_wheel(tmp_path / WHEEL_NAME) + artifact = kernel_build_cache._wheel_artifact(wheel) + if corruption == "digest": + artifact["sha256"] = "0" * 64 + elif corruption == "size": + artifact["size_bytes"] = int(artifact["size_bytes"]) + 1 + elif corruption == "archive": + wheel.write_bytes(b"not a wheel") + else: + wheel = _write_wheel(tmp_path / "unexpected.whl", distribution="other-package") + artifact = {**artifact, "wheel_name": wheel.name} + + valid, reason = kernel_build_cache._validate_wheel(wheel, artifact) + + assert not valid + assert reason + + +def test_store_cache_entry_handles_concurrent_valid_entry(monkeypatch, tmp_path) -> None: + cache_root = tmp_path / "cache" + cache_root.mkdir() + wheel = _write_wheel(tmp_path / WHEEL_NAME) + rename_calls = [] + + def fake_rename(self, target): + rename_calls.append((self, target)) + _write_cache_entry(target.parent, target.name) + raise OSError(errno.ENOTEMPTY, "Directory not empty") + + monkeypatch.setattr(Path, "rename", fake_rename) + + cache_entry = kernel_build_cache._store_cache_entry(cache_root, _metadata(), wheel) + + assert cache_entry == cache_root / "cache-key" + assert rename_calls + assert not list(cache_root.glob(".cache-key.tmp-*")) + + +def test_store_cache_entry_evicts_existing_invalid_entry(tmp_path) -> None: + cache_root = tmp_path / "cache" + cache_root.mkdir() + invalid_entry = cache_root / "cache-key" + invalid_entry.mkdir() + sentinel = invalid_entry / "stale" + sentinel.write_text("invalid") + wheel = _write_wheel(tmp_path / WHEEL_NAME) + + cache_entry = kernel_build_cache._store_cache_entry(cache_root, _metadata(), wheel) + cached_wheel, reason = kernel_build_cache._cache_entry_wheel(cache_entry, "cache-key") + + assert cached_wheel == cache_entry / WHEEL_NAME + assert reason == "" + assert not sentinel.exists() + + +def test_store_cache_entry_raises_when_renamed_entry_fails_validation(monkeypatch, tmp_path) -> None: + cache_root = tmp_path / "cache" + cache_root.mkdir() + wheel = _write_wheel(tmp_path / WHEEL_NAME) + original_rename = Path.rename + + def fake_rename(self, target): + original_rename(self, target) + metadata_path = target / kernel_build_cache.METADATA_FILE + payload = json.loads(metadata_path.read_text(encoding="utf-8")) + payload["cache_key"] = "other-key" + metadata_path.write_text(json.dumps(payload), encoding="utf-8") + + monkeypatch.setattr(Path, "rename", fake_rename) + + with pytest.raises(RuntimeError, match="could not be validated"): + kernel_build_cache._store_cache_entry(cache_root, _metadata(), wheel) + + assert not (cache_root / "cache-key").exists() + + +def test_consumer_installs_valid_cache_hit_without_build(monkeypatch, tmp_path) -> None: + cache_root = tmp_path / "cache" + cache_root.mkdir() + cache_entry = _write_cache_entry(cache_root, "cache-key") + installed = [] + monkeypatch.setattr(kernel_build_cache, "_build_metadata", lambda repo_root: _metadata()) + monkeypatch.setattr(kernel_build_cache, "_try_install_prebuilt", lambda metadata, path: False) + monkeypatch.setattr(kernel_build_cache, "_install_wheel", installed.append) + monkeypatch.setattr( + kernel_build_cache, + "_build_and_install_local", + lambda repo_root, metadata: pytest.fail("cache hit unexpectedly built locally"), + ) + + kernel_build_cache.install_cached_or_build(tmp_path, cache_root, tmp_path / "missing-prebuilt.json") + + assert installed == [cache_entry / WHEEL_NAME] + + +def test_consumer_rejects_corrupt_cache_without_mutating_it(monkeypatch, tmp_path) -> None: + cache_root = tmp_path / "cache" + cache_root.mkdir() + cache_entry = _write_cache_entry(cache_root, "cache-key") + wheel = cache_entry / WHEEL_NAME + wheel.write_bytes(b"corrupt") + before_metadata = (cache_entry / kernel_build_cache.METADATA_FILE).read_bytes() + built = [] + monkeypatch.setattr(kernel_build_cache, "_build_metadata", lambda repo_root: _metadata()) + monkeypatch.setattr(kernel_build_cache, "_try_install_prebuilt", lambda metadata, path: False) + monkeypatch.setattr(kernel_build_cache, "_build_and_install_local", lambda repo_root, metadata: built.append(True)) + + kernel_build_cache.install_cached_or_build(tmp_path, cache_root, tmp_path / "missing-prebuilt.json") + + assert built == [True] + assert wheel.read_bytes() == b"corrupt" + assert (cache_entry / kernel_build_cache.METADATA_FILE).read_bytes() == before_metadata + + +def test_consumer_falls_back_when_cached_install_fails(monkeypatch, tmp_path) -> None: + cache_root = tmp_path / "cache" + cache_root.mkdir() + _write_cache_entry(cache_root, "cache-key") + built = [] + monkeypatch.setattr(kernel_build_cache, "_build_metadata", lambda repo_root: _metadata()) + monkeypatch.setattr(kernel_build_cache, "_try_install_prebuilt", lambda metadata, path: False) + monkeypatch.setattr( + kernel_build_cache, + "_install_wheel", + lambda wheel: (_ for _ in ()).throw(subprocess.CalledProcessError(1, ["uv"])), + ) + monkeypatch.setattr(kernel_build_cache, "_build_and_install_local", lambda repo_root, metadata: built.append(True)) + + kernel_build_cache.install_cached_or_build(tmp_path, cache_root, tmp_path / "missing-prebuilt.json") + + assert built == [True] + + +def test_normal_l40s_pr_ssim_install_uses_docker_prebuilt_without_build(monkeypatch, tmp_path, capsys) -> None: + prebuilt_root = tmp_path / "prebuilt" + _write_prebuilt_artifact(prebuilt_root, "9.0a", "hopper-key") + l40s_wheel = _write_prebuilt_artifact(prebuilt_root, "8.9", "l40s-key") + installed = [] + monkeypatch.setattr(kernel_build_cache, "_build_metadata", lambda repo_root: _metadata("l40s-key")) + monkeypatch.setattr(kernel_build_cache, "_install_wheel", installed.append) + monkeypatch.setattr( + kernel_build_cache, + "_build_and_install_local", + lambda repo_root, metadata: pytest.fail("L40S prebuilt hit unexpectedly called build.sh"), + ) + + kernel_build_cache.install_cached_or_build( + tmp_path, + cache_root=None, + prebuilt_info_path=prebuilt_root, + ) + + output = capsys.readouterr().out + assert installed == [l40s_wheel] + assert "Docker-prebuilt cache hit for key l40s-key" in output + assert "build.sh" not in output + + +def test_prebuilt_requires_validated_artifact(monkeypatch, tmp_path) -> None: + wheel = _write_wheel(tmp_path / WHEEL_NAME) + prebuilt_info = tmp_path / "prebuilt.json" + payload = { + **_metadata(), + "artifact": kernel_build_cache._wheel_artifact(wheel), + "wheel_path": str(wheel), + } + prebuilt_info.write_text(json.dumps(payload), encoding="utf-8") + installed = [] + monkeypatch.setattr(kernel_build_cache, "_install_wheel", installed.append) + + assert kernel_build_cache._try_install_prebuilt(_metadata(), prebuilt_info) + assert installed == [wheel] + + installed.clear() + wheel.write_bytes(b"corrupt") + assert not kernel_build_cache._try_install_prebuilt(_metadata(), prebuilt_info) + assert installed == [] + + +def test_producer_repairs_invalid_entry_and_stores_valid_wheel(monkeypatch, tmp_path) -> None: + cache_root = tmp_path / "cache" + cache_root.mkdir() + invalid_entry = cache_root / "cache-key" + invalid_entry.mkdir() + (invalid_entry / kernel_build_cache.METADATA_FILE).write_text("not json", encoding="utf-8") + built_dir = tmp_path / "built" + built_dir.mkdir() + built_wheel = _write_wheel(built_dir / WHEEL_NAME) + installed = [] + monkeypatch.setattr(kernel_build_cache, "_build_metadata", lambda repo_root: _metadata()) + monkeypatch.setattr(kernel_build_cache, "_build_wheel", lambda repo_root, metadata: built_wheel) + monkeypatch.setattr(kernel_build_cache, "_install_wheel", installed.append) + + kernel_build_cache.produce_cached_wheel(tmp_path, cache_root) + + cached_wheel, reason = kernel_build_cache._cache_entry_wheel(cache_root / "cache-key", "cache-key") + assert installed == [built_wheel] + assert cached_wheel == cache_root / "cache-key" / WHEEL_NAME + assert reason == "" diff --git a/fastvideo/tests/modal/test_pr_test.py b/fastvideo/tests/modal/test_pr_test.py index 0315043a3..7190d0bd0 100644 --- a/fastvideo/tests/modal/test_pr_test.py +++ b/fastvideo/tests/modal/test_pr_test.py @@ -265,13 +265,61 @@ def test_run_test_command_composes_valid_post_checkout_shell( install_command=install_command, ) - assert [event for event, _ in events] == ["checkout", "run"] - shell_command = events[1][1][-1] - assert "cd /FastVideo" in shell_command - assert "git clone" not in shell_command - assert ("./build.sh" in shell_command) is build_kernel + assert [event for event, _ in events] == ["checkout", "run", "run"] + setup_command = events[1][1][-1] + test_command = events[2][1][-1] + assert "cd /FastVideo" in setup_command + assert "git clone" not in setup_command + assert ("kernel_build_cache.py install" in setup_command) is build_kernel + assert "pytest fastvideo/tests/api -q" not in setup_command + assert "kernel_build_cache.py install" not in test_command + assert "pytest fastvideo/tests/api -q" in test_command if install_command: - assert install_command in shell_command + assert install_command in setup_command else: - assert "uv pip install -e" not in shell_command - real_run(["/bin/bash", "-n"], input=shell_command, text=True, check=True) + assert "uv pip install -e" not in setup_command + for shell_command in (setup_command, test_command): + real_run(["/bin/bash", "-n"], input=shell_command, text=True, check=True) + + +def test_run_test_command_uses_nonshared_kernel_install_before_tests(monkeypatch): + module = _load_pr_test_module(monkeypatch) + commands = [] + + def fake_run(args, **_kwargs): + commands.append(args[-1]) + return types.SimpleNamespace(returncode=0) + + monkeypatch.setenv("BUILDKITE_REPO", "https://example.com/FastVideo.git") + monkeypatch.setenv("BUILDKITE_COMMIT", "0123456789abcdef") + monkeypatch.setenv("BUILDKITE_PULL_REQUEST", "false") + monkeypatch.setattr(module, "_checkout_repository", + lambda *_args: None) + monkeypatch.setattr(subprocess, "run", fake_run) + + module.run_test_command("pytest fastvideo/tests/api -q", build_kernel=True) + + assert len(commands) == 2 + setup_command, test_command = commands + assert "kernel_build_cache.py install" in setup_command + assert "--cache-root" not in setup_command + assert "pytest fastvideo/tests/api -q" not in setup_command + assert "kernel_build_cache.py install" not in test_command + assert "pytest fastvideo/tests/api -q" in test_command + assert not hasattr(module, "kernel_cache_vol") + + +def test_run_unit_test_collects_modal_cache_runner_tests(monkeypatch): + module = _load_pr_test_module(monkeypatch) + commands = [] + monkeypatch.setattr(module, "run_test", commands.append) + + module.run_unit_test() + + assert len(commands) == 1 + for test_path in ( + "./fastvideo/tests/modal/test_kernel_build_cache.py", + "./fastvideo/tests/modal/test_pr_test.py", + "./fastvideo/tests/modal/test_ssim_test.py", + ): + assert test_path in commands[0] diff --git a/fastvideo/tests/modal/test_ssim_test.py b/fastvideo/tests/modal/test_ssim_test.py index 05cec7d76..574307740 100644 --- a/fastvideo/tests/modal/test_ssim_test.py +++ b/fastvideo/tests/modal/test_ssim_test.py @@ -99,3 +99,36 @@ def test_spawn_ssim_task_passes_hf_api_key_to_pytest_env(monkeypatch, tmp_path): assert env["FASTVIDEO_SSIM_MODEL_ID"] == "model-a" assert running_task.process.__class__ is FakeProcess running_task.log_handle.close() + + +def test_prepare_workspace_uses_nonshared_kernel_install(monkeypatch): + module = _load_ssim_test_module(monkeypatch) + captured = {} + task = module.SSIMTask( + task_id=0, + test_file="fastvideo/tests/ssim/test_example.py", + required_gpus=1, + ) + + def fake_run(args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + return types.SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(module.subprocess, "run", fake_run) + monkeypatch.setattr(module, "_discover_ssim_tasks", lambda *_args, **_kwargs: [task]) + + repo_root, tasks = module._prepare_ssim_workspace( + git_repo="https://example.com/FastVideo.git", + git_commit="0123456789abcdef", + pr_number="false", + hf_api_key="hf_test_token", + ) + + setup_command = captured["args"][2] + assert "kernel_build_cache.py install" in setup_command + assert "--cache-root" not in setup_command + assert repo_root == "/FastVideo" + assert tasks == [task] + assert module.SSIM_COMMON_KWARGS["volumes"] == {"/root/data": module.model_vol} + assert not hasattr(module, "kernel_cache_vol")