python3Packages.vllm: 0.3.3->0.5.2

pkgs/development/python-modules/lm-format-enforcer/default.nix

@@ -0,0 +1,42 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+  pydantic,
+  interegular,
+  pyyaml,
+  poetry-core,
+}:
+
+buildPythonPackage rec {
+  pname = "lm-format-enforcer";
+  version = "0.10.4";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "noamgat";
+    repo = pname;
+    rev = "refs/tags/v${version}";
+    hash = "sha256-GOnMj910rgzYeIeN2yLcXZDDel/Hu6nv7ov5BrlHJLg=";
+  };
+
+  build-system = [ poetry-core ];
+
+  dependencies = [
+    interegular
+    pydantic
+    pyyaml
+  ];
+
+  doCheck = false; # most tests require internet access
+
+  pythonImportsCheck = [ "lmformatenforcer" ];
+
+  meta = with lib; {
+    description = "Enforce the output format (JSON Schema, Regex etc) of a language model";
+    changelog = "https://github.com/noamgat/lm-format-enforcer/releases/tag/v${version}";
+    homepage = "https://github.com/noamgat/lm-format-enforcer";
+    license = licenses.mit;
+    maintainers = with maintainers; [ cfhammill ];
+  };
+}

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

@@ -6,9 +6,12 @@
   setuptools-scm,
   interegular,
   cloudpickle,
+  datasets,
   diskcache,
   joblib,
   jsonschema,
+  pyairports,
+  pycountry,
   pydantic,
   lark,
   nest-asyncio,
@@ -38,6 +41,7 @@ buildPythonPackage rec {
   propagatedBuildInputs = [
     interegular
     cloudpickle
+    datasets
     diskcache
     joblib
     jsonschema
@@ -48,16 +52,19 @@ buildPythonPackage rec {
     scipy
     torch
     transformers
+    pycountry
+    pyairports
   ];
 
-  pythonImportsCheck = [ "outlines" ];
+  checkPhase = ''
+    export HOME=$(mktemp -d)
+    python3 -c 'import outlines'
+  '';
 
   meta = with lib; {
     description = "Structured text generation";
     homepage = "https://github.com/outlines-dev/outlines";
     license = licenses.asl20;
     maintainers = with maintainers; [ lach ];
-    # Missing dependencies since the last update
-    broken = true;
   };
 }

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

@@ -0,0 +1,30 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchPypi,
+  setuptools,
+}:
+
+buildPythonPackage rec {
+  pname = "pyairports";
+  version = "2.1.1";
+  pyproject = true;
+
+  src = fetchPypi {
+    inherit pname version;
+    hash = "sha256-PWCnJ/zk2oG5xjk+qK4LM9Z7N+zjRN/8hj90njrWK80=";
+  };
+
+  build-system = [ setuptools ];
+
+  doCheck = false;
+
+  pythonImportChecks = [ "pyairports" ];
+
+  meta = with lib; {
+    description = "pyairports is a package which enables airport lookup by 3-letter IATA code.";
+    homepage = "https://github.com/ozeliger/pyairports";
+    license = licenses.asl20;
+    maintainers = with maintainers; [ cfhammill ];
+  };
+}

pkgs/development/python-modules/vllm/0001-setup.py-don-t-ask-for-hipcc-version.patch

@@ -0,0 +1,24 @@
+From f6a7748bee79fc2e1898968fef844daacfa7860b Mon Sep 17 00:00:00 2001
+From: SomeoneSerge <else@someonex.net>
+Date: Wed, 31 Jul 2024 12:02:53 +0000
+Subject: [PATCH 1/2] setup.py: don't ask for hipcc --version
+
+---
+ setup.py | 1 +
+ 1 file changed, 1 insertion(+)
+
+diff --git a/setup.py b/setup.py
+index 72ef26f1..01e006f9 100644
+--- a/setup.py
++++ b/setup.py
+@@ -279,6 +279,7 @@ def _install_punica() -> bool:
+
+
+ def get_hipcc_rocm_version():
++    return "0.0" # `hipcc --version` misbehaves ("unresolved paths") inside the nix sandbox
+     # Run the hipcc --version command
+     result = subprocess.run(['hipcc', '--version'],
+                             stdout=subprocess.PIPE,
+-- 
+2.45.1
+

pkgs/development/python-modules/vllm/0002-setup.py-nix-support-respect-cmakeFlags.patch

@@ -0,0 +1,40 @@
+From 10b7e8330bdba319a4162cceb8e5dd4280215b04 Mon Sep 17 00:00:00 2001
+From: SomeoneSerge <else@someonex.net>
+Date: Wed, 31 Jul 2024 12:06:15 +0000
+Subject: [PATCH 2/2] setup.py: nix-support (respect cmakeFlags)
+
+---
+ setup.py | 10 ++++++++++
+ 1 file changed, 10 insertions(+)
+
+diff --git a/setup.py b/setup.py
+index 01e006f9..14762146 100644
+--- a/setup.py
++++ b/setup.py
+@@ -15,6 +15,15 @@ from setuptools import Extension, find_packages, setup
+ from setuptools.command.build_ext import build_ext
+ from torch.utils.cpp_extension import CUDA_HOME
+
++import os
++import json
++
++if "NIX_ATTRS_JSON_FILE" in os.environ:
++    with open(os.environ["NIX_ATTRS_JSON_FILE"], "r") as f:
++        NIX_ATTRS = json.load(f)
++else:
++    NIX_ATTRS = { "cmakeFlags": os.environ.get("cmakeFlags", "").split() }
++
+
+ def load_module_from_path(module_name, path):
+     spec = importlib.util.spec_from_file_location(module_name, path)
+@@ -159,6 +168,7 @@ class cmake_build_ext(build_ext):
+             '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={}'.format(outdir),
+             '-DCMAKE_ARCHIVE_OUTPUT_DIRECTORY={}'.format(self.build_temp),
+             '-DVLLM_TARGET_DEVICE={}'.format(VLLM_TARGET_DEVICE),
++            *NIX_ATTRS["cmakeFlags"],
+         ]
+
+         verbose = envs.VERBOSE
+-- 
+2.45.1
+

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

@@ -1,10 +1,13 @@
 {
   lib,
   stdenv,
+  python,
   buildPythonPackage,
+  pythonRelaxDepsHook,
   fetchFromGitHub,
   which,
   ninja,
+  cmake,
   packaging,
   setuptools,
   torch,
@@ -23,6 +26,13 @@
   pydantic,
   aioprometheus,
   pynvml,
+  openai,
+  pyzmq,
+  tiktoken,
+  torchvision,
+  py-cpuinfo,
+  lm-format-enforcer,
+  prometheus-fastapi-instrumentator,
   cupy,
   writeShellScript,
 
@@ -31,83 +41,73 @@
   cudaSupport ? config.cudaSupport,
   cudaPackages ? { },
 
-  rocmSupport ? config.rocmSupport,
+  # Has to be either rocm or cuda, default to the free one
+  rocmSupport ? !config.cudaSupport,
   rocmPackages ? { },
   gpuTargets ? [ ],
-}:
+}@args:
 
 let
-  stdenv_pkg = stdenv;
+  cutlass = fetchFromGitHub {
+    owner = "NVIDIA";
+    repo = "cutlass";
+    rev = "refs/tags/v3.5.0";
+    sha256 = "sha256-D/s7eYsa5l/mfx73tE4mnFcTQdYqGmXa9d9TCryw4e4=";
+  };
 in
 
 buildPythonPackage rec {
   pname = "vllm";
-  version = "0.3.3";
-  format = "pyproject";
+  version = "0.5.3.post1";
+  pyproject = true;
+
+  stdenv = if cudaSupport then cudaPackages.backendStdenv else args.stdenv;
 
   src = fetchFromGitHub {
     owner = "vllm-project";
     repo = pname;
-    rev = "v${version}";
-    hash = "sha256-LU5pCPVv+Ws9dL8oWL1sJGzwQKI1IFk2A1I6TP9gXL4=";
+    rev = "refs/tags/v${version}";
+    hash = "sha256-++DK2Y2zz+1KrEcdQc5XFrSjc7fCwMD2DQ/RqY7PoFU=";
   };
 
-  # Otherwise it tries to enumerate host supported ROCM gfx archs, and that is not possible due to sandboxing.
-  PYTORCH_ROCM_ARCH = lib.optionalString rocmSupport (
-    lib.strings.concatStringsSep ";" rocmPackages.clr.gpuTargets
-  );
-
-  # cupy-cuda12x is the same wheel as cupy, but built with cuda dependencies, we already have it set up
-  # like that in nixpkgs. Version upgrade is due to upstream shenanigans
-  # https://github.com/vllm-project/vllm/pull/2845/commits/34a0ad7f9bb7880c0daa2992d700df3e01e91363
-  #
-  # hipcc --version works badly on NixOS due to unresolved paths.
-  # Unclear why pythonRelaxDeps doesn't work here, but on last attempt, it didn't.
-  postPatch =
-    ''
-      substituteInPlace requirements.txt \
-        --replace "xformers == 0.0.23.post1" "xformers"
-      substituteInPlace requirements.txt \
-        --replace "cupy-cuda12x == 12.1.0" "cupy"
-      substituteInPlace requirements-build.txt \
-        --replace "torch==2.1.2" "torch"
-      substituteInPlace pyproject.toml \
-        --replace "torch == 2.1.2" "torch"
-      substituteInPlace requirements.txt \
-        --replace "torch == 2.1.2" "torch"
-    ''
-    + lib.optionalString rocmSupport ''
-      substituteInPlace setup.py \
-        --replace "'hipcc', '--version'" "'${writeShellScript "hipcc-version-stub" "echo HIP version: 0.0"}'"
-    '';
-
-  preBuild =
-    lib.optionalString cudaSupport ''
-      export CUDA_HOME=${cudaPackages.cuda_nvcc}
-    ''
-    + lib.optionalString rocmSupport ''
-      export ROCM_HOME=${rocmPackages.clr}
-      export PATH=$PATH:${rocmPackages.hipcc}
-    '';
+  patches = [
+    ./0001-setup.py-don-t-ask-for-hipcc-version.patch
+    ./0002-setup.py-nix-support-respect-cmakeFlags.patch
+  ];
+
+  # Ignore the python version check because it hard-codes minor versions and
+  # lags behind `ray`'s python interpreter support
+  postPatch = ''
+    substituteInPlace CMakeLists.txt \
+      --replace-fail \
+        'set(PYTHON_SUPPORTED_VERSIONS' \
+        'set(PYTHON_SUPPORTED_VERSIONS "${lib.versions.majorMinor python.version}"'
+  '';
 
   nativeBuildInputs = [
+    cmake
     ninja
+    pythonRelaxDepsHook
+    which
+  ] ++ lib.optionals rocmSupport [ rocmPackages.hipcc ];
+
+  build-system = [
     packaging
     setuptools
-    torch
     wheel
-    which
-  ] ++ lib.optionals rocmSupport [ rocmPackages.hipcc ];
+  ];
 
   buildInputs =
     (lib.optionals cudaSupport (
       with cudaPackages;
       [
         cuda_cudart # cuda_runtime.h, -lcudart
-        cuda_cccl # <thrust/*>
+        cuda_cccl
         libcusparse # cusparse.h
-        libcublas # cublas_v2.h
         libcusolver # cusolverDn.h
+        cuda_nvcc
+        cuda_nvtx
+        libcublas
       ]
     ))
     ++ (lib.optionals rocmSupport (
@@ -121,31 +121,49 @@ buildPythonPackage rec {
       ]
     ));
 
-  propagatedBuildInputs =
+  dependencies =
     [
-      psutil
-      ray
+      aioprometheus
+      fastapi
+      lm-format-enforcer
+      numpy
+      openai
+      outlines
       pandas
+      prometheus-fastapi-instrumentator
+      psutil
+      py-cpuinfo
       pyarrow
+      pydantic
+      pyzmq
+      ray
       sentencepiece
-      numpy
+      tiktoken
       torch
+      torchvision
       transformers
-      outlines
-      xformers
-      fastapi
       uvicorn
-      pydantic
-      aioprometheus
+      xformers
     ]
     ++ uvicorn.optional-dependencies.standard
     ++ aioprometheus.optional-dependencies.starlette
     ++ lib.optionals cudaSupport [
-      pynvml
       cupy
+      pynvml
     ];
 
-  stdenv = if cudaSupport then cudaPackages.backendStdenv else stdenv_pkg;
+  dontUseCmakeConfigure = true;
+  cmakeFlags = [ (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${lib.getDev cutlass}") ];
+
+  env =
+    lib.optionalAttrs cudaSupport { CUDA_HOME = "${lib.getDev cudaPackages.cuda_nvcc}"; }
+    // lib.optionalAttrs rocmSupport {
+      # Otherwise it tries to enumerate host supported ROCM gfx archs, and that is not possible due to sandboxing.
+      PYTORCH_ROCM_ARCH = lib.strings.concatStringsSep ";" rocmPackages.clr.gpuTargets;
+      ROCM_HOME = "${rocmPackages.clr}";
+    };
+
+  pythonRelaxDeps = true;
 
   pythonImportsCheck = [ "vllm" ];