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 111cfa97f1ebf07a598386e0fba915695f65d30b.

* move transformers and sentencetransformers to common env

* Share also transformers-musicgen

.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..."
+	$(MAKE) -C ../common-env/transformers
-	@conda env create --name ttsbark --file ttsbark.yml
-	@echo "Virtual environment created."
 
 .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/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/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
+      - rouge==1.0.1
-      - scikit-learn==1.3.1
+      - s3transfer==0.7.0
+      - safetensors==0.3.3
       - scipy==1.11.3
-      - sentence-transformers==2.2.2
+      - six==1.16.0
-      - sentencepiece==0.1.99
       - sympy==1.12
-      - threadpoolctl==3.2.0
+      - tokenizers==0.14.0
-      - tokenizers==0.14.1
       - 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..."
+	$(MAKE) -C ../common-env/transformers
-	@conda env create --name petals --file petals.yml
-	@echo "Virtual environment created."
 
 .PHONY: run
 run:
 	@echo "Running petals..."
 	bash run.sh
 	@echo "petals run."
+
+.PHONY: test
+test:
+	@echo "Testing petals..."
+	bash test.sh
+	@echo "petals tested."

backend/python/petals/run.sh

@@ -9,10 +9,10 @@ export PATH=$PATH:/opt/conda/bin
 # if source is available use it, or use conda
 #
 if [ -f /opt/conda/bin/activate ]; then
-    source activate petals
+    source activate transformers
 else
     eval "$(conda shell.bash hook)"
-    conda activate petals
+    conda activate transformers
 fi
 
 # get the directory where the bash script is located

backend/python/petals/test.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+##
+## A bash script wrapper that runs the transformers server with conda
+
+# Activate conda environment
+source activate transformers
+
+# get the directory where the bash script is located
+DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+
+python -m unittest $DIR/test_petals.py

backend/python/petals/test_petals.py

@@ -0,0 +1,58 @@
+import unittest
+import subprocess
+import time
+import backend_pb2
+import backend_pb2_grpc
+
+import grpc
+
+import unittest
+import subprocess
+import time
+import grpc
+import backend_pb2_grpc
+import backend_pb2
+
+class TestBackendServicer(unittest.TestCase):
+    """
+    TestBackendServicer is the class that tests the gRPC service.
+
+    This class contains methods to test the startup and shutdown of the gRPC service.
+    """
+    def setUp(self):
+        self.service = subprocess.Popen(["python", "backend_petals.py", "--addr", "localhost:50051"])
+        time.sleep(10)
+
+    def tearDown(self) -> None:
+        self.service.terminate()
+        self.service.wait()
+
+    def test_server_startup(self):
+        try:
+            self.setUp()
+            with grpc.insecure_channel("localhost:50051") as channel:
+                stub = backend_pb2_grpc.BackendStub(channel)
+                response = stub.Health(backend_pb2.HealthMessage())
+                self.assertEqual(response.message, b'OK')
+        except Exception as err:
+            print(err)
+            self.fail("Server failed to start")
+        finally:
+            self.tearDown()
+    def test_load_model(self):
+        """
+        This method tests if the model is loaded successfully
+        """
+        try:
+            self.setUp()
+            with grpc.insecure_channel("localhost:50051") as channel:
+                stub = backend_pb2_grpc.BackendStub(channel)
+                response = stub.LoadModel(backend_pb2.ModelOptions(Model="bigscience/bloom-560m"))
+                print(response)
+                self.assertTrue(response.success)
+                self.assertEqual(response.message, "Model loaded successfully")
+        except Exception as err:
+            print(err)
+            self.fail("LoadModel service failed")
+        finally:
+            self.tearDown()

backend/python/sentencetransformers/Makefile

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

backend/python/sentencetransformers/run.sh

@@ -6,7 +6,7 @@
 export PATH=$PATH:/opt/conda/bin
 
 # Activate conda environment
-source activate sentencetransformers
+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/sentencetransformers/sentencetransformers.yml

@@ -1,77 +0,0 @@
-name: sentencetransformers
-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:
-      - certifi==2023.7.22
-      - charset-normalizer==3.3.0
-      - click==8.1.7
-      - filelock==3.12.4
-      - fsspec==2023.9.2
-      - grpcio==1.59.0
-      - huggingface-hub==0.17.3
-      - idna==3.4
-      - install==1.3.5
-      - jinja2==3.1.2
-      - joblib==1.3.2
-      - markupsafe==2.1.3
-      - mpmath==1.3.0
-      - 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
-      - protobuf==4.24.4
-      - pyyaml==6.0.1
-      - regex==2023.10.3
-      - requests==2.31.0
-      - safetensors==0.4.0
-      - scikit-learn==1.3.1
-      - scipy==1.11.3
-      - sentence-transformers==2.2.2
-      - sentencepiece==0.1.99
-      - sympy==1.12
-      - threadpoolctl==3.2.0
-      - tokenizers==0.14.1
-      - torch==2.1.0
-      - torchvision==0.16.0
-      - tqdm==4.66.1
-      - transformers==4.34.0
-      - triton==2.1.0
-      - typing-extensions==4.8.0
-      - urllib3==2.0.6
-prefix: /opt/conda/envs/sentencetransformers

backend/python/sentencetransformers/test.sh

@@ -3,7 +3,7 @@
 ## A bash script wrapper that runs the sentencetransformers server with conda
 
 # Activate conda environment
-source activate sentencetransformers
+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/transformers-musicgen/Makefile

@@ -1,15 +1,7 @@
 
-TRANSFORMERS_MUSICGEN_CONDA_PATH = "transformers-musicgen.yml"
-
-ifeq ($(BUILD_TYPE), cublas)
-	TRANSFORMERS_MUSICGEN_CONDA_PATH = "transformers-musicgen-nvidia.yml"
-endif
-
 .PHONY: transformers-musicgen
 transformers-musicgen:
-	@echo "Creating virtual environment..."
+	$(MAKE) -C ../common-env/transformers
-	@conda env create --name transformers-musicgen --file $(TRANSFORMERS_MUSICGEN_CONDA_PATH)
-	@echo "Virtual environment created."
 
 .PHONY: run
 run:

backend/python/transformers-musicgen/test.sh

@@ -3,7 +3,7 @@
 ## A bash script wrapper that runs the transformers server with conda
 
 # Activate conda environment
-source activate transformers-musicgen
+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/transformers-musicgen/transformers-musicgen-nvidia.yml

@@ -1,71 +0,0 @@
-name: transformers-musicgen
-channels:
-  - defaults
-dependencies:
-  - bzip2=1.0.8
-  - ca-certificates=2023.08.22
-  - libffi=3.4.4
-  - libuuid=1.41.5
-  - ncurses=6.4
-  - openssl=3.0.11
-  - pip=23.2.1
-  - python=3.11.5
-  - readline=8.2
-  - setuptools=68.0.0
-  - sqlite=3.41.2
-  - tk=8.6.12
-  - tzdata=2023c
-  - wheel=0.41.2
-  - xz=5.4.2
-  - zlib=1.2.13
-  - pip:
-      - certifi==2023.7.22
-      - charset-normalizer==3.3.0
-      - click==8.1.7
-      - filelock==3.12.4
-      - fsspec==2023.9.2
-      - grpcio==1.59.0
-      - huggingface-hub==0.17.3
-      - idna==3.4
-      - install==1.3.5
-      - jinja2==3.1.2
-      - joblib==1.3.2
-      - markupsafe==2.1.3
-      - mpmath==1.3.0
-      - 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
-      - protobuf==4.24.4
-      - pyyaml==6.0.1
-      - regex==2023.10.3
-      - requests==2.31.0
-      - safetensors==0.4.0
-      - scikit-learn==1.3.1
-      - scipy==1.11.3
-      - sentence-transformers==2.2.2
-      - sentencepiece==0.1.99
-      - sympy==1.12
-      - threadpoolctl==3.2.0
-      - tokenizers==0.14.1
-      - torch==2.1.0
-      - torchvision==0.16.0
-      - tqdm==4.66.1
-      - transformers==4.34.0
-      - triton==2.1.0
-      - typing-extensions==4.8.0
-      - urllib3==2.0.6
-prefix: /opt/conda/envs/transformers-musicgen

backend/python/transformers-musicgen/transformers-musicgen.yml

@@ -1,58 +0,0 @@
-name: transformers-musicgen
-channels:
-  - defaults
-dependencies:
-  - bzip2=1.0.8
-  - ca-certificates=2023.08.22
-  - libffi=3.4.4
-  - libuuid=1.41.5
-  - ncurses=6.4
-  - openssl=3.0.11
-  - pip=23.2.1
-  - python=3.11.5
-  - readline=8.2
-  - setuptools=68.0.0
-  - sqlite=3.41.2
-  - tk=8.6.12
-  - tzdata=2023c
-  - wheel=0.41.2
-  - xz=5.4.2
-  - zlib=1.2.13
-  - pip:
-      - certifi==2023.7.22
-      - charset-normalizer==3.3.0
-      - click==8.1.7
-      - filelock==3.12.4
-      - fsspec==2023.9.2
-      - grpcio==1.59.0
-      - huggingface-hub==0.17.3
-      - idna==3.4
-      - install==1.3.5
-      - jinja2==3.1.2
-      - joblib==1.3.2
-      - markupsafe==2.1.3
-      - mpmath==1.3.0
-      - networkx==3.1
-      - nltk==3.8.1
-      - numpy==1.26.0
-      - packaging==23.2
-      - pillow==10.0.1
-      - protobuf==4.24.4
-      - pyyaml==6.0.1
-      - regex==2023.10.3
-      - requests==2.31.0
-      - safetensors==0.4.0
-      - scikit-learn==1.3.1
-      - scipy==1.11.3
-      - sentence-transformers==2.2.2
-      - sentencepiece==0.1.99
-      - sympy==1.12
-      - threadpoolctl==3.2.0
-      - tokenizers==0.14.1
-      - torch==2.1.0
-      - torchvision==0.16.0
-      - tqdm==4.66.1
-      - transformers==4.34.0
-      - typing-extensions==4.8.0
-      - urllib3==2.0.6
-prefix: /opt/conda/envs/transformers-musicgen

backend/python/transformers/Makefile

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

backend/python/vllm/test_backend_vllm.py

@@ -66,7 +66,7 @@ class TestBackendServicer(unittest.TestCase):
                 stub = backend_pb2_grpc.BackendStub(channel)
                 response = stub.LoadModel(backend_pb2.ModelOptions(Model="facebook/opt-125m"))
                 self.assertTrue(response.success)
-                req = backend_pb2.PredictOptions(prompt="The capital of France is")
+                req = backend_pb2.PredictOptions(Prompt="The capital of France is")
                 resp = stub.Predict(req)
                 self.assertIsNotNone(resp.message)
         except Exception as err: