Fixes multiregion code for autoscaler lambdas (#6231)

pytorch · Feb 4, 2025 · 4e1f892 · 4e1f892
1 parent 7290f72
commit 4e1f892
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 17 deletions.
diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/Makefile b/terraform-aws-github-runner/modules/runners/lambdas/runners/Makefile
@@ -2,7 +2,7 @@ SHELL=/bin/bash -o pipefail
 
 .PHONY: clean
 clean:
-	rm -rf dist node_modules
+	rm -rf dist node_modules coverage
 	rm runners.zip
 
 .PHONY: build

diff --git a/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/config.ts b/terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/config.ts
@@ -137,13 +137,9 @@ export class Config {
   }
 
   get shuffledAwsRegionInstances(): string[] {
-    let arr: string[];
-    if (this.awsRegionsToVpcIds.size > 0) {
-      arr = Array.from(this.awsRegionsToVpcIds.keys());
-    } else {
-      arr = [...this.awsRegionInstances];
-    }
-    return shuffleArrayInPlace(arr);
+    const regions: Set<string> = new Set(this.awsRegionsToVpcIds.keys());
+    this.awsRegionInstances.forEach((region) => regions.add(region));
+    return shuffleArrayInPlace(Array.from(regions));
   }
 
   get ghesUrlApi(): undefined | string {

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
@@ -146,10 +146,11 @@ export async function listRunners(
           ec2Filters.push({ Name: tags[attr as keyof typeof tags], Values: [filters[attr] as string] }),
         );
     }
-    console.debug(`[listRunners]: REGIONS ${Config.Instance.shuffledAwsRegionInstances}`);
+    const awsRegionsInstances = Config.Instance.shuffledAwsRegionInstances;
+    console.debug(`[listRunners]: REGIONS ${awsRegionsInstances}`);
     const runningInstances = (
       await Promise.all(
-        Config.Instance.shuffledAwsRegionInstances
+        awsRegionsInstances
           .filter((r) => regions?.has(r) ?? true)
           .map((awsRegion) => {
             console.debug(`[listRunners]: Running for region ${awsRegion}`);
@@ -163,11 +164,23 @@ export async function listRunners(
                     .describeInstances({ Filters: ec2Filters })
                     .promise()
                     .then((describeInstanceResult): DescribeInstancesResultRegion => {
+                      const listOfRunnersIdType: string[] = (
+                        describeInstanceResult?.Reservations?.flatMap((reservation) => {
+                          return (
+                            reservation.Instances?.map((instance) => {
+                              return `${instance.InstanceId} - ${
+                                instance.Tags?.find((e) => e.Key === 'RunnerType')?.Value
+                              }`;
+                            }) ?? []
+                          );
+                        }) ?? []
+                      ).filter((desc): desc is string => desc !== undefined);
                       console.debug(
                         `[listRunners]: Result for EC2({ region: ${awsRegion} })` +
-                          `.describeInstances({ Filters: ${ec2Filters} }) = ` +
+                          `.describeInstances({ Filters: ${JSON.stringify(ec2Filters)} }) = ` +
                           `${describeInstanceResult?.Reservations?.length ?? 'UNDEF'}`,
                       );
+                      console.debug(`[listRunners]: ${listOfRunnersIdType.join('\n ')}`);
                       return { describeInstanceResult, awsRegion };
                     });
                 },
@@ -482,8 +495,8 @@ export async function createRunner(runnerParameters: RunnerInputParameters, metr
       runnerParameters.runnerType.labels ? ' [' + runnerParameters.runnerType.labels.join(',') + ']' : ''
     }`;
 
-    const shuffledAwsRegionInstances = Config.Instance.shuffledAwsRegionInstances;
-    for (const [awsRegionIdx, awsRegion] of shuffledAwsRegionInstances.entries()) {
+    const awsRegionsInstances = Config.Instance.shuffledAwsRegionInstances;
+    for (const [awsRegionIdx, awsRegion] of awsRegionsInstances.entries()) {
       const runnerSubnetSequence = await getCreateRunnerSubnetSequence(runnerParameters, awsRegion, metrics);
 
       const ec2 = new EC2({ region: awsRegion });
@@ -582,7 +595,7 @@ export async function createRunner(runnerParameters: RunnerInputParameters, metr
           const msg =
             `[${subnetIdx}/${subnets.length} - ${subnet}] ` +
             `[${vpcId}] ` +
-            `[${awsRegionIdx}/${shuffledAwsRegionInstances.length} - ${awsRegion}] Issue creating instance ` +
+            `[${awsRegionIdx}/${awsRegionsInstances.length} - ${awsRegion}] Issue creating instance ` +
             `${runnerParameters.runnerType.instance_type}${labelsStrLog}: ${e}`;
           errors.push([msg, e, awsRegion]);
           console.warn(msg);

diff --git a/...form-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/scale-up.test.ts b/...form-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/scale-up.test.ts
@@ -1083,9 +1083,6 @@ describe('_calculateScaleUpAmount', () => {
           availableCount,
         );
 
-        const availAfterHandinglingRequest = availableCount - requestedCount;
-        const amtBelowMin = minRunners - availAfterHandinglingRequest;
-
         // We were above min runners before, and we should scale up enough to not dip below min runners
         expect(scaleUpAmount).toEqual(3);
       }