File tree Expand file tree Collapse file tree 2 files changed +17
-2
lines changed Expand file tree Collapse file tree 2 files changed +17
-2
lines changed Original file line number Diff line number Diff line change 3
3
import sys
4
4
import os
5
5
import ctypes
6
- import functools
7
6
import pathlib
7
+ import functools
8
+ import contextlib
8
9
9
10
from typing import (
10
11
Any ,
22
23
23
24
# Load the library
24
25
def _load_shared_library (lib_base_name : str ):
26
+ _exit_stack = contextlib .ExitStack ()
25
27
# Construct the paths to the possible shared library names
26
28
_base_path = pathlib .Path (os .path .abspath (os .path .dirname (__file__ ))) / "lib"
27
29
# 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):
56
58
# Add the library directory to the DLL search path on Windows (if needed)
57
59
if sys .platform == "win32" :
58
60
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
59
65
os .environ ['PATH' ] = str (_base_path ) + os .pathsep + os .environ ['PATH' ]
66
+ _exit_stack .callback (restore_path )
60
67
61
68
if sys .platform == "win32" and sys .version_info >= (3 , 8 ):
62
69
os .add_dll_directory (str (_base_path ))
Original file line number Diff line number Diff line change 1
1
from __future__ import annotations
2
2
3
- import sys
4
3
import os
4
+ import sys
5
5
import ctypes
6
6
import functools
7
+ import contextlib
7
8
from ctypes import (
8
9
c_bool ,
9
10
c_char_p ,
34
35
35
36
# Load the library
36
37
def _load_shared_library (lib_base_name : str ):
38
+ _exit_stack = contextlib .ExitStack ()
37
39
# Construct the paths to the possible shared library names
38
40
_base_path = pathlib .Path (os .path .abspath (os .path .dirname (__file__ ))) / "lib"
39
41
# 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):
71
73
os .add_dll_directory (os .path .join (os .environ ["CUDA_PATH" ], "bin" ))
72
74
os .add_dll_directory (os .path .join (os .environ ["CUDA_PATH" ], "lib" ))
73
75
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 )
74
82
75
83
# Try to load the shared library, handling potential errors
76
84
for _lib_path in _lib_paths :
You can’t perform that action at this time.
0 commit comments