fix(core.utils): reject bytearray in NVRTC name_expressions cache keys

Author: cpcloud

cuda_core/cuda/core/utils/_program_cache.py

@@ -483,11 +483,12 @@ def make_program_cache_key(
         One of ``"ptx"``, ``"cubin"``, ``"ltoir"``.
     name_expressions:
         Optional iterable of mangled-name lookups. Order is not significant.
-        Elements may be ``str``, ``bytes``, or ``bytearray``; ``"foo"`` and
-        ``b"foo"`` produce distinct keys because ``Program.compile`` records
-        the original Python object as the ``ObjectCode.symbol_mapping`` key,
-        and ``get_kernel`` lookups must use the same type the cache key
-        recorded.
+        Elements may be ``str`` or ``bytes``; ``"foo"`` and ``b"foo"`` produce
+        distinct keys because ``Program.compile`` records the original Python
+        object as the ``ObjectCode.symbol_mapping`` key, and ``get_kernel``
+        lookups must use the same type the cache key recorded. ``bytearray``
+        is rejected because ``Program.compile`` stores each element as a
+        dict key and ``bytearray`` is unhashable.
     extra_digest:
         Caller-supplied bytes mixed into the key. Required whenever
         :class:`cuda.core.ProgramOptions` sets any option that pulls in
@@ -681,15 +682,25 @@ def make_program_cache_key(
     # because Program.compile records the original Python object as the
     # ObjectCode.symbol_mapping key (_program.pyx:759), so a cached
     # ObjectCode whose mapping-key type differs from what the caller's
-    # later ``get_kernel`` passes would silently miss.
+    # later ``get_kernel`` passes would silently miss. Reject
+    # ``bytearray`` because Program.compile also uses the raw element as a
+    # dict key -- bytearray is unhashable, so a cache miss would compile
+    # then crash in ``symbol_mapping[n] = ...``. Accepting it here would
+    # let the cache serve hits for inputs the uncached path can't handle.
     if backend == "nvrtc":
 
         def _tag_name(n):
-            if isinstance(n, (bytes, bytearray)):
-                return b"b:" + bytes(n)
+            if isinstance(n, bytes):
+                return b"b:" + n
             if isinstance(n, str):
                 return b"s:" + n.encode("utf-8")
-            raise TypeError(f"name_expressions elements must be str, bytes, or bytearray; got {type(n).__name__}")
+            if isinstance(n, bytearray):
+                raise TypeError(
+                    "name_expressions elements must be str or bytes; "
+                    "bytearray is not accepted because Program.compile uses "
+                    "each element as a dict key and bytearray is unhashable."
+                )
+            raise TypeError(f"name_expressions elements must be str or bytes; got {type(n).__name__}")
 
         names = tuple(sorted(_tag_name(n) for n in name_expressions))
     else:

cuda_core/tests/test_program_cache.py

@@ -306,6 +306,16 @@ def test_make_program_cache_key_name_expressions_str_bytes_distinct():
     assert _make_key(name_expressions=("foo",)) != _make_key(name_expressions=(b"foo",))
 
 
+def test_make_program_cache_key_rejects_bytearray_in_name_expressions():
+    """``bytearray`` is unhashable, and ``Program.compile`` stores each
+    element of ``name_expressions`` as a dict key
+    (``symbol_mapping[n] = ...`` in ``_program.pyx``). Accepting it in the
+    cache helper would mean hits served for inputs the uncached compile
+    path crashes on -- so reject up front."""
+    with pytest.raises(TypeError, match="bytearray"):
+        _make_key(name_expressions=("ok", bytearray(b"bad")))
+
+
 @pytest.mark.parametrize(
     "code_type, target_type",
     [