#!/bin/bash
# vllm-omni customized version
# Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/bootstrap-amd.sh
# Last synced: 2025-12-15
# Modifications: Use local template file instead of downloading from ci-infra

set -euo pipefail

if [[ -z "${RUN_ALL:-}" ]]; then
    RUN_ALL=0
fi

if [[ -z "${NIGHTLY:-}" ]]; then
    NIGHTLY=0
fi

if [[ -z "${VLLM_CI_BRANCH:-}" ]]; then
    VLLM_CI_BRANCH="main"
fi

if [[ -z "${AMD_MIRROR_HW:-}" ]]; then
    AMD_MIRROR_HW="amdproduction"
fi

if [[ -z "${DOCS_ONLY_DISABLE:-}" ]]; then
    DOCS_ONLY_DISABLE=0
fi

fail_fast() {
    DISABLE_LABEL="ci-no-fail-fast"
    # If BUILDKITE_PULL_REQUEST != "false", then we check the PR labels using curl and jq
    if [ "$BUILDKITE_PULL_REQUEST" != "false" ]; then
        PR_LABELS=$(curl -s "https://api.github.com/repos/vllm-project/vllm-omni/pulls/$BUILDKITE_PULL_REQUEST" | jq -r '.labels[].name')
        if [[ $PR_LABELS == *"$DISABLE_LABEL"* ]]; then
            echo false
        else
            echo true
        fi
    else
        echo false  # not a PR or BUILDKITE_PULL_REQUEST not set
    fi
}

check_run_all_label() {
    RUN_ALL_LABEL="ready-run-all-tests"
    # If BUILDKITE_PULL_REQUEST != "false", then we check the PR labels using curl and jq
    if [ "$BUILDKITE_PULL_REQUEST" != "false" ]; then
        PR_LABELS=$(curl -s "https://api.github.com/repos/vllm-project/vllm-omni/pulls/$BUILDKITE_PULL_REQUEST" | jq -r '.labels[].name')
        if [[ $PR_LABELS == *"$RUN_ALL_LABEL"* ]]; then
            echo true
        else
            echo false
        fi
    else
        echo false  # not a PR or BUILDKITE_PULL_REQUEST not set
    fi
}

is_docs_only_change() {
    local file_path
    local has_any=0

    while IFS= read -r file_path; do
        [[ -z "${file_path}" ]] && continue
        has_any=1

        if [[ "${file_path}" == docs/* ]]; then
            continue
        fi
        if [[ "${file_path}" == *.md ]]; then
            continue
        fi
        if [[ "${file_path}" == "mkdocs.yml" ]]; then
            continue
        fi
        return 1
    done

    [[ "${has_any}" -eq 1 ]]
}

resolve_skip_ci() {
    local is_pr_build=0
    local files
    local base_branch base_ref

    if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
        is_pr_build=1
    fi

    if [[ "${is_pr_build}" -eq 1 ]]; then
        base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
        if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
            echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
            git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
        fi

        base_ref=""
        if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
            base_ref="origin/${base_branch}"
        elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
            base_ref="${base_branch}"
        else
            echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
            echo -n 0
            return 0
        fi

        if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
            echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
            echo -n 0
            return 0
        fi
    elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
        if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
            echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
            echo -n 0
            return 0
        fi
        if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
            echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
            echo -n 0
            return 0
        fi
    else
        echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
        echo -n 0
        return 0
    fi

    if is_docs_only_change <<< "${files}"; then
        echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
        echo -n 1
        return 0
    fi

    echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
    echo -n 0
}

if [[ -z "${COV_ENABLED:-}" ]]; then
    COV_ENABLED=0
fi

upload_pipeline() {
    echo "Uploading pipeline..."
    # Install minijinja
    ls .buildkite || buildkite-agent annotate --style error 'Please merge upstream main branch for buildkite CI'
    curl -sSfL https://github.com/mitsuhiko/minijinja/releases/download/2.3.1/minijinja-cli-installer.sh | sh
    source /var/lib/buildkite-agent/.cargo/env

    if [[ $BUILDKITE_PIPELINE_SLUG == "fastcheck" ]]; then
        AMD_MIRROR_HW="amdtentative"
    fi

    # Use local template file for vllm-omni
    cp .buildkite/test-template-amd-omni.j2 .buildkite/test-template.j2


    # (WIP) Use pipeline generator instead of jinja template
    if [ -e ".buildkite/pipeline_generator/pipeline_generator.py" ]; then
        python -m pip install click pydantic
        python .buildkite/pipeline_generator/pipeline_generator.py --run_all=$RUN_ALL --list_file_diff="$LIST_FILE_DIFF" --nightly="$NIGHTLY" --mirror_hw="$AMD_MIRROR_HW"
        buildkite-agent pipeline upload .buildkite/pipeline.yaml
        exit 0
    fi
    echo "List file diff: $LIST_FILE_DIFF"
    echo "Run all: $RUN_ALL"
    echo "Nightly: $NIGHTLY"
    echo "AMD Mirror HW: $AMD_MIRROR_HW"

    FAIL_FAST=$(fail_fast)

    cd .buildkite

    # Select test definition file: merge suite for main, ready suite for PRs.
    # For debugging, DEBUG_TEST_YAML can override the selection — accepts
    # "merge" or "ready" (case-insensitive).
    if [[ -n "${DEBUG_TEST_YAML:-}" ]]; then
        case "${DEBUG_TEST_YAML,,}" in
            merge)
                TEST_YAML="test-amd-merge.yml"
                ;;
            ready)
                TEST_YAML="test-amd-ready.yaml"
                ;;
            *)
                echo "ERROR: DEBUG_TEST_YAML must be 'merge' or 'ready', got '$DEBUG_TEST_YAML'" >&2
                exit 1
                ;;
        esac
        echo "DEBUG_TEST_YAML override: using $TEST_YAML"
    elif [[ $BUILDKITE_BRANCH == "main" ]]; then
        TEST_YAML="test-amd-merge.yml"
    else
        TEST_YAML="test-amd-ready.yaml"
    fi

    (
        set -x
        # Output pipeline.yaml with all blank lines removed
        minijinja-cli test-template.j2 "$TEST_YAML" \
            -D branch="$BUILDKITE_BRANCH" \
            -D list_file_diff="$LIST_FILE_DIFF" \
            -D run_all="$RUN_ALL" \
            -D nightly="$NIGHTLY" \
            -D mirror_hw="$AMD_MIRROR_HW" \
            -D fail_fast="$FAIL_FAST" \
            -D vllm_use_precompiled="$VLLM_USE_PRECOMPILED" \
            -D vllm_merge_base_commit="$(git merge-base origin/main HEAD)" \
            -D cov_enabled="$COV_ENABLED" \
            -D vllm_ci_branch="$VLLM_CI_BRANCH" \
            | sed '/^[[:space:]]*$/d' \
            > pipeline.yaml
    )
    cat pipeline.yaml
    buildkite-agent artifact upload pipeline.yaml
    buildkite-agent pipeline upload pipeline.yaml
    exit 0
}

get_diff() {
    $(git add .)
    echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
}

get_diff_main() {
    $(git add .)
    echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
}

file_diff=$(get_diff)
if [[ $BUILDKITE_BRANCH == "main" ]]; then
    file_diff=$(get_diff_main)
fi

# ----------------------------------------------------------------------
# Early exit start: skip pipeline if conditions are met
# ----------------------------------------------------------------------

# Match CUDA Buildkite skip-ci behavior for docs/markdown-only changes.
if [[ "${DOCS_ONLY_DISABLE}" != "1" ]]; then
    SKIP_CI="$(resolve_skip_ci)"
    if [[ "${SKIP_CI}" == "1" ]]; then
        echo "[docs-only] Docs/markdown-only changes detected. Exiting before pipeline upload."
        exit 0
    fi
fi

# ----------------------------------------------------------------------
# Early exit end
# ----------------------------------------------------------------------

patterns=(
    "docker/Dockerfile"
    "CMakeLists.txt"
    "requirements/common.txt"
    "requirements/cuda.txt"
    "requirements/build.txt"
    "requirements/test.txt"
    "setup.py"
    "csrc/"
    "cmake/"
)

ignore_patterns=(
    "docker/Dockerfile."
    "csrc/cpu"
    "csrc/rocm"
    "cmake/hipify.py"
    "cmake/cpu_extension.cmake"
)

for file in $file_diff; do
    # First check if file matches any pattern
    matches_pattern=0
    for pattern in "${patterns[@]}"; do
        if [[ $file == $pattern* ]] || [[ $file == $pattern ]]; then
            matches_pattern=1
            break
        fi
    done

    # If file matches pattern, check it's not in ignore patterns
    if [[ $matches_pattern -eq 1 ]]; then
        matches_ignore=0
        for ignore in "${ignore_patterns[@]}"; do
            if [[ $file == $ignore* ]] || [[ $file == $ignore ]]; then
                matches_ignore=1
                break
            fi
        done

        if [[ $matches_ignore -eq 0 ]]; then
            RUN_ALL=1
            echo "Found changes: $file. Run all tests"
            break
        fi
    fi
done

# Check for ready-run-all-tests label
LABEL_RUN_ALL=$(check_run_all_label)
if [[ $LABEL_RUN_ALL == true ]]; then
    RUN_ALL=1
    NIGHTLY=1
    echo "Found 'ready-run-all-tests' label. Running all tests including optional tests."
fi

# Decide whether to use precompiled wheels
# Relies on existing patterns array as a basis.
if [[ -n "${VLLM_USE_PRECOMPILED:-}" ]]; then
    echo "VLLM_USE_PRECOMPILED is already set to: $VLLM_USE_PRECOMPILED"
elif [[ $RUN_ALL -eq 1 ]]; then
    export VLLM_USE_PRECOMPILED=0
    echo "Detected critical changes, building wheels from source"
else
    export VLLM_USE_PRECOMPILED=1
    echo "No critical changes, using precompiled wheels"
fi


LIST_FILE_DIFF=$(get_diff | tr ' ' '|')
if [[ $BUILDKITE_BRANCH == "main" ]]; then
    LIST_FILE_DIFF=$(get_diff_main | tr ' ' '|')
fi
upload_pipeline