setup.sh

#!/bin/bash
# =============================================================================
# Lambda Labs CUDA Development Environment Setup (Native)
# =============================================================================
# No Docker required - Lambda has CUDA pre-installed.
# This script sets up environment variables and installs useful tools.
# =============================================================================

set -e

echo "=============================================="
echo "  Lambda Labs CUDA Environment Setup"
echo "=============================================="
echo ""

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# -----------------------------------------------------------------------------
# Step 1: Detect CUDA installation
# -----------------------------------------------------------------------------
echo "[1/4] Detecting CUDA installation..."

if command -v nvcc &> /dev/null; then
    NVCC_VERSION=$(nvcc --version | grep "release" | awk '{print $6}' | cut -d',' -f1)
    echo "      Found CUDA $NVCC_VERSION"
else
    echo "      ERROR: nvcc not found!"
    exit 1
fi

# Find CUDA home
if [ -d "/usr/local/cuda" ]; then
    CUDA_HOME="/usr/local/cuda"
elif [ -d "/usr/lib/cuda" ]; then
    CUDA_HOME="/usr/lib/cuda"
else
    CUDA_HOME=$(dirname $(dirname $(which nvcc)))
fi
echo "      CUDA_HOME: $CUDA_HOME"

# -----------------------------------------------------------------------------
# Step 2: Create shell config additions
# -----------------------------------------------------------------------------
echo "[2/4] Setting up environment variables..."

ENV_FILE="$SCRIPT_DIR/.cuda_env"
cat > "$ENV_FILE" << EOF
# CUDA Environment (generated by setup.sh)
export CUDA_HOME=$CUDA_HOME
export PATH=\$CUDA_HOME/bin:\$PATH
export LD_LIBRARY_PATH=\$CUDA_HOME/lib64:\$LD_LIBRARY_PATH

# Convenience alias
alias cdw='cd $SCRIPT_DIR/workspace'
EOF

# Add to .bashrc if not already there
if ! grep -q "lambda-dev-env/.cuda_env" ~/.bashrc 2>/dev/null; then
    echo "" >> ~/.bashrc
    echo "# CUDA development environment" >> ~/.bashrc
    echo "[ -f $ENV_FILE ] && source $ENV_FILE" >> ~/.bashrc
    echo "      Added to ~/.bashrc"
else
    echo "      Already in ~/.bashrc"
fi

# Source it now
source "$ENV_FILE"

# -----------------------------------------------------------------------------
# Step 3: Install useful packages (if not present)
# -----------------------------------------------------------------------------
echo "[3/4] Checking development tools..."

MISSING_PKGS=""
for pkg in build-essential cmake git python3-pip; do
    if ! dpkg -l | grep -q "^ii  $pkg "; then
        MISSING_PKGS="$MISSING_PKGS $pkg"
    fi
done

if [ -n "$MISSING_PKGS" ]; then
    echo "      Installing:$MISSING_PKGS"
    sudo apt-get update -qq
    sudo apt-get install -y -qq $MISSING_PKGS
else
    echo "      All tools present"
fi

# -----------------------------------------------------------------------------
# Step 4: Create directories and test
# -----------------------------------------------------------------------------
echo "[4/4] Creating workspace..."

mkdir -p workspace scripts
echo "      Created ./workspace and ./scripts"

# Quick GPU test
echo ""
echo "=============================================="
echo "  Setup Complete!"
echo "=============================================="
echo ""
nvidia-smi --query-gpu=name,compute_cap,memory.total --format=csv
echo ""
echo "Usage:"
echo "  source ~/.bashrc     # Load environment (or open new terminal)"
echo "  cd workspace         # Your code goes here"
echo "  nvcc -O3 -arch=sm_90 -o prog prog.cu"
echo "  ./validate.sh        # Test the environment"
echo ""