diff --git a/runner_manager/models/backend.py b/runner_manager/models/backend.py index 599d7ac2..ce75bda7 100644 --- a/runner_manager/models/backend.py +++ b/runner_manager/models/backend.py @@ -3,11 +3,16 @@ from string import Template from typing import Dict, List, Literal, Optional, Sequence, TypedDict -from mypy_boto3_ec2.literals import InstanceTypeType, VolumeTypeType +from mypy_boto3_ec2.literals import ( + InstanceMetadataTagsStateType, + InstanceTypeType, + VolumeTypeType, +) from mypy_boto3_ec2.type_defs import ( BlockDeviceMappingTypeDef, EbsBlockDeviceTypeDef, IamInstanceProfileTypeDef, + InstanceMetadataOptionsRequestTypeDef, TagSpecificationTypeDef, TagTypeDef, ) @@ -141,6 +146,7 @@ class AWSConfig(BackendConfig): "MaxCount": int, "MinCount": int, "IamInstanceProfile": IamInstanceProfileTypeDef, + "MetadataOptions": InstanceMetadataOptionsRequestTypeDef, }, ) @@ -159,6 +165,7 @@ class AWSInstanceConfig(InstanceConfig): volume_type: VolumeTypeType = "gp3" disk_size_gb: int = 20 iam_instance_profile_arn: str = "" + instance_metadata_tags: InstanceMetadataTagsStateType = "disabled" def configure_instance(self, runner: Runner) -> AwsInstance: """Configure instance.""" @@ -197,6 +204,11 @@ def configure_instance(self, runner: Runner) -> AwsInstance: iam_instance_profile: IamInstanceProfileTypeDef = IamInstanceProfileTypeDef( Arn=self.iam_instance_profile_arn ) + instance_metadata_options: InstanceMetadataOptionsRequestTypeDef = ( + InstanceMetadataOptionsRequestTypeDef( + InstanceMetadataTags=self.instance_metadata_tags, + ) + ) return AwsInstance( ImageId=self.image, InstanceType=self.instance_type, @@ -208,6 +220,7 @@ def configure_instance(self, runner: Runner) -> AwsInstance: MinCount=self.min_count, BlockDeviceMappings=block_device_mappings, IamInstanceProfile=iam_instance_profile, + MetadataOptions=instance_metadata_options, ) diff --git a/tests/unit/backend/test_aws.py b/tests/unit/backend/test_aws.py index 61310dbc..b5633411 100644 --- a/tests/unit/backend/test_aws.py +++ b/tests/unit/backend/test_aws.py @@ -49,6 +49,7 @@ def test_aws_instance_config(runner: Runner): tags={"test": "test"}, subnet_id="i-0f9b0a3b7b3b3b3b3", iam_instance_profile_arn="test", + instance_metadata_tags="enabled", ) instance: AwsInstance = instance_config.configure_instance(runner) assert instance["ImageId"] == instance_config.image @@ -57,6 +58,10 @@ def test_aws_instance_config(runner: Runner): instance["IamInstanceProfile"]["Arn"] == instance_config.iam_instance_profile_arn ) + assert ( + instance["MetadataOptions"]["InstanceMetadataTags"] + == instance_config.instance_metadata_tags + ) assert runner.name in instance["UserData"] tags = instance["TagSpecifications"][0]["Tags"] assert TagTypeDef(Key="test", Value="test") in tags