fixup! feat(core.utils): couple key schema with backend schema; document pickle compat; add usage example

Author: cpcloud

cuda_core/cuda/core/utils/_program_cache.py

@@ -508,6 +508,34 @@ def make_program_cache_key(
         If ``extra_digest`` is ``None`` while ``options`` sets any option whose
         compilation effect depends on external file content that the key
         cannot otherwise observe.
+
+    Examples
+    --------
+    Wiring a cache around :class:`~cuda.core.Program` compile::
+
+        from cuda.core import Program, ProgramOptions
+        from cuda.core.utils import FileStreamProgramCache, make_program_cache_key
+
+        source = "extern \"C\" __global__ void k(int *a){ *a = 1; }"
+        options = ProgramOptions(arch="sm_80")
+
+        with FileStreamProgramCache("/var/cache/myapp/cuda") as cache:
+            key = make_program_cache_key(
+                code=source, code_type="c++", options=options, target_type="cubin"
+            )
+            obj = cache.get(key)
+            if obj is None:
+                obj = Program(source, "c++", options=options).compile("cubin")
+                cache[key] = obj
+
+    Options that read external files (``include_path``, ``pre_include``,
+    ``pch``, ``use_pch``, ``pch_dir``; and ``use_libdevice=True`` on the
+    NVVM path) require ``extra_digest`` -- fingerprint the bytes the
+    compiler will pull in and pass that digest so changes to those files
+    force a cache miss. Options that have compile-time side effects
+    (``create_pch``, ``time``, ``fdevice_time_trace``) cannot be cached
+    and raise ``ValueError``; compile directly, or disable the flag, for
+    those cases.
     """
     # Mirror Program.compile (_program.pyx Program_init lowercases code_type
     # before dispatch); a caller that passes "PTX" or "C++" must get the
@@ -764,7 +792,12 @@ def _probe(label: str, fn):
 # ---------------------------------------------------------------------------
 
 
-_SQLITE_SCHEMA_VERSION = "1"
+# Composite of (backend-storage schema, key schema): a bump in either one
+# forces a wipe-on-open. This keeps ``_KEY_SCHEMA_VERSION`` bumps from
+# leaving unreachable entries on disk (they would be orphaned forever,
+# since the new hash never collides with the old one).
+_SQLITE_BACKEND_SCHEMA = "1"
+_SQLITE_SCHEMA_VERSION = f"{_SQLITE_BACKEND_SCHEMA}.{_KEY_SCHEMA_VERSION}"
 
 
 class SQLiteProgramCache(ProgramCacheResource):
@@ -1066,7 +1099,9 @@ def _enforce_size_cap(self) -> None:
 # ---------------------------------------------------------------------------
 
 
-_FILESTREAM_SCHEMA_VERSION = 2
+# Composite of (backend-storage schema, key schema) -- see SQLite comment.
+_FILESTREAM_BACKEND_SCHEMA = 2
+_FILESTREAM_SCHEMA_VERSION = f"{_FILESTREAM_BACKEND_SCHEMA}.{_KEY_SCHEMA_VERSION}"
 _ENTRIES_SUBDIR = "entries"
 _TMP_SUBDIR = "tmp"
 _SCHEMA_FILE = "SCHEMA_VERSION"
@@ -1197,17 +1232,25 @@ class FileStreamProgramCache(ProgramCacheResource):
 
     .. note:: **Cross-version sharing.**
 
-        ``_FILESTREAM_SCHEMA_VERSION`` guards on-disk format changes: a
-        cache written by an incompatible version is wiped on open. Within
-        a single schema version, the cache is safe to share across
-        ``cuda.core`` patch releases because every entry's key encodes
-        the relevant backend/compiler/runtime fingerprints for its
-        compilation path (NVRTC entries pin the NVRTC version, NVVM
-        entries pin the libNVVM library and IR versions, PTX/linker
-        entries pin the chosen linker backend and its version -- and,
-        when the cuLink/driver backend is selected, the driver version
-        too; nvJitLink-backed PTX entries are deliberately driver-version
-        independent).
+        ``_FILESTREAM_SCHEMA_VERSION`` encodes both the on-disk storage
+        format and the key-schema version, so a cache written by an
+        incompatible version is wiped on open (bumping either
+        ``_KEY_SCHEMA_VERSION`` or ``_FILESTREAM_BACKEND_SCHEMA`` forces
+        cleanup instead of leaving orphaned entries on disk).
+
+        Within a single schema version the cache is safe to share across
+        ``cuda.core`` patch releases on a best-effort basis: every entry's
+        key encodes the relevant backend/compiler/runtime fingerprints for
+        its compilation path (NVRTC entries pin the NVRTC version, NVVM
+        entries pin the libNVVM library and IR versions, PTX/linker entries
+        pin the chosen linker backend and its version -- and, when the
+        cuLink/driver backend is selected, the driver version too;
+        nvJitLink-backed PTX entries are deliberately driver-version
+        independent). Entries are stored as pickled :class:`ObjectCode`,
+        so the sharing guarantee also assumes ``ObjectCode``'s pickle
+        representation stays compatible across the patch releases in
+        question; a change to its ``__reduce__`` protocol would require
+        bumping the schema version.
 
     Parameters
     ----------

cuda_core/tests/test_program_cache.py

@@ -1481,6 +1481,37 @@ def test_sqlite_cache_wipes_on_schema_mismatch(tmp_path):
         assert b"k" not in cache
 
 
+@needs_sqlite3
+def test_sqlite_cache_schema_version_encodes_key_schema(tmp_path, monkeypatch):
+    """Bumping ``_KEY_SCHEMA_VERSION`` must invalidate existing on-disk
+    entries even when the backend storage layout is unchanged -- otherwise
+    old rows linger unreachable after a key-hash format change."""
+    import sqlite3
+
+    from cuda.core.utils import SQLiteProgramCache, _program_cache
+
+    db = tmp_path / "cache.db"
+    with SQLiteProgramCache(db) as cache:
+        cache[b"k"] = _fake_object_code(b"old-payload")
+
+    # Bump just the key-schema version; the backend version stays the same.
+    monkeypatch.setattr(_program_cache, "_KEY_SCHEMA_VERSION", _program_cache._KEY_SCHEMA_VERSION + 1)
+    monkeypatch.setattr(
+        _program_cache,
+        "_SQLITE_SCHEMA_VERSION",
+        f"{_program_cache._SQLITE_BACKEND_SCHEMA}.{_program_cache._KEY_SCHEMA_VERSION}",
+    )
+
+    with SQLiteProgramCache(db) as cache:
+        assert len(cache) == 0
+        assert b"k" not in cache
+
+    # The old row must be physically gone, not just invisible.
+    with sqlite3.connect(db) as conn:
+        row_count = conn.execute("SELECT COUNT(*) FROM entries").fetchone()[0]
+    assert row_count == 0
+
+
 @needs_sqlite3
 def test_sqlite_cache_drops_tables_on_schema_mismatch(tmp_path):
     """On a schema mismatch the cache must DROP the old tables, not just
@@ -2040,6 +2071,31 @@ def test_filestream_cache_wipes_on_schema_mismatch(tmp_path):
     assert (root / "SCHEMA_VERSION").read_text().strip() != "0"
 
 
+def test_filestream_cache_schema_version_encodes_key_schema(tmp_path, monkeypatch):
+    """As with the SQLite backend, bumping ``_KEY_SCHEMA_VERSION`` alone
+    must invalidate the on-disk cache so orphaned entries from the old
+    key-hash format do not linger after an upgrade."""
+    from cuda.core.utils import FileStreamProgramCache, _program_cache
+
+    root = tmp_path / "fc"
+    with FileStreamProgramCache(root) as cache:
+        cache[b"k"] = _fake_object_code(b"old-payload")
+        path = cache._path_for_key(b"k")
+    assert path.exists()
+
+    monkeypatch.setattr(_program_cache, "_KEY_SCHEMA_VERSION", _program_cache._KEY_SCHEMA_VERSION + 1)
+    monkeypatch.setattr(
+        _program_cache,
+        "_FILESTREAM_SCHEMA_VERSION",
+        f"{_program_cache._FILESTREAM_BACKEND_SCHEMA}.{_program_cache._KEY_SCHEMA_VERSION}",
+    )
+
+    with FileStreamProgramCache(root) as cache:
+        assert len(cache) == 0
+        assert b"k" not in cache
+    assert not path.exists()
+
+
 # ---------------------------------------------------------------------------
 # End-to-end: real NVRTC compilation through persistent cache
 # ---------------------------------------------------------------------------