Clean up Datastore YCSB benchmark since multiple databases per project are now supported.

PiperOrigin-RevId: 720665610

Author: bvliu

perfkitbenchmarker/linux_benchmarks/cloud_datastore_ycsb_benchmark.py

@@ -25,19 +25,13 @@
 to test Datastore.
 """
 
-import concurrent.futures
 import logging
-from multiprocessing import pool as pool_lib
-import time
 
 from absl import flags
-from google.cloud import datastore
-from google.oauth2 import service_account
 from perfkitbenchmarker import background_tasks
 from perfkitbenchmarker import configs
 from perfkitbenchmarker import errors
 from perfkitbenchmarker import virtual_machine
-from perfkitbenchmarker import vm_util
 from perfkitbenchmarker.linux_packages import ycsb
 from perfkitbenchmarker.providers.gcp import util
 
@@ -64,38 +58,15 @@
 # https://github.com/brianfrankcooper/YCSB/tree/master/googledatastore
 _YCSB_COLLECTIONS = ['usertable']
 
-# constants used for cleaning up database
-_CLEANUP_THREAD_POOL_WORKERS = 30
-_CLEANUP_KIND_READ_BATCH_SIZE = 12000
-_CLEANUP_KIND_DELETE_BATCH_SIZE = 6000
-_CLEANUP_KIND_DELETE_PER_THREAD_BATCH_SIZE = 3000
-_CLEANUP_KIND_DELETE_OP_BATCH_SIZE = 500
-
 FLAGS = flags.FLAGS
-_KEYFILE = flags.DEFINE_string(
-    'google_datastore_keyfile',
-    None,
-    'The path to Google API JSON private key file',
-)
-_PROJECT_ID = flags.DEFINE_string(
-    'google_datastore_projectId',
-    None,
-    'The project ID that has Cloud Datastore service',
-)
 _DATASET_ID = flags.DEFINE_string(
     'google_datastore_datasetId',
     None,
-    'The database ID that has Cloud Datastore service',
+    'The ID of the database to use for benchmarking',
 )
 _DEBUG = flags.DEFINE_string(
     'google_datastore_debug', 'false', 'The logging level when running YCSB'
 )
-_REPOPULATE = flags.DEFINE_boolean(
-    'google_datastore_repopulate',
-    False,
-    'If True, empty database & repopulate with new data.'
-    'By default, tests are run with pre-populated data.',
-)
 _TARGET_LOAD_QPS = flags.DEFINE_integer(
     'google_datastore_target_load_qps',
     500,
@@ -104,9 +75,7 @@
     ' for more info.',
 )
 
-_KEYFILE_LOCAL_PATH = '/tmp/key.json'
 _INSERTION_RETRY_LIMIT = 100
-_SLEEP_AFTER_LOADING_SECS = 30 * 60
 
 
 def GetConfig(user_config):
@@ -128,20 +97,10 @@ def _Install(vm):
   """Installs YCSB benchmark & copies datastore keyfile to client vm."""
   vm.Install('ycsb')
 
-  # Copy private key file to VM
-  if _KEYFILE.value:
-    if _KEYFILE.value.startswith('gs://'):
-      vm.Install('google_cloud_sdk')
-      vm.RemoteCommand(f'gsutil cp {_KEYFILE.value} {_KEYFILE_LOCAL_PATH}')
-    else:
-      vm.RemoteCopy(_KEYFILE.value, _KEYFILE_LOCAL_PATH)
-
 
 def _GetCommonYcsbArgs():
   """Returns common YCSB args."""
-  project = _PROJECT_ID.value
-  if project is None:
-    project = FLAGS.project or util.GetDefaultProject()
+  project = FLAGS.project or util.GetDefaultProject()
   args = {
       'googledatastore.projectId': project,
       'googledatastore.debug': _DEBUG.value,
@@ -153,10 +112,7 @@ def _GetCommonYcsbArgs():
 
 
 def _GetYcsbExecutor():
-  env = {}
-  if _KEYFILE.value:
-    env = {'GOOGLE_APPLICATION_CREDENTIALS': _KEYFILE_LOCAL_PATH}
-  return ycsb.YCSBExecutor('googledatastore', environment=env)
+  return ycsb.YCSBExecutor('googledatastore')
 
 
 def RampUpLoad(
@@ -205,9 +161,6 @@ def Prepare(benchmark_spec):
     benchmark_spec: The benchmark specification. Contains all data that is
       required to run the benchmark.
   """
-  if _REPOPULATE.value:
-    EmptyDatabase()
-
   vms = benchmark_spec.vms
 
   # Install required packages and copy credential files
@@ -232,10 +185,6 @@ def Run(benchmark_spec):
   Returns:
     A list of sample.Sample instances.
   """
-  if _REPOPULATE.value and not FLAGS.ycsb_skip_run_stage:
-    logging.info('Sleeping 30 minutes to allow for compaction.')
-    time.sleep(_SLEEP_AFTER_LOADING_SECS)
-
   executor = _GetYcsbExecutor()
   vms = benchmark_spec.vms
   run_kwargs = _GetCommonYcsbArgs()
@@ -250,247 +199,3 @@ def Run(benchmark_spec):
 
 def Cleanup(_):
   pass
-
-
-def EmptyDatabase():
-  """Deletes all entries in a datastore database."""
-  if _KEYFILE.value:
-    dataset_id = _DATASET_ID.value
-    executor = concurrent.futures.ThreadPoolExecutor(
-        max_workers=_CLEANUP_THREAD_POOL_WORKERS
-    )
-
-    logging.info('Attempting to delete all data in %s', dataset_id)
-
-    credentials = GetDatastoreDeleteCredentials()
-
-    futures = []
-    for kind in _YCSB_COLLECTIONS:
-      futures.append(
-          executor.submit(
-              _ReadAndDeleteAllEntities(dataset_id, credentials, kind)
-          )
-      )
-
-    concurrent.futures.wait(
-        futures, timeout=None, return_when=concurrent.futures.ALL_COMPLETED
-    )
-    logging.info('Deleted all data for %s', dataset_id)
-
-  else:
-    logging.warning('Manually delete all the entries via GCP portal.')
-
-
-def GetDatastoreDeleteCredentials():
-  """Returns credentials to datastore db."""
-  if _KEYFILE.value is not None and _KEYFILE.value.startswith('gs://'):
-    # Copy private keyfile to local disk
-    cp_cmd = [
-        'gsutil',
-        'cp',
-        _KEYFILE.value,
-        _KEYFILE_LOCAL_PATH,
-    ]
-    vm_util.IssueCommand(cp_cmd)
-    credentials_path = _KEYFILE_LOCAL_PATH
-  else:
-    credentials_path = _KEYFILE.value
-
-  if credentials_path is None:
-    raise errors.Benchmarks.RunError('Credentials path is None')
-
-  credentials = service_account.Credentials.from_service_account_file(
-      credentials_path,
-      scopes=datastore.client.Client.SCOPE,
-  )
-
-  return credentials
-
-
-def _ReadAndDeleteAllEntities(dataset_id, credentials, kind):
-  """Reads and deletes all kind entries in a datastore database.
-
-  Args:
-    dataset_id: Cloud Datastore client dataset id.
-    credentials: Cloud Datastore client credentials.
-    kind: Kind for which entities will be deleted.
-
-  Raises:
-    ValueError: In case of delete failures.
-  """
-  task_id = 1
-  start_cursor = None
-  pool = pool_lib.ThreadPool(processes=_CLEANUP_THREAD_POOL_WORKERS)
-
-  # We use a cursor to fetch entities in larger read batches and submit delete
-  # tasks to delete them in smaller delete batches (500 at a time) due to
-  # datastore single operation restriction.
-  entity_read_count = 0
-  total_entity_count = 0
-  delete_keys = []
-  while True:
-    query = _CreateClient(dataset_id, credentials).query(kind=kind)
-    query.keys_only()
-    query_iter = query.fetch(
-        start_cursor=start_cursor, limit=_CLEANUP_KIND_READ_BATCH_SIZE
-    )
-
-    for current_entities in query_iter.pages:
-      delete_keys.extend([entity.key for entity in current_entities])
-      entity_read_count = len(delete_keys)
-      # logging.debug('next batch of entities for %s - total = %d', kind,
-      #              entity_read_count)
-      if entity_read_count >= _CLEANUP_KIND_DELETE_BATCH_SIZE:
-        total_entity_count += entity_read_count
-        logging.info('Creating tasks...Read %d in total', total_entity_count)
-        while delete_keys:
-          delete_chunk = delete_keys[
-              :_CLEANUP_KIND_DELETE_PER_THREAD_BATCH_SIZE
-          ]
-          delete_keys = delete_keys[_CLEANUP_KIND_DELETE_PER_THREAD_BATCH_SIZE:]
-          # logging.debug(
-          #    'Creating new Task %d - Read %d entities for %s kind , Read %d '
-          #    + 'in total.',
-          #    task_id, entity_read_count, kind, total_entity_count)
-          deletion_task = _DeletionTask(kind, task_id)
-          pool.apply_async(
-              deletion_task.DeleteEntities,
-              (
-                  dataset_id,
-                  credentials,
-                  delete_chunk,
-              ),
-          )
-          task_id += 1
-
-        # Reset delete batch,
-        entity_read_count = 0
-        delete_keys = []
-
-    # Read this after the pages are retrieved otherwise it will be set to None.
-    start_cursor = query_iter.next_page_token
-    if start_cursor is None:
-      logging.info('Read all existing records for %s', kind)
-      if delete_keys:
-        logging.info(
-            'Entities batch is not empty %d, submitting new tasks',
-            len(delete_keys),
-        )
-        while delete_keys:
-          delete_chunk = delete_keys[
-              :_CLEANUP_KIND_DELETE_PER_THREAD_BATCH_SIZE
-          ]
-          delete_keys = delete_keys[_CLEANUP_KIND_DELETE_PER_THREAD_BATCH_SIZE:]
-          logging.debug(
-              (
-                  'Creating new Task %d - Read %d entities for %s kind , Read'
-                  ' %d in total.'
-              ),
-              task_id,
-              entity_read_count,
-              kind,
-              total_entity_count,
-          )
-          deletion_task = _DeletionTask(kind, task_id)
-          pool.apply_async(
-              deletion_task.DeleteEntities,
-              (
-                  dataset_id,
-                  credentials,
-                  delete_chunk,
-              ),
-          )
-          task_id += 1
-      break
-
-  logging.info('Waiting for all tasks - %d to complete...', task_id)
-  time.sleep(60)
-  pool.close()
-  pool.join()
-
-  # Rerun the query and delete any leftovers to make sure that all records
-  # are deleted as intended.
-  client = _CreateClient(dataset_id, credentials)
-  query = client.query(kind=kind)
-  query.keys_only()
-  entities = list(query.fetch(limit=20000))
-  if entities:
-    logging.info('Deleting leftover %d entities for %s', len(entities), kind)
-    total_entity_count += len(entities)
-    deletion_task = _DeletionTask(kind, task_id)
-    delete_keys = []
-    delete_keys.extend([entity.key for entity in entities])
-    deletion_task.DeleteEntities(dataset_id, credentials, delete_keys)
-
-  logging.info(
-      'Deleted all data for %s - %s - %d', dataset_id, kind, total_entity_count
-  )
-
-
-def _CreateClient(dataset_id, credentials):
-  """Creates a datastore client for the dataset using the credentials.
-
-  Args:
-    dataset_id: Cloud Datastore client dataset id.
-    credentials: Cloud Datastore client credentials.
-
-  Returns:
-    Datastore client.
-  """
-  return datastore.Client(project=dataset_id, credentials=credentials)
-
-
-class _DeletionTask:
-  """Represents a cleanup deletion task.
-
-  Attributes:
-    kind: Datastore kind to be deleted.
-    task_id: Task id
-    entity_deletion_count: No of entities deleted.
-    deletion_error: Set to true if deletion fails with an error.
-  """
-
-  def __init__(self, kind, task_id):
-    self.kind = kind
-    self.task_id = task_id
-    self.entity_deletion_count = 0
-    self.deletion_error = False
-
-  def DeleteEntities(self, dataset_id, credentials, delete_entities):
-    """Deletes entities in a datastore database in batches.
-
-    Args:
-      dataset_id: Cloud Datastore client dataset id.
-      credentials: Cloud Datastore client credentials.
-      delete_entities: Entities to delete.
-
-    Returns:
-      number of records deleted.
-    Raises:
-      ValueError: In case of delete failures.
-    """
-    try:
-      client = datastore.Client(project=dataset_id, credentials=credentials)
-      logging.info('Task %d - Started deletion for %s', self.task_id, self.kind)
-      while delete_entities:
-        chunk = delete_entities[:_CLEANUP_KIND_DELETE_OP_BATCH_SIZE]
-        delete_entities = delete_entities[_CLEANUP_KIND_DELETE_OP_BATCH_SIZE:]
-        client.delete_multi(chunk)
-        self.entity_deletion_count += len(chunk)
-
-      logging.info(
-          'Task %d - Completed deletion for %s - %d',
-          self.task_id,
-          self.kind,
-          self.entity_deletion_count,
-      )
-      return self.entity_deletion_count
-    except ValueError as error:
-      logging.exception(
-          'Task %d - Delete entities for %s failed due to %s',
-          self.task_id,
-          self.kind,
-          error,
-      )
-      self.deletion_error = True
-      raise error

requirements.txt

@@ -18,7 +18,6 @@ colorlog[windows]==2.6.0
 contextlib2>=0.5.1
 google-cloud-aiplatform
 google-cloud-core
-google-cloud-datastore
 google-cloud-monitoring>=2.0.0
 immutabledict
 jinja2>=2.10.2