fix upload alerts to rockset (#103995)

Testing is the CI of https://github.com/pytorch/pytorch/pull/103996

Pull Request resolved: https://github.com/pytorch/pytorch/pull/103995
Approved by: https://github.com/huydhn

.github/workflows/nightly-rockset-uploads.yml

@@ -4,9 +4,14 @@ on:
   schedule:
     # Choose a random time near midnight PST because it may be delayed if there are high loads
     - cron:  37 7 * * *
+  pull_request:
+    paths:
+      - 'tools/stats/upload_external_contrib_stats.py'
+      - 'tools/stats/upload_test_stat_aggregates.py'
+      - '.github/workflows/nightly-rockset-uploads.yml'
 
 concurrency:
-  group: ${{ github.workflow }}-${{ github.event_name == 'schedule' || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
   cancel-in-progress: true
 
 jobs:

.github/workflows/upload-alerts.yml

@@ -5,7 +5,10 @@ name: Upload Alerts to AWS/Rockset
 on:
   schedule:
     - cron: '*/10 * * * *'
-  workflow_dispatch:
+  pull_request:
+    paths:
+      - 'tools/alerts/create_alerts.py'
+      - '.github/workflows/upload-alerts.yml'
 
 jobs:
   upload-alerts:
@@ -25,7 +28,7 @@ jobs:
 
       - name: Create alerts
         run: |
-          output=$(python3 "tools/alerts/create_alerts.py")
+          output=$(PYTHONPATH=$PYTHONPATH:$(pwd) python3 "tools/alerts/create_alerts.py")
           echo "uploading following alerts"
           echo "$output"
           echo "script-output=$output" >> "$GITHUB_OUTPUT"

tools/github/github_utils.py

@@ -0,0 +1,81 @@
+"""GitHub Utilities"""
+
+import json
+import os
+
+from typing import Any, Callable, cast, Dict, Optional, Tuple
+
+from urllib.error import HTTPError
+from urllib.parse import quote
+from urllib.request import Request, urlopen
+
+
+def gh_fetch_url_and_headers(
+    url: str,
+    *,
+    headers: Optional[Dict[str, str]] = None,
+    data: Optional[Dict[str, Any]] = None,
+    method: Optional[str] = None,
+    reader: Callable[[Any], Any] = lambda x: x.read(),
+) -> Tuple[Any, Any]:
+    if headers is None:
+        headers = {}
+    token = os.environ.get("GITHUB_TOKEN")
+    if token is not None and url.startswith("https://api.github.com/"):
+        headers["Authorization"] = f"token {token}"
+    data_ = json.dumps(data).encode() if data is not None else None
+    try:
+        with urlopen(Request(url, headers=headers, data=data_, method=method)) as conn:
+            return conn.headers, reader(conn)
+    except HTTPError as err:
+        if err.code == 403 and all(
+            key in err.headers for key in ["X-RateLimit-Limit", "X-RateLimit-Used"]
+        ):
+            print(
+                f"""Rate limit exceeded:
+                Used: {err.headers['X-RateLimit-Used']}
+                Limit: {err.headers['X-RateLimit-Limit']}
+                Remaining: {err.headers['X-RateLimit-Remaining']}
+                Resets at: {err.headers['x-RateLimit-Reset']}"""
+            )
+        raise
+
+
+def gh_fetch_url(
+    url: str,
+    *,
+    headers: Optional[Dict[str, str]] = None,
+    data: Optional[Dict[str, Any]] = None,
+    method: Optional[str] = None,
+    reader: Callable[[Any], Any] = lambda x: x.read(),
+) -> Any:
+    return gh_fetch_url_and_headers(
+        url, headers=headers, data=data, reader=json.load, method=method
+    )[1]
+
+
+def _gh_fetch_json_any(
+    url: str,
+    params: Optional[Dict[str, Any]] = None,
+    data: Optional[Dict[str, Any]] = None,
+) -> Any:
+    headers = {"Accept": "application/vnd.github.v3+json"}
+    if params is not None and len(params) > 0:
+        url += "?" + "&".join(
+            f"{name}={quote(str(val))}" for name, val in params.items()
+        )
+    return gh_fetch_url(url, headers=headers, data=data, reader=json.load)
+
+
+def gh_fetch_json_dict(
+    url: str,
+    params: Optional[Dict[str, Any]] = None,
+    data: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    return cast(Dict[str, Any], _gh_fetch_json_any(url, params, data))
+
+
+def gh_fetch_commit(org: str, repo: str, sha: str) -> Dict[str, Any]:
+    return gh_fetch_json_dict(
+        f"https://api.github.com/repos/{org}/{repo}/commits/{sha}"
+    )