def _get_default_subnet_id(zone: str) -> str:
"""Returns the default subnet ID for the zone.
Creates a default subnet if one did not exist previously
Args:
zone: Zone of the default subnet.
Returns: Default Subnet ID
Raises:
UnsupportedConfigError: When default subnet does not exist and cannot be
created.
"""
region = util.GetRegionFromZone(zone)
subnet_cmd = util.AWS_PREFIX + [
'ec2', 'describe-subnets', '--region', region, '--filter',
f'Name=availabilityZone,Values={zone}', 'Name=defaultForAz,Values=true'
]
stdout, _ = vm_util.IssueRetryableCommand(subnet_cmd)
subnets = json.loads(stdout)['Subnets']
if subnets:
return subnets[0]['SubnetId']
create_cmd = util.AWS_PREFIX + [
'ec2', 'create-default-subnet', '--region', region,
'--availability-zone', zone
]
stdout, _, ret = vm_util.IssueCommand(create_cmd, raise_on_failure=False)
if ret:
raise errors.Benchmarks.UnsupportedConfigError(
f'AWS default subnet does not exist for zone {zone}.')
return json.loads(stdout)['Subnet']['SubnetId']
def __init__(self, vm_spec):
"""Initialize a AWS virtual machine.
Args:
vm_spec: virtual_machine.BaseVirtualMachineSpec object of the vm.
"""
super(AwsVirtualMachine, self).__init__(vm_spec)
self.region = util.GetRegionFromZone(self.zone)
self.user_name = FLAGS.aws_user_name
if self.machine_type in NUM_LOCAL_VOLUMES:
self.max_local_disks = NUM_LOCAL_VOLUMES[self.machine_type]
self.user_data = None
self.network = aws_network.AwsNetwork.GetNetwork(self)
self.firewall = aws_network.AwsFirewall.GetFirewall()
self.use_dedicated_host = vm_spec.use_dedicated_host
self.use_spot_instance = vm_spec.use_spot_instance
self.spot_price = vm_spec.spot_price
self.boot_disk_size = vm_spec.boot_disk_size
self.client_token = str(uuid.uuid4())
self.host = None
self.id = None
if self.use_dedicated_host and util.IsRegion(self.zone):
raise ValueError(
'In order to use dedicated hosts, you must specify an availability '
'zone, not a region ("zone" was %s).' % self.zone)
if self.use_spot_instance and self.spot_price <= 0.0:
raise ValueError(
'In order to use spot instances you must specify a spot price '
'greater than 0.0.')
def __init__(self, disk_spec, zone):
super(AwsNfsService, self).__init__(disk_spec, zone)
self.region = util.GetRegionFromZone(self.zone)
self.aws_commands = AwsEfsCommands(self.region)
self.disk_spec.disk_size = 0
self.filer_id = None
self.mount_id = None
def __init__(self, managed_relational_db_spec):
super(AwsManagedRelationalDb,
self).__init__(managed_relational_db_spec)
self.spec = managed_relational_db_spec
self.instance_id = 'pkb-db-instance-' + FLAGS.run_uri
self.zone = self.spec.vm_spec.zone
self.region = util.GetRegionFromZone(self.zone)
def __init__(self, spec):
"""Initializes AwsNetwork instances.
Args:
spec: An AwsNetworkSpec object.
"""
super(AwsNetwork, self).__init__(spec)
self.region = util.GetRegionFromZone(spec.zone)
self.regional_network = _AwsRegionalNetwork.GetForRegion(
self.region, spec.vpc_id)
self.subnet = None
self.vpc_peering = None
if (FLAGS.placement_group_style ==
placement_group.PLACEMENT_GROUP_NONE):
self.placement_group = None
else:
placement_group_spec = aws_placement_group.AwsPlacementGroupSpec(
'AwsPlacementGroupSpec', flag_values=FLAGS, zone=spec.zone)
self.placement_group = aws_placement_group.AwsPlacementGroup(
placement_group_spec)
self.is_static = False
if spec.vpc_id:
self.is_static = True
self.subnet = AwsSubnet(
self.zone,
spec.vpc_id,
cidr_block=self.regional_network.cidr_block,
subnet_id=spec.subnet_id)
def __init__(self, table_name, **kwargs):
super(AwsDynamoDBInstance, self).__init__(**kwargs)
self.zone = FLAGS.zones[0] if FLAGS.zones else FLAGS.zone[0]
self.region = util.GetRegionFromZone(self.zone)
self.primary_key = (
'{{\"AttributeName\": \"{0}\",\"KeyType\": \"HASH\"}}'.format(
FLAGS.aws_dynamodb_primarykey))
self.sort_key = (
'{{\"AttributeName\": \"{0}\",\"KeyType\": \"RANGE\"}}'.format(
FLAGS.aws_dynamodb_sortkey))
self.part_attributes = ('{{\"AttributeName\": \"{0}\",'
'\"AttributeType\": \"{1}\"}}'.format(
FLAGS.aws_dynamodb_primarykey,
FLAGS.aws_dynamodb_attributetype))
self.sort_attributes = ('{{\"AttributeName\": \"{0}\",'
'\"AttributeType\": \"{1}\"}}'.format(
FLAGS.aws_dynamodb_sortkey,
FLAGS.aws_dynamodb_attributetype))
self.table_name = table_name
self.rcu = FLAGS.aws_dynamodb_read_capacity
self.wcu = FLAGS.aws_dynamodb_write_capacity
self.throughput = 'ReadCapacityUnits={read},WriteCapacityUnits={write}'.format(
read=self.rcu, write=self.wcu)
self.lsi_indexes = _GetIndexes().CreateLocalSecondaryIndex()
self.gsi_indexes = _GetIndexes().CreateGlobalSecondaryIndex()
self.resource_arn: str = None # Set during the _Exists() call.
def __init__(self, edw_service_spec):
super(Redshift, self).__init__(edw_service_spec)
# pkb setup attribute
self.project = None
self.cmd_prefix = list(util.AWS_PREFIX)
if FLAGS.zones:
self.zone = FLAGS.zones[0]
self.region = util.GetRegionFromZone(self.zone)
else:
self.region = GetDefaultRegion()
self.cmd_prefix += ['--region', self.region]
# Redshift specific attribute (see if they can be set)
self.cluster_subnet_group = None
self.cluster_parameter_group = None
self.arn = ''
self.cluster_subnet_group = aws_cluster_subnet_group.RedshiftClusterSubnetGroup(
self.cmd_prefix)
self.cluster_parameter_group = aws_cluster_parameter_group.RedshiftClusterParameterGroup(
self.cmd_prefix)
if self.db is None:
self.db = DEFAULT_DATABASE_NAME
self.client_interface = GetRedshiftClientInterface(
self.db, self.user, self.password)
def __init__(self,
table_name: Optional[str] = None,
zone: Optional[str] = None,
rcu: Optional[int] = None,
wcu: Optional[int] = None,
primary_key: Optional[str] = None,
sort_key: Optional[str] = None,
attribute_type: Optional[str] = None,
lsi_count: Optional[int] = None,
gsi_count: Optional[int] = None,
use_sort: Optional[bool] = None,
**kwargs):
super(AwsDynamoDBInstance, self).__init__(**kwargs)
self.table_name = table_name or f'pkb-{FLAGS.run_uri}'
self.zone = zone or _DEFAULT_ZONE
self.region = util.GetRegionFromZone(self.zone)
self.resource_arn: str = None # Set during the _Exists() call.
self.rcu = rcu or _FREE_TIER_RCU
self.wcu = wcu or _FREE_TIER_WCU
self.throughput = (
f'ReadCapacityUnits={self.rcu},WriteCapacityUnits={self.wcu}')
self.primary_key = primary_key or 'primary_key'
self.sort_key = sort_key or 'sort_key'
self.use_sort = use_sort or False
self.attribute_type = attribute_type or 'S'
self.lsi_count = lsi_count or 0
self.lsi_indexes = self._CreateLocalSecondaryIndex()
self.gsi_count = gsi_count or 0
self.gsi_indexes = self._CreateGlobalSecondaryIndex()
def __init__(self, machine_type, zone): super(AwsDedicatedHost, self).__init__() self.machine_type = machine_type self.zone = zone self.region = util.GetRegionFromZone(self.zone) self.client_token = str(uuid.uuid4()) self.id = None
def __init__(self, dpb_service_spec):
super(UnmanagedDpbService, self).__init__(dpb_service_spec)
# Dictionary to hold the cluster vms.
self.vms = {}
self.cloud = dpb_service_spec.worker_group.cloud
if not self.dpb_service_zone:
raise errors.Setup.InvalidSetupError(
'dpb_service_zone must be provided, for provisioning.')
if self.cloud == 'GCP':
self.region = gcp_util.GetRegionFromZone(FLAGS.dpb_service_zone)
self.storage_service = gcs.GoogleCloudStorageService()
self.persistent_fs_prefix = 'gs://'
elif self.cloud == 'AWS':
self.region = aws_util.GetRegionFromZone(FLAGS.dpb_service_zone)
self.storage_service = s3.S3Service()
self.persistent_fs_prefix = 's3://'
else:
self.region = None
self.storage_service = None
self.persistent_fs_prefix = None
self.manage_bucket = False
logging.warning(
'Cloud provider %s does not support object storage. '
'Some benchmarks will not work.', self.cloud)
if self.storage_service:
self.storage_service.PrepareService(location=self.region)
# set in _Create of derived classes
self.leader = None
def __init__(self, dpb_service_spec):
super(UnmanagedDpbService, self).__init__(dpb_service_spec)
# Dictionary to hold the cluster vms.
self.vms = {}
self.cloud = dpb_service_spec.worker_group.cloud
if not self.dpb_service_zone:
raise errors.Setup.InvalidSetupError(
'dpb_service_zone must be provided, for provisioning.')
if self.cloud == 'GCP':
self.region = gcp_util.GetRegionFromZone(FLAGS.dpb_service_zone)
self.storage_service = gcs.GoogleCloudStorageService()
self.persistent_fs_prefix = 'gs://'
elif self.cloud == 'AWS':
self.region = aws_util.GetRegionFromZone(FLAGS.dpb_service_zone)
self.storage_service = s3.S3Service()
self.persistent_fs_prefix = 's3://'
else:
raise errors.Config.InvalidValue(
f'Unsupported Cloud provider {self.cloud}')
if self.storage_service:
self.storage_service.PrepareService(location=self.region)
# set in _Create of derived classes
self.leader = None
def __init__(self, spec):
super(ElastiCacheMemcached, self).__init__(spec)
self.subnet_group_name = 'subnet-%s' % self.name
self.zone = self.spec.vms[0].zone
self.region = util.GetRegionFromZone(self.zone)
self.node_type = FLAGS.cache_node_type
self.version = FLAGS.managed_memory_store_version
def _BuildContext(launcher_vm, booter_template_vm):
"""Returns the context variables for Jinja2 template during rendering."""
context = {
'boot_machine_type': booter_template_vm.machine_type,
'cloud': FLAGS.cloud,
'contact_launcher': FLAGS.vms_contact_launcher,
'launcher_vm_name': launcher_vm.name,
'os_type': 'linux' if _IsLinux() else 'windows',
'server_ip': launcher_vm.internal_ip,
'server_port': _PORT,
'start_time_file': _START_TIME_FILE_PATH,
'timeout': _TIMEOUT_SECONDS,
'vm_count': FLAGS.boots_per_launcher,
'zone': launcher_vm.zone,
}
cloud = FLAGS.cloud
if cloud == 'GCP':
context.update({
'boot_disk_size':
booter_template_vm.boot_disk_size,
'boot_vm_name_prefix':
_BOOT_VM_NAME_PREFIX.format(launcher_name=launcher_vm.name),
'image_family':
booter_template_vm.image_family,
'image_project':
booter_template_vm.image_project,
'gcloud_path':
FLAGS.gcloud_path,
'project':
FLAGS.project,
})
elif cloud == 'AWS':
context.update({
'group_name': booter_template_vm.placement_group.name,
'image': booter_template_vm.image,
'key_name': 'perfkit-key-{0}'.format(FLAGS.run_uri),
'region': aws_util.GetRegionFromZone(launcher_vm.zone),
'subnet_id': booter_template_vm.network.subnet.id,
})
elif cloud == 'Azure':
context.update({
'boot_vm_name_prefix':
launcher_vm.name.split('-', 1)[1],
'location':
launcher_vm.location,
'image':
booter_template_vm.image,
'storage_sku':
booter_template_vm.os_disk.disk_type,
'resource_group':
launcher_vm.resource_group.name,
'nic':
_BOOT_NIC_NAME_PREFIX.format(run_uri=FLAGS.run_uri),
'password':
booter_template_vm.password,
'start_id':
GetAzBootVMStartIdByLauncher(launcher_vm.name),
})
return context
def __init__(self, zone, vpc_id, cidr_block='10.0.0.0/24', subnet_id=None): super(AwsSubnet, self).__init__(subnet_id is not None) self.zone = zone self.region = util.GetRegionFromZone(zone) self.vpc_id = vpc_id self.id = subnet_id self.cidr_block = cidr_block
def __init__(self, spec):
super(AwsKopsCluster, self).__init__(spec)
self.name += '.k8s.local'
self.config_bucket = 'kops-%s-%s' % (FLAGS.run_uri, str(uuid.uuid4()))
self.region = util.GetRegionFromZone(self.zone)
self.s3_service = s3.S3Service()
self.s3_service.PrepareService(self.region)
def __init__(self, cluster_spec):
super(FargateCluster, self).__init__(cluster_spec)
self.region = util.GetRegionFromZone(self.zone)
self.network = aws_network.AwsNetwork.GetNetwork(self)
self.firewall = aws_network.AwsFirewall.GetFirewall()
self.name = 'pkb-%s' % FLAGS.run_uri
self.task_defs = {}
self.arn = None
def __init__(self, disk_spec, zone): super(AwsNfsService, self).__init__(disk_spec, zone) self.region = util.GetRegionFromZone(self.zone) self.aws_commands = AwsEfsCommands(self.region) self.disk_spec.disk_size = 0 self.filer_id = None self.mount_id = None self.throughput_mode = FLAGS.efs_throughput_mode self.provisioned_throughput = FLAGS.efs_provisioned_throughput
def __init__(self, table_name):
super(AwsDynamoDBInstance, self).__init__()
self.zone = FLAGS.zones[0]
self.region = util.GetRegionFromZone(self.zone)
self.primary_key = 'AttributeName={0},KeyType=HASH'.format(FLAGS.aws_dynamodb_primarykey)
self.attributes = 'AttributeName={0},AttributeType={1}'\
.format(FLAGS.aws_dynamodb_primarykey, FLAGS.aws_dynamodb_attributetype)
self.table_name = table_name
self.throughput = FLAGS.aws_dynamodb_capacity
def __init__(self, vm_spec):
"""Initialize a AWS virtual machine.
Args:
vm_spec: virtual_machine.BaseVirtualMachineSpec object of the vm.
Raises:
ValueError: If an incompatible vm_spec is passed.
"""
super(AwsVirtualMachine, self).__init__(vm_spec)
self.region = util.GetRegionFromZone(self.zone)
self.user_name = FLAGS.aws_user_name or self.DEFAULT_USER_NAME
if self.machine_type in aws_disk.NUM_LOCAL_VOLUMES:
self.max_local_disks = aws_disk.NUM_LOCAL_VOLUMES[self.machine_type]
self.user_data = None
self.network = aws_network.AwsNetwork.GetNetwork(self)
self.placement_group = getattr(vm_spec, 'placement_group',
self.network.placement_group)
self.firewall = aws_network.AwsFirewall.GetFirewall()
self.use_dedicated_host = vm_spec.use_dedicated_host
self.num_vms_per_host = vm_spec.num_vms_per_host
self.use_spot_instance = vm_spec.use_spot_instance
self.spot_price = vm_spec.spot_price
self.spot_block_duration_minutes = vm_spec.spot_block_duration_minutes
self.boot_disk_size = vm_spec.boot_disk_size
self.client_token = str(uuid.uuid4())
self.host = None
self.id = None
self.metadata.update({
'spot_instance':
self.use_spot_instance,
'spot_price':
self.spot_price,
'spot_block_duration_minutes':
self.spot_block_duration_minutes,
'placement_group_strategy':
self.placement_group.strategy
if self.placement_group else placement_group.PLACEMENT_GROUP_NONE,
'aws_credit_specification':
FLAGS.aws_credit_specification
if FLAGS.aws_credit_specification else 'none'
})
self.spot_early_termination = False
self.spot_status_code = None
# See:
# https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/enhanced-networking-os.html
self._smp_affinity_script = 'smp_affinity.sh'
if self.use_dedicated_host and util.IsRegion(self.zone):
raise ValueError(
'In order to use dedicated hosts, you must specify an availability '
'zone, not a region ("zone" was %s).' % self.zone)
if self.use_dedicated_host and self.use_spot_instance:
raise ValueError(
'Tenancy=host is not supported for Spot Instances')
def __init__(self, spec):
"""Initializes AwsNetwork instances.
Args:
spec: A BaseNetworkSpec object.
"""
super(AwsNetwork, self).__init__(spec)
self.region = util.GetRegionFromZone(spec.zone)
self.regional_network = _AwsRegionalNetwork.GetForRegion(self.region)
self.subnet = None
self.placement_group = AwsPlacementGroup(self.region)
def testDeleteBucket(self):
local_emr = LocalAwsDpbEmr()
with mock.patch(vm_util.__name__ + '.IssueCommand',
return_value=('out_', 'err_', 0)) as mock_issue:
local_emr.DeleteBucket('foo')
self.assertEqual(mock_issue.call_count, 1)
call_arg_list, _ = mock_issue.call_args
self.assertListEqual([
'aws', 's3', 'rb', 's3://{}'.format('foo'), '--region',
util.GetRegionFromZone(FLAGS.dpb_service_zone), '--force'
], call_arg_list[0])
def __init__(self, edw_service_spec):
super(Athena, self).__init__(edw_service_spec)
self.output_location = FLAGS.athena_output_location
self.region = util.GetRegionFromZone(FLAGS.zones[0])
self.client_interface = GetAthenaClientInterface(self.cluster_identifier)
if FLAGS.provision_athena:
self.data_bucket = 'pkb' + self.cluster_identifier.replace('_', '')
self.tables = (
TPC_H_TABLES if FLAGS.edw_tpc_dsb_type == 'tpc_h' else TPC_DS_TABLES)
self.athena_db_create_time = 0
self.athena_table_create_time = 0
def __init__(self, aws_placement_group_spec):
"""Init method for AwsPlacementGroup.
Args:
aws_placement_group_spec: Object containing the
information needed to create an AwsPlacementGroup.
"""
super(AwsPlacementGroup, self).__init__(aws_placement_group_spec)
self.name = ('perfkit-%s-%s' %
(FLAGS.run_uri, str(uuid.uuid4())[-12:]))
self.region = util.GetRegionFromZone(self.zone)
self.strategy = aws_placement_group_spec.aws_placement_group_style
def __init__(self, edw_service_spec):
super(Athena, self).__init__(edw_service_spec)
self.output_location = FLAGS.athena_output_location
self.region = util.GetRegionFromZone(FLAGS.zones[0])
if FLAGS.provision_athena:
self.db = '_'.join([
FLAGS.edw_tpc_dsb_type,
str(FLAGS.edw_tpc_dataset_size_in_GB)
])
self.data_bucket = self.db.replace('_', '')
self.athena_db_create_time = 0
self.athena_table_create_time = 0
def __init__(self, vm_spec):
"""Initialize a AWS virtual machine.
Args:
vm_spec: virtual_machine.BaseVirtualMachineSpec object of the vm.
Raises:
ValueError: If an incompatible vm_spec is passed.
"""
super(AwsVirtualMachine, self).__init__(vm_spec)
self.region = util.GetRegionFromZone(self.zone)
self.user_name = FLAGS.aws_user_name
if self.machine_type in NUM_LOCAL_VOLUMES:
self.max_local_disks = NUM_LOCAL_VOLUMES[self.machine_type]
self.user_data = None
self.network = aws_network.AwsNetwork.GetNetwork(self)
self.placement_group = getattr(vm_spec, 'placement_group',
self.network.placement_group)
self.firewall = aws_network.AwsFirewall.GetFirewall()
self.use_dedicated_host = vm_spec.use_dedicated_host
# TODO(buggay): implement num_vms_per_host functionality
self.num_vms_per_host = vm_spec.num_vms_per_host
if self.num_vms_per_host:
raise NotImplementedError(
'Num vms per host for AWS is not supported.')
self.use_spot_instance = vm_spec.use_spot_instance
self.spot_price = vm_spec.spot_price
self.boot_disk_size = vm_spec.boot_disk_size
self.client_token = str(uuid.uuid4())
self.host = None
self.id = None
self.metadata.update({
'spot_instance':
self.use_spot_instance,
'spot_price':
self.spot_price,
'placement_group_strategy':
self.placement_group.strategy if self.placement_group else 'none',
'aws_credit_specification':
FLAGS.aws_credit_specification
if FLAGS.aws_credit_specification else 'none'
})
self.early_termination = False
self.spot_status_code = None
if self.use_dedicated_host and util.IsRegion(self.zone):
raise ValueError(
'In order to use dedicated hosts, you must specify an availability '
'zone, not a region ("zone" was %s).' % self.zone)
if self.use_dedicated_host and self.use_spot_instance:
raise ValueError(
'Tenancy=host is not supported for Spot Instances')
def __init__(self, vm_group):
if not vm_group:
raise InvalidVmGroupSizeError(
'AwsCapacityReservation must be initialized with at least one '
'VM in the vm_group.')
super(AwsCapacityReservation, self).__init__(vm_group)
self.zone_or_region = vm_group[0].zone
self.region = util.GetRegionFromZone(self.zone_or_region)
self.machine_type = vm_group[0].machine_type
self.os_type = vm_group[0].OS_TYPE
self.vm_count = len(vm_group)
def __init__(self, dpb_service_spec):
super(AwsDpbEmr, self).__init__(dpb_service_spec)
self.project = None
self.cmd_prefix = list(util.AWS_PREFIX)
if FLAGS.zones:
self.zone = FLAGS.zones[0]
region = util.GetRegionFromZone(self.zone)
self.cmd_prefix += ['--region', region]
self.network = aws_network.AwsNetwork.GetNetwork(self)
self.bucket_to_delete = None
self.emr_release_label = FLAGS.dpb_emr_release_label
def _AssertClientAndDbInSameRegion(self):
"""Asserts that the client vm is in the same region requested by the server.
Raises:
AwsManagedRelationalDbCrossRegionException if the client vm is in a
different region that is requested by the server.
"""
client_region = self.client_vm.region
db_region = util.GetRegionFromZone(self.zone)
if client_region != db_region:
raise AwsManagedRelationalDbCrossRegionException(
('client_vm and managed_relational_db server '
'must be in the same region'))
def __init__(self, spec):
super(EksCluster, self).__init__(spec)
self.zones = sorted(FLAGS.eks_zones)
self.zone = ','.join(self.zones)
self.region = util.GetRegionFromZone(self.zones[0])
self.networks = [
aws_network.AwsNetwork.GetNetworkFromNetworkSpec(
aws_network.AwsNetworkSpec(zone)) for zone in self.zones
]
self.eks_control_plane = _EksControlPlane(self.name, self.region,
self.networks)
self.eks_workers = _EksWorkers('%s-worker-nodes' % self.name,
self.region)
def __init__(self, disk_spec, zone, machine_type):
super(AwsDisk, self).__init__(disk_spec)
self.iops = disk_spec.iops
self.id = None
self.zone = zone
self.region = util.GetRegionFromZone(zone)
self.device_letter = None
self.attached_vm_id = None
if self.disk_type != disk.LOCAL:
self.metadata = DISK_METADATA[self.disk_type]
else:
self.metadata = (LOCAL_HDD_METADATA if LocalDiskIsHDD(machine_type)
else LOCAL_SSD_METADATA)
