#!/usr/bin/env bash
# Core training/eval deps for the Phase-0 go/no-go (NO flash-attn build, NO vllm).
# Run on the GPU host from the repo root. vLLM is installed separately afterward
# so a vLLM dependency conflict can't invalidate the Phase-0 result.
set -euo pipefail

export HF_HOME=${HF_HOME:-/workspace/hf-cache}
export HF_HUB_CACHE=${HF_HUB_CACHE:-/workspace/hf-cache/hub}
export PIP_ROOT_USER_ACTION=ignore
# Image ships torch in the system (PEP-668 externally-managed) Python; allow installs there.
export PIP_BREAK_SYSTEM_PACKAGES=1
mkdir -p "$HF_HOME" /workspace/reports /workspace/tmp

python -m pip install --upgrade pip setuptools wheel packaging >/dev/null

echo "== installing core deps =="
python -m pip install -U \
  "transformers>=5.12.1" "accelerate>=1.14.0" "datasets>=5.0.0" "peft>=0.19.1" \
  "huggingface_hub>=1.1.0" "safetensors>=0.6.0" "tokenizers>=0.22.0" \
  "sentencepiece>=0.2.0" "protobuf>=6.0.0" "einops>=0.8.0" "pyyaml>=6.0.0"

# Hybrid (qwen3_5 / Gated-DeltaNet) kernels. Triton-based FLA + causal-conv1d are
# usually wheels; tolerate failure and fall back to sdpa for the smoke.
echo "== installing hybrid-arch kernels (best effort) =="
python -m pip install -U "causal-conv1d>=1.6.2.post1" || echo "WARN: causal-conv1d failed"
python -m pip install -U "flash-linear-attention>=0.5.1" || echo "WARN: flash-linear-attention failed"

echo "== versions =="
python - <<'PY'
import importlib
for m in ["torch","transformers","accelerate","datasets","peft","tokenizers"]:
    try:
        mod=importlib.import_module(m); print(m, getattr(mod,"__version__","?"))
    except Exception as e:
        print(m, "MISSING", repr(e))
import torch
print("cuda_available", torch.cuda.is_available(), "devices", torch.cuda.device_count())
if torch.cuda.is_available():
    print("gpu0", torch.cuda.get_device_name(0))
PY
echo "core setup done."