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

.github/workflows/github-status-test-lambda.yml

@@ -1,4 +1,4 @@
-name: Deploy github-status-test
+name: Test and deploy github-status-test
 
 on:
   push:
@@ -17,7 +17,19 @@
     working-directory: aws/lambda/github-status-test/
 
 jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: '3.11'
+          cache: pip
+      - run: pip3 install -r requirements.txt pytest
+      - run: pytest -v test_lambda_function.py
+
   deploy:
+    needs: test
     runs-on: ubuntu-latest
     if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }}
     permissions:

aws/lambda/github-status-test/Makefile

@@ -1,11 +1,34 @@
-prepare:
+ZIP := github-status-test-deployment.zip
+# Third-party modules lambda_function.py imports. The built zip is checked for
+# these before it can be deployed: a package missing a dependency fails at
+# import, which drops every event the webhook sends -- not just the log
+# download -- and update-function-code publishes straight to $LATEST, which is
+# what API Gateway invokes.
+VENDORED := boto3 requests github
+
+# 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.
+# Starts from clean so a stale packages/ or zip can't leak into the artifact.
+prepare: clean
 	mkdir -p ./packages
-	pip install --target ./packages -r requirements.txt
-	cd packages && zip -r ../github-status-test-deployment.zip .
-	zip -g github-status-test-deployment.zip lambda_function.py
+	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 ../$(ZIP) .
+	zip -g $(ZIP) lambda_function.py
+	$(MAKE) verify
+
+verify:
+	@for m in $(VENDORED); do \
+		unzip -l $(ZIP) | grep -qE " $$m/__init__\.py$$" \
+			|| { echo "ERROR: '$$m' missing from $(ZIP), refusing to deploy"; exit 1; }; \
+	done
+	@echo "verified: $(ZIP) contains $(VENDORED)"
 
 deploy: prepare
-	aws lambda update-function-code --function-name github-status-test --zip-file fileb://github-status-test-deployment.zip
+	aws lambda update-function-code --function-name github-status-test --zip-file fileb://$(ZIP)
 
 clean:
-	rm -rf github-status-test-deployment.zip packages
+	rm -rf $(ZIP) packages

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

@@ -1,8 +1,43 @@
 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
 
+> **`make deploy` is immediately live in production.** The API Gateway integration currently points at
+> the unqualified function (`:function:github-status-test/invocations`), so `update-function-code` puts
+> the new code on `$LATEST` and every webhook hits it right away. The publish-a-version steps below are
+> stale — they describe pinning the integration to a numbered version, which is not how it is wired
+> today, and the resource id is now `xtmtzj` rather than `clc02o`. Until that is fixed, treat any deploy
+> as a direct production change: `make prepare` verifies the zip contains every vendored module before
+> `make deploy` will run, but there is no staged rollout behind it.
+
 A new version of the lambda can be deployed using `make deploy` and it will be done so automatically by the workflow
 `github-status-test-lambda` when a change is committed to main. We have limited capacity for testing this lambda at
 the moment, so additional verification steps are needed to get the new deployed version to prod. More tests and guardrails

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

@@ -1,34 +1,151 @@
 # Copyright (c) 2019-present, Facebook, Inc.
 
+import base64
+import contextlib
 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 a repo 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)
+
+# Keyed by "owner/repo", not owner: get_repo_installation() resolves per repo,
+# so a "Selected repositories" install can cover one repo of an owner and not
+# its sibling. Sharing an owner's entry would hand a repo a token minted for a
+# different one, or let one repo's "not installed" result mask another's.
+# full_name -> (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:
+        with contextlib.suppress(ValueError):
+            return float(reset)
+    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 repo, 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, or None if unavailable."""
+    if not GITHUB_APP_ID or not GITHUB_APP_PRIVATE_KEY:
+        return None
+
+    cached = _token_cache.get(full_name)
+    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 {full_name}: {err}")
+        return None
+
+    _token_cache[full_name] = (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.
+            _token_cache[full_name] = (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 +178,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