#!/bin/bash

# Directorio de trabajo persistente en RunPod
workspace="workspace"

# Modelo
model="neverSleep/Noromaid-20b-v0.1.1"

# Longitud máxima de contexto en tokens
# 4096 = límite de entrenamiento del modelo (máxima calidad)
# No subir más para evitar degradación de respuestas
maxModelLen=4096

# Tipo de datos para KV Cache (Key-Value cache)
# fp8_e4m3 = 8-bit floating point con mejor precisión
# Reduce uso de memoria ~50% vs fp16, ideal para H100
# Permite ~2x más peticiones concurrentes
kvCacheDtype="auto"

# Tamaño de bloques contiguos en el cache
# 32 = bloques grandes, eficiente para GPUs potentes como H100
# Valores menores (16) = más granular pero menos eficiente
# Valores mayores (64+) = no soportados en CUDA
blockSize=32

# IP de bind del servidor
host="0.0.0.0"

# Puerto del servidor HTTP
port=2300

# Tipo de datos para los pesos del modelo
# bfloat16 = 16-bit brain floating point
# Optimizado para hardware Hopper (H100), mejor que fp16
# Menor uso de memoria que fp32 sin perder mucha precisión
dtype="bfloat16"

# Api Key
apiKey="nk-uZ6LwMLN82d5xulWgawal1H3noSpgfj9hh3ZDjoNZp1j"


echo "=== vLLM Setup Script ==="

# Install uv if not present
# if ! command -v uv &> /dev/null; then
echo ">> Installing uv..."
curl -LsSf https://astral.sh/uv/install.sh | sh
source $HOME/.local/bin/env 
# else
#     echo ">> uv is already installed."
# fi

# Verify if tmux is installed
if ! command -v tmux &> /dev/null; then
    echo ">> tmux is not installed. Proceeding with installation..."
    apt update
    apt install tmux -y
else
    echo ">> tmux is already installed."
fi

# Create necessary directories
echo ">> Creating required directories..."
mkdir -p "/$workspace/models"
mkdir -p "/$workspace/.cache/"{vllm,huggingface,torch}

# Navigate to the workspace
echo ">> Navigating to the workspace directory..."
cd "/$workspace" || exit

# Create virtual environment if it doesn't exist
if [ -d ".venv" ]; then
    echo ">> Fixing existing virtual environment..."
    # This will skip the prompt and repair the venv
    UV_VENV_CLEAR=0 uv venv --python 3.12 --seed 2>/dev/null || true
else
    echo ">> Creating new virtual environment..."
    uv venv --python 3.12 --seed
fi

# Activate virtual environment
echo ">> Activating the virtual environment..."
source .venv/bin/activate

# Verify virtual environment activation
if [[ "$VIRTUAL_ENV" != "" ]]; then
    echo ">> Virtual environment activated successfully."
    
    # Check if vllm is already installed
    if ! python -c "import vllm" 2>/dev/null; then
        echo ">> Installing vLLM..."
        uv pip install vllm --torch-backend=auto
    else
        echo ">> vLLM is already installed."
    fi
    
    # Kill existing tmux session if it exists
    tmux kill-session -t vllm 2>/dev/null || true
    
    echo ">> Starting vLLM server in tmux session..."

    # Download chat template if it doesn't exist
    if [ ! -f "/$workspace/alpaca.jinja" ]; then
        wget repo.zarat.dev/alpaca.txt -O "/$workspace/alpaca.jinja"
    fi

    # Start vLLM in tmux with proper environment variables
    tmux new-session -s vllm -d bash -c "
        cd /workspace && \
        source .venv/bin/activate && \
        export VLLM_CACHE_ROOT=/workspace/.cache/vllm && \
        export HF_HOME=/workspace/.cache/huggingface && \
        export TORCH_HOME=/workspace/.cache/torch && \
        export CUDA_VISIBLE_DEVICES=0 && \
        vllm serve $model \
            --gpu-memory-utilization 0.85 \
            --max-model-len $maxModelLen \
            --download-dir /$workspace/models \
            --enable-prefix-caching \
            --kv-cache-dtype $kvCacheDtype \
            --block-size $blockSize \
            --host $host \
            --port $port \
            --dtype $dtype \
            --api-key $apiKey \
            --chat-template /$workspace/alpaca.jinja
    "
    
    echo ">> vLLM server started in tmux session 'vllm'"
    echo ">> Use 'tmux attach -t vllm' to view the session"
    echo ">> Use 'tmux kill-session -t vllm' to stop the server"
    echo ">> Use 'tmux a' to attach to the last session if you don't specify a session name"
    
    # Wait a moment and show session status
    sleep 2
    echo ">> Current tmux sessions:"
    tmux list-sessions 2>/dev/null || echo "No active sessions"

else
    echo "!! Failed to activate virtual environment."
    exit 1
fi

echo "=== Setup completed ==="