vLLM/.github/workflows/scripts/run_ci_command.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import http.client
import json
import os
import random
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from collections.abc import Callable, Mapping, Sequence
from typing import Any

COMMAND_RUN_CI = "/ci run"
COMMAND_RUN_CI_ALL = "/ci run all"
COMMAND_RUN_CI_NIGHTLY = "/ci run nightly"
COMMAND_RETRY_FAILED = "/ci retry"
COMMAND_CANCEL_CI = "/ci cancel"
COMMAND_RUN_AMD_CI = "/amd-ci run"
COMMAND_RUN_AMD_CI_ALL = "/amd-ci run all"
COMMAND_RUN_AMD_CI_NIGHTLY = "/amd-ci run nightly"
COMMAND_RETRY_AMD_FAILED = "/amd-ci retry"
COMMAND_CANCEL_AMD_CI = "/amd-ci cancel"
ALLOW_STALE_SUFFIX = " --allow-stale"
RUN_CI_COMMAND_ENV = {
    COMMAND_RUN_CI: {},
    COMMAND_RUN_CI_ALL: {"RUN_ALL": "1"},
    COMMAND_RUN_CI_NIGHTLY: {"RUN_ALL": "1", "NIGHTLY": "1"},
    COMMAND_RUN_AMD_CI: {},
    COMMAND_RUN_AMD_CI_ALL: {"RUN_ALL": "1"},
    COMMAND_RUN_AMD_CI_NIGHTLY: {"RUN_ALL": "1", "NIGHTLY": "1"},
}
ALLOW_STALE_COMMANDS = frozenset(
    f"{command}{ALLOW_STALE_SUFFIX}" for command in RUN_CI_COMMAND_ENV
)
RUN_CI_COMMAND_ENV.update(
    {
        command: RUN_CI_COMMAND_ENV[command.removesuffix(ALLOW_STALE_SUFFIX)]
        for command in ALLOW_STALE_COMMANDS
    }
)
UPSTREAM_CI_COMMANDS = frozenset(
    {
        *(command for command in RUN_CI_COMMAND_ENV if command.startswith("/ci ")),
        COMMAND_RETRY_FAILED,
        COMMAND_CANCEL_CI,
    }
)
AMD_CI_COMMANDS = frozenset(
    {
        *(command for command in RUN_CI_COMMAND_ENV if command.startswith("/amd-ci ")),
        COMMAND_RETRY_AMD_FAILED,
        COMMAND_CANCEL_AMD_CI,
    }
)
RETRY_COMMANDS = frozenset({COMMAND_RETRY_FAILED, COMMAND_RETRY_AMD_FAILED})
CANCEL_COMMANDS = frozenset({COMMAND_CANCEL_CI, COMMAND_CANCEL_AMD_CI})
ALL_CI_COMMANDS = UPSTREAM_CI_COMMANDS | AMD_CI_COMMANDS
CI_COMMAND_PREFIXES = frozenset({"/ci", "/amd-ci"})
CI_AUTHORIZED_COMMENT_MARKER = "<!-- vllm-ci-authorized -->"
READY_LABELS = {"ready", "ready-run-all-tests"}
TRUSTED_PERMISSIONS = {"admin", "maintain", "write"}
BRANCH_UPDATE_POLL_SECONDS = 3
BRANCH_UPDATE_POLL_ATTEMPTS = 40
ACTIVE_BUILD_STATES = {
    "blocked",
    "creating",
    "scheduled",
    "running",
    "failing",
    "canceling",
    "waiting",
    "waiting_failed",
}
RETRY_STATES = "failed,timed_out,expired"
CANCELABLE_BUILD_STATES = ("scheduled", "running", "failing")
SETUP_STEP_KEYS = {
    "ensure-ci-base-amd",
    "pre-commit",
    "refresh-rocm-base-amd",
}


class ApiError(RuntimeError):
    def __init__(self, status: int | None, message: str) -> None:
        super().__init__(message)
        self.status = status


class CiPreparationError(RuntimeError):
    pass


def rate_limit_jitter() -> float:
    return random.uniform(1, 5)


class HttpTransport:
    def __init__(
        self,
        *,
        max_retries: int = 3,
        jitter: Callable[[], float] = rate_limit_jitter,
        sleep: Callable[[float], None] = time.sleep,
    ) -> None:
        self.max_retries = max_retries
        self.jitter = jitter
        self.sleep = sleep

    def request(
        self,
        url: str,
        *,
        body: Mapping[str, Any] | None = None,
        headers: Mapping[str, str] | None = None,
        method: str = "GET",
    ) -> Any:
        data = None if body is None else json.dumps(body).encode()
        request = urllib.request.Request(
            url,
            data=data,
            headers=dict(headers or {}),
            method=method,
        )
        for attempt in range(self.max_retries + 1):
            try:
                with urllib.request.urlopen(request, timeout=30) as response:
                    response_body = response.read().decode()
                break
            except urllib.error.HTTPError as error:
                response_body = error.read().decode()
                if error.code == 429 and attempt < self.max_retries:
                    delay = self._rate_limit_delay(error, response_body)
                    print(
                        "API rate limit reached; "
                        f"retry {attempt + 1}/{self.max_retries} "
                        f"in {delay:g} seconds.",
                        file=sys.stderr,
                    )
                    self.sleep(delay)
                    continue
                message = self._error_message(response_body, error.reason)
                raise ApiError(
                    error.code,
                    f"API returned {error.code}: {message}",
                ) from error
            except urllib.error.URLError as error:
                raise ApiError(
                    None,
                    f"API request failed: {error.reason}",
                ) from error
            except (OSError, http.client.HTTPException) as error:
                raise ApiError(None, f"API request failed: {error}") from error

        if not response_body:
            return None
        try:
            return json.loads(response_body)
        except json.JSONDecodeError as error:
            raise ApiError(None, "API returned a non-JSON response.") from error

    @staticmethod
    def _error_message(response_body: str, fallback: str) -> str:
        try:
            parsed = json.loads(response_body)
        except json.JSONDecodeError:
            return fallback
        return str(parsed.get("message", fallback))

    def _rate_limit_delay(
        self,
        error: urllib.error.HTTPError,
        response_body: str,
    ) -> float:
        try:
            parsed = json.loads(response_body)
        except json.JSONDecodeError:
            parsed = {}

        scope = parsed.get("scope")
        reset_header = (
            "RateLimit-User-Reset" if scope == "rest_user" else "RateLimit-Reset"
        )
        candidates = [
            error.headers.get(reset_header),
            parsed.get("reset"),
        ]
        for candidate in candidates:
            try:
                delay = float(candidate)
            except (TypeError, ValueError):
                continue
            if delay >= 0:
                return delay + self.jitter()
        return 60 + self.jitter()


class GitHubClient:
    def __init__(
        self,
        token: str,
        repository: str,
        transport: HttpTransport | None = None,
    ) -> None:
        if not token:
            raise RuntimeError("GH_TOKEN is not set.")
        self.owner, self.repo = repository.split("/", maxsplit=1)
        self.transport = transport or HttpTransport()
        self.headers = {
            "Accept": "application/vnd.github+json",
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
            "User-Agent": "vllm-ci-command",
            "X-GitHub-Api-Version": "2022-11-28",
        }

    def _request(
        self,
        path: str,
        *,
        body: Mapping[str, Any] | None = None,
        method: str = "GET",
    ) -> Any:
        return self.transport.request(
            f"https://api.github.com{path}",
            body=body,
            headers=self.headers,
            method=method,
        )

    def _repo_path(self, suffix: str) -> str:
        owner = urllib.parse.quote(self.owner, safe="")
        repo = urllib.parse.quote(self.repo, safe="")
        return f"/repos/{owner}/{repo}{suffix}"

    def _paginate(self, path: str) -> list[dict[str, Any]]:
        results: list[dict[str, Any]] = []
        separator = "&" if "?" in path else "?"
        for page in range(1, 101):
            response = self._request(f"{path}{separator}per_page=100&page={page}")
            if not isinstance(response, list):
                raise ApiError(None, "GitHub API returned an invalid list response.")
            results.extend(response)
            if len(response) < 100:
                return results
        raise ApiError(None, "GitHub API pagination exceeded 10,000 results.")

    def get_pr(self, number: int) -> dict[str, Any]:
        return self._request(self._repo_path(f"/pulls/{number}"))

    def get_commits_behind_base(self, base_ref: str, head_sha: str) -> int:
        base_ref = urllib.parse.quote(f"refs/heads/{base_ref}", safe="")
        head_sha = urllib.parse.quote(head_sha, safe="")
        response = self._request(
            self._repo_path(f"/compare/{base_ref}...{head_sha}?per_page=1")
        )
        behind = response.get("behind_by") if isinstance(response, dict) else None
        if type(behind) is not int or behind < 0:
            raise ApiError(None, "GitHub returned an invalid commits-behind count.")
        return behind

    def update_branch(self, number: int, expected_head_sha: str) -> None:
        self._request(
            self._repo_path(f"/pulls/{number}/update-branch"),
            body={"expected_head_sha": expected_head_sha},
            method="PUT",
        )

    def list_pulls_for_commit(self, commit: str) -> list[dict[str, Any]]:
        commit = urllib.parse.quote(commit, safe="")
        response = self._request(self._repo_path(f"/commits/{commit}/pulls"))
        if not isinstance(response, list):
            raise ApiError(None, "GitHub API returned an invalid pull request list.")
        return response

    def get_permission(self, actor: str) -> str:
        username = urllib.parse.quote(actor, safe="")
        try:
            response = self._request(
                self._repo_path(f"/collaborators/{username}/permission")
            )
        except ApiError as error:
            if error.status == 404:
                return "none"
            raise
        return str(response["permission"])

    def get_review_decision(self, number: int) -> str | None:
        query = """
          query($owner: String!, $repo: String!, $number: Int!) {
            repository(owner: $owner, name: $repo) {
              pullRequest(number: $number) {
                reviewDecision
              }
            }
          }
        """
        response = self._request(
            "/graphql",
            body={
                "query": query,
                "variables": {
                    "number": number,
                    "owner": self.owner,
                    "repo": self.repo,
                },
            },
            method="POST",
        )
        return response["data"]["repository"]["pullRequest"]["reviewDecision"]

    def list_reviews(self, number: int) -> list[dict[str, Any]]:
        return self._paginate(self._repo_path(f"/pulls/{number}/reviews"))

    def list_issue_comments(self, number: int) -> list[dict[str, Any]]:
        return self._paginate(self._repo_path(f"/issues/{number}/comments"))

    def list_reactions(self, comment_id: int) -> list[dict[str, Any]]:
        return self._paginate(
            self._repo_path(f"/issues/comments/{comment_id}/reactions")
        )

    def add_reaction(self, comment_id: int, content: str) -> None:
        self._request(
            self._repo_path(f"/issues/comments/{comment_id}/reactions"),
            body={"content": content},
            method="POST",
        )

    def add_comment(self, issue_number: int, body: str) -> None:
        self._request(
            self._repo_path(f"/issues/{issue_number}/comments"),
            body={"body": body},
            method="POST",
        )


class BuildkiteClient:
    def __init__(
        self,
        token: str,
        organization: str,
        pipeline: str,
        transport: HttpTransport | None = None,
    ) -> None:
        self.token = token
        self.transport = transport or HttpTransport()
        organization = urllib.parse.quote(organization, safe="")
        pipeline = urllib.parse.quote(pipeline, safe="")
        self.base_url = (
            "https://api.buildkite.com/v2/organizations/"
            f"{organization}/pipelines/{pipeline}/builds"
        )

    def _headers(self) -> dict[str, str]:
        if not self.token:
            raise RuntimeError("The BUILDKITE_API_TOKEN repository secret is not set.")
        return {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json",
            "User-Agent": "vllm-ci-command",
        }

    def _request_url(
        self,
        url: str,
        *,
        body: Mapping[str, Any] | None = None,
        method: str = "GET",
    ) -> Any:
        if (
            url != self.base_url
            and not url.startswith(f"{self.base_url}?")
            and not url.startswith(f"{self.base_url}/")
        ):
            raise ApiError(None, "Buildkite API returned an invalid pagination URL.")
        return self.transport.request(
            url,
            body=body,
            headers=self._headers(),
            method=method,
        )

    def _request(
        self,
        *,
        body: Mapping[str, Any] | None = None,
        method: str = "GET",
        path: str = "",
        query: Sequence[tuple[str, str]] = (),
    ) -> Any:
        url = f"{self.base_url}{path}"
        if query:
            url = f"{url}?{urllib.parse.urlencode(query)}"
        return self._request_url(url, body=body, method=method)

    def list_builds(
        self,
        commit: str | None,
        *,
        branch: str | None = None,
        metadata: tuple[str, str] | None = None,
        states: Sequence[str] = (),
    ) -> list[dict[str, Any]]:
        query = [
            ("exclude_jobs", "true"),
            ("exclude_pipeline", "true"),
            ("per_page", "100"),
        ]
        if commit:
            query.append(("commit", commit))
        if branch:
            query.append(("branch", branch))
        if metadata:
            key, value = metadata
            query.append((f"meta_data[{key}]", value))
        query.extend(("state[]", state) for state in states)
        response = self._request(query=query)
        if not isinstance(response, list):
            raise ApiError(None, "Buildkite API returned an invalid build list.")
        return response

    def create_build(self, body: Mapping[str, Any]) -> dict[str, Any]:
        return self._request(body=body, method="POST")

    def retry_failed_jobs(
        self,
        build_number: int,
        states: str,
    ) -> dict[str, Any]:
        number = urllib.parse.quote(str(build_number), safe="")
        return self._request(
            body={"states": states},
            method="PUT",
            path=f"/{number}/retry_failed_jobs",
        )

    def cancel_build(self, build_number: int) -> dict[str, Any]:
        number = urllib.parse.quote(str(build_number), safe="")
        return self._request(method="PUT", path=f"/{number}/cancel")

    def list_failed_jobs(self, build_number: int) -> list[dict[str, Any]]:
        number = urllib.parse.quote(str(build_number), safe="")
        query = [
            ("state[]", "failed"),
            ("state[]", "timed_out"),
            ("state[]", "expired"),
            ("include_retried_jobs", "false"),
            ("per_page", "100"),
        ]
        url = f"{self.base_url}/{number}/jobs?{urllib.parse.urlencode(query)}"
        jobs: list[dict[str, Any]] = []
        while url:
            response = self._request_url(url)
            if not isinstance(response, Mapping):
                raise ApiError(None, "Buildkite API returned an invalid job list.")
            items = response.get("items")
            links = response.get("links")
            if not isinstance(items, list) or not isinstance(links, Mapping):
                raise ApiError(None, "Buildkite API returned an invalid job list.")
            jobs.extend(items)
            next_url = links.get("next")
            if next_url is not None and not isinstance(next_url, str):
                raise ApiError(
                    None, "Buildkite API returned an invalid pagination URL."
                )
            url = next_url
        return jobs


def parse_command(body: str) -> str | None:
    body = body.strip()
    if body in ALL_CI_COMMANDS:
        return body
    return None


def is_ci_command_attempt(body: str) -> bool:
    words = body.split(maxsplit=1)
    return bool(words) and words[0].casefold() in CI_COMMAND_PREFIXES


def pipeline_for_command(
    command: str,
    *,
    amd_ci_pipeline: str = "amd-ci",
    upstream_ci_pipeline: str = "ci",
) -> str:
    if command in AMD_CI_COMMANDS:
        return amd_ci_pipeline
    if command in UPSTREAM_CI_COMMANDS:
        return upstream_ci_pipeline
    raise ValueError(f"Unsupported CI command: {command}")


def ci_name_for_command(command: str) -> str:
    if command in AMD_CI_COMMANDS:
        return "AMD CI"
    if command in UPSTREAM_CI_COMMANDS:
        return "CI"
    raise ValueError(f"Unsupported CI command: {command}")


def run_command_for_command(command: str) -> str:
    if command in AMD_CI_COMMANDS:
        return COMMAND_RUN_AMD_CI
    if command in UPSTREAM_CI_COMMANDS:
        return COMMAND_RUN_CI
    raise ValueError(f"Unsupported CI command: {command}")


def retry_command_for_command(command: str) -> str:
    if command in AMD_CI_COMMANDS:
        return COMMAND_RETRY_AMD_FAILED
    if command in UPSTREAM_CI_COMMANDS:
        return COMMAND_RETRY_FAILED
    raise ValueError(f"Unsupported CI command: {command}")


def parse_trusted_users(value: str = "") -> set[str]:
    return {
        user.casefold() for item in value.split(",") for user in item.split() if user
    }


def has_ready_label(pr: Mapping[str, Any]) -> bool:
    return any(label["name"] in READY_LABELS for label in pr["labels"])


def is_trusted_permission(permission: str) -> bool:
    return permission in TRUSTED_PERMISSIONS


def authorize(
    *,
    actor: str,
    permission: str,
    pr: Mapping[str, Any],
    run_command: str = COMMAND_RUN_CI,
    trusted_approval: bool = False,
    trusted_users: set[str] | None = None,
) -> tuple[bool, str]:
    trusted_users = trusted_users or set()
    if is_trusted_permission(permission):
        return True, f"repository {permission} permission"
    if actor.casefold() in trusted_users:
        return True, "configured trusted contributor"
    if actor.casefold() != pr["user"]["login"].casefold():
        return (
            False,
            "Only reviewers with write access can use CI commands before CI is "
            "delegated to the PR author.",
        )
    if pr["draft"]:
        return False, "PR authors cannot run CI while the PR is a draft."
    if has_ready_label(pr):
        return True, "ready label"
    if trusted_approval:
        return True, "approval from a trusted reviewer"
    return (
        False,
        f"A reviewer with write access must run `{run_command}`, approve the PR, "
        "or add the `ready` label first.",
    )


def has_trusted_approval(
    github: GitHubClient,
    number: int,
    trusted_users: set[str],
) -> bool:
    if github.get_review_decision(number) != "APPROVED":
        return False

    latest_review_states: dict[str, tuple[str, str]] = {}
    for review in github.list_reviews(number):
        user = review.get("user") or {}
        login = user.get("login")
        state = review.get("state")
        if login and state in {"APPROVED", "CHANGES_REQUESTED", "DISMISSED"}:
            latest_review_states[login.casefold()] = (login, state)

    for login, state in latest_review_states.values():
        if state != "APPROVED":
            continue
        if login.casefold() in trusted_users:
            return True
        if is_trusted_permission(github.get_permission(login)):
            return True
    return False


def is_build_for_pr(build: Mapping[str, Any], pr_number: int) -> bool:
    pull_request = build.get("pull_request")
    if isinstance(pull_request, Mapping):
        build_pr_number = pull_request.get("id", pull_request.get("number"))
        if build_pr_number is not None:
            return str(build_pr_number) == str(pr_number)
    metadata = build.get("meta_data") or {}
    return str(metadata.get("github-pr-number")) == str(pr_number)


def is_active_build(build: Mapping[str, Any]) -> bool:
    return bool(build.get("blocked")) or build.get("state") in ACTIVE_BUILD_STATES


def is_comment_triggered_build(build: Mapping[str, Any]) -> bool:
    metadata = build.get("meta_data") or {}
    return bool(metadata.get("github-comment-id"))


def blocks_new_run(command: str, build: Mapping[str, Any]) -> bool:
    if command in AMD_CI_COMMANDS:
        return (
            is_comment_triggered_build(build)
            and build.get("state") in ACTIVE_BUILD_STATES
            and build.get("state") != "blocked"
        )
    return is_active_build(build)


def select_latest_build(
    builds: Sequence[dict[str, Any]],
    pr_number: int,
) -> dict[str, Any] | None:
    matching = [build for build in builds if is_build_for_pr(build, pr_number)]
    return max(matching, key=lambda build: build.get("created_at", ""), default=None)


def create_build_payload(
    *,
    actor: str,
    comment_id: int,
    command: str = COMMAND_RUN_CI,
    pr: Mapping[str, Any],
) -> dict[str, Any]:
    if command not in RUN_CI_COMMAND_ENV:
        raise ValueError(f"Unsupported run command: {command}")

    env = {
        "VLLM_CI_GITHUB_COMMENT_ID": str(comment_id),
        "VLLM_CI_TRIGGERED_BY": actor,
        **RUN_CI_COMMAND_ENV[command],
    }
    return {
        "commit": pr["head"]["sha"],
        # Use the head label (owner:branch) rather than the bare ref so a PR
        # whose fork branch is named e.g. "main" does not collide with the
        # upstream branch of the same name.
        "branch": pr["head"].get("label") or pr["head"]["ref"],
        "message": f"PR #{pr['number']} {command} by @{actor}",
        "pull_request_id": pr["number"],
        "pull_request_base_branch": pr["base"]["ref"],
        "pull_request_repository": pr["head"]["repo"]["clone_url"],
        "pull_request_labels": [label["name"] for label in pr["labels"]],
        "ignore_pipeline_branch_filters": True,
        "env": env,
        "meta_data": {
            "github-comment-id": str(comment_id),
            "github-pr-number": str(pr["number"]),
            "github-triggered-by": actor,
        },
    }


def create_retry_build_payload(
    *,
    actor: str,
    comment_id: int,
    command: str = COMMAND_RETRY_FAILED,
    pr: Mapping[str, Any],
    source_build: Mapping[str, Any],
    step_keys: Sequence[str],
) -> dict[str, Any]:
    if command not in RETRY_COMMANDS:
        raise ValueError(f"Unsupported retry command: {command}")
    payload = create_build_payload(
        actor=actor,
        comment_id=comment_id,
        command=run_command_for_command(command),
        pr=pr,
    )
    source_number = str(source_build["number"])
    payload["message"] = f"PR #{pr['number']} {command} by @{actor}"
    payload["env"]["VLLM_CI_ONLY_STEP_KEYS"] = json.dumps(
        step_keys, separators=(",", ":")
    )
    payload["meta_data"].update(
        {
            "github-retry-source-build": source_number,
            "github-retry-source-commit": str(source_build.get("commit", "")),
        }
    )
    return payload


def add_reaction_safely(
    github: GitHubClient,
    comment_id: int,
    content: str,
) -> None:
    try:
        github.add_reaction(comment_id, content)
    except Exception as error:
        print(f"Could not add {content} reaction: {error}", file=sys.stderr)


def command_comment_marker(comment_id: int) -> str:
    return f"<!-- vllm-ci-command:{comment_id} -->"


def has_bot_comment_marker(
    github: GitHubClient,
    issue_number: int,
    marker: str,
) -> bool:
    return any(
        marker in str(comment.get("body", ""))
        and (comment.get("user") or {}).get("login") == "github-actions[bot]"
        for comment in github.list_issue_comments(issue_number)
    )


def is_already_handled(
    github: GitHubClient,
    issue_number: int,
    comment_id: int,
) -> bool:
    terminal_reaction = any(
        reaction.get("content") in {"rocket", "-1"}
        and (reaction.get("user") or {}).get("login") == "github-actions[bot]"
        for reaction in github.list_reactions(comment_id)
    )
    if terminal_reaction:
        return True
    return has_bot_comment_marker(
        github,
        issue_number,
        command_comment_marker(comment_id),
    )


def notify_authorized(
    event: Mapping[str, Any],
    github: GitHubClient,
    trusted_users_value: str = "",
) -> None:
    pr = event["pull_request"]
    if pr["state"] != "open" or pr["draft"]:
        return

    trusted_users = parse_trusted_users(trusted_users_value)
    author = pr["user"]["login"]
    author_permission = github.get_permission(author)
    if (
        is_trusted_permission(author_permission)
        or author.casefold() in trusted_users
        or has_bot_comment_marker(
            github,
            pr["number"],
            CI_AUTHORIZED_COMMENT_MARKER,
        )
    ):
        return

    if "label" in event:
        if (
            event.get("action") != "labeled"
            or event["label"]["name"] not in READY_LABELS
        ):
            return
    elif "review" in event:
        if (
            event.get("action") != "submitted"
            or str(event["review"].get("state", "")).casefold() != "approved"
            or has_ready_label(pr)
            or not has_trusted_approval(github, pr["number"], trusted_users)
        ):
            return
    else:
        return

    github.add_comment(
        pr["number"],
        (
            f"✅ @{author}, CI is now available for this PR.\n\n"
            "- `/ci run` starts upstream CI; `/amd-ci run` starts AMD CI only.\n"
            "- Run commands test your branch with every commit currently on its "
            "upstream target branch. If it is behind, the bot merges the target "
            "branch into it first (pull before pushing again). Append "
            "`--allow-stale` to a run command to test an outdated branch as is, "
            "at your own risk.\n"
            "- `/ci retry` retries failed jobs in the CI build for the current "
            "PR head. If the current head has no CI build, it starts a new CI "
            "build for the current head containing only jobs that failed in "
            "the latest earlier CI build for this PR.\n"
            "- `/amd-ci retry` retries failed jobs in AMD CI for the current PR "
            "head. Use `/amd-ci run` when the current head has no AMD CI build.\n"
            "- `/ci cancel` cancels scheduled or running CI builds for this PR "
            "branch; `/amd-ci cancel` does the same for AMD CI only.\n\n"
            f"{CI_AUTHORIZED_COMMENT_MARKER}"
        ),
    )


def resolve_workflow_run_pr(
    workflow_run: Mapping[str, Any],
    github: GitHubClient,
) -> dict[str, Any] | None:
    head_sha = str(workflow_run.get("head_sha", ""))
    if not head_sha:
        return None

    seen: set[int] = set()

    def find_matching_pr(
        candidates: Sequence[Mapping[str, Any]],
    ) -> dict[str, Any] | None:
        for candidate in candidates:
            candidate_number = candidate.get("number")
            if not isinstance(candidate_number, int) or candidate_number in seen:
                continue
            seen.add(candidate_number)
            try:
                pr = github.get_pr(candidate_number)
            except ApiError as error:
                if error.status == 404:
                    continue
                raise
            if pr["state"] == "open" and pr["head"]["sha"] == head_sha:
                return pr
        return None

    associated_pr = find_matching_pr(workflow_run.get("pull_requests") or [])
    if associated_pr is not None:
        return associated_pr
    return find_matching_pr(github.list_pulls_for_commit(head_sha))


def merge_base_into_pr(
    github: GitHubClient,
    branch_updater: GitHubClient,
    pr: Mapping[str, Any],
    command: str,
    lag: str,
) -> tuple[dict[str, Any], str]:
    base_ref = pr["base"]["ref"]

    def merge_failed(reason: str) -> CiPreparationError:
        return CiPreparationError(
            f"{lag} The bot could not merge upstream `{base_ref}` into this "
            f"branch ({reason}). No new CI build was started. Either merge or "
            f"rebase onto the latest `{base_ref}` yourself and rerun `{command}`, "
            f"or use `{command}{ALLOW_STALE_SUFFIX}` to test this branch as is."
        )

    try:
        branch_updater.update_branch(pr["number"], pr["head"]["sha"])
        for _ in range(BRANCH_UPDATE_POLL_ATTEMPTS):
            time.sleep(BRANCH_UPDATE_POLL_SECONDS)
            current_pr = github.get_pr(pr["number"])
            if current_pr["head"]["sha"] != pr["head"]["sha"]:
                break
        else:
            raise merge_failed("the merge commit did not appear in time")
        if (
            current_pr["state"] != "open"
            or current_pr["base"]["ref"] != base_ref
            or github.get_commits_behind_base(base_ref, current_pr["head"]["sha"])
        ):
            raise merge_failed("the PR changed during the merge")
    except ApiError as error:
        raise merge_failed(str(error)) from error

    note = (
        f"ℹ️ {lag} Merged upstream `{base_ref}` into this branch as "
        f"`{current_pr['head']['sha'][:12]}` before starting CI. "
        "Pull the branch before pushing more changes."
    )
    print(note)
    return current_pr, note


def prepare_pr_for_ci(
    github: GitHubClient,
    pr: Mapping[str, Any],
    command: str,
    branch_updater: GitHubClient | None = None,
) -> tuple[dict[str, Any], str]:
    base_ref = pr["base"]["ref"]
    no_action = "No new CI build was started."
    try:
        behind = github.get_commits_behind_base(base_ref, pr["head"]["sha"])
        current_pr = github.get_pr(pr["number"])
    except ApiError as error:
        raise CiPreparationError(
            f"Could not check the PR against upstream `{base_ref}`. "
            f"{no_action} Comment `{command}` again. {error}"
        ) from error
    if (
        current_pr["state"] != "open"
        or current_pr["head"]["sha"] != pr["head"]["sha"]
        or current_pr["base"]["ref"] != pr["base"]["ref"]
    ):
        raise CiPreparationError(
            f"The PR changed while checking its branch. {no_action} "
            f"Comment `{command}` again."
        )
    commits = "commit" if behind == 1 else "commits"
    lag = f"This PR is {behind} {commits} behind upstream `{base_ref}`."
    if command in ALLOW_STALE_COMMANDS:
        warning = ""
        if behind:
            run_command = command.removesuffix(ALLOW_STALE_SUFFIX)
            warning = (
                f"⚠️ {lag} Running CI at your own risk because "
                "`--allow-stale` was requested; outdated CI configuration may "
                "cause failures. Before merging, merge or rebase onto the latest "
                f"`{base_ref}`, then rerun `{run_command}` on the latest PR commit."
            )
            print(warning)
        return current_pr, warning

    if behind and branch_updater is not None:
        return merge_base_into_pr(github, branch_updater, current_pr, command, lag)
    if behind:
        raise CiPreparationError(
            f"{lag} Your branch must contain every commit currently on upstream "
            f"`{base_ref}`. {no_action} "
            f"Merge or rebase onto the latest `{base_ref}`, then rerun `{command}`."
            " To test this branch at your own risk, "
            f"use `{command}{ALLOW_STALE_SUFFIX}`."
        )
    return current_pr, ""


def handle_run_ci(
    *,
    actor: str,
    buildkite: BuildkiteClient,
    comment_id: int,
    command: str,
    github: GitHubClient,
    pr: Mapping[str, Any],
    branch_updater: GitHubClient | None = None,
    allow_repeat_run: bool = False,
) -> str:
    ci_name = ci_name_for_command(command)
    duplicate_builds = buildkite.list_builds(
        pr["head"]["sha"],
        metadata=("github-comment-id", str(comment_id)),
    )
    duplicate = select_latest_build(duplicate_builds, pr["number"])
    if duplicate:
        return (
            f"{ci_name} was already requested by this comment: {duplicate['web_url']}"
        )

    current_builds = buildkite.list_builds(pr["head"]["sha"])
    active_build = next(
        (
            build
            for build in current_builds
            if is_build_for_pr(build, pr["number"]) and blocks_new_run(command, build)
        ),
        None,
    )
    if active_build:
        return (
            f"{ci_name} is already running for this commit: {active_build['web_url']}"
        )

    tested_build = select_latest_build(
        [build for build in current_builds if build.get("state") != "canceled"],
        pr["number"],
    )
    if tested_build and not allow_repeat_run:
        raise CiPreparationError(
            f"{ci_name} already ran on this commit: {tested_build['web_url']}. "
            "No new CI build was started. Push new changes to run CI again, use "
            f"`{retry_command_for_command(command)}` to rerun failed jobs, or ask "
            "a reviewer with write access to rerun it."
        )

    current_pr, warning = prepare_pr_for_ci(github, pr, command, branch_updater)

    build = buildkite.create_build(
        create_build_payload(
            actor=actor,
            comment_id=comment_id,
            command=command,
            pr=current_pr,
        )
    )
    return (
        f"Triggered [Buildkite {ci_name} #{build['number']}]({build['web_url']}) "
        f"for commit `{current_pr['head']['sha'][:12]}`."
        + (f"\n\n{warning}" if warning else "")
    )


def handle_retry_failed(
    *,
    actor: str,
    buildkite: BuildkiteClient,
    comment_id: int,
    command: str,
    github: GitHubClient,
    pr: Mapping[str, Any],
) -> str:
    ci_name = ci_name_for_command(command)
    run_command = run_command_for_command(command)
    retry_command = retry_command_for_command(command)
    builds = buildkite.list_builds(pr["head"]["sha"])
    build = select_latest_build(builds, pr["number"])
    if build:
        metadata = build.get("meta_data") or {}
        if str(metadata.get("github-comment-id")) == str(comment_id):
            return (
                f"{ci_name} was already requested by this comment: {build['web_url']}"
            )

        retried = buildkite.retry_failed_jobs(build["number"], RETRY_STATES)
        if retried["retried_jobs_count"] == 0:
            return (
                "No failed, timed-out, or expired jobs need retrying: "
                f"{build['web_url']}"
            )
        return (
            f"Queued {retried['retried_jobs_count']} failed job(s) for retry in "
            f"[Buildkite {ci_name} #{build['number']}]({build['web_url']})."
        )

    if command == COMMAND_RETRY_AMD_FAILED:
        return (
            "No AMD CI build exists for the current PR head. "
            f"Use `{run_command}` first."
        )

    previous_builds = buildkite.list_builds(
        None,
        metadata=("github-pr-number", str(pr["number"])),
    )
    previous_builds = [
        candidate
        for candidate in previous_builds
        if candidate.get("commit") != pr["head"]["sha"]
    ]
    source_build = select_latest_build(previous_builds, pr["number"])
    if not source_build:
        return (
            f"No earlier {ci_name} build exists for this PR. Use `{run_command}` first."
        )
    if not source_build.get("finished_at") or is_active_build(source_build):
        return (
            f"The previous {ci_name} build is still running: {source_build['web_url']}"
        )

    failed_jobs = buildkite.list_failed_jobs(source_build["number"])
    failed_script_jobs = [job for job in failed_jobs if job.get("type") == "script"]
    missing_step_keys = [job for job in failed_script_jobs if not job.get("step_key")]
    if missing_step_keys:
        return (
            f"[Buildkite {ci_name} #{source_build['number']}]"
            f"({source_build['web_url']}) has failed jobs without stable step "
            "keys, so they cannot be retried on a new commit. "
            f"Use `{run_command}`."
        )

    failed_step_keys = {str(job["step_key"]) for job in failed_script_jobs}
    setup_failures = sorted(
        step_key
        for step_key in failed_step_keys
        if step_key.startswith("image-build") or step_key in SETUP_STEP_KEYS
    )
    if setup_failures:
        return (
            f"[Buildkite {ci_name} #{source_build['number']}]"
            f"({source_build['web_url']}) failed during CI setup, so its test "
            f"failure set is incomplete. Use `{run_command}` for the new commit."
        )

    step_keys = sorted(failed_step_keys)
    if not step_keys:
        return (
            "No failed, timed-out, or expired jobs need retrying in "
            f"[Buildkite {ci_name} #{source_build['number']}]"
            f"({source_build['web_url']})."
        )

    current_pr = github.get_pr(pr["number"])
    if current_pr["state"] != "open" or current_pr["head"]["sha"] != pr["head"]["sha"]:
        return (
            "The PR head changed while processing the command. "
            f"Comment `{retry_command}` again."
        )

    retry_build = buildkite.create_build(
        create_retry_build_payload(
            actor=actor,
            comment_id=comment_id,
            command=command,
            pr=current_pr,
            source_build=source_build,
            step_keys=step_keys,
        )
    )
    return (
        f"Triggered [Buildkite {ci_name} #{retry_build['number']}]"
        f"({retry_build['web_url']}) for commit "
        f"`{current_pr['head']['sha'][:12]}`, running {len(step_keys)} failed "
        f"step(s) from [Buildkite {ci_name} #{source_build['number']}]"
        f"({source_build['web_url']})."
    )


def handle_cancel_ci(
    *,
    buildkite: BuildkiteClient,
    command: str = COMMAND_CANCEL_CI,
    pr: Mapping[str, Any],
) -> str:
    ci_name = ci_name_for_command(command)
    branch = pr["head"]["ref"]
    # Builds are created under the head label (owner:branch); also match the
    # bare ref so builds created before that change can still be cancelled.
    branches = [branch]
    head_label = pr["head"].get("label")
    if head_label and head_label not in branches:
        branches.insert(0, str(head_label))
    states = CANCELABLE_BUILD_STATES

    builds_by_number: dict[int, dict[str, Any]] = {}
    for candidate_branch in branches:
        for build in buildkite.list_builds(
            None,
            branch=candidate_branch,
            states=states,
        ):
            builds_by_number[build["number"]] = build

    cancelable_builds = [
        build
        for build in builds_by_number.values()
        if build.get("branch") in branches
        and is_build_for_pr(build, pr["number"])
        and build.get("state") in states
    ]
    if not cancelable_builds:
        return f"No cancelable {ci_name} build is running for branch `{branch}`."

    for build in cancelable_builds:
        buildkite.cancel_build(build["number"])

    links = ", ".join(
        f"[#{build['number']}]({build['web_url']})" for build in cancelable_builds
    )
    count = len(cancelable_builds)
    noun = "build" if count == 1 else "builds"
    return (
        f"Requested cancellation of {count} {ci_name} {noun} for `{branch}`: {links}."
    )


def reply_unrecognized_command(
    event: Mapping[str, Any],
    github: GitHubClient,
) -> None:
    if "pull_request" not in event["issue"]:
        return
    issue_number = event["issue"]["number"]
    comment_id = event["comment"]["id"]
    actor = event["comment"]["user"]["login"]
    if is_already_handled(github, issue_number, comment_id):
        print(f"Comment {comment_id} was already handled.")
        return
    add_reaction_safely(github, comment_id, "-1")
    commands = ", ".join(f"`{command}`" for command in sorted(ALL_CI_COMMANDS))
    github.add_comment(
        issue_number,
        f"❌ @{actor}, that is not a recognized CI command, so no CI was "
        "started. Post a new comment containing only one of: "
        f"{commands}.\n\n{command_comment_marker(comment_id)}",
    )


def run(
    event: Mapping[str, Any],
    github: GitHubClient,
    buildkite: BuildkiteClient,
    trusted_users_value: str = "",
    branch_updater: GitHubClient | None = None,
) -> None:
    command = parse_command(event["comment"]["body"])
    if not command or "pull_request" not in event["issue"]:
        return

    issue_number = event["issue"]["number"]
    comment_id = event["comment"]["id"]
    actor = event["comment"]["user"]["login"]

    if is_already_handled(github, issue_number, comment_id):
        print(f"Comment {comment_id} was already handled.")
        return
    add_reaction_safely(github, comment_id, "eyes")

    try:
        pr = github.get_pr(issue_number)
        permission = github.get_permission(actor)
        if pr["state"] != "open":
            github.add_comment(
                issue_number,
                "❌ CI commands require an open PR.\n\n"
                f"{command_comment_marker(comment_id)}",
            )
            return

        trusted_users = parse_trusted_users(trusted_users_value)
        should_check_approval = (
            not is_trusted_permission(permission)
            and actor.casefold() not in trusted_users
            and actor.casefold() == pr["user"]["login"].casefold()
            and not pr["draft"]
            and not has_ready_label(pr)
        )
        trusted_approval = should_check_approval and has_trusted_approval(
            github,
            issue_number,
            trusted_users,
        )
        allowed, reason = authorize(
            actor=actor,
            permission=permission,
            pr=pr,
            run_command=run_command_for_command(command),
            trusted_approval=trusted_approval,
            trusted_users=trusted_users,
        )
        if not allowed:
            github.add_comment(
                issue_number,
                f"❌ @{actor}, {reason}\n\n{command_comment_marker(comment_id)}",
            )
            return

        print(f"Authorized @{actor}: {reason}")
        if command in RUN_CI_COMMAND_ENV:
            message = handle_run_ci(
                actor=actor,
                buildkite=buildkite,
                comment_id=comment_id,
                command=command,
                github=github,
                pr=pr,
                branch_updater=branch_updater,
                allow_repeat_run=is_trusted_permission(permission),
            )
        elif command in RETRY_COMMANDS:
            message = handle_retry_failed(
                actor=actor,
                buildkite=buildkite,
                comment_id=comment_id,
                command=command,
                github=github,
                pr=pr,
            )
        elif command in CANCEL_COMMANDS:
            message = handle_cancel_ci(
                buildkite=buildkite,
                command=command,
                pr=pr,
            )
        else:
            raise ValueError(f"Unsupported CI command: {command}")
        add_reaction_safely(github, comment_id, "rocket")
        github.add_comment(issue_number, f"✅ {message}")
    except CiPreparationError as error:
        add_reaction_safely(github, comment_id, "-1")
        github.add_comment(
            issue_number,
            f"❌ {error}\n\n{command_comment_marker(comment_id)}",
        )
    except Exception:
        add_reaction_safely(github, comment_id, "confused")
        raise


def main() -> None:
    event_path = os.environ["GITHUB_EVENT_PATH"]
    with open(event_path, encoding="utf-8") as event_file:
        event = json.load(event_file)

    github = GitHubClient(
        os.environ.get("GH_TOKEN", ""),
        os.environ["GITHUB_REPOSITORY"],
    )
    event_name = os.environ.get("GITHUB_EVENT_NAME", "issue_comment")
    if event_name == "pull_request_target":
        notify_authorized(
            event,
            github,
            os.environ.get("CI_TRUSTED_USERS", ""),
        )
        return
    if event_name == "workflow_run":
        pr = resolve_workflow_run_pr(event["workflow_run"], github)
        if pr is None:
            print("Could not resolve an open PR for the approval workflow run.")
            return
        notify_authorized(
            {
                "action": "submitted",
                "pull_request": pr,
                "review": {"state": "approved"},
            },
            github,
            os.environ.get("CI_TRUSTED_USERS", ""),
        )
        return

    body = event["comment"]["body"]
    command = parse_command(body)
    if not command:
        if is_ci_command_attempt(body):
            reply_unrecognized_command(event, github)
        return

    pipeline = pipeline_for_command(
        command,
        amd_ci_pipeline=os.environ.get("BUILDKITE_AMD_PIPELINE", "amd-ci"),
        upstream_ci_pipeline=os.environ.get("BUILDKITE_PIPELINE", "ci"),
    )
    buildkite = BuildkiteClient(
        os.environ.get("BUILDKITE_API_TOKEN", ""),
        os.environ.get("BUILDKITE_ORGANIZATION", "vllm"),
        pipeline,
    )
    branch_update_token = os.environ.get("GH_BRANCH_UPDATE_TOKEN", "")
    run(
        event,
        github,
        buildkite,
        os.environ.get("CI_TRUSTED_USERS", ""),
        GitHubClient(branch_update_token, os.environ["GITHUB_REPOSITORY"])
        if branch_update_token
        else None,
    )


if __name__ == "__main__":
    main()