diff --git a/pkgs/development/python-modules/flashinfer/default.nix b/pkgs/development/python-modules/flashinfer/default.nix index 35a729767ce8..a0a5a107d0ec 100644 --- a/pkgs/development/python-modules/flashinfer/default.nix +++ b/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 + ]; }; } 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 ];