github-status-test: use a GitHub App for log downloads, with PAT fallback

Job logs are currently downloaded with a randomly chosen PAT from the
GITHUB_TOKENS pool. PATs are tied to a personal account, need `repo` scope
to reach private repos, and rotate whenever their owner does.

Download logs with a GitHub App installation token instead, falling back to
the existing PAT pool when the app is rate limited, rejected, or not
installed on the repo's owner. Both app env vars unset means the previous
behaviour exactly, so the app can be rolled back by clearing them.

- GITHUB_APP_ID / GITHUB_APP_PRIVATE_KEY (base64 PEM, raw PEM also accepted)
- Tokens are minted with PyGithub's GithubIntegration, matching
  cross_repo_ci_relay, rather than hand-rolling the app JWT
- Installation tokens are cached per repo owner in module scope and
  refreshed 5 minutes early, so warm invocations don't mint one per job
- Fallback triggers on 401/403/429. On a rate limit the app is parked until
  x-ratelimit-reset so later jobs skip it rather than each paying for a
  doomed request first
- Repos outside the installation (e.g. vllm-project/vllm) resolve to no
  installation and go straight to the pool, cached briefly to avoid a
  lookup per job

Two existing bugs fixed, both needed for the fallback to work:

- Non-OK responses were archived as though they were logs, so a rate-limited
  403 body would land in S3 as the job's log and be sent to the classifier.
  Detecting that condition is also what fallback depends on.
- `except HTTPError` caught urllib's, never requests'. It now catches
  requests.RequestException too.

The lambda runs on python3.9 but the deploy workflow's setup-python pins no
version, so `pip install` used the runner's default. cryptography (via
PyGithub) ships compiled wheels, so the Makefile now pins the target platform
and python version; otherwise the zip gets wheels the runtime can't import.

Author: huydhn

aws/lambda/github-status-test/Makefile

@@ -1,6 +1,12 @@
+# The lambda runs on python3.9/x86_64. cryptography ships compiled wheels, so
+# pin the target platform rather than inheriting whatever python the CI runner
+# happens to default to -- otherwise the zip gets wheels the runtime can't load.
 prepare:
 	mkdir -p ./packages
-	pip install --target ./packages -r requirements.txt
+	pip install --target ./packages \
+		--platform manylinux2014_x86_64 --python-version 3.9 \
+		--implementation cp --only-binary=:all: --no-compile \
+		-r requirements.txt
 	cd packages && zip -r ../github-status-test-deployment.zip .
 	zip -g github-status-test-deployment.zip lambda_function.py
 

aws/lambda/github-status-test/README.md

@@ -1,6 +1,33 @@
 Despite the name, this is the lambda used to write GitHub webhook payloads to S3 as mentioned
 in https://github.com/pytorch/test-infra/blob/main/torchci/docs/architecture.md
 
+### GitHub credentials
+
+Job logs are downloaded with a GitHub App installation token, falling back to the `GITHUB_TOKENS`
+PAT pool when the app is rate limited, rejected, or not installed on the repo's owner.
+
+| Env var | Required | Purpose |
+| --- | --- | --- |
+| `GITHUB_APP_ID` | no | App id used to mint installation tokens (e.g. `4550824`, `pytorch-bot-preview`) |
+| `GITHUB_APP_PRIVATE_KEY` | no | The app's private key, base64-encoded PEM (same encoding torchci uses) |
+| `GITHUB_TOKENS` | yes | Comma-separated PAT pool, used as the fallback and when no app is configured |
+
+With both app vars unset the lambda behaves exactly as before and only uses `GITHUB_TOKENS`, so
+the app can be rolled back by clearing the env vars — no code change or redeploy needed.
+
+Notes on the app path:
+
+- Installation tokens last an hour and are cached per repo owner in module scope, so a warm
+  invocation reuses one rather than minting a token per job.
+- The app's rate limit is per installation. `pytorch` is enterprise-owned, so its installation
+  gets 15,000 requests/hour, independent of any other app's quota. Use a dedicated app rather
+  than the shared `pytorch-bot` installation, whose quota Dr. CI and the HUD already draw on.
+- Repos outside the installation (e.g. `vllm-project/vllm`) resolve to no installation and go
+  straight to the PAT pool; that negative result is cached briefly to avoid a lookup per job.
+- Downloading job logs is documented as needing the `actions: read` permission. It currently
+  works without it because pytorch repos are public, but the permission should be granted so the
+  dependency is explicit and private repos keep working.
+
 ### Deployment
 
 A new version of the lambda can be deployed using `make deploy` and it will be done so automatically by the workflow

aws/lambda/github-status-test/lambda_function.py

@@ -1,34 +1,150 @@
 # Copyright (c) 2019-present, Facebook, Inc.
 
+import base64
 import gzip
 import json
 import os
 import random
+import time
 from urllib.error import HTTPError
 from urllib.request import urlopen
 from uuid import uuid4
 
 import boto3
 import requests
+from github import Auth, GithubIntegration
+from github.GithubException import UnknownObjectException
 
 
 s3 = boto3.resource("s3")
 GITHUB_TOKENS = os.environ.get("GITHUB_TOKENS")
+GITHUB_APP_ID = os.environ.get("GITHUB_APP_ID")
+# Base64-encoded PEM, the same encoding torchci uses for its app key
+GITHUB_APP_PRIVATE_KEY = os.environ.get("GITHUB_APP_PRIVATE_KEY")
 BUCKET_NAME = "ossci-raw-job-status"
 
+GITHUB_API_URL = "https://api.github.com"
+# Installation tokens last an hour. Refresh early so a warm invocation never
+# signs a request with a token that expires mid-flight.
+TOKEN_EXPIRY_MARGIN = 300
+# How long to remember that an owner has no app installation, so repos outside
+# the installation don't trigger a lookup for every job.
+NO_INSTALLATION_TTL = 900
+# Used when a credential is rejected without telling us when it recovers.
+DEFAULT_COOL_OFF = 60
+# Statuses meaning "this credential can't do it, try the next one": rate limited
+# (403 or 429) or rejected outright (401).
+FALLBACK_STATUSES = (401, 403, 429)
+
+# owner -> (installation token or None, epoch seconds the entry goes stale)
+_token_cache = {}
+
 
 def json_dumps(obj):
     return json.dumps(obj, sort_keys=True, indent=4, separators=(",", ": "))
 
 
-def download_log(full_name, conclusion, job_id):
-    url = f"https://api.github.com/repos/{full_name}/actions/jobs/{job_id}/logs"
+def app_private_key():
+    key = GITHUB_APP_PRIVATE_KEY
+    if "PRIVATE KEY" not in key:
+        key = base64.b64decode(key).decode("utf-8")
+    return key
+
+
+def cool_off_until(response):
+    reset = response.headers.get("x-ratelimit-reset")
+    if reset:
+        try:
+            return float(reset)
+        except ValueError:
+            pass
+    return time.time() + DEFAULT_COOL_OFF
+
+
+def fetch_installation_token(full_name):
+    """Mint an installation token for the app installation covering full_name.
+
+    Returns (None, expiry) when the app isn't installed on that owner, so the
+    caller falls back to a PAT instead of retrying the lookup for every job.
+    """
+    owner, repo = full_name.split("/", 1)
+    integration = GithubIntegration(
+        auth=Auth.AppAuth(int(GITHUB_APP_ID), app_private_key())
+    )
+
+    try:
+        installation = integration.get_repo_installation(owner, repo)
+    except UnknownObjectException:
+        return None, time.time() + NO_INSTALLATION_TTL
+
+    token = integration.get_access_token(installation.id)
+    return token.token, token.expires_at.timestamp() - TOKEN_EXPIRY_MARGIN
+
+
+def installation_token(full_name):
+    """Cached installation token for full_name's owner, or None if unavailable."""
+    if not GITHUB_APP_ID or not GITHUB_APP_PRIVATE_KEY:
+        return None
+
+    owner = full_name.split("/")[0]
+    cached = _token_cache.get(owner)
+    if cached and time.time() < cached[1]:
+        return cached[0]
+
+    try:
+        token, expires_at = fetch_installation_token(full_name)
+    except Exception as err:
+        # Deliberately broad: a bad app id, an unparseable private key or a
+        # GitHub blip must degrade to the PAT pool, never fail the webhook and
+        # lose the payload archiving that happens after the log download.
+        # Not cached either, so the next invocation retries.
+        print(f"ERROR minting installation token for {owner}: {err}")
+        return None
+
+    _token_cache[owner] = (token, expires_at)
+    return token
+
+
+def fetch_log(full_name, job_id, token):
+    url = f"{GITHUB_API_URL}/repos/{full_name}/actions/jobs/{job_id}/logs"
     headers = {
         "Accept": "application/vnd.github.v3+json",
-        "Authorization": "token " + random.choice(GITHUB_TOKENS.split(",")),
+        "Authorization": "token " + token,
     }
-    r = requests.get(url, headers=headers)
-    log_data = r.content
+    return requests.get(url, headers=headers, timeout=30)
+
+
+def download_log(full_name, conclusion, job_id):
+    response = None
+
+    app_token = installation_token(full_name)
+    if app_token:
+        response = fetch_log(full_name, job_id, app_token)
+        if response.status_code in FALLBACK_STATUSES:
+            # Stop using the app until its window resets, otherwise every job
+            # for the rest of the hour pays for a doomed request first.
+            owner = full_name.split("/")[0]
+            _token_cache[owner] = (None, cool_off_until(response))
+            print(
+                f"App auth returned {response.status_code} for {full_name} "
+                f"job {job_id}, falling back to a PAT"
+            )
+            response = None
+
+    if response is None:
+        if not GITHUB_TOKENS:
+            print(f"ERROR no usable credential for {full_name} job {job_id}")
+            return
+        response = fetch_log(full_name, job_id, random.choice(GITHUB_TOKENS.split(",")))
+
+    if not response.ok:
+        # Bail out rather than archive the API error body as if it were the log
+        print(
+            f"ERROR {response.status_code} downloading log for {full_name} job {job_id}"
+        )
+        return
+
+    log_data = response.content
 
     object_path = f"log/{job_id}"
     if full_name != "pytorch/pytorch":
@@ -61,7 +177,7 @@ def lambda_handler(event, context):
             conclusion = body[event_type]["conclusion"]
             job_id = body[event_type]["id"]
             download_log(full_name, conclusion, job_id)
-        except HTTPError as err:
+        except (HTTPError, requests.RequestException) as err:
             # Just eat the error as logs are optional.
             print("ERROR", err)
             pass

aws/lambda/github-status-test/requirements.txt

@@ -1,2 +1,4 @@
 boto3==1.24.59
 requests==2.32.2
+# Mints the app installation token. Same version cross_repo_ci_relay uses.
+PyGithub==2.9.0

aws/lambda/github-status-test/test_lambda_function.py

@@ -0,0 +1,178 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+import lambda_function
+from lambda_function import download_log, installation_token
+
+
+def make_response(status_code, content=b"log data", headers=None):
+    response = MagicMock()
+    response.status_code = status_code
+    response.ok = 200 <= status_code < 300
+    response.content = content
+    response.headers = headers or {}
+    return response
+
+
+class TestInstallationToken(unittest.TestCase):
+    def setUp(self):
+        lambda_function._token_cache.clear()
+
+    @patch.object(lambda_function, "GITHUB_APP_ID", None)
+    @patch.object(lambda_function, "GITHUB_APP_PRIVATE_KEY", None)
+    def test_returns_none_without_app_credentials(self):
+        with patch.object(lambda_function, "fetch_installation_token") as fetch:
+            self.assertIsNone(installation_token("pytorch/pytorch"))
+            fetch.assert_not_called()
+
+    @patch.object(lambda_function, "GITHUB_APP_ID", "4550824")
+    @patch.object(lambda_function, "GITHUB_APP_PRIVATE_KEY", "key")
+    def test_token_is_cached_per_owner(self):
+        with patch.object(
+            lambda_function,
+            "fetch_installation_token",
+            return_value=("tok", 2**31),
+        ) as fetch:
+            self.assertEqual(installation_token("pytorch/pytorch"), "tok")
+            # Same owner, different repo: served from cache, no second mint
+            self.assertEqual(installation_token("pytorch/executorch"), "tok")
+            fetch.assert_called_once()
+
+    @patch.object(lambda_function, "GITHUB_APP_ID", "4550824")
+    @patch.object(lambda_function, "GITHUB_APP_PRIVATE_KEY", "key")
+    def test_expired_cache_entry_is_refreshed(self):
+        lambda_function._token_cache["pytorch"] = ("stale", 0)
+        with patch.object(
+            lambda_function,
+            "fetch_installation_token",
+            return_value=("fresh", 2**31),
+        ):
+            self.assertEqual(installation_token("pytorch/pytorch"), "fresh")
+
+    @patch.object(lambda_function, "GITHUB_APP_ID", "4550824")
+    @patch.object(lambda_function, "GITHUB_APP_PRIVATE_KEY", "key")
+    def test_uninstalled_owner_is_cached_as_none(self):
+        with patch.object(
+            lambda_function,
+            "fetch_installation_token",
+            return_value=(None, 2**31),
+        ) as fetch:
+            self.assertIsNone(installation_token("vllm-project/vllm"))
+            self.assertIsNone(installation_token("vllm-project/vllm"))
+            fetch.assert_called_once()
+
+    @patch.object(lambda_function, "GITHUB_APP_ID", "4550824")
+    @patch.object(lambda_function, "GITHUB_APP_PRIVATE_KEY", "key")
+    def test_mint_failure_is_not_cached(self):
+        with patch.object(
+            lambda_function,
+            "fetch_installation_token",
+            side_effect=lambda_function.requests.RequestException("boom"),
+        ):
+            self.assertIsNone(installation_token("pytorch/pytorch"))
+        self.assertNotIn("pytorch", lambda_function._token_cache)
+
+    @patch.object(lambda_function, "GITHUB_APP_ID", "4550824")
+    @patch.object(lambda_function, "GITHUB_APP_PRIVATE_KEY", "bm90LWEta2V5")
+    def test_unparseable_private_key_degrades_to_none(self):
+        # A misconfigured key must not escape as an exception: the webhook still
+        # has to archive its payload after the log download is skipped.
+        self.assertIsNone(installation_token("pytorch/pytorch"))
+
+
+@patch.object(lambda_function, "GITHUB_TOKENS", "pat1,pat2")
+@patch.object(lambda_function, "s3")
+@patch.object(lambda_function, "urlopen")
+class TestDownloadLog(unittest.TestCase):
+    def setUp(self):
+        lambda_function._token_cache.clear()
+
+    def test_uses_app_token_when_available(self, urlopen, s3):
+        with patch.object(
+            lambda_function, "installation_token", return_value="app-token"
+        ), patch.object(
+            lambda_function, "fetch_log", return_value=make_response(200)
+        ) as fetch_log:
+            download_log("pytorch/pytorch", "failure", 123)
+
+        fetch_log.assert_called_once_with("pytorch/pytorch", 123, "app-token")
+        s3.Object.assert_called_once_with("ossci-raw-job-status", "log/123")
+        urlopen.assert_called_once()
+
+    def test_falls_back_to_pat_when_rate_limited(self, urlopen, s3):
+        responses = [
+            make_response(403, headers={"x-ratelimit-remaining": "0"}),
+            make_response(200),
+        ]
+        with patch.object(
+            lambda_function, "installation_token", return_value="app-token"
+        ), patch.object(
+            lambda_function, "fetch_log", side_effect=responses
+        ) as fetch_log:
+            download_log("pytorch/pytorch", "failure", 123)
+
+        self.assertEqual(fetch_log.call_count, 2)
+        self.assertIn(fetch_log.call_args_list[1][0][2], ("pat1", "pat2"))
+        # The log still gets archived via the fallback credential
+        s3.Object.assert_called_once_with("ossci-raw-job-status", "log/123")
+
+    def test_rate_limited_app_is_skipped_until_reset(self, urlopen, s3):
+        reset_at = 2**31
+        responses = [
+            make_response(429, headers={"x-ratelimit-reset": str(reset_at)}),
+            make_response(200),
+            make_response(200),
+        ]
+        with patch.object(
+            lambda_function, "fetch_installation_token", return_value=("app", 2**31)
+        ), patch.object(lambda_function, "GITHUB_APP_ID", "4550824"), patch.object(
+            lambda_function, "GITHUB_APP_PRIVATE_KEY", "key"
+        ), patch.object(
+            lambda_function, "fetch_log", side_effect=responses
+        ) as fetch_log:
+            download_log("pytorch/pytorch", "failure", 123)
+            download_log("pytorch/pytorch", "failure", 456)
+
+        # 2 calls for the first job (app then PAT), 1 for the second (PAT only)
+        self.assertEqual(fetch_log.call_count, 3)
+        self.assertEqual(
+            lambda_function._token_cache["pytorch"], (None, float(reset_at))
+        )
+
+    def test_error_response_is_not_archived(self, urlopen, s3):
+        with patch.object(
+            lambda_function, "installation_token", return_value=None
+        ), patch.object(
+            lambda_function,
+            "fetch_log",
+            return_value=make_response(404, content=b'{"message": "Not Found"}'),
+        ):
+            download_log("pytorch/pytorch", "skipped", 123)
+
+        s3.Object.assert_not_called()
+        urlopen.assert_not_called()
+
+    def test_non_pytorch_repo_is_prefixed(self, urlopen, s3):
+        with patch.object(
+            lambda_function, "installation_token", return_value=None
+        ), patch.object(lambda_function, "fetch_log", return_value=make_response(200)):
+            download_log("pytorch/executorch", "success", 999)
+
+        s3.Object.assert_called_once_with(
+            "ossci-raw-job-status", "log/pytorch/executorch/999"
+        )
+
+    def test_no_credentials_at_all_is_a_noop(self, urlopen, s3):
+        with patch.object(
+            lambda_function, "installation_token", return_value=None
+        ), patch.object(lambda_function, "GITHUB_TOKENS", None), patch.object(
+            lambda_function, "fetch_log"
+        ) as fetch_log:
+            download_log("pytorch/pytorch", "failure", 123)
+
+        fetch_log.assert_not_called()
+        s3.Object.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()