From 7c9dec453c7d3dd8bbf7119b86d85a9e3d22d2f7 Mon Sep 17 00:00:00 2001
From: Daniel Fahey <daniel.fahey@pm.me>
Date: Mon, 10 Nov 2025 15:13:04 +0000
Subject: [PATCH] python3Packages.vllm: unvendor `pynvml`

- Replaces `pynvml` dependency with `nvidia-ml-py`
- Patches vLLM to use system `pynvml` from `nvidia-ml-py` instead of its vendored copy

Removes need to manually set `LD_LIBRARY_PATH=/run/opengl-driver/lib` when running `vllm` commands, as `nvidia-ml-py` includes patches to locate `libnvidia-ml.so.1` in Nix environments.
---
 pkgs/development/python-modules/vllm/default.nix | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/pkgs/development/python-modules/vllm/default.nix b/pkgs/development/python-modules/vllm/default.nix
index ca759392b585..b8b2e7f59d3b 100644
--- a/pkgs/development/python-modules/vllm/default.nix
+++ b/pkgs/development/python-modules/vllm/default.nix
@@ -34,7 +34,7 @@
   uvicorn,
   pydantic,
   aioprometheus,
-  pynvml,
+  nvidia-ml-py,
   openai,
   pyzmq,
   tiktoken,
@@ -314,6 +314,13 @@ buildPythonPackage rec {
   ];
 
   postPatch = ''
+    # Remove vendored pynvml entirely
+    rm vllm/third_party/pynvml.py
+    substituteInPlace tests/utils.py \
+      --replace-fail "from vllm.third_party.pynvml import" "from pynvml import"
+    substituteInPlace vllm/utils/__init__.py \
+      --replace-fail "import vllm.third_party.pynvml" "import pynvml"
+
     # pythonRelaxDeps does not cover build-system
     substituteInPlace pyproject.toml \
       --replace-fail "torch ==" "torch >=" \
@@ -442,7 +449,7 @@ buildPythonPackage rec {
   ]
   ++ lib.optionals cudaSupport [
     cupy
-    pynvml
+    nvidia-ml-py
     flashinfer
   ];