|
| 1 | +# Copyright (c) 2019-present, Facebook, Inc. |
| 2 | + |
| 3 | +"""Download a completed GitHub Actions job log and archive it to S3. |
| 4 | +
|
| 5 | +Invoked asynchronously (``InvocationType: "Event"``) by the PyTorch bot's |
| 6 | +``workflow_job`` handler in torchci, and by torchci's backfill route. There is no |
| 7 | +API Gateway integration and no Lambda function URL: the only way in is |
| 8 | +``lambda:InvokeFunction``, which is IAM-authenticated. |
| 9 | +
|
| 10 | +Classification is *not* triggered from here. An S3 ObjectCreated notification on |
| 11 | +the ``log/`` prefix drives it, so any path that lands a log gets classified and |
| 12 | +this function never blocks on the classifier finishing. |
| 13 | +""" |
| 14 | + |
| 15 | +import base64 |
| 16 | +import contextlib |
| 17 | +import gzip |
| 18 | +import os |
| 19 | +import random |
| 20 | +import time |
| 21 | + |
| 22 | +import boto3 |
| 23 | +import requests |
| 24 | +from github import Auth, GithubIntegration |
| 25 | +from github.GithubException import UnknownObjectException |
| 26 | + |
| 27 | + |
| 28 | +s3 = boto3.resource("s3") |
| 29 | +GITHUB_TOKENS = os.environ.get("GITHUB_TOKENS") |
| 30 | +GITHUB_APP_ID = os.environ.get("GITHUB_APP_ID") |
| 31 | +# Base64-encoded PEM, the same encoding torchci uses for its app key |
| 32 | +GITHUB_APP_PRIVATE_KEY = os.environ.get("GITHUB_APP_PRIVATE_KEY") |
| 33 | +BUCKET_NAME = "ossci-raw-job-status" |
| 34 | + |
| 35 | +GITHUB_API_URL = "https://api.github.com" |
| 36 | +# Installation tokens last an hour. Refresh early so a warm invocation never |
| 37 | +# signs a request with a token that expires mid-flight. |
| 38 | +TOKEN_EXPIRY_MARGIN = 300 |
| 39 | +# How long to remember that a repo has no app installation, so repos outside |
| 40 | +# the installation don't trigger a lookup for every job. |
| 41 | +NO_INSTALLATION_TTL = 900 |
| 42 | +# Used when a credential is rejected without telling us when it recovers. |
| 43 | +DEFAULT_COOL_OFF = 60 |
| 44 | +# Statuses meaning "this credential can't do it, try the next one": rate limited |
| 45 | +# (403 or 429) or rejected outright (401). |
| 46 | +FALLBACK_STATUSES = (401, 403, 429) |
| 47 | + |
| 48 | +# Keyed by "owner/repo", not owner: get_repo_installation() resolves per repo, |
| 49 | +# so a "Selected repositories" install can cover one repo of an owner and not |
| 50 | +# its sibling. Sharing an owner's entry would hand a repo a token minted for a |
| 51 | +# different one, or let one repo's "not installed" result mask another's. |
| 52 | +# full_name -> (installation token or None, epoch seconds the entry goes stale) |
| 53 | +_token_cache = {} |
| 54 | + |
| 55 | + |
| 56 | +def app_private_key(): |
| 57 | + key = GITHUB_APP_PRIVATE_KEY |
| 58 | + if "PRIVATE KEY" not in key: |
| 59 | + key = base64.b64decode(key).decode("utf-8") |
| 60 | + return key |
| 61 | + |
| 62 | + |
| 63 | +def cool_off_until(response): |
| 64 | + reset = response.headers.get("x-ratelimit-reset") |
| 65 | + if reset: |
| 66 | + with contextlib.suppress(ValueError): |
| 67 | + return float(reset) |
| 68 | + return time.time() + DEFAULT_COOL_OFF |
| 69 | + |
| 70 | + |
| 71 | +def fetch_installation_token(full_name): |
| 72 | + """Mint an installation token for the app installation covering full_name. |
| 73 | +
|
| 74 | + Returns (None, expiry) when the app isn't installed on that repo, so the |
| 75 | + caller falls back to a PAT instead of retrying the lookup for every job. |
| 76 | + """ |
| 77 | + owner, repo = full_name.split("/", 1) |
| 78 | + integration = GithubIntegration( |
| 79 | + auth=Auth.AppAuth(int(GITHUB_APP_ID), app_private_key()) |
| 80 | + ) |
| 81 | + |
| 82 | + try: |
| 83 | + installation = integration.get_repo_installation(owner, repo) |
| 84 | + except UnknownObjectException: |
| 85 | + return None, time.time() + NO_INSTALLATION_TTL |
| 86 | + |
| 87 | + token = integration.get_access_token(installation.id) |
| 88 | + return token.token, token.expires_at.timestamp() - TOKEN_EXPIRY_MARGIN |
| 89 | + |
| 90 | + |
| 91 | +def installation_token(full_name): |
| 92 | + """Cached installation token for full_name, or None if unavailable.""" |
| 93 | + if not GITHUB_APP_ID or not GITHUB_APP_PRIVATE_KEY: |
| 94 | + return None |
| 95 | + |
| 96 | + cached = _token_cache.get(full_name) |
| 97 | + if cached and time.time() < cached[1]: |
| 98 | + return cached[0] |
| 99 | + |
| 100 | + try: |
| 101 | + token, expires_at = fetch_installation_token(full_name) |
| 102 | + except Exception as err: |
| 103 | + # Deliberately broad: a bad app id, an unparseable private key or a |
| 104 | + # GitHub blip must degrade to the PAT pool rather than raise. Not cached |
| 105 | + # either, so the next invocation retries. |
| 106 | + print(f"ERROR minting installation token for {full_name}: {err}") |
| 107 | + return None |
| 108 | + |
| 109 | + _token_cache[full_name] = (token, expires_at) |
| 110 | + return token |
| 111 | + |
| 112 | + |
| 113 | +def fetch_log(full_name, job_id, token): |
| 114 | + url = f"{GITHUB_API_URL}/repos/{full_name}/actions/jobs/{job_id}/logs" |
| 115 | + headers = { |
| 116 | + "Accept": "application/vnd.github.v3+json", |
| 117 | + "Authorization": "token " + token, |
| 118 | + } |
| 119 | + return requests.get(url, headers=headers, timeout=30) |
| 120 | + |
| 121 | + |
| 122 | +def log_object_path(full_name, job_id): |
| 123 | + """S3 key for a job's log. |
| 124 | +
|
| 125 | + pytorch/pytorch is unprefixed for historical reasons and must stay that way: |
| 126 | + default.workflow_job's `log_url` ALIAS in ClickHouse derives the URL from |
| 127 | + this exact scheme. |
| 128 | + """ |
| 129 | + if full_name == "pytorch/pytorch": |
| 130 | + return f"log/{job_id}" |
| 131 | + return f"log/{full_name}/{job_id}" |
| 132 | + |
| 133 | + |
| 134 | +def download_log(full_name, conclusion, job_id): |
| 135 | + """Fetch a job log from GitHub and archive it. Returns True when stored.""" |
| 136 | + response = None |
| 137 | + |
| 138 | + app_token = installation_token(full_name) |
| 139 | + if app_token: |
| 140 | + response = fetch_log(full_name, job_id, app_token) |
| 141 | + if response.status_code in FALLBACK_STATUSES: |
| 142 | + # Stop using the app until its window resets, otherwise every job |
| 143 | + # for the rest of the hour pays for a doomed request first. |
| 144 | + _token_cache[full_name] = (None, cool_off_until(response)) |
| 145 | + print( |
| 146 | + f"App auth returned {response.status_code} for {full_name} " |
| 147 | + f"job {job_id}, falling back to a PAT" |
| 148 | + ) |
| 149 | + response = None |
| 150 | + |
| 151 | + if response is None: |
| 152 | + if not GITHUB_TOKENS: |
| 153 | + print(f"ERROR no usable credential for {full_name} job {job_id}") |
| 154 | + return False |
| 155 | + response = fetch_log(full_name, job_id, random.choice(GITHUB_TOKENS.split(","))) |
| 156 | + |
| 157 | + if not response.ok: |
| 158 | + # Bail out rather than archive the API error body as if it were the log |
| 159 | + print( |
| 160 | + f"ERROR {response.status_code} downloading log for {full_name} job {job_id}" |
| 161 | + ) |
| 162 | + return False |
| 163 | + |
| 164 | + # Note: brotli would compress better, but is annoying to add as a dep |
| 165 | + # If space becomes a problem it's roughly ~2x better in TEXT_MODE |
| 166 | + s3.Object(BUCKET_NAME, log_object_path(full_name, job_id)).put( |
| 167 | + Body=gzip.compress(response.content), |
| 168 | + ContentType="text/plain", |
| 169 | + ContentEncoding="gzip", |
| 170 | + Metadata={"conclusion": conclusion or ""}, |
| 171 | + ) |
| 172 | + return True |
| 173 | + |
| 174 | + |
| 175 | +def parse_event(event): |
| 176 | + """Validate the invoke payload, returning (full_name, conclusion, job_id). |
| 177 | +
|
| 178 | + Raises ValueError on anything malformed. The caller is an async invoke, so a |
| 179 | + raised error is retried twice by Lambda and then lands in the DLQ, which is |
| 180 | + what we want for a payload we can't interpret. |
| 181 | + """ |
| 182 | + if not isinstance(event, dict): |
| 183 | + raise ValueError(f"expected a JSON object, got {type(event).__name__}") |
| 184 | + |
| 185 | + full_name = event.get("repo") |
| 186 | + if not full_name or "/" not in full_name: |
| 187 | + raise ValueError(f"missing or malformed 'repo': {full_name!r}") |
| 188 | + |
| 189 | + try: |
| 190 | + job_id = int(event["job_id"]) |
| 191 | + except (KeyError, TypeError, ValueError): |
| 192 | + raise ValueError(f"missing or non-numeric 'job_id': {event.get('job_id')!r}") |
| 193 | + |
| 194 | + return full_name, event.get("conclusion"), job_id |
| 195 | + |
| 196 | + |
| 197 | +def lambda_handler(event, context): |
| 198 | + full_name, conclusion, job_id = parse_event(event) |
| 199 | + |
| 200 | + try: |
| 201 | + stored = download_log(full_name, conclusion, job_id) |
| 202 | + except requests.RequestException as err: |
| 203 | + # A GitHub blip is worth a Lambda retry, so let it propagate. |
| 204 | + print(f"ERROR downloading log for {full_name} job {job_id}: {err}") |
| 205 | + raise |
| 206 | + |
| 207 | + return {"repo": full_name, "job_id": job_id, "stored": stored} |
0 commit comments