Automatically create and delete cluster

perfkitbenchmarker/linux_benchmarks/cloud_bigtable_ycsb_benchmark.py

@@ -42,6 +42,7 @@
     as hbase_ycsb
 from perfkitbenchmarker.linux_packages import hbase
 from perfkitbenchmarker.linux_packages import ycsb
+from perfkitbenchmarker.providers.gcp import gce_bigtable
 
 FLAGS = flags.FLAGS
 
@@ -91,6 +92,9 @@
 # TODO(connormccoy): Make table parameters configurable.
 COLUMN_FAMILY = 'cf'
 
+# Only used when we need to create the cluster.
+CLUSTER_SIZE = 3
+
 
 def GetConfig(user_config):
   return configs.LoadConfig(BENCHMARK_CONFIG, user_config, BENCHMARK_NAME)
@@ -113,14 +117,12 @@ def CheckPrerequisites():
       raise ValueError('Scope {0} required.'.format(scope))
 
   # TODO: extract from gcloud config if available.
-  if not FLAGS.google_bigtable_cluster_name:
-    raise ValueError('Missing --google_bigtable_cluster_name')
-  if not FLAGS.google_bigtable_zone_name:
-    raise ValueError('Missing --google_bigtable_zone_name')
-  cluster = _GetClusterDescription(FLAGS.project or _GetDefaultProject(),
-                                   FLAGS.google_bigtable_zone_name,
-                                   FLAGS.google_bigtable_cluster_name)
-  logging.info('Found cluster: %s', cluster)
+  if FLAGS.google_bigtable_cluster_name:
+    cluster = _GetClusterDescription(FLAGS.project or _GetDefaultProject(),
+                                     FLAGS.google_bigtable_zone_name,
+                                     FLAGS.google_bigtable_cluster_name)
+    logging.info('Found cluster: %s', cluster)
+  logging.info('No cluster; will create in Prepare.')
 
 
 def _GetClusterDescription(project, zone, cluster_name):
@@ -180,6 +182,8 @@ def _Install(vm):
   vm.Install('ycsb')
   vm.Install('curl')
 
+  cluster_name = (FLAGS.google_bigtable_cluster_name or
+                  'pkb-bigtable-{0}'.format(FLAGS.run_uri))
   hbase_lib = posixpath.join(hbase.HBASE_DIR, 'lib')
   for url in [FLAGS.google_bigtable_hbase_jar_url, TCNATIVE_BORINGSSL_URL]:
     jar_name = os.path.basename(url)
@@ -194,7 +198,7 @@ def _Install(vm):
       'google_bigtable_endpoint': FLAGS.google_bigtable_endpoint,
       'google_bigtable_admin_endpoint': FLAGS.google_bigtable_admin_endpoint,
       'project': FLAGS.project or _GetDefaultProject(),
-      'cluster': FLAGS.google_bigtable_cluster_name,
+      'cluster': cluster_name,
       'zone': FLAGS.google_bigtable_zone_name,
   }
 
@@ -218,6 +222,23 @@ def Prepare(benchmark_spec):
   benchmark_spec.always_call_cleanup = True
   vms = benchmark_spec.vms
 
+  # TODO: in the future, it might be nice to chagne this so that
+  # a gce_bigtable.GceBigtableCluster can be created with an
+  # flag that says don't create/delete the cluster.  That would
+  # reduce the code paths here.
+  if FLAGS.google_bigtable_cluster_name is None:
+    cluster_name = 'pkb-bigtable-{0}'.format(FLAGS.run_uri)
+    project = FLAGS.project or _GetDefaultProject()
+    logging.info('Creating bigtable cluster %s', cluster_name)
+    zone = FLAGS.google_bigtable_zone_name
+    benchmark_spec.bigtable_cluster = gce_bigtable.GceBigtableCluster(
+        cluster_name, CLUSTER_SIZE, project, zone)
+    benchmark_spec.bigtable_cluster.Create()
+    cluster = _GetClusterDescription(project,
+                                     FLAGS.google_bigtable_zone_name,
+                                     cluster_name)
+    logging.info('Cluster %s created successfully', cluster)
+
   vm_util.RunThreaded(_Install, vms)
 
   # Create table
@@ -248,9 +269,11 @@ def Run(benchmark_spec):
                     'table': table_name}
 
   executor = ycsb.YCSBExecutor('hbase-10', **executor_flags)
+  cluster_name = (FLAGS.google_bigtable_cluster_name or
+                  'pkb-bigtable-{0}'.format(FLAGS.run_uri))
   cluster_info = _GetClusterDescription(FLAGS.project or _GetDefaultProject(),
                                         FLAGS.google_bigtable_zone_name,
-                                        FLAGS.google_bigtable_cluster_name)
+                                        cluster_name)
 
   metadata = {'ycsb_client_vms': len(vms),
               'bigtable_nodes': cluster_info.get('serveNodes')}
@@ -285,8 +308,12 @@ def Cleanup(benchmark_spec):
     benchmark_spec: The benchmark specification. Contains all data that is
         required to run the benchmark.
   """
-  vm = benchmark_spec.vms[0]
-  # Delete table
-  command = ("""echo 'disable "{0}"; drop "{0}"; exit' | """
-             """{1}/hbase shell""").format(_GetTableName(), hbase.HBASE_BIN)
-  vm.RemoteCommand(command, should_log=True, ignore_failure=True)
+ # Delete table
+  if FLAGS.google_bigtable_cluster_name is None:
+    benchmark_spec.bigtable_cluster.Delete()
+  else:
+    # Only need to drop the tables if we're not deleting the cluster.
+    vm = benchmark_spec.vms[0]
+    command = ("""echo 'disable "{0}"; drop "{0}"; exit' | """
+               """{1}/hbase shell""").format(_GetTableName(), hbase.HBASE_BIN)
+    vm.RemoteCommand(command, should_log=True, ignore_failure=True)

perfkitbenchmarker/providers/gcp/gce_bigtable.py

@@ -0,0 +1,84 @@
+# Copyright 2014 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.
+"""Module containing class for GCP's bigtable clusters.
+
+Clusters can be created and deleted.
+"""
+
+import json
+
+from perfkitbenchmarker import flags
+from perfkitbenchmarker import resource
+from perfkitbenchmarker.providers.gcp import util
+
+FLAGS = flags.FLAGS
+
+
+class GceBigtableCluster(resource.BaseResource):
+  """Object representing an GCE BigTable Cluster.
+
+  Attributes:
+    name: Cluster name.
+    num_nodes: Number of nodes in the cluster.
+    project: Enclosing project for the cluster.
+    zone: zone of the cluster.
+  """
+
+  def __init__(self, name, num_nodes, project, zone):
+    super(GceBigtableCluster, self).__init__()
+    self.num_nodes = num_nodes
+    self.name = name
+    self.zone = zone
+    self.project = project
+
+
+  def _Create(self):
+    """Creates the cluster."""
+    cmd = util.GcloudCommand(self, 'alpha', 'bigtable', 'clusters', 'create',
+                             self.name)
+    cmd.flags['description'] = 'PkbCreatedCluster'
+    cmd.flags['nodes'] = str(self.num_nodes)
+    cmd.flags['zone'] = self.zone
+    cmd.flags['project'] = self.project
+    cmd.Issue()
+
+
+  def _Delete(self):
+    """Deletes the cluster."""
+    cmd = util.GcloudCommand(self, 'alpha', 'bigtable', 'clusters', 'delete',
+                             self.name)
+    cmd.Issue()
+
+
+  def _Exists(self):
+    """Returns true if the cluster exists."""
+    cmd = util.GcloudCommand(self, 'alpha', 'bigtable', 'clusters', 'list')
+    cmd.flags['format'] = 'json'
+    cmd.flags['project'] = self.project
+    # The zone flag makes this command fail.
+    cmd.flags['zone'] = []
+    stdout, stderr, _ = cmd.Issue(suppress_warning=True)
+    try:
+      result = json.loads(stdout)
+      clusters = {cluster['name']: cluster for cluster in result['clusters']}
+      expected_cluster_name = 'projects/{0}/zones/{1}/clusters/{2}'.format(
+          self.project, self.zone, self.name)
+      try:
+        clusters[expected_cluster_name]
+        return True
+      except KeyError:
+        return False
+    except ValueError:
+      return False
+    return True