cat <<'__RUN_SBATCH__' > run_deepseek.sbatch
#!/bin/bash
#SBATCH --job-name=qwen3_8gpu
#SBATCH --partition=P10
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=240
#SBATCH --mem=0
#SBATCH --time=12: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=ALL,OPENAI_API_KEY="openai_api_keyをここに"

#=============== モード切替（--skip-env対応） ===============
MODE_ARG="${1:-}"
case "$MODE_ARG" in
  --skip-env) SKIP_ENV_SETUP=1 ;;
  --setup-env) SKIP_ENV_SETUP=0 ;;
esac
: "${SKIP_ENV_SETUP:=0}"                  # 1 で conda 作成/インストールをスキップ
: "${DEEPSEEK_ENV:=$HOME/conda_env_llm}"  # 既存/新規 conda 環境のパス

#=============== 前処理 / 片付け ===============
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_vllm_pretty:-}" ]] && kill "${pid_vllm_pretty}" 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 "'"$DEEPSEEK_ENV"'" || true
    timeout 10s ray stop --force || true
    pkill -9 ray || true
    pkill -9 gcs_server || true
    # 自分の所有物のみ削除（他ユーザの残骸は触らない）
    find /tmp/ray -maxdepth 1 -type d -name "session_*" -user "$USER" -print0 2>/dev/null | xargs -0r rm -rf
    find /tmp -maxdepth 1 -type d -name "ray_tmp_*" -user "$USER" -print0 2>/dev/null | xargs -0r rm -rf
    rm -rf /dev/shm/ray* 2>/dev/null || 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

if [[ "$SKIP_ENV_SETUP" -eq 1 ]]; then
  echo "[ENV] SKIP_ENV_SETUP=1: conda作成/インストールをスキップします。環境を有効化: $DEEPSEEK_ENV"
  conda activate "$DEEPSEEK_ENV" || echo "[ENV][WARN] conda activate 失敗: $DEEPSEEK_ENV（既存PATHで続行）"
else
  echo "[ENV] SKIP_ENV_SETUP=0: 新規に環境を構築します（従来どおり）"
  export DEEPSEEK_ENV="$DEEPSEEK_ENV"
  [ -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
fi

#=============== Ray の前処理（自分の所有物だけ掃除） ===============
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 "'"$DEEPSEEK_ENV"'" || true
  timeout 10s ray stop --force || true
  pkill -9 ray || true
  pkill -9 gcs_server || true
  find /tmp/ray -maxdepth 1 -type d -name "session_*" -user "$USER" -print0 2>/dev/null | xargs -0r rm -rf
  find /tmp -maxdepth 1 -type d -name "ray_tmp_*" -user "$USER" -print0 2>/dev/null | xargs -0r rm -rf
  df -h /dev/shm || true
  ls -la /dev/shm/ || true
  rm -rf /dev/shm/ray* 2>/dev/null || true
'

# 割り当てノード一覧
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 RAY_TMPDIR="${SLURM_TMPDIR:-/tmp}/ray-${SLURM_JOB_ID:-manual}"
mkdir -p "$RAY_TMPDIR"

# ヘッドノードでRay起動（--temp-dir 追加）
export VLLM_HOST_IP=$HEAD_IP
export NCCL_SOCKET_IFNAME=enp25s0np0
export NCCL_IB_DISABLE=0
ray start --head --node-ip-address="$HEAD_IP" --port=6379 \
  --dashboard-host=0.0.0.0 --dashboard-port=8265 \
  --temp-dir "$RAY_TMPDIR"
export RAY_ADDRESS=$HEAD_IP:6379
echo "✅ Rayヘッドノード起動完了: $RAY_ADDRESS"
echo "📁 RAY_TMPDIR: $RAY_TMPDIR"

# 残りノードでワーカー起動（--temp-dir 追加、--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"'" || true
      export NCCL_SOCKET_IFNAME=enp25s0np0
      export NCCL_IB_DISABLE=0
      export RAY_TMPDIR="${SLURM_TMPDIR:-/tmp}/ray-${SLURM_JOB_ID:-manual}"
      mkdir -p "$RAY_TMPDIR"
      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" \
        --temp-dir "$RAY_TMPDIR" --block
    ' &
  done
  # Nodes が揃うまでヘッドでポーリング（timeout付きでハング回避）
  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"'" || true
      RAY_ADDRESS="'"$RAY_ADDRESS"'" timeout 6s ray status 2>/dev/null | 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 "${VLLM_HOST_IP:-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"

# 見やすい整形ログ: タイムスタンプ＋プレフィックスで別ファイルに並行生成
stdbuf -oL -eL tail -n +1 -F vllm_server.log \
  | awk '{ print strftime("[%Y-%m-%d %H:%M:%S]"), "[vllm]", $0; fflush(); }' \
  > vllm_server.pretty.log &
pid_vllm_pretty=$!

# 起動待ち（ポート 8000 が開くまで）
echo "Waiting for vLLM server to become ready on :8000 ..."
for i in $(seq 1 60); do
  python - <<'PY' || true
import socket, sys, os
s=socket.socket()
s.settimeout(2)
try:
    head_ip = os.environ.get("VLLM_HOST_IP", "127.0.0.1")
    s.connect((head_ip, 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_vllm_pretty}" 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 "'"$DEEPSEEK_ENV"'" || true
  timeout 10s ray stop --force || true
  pkill -9 ray || true
  pkill -9 gcs_server || true
  find /tmp/ray -maxdepth 1 -type d -name "session_*" -user "$USER" -print0 2>/dev/null | xargs -0r rm -rf
  find /tmp -maxdepth 1 -type d -name "ray_tmp_*" -user "$USER" -print0 2>/dev/null | xargs -0r rm -rf
  rm -rf /dev/shm/ray* 2>/dev/null || true
' || true
echo "Done."
__RUN_SBATCH__

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