Add the log uploader bot

torchci/lib/bot/index.ts

@@ -8,6 +8,7 @@ import ciflowPushTrigger from "./ciflowPushTrigger";
 import codevNoWritePerm from "./codevNoWritePermBot";
 import crcrOncallBot from "./crcrOncallBot";
 import drciBot from "./drciBot";
+import logUploader from "./logUploader";
 import nitpickBot from "./nitpickBot";
 import pytorchBot from "./pytorchBot";
 import retryBot from "./retryBot";
@@ -25,6 +26,7 @@ export default function bot(app: Probot) {
   codevNoWritePerm(app);
   crcrOncallBot(app);
   drciBot(app);
+  logUploader(app);
   nitpickBot(app);
   pytorchBot(app);
   retryBot(app);

torchci/lib/bot/logUploader.ts

@@ -0,0 +1,71 @@
+import { invokeLogUploader } from "lib/lambda";
+import { Context, Probot } from "probot";
+import { isPyTorchbotSupportedOrg, isVLLM } from "./utils";
+
+/**
+ * Comma-separated list of repos whose logs this handler uploads. Entries are
+ * either `owner/repo` or `owner/*`; an empty or unset value disables the
+ * handler entirely.
+ *
+ * This exists so the cutover from the github-status-test webhook can be done a
+ * repo at a time. While a repo still has that webhook, both paths write the same
+ * S3 key -- harmless, but it doubles classifier calls, so the allowlist is what
+ * bounds the overlap. Rolling back is editing this variable.
+ */
+export function parseRepoAllowlist(raw: string | undefined): Set<string> {
+  return new Set(
+    (raw ?? "")
+      .split(",")
+      .map((entry) => entry.trim().toLowerCase())
+      .filter((entry) => entry.length > 0)
+  );
+}
+
+export function isRepoEnabled(
+  allowlist: Set<string>,
+  owner: string,
+  repo: string
+): boolean {
+  return (
+    allowlist.has(`${owner}/${repo}`.toLowerCase()) ||
+    allowlist.has(`${owner.toLowerCase()}/*`)
+  );
+}
+
+async function handleCompletedJob(event: Context<"workflow_job">) {
+  if (event.payload.action !== "completed") {
+    return;
+  }
+
+  const owner = event.payload.repository.owner.login;
+  const repo = event.payload.repository.name;
+  if (!isPyTorchbotSupportedOrg(owner) && !isVLLM(owner)) {
+    event.log(`${__filename} isn't enabled on ${owner}'s repos`);
+    return;
+  }
+
+  const allowlist = parseRepoAllowlist(process.env.LOG_UPLOADER_REPOS);
+  if (!isRepoEnabled(allowlist, owner, repo)) {
+    return;
+  }
+
+  try {
+    await invokeLogUploader({
+      repo: event.payload.repository.full_name,
+      job_id: event.payload.workflow_job.id,
+      conclusion: event.payload.workflow_job.conclusion,
+    });
+  } catch (error) {
+    // Never fail the webhook over a log. GitHub would redeliver the whole event,
+    // re-running every other handler, to retry something Dr.CI already repairs
+    // on its own via backfillMissingLog.
+    event.log.error(
+      `Failed to queue a log upload for ${event.payload.repository.full_name} ` +
+        `job ${event.payload.workflow_job.id}: ${error}`
+    );
+  }
+}
+
+export default function logUploader(app: Probot) {
+  app.on("workflow_job", handleCompletedJob);
+}

torchci/lib/lambda.ts

@@ -0,0 +1,72 @@
+import {
+  InvocationType,
+  InvokeCommand,
+  LambdaClient,
+} from "@aws-sdk/client-lambda";
+
+export const GHA_LOG_UPLOADER_FUNCTION = "gha-log-uploader";
+
+export class MissingAwsCredentialsError extends Error {
+  constructor() {
+    // The SDK's own failure here is "Resolved credential object is not valid",
+    // which says nothing about which variable to go set.
+    super(
+      "OUR_AWS_ACCESS_KEY_ID / OUR_AWS_SECRET_ACCESS_KEY are not set, " +
+        "cannot reach the log uploader"
+    );
+    this.name = "MissingAwsCredentialsError";
+  }
+}
+
+export function getLambdaClient(): LambdaClient {
+  const accessKeyId = process.env.OUR_AWS_ACCESS_KEY_ID;
+  const secretAccessKey = process.env.OUR_AWS_SECRET_ACCESS_KEY;
+  if (!accessKeyId || !secretAccessKey) {
+    throw new MissingAwsCredentialsError();
+  }
+
+  return new LambdaClient({
+    region: "us-east-1",
+    credentials: { accessKeyId, secretAccessKey },
+    // This call sits on the webhook ack path, so it has to fail fast rather than
+    // fail well. The SDK defaults to 3 attempts and no socket timeout, which on a
+    // black-holed endpoint hangs for the OS default (~75s) per attempt; measured,
+    // a connectionTimeout of 1s aborts at ~1.5s instead. Worst case here is
+    // roughly 4s, against a 10s GitHub webhook timeout that every other handler
+    // also has to fit inside.
+    maxAttempts: 2,
+    requestHandler: { connectionTimeout: 1000, requestTimeout: 2000 },
+  });
+}
+
+export interface LogUploadRequest {
+  repo: string;
+  job_id: number;
+  conclusion?: string | null;
+}
+
+/**
+ * Ask gha-log-uploader to archive a job's log to S3.
+ *
+ * Invoked with InvocationType Event, so this returns as soon as Lambda accepts
+ * the payload rather than waiting on the GitHub download. That matters because
+ * the caller is a Probot webhook handler: Probot only acks GitHub once every
+ * handler resolves, and nothing runs after a Vercel function returns, so the
+ * handoff has to be both awaited and bounded.
+ *
+ * Delivery failures are Lambda's problem from here -- it retries twice and then
+ * DLQs. Failures to hand off at all reject, and every caller treats that as
+ * non-fatal: Dr.CI re-requests a missing log through backfillMissingLog on its
+ * next run, so losing a handoff costs a log, never a webhook.
+ */
+export async function invokeLogUploader(
+  request: LogUploadRequest
+): Promise<void> {
+  await getLambdaClient().send(
+    new InvokeCommand({
+      FunctionName: GHA_LOG_UPLOADER_FUNCTION,
+      InvocationType: InvocationType.Event,
+      Payload: Buffer.from(JSON.stringify(request)),
+    })
+  );
+}

torchci/package.json

@@ -15,6 +15,7 @@
   },
   "dependencies": {
     "@aws-sdk/client-dynamodb": "^3.347.1",
+    "@aws-sdk/client-lambda": "^3.347.1",
     "@aws-sdk/client-s3": "^3.347.1",
     "@aws-sdk/lib-dynamodb": "^3.72.0",
     "@clickhouse/client": "^1.11.1",

torchci/test/lambdaClient.test.ts

@@ -0,0 +1,65 @@
+import {
+  getLambdaClient,
+  invokeLogUploader,
+  MissingAwsCredentialsError,
+} from "lib/lambda";
+
+describe("getLambdaClient", () => {
+  const saved = {
+    id: process.env.OUR_AWS_ACCESS_KEY_ID,
+    secret: process.env.OUR_AWS_SECRET_ACCESS_KEY,
+  };
+
+  afterEach(() => {
+    process.env.OUR_AWS_ACCESS_KEY_ID = saved.id;
+    process.env.OUR_AWS_SECRET_ACCESS_KEY = saved.secret;
+  });
+
+  function setCredentials(id?: string, secret?: string) {
+    if (id === undefined) {
+      delete process.env.OUR_AWS_ACCESS_KEY_ID;
+    } else {
+      process.env.OUR_AWS_ACCESS_KEY_ID = id;
+    }
+    if (secret === undefined) {
+      delete process.env.OUR_AWS_SECRET_ACCESS_KEY;
+    } else {
+      process.env.OUR_AWS_SECRET_ACCESS_KEY = secret;
+    }
+  }
+
+  test.each([
+    ["neither is set", undefined, undefined],
+    ["only the key id is set", "AKIA", undefined],
+    ["only the secret is set", undefined, "shh"],
+    ["the key id is empty", "", "shh"],
+  ])("throws a named error when %s", (_label, id, secret) => {
+    setCredentials(id, secret);
+    expect(() => getLambdaClient()).toThrow(MissingAwsCredentialsError);
+  });
+
+  test("the error names the variables to set", () => {
+    setCredentials(undefined, undefined);
+    // The SDK's own message is "Resolved credential object is not valid", which
+    // gives whoever is paged nothing to act on.
+    expect(() => getLambdaClient()).toThrow(/OUR_AWS_ACCESS_KEY_ID/);
+  });
+
+  test("builds a client when both are set", () => {
+    setCredentials("AKIA", "shh");
+    expect(getLambdaClient()).toBeDefined();
+  });
+
+  test("bounds retries and socket waits", async () => {
+    setCredentials("AKIA", "shh");
+    const config = getLambdaClient().config;
+    expect(await config.maxAttempts()).toBe(2);
+  });
+
+  test("invokeLogUploader rejects rather than hanging with no credentials", async () => {
+    setCredentials(undefined, undefined);
+    await expect(
+      invokeLogUploader({ repo: "pytorch/pytorch", job_id: 1 })
+    ).rejects.toThrow(MissingAwsCredentialsError);
+  });
+});

torchci/test/logUploader.test.ts

@@ -0,0 +1,155 @@
+import logUploader, {
+  isRepoEnabled,
+  parseRepoAllowlist,
+} from "lib/bot/logUploader";
+import * as lambda from "lib/lambda";
+import nock from "nock";
+import { Probot } from "probot";
+import * as utils from "./utils";
+
+nock.disableNetConnect();
+
+function workflowJobPayload({
+  owner = "meta-pytorch",
+  name = "torchcomms",
+  action = "completed",
+  id = 12345,
+  conclusion = "failure" as string | null,
+} = {}) {
+  return {
+    action,
+    workflow_job: { id, conclusion },
+    repository: {
+      full_name: `${owner}/${name}`,
+      name,
+      owner: { login: owner },
+    },
+  };
+}
+
+describe("parseRepoAllowlist", () => {
+  test("an unset value disables the handler", () => {
+    expect(parseRepoAllowlist(undefined).size).toBe(0);
+    expect(parseRepoAllowlist("").size).toBe(0);
+  });
+
+  test("entries are trimmed and lowercased", () => {
+    expect(parseRepoAllowlist(" Pytorch/PyTorch , meta-pytorch/* ")).toEqual(
+      new Set(["pytorch/pytorch", "meta-pytorch/*"])
+    );
+  });
+
+  test("empty entries from a trailing comma are dropped", () => {
+    expect(parseRepoAllowlist("pytorch/pytorch,,").size).toBe(1);
+  });
+});
+
+describe("isRepoEnabled", () => {
+  test("matches an exact repo", () => {
+    const allowlist = parseRepoAllowlist("pytorch/pytorch");
+    expect(isRepoEnabled(allowlist, "pytorch", "pytorch")).toBe(true);
+    expect(isRepoEnabled(allowlist, "pytorch", "executorch")).toBe(false);
+  });
+
+  test("matches a whole org via a wildcard", () => {
+    const allowlist = parseRepoAllowlist("meta-pytorch/*");
+    expect(isRepoEnabled(allowlist, "meta-pytorch", "torchcomms")).toBe(true);
+    expect(isRepoEnabled(allowlist, "meta-pytorch", "monarch")).toBe(true);
+    expect(isRepoEnabled(allowlist, "pytorch", "pytorch")).toBe(false);
+  });
+
+  test("an org wildcard does not leak into a similarly named org", () => {
+    const allowlist = parseRepoAllowlist("pytorch/*");
+    expect(isRepoEnabled(allowlist, "meta-pytorch", "torchcomms")).toBe(false);
+  });
+
+  test("comparison is case insensitive", () => {
+    const allowlist = parseRepoAllowlist("PyTorch/PyTorch");
+    expect(isRepoEnabled(allowlist, "pytorch", "pytorch")).toBe(true);
+  });
+});
+
+describe("logUploader", () => {
+  let probot: Probot;
+  let invoke: jest.SpyInstance;
+
+  beforeEach(() => {
+    probot = utils.testProbot();
+    probot.load(logUploader);
+    invoke = jest.spyOn(lambda, "invokeLogUploader").mockResolvedValue();
+    process.env.LOG_UPLOADER_REPOS = "meta-pytorch/*,pytorch/pytorch";
+  });
+
+  afterEach(() => {
+    jest.restoreAllMocks();
+    nock.cleanAll();
+    delete process.env.LOG_UPLOADER_REPOS;
+  });
+
+  async function receive(payload: any) {
+    await probot.receive({ name: "workflow_job", payload, id: "1" } as any);
+  }
+
+  test("queues an upload for a completed job on an allowlisted repo", async () => {
+    await receive(workflowJobPayload());
+
+    expect(invoke).toHaveBeenCalledWith({
+      repo: "meta-pytorch/torchcomms",
+      job_id: 12345,
+      conclusion: "failure",
+    });
+  });
+
+  test("ignores anything but the completed action", async () => {
+    // workflow_job also fires on queued and in_progress, where there is no log
+    // to fetch yet. Uploading then would archive a truncated log.
+    await receive(workflowJobPayload({ action: "queued" }));
+    await receive(workflowJobPayload({ action: "in_progress" }));
+
+    expect(invoke).not.toHaveBeenCalled();
+  });
+
+  test("skips a repo that is not on the allowlist", async () => {
+    await receive(workflowJobPayload({ owner: "pytorch", name: "executorch" }));
+
+    expect(invoke).not.toHaveBeenCalled();
+  });
+
+  test("skips an org the bot does not serve, even if allowlisted", async () => {
+    // The allowlist narrows the org gate, it must not widen it.
+    process.env.LOG_UPLOADER_REPOS = "someoneelse/*";
+    await receive(workflowJobPayload({ owner: "someoneelse", name: "repo" }));
+
+    expect(invoke).not.toHaveBeenCalled();
+  });
+
+  test("does nothing when the allowlist is unset", async () => {
+    delete process.env.LOG_UPLOADER_REPOS;
+    await receive(workflowJobPayload());
+
+    expect(invoke).not.toHaveBeenCalled();
+  });
+
+  test("passes a null conclusion through rather than dropping the job", async () => {
+    await receive(workflowJobPayload({ conclusion: null }));
+
+    expect(invoke).toHaveBeenCalledWith(
+      expect.objectContaining({ conclusion: null })
+    );
+  });
+
+  test.each([
+    ["the role cannot invoke the function", "AccessDeniedException"],
+    ["credentials are missing entirely", "MissingAwsCredentialsError"],
+    ["the Lambda API is unreachable", "TimeoutError"],
+    ["Lambda throttles us", "TooManyRequestsException"],
+  ])("a failed invoke does not fail the webhook when %s", async (_l, name) => {
+    // Throwing here would make GitHub redeliver the event and re-run every other
+    // handler, to retry something Dr.CI repairs on its own.
+    const error = new Error(name);
+    error.name = name;
+    invoke.mockRejectedValue(error);
+
+    await expect(receive(workflowJobPayload())).resolves.not.toThrow();
+  });
+});