cat <<'__RUN_SBATCH__' > run_deepseek.sbatch
#!/bin/bash
#SBATCH --job-name=qwen3_8gpu
#SBATCH --partition=P01
#SBATCH --nodes=2
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=64
#SBATCH --time=04:00:00
#SBATCH --output=/home/Competition2025/P10/shareP10/l/logs/%x-%j.out
#SBATCH --error=/home/Competition2025/P10/shareP10/l/logs/%x-%j.err
#SBATCH --export=OPENAI_API_KEY="openai_api_keyをここに"

#=============== 前処理 / 片付け ===============
mkdir -p /home/Competition2025/P10/shareP10/l/logs

cleanup () {
  echo "[CLEANUP] stopping background processes..."
  [[ -n "${pid_vllm:-}"  ]] && kill "${pid_vllm}"  2>/dev/null || true
  [[ -n "${pid_nvsmi:-}" ]] && kill "${pid_nvsmi}" 2>/dev/null || true
  srun -N ${SLURM_NNODES:-1} -n ${SLURM_NNODES:-1} /bin/bash -lc '
    source /home/appli/miniconda3/24.7.1-py311/etc/profile.d/conda.sh
    conda activate "$HOME/conda_env_llm" || true
    ray stop --force || true
    pkill -9 ray || true
    pkill -9 gcs_server || true
    rm -rf /tmp/ray{,_tmp_*} || true
    rm -rf /dev/shm/ray* || true
  ' >/dev/null 2>&1 || true
}
trap cleanup EXIT ERR SIGINT SIGTERM SIGQUIT

#=============== Hugging Face 認証とキャッシュ ===============
export HF_TOKEN="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"
export HF_HOME=${SLURM_TMPDIR:-$HOME}/.hf_cache
export TRANSFORMERS_CACHE=$HF_HOME
export HUGGINGFACE_HUB_TOKEN=$HF_TOKEN
mkdir -p "$HF_HOME"
echo "HF cache dir : $HF_HOME"

#=============== GPU 監視 ======================
nvidia-smi -i 0,1,2,3,4,5,6,7 -l 3 > nvidia-smi.log &
pid_nvsmi=$!

#=============== モジュールと Conda 環境 ===============
module reset
module load nccl/2.24.3
module load hpcx/2.18.1-gcc-cuda12/hpcx-mt
module load cudnn/9.4.0
module load miniconda/24.7.1-py311
source /home/appli/miniconda3/24.7.1-py311/etc/profile.d/conda.sh

export DEEPSEEK_ENV="$HOME/conda_env_llm"
[ -d "$DEEPSEEK_ENV" ] && conda env remove --prefix "$DEEPSEEK_ENV" -y && rm -rf "$DEEPSEEK_ENV"
conda create --prefix "$DEEPSEEK_ENV" python=3.11 -y
conda activate "$DEEPSEEK_ENV"

conda install cuda-toolkit=12.6 -c conda-forge -y
conda install gcc_linux-64 gxx_linux-64 -y
pip install --upgrade pip wheel cmake ninja
pip install huggingface-hub
pip install hydra-core
pip cache purge && rm -rf ~/.cache/pip/*
pip install vllm==0.9.2
pip install --index-url https://download.pytorch.org/whl/cu126 torch==2.7.0 torchvision==0.22.0+cu126 torchaudio==2.7.0
pip install "ray[cgraph,data,train,tune,serve]==2.43.0"
pip install transformers==4.53.2 accelerate==1.2.1 sentencepiece==0.2.0 \
           protobuf==4.25.5 pydantic==2.10.3 aiohttp==3.11.10 \
           numpy==2.2.6 scipy==1.14.1
pip install datasets requests tqdm psutil gputil
pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.2/flash_attn-2.8.2+cu12torch2.7cxx11abiTRUE-cp311-cp311-linux_x86_64.whl
pip install flashinfer-python

cat >> ~/.bashrc << 'EOF'
# DeepSeek環境用ライブラリパス
export LD_LIBRARY_PATH="/home/Competition2025/P10/P10U017/.local/lib/python3.11/site-packages/nvidia/cusparselt/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="/home/Competition2025/P10/P10U017/.local/lib/python3.11/site-packages/nvidia/cublas/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="/home/Competition2025/P10/P10U017/.local/lib/python3.11/site-packages/nvidia/cusparse/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="/home/Competition2025/P10/P10U017/.local/lib/python3.11/site-packages/nvidia/curand/lib:$LD_LIBRARY_PATH"
EOF
source ~/.bashrc

#=============== Ray の前処理 ===============
srun -N ${SLURM_NNODES:-1} -n ${SLURM_NNODES:-1} /bin/bash -lc '
  ray stop --force || true
  pkill -9 ray || true
  pkill -9 gcs_server || true
  rm -rf /tmp/ray{,_tmp_*}
  df -h /dev/shm
  ls -la /dev/shm/ || true
  rm -rf /dev/shm/ray*
'

# 割り当てノード一覧
HOSTS=($(scontrol show hostnames "$SLURM_JOB_NODELIST"))
HEAD_NODE=${HOSTS[0]}
echo "🖥️ 割り当てノード: ${HOSTS[*]}"
echo "🧠 ヘッドノード: $HEAD_NODE"

# NIC 固定でIP取得
HEAD_IF=enp25s0np0
HEAD_IP=$(ip -o -4 addr show dev "$HEAD_IF" | awk '{print $4}' | cut -d/ -f1)
if [[ -z "$HEAD_IP" ]]; then HEAD_IP=$(hostname -i | awk '{print $1}'); fi
echo "🖥️ ヘッドノードIP: $HEAD_IP"

# ヘッドノードでRay起動
export VLLM_HOST_IP=$HEAD_IP
ray start --head --node-ip-address=$HEAD_IP --port=6379 \
  --dashboard-host=0.0.0.0 --dashboard-port=8265
export RAY_ADDRESS=$HEAD_IP:6379
echo "✅ Rayヘッドノード起動完了: $RAY_ADDRESS"

# 残りノードでワーカー起動（--block で存続）
if [[ ${#HOSTS[@]} -gt 1 ]]; then
  for node in "${HOSTS[@]:1}"; do
    echo "➡️  ワーカーノード起動: $node"
    srun -N1 -n1 -w "$node" /bin/bash -lc '
      source /home/appli/miniconda3/24.7.1-py311/etc/profile.d/conda.sh
      conda activate "'"$DEEPSEEK_ENV"'"
      IF=enp25s0np0
      WORKER_IP=$(ip -o -4 addr show dev "$IF" | awk "{print \$4}" | cut -d/ -f1)
      if [[ -z "$WORKER_IP" ]]; then WORKER_IP=$(hostname -i | awk "{print \$1}"); fi
      export VLLM_HOST_IP=$WORKER_IP
      ray start --address="'"$RAY_ADDRESS"'" --node-ip-address=$WORKER_IP --block
    ' &
  done
  # Nodes が揃うまでヘッドでポーリング
  for i in {1..60}; do
    ok=$(srun -N1 -n1 -w "$HEAD_NODE" /bin/bash -lc '
      source /home/appli/miniconda3/24.7.1-py311/etc/profile.d/conda.sh
      conda activate "'"$DEEPSEEK_ENV"'"
      ray status | awk "/Nodes:/ {print \$2}"
    ' 2>/dev/null | tail -n1)
    if [[ "$ok" == "${SLURM_NNODES}" ]]; then
      echo "✅ すべてのワーカーノードが参加しました"
      break
    fi
    sleep 2
  done
fi

#=============== vLLM サーバ起動（バックグラウンド） ===============
cat > deepseek_optimized_server_v3.sh << 'EOF'
#!/bin/bash
export MODEL_PATH="/home/Competition2025/P10/shareP10/llm_project/models/deepseek/DeepSeek-R1-0528"
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export NCCL_SOCKET_IFNAME=enp25s0np0
export NCCL_IB_DISABLE=0
export NCCL_DEBUG=WARN
export RAY_ADDRESS=${RAY_ADDRESS:-"auto"}
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
export CUDA_LAUNCH_BLOCKING=1

echo "========================================"
echo "DeepSeek-R1-0528 vLLM Server (Performance Optimized v3)"
echo "========================================"
echo "Model         : $MODEL_PATH"
echo "Configuration : 2 × 8 GPU (TP=16)"
echo "Performance   : ~1000 token/sec (validated)"
echo "Memory Profile: 90% GPU utilization, KV cache 90-95%"
echo "vLLM          : $(python -c 'import vllm; print(vllm.__version__)')"
echo "VLLM_HOST_IP  : ${VLLM_HOST_IP:-未設定}"

echo "========================================"
echo "🔍 GPU Memory Check:"
nvidia-smi --query-gpu=index,memory.used,memory.total --format=csv,noheader,nounits | while read line; do
    echo "  GPU $line"
done
echo "========================================"
echo "🚀 Starting server with optimized parameters..."
echo "   - Max sequences: 80"
echo "   - Batch tokens: 32768"
echo "   - Expert parallel: ENABLED"
echo "   - Swap space: 16GB"
echo "========================================"

vllm serve "$MODEL_PATH" \
  --served-model-name deepseek-r1 \
  --tensor-parallel-size 16 \
  --enable-expert-parallel \
  --distributed-executor-backend ray \
  --trust-remote-code \
  --dtype auto \
  --gpu-memory-utilization 0.90 \
  --max-model-len 8192 \
  --max-num-seqs 80 \
  --max-num-batched-tokens 32768 \
  --host 0.0.0.0 --port 8000 \
  --api-key p10-deepseek-key \
  --disable-log-requests \
  --swap-space 16
EOF
chmod +x deepseek_optimized_server_v3.sh

./deepseek_optimized_server_v3.sh > vllm_server.log 2>&1 &
pid_vllm=$!
echo "vLLM server PID: $pid_vllm"

# 起動待ち（ポート 8000 が開くまで）
echo "Waiting for vLLM server to become ready on :8000 ..."
for i in $(seq 1 60); do
  python - <<'PY' || true
import socket, sys
s=socket.socket()
s.settimeout(2)
try:
    s.connect(("127.0.0.1", 8000))
    s.close()
    sys.exit(0)
except Exception:
    sys.exit(1)
PY
  rc=$?
  if [[ ${rc:-1} -eq 0 ]]; then
    echo "vLLM server is up."
    break
  fi
  sleep 5
done

#=============== 必要なディレクトリ ===============
cd llm_bridge_prod/eval_hle/ || exit 0
mkdir -p predictions judged

#=============== 推論・評価 ===============
python predict.py > predict.log 2>&1
python judge.py > judge.log 2>&1

#=============== 後片付け（明示） ===============
kill "${pid_vllm}"  2>/dev/null || true
kill "${pid_nvsmi}" 2>/dev/null || true
srun -N ${SLURM_NNODES:-1} -n ${SLURM_NNODES:-1} /bin/bash -lc '
  source /home/appli/miniconda3/24.7.1-py311/etc/profile.d/conda.sh
  conda activate "$HOME/conda_env_llm" || true
  ray stop --force || true
' || true
echo "Done."
__RUN_SBATCH__

echo "Saved to run_deepseek.sbatch (submit with: sbatch run_deepseek.sbatch)"