# 定义流水线要使用的 Docker 镜像
image: image.sourcefind.cn:5000/dcu/admin/base/vllm:0.9.2-ubuntu22.04-dtk25.04.1-rc5-rocblas101839-0811-das1.6-py3.10-20250908-rc1

# 只定义一个阶段
stages:
  - test

# === 唯一的作业：完整前置操作 + for 循环串行测试 ===
benchmark_hunyuan-dit:
  stage: test
  tags:
    - demos

  script:
    - echo "========================================="
    - echo "Step 1:Setting up the environment from /workspace/packages/hunyuan-dit (One-time setup)"
    - echo "========================================="
    # 定义预置文件的根目录
    - export PACKAGE_DIR="/workspace/packages/hunyuan-dit"
    # 从预置目录解压库文件 (这些操作只做一次！)
    - tar -xzf "$PACKAGE_DIR/hipblaslt-install0925.tar.gz"
    - tar -xzf "$PACKAGE_DIR/package_0915_ubuntu.tar.gz"

    - echo "========================================="
    - echo "Step 2:Installing Python packages from /workspace/packages/hunyuan-dit (One-time setup)"
    - echo "========================================="
    # 从预置目录安装 Python 包 (这些操作也只做一次！)
    - pip install "$PACKAGE_DIR/apex-1.5.0+das.opt1.dtk25041-cp310-cp310-linux_x86_64.whl" "$PACKAGE_DIR/deepspeed-0.14.2+das.opt1.dtk25041-cp310-cp310-manylinux_2_28_x86_64.whl" "$PACKAGE_DIR/lightop-0.5.0+das.dtk25041.unknown-cp310-cp310-linux_x86_64.whl" 
    - pip install -r requirements.txt   -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple

    - echo "========================================="
    - echo "Step 3:Starting serial benchmark with a for-loop"
    - echo "========================================="

    # 【核心】定义要测试的批次大小列表
    - BATCH_SIZES=(1 2 4 8)

    # 创建一个最终的产物目录，用于存放所有批次的测试结果
    - mkdir -p all_results

    # 【核心】使用 for 循环串行执行
    - |
      for bs in "${BATCH_SIZES[@]}"; do
        echo "============================================================"
        echo ">>> Running benchmark for BATCH_SIZE = $bs"
        echo "============================================================"
        
        # 设置本次运行的环境变量 (这些变量在每次循环中都需要)
        export LD_LIBRARY_PATH=/workspace/packages/hunyuan-dit/hipblaslt-install/lib/:$LD_LIBRARY_PATH
        export LD_LIBRARY_PATH=/workspace/packages/hunyuan-dit/package/miopen/lib/:$LD_LIBRARY_PATH
        
        # 运行测试脚本，将 $bs 作为 batch-size 参数
        python sample_t2i_dcu.py \
          --model-root /workspace/packages/hunyuan-dit/HunyuanDiT-v1.2/ \
          --batch-size $bs \
          --infer-mode fa \
          --prompt "青花瓷风格，一只可爱的哈士奇" \
          --no-enhance \
          --load-key module \
          --image-size 1024 1024 \
          --infer-steps 20

        # 【关键】整理本次循环的产物
        # 检查 results 目录是否存在且不为空
        if [ -d "results" ] && [ "$(ls -A results)" ]; then
          echo ">>> Saving results for batch size $bs..."
          # 将本次的 results 目录重命名并移动到 all_results 中
          mv results "all_results/results_bs_$bs"
        else
          echo ">>> No results found for batch size $bs, skipping."
        fi
        
        # 清理环境，为下一次循环做准备（可选，但推荐）
        rm -rf results # 确保没有残留的空目录
        echo ">>> Benchmark for batch size $bs finished."
        echo ""
      done

    - echo "========================================="
    - echo "All benchmarks finished."
    - echo "Final result directory structure:"
    - ls -R all_results # 打印最终的目录结构，方便检查
    - echo "========================================="

  # 【优化】定义最终的产物
  artifacts:
    name: "benchmark-hunyuan-dit-serial-results"
    # 只上传包含所有批次结果的最终目录
    paths:
      - all_results/
    expire_in: 1 week