def ConstructDpbService(self):
"""Create the dpb_service object and create groups for its vms."""
if self.config.dpb_service is None:
return
dpb_service_spec = self.config.dpb_service
dpb_service_cloud = dpb_service_spec.worker_group.cloud
dpb_service_spec.worker_group.vm_count = dpb_service_spec.worker_count
providers.LoadProvider(dpb_service_cloud)
dpb_service_type = dpb_service_spec.service_type
dpb_service_class = dpb_service.GetDpbServiceClass(dpb_service_type)
self.dpb_service = dpb_service_class(dpb_service_spec)
# If the dpb service is un-managed, the provisioning needs to be handed
# over to the vm creation module.
if dpb_service_type == dpb_service.UNMANAGED_DPB_SVC_YARN_CLUSTER:
# Ensure non cluster vms are not present in the spec.
if self.vms_to_boot:
raise Exception('Invalid Non cluster vm group {0} when benchmarking '
'unmanaged dpb service'.format(self.vms_to_boot))
base_vm_spec = dpb_service_spec.worker_group
base_vm_spec.vm_spec.zone = self.dpb_service.dpb_service_zone
if dpb_service_spec.worker_count:
self.vms_to_boot['worker_group'] = dpb_service_spec.worker_group
# else we have a single node cluster.
master_group_spec = copy.copy(base_vm_spec)
master_group_spec.vm_count = 1
self.vms_to_boot['master_group'] = master_group_spec
def __init__(self, component_full_name, flag_values=None, **kwargs):
super(_VmGroupSpec, self).__init__(component_full_name,
flag_values=flag_values,
**kwargs)
providers.LoadProvider(self.cloud.lower())
if self.disk_spec:
disk_config = getattr(self.disk_spec, self.cloud, None)
if disk_config is None:
raise errors.Config.MissingOption(
'{0}.cloud is "{1}", but {0}.disk_spec does not contain a '
'configuration for "{1}".'.format(component_full_name,
self.cloud))
disk_spec_class = disk.GetDiskSpecClass(self.cloud)
self.disk_spec = disk_spec_class('{0}.disk_spec.{1}'.format(
component_full_name, self.cloud),
flag_values=flag_values,
**disk_config)
vm_config = getattr(self.vm_spec, self.cloud, None)
if vm_config is None:
raise errors.Config.MissingOption(
'{0}.cloud is "{1}", but {0}.vm_spec does not contain a '
'configuration for "{1}".'.format(component_full_name,
self.cloud))
vm_spec_class = virtual_machine.GetVmSpecClass(self.cloud)
self.vm_spec = vm_spec_class('{0}.vm_spec.{1}'.format(
component_full_name, self.cloud),
flag_values=flag_values,
**vm_config)
def ConstructNfsService(self):
"""Construct the NFS service object.
Creates an NFS Service only if an NFS disk is found in the disk_specs.
"""
if self.nfs_service:
logging.info('NFS service already created: %s', self.nfs_service)
return
for group_spec in self.vms_to_boot.values():
if not group_spec.disk_spec or not group_spec.vm_count:
continue
disk_spec = group_spec.disk_spec
if disk_spec.disk_type != disk.NFS:
continue
# Choose which nfs_service to create.
if disk_spec.nfs_ip_address:
self.nfs_service = nfs_service.StaticNfsService(disk_spec)
elif disk_spec.nfs_managed:
cloud = group_spec.cloud
providers.LoadProvider(cloud)
nfs_class = nfs_service.GetNfsServiceClass(cloud)
self.nfs_service = nfs_class(disk_spec, group_spec.vm_spec.zone)
else:
self.nfs_service = nfs_service.UnmanagedNfsService(disk_spec,
self.vms[0])
logging.debug('NFS service %s', self.nfs_service)
break
def Prepare(benchmark_spec):
"""Prepare vm with cloud provider tool and prepare vm with data file.
Args:
benchmark_spec: The benchmark specification. Contains all data that is
required to run the benchmark.
"""
providers.LoadProvider(FLAGS.storage)
service = object_storage_service.GetObjectStorageClass(FLAGS.storage)()
service.PrepareService(FLAGS.object_storage_region)
vms = benchmark_spec.vms
vm_util.RunThreaded(lambda vm: PrepareVM(vm, service), vms)
# We would like to always cleanup server side states when exception happens.
benchmark_spec.always_call_cleanup = True
# Make the bucket(s)
bucket_name = 'pkb%s' % FLAGS.run_uri
if FLAGS.storage != 'GCP' or not FLAGS.object_storage_gcs_multiregion:
service.MakeBucket(bucket_name)
else:
# Use a GCS multiregional bucket
multiregional_service = gcs.GoogleCloudStorageService()
multiregional_service.PrepareService(
FLAGS.object_storage_gcs_multiregion or DEFAULT_GCS_MULTIREGION)
multiregional_service.MakeBucket(bucket_name)
# Save the service and the buckets for later
benchmark_spec.service = service
benchmark_spec.buckets = [bucket_name]
def testPingSupported(self):
for cloud in benchmark_spec.VALID_CLOUDS:
providers.LoadProvider(cloud.lower())
ThisProviderInfoClass = provider_info.GetProviderInfoClass(cloud)
self.assertTrue(
ThisProviderInfoClass.IsBenchmarkSupported('iperf'),
'provider {0} does not support iperf'.format(cloud))
def ConstructSparkService(self):
"""Create the spark_service object and create groups for its vms."""
if self.config.spark_service is None:
return
spark_spec = self.config.spark_service
# Worker group is required, master group is optional
cloud = spark_spec.worker_group.cloud
if spark_spec.master_group:
cloud = spark_spec.master_group.cloud
providers.LoadProvider(cloud)
service_type = spark_spec.service_type
spark_service_class = spark_service.GetSparkServiceClass(
cloud, service_type)
self.spark_service = spark_service_class(spark_spec)
# If this is Pkb managed, the benchmark spec needs to adopt vms.
if service_type == spark_service.PKB_MANAGED:
for name, group_spec in [('master_group', spark_spec.master_group),
('worker_group', spark_spec.worker_group)
]:
if name in self.vms_to_boot:
raise Exception(
'Cannot have a vm group {0} with a {1} spark '
'service'.format(name, spark_service.PKB_MANAGED))
self.vms_to_boot[name] = group_spec
def _VerifyProviderBenchmarkSupport(self, cloud, benchmark, support_expected):
providers.LoadProvider(cloud)
provider_info_class = provider_info.GetProviderInfoClass(cloud)
supported = provider_info_class.IsBenchmarkSupported(benchmark)
fmt_args = ('', ' not') if support_expected else (' not', '')
self.assertEqual(supported, support_expected, (
'Expected provider {provider} {0}to support benchmark {benchmark}, but '
'it did{1}.'.format(*fmt_args, provider=cloud, benchmark=benchmark)))
def ConstructCloudTpu(self):
"""Constructs the BenchmarkSpec's cloud TPU objects."""
if self.config.cloud_tpu is None:
return
cloud = self.config.cloud_tpu.cloud
providers.LoadProvider(cloud)
cloud_tpu_class = cloud_tpu.GetCloudTpuClass(cloud)
self.cloud_tpu = cloud_tpu_class(self.config.cloud_tpu)
def ConstructRelationalDb(self):
"""Create the relational db and create groups for its vms."""
if self.config.relational_db is None:
return
cloud = self.config.relational_db.cloud
providers.LoadProvider(cloud)
relational_db_class = (relational_db.GetRelationalDbClass(cloud))
self.relational_db = relational_db_class(self.config.relational_db)
def ConstructDpbService(self):
"""Create the dpb_service object and create groups for its vms."""
if self.config.dpb_service is None:
return
providers.LoadProvider(self.config.dpb_service.worker_group.cloud)
dpb_service_class = dpb_service.GetDpbServiceClass(
self.config.dpb_service.service_type)
self.dpb_service = dpb_service_class(self.config.dpb_service)
def ConstructTpuGroup(self, group_spec):
"""Constructs the BenchmarkSpec's cloud TPU objects."""
if group_spec is None:
return
cloud = group_spec.cloud
providers.LoadProvider(cloud)
tpu_class = cloud_tpu.GetTpuClass(cloud)
return tpu_class(group_spec)
def ConstructCloudRedis(self):
"""Create the cloud_redis object."""
if self.config.cloud_redis is None:
return
cloud = self.config.cloud_redis.cloud
providers.LoadProvider(cloud)
cloud_redis_class = cloud_redis.GetCloudRedisClass(cloud)
self.cloud_redis = cloud_redis_class(self.config.cloud_redis)
def ConstructContainerCluster(self):
"""Create the container cluster."""
if self.config.container_cluster is None:
return
cloud = self.config.container_cluster.cloud
providers.LoadProvider(cloud)
container_cluster_class = container_service.GetContainerClusterClass(cloud)
self.container_cluster = container_cluster_class(
self.config.container_cluster)
def _GetService() -> object_storage_service.ObjectStorageService:
"""Get a ready to use instance of ObjectStorageService."""
# TODO(pclay): consider using FLAGS.storage to allow cross cloud testing?
cloud = FLAGS.cloud
providers.LoadProvider(cloud)
service = object_storage_service.GetObjectStorageClass(cloud)()
# This method is idempotent with default args and safe to call in each phase.
service.PrepareService(FLAGS.object_storage_region)
return service
def ConstructContainerRegistry(self):
"""Create the container registry."""
if self.config.container_registry is None:
return
cloud = self.config.container_registry.cloud
providers.LoadProvider(cloud)
container_registry_class = container_service.GetContainerRegistryClass(
cloud)
self.container_registry = container_registry_class(
self.config.container_registry)
def ConstructSparkService(self):
if self.config.spark_service is None:
return
providers.LoadProvider(self.config.spark_service.cloud)
spark_service_spec = self.config.spark_service
service_type = spark_service_spec.service_type
spark_service_class = spark_service.GetSparkServiceClass(
spark_service_spec.cloud, service_type)
self.spark_service = spark_service_class(spark_service_spec)
def ConstructManagedRelationalDb(self):
"""Create the managed relational db and create groups for its vms."""
if self.config.managed_relational_db is None:
return
cloud = self.config.managed_relational_db.cloud
providers.LoadProvider(cloud)
managed_relational_db_class = (
managed_relational_db.GetManagedRelationalDbClass(cloud))
self.managed_relational_db = managed_relational_db_class(
self.config.managed_relational_db)
def ConstructCapacityReservations(self):
"""Construct capacity reservations for each VM group."""
if not FLAGS.use_capacity_reservations:
return
for vm_group in six.itervalues(self.vm_groups):
cloud = vm_group[0].CLOUD
providers.LoadProvider(cloud)
capacity_reservation_class = capacity_reservation.GetResourceClass(
cloud)
self.capacity_reservations.append(
capacity_reservation_class(vm_group))
def ConstructDataDiscoveryService(self):
"""Create the data_discovery_service object."""
if not self.config.data_discovery_service:
return
cloud = self.config.data_discovery_service.cloud
service_type = self.config.data_discovery_service.service_type
providers.LoadProvider(cloud)
data_discovery_service_class = (
data_discovery_service.GetDataDiscoveryServiceClass(cloud, service_type)
)
self.data_discovery_service = data_discovery_service_class.FromSpec(
self.config.data_discovery_service)
def __init__(self, component_full_name, flag_values=None, **kwargs):
super(_ManagedRelationalDbSpec, self).__init__(component_full_name,
flag_values=flag_values,
**kwargs)
# TODO(ferneyhough): This is a lot of boilerplate, and is repeated
# below in VmGroupSpec. See if some can be consolidated. Maybe we can
# specify a VmGroupSpec instead of both vm_spec and disk_spec.
ignore_package_requirements = (getattr(
flag_values, 'ignore_package_requirements', True)
if flag_values else True)
providers.LoadProvider(self.cloud, ignore_package_requirements)
if self.disk_spec:
disk_config = getattr(self.disk_spec, self.cloud, None)
if disk_config is None:
raise errors.Config.MissingOption(
'{0}.cloud is "{1}", but {0}.disk_spec does not contain a '
'configuration for "{1}".'.format(component_full_name,
self.cloud))
disk_spec_class = disk.GetDiskSpecClass(self.cloud)
self.disk_spec = disk_spec_class('{0}.disk_spec.{1}'.format(
component_full_name, self.cloud),
flag_values=flag_values,
**disk_config)
vm_config = getattr(self.vm_spec, self.cloud, None)
if vm_config is None:
raise errors.Config.MissingOption(
'{0}.cloud is "{1}", but {0}.vm_spec does not contain a '
'configuration for "{1}".'.format(component_full_name,
self.cloud))
vm_spec_class = virtual_machine.GetVmSpecClass(self.cloud)
self.vm_spec = vm_spec_class('{0}.vm_spec.{1}'.format(
component_full_name, self.cloud),
flag_values=flag_values,
**vm_config)
# Set defaults that were not able to be set in
# GetOptionDecoderConstructions()
if not self.engine_version:
managed_db_class = managed_relational_db.GetManagedRelationalDbClass(
self.cloud)
self.engine_version = managed_db_class.GetDefaultEngineVersion(
self.engine)
if not self.database_name:
self.database_name = 'pkb-db-%s' % flag_values.run_uri
if not self.database_username:
self.database_username = '******' % flag_values.run_uri
if not self.database_password:
self.database_password = managed_relational_db.GenerateRandomDbPassword(
)
def testMYSQLSupport(self):
for cloud in benchmark_spec.VALID_CLOUDS:
providers.LoadProvider(cloud.lower())
ThisProviderInfoClass = provider_info.GetProviderInfoClass(cloud)
if (cloud == benchmark_spec.AWS or cloud == benchmark_spec.GCP):
self.assertTrue(
ThisProviderInfoClass.IsBenchmarkSupported(
'mysql_service'))
else:
self.assertFalse(
ThisProviderInfoClass.IsBenchmarkSupported(
'mysql_service'),
'Cloud {0} is not supposed to support mysql_service {1}'.
format(cloud, ThisProviderInfoClass.__name__))
def ConstructMessagingService(self):
"""Create the messaging_service object.
Assumes VMs are already constructed.
"""
if self.config.messaging_service is None:
return
cloud = self.config.messaging_service.cloud
delivery = self.config.messaging_service.delivery
providers.LoadProvider(cloud)
messaging_service_class = messaging_service.GetMessagingServiceClass(
cloud, delivery)
self.messaging_service = messaging_service_class.FromSpec(
self.config.messaging_service)
self.messaging_service.setVms(self.vm_groups)
def ConstructEdwService(self):
"""Create the edw_service object."""
if self.config.edw_service is None:
return
# Load necessary modules from the provider to account for dependencies
providers.LoadProvider(
edw_service.TYPE_2_PROVIDER.get(self.config.edw_service.type))
# Load the module for the edw service based on type
edw_service_module = importlib.import_module(
edw_service.TYPE_2_MODULE.get(self.config.edw_service.type))
edw_service_class = getattr(
edw_service_module, self.config.edw_service.type[0].upper() +
self.config.edw_service.type[1:])
# Check if a new instance needs to be created or restored from snapshot
self.edw_service = edw_service_class(self.config.edw_service)
def ConstructSparkService(self):
if self.config.spark_service is None:
return
providers.LoadProvider(self.config.spark_service.cloud)
spark_service_spec = self.config.spark_service
service_type = spark_service_spec.spark_service_type
spark_service_class = spark_service.GetSparkServiceClass(
spark_service_spec.cloud, service_type)
if self.config.spark_service.static_cluster_name:
name = self.config.spark_service.static_cluster_name
static_cluster = True
else:
name = 'pkb-' + FLAGS.run_uri
static_cluster = False
self.spark_service = spark_service_class(name, static_cluster,
spark_service_spec)
def Prepare(benchmark_spec):
"""Prepare vm with cloud provider tool and prepare vm with data file.
Args:
benchmark_spec: The benchmark specification. Contains all data that is
required to run the benchmark.
"""
providers.LoadProvider(FLAGS.storage)
service = object_storage_service.GetObjectStorageClass(FLAGS.storage)()
service.PrepareService(FLAGS.object_storage_region)
vms = benchmark_spec.vms
for vm in vms:
PrepareVM(vm, service)
service.PrepareVM(vm)
# We would like to always cleanup server side states when exception happens.
benchmark_spec.always_call_cleanup = True
# Make the bucket(s)
bucket_name = 'pkb%s' % FLAGS.run_uri
if FLAGS.storage != 'GCP':
service.MakeBucket(bucket_name)
buckets = [bucket_name]
else:
# TODO(nlavine): make GCP bucket name handling match other
# providers. Leaving it inconsistent for now to match previous
# behavior, but should change it after a reasonable deprecation
# period.
multiregional_service = gcs.GoogleCloudStorageService()
multiregional_service.PrepareService(
FLAGS.object_storage_gcs_multiregion or DEFAULT_GCS_MULTIREGION)
multiregional_service.MakeBucket(bucket_name)
region = FLAGS.object_storage_region or gcs.DEFAULT_GCP_REGION
regional_bucket_name = 'pkb%s-%s' % (FLAGS.run_uri, region)
regional_service = gcs.GoogleCloudStorageService()
regional_service.PrepareService(region)
regional_service.MakeBucket(regional_bucket_name)
buckets = [bucket_name, regional_bucket_name]
# Save the service and the buckets for later
benchmark_spec.service = service
benchmark_spec.buckets = buckets
def ConstructNfsService(self):
"""Construct the NFS service object.
Creates an NFS Service only if an NFS disk is found in the disk_specs.
"""
if self.nfs_service:
logging.info('NFS service already created: %s', self.nfs_service)
return
for group_spec in self.config.vm_groups.values():
if not group_spec.disk_spec:
continue
disk_spec = group_spec.disk_spec
if disk_spec.disk_type != disk.NFS:
continue
cloud = group_spec.cloud
providers.LoadProvider(cloud)
nfs_class = nfs_service.GetNfsServiceClass(cloud)
self.nfs_service = nfs_class(disk_spec, group_spec.vm_spec.zone)
logging.info('NFS service %s', self.nfs_service)
break
def ConstructSmbService(self):
"""Construct the SMB service object.
Creates an SMB Service only if an SMB disk is found in the disk_specs.
"""
if self.smb_service:
logging.info('SMB service already created: %s', self.smb_service)
return
for group_spec in self.vms_to_boot.values():
if not group_spec.disk_spec or not group_spec.vm_count:
continue
disk_spec = group_spec.disk_spec
if disk_spec.disk_type != disk.SMB:
continue
cloud = group_spec.cloud
providers.LoadProvider(cloud)
smb_class = smb_service.GetSmbServiceClass(cloud)
self.smb_service = smb_class(disk_spec, group_spec.vm_spec.zone)
logging.debug('SMB service %s', self.smb_service)
break
def ConstructEdwService(self):
"""Create the edw_service object."""
if self.config.edw_service is None:
return
# Load necessary modules from the provider to account for dependencies
# TODO(saksena): Replace with
# providers.LoadProvider(string.lower(FLAGS.cloud))
providers.LoadProvider(
edw_service.TYPE_2_PROVIDER.get(self.config.edw_service.type))
# Load the module for the edw service based on type
edw_service_type = self.config.edw_service.type
edw_service_module = importlib.import_module(
edw_service.TYPE_2_MODULE.get(edw_service_type))
# The edw_service_type in certain cases may be qualified with a hosting
# cloud eg. snowflake_aws,snowflake_gcp, etc.
# However the edw_service_class_name in all cases will still be cloud
# agnostic eg. Snowflake.
edw_service_class_name = edw_service_type.split('_')[0]
edw_service_class = getattr(
edw_service_module,
edw_service_class_name[0].upper() + edw_service_class_name[1:])
# Check if a new instance needs to be created or restored from snapshot
self.edw_service = edw_service_class(self.config.edw_service)
def testLoadProviderChecksRequirements(self):
with mock.patch(requirements.__name__ + '.CheckProviderRequirements'):
providers.LoadProvider('GCP', ignore_package_requirements=False)
requirements.CheckProviderRequirements.assert_called_once_with(
'gcp')
def testLoadProviderIgnoresRequirements(self):
with mock.patch(requirements.__name__ + '.CheckProviderRequirements'):
providers.LoadProvider('GCP')
requirements.CheckProviderRequirements.assert_not_called()
