Migrate update-viablestrict to test-infra as a generic GHA

.github/actions/update-viablestrict/action.yml

@@ -0,0 +1,96 @@
+name: Update viable/strict to the latest green commit from a PyTorch repo
+
+description: |
+  The latest green commit is the newest commit that have passed all the required
+  test in trunk and can be consider stable. This GHA only works for repos that
+  have been onboard to PyTorch Dev Infra.
+
+inputs:
+  repository:
+    description: The repository name, i.e. pytorch/pytorch
+    required: true
+    type: string
+  stable-branch:
+    description: The name of the stable branch to push to
+    required: false
+    default: 'viable/strict'
+    type: string
+  requires:
+    description: |
+      The list of required jobs that need to pass before the commit can be
+      considered stable
+    required: true
+    type: string
+  secret-bot-token:
+    description: The token to use to push to the stable protected branch
+    required: true
+    type: string
+  rockset-api-key:
+    description: The API key to query Rockset, read-only
+    required: true
+    type: string
+  test-infra-repository:
+    description: Test infra repository to use
+    default: 'pytorch/test-infra'
+    type: string
+  test-infra-ref:
+    description: Test infra reference to use
+    default: ''
+    type: string
+
+runs:
+  using: composite
+  steps:
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.11'
+
+    - name: Checkout test-infra for the fetch_latest_green_commit scripts
+      uses: actions/checkout@v3
+      with:
+        repository: ${{ inputs.test-infra-repository }}
+        ref: ${{ inputs.test-infra-ref }}
+        path: test-infra
+
+    - uses: actions/checkout@v3
+      with:
+        repository: ${{ inputs.repository }}
+        token: ${{ inputs.secret-bot-token }}
+        path: ${{ inputs.repository }}
+        # Need the whole history here to get the lastest green commit, which is
+        # usually few hours behind main
+        fetch-depth: 0
+
+    - name: Install Python Packages
+      shell: bash
+      run: |
+        pip install rockset==1.0.3 boto3==1.19.12
+
+    - name: Get latest viable commit
+      id: get-latest-commit
+      working-directory: ${{ inputs.repository }}
+      env:
+        ROCKSET_API_KEY: ${{ inputs.rockset-api-key }}
+      shell: bash
+      run: |
+        set -ex
+
+        output=$(python ${GITHUB_WORKSPACE}/test-infra/.github/scripts/fetch_latest_green_commit.py --requires "${{ inputs.requires }}")
+        echo "latest_viable_sha=$output" >> "${GITHUB_OUTPUT}"
+
+    - name: Push SHA to viable/strict branch
+      if: steps.get-latest-commit.outputs.latest_viable_sha != 'None'
+      working-directory: ${{ inputs.repository }}
+      env:
+        GITHUB_TOKEN: ${{ inputs.secret-bot-token }}
+        STABLE_BRANCH: ${{ inputs.stable-branch }}
+        LATEST_VIABLE_SHA: ${{ steps.get-latest-commit.outputs.latest_viable_sha }}
+      shell: bash
+      run: |
+        git config --global user.email "pytorchmergebot@users.noreply.github.com"
+        git config --global user.name "PyTorch MergeBot"
+
+        echo "Set the latest sha variable to be ${LATEST_VIABLE_SHA}"
+        # Pushing an older green commit here will fail because it's non-fast-forward, which is ok
+        # to ignore because we already have the later green commit in visable/strict
+        git push origin "${LATEST_VIABLE_SHA}":"${STABLE_BRANCH}" || true

.github/scripts/fetch_latest_green_commit.py

@@ -0,0 +1,156 @@
+import json
+import os
+import re
+import sys
+from typing import Any, cast, Dict, List, NamedTuple, Optional, Tuple
+
+import rockset  # type: ignore[import]
+from gitutils import _check_output
+
+
+def eprint(msg: str) -> None:
+    print(msg, file=sys.stderr)
+
+
+class WorkflowCheck(NamedTuple):
+    workflowName: str
+    name: str
+    jobName: str
+    conclusion: str
+
+
+def get_latest_commits() -> List[str]:
+    latest_viable_commit = _check_output(
+        [
+            "git",
+            "log",
+            "-n",
+            "1",
+            "--pretty=format:%H",
+            "origin/viable/strict",
+        ],
+        encoding="ascii",
+    )
+    commits = _check_output(
+        [
+            "git",
+            "rev-list",
+            f"{latest_viable_commit}^..HEAD",
+            "--remotes=*origin/main",
+        ],
+        encoding="ascii",
+    ).splitlines()
+
+    return commits
+
+
+def query_commits(commits: List[str]) -> List[Dict[str, Any]]:
+    rs = rockset.RocksetClient(
+        host="api.usw2a1.rockset.com", api_key=os.environ["ROCKSET_API_KEY"]
+    )
+    params = [{"name": "shas", "type": "string", "value": ",".join(commits)}]
+    res = rs.QueryLambdas.execute_query_lambda(
+        # https://console.rockset.com/lambdas/details/commons.commit_jobs_batch_query
+        query_lambda="commit_jobs_batch_query",
+        version="19c74e10819104f9",
+        workspace="commons",
+        parameters=params,
+    )
+
+    return cast(List[Dict[str, Any]], res.results)
+
+
+def print_commit_status(commit: str, results: Dict[str, Any]) -> None:
+    print(commit)
+    for check in results["results"]:
+        if check["sha"] == commit:
+            print(f"\t{check['conclusion']:>10}: {check['name']}")
+
+
+def get_commit_results(
+    commit: str, results: List[Dict[str, Any]]
+) -> List[Dict[str, Any]]:
+    workflow_checks = []
+    for check in results:
+        if check["sha"] == commit:
+            workflow_checks.append(
+                WorkflowCheck(
+                    workflowName=check["workflowName"],
+                    name=check["name"],
+                    jobName=check["jobName"],
+                    conclusion=check["conclusion"],
+                )._asdict()
+            )
+    return workflow_checks
+
+
+def is_green(
+    commit: str, requires: List[str], results: List[Dict[str, Any]]
+) -> Tuple[bool, str]:
+    workflow_checks = get_commit_results(commit, results)
+
+    regex = {name: False for name in requires}
+
+    for check in workflow_checks:
+        jobName = check["jobName"]
+        # Ignore result from unstable job, be it success or failure
+        if "unstable" in jobName:
+            continue
+
+        workflowName = check["workflowName"]
+        conclusion = check["conclusion"]
+        for required_check in regex:
+            if re.match(required_check, workflowName, flags=re.IGNORECASE):
+                if conclusion not in ["success", "skipped"]:
+                    return (False, workflowName + " checks were not successful")
+                else:
+                    regex[required_check] = True
+
+    missing_workflows = [x for x in regex.keys() if not regex[x]]
+    if len(missing_workflows) > 0:
+        return (False, "missing required workflows: " + ", ".join(missing_workflows))
+
+    return (True, "")
+
+
+def get_latest_green_commit(
+    commits: List[str], requires: List[str], results: List[Dict[str, Any]]
+) -> Optional[str]:
+    for commit in commits:
+        eprint(f"Checking {commit}")
+        green, msg = is_green(commit, requires, results)
+        if green:
+            eprint("GREEN")
+            return commit
+        else:
+            eprint("RED: " + msg)
+    return None
+
+
+def parse_args() -> Any:
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser("Return the latest green commit from a PyTorch repo")
+    parser.add_argument(
+        "--requires",
+        type=str,
+        required=True,
+        help="the JSON list of required jobs that need to pass for the commit to be green",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    commits = get_latest_commits()
+    results = query_commits(commits)
+
+    latest_viable_commit = get_latest_green_commit(
+        commits, json.loads(args.requires), results
+    )
+    print(latest_viable_commit)
+
+
+if __name__ == "__main__":
+    main()