feat(venv): shared env (#3195)

* feat(venv): allow to share veenvs

Signed-off-by: Ettore Di Giacinto <[email protected]>

* fix(vllm): add back flash-attn

Signed-off-by: Ettore Di Giacinto <[email protected]>

---------

Signed-off-by: Ettore Di Giacinto <[email protected]>

Author: mudler

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

backend/python/vllm/requirements-cublas11-after.txt

@@ -0,0 +1 @@
+flash-attn

backend/python/vllm/requirements-cublas12-after.txt

@@ -0,0 +1 @@
+flash-attn