Coverage for cuda / core / utils / _program_cache / _abc.py: 100.00%
49 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-22 01:37 +0000
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-22 01:37 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
5"""Abstract base class and value-coercion helpers shared by every cache backend."""
7from __future__ import annotations
9import abc
10import collections.abc
11from pathlib import Path
13from cuda.core._module import ObjectCode
16def _extract_bytes(value: object) -> bytes:
17 """Return the raw binary bytes to store on disk.
19 Accepts ``bytes``, ``bytearray``, ``memoryview``, or any
20 :class:`ObjectCode`. Path-backed ``ObjectCode`` (created via
21 ``ObjectCode.from_cubin('/path')`` etc.) is read from the filesystem
22 at write time so the cached entry is the binary content itself, not
23 a path that could later be moved or modified.
24 """
25 if isinstance(value, (bytes, bytearray, memoryview)): 1cVWpdqewxyzfXYgbAZrBhiCDjEF0%G1okHs2*3IJl45K67+,NtOu'8P-Q;RLST.Uv9mn
26 return bytes(value) 1VWqXYbZ01s2*34567+,O'P-QRST.Uv9
27 if isinstance(value, ObjectCode): 1cpdewxyzfgbArBhiCDjEF%GokHIJlKNtu8;Lmn
28 code = value.code 1cpdewxyzfgbArBhiCDjEFGokHIJlKNtu8Lmn
29 if isinstance(code, str): 1cpdewxyzfgbArBhiCDjEFGokHIJlKNtu8Lmn
30 try: 1pbrN
31 return Path(code).read_bytes() 1pbrN
32 except FileNotFoundError: 1r
33 # Avoid ``{code!r}``: on Windows it doubles every backslash in
34 # the path, which breaks naive ``str(path) in str(exc)`` checks
35 # callers (and tests) write against the raw filesystem string.
36 raise FileNotFoundError( 1r
37 f"cannot store path-backed ObjectCode in cache: {code} no longer exists"
38 ) from None
39 return bytes(code) 1cdewxyzfgbABhiCDjEFGokHIJlKtu8Lmn
40 raise TypeError(f"cache values must be bytes-like or ObjectCode, got {type(value).__name__}") 1%;
43def _as_key_bytes(key: object) -> bytes:
44 if isinstance(key, (bytes, bytearray)): 1cVWpdqewxyzfMXYgbAZBhiCDjEF0G1okHs2*3IJl45K67+,NtOu()'8P-QRLST.Uv!9#$mn
45 return bytes(key) 1cVWpdqewxyzfMXYgbAZBhiCDjEF0G1okHs2*3IJl45K67+,NtOu()'8P-QLST.Uv!9#$mn
46 if isinstance(key, str): 1dgRT
47 return key.encode("utf-8") 1dgT
48 raise TypeError(f"cache keys must be bytes or str, got {type(key).__name__}") 1R
51class ProgramCacheResource(abc.ABC):
52 """Abstract base class for compiled-program caches.
54 Concrete implementations store and retrieve **raw binary bytes** keyed
55 by ``bytes`` or ``str``. A ``str`` key is encoded as UTF-8 before
56 being used, so ``"k"`` and ``b"k"`` refer to the same entry. A typical
57 key is produced by :func:`make_program_cache_key`, which returns
58 ``bytes``.
60 The values written are the compiled program bytes themselves --
61 cubin, PTX, LTO-IR, etc. Reads return raw bytes so cache files
62 remain consumable by external NVIDIA tools (``cuobjdump``,
63 ``nvdisasm``, ``cuda-gdb``, ...).
65 Most callers don't interact with this object directly. The
66 recommended usage is :meth:`cuda.core.Program.compile`'s ``cache=``
67 keyword, which derives the key, returns a fresh
68 :class:`~cuda.core.ObjectCode` on hit, and stores the compile
69 result on miss::
71 with FileStreamProgramCache() as cache:
72 obj = program.compile("cubin", cache=cache)
74 The escape hatch -- only needed when the compile inputs require an
75 ``extra_digest`` (header / PCH content fingerprints, NVVM
76 libdevice) -- is to call :func:`make_program_cache_key` yourself
77 and use the cache as a plain ``bytes`` mapping::
79 from cuda.core import ObjectCode
81 key = make_program_cache_key(
82 code=source,
83 code_type="c++",
84 options=options,
85 target_type="cubin",
86 extra_digest=header_fingerprint(),
87 )
88 data = cache.get(key)
89 if data is None:
90 obj = program.compile("cubin")
91 cache[key] = obj # extracts bytes(obj.code)
92 else:
93 obj = ObjectCode.from_cubin(data)
95 The cache layer does no payload validation; bytes go in and come
96 back out unchanged. Symbol-mapping metadata that
97 :class:`~cuda.core.ObjectCode` carries when produced with NVRTC
98 name expressions is **not** preserved across a cache round-trip --
99 the binary alone is stored. Callers that need ``symbol_mapping``
100 for ``get_kernel(name_expression)`` should compile fresh, or look
101 the mangled symbol up by hand.
103 .. note:: **Concurrent-access idiom.**
105 Use :meth:`get` (or ``data = cache[key]`` inside a ``try /
106 except KeyError``) for lookups. There is intentionally no
107 ``__contains__``: the obvious ``if key in cache: data =
108 cache[key]`` idiom is racy across processes (another writer
109 can ``os.replace`` over the entry, or eviction can unlink
110 it, between the check and the read), and exposing
111 ``__contains__`` invites that pattern. ``get`` answers
112 both questions in one filesystem-level operation, so a
113 successful return always carries the bytes.
114 """
116 @abc.abstractmethod
117 def __getitem__(self, key: bytes | str) -> bytes:
118 """Retrieve the cached binary bytes.
120 Raises
121 ------
122 KeyError
123 If ``key`` is not in the cache.
124 """
126 @abc.abstractmethod
127 def __setitem__(self, key: bytes | str, value: bytes | bytearray | memoryview | ObjectCode) -> None:
128 """Store ``value`` under ``key``.
130 Path-backed :class:`~cuda.core.ObjectCode` is read from disk at
131 write time so the cached entry holds the bytes, not a path.
132 """
134 @abc.abstractmethod
135 def __delitem__(self, key: bytes | str) -> None:
136 """Remove the entry associated with ``key``.
138 Raises
139 ------
140 KeyError
141 If ``key`` is not in the cache.
142 """
144 @abc.abstractmethod
145 def __len__(self) -> int:
146 """Return the number of entries currently in the cache.
148 Implementations that store entries on disk by hashed key may
149 count orphaned files (entries from a previous
150 ``_KEY_SCHEMA_VERSION`` that are still on disk but no longer
151 reachable by post-bump lookups) until eviction reaps them.
152 Callers that need an exact count of live entries should not
153 rely on ``__len__`` across schema bumps.
154 """
156 @abc.abstractmethod
157 def clear(self) -> None:
158 """Remove every entry from the cache."""
160 def get(self, key: bytes | str, default: bytes | None = None) -> bytes | None:
161 """Return ``self[key]`` or ``default`` if absent."""
162 try: 1cdqefMghijoksltOu()'PQRLST/:!#$mn
163 return self[key] 1cdqefMghijoksltOu()'PQRLST/:!#$mn
164 except KeyError: 1cqefMhijksltOu()PQRS/:mn
165 return default 1cqefMhijksltOu()PQS/:mn
167 def update(
168 self,
169 items: (
170 collections.abc.Mapping[bytes | str, bytes | bytearray | memoryview | ObjectCode]
171 | collections.abc.Iterable[tuple[bytes | str, bytes | bytearray | memoryview | ObjectCode]]
172 ),
173 /,
174 ) -> None:
175 """Bulk ``__setitem__``.
177 Accepts a mapping or an iterable of ``(key, value)`` pairs. Each
178 write goes through ``__setitem__`` so backend-specific value
179 coercion (e.g. extracting bytes from an :class:`~cuda.core.ObjectCode`)
180 and size-cap enforcement run on every entry. Not transactional --
181 a failure mid-iteration leaves earlier writes committed.
182 """
183 if isinstance(items, collections.abc.Mapping): 1bUv
184 items = items.items() 1bUv
185 for key, value in items: 1bUv
186 self[key] = value 1bUv
188 def close(self) -> None: # noqa: B027 1acVWpdqewxyzfMXYgbAZrBhiCDjEF0%G1okHs=23IJl45K67@v!9#$mn
189 """Release backend resources.
191 The default implementation does nothing. Subclasses that hold
192 long-lived state (open file handles, database connections,
193 network sockets, ...) should override this to release them.
195 Callers should use the context-manager form (``with cache:``)
196 or call :meth:`close` explicitly when finished, so code stays
197 portable across backends that do hold resources.
198 """
200 def __enter__(self) -> ProgramCacheResource:
201 return self 1cVWpdqewxyzfMXYgbAZrBhiCDjEF0%G1okHs=23IJl45K67?v!9#$mn
203 def __exit__(self, exc_type, exc_value, traceback) -> None:
204 self.close() 1cVWpdqewxyzfMXYgbAZrBhiCDjEF0%G1okHs=23IJl45K67?v!9#$mn