Skip to content

Commit d53bce2

Browse files
committed
restore PATH after loading dlls
1 parent 62804ee commit d53bce2

File tree

2 files changed

+17
-2
lines changed

2 files changed

+17
-2
lines changed

llama_cpp/llama_cpp.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,9 @@
33
import sys
44
import os
55
import ctypes
6-
import functools
76
import pathlib
7+
import functools
8+
import contextlib
89

910
from typing import (
1011
Any,
@@ -22,6 +23,7 @@
2223

2324
# Load the library
2425
def _load_shared_library(lib_base_name: str):
26+
_exit_stack = contextlib.ExitStack()
2527
# Construct the paths to the possible shared library names
2628
_base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib"
2729
# Searching for the library in the current directory under the name "libllama" (default name
@@ -56,7 +58,12 @@ def _load_shared_library(lib_base_name: str):
5658
# Add the library directory to the DLL search path on Windows (if needed)
5759
if sys.platform == "win32":
5860
os.add_dll_directory(str(_base_path))
61+
old_path = os.environ.get("PATH", "")
62+
def restore_path():
63+
os.environ["PATH"] = old_path
64+
# required to restore the PATH in case of an exception
5965
os.environ['PATH'] = str(_base_path) + os.pathsep + os.environ['PATH']
66+
_exit_stack.callback(restore_path)
6067

6168
if sys.platform == "win32" and sys.version_info >= (3, 8):
6269
os.add_dll_directory(str(_base_path))

llama_cpp/llava_cpp.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
from __future__ import annotations
22

3-
import sys
43
import os
4+
import sys
55
import ctypes
66
import functools
7+
import contextlib
78
from ctypes import (
89
c_bool,
910
c_char_p,
@@ -34,6 +35,7 @@
3435

3536
# Load the library
3637
def _load_shared_library(lib_base_name: str):
38+
_exit_stack = contextlib.ExitStack()
3739
# Construct the paths to the possible shared library names
3840
_base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib"
3941
# Searching for the library in the current directory under the name "libllama" (default name
@@ -71,6 +73,12 @@ def _load_shared_library(lib_base_name: str):
7173
os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "bin"))
7274
os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "lib"))
7375
cdll_args["winmode"] = ctypes.RTLD_GLOBAL
76+
old_path = os.environ.get("PATH", "")
77+
def restore_path():
78+
os.environ["PATH"] = old_path
79+
# required to restore the PATH in case of an exception
80+
os.environ['PATH'] = str(_base_path) + os.pathsep + os.environ['PATH']
81+
_exit_stack.callback(restore_path)
7482

7583
# Try to load the shared library, handling potential errors
7684
for _lib_path in _lib_paths:

0 commit comments

Comments
 (0)