Introduce runner variants to help test changes (#5531)

pytorch · Aug 2, 2024 · f3c58fe · f3c58fe
1 parent 2f74a6e
commit f3c58fe
Show file tree

Hide file tree

Showing 3 changed files with 132 additions and 51 deletions.
diff --git a/...rm-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/gh-runners.test.ts b/...rm-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/gh-runners.test.ts
@@ -669,7 +669,79 @@ runner_types:
       os: linux
       max_available: 1
       disk_size: 150
-      is_ephemeral: false`;
+      is_ephemeral: false
+    linux.4xlarge:
+      instance_type: c5.2xlarge
+      os: linux
+      max_available: 1
+      disk_size: 150
+      is_ephemeral: false
+      variants:
+        ephemeral:
+          is_ephemeral: true
+        largedisk:
+          disk_size: 300
+        ami123:
+          ami: ami-123`;
+
+  const getRunnerTypeResponse = new Map([
+    [
+      'linux.2xlarge',
+      {
+        runnerTypeName: 'linux.2xlarge',
+        instance_type: 'c5.2xlarge',
+        os: 'linux',
+        max_available: 1,
+        disk_size: 150,
+        is_ephemeral: false,
+      },
+    ],
+    [
+      'linux.4xlarge',
+      {
+        runnerTypeName: 'linux.4xlarge',
+        instance_type: 'c5.2xlarge',
+        os: 'linux',
+        max_available: 1,
+        disk_size: 150,
+        is_ephemeral: false,
+      },
+    ],
+    [
+      'ephemeral.linux.4xlarge',
+      {
+        runnerTypeName: 'ephemeral.linux.4xlarge',
+        instance_type: 'c5.2xlarge',
+        os: 'linux',
+        max_available: 1,
+        disk_size: 150,
+        is_ephemeral: true,
+      },
+    ],
+    [
+      'largedisk.linux.4xlarge',
+      {
+        runnerTypeName: 'largedisk.linux.4xlarge',
+        instance_type: 'c5.2xlarge',
+        os: 'linux',
+        max_available: 1,
+        disk_size: 300,
+        is_ephemeral: false,
+      },
+    ],
+    [
+      'ami123.linux.4xlarge',
+      {
+        runnerTypeName: 'ami123.linux.4xlarge',
+        instance_type: 'c5.2xlarge',
+        os: 'linux',
+        max_available: 1,
+        disk_size: 150,
+        is_ephemeral: false,
+        ami: 'ami-123',
+      },
+    ],
+  ]);
 
   it('gets the contents, twice', async () => {
     const repo = { owner: 'owner', repo: 'repo' };
@@ -697,36 +769,8 @@ runner_types:
     mockCreateOctoClient.mockReturnValueOnce(mockedOctokit as unknown as Octokit);
 
     await resetGHRunnersCaches();
-    expect(await getRunnerTypes(repo, metrics)).toEqual(
-      new Map([
-        [
-          'linux.2xlarge',
-          {
-            runnerTypeName: 'linux.2xlarge',
-            instance_type: 'c5.2xlarge',
-            os: 'linux',
-            max_available: 1,
-            disk_size: 150,
-            is_ephemeral: false,
-          },
-        ],
-      ]),
-    );
-    expect(await getRunnerTypes(repo, metrics)).toEqual(
-      new Map([
-        [
-          'linux.2xlarge',
-          {
-            runnerTypeName: 'linux.2xlarge',
-            instance_type: 'c5.2xlarge',
-            os: 'linux',
-            max_available: 1,
-            disk_size: 150,
-            is_ephemeral: false,
-          },
-        ],
-      ]),
-    );
+    expect(await getRunnerTypes(repo, metrics)).toEqual(getRunnerTypeResponse);
+    expect(await getRunnerTypes(repo, metrics)).toEqual(getRunnerTypeResponse);
 
     expect(mockCreateGithubAuth).toBeCalledTimes(2);
     expect(mockCreateGithubAuth).toBeCalledWith(undefined, 'app', Config.Instance.ghesUrlApi, metrics);

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
@@ -1,5 +1,5 @@
 import { Repo, getRepoKey, expBackOff } from './utils';
-import { RunnerType } from './runners';
+import { RunnerType, RunnerTypeScaleConfig } from './runners';
 import { createGithubAuth, createOctoClient } from './gh-auth';
 import { locallyCached, redisCached, clearLocalCacheNamespace, redisClearCacheKeyPattern } from './cache';
 
@@ -340,7 +340,7 @@ export async function getRunnerTypes(
       console.debug(`'${filepath}' contents: ${configYml}`);
 
       const config = YAML.parse(configYml);
-      const result: Map<string, RunnerType> = new Map(
+      const result: Map<string, RunnerTypeScaleConfig> = new Map(
         /* eslint-disable-next-line @typescript-eslint/no-explicit-any */
         (Object.entries(config.runner_types) as [string, any][]).map(([prop, runner_type]) => [
           prop,
@@ -358,33 +358,58 @@ export async function getRunnerTypes(
             max_available: runner_type.max_available,
             os: runner_type.os,
             runnerTypeName: prop,
+            variants: new Map(Object.entries(runner_type.variants || {})),
           },
         ]),
       );
 
+      Array.from(result.keys()).forEach((key) => {
+        const runnerType = result.get(key);
+        if (runnerType?.variants === undefined) {
+          return;
+        }
+
+        if (runnerType.variants.size > 0) {
+          Array.from(runnerType.variants.keys()).forEach((variant) => {
+            const variantType = runnerType.variants?.get(variant);
+            if (!variantType) {
+              return;
+            }
+
+            result.set(`${variant}.${key}`, { ...runnerType, ...variantType, runnerTypeName: `${variant}.${key}` });
+          });
+        }
+      });
+
       const filteredResult: Map<string, RunnerType> = new Map(
-        [...result.entries()].filter(
-          ([, runnerType]) =>
-            typeof runnerType.runnerTypeName === 'string' &&
-            alphaNumericStr.test(runnerType.runnerTypeName) &&
-            typeof runnerType.instance_type === 'string' &&
-            alphaNumericStr.test(runnerType.instance_type) &&
-            ['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.max_available === 'number' || runnerType.max_available === undefined) &&
-            (typeof runnerType.ami === 'string' || runnerType.ami === undefined) &&
-            (typeof runnerType.ami_experiment?.ami === 'string' || runnerType.ami_experiment === undefined) &&
-            (typeof runnerType.ami_experiment?.percentage === 'number' || runnerType.ami_experiment === undefined),
-        ),
+        [...result.entries()]
+          .filter(
+            ([, runnerType]) =>
+              typeof runnerType.runnerTypeName === 'string' &&
+              alphaNumericStr.test(runnerType.runnerTypeName) &&
+              typeof runnerType.instance_type === 'string' &&
+              alphaNumericStr.test(runnerType.instance_type) &&
+              ['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.max_available === 'number' || runnerType.max_available === undefined) &&
+              (typeof runnerType.ami === 'string' || runnerType.ami === undefined) &&
+              (typeof runnerType.ami_experiment?.ami === 'string' || runnerType.ami_experiment === undefined) &&
+              (typeof runnerType.ami_experiment?.percentage === 'number' || runnerType.ami_experiment === undefined),
+          )
+          .map(([key, runnerType]) => {
+            const rt: RunnerTypeScaleConfig = { ...runnerType };
+            delete rt.variants;
+            return [key, rt];
+          }),
       );
 
       if (result.size != filteredResult.size) {
         console.error(
           `Some runner types were filtered out due to invalid values: ${result.size} -> ${filteredResult.size}`,
         );
-        console.error(`Original runner types: ${JSON.stringify(result)}`);
-        console.error(`Filtered runner types: ${JSON.stringify(filteredResult)}`);
+        console.error(`Original runner types: ${JSON.stringify(Array.from(result.keys()).sort())}`);
+        console.error(`Filtered runner types: ${JSON.stringify(Array.from(filteredResult.keys()).sort())}`);
       }
 
       status = 'success';

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
@@ -27,18 +27,30 @@ export interface AmiExpermient {
   percentage: number;
 }
 
-export interface RunnerType {
+export interface RunnerTypeOptional {
   ami_experiment?: AmiExpermient;
   ami?: string;
+  disk_size?: number;
+  instance_type?: string;
+  is_ephemeral?: boolean;
+  labels?: Array<string>;
+  max_available?: number;
+  os?: string;
+}
+
+export interface RunnerType extends RunnerTypeOptional {
   disk_size: number;
   instance_type: string;
   is_ephemeral: boolean;
-  labels?: Array<string>;
   max_available: number;
   os: string;
   runnerTypeName: string;
 }
 
+export interface RunnerTypeScaleConfig extends RunnerType {
+  variants?: Map<string, RunnerTypeOptional>;
+}
+
 export interface DescribeInstancesResultRegion {
   awsRegion: string;
   describeInstanceResult: PromiseResult<EC2.Types.DescribeInstancesResult, AWS.AWSError>;