python3Packages.vllm: enhance CLI UX (#460701)

pkgs/development/python-modules/flashinfer/default.nix

@@ -9,6 +9,7 @@
   config,
   buildPythonPackage,
   fetchFromGitHub,
+  fetchpatch2,
 
   # build-system
   setuptools,
@@ -22,7 +23,7 @@
   click,
   einops,
   numpy,
-  pynvml,
+  nvidia-ml-py,
   tabulate,
   torch,
   tqdm,
@@ -41,6 +42,15 @@ buildPythonPackage rec {
     hash = "sha256-e9PfLfU0DdoLKlXiHylCbGd125c7Iw9y4NDIOAP0xHs=";
   };
 
+  patches = [
+    # TODO: remove patch with update to v0.5.2+
+    # Switch pynvml to nvidia-ml-py
+    (fetchpatch2 {
+      url = "https://github.com/flashinfer-ai/flashinfer/commit/a42f99255d68d1a54b689bd4985339c6b44963a6.patch?full_index=1";
+      hash = "sha256-3XJFcdQeZ/c5fwiQvd95z4p9BzTn8pjle21WzeBxUgk=";
+    })
+  ];
+
   build-system = [ setuptools ];
 
   nativeBuildInputs = [
@@ -86,7 +96,7 @@ buildPythonPackage rec {
     click
     einops
     numpy
-    pynvml
+    nvidia-ml-py
     tabulate
     torch
     tqdm
@@ -104,6 +114,9 @@ buildPythonPackage rec {
       scenarios.
     '';
     license = lib.licenses.asl20;
-    maintainers = with lib.maintainers; [ breakds ];
+    maintainers = with lib.maintainers; [
+      breakds
+      daniel-fahey
+    ];
   };
 }

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
   ];