Set client_vm based on benchmark_spec.vm_groups['clients'].

PiperOrigin-RevId: 692277086

perfkitbenchmarker/linux_benchmarks/hammerdbcli_benchmark.py

@@ -127,7 +127,9 @@ def Prepare(benchmark_spec: bm_spec.BenchmarkSpec):
     benchmark_spec: The benchmark specification. Contains all data that is
       required to run the benchmark.
   """
-  vm = benchmark_spec.vms[0]
+  client_vms = benchmark_spec.vm_groups['clients']
+  assert len(client_vms) == 1
+  client_vm = client_vms[0]
   relational_db = benchmark_spec.relational_db
   db_engine = relational_db.engine
   num_cpus = None
@@ -142,7 +144,7 @@ def Prepare(benchmark_spec: bm_spec.BenchmarkSpec):
     # Use 28800, the default mysql timeout
     relational_db.SetDbConfiguration('wait_timeout', '28800')
 
-  vm.Install('hammerdb')
+  client_vm.Install('hammerdb')
   optimized_server_config = (
       hammerdb.HAMMERDB_OPTIMIZED_SERVER_CONFIGURATION.value
   )
@@ -165,7 +167,7 @@ def Prepare(benchmark_spec: bm_spec.BenchmarkSpec):
           custom_server_config,
       )
   hammerdb.SetupConfig(
-      vm=vm,
+      vm=client_vm,
       db_engine=db_engine,
       hammerdb_script=hammerdb.HAMMERDB_SCRIPT.value,
       ip=relational_db.endpoint,
@@ -341,7 +343,9 @@ def Run(benchmark_spec: bm_spec.BenchmarkSpec) -> List[sample.Sample]:
   Raises:
     Exception: if the script is unknown
   """
-  vm = benchmark_spec.vms[0]
+  client_vms = benchmark_spec.vm_groups['clients']
+  assert len(client_vms) == 1
+  client_vm = client_vms[0]
   relational_db = benchmark_spec.relational_db
   num_cpus = None
   if hasattr(relational_db, 'server_vm'):
@@ -358,8 +362,10 @@ def Run(benchmark_spec: bm_spec.BenchmarkSpec) -> List[sample.Sample]:
   samples = []
   for i in range(1, 1 + hammerdb.NUM_RUN.value):
     metadata['run_iteration'] = i
-    stdout = hammerdb.Run(vm, db_engine, script, timeout=timeout)
-    current_samples = hammerdb.ParseResults(script=script, stdout=stdout, vm=vm)
+    stdout = hammerdb.Run(client_vm, db_engine, script, timeout=timeout)
+    current_samples = hammerdb.ParseResults(
+        script=script, stdout=stdout, vm=client_vm
+    )
     if (
         db_engine == sql_engine_utils.ALLOYDB
         and relational_db.enable_columnar_engine_recommendation
@@ -373,9 +379,9 @@ def Run(benchmark_spec: bm_spec.BenchmarkSpec) -> List[sample.Sample]:
       )
       relational_db.WaitColumnarEnginePopulates(database_name)
       # Another prewarm
-      stdout = hammerdb.Run(vm, db_engine, script, timeout=timeout)
+      stdout = hammerdb.Run(client_vm, db_engine, script, timeout=timeout)
       current_samples = hammerdb.ParseResults(
-          script=script, stdout=stdout, vm=vm
+          script=script, stdout=stdout, vm=client_vm
       )
 
     for s in current_samples:

perfkitbenchmarker/linux_benchmarks/pgbench_benchmark.py

@@ -194,8 +194,10 @@ def Prepare(benchmark_spec):
     benchmark_spec: benchmark_spec object which contains the database server and
       client_vm
   """
-  vm = benchmark_spec.vms[0]
-  vm.Install('pgbench')
+  client_vms = benchmark_spec.vm_groups['clients']
+  assert len(client_vms) == 1
+  client_vm = client_vms[0]
+  client_vm.Install('pgbench')
 
   UpdateBenchmarkSpecWithPrepareStageFlags(benchmark_spec)
 
@@ -235,15 +237,15 @@ def Prepare(benchmark_spec):
                 filler varchar(84)  NULL \\
             ); \\
             RUN BATCH;\"""")
-    vm.RobustRemoteCommand(
+    client_vm.RobustRemoteCommand(
         f"""pgbench "host=/tmp port=5432 dbname={TEST_DB_NAME} \\
         options='-c spanner.force_autocommit=on -c \\
         spanner.copy_max_parallelism=200 -c spanner.autocommit_dml_mode=\\'partitioned_non_atomic\\''" \\
         -i -Ig \\
         --scale={benchmark_spec.scale_factor}"""
     )
   else:
-    vm.RobustRemoteCommand(
+    client_vm.RobustRemoteCommand(
         f'pgbench {connection_string} -i -s {benchmark_spec.scale_factor}'
     )
 
@@ -375,5 +377,7 @@ def GetMetaData(db_size, benchmark_spec):
 
 def Cleanup(benchmark_spec):
   """Uninstalls pgbench from the client vm."""
-  vm = benchmark_spec.vms[0]
-  vm.Uninstall('pgbench')
+  client_vms = benchmark_spec.vm_groups['clients']
+  assert len(client_vms) == 1
+  client_vm = client_vms[0]
+  client_vm.Uninstall('pgbench')

perfkitbenchmarker/linux_benchmarks/sysbench_benchmark.py

@@ -826,7 +826,7 @@ def Run(benchmark_spec):
   """
   logging.info('Start benchmarking, Cloud Provider is %s.', FLAGS.cloud)
   results = []
-  client_vms = benchmark_spec.vms
+  client_vms = benchmark_spec.vm_groups['clients']
   db = benchmark_spec.relational_db
 
   for thread_count in FLAGS.sysbench_run_threads: