fixup! adapt Program.compile(cache=...) wrapper + tests for bytes API

Salvage the high-level wrapper from the diverged origin branch and adapt
to the raw-bytes cache API:

* Wrapper rebuilds an ObjectCode on hit via
  ObjectCode._init(hit_bytes, target_type, name=self._options.name) since
  the cache returns bytes, not an ObjectCode.
* Documented in the Program.compile docstring that symbol_mapping is
  not preserved across a cache round-trip.
* Test stub _RecordingCache mirrors the FileStream contract: __setitem__
  accepts ObjectCode and stores its bytes, get returns bytes (or None).
* Drop SQLite/InMemory imports and tests left over from the cherry-pick
  (those backends are gone in this branch).
* Fix import order in _program_cache.py (cherry-pick interleaved
  platformdirs between pathlib and typing imports).

Author: cpcloud

cuda_core/cuda/core/_program.pyx

@@ -104,11 +104,17 @@ cdef class Program:
         logs : object, optional
             Object with a ``write`` method to receive compilation logs.
         cache : :class:`~cuda.core.utils.ProgramCacheResource`, optional
-            If provided, the compiled :class:`~cuda.core.ObjectCode` is looked
-            up in ``cache`` via a key derived from the program's code, options,
-            ``target_type`` and ``name_expressions``. On a hit the cached
-            ``ObjectCode`` is returned without re-compiling; on a miss the
-            fresh compile result is stored. Options that require an
+            If provided, the compiled binary is looked up in ``cache`` via a
+            key derived from the program's code, options, ``target_type`` and
+            ``name_expressions``. On a hit the cached bytes are wrapped in a
+            fresh :class:`~cuda.core.ObjectCode` (with the same ``target_type``
+            and ``ProgramOptions.name``) and returned without re-compiling;
+            on a miss the compile output is stored as raw bytes (the cache
+            extracts ``bytes(object_code.code)``). Note that
+            ``ObjectCode.symbol_mapping`` is not preserved across a cache
+            round-trip -- callers using ``name_expressions`` who need
+            ``get_kernel(name_expression)`` after a hit must compile fresh
+            or look the mangled symbol up by hand. Options that require an
             ``extra_digest`` (``include_path``, ``pre_include``, ``pch``,
             ``use_pch``, ``pch_dir``, NVVM ``use_libdevice=True``, or NVRTC
             ``options.name`` with a directory component) raise ``ValueError``
@@ -146,9 +152,9 @@ cdef class Program:
             target_type=target_type,
             name_expressions=name_expressions,
         )
-        hit = cache.get(key)
-        if hit is not None:
-            return hit
+        hit_bytes = cache.get(key)
+        if hit_bytes is not None:
+            return ObjectCode._init(hit_bytes, target_type, name=self._options.name)
         compiled = _program_compile_uncached(self, target_type, name_expressions, logs)
         cache[key] = compiled
         return compiled

cuda_core/cuda/core/utils/_program_cache.py

@@ -23,9 +23,9 @@
 import threading
 import time
 from pathlib import Path
+from typing import Iterable, Sequence
 
 import platformdirs
-from typing import Iterable, Sequence
 
 from cuda.core._module import ObjectCode
 from cuda.core._program import ProgramOptions

cuda_core/tests/test_program_compile_cache.py

@@ -1,5 +1,4 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-#
 # SPDX-License-Identifier: Apache-2.0
 """Tests for the ``Program.compile(cache=...)`` convenience integration."""
 
@@ -12,34 +11,37 @@
 from cuda.core._module import ObjectCode
 from cuda.core.utils import (
     FileStreamProgramCache,
-    InMemoryProgramCache,
-    SQLiteProgramCache,
     make_program_cache_key,
 )
 
-try:
-    import sqlite3  # noqa: F401
-
-    _has_sqlite3 = True
-except ImportError:
-    _has_sqlite3 = False
-
-needs_sqlite3 = pytest.mark.skipif(not _has_sqlite3, reason="libsqlite3 not available")
-
 
 class _RecordingCache:
-    """Minimal recording stub for the two-method cache protocol.
+    """Minimal recording stub for the bytes-in / bytes-out cache protocol.
+
+    Mirrors :class:`FileStreamProgramCache`'s contract: ``__setitem__``
+    accepts bytes-like or :class:`ObjectCode` (extracts bytes), and
+    ``get`` returns the stored bytes (or ``None``).
 
     Intentionally does NOT subclass ``ProgramCacheResource`` -- the wrapper
     should be duck-typed, so we test the duck-typed surface directly.
     """
 
     def __init__(self, preseed=None):
-        self._store = dict(preseed or {})
-        self.get_calls = []
-        self.set_calls = []
-        self.get_side_effect = None
-        self.set_side_effect = None
+        self._store: dict[bytes, bytes] = {}
+        for k, v in (preseed or {}).items():
+            self._store[k] = self._extract(v)
+        self.get_calls: list[bytes] = []
+        self.set_calls: list[tuple[bytes, bytes]] = []
+        self.get_side_effect: BaseException | None = None
+        self.set_side_effect: BaseException | None = None
+
+    @staticmethod
+    def _extract(value) -> bytes:
+        if isinstance(value, ObjectCode):
+            return bytes(value.code)
+        if isinstance(value, (bytes, bytearray, memoryview)):
+            return bytes(value)
+        raise TypeError(f"unexpected value type: {type(value).__name__}")
 
     def get(self, key, default=None):
         self.get_calls.append(key)
@@ -48,22 +50,24 @@ def get(self, key, default=None):
         return self._store.get(key, default)
 
     def __setitem__(self, key, value):
-        self.set_calls.append((key, value))
+        data = self._extract(value)
+        self.set_calls.append((key, data))
         if self.set_side_effect is not None:
             raise self.set_side_effect
-        self._store[key] = value
+        self._store[key] = data
 
 
 _KERNEL = 'extern "C" __global__ void k() {}'
+_SENTINEL_BYTES = b"sentinel-cubin-bytes"
 
 
 def _make_sentinel_object_code():
     """Construct a cache-safe ``ObjectCode`` that doesn't require compilation."""
-    return ObjectCode._init(b"sentinel-cubin-bytes", "cubin", name="sentinel")
+    return ObjectCode._init(_SENTINEL_BYTES, "cubin", name="sentinel")
 
 
 def test_cache_miss_runs_compile_then_stores(monkeypatch):
-    """On cache miss: get(key) once, _program_compile_uncached once, __setitem__(key, obj) once."""
+    """On cache miss: get(key) once, _program_compile_uncached once, __setitem__ once."""
     sentinel = _make_sentinel_object_code()
     program = Program(_KERNEL, "c++", ProgramOptions(arch="sm_80"))
 
@@ -75,16 +79,19 @@ def _return_sentinel(_program, *_args, **_kwargs):
 
     result = program.compile("cubin", cache=cache)
 
+    # On miss the wrapper returns the freshly-compiled ObjectCode unchanged.
     assert result is sentinel
     assert len(cache.get_calls) == 1
     assert len(cache.set_calls) == 1
-    assert cache.set_calls[0][1] is sentinel
+    # The cache stored the binary bytes extracted from the ObjectCode.
+    assert cache.set_calls[0][1] == _SENTINEL_BYTES
 
 
-def test_cache_hit_short_circuits_compile(monkeypatch):
-    """On cache hit: get(key) returns, _program_compile_uncached is not called, no __setitem__."""
-    sentinel = _make_sentinel_object_code()
-    options = ProgramOptions(arch="sm_80")
+def test_cache_hit_returns_object_code_reconstructed_from_bytes(monkeypatch):
+    """On hit: get(key) returns bytes, the wrapper rebuilds an ObjectCode with
+    the same code_type and ProgramOptions.name. _program_compile_uncached is
+    NOT called and there is no __setitem__."""
+    options = ProgramOptions(arch="sm_80", name="my_program")
     program = Program(_KERNEL, "c++", options)
     key = make_program_cache_key(
         code=_KERNEL,
@@ -97,11 +104,14 @@ def _explode(_program, *_args, **_kwargs):
         raise AssertionError("_program_compile_uncached must not be called on cache hit")
 
     monkeypatch.setattr(_program_module, "_program_compile_uncached", _explode)
-    cache = _RecordingCache(preseed={key: sentinel})
+    cache = _RecordingCache(preseed={key: _SENTINEL_BYTES})
 
     result = program.compile("cubin", cache=cache)
 
-    assert result is sentinel
+    assert isinstance(result, ObjectCode)
+    assert bytes(result.code) == _SENTINEL_BYTES
+    assert result.code_type == "cubin"
+    assert result.name == "my_program"
     assert cache.get_calls == [key]
     assert cache.set_calls == []
 
@@ -110,7 +120,9 @@ def test_name_expressions_affects_cache_key(monkeypatch):
     """Different ``name_expressions`` must produce different cache keys."""
     sentinel = _make_sentinel_object_code()
     program = Program(_KERNEL, "c++", ProgramOptions(arch="sm_80"))
-    monkeypatch.setattr(_program_module, "_program_compile_uncached", lambda _self, *_args, **_kwargs: sentinel)
+    monkeypatch.setattr(
+        _program_module, "_program_compile_uncached", lambda _self, *_args, **_kwargs: sentinel
+    )
     cache = _RecordingCache()
 
     program.compile("cubin", name_expressions=("foo",), cache=cache)
@@ -188,7 +200,9 @@ def test_cache_write_exception_propagates(monkeypatch):
     """Exceptions from cache.__setitem__ propagate after compile runs."""
     sentinel = _make_sentinel_object_code()
     program = Program(_KERNEL, "c++", ProgramOptions(arch="sm_80"))
-    monkeypatch.setattr(_program_module, "_program_compile_uncached", lambda _self, *_args, **_kwargs: sentinel)
+    monkeypatch.setattr(
+        _program_module, "_program_compile_uncached", lambda _self, *_args, **_kwargs: sentinel
+    )
     cache = _RecordingCache()
     cache.set_side_effect = RuntimeError("disk full")
 
@@ -203,7 +217,9 @@ def test_no_cache_kwarg_does_not_derive_key(monkeypatch):
     """Without cache=, no cache-module functions run; compile goes straight through."""
     sentinel = _make_sentinel_object_code()
     program = Program(_KERNEL, "c++", ProgramOptions(arch="sm_80"))
-    monkeypatch.setattr(_program_module, "_program_compile_uncached", lambda _self, *_args, **_kwargs: sentinel)
+    monkeypatch.setattr(
+        _program_module, "_program_compile_uncached", lambda _self, *_args, **_kwargs: sentinel
+    )
 
     # If the implementation accidentally derived a key, it would call
     # make_program_cache_key. Replace it with a raising stub to catch that.
@@ -219,35 +235,9 @@ def _cache_path_must_not_run(*_args, **_kwargs):
     assert result is sentinel
 
 
-def test_inmemory_hit_returns_same_instance(init_cuda):
-    """InMemoryProgramCache stores by reference; a hit returns the same object."""
-    cache = InMemoryProgramCache()
-    program = Program(_KERNEL, "c++", ProgramOptions(arch="sm_80"))
-
-    first = program.compile("cubin", cache=cache)
-    second = program.compile("cubin", cache=cache)
-
-    assert second is first
-
-
-@needs_sqlite3
-def test_sqlite_hit_returns_byte_equal_object_code(init_cuda, tmp_path):
-    """SQLiteProgramCache pickle-roundtrips; a hit returns byte-identical code."""
-    program = Program(_KERNEL, "c++", ProgramOptions(arch="sm_80"))
-    db_path = tmp_path / "cache.sqlite"
-
-    with SQLiteProgramCache(db_path) as cache:
-        first = program.compile("cubin", cache=cache)
-
-    with SQLiteProgramCache(db_path) as cache:
-        second = program.compile("cubin", cache=cache)
-
-    assert second.code == first.code
-    assert second.code_type == first.code_type
-
-
 def test_filestream_hit_returns_byte_equal_object_code(init_cuda, tmp_path):
-    """FileStreamProgramCache pickle-roundtrips; a hit returns byte-identical code."""
+    """End-to-end: real compile, FileStreamProgramCache roundtrip, second
+    compile returns an ObjectCode whose bytes match the first compile."""
     program = Program(_KERNEL, "c++", ProgramOptions(arch="sm_80"))
     cache_dir = tmp_path / "fc"
 
@@ -257,42 +247,31 @@ def test_filestream_hit_returns_byte_equal_object_code(init_cuda, tmp_path):
     with FileStreamProgramCache(cache_dir) as cache:
         second = program.compile("cubin", cache=cache)
 
-    assert second.code == first.code
+    assert bytes(second.code) == bytes(first.code)
     assert second.code_type == first.code_type
 
 
-@pytest.mark.parametrize(
-    "backend",
-    [
-        pytest.param("sqlite", marks=needs_sqlite3),
-        "filestream",
-    ],
-)
-def test_cache_kwarg_roundtrip_across_reopen(init_cuda, tmp_path, backend, monkeypatch):
+def test_cache_kwarg_roundtrip_across_reopen(init_cuda, tmp_path, monkeypatch):
     """Compile with cache= in one 'session', reopen and fetch via cache= again."""
-    program_init_args = (_KERNEL, "c++", ProgramOptions(arch="sm_80", name="cached_kernel"))
-
-    if backend == "sqlite":
-        cache_path = tmp_path / "cache.sqlite"
-        cache_cls = SQLiteProgramCache
-    else:
-        cache_path = tmp_path / "fc"
-        cache_cls = FileStreamProgramCache
+    init_args = (_KERNEL, "c++", ProgramOptions(arch="sm_80", name="cached_kernel"))
+    cache_path = tmp_path / "fc"
 
-    with cache_cls(cache_path) as cache:
-        program = Program(*program_init_args)
+    with FileStreamProgramCache(cache_path) as cache:
+        program = Program(*init_args)
         first = program.compile("cubin", cache=cache)
 
     # Fresh process / fresh Program and cache-handle -- same cache path.
-    with cache_cls(cache_path) as cache:
-        program = Program(*program_init_args)
+    with FileStreamProgramCache(cache_path) as cache:
+        program = Program(*init_args)
 
-        # Monkeypatch the module-level compile seam so a cache miss would raise --
-        # the only way this succeeds is if the reopen finds the prior entry.
+        # If the reopened cache misses, the wrapper would fall back to
+        # _program_compile_uncached -- replace it with a raising stub so
+        # the test can only succeed via a hit.
         def _must_not_recompile(_program, *_args, **_kwargs):
             raise AssertionError("cache miss: reopened cache didn't serve entry")
 
         monkeypatch.setattr(_program_module, "_program_compile_uncached", _must_not_recompile)
         second = program.compile("cubin", cache=cache)
 
-    assert second.code == first.code
+    assert bytes(second.code) == bytes(first.code)
+    assert second.name == "cached_kernel"