Adds scaleUpHealing chron (#6412)

# TLDR 

This change introduces a new lambda `${var.environment}-scale-up-chron`.
With all the typescript code and required terraform changes.

# What is chaning?

This PR introduces the typescript code for the new lambda, and the
related terraform changes to run the lambda every 30 minutes. The lambda
should timeout in 15 minutes. Its permissions and access should be the
same as the one in scaleUp.

It goes to hud in a URL specified in a user configuration
`retry_scale_up_chron_hud_query_url` and gets a list of instance types
and number of jobs enqueued. It then synchronously tries to deploy those
runners.

It introduces 2 new parameters in the main module:
* `retry_scale_up_chron_hud_query_url` that for now should point to
https://hud.pytorch.org/api/clickhouse/queued_jobs_aggregate?parameters=%5B%5D
only in the installations that will benefit from it (both meta and linux
foundation PROD clusters, NOT canary) as when this variable is set to
empty string (default) the installation of this cron is not performed.
* `scale_config_org` that should point to the org where scale-config
files are defined. In our case it is `pytorch`.

[example of the
change](https://github.com/pytorch-labs/pytorch-gha-infra/pull/622/files)

# Why are we changing this?

We're introducing this change in order to provide a solution to help
recover lost requests for infra scaling. Its been proven for a while
that when there are github API outages we fail to get new jobs webhook
or fail to provision new runners. Most of the time our retry mechanism
is capable of dealing with the situation. But, in cases where we are not
receiving webhooks or other more esoteric problems, there is no way to
recover.

With this change, every 30 minutes, jobs enqueued for longer than 30
minutes for one of the autoscaled instance types, will trigger the
creation of those instances.

A few design decisions:
1 - Why rely on hud?
Hud currently already have these informations, so it should be simple to
just get it from there;

2 - Why not send a scale message and allow scaleUp to handle it?
We want to have isolation, in a way that we can easily circuit-break the
creation of enqueued instances. This also includes the isolation that
guarantees that if scaler is failing to deploy given instance type, this
mechanism won;t risk flood/overflow the main scaler that have to deal
with all other ones.

3 - why randomise the instance creation order?
So if some instance type is problematic, we are not absolutely
preventing the recovery of other instances types (just interfering).
Also we gain some time between instances creations of the same type,
allowing for a smoother operation.

4 - why a new lambda?
check number 2

# If something goes wrong?
Given we introduced as much as possible work to make sure there are
maximal isolation between the regular scaler and the cron recovery
scaler that we're introducing, we;re not foreseeing any potential gaps
that could break the main scaler and as a consequence introduce system
breakages.

Having said that, if you need to revert those changes from production,
just follow the steps:
https://docs.google.com/document/d/1nq3dx-_8wasii1koCkXJDSo3uz_0Ee8DzIS2-j2TOpA/edit?tab=t.0#heading=h.jwflgevrww4j

---------

Co-authored-by: Zain Rizvi <[email protected]>
Co-authored-by: Camyll Harajli <[email protected]>

Author: 3 people

terraform-aws-github-runner/main.tf

@@ -104,6 +104,7 @@ module "runners" {
   environment          = var.environment
   tags                 = local.tags
 
+  scale_config_org       = var.scale_config_org
   scale_config_repo      = var.scale_config_repo
   scale_config_repo_path = var.scale_config_repo_path
 
@@ -112,6 +113,8 @@ module "runners" {
     encrypt    = var.encrypt_secrets
   }
 
+  retry_scale_up_chron_hud_query_url = var.retry_scale_up_chron_hud_query_url
+
   must_have_issues_labels = var.must_have_issues_labels
   cant_have_issues_labels = var.cant_have_issues_labels
 

terraform-aws-github-runner/modules/runners/lambdas/runners/jest.config.js

@@ -11,5 +11,8 @@ module.exports = {
       lines: 93,
       statements: 94
     }
-  }
+  },
+  moduleNameMapper: {
+    axios: 'axios/dist/node/axios.cjs', // Allow axios to work in tests
+  },
 };

terraform-aws-github-runner/modules/runners/lambdas/runners/package.json

@@ -42,6 +42,7 @@
     "@types/uuid": "^9.0.1",
     "async-mutex": "^0.4.0",
     "aws-sdk": "^2.863.0",
+    "axios": "^1.7.7",
     "cron-parser": "^3.3.0",
     "generic-pool": "^3.9.0",
     "lru-cache": "^6.0.0",

terraform-aws-github-runner/modules/runners/lambdas/runners/src/lambda.test.ts

@@ -1,11 +1,12 @@
-import { scaleDown as scaleDownL, scaleUp as scaleUpL } from './lambda';
+import { scaleDown as scaleDownL, scaleUp as scaleUpL, scaleUpChron as scaleUpChronL } from './lambda';
 
 import nock from 'nock';
 import { Config } from './scale-runners/config';
 import { Context, SQSEvent, ScheduledEvent } from 'aws-lambda';
 import { mocked } from 'ts-jest/utils';
 import { scaleDown } from './scale-runners/scale-down';
 import { scaleUp, RetryableScalingError } from './scale-runners/scale-up';
+import { scaleUpChron } from './scale-runners/scale-up-chron';
 import { sqsSendMessages, sqsDeleteMessageBatch } from './scale-runners/sqs';
 import * as MetricsModule from './scale-runners/metrics';
 
@@ -21,8 +22,10 @@ jest.mock('aws-sdk', () => ({
 jest.mock('./scale-runners/scale-down');
 jest.mock('./scale-runners/scale-up');
 jest.mock('./scale-runners/sqs');
+jest.mock('./scale-runners/scale-up-chron');
 
-const metrics = new MetricsModule.ScaleUpMetrics();
+const mockScaleUpMetrics = new MetricsModule.ScaleUpMetrics();
+const mockScaleUpChronMetrics = new MetricsModule.ScaleUpChronMetrics();
 
 beforeEach(() => {
   jest.resetModules();
@@ -34,7 +37,7 @@ beforeEach(() => {
 describe('scaleUp', () => {
   beforeEach(() => {
     jest.spyOn(global.Math, 'random').mockReturnValue(1.0);
-    jest.spyOn(MetricsModule, 'ScaleUpMetrics').mockReturnValue(metrics);
+    jest.spyOn(MetricsModule, 'ScaleUpMetrics').mockReturnValue(mockScaleUpMetrics);
   });
 
   afterEach(() => {
@@ -55,8 +58,8 @@ describe('scaleUp', () => {
       callback,
     );
     expect(mockedScaleUp).toBeCalledTimes(2);
-    expect(mockedScaleUp).toBeCalledWith('aws:sqs', { id: 1 }, metrics);
-    expect(mockedScaleUp).toBeCalledWith('aws:sqs', { id: 2 }, metrics);
+    expect(mockedScaleUp).toBeCalledWith('aws:sqs', { id: 1 }, mockScaleUpMetrics);
+    expect(mockedScaleUp).toBeCalledWith('aws:sqs', { id: 2 }, mockScaleUpMetrics);
     expect(callback).toBeCalledTimes(1);
     expect(callback).toBeCalledWith(null);
   });
@@ -88,12 +91,12 @@ describe('scaleUp', () => {
       callback,
     );
     expect(mockedScaleUp).toBeCalledTimes(1);
-    expect(mockedScaleUp).toBeCalledWith('aws:sqs', { id: 1 }, metrics);
+    expect(mockedScaleUp).toBeCalledWith('aws:sqs', { id: 1 }, mockScaleUpMetrics);
     expect(callback).toBeCalledTimes(1);
     expect(callback).toBeCalledWith('Failed handling SQS event');
 
     expect(sqsDeleteMessageBatch).toBeCalledTimes(1);
-    expect(sqsDeleteMessageBatch).toBeCalledWith(metrics, evts);
+    expect(sqsDeleteMessageBatch).toBeCalledWith(mockScaleUpMetrics, evts);
   });
 
   it('stochasticOvershoot when retryCount > 5', async () => {
@@ -137,7 +140,7 @@ describe('scaleUp', () => {
       },
     ];
     expect(sqsSendMessages).toBeCalledTimes(1);
-    expect(sqsSendMessages).toBeCalledWith(metrics, expected, 'asdf');
+    expect(sqsSendMessages).toBeCalledWith(mockScaleUpMetrics, expected, 'asdf');
 
     expect(sqsDeleteMessageBatch).toBeCalledTimes(0);
   });
@@ -205,10 +208,10 @@ describe('scaleUp', () => {
       },
     ];
     expect(sqsSendMessages).toBeCalledTimes(1);
-    expect(sqsSendMessages).toBeCalledWith(metrics, expected, 'asdf');
+    expect(sqsSendMessages).toBeCalledWith(mockScaleUpMetrics, expected, 'asdf');
 
     expect(sqsDeleteMessageBatch).toBeCalledTimes(1);
-    expect(sqsDeleteMessageBatch).toBeCalledWith(metrics, records);
+    expect(sqsDeleteMessageBatch).toBeCalledWith(mockScaleUpMetrics, records);
   });
 });
 
@@ -231,3 +234,28 @@ describe('scaleDown', () => {
     expect(callback).toBeCalledWith('Failed');
   });
 });
+
+describe('scaleUpChron', () => {
+  beforeEach(() => {
+    jest.spyOn(MetricsModule, 'ScaleUpChronMetrics').mockReturnValue(mockScaleUpChronMetrics);
+  });
+
+  it('succeeds', async () => {
+    const mockedScaleUpChron = mocked(scaleUpChron).mockResolvedValue(undefined);
+    const callback = jest.fn();
+    await scaleUpChronL({} as unknown as ScheduledEvent, {} as unknown as Context, callback);
+    expect(mockedScaleUpChron).toBeCalledTimes(1);
+    expect(mockedScaleUpChron).toBeCalledWith(mockScaleUpChronMetrics);
+    expect(callback).toBeCalledTimes(1);
+    expect(callback).toBeCalledWith(null);
+  });
+
+  it('fails', async () => {
+    const mockedScaleUpChron = mocked(scaleUpChron).mockRejectedValue(Error('error'));
+    const callback = jest.fn();
+    await scaleUpChronL({} as unknown as ScheduledEvent, {} as unknown as Context, callback);
+    expect(mockedScaleUpChron).toBeCalledTimes(1);
+    expect(callback).toBeCalledTimes(1);
+    expect(callback).toBeCalledWith('Failed');
+  });
+});

terraform-aws-github-runner/modules/runners/lambdas/runners/src/lambda.ts

@@ -2,9 +2,15 @@ import { ActionRequestMessage, RetryableScalingError, scaleUp as scaleUpR } from
 import { Context, SQSEvent, SQSRecord, ScheduledEvent } from 'aws-lambda';
 
 import { Config } from './scale-runners/config';
-import { ScaleUpMetrics, sendMetricsAtTimeout, sendMetricsTimeoutVars } from './scale-runners/metrics';
+import {
+  ScaleUpMetrics,
+  ScaleUpChronMetrics,
+  sendMetricsAtTimeout,
+  sendMetricsTimeoutVars,
+} from './scale-runners/metrics';
 import { getDelayWithJitterRetryCount, stochaticRunOvershoot } from './scale-runners/utils';
 import { scaleDown as scaleDownR } from './scale-runners/scale-down';
+import { scaleUpChron as scaleUpChronR } from './scale-runners/scale-up-chron';
 import { sqsSendMessages, sqsDeleteMessageBatch } from './scale-runners/sqs';
 
 async function sendRetryEvents(evtFailed: Array<[SQSRecord, boolean, number]>, metrics: ScaleUpMetrics) {
@@ -155,3 +161,35 @@ export async function scaleDown(event: ScheduledEvent, context: Context, callbac
     return callback('Failed');
   }
 }
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export async function scaleUpChron(event: ScheduledEvent, context: Context, callback: any) {
+  // we mantain open connections to redis, so the event pool is only cleaned when the SIGTERM is sent
+  context.callbackWaitsForEmptyEventLoop = false;
+
+  const metrics = new ScaleUpChronMetrics();
+  const sndMetricsTimout: sendMetricsTimeoutVars = {
+    metrics: metrics,
+  };
+  sndMetricsTimout.setTimeout = setTimeout(
+    sendMetricsAtTimeout(sndMetricsTimout),
+    (Config.Instance.lambdaTimeout - 10) * 1000,
+  );
+
+  try {
+    await scaleUpChronR(metrics);
+    return callback(null);
+  } catch (e) {
+    console.error(e);
+    return callback('Failed');
+  } finally {
+    try {
+      clearTimeout(sndMetricsTimout.setTimeout);
+      sndMetricsTimout.metrics = undefined;
+      sndMetricsTimout.setTimeout = undefined;
+      await metrics.sendMetrics();
+    } catch (e) {
+      console.error(`Error sending metrics: ${e}`);
+    }
+  }
+}

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/config.ts

@@ -36,8 +36,11 @@ export class Config {
   readonly retryScaleUpRecordQueueUrl: string | undefined;
   readonly runnerGroupName: string | undefined;
   readonly runnersExtraLabels: undefined | string;
+  readonly scaleConfigOrg: string;
   readonly scaleConfigRepo: string;
   readonly scaleConfigRepoPath: string;
+  readonly scaleUpMinQueueTimeMinutes: number;
+  readonly scaleUpChronRecordQueueUrl: string | undefined;
   readonly secretsManagerSecretsId: string | undefined;
   readonly sSMParamCleanupAgeDays: number;
   readonly sSMParamMaxCleanupAllowance: number;
@@ -93,9 +96,14 @@ export class Config {
     this.retryScaleUpRecordDelayS = Number(process.env.RETRY_SCALE_UP_RECORD_DELAY_S || '0');
     /* istanbul ignore next */
     this.retryScaleUpRecordJitterPct = Number(process.env.RETRY_SCALE_UP_RECORD_JITTER_PCT || '0');
-    this.retryScaleUpRecordQueueUrl = process.env.RETRY_SCALE_UP_RECORD_QUEUE_URL;
+    this.retryScaleUpRecordQueueUrl = process.env.RETRY_SCALE_UP_CHRON_RECORD_QUEUE_URL;
+    this.scaleUpChronRecordQueueUrl = process.env.SCALE_UP_CHRON_HUD_QUERY_URL;
+    this.scaleUpMinQueueTimeMinutes = process.env.SCALE_UP_MIN_QUEUE_TIME_MINUTES
+      ? Number(process.env.SCALE_UP_MIN_QUEUE_TIME_MINUTES)
+      : 30;
     this.runnerGroupName = process.env.RUNNER_GROUP_NAME;
     this.runnersExtraLabels = process.env.RUNNER_EXTRA_LABELS;
+    this.scaleConfigOrg = process.env.SCALE_CONFIG_ORG || '';
     /* istanbul ignore next */
     this.scaleConfigRepo = process.env.SCALE_CONFIG_REPO || '';
     if (this.enableOrganizationRunners && !this.scaleConfigRepo) {

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/metrics.ts

@@ -913,8 +913,8 @@ export class Metrics {
 }
 
 export class ScaleUpMetrics extends Metrics {
-  constructor() {
-    super('scaleUp');
+  constructor(lambdaName: string | undefined = undefined) {
+    super(lambdaName || 'scaleUp');
   }
 
   /* istanbul ignore next */
@@ -1630,6 +1630,64 @@ export class ScaleDownMetrics extends Metrics {
   }
 }
 
+export class ScaleUpChronMetrics extends ScaleUpMetrics {
+  constructor() {
+    super('scaleUpChron');
+  }
+
+  queuedRunnerStats(org: string, runnerType: string, numQueuedJobs: number) {
+    const dimensions = new Map([
+      ['Org', org],
+      ['RunnerType', runnerType],
+      ['numQueuedJobs', numQueuedJobs.toString()],
+    ]);
+    this.addEntry('gh.scaleupchron.queuedRunners', 3, dimensions);
+  }
+
+  queuedRunnerFailure(error: string) {
+    const dimensions = new Map([['error', error]]);
+    this.countEntry('gh.scaleupchron.queuedRunners.failure', 1, dimensions);
+  }
+
+  /* istanbul ignore next */
+  getQueuedJobsEndpointSuccess(ms: number) {
+    this.countEntry(`gh.calls.total`, 1);
+    this.countEntry(`gh.calls.getQueuedJobsEndpoint.count`, 1);
+    this.countEntry(`gh.calls.getQueuedJobsEndpoint.success`, 1);
+    this.addEntry(`gh.calls.getQueuedJobsEndpoint.wallclock`, ms);
+  }
+
+  /* istanbul ignore next */
+  getQueuedJobsEndpointFailure(ms: number) {
+    this.countEntry(`gh.calls.total`, 1);
+    this.countEntry(`gh.calls.getQueuedJobsEndpoint.count`, 1);
+    this.countEntry(`gh.calls.getQueuedJobsEndpoint.failure`, 1);
+    this.addEntry(`gh.calls.getQueuedJobsEndpoint.wallclock`, ms);
+  }
+
+  scaleUpInstanceSuccess() {
+    this.scaleUpSuccess();
+    this.countEntry('run.scaleupchron.success');
+  }
+
+  scaleUpInstanceFailureNonRetryable(error: string) {
+    const dimensions = new Map([['error', error]]);
+    // should we add more information about this or do we not care since  it'll be requeued?
+    this.countEntry('run.scaleupchron.failure.nonRetryable', 1, dimensions);
+  }
+
+  scaleUpInstanceFailureRetryable(error: string) {
+    const dimensions = new Map([['error', error]]);
+
+    // should we add more information about this or do we not care since  it'll be requeued?
+    this.countEntry('run.scaleupchron.failure.retryable', 1, dimensions);
+  }
+
+  scaleUpInstanceNoOp() {
+    this.countEntry('run.scaleupchron.noop');
+  }
+}
+
 export interface sendMetricsTimeoutVars {
   metrics?: Metrics;
   setTimeout?: ReturnType<typeof setTimeout>;