# 定义流水线要使用的 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 "========================================="
    - 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
    - bash run-main.sh > logs/comfyui.log 2>&1 &
    - export COMFYUI_PID=$!
    - echo "ComfyUI server started with PID $COMFYUI_PID"
    
    - echo "========================================="
    - echo "Step 4 Waiting for ComfyUI server to be ready"
    - echo "========================================="
    - 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
    
    - echo "Running first test with seed=2675441231"
    - cp wan_t2v_14B_1_gpu.py temp_test1.py
    - 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 "SUCCESS: Seed 2675441231 set"; grep -n '"seed": 2675441231,' temp_test1.py; else echo "ERROR: Failed to set seed 2675441231"; 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..."
    - 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 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"
    
    - echo "Running second test with seed=1234567890"
    - cp wan_t2v_14B_1_gpu.py temp_test2.py
    - 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 "SUCCESS: Seed 1234567890 set"; grep -n '"seed": 1234567890,' temp_test2.py; else echo "ERROR: Failed to set seed 1234567890"; echo "Using Python fallback method"; python3 -c "import re; content=open('temp_test2.py').read(); content=re.sub(r'(\\\"seed\\\":\\s*)\\d+', r'\\g<1>1234567890', content); open('temp_test2.py','w').write(content)"; if grep -q '"seed": 1234567890,' temp_test2.py; then echo "SUCCESS: Seed set using Python method"; grep -n '"seed": 1234567890,' temp_test2.py; else echo "ERROR: All methods failed"; 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..."
    - 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 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"
    - echo "========================================="
    - 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
    - find output -name "*.webp" -type f -exec ls -lt {} + | head -n 2 > results/latest_files.txt
    - cat results/latest_files.txt | awk '{print $NF}' | while read filepath; do cp "$filepath" results/ && echo "Copied: $filepath"; done
    - 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 "No execution time found in test1" >> 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 "No execution time found in test2" >> results/summary.txt
    - echo "" >> results/summary.txt
    - echo "Generated files:" >> results/summary.txt
    - ls results/*.webp >> results/summary.txt || echo "No webp files found in results" >> results/summary.txt
    - echo "All tests completed successfully!"
    
    - echo "========================================="
    - echo "Step 7 Cleaning up"
    - echo "========================================="
    - 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/"
      - "logs/"
    expire_in: 1 week