Add metric for GitHub API Rate Limit (#5654)

This captures the rate limit values from GitHub for the ALI account user/process. We can use this to graph and track how much of the API rate limit is used by the CI infrastructure and flag when we are getting too close to the overall limit. Closes: pytorch/ci-infra#273 Signed-off-by: Thanh Ha <[email protected]> Co-authored-by: Jean Schmidt <[email protected]>
pytorch · Sep 24, 2024 · 2791b5a · 2791b5a
1 parent 68d2cf7
commit 2791b5a
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 0 deletions.
diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/gh-runners.ts b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/gh-runners.ts
@@ -514,3 +514,33 @@ export async function createRegistrationTokenOrg(
     throw e;
   }
 }
+
+export async function getGitHubRateLimit(repo: Repo, installationId: number, metrics: Metrics): Promise<void> {
+  try {
+    const { used, limit, remaining } = await locallyCached('ghRunners', 'getGitHubRateLimit', 10, async () => {
+      try {
+        const client = await createGitHubClientForRunnerRepo(repo, metrics);
+
+        const rateLimit = await expBackOff(() => {
+          return metrics.trackRequest(metrics.getGitHubRateLimitSuccess, metrics.getGitHubRateLimitFailure, () => {
+            return client.rateLimit.get();
+          });
+        });
+
+        const limit = Number(rateLimit.headers['x-ratelimit-limit']);
+        const remaining = Number(rateLimit.headers['x-ratelimit-remaining']);
+        const used = Number(rateLimit.headers['x-ratelimit-used']);
+
+        return { used, limit, remaining };
+      } catch (e) {
+        console.error(`[getGitHubRateLimit]: <anonymous> ${e}`);
+        throw e;
+      }
+    });
+
+    metrics.gitHubRateLimitStats(limit, remaining, used);
+  } catch (e) {
+    console.error(`[getGitHubRateLimit]: ${e}`);
+    throw e;
+  }
+}
diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/metrics.ts b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/metrics.ts
@@ -241,6 +241,29 @@ export class Metrics {
   }
 
   // GitHub API CALLS
+  /* istanbul ignore next */
+  getGitHubRateLimitSuccess(ms: number) {
+    this.countEntry(`gh.calls.total`, 1);
+    this.countEntry(`gh.calls.getGitHubRateLimit.count`, 1);
+    this.countEntry(`gh.calls.getGitHubRateLimit.success`, 1);
+    this.countEntry(`gh.calls.getGitHubRateLimit.wallclock`, ms);
+  }
+
+  /* istanbul ignore next */
+  getGitHubRateLimitFailure(ms: number) {
+    this.countEntry(`gh.calls.total`, 1);
+    this.countEntry(`gh.calls.getGitHubRateLimit.count`, 1);
+    this.countEntry(`gh.calls.getGitHubRateLimit.failure`, 1);
+    this.countEntry(`gh.calls.getGitHubRateLimit.wallclock`, ms);
+  }
+
+  /* istanbul ignore next */
+  gitHubRateLimitStats(limit: number, remaining: number, used: number) {
+    this.addEntry(`gh.calls.ratelimit.limit`, limit);
+    this.addEntry(`gh.calls.ratelimit.remaining`, remaining);
+    this.addEntry(`gh.calls.ratelimit.used`, used);
+  }
+
   /* istanbul ignore next */
   createAppAuthGHCallSuccess(ms: number) {
     this.countEntry(`gh.calls.total`, 1);

diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/scale-up.ts b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/scale-up.ts
@@ -7,6 +7,7 @@ import {
   getRunnerTypes,
   listGithubRunnersOrg,
   listGithubRunnersRepo,
+  getGitHubRateLimit,
 } from './gh-runners';
 
 import { Config } from './config';
@@ -54,6 +55,8 @@ export async function scaleUp(
   metrics.runRepo(repo);
   metrics.run();
 
+  getGitHubRateLimit(repo, Number(payload.installationId), metrics);
+
   const runnerTypes = await getRunnerTypes(
     {
       owner: repo.owner,