diff --git a/perfkitbenchmarker/linux_benchmarks/cloud_bigtable_ycsb_benchmark.py b/perfkitbenchmarker/linux_benchmarks/cloud_bigtable_ycsb_benchmark.py index 39038bc153..47b2e23f03 100644 --- a/perfkitbenchmarker/linux_benchmarks/cloud_bigtable_ycsb_benchmark.py +++ b/perfkitbenchmarker/linux_benchmarks/cloud_bigtable_ycsb_benchmark.py @@ -357,19 +357,6 @@ def _GetYcsbExecutor( return ycsb.YCSBExecutor(FLAGS.hbase_binding, **executor_flags) -def _LoadDatabase(executor: ycsb.YCSBExecutor, - bigtable: gcp_bigtable.GcpBigtableInstance, - vms: list[virtual_machine.VirtualMachine], - load_kwargs: dict[str, Any]) -> list[sample.Sample]: - """Loads the database with the specified infrastructure capacity.""" - if bigtable.restored or ycsb.SKIP_LOAD_STAGE.value: - return [] - bigtable.UpdateCapacityForLoad() - results = list(executor.Load(vms, load_kwargs=load_kwargs)) - bigtable.UpdateCapacityForRun() - return results - - def Run(benchmark_spec: bm_spec.BenchmarkSpec) -> List[sample.Sample]: """Spawn YCSB and gather the results. @@ -391,7 +378,8 @@ def Run(benchmark_spec: bm_spec.BenchmarkSpec) -> List[sample.Sample]: load_kwargs = _GenerateLoadKwargs(instance) run_kwargs = _GenerateRunKwargs(instance) executor: ycsb.YCSBExecutor = _GetYcsbExecutor(vms) - samples += _LoadDatabase(executor, instance, vms, load_kwargs) + if not instance.restored: + samples += list(executor.Load(vms, load_kwargs=load_kwargs)) samples += list(executor.Run(vms, run_kwargs=run_kwargs)) # Optionally add new samples for cluster cpu utilization. diff --git a/perfkitbenchmarker/providers/gcp/gcp_bigtable.py b/perfkitbenchmarker/providers/gcp/gcp_bigtable.py index 648def6a86..4536e3e13d 100644 --- a/perfkitbenchmarker/providers/gcp/gcp_bigtable.py +++ b/perfkitbenchmarker/providers/gcp/gcp_bigtable.py @@ -18,7 +18,6 @@ import json import logging -import time from typing import Any, Dict, List, Optional from absl import flags @@ -53,13 +52,6 @@ 'Ignored if --bigtable_autoscaling_min_nodes is set.' ), ) -_LOAD_NODES = flags.DEFINE_integer( - 'bigtable_load_node_count', - None, - 'The number of nodes for the Bigtable instance to use for the load' - ' phase. Assumes that the benchmark calls UpdateRunCapacity to set the ' - ' correct node count manually before the run phase.', -) _AUTOSCALING_MIN_NODES = flags.DEFINE_integer( 'bigtable_autoscaling_min_nodes', None, 'Minimum number of nodes for autoscaling.') @@ -97,7 +89,6 @@ class BigtableSpec(non_relational_db.BaseNonRelationalDbSpec): zone: str project: str node_count: int - load_node_count: int storage_type: str replication_cluster: bool replication_cluster_zone: str @@ -119,7 +110,6 @@ def _GetOptionDecoderConstructions(cls): 'zone': (option_decoders.StringDecoder, none_ok), 'project': (option_decoders.StringDecoder, none_ok), 'node_count': (option_decoders.IntDecoder, none_ok), - 'load_node_count': (option_decoders.IntDecoder, none_ok), 'storage_type': (option_decoders.StringDecoder, none_ok), 'replication_cluster': (option_decoders.BooleanDecoder, none_ok), 'replication_cluster_zone': (option_decoders.StringDecoder, none_ok), @@ -157,7 +147,6 @@ def _ApplyFlags(cls, config_values, flag_values) -> None: 'google_bigtable_zone': 'zone', 'bigtable_storage_type': 'storage_type', 'bigtable_node_count': 'node_count', - 'bigtable_load_node_count': 'load_node_count', 'bigtable_replication_cluster': 'replication_cluster', 'bigtable_replication_cluster_zone': 'replication_cluster_zone', 'bigtable_multicluster_routing': 'multicluster_routing', @@ -206,7 +195,6 @@ def __init__(self, project: Optional[str], zone: Optional[str], node_count: Optional[int], - load_node_count: Optional[int], storage_type: Optional[str], replication_cluster: Optional[bool], replication_cluster_zone: Optional[str], @@ -222,7 +210,6 @@ def __init__(self, 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._load_node_count = load_node_count or self.node_count self.storage_type: str = storage_type or _DEFAULT_STORAGE_TYPE self.replication_cluster: bool = replication_cluster or False self.replication_cluster_zone: Optional[str] = ( @@ -240,7 +227,6 @@ def FromSpec(cls, spec: BigtableSpec) -> 'GcpBigtableInstance': zone=spec.zone, project=spec.project, node_count=spec.node_count, - load_node_count=spec.load_node_count, storage_type=spec.storage_type, replication_cluster=spec.replication_cluster, replication_cluster_zone=spec.replication_cluster_zone, @@ -401,29 +387,13 @@ def _GetClusters(self) -> List[Dict[str, Any]]: return json.loads(stdout) def _UpdateNodes(self, nodes: int) -> None: - """Updates clusters to the specified node count and waits until ready.""" clusters = self._GetClusters() for i in range(len(clusters)): - # Do nothing if the node count is already equal to what we want. - if clusters[i]['serveNodes'] == nodes: - continue cmd = _GetBigtableGcloudCommand(self, 'bigtable', 'clusters', 'update', f'{self.name}-{i}') cmd.flags['instance'] = self.name cmd.flags['num-nodes'] = nodes or self.node_count cmd.Issue() - # Note that Exists is implemented like IsReady, but should likely be - # refactored. - while not self._Exists(): - time.sleep(10) - - def UpdateCapacityForLoad(self) -> None: - """See base class.""" - self._UpdateNodes(self._load_node_count) - - def UpdateCapacityForRun(self) -> None: - """See base class.""" - self._UpdateNodes(self.node_count) def _Freeze(self) -> None: self._UpdateNodes(_FROZEN_NODE_COUNT) diff --git a/tests/linux_benchmarks/cloud_bigtable_ycsb_benchmark_test.py b/tests/linux_benchmarks/cloud_bigtable_ycsb_benchmark_test.py deleted file mode 100644 index 69109fb9c2..0000000000 --- a/tests/linux_benchmarks/cloud_bigtable_ycsb_benchmark_test.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright 2023 PerfKitBenchmarker Authors. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Tests for cloud_bigtable_ycsb_benchmark.""" -import textwrap -import unittest -from absl import flags -from absl.testing import flagsaver -import mock -from perfkitbenchmarker.linux_benchmarks import cloud_bigtable_ycsb_benchmark -from perfkitbenchmarker.linux_packages import ycsb -from perfkitbenchmarker.providers.gcp import gcp_bigtable -from tests import pkb_common_test_case - -FLAGS = flags.FLAGS - - -class CloudBigtableYcsbBenchmarkTest(pkb_common_test_case.PkbCommonTestCase): - - def _CreateMockBigtable(self, node_count, load_node_count): - mock_spec = textwrap.dedent(f""" - cloud_bigtable_ycsb: - non_relational_db: - service_type: bigtable - node_count: {node_count} - load_node_count: {load_node_count} - """) - self.mock_bm_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml( - mock_spec, 'cloud_bigtable_ycsb' - ) - self.mock_bm_spec.ConstructNonRelationalDb() - self.mock_update_nodes = self.enter_context( - mock.patch.object(self.mock_bm_spec.non_relational_db, '_UpdateNodes') - ) - return self.mock_bm_spec.non_relational_db - - def setUp(self): - super().setUp() - FLAGS.run_uri = 'test_uri' - self.executor = ycsb.YCSBExecutor('cloudbigtable') - self.mock_load = self.enter_context( - mock.patch.object(self.executor, 'Load') - ) - - def testLoadDatabaseIncreasedCapacity(self): - test_instance: gcp_bigtable.GcpBigtableInstance = self._CreateMockBigtable( - node_count=3, load_node_count=6 - ) - cloud_bigtable_ycsb_benchmark._LoadDatabase( - self.executor, test_instance, [], {} - ) - - self.mock_update_nodes.assert_has_calls([mock.call(6), mock.call(3)]) - self.mock_load.assert_called_once() - - def testLoadDatabaseNotCalledRestored(self): - test_instance: gcp_bigtable.GcpBigtableInstance = self._CreateMockBigtable( - node_count=3, load_node_count=6 - ) - test_instance.restored = True - - cloud_bigtable_ycsb_benchmark._LoadDatabase( - self.executor, test_instance, [], {} - ) - - self.mock_update_nodes.assert_not_called() - self.mock_load.assert_not_called() - - @flagsaver.flagsaver(ycsb_skip_load_stage=True) - def testLoadDatabaseNotCalledSkipFlag(self): - test_instance: gcp_bigtable.GcpBigtableInstance = self._CreateMockBigtable( - node_count=3, load_node_count=6 - ) - - cloud_bigtable_ycsb_benchmark._LoadDatabase( - self.executor, test_instance, [], {} - ) - - self.mock_update_nodes.assert_not_called() - self.mock_load.assert_not_called() - - -if __name__ == '__main__': - unittest.main() diff --git a/tests/providers/gcp/gcp_bigtable_test.py b/tests/providers/gcp/gcp_bigtable_test.py index a1a119201e..831910b78e 100644 --- a/tests/providers/gcp/gcp_bigtable_test.py +++ b/tests/providers/gcp/gcp_bigtable_test.py @@ -18,11 +18,10 @@ from absl import flags from absl.testing import flagsaver import mock + from perfkitbenchmarker import errors -from perfkitbenchmarker import vm_util from perfkitbenchmarker.providers.gcp import gcp_bigtable from perfkitbenchmarker.providers.gcp import util -from tests import matchers from tests import pkb_common_test_case import requests @@ -40,23 +39,6 @@ }} """ -_TEST_LIST_CLUSTERS_OUTPUT = [ - { - 'defaultStorageType': 'SSD', - 'location': 'projects/pkb-test/locations/us-west1-a', - 'name': 'projects/pkb-test/instances/pkb-bigtable-730b1e6b/clusters/pkb-bigtable-730b1e6b-1', - 'serveNodes': 3, - 'state': 'READY', - }, - { - 'defaultStorageType': 'SSD', - 'location': 'projects/pkb-test/locations/us-west1-c', - 'name': 'projects/pkb-test/instances/pkb-bigtable-730b1e6b/clusters/pkb-bigtable-730b1e6b-0', - 'serveNodes': 3, - 'state': 'READY', - }, -] - OUT_OF_QUOTA_STDERR = """ ERROR: (gcloud.beta.bigtable.instances.create) Operation successfully rolled @@ -300,52 +282,6 @@ def testBigtableGcloudCommand(self): self.assertEqual(cmd.flags['project'], PROJECT) self.assertEmpty(cmd.flags['zone']) - def testSetNodes(self): - test_instance = GetTestBigtableInstance() - self.enter_context( - mock.patch.object( - test_instance, - '_GetClusters', - return_value=_TEST_LIST_CLUSTERS_OUTPUT, - ) - ) - self.enter_context( - mock.patch.object(test_instance, '_Exists', return_value=True) - ) - cmd = self.enter_context( - mock.patch.object(vm_util, 'IssueCommand', return_value=[None, None, 0]) - ) - - test_instance._UpdateNodes(6) - - self.assertSequenceEqual( - [ - mock.call(matchers.HASALLOF('--num-nodes', '6')), - mock.call(matchers.HASALLOF('--num-nodes', '6')), - ], - cmd.mock_calls, - ) - - def testSetNodesSkipsIfCountAlreadyCorrect(self): - test_instance = GetTestBigtableInstance() - self.enter_context( - mock.patch.object( - test_instance, - '_GetClusters', - return_value=_TEST_LIST_CLUSTERS_OUTPUT, - ) - ) - self.enter_context( - mock.patch.object(test_instance, '_Exists', return_value=True) - ) - cmd = self.enter_context( - mock.patch.object(vm_util, 'IssueCommand', return_value=[None, None, 0]) - ) - - test_instance._UpdateNodes(3) - - cmd.assert_not_called() - if __name__ == '__main__': unittest.main()