SDAFOASDFSAKFSDA

app.py

@@ -11,6 +11,7 @@ def ensure_cuda_runtime():
     and restarts the script with the correct LD_LIBRARY_PATH.
     """
     missing_libs = []
+    lib_paths = []
     
     # 1. Check for libcudart (CUDA Runtime)
     try:
@@ -26,17 +27,47 @@ def ensure_cuda_runtime():
         print("libcublas.so.12 not found.")
         missing_libs.append("nvidia-cublas-cu12")
 
-    if not missing_libs:
-        return  # All libraries found
-
-    print(f"Missing CUDA libraries. Installing: {', '.join(missing_libs)}...")
+    # 3. Check for libcuda.so.1 (NVIDIA Driver)
+    # This CANNOT be installed via pip. It must be mounted from the host.
+    # If missing, we must fall back to CPU mode.
+    driver_found = False
+    try:
+        ctypes.CDLL("libcuda.so.1")
+        driver_found = True
+    except OSError:
+        print("libcuda.so.1 (NVIDIA Driver) not found via dlopen.")
+        # Search common paths
+        search_paths = [
+            "/usr/lib/x86_64-linux-gnu",
+            "/usr/lib64",
+            "/usr/lib",
+            "/usr/local/cuda/lib64",
+            "/usr/lib/wsl/lib"
+        ]
+        for path in search_paths:
+            if os.path.exists(os.path.join(path, "libcuda.so.1")):
+                print(f"Found libcuda.so.1 manually at {path}")
+                lib_paths.append(path)
+                driver_found = True
+                break
     
-    # 3. Install missing packages
-    subprocess.check_call([sys.executable, "-m", "pip", "install"] + missing_libs)
+    if not driver_found and not missing_libs:
+        # If we have libs but no driver, we can't run GPU.
+        # We will set an env var to force CPU install later.
+        print("CRITICAL: NVIDIA Driver (libcuda.so.1) not found. GPU acceleration will fail.")
+        print("Switching to CPU-only mode for this run.")
+        os.environ["FORCE_CPU_MODE"] = "1"
+
+    if not missing_libs and driver_found:
+        return  # All libraries and driver found
+
+    if missing_libs:
+        print(f"Missing CUDA libraries. Installing: {', '.join(missing_libs)}...")
+        # 3. Install missing packages
+        subprocess.check_call([sys.executable, "-m", "pip", "install"] + missing_libs)
 
     # 4. Find library paths
     import site
-    lib_paths = []
     for sp in site.getsitepackages():
         # Runtime libs
         rt_path = os.path.join(sp, "nvidia", "cuda_runtime", "lib")
@@ -48,7 +79,7 @@ def ensure_cuda_runtime():
         if os.path.isdir(cublas_path):
             lib_paths.append(cublas_path)
 
-    if not lib_paths:
+    if not lib_paths and missing_libs:
         print("Warning: Could not find nvidia lib paths after installation.")
         return
 
@@ -68,8 +99,12 @@ def ensure_cuda_runtime():
             new_ld = f"{new_ld}{os.pathsep}{current_ld}"
             
         os.environ["LD_LIBRARY_PATH"] = new_ld
+        
+        # Pass the FORCE_CPU_MODE flag to the restarted process if set
+        env = os.environ.copy()
+        
         # Re-execute the current script with the new environment
-        os.execv(sys.executable, [sys.executable] + sys.argv)
+        os.execvpe(sys.executable, [sys.executable] + sys.argv, env)
 
 ensure_cuda_runtime()
 # --- CUDA DEPENDENCY FIX END ---
@@ -84,14 +119,24 @@ try:
     print("llama-cpp-python is correctly installed.")
 except (ImportError, RuntimeError, OSError) as e:
     print(f"llama-cpp-python needs installation or repair: {e}")
-    print("Installing llama-cpp-python from pre-built wheel index...")
-    subprocess.check_call([
-        sys.executable, "-m", "pip", "install", 
-        "llama-cpp-python", 
-        "--extra-index-url", "https://abetlen.github.io/llama-cpp-python/whl/cu121",
-        "--prefer-binary",
-        "--force-reinstall"
-    ])
+    
+    if os.environ.get("FORCE_CPU_MODE") == "1":
+        print("Installing CPU-only llama-cpp-python (Fallback)...")
+        subprocess.check_call([
+            sys.executable, "-m", "pip", "install", 
+            "llama-cpp-python", 
+            "--prefer-binary",
+            "--force-reinstall"
+        ])
+    else:
+        print("Installing llama-cpp-python from pre-built wheel index (CUDA)...")
+        subprocess.check_call([
+            sys.executable, "-m", "pip", "install", 
+            "llama-cpp-python", 
+            "--extra-index-url", "https://abetlen.github.io/llama-cpp-python/whl/cu121",
+            "--prefer-binary",
+            "--force-reinstall"
+        ])
     print("Installation complete.")
 
 # --- IMPORTS AFTER INSTALL ---