# 定义流水线要使用的 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 precise seed replacement"
    - echo "========================================="
    # 创建结果目录
    - mkdir -p results
    # 创建性能日志文件
    - touch results/performance.log
    
    # 运行第一次测试（seed=2675441231）
    - echo "Running first test with seed=2675441231"
    - |
      # 创建第一个测试文件
      cp wan_t2v_14B_1_gpu.py temp_test1.py
      
      # 精确替换JSON中的seed值
      sed -i 's/"seed": [0-9]\+,/"seed": 2675441231,/' temp_test1.py
      
      # 验证修改是否成功
      echo "=== Verifying temp_test1.py seed modification ==="
      if grep -q '"seed": 2675441231,' temp_test1.py; then
        echo "✓ Seed 2675441231 successfully set in temp_test1.py"
        grep -n '"seed": 2675441231,' temp_test1.py
      else
        echo "✗ Failed to set seed 2675441231 in temp_test1.py"
        echo "Current seed lines:"
        grep -n "seed" temp_test1.py
        exit 1
      fi
      
      # 记录开始时间
      start_time=$(date +%s)
      
      # 运行测试
      python temp_test1.py > results/test1_output.log 2>&1 &
      export TEST1_PID=$!
      
      echo "Waiting for first webp file to be generated..."
      # 等待第一个webp文件生成，最多等待10分钟
      timeout=600
      count=0
      while [ $count -lt $timeout ]; do
        webp_count=$(find output -name "*.webp" -type f | wc -l)
        if [ $webp_count -gt 0 ]; then
          echo "First webp file detected! Count: $webp_count"
          break
        fi
        sleep 2
        count=$((count+2))
        echo "Waiting... ${count}s elapsed, webp count: $webp_count"
      done
      
      if [ $count -eq $timeout ]; then
        echo "Timeout waiting for first webp file"
        exit 1
      fi
      
      # 记录结束时间和计算执行时间
      end_time=$(date +%s)
      duration=$((end_time - start_time))
      echo "Test 1 execution time $duration seconds" >> results/performance.log
      echo "Test 1 completed in $duration seconds"
    
    # 运行第二次测试（seed=1234567890）
    - echo "Running second test with seed=1234567890"
    - |
      # 创建第二个测试文件
      cp wan_t2v_14B_1_gpu.py temp_test2.py
      
      # 精确替换JSON中的seed值
      sed -i 's/"seed": [0-9]\+,/"seed": 1234567890,/' temp_test2.py
      
      # 验证修改是否成功
      echo "=== Verifying temp_test2.py seed modification ==="
      if grep -q '"seed": 1234567890,' temp_test2.py; then
        echo "✓ Seed 1234567890 successfully set in temp_test2.py"
        grep -n '"seed": 1234567890,' temp_test2.py
      else
        echo "✗ Failed to set seed 1234567890 in temp_test2.py"
        echo "Current seed lines:"
        grep -n "seed" temp_test2.py
        echo "Attempting alternative replacement method..."
        
        # 备用方法：使用Python直接修改
        python3 -c "
      import re
      with open('temp_test2.py', 'r') as f:
          content = f.read()
          
      # 使用正则表达式替换seed值
      content = re.sub(r'(\"seed\":\s*)\d+', r'\g<1>1234567890', content)

      with open('temp_test2.py', 'w') as f:
          f.write(content)
      "
        
        # 再次验证
        if grep -q '"seed": 1234567890,' temp_test2.py; then
          echo "✓ Seed 1234567890 successfully set using Python method"
          grep -n '"seed": 1234567890,' temp_test2.py
        else
          echo "✗ All methods failed to set seed 1234567890"
          echo "Full file content for debugging:"
          cat temp_test2.py
          exit 1
        fi
      fi
      
      # 记录开始时间
      start_time=$(date +%s)
      
      # 运行测试
      python temp_test2.py > results/test2_output.log 2>&1 &
      export TEST2_PID=$!
      
      echo "Waiting for second webp file to be generated..."
      # 等待第二个webp文件生成，最多等待10分钟
      timeout=600
      count=0
      while [ $count -lt $timeout ]; do
        webp_count=$(find output -name "*.webp" -type f | wc -l)
        if [ $webp_count -gt 1 ]; then
          echo "Second webp file detected! Total count: $webp_count"
          break
        fi
        sleep 2
        count=$((count+2))
        echo "Waiting... ${count}s elapsed, webp count: $webp_count"
      done
      
      if [ $count -eq $timeout ]; then
        echo "Timeout waiting for second webp file"
        exit 1
      fi
      
      # 记录结束时间和计算执行时间
      end_time=$(date +%s)
      duration=$((end_time - start_time))
      echo "Test 2 execution time $duration seconds" >> results/performance.log
      echo "Test 2 completed in $duration seconds"
    
    - echo "========================================="
    - echo "Step 6 Verifying and collecting results (robust version)"
    - echo "========================================="
    - |
      # 验证两个webp文件都已生成
      webp_count=$(find output -name "*.webp" -type f | wc -l)
      echo "Total webp files found: $webp_count"
      
      if [ $webp_count -lt 2 ]; then
        echo "ERROR: Expected 2 webp files, but found $webp_count"
        find output -name "*.webp" -type f -ls
        exit 1
      fi
      
      # 查找并复制最新的两个webp文件
      find output -name "*.webp" -type f -exec ls -lt {} + | head -n 2 > results/latest_files.txt
      
      # 复制文件到结果目录，添加文件存在性检查
      while read line; do
        filepath=$(echo "$line" | awk '{print $NF}')
        if [ -f "$filepath" ]; then
          cp "$filepath" results/
          echo "Copied: $filepath"
        else
          echo "Warning: File not found, skipping: $filepath"
        fi
      done < results/latest_files.txt
      
      # 创建结果摘要，使用不会失败的方法
      echo "=== Test Results Summary ===" > results/summary.txt
      
      echo "Test 1 (seed=2675441231):" >> results/summary.txt
      test1_time=$(grep "Prompt executed in" results/test1_output.log || echo "Execution time not found in log")
      echo "$test1_time" >> results/summary.txt
      
      echo "" >> results/summary.txt
      echo "Test 2 (seed=1234567890):" >> results/summary.txt
      test2_time=$(grep "Prompt executed in" results/test2_output.log || echo "Execution time not found in log")
      echo "$test2_time" >> results/summary.txt
      
      echo "" >> results/summary.txt
      echo "Performance data from script timing:" >> results/summary.txt
      cat results/performance.log >> results/summary.txt
      
      echo "" >> results/summary.txt
      echo "Generated files:" >> results/summary.txt
      ls results/*.webp >> results/summary.txt 2>/dev/null || echo "No .webp files found in results directory" >> results/summary.txt
      
      echo "All tests completed successfully!"
    
    - echo "========================================="
    - echo "Step 7 Cleaning up"
    - echo "========================================="
    # 停止ComfyUI服务器
    - kill $COMFYUI_PID || true
    # 清理临时文件
    - rm -f temp_test1.py temp_test2.py
    - echo "Cleanup completed"

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