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)
- Installation tokens are cached per repo owner in module scope and
  refreshed 5 minutes early, so warm invocations don't mint one per job
- The JWT's iat is backdated 60s; GitHub rejects a future-dated iat
- 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 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,186 @@
 # Copyright (c) 2019-present, Facebook, Inc.
 
+import base64
 import gzip
 import json
 import os
 import random
+import time
+from datetime import datetime, timezone
 from urllib.error import HTTPError
 from urllib.request import urlopen
 from uuid import uuid4
 
 import boto3
+import jwt
 import requests
 
 
 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"
+# GitHub rejects a JWT whose iat is in the future, so backdate it to absorb any
+# clock skew between the lambda and GitHub.
+JWT_ISSUED_SKEW = 60
+JWT_EXPIRY = 540
+# 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 app_jwt():
+    now = int(time.time())
+    return jwt.encode(
+        {
+            "iat": now - JWT_ISSUED_SKEW,
+            "exp": now + JWT_EXPIRY,
+            "iss": GITHUB_APP_ID,
+        },
+        app_private_key(),
+        algorithm="RS256",
+    )
+
+
+def token_expiry(expires_at):
+    if expires_at:
+        try:
+            parsed = datetime.strptime(expires_at, "%Y-%m-%dT%H:%M:%SZ")
+            return parsed.replace(tzinfo=timezone.utc).timestamp() - TOKEN_EXPIRY_MARGIN
+        except ValueError:
+            pass
+    return time.time() + 3600 - TOKEN_EXPIRY_MARGIN
+
+
+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.
+    """
     headers = {
         "Accept": "application/vnd.github.v3+json",
-        "Authorization": "token " + random.choice(GITHUB_TOKENS.split(",")),
+        "Authorization": "Bearer " + app_jwt(),
     }
-    r = requests.get(url, headers=headers)
-    log_data = r.content
+
+    r = requests.get(
+        f"{GITHUB_API_URL}/repos/{full_name}/installation", headers=headers, timeout=10
+    )
+    if r.status_code == 404:
+        return None, time.time() + NO_INSTALLATION_TTL
+    r.raise_for_status()
+
+    installation_id = r.json()["id"]
+    r = requests.post(
+        f"{GITHUB_API_URL}/app/installations/{installation_id}/access_tokens",
+        headers=headers,
+        timeout=10,
+    )
+    r.raise_for_status()
+    body = r.json()
+    return body["token"], token_expiry(body.get("expires_at"))
+
+
+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 " + token,
+    }
+    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 +213,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,6 @@
 boto3==1.24.59
 requests==2.32.2
+PyJWT==2.13.0
+# Needed by PyJWT to sign the app JWT with RS256. 43.0.3 is the last release
+# publishing a cp39 wheel, which is what the python3.9 runtime needs.
+cryptography==43.0.3