Merge pull request #46 from ethanholz/feat/storage

feat: add support for custom sized root block devices.

README.md

@@ -33,6 +33,7 @@ For more information see the [self-hosted runner security docs](https://docs.git
 | 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_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. Will not start if not specified.                                       | true               |         |
+| aws_root_device_size  | The root device size in GB to use for your runner.                                                                 | false              | The default AMI root device size |
 | aws_security_group_id | The AWS security group ID to use for your runner. Will use the account default security group if not specified.    | false              | The default AWS security group |
 | aws_subnet_id         | The AWS subnet ID to use for your runner. Will use the account default subnet if not specified.                    | false              | The default AWS subnet ID |
 | aws_tags              | The AWS tags to use for your runner, formatted as a JSON list. See `README` for more details.                      | false              |         |

action.yml

@@ -16,6 +16,9 @@ inputs:
   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."
     required: false
+  aws_root_device_size:
+    description: "The root device size in GB to use for your runner. Optional, defaults to the AMI default."
+    required: false
   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: false

docs/aws.md

@@ -29,7 +29,8 @@ The goal of this document is to provide a guide on how to set up the GitHub Acti
                 "ec2:RunInstances",
                 "ec2:TerminateInstances",
                 "ec2:DescribeInstances",
-                "ec2:DescribeInstanceStatus"
+                "ec2:DescribeInstanceStatus",
+                "ec2:DescribeImages"
               ],
               "Resource": "*"
             }

src/gha_runner/__main__.py

@@ -24,6 +24,10 @@ def parse_aws_params() -> dict:
     instance_type = os.environ.get("INPUT_AWS_INSTANCE_TYPE")
     if instance_type is not None:
         params["instance_type"] = instance_type
+    root_device_size = os.environ.get("INPUT_AWS_ROOT_DEVICE_SIZE")
+    # We need to convert this to an integer, but it is not required to start
+    if root_device_size is not None and root_device_size != "":
+        params["root_device_size"] = int(root_device_size)
     params = _env_parse_helper(params, "INPUT_AWS_SUBNET_ID", "subnet_id")
     params = _env_parse_helper(
         params, "INPUT_AWS_SECURITY_GROUP_ID", "security_group_id"

src/gha_runner/clouddeployment.py

@@ -1,8 +1,10 @@
+from copy import deepcopy
 from abc import ABC, abstractmethod
 from gha_runner.gh import GitHubInstance
 from dataclasses import dataclass, field
 import importlib.resources
 import boto3
+from botocore.exceptions import ClientError
 from string import Template
 
 
@@ -93,6 +95,7 @@ class AWS(CloudDeployment):
     runner_release: str = ""
     tags: list[dict[str, str]] = field(default_factory=list)
     gh_runner_tokens: list[str] = field(default_factory=list)
+    root_device_size: int = 0
     labels: str = ""
     subnet_id: str = ""
     security_group_id: str = ""
@@ -134,6 +137,48 @@ def _build_aws_params(self, user_data_params: dict) -> dict:
 
         return params
 
+    def _modify_root_disk_size(self, client, params) -> dict:
+        """Modify the root disk size of the instance.
+
+        Parameters
+        ----------
+        client
+            The boto3 client to use for the API call.
+        params
+            The parameters for the create_instances AWS API call.
+
+        Returns
+        -------
+        dict
+            The modified parameters for the AWS API call.
+
+        Raises
+        ------
+        ClientError
+            If the user does not have permissions to describe images.
+
+        """
+        try:
+            # Check if we have permissions to describe images
+            client.describe_images(ImageIds=[self.image_id], DryRun=True)
+        except ClientError as e:
+            # This is the case where we do have permissions
+            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"])
+                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
+                            params["BlockDeviceMappings"] = block_devices
+                        break
+            else:
+                # If not, we should receive an UnauthorizedOperation error
+                raise e
+        return params
+
+
     def create_instances(self) -> dict[str, str]:
         if not self.gh_runner_tokens:
             raise ValueError(
@@ -172,6 +217,8 @@ def create_instances(self) -> dict[str, str]:
                 "labels": labels,
             }
             params = self._build_aws_params(user_data_params)
+            if self.root_device_size > 0:
+                params = self._modify_root_disk_size(ec2, params)
             result = ec2.run_instances(**params)
             instances = result["Instances"]
             id = instances[0]["InstanceId"]