feat(conda): share envs with transformer-based backends (#1465)

* feat(conda): share env between diffusers and bark

* Detect if env already exists

* share diffusers and petals

* tests: add petals

* Use smaller model for tests with petals

* test only model load on petals

* tests(petals): run only load model tests

* Revert "test only model load on petals"

This reverts commit 111cfa9.

* move transformers and sentencetransformers to common env

* Share also transformers-musicgen

Author: mudler

.github/workflows/test-extra.yml

@@ -133,6 +133,37 @@ jobs:
 
 
 
+  tests-petals:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Clone
+        uses: actions/checkout@v4
+        with: 
+          submodules: true
+      - name: Dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install build-essential ffmpeg
+          curl https://repo.anaconda.com/pkgs/misc/gpgkeys/anaconda.asc | gpg --dearmor > conda.gpg && \
+             sudo 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 && \
+             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" > /etc/apt/sources.list.d/conda.list' && \
+             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
+          sudo apt-get install -y ca-certificates cmake curl patch
+          sudo apt-get install -y libopencv-dev && sudo ln -s /usr/include/opencv4/opencv2 /usr/include/opencv2
+          
+          sudo rm -rfv /usr/bin/conda || true
+
+      - name: Test petals
+        run: |
+           export PATH=$PATH:/opt/conda/bin
+           make -C backend/python/petals
+           make -C backend/python/petals test
+
+           
+
   tests-bark:
     runs-on: ubuntu-latest
     steps:

backend/python/bark/Makefile

@@ -1,8 +1,6 @@
 .PHONY: ttsbark
 ttsbark:
-	@echo "Creating virtual environment..."
-	@conda env create --name ttsbark --file ttsbark.yml
-	@echo "Virtual environment created."
+	$(MAKE) -C ../common-env/transformers
 
 .PHONY: run
 run:

backend/python/bark/run.sh

@@ -6,7 +6,7 @@
 export PATH=$PATH:/opt/conda/bin
 
 # Activate conda environment
-source activate ttsbark
+source activate transformers
 
 # get the directory where the bash script is located
 DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"

backend/python/bark/test.sh

@@ -3,7 +3,7 @@
 ## A bash script wrapper that runs the bark server with conda
 
 # Activate conda environment
-source activate ttsbark
+source activate transformers
 
 # get the directory where the bash script is located
 DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"

backend/python/common-env/transformers/Makefile

@@ -0,0 +1,10 @@
+CONDA_ENV_PATH = "transformers.yml"
+
+ifeq ($(BUILD_TYPE), cublas)
+	CONDA_ENV_PATH = "transformers-nvidia.yml"
+endif
+
+.PHONY: transformers
+transformers:
+	@echo "Installing $(CONDA_ENV_PATH)..."
+	bash install.sh $(CONDA_ENV_PATH)

backend/python/common-env/transformers/install.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+set -ex
+
+# Check if environment exist
+conda_env_exists(){
+    ! conda list --name "${@}" >/dev/null 2>/dev/null
+}
+
+if conda_env_exists "transformers" ; then
+    echo "Creating virtual environment..."
+    conda env create --name transformers --file $1
+    echo "Virtual environment created."
+else 
+    echo "Virtual environment already exists."
+fi

backend/python/bark/ttsbark.yml renamed to backend/python/common-env/transformers/transformers-nvidia.yml

@@ -1,4 +1,4 @@
-name: bark
+name: transformers
 channels:
   - defaults
 dependencies:
@@ -35,6 +35,8 @@ dependencies:
       - certifi==2023.7.22
       - charset-normalizer==3.3.0
       - datasets==2.14.5
+      - sentence-transformers==2.2.2
+      - sentencepiece==0.1.99
       - dill==0.3.7
       - einops==0.7.0
       - encodec==0.1.1
@@ -68,6 +70,7 @@ dependencies:
       - packaging==23.2
       - pandas==2.1.1
       - peft==0.5.0
+      - git+https://github.com/bigscience-workshop/petals
       - protobuf==4.24.4
       - psutil==5.9.5
       - pyarrow==13.0.0
@@ -93,4 +96,4 @@ dependencies:
       - urllib3==1.26.17
       - xxhash==3.4.1
       - yarl==1.9.2
-prefix: /opt/conda/envs/bark
+prefix: /opt/conda/envs/transformers

backend/python/transformers/transformers.yml renamed to backend/python/common-env/transformers/transformers.yml

@@ -20,58 +20,68 @@ dependencies:
   - 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:
+      - accelerate==0.23.0
+      - aiohttp==3.8.5
+      - aiosignal==1.3.1
+      - async-timeout==4.0.3
+      - attrs==23.1.0
+      - bark==0.1.5
+      - boto3==1.28.61
+      - botocore==1.31.61
       - certifi==2023.7.22
       - charset-normalizer==3.3.0
-      - click==8.1.7
+      - datasets==2.14.5
+      - sentence-transformers==2.2.2
+      - sentencepiece==0.1.99
+      - dill==0.3.7
+      - einops==0.7.0
+      - encodec==0.1.1
       - filelock==3.12.4
-      - fsspec==2023.9.2
+      - frozenlist==1.4.0
+      - fsspec==2023.6.0
+      - funcy==2.0
       - grpcio==1.59.0
-      - huggingface-hub==0.17.3
+      - huggingface-hub==0.16.4
       - idna==3.4
-      - install==1.3.5
       - jinja2==3.1.2
-      - joblib==1.3.2
+      - jmespath==1.0.1
       - markupsafe==2.1.3
       - mpmath==1.3.0
+      - multidict==6.0.4
+      - multiprocess==0.70.15
       - networkx==3.1
-      - nltk==3.8.1
       - numpy==1.26.0
-      - 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
       - packaging==23.2
-      - pillow==10.0.1
+      - pandas==2.1.1
+      - peft==0.5.0
+      - git+https://github.com/bigscience-workshop/petals
       - protobuf==4.24.4
+      - psutil==5.9.5
+      - pyarrow==13.0.0
+      - python-dateutil==2.8.2
+      - pytz==2023.3.post1
       - pyyaml==6.0.1
       - regex==2023.10.3
       - requests==2.31.0
-      - safetensors==0.4.0
-      - scikit-learn==1.3.1
+      - rouge==1.0.1
+      - s3transfer==0.7.0
+      - safetensors==0.3.3
       - scipy==1.11.3
-      - sentence-transformers==2.2.2
-      - sentencepiece==0.1.99
+      - six==1.16.0
       - sympy==1.12
-      - threadpoolctl==3.2.0
-      - tokenizers==0.14.1
+      - tokenizers==0.14.0
       - torch==2.1.0
-      - torchvision==0.16.0
+      - torchaudio==2.1.0
       - tqdm==4.66.1
       - transformers==4.34.0
       - triton==2.1.0
       - typing-extensions==4.8.0
-      - urllib3==2.0.6
+      - tzdata==2023.3
+      - urllib3==1.26.17
+      - xxhash==3.4.1
+      - yarl==1.9.2
 prefix: /opt/conda/envs/transformers

backend/python/diffusers/diffusers.yml

@@ -70,4 +70,4 @@ dependencies:
       - typing-extensions==4.8.0
       - urllib3==2.0.6
       - zipp==3.17.0
-prefix: /opt/conda/envs/diffusers
+prefix: /opt/conda/envs/diffusers

backend/python/petals/Makefile

@@ -1,11 +1,15 @@
 .PHONY: petals
 petals:
-	@echo "Creating virtual environment..."
-	@conda env create --name petals --file petals.yml
-	@echo "Virtual environment created."
+	$(MAKE) -C ../common-env/transformers
 
 .PHONY: run
 run:
 	@echo "Running petals..."
 	bash run.sh
 	@echo "petals run."
+
+.PHONY: test
+test:
+	@echo "Testing petals..."
+	bash test.sh
+	@echo "petals tested."