#!/usr/bin/env bash
set -euo pipefail
REPO="${REPO:-/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/huiyingl/sd/speculators}"
WORKSPACE="$(dirname "$REPO")"
export UV_CACHE_DIR="$WORKSPACE/.cache/uv-aarch64-py312"
export UV_LINK_MODE=copy
export UV_HTTP_TIMEOUT=60
export PIP_DISABLE_PIP_VERSION_CHECK=1
export PIP_CACHE_DIR="$WORKSPACE/.cache/pip-aarch64-py312"
export PIP_PROGRESS_BAR=off
if [[ -f "$WORKSPACE/.cache/seed-wheels/seed-constraints.txt" ]]; then
    export PIP_FIND_LINKS="$WORKSPACE/.cache/seed-wheels"
    export PIP_CONSTRAINT="$WORKSPACE/.cache/seed-wheels/seed-constraints.txt"
fi
export VLLM_PLUGINS=""
cd "$REPO"
if ! command -v git >/dev/null; then
    apt-get update -qq
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends git rsync
fi
if ! command -v g++-12 >/dev/null; then
    apt-get update -qq
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends gcc-12 g++-12
fi
update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 120
update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-12 120
git config --global --add safe.directory "$REPO"
UV="${UV:-$(command -v uv)}"
# The official vLLM ARM64 image supplies the matching CUDA/PyTorch binaries.
# Both venvs inherit those immutable base packages and isolate added packages.
/usr/bin/python3 -c 'import torch,vllm; assert vllm.__version__ == "0.27.1"; assert torch.__version__.startswith("2.13.0")'
if [[ ! -x /opt/vllm_venv/bin/python ]]; then
    "$UV" venv --system-site-packages --python /usr/bin/python3 /opt/vllm_venv
fi
if [[ ! -x /opt/speculators_venv/bin/python ]]; then
    "$UV" venv --system-site-packages --python /usr/bin/python3 /opt/speculators_venv
fi
if [[ -f "$WORKSPACE/.cache/seed-wheels/seed-constraints.txt" ]]; then
    for environment in /opt/speculators_venv /opt/vllm_venv; do
        /usr/bin/python3 -m pip --python "$environment/bin/python" install \
            --no-index --no-deps "$WORKSPACE/.cache/seed-wheels/"*.whl
    done
fi
/usr/bin/python3 -m pip --python /opt/speculators_venv/bin/python install \
    -e . 'torch==2.13.0' wandb pytest pytest-asyncio pytest-mock \
    'guidellm==0.7.3' matplotlib scikit-learn scipy
touch logs/ngram_setup/train-env-ready
/usr/bin/python3 -m pip --python /opt/vllm_venv/bin/python install 'vllm==0.27.1' -e . wandb matplotlib
/usr/bin/python3 scripts/cluster/check_environment_dependencies.py
/opt/speculators_venv/bin/python -c 'import torch, speculators; print("TRAIN", torch.__version__, torch.version.cuda, torch.cuda.get_device_name())'
/opt/vllm_venv/bin/python -c 'import torch, vllm; print("SERVE", vllm.__version__, torch.__version__, torch.version.cuda, torch.cuda.get_device_name())'
/usr/bin/python3 -m pip --python /opt/speculators_venv/bin/python freeze > logs/ngram_setup/speculators-requirements.txt
/usr/bin/python3 -m pip --python /opt/vllm_venv/bin/python freeze > logs/ngram_setup/vllm-requirements.txt