Fixed so that the workload to be tuned can be specified by the following two parameters.

1. data_load_command parameter
2. run_workload_command parameter

Author: ssl-oyamata

conf/postgres_opttune.conf

@@ -19,12 +19,12 @@ ssh_port = 22 # ssh port
 ssh_password = postgres # pg_os_user's ssh password
 
 [turning]
-study_name = pgbench_study # study name
+study_name = my_workload_study # study name
 required_recovery_time_second = 0
 # The maximum recovery time allowed by the user in case of a PostgreSQL crash,
 # which is used to estimate the wax_wal_size parameter.
 # Note: The default value of 0 does not perform the estimation of the wax_wal_size parameter.
-benchmark = pgbench # Benchmark tool name('sampled_workload' or pgbench' or 'oltpbench' or 'star_schema_benchmark')
+benchmark = my_workload # Benchmark tool name('my_workload' or 'sampled_workload' or pgbench' or 'oltpbench' or 'star_schema_benchmark')
 parameter_json_dir = ./conf/
 number_trail = 100 # Number of benchmarks to run for turning
 data_load_interval = 10 # Specify the data load interval by the number of benchmarks
@@ -34,10 +34,26 @@ save_study_history = True # Whether to save study history
 load_study_history = True # Whether to load study history if a study name already exists.
 history_database_url = sqlite:///study-history.db   # Example PostgreSQL. postgresql://postgres@localhost/study_history
 
-[sampled_workload]
-sampled_workload_save_file = workload_data/2020-09-13_202209.011708-2020-09-13_202239.011973.pkl
+[my-workload]
+data_load_command = /usr/pgsql-12/bin/pgbench -i -s 10 tpcc
+run_workload_command = /usr/pgsql-12/bin/pgbench tpcc -T 1200
+
+[sampled-workload]
 # File saved using workload_sampler.py
+sampled_workload_save_file = workload_data/2020-09-13_202209.011708-2020-09-13_202239.011973.pkl
+
+[workload-sampling]
+workload_sampling_time_second = 30
+# Time (in seconds) to sample the workload running on the database in the [PostgreSQL] section
+my_workload_save_dir = ./workload_data/ # workload save directory
+# Database settings to temporarily store workload information
+pghost = localhost # PostgreSQL server host
+pgport = 5432 # PostgreSQL server port
+pguser = postgres # PostgreSQL user name(Database user)
+pgpassword = postgres12 # PostgreSQL user password(Database user)
+pgdatabase = sampling # PostgreSQL Database
 
+##### Experimental settings
 [pgbench]
 scale_factor = 10 # pgbench scale factor
 clients = 10 # Number of clients
@@ -60,13 +76,3 @@ sql_key = Q1.1, Q2.1, Q3.1
 # Please specify the name of the file in pgopttune/workload/star_schema_sql/ directory(sql_file_path parameter diretory)
 # (e.g., Q1.1,Q2,1).
 
-[workload-sampling]
-workload_sampling_time_second = 30
-# Time (in seconds) to sample the workload running on the database in the [PostgreSQL] section
-my_workload_save_dir = ./workload_data/ # workload save directory
-# Database settings to temporarily store workload information
-pghost = localhost # PostgreSQL server host
-pgport = 5432 # PostgreSQL server port
-pguser = postgres # PostgreSQL user name(Database user)
-pgpassword = postgres12 # PostgreSQL user password(Database user)
-pgdatabase = sampling # PostgreSQL Database

pgopttune/config/my_workload_config.py

@@ -0,0 +1,15 @@
+from pgopttune.config.config import Config
+
+
+class MyWorkloadConfig(Config):
+    def __init__(self, conf_path, section='my-workload'):
+        super().__init__(conf_path)
+        self.config_dict = dict(self.config.items(section))
+
+    @property
+    def data_load_command(self):
+        return self.get_parameter_value('data_load_command')
+
+    @property
+    def run_workload_command(self):
+        return self.get_parameter_value('run_workload_command')

pgopttune/config/sampled_workload_config.py

@@ -3,7 +3,7 @@
 
 
 class SampledWorkloadConfig(Config):
-    def __init__(self, conf_path, section='sampled_workload'):
+    def __init__(self, conf_path, section='sampled-workload'):
         super().__init__(conf_path)
         self.conf_path = conf_path
         self.config_dict = dict(self.config.items(section))

pgopttune/objective/objective_my_workload.py

@@ -0,0 +1,18 @@
+from logging import getLogger
+from pgopttune.workload.my_workload import MyWorkLoad
+from pgopttune.objective.objective import Objective
+from pgopttune.config.postgres_server_config import PostgresServerConfig
+from pgopttune.config.tune_config import TuneConfig
+from pgopttune.config.my_workload_config import MyWorkloadConfig
+
+logger = getLogger(__name__)
+
+
+class ObjectiveMyWorkload(Objective):
+
+    def __init__(self,
+                 postgres_server_config: PostgresServerConfig,
+                 tune_config: TuneConfig,
+                 my_workload_config: MyWorkloadConfig):
+        super().__init__(postgres_server_config, tune_config)
+        self.workload = MyWorkLoad(postgres_server_config, my_workload_config)

pgopttune/objective/objective_sampled_workload.py

@@ -13,7 +13,7 @@ class ObjectiveSampledWorkload(Objective):
     def __init__(self,
                  postgres_server_config: PostgresServerConfig,
                  tune_config: TuneConfig,
-                 my_workload_config: SampledWorkloadConfig):
+                 sampled_workload_config: SampledWorkloadConfig):
         super().__init__(postgres_server_config, tune_config)
-        self.workload = SampledWorkload.load_sampled_workload(my_workload_config.my_workload_save_file,
+        self.workload = SampledWorkload.load_sampled_workload(sampled_workload_config.my_workload_save_file,
                                                               postgres_server_config=postgres_server_config)

pgopttune/workload/my_workload.py

@@ -0,0 +1,46 @@
+import os
+import time
+from logging import getLogger
+from .workload import Workload
+from pgopttune.utils.command import run_command
+from pgopttune.config.postgres_server_config import PostgresServerConfig
+from pgopttune.config.my_workload_config import MyWorkloadConfig
+
+logger = getLogger(__name__)
+
+
+class MyWorkLoad(Workload):
+    def __init__(self, postgres_server_config: PostgresServerConfig, my_workload_config: MyWorkloadConfig):
+        super().__init__(postgres_server_config)
+        self.my_workload_config = my_workload_config
+        os.environ['PGHOST'] = postgres_server_config.host
+        os.environ['PGPORT'] = postgres_server_config.port
+        os.environ['PGDATABASE'] = postgres_server_config.database
+        os.environ['PGUSER'] = postgres_server_config.user
+        os.environ['PGPASSWORD'] = postgres_server_config.password
+
+    def data_load(self):
+        data_load_cmd = self.my_workload_config.data_load_command
+        logger.debug('run my workload data load command : {}'.format(data_load_cmd))
+        run_command(data_load_cmd)
+        self.vacuum_database()  # vacuum analyze
+
+    def run(self):
+        run_workload_command = self.my_workload_config.run_workload_command
+        start_number_of_xact_commit = self.get_number_of_xact_commit()
+        workload_start_time = time.time()  # start measurement time
+        run_command(run_workload_command)  # run workload
+        # command_result = run_command(run_workload_command)  # run workload
+        # logger.info(command_result.stdout.decode("utf8"))
+        workload_elapsed_times = time.time() - workload_start_time
+        # logger.info(workload_elapsed_times)
+        time.sleep(1)  # default PGSTAT_STAT_INTERVAL(500ms)
+        workload_number_of_xact_commit = self.get_number_of_xact_commit() - start_number_of_xact_commit
+        # logger.info(workload_number_of_xact_commit)
+        tps = self.calculate_transaction_per_second(workload_number_of_xact_commit, workload_elapsed_times)
+        # logger.info("tps : {}".format(tps))
+        return tps
+
+    @staticmethod
+    def calculate_transaction_per_second(number_of_xact_commit, elapsed_seconds):
+        return round(number_of_xact_commit / elapsed_seconds, 6)

pgopttune/workload/workload.py

@@ -1,4 +1,5 @@
 from logging import getLogger
+from psycopg2.extras import DictCursor
 from pgopttune.utils.pg_connect import get_pg_connection
 from pgopttune.config.postgres_server_config import PostgresServerConfig
 
@@ -26,3 +27,12 @@ def execute_sql_file(self, sql_filepath):
             with conn.cursor() as cur:
                 cur.execute(open(sql_filepath, "r").read())
         logger.debug("finish execute {}".format(sql_filepath))
+
+    def get_number_of_xact_commit(self):
+        get_number_of_xact_commit_sql = "SELECT xact_commit FROM pg_stat_database WHERE datname = %s"
+        with get_pg_connection(dsn=self.postgres_server_config.dsn) as conn:
+            conn.set_session(autocommit=True)
+            with conn.cursor(cursor_factory=DictCursor) as cur:
+                cur.execute(get_number_of_xact_commit_sql, (self.postgres_server_config.database,))
+                xact_commit = cur.fetchone()["xact_commit"]
+        return xact_commit

sampling_workload.py

@@ -30,7 +30,7 @@ def main(
         "[turning]\n"
         "benchmark = sampled_workload \n"
         ":\n"
-        "[sampled_workload]\n"
+        "[sampled-workload]\n"
         "sampled_workload_save_file = {}".format(workload_save_file_path))
 
 

tune.py

@@ -7,6 +7,7 @@
 from pgopttune.utils.logger import logging_dict
 from pgopttune.sampler.sampler import get_sampler
 from pgopttune.study.study import create_study
+from pgopttune.objective.objective_my_workload import ObjectiveMyWorkload
 from pgopttune.objective.objective_pgbench import ObjectivePgbench
 from pgopttune.objective.objective_oltpbench import ObjectiveOltpbench
 from pgopttune.objective.objective_star_schema_benchmark import ObjectiveStarSchemaBenchmark
@@ -16,6 +17,7 @@
 from pgopttune.recovery.pg_recovery import Recovery
 from pgopttune.config.postgres_server_config import PostgresServerConfig
 from pgopttune.config.tune_config import TuneConfig
+from pgopttune.config.my_workload_config import MyWorkloadConfig
 from pgopttune.config.pgbench_config import PgbenchConfig
 from pgopttune.config.oltpbench_config import OltpbenchConfig
 from pgopttune.config.star_schema_benchmark_config import StarSchemaBenchmarkConfig
@@ -36,8 +38,12 @@ def main(
     optuna.logging.disable_default_handler()  # Stop showing logs in sys.stderr.
 
     # set objective
+    # my workload
+    if tune_config.benchmark == 'my_workload':
+        my_workload_config = MyWorkloadConfig(conf_path)  # pgbench config
+        objective = ObjectiveMyWorkload(postgres_server_config, tune_config, my_workload_config)
     # pgbench
-    if tune_config.benchmark == 'pgbench':
+    elif tune_config.benchmark == 'pgbench':
         pgbench_config = PgbenchConfig(conf_path)  # pgbench config
         objective = ObjectivePgbench(postgres_server_config, tune_config, pgbench_config)
     # oltpbench
@@ -48,7 +54,7 @@ def main(
     elif tune_config.benchmark == 'star_schema_benchmark':
         star_schema_benchmark_config = StarSchemaBenchmarkConfig(conf_path)  # star schema benchmark config
         objective = ObjectiveStarSchemaBenchmark(postgres_server_config, tune_config, star_schema_benchmark_config)
-    # my workload (save using sampling_workload.py)
+    # sampled workload (save using sampling_workload.py)
     elif tune_config.benchmark == 'sampled_workload':
         sampled_workload_config = SampledWorkloadConfig(conf_path)  # my workload sampled config
         objective = ObjectiveSampledWorkload(postgres_server_config, tune_config, sampled_workload_config)