Added pggrill workload.

bench/ibench/autovac_driver.py

@@ -1,3 +1,4 @@
+import importlib
 import os
 import argparse
 import sys
@@ -6,14 +7,13 @@
 from learning.rl_glue import RLGlue
 from learning.rl import Agent, default_network_arch
 
-from workloads.iibench_driver import run_with_params
-
 from tqdm.auto import tqdm
 
 from executors.simulated_vacuum import SimulatedVacuum
 from executors.pg_stat_and_vacuum import PGStatAndVacuum
 
-def benchmark(resume_id, experiment_duration, model_type, model1_filename, model2_filename, instance_url, instance_user, instance_password, instance_dbname):
+def benchmark(workload_driver, resume_id, experiment_duration, model_type, model1_filename, model2_filename, instance_url, instance_user, instance_password, instance_dbname):
+    run_with_params = workload_driver.run_with_params
     id = 0
     for initial_size in tqdm([10000, 100000, 1000000]):
         for update_speed in tqdm([500, 1000, 2000, 4000, 8000, 16000, 32000, 64000]):
@@ -34,27 +34,28 @@ def benchmark(resume_id, experiment_duration, model_type, model1_filename, model
                             initial_size, update_speed, 5, experiment_duration, True, False, False,
                             model1_filename, True)
 
-            # Control with RL model #1
-            run_with_params(False, tag2, instance_url, instance_user, instance_password, instance_dbname,
-                            initial_size, update_speed, 5, experiment_duration, True, False, False,
-                            model2_filename, True)
+            # # Control with RL model #1
+            # if model2_filename and model2_filename != "":
+            #     run_with_params(False, tag2, instance_url, instance_user, instance_password, instance_dbname,
+            #                     initial_size, update_speed, 5, experiment_duration, True, False, False,
+            #                     model2_filename, True)
 
-            # Control with PID
-            run_with_params(False, tag3, instance_url, instance_user, instance_password, instance_dbname,
-                            initial_size, update_speed, 5, experiment_duration, True, True, False,
-                            "", True)
+            # # Control with PID
+            # run_with_params(False, tag3, instance_url, instance_user, instance_password, instance_dbname,
+            #                 initial_size, update_speed, 5, experiment_duration, True, True, False,
+            #                 "", True)
 
-            # Control with default autovacuum
-            run_with_params(False, tag4, instance_url, instance_user, instance_password, instance_dbname,
-                            initial_size, update_speed, 5, experiment_duration, False, False, False,
-                            "", True)
+            # # Control with default autovacuum
+            # run_with_params(False, tag4, instance_url, instance_user, instance_password, instance_dbname,
+            #                 initial_size, update_speed, 5, experiment_duration, False, False, False,
+            #                 "", True)
 
             gnuplot_cmd = ("gnuplot -e \"outfile='graph%s.png'; titlestr='Query latency graph (%s)'; filename1='%s_latencies.txt'; filename2='%s_latencies.txt'; filename3='%s_latencies.txt'; filename4='%s_latencies.txt'\" gnuplot_script.txt"
                            % (tag_suffix, tag_suffix, tag1, tag2, tag3, tag4))
             print("Gnuplot command: ", gnuplot_cmd)
             os.system(gnuplot_cmd)
 
-def learn(resume_id, experiment_duration, model_type, model1_filename, model2_filename, instance_url, instance_user, instance_password, instance_dbname):
+def learn(workload_driver, resume_id, experiment_duration, model_type, model1_filename, model2_filename, instance_url, instance_user, instance_password, instance_dbname):
     agent_configs = {
         'network_arch': default_network_arch,
 
@@ -91,7 +92,17 @@ def learn(resume_id, experiment_duration, model_type, model1_filename, model2_fi
         'approx_bytes_per_tuple': 100,
         'is_replay': is_replay,
         'replay_filename_mask': 'replay_n%d.txt',
-        'state_history_length': 10
+        'state_history_length': 10,
+        'workload_driver': workload_driver,
+        # START pggrill only
+        'initial_size_range': [1000_000, 1000_000],
+        'update_speed_range': [100, 100_000],
+        'num_cols_range': [0, 0],
+        'num_indexes_range': [0, 0],
+        'num_partitions_range': [0, 0],
+        'updated_percentage_range': [1, 50],
+        'num_workers_range': [1, 50],
+        # END pggrill only
     }
 
     experiment_configs = {
@@ -107,6 +118,7 @@ def learn(resume_id, experiment_duration, model_type, model1_filename, model2_fi
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description="Run the AutoVacuum reinforcement learning driver.")
     parser.add_argument('--cmd', type=str, choices=['benchmark', 'learn'], help='Command to execute (benchmark or learn)')
+    parser.add_argument('--workload-type', type=str, choices=['pggrill', 'iibench'], required=True, help='Type of the workload driver (pggrill or iibench)')
     parser.add_argument('--max-episodes', type=int, default=100, help='Maximum number of episodes for the experiment')
     parser.add_argument('--resume-id', type=int, default=0, help='Identifier to resume from a previous state')
     parser.add_argument('--experiment-duration', type=int, default=120, help='Duration of the experiment in seconds')
@@ -121,6 +133,9 @@ def learn(resume_id, experiment_duration, model_type, model1_filename, model2_fi
     args = parser.parse_args()
 
     cmd = args.cmd
+
+    # Load the specified workload driver dynamically
+    workload_driver = importlib.import_module(f"workloads.{args.workload_type}_driver")
 
     max_episodes = args.max_episodes
     resume_id = args.resume_id
@@ -134,8 +149,8 @@ def learn(resume_id, experiment_duration, model_type, model1_filename, model2_fi
     instance_dbname = args.instance_dbname
 
     if cmd == "benchmark":
-        benchmark(resume_id, experiment_duration, model_type, model1_filename, model2_filename, instance_url, instance_user, instance_password, instance_dbname)
+        benchmark(workload_driver, resume_id, experiment_duration, model_type, model1_filename, model2_filename, instance_url, instance_user, instance_password, instance_dbname)
     elif cmd == "learn":
-        learn(resume_id, experiment_duration, model_type, model1_filename, model2_filename, instance_url, instance_user, instance_password, instance_dbname)
+        learn(workload_driver, resume_id, experiment_duration, model_type, model1_filename, model2_filename, instance_url, instance_user, instance_password, instance_dbname)
     else:
         print("Invalid command")

bench/ibench/executors/pg_stat_and_vacuum.py

@@ -1,6 +1,5 @@
 import time
 import psycopg2
-from workloads.iibench_driver import run_with_default_settings
 from multiprocessing import Barrier, Process
 from executors.vacuum_experiment import VacuumExperiment
 
@@ -11,11 +10,16 @@ def startExp(self, env_info):
         self.db_host = env_info['db_host']
         self.db_user = env_info['db_user']
         self.db_pwd = env_info['db_pwd']
-        self.table_name = env_info['table_name']
+        self.workload_driver = env_info['workload_driver']
 
-        #print("Environment info (for PGStatAndVacuum):")
-        #for x in self.env_info:
-        #    print ('\t', x, ':', self.env_info[x])
+        params = self.workload_driver.collectExperimentParams(self.env_info)
+        self.table_name = params['table_name']
+
+        # print("Environment info (for PGStatAndVacuum):")
+        # for x in self.env_info:
+        #     print ('\t', x, ':', self.env_info[x])
+        # for x in params:
+        #     print ('\t', x, ':', params[x])
 
         self.is_replay = env_info['is_replay']
         self.replay_filename = env_info['replay_filename_mask'] % env_info['experiment_id']
@@ -28,7 +32,7 @@ def startExp(self, env_info):
                   % (len(self.replay_buffer), self.replay_filename))
         else:
             barrier = Barrier(2)
-            self.workload_thread = Process(target=run_with_default_settings, args=(barrier, self.env_info))
+            self.workload_thread = Process(target=self.workload_driver.run_with_default_settings, args=(barrier, self.env_info))
             self.workload_thread.start()
             # We wait until the workload is initialized and ready to start
             barrier.wait()
@@ -37,6 +41,9 @@ def startExp(self, env_info):
             self.conn.set_session(autocommit=True)
             self.cursor = self.conn.cursor()
 
+            print("Resetting stats...")
+            self.cursor.execute("SELECT pg_stat_reset()")
+
             print("Disabling autovacuum...")
             self.cursor.execute("alter table %s set ("
                                 "autovacuum_enabled = off,"

bench/ibench/executors/simulated_vacuum.py

@@ -1,10 +1,12 @@
-from workloads.iibench_driver import collectExperimentParams
 from executors.vacuum_experiment import VacuumExperiment
 
 class SimulatedVacuum(VacuumExperiment):
     def startExp(self, env_info):
         self.env_info = env_info
-        self.initial_size, self.update_speed = collectExperimentParams(self.env_info)
+        self.workload_driver = self.env_info['workload_driver']
+        params = self.workload_driver.collectExperimentParams(self.env_info)
+        self.initial_size = params['initial_size']
+        self.update_speed = params['update_speed']
 
         #print("Environment info (for SimulatedVacuum):")
         #for x in self.env_info: