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

1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5"""Abstract base class and value-coercion helpers shared by every cache backend.""" 

6 

7from __future__ import annotations 

8 

9import abc 

10import collections.abc 

11from pathlib import Path 

12 

13from cuda.core._module import ObjectCode 

14 

15 

16def _extract_bytes(value: object) -> bytes: 

17 """Return the raw binary bytes to store on disk. 

18 

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%;

41 

42 

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

49 

50 

51class ProgramCacheResource(abc.ABC): 

52 """Abstract base class for compiled-program caches. 

53 

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``. 

59 

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``, ...). 

64 

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:: 

70 

71 with FileStreamProgramCache() as cache: 

72 obj = program.compile("cubin", cache=cache) 

73 

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:: 

78 

79 from cuda.core import ObjectCode 

80 

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) 

94 

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. 

102 

103 .. note:: **Concurrent-access idiom.** 

104 

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 """ 

115 

116 @abc.abstractmethod 

117 def __getitem__(self, key: bytes | str) -> bytes: 

118 """Retrieve the cached binary bytes. 

119 

120 Raises 

121 ------ 

122 KeyError 

123 If ``key`` is not in the cache. 

124 """ 

125 

126 @abc.abstractmethod 

127 def __setitem__(self, key: bytes | str, value: bytes | bytearray | memoryview | ObjectCode) -> None: 

128 """Store ``value`` under ``key``. 

129 

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 """ 

133 

134 @abc.abstractmethod 

135 def __delitem__(self, key: bytes | str) -> None: 

136 """Remove the entry associated with ``key``. 

137 

138 Raises 

139 ------ 

140 KeyError 

141 If ``key`` is not in the cache. 

142 """ 

143 

144 @abc.abstractmethod 

145 def __len__(self) -> int: 

146 """Return the number of entries currently in the cache. 

147 

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 """ 

155 

156 @abc.abstractmethod 

157 def clear(self) -> None: 

158 """Remove every entry from the cache.""" 

159 

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

166 

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__``. 

176 

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

187 

188 def close(self) -> None: # noqa: B027 1acVWpdqewxyzfMXYgbAZrBhiCDjEF0%G1okHs=23IJl45K67@v!9#$mn

189 """Release backend resources. 

190 

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. 

194 

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 """ 

199 

200 def __enter__(self) -> ProgramCacheResource: 

201 return self 1cVWpdqewxyzfMXYgbAZrBhiCDjEF0%G1okHs=23IJl45K67?v!9#$mn

202 

203 def __exit__(self, exc_type, exc_value, traceback) -> None: 

204 self.close() 1cVWpdqewxyzfMXYgbAZrBhiCDjEF0%G1okHs=23IJl45K67?v!9#$mn