Improves visibility for AMI experiments (#5421)

pytorch · Jul 12, 2024 · e523942 · e523942
1 parent 4e860d3
commit e523942
Show file tree

Hide file tree

Showing 5 changed files with 69 additions and 23 deletions.
diff --git a/terraform-aws-github-runner/modules/runners-instances/templates/install-config-runner.sh b/terraform-aws-github-runner/modules/runners-instances/templates/install-config-runner.sh
@@ -51,6 +51,11 @@ function metric_report () {
     if [ ! -z "$OS_ID" ]; then
         aws cloudwatch put-metric-data --metric-name "\$metric_name" --namespace "\$namespace" --value \$value --region us-east-1 --dimensions "os=$OS_ID" || true
     fi
+    while read line ; do
+        aws cloudwatch put-metric-data --metric-name "\$metric_name" --namespace "\$namespace" --value \$value --region us-east-1 --dimensions "label=\$line,OnAMIExperiment=\$on_ami_experiment" || true
+        aws cloudwatch put-metric-data --metric-name "\$metric_name" --namespace "\$namespace" --value \$value --region us-east-1 --dimensions "label=\$line" || true
+    done < /home/$USER_NAME/runner-labels
+
     aws cloudwatch put-metric-data --metric-name "\$metric_name" --namespace "\$namespace" --value \$value --region us-east-1 --dimensions "OnAMIExperiment=\$on_ami_experiment" || true
 }
 EOF
@@ -167,6 +172,16 @@ fallback_to_node16() {
   echo $FALLBACK_TO_NODE16 >> $RUNNER_ENV
 }
 
+get_labels_from_config() {
+  while [[ "$#" -gt 0 ]]; do
+    case $1 in
+      --labels) target="$2"; shift ;;
+    esac
+    shift
+  done
+  echo $target | sed 's/,/\n/g'
+}
+
 cd /home/$USER_NAME
 mkdir actions-runner && cd actions-runner
 
@@ -191,13 +206,15 @@ while [[ $(aws ssm get-parameters --names ${environment}-$INSTANCE_ID --with-dec
     fi
 done
 CONFIG=$(aws ssm get-parameters --names ${environment}-$INSTANCE_ID --with-decryption --region $REGION | jq -r ".Parameters | .[0] | .Value")
-echo "Configuration received: '$CONFIG'"
+retry aws ssm delete-parameter --name ${environment}-$INSTANCE_ID --region $REGION
 
+echo "Configuration received: '$CONFIG'"
 if echo "$CONFIG" | grep -q "#ON_AMI_EXPERIMENT"; then
   CONFIG=$(echo "$CONFIG" | sed 's/ #ON_AMI_EXPERIMENT//g')
   touch /home/$USER_NAME/on-ami-experiment
 fi
-retry aws ssm delete-parameter --name ${environment}-$INSTANCE_ID --region $REGION
+
+get_labels_from_config $CONFIG > /home/$USER_NAME/runner-labels
 
 export RUNNER_ALLOW_RUNASROOT=1
 os_id=$(awk -F= '/^ID/{print $2}' /etc/os-release)

diff --git a/...aform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/runners.test.ts b/...aform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/runners.test.ts
@@ -535,7 +535,7 @@ describe('createRunner', () => {
       expect(mockEC2.runInstances).toHaveBeenCalledTimes(1);
       expect(mockEC2.runInstances).toBeCalledWith(createExpectedRunInstancesLinux(runnerParameters, 0));
       expect(runnerConfigFn).toBeCalledTimes(1);
-      expect(runnerConfigFn).toBeCalledWith(config.awsRegion);
+      expect(runnerConfigFn).toBeCalledWith(config.awsRegion, false);
     });
 
     it('calls run instances with the correct config for repo && linux && organization', async () => {
@@ -559,7 +559,7 @@ describe('createRunner', () => {
       expect(mockEC2.runInstances).toHaveBeenCalledTimes(1);
       expect(mockEC2.runInstances).toBeCalledWith(createExpectedRunInstancesLinux(runnerParameters, 0, true));
       expect(runnerConfigFn).toBeCalledTimes(1);
-      expect(runnerConfigFn).toBeCalledWith(config.awsRegion);
+      expect(runnerConfigFn).toBeCalledWith(config.awsRegion, false);
     });
 
     it('calls run instances with the correct config for repo && windows', async () => {
@@ -581,7 +581,7 @@ describe('createRunner', () => {
       await createRunner(runnerParameters, metrics);
 
       expect(runnerConfigFn).toBeCalledTimes(1);
-      expect(runnerConfigFn).toBeCalledWith(config.awsRegion);
+      expect(runnerConfigFn).toBeCalledWith(config.awsRegion, false);
       expect(mockEC2.runInstances).toHaveBeenCalledTimes(1);
       const secGroup = Config.Instance.vpcIdToSecurityGroupIds.get('vpc-agdgaduwg113') || [];
       expect(mockEC2.runInstances).toBeCalledWith({
@@ -647,7 +647,7 @@ describe('createRunner', () => {
         metrics,
       );
       expect(runnerConfigFn).toBeCalledTimes(1);
-      expect(runnerConfigFn).toBeCalledWith(config.awsRegion);
+      expect(runnerConfigFn).toBeCalledWith(config.awsRegion, false);
       expect(mockSSM.putParameter).toBeCalledTimes(1);
       expect(mockSSM.putParameter).toBeCalledWith({
         Name: 'unit-test-env-i-1234',
@@ -903,7 +903,7 @@ describe('createRunner', () => {
       expect(mockEC2.runInstances).toHaveBeenCalledTimes(1);
       expect(mockEC2.runInstances).toBeCalledWith(createExpectedRunInstancesLinux(runnerParameters, 2));
       expect(runnerConfigFn).toBeCalledTimes(1);
-      expect(runnerConfigFn).toBeCalledWith(config.shuffledAwsRegionInstances[0]);
+      expect(runnerConfigFn).toBeCalledWith(config.shuffledAwsRegionInstances[0], false);
     });
 
     it('succeed, 2nd subnet and 1st region', async () => {
@@ -935,7 +935,7 @@ describe('createRunner', () => {
         createExpectedRunInstancesLinux(runnerParameters, 4, false, 'vpc-agdgaduwg113-12'),
       );
       expect(runnerConfigFn).toBeCalledTimes(1);
-      expect(runnerConfigFn).toBeCalledWith(config.shuffledAwsRegionInstances[0]);
+      expect(runnerConfigFn).toBeCalledWith(config.shuffledAwsRegionInstances[0], false);
     });
 
     it('succeed, 1nd subnet and 2nd region', async () => {
@@ -983,7 +983,7 @@ describe('createRunner', () => {
         ),
       );
       expect(runnerConfigFn).toBeCalledTimes(1);
-      expect(runnerConfigFn).toBeCalledWith(config.shuffledAwsRegionInstances[1]);
+      expect(runnerConfigFn).toBeCalledWith(config.shuffledAwsRegionInstances[1], false);
     });
 
     it('fails, everywere', async () => {

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
@@ -15,7 +15,7 @@ export interface ListRunnerFilters {
 }
 
 export interface RunnerInputParameters {
-  runnerConfig: (awsRegion: string) => Promise<string>;
+  runnerConfig: (awsRegion: string, experimentalRunner: boolean) => Promise<string>;
   environment: string;
   repoName?: string;
   orgName?: string;
@@ -319,7 +319,7 @@ async function addSSMParameterRunnerConfig(
     return;
   }
 
-  let runnerConfig = await runnerParameters.runnerConfig(awsRegion);
+  let runnerConfig = await runnerParameters.runnerConfig(awsRegion, customAmiExperiment);
   if (customAmiExperiment) {
     runnerConfig = `${runnerConfig} #ON_AMI_EXPERIMENT`;
   }

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
@@ -275,11 +275,15 @@ describe('scaleUp', () => {
       metrics,
     );
 
-    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion)).toEqual(
+    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion, false)).toEqual(
       `--url ${config.ghesUrlHost}/owner --token ${token} --labels AWS:${config.awsRegion},linux.2xlarge,` +
         `extra-label  --runnergroup group_one`,
     );
-    expect(mockedCreateRegistrationTokenForOrg).toBeCalledTimes(1);
+    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion, true)).toEqual(
+      `--url ${config.ghesUrlHost}/owner --token ${token} --labels AWS:${config.awsRegion},linux.2xlarge,` +
+        `experimental.ami,extra-label --ephemeral --runnergroup group_one`,
+    );
+    expect(mockedCreateRegistrationTokenForOrg).toBeCalledTimes(2);
     expect(mockedCreateRegistrationTokenForOrg).toBeCalledWith(repo.owner, metrics, 2);
   });
 
@@ -357,11 +361,15 @@ describe('scaleUp', () => {
       metrics,
     );
 
-    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion)).toEqual(
+    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion, false)).toEqual(
       `--url ${config.ghesUrlHost}/owner/repo --token ${token} --labels AWS:${config.awsRegion},linux.2xlarge,` +
         `extra-label `,
     );
-    expect(mockedCreateRegistrationTokenForRepo).toBeCalledTimes(1);
+    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion, true)).toEqual(
+      `--url ${config.ghesUrlHost}/owner/repo --token ${token} --labels AWS:${config.awsRegion},linux.2xlarge,` +
+        `experimental.ami,extra-label --ephemeral`,
+    );
+    expect(mockedCreateRegistrationTokenForRepo).toBeCalledTimes(2);
     expect(mockedCreateRegistrationTokenForRepo).toBeCalledWith(repo, metrics, 2);
   });
 
@@ -439,11 +447,15 @@ describe('scaleUp', () => {
       metrics,
     );
 
-    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion)).toEqual(
+    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion, false)).toEqual(
       `--url ${config.ghesUrlHost}/owner/repo --token ${token} --labels AWS:${config.awsRegion},linux.2xlarge,` +
         `extra-label `,
     );
-    expect(mockedCreateRegistrationTokenForRepo).toBeCalledTimes(1);
+    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion, true)).toEqual(
+      `--url ${config.ghesUrlHost}/owner/repo --token ${token} --labels AWS:${config.awsRegion},linux.2xlarge,` +
+        `experimental.ami,extra-label --ephemeral`,
+    );
+    expect(mockedCreateRegistrationTokenForRepo).toBeCalledTimes(2);
     expect(mockedCreateRegistrationTokenForRepo).toBeCalledWith(repo, metrics, 0);
   });
 
@@ -521,11 +533,15 @@ describe('scaleUp', () => {
       metrics,
     );
 
-    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion)).toEqual(
+    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion, false)).toEqual(
       `--url ${config.ghesUrlHost}/owner/repo --token ${token} --labels AWS:${config.awsRegion},linux.2xlarge,` +
         `extra-label `,
     );
-    expect(mockedCreateRegistrationTokenForRepo).toBeCalledTimes(1);
+    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion, true)).toEqual(
+      `--url ${config.ghesUrlHost}/owner/repo --token ${token} --labels AWS:${config.awsRegion},linux.2xlarge,` +
+        `experimental.ami,extra-label --ephemeral`,
+    );
+    expect(mockedCreateRegistrationTokenForRepo).toBeCalledTimes(2);
     expect(mockedCreateRegistrationTokenForRepo).toBeCalledWith(repo, metrics, undefined);
   });
 
@@ -636,10 +652,14 @@ describe('scaleUp', () => {
       metrics,
     );
 
-    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion)).toEqual(
+    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion, false)).toEqual(
       `--url ${config.ghesUrlHost}/owner/repo --token ${token} --labels AWS:${config.awsRegion},linux.2xlarge` +
         ` --ephemeral`,
     );
+    expect(await mockedCreateRunner.mock.calls[0][0].runnerConfig(config.awsRegion, true)).toEqual(
+      `--url ${config.ghesUrlHost}/owner/repo --token ${token} --labels AWS:${config.awsRegion},linux.2xlarge,` +
+        `experimental.ami --ephemeral`,
+    );
   });
 
   it('dont have mustHaveIssuesLabels', async () => {

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
@@ -83,8 +83,15 @@ export async function scaleUp(
       try {
         const createRunnerParams: RunnerInputParameters = {
           environment: Config.Instance.environment,
-          runnerConfig: (awsRegion: string) => {
-            return createRunnerConfigArgument(runnerType, repo, payload.installationId, metrics, awsRegion);
+          runnerConfig: (awsRegion: string, experimentalRunner: boolean) => {
+            return createRunnerConfigArgument(
+              runnerType,
+              repo,
+              payload.installationId,
+              metrics,
+              awsRegion,
+              experimentalRunner,
+            );
           },
           runnerType: runnerType,
         };
@@ -132,11 +139,13 @@ async function createRunnerConfigArgument(
   installationId: number | undefined,
   metrics: Metrics,
   awsRegion: string,
+  experimentalRunner: boolean,
 ): Promise<string> {
-  const ephemeralArgument = runnerType.is_ephemeral ? '--ephemeral' : '';
+  const ephemeralArgument = runnerType.is_ephemeral || experimentalRunner ? '--ephemeral' : '';
   const labelsArgument = [
     `AWS:${awsRegion}`,
     `${runnerType.runnerTypeName}`,
+    ...(experimentalRunner ? ['experimental.ami'] : []),
     ...(Config.Instance.runnersExtraLabels ? Config.Instance.runnersExtraLabels.split(',') : []),
     ...(runnerType.labels ?? []),
   ].join(',');