From f44400fd84ba96e90641eb68aa51c4a65f30fda3 Mon Sep 17 00:00:00 2001
From: barshaul <barshaul@amazon.com>
Date: Sat, 16 Mar 2024 09:23:00 +0000
Subject: [PATCH] CD: Automate the start&stop of the self hosted runner

---
 .github/workflows/pypi-cd.yml | 56 +++++++++++++++++++++++++++++++++--
 1 file changed, 54 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/pypi-cd.yml b/.github/workflows/pypi-cd.yml
index 8e0b81c8be..81bd8c165f 100644
--- a/.github/workflows/pypi-cd.yml
+++ b/.github/workflows/pypi-cd.yml
@@ -14,12 +14,43 @@ on:
 concurrency:
     group: pypi-${{ github.head_ref || github.ref }}
     cancel-in-progress: true
+permissions: write-all
 
 jobs:
+    start-self-hosted-runner:
+      name: Start self hosted EC2 runner
+      runs-on: ubuntu-latest
+      outputs:
+        label: ${{ steps.start-ec2-runner.outputs.label }}
+        ec2-instance-id: ${{ steps.start-ec2-runner.outputs.ec2-instance-id }}
+      steps:
+        - name: Configure AWS Credentials
+          uses: aws-actions/configure-aws-credentials@v4
+          with:
+            aws-access-key-id: ${{ secrets.AWS_EC2_ACCESS_KEY_ID }}
+            aws-secret-access-key: ${{ secrets.AWS_EC2_SECRET_ACCESS_KEY }}
+            aws-region: ${{ secrets.AWS_REGION }}
+        - name: Start EC2 runner
+          id: start-ec2-runner
+          uses: machulav/ec2-github-runner@v2
+          with:
+            mode: start
+            github-token: ${{ secrets.GITHUB_TOKEN }}
+            ec2-image-id: ami-03831836fb1f65aac
+            ec2-instance-type: m6g.xlarge
+            subnet-id: subnet-0a88f7abc7e1490e1
+            security-group-id: sg-05024f40170e648f0
+            iam-role-name: github-actions-role 
+            aws-resource-tags: >
+              [
+                {"Key": "Name", "Value": "ec2-github-runner"},
+                {"Key": "GitHubRepository", "Value": "${{ github.repository }}"}  
+              ]
     publish-binaries:
+        needs: start-self-hosted-runner
         if: github.repository_owner == 'aws'
         name: Publish packages to PyPi
-        runs-on: ${{ matrix.build.RUNNER }}
+        runs-on: ${{ matrix.build.RUNNER == 'self-hosted' && needs.start-self-hosted-runner.outputs.label || matrix.build.RUNNER }}
         strategy:
             fail-fast: false
             matrix:
@@ -34,7 +65,7 @@ jobs:
                     - {
                           OS: ubuntu-latest,
                           NAMED_OS: linux,
-                          RUNNER: [self-hosted, Linux, ARM64],
+                          RUNNER: self-hosted,
                           ARCH: arm64,
                           TARGET: aarch64-unknown-linux-gnu,
                           CONTAINER: "2_28",
@@ -170,3 +201,24 @@ jobs:
               with:
                   command: upload
                   args: --skip-existing python/wheels/*
+    stop-runner:
+      name: Stop self-hosted EC2 runner
+      needs:
+        - start-self-hosted-runner # required to get output from the start-runner job
+        - publish-to-pypi # required to wait when the main job is done
+      runs-on: ubuntu-latest
+      if: ${{ always() }} # required to stop the runner even if the error happened in the previous jobs
+      steps:
+        - name: Configure AWS Credentials
+          uses: aws-actions/configure-aws-credentials@v4
+          with:
+            aws-access-key-id: ${{ secrets.AWS_EC2_ACCESS_KEY_ID }}
+            aws-secret-access-key: ${{ secrets.AWS_EC2_SECRET_ACCESS_KEY }}
+            aws-region: ${{ secrets.AWS_REGION }}
+        - name: Stop EC2 runner
+          uses: machulav/ec2-github-runner@v2
+          with:
+            mode: stop
+            github-token: ${{ secrets.GITHUB_TOKEN }}
+            label: ${{ needs.start-self-hosted-runner.outputs.label }}
+            ec2-instance-id: ${{ needs.start-self-hosted-runner.outputs.ec2-instance-id }}