feat: migrate python backends from conda to uv (#2215)

* feat: migrate diffusers backend from conda to uv - replace conda with UV for diffusers install (prototype for all extras backends) - add ability to build docker with one/some/all extras backends instead of all or nothing Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: migrate autogtpq bark coqui from conda to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: convert exllama over to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: migrate exllama2 to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: migrate mamba to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: migrate parler to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: migrate petals to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: fix tests Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: migrate rerankers to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: migrate sentencetransformers to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: install uv for tests-linux Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: make sure file exists before installing on intel images Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: migrate transformers backend to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: migrate transformers-musicgen to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: migrate vall-e-x to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: migrate vllm to uv Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: add uv install to the rest of test-extra.yml Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: adjust file perms on all install/run/test scripts Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: add missing acclerate dependencies Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: add some more missing dependencies to python backends Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: parler tests venv py dir fix Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: correct filename for transformers-musicgen tests Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: adjust the pwd for valle tests Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: cleanup and optimization work for uv migration Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: add setuptools to requirements-install for mamba Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: more size optimization work Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * feat: make installs and tests more consistent, cleanup some deps Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: cleanup Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: mamba backend is cublas only Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> * fix: uncomment lines in makefile Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com> --------- Signed-off-by: Chris Jowett <421501+cryptk@users.noreply.github.com>
2025-05-20 10:35:01 +00:00 · 2024-05-10 08:08:08 -05:00 · 2024-05-10 08:08:08 -05:00 · 28a421cb1d
commit 28a421cb1d
parent e6768097f4
101 changed files with 988 additions and 740 deletions
--- a/backend/python/exllama2/.gitignore
+++ b/backend/python/exllama2/.gitignore
@ -0,0 +1 @@
+source
--- a/backend/python/exllama2/Makefile
+++ b/backend/python/exllama2/Makefile
@ -1,6 +1,5 @@
 .PHONY: exllama2
 exllama2: protogen
-	$(MAKE) -C ../common-env/transformers
 	bash install.sh

 .PHONY: run
@ -17,4 +16,8 @@ protogen-clean:
 	$(RM) backend_pb2_grpc.py backend_pb2.py

 backend_pb2_grpc.py backend_pb2.py:
-	python3 -m grpc_tools.protoc -I../.. --python_out=. --grpc_python_out=. backend.proto
+	python3 -m grpc_tools.protoc -I../.. --python_out=. --grpc_python_out=. backend.proto
+
+.PHONY: clean
+clean: protogen-clean
+	$(RM) -r venv source
--- a/backend/python/exllama2/exllama2.yml
+++ b/backend/python/exllama2/exllama2.yml
@ -1,57 +0,0 @@
-name: exllama2
-channels:
-  - defaults
-dependencies:
-  - _libgcc_mutex=0.1=main
-  - _openmp_mutex=5.1=1_gnu
-  - bzip2=1.0.8=h7b6447c_0
-  - ca-certificates=2023.08.22=h06a4308_0
-  - ld_impl_linux-64=2.38=h1181459_1
-  - libffi=3.4.4=h6a678d5_0
-  - libgcc-ng=11.2.0=h1234567_1
-  - libgomp=11.2.0=h1234567_1
-  - libstdcxx-ng=11.2.0=h1234567_1
-  - libuuid=1.41.5=h5eee18b_0
-  - ncurses=6.4=h6a678d5_0
-  - openssl=3.0.11=h7f8727e_2
-  - pip=23.2.1=py311h06a4308_0
-  - python=3.11.5=h955ad1f_0
-  - readline=8.2=h5eee18b_0
-  - setuptools=68.0.0=py311h06a4308_0
-  - sqlite=3.41.2=h5eee18b_0
-  - tk=8.6.12=h1ccaba5_0
-  - tzdata=2023c=h04d1e81_0
-  - wheel=0.41.2=py311h06a4308_0
-  - xz=5.4.2=h5eee18b_0
-  - zlib=1.2.13=h5eee18b_0
-  - pip:
-      - filelock==3.12.4
-      - fsspec==2023.9.2
-      - grpcio==1.63.0
-      - markupsafe==2.1.3
-      - mpmath==1.3.0
-      - networkx==3.1
-      - protobuf==4.24.4
-      - nvidia-cublas-cu12==12.1.3.1
-      - nvidia-cuda-cupti-cu12==12.1.105
-      - nvidia-cuda-nvrtc-cu12==12.1.105
-      - nvidia-cuda-runtime-cu12==12.1.105
-      - nvidia-cudnn-cu12==8.9.2.26
-      - nvidia-cufft-cu12==11.0.2.54
-      - nvidia-curand-cu12==10.3.2.106
-      - nvidia-cusolver-cu12==11.4.5.107
-      - nvidia-cusparse-cu12==12.1.0.106
-      - nvidia-nccl-cu12==2.18.1
-      - nvidia-nvjitlink-cu12==12.2.140
-      - nvidia-nvtx-cu12==12.1.105
-      - pandas
-      - numpy
-      - ninja
-      - fastparquet
-      - torch>=2.1.0
-      - safetensors>=0.3.2
-      - sentencepiece>=0.1.97
-      - pygments
-      - websockets
-      - regex
-prefix: /opt/conda/envs/exllama2
--- a/backend/python/exllama2/install.sh
+++ b/backend/python/exllama2/install.sh
@ -2,30 +2,42 @@
 set -e
 ##
 ## A bash script installs the required dependencies of VALL-E-X and prepares the environment
-export SHA=c0ddebaaaf8ffd1b3529c2bb654e650bce2f790f
+EXLLAMA2_VERSION=c0ddebaaaf8ffd1b3529c2bb654e650bce2f790f
+
+BUILD_ISOLATION_FLAG=""

 if [ "$BUILD_TYPE" != "cublas" ]; then
-    echo "[exllamav2] Attention!!! Nvidia GPU is required - skipping installation"
+    echo "[exllama] Attention!!! Nvidia GPU is required - skipping installation"
    exit 0
 fi

-export PATH=$PATH:/opt/conda/bin
-source activate transformers
+MY_DIR="$(dirname -- "${BASH_SOURCE[0]}")"

-echo $CONDA_PREFIX
+uv venv ${MY_DIR}/venv
+source ${MY_DIR}/venv/bin/activate

-git clone https://github.com/turboderp/exllamav2 $CONDA_PREFIX/exllamav2
+if [ -f "requirements-install.txt" ]; then
+    # If we have a requirements-install.txt, it means that a package does not properly declare it's build time
+    # dependencies per PEP-517, so we have to set up the proper build environment ourselves, and then install
+    # the package without build isolation
+    BUILD_ISOLATION_FLAG="--no-build-isolation"
+    uv pip install --requirement ${MY_DIR}/requirements-install.txt
+fi

-pushd $CONDA_PREFIX/exllamav2
+uv pip install ${BUILD_ISOLATION_FLAG} --requirement ${MY_DIR}/requirements.txt

-git checkout -b build $SHA
+if [ -f "requirements-${BUILD_TYPE}.txt" ]; then
+    uv pip install ${BUILD_ISOLATION_FLAG} --requirement ${MY_DIR}/requirements-${BUILD_TYPE}.txt
+fi

-# TODO: this needs to be pinned within the conda environments
-pip install -r requirements.txt
+git clone https://github.com/turboderp/exllamav2 $MY_DIR/source
+pushd ${MY_DIR}/source && git checkout -b build ${EXLLAMA2_VERSION} && popd

-popd
+uv pip install ${BUILD_ISOLATION_FLAG} --requirement ${MY_DIR}/source/requirements.txt
+# This installs exllamav2 in JIT mode so it will compile the appropriate torch extension at runtime
+EXLLAMA_NOCOMPILE= uv pip install ${BUILD_ISOLATION_FLAG} ${MY_DIR}/source/

-cp -rfv $CONDA_PREFIX/exllamav2/* ./  
+cp -rfv ./*py $MY_DIR/source/

 if [ "$PIP_CACHE_PURGE" = true ] ; then
    pip cache purge
--- a/backend/python/exllama2/requirements-install.txt
+++ b/backend/python/exllama2/requirements-install.txt
@ -0,0 +1,4 @@
+# This is here to trigger the install script to add --no-build-isolation to the uv pip install commands
+# exllama2 does not specify it's build requirements per PEP517, so we need to provide some things ourselves
+wheel
+setuptools
--- a/backend/python/exllama2/requirements.txt
+++ b/backend/python/exllama2/requirements.txt
@ -0,0 +1,7 @@
+accelerate
+grpcio==1.63.0
+protobuf
+certifi
+torch
+wheel
+setuptools
--- a/backend/python/exllama2/run.sh
+++ b/backend/python/exllama2/run.sh
@ -1,16 +1,10 @@
 #!/bin/bash

 ##
-## A bash script wrapper that runs the exllama server with conda
+## A bash script wrapper that runs the exllama2 server

-export PATH=$PATH:/opt/conda/bin
+MY_DIR="$(dirname -- "${BASH_SOURCE[0]}")"

-# Activate conda environment
-source activate transformers
+source $MY_DIR/venv/bin/activate

-# get the directory where the bash script is located
-DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
-
-cd $DIR
-
-python $DIR/exllama2_backend.py $@
+python $MY_DIR/source/exllama2_backend.py $@
--- a/backend/python/exllama2/test.sh
+++ b/backend/python/exllama2/test.sh
@ -0,0 +1,16 @@
+#!/bin/bash
+##
+## A bash script wrapper that runs python unittests
+
+MY_DIR="$(dirname -- "${BASH_SOURCE[0]}")"
+
+source $MY_DIR/venv/bin/activate
+
+if [ -f "${MY_DIR}/test.py" ]; then
+    pushd ${MY_DIR}
+    python -m unittest test.py
+    popd
+else
+    echo "ERROR: No tests defined for backend!"
+    exit 1
+fi