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

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

# === 唯一的作业：使用预置文件运行基准测试 ===
benchmark_wan2.1:
  stage: test
  tags:
    - demos
  script:
    - echo "========================================="
    - echo "Step 1 Setting up the environment from /workspace/packages/wan2.1"
    - echo "========================================="
    # 定义预置文件的根目录
    - export PACKAGE_DIR="/workspace/packages/wan2.1"
    # 从预置目录解压库文件
    - tar -xzf "$PACKAGE_DIR/rocblas-install-0910-bug.tar.gz"
    # 设置环境变量，模型路径指向预置目录
    - export LD_LIBRARY_PATH="$CI_PROJECT_DIR/rocblas-install/lib/:$LD_LIBRARY_PATH"
    - pip install -r requirements.txt -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
    - ln -s /workspace/packages/wan2.1/Wan_2.1/models/ ./models
    
    - echo "========================================="
    - echo "Step 2 Setting up ComfyUI directories"
    - echo "========================================="
    # 创建ComfyUI需要的目录结构
    - mkdir -p custom_nodes
    - mkdir -p input
    - mkdir -p output
    - mkdir -p temp
    - echo "Created ComfyUI directory structure"
    - ls -la
    
    - echo "========================================="
    - echo "Step 3 Starting ComfyUI server"
    - echo "========================================="
    # 创建日志目录
    - mkdir -p logs
    # 后台启动ComfyUI服务器并记录日志
    - bash run-main.sh > logs/comfyui.log 2>&1 &
    # 获取ComfyUI进程ID
    - export COMFYUI_PID=$!
    - echo "ComfyUI server started with PID $COMFYUI_PID"
    
    - echo "========================================="
    - echo "Step 4 Waiting for ComfyUI server to be ready"
    - echo "========================================="
    # 等待服务器启动，最多等待5分钟
    - |
      timeout=300
      count=0
      while [ $count -lt $timeout ]; do
        if grep -q "http://127.0.0.1:8188" logs/comfyui.log; then
          echo "ComfyUI server is ready!"
          break
        fi
        sleep 1
        count=$((count+1))
      done
      
      if [ $count -eq $timeout ]; then
        echo "Timeout waiting for ComfyUI server to start"
        exit 1
      fi
    
    - echo "========================================="
    - echo "Step 5 Running benchmark tests with different seeds"
    - echo "========================================="
    # 创建结果目录
    - mkdir -p results
    # 创建性能日志文件
    - touch results/performance.log
    
    # 运行第一次测试（seed=2675441231）
    - echo "Running first test with seed=2675441231"
    - |
      # 修改Python脚本以使用第一个seed
      sed 's/prompt\["3"\]\["inputs"\]\["seed"\] = args.seed/prompt["3"]["inputs"]["seed"] = 2675441231/' wan_t2v_14B_1_gpu.py > temp_test1.py
      
      # 记录开始时间
      start_time=$(date +%s)
      
      # 运行测试并捕获输出
      python temp_test1.py 2>&1 | tee results/test1_output.log
      
      # 记录结束时间和计算执行时间
      end_time=$(date +%s)
      duration=$((end_time - start_time))
      echo "Test 1 execution time $duration seconds" >> results/performance.log
    
    # 等待一段时间确保第一次测试完成
    - sleep 10
    
    # 运行第二次测试（seed=1234567890）
    - echo "Running second test with seed=1234567890"
    - |
      # 修改Python脚本以使用第二个seed
      sed 's/prompt\["3"\]\["inputs"\]\["seed"\] = args.seed/prompt["3"]["inputs"]["seed"] = 1234567890/' wan_t2v_14B_1_gpu.py > temp_test2.py
      
      # 记录开始时间
      start_time=$(date +%s)
      
      # 运行测试并捕获输出
      python temp_test2.py 2>&1 | tee results/test2_output.log
      
      # 记录结束时间和计算执行时间
      end_time=$(date +%s)
      duration=$((end_time - start_time))
      echo "Test 2 execution time $duration seconds" >> results/performance.log
    
    - echo "========================================="
    - echo "Step 6 Collecting results"
    - echo "========================================="
    # 查找并复制输出文件
    - |
      # 查找最新的.webp文件
      find output -name "*.webp" -type f -exec ls -lt {} + | head -n 2 > results/latest_files.txt
      
      # 复制最新的两个.webp文件到结果目录
      while read line; do
        filepath=$(echo "$line" | awk '{print $NF}')
        cp "$filepath" results/
      done < results/latest_files.txt
      
      # 创建结果摘要
      echo "=== Test Results Summary ===" > results/summary.txt
      echo "Test 1 (seed=2675441231):" >> results/summary.txt
      grep "Prompt executed in" results/test1_output.log >> results/summary.txt
      echo "" >> results/summary.txt
      echo "Test 2 (seed=1234567890):" >> results/summary.txt
      grep "Prompt executed in" results/test2_output.log >> results/summary.txt
    
    - echo "========================================="
    - echo "Step 7 Cleaning up"
    - echo "========================================="
    # 停止ComfyUI服务器
    - kill $COMFYUI_PID || true
    # 清理临时文件
    - rm -f temp_test1.py temp_test2.py

  # 定义产物，保留测试生成的文件
  artifacts:
    name: "$CI_JOB_NAME-$CI_COMMIT_REF_NAME"
    paths:
      - "results/*.webp"
      - "results/*.log"
      - "results/*.txt"
      - "logs/comfyui.log"
    expire_in: 1 week