[CUDA][Cache] Publish immutable binary cache directories (#3177)

* [CUDA][Cache] Publish immutable binary cache directories * [Example] Drop stale torch JIT build locks in minference example The vertical_slash_index extension builds into a cache directory shared across CI runs. torch's FileBaton lock file is not tied to the owning process, so a run cancelled mid-compile leaves it behind and every later load() polls on it forever (hung CI at test_vs_sparse_attn). Guard the build with an OS-level flock, which dies with its holder, and remove any leftover FileBaton lock before invoking torch. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-10-26 07:22:54 +08:00 · 2026-09-08 02:46:47 +08:00
parent 62bba8d20d
commit 7525499330
6 changed files with 365 additions and 76 deletions
@@ -1,9 +1,19 @@
 from __future__ import annotations

+import builtins
+import errno
+import io
+import json
 import os
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from hashlib import sha256
+from pathlib import Path

+import pytest
 import tilelang
 from tilelang import tvm
+import tilelang.cache.cuda_binary_cache as cuda_binary_cache_mod
 import tilelang.cache.kernel_cache as kernel_cache_mod
 from tilelang.backend import create_backend_context
 from tilelang.cache.cuda_binary_cache import CUDABinaryCache
@@ -69,8 +79,11 @@ def test_cuda_binary_cache_hit_skips_nvcc_compile(monkeypatch, tmp_path):
    # first compiles, second hits; third compiles (new options), fourth hits
    assert len(compile_calls) == 2
    assert compile_calls[0][3] != compile_calls[1][3]
-    cache_files = list((tmp_path / "cache").glob("*/cuda-binaries/*.cubin"))
+    cache_files = list((tmp_path / "cache").glob("*/cuda-binaries/*/kernel.cubin"))
    assert len(cache_files) == 2
+    for cache_file in cache_files:
+        assert cache_file.read_bytes() == b"fake-cubin"
+        assert (cache_file.parent / "metadata.json").is_file()


 def test_cuda_binary_cache_corrupted_entry_recompiles(monkeypatch, tmp_path):
@@ -93,30 +106,281 @@ def test_cuda_binary_cache_corrupted_entry_recompiles(monkeypatch, tmp_path):
    cuda_backend.tilelang_callback_cuda_compile(source, target)
    assert len(compile_calls) == 1

-    [cache_file] = (tmp_path / "cache").glob("*/cuda-binaries/*.cubin")
-    assert cache_file.with_name(cache_file.name + ".sha256").exists()
+    [cache_file] = (tmp_path / "cache").glob("*/cuda-binaries/*/kernel.cubin")
    # Same-size corruption, as left behind by a crashed writer/filesystem client.
-    cache_file.write_bytes(b"\x00" * len(b"fake-cubin"))
+    corrupted = b"\x00" * cache_file.stat().st_size
+    cache_file.write_bytes(corrupted)
+    metadata_path = cache_file.parent / "metadata.json"
+    original_metadata = metadata_path.read_bytes()

    recompiled = cuda_backend.tilelang_callback_cuda_compile(source, target)
    assert bytes(recompiled) == b"fake-cubin"
    assert len(compile_calls) == 2

-    # The corrupted entry was rewritten, so the next call hits the cache again.
-    cuda_backend.tilelang_callback_cuda_compile(source, target)
-    assert len(compile_calls) == 2
+    # Shared entries stay immutable even after a miss. Recompilation is usable
+    # for this call, but repairing the on-disk entry needs offline cleanup.
+    assert bytes(cuda_backend.tilelang_callback_cuda_compile(source, target)) == b"fake-cubin"
+    assert len(compile_calls) == 3
+    assert cache_file.read_bytes() == corrupted
+    assert metadata_path.read_bytes() == original_metadata


-def test_cuda_binary_cache_accepts_legacy_entry_without_sidecar(monkeypatch, tmp_path):
+def test_cuda_binary_cache_directory_coexists_with_legacy_sidecar_entry(monkeypatch, tmp_path):
    _set_cache_dirs(monkeypatch, tmp_path)

    key = "legacy-key"
-    path = CUDABinaryCache.get_path(key, "cubin")
-    os.makedirs(os.path.dirname(path), exist_ok=True)
-    with open(path, "wb") as f:
+    cache_root = CUDABinaryCache._get_cache_root()
+    legacy_path = os.path.join(cache_root, f"{key}.cubin")
+    os.makedirs(cache_root, exist_ok=True)
+    with open(legacy_path, "wb") as f:
        f.write(b"legacy-cubin")
+    with open(legacy_path + ".sha256", "w") as f:
+        f.write(sha256(b"legacy-cubin").hexdigest())

-    assert CUDABinaryCache.load(key, "cubin") == b"legacy-cubin"
+    assert CUDABinaryCache.load(key, "cubin") is None
+    assert os.path.exists(legacy_path)
+
+    CUDABinaryCache.save(key, "cubin", b"new-cubin")
+
+    assert CUDABinaryCache.load(key, "cubin") == b"new-cubin"
+    assert Path(legacy_path).read_bytes() == b"legacy-cubin"
+    assert Path(legacy_path + ".sha256").read_text() == sha256(b"legacy-cubin").hexdigest()
+
+
+@pytest.mark.parametrize("failure", ["missing", "empty", "malformed", "invalid-encoding", "wrong-type", "unreadable"])
+def test_cuda_binary_cache_rejects_bad_metadata(monkeypatch, tmp_path, failure):
+    _set_cache_dirs(monkeypatch, tmp_path)
+
+    key = "bad-metadata-key"
+    CUDABinaryCache.save(key, "cubin", b"valid-cubin")
+    path = Path(CUDABinaryCache.get_path(key, "cubin"))
+    metadata_path = path.parent / "metadata.json"
+    metadata_path.unlink()
+    if failure == "unreadable":
+        metadata_path.mkdir()
+    elif failure != "missing":
+        contents = {"empty": b"", "malformed": b"{", "invalid-encoding": b"\xff", "wrong-type": b"[]"}
+        metadata_path.write_bytes(contents[failure])
+
+    assert CUDABinaryCache.load(key, "cubin") is None
+    assert path.read_bytes() == b"valid-cubin"
+
+    # A later writer must leave even an invalid published directory alone.
+    CUDABinaryCache.save(key, "cubin", b"new-cubin")
+    assert path.read_bytes() == b"valid-cubin"
+    assert CUDABinaryCache.load(key, "cubin") is None
+
+
+@pytest.mark.parametrize(
+    "field,value",
+    [
+        ("size", None),
+        ("size", True),
+        ("size", 0),
+        ("size", "11"),
+        ("sha256", None),
+        ("sha256", ""),
+        ("sha256", "x" * 64),
+        ("sha256", " " * 64),
+        ("format", "tilelang.cuda-binary-cache.v2"),
+    ],
+)
+def test_cuda_binary_cache_rejects_invalid_metadata_fields(monkeypatch, tmp_path, field, value):
+    _set_cache_dirs(monkeypatch, tmp_path)
+
+    key = "invalid-metadata-key"
+    CUDABinaryCache.save(key, "cubin", b"valid-cubin")
+    path = Path(CUDABinaryCache.get_path(key, "cubin"))
+    metadata_path = path.parent / "metadata.json"
+    metadata = json.loads(metadata_path.read_text())
+    if value is None:
+        del metadata[field]
+    else:
+        metadata[field] = value
+    metadata_path.write_text(json.dumps(metadata))
+
+    assert CUDABinaryCache.load(key, "cubin") is None
+    assert path.exists()
+    assert json.loads(metadata_path.read_text()) == metadata
+
+
+def test_cuda_binary_cache_rejects_empty_entry(monkeypatch, tmp_path):
+    _set_cache_dirs(monkeypatch, tmp_path)
+
+    key = "empty-key"
+    CUDABinaryCache.save(key, "fatbin", b"valid-fatbin")
+    path = Path(CUDABinaryCache.get_path(key, "fatbin"))
+    path.write_bytes(b"")
+    metadata_path = path.parent / "metadata.json"
+    metadata = json.loads(metadata_path.read_text())
+    metadata.update(size=0, sha256=sha256(b"").hexdigest())
+    metadata_path.write_text(json.dumps(metadata))
+
+    assert CUDABinaryCache.load(key, "fatbin") is None
+    assert path.exists()
+
+
+@pytest.mark.parametrize("missing_metadata", [False, True])
+def test_cuda_binary_cache_empty_read_is_miss(monkeypatch, tmp_path, missing_metadata):
+    _set_cache_dirs(monkeypatch, tmp_path)
+
+    key = "empty-read-key"
+    CUDABinaryCache.save(key, "fatbin", b"valid-fatbin")
+    path = Path(CUDABinaryCache.get_path(key, "fatbin"))
+    metadata_path = path.parent / "metadata.json"
+
+    def empty_read(file, *args, **kwargs):
+        if Path(file) == path:
+            # Model a 3FS fd returning EOF despite a nonempty cached binary.
+            return io.BytesIO(b"")
+        if missing_metadata and Path(file) == metadata_path:
+            raise FileNotFoundError(errno.ENOENT, "metadata disappeared")
+        return builtins.open(file, *args, **kwargs)
+
+    monkeypatch.setattr(cuda_binary_cache_mod, "open", empty_read, raising=False)
+
+    assert CUDABinaryCache.load(key, "fatbin") is None
+    assert path.read_bytes() == b"valid-fatbin"
+    assert metadata_path.exists()
+
+
+def test_cuda_binary_cache_hash_mismatch_does_not_delete_entry(monkeypatch, tmp_path):
+    _set_cache_dirs(monkeypatch, tmp_path)
+
+    key = "mismatched-key"
+    CUDABinaryCache.save(key, "cubin", b"valid-cubin")
+    path = Path(CUDABinaryCache.get_path(key, "cubin"))
+    corrupted = b"other-cubin"
+    assert len(corrupted) == path.stat().st_size
+    path.write_bytes(corrupted)
+
+    assert CUDABinaryCache.load(key, "cubin") is None
+    assert path.read_bytes() == corrupted
+    assert (path.parent / "metadata.json").exists()
+
+
+def test_cuda_binary_cache_rejects_truncated_binary(monkeypatch, tmp_path):
+    _set_cache_dirs(monkeypatch, tmp_path)
+
+    key = "truncated-binary-key"
+    CUDABinaryCache.save(key, "cubin", b"valid-cubin")
+    path = Path(CUDABinaryCache.get_path(key, "cubin"))
+    path.write_bytes(path.read_bytes()[:-1])
+
+    assert CUDABinaryCache.load(key, "cubin") is None
+    assert path.exists()
+
+
+@pytest.mark.parametrize("compile_format", ["cubin", "fatbin"])
+def test_cuda_binary_cache_first_valid_writer_wins(monkeypatch, tmp_path, compile_format):
+    _set_cache_dirs(monkeypatch, tmp_path)
+
+    key = "first-writer-key"
+    path = Path(CUDABinaryCache.get_path(key, compile_format))
+    CUDABinaryCache.save(key, compile_format, b"first-binary")
+    first_inode = path.stat().st_ino
+    with path.open("rb") as reader:
+        CUDABinaryCache.save(key, compile_format, b"second-binary")
+        assert reader.read() == b"first-binary"
+
+    assert CUDABinaryCache.load(key, compile_format) == b"first-binary"
+    assert path.stat().st_ino == first_inode
+    assert sorted(p.name for p in path.parent.iterdir()) == [f"kernel.{compile_format}", "metadata.json"]
+
+
+def test_cuda_binary_cache_directory_publication_is_atomic(monkeypatch, tmp_path):
+    _set_cache_dirs(monkeypatch, tmp_path)
+
+    key = "atomic-key"
+    path = Path(CUDABinaryCache.get_path(key, "fatbin"))
+    synced_inodes = set()
+    real_fsync = os.fsync
+    real_rename = os.rename
+    publications = []
+
+    def track_fsync(fd):
+        real_fsync(fd)
+        stat = os.fstat(fd)
+        synced_inodes.add((stat.st_dev, stat.st_ino))
+
+    def check_publication(src, dst):
+        staging = Path(src)
+        assert staging.parent == Path(CUDABinaryCache._get_staging_root())
+        assert staging.stat().st_dev == path.parent.parent.stat().st_dev
+        assert not path.parent.exists()
+        assert CUDABinaryCache.load(key, "fatbin") is None
+        assert sorted(p.name for p in staging.iterdir()) == ["kernel.fatbin", "metadata.json"]
+        assert (staging / "kernel.fatbin").read_bytes() == b"valid-fatbin"
+        for entry in [staging, *staging.iterdir()]:
+            stat = entry.stat()
+            assert (stat.st_dev, stat.st_ino) in synced_inodes
+        real_rename(src, dst)
+        assert CUDABinaryCache.load(key, "fatbin") == b"valid-fatbin"
+        publications.append(dst)
+
+    monkeypatch.setattr(os, "fsync", track_fsync)
+    monkeypatch.setattr(os, "rename", check_publication)
+    CUDABinaryCache.save(key, "fatbin", b"valid-fatbin")
+
+    assert publications == [str(path.parent)]
+    assert not list(Path(CUDABinaryCache._get_staging_root()).iterdir())
+
+
+def test_cuda_binary_cache_concurrent_publishers_do_not_replace_winner(monkeypatch, tmp_path):
+    _set_cache_dirs(monkeypatch, tmp_path)
+
+    writer_count = 16
+    barrier = threading.Barrier(writer_count)
+    payloads = [f"cubin-{index}".encode() for index in range(writer_count)]
+    real_rename = os.rename
+    winners = []
+
+    def concurrent_rename(src, dst):
+        # Force every writer to stage its files before any can publish.
+        barrier.wait(timeout=10)
+        real_rename(src, dst)
+        winners.append((Path(dst) / "kernel.cubin").read_bytes())
+
+    monkeypatch.setattr(os, "rename", concurrent_rename)
+
+    with ThreadPoolExecutor(max_workers=writer_count) as executor:
+        list(executor.map(lambda data: CUDABinaryCache.save("concurrent-key", "cubin", data), payloads))
+
+    assert len(winners) == 1
+    assert winners[0] in payloads
+    assert CUDABinaryCache.load("concurrent-key", "cubin") == winners[0]
+    cache_entries = os.listdir(CUDABinaryCache._get_cache_root())
+    assert cache_entries == ["concurrent-key"]
+    assert not list(Path(CUDABinaryCache._get_staging_root()).iterdir())
+
+
+@pytest.mark.parametrize("operation", ["fsync", "rename"])
+def test_cuda_binary_cache_failed_publish_cleans_only_own_staging(monkeypatch, tmp_path, operation):
+    _set_cache_dirs(monkeypatch, tmp_path)
+
+    staging_root = Path(CUDABinaryCache._get_staging_root())
+    other_writer = staging_root / "other-writer"
+    other_writer.mkdir(parents=True)
+    (other_writer / "kernel.cubin").write_bytes(b"other-cubin")
+
+    def fail(*args, **kwargs):
+        raise OSError(errno.EIO, "injected I/O failure")
+
+    monkeypatch.setattr(os, operation, fail)
+    with pytest.raises(OSError, match="injected I/O failure"):
+        CUDABinaryCache.save("failed-publish-key", "cubin", b"valid-cubin")
+
+    assert CUDABinaryCache.load("failed-publish-key", "cubin") is None
+    assert not Path(CUDABinaryCache.get_path("failed-publish-key", "cubin")).parent.exists()
+    assert list(staging_root.iterdir()) == [other_writer]
+    assert (other_writer / "kernel.cubin").read_bytes() == b"other-cubin"
+
+
+def test_cuda_binary_cache_rejects_empty_save(monkeypatch, tmp_path):
+    _set_cache_dirs(monkeypatch, tmp_path)
+
+    with pytest.raises(ValueError, match="empty CUDA binary"):
+        CUDABinaryCache.save("empty-save-key", "fatbin", b"")


 def test_disk_cache_load_failure_is_cache_miss(monkeypatch, tmp_path):