tilelang/testing/python/cache/test_tilelang_kernel_cache_integrity.py

"""Regression tests for cache-entry integrity verification.

A crashed writer (or crashed network-filesystem client) can publish a cache
entry whose kernel_lib.so is truncated or partially written. Such a binary can
dlopen fine and only fail at first kernel launch with CUDA_ERROR_INVALID_IMAGE,
where nothing repairs the entry. The kernel cache therefore records a manifest
of file sizes/hashes on save and verifies it on load, treating mismatches as
cache misses.
"""

import json
from hashlib import sha256
from pathlib import Path

import pytest

from tilelang import tvm

import tilelang.cache.kernel_cache as kernel_cache_mod
from tilelang.backend import create_backend_context
from tilelang.cache.kernel_cache import KernelCache
from tilelang.engine.param import KernelParam, dump_kernel_params
from tilelang.env import env


class _FakeAdapter:
    def __init__(self, libpath: str):
        self.libpath = libpath

    def get_kernel_source(self):
        return "// host kernel"


class _FakeKernel:
    def __init__(self, libpath: str):
        self.adapter = _FakeAdapter(libpath)
        self.kernel_source = "// device kernel"
        self.params = [KernelParam(tvm.DataType("float32"), [4])]


@pytest.fixture
def cache_dirs(tmp_path, monkeypatch):
    cache_dir = tmp_path / "cache"
    cache_dir.mkdir()
    monkeypatch.setattr(env, "TILELANG_CACHE_DIR", str(cache_dir))
    return cache_dir


def _make_fake_kernel(tmp_path):
    lib_path = tmp_path / "kernel_lib.so"
    lib_path.write_bytes(b"fake-so-contents")
    return _FakeKernel(str(lib_path))


def _load_expecting_no_build(cache: KernelCache, key: str, monkeypatch):
    def fail_from_database(cls, **kwargs):
        raise AssertionError("corrupted cache entries must miss before rebuilding from database")

    monkeypatch.setattr(kernel_cache_mod.JITKernel, "from_database", classmethod(fail_from_database))
    return cache._load_kernel_from_disk(
        key,
        backend_context=create_backend_context("cuda", execution_backend="tvm_ffi"),
        out_idx=[0],
        pass_configs=None,
        compile_flags=None,
        func=None,
    )


def _load_expecting_hit(cache: KernelCache, key: str, monkeypatch):
    sentinel = object()
    monkeypatch.setattr(kernel_cache_mod.JITKernel, "from_database", classmethod(lambda cls, **kwargs: sentinel))
    loaded = cache._load_kernel_from_disk(
        key,
        backend_context=create_backend_context("cuda", execution_backend="tvm_ffi"),
        out_idx=[0],
        pass_configs=None,
        compile_flags=None,
        func=None,
    )
    return loaded is sentinel


def test_save_writes_manifest_covering_all_files(cache_dirs, tmp_path):
    cache = KernelCache()
    key = "manifest-on-save"
    cache._save_kernel_to_disk(key, _make_fake_kernel(tmp_path))

    cache_path = Path(cache._get_cache_path(key))
    manifest = json.loads((cache_path / cache.manifest_path).read_text())

    expected_files = {
        cache.device_kernel_path,
        cache.host_kernel_path,
        cache.kernel_lib_path,
        cache.params_path,
    }
    assert set(manifest["files"]) == expected_files
    lib_bytes = (cache_path / cache.kernel_lib_path).read_bytes()
    assert manifest["files"][cache.kernel_lib_path]["size"] == len(lib_bytes)
    assert manifest["files"][cache.kernel_lib_path]["sha256"] == sha256(lib_bytes).hexdigest()


def test_load_removes_entry_with_truncated_kernel_lib(cache_dirs, tmp_path, monkeypatch):
    cache = KernelCache()
    key = "truncated-lib"
    cache._save_kernel_to_disk(key, _make_fake_kernel(tmp_path))
    cache_path = Path(cache._get_cache_path(key))

    lib_file = cache_path / cache.kernel_lib_path
    lib_file.write_bytes(lib_file.read_bytes()[: len(b"fake-so-contents") // 2])

    assert _load_expecting_no_build(cache, key, monkeypatch) is None
    assert not cache_path.exists()


def test_load_removes_entry_with_same_size_corruption(cache_dirs, tmp_path, monkeypatch):
    cache = KernelCache()
    key = "bitflipped-lib"
    cache._save_kernel_to_disk(key, _make_fake_kernel(tmp_path))
    cache_path = Path(cache._get_cache_path(key))

    lib_file = cache_path / cache.kernel_lib_path
    original = lib_file.read_bytes()
    lib_file.write_bytes(b"\x00" * len(original))

    assert _load_expecting_no_build(cache, key, monkeypatch) is None
    assert not cache_path.exists()


def test_source_files_are_size_checked_but_not_hashed(cache_dirs, tmp_path, monkeypatch):
    # Hashing sources on load would regress lazy source loading; only their
    # size is checked, so a same-size rewrite of a source file must not
    # invalidate the entry.
    cache = KernelCache()
    key = "source-not-hashed"
    cache._save_kernel_to_disk(key, _make_fake_kernel(tmp_path))
    cache_path = Path(cache._get_cache_path(key))

    device_file = cache_path / cache.device_kernel_path
    device_file.write_text("X" * len(device_file.read_text()))

    assert _load_expecting_hit(cache, key, monkeypatch)


def test_legacy_entry_without_manifest_misses_and_save_repairs(cache_dirs, tmp_path, monkeypatch):
    cache = KernelCache()
    key = "legacy-no-manifest"
    cache_path = Path(cache._get_cache_path(key))
    cache_path.mkdir(parents=True)
    (cache_path / cache.device_kernel_path).write_text("// device kernel")
    (cache_path / cache.host_kernel_path).write_text("// host kernel")
    (cache_path / cache.kernel_lib_path).write_bytes(b"legacy-truncated")
    (cache_path / cache.params_path).write_text(dump_kernel_params([KernelParam(tvm.DataType("float32"), [4])]))

    assert _load_expecting_no_build(cache, key, monkeypatch) is None

    cache._save_kernel_to_disk(key, _make_fake_kernel(tmp_path))
    assert (cache_path / cache.manifest_path).exists()
    assert (cache_path / cache.kernel_lib_path).read_bytes() == b"fake-so-contents"
    assert _load_expecting_hit(cache, key, monkeypatch)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])