feat: migrate python backends from conda to uv

.dockerignore

@@ -6,6 +6,11 @@ examples/chatbot-ui/models
 examples/rwkv/models
 examples/**/models
 Dockerfile*
+__pycache__
 
 # SonarQube
-.scannerwork
+.scannerwork
+
+# backend virtual environments
+**/venv
+backend/python/**/source

.github/workflows/test.yml

@@ -78,6 +78,8 @@ jobs:
              sudo /bin/bash -c 'echo "deb [arch=amd64 signed-by=/usr/share/keyrings/conda-archive-keyring.gpg] https://repo.anaconda.com/pkgs/misc/debrepo/conda stable main" | tee -a /etc/apt/sources.list.d/conda.list' && \
              sudo apt-get update && \
              sudo apt-get install -y conda
+          # Install UV
+          curl -LsSf https://astral.sh/uv/install.sh | sh
           sudo apt-get install -y ca-certificates cmake patch python3-pip unzip
           sudo apt-get install -y libopencv-dev
 

.gitignore

@@ -46,4 +46,7 @@ prepare
 *pb2_grpc.py
 
 # SonarQube
-.scannerwork
+.scannerwork
+
+# backend virtual environments
+**/venv

Dockerfile

@@ -76,26 +76,16 @@ RUN test -n "$TARGETARCH" \
 # The requirements-extras target is for any builds with IMAGE_TYPE=extras. It should not be placed in this target unless every IMAGE_TYPE=extras build will use it
 FROM requirements-core AS requirements-extras
 
-RUN apt-get update && \
-    apt-get install -y --no-install-recommends gpg && \
-    curl https://repo.anaconda.com/pkgs/misc/gpgkeys/anaconda.asc | gpg --dearmor > conda.gpg && \
-    install -o root -g root -m 644 conda.gpg /usr/share/keyrings/conda-archive-keyring.gpg && \
-    gpg --keyring /usr/share/keyrings/conda-archive-keyring.gpg --no-default-keyring --fingerprint 34161F5BF5EB1D4BFBBB8F0A8AEB4F8B29D82806 && \
-    echo "deb [arch=amd64 signed-by=/usr/share/keyrings/conda-archive-keyring.gpg] https://repo.anaconda.com/pkgs/misc/debrepo/conda stable main" > /etc/apt/sources.list.d/conda.list && \
-    echo "deb [arch=amd64 signed-by=/usr/share/keyrings/conda-archive-keyring.gpg] https://repo.anaconda.com/pkgs/misc/debrepo/conda stable main" | tee -a /etc/apt/sources.list.d/conda.list && \
-    apt-get update && \
-    apt-get install -y --no-install-recommends \
-        conda && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
-
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh
 ENV PATH="/root/.cargo/bin:${PATH}"
 
 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
         espeak-ng \
-        espeak && \
+        espeak \
+        python3-dev \
+        python3-venv && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
@@ -246,6 +236,7 @@ ARG FFMPEG
 ARG BUILD_TYPE
 ARG TARGETARCH
 ARG IMAGE_TYPE=extras
+ARG EXTRA_BACKENDS
 ARG MAKEFLAGS
 
 ENV BUILD_TYPE=${BUILD_TYPE}
@@ -257,7 +248,6 @@ ARG CUDA_MAJOR_VERSION=11
 ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
 ENV NVIDIA_REQUIRE_CUDA="cuda>=${CUDA_MAJOR_VERSION}.0"
 ENV NVIDIA_VISIBLE_DEVICES=all
-ENV PIP_CACHE_PURGE=true
 
 # Add FFmpeg
 RUN if [ "${FFMPEG}" = "true" ]; then \
@@ -290,51 +280,58 @@ COPY --from=builder /build/sources/go-piper/piper-phonemize/pi/lib/* /usr/lib/
 # do not let stablediffusion rebuild (requires an older version of absl)
 COPY --from=builder /build/backend-assets/grpc/stablediffusion ./backend-assets/grpc/stablediffusion
 
-## Duplicated from Makefile to avoid having a big layer that's hard to push
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/autogptq \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/bark \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/diffusers \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/vllm \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/mamba \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/sentencetransformers \
+# Change the shell to bash so we can use [[ tests below
+SHELL ["/bin/bash", "-c"]
+# We try to strike a balance between individual layer size (as that affects total push time) and total image size
+# Splitting the backends into more groups with fewer items results in a larger image, but a smaller size for the largest layer
+# Splitting the backends into fewer groups with more items results in a smaller image, but a larger size for the largest layer
+
+RUN if [[ ( "${EXTRA_BACKENDS}" =~ "coqui" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/coqui \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "parler-tts" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/parler-tts \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "diffusers" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/diffusers \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "transformers-musicgen" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/transformers-musicgen \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "exllama1" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/exllama \
     ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/rerankers \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/transformers \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/vall-e-x \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/exllama \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/exllama2 \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/petals \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/transformers-musicgen \
-    ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/parler-tts \
+
+RUN if [[ ( "${EXTRA_BACKENDS}" =~ "vall-e-x" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/vall-e-x \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "petals" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/petals \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "sentencetransformers" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/sentencetransformers \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "exllama2" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/exllama2 \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "transformers" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/transformers \
     ; fi
-RUN if [ "${IMAGE_TYPE}" = "extras" ]; then \
-    make -C backend/python/coqui \
+
+RUN if [[ ( "${EXTRA_BACKENDS}" =~ "vllm" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/vllm \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "autogptq" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/autogptq \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "bark" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/bark \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "rerankers" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/rerankers \
+    ; fi && \
+    if [[ ( "${EXTRA_BACKENDS}" =~ "mamba" || -z "${EXTRA_BACKENDS}" ) && "$IMAGE_TYPE" == "extras" ]]; then \
+        make -C backend/python/mamba \
     ; fi
 
 # Make sure the models directory exists

Makefile

@@ -574,20 +574,20 @@ vllm-protogen-clean:
 # Note: it is duplicated in the Dockerfile
 prepare-extra-conda-environments: protogen-python
 	$(MAKE) -C backend/python/autogptq
-	$(MAKE) -C backend/python/bark
-	$(MAKE) -C backend/python/coqui
-	$(MAKE) -C backend/python/diffusers
-	$(MAKE) -C backend/python/vllm
-	$(MAKE) -C backend/python/mamba
-	$(MAKE) -C backend/python/sentencetransformers
-	$(MAKE) -C backend/python/rerankers
-	$(MAKE) -C backend/python/transformers
-	$(MAKE) -C backend/python/transformers-musicgen
-	$(MAKE) -C backend/python/parler-tts
-	$(MAKE) -C backend/python/vall-e-x
-	$(MAKE) -C backend/python/exllama
-	$(MAKE) -C backend/python/petals
-	$(MAKE) -C backend/python/exllama2
+#	$(MAKE) -C backend/python/bark
+#	$(MAKE) -C backend/python/coqui
+#	$(MAKE) -C backend/python/diffusers
+#	$(MAKE) -C backend/python/vllm
+#	$(MAKE) -C backend/python/mamba
+#	$(MAKE) -C backend/python/sentencetransformers
+#	$(MAKE) -C backend/python/rerankers
+#	$(MAKE) -C backend/python/transformers
+#	$(MAKE) -C backend/python/transformers-musicgen
+#	$(MAKE) -C backend/python/parler-tts
+#	$(MAKE) -C backend/python/vall-e-x
+#	$(MAKE) -C backend/python/exllama
+#	$(MAKE) -C backend/python/petals
+#	$(MAKE) -C backend/python/exllama2
 
 prepare-test-extra: protogen-python
 	$(MAKE) -C backend/python/transformers

backend/cpp/llama/Makefile

@@ -62,8 +62,8 @@ grpc-server: llama.cpp llama.cpp/examples/grpc-server
 	@echo "Building grpc-server with $(BUILD_TYPE) build type and $(CMAKE_ARGS)"
 ifneq (,$(findstring sycl,$(BUILD_TYPE)))
 	bash -c "source $(ONEAPI_VARS); \
-	cd llama.cpp && mkdir -p build && cd build && cmake .. $(CMAKE_ARGS) && cmake --build . --config Release"	
+	cd llama.cpp && mkdir -p build && cd build && cmake .. $(CMAKE_ARGS) && $(MAKE)"
 else
-	cd llama.cpp && mkdir -p build && cd build && cmake .. $(CMAKE_ARGS) && cmake --build . --config Release
+	cd llama.cpp && mkdir -p build && cd build && cmake .. $(CMAKE_ARGS) && $(MAKE)
 endif
 	cp llama.cpp/build/bin/grpc-server .

backend/python/autogptq/Makefile

@@ -1,6 +1,6 @@
 .PHONY: autogptq
 autogptq: protogen
-	$(MAKE) -C ../common-env/transformers
+	bash install.sh
 
 .PHONY: protogen
 protogen: backend_pb2_grpc.py backend_pb2.py
@@ -10,4 +10,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 -rf venv

backend/python/autogptq/install.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+set -ex
+
+BUILD_ISOLATION_FLAG=""
+
+MY_DIR="$(dirname -- "${BASH_SOURCE[0]}")"
+
+uv venv ${MY_DIR}/venv
+source ${MY_DIR}/venv/bin/activate
+
+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
+uv pip install ${BUILD_ISOLATION_FLAG} --requirement ${MY_DIR}/requirements.txt
+
+if [ -f "requirements-${BUILD_TYPE}.txt" ]; then
+    uv pip install ${BUILD_ISOLATION_FLAG}  --requirement ${MY_DIR}/requirements-${BUILD_TYPE}.txt
+fi
+
+if [ -d "/opt/intel" ]; then
+    # Intel GPU: If the directory exists, we assume we are using the Intel image
+    # https://github.com/intel/intel-extension-for-pytorch/issues/538
+    if [ -f "requirements-intel.txt" ]; then
+        uv pip install ${BUILD_ISOLATION_FLAG}  --index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/ --requirement ${MY_DIR}/requirements-intel.txt
+    fi
+fi
+
+if [ "$PIP_CACHE_PURGE" = true ] ; then
+    pip cache purge
+fi

backend/python/autogptq/requirements.txt

@@ -0,0 +1,7 @@
+accelerate
+auto-gptq==0.7.1
+grpcio==1.63.0
+protobuf
+torch
+certifi
+transformers

backend/python/autogptq/run.sh

@@ -1,14 +1,10 @@
 #!/bin/bash
 
 ##
-## A bash script wrapper that runs the autogptq server with conda
+## A bash script wrapper that runs the autogptq 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 )"
-
-python $DIR/autogptq.py $@
+python $MY_DIR/autogptq.py $@

backend/python/autogptq/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

backend/python/bark/Makefile

@@ -1,6 +1,6 @@
 .PHONY: ttsbark
 ttsbark: protogen
-	$(MAKE) -C ../common-env/transformers
+	bash install.sh
 
 .PHONY: run
 run: protogen
@@ -22,4 +22,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 -rf venv

backend/python/bark/install.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+set -ex
+
+BUILD_ISOLATION_FLAG=""
+
+MY_DIR="$(dirname -- "${BASH_SOURCE[0]}")"
+
+uv venv ${MY_DIR}/venv
+source ${MY_DIR}/venv/bin/activate
+
+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
+uv pip install ${BUILD_ISOLATION_FLAG} --requirement ${MY_DIR}/requirements.txt
+
+if [ -f "requirements-${BUILD_TYPE}.txt" ]; then
+    uv pip install ${BUILD_ISOLATION_FLAG}  --requirement ${MY_DIR}/requirements-${BUILD_TYPE}.txt
+fi
+
+if [ -d "/opt/intel" ]; then
+    # Intel GPU: If the directory exists, we assume we are using the Intel image
+    # https://github.com/intel/intel-extension-for-pytorch/issues/538
+    if [ -f "requirements-intel.txt" ]; then
+        uv pip install ${BUILD_ISOLATION_FLAG}  --index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/ --requirement ${MY_DIR}/requirements-intel.txt
+    fi
+fi
+
+if [ "$PIP_CACHE_PURGE" = true ] ; then
+    pip cache purge
+fi

backend/python/bark/requirements.txt

@@ -0,0 +1,6 @@
+accelerate
+bark==0.1.5
+grpcio==1.63.0
+protobuf
+certifi
+transformers