jobName: "sft-train-r{RUNDIC}"
description: "AutoResearch R{RUNDIC} SFT training (40k+ samples, qwen3-4B base, 3 epochs FSDP2)"
accessType: PUBLIC

imageConfig:
  imageUrl: micr.cloud.mioffice.cn/vllm-image/ai-arch-llm-prod:vllm-v0.12.0-f098b188
  imageCommand: |-
    set -ex
    cd {AUTORESEARCH_ROOT}

    # 0a. 验环境 + 装缺的包（vllm 镜像 py3.12 + torch 已含）
    python3 --version
    python3 -c "import torch; print(f'torch={torch.__version__}')"
    pip install -q --no-deps "accelerate==1.7.0" 2>&1 | tail -3
    pip install -q --ignore-installed blinker 2>&1 | tail -2
    pip install -q peft 2>&1 | tail -3
    pip install -q wandb bitsandbytes 2>&1 | tail -3
    pip install -q luigi mlflow scikit-learn openpyxl pyyaml sentencepiece tiktoken protobuf pynvml datasets 2>&1 | tail -3
    pip install -q "transformers>=4.45" 2>&1 | tail -3
    python3 -c 'import torch, accelerate, transformers, peft; print(torch.__version__, accelerate.__version__, transformers.__version__, peft.__version__)'

    # 0c. 把 HF cache 重定向到容器本地大盘（默认在 ~/.cache 可能是 juicefs，mmap 易 SIGBUS）
    export HF_HOME=/tmp/hf_cache
    export HF_DATASETS_CACHE=/tmp/hf_cache/datasets
    export TRANSFORMERS_CACHE=/tmp/hf_cache/transformers
    mkdir -p /tmp/hf_cache/datasets /tmp/hf_cache/transformers
    df -h /tmp /dev/shm 2>/dev/null || true

    # 强制 datasets 加载进内存而不是 arrow mmap（避免 /dev/shm 溢出 SIGBUS）
    export HF_DATASETS_IN_MEMORY_MAX_SIZE=20000000000
    export HF_DATASETS_NUM_PROC=1
    export TOKENIZERS_PARALLELISM=false
    export OMP_NUM_THREADS=1
    export MKL_NUM_THREADS=1
    # 关闭 NCCL 用 shm 通信（改 socket 通道）
    export NCCL_SHM_DISABLE=1
    export NCCL_P2P_DISABLE=0
    export NCCL_DEBUG=WARN

    # 0b. 把上一轮产出搬走（cml job 内幂等，本地搬过就跳过）
    if [ -d sft_output ] && [ ! -d sft_output_r{PREV_RUNDIC} ]; then
      mv sft_output sft_output_r{PREV_RUNDIC}
    fi
    rm -rf sft_output

    # 1. 组装数据
    python3 {AUTORESEARCH_ROOT}/prepare_and_train_sft.py prepare \
      --output_dir {AUTORESEARCH_ROOT}/sft_data

    # 2. 启动训练（绝对路径，避免相对路径在容器内找不到 zk_trainer）
    python3 {AUTORESEARCH_ROOT}/prepare_and_train_sft.py train \
      --data_dir {AUTORESEARCH_ROOT}/sft_data \
      --model_path /mnt/wangsenhao/verl_zk/Qwen3-4B-Instruct-2507 \
      --model_type qwen3 \
      --train_output {AUTORESEARCH_ROOT}/sft_output \
      --epochs 3 --lr 1e-5

    # 3. 训练成功标记（被 watcher 检测）
    [ -f {AUTORESEARCH_ROOT}/sft_output/config.json ] && \
      touch {AUTORESEARCH_ROOT}/sft_output/_SUCCESS

# 挂载 wangsenhao + xiaoai-zk-model-train-tj5 + verl_zk 所在卷
juiceFsMountConfigs:
  - volume: wangsenhao
    juiceFsCluster: tj5-common
    subPath: /
    mountPath: /mnt/wangsenhao
    readOnly: false
  - volume: xiaoai-zk-model-train-tj5
    juiceFsCluster: tj5-common
    subPath: /
    mountPath: /mnt/xiaoai-zk-model-train-tj5
    readOnly: false

envConfigs:
  - key: PYTHONPATH
    value: {AUTORESEARCH_ROOT}/zk_trainer
  - key: WANDB_DISABLED
    value: "true"
  - key: WANDB_MODE
    value: offline
  - key: VOLUME_PREFIX
    value: /mnt/xiaoai-zk-model-train-tj5
  - key: HF_HOME
    value: /mnt/wangsenhao/.hf_cache

# h20-96g 8 卡 FSDP2
queueId: "6052"
priority: 5
preemptible: false
framework: pytorch
resourceConfigs:
  - nodeRole: worker
    nodeNumber: 1
    perNodeResourceSpec:
      resourcePriority: GUARANTEED
      resourceName: cloudml.ng2h20-8-8.20-199
      resourceNumber: 8

# 故障自动重试（节点级失败）
retryConfig:
  enableRetry: true
  maxRetryTimes: 2
  policySets:
    - NodeFailure

# 失败/完成飞书告警
alertConfig:
  enableAlert: true
  alertItems:
    - alertConditions:
        - FAILED
        - SUCCEED
      alertLevel: P2
      alertReceivers:
        persons:
          - wangsenhao
