Merge pull request #4 from omsf/feat/auto-updating-ami

feat: add support for pulling the latest AMI

Author: dwhswenson

README.md

@@ -6,7 +6,8 @@ This repository contains the code to start a GitHub Actions runner on an AWS EC2
 | arch                  | The AMI architecture                                                                                               | true               | x64     |
 | aws_home_dir          | The AWS AMI home directory to use for your runner. Will not start if not specified.                                | true               |         |
 | aws_iam_role          | The optional AWS IAM role to assume for provisioning your runner.                                                  | false              |         |
-| aws_image_id          | The machine AMI to use for your runner. This AMI can be a default but should have docker installed in the AMI.     | true               |         |
+| aws_image_id          | The machine AMI to use for your runner. This AMI can be a default but should have docker installed in the AMI. If set to `latest`, aws_image_name is required     | true               |         |
+| aws_image_name        | The name of AMI you want to use, only required if you don't specify `aws_image_id`                                 | false              |         |
 | aws_instance_type     | The type of instance to use for your runner. For example: t2.micro, t4g.nano, etc. Will not start if not specified.| true               |         |
 | aws_region_name       | The AWS region name to use for your runner. Defaults to AWS_REGION                                                 | true               |         |
 | aws_root_device_size  | The root device size in GB to use for your runner.                                                                 | false              | The AMI default root disk size |
@@ -46,7 +47,8 @@ jobs:
         id: aws-start
         uses: omsf/start-aws-gha-runner@v1.0.0
         with:
-          aws_image_id: ami-0f7c4a792e3fb63c8
+          aws_image_name: "Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 24.04)"
+          aws_image_id: latest
           aws_instance_type: g4dn.xlarge
           aws_home_dir: /home/ubuntu
         env:

action.yml

@@ -11,8 +11,10 @@ inputs:
     description: "The optional AWS IAM role to assume for provisioning your runner."
     required: false
   aws_image_id:
-    description: "The machine AMI to use for your runner. This AMI can be a default but should have docker installed in the AMI. Will not start if not specified."
+    description: "The machine AMI to use for your runner. This AMI can be a default but should have docker installed in the AMI. Will not start if not specified. If set to `latest`, `aws_image_name` is required"
     required: true
+  aws_image_name:
+    description: "The AMI name. Only required if `aws_image_id` is set to latest. Be as specific as possible. For example: Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 24.04)"
   aws_instance_type:
     description: "The type of instance to use for your runner. For example: t2.micro, t4g.nano, etc. Will not start if not specified."
     required: true

src/start_aws_gha_runner/__main__.py

@@ -20,6 +20,7 @@ def main():
     builder = (
         EnvVarBuilder(env)
         .update_state("INPUT_AWS_IMAGE_ID", "image_id")
+        .update_state("INPUT_AWS_IMAGE_NAME", "image_name")
         .update_state("INPUT_AWS_INSTANCE_TYPE", "instance_type")
         .update_state("INPUT_AWS_SUBNET_ID", "subnet_id")
         .update_state("INPUT_AWS_SECURITY_GROUP_ID", "security_group_id")

src/start_aws_gha_runner/start.py

@@ -52,6 +52,7 @@ class StartAWS(CreateCloudInstance):
     repo: str
     region_name: str
     runner_release: str = ""
+    image_name: str = ""
     tags: list[dict[str, str]] = field(default_factory=list)
     gh_runner_tokens: list[str] = field(default_factory=list)
     root_device_size: int = 0
@@ -119,8 +120,28 @@ def _build_user_data(self, **kwargs) -> str:
             except Exception as e:
                 raise Exception(f"Error parsing user data template: {e}")
 
+    def _fetch_latest_ami(
+        self, client, ami_name: str, owner: str = "amazon"
+    ) -> str:
+        out = client.describe_images(
+            Owners=[owner],
+            Filters=[
+                {
+                    "Name": "name",
+                    "Values": [f"{ami_name}*"],
+                },
+                {"Name": "state", "Values": ["available"]},
+            ],
+        )
+        images = out.get("Images", [])
+        newest = sorted(images, key=lambda i: i["CreationDate"], reverse=True)[
+            0
+        ]
+
+        return newest["ImageId"]
+
     def _modify_root_disk_size(self, client, params: dict) -> dict:
-        """ Modify the root disk size of the instance.
+        """Modify the root disk size of the instance.
 
         Parameters
         ----------
@@ -146,11 +167,15 @@ def _modify_root_disk_size(self, client, params: dict) -> dict:
             if "DryRunOperation" in str(e):
                 image_options = client.describe_images(ImageIds=[self.image_id])
                 root_device_name = image_options["Images"][0]["RootDeviceName"]
-                block_devices = deepcopy(image_options["Images"][0]["BlockDeviceMappings"])
+                block_devices = deepcopy(
+                    image_options["Images"][0]["BlockDeviceMappings"]
+                )
                 for idx, block_device in enumerate(block_devices):
                     if block_device["DeviceName"] == root_device_name:
                         if self.root_device_size > 0:
-                            block_devices[idx]["Ebs"]["VolumeSize"] = self.root_device_size
+                            block_devices[idx]["Ebs"]["VolumeSize"] = (
+                                self.root_device_size
+                            )
                             params["BlockDeviceMappings"] = block_devices
                         break
             else:
@@ -206,6 +231,14 @@ def create_instances(self) -> dict[str, str]:
                 "runner_release": self.runner_release,
                 "labels": labels,
             }
+            # We need to handle the case where someone wants to always use latest
+            if self.image_id == "latest":
+                if not self.image_name:
+                    raise ValueError(
+                        "Looking for latest image but name not provided"
+                    )
+                # This updates the image ID to the latest, will fail if image does not exist
+                self.image_id = self._fetch_latest_ami(ec2, self.image_name)
             params = self._build_aws_params(user_data_params)
             if self.root_device_size > 0:
                 params = self._modify_root_disk_size(ec2, params)

tests/test_start.py

@@ -1,5 +1,6 @@
 import pytest
 from moto import mock_aws
+from moto.ec2.models import ec2_backends
 import boto3
 from unittest.mock import call, patch, mock_open, Mock
 from start_aws_gha_runner.start import StartAWS
@@ -21,6 +22,26 @@ def aws():
         yield StartAWS(**params)
 
 
+@pytest.fixture(scope="function")
+def aws_latest_ami():
+    with mock_aws():
+        params = {
+            "image_id": "latest",
+            # This comes from https://github.com/getmoto/moto/blob/master/moto/ec2/resources/amis.json
+            # These are AMIs used to mock out data.
+            # For more info see here:
+            # https://docs.getmoto.org/en/latest/docs/services/ec2.html
+            "image_name": "Ubuntu CUDA9 DLAMI",
+            "instance_type": "t2.micro",
+            "region_name": "us-east-1",
+            "gh_runner_tokens": ["testing"],
+            "home_dir": "/home/ec2-user",
+            "runner_release": "testing",
+            "repo": "omsf-eco-infra/awsinfratesting",
+        }
+        yield StartAWS(**params)
+
+
 def test_build_user_data(aws):
     params = {
         "homedir": "/home/ec2-user",
@@ -57,6 +78,7 @@ def test_build_user_data_missing_params(aws):
     with pytest.raises(Exception):
         aws._build_user_data(**params)
 
+
 @pytest.fixture(scope="function")
 def complete_params():
     params = {
@@ -74,10 +96,11 @@ def complete_params():
         "subnet_id": "test",
         "security_group_id": "test",
         "iam_role": "test",
-        "root_device_size": 100
+        "root_device_size": 100,
     }
     yield params
 
+
 def test_build_aws_params(complete_params):
     user_data_params = {
         "token": "test",
@@ -119,40 +142,37 @@ def test_build_aws_params(complete_params):
         ],
     }
 
+
 def test_modify_root_disk_size(complete_params):
     mock_client = Mock()
 
     # Mock image data with all device mappings
     mock_image_data = {
-        "Images": [{
-            "RootDeviceName": "/dev/sda1",
-            "BlockDeviceMappings": [
-                {
-                    "Ebs": {
-                        "DeleteOnTermination": True,
-                        "VolumeSize": 50,
-                        "VolumeType": "gp3",
-                        "Encrypted": False
+        "Images": [
+            {
+                "RootDeviceName": "/dev/sda1",
+                "BlockDeviceMappings": [
+                    {
+                        "Ebs": {
+                            "DeleteOnTermination": True,
+                            "VolumeSize": 50,
+                            "VolumeType": "gp3",
+                            "Encrypted": False,
+                        },
+                        "DeviceName": "/dev/sda1",
                     },
-                    "DeviceName": "/dev/sda1"
-                },
-                {
-                    "DeviceName": "/dev/sdb",
-                    "VirtualName": "ephemeral0"
-                },
-                {
-                    "DeviceName": "/dev/sdc",
-                    "VirtualName": "ephemeral1"
-                }
-            ]
-        }]
+                    {"DeviceName": "/dev/sdb", "VirtualName": "ephemeral0"},
+                    {"DeviceName": "/dev/sdc", "VirtualName": "ephemeral1"},
+                ],
+            }
+        ]
     }
 
     def mock_describe_images(**kwargs):
-        if kwargs.get('DryRun', False):
+        if kwargs.get("DryRun", False):
             raise ClientError(
                 error_response={"Error": {"Code": "DryRunOperation"}},
-                operation_name="DescribeImages"
+                operation_name="DescribeImages",
             )
         return mock_image_data
 
@@ -168,65 +188,57 @@ def mock_describe_images(**kwargs):
                     "DeleteOnTermination": True,
                     "VolumeSize": 100,
                     "VolumeType": "gp3",
-                    "Encrypted": False
-                }
-            },
-            {
-                "DeviceName": "/dev/sdb",
-                "VirtualName": "ephemeral0"
+                    "Encrypted": False,
+                },
             },
-            {
-                "DeviceName": "/dev/sdc",
-                "VirtualName": "ephemeral1"
-            }
+            {"DeviceName": "/dev/sdb", "VirtualName": "ephemeral0"},
+            {"DeviceName": "/dev/sdc", "VirtualName": "ephemeral1"},
         ]
     }
     assert out == expected_output
 
+
 def test_modify_root_disk_size_permission_error(complete_params):
     mock_client = Mock()
 
     # Mock permission denied error
     mock_client.describe_images.side_effect = ClientError(
-        error_response={'Error': {'Code': 'AccessDenied'}},
-        operation_name='DescribeImages'
+        error_response={"Error": {"Code": "AccessDenied"}},
+        operation_name="DescribeImages",
     )
 
     aws = StartAWS(**complete_params)
 
     with pytest.raises(ClientError) as exc_info:
         aws._modify_root_disk_size(mock_client, {})
 
-    assert 'AccessDenied' in str(exc_info.value)
+    assert "AccessDenied" in str(exc_info.value)
+
 
 def test_modify_root_disk_size_no_change(complete_params):
     mock_client = Mock()
     complete_params["root_device_size"] = 0
 
     mock_image_data = {
-        "Images": [{
-            "RootDeviceName": "/dev/sda1",
-            "BlockDeviceMappings": [
-                {
-                    "DeviceName": "/dev/sda1",
-                    "Ebs": {
-                        "VolumeSize": 50,
-                        "VolumeType": "gp3"
-                    }
-                },
-                {
-                    "DeviceName": "/dev/sdb",
-                    "VirtualName": "ephemeral0"
-                }
-            ]
-        }]
+        "Images": [
+            {
+                "RootDeviceName": "/dev/sda1",
+                "BlockDeviceMappings": [
+                    {
+                        "DeviceName": "/dev/sda1",
+                        "Ebs": {"VolumeSize": 50, "VolumeType": "gp3"},
+                    },
+                    {"DeviceName": "/dev/sdb", "VirtualName": "ephemeral0"},
+                ],
+            }
+        ]
     }
 
     def mock_describe_images(**kwargs):
-        if kwargs.get('DryRun', False):
+        if kwargs.get("DryRun", False):
             raise ClientError(
-                error_response={'Error': {'Code': 'DryRunOperation'}},
-                operation_name='DescribeImages'
+                error_response={"Error": {"Code": "DryRunOperation"}},
+                operation_name="DescribeImages",
             )
         return mock_image_data
 
@@ -239,6 +251,59 @@ def mock_describe_images(**kwargs):
     # With root_device_size = 0, no modifications should be made
     assert result == input_params
 
+
+@pytest.fixture(scope="function")
+def complete_params_latest():
+    params = {
+        "image_name": "Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 22.04)",
+        "image_id": "latest",
+        "instance_type": "t2.micro",
+        "tags": [
+            {"Key": "Name", "Value": "test"},
+            {"Key": "Owner", "Value": "test"},
+        ],
+        "region_name": "us-east-1",
+        "gh_runner_tokens": ["testing"],
+        "home_dir": "/home/ec2-user",
+        "runner_release": "testing",
+        "repo": "omsf-eco-infra/awsinfratesting",
+        "subnet_id": "test",
+        "security_group_id": "test",
+        "iam_role": "test",
+        "root_device_size": 100,
+    }
+    yield params
+
+
+def test_fetch_latest_ami(complete_params_latest):
+    mock_client = Mock()
+
+    mock_image_data = {
+        "Images": [
+            {"CreationDate": "2025-08-03", "ImageId": "ami-12345678"},
+            {"CreationDate": "2025-08-05", "ImageId": "ami-89123456"},
+            {"CreationDate": "2025-09-05", "ImageId": "ami-89121111"},
+        ]
+    }
+    mock_client.describe_images.return_value = mock_image_data
+    aws = StartAWS(**complete_params_latest)
+    result = aws._fetch_latest_ami(mock_client, "Test")
+    assert result == "ami-89121111"
+
+
+def test_create_instances_latest(aws_latest_ami):
+    ids = aws_latest_ami.create_instances()
+    assert len(ids) == 1
+
+
+def test_create_instatnces_latest_no_name(aws_latest_ami):
+    aws_latest_ami.image_name = ""
+    with pytest.raises(
+        ValueError, match="Looking for latest image but name not provided"
+    ):
+        aws_latest_ami.create_instances()
+
+
 def test_create_instance_with_labels(aws):
     aws.labels = "test"
     ids = aws.create_instances()