optimized setup.sh

setup.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-echo "Starting training..."
+echo "Starting building the environment..."
 uv sync --offline
 echo "Activated virtual environment"
 source .venv/bin/activate
@@ -10,20 +10,27 @@ cp debug/modeling_llama.py .venv/lib/python3.10/site-packages/transformers/model
 export LLM_PATH=meta-llama/Llama-2-7b-hf
 echo "LLM_PATH: $LLM_PATH"
 
-echo "Setting up CUDA..."
+if [ -n "$LD_LIBRARY_PATH" ]; then
+    echo "LD_LIBRARY_PATH is already set. Skipping CUDA setup..."
+else
+    echo "Setting up CUDA..."
+fi
 cuda_dirs=(/usr/local/cuda-*/)
 if [ ${#cuda_dirs[@]} -eq 0 ]; then
+
     echo "Can not find any cuda folder in /usr/local"
     exit 1
-fi
 
-# Sort the cuda folders and get the latest one
-latest_cuda=$(for dir in "${cuda_dirs[@]}"; do echo "${dir%/}"; done | sort -V | tail -n 1)
+    # Sort the cuda folders and get the latest one
+    latest_cuda=$(for dir in "${cuda_dirs[@]}"; do echo "${dir%/}"; done | sort -V | tail -n 1)
 
-echo "The latest CUDA folder is: ${latest_cuda}"
+    echo "The latest CUDA folder is: ${latest_cuda}"
 
-export CUDA_HOME="${latest_cuda}"
-export PATH="${latest_cuda}/bin:${PATH}"
-export LD_LIBRARY_PATH="${latest_cuda}/lib64/stubs:${latest_cuda}/lib64:${latest_cuda}/cudnn/lib:${LD_LIBRARY_PATH}"
+    export CUDA_HOME="${latest_cuda}"
+    export PATH="${latest_cuda}/bin:${PATH}"
+    export LD_LIBRARY_PATH="${latest_cuda}/lib64/stubs:${latest_cuda}/lib64:${latest_cuda}/cudnn/lib:${LD_LIBRARY_PATH}"
 
+fi
+
+echo "Starting training..."
 sh train_movielens.sh