restore PATH after loading dlls

abetlen · abetlen · commit d53bce2d9541 · 2024-07-05T21:44:03.000-04:00
diff --git a/llama_cpp/llama_cpp.py b/llama_cpp/llama_cpp.py
@@ -3,8 +3,9 @@
 import sys
 import os
 import ctypes
-import functools
 import pathlib
+import functools
+import contextlib
 
 from typing import (
     Any,
@@ -22,6 +23,7 @@
 
 # Load the library
 def _load_shared_library(lib_base_name: str):
+    _exit_stack = contextlib.ExitStack()
     # Construct the paths to the possible shared library names
     _base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib"
     # 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):
     # Add the library directory to the DLL search path on Windows (if needed)
     if sys.platform == "win32":
         os.add_dll_directory(str(_base_path))
+        old_path = os.environ.get("PATH", "")
+        def restore_path():
+            os.environ["PATH"] = old_path
+        # required to restore the PATH in case of an exception
         os.environ['PATH'] = str(_base_path) + os.pathsep + os.environ['PATH']
+        _exit_stack.callback(restore_path)
 
     if sys.platform == "win32" and sys.version_info >= (3, 8):
         os.add_dll_directory(str(_base_path))
diff --git a/llama_cpp/llava_cpp.py b/llama_cpp/llava_cpp.py
@@ -1,9 +1,10 @@
 from __future__ import annotations
 
-import sys
 import os
+import sys
 import ctypes
 import functools
+import contextlib
 from ctypes import (
     c_bool,
     c_char_p,
@@ -34,6 +35,7 @@
 
 # Load the library
 def _load_shared_library(lib_base_name: str):
+    _exit_stack = contextlib.ExitStack()
     # Construct the paths to the possible shared library names
     _base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib"
     # 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):
             os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "bin"))
             os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "lib"))
         cdll_args["winmode"] = ctypes.RTLD_GLOBAL
+        old_path = os.environ.get("PATH", "")
+        def restore_path():
+            os.environ["PATH"] = old_path
+        # required to restore the PATH in case of an exception
+        os.environ['PATH'] = str(_base_path) + os.pathsep + os.environ['PATH']
+        _exit_stack.callback(restore_path)
 
     # Try to load the shared library, handling potential errors
     for _lib_path in _lib_paths:
