From 23beb40a0d687a730952dd2e64384428ac955b44 Mon Sep 17 00:00:00 2001
From: cfhammill <christopher.hammill@unityhealth.to>
Date: Tue, 9 Jul 2024 16:15:47 -0400
Subject: [PATCH] python3Packages.vllm: relax deps, add cuda stdenv

---
 .../python-modules/vllm/default.nix           | 20 ++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/pkgs/development/python-modules/vllm/default.nix b/pkgs/development/python-modules/vllm/default.nix
index 7ccab0a23129..23717ec9570a 100644
--- a/pkgs/development/python-modules/vllm/default.nix
+++ b/pkgs/development/python-modules/vllm/default.nix
@@ -1,5 +1,6 @@
 {
   lib,
+  stdenv,
   buildPythonPackage,
   fetchFromGitHub,
   which,
@@ -35,6 +36,10 @@
   gpuTargets ? [ ],
 }:
 
+let
+  stdenv_pkg = stdenv;
+in
+
 buildPythonPackage rec {
   pname = "vllm";
   version = "0.3.3";
@@ -52,25 +57,24 @@ buildPythonPackage rec {
     lib.strings.concatStringsSep ";" rocmPackages.clr.gpuTargets
   );
 
-  # xformers 0.0.23.post1 github release specifies its version as 0.0.24
-  #
   # 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 == 0.0.24"
+        --replace "xformers == 0.0.23.post1" "xformers"
       substituteInPlace requirements.txt \
-        --replace "cupy-cuda12x == 12.1.0" "cupy == 12.3.0"
+        --replace "cupy-cuda12x == 12.1.0" "cupy"
       substituteInPlace requirements-build.txt \
-        --replace "torch==2.1.2" "torch == 2.2.1"
+        --replace "torch==2.1.2" "torch"
       substituteInPlace pyproject.toml \
-        --replace "torch == 2.1.2" "torch == 2.2.1"
+        --replace "torch == 2.1.2" "torch"
       substituteInPlace requirements.txt \
-        --replace "torch == 2.1.2" "torch == 2.2.1"
+        --replace "torch == 2.1.2" "torch"
     ''
     + lib.optionalString rocmSupport ''
       substituteInPlace setup.py \
@@ -141,6 +145,8 @@ buildPythonPackage rec {
       cupy
     ];
 
+  stdenv = if cudaSupport then cudaPackages.backendStdenv else stdenv_pkg;
+
   pythonImportsCheck = [ "vllm" ];
 
   meta = with lib; {