pathfinder: simplify NVML init/shutdown handling and document ref-count contract

rwgk · cursoragent · rwgk · commit fe4b2fe47e93 · 2026-05-04T08:49:49.000-07:00
Collapse the duplicated ``nvmlShutdown()`` calls in ``_query_driver_release_version_text`` into a single ``try/finally`` so the cleanup always runs in one place. The asymmetric error-precedence rule is preserved via ``sys.exc_info()[1]``: when both the NVML body and shutdown fail, the body's error wins (Python keeps the shutdown error on ``__context__`` for debugging); when only shutdown fails, the shutdown error surfaces. Add comments above the matched ``nvmlInit_v2()`` / ``nvmlShutdown()`` pair noting that NVML's init/shutdown is reference-counted, so this balanced pair is safe even when the caller has already initialized NVML elsewhere in the process. Pre-empts a question raised in review on PR #2000. Add two focused tests filling out the cleanup matrix: - ``test_query_driver_release_version_text_raises_when_only_shutdown_fails`` asserts a non-zero shutdown status surfaces when the body succeeded. - ``test_query_driver_release_version_text_body_error_wins_when_both_fail`` locks in the body-error-wins precedence when both calls fail. Co-authored-by: Cursor <cursoragent@cursor.com>
diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/driver_info.py b/cuda_pathfinder/cuda/pathfinder/_utils/driver_info.py
@@ -6,6 +6,7 @@
 import ctypes
 import functools
 import re
+import sys
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import cast
@@ -140,6 +141,9 @@ def _query_driver_release_version_text() -> str:
     nvml_shutdown.argtypes = []
     nvml_shutdown.restype = ctypes.c_int
 
+    # NVML's init/shutdown pair is reference-counted (see "Initialization and
+    # Cleanup" in the NVML API docs), so this balanced pair is safe even when
+    # the caller has already initialized NVML elsewhere in the process.
     init_status = nvml_init_v2()
     if init_status != _NVML_SUCCESS:
         raise RuntimeError(f"Failed to initialize NVML via nvmlInit_v2() (status={init_status}).")
@@ -152,13 +156,13 @@ def _query_driver_release_version_text() -> str:
                 f"Failed to query driver release version via nvmlSystemGetDriverVersion() (status={status})."
             )
         release_version = version_buffer.value.decode()
-    except BaseException as exc:
+    finally:
+        # Balance the init_v2() above unconditionally. If the body already
+        # raised, let that error win; a non-zero shutdown status here would
+        # only mask the more useful root cause (Python keeps it on
+        # ``__context__`` for debugging). ``sys.exc_info()[1]`` is the
+        # currently-propagating exception inside the finally, or None.
         shutdown_status = nvml_shutdown()
-        if shutdown_status != _NVML_SUCCESS:
-            raise RuntimeError(f"Failed to shut down NVML via nvmlShutdown() (status={shutdown_status}).") from exc
-        raise
-
-    shutdown_status = nvml_shutdown()
-    if shutdown_status != _NVML_SUCCESS:
-        raise RuntimeError(f"Failed to shut down NVML via nvmlShutdown() (status={shutdown_status}).")
+        if shutdown_status != _NVML_SUCCESS and sys.exc_info()[1] is None:
+            raise RuntimeError(f"Failed to shut down NVML via nvmlShutdown() (status={shutdown_status}).")
     return release_version
diff --git a/cuda_pathfinder/tests/test_utils_driver_info.py b/cuda_pathfinder/tests/test_utils_driver_info.py
@@ -155,6 +155,36 @@ def test_query_driver_release_version_text_raises_when_nvml_call_fails(monkeypat
     assert fake_nvml_lib.shutdown_calls == 1
 
 
+def test_query_driver_release_version_text_raises_when_only_shutdown_fails(monkeypatch):
+    fake_nvml_lib = _FakeNvmlLib(shutdown_statuses=(2,))
+
+    monkeypatch.setattr(
+        driver_info,
+        "_load_nvidia_dynamic_lib",
+        lambda _libname: _loaded_nvml("/usr/lib/libnvidia-ml.so.1"),
+    )
+    monkeypatch.setattr(driver_info.ctypes, "CDLL", lambda _abs_path: fake_nvml_lib)
+
+    with pytest.raises(RuntimeError, match=r"nvmlShutdown\(\) \(status=2\)"):
+        driver_info._query_driver_release_version_text()
+    assert fake_nvml_lib.shutdown_calls == 1
+
+
+def test_query_driver_release_version_text_body_error_wins_when_both_fail(monkeypatch):
+    fake_nvml_lib = _FakeNvmlLib(query_status=1, shutdown_statuses=(2,))
+
+    monkeypatch.setattr(
+        driver_info,
+        "_load_nvidia_dynamic_lib",
+        lambda _libname: _loaded_nvml("/usr/lib/libnvidia-ml.so.1"),
+    )
+    monkeypatch.setattr(driver_info.ctypes, "CDLL", lambda _abs_path: fake_nvml_lib)
+
+    with pytest.raises(RuntimeError, match=r"nvmlSystemGetDriverVersion\(\) \(status=1\)"):
+        driver_info._query_driver_release_version_text()
+    assert fake_nvml_lib.shutdown_calls == 1
+
+
 def test_query_driver_cuda_version_uses_windll_on_windows(monkeypatch):
     fake_driver_lib = _FakeDriverLib(status=0, version=12080)
     loaded_paths: list[str] = []
