From 55e79904caa4168d8806b33de91be3a34fab07ca Mon Sep 17 00:00:00 2001 From: Huy Do Date: Thu, 20 Aug 2026 12:36:24 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- aws/lambda/call-log-classifier/Makefile | 17 +++ aws/lambda/call-log-classifier/README.md | 68 +++++++++++ .../call-log-classifier/lambda_function.py | 114 +++++++++++++++++ .../test_lambda_function.py | 115 ++++++++++++++++++ 4 files changed, 314 insertions(+) create mode 100644 aws/lambda/call-log-classifier/Makefile create mode 100644 aws/lambda/call-log-classifier/README.md create mode 100644 aws/lambda/call-log-classifier/lambda_function.py create mode 100644 aws/lambda/call-log-classifier/test_lambda_function.py diff --git a/aws/lambda/call-log-classifier/Makefile b/aws/lambda/call-log-classifier/Makefile new file mode 100644 index 0000000000..1fdf10f11d --- /dev/null +++ b/aws/lambda/call-log-classifier/Makefile @@ -0,0 +1,17 @@ +ZIP := deployment.zip +FUNCTION := call-log-classifier + +# boto3 ships in the Lambda python runtime and nothing else is imported, so the +# package is just the handler. +prepare: clean + mkdir -p deployment + cp lambda_function.py ./deployment/. + cd ./deployment && zip -q -r ../$(ZIP) . + +deploy: prepare + aws lambda update-function-code --function-name $(FUNCTION) --zip-file fileb://$(ZIP) + +clean: + rm -rf deployment $(ZIP) + +.PHONY: prepare deploy clean diff --git a/aws/lambda/call-log-classifier/README.md b/aws/lambda/call-log-classifier/README.md new file mode 100644 index 0000000000..9919a5aeb4 --- /dev/null +++ b/aws/lambda/call-log-classifier/README.md @@ -0,0 +1,68 @@ +# call-log-classifier + +Triggers log classification when a job log lands in +`s3://ossci-raw-job-status/log/`. Wired to `s3:ObjectCreated:*` on that prefix, so +whatever wrote the log — `gha-log-uploader`, torchci's backfill route, or a person +running `aws s3 cp` — gets it classified. + +This exists so uploading and classifying are decoupled. `github-status-test` +called the classifier inline with an untimed `urlopen` and blocked until it +finished, which is what produced its 274s/344s/400s/900s duration tails. + +## Key shapes + +| Key | Repo | +| --- | --- | +| `log/` | `pytorch/pytorch` (unprefixed for historical reasons) | +| `log///` | `/` | + +Anything else is skipped and logged. `log//` in particular cannot be +attributed to a repo, so it is ignored rather than guessed at. + +## Why it does not use the classifier's function URL + +`log_classifier` has a Lambda function URL with `AuthType: NONE`. This function +reaches it through `lambda:InvokeFunction` instead, so the path from S3 to +classification never crosses a public endpoint. + +`log_classifier` is built on `lambda_http` with only the `apigw_http` feature, so +it expects an API Gateway HTTP API v2.0 request. `classifier_payload()` reproduces +that shape. Verified against the deployed function: a payload with no `job_id` +returns its 400 `no job id provided` branch, and one with a non-numeric `job_id` +fails inside its `parse::()` — together showing that both the envelope and +the query string are read from a direct invoke. + +The invoke is asynchronous. Classification can take minutes and nothing here reads +the result. + +## Relationship to keep-going-call-log-classifier + +`keep-going-call-log-classifier` does the same job for the `temp_logs/` prefix on +`gha-artifacts`, and still calls the classifier over its public function URL. It +is deliberately left alone during the `github-status-test` cutover. Folding the +two together, and then removing the public function URL, is cleanup for +afterwards. + +## One-time AWS setup + +1. Create the function: python3.12, handler `lambda_function.lambda_handler`. Only + boto3 is imported, so the package is just the handler. +2. Give its execution role `lambda:InvokeFunction` on + `arn:aws:lambda:us-east-1:308535385114:function:log_classifier`, plus the usual + CloudWatch Logs permissions. +3. Let S3 invoke it: + ``` + aws lambda add-permission --function-name call-log-classifier \ + --statement-id s3-ossci-raw-job-status --action lambda:InvokeFunction \ + --principal s3.amazonaws.com \ + --source-arn arn:aws:s3:::ossci-raw-job-status \ + --source-account 308535385114 + ``` +4. Add an `s3:ObjectCreated:*` notification on `ossci-raw-job-status` filtered to + prefix `log/`. **Read the existing configuration and add to it** — + `put-bucket-notification-configuration` replaces the whole document, and that + bucket already carries twelve `clickhouse-replicator-s3` rules that must + survive. None of them overlap `log/`. + +Do not enable the notification before `gha-log-uploader` is live, or every log the +old lambda writes gets classified twice. diff --git a/aws/lambda/call-log-classifier/lambda_function.py b/aws/lambda/call-log-classifier/lambda_function.py new file mode 100644 index 0000000000..3e22d887e1 --- /dev/null +++ b/aws/lambda/call-log-classifier/lambda_function.py @@ -0,0 +1,114 @@ +"""Kick off log classification when a job log lands in S3. + +Wired to `s3:ObjectCreated:*` on the `log/` prefix of `ossci-raw-job-status`, so +whatever put the log there -- gha-log-uploader, torchci's backfill route, or a +person running `aws s3 cp` -- gets it classified. Decoupling this from the +uploader is what keeps a slow classification from showing up as uploader latency. + +`log_classifier` is reached through `lambda:InvokeFunction` rather than its public +function URL. It is built on lambda_http with only the `apigw_http` feature, so it +expects an API Gateway HTTP API v2.0 request; PAYLOAD_TEMPLATE reproduces that +shape. Verified against the deployed function: a payload with no `job_id` returns +its 400 "no job id provided" branch, and a non-numeric one fails in its +`parse::()`, which together show both the envelope and the query string +are read. + +Sibling `keep-going-call-log-classifier` does the same job for the `temp_logs/` +prefix on `gha-artifacts`. The two should be folded together once the +github-status-test cutover is finished. +""" + +import json +from typing import Any, Optional, Tuple + +import boto3 + + +LOG_CLASSIFIER_FUNCTION = "log_classifier" +DEFAULT_REPO = "pytorch/pytorch" +LOG_PREFIX = "log/" + +lambda_client = boto3.client("lambda") + + +def parse_key(key: str) -> Optional[Tuple[str, int]]: + """Map an S3 key under `log/` to (repo, job_id), or None if it isn't one. + + pytorch/pytorch is unprefixed for historical reasons, so `log/` means + pytorch/pytorch and `log///` names its repo explicitly. + """ + if not key.startswith(LOG_PREFIX): + return None + + parts = key[len(LOG_PREFIX) :].split("/") + if len(parts) == 1: + repo = DEFAULT_REPO + elif len(parts) == 3: + repo = f"{parts[0]}/{parts[1]}" + else: + # Neither shape. Includes `log//` and anything deeper, which + # we have no way to attribute to a repo. + return None + + try: + return repo, int(parts[-1]) + except ValueError: + return None + + +def classifier_payload(repo: str, job_id: int) -> dict: + """An API Gateway HTTP API v2.0 request, which is what lambda_http parses.""" + query = {"job_id": str(job_id), "repo": repo} + return { + "version": "2.0", + "routeKey": "$default", + "rawPath": "/", + "rawQueryString": f"job_id={job_id}&repo={repo}", + "headers": {}, + "queryStringParameters": query, + "requestContext": { + "accountId": "308535385114", + "apiId": "call-log-classifier", + "domainName": "lambda-invoke", + "domainPrefix": "lambda-invoke", + "http": { + "method": "GET", + "path": "/", + "protocol": "HTTP/1.1", + "sourceIp": "127.0.0.1", + "userAgent": "call-log-classifier", + }, + "requestId": f"call-log-classifier-{job_id}", + "routeKey": "$default", + "stage": "$default", + "time": "01/Jan/1970:00:00:00 +0000", + "timeEpoch": 0, + }, + "isBase64Encoded": False, + } + + +def lambda_handler(event: Any, context: Any) -> None: + for record in event.get("Records", []): + key = record.get("s3", {}).get("object", {}).get("key", "") + + parsed = parse_key(key) + if parsed is None: + print(f"Skipping key that isn't a job log: key={key}") + continue + + repo, job_id = parsed + try: + # Async: classification can take minutes and nothing here reads the + # result, so blocking on it would only burn this function's runtime. + lambda_client.invoke( + FunctionName=LOG_CLASSIFIER_FUNCTION, + InvocationType="Event", + Payload=json.dumps(classifier_payload(repo, job_id)).encode(), + ) + except Exception as error: + # One bad key must not strand the rest of the batch. + print( + f"Failed to call log classifier for job_id={job_id}, " + f"repo={repo}, key={key}, error={error}" + ) diff --git a/aws/lambda/call-log-classifier/test_lambda_function.py b/aws/lambda/call-log-classifier/test_lambda_function.py new file mode 100644 index 0000000000..5196c70d03 --- /dev/null +++ b/aws/lambda/call-log-classifier/test_lambda_function.py @@ -0,0 +1,115 @@ +import json +import unittest +from unittest.mock import patch + +from lambda_function import classifier_payload, lambda_handler, parse_key + + +def s3_event(*keys: str) -> dict: + return { + "Records": [ + { + "s3": { + "bucket": {"name": "ossci-raw-job-status"}, + "object": {"key": key}, + } + } + for key in keys + ] + } + + +class TestParseKey(unittest.TestCase): + def test_bare_id_is_pytorch_pytorch(self): + self.assertEqual(parse_key("log/123345"), ("pytorch/pytorch", 123345)) + + def test_prefixed_key_names_its_repo(self): + self.assertEqual( + parse_key("log/pytorch/executorch/999"), ("pytorch/executorch", 999) + ) + + def test_meta_pytorch_repo(self): + self.assertEqual( + parse_key("log/meta-pytorch/torchcomms/42"), ("meta-pytorch/torchcomms", 42) + ) + + def test_ignores_other_prefixes(self): + # The notification is filtered to `log/`, but `classification/` and + # `logs_something/` must not be mistaken for it if that ever changes. + self.assertIsNone(parse_key("classification/123")) + self.assertIsNone(parse_key("log_archive/123")) + + def test_ignores_a_non_numeric_id(self): + self.assertIsNone(parse_key("log/not-a-number")) + self.assertIsNone(parse_key("log/pytorch/executorch/not-a-number")) + + def test_ignores_an_unattributable_depth(self): + # `log//` gives no repo, and anything deeper is not ours. + self.assertIsNone(parse_key("log/pytorch/123")) + self.assertIsNone(parse_key("log/a/b/c/123")) + + def test_ignores_a_directory_marker(self): + self.assertIsNone(parse_key("log/")) + + +class TestClassifierPayload(unittest.TestCase): + def test_is_a_v2_request_the_classifier_can_parse(self): + payload = classifier_payload("pytorch/executorch", 999) + # log_classifier builds on lambda_http with only the apigw_http feature, + # so version 2.0 and requestContext.http are what make it deserialize. + self.assertEqual(payload["version"], "2.0") + self.assertIn("http", payload["requestContext"]) + self.assertEqual( + payload["queryStringParameters"], + {"job_id": "999", "repo": "pytorch/executorch"}, + ) + self.assertEqual( + payload["rawQueryString"], "job_id=999&repo=pytorch/executorch" + ) + + def test_is_json_serializable(self): + json.dumps(classifier_payload("pytorch/pytorch", 1)) + + +class TestLambdaHandler(unittest.TestCase): + def test_invokes_the_classifier_asynchronously(self): + with patch("lambda_function.lambda_client") as client: + lambda_handler(s3_event("log/123345"), None) + + kwargs = client.invoke.call_args.kwargs + self.assertEqual(kwargs["FunctionName"], "log_classifier") + # Event, not RequestResponse: nothing here reads the classification. + self.assertEqual(kwargs["InvocationType"], "Event") + self.assertEqual( + json.loads(kwargs["Payload"])["queryStringParameters"], + {"job_id": "123345", "repo": "pytorch/pytorch"}, + ) + + def test_handles_every_record_in_a_batch(self): + with patch("lambda_function.lambda_client") as client: + lambda_handler(s3_event("log/1", "log/pytorch/rl/2"), None) + + self.assertEqual(client.invoke.call_count, 2) + + def test_skips_a_key_it_cannot_attribute(self): + with patch("lambda_function.lambda_client") as client: + lambda_handler(s3_event("log/not-a-number"), None) + + client.invoke.assert_not_called() + + def test_one_failure_does_not_strand_the_batch(self): + with patch("lambda_function.lambda_client") as client: + client.invoke.side_effect = [RuntimeError("throttled"), None] + lambda_handler(s3_event("log/1", "log/2"), None) + + self.assertEqual(client.invoke.call_count, 2) + + def test_an_empty_event_is_a_noop(self): + with patch("lambda_function.lambda_client") as client: + lambda_handler({}, None) + + client.invoke.assert_not_called() + + +if __name__ == "__main__": + unittest.main()