From 140d43b4199c662b790c3a7f019fb8d3468ce89d Mon Sep 17 00:00:00 2001
From: Thanh Ha <thanh.ha@linuxfoundation.org>
Date: Wed, 25 Sep 2024 12:47:39 -0400
Subject: [PATCH] Add min-available runners to scale-config

This change allows the min-available runners to be configured more
specifically by the runner type rather than as a single global setting.

Issue: pytorch/ci-infra#275
Signed-off-by: Thanh Ha <thanh.ha@linuxfoundation.org>
---
 .../runners/src/scale-runners/gh-runners.ts   |  2 ++
 .../runners/src/scale-runners/runners.ts      |  1 +
 .../runners/src/scale-runners/scale-down.ts   | 22 ++++++++++++++++++-
 3 files changed, 24 insertions(+), 1 deletion(-)

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
index 07e529903b..9185e394b5 100644
--- 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
@@ -355,6 +355,7 @@ export async function getRunnerTypes(
             is_ephemeral: runner_type.is_ephemeral || false,
             /* istanbul ignore next */
             labels: runner_type.labels?.map((label: string) => label.trim()),
+            min_available: runner_type.min_available || Config.Instance.minAvailableRunners,
             max_available: runner_type.max_available,
             os: runner_type.os,
             runnerTypeName: prop,
@@ -401,6 +402,7 @@ export async function getRunnerTypes(
               ['linux', 'windows'].includes(runnerType.os) &&
               (runnerType.labels?.every((label) => typeof label === 'string' && alphaNumericStr.test(label)) ?? true) &&
               (typeof runnerType.disk_size === 'number' || runnerType.disk_size === undefined) &&
+              (typeof runnerType.min_available === 'number' || runnerType.min_available === undefined) &&
               (typeof runnerType.max_available === 'number' || runnerType.max_available === undefined) &&
               (typeof runnerType.ami === 'string' || runnerType.ami === undefined) &&
               (typeof runnerType.ami_experiment?.ami === 'string' || runnerType.ami_experiment === undefined) &&
diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/runners.ts b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/runners.ts
index 3f06d50c53..b4f689dffe 100644
--- a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/runners.ts
+++ b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/runners.ts
@@ -34,6 +34,7 @@ export interface RunnerTypeOptional {
   instance_type?: string;
   is_ephemeral?: boolean;
   labels?: Array<string>;
+  min_available?: number;
   max_available?: number;
   os?: string;
 }
diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/scale-down.ts b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/scale-down.ts
index 8d5c6a7410..73d82550ad 100644
--- a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/scale-down.ts
+++ b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/scale-down.ts
@@ -103,7 +103,7 @@ export async function scaleDown(): Promise<void> {
         // We only limit the number of removed instances here for the reason: while sorting and getting info
         // on getRunner[Org|Repo] we send statistics that are relevant for monitoring
         if (
-          ghRunnersRemovable.length - removedRunners <= Config.Instance.minAvailableRunners &&
+          ghRunnersRemovable.length - removedRunners <= (await minRunners(ec2runner, metrics)) &&
           ghRunner !== undefined &&
           ec2runner.applicationDeployDatetime == Config.Instance.datetimeDeploy &&
           !(await isEphemeralRunner(ec2runner, metrics))
@@ -387,6 +387,26 @@ export async function isEphemeralRunner(ec2runner: RunnerInfo, metrics: ScaleDow
   return runnerTypes.get(ec2runner.runnerType)?.is_ephemeral ?? false;
 }
 
+export async function minRunners(ec2runner: RunnerInfo, metrics: ScaleDownMetrics): Promise<number> {
+  if (ec2runner.runnerType === undefined) {
+    return Config.Instance.minAvailableRunners;
+  }
+
+  const repo: Repo = (() => {
+    if (Config.Instance.enableOrganizationRunners) {
+      return {
+        owner: ec2runner.org !== undefined ? (ec2runner.org as string) : getRepo(ec2runner.repo as string).owner,
+        repo: Config.Instance.scaleConfigRepo,
+      };
+    }
+    return getRepo(ec2runner.repo as string);
+  })();
+
+  const runnerTypes = await getRunnerTypes(repo, metrics);
+
+  return runnerTypes.get(ec2runner.runnerType)?.min_available ?? Config.Instance.minAvailableRunners;
+}
+
 export function isRunnerRemovable(
   ghRunner: GhRunner | undefined,
   ec2runner: RunnerInfo,