Unverified Commit 1477b44f authored by Yt's avatar Yt Committed by GitHub
Browse files

python3Packages.vllm: enhance CLI UX (#460701)

parents 9bf27aec daaf408d

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

+16 −3
Original line number Diff line number Diff line
@@ -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

+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

  ];