def __init__(self, name: Optional[str], project: Optional[str],
zone: Optional[str], node_count: Optional[int],
storage_type: Optional[str],
replication_cluster: Optional[bool],
replication_cluster_zone: Optional[str],
multicluster_routing: Optional[bool],
autoscaling_min_nodes: Optional[int],
autoscaling_max_nodes: Optional[int],
autoscaling_cpu_target: Optional[int]):
super(GcpBigtableInstance, self).__init__()
if name is not None:
self.user_managed = True
self.name: str = name or f'pkb-bigtable-{FLAGS.run_uri}'
self.zone: str = zone or FLAGS.google_bigtable_zone
self.project: str = project or FLAGS.project or util.GetDefaultProject(
)
self.node_count: int = node_count or _DEFAULT_NODE_COUNT
self.storage_type: str = storage_type or _DEFAULT_STORAGE_TYPE
self.replication_cluster: bool = replication_cluster or False
self.replication_cluster_zone: str = (replication_cluster_zone
or _DEFAULT_REPLICATION_ZONE)
self.multicluster_routing: bool = multicluster_routing or False
self.autoscaling_min_nodes: Optional[
int] = autoscaling_min_nodes or None
self.autoscaling_max_nodes: Optional[
int] = autoscaling_max_nodes or None
self.autoscaling_cpu_target: Optional[
int] = autoscaling_cpu_target or None
def __init__(self, vm_spec):
"""Initialize a GCE virtual machine.
Args:
vm_spec: virtual_machine.BaseVirtualMachineSpec object of the vm.
Raises:
errors.Config.MissingOption: If the spec does not include a "machine_type"
or both "cpus" and "memory".
errors.Config.InvalidValue: If the spec contains both "machine_type" and
at least one of "cpus" or "memory".
"""
super(GceVirtualMachine, self).__init__(vm_spec)
self.boot_metadata = {}
self.cpus = vm_spec.cpus
self.image = self.image or self.DEFAULT_IMAGE
self.max_local_disks = vm_spec.num_local_ssds
self.memory_mib = vm_spec.memory
self.preemptible = vm_spec.preemptible
self.project = vm_spec.project or util.GetDefaultProject()
self.image_project = vm_spec.image_project
self.network = gce_network.GceNetwork.GetNetwork(self)
self.firewall = gce_network.GceFirewall.GetFirewall()
self.boot_disk_size = vm_spec.boot_disk_size
self.boot_disk_type = vm_spec.boot_disk_type
self.id = None
def Run(benchmark_spec):
"""Spawn YCSB and gather the results.
Args:
benchmark_spec: The benchmark specification. Contains all data that is
required to run the benchmark.
Returns:
A list of sample.Sample instances.
"""
vms = benchmark_spec.vms
run_kwargs = {
'table': BENCHMARK_TABLE,
'zeropadding': BENCHMARK_ZERO_PADDING,
'cloudspanner.instance': BENCHMARK_INSTANCE_PREFIX + FLAGS.run_uri,
'cloudspanner.database': BENCHMARK_DATABASE,
'cloudspanner.readmode': FLAGS.cloud_spanner_ycsb_readmode,
'cloudspanner.boundedstaleness':
FLAGS.cloud_spanner_ycsb_boundedstaleness,
'cloudspanner.batchinserts': FLAGS.cloud_spanner_ycsb_batchinserts,
}
if FLAGS.cloud_spanner_ycsb_client_type == 'go':
run_kwargs['cloudspanner.project'] = util.GetDefaultProject()
load_kwargs = run_kwargs.copy()
samples = list(benchmark_spec.executor.LoadAndRun(
vms, load_kwargs=load_kwargs, run_kwargs=run_kwargs))
metadata = {'ycsb_client_type': FLAGS.cloud_spanner_ycsb_client_type}
for sample in samples:
sample.metadata.update(metadata)
return samples
def __init__(self, relational_db_spec):
super(GCPRelationalDb, self).__init__(relational_db_spec)
self.project = FLAGS.project or util.GetDefaultProject()
self.instance_id = 'pkb-db-instance-' + FLAGS.run_uri
if self.is_managed_db:
self.unmanaged_db_exists = False
def __init__(self, vm_spec):
"""Initialize a GCE virtual machine.
Args:
vm_spec: virtual_machine.BaseVmSpec object of the vm.
Raises:
errors.Config.MissingOption: If the spec does not include a "machine_type"
or both "cpus" and "memory".
errors.Config.InvalidValue: If the spec contains both "machine_type" and
at least one of "cpus" or "memory".
"""
super(GceVirtualMachine, self).__init__(vm_spec)
self.boot_metadata = {}
self.cpus = vm_spec.cpus
self.image = self.image or self.DEFAULT_IMAGE
self.max_local_disks = vm_spec.num_local_ssds
self.memory_mib = vm_spec.memory
self.preemptible = vm_spec.preemptible
self.project = vm_spec.project or util.GetDefaultProject()
self.image_family = vm_spec.image_family or self.DEFAULT_IMAGE_FAMILY
self.image_project = vm_spec.image_project or self.DEFAULT_IMAGE_PROJECT
self.backfill_image = False
self.network = gce_network.GceNetwork.GetNetwork(self)
self.firewall = gce_network.GceFirewall.GetFirewall()
self.boot_disk_size = vm_spec.boot_disk_size
self.boot_disk_type = vm_spec.boot_disk_type
self.id = None
self.node_type = vm_spec.node_type
self.node_group = None
self.use_dedicated_host = vm_spec.use_dedicated_host
self.num_vms_per_host = vm_spec.num_vms_per_host
self.min_cpu_platform = vm_spec.min_cpu_platform
self.gce_remote_access_firewall_rule = FLAGS.gce_remote_access_firewall_rule
self.gce_accelerator_type_override = FLAGS.gce_accelerator_type_override
def __init__(self, managed_relational_db_spec):
super(GCPManagedRelationalDb,
self).__init__(managed_relational_db_spec)
self.spec = managed_relational_db_spec
self.project = FLAGS.project or util.GetDefaultProject()
self.instance_id = 'pkb-db-instance-' + FLAGS.run_uri
if (self.spec.engine == managed_relational_db.POSTGRES
and self.spec.high_availability):
raise Exception(
('High availability configuration is not supported on '
'CloudSQL PostgreSQL'))
def Run(benchmark_spec):
"""Spawn YCSB and gather the results.
Args:
benchmark_spec: The benchmark specification. Contains all data that is
required to run the benchmark.
Returns:
A list of sample.Sample instances.
"""
vms = benchmark_spec.vms
spanner: gcp_spanner.GcpSpannerInstance = benchmark_spec.spanner
run_kwargs = {
'table': BENCHMARK_TABLE,
'zeropadding': BENCHMARK_ZERO_PADDING,
'cloudspanner.instance': benchmark_spec.spanner.name,
'cloudspanner.database': benchmark_spec.spanner.database,
'cloudspanner.readmode': FLAGS.cloud_spanner_ycsb_readmode,
'cloudspanner.boundedstaleness':
FLAGS.cloud_spanner_ycsb_boundedstaleness,
'cloudspanner.batchinserts': FLAGS.cloud_spanner_ycsb_batchinserts,
}
# Uses overridden cloud spanner endpoint in gcloud configuration
end_point = benchmark_spec.spanner.GetEndPoint()
if end_point:
run_kwargs['cloudspanner.host'] = end_point
if FLAGS.cloud_spanner_ycsb_client_type == 'go':
run_kwargs['cloudspanner.project'] = util.GetDefaultProject()
load_kwargs = run_kwargs.copy()
samples = []
metadata = {'ycsb_client_type': FLAGS.cloud_spanner_ycsb_client_type}
if not benchmark_spec.spanner.restored:
samples += list(
benchmark_spec.executor.Load(vms, load_kwargs=load_kwargs))
if _CPU_OPTIMIZATION.value:
samples += CpuUtilizationRun(benchmark_spec.executor, spanner, vms,
run_kwargs)
metadata.update(_GetCpuOptimizationMetadata())
else:
samples += list(benchmark_spec.executor.Run(vms,
run_kwargs=run_kwargs))
for result in samples:
result.metadata.update(metadata)
result.metadata.update(spanner.GetResourceMetadata())
return samples
def _RunBenchmark(vm, streams=1, batch_size=2048, data_format='AVRO'):
"""Run a benchmark python script on a VM for given arguments.
Args:
vm: GCE VM instance where to run the benchmark.
streams: Number of BigQuery streams to use.
batch_size: batch size to use.
data_format: serialization data format to use (AVRO or ARROW).
Returns:
Benchmark result.
"""
project_id = FLAGS.project or gcp_util.GetDefaultProject()
options = (f' --project_id="{project_id}"'
f' --dataset_project_id="{FLAGS.dataset_project_id}"'
f' --dataset_id="{FLAGS.dataset_id}"'
f' --table_id="{FLAGS.table_id}"'
f' --requested_streams="{streams}"'
f' --batch_size="{batch_size}"'
f' --format="{data_format}"')
cmd = f'{DLVM_PYTHON} {REMOTE_SCRIPT} {options}'
logging.info(cmd)
stdout, stderr = vm.RemoteCommand(cmd)
logging.info(stdout)
logging.info(stderr)
result = ''
for line in stdout.split('\n'):
m = re.match(RESULT_MATCH_REGEX, line)
if m:
result = m[1]
break
if not result:
return []
metadata = {
'bq_requested_streams': streams,
'bq_batch_size': batch_size,
'bq_data_format': data_format,
'dataset_project_id': FLAGS.dataset_project_id,
'dataset_id': FLAGS.dataset_id,
'table_id': FLAGS.table_id,
}
return [
sample.Sample(
'BigQuery TensorFlow connector read throughput',
result,
UNIT,
metadata)
]
def __init__(self, name, description, database, ddl):
super(GcpSpannerInstance, self).__init__()
self._name = name
self._description = description
self._database = database
self._ddl = ddl
self._config = FLAGS.cloud_spanner_config
self._nodes = FLAGS.cloud_spanner_nodes
self._end_point = None
# Cloud Spanner may not explicitly set the following common flags.
self.project = FLAGS.cloud_spanner_project or util.GetDefaultProject()
self.zone = None
def __init__(self, vm_spec):
"""Initialize a GCE virtual machine.
Args:
vm_spec: virtual_machine.BaseVmSpec object of the vm.
Raises:
errors.Config.MissingOption: If the spec does not include a "machine_type"
or both "cpus" and "memory".
errors.Config.InvalidValue: If the spec contains both "machine_type" and
at least one of "cpus" or "memory".
"""
super(GceVirtualMachine, self).__init__(vm_spec)
self.boot_metadata = {}
self.cpus = vm_spec.cpus
self.image = self.image or self.DEFAULT_IMAGE
self.max_local_disks = vm_spec.num_local_ssds
self.memory_mib = vm_spec.memory
self.preemptible = vm_spec.preemptible
self.early_termination = False
self.preemptible_status_code = None
self.project = vm_spec.project or util.GetDefaultProject()
self.image_family = vm_spec.image_family or self.DEFAULT_IMAGE_FAMILY
self.image_project = vm_spec.image_project or self.DEFAULT_IMAGE_PROJECT
self.backfill_image = False
self.network = self._GetNetwork()
self.firewall = gce_network.GceFirewall.GetFirewall()
self.boot_disk_size = vm_spec.boot_disk_size or self.BOOT_DISK_SIZE_GB
self.boot_disk_type = vm_spec.boot_disk_type or self.BOOT_DISK_TYPE
self.id = None
self.node_type = vm_spec.node_type
self.node_group = None
self.use_dedicated_host = vm_spec.use_dedicated_host
self.num_vms_per_host = vm_spec.num_vms_per_host
self.min_cpu_platform = vm_spec.min_cpu_platform
self.gce_remote_access_firewall_rule = FLAGS.gce_remote_access_firewall_rule
self.gce_accelerator_type_override = FLAGS.gce_accelerator_type_override
self.gce_tags = vm_spec.gce_tags
self.gce_network_tier = FLAGS.gce_network_tier
self.gce_shielded_secure_boot = FLAGS.gce_shielded_secure_boot
# We don't want boot time samples to be affected from retrying, so don't
# retry cluster_boot when rate limited.
if 'cluster_boot' in FLAGS.benchmarks:
FLAGS.gcp_retry_on_rate_limited = False
def __init__(self,
name: Optional[str] = None,
description: Optional[str] = None,
database: Optional[str] = None,
ddl: Optional[str] = None,
config: Optional[str] = None,
nodes: Optional[int] = None,
project: Optional[str] = None,
**kwargs):
super(GcpSpannerInstance, self).__init__(**kwargs)
self.name = name or f'pkb-instance-{FLAGS.run_uri}'
self.database = database or f'pkb-database-{FLAGS.run_uri}'
self._description = description or _DEFAULT_DESCRIPTION
self._ddl = ddl or _DEFAULT_DDL
self._config = config or self._GetDefaultConfig()
self._nodes = nodes or _DEFAULT_NODES
self._end_point = None
# Cloud Spanner may not explicitly set the following common flags.
self.project = (project or FLAGS.project or util.GetDefaultProject())
self.zone = None
def Run(benchmark_spec):
"""Spawn YCSB and gather the results.
Args:
benchmark_spec: The benchmark specification. Contains all data that is
required to run the benchmark.
Returns:
A list of sample.Sample instances.
"""
vms = benchmark_spec.vms
run_kwargs = {
'table': BENCHMARK_TABLE,
'zeropadding': BENCHMARK_ZERO_PADDING,
'cloudspanner.instance': benchmark_spec.spanner.name,
'cloudspanner.database': benchmark_spec.spanner.database,
'cloudspanner.readmode': FLAGS.cloud_spanner_ycsb_readmode,
'cloudspanner.boundedstaleness':
FLAGS.cloud_spanner_ycsb_boundedstaleness,
'cloudspanner.batchinserts': FLAGS.cloud_spanner_ycsb_batchinserts,
}
if FLAGS.cloud_spanner_ycsb_client_type == 'go':
run_kwargs['cloudspanner.project'] = util.GetDefaultProject()
load_kwargs = run_kwargs.copy()
samples = list(
benchmark_spec.executor.LoadAndRun(vms,
load_kwargs=load_kwargs,
run_kwargs=run_kwargs))
metadata = {'ycsb_client_type': FLAGS.cloud_spanner_ycsb_client_type}
for sample in samples:
# YCSB writes error samples, there is no direct output to parse.
if 'Return=ERROR' in sample.metric:
raise errors.Benchmarks.RunError(
'Error running YCSB, please check the output log.')
sample.metadata.update(metadata)
return samples
def __init__(self, tpu_spec):
super(GcpTpu, self).__init__(tpu_spec)
self.spec = tpu_spec
self.project = FLAGS.project or util.GetDefaultProject()
def __init__(self, cloud_tpu_spec):
super(GcpCloudTpu, self).__init__(cloud_tpu_spec)
self.spec = cloud_tpu_spec
self.project = FLAGS.project or util.GetDefaultProject()
def __init__(self, registry_spec):
super(GoogleContainerRegistry, self).__init__(registry_spec)
self.project = self.project or util.GetDefaultProject()
def __init__(self, managed_relational_db_spec):
super(GCPManagedRelationalDb,
self).__init__(managed_relational_db_spec)
self.spec = managed_relational_db_spec
self.project = FLAGS.project or util.GetDefaultProject()
self.instance_id = 'pkb-db-instance-' + FLAGS.run_uri
def __init__(self):
super().__init__()
self.project = FLAGS.project or util.GetDefaultProject()
self.pubsub_topic = 'pkb-topic-{0}'.format(FLAGS.run_uri)
self.pubsub_subscription = 'pkb-subscription-{0}'.format(FLAGS.run_uri)
def __init__(self, relational_db_spec):
super(GCPRelationalDb, self).__init__(relational_db_spec)
self.project = FLAGS.project or util.GetDefaultProject()
self.unmanaged_db_exists = None if self.is_managed_db else False
