# 定义流水线要使用的 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

# === 唯一的作业：使用预置文件运行基准测试 ===
benchmark_hunyuan-dit:
  stage: test
  tags:
    - demos

  # 【核心修改 1】使用矩阵策略，自动创建多个并行 Job
  # GitLab 会为 BATCH_SIZE 列表中的每个值创建一个独立的 Job
  parallel:
    matrix:
      - BATCH_SIZE: [1, 2, 4, 8]

  script:
    - echo "========================================="
    - echo "Step 1: Setting up the environment for BATCH_SIZE = $BATCH_SIZE"
    - 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"
    - 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" 

    - echo "========================================="
    - echo "Step 3: Running the benchmark with BATCH_SIZE = $BATCH_SIZE"
    - 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
    
    # 【核心修改 2】使用矩阵变量 $BATCH_SIZE
    # 运行测试脚本，batch-size 从固定的 4 改为动态的 $BATCH_SIZE
    - python sample_t2i_dcu.py --model-root /workspace/OEM_ADVTG_TEST/hunyuan/HunyuanDiT-v1.2/ --batch-size $BATCH_SIZE --infer-mode fa --prompt "青花瓷风格，一只可爱的哈士奇" --no-enhance --load-key module --image-size 1024 1024 --infer-steps 20

  # 【核心修改 3】优化产物定义
  artifacts:
    # 为产物命名，包含批次大小，方便在 GitLab UI 上区分
    name: "benchmark-hunyuan-dit-bs-$BATCH_SIZE"
    # 【优化】指定 'results/' 目录，因为你的 Python 脚本把图片保存在这里
    # 这比 "*.png" 更精确，能确保只收集我们想要的文件
    paths:
      - results/
    expire_in: 1 week