Skip to content

Commit 01c7607

Browse files
abetlenabetlen
authored
feat: Expose libggml in internal APIs (#1761)
* Expose libggml and refactor ctypes extension * Only expose libggml * Use ctypes_extensions module for libllama and libllava
1 parent dca0c9a commit 01c7607

File tree

4 files changed

+171
-225
lines changed

4 files changed

+171
-225
lines changed

Diff for: ‎llama_cpp/_ctypes_extensions.py

+131
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
from __future__ import annotations
2+
3+
import sys
4+
import os
5+
import ctypes
6+
import functools
7+
import pathlib
8+
9+
from typing import (
10+
Any,
11+
Callable,
12+
List,
13+
Union,
14+
Optional,
15+
TYPE_CHECKING,
16+
TypeVar,
17+
Generic,
18+
)
19+
from typing_extensions import TypeAlias
20+
21+
22+
# Load the library
23+
def load_shared_library(lib_base_name: str, base_path: pathlib.Path):
24+
"""Platform independent shared library loader"""
25+
# Searching for the library in the current directory under the name "libllama" (default name
26+
# for llamacpp) and "llama" (default name for this repo)
27+
lib_paths: List[pathlib.Path] = []
28+
# Determine the file extension based on the platform
29+
if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"):
30+
lib_paths += [
31+
base_path / f"lib{lib_base_name}.so",
32+
]
33+
elif sys.platform == "darwin":
34+
lib_paths += [
35+
base_path / f"lib{lib_base_name}.so",
36+
base_path / f"lib{lib_base_name}.dylib",
37+
]
38+
elif sys.platform == "win32":
39+
lib_paths += [
40+
base_path / f"{lib_base_name}.dll",
41+
base_path / f"lib{lib_base_name}.dll",
42+
]
43+
else:
44+
raise RuntimeError("Unsupported platform")
45+
46+
cdll_args = dict() # type: ignore
47+
48+
# Add the library directory to the DLL search path on Windows (if needed)
49+
if sys.platform == "win32":
50+
os.add_dll_directory(str(base_path))
51+
os.environ["PATH"] = str(base_path) + os.pathsep + os.environ["PATH"]
52+
53+
if sys.platform == "win32" and sys.version_info >= (3, 8):
54+
os.add_dll_directory(str(base_path))
55+
if "CUDA_PATH" in os.environ:
56+
os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "bin"))
57+
os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "lib"))
58+
if "HIP_PATH" in os.environ:
59+
os.add_dll_directory(os.path.join(os.environ["HIP_PATH"], "bin"))
60+
os.add_dll_directory(os.path.join(os.environ["HIP_PATH"], "lib"))
61+
cdll_args["winmode"] = ctypes.RTLD_GLOBAL
62+
63+
# Try to load the shared library, handling potential errors
64+
for lib_path in lib_paths:
65+
if lib_path.exists():
66+
try:
67+
return ctypes.CDLL(str(lib_path), **cdll_args) # type: ignore
68+
except Exception as e:
69+
raise RuntimeError(f"Failed to load shared library '{lib_path}': {e}")
70+
71+
raise FileNotFoundError(
72+
f"Shared library with base name '{lib_base_name}' not found"
73+
)
74+
75+
76+
# ctypes sane type hint helpers
77+
#
78+
# - Generic Pointer and Array types
79+
# - PointerOrRef type with a type hinted byref function
80+
#
81+
# NOTE: Only use these for static type checking not for runtime checks
82+
# no good will come of that
83+
84+
if TYPE_CHECKING:
85+
CtypesCData = TypeVar("CtypesCData", bound=ctypes._CData) # type: ignore
86+
87+
CtypesArray: TypeAlias = ctypes.Array[CtypesCData] # type: ignore
88+
89+
CtypesPointer: TypeAlias = ctypes._Pointer[CtypesCData] # type: ignore
90+
91+
CtypesVoidPointer: TypeAlias = ctypes.c_void_p
92+
93+
class CtypesRef(Generic[CtypesCData]):
94+
pass
95+
96+
CtypesPointerOrRef: TypeAlias = Union[
97+
CtypesPointer[CtypesCData], CtypesRef[CtypesCData]
98+
]
99+
100+
CtypesFuncPointer: TypeAlias = ctypes._FuncPointer # type: ignore
101+
102+
F = TypeVar("F", bound=Callable[..., Any])
103+
104+
105+
def ctypes_function_for_shared_library(lib: ctypes.CDLL):
106+
"""Decorator for defining ctypes functions with type hints"""
107+
108+
def ctypes_function(
109+
name: str, argtypes: List[Any], restype: Any, enabled: bool = True
110+
):
111+
def decorator(f: F) -> F:
112+
if enabled:
113+
func = getattr(lib, name)
114+
func.argtypes = argtypes
115+
func.restype = restype
116+
functools.wraps(f)(func)
117+
return func
118+
else:
119+
return f
120+
121+
return decorator
122+
123+
return ctypes_function
124+
125+
126+
def _byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCData]:
127+
"""Type-annotated version of ctypes.byref"""
128+
...
129+
130+
131+
byref = _byref if TYPE_CHECKING else ctypes.byref

Diff for: ‎llama_cpp/_ggml.py

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
"""Internal module use at your own risk
2+
3+
This module provides a minimal interface for working with ggml tensors from llama-cpp-python
4+
"""
5+
import os
6+
import pathlib
7+
8+
import llama_cpp._ctypes_extensions as ctypes_ext
9+
10+
libggml_base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib"
11+
libggml = ctypes_ext.load_shared_library("ggml", libggml_base_path)
12+

Diff for: ‎llama_cpp/llama_cpp.py

+17-123
Original file line numberDiff line numberDiff line change
@@ -1,151 +1,45 @@
11
from __future__ import annotations
22

3-
import sys
43
import os
54
import ctypes
6-
import functools
75
import pathlib
86

97
from typing import (
10-
Any,
118
Callable,
12-
List,
139
Union,
1410
NewType,
1511
Optional,
1612
TYPE_CHECKING,
17-
TypeVar,
18-
Generic,
1913
)
20-
from typing_extensions import TypeAlias
2114

15+
from llama_cpp._ctypes_extensions import (
16+
load_shared_library,
17+
byref,
18+
ctypes_function_for_shared_library,
19+
)
2220

23-
# Load the library
24-
def _load_shared_library(lib_base_name: str):
25-
# Construct the paths to the possible shared library names
26-
_base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib"
27-
# Searching for the library in the current directory under the name "libllama" (default name
28-
# for llamacpp) and "llama" (default name for this repo)
29-
_lib_paths: List[pathlib.Path] = []
30-
# Determine the file extension based on the platform
31-
if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"):
32-
_lib_paths += [
33-
_base_path / f"lib{lib_base_name}.so",
34-
]
35-
elif sys.platform == "darwin":
36-
_lib_paths += [
37-
_base_path / f"lib{lib_base_name}.so",
38-
_base_path / f"lib{lib_base_name}.dylib",
39-
]
40-
elif sys.platform == "win32":
41-
_lib_paths += [
42-
_base_path / f"{lib_base_name}.dll",
43-
_base_path / f"lib{lib_base_name}.dll",
44-
]
45-
else:
46-
raise RuntimeError("Unsupported platform")
47-
48-
if "LLAMA_CPP_LIB" in os.environ:
49-
lib_base_name = os.environ["LLAMA_CPP_LIB"]
50-
_lib = pathlib.Path(lib_base_name)
51-
_base_path = _lib.parent.resolve()
52-
_lib_paths = [_lib.resolve()]
53-
54-
cdll_args = dict() # type: ignore
55-
56-
# Add the library directory to the DLL search path on Windows (if needed)
57-
if sys.platform == "win32":
58-
os.add_dll_directory(str(_base_path))
59-
os.environ["PATH"] = str(_base_path) + os.pathsep + os.environ["PATH"]
60-
61-
if sys.platform == "win32" and sys.version_info >= (3, 8):
62-
os.add_dll_directory(str(_base_path))
63-
if "CUDA_PATH" in os.environ:
64-
os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "bin"))
65-
os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "lib"))
66-
if "HIP_PATH" in os.environ:
67-
os.add_dll_directory(os.path.join(os.environ["HIP_PATH"], "bin"))
68-
os.add_dll_directory(os.path.join(os.environ["HIP_PATH"], "lib"))
69-
cdll_args["winmode"] = ctypes.RTLD_GLOBAL
70-
71-
# Try to load the shared library, handling potential errors
72-
for _lib_path in _lib_paths:
73-
if _lib_path.exists():
74-
try:
75-
return ctypes.CDLL(str(_lib_path), **cdll_args) # type: ignore
76-
except Exception as e:
77-
raise RuntimeError(f"Failed to load shared library '{_lib_path}': {e}")
78-
79-
raise FileNotFoundError(
80-
f"Shared library with base name '{lib_base_name}' not found"
21+
if TYPE_CHECKING:
22+
from llama_cpp._ctypes_extensions import (
23+
CtypesCData,
24+
CtypesArray,
25+
CtypesPointer,
26+
CtypesVoidPointer,
27+
CtypesRef,
28+
CtypesPointerOrRef,
29+
CtypesFuncPointer,
8130
)
8231

8332

8433
# Specify the base name of the shared library to load
8534
_lib_base_name = "llama"
86-
35+
_override_base_path = os.environ.get("LLAMA_CPP_LIB_PATH")
36+
_base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib" if _override_base_path is None else pathlib.Path(_override_base_path)
8737
# Load the library
88-
_lib = _load_shared_library(_lib_base_name)
89-
90-
91-
# ctypes sane type hint helpers
92-
#
93-
# - Generic Pointer and Array types
94-
# - PointerOrRef type with a type hinted byref function
95-
#
96-
# NOTE: Only use these for static type checking not for runtime checks
97-
# no good will come of that
98-
99-
if TYPE_CHECKING:
100-
CtypesCData = TypeVar("CtypesCData", bound=ctypes._CData) # type: ignore
101-
102-
CtypesArray: TypeAlias = ctypes.Array[CtypesCData] # type: ignore
103-
104-
CtypesPointer: TypeAlias = ctypes._Pointer[CtypesCData] # type: ignore
105-
106-
CtypesVoidPointer: TypeAlias = ctypes.c_void_p
107-
108-
class CtypesRef(Generic[CtypesCData]):
109-
pass
110-
111-
CtypesPointerOrRef: TypeAlias = Union[
112-
CtypesPointer[CtypesCData], CtypesRef[CtypesCData]
113-
]
114-
115-
CtypesFuncPointer: TypeAlias = ctypes._FuncPointer # type: ignore
116-
117-
F = TypeVar("F", bound=Callable[..., Any])
118-
119-
120-
def ctypes_function_for_shared_library(lib: ctypes.CDLL):
121-
def ctypes_function(
122-
name: str, argtypes: List[Any], restype: Any, enabled: bool = True
123-
):
124-
def decorator(f: F) -> F:
125-
if enabled:
126-
func = getattr(lib, name)
127-
func.argtypes = argtypes
128-
func.restype = restype
129-
functools.wraps(f)(func)
130-
return func
131-
else:
132-
return f
133-
134-
return decorator
135-
136-
return ctypes_function
137-
38+
_lib = load_shared_library(_lib_base_name, _base_path)
13839

13940
ctypes_function = ctypes_function_for_shared_library(_lib)
14041

14142

142-
def byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCData]:
143-
"""Type-annotated version of ctypes.byref"""
144-
...
145-
146-
147-
byref = ctypes.byref # type: ignore
148-
14943
# from ggml.h
15044
# // NOTE: always add types at the end of the enum to keep backward compatibility
15145
# enum ggml_type {

0 commit comments

Comments
 (0)