def _SetThroughput(self, rcu: int, wcu: int) -> None:
"""Updates the table's rcu and wcu."""
cmd = util.AWS_PREFIX + [
'dynamodb', 'update-table',
'--table-name', self.table_name,
'--region', self.region,
'--provisioned-throughput',
f'ReadCapacityUnits={rcu},WriteCapacityUnits={wcu}',
]
logging.info('Setting %s table provisioned throughput to %s rcu and %s wcu',
self.table_name, rcu, wcu)
util.IssueRetryableCommand(cmd)
def _Exists(self):
"""Returns true if the internet gateway exists."""
describe_cmd = util.AWS_PREFIX + [
'ec2', 'describe-internet-gateways',
'--region=%s' % self.region,
'--filter=Name=internet-gateway-id,Values=%s' % self.id
]
stdout, _ = util.IssueRetryableCommand(describe_cmd)
response = json.loads(stdout)
internet_gateways = response['InternetGateways']
assert len(internet_gateways) < 2, 'Too many internet gateways.'
return len(internet_gateways) > 0
def Attach(self, vpc_id):
"""Attaches the internetgateway to the VPC."""
if not self.attached:
self.vpc_id = vpc_id
attach_cmd = util.AWS_PREFIX + [
'ec2', 'attach-internet-gateway',
'--region=%s' % self.region,
'--internet-gateway-id=%s' % self.id,
'--vpc-id=%s' % self.vpc_id
]
util.IssueRetryableCommand(attach_cmd)
self.attached = True
def _Exists(self):
"""Returns true if the VPC exists."""
describe_cmd = util.AWS_PREFIX + [
'ec2', 'describe-vpcs',
'--region=%s' % self.region,
'--filter=Name=vpc-id,Values=%s' % self.id
]
stdout, _ = util.IssueRetryableCommand(describe_cmd)
response = json.loads(stdout)
vpcs = response['Vpcs']
assert len(vpcs) < 2, 'Too many VPCs.'
return len(vpcs) > 0
def _Exists(self):
"""Returns true if the subnet exists."""
describe_cmd = util.AWS_PREFIX + [
'ec2', 'describe-subnets',
'--region=%s' % self.region,
'--filter=Name=subnet-id,Values=%s' % self.id
]
stdout, _ = util.IssueRetryableCommand(describe_cmd)
response = json.loads(stdout)
subnets = response['Subnets']
assert len(subnets) < 2, 'Too many subnets.'
return len(subnets) > 0
def _Exists(self):
"""Returns true if the Placement Group exists."""
describe_cmd = util.AWS_PREFIX + [
'ec2',
'describe-placement-groups',
'--region=%s' % self.region,
'--filter=Name=group-name,Values=%s' % self.name]
stdout, _ = util.IssueRetryableCommand(describe_cmd)
response = json.loads(stdout)
placement_groups = response['PlacementGroups']
assert len(placement_groups) < 2, 'Too many placement groups.'
return bool(placement_groups)
def _RunCommand(self, cmds):
"""Runs the AWS ec2 command in the defined region.
Args:
cmds: List of AWS ec2 commands to run, example: ['describe-route-tables']
Returns:
Dict of the AWS response.
"""
cmd = util.AWS_PREFIX + ['ec2', '--region=%s' % self.region] + list(cmds)
stdout, _ = util.IssueRetryableCommand(cmd)
return json.loads(stdout)
def _GetImageId(self):
"""Gets the current EKS worker image for the region."""
describe_cmd = util.AWS_PREFIX + [
'--region=%s' % self.region,
'ec2',
'describe-images',
'--query', 'Images[*].{Name:Name,ImageId:ImageId}',
'--filters',
'Name=name,Values=eks-worker-v*',
]
stdout, _ = util.IssueRetryableCommand(describe_cmd)
return max(json.loads(stdout), key=lambda image: image['Name'])['ImageId']
def AllowPortInSecurityGroup(self,
region,
security_group,
start_port,
end_port=None,
source_range=None):
"""Opens a port on the firewall for a security group.
Args:
region: The region of the security group
security_group: The security group in which to open the ports
start_port: The first local port to open in a range.
end_port: The last local port to open in a range. If None, only start_port
will be opened.
source_range: List of source CIDRs to allow for this port.
"""
end_port = end_port or start_port
source_range = source_range or ['0.0.0.0/0']
for source in source_range:
entry = (start_port, end_port, region, security_group, source)
if entry in self.firewall_set:
continue
if self._RuleExists(region, security_group, start_port, end_port,
source):
self.firewall_set.add(entry)
continue
with self._lock:
if entry in self.firewall_set:
continue
authorize_cmd = util.AWS_PREFIX + [
'ec2',
'authorize-security-group-ingress',
'--region=%s' % region,
'--group-id=%s' % security_group,
'--port=%s-%s' % (start_port, end_port),
'--cidr=%s' % source,
]
util.IssueRetryableCommand(authorize_cmd + ['--protocol=tcp'])
util.IssueRetryableCommand(authorize_cmd + ['--protocol=udp'])
self.firewall_set.add(entry)
def _Exists(self):
"""Returns whether the VM exists.
This method waits until the VM is no longer pending.
Returns:
Whether the VM exists.
Raises:
AwsUnknownStatusError: If an unknown status is returned from AWS.
AwsTransitionalVmRetryableError: If the VM is pending. This is retried.
"""
describe_cmd = util.AWS_PREFIX + [
'ec2', 'describe-instances',
'--region=%s' % self.region,
'--filter=Name=client-token,Values=%s' % self.client_token
]
stdout, _ = util.IssueRetryableCommand(describe_cmd)
response = json.loads(stdout)
reservations = response['Reservations']
assert len(reservations) < 2, 'Too many reservations.'
if not reservations:
if not self.create_start_time:
return False
logging.info('No reservation returned by describe-instances. This '
'sometimes shows up immediately after a successful '
'run-instances command. Retrying describe-instances '
'command.')
raise AwsTransitionalVmRetryableError()
instances = reservations[0]['Instances']
assert len(instances) == 1, 'Wrong number of instances.'
status = instances[0]['State']['Name']
self.id = instances[0]['InstanceId']
if self.use_spot_instance:
self.spot_instance_request_id = instances[0][
'SpotInstanceRequestId']
if status not in INSTANCE_KNOWN_STATUSES:
raise AwsUnknownStatusError('Unknown status %s' % status)
if status in INSTANCE_TRANSITIONAL_STATUSES:
logging.info(
'VM has status %s; retrying describe-instances command.',
status)
raise AwsTransitionalVmRetryableError()
# In this path run-instances succeeded, a pending instance was created, but
# not fulfilled so it moved to terminated.
if (status == TERMINATED and instances[0]['StateReason']['Code']
== 'Server.InsufficientInstanceCapacity'):
raise errors.Benchmarks.InsufficientCapacityCloudFailure(
instances[0]['StateReason']['Message'])
return status in INSTANCE_EXISTS_STATUSES
def CreateRoute(self, internet_gateway_id):
"""Adds a route to the internet gateway."""
if self.RouteExists():
logging.info('Internet route already exists.')
return
create_cmd = util.AWS_PREFIX + [
'ec2',
'create-route',
'--region=%s' % self.region,
'--route-table-id=%s' % self.id,
'--gateway-id=%s' % internet_gateway_id,
'--destination-cidr-block=0.0.0.0/0']
util.IssueRetryableCommand(create_cmd)
def DeleteKeyfile(cls, region):
"""Deletes the imported keyfile for a region."""
with cls._lock:
if _GetKeyfileSetKey(region) in cls.deleted_keyfile_set:
return
delete_cmd = util.AWS_PREFIX + [
'ec2', '--region=%s' % region,
'delete-key-pair',
'--key-name=%s' % cls.GetKeyNameForRun()]
util.IssueRetryableCommand(delete_cmd)
cls.deleted_keyfile_set.add(_GetKeyfileSetKey(region))
if _GetKeyfileSetKey(region) in cls.imported_keyfile_set:
cls.imported_keyfile_set.remove(_GetKeyfileSetKey(region))
def DeleteKeyfile(self):
"""Deletes the imported keyfile for a region."""
with self._lock:
if self.region in self.deleted_keyfile_set:
return
delete_cmd = util.AWS_PREFIX + [
'ec2', '--region=%s' % self.region,
'delete-key-pair',
'--key-name=%s' % 'perfkit-key-%s' % FLAGS.run_uri]
util.IssueRetryableCommand(delete_cmd)
self.deleted_keyfile_set.add(self.region)
if self.region in self.imported_keyfile_set:
self.imported_keyfile_set.remove(self.region)
def _EnableDnsHostnames(self):
"""Sets the enableDnsHostnames attribute of this VPC to True.
By default, instances launched in non-default VPCs are assigned an
unresolvable hostname. This breaks the hadoop benchmark. Setting the
enableDnsHostnames attribute to 'true' on the VPC resolves this. See:
http://docs.aws.amazon.com/AmazonVPC/latest/UserGuide/VPC_DHCP_Options.html
"""
enable_hostnames_command = util.AWS_PREFIX + [
'ec2', 'modify-vpc-attribute',
'--region=%s' % self.region, '--vpc-id', self.id,
'--enable-dns-hostnames', '{ "Value": true }'
]
util.IssueRetryableCommand(enable_hostnames_command)
def Detach(self):
"""Detaches the disk from a VM."""
detach_cmd = util.AWS_PREFIX + [
'ec2', 'detach-volume',
'--region=%s' % self.region,
'--instance-id=%s' % self.attached_vm_id,
'--volume-id=%s' % self.id
]
util.IssueRetryableCommand(detach_cmd)
with self._lock:
assert self.attached_vm_id in AwsDisk.vm_devices
AwsDisk.vm_devices[self.attached_vm_id].add(self.device_letter)
self.attached_vm_id = None
self.device_letter = None
def AllowPort(self, vm, port):
"""Opens a port on the firewall.
Args:
vm: The BaseVirtualMachine object to open the port for.
port: The local port to open.
"""
if vm.is_static:
return
entry = (port, vm.group_id)
if entry in self.firewall_set:
return
with self._lock:
if entry in self.firewall_set:
return
authorize_cmd = util.AWS_PREFIX + [
'ec2', 'authorize-security-group-ingress',
'--region=%s' % vm.region,
'--group-id=%s' % vm.group_id,
'--port=%s' % port, '--cidr=0.0.0.0/0'
]
util.IssueRetryableCommand(authorize_cmd + ['--protocol=tcp'])
util.IssueRetryableCommand(authorize_cmd + ['--protocol=udp'])
self.firewall_set.add(entry)
def _Exists(self):
"""Returns true if the disk exists."""
describe_cmd = util.AWS_PREFIX + [
'ec2', 'describe-volumes',
'--region=%s' % self.region,
'--filter=Name=volume-id,Values=%s' % self.id
]
stdout, _ = util.IssueRetryableCommand(describe_cmd)
response = json.loads(stdout)
volumes = response['Volumes']
assert len(volumes) < 2, 'Too many volumes.'
if not volumes:
return False
status = volumes[0]['State']
assert status in VOLUME_KNOWN_STATUSES, status
return status in VOLUME_EXISTS_STATUSES
def _RuleExists(self, region, security_group, start_port, end_port, source):
"""Whether the firewall rule exists in the VPC."""
query_cmd = util.AWS_PREFIX + [
'ec2',
'describe-security-groups',
'--region=%s' % region,
'--group-ids=%s' % security_group,
'--filters',
'Name=ip-permission.cidr,Values={}'.format(source),
'Name=ip-permission.from-port,Values={}'.format(start_port),
'Name=ip-permission.to-port,Values={}'.format(end_port),
]
stdout, _ = util.IssueRetryableCommand(query_cmd)
# "groups" will be an array of all the matching firewall rules
groups = json.loads(stdout)['SecurityGroups']
return bool(groups)
def ImportKeyfile(self):
"""Imports the public keyfile to AWS."""
with self._lock:
if self.region in self.imported_keyfile_set:
return
cat_cmd = ['cat',
vm_util.GetPublicKeyPath()]
keyfile, _ = vm_util.IssueRetryableCommand(cat_cmd)
import_cmd = util.AWS_PREFIX + [
'ec2', '--region=%s' % self.region,
'import-key-pair',
'--key-name=%s' % 'perfkit-key-%s' % FLAGS.run_uri,
'--public-key-material=%s' % keyfile]
util.IssueRetryableCommand(import_cmd)
self.imported_keyfile_set.add(self.region)
if self.region in self.deleted_keyfile_set:
self.deleted_keyfile_set.remove(self.region)
def ImportKeyfile(cls, region):
"""Imports the public keyfile to AWS."""
with cls._lock:
if _GetKeyfileSetKey(region) in cls.imported_keyfile_set:
return
cat_cmd = ['cat', vm_util.GetPublicKeyPath()]
keyfile, _ = vm_util.IssueRetryableCommand(cat_cmd)
import_cmd = util.AWS_PREFIX + [
'ec2',
'--region=%s' % region, 'import-key-pair',
'--key-name=%s' % cls.GetKeyNameForRun(),
'--public-key-material=%s' % keyfile
]
util.IssueRetryableCommand(import_cmd)
cls.imported_keyfile_set.add(_GetKeyfileSetKey(region))
if _GetKeyfileSetKey(region) in cls.deleted_keyfile_set:
cls.deleted_keyfile_set.remove(_GetKeyfileSetKey(region))
def _GetDefaultImage(cls, machine_type, region):
"""Returns the default image given the machine type and region.
If specified, aws_image_name_filter will override os_type defaults.
If no default is configured, this will return None.
Args:
machine_type: The machine_type of the VM, used to determine virtualization
type.
region: The region of the VM, as images are region specific.
"""
if cls.IMAGE_NAME_FILTER is None:
return None
if FLAGS.aws_image_name_filter:
cls.IMAGE_NAME_FILTER = FLAGS.aws_image_name_filter
prefix = machine_type.split('.')[0]
virt_type = 'paravirtual' if prefix in NON_HVM_PREFIXES else 'hvm'
describe_cmd = util.AWS_PREFIX + [
'--region=%s' % region,
'ec2',
'describe-images',
'--query', 'Images[*].{Name:Name,ImageId:ImageId}',
'--filters',
'Name=name,Values=%s' % cls.IMAGE_NAME_FILTER,
'Name=block-device-mapping.volume-type,Values=%s' %
cls.DEFAULT_ROOT_DISK_TYPE,
'Name=virtualization-type,Values=%s' % virt_type]
if cls.IMAGE_PRODUCT_CODE_FILTER:
describe_cmd.extend(['Name=product-code,Values=%s' %
cls.IMAGE_PRODUCT_CODE_FILTER])
if cls.IMAGE_OWNER:
describe_cmd.extend(['--owners', cls.IMAGE_OWNER])
stdout, _ = util.IssueRetryableCommand(describe_cmd)
if not stdout:
return None
images = json.loads(stdout)
# We want to return the latest version of the image, and since the wildcard
# portion of the image name is the image's creation date, we can just take
# the image with the 'largest' name.
return max(images, key=lambda image: image['Name'])['ImageId']
def _Exists(self):
"""Returns true if the VM exists."""
describe_cmd = util.AWS_PREFIX + [
'ec2', 'describe-instances',
'--region=%s' % self.region,
'--filter=Name=instance-id,Values=%s' % self.id
]
stdout, _ = util.IssueRetryableCommand(describe_cmd)
response = json.loads(stdout)
reservations = response['Reservations']
assert len(reservations) < 2, 'Too many reservations.'
if not reservations:
return False
instances = reservations[0]['Instances']
assert len(instances) == 1, 'Wrong number of instances.'
status = instances[0]['State']['Name']
assert status in INSTANCE_KNOWN_STATUSES, status
return status in INSTANCE_EXISTS_STATUSES
def _GetDecodedPasswordData(self):
# Retrieve a base64 encoded, encrypted password for the VM.
get_password_cmd = util.AWS_PREFIX + [
'ec2',
'get-password-data',
'--region=%s' % self.region,
'--instance-id=%s' % self.id]
stdout, _ = util.IssueRetryableCommand(get_password_cmd)
response = json.loads(stdout)
password_data = response['PasswordData']
# AWS may not populate the password data until some time after
# the VM shows as running. Simply retry until the data shows up.
if not password_data:
raise ValueError('No PasswordData in response.')
# Decode the password data.
return base64.b64decode(password_data)
def Detach(self):
"""Detaches the internet gateway from the VPC."""
def _suppress_failure(stdout, stderr, retcode):
"""Suppresses Detach failure when internet gateway is in a bad state."""
del stdout # unused
if retcode and ('InvalidInternetGatewayID.NotFound' in stderr or
'Gateway.NotAttached' in stderr):
return True
return False
if self.attached and not self.user_managed:
detach_cmd = util.AWS_PREFIX + [
'ec2',
'detach-internet-gateway',
'--region=%s' % self.region,
'--internet-gateway-id=%s' % self.id,
'--vpc-id=%s' % self.vpc_id]
util.IssueRetryableCommand(detach_cmd, suppress_failure=_suppress_failure)
self.attached = False
def _PostCreate(self):
"""Get the instance's data and tag it."""
describe_cmd = util.AWS_PREFIX + [
'ec2',
'describe-instances',
'--region=%s' % self.region,
'--instance-ids=%s' % self.id]
logging.info('Getting instance %s public IP. This will fail until '
'a public IP is available, but will be retried.', self.id)
stdout, _ = util.IssueRetryableCommand(describe_cmd)
response = json.loads(stdout)
instance = response['Reservations'][0]['Instances'][0]
self.ip_address = instance['PublicIpAddress']
self.internal_ip = instance['PrivateIpAddress']
if util.IsRegion(self.zone):
self.zone = str(instance['Placement']['AvailabilityZone'])
util.AddDefaultTags(self.id, self.region)
assert self.group_id == instance['SecurityGroups'][0]['GroupId'], (
self.group_id, instance['SecurityGroups'][0]['GroupId'])
def SetThroughput(self,
rcu: Optional[int] = None,
wcu: Optional[int] = None) -> None:
"""Updates the table's rcu and wcu."""
if not rcu:
rcu = self.rcu
if not wcu:
wcu = self.wcu
cmd = util.AWS_PREFIX + [
'dynamodb', 'update-table',
'--table-name', self.table_name,
'--region', self.region,
'--provisioned-throughput',
f'ReadCapacityUnits={rcu},WriteCapacityUnits={wcu}',
]
logging.info('Setting %s table provisioned throughput to %s rcu and %s wcu',
self.table_name, rcu, wcu)
util.IssueRetryableCommand(cmd)
while not self._IsReady():
continue
def GetDict(self):
"""The 'aws ec2 describe-subnets' for this VPC / subnet id.
Returns:
A dict of the single subnet or an empty dict if there are no subnets.
Raises:
AssertionError: If there is more than one subnet.
"""
describe_cmd = util.AWS_PREFIX + [
'ec2', 'describe-subnets',
'--region=%s' % self.region,
'--filter=Name=vpc-id,Values=%s' % self.vpc_id
]
if self.id:
describe_cmd.append('--filter=Name=subnet-id,Values=%s' % self.id)
stdout, _ = util.IssueRetryableCommand(describe_cmd)
response = json.loads(stdout)
subnets = response['Subnets']
assert len(subnets) < 2, 'Too many subnets.'
return subnets[0] if subnets else {}
def Attach(self, vm):
"""Attaches the disk to a VM.
Args:
vm: The AwsVirtualMachine instance to which the disk will be attached.
"""
with self._lock:
self.attached_vm_id = vm.id
if self.attached_vm_id not in AwsDisk.vm_devices:
AwsDisk.vm_devices[self.attached_vm_id] = set(
string.ascii_lowercase)
self.device_letter = min(AwsDisk.vm_devices[self.attached_vm_id])
AwsDisk.vm_devices[self.attached_vm_id].remove(self.device_letter)
attach_cmd = util.AWS_PREFIX + [
'ec2',
'attach-volume',
'--region=%s' % self.region,
'--instance-id=%s' % self.attached_vm_id,
'--volume-id=%s' % self.id,
'--device=%s' % self.GetDevicePath()]
logging.info('Attaching AWS volume %s. This may fail if the disk is not '
'ready, but will be retried.', self.id)
util.IssueRetryableCommand(attach_cmd)
def _PostCreate(self):
"""Get the instance's data and tag it."""
describe_cmd = util.AWS_PREFIX + [
'ec2',
'describe-instances',
'--region=%s' % self.region,
'--instance-ids=%s' % self.id]
logging.info('Getting instance %s public IP. This will fail until '
'a public IP is available, but will be retried.', self.id)
stdout, _ = util.IssueRetryableCommand(describe_cmd)
response = json.loads(stdout)
instance = response['Reservations'][0]['Instances'][0]
self.ip_address = instance['PublicIpAddress']
self.internal_ip = instance['PrivateIpAddress']
if util.IsRegion(self.zone):
self.zone = str(instance['Placement']['AvailabilityZone'])
assert self.group_id == instance['SecurityGroups'][0]['GroupId'], (
self.group_id, instance['SecurityGroups'][0]['GroupId'])
if FLAGS.aws_efa:
self.InstallPackages('curl')
url = _EFA_URL.format(version=FLAGS.aws_efa_version)
self.RemoteCommand(
_EFA_INSTALL_CMD.format(url=url, tarfile=posixpath.basename(url)))
def UpdateWithDefaultTags(self) -> None:
"""Adds default tags to the table."""
tags = util.MakeFormattedDefaultTags()
cmd = self._GetTagResourceCommand(tags)
logging.info('Setting default tags on table %s', self.table_name)
util.IssueRetryableCommand(cmd)
