Add the call-log-classifier lambda

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

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/<job_id>` | `pytorch/pytorch` (unprefixed for historical reasons) |
+| `log/<owner>/<repo>/<job_id>` | `<owner>/<repo>` |
+
+Anything else is skipped and logged. `log/<owner>/<job_id>` 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::<usize>()` — 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.

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::<usize>()`, 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/<id>` means
+    pytorch/pytorch and `log/<owner>/<repo>/<id>` 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/<owner>/<id>` 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}"
+            )

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/<owner>/<id>` 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()