#!/bin/bash
# CoT 数据生成 - Nebula 提交（本地大 Qwen-VL 模型 vLLM 离线批推）
# Usage: sh scripts/submit/submit_cot_gen.sh
# 远程机器上 OSS 挂载在 /data/oss_bucket_0/，模型/数据/输出都直接走该路径
# entry 直接指向生成脚本，nebula 每卡拉起一个进程，脚本内部按 RANK/WORLD_SIZE 分片，
# 每个进程用 tp=1 独占一张卡跑 vLLM，各分片写各自的 .shardN-M.jsonl，跑完再合并

## 集群资源配置（gpu:100 表示 1 张卡，27B 模型 bf16 单卡 A800-80G 放得下）
CLUSTER_CONFIG='./cluster.json'
cat <<EOT > ${CLUSTER_CONFIG}
{
    "worker": {
        "gpu": 100,
        "cpu": 1200,
        "memory": 50000
    }
}
EOT
## worker 数即分片数，8 个 worker 各跑 1/8 数据
WORKER_COUNT=8

## ============ 模型设置 ============
MODEL_PATH=/data/oss_bucket_0/planck/hf_models/Qwen3.6-27B

## ============ 数据集设置 ============
TRAIN_DATA_PATH=/data/oss_bucket_0/siku/Last_reranker/dataset/all_data_wo_q_bbox/00-07/with_bbox/train_format/eval.jsonl
NUM_DOCS=10
MAX_CROP_TOKENS=80
PATCH_SIZE=16

## ============ 生成设置 ============
MAX_MODEL_LEN=16384
MAX_NEW_TOKENS=4096
MAX_RETRIES=3
RETRY_TEMPERATURE=0.7
BATCH_SIZE=64
SEED=42

## ============ 保存设置 ============
RUN_ID=qwen36-27b-cot-gen-eval
OUTPUT_PATH=/data/oss_bucket_0/liqi/dataset/reranker/cot_data_eval/eval_cot.jsonl

## 试跑：设为 "--test 20" 时每个分片只跑前 20 条、输出 train_cot.test.shardN-M.jsonl
## 正式全量生成时留空
TEST=""

## ============ 拼装参数（user_params 会原样传给 entry 脚本） ============
## 不传 --shard-id/--num-shards，脚本自动取 RANK/WORLD_SIZE；分片模式下 tp 强制为 1
args="\
    --model-path $MODEL_PATH \
    --train-data-path $TRAIN_DATA_PATH \
    --output-path $OUTPUT_PATH \
    --num-docs $NUM_DOCS \
    --seed $SEED \
    --max-crop-tokens $MAX_CROP_TOKENS \
    --patch-size $PATCH_SIZE \
    --tensor-parallel-size 1 \
    --max-model-len $MAX_MODEL_LEN \
    --max-new-tokens $MAX_NEW_TOKENS \
    --max-retries $MAX_RETRIES \
    --retry-temperature $RETRY_TEMPERATURE \
    --batch-size $BATCH_SIZE \
    $TEST"

nebulactl run mdl \
          --engine=xdl \
          --ignore=venv/*,data/*,.git/*,cot_data/* \
          --job_name=${RUN_ID} \
          --queue="pailitao_algo_high_2" \
          --entry="cot_gen/gen_cot_data.py" \
          --nebula_project=pailitao_v100 \
          --algo_name=pytorch2100 \
          --worker_count=${WORKER_COUNT} \
          --file.cluster_file="./cluster.json" \
          --user_params="$args" \
          --nas_file_system_id=98d384a407-okl3.cn-zhangjiakou.nas.aliyuncs.com \
          --oss_access_id=LTAI5tMZb5HHH6HHgAExAyiL \
          --oss_access_key=6MVP8JnTE01t4TLUHhPoFV0swZNTrb \
          --oss_bucket=pailitao-ai \
          --oss_endpoint=oss-cn-zhangjiakou.aliyuncs.com \
          --requirements_file_name=requirements_cot_gen.txt \
          --force \
          --max_failover_times=2

rm ${CLUSTER_CONFIG}
# 分片数由 WORKER_COUNT 决定，各分片输出：
#   ${OUTPUT_PATH%.jsonl}.shard{0..7}-8.jsonl        成功样本
#   ${OUTPUT_PATH%.jsonl}.shard{0..7}-8.jsonl.failed 重试仍失败的样本（含 raw_output 便于排查）
# 全部跑完后合并：
#   cat ${OUTPUT_PATH%.jsonl}.shard*-${WORKER_COUNT}.jsonl > ${OUTPUT_PATH}
# 各分片支持断点续跑：重新提交同一 job 会跳过已写入的 sample_idx