From 66cf38b0b36f46d915a79b0e8d2ceae90614f6bb Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@users.noreply.github.com>
Date: Wed, 7 Aug 2024 19:45:14 +0200
Subject: [PATCH] feat(venv): shared env (#3195)

* feat(venv): allow to share veenvs

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(vllm): add back flash-attn

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
---
 backend/python/common/libbackend.sh           | 40 ++++++++++++++-----
 .../vllm/requirements-cublas11-after.txt      |  1 +
 .../vllm/requirements-cublas12-after.txt      |  1 +
 3 files changed, 31 insertions(+), 11 deletions(-)
 create mode 100644 backend/python/vllm/requirements-cublas11-after.txt
 create mode 100644 backend/python/vllm/requirements-cublas12-after.txt

diff --git a/backend/python/common/libbackend.sh b/backend/python/common/libbackend.sh
index 7287fb95..934b1fd3 100644
--- a/backend/python/common/libbackend.sh
+++ b/backend/python/common/libbackend.sh
@@ -18,10 +18,23 @@
 # source $(dirname $0)/../common/libbackend.sh
 #
 function init() {
+    # Name of the backend (directory name)
     BACKEND_NAME=${PWD##*/}
+
+    # Path where all backends files are
     MY_DIR=$(realpath `dirname $0`)
+
+    # Build type
     BUILD_PROFILE=$(getBuildProfile)
 
+    # Environment directory
+    EDIR=${MY_DIR}
+
+    # Allow to specify a custom env dir for shared environments
+    if [ "x${ENV_DIR}" != "x" ]; then
+        EDIR=${ENV_DIR}
+    fi
+
     # If a backend has defined a list of valid build profiles...
     if [ ! -z "${LIMIT_TARGETS}" ]; then
         isValidTarget=$(checkTargets ${LIMIT_TARGETS})
@@ -74,13 +87,14 @@ function getBuildProfile() {
 # This function is idempotent, so you can call it as many times as you want and it will
 # always result in an activated virtual environment
 function ensureVenv() {
-    if [ ! -d "${MY_DIR}/venv" ]; then
-        uv venv ${MY_DIR}/venv
+    if [ ! -d "${EDIR}/venv" ]; then
+        uv venv ${EDIR}/venv
         echo "virtualenv created"
     fi
-    
-    if [ "x${VIRTUAL_ENV}" != "x${MY_DIR}/venv" ]; then
-        source ${MY_DIR}/venv/bin/activate
+
+    # Source if we are not already in a Virtual env
+    if [ "x${VIRTUAL_ENV}" != "x${EDIR}/venv" ]; then
+        source ${EDIR}/venv/bin/activate
         echo "virtualenv activated"
     fi
 
@@ -113,21 +127,25 @@ function installRequirements() {
 
     # These are the requirements files we will attempt to install, in order
     declare -a requirementFiles=(
-        "${MY_DIR}/requirements-install.txt"
-        "${MY_DIR}/requirements.txt"
-        "${MY_DIR}/requirements-${BUILD_TYPE}.txt"
+        "${EDIR}/requirements-install.txt"
+        "${EDIR}/requirements.txt"
+        "${EDIR}/requirements-${BUILD_TYPE}.txt"
     )
 
     if [ "x${BUILD_TYPE}" != "x${BUILD_PROFILE}" ]; then
-        requirementFiles+=("${MY_DIR}/requirements-${BUILD_PROFILE}.txt")
+        requirementFiles+=("${EDIR}/requirements-${BUILD_PROFILE}.txt")
     fi
 
     # if BUILD_TYPE is empty, we are a CPU build, so we should try to install the CPU requirements
     if [ "x${BUILD_TYPE}" == "x" ]; then
-        requirementFiles+=("${MY_DIR}/requirements-cpu.txt")
+        requirementFiles+=("${EDIR}/requirements-cpu.txt")
     fi
 
-    requirementFiles+=("${MY_DIR}/requirements-after.txt")
+    requirementFiles+=("${EDIR}/requirements-after.txt")
+
+    if [ "x${BUILD_TYPE}" != "x${BUILD_PROFILE}" ]; then
+        requirementFiles+=("${EDIR}/requirements-${BUILD_PROFILE}-after.txt")
+    fi
 
     for reqFile in ${requirementFiles[@]}; do
         if [ -f ${reqFile} ]; then
diff --git a/backend/python/vllm/requirements-cublas11-after.txt b/backend/python/vllm/requirements-cublas11-after.txt
new file mode 100644
index 00000000..7bfe8efe
--- /dev/null
+++ b/backend/python/vllm/requirements-cublas11-after.txt
@@ -0,0 +1 @@
+flash-attn
\ No newline at end of file
diff --git a/backend/python/vllm/requirements-cublas12-after.txt b/backend/python/vllm/requirements-cublas12-after.txt
new file mode 100644
index 00000000..7bfe8efe
--- /dev/null
+++ b/backend/python/vllm/requirements-cublas12-after.txt
@@ -0,0 +1 @@
+flash-attn
\ No newline at end of file