Commit 7c9dec45 authored by Daniel Fahey's avatar Daniel Fahey
Browse files

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.
parent 54e64edf

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

+9 −2
Original line number Diff line number Diff line
@@ -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

  ];