From aef80e0fc1c6d375cbc58e33af6045171c908d4b Mon Sep 17 00:00:00 2001
From: Abhinav Rao <39939017+AetherPrior@users.noreply.github.com>
Date: Sun, 10 Mar 2024 10:50:51 -0400
Subject: [PATCH] Upgrade rustc deps, use environment cargo path

---
 setup_scripts/conda_server.sh | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)
 mode change 100644 => 100755 setup_scripts/conda_server.sh

diff --git a/setup_scripts/conda_server.sh b/setup_scripts/conda_server.sh
old mode 100644
new mode 100755
index 2da4e053e16..d8a6f0b1760
--- a/setup_scripts/conda_server.sh
+++ b/setup_scripts/conda_server.sh
@@ -7,6 +7,11 @@
 DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 DIR="$(dirname "$DIR")"
 
+# define colors for output
+COLOR_GREEN='\033[0;32m'
+COLOR_NC='\033[0m'
+COLOR_RED='\033[0;31m'
+
 # currently can only build in TIR without extensions
 # seems un-important, as it only affects BLOOM/NEOX
 ENV_NAME=tgi-env
@@ -90,7 +95,7 @@ conda activate ${ENV_NAME}
 # # Install dependencies
 mamba install -y -c conda-forge coreutils "gxx<12.0" 
 mamba install -y -c conda-forge curl git tar
-mamba install -y -c conda-forge "rust>=1.65.0"
+mamba install -y -c conda-forge "rust>=1.74"
 mamba install -y -c conda-forge openssh 
 mamba install -y -c "nvidia/label/cuda-11.8.0" cuda-toolkit
 # pin pytorch due to some cuda-issue in pytorch==2.1.0 / something with vllm
@@ -102,9 +107,6 @@ export LD_LIBRARY_PATH=${CONDA_HOME}/envs/${ENV_NAME}/lib:$LD_LIBRARY_PATH
 export PATH=${CONDA_HOME}/envs/${ENV_NAME}/bin:$PATH
 export CUDA_HOME=${CONDA_HOME}/envs/${ENV_NAME}
 
-# add cargo bin
-export PATH=~/.cargo/bin:$PATH
-
 # add protoc
 export PROTOC_ZIP=protoc-21.12-linux-x86_64.zip
 mkdir -p /tmp/protoc
@@ -154,6 +156,8 @@ if [ "$BUILD_VLLM" = true ] ; then
     rm -rf workdir/*
 fi
 
+echo -e "${COLOR_GREEN}DONE INSTALLING VLLM${COLOR_NC}"
+
 # install base package
 cd ${DIR}
 OPENSSL_DIR=${DIR}/.openssl \
@@ -162,6 +166,7 @@ OPENSSL_INCLUDE_DIR=${DIR}/.openssl/include \
 BUILD_EXTENSIONS=$BUILD_EXTENSIONS \
     make install
 
+echo -e "${COLOR_GREEN}DONE INSTALLING BASE PACKAGE${COLOR_NC}"
 
 # install flash attention
 if [ "$BUILD_FLASHATTN" = true ] ; then
@@ -177,6 +182,8 @@ fi
 
 rm -rf workdir
 
+echo -e "${COLOR_GREEN}DONE INSTALLING FLASH ATTENTION${COLOR_NC}"
+
 # override protobuf
 pip install 'protobuf<3.21'