Repository navigation
Expand file tree
/
Copy path_ctypes_extensions.py
More file actions
320 lines (258 loc) · 10.5 KB
/
Copy path_ctypes_extensions.py
File metadata and controls
320 lines (258 loc) · 10.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
from __future__ import annotations
import sys
import os
import ctypes
import functools
import pathlib
import importlib.metadata
from ctypes.util import find_library
from typing import (
Any,
Callable,
Iterable,
List,
Union,
Optional,
TYPE_CHECKING,
TypeVar,
Generic,
)
from typing_extensions import TypeAlias
def _version_at_least(version: str) -> bool:
"""Check whether installed llama-cpp-python version meets requirement."""
try:
current = importlib.metadata.version("llama-cpp-python")
from packaging.version import Version
return Version(current) >= Version(version)
except Exception:
return False
def _format_library_dir_contents(base_paths: list[pathlib.Path]) -> str:
"""Format directory contents for diagnostics after library loading fails."""
sections = []
for base_path in base_paths:
p = pathlib.Path(base_path)
if not p.exists():
sections.append(f"{p}: <not found>")
continue
if not p.is_dir():
sections.append(f"{p}: <not a directory>")
continue
try:
# Only list files when reporting a final loading failure.
files = sorted(x.name for x in p.iterdir())
except Exception as e:
sections.append(f"{p}: <failed to list: {e}>")
continue
if files:
sections.append(
f"{p}:\n"
+ "\n".join(f" - {name}" for name in files)
)
else:
sections.append(f"{p}: <empty>")
return "\n".join(sections)
# Load the library
def load_shared_library(lib_base_name: str, base_paths: Union[pathlib.Path, list[pathlib.Path]]):
if isinstance(base_paths, pathlib.Path):
base_paths = [base_paths]
lib_names = []
if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"):
lib_names = [f"lib{lib_base_name}.so"]
base_paths.extend([
"/usr/local/lib",
"/usr/lib",
"/usr/lib64",
])
elif sys.platform == "darwin":
lib_names = [
f"lib{lib_base_name}.dylib",
f"lib{lib_base_name}.so",
]
base_paths.extend([
"/usr/local/lib",
"/opt/homebrew/lib",
"/usr/lib",
])
elif sys.platform == "win32":
lib_names = [
f"{lib_base_name}.dll",
f"lib{lib_base_name}.dll",
]
else:
raise RuntimeError("Unsupported platform")
cdll_args = dict() # type: ignore
# Add the library directory to the DLL search path on Windows (if needed)
if sys.platform == "win32":
# Add CUDA runtime DLL directories if CUDA is available.
if "CUDA_PATH" in os.environ:
cuda_path = os.environ["CUDA_PATH"]
sub_dirs_to_add = [
"bin",
os.path.join("bin", "x64"), # CUDA 13.0+
"lib",
os.path.join("lib", "x64")
]
for sub_dir in sub_dirs_to_add:
full_path = os.path.join(cuda_path, sub_dir)
if os.path.exists(full_path):
os.add_dll_directory(full_path)
# Add HIP runtime DLL directories when HIP backend is available.
if "HIP_PATH" in os.environ:
os.add_dll_directory(os.path.join(os.environ["HIP_PATH"], "bin"))
os.add_dll_directory(os.path.join(os.environ["HIP_PATH"], "lib"))
# Add Vulkan SDK DLL directories when Vulkan backend is enabled.
if "VULKAN_SDK" in os.environ:
os.add_dll_directory(os.path.join(os.environ["VULKAN_SDK"], "Bin"))
os.add_dll_directory(os.path.join(os.environ["VULKAN_SDK"], "Lib"))
# Add package-provided library directories.
#
# The paths are added in reverse order intentionally.
# This ensures that the first entry in base_paths gets prepended
# to PATH last, making it the highest priority search location.
#
# Example:
# base_paths = [
# package/lib,
# package/bin,
# ]
#
# After reversed iteration:
# PATH = package/lib;package/bin;...
for base_path in reversed(base_paths):
p = pathlib.Path(base_path)
if p.exists() and p.is_dir():
os.add_dll_directory(str(p))
os.environ["PATH"] = str(p) + os.pathsep + os.environ["PATH"]
cdll_args["winmode"] = ctypes.RTLD_GLOBAL
errors = []
# First, try to find an available library through the system
lib_path = find_library(lib_base_name)
if lib_path:
try:
lib = ctypes.CDLL(lib_path, **cdll_args)
print(f"[llama-cpp-python].find_library: loaded library from {lib_path}")
return lib
except Exception as e:
errors.append(f"{lib_path}: {e}")
# Then fallback to manually checking the list of paths.
for base_path in base_paths:
for lib_name in lib_names:
lib_path = pathlib.Path(base_path) / lib_name
if lib_path.exists():
try:
lib = ctypes.CDLL(str(lib_path), **cdll_args)
print(f"[llama-cpp-python].provided_path: loaded library from {lib_path}")
return lib
except Exception as e:
errors.append(f"{lib_path}: {e}")
# Include directory contents only in the failure path to avoid extra work during successful imports.
raise RuntimeError(
f"Failed to load '{lib_base_name}' from {base_paths}\n"
+ "\n".join(errors)
+ "\nLibrary search path contents:\n"
+ _format_library_dir_contents(base_paths)
)
# ctypes sane type hint helpers
#
# - Generic Pointer and Array types
# - PointerOrRef type with a type hinted byref function
#
# NOTE: Only use these for static type checking not for runtime checks
# no good will come of that
if TYPE_CHECKING:
CtypesCData = TypeVar("CtypesCData", bound=ctypes._CData) # type: ignore
CtypesArray: TypeAlias = ctypes.Array[CtypesCData] # type: ignore
CtypesPointer: TypeAlias = ctypes._Pointer[CtypesCData] # type: ignore
CtypesVoidPointer: TypeAlias = ctypes.c_void_p
class CtypesRef(Generic[CtypesCData]):
pass
CtypesPointerOrRef: TypeAlias = Union[
CtypesPointer[CtypesCData], CtypesRef[CtypesCData]
]
CtypesFuncPointer: TypeAlias = ctypes._FuncPointer # type: ignore
F = TypeVar("F", bound=Callable[..., Any])
def ctypes_function_for_shared_library(lib: ctypes.CDLL):
"""Create a decorator used to bind typed Python declarations to C symbols.
The returned decorator accepts either a single exported symbol name or an
iterable of ABI-compatible aliases. When aliases are provided, they are
checked in order and the first available symbol is selected.
"""
def ctypes_function(
name: Union[str, Iterable[str]],
argtypes: List[Any],
restype: Any,
enabled: bool = True,
required: bool = True,
):
"""Bind a Python declaration to one of the requested C symbols.
Args:
name: A symbol name or an ordered iterable of compatible aliases.
argtypes: The ctypes argument types assigned to the C function.
restype: The ctypes return type assigned to the C function.
enabled: Return the original Python declaration when disabled.
required: Raise if symbol is missing. If False, create a runtime unavailable stub.
Raises:
ValueError: If no symbol names are provided.
AttributeError: If none of the requested symbols exist in the
shared library.
"""
symbol_names = (name,) if isinstance(name, str) else tuple(name)
if not symbol_names:
raise ValueError("At least one shared library symbol name is required")
def decorator(f: F) -> F:
if not enabled:
return f
for symbol_name in symbol_names:
try:
func = getattr(lib, symbol_name)
except AttributeError:
continue
# Validate ctypes argument declarations before assigning them.
# ctypes requires every argtype to provide from_param().
for index, argtype in enumerate(argtypes):
if not hasattr(argtype, "from_param"):
raise TypeError(
"Invalid ctypes argument type:\n"
f" function: {f.__name__}\n"
f" symbol: {symbol_name}\n"
f" arg index: {index}\n"
f" arg type: {argtype!r}\n"
f" expected: a ctypes type with from_param()"
)
func.argtypes = argtypes
func.restype = restype
functools.update_wrapper(func, f)
# Preserve the actual exported symbol selected at runtime for
# diagnostics, especially when ABI aliases are being used.
func.__ctypes_symbol_name__ = symbol_name
return func
message = (
"None of the shared library symbols were found: "
+ ", ".join(symbol_names)
)
if required:
raise AttributeError(message)
# Optional extension API.
# Keep import working when the symbol is unavailable.
print(
"[llama-cpp-python].ctypes_function: WARNING! optional API unavailable\n"
f" symbols: {', '.join(symbol_names)}\n"
f" library: {getattr(lib, '_name', '<unknown>')}"
)
def unavailable(*args, **kwargs):
raise RuntimeError(
"This llama.cpp extension API is unavailable.\n"
f"Required symbol(s): {', '.join(symbol_names)}\n"
f"Library: {getattr(lib, '_name', '<unknown>')}"
)
functools.update_wrapper(unavailable, f)
# Mark unavailable extension API.
unavailable.__ctypes_symbol_name__ = None
unavailable.__ctypes_optional__ = True
return unavailable
return decorator
return ctypes_function
def _byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCData]:
"""Type-annotated version of ctypes.byref"""
...
byref = _byref if TYPE_CHECKING else ctypes.byref