From c0552bab9222209fcff4958ab88831e70054c7a9 Mon Sep 17 00:00:00 2001
From: Yanlin-Duanmu <68236504+Yanlin-Duanmu@users.noreply.github.com>
Date: Sun, 2 Mar 2025 20:29:21 +0800
Subject: [PATCH] Fix: install.py - Prevent nvmlShutdown from being called when
 uninitialized

Problem:
`check_nvidia_gpu()` crashes with `pynvml.NVMLError_Uninitialized` if NVIDIA GPU is not detected or drivers are missing. This happens because `pynvml.nvmlShutdown()` is called in the `finally` block even if `pynvml.nvmlInit()` failed.

Cause:
`pynvml.nvmlShutdown()` must only be called after successful `pynvml.nvmlInit()`. When `nvmlInit()` fails, the library is uninitialized, and calling `nvmlShutdown()` leads to the error. The `finally` block unconditionally executes `nvmlShutdown()`.

Solution:
Introduce an `initialized` flag. Set it to `True` after successful `pynvml.nvmlInit()`.  Conditionally call `pynvml.nvmlShutdown()` in the `finally` block only when `initialized` is `True`.

Impact:
Prevents crash when NVIDIA GPU is not detected or drivers are missing. Improves error handling and robustness.

No new tests needed as it fixes error flow.
---
 install.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/install.py b/install.py
index 551bca6..1ff0430 100644
--- a/install.py
+++ b/install.py
@@ -19,8 +19,10 @@ def check_nvidia_gpu():
     install_package("pynvml")
     import pynvml
     from translations.translations import translate as t
+    initialized = False
     try:
         pynvml.nvmlInit()
+        initialized = True
         device_count = pynvml.nvmlDeviceGetCount()
         if device_count > 0:
             print(t("Detected NVIDIA GPU(s)"))
@@ -36,7 +38,8 @@ def check_nvidia_gpu():
         print(t("No NVIDIA GPU detected or NVIDIA drivers not properly installed"))
         return False
     finally:
-        pynvml.nvmlShutdown()
+        if initialized:
+            pynvml.nvmlShutdown()
 
 def check_ffmpeg():
     from rich.console import Console