Typically, database tuning tools focus on one type of decision. Some tools recommend database indexes, others optimize system configuration parameters such as the buffer pool size. With UDO (Universal Database Optimizer), we present a tool that optimizes all of those tuning choices by one unified approach.
10 | 11 |UDO uses reinforcement learning to converge to optimal tuning choices. Given a list of index candidates and tuning parameters, it explores interesting options and evaluates performance on an example workload. We completely bypass simplifying cost models and rely on observations, collected during optimization, alone. In doing so, we prioritize optimization quality over optimization speed: UDO may take hours to optimize a given workload but it converges to optimal results.
12 | 13 |In the current version, UDO tunes Postgres for given SQL workloads. It can be installed via pip (see here). See the README file for a quick intro (using UDO to optimize for TPC-H). 14 | 15 |
UDO is supported by NSF Grant 1910830.
16 | 17 | 18 | 19 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | setup(name='udo_db', 4 | version='0.20', 5 | description='a universal optimizer for database system', 6 | author='Junxiong Wang', 7 | author_email='chuangzhetianxia@gmail.coma', 8 | url='https://ovss.github.io/UDO/', 9 | download_url='https://github.com/OVSS/UDO/archive/refs/tags/0.01.tar.gz', 10 | keywords=['Database Optimization', 'OLAP', 'Index selection', 'System parameters'], 11 | packages=find_packages(), 12 | license='MIT', 13 | install_requires=[ # I get to this in a second 14 | 'numpy==1.19.2', 15 | 'h5py==2.10.0', 16 | 'testresources>=2.0.1', 17 | 'gym>=0.17.3', 18 | 'udo_optimization>=0.0.3', 19 | 'tensorflow>=2.4.2', 20 | 'keras>=2.4.3', 21 | 'keras-rl>=0.4.2', 22 | 'mysqlclient>=2.0.3', 23 | 'psycopg2>=2.8.6' 24 | ], 25 | classifiers=[ 26 | 'Development Status :: 3 - Alpha', 27 | # Chose either "3 - Alpha", "4 - Beta" or "5 - Production/Stable" as the current state of your package 28 | 'Intended Audience :: Developers', # Define that your audience are developers 29 | 'Topic :: Database', 30 | 'License :: OSI Approved :: MIT License', # Again, pick a license 31 | 'Programming Language :: Python :: 3', # Specify which python versions that you want to support 32 | 'Programming Language :: Python :: 3.4', 33 | 'Programming Language :: Python :: 3.5', 34 | 'Programming Language :: Python :: 3.6', 35 | ]) 36 | -------------------------------------------------------------------------------- /udo/mcts/mcts_node.py: -------------------------------------------------------------------------------- 1 | # ----------------------------------------------------------------------- 2 | # Copyright (c) 2021 Cornell Database Group 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining 5 | # a copy of this software and associated documentation files (the 6 | # "Software"), to deal in the Software without restriction, including 7 | # without limitation the rights to use, copy, modify, merge, publish, 8 | # distribute, sublicense, and/or sell copies of the Software, and to 9 | # permit persons to whom the Software is furnished to do so, subject to 10 | # the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be 13 | # included in all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 17 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT 18 | # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 19 | # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 20 | # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 21 | # OTHER DEALINGS IN THE SOFTWARE. 22 | # ----------------------------------------------------------------------- 23 | 24 | from enum import Enum 25 | 26 | 27 | class SelectionPolicy(Enum): 28 | UCB1 = 1 29 | UCBV = 2 30 | 31 | 32 | class UpdatePolicy(Enum): 33 | MCTS = 1 34 | RAVE = 2 35 | 36 | 37 | class SpaceType(Enum): 38 | Light = 1 39 | Heavy = 2 40 | All = 3 41 | 42 | 43 | class mcts_node(object): 44 | 45 | def select_action(self): 46 | """select actions""" 47 | return None 48 | 49 | def playout(self, current_level_action): 50 | """playout from a leaf node""" 51 | return None 52 | 53 | def sample(self, round): 54 | """sample from current node""" 55 | return None 56 | 57 | def opt_policy(self): 58 | """retrieve the current optimal policy""" 59 | return None 60 | 61 | ## MCTS 62 | -------------------------------------------------------------------------------- /udo/agent/sarsa_agent.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | import gym 4 | import udo_optimization 5 | 6 | from rl.agents import SARSAAgent 7 | from rl.policy import BoltzmannQPolicy 8 | from tensorflow.keras.layers import Dense, Activation, Flatten 9 | from tensorflow.keras.models import Sequential 10 | from tensorflow.keras.optimizers import Adam 11 | 12 | def run_sarsa_agent(driver, queries, candidate_indices, tuning_config): 13 | # Get the environment and extract the number of actions. 14 | env = gym.make("udo_optimization-v0", driver=driver, queries=queries, candidate_indices=candidate_indices, 15 | config=tuning_config) 16 | env.horizon = tuning_config['horizon'] 17 | 18 | nb_actions = env.action_space.n 19 | logging.info(f"nr action: {nb_actions}") 20 | logging.info(f"observation space: {env.observation_space.shape}") 21 | 22 | # Next, we build a very simple model. 23 | model = Sequential() 24 | model.add(Flatten(input_shape=(1,) + env.observation_space.shape)) 25 | model.add(Dense(52)) 26 | model.add(Activation('relu')) 27 | model.add(Dense(252)) 28 | model.add(Activation('relu')) 29 | model.add(Dense(526)) 30 | model.add(Activation('relu')) 31 | model.add(Dense(252)) 32 | model.add(Activation('relu')) 33 | model.add(Dense(nb_actions)) 34 | model.add(Activation('linear')) 35 | 36 | logging.info(model.summary()) 37 | 38 | # SARSA does not require a memory. 39 | policy = BoltzmannQPolicy() 40 | # policy.select_action() 41 | sarsa = SARSAAgent(model=model, nb_actions=nb_actions, nb_steps_warmup=10, policy=policy) 42 | sarsa.compile(Adam(lr=1e-3), metrics=['mae']) 43 | 44 | # Okay, now it's time to learn something! We visualize the training here for show, but this 45 | # slows down training quite a lot. You can always safely abort the training prematurely using 46 | # Ctrl + C. 47 | sarsa.fit(env, nb_steps=500, visualize=False, verbose=2) 48 | 49 | # After training is done, we save the final weights. 50 | # sarsa.save_weights('sarsa_{}_weights.h5f'.format(udo_optimization-v0), overwrite=True) 51 | 52 | # Finally, evaluate our algorithm for 5 episodes. 53 | sarsa.test(env, nb_episodes=5, visualize=False) 54 | env.print_state_summary(env.best_state) 55 | -------------------------------------------------------------------------------- /udo/extract_index.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # ----------------------------------------------------------------------- 3 | # Copyright (c) 2021 Cornell Database Group 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining 6 | # a copy of this software and associated documentation files (the 7 | # "Software"), to deal in the Software without restriction, including 8 | # without limitation the rights to use, copy, modify, merge, publish, 9 | # distribute, sublicense, and/or sell copies of the Software, and to 10 | # permit persons to whom the Software is furnished to do so, subject to 11 | # the following conditions: 12 | # 13 | # The above copyright notice and this permission notice shall be 14 | # included in all copies or substantial portions of the Software. 15 | # 16 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT 19 | # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 20 | # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 21 | # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 22 | # OTHER DEALINGS IN THE SOFTWARE. 23 | # ----------------------------------------------------------------------- 24 | 25 | import argparse 26 | import json 27 | import os 28 | import re 29 | import itertools 30 | 31 | udo_parser = argparse.ArgumentParser(description='UDO index candidate generator.') 32 | 33 | udo_parser.add_argument('-db_schema', help='the database schema to optimizes') 34 | udo_parser.add_argument('-queries', help='queries') 35 | 36 | args = udo_parser.parse_args() 37 | # change 38 | args = vars(args) 39 | 40 | with open(args['db_schema']) as f: 41 | db_schema = json.load(f) 42 | 43 | if args['queries']: 44 | queries = dict() 45 | for file_name in os.listdir(args['queries']): 46 | if file_name.endswith(".sql"): 47 | with open(os.path.join(args['queries'], file_name)) as f: 48 | content = f.read() 49 | queries[file_name] = content 50 | 51 | table_pattern = re.compile("FROM(.*)WHERE") 52 | where_pattern = re.compile("WHERE(.*)") 53 | collect_indices = dict() 54 | for query_id, query in queries.items(): 55 | matches = table_pattern.findall(query) 56 | join_elements = matches[0].split(',') 57 | rename_table_dict = {} 58 | for join_element in join_elements: 59 | rename_pos = join_element.index("AS") 60 | if join_element[:rename_pos].strip() not in rename_table_dict: 61 | rename_table_dict[join_element[:rename_pos].strip()] = [join_element[rename_pos + 2:].strip()] 62 | else: 63 | rename_table_dict[join_element[:rename_pos].strip()].append(join_element[rename_pos + 2:].strip()) 64 | print(rename_table_dict) 65 | where_clause = where_pattern.findall(query)[0] 66 | for table, columns in db_schema.items(): 67 | if table in rename_table_dict: 68 | for rename_table in rename_table_dict[table]: 69 | for column in columns: 70 | if rename_table + '.' + column in where_clause: 71 | if table not in collect_indices: 72 | collect_indices[table] = [column] 73 | elif column not in collect_indices[table]: 74 | collect_indices[table].append(column) 75 | print(collect_indices) 76 | 77 | for k in collect_indices.keys(): 78 | all_candidate = collect_indices[k] 79 | v = [] 80 | for i in range(1, len(all_candidate) + 1): 81 | v += [','.join(comb) for comb in (itertools.combinations(all_candidate, i))] 82 | for i in range(len(v)): 83 | print("('IDX_%s_%d','%s','%s')," % (k, i, k, v[i])) 84 | -------------------------------------------------------------------------------- /udo/agent/udo_simplifed_agent.py: -------------------------------------------------------------------------------- 1 | # ----------------------------------------------------------------------- 2 | # Copyright (c) 2021 Cornell Database Group 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining 5 | # a copy of this software and associated documentation files (the 6 | # "Software"), to deal in the Software without restriction, including 7 | # without limitation the rights to use, copy, modify, merge, publish, 8 | # distribute, sublicense, and/or sell copies of the Software, and to 9 | # permit persons to whom the Software is furnished to do so, subject to 10 | # the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be 13 | # included in all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 17 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT 18 | # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 19 | # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 20 | # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 21 | # OTHER DEALINGS IN THE SOFTWARE. 22 | # ----------------------------------------------------------------------- 23 | 24 | import logging 25 | import time 26 | 27 | import gym 28 | import udo_optimization 29 | 30 | from udo.mcts.mcts_node import SpaceType 31 | from udo.mcts.uct_node import uct_node 32 | 33 | def run_simplifed_udo_agent(driver, queries, candidate_indices, tuning_config): 34 | start_tune_time = time.time() 35 | env = gym.make('udo_optimization-v0', driver=driver, queries=queries, candidate_indices=candidate_indices, 36 | config=tuning_config) 37 | env.reset() 38 | duration_in_seconds = tuning_config['duration'] * 3600 39 | init_state = env.state_decoder(0) 40 | root = uct_node(0, 0, tuning_config['horizon'], init_state, env, space_type=SpaceType.All) 41 | 42 | start_time = time.time() 43 | current_time = time.time() 44 | query_keys = range(len(queries)) 45 | 46 | ep = 0 47 | logging.debug(f"start: {time.time()}") 48 | while (current_time - start_time) < duration_in_seconds: 49 | # for the micro episode 50 | selected_actions = root.sample(ep) 51 | logging.debug(f"selected_actions: {selected_actions}") 52 | # evaluate the light actions 53 | env.reset() 54 | for selected_action in selected_actions: 55 | # move to next state 56 | state = env.step_without_evaluation(selected_action) 57 | # obtain run time info by running queries within timeout 58 | run_time = env.evaluate(query_keys) 59 | # the total time of sampled queries 60 | total_run_time = sum(run_time) 61 | # the default time of sampled queries 62 | default_time = sum(env.default_runtime) 63 | # the relative ration of the improvement, the less of total_run_time, the better 64 | # light_reward = default_time / total_run_time 65 | light_reward = max(default_time - total_run_time, 0) 66 | logging.debug(f"light reward: {light_reward}") 67 | root.update_statistics_with_mcts_reward(light_reward, selected_actions) 68 | root.print_reward_info() 69 | current_best = root.best_actions() 70 | logging.debug(f"current best action: {current_best}") 71 | logging.debug(f"runtime: {total_run_time}") 72 | logging.info(f"current best configurations:") 73 | env.print_action_summary(root.best_actions()) 74 | # add more episode 75 | ep += 1 76 | final_tune_time = time.time() 77 | logging.debug(f"end: {final_tune_time}") 78 | # tuning summary 79 | best_actions = root.best_actions() 80 | logging.info(f"Summary: Total Tuning Time {(final_tune_time - start_tune_time) / 3600} hours") 81 | env.print_action_summary(best_actions) 82 | -------------------------------------------------------------------------------- /udo/drivers/abstractdriver.py: -------------------------------------------------------------------------------- 1 | # ----------------------------------------------------------------------- 2 | # Copyright (c) 2021 Cornell Database Group 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining 5 | # a copy of this software and associated documentation files (the 6 | # "Software"), to deal in the Software without restriction, including 7 | # without limitation the rights to use, copy, modify, merge, publish, 8 | # distribute, sublicense, and/or sell copies of the Software, and to 9 | # permit persons to whom the Software is furnished to do so, subject to 10 | # the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be 13 | # included in all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 17 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT 18 | # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 19 | # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 20 | # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 21 | # OTHER DEALINGS IN THE SOFTWARE. 22 | # ----------------------------------------------------------------------- 23 | 24 | import logging 25 | 26 | class AbstractDriver(object): 27 | def __init__(self, driver_name, conf, sys_params): 28 | self.driver_name = driver_name 29 | self.config = conf 30 | self.sys_params = sys_params 31 | self.sys_params_type = len(self.sys_params) 32 | self.sys_params_space = [len(specific_parameter) for specific_parameter in self.sys_params] 33 | 34 | def __str__(self): 35 | return self.driver_name 36 | 37 | def connect(self): 38 | """connect to a DBMS""" 39 | raise NotImplementedError("%s does not implement connect function" % (self.driver_name)) 40 | 41 | def run_queries_with_timeout(self, query_list, timeout): 42 | """run queries with specific timeout""" 43 | raise NotImplementedError("%s does not implement run queries with timeout function" % (self.driver_name)) 44 | 45 | def run_queries_without_timeout(self, query_list): 46 | """run queries without specific timeout""" 47 | raise NotImplementedError("%s does not implement run queries without timeout function" % (self.driver_name)) 48 | 49 | def build_index(self, index_creation_sql): 50 | """build index via SQL""" 51 | raise NotImplementedError("%s does not implement build index function" % (self.driver_name)) 52 | 53 | def drop_index(self, index_drop_sql): 54 | """drop index via SQL""" 55 | raise NotImplementedError("%s does not implement drop index function" % (self.driver_name)) 56 | 57 | def build_index_command(self, index_to_create): 58 | """build index command""" 59 | return None 60 | 61 | def set_system_parameter(self, parameter_sql): 62 | """switch system parameters via SQL""" 63 | raise NotImplementedError("%s does not implement set system parameter function" % (self.driver_name)) 64 | 65 | def change_system_parameter(self, parameter_choices): 66 | """change system parameter values using the input parameter choices""" 67 | for i in range(self.sys_params_type): 68 | parameter_choice = int(parameter_choices[i]) 69 | parameter_change_sql = self.sys_params[i][parameter_choice] 70 | logging.info(f"{parameter_change_sql}") 71 | self.set_system_parameter(parameter_change_sql) 72 | 73 | def get_system_parameter_command(self, parameter_type, parameter_value): 74 | """get the system parameter command based on the parameter_type and its value""" 75 | return self.sys_params[parameter_type][parameter_value] 76 | 77 | def get_system_parameter_space(self): 78 | """return the search space of system parameter""" 79 | return self.sys_params_space 80 | 81 | ## CLASS 82 | -------------------------------------------------------------------------------- /udo/optimizer/order_optimizer.py: -------------------------------------------------------------------------------- 1 | # ----------------------------------------------------------------------- 2 | # Copyright (c) 2021 Cornell Database Group 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining 5 | # a copy of this software and associated documentation files (the 6 | # "Software"), to deal in the Software without restriction, including 7 | # without limitation the rights to use, copy, modify, merge, publish, 8 | # distribute, sublicense, and/or sell copies of the Software, and to 9 | # permit persons to whom the Software is furnished to do so, subject to 10 | # the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be 13 | # included in all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 17 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT 18 | # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 19 | # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 20 | # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 21 | # OTHER DEALINGS IN THE SOFTWARE. 22 | # ----------------------------------------------------------------------- 23 | 24 | import logging 25 | 26 | class OrderOptimizer: 27 | """ 28 | Optimize the evaluation order using the estimated cost 29 | """ 30 | 31 | def __init__(self, index_card_info): 32 | """cardinality information""" 33 | self.index_card_info = index_card_info 34 | 35 | def index_build_cost(self, current_index_list): 36 | """index build cost according to table cardinality""" 37 | return sum(self.index_card_info[current_index] for current_index in current_index_list) 38 | 39 | def greedy_min_cost_order(self, selected_action_batch): 40 | """determine the order to evaluate those actions of the given batch using greedy""" 41 | batch_size = len(selected_action_batch) 42 | if batch_size < 3: 43 | return range(batch_size) 44 | else: 45 | selected_action_batch_set = [set(batch) for batch in selected_action_batch] 46 | # take the first 47 | first_action = selected_action_batch_set[0] 48 | second_action = selected_action_batch_set[1] 49 | # cost for the first two elements 50 | common = first_action.intersection(second_action) 51 | cost = self.index_build_cost(first_action) + self.index_build_cost(second_action) - self.index_build_cost( 52 | common) 53 | order = [0, 1] 54 | for selected_action_idx in range(2, len(selected_action_batch_set)): 55 | # obtain select action 56 | selected_action = selected_action_batch_set[selected_action_idx] 57 | # consider to insert the new action to the first or last element 58 | current_cost1 = cost + self.index_build_cost(selected_action) - self.index_build_cost( 59 | selected_action.intersection(selected_action_batch_set[order[0]])) 60 | current_cost2 = cost + self.index_build_cost(selected_action) - self.index_build_cost( 61 | selected_action.intersection(selected_action_batch_set[order[-1]])) 62 | if current_cost2 > current_cost1: 63 | min_cost = current_cost1 64 | min_pos = -1 65 | else: 66 | min_cost = current_cost2 67 | min_pos = len(order) - 1 68 | # consider to insert the new action to the middle element 69 | for insert_pos in range(0, len(order) - 1): 70 | prev_pos = order[insert_pos] 71 | next_pos = order[insert_pos + 1] 72 | current_cost = cost + self.index_build_cost(selected_action_batch_set[prev_pos].intersection( 73 | selected_action_batch_set[next_pos])) - self.index_build_cost( 74 | selected_action.intersection(selected_action_batch_set[prev_pos])) - self.index_build_cost( 75 | selected_action.intersection(selected_action_batch_set[next_pos])) + self.index_build_cost( 76 | selected_action) 77 | if current_cost < min_cost: 78 | min_cost = current_cost 79 | min_pos = insert_pos 80 | # update the order 81 | order.insert(min_pos + 1, selected_action_idx) 82 | logging.debug(f"min cost order: {order}") 83 | return order 84 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # UDO Quickstart 2 | 3 | The following installation procedure was tested on Ubuntu 20.04 with Python 3. More precisely, we used a t3.medium EC2 instance and the "Ubuntu Server 20.04 LTS (HVM), SSD Volume Type" AMI. 4 | 5 | ## Installation from GitHub 6 | 7 | 1. Download UDO package from UDO repository and switch to UDO directory. 8 | 9 | ``` 10 | git clone https://[username]@github.com/OVSS/UDO.git 11 | cd UDO 12 | ``` 13 | 14 | 2. Install DBMS packages. `bash ./install.sh` 15 | 16 | 3. Install UDO packages. 17 | 18 | ``` 19 | cd udo 20 | python3 -m pip install -r requirement.txt` 21 | ``` 22 | 23 | 4. Install UDO Gym environment. 24 | 25 | ``` 26 | cd udo-optimization/ 27 | python3 -m pip install -e . 28 | ``` 29 | 30 | ## Installation via PIP 31 | 32 | 1. Install DBMS packages via `bash ./install.sh` if Postgres and MySQL are not installed. 33 | 34 | 2. Use `python3 -m pip install UDO-DB` to install packages. 35 | 36 | ## Prepare TPC-H Database 37 | 38 | The TPC-H schema, dataset, and queries are available at https://drive.google.com/drive/folders/123pwHaoz8C1dakvUef8AjKqci3_JNG47. 39 | 40 | 1. To download data from Google Drive, install gdown via `python3 -m pip install gdown`. 41 | 42 | 2. Download TPC-H .zip file using `/home/ubuntu/.local/bin/gdown https://drive.google.com/uc?id=1IgzHMOc75Km9h-FLMepV-t9lrQhWGTwt`. 43 | 44 | 3. Install unzip via `sudo apt install unzip` and use it to extract files via `unzip TPC-H.zip` 45 | 46 | 4. Create TPC-H database via `sudo -u postgres createdb tpch_sf10`. 47 | 48 | 5. Create TPC-H database schema via `sudo -u postgres psql tpch_sf10 < tpch_schema.sql`. 49 | 50 | 6. Load TPC-H data via `sudo -u postgres psql tpch_sf10 < tpch_sf10_data.sql`. 51 | 52 | ## Use UDO to Tune Postgres for TPC-H 53 | 54 | 1. Optional, using the UDO tool to extract indexes. The output format should be `index name;table;columns`, which is the index format required by UDO. 55 | 56 | ``` 57 | usage: extract_index.py [-h] [-db_schema DB_SCHEMA] [-queries QUERIES] 58 | 59 | UDO index candidate generator. 60 | 61 | optional arguments: 62 | -h, --help show this help message and exit 63 | -db_schema DB_SCHEMA the database schmea to optimizes 64 | -queries QUERIES queries 65 | ``` 66 | 67 | 2. ```echo "PATH=$PATH:/home/ubuntu/.local/bin" >> ~/.bashrc``` 68 | 69 | 3. Run agents after installing udo from pip 70 | 71 | ``` 72 | usage: python3 -m udo [-h] [-system {mysql,postgres}] [-db DB] [-username USERNAME] [-password PASSWORD] [-queries QUERIES] 73 | [-indices INDICES] [-sys_params SYS_PARAMS] [-duration DURATION] [-agent {udo,udo-s,ddpg,sarsa}] 74 | [-horizon HORIZON] [-heavy_horizon HEAVY_HORIZON] [-rl_update {RAVE,MCTS}] [-rl_select {UCB1,UCBV}] 75 | [-rl_reward {delta,accumulate}] [-rl_delay {UCB,Exp3}] [-rl_max_delay_time RL_MAX_DELAY_TIME] 76 | [-sample_rate SAMPLE_RATE] [-default_query_time_out DEFAULT_QUERY_TIME_OUT] [-time_out_ratio TIME_OUT_RATIO] 77 | [--load_json LOAD_JSON] 78 | 79 | UDO optimizer. 80 | 81 | optional arguments: 82 | -h, --help show this help message and exit 83 | -system {mysql,postgres} 84 | Target system driver 85 | -db DB the database to optimizes 86 | -username USERNAME username 87 | -password PASSWORD password 88 | -queries QUERIES the input query file 89 | -indices INDICES the input query file 90 | -sys_params SYS_PARAMS 91 | the input system params json file 92 | -duration DURATION time for tuning in hours 93 | -agent {udo,udo-s,ddpg,sarsa} 94 | reinforcement learning agent 95 | -horizon HORIZON the number horizon for reinforcement agent 96 | -heavy_horizon HEAVY_HORIZON 97 | the number horizon for heavy parameters in UDO 98 | -rl_update {RAVE,MCTS} 99 | the update policy of UDO tree search 100 | -rl_select {UCB1,UCBV} 101 | the selection policy of UDO tree search 102 | -rl_reward {delta,accumulate} 103 | the reward of reinforcement learning agent 104 | -rl_delay {UCB, Exp3} the delay selection policy 105 | -rl_max_delay_time RL_MAX_DELAY_TIME 106 | the delay selection policy 107 | -sample_rate SAMPLE_RATE 108 | sampled rate from workload 109 | -default_query_time_out DEFAULT_QUERY_TIME_OUT 110 | default timeout in seconds for each query 111 | -time_out_ratio TIME_OUT_RATIO 112 | timeout ratio respect to default time 113 | --load_json LOAD_JSON 114 | Load settings from file in json format. Command line options override values in file. 115 | ``` 116 | 117 | For example 118 | 119 | for TPC-H with scaling factor 1 120 | 121 | ``` 122 | python3 -m udo -system postgres -db tpch_sf1 -username postgres -queries tpch_queries -indices tpch_index.txt -sys_params postgressysparams.json -duration 5 -agent udo -horizon 8 -heavy_horizon 3 -rl_max_delay_time 5 -default_query_time_out 6 123 | ``` 124 | 125 | for TPC-H with scaling factor 10 126 | 127 | ``` 128 | python3 -m udo -system postgres -db tpch_sf10 -username postgres -queries tpch_queries -indices tpch_index.txt -sys_params postgressysparams.json -duration 5 -agent udo -horizon 8 -heavy_horizon 3 -rl_max_delay_time 5 -default_query_time_out 18 129 | ``` 130 | 131 | # Citation 132 | 133 | If you use our repository, please cite the following papers. 134 | 135 | ``` 136 | @article{wang2021udo, 137 | title={UDO: universal database optimization using reinforcement learning}, 138 | author={Wang, Junxiong and Trummer, Immanuel and Basu, Debabrota}, 139 | journal={Proceedings of the VLDB Endowment}, 140 | volume={14}, 141 | number={13}, 142 | pages={3402--3414}, 143 | year={2021}, 144 | publisher={VLDB Endowment} 145 | } 146 | 147 | @inproceedings{wang2021demonstrating, 148 | title={Demonstrating UDO: A Unified Approach for Optimizing Transaction Code, Physical Design, and System Parameters via Reinforcement Learning}, 149 | author={Wang, Junxiong and Trummer, Immanuel and Basu, Debabrota}, 150 | booktitle={Proceedings of the 2021 International Conference on Management of Data}, 151 | pages={2794--2797}, 152 | year={2021} 153 | } 154 | ``` 155 | -------------------------------------------------------------------------------- /udo/drivers/mysqldriver.py: -------------------------------------------------------------------------------- 1 | # ----------------------------------------------------------------------- 2 | # Copyright (c) 2021 Cornell Database Group 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining 5 | # a copy of this software and associated documentation files (the 6 | # "Software"), to deal in the Software without restriction, including 7 | # without limitation the rights to use, copy, modify, merge, publish, 8 | # distribute, sublicense, and/or sell copies of the Software, and to 9 | # permit persons to whom the Software is furnished to do so, subject to 10 | # the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be 13 | # included in all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 17 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT 18 | # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 19 | # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 20 | # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 21 | # OTHER DEALINGS IN THE SOFTWARE. 22 | # ----------------------------------------------------------------------- 23 | 24 | import json 25 | import re 26 | import time 27 | 28 | import MySQLdb 29 | 30 | from .abstractdriver import * 31 | 32 | 33 | class MysqlDriver(AbstractDriver): 34 | """the DBMS driver for MySQL""" 35 | 36 | def __init__(self, conf, sys_params): 37 | super(MysqlDriver, self).__init__("mysql", conf, sys_params) 38 | 39 | def connect(self): 40 | """connect to a database""" 41 | self.conn = MySQLdb.connect(self.config['host'], self.config['user'], self.config['passwd'], self.config['db']) 42 | self.cursor = self.conn.cursor() 43 | self.index_creation_format = "CREATE INDEX %s ON %s (%s) USING BTREE;" 44 | self.index_drop_format = "ALTER TABLE %s drop index %s;" 45 | # self.is_cluster = True 46 | self.sys_params_type = len(self.sys_params) 47 | self.sys_params_space = [len(specific_parameter) for specific_parameter in self.sys_params] 48 | self.retrieve_table_name_sql = "show tables;" 49 | self.cardinality_format = "select count(*) from %s;" 50 | self.innodb_buffer_pattern = re.compile("set global innodb_buffer_pool_size = (.*);") 51 | self.select_innodb_buffer_value_sql = "select @@innodb_buffer_pool_size;" 52 | 53 | def cardinalities(self): 54 | """get cardinality of the connected database""" 55 | self.cursor.execute(self.retrieve_table_name_sql) 56 | dbms_tables = [] 57 | cardinality_info = {} 58 | for (table_name,) in self.cursor: 59 | dbms_tables.append(table_name) 60 | for table in dbms_tables: 61 | self.cursor.execute(self.cardinality_format % table) 62 | result = self.cursor.fetchone() 63 | cardinality = result[0] 64 | cardinality_info[table] = cardinality 65 | return cardinality_info 66 | 67 | def run_queries_with_timeout(self, query_list, timeout): 68 | """run queries with a timeout""" 69 | run_time = [] 70 | for query_sql, current_timeout in zip(query_list, timeout): 71 | try: 72 | current_timeout = current_timeout 73 | self.cursor.execute("SET SESSION MAX_EXECUTION_TIME=%d" % (current_timeout * 1000)) 74 | start_time = time.time() 75 | self.cursor.execute(query_sql) 76 | finish_time = time.time() 77 | duration = finish_time - start_time 78 | print("query run:%s" % duration) 79 | except MySQLdb.OperationalError as oe: 80 | # print(oe) 81 | # print("timeout") 82 | duration = current_timeout 83 | run_time.append(duration) 84 | print(run_time) 85 | # reset back to default configuraiton 86 | try: 87 | self.cursor.execute("SET SESSION MAX_EXECUTION_TIME=0") 88 | self.cursor.execute("drop view if exists REVENUE0;") 89 | except MySQLdb.OperationalError as oe: 90 | print(oe) 91 | return run_time 92 | 93 | def run_queries_without_timeout(self, query_list): 94 | """run queries without timeout""" 95 | return self.run_queries_with_timeout(query_list, [0 for query in query_list]) 96 | 97 | def build_index(self, index_to_create): 98 | """build index""" 99 | index_sql = self.index_creation_format % (index_to_create[0], index_to_create[1], index_to_create[2]) 100 | logging.debug(f"create index {index_sql}") 101 | self.cursor.execute(index_sql) 102 | self.conn.commit() 103 | 104 | def build_index_command(self, index_to_create): 105 | """build index command""" 106 | index_sql = self.index_creation_format % (index_to_create[0], index_to_create[1], index_to_create[2]) 107 | return index_sql 108 | 109 | def drop_index(self, index_to_drop): 110 | """drop index""" 111 | index_sql = self.index_drop_format % (index_to_drop[1], index_to_drop[0]) 112 | logging.debug(f"drop index {index_sql}") 113 | self.cursor.execute(index_sql) 114 | self.conn.commit() 115 | 116 | def change_system_parameter(self, parameter_choices): 117 | """change system parameter values using the input parameter choices""" 118 | for i in range(self.sys_params_type): 119 | parameter_choice = int(parameter_choices[i]) 120 | parameter_change_sql = self.sys_params[i][parameter_choice] 121 | self.set_system_parameter(parameter_change_sql) 122 | if "innodb_buffer_pool_size" in parameter_change_sql: 123 | # wait finish until it get effective 124 | value_to_change = int(self.innodb_buffer_pattern.findall(parameter_change_sql)[0]) 125 | # verify whether the value is changed or not 126 | while True: 127 | self.cursor.execute(self.select_innodb_buffer_value_sql) 128 | current_value = (self.cursor.fetchall())[0][0] 129 | if value_to_change == current_value: 130 | break 131 | else: 132 | time.sleep(3) 133 | self.set_system_parameter(parameter_change_sql) 134 | 135 | def set_system_parameter(self, parameter_sql): 136 | """switch system parameters""" 137 | try: 138 | logging.debug(f"set system parameter {parameter_sql}") 139 | self.cursor.execute(parameter_sql) 140 | self.conn.commit() 141 | except MySQLdb.OperationalError as oe: 142 | print(oe) 143 | 144 | def analyze_queries_cost(self, query_sqls): 145 | """analyze cost of queries""" 146 | total_cost = [] 147 | for analyze_sql in query_sqls: 148 | analyze_sql = f"EXPLAIN FORMAT=JSON {analyze_sql}" 149 | self.cursor.execute(analyze_sql) 150 | plan = json.loads((self.cursor.fetchone())[0]) 151 | total_cost.append(float(plan['query_block']['cost_info']['query_cost'])) 152 | logging.info(f"estimate costs {total_cost}") 153 | return total_cost 154 | 155 | def close(self): 156 | """close the connection""" 157 | self.cursor.close() 158 | self.conn.close() 159 | 160 | ## CLASS 161 | -------------------------------------------------------------------------------- /udo/drivers/postgresdriver.py: -------------------------------------------------------------------------------- 1 | # ----------------------------------------------------------------------- 2 | # Copyright (c) 2021 Cornell Database Group 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining 5 | # a copy of this software and associated documentation files (the 6 | # "Software"), to deal in the Software without restriction, including 7 | # without limitation the rights to use, copy, modify, merge, publish, 8 | # distribute, sublicense, and/or sell copies of the Software, and to 9 | # permit persons to whom the Software is furnished to do so, subject to 10 | # the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be 13 | # included in all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 17 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT 18 | # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 19 | # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 20 | # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 21 | # OTHER DEALINGS IN THE SOFTWARE. 22 | # ----------------------------------------------------------------------- 23 | 24 | import logging 25 | import time 26 | 27 | import psycopg2 28 | from psycopg2._psycopg import InternalError 29 | from psycopg2._psycopg import QueryCanceledError 30 | 31 | from .abstractdriver import * 32 | 33 | class PostgresDriver(AbstractDriver): 34 | """the DBMS driver for Postgres""" 35 | 36 | def __init__(self, conf, sys_params): 37 | super(PostgresDriver, self).__init__("postgres", conf, sys_params) 38 | 39 | def connect(self): 40 | """connect to a database""" 41 | self.conn = psycopg2.connect("dbname='%s' user='%s'" % (self.config["db"], self.config["user"])) 42 | self.conn.autocommit = True 43 | self.cursor = self.conn.cursor() 44 | self.index_creation_format = "CREATE INDEX %s ON %s (%s);" 45 | self.index_drop_format = "drop index %s;" 46 | self.is_cluster = False 47 | self.retrieve_table_name_sql = "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' ORDER BY table_name;" 48 | self.cardinality_format = "select count(*) from %s;" 49 | self.cluster_indices_format = "CLUSTER %s ON %s;" 50 | self.enable_index_format = "update pg_index set indisvalid = true where indexrelid = '%s'::regclass;" 51 | self.disable_index_format = "update pg_index set indisvalid = false where indexrelid = '%s'::regclass;" 52 | 53 | def cardinalities(self): 54 | """get cardinality of the connected database""" 55 | self.cursor.execute(self.retrieve_table_name_sql) 56 | dbms_tables = [] 57 | cardinality_info = {} 58 | for table in self.cursor.fetchall(): 59 | dbms_tables.append(table) 60 | for table in dbms_tables: 61 | self.cursor.execute(self.cardinality_format % table) 62 | result = self.cursor.fetchone() 63 | cardinality = result[0] 64 | cardinality_info[table[0].lower()] = cardinality 65 | return cardinality_info 66 | 67 | def run_queries_with_timeout(self, query_list, timeout): 68 | """run queries with a timeout""" 69 | run_time = [] 70 | for query_sql, current_timeout in zip(query_list, timeout): 71 | try: 72 | # logging.debug(f"query sql: {query_sql}") 73 | logging.debug(f"current timeout: {current_timeout}") 74 | self.cursor.execute("set statement_timeout = %d" % (current_timeout * 1000)) 75 | start_time = time.time() 76 | self.cursor.execute(query_sql) 77 | finish_time = time.time() 78 | duration = finish_time - start_time 79 | except QueryCanceledError: 80 | duration = current_timeout 81 | except InternalError: 82 | # error to run the query, set duration to a large number 83 | logging.debug(f"Internal Error for query {query_sql}") 84 | duration = current_timeout * 1000 85 | logging.debug(f"duration: {duration}") 86 | run_time.append(duration) 87 | # reset the timeout to the default configuration 88 | self.cursor.execute("set statement_timeout=0;") 89 | self.cursor.execute("drop view if exists REVENUE0;") 90 | return run_time 91 | 92 | def run_queries_with_total_timeout(self, query_list, timeout): 93 | """run queries with a timeout""" 94 | run_time = [] 95 | current_timeout = timeout 96 | total_runtime = 0 97 | for query_sql in query_list: 98 | try: 99 | logging.debug(f"current timeout: {current_timeout}") 100 | self.cursor.execute("set statement_timeout = %d" % (current_timeout * 1000)) 101 | start_time = time.time() 102 | self.cursor.execute(query_sql) 103 | finish_time = time.time() 104 | duration = finish_time - start_time 105 | except QueryCanceledError: 106 | total_runtime = timeout 107 | break 108 | except InternalError: 109 | # error to run the query, set duration to a large number 110 | logging.debug(f"Internal Error for query {query_sql}") 111 | total_runtime = timeout * 1000 112 | break 113 | logging.debug(f"duration: {duration}") 114 | run_time.append(duration) 115 | current_timeout = current_timeout - duration 116 | total_runtime += duration 117 | # reset the timeout to the default configuration 118 | self.cursor.execute("set statement_timeout=0;") 119 | self.cursor.execute("drop view if exists REVENUE0;") 120 | logging.debug(f"runtime {run_time}") 121 | return total_runtime 122 | 123 | def run_queries_without_timeout(self, query_list): 124 | """run queries without timeout""" 125 | return self.run_queries_with_timeout(query_list, [0] * len(query_list)) 126 | 127 | def build_index(self, index_to_create): 128 | """build index""" 129 | index_sql = self.index_creation_format % (index_to_create[0], index_to_create[1], index_to_create[2]) 130 | logging.debug(f"create index {index_sql}") 131 | self.cursor.execute(index_sql) 132 | # if we consider the cluster indices 133 | if self.is_cluster: 134 | self.cursor.execute(self.cluster_indices_format % (index_to_create[0], index_to_create[1])) 135 | # self.conn.commit() 136 | 137 | def build_index_command(self, index_to_create): 138 | """build index command""" 139 | index_sql = self.index_creation_format % (index_to_create[0], index_to_create[1], index_to_create[2]) 140 | if self.is_cluster: 141 | cluster_sql = self.cluster_indices_format % (index_to_create[0], index_to_create[1]) 142 | return f"{index_sql} \n {cluster_sql}" 143 | else: 144 | return index_sql 145 | 146 | def drop_index(self, index_to_drop): 147 | """drop index""" 148 | index_sql = self.index_drop_format % (index_to_drop[0]) 149 | logging.debug(f"drop index {index_sql}") 150 | self.cursor.execute(index_sql) 151 | # self.conn.commit() 152 | 153 | def set_system_parameter(self, parameter_sql): 154 | """switch system parameters""" 155 | logging.info(f"{parameter_sql}") 156 | self.cursor.execute(parameter_sql) 157 | # self.conn.commit() 158 | 159 | def analyze_queries_cost(self, query_sqls): 160 | """analyze cost of queries""" 161 | total_cost = [] 162 | for analyze_sql in query_sqls: 163 | analyze_sql = f"explain (format json) {analyze_sql}" 164 | self.cursor.execute(analyze_sql) 165 | total_cost.append((self.cursor.fetchone())[0][0]['Plan']['Total Cost']) 166 | logging.info(f"estimate costs {total_cost}") 167 | return total_cost 168 | 169 | def enable_indices(self, index_to_enable): 170 | """enable indices""" 171 | index_sql = self.enable_index_format % (index_to_enable[0]) 172 | logging.info(f"enable: {index_sql}") 173 | self.cursor.execute(index_sql) 174 | 175 | def disable_indices(self, index_to_disable): 176 | """disable indices""" 177 | index_sql = self.disable_index_format % (index_to_disable[0]) 178 | logging.info(f"disable: {index_sql}") 179 | self.cursor.execute(index_sql) 180 | 181 | def close(self): 182 | """close the connection""" 183 | self.cursor.close() 184 | self.conn.close() 185 | 186 | ## CLASS 187 | -------------------------------------------------------------------------------- /udo/__main__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # ----------------------------------------------------------------------- 3 | # Copyright (c) 2021 Cornell Database Group 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining 6 | # a copy of this software and associated documentation files (the 7 | # "Software"), to deal in the Software without restriction, including 8 | # without limitation the rights to use, copy, modify, merge, publish, 9 | # distribute, sublicense, and/or sell copies of the Software, and to 10 | # permit persons to whom the Software is furnished to do so, subject to 11 | # the following conditions: 12 | # 13 | # The above copyright notice and this permission notice shall be 14 | # included in all copies or substantial portions of the Software. 15 | # 16 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT 19 | # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 20 | # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 21 | # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 22 | # OTHER DEALINGS IN THE SOFTWARE. 23 | # ----------------------------------------------------------------------- 24 | 25 | import argparse 26 | import json 27 | import logging 28 | import os 29 | import random 30 | import sys 31 | 32 | import numpy as np 33 | 34 | from udo.agent.ddpg_agent import run_ddpg_agent 35 | from udo.agent.sarsa_agent import run_sarsa_agent 36 | from udo.agent.udo_agent import run_udo_agent 37 | from udo.agent.udo_simplifed_agent import run_simplifed_udo_agent 38 | from udo.drivers.mysqldriver import MysqlDriver 39 | from udo.drivers.postgresdriver import PostgresDriver 40 | 41 | if __name__ == "__main__": 42 | logging.basicConfig(level=logging.INFO, 43 | format="%(asctime)s [%(funcName)s:%(lineno)03d] %(levelname)-5s: %(message)s", 44 | datefmt="%m-%d-%Y %H:%M:%S", 45 | stream=sys.stdout) 46 | 47 | udo_parser = argparse.ArgumentParser(description='UDO optimizer.') 48 | # database setting 49 | udo_parser.add_argument('-system', choices=('mysql', 'postgres'), 50 | help='Target system driver') 51 | udo_parser.add_argument('-db', default="tpch", 52 | help='the database to optimizes') 53 | udo_parser.add_argument('-username', default="udo", 54 | help='username') 55 | udo_parser.add_argument('-password', default="udo", 56 | help='password') 57 | udo_parser.add_argument('-queries', default="tpch_query.sql", 58 | help='the input query file') 59 | udo_parser.add_argument('-indices', 60 | help='the input query file') 61 | udo_parser.add_argument('-sys_params', 62 | help='the input system params json file') 63 | # tuning time 64 | udo_parser.add_argument('-duration', default=5, type=float, 65 | help='time for tuning in hours') 66 | 67 | # rl algorithm 68 | udo_parser.add_argument('-agent', default='udo', choices=('udo', 'udo-s', 'ddpg', 'sarsa'), 69 | help='reinforcement learning agent') 70 | udo_parser.add_argument('-horizon', default=5, type=int, 71 | help='the number horizon for reinforcement agent') 72 | udo_parser.add_argument('-heavy_horizon', default=3, type=int, 73 | help='the number horizon for heavy parameters in UDO') 74 | # mcts algorithm 75 | udo_parser.add_argument('-rl_update', choices=('RAVE', 'MCTS'), default='RAVE', 76 | help='the update policy of UDO tree search') 77 | udo_parser.add_argument('-rl_select', choices=('UCB1', 'UCBV'), default='UCBV', 78 | help='the selection policy of UDO tree search') 79 | udo_parser.add_argument('-rl_reward', choices=('delta', 'accumulate'), 80 | help='the reward of reinforcement learning agent') 81 | udo_parser.add_argument('-rl_delay', choices=('UCB', 'Exp3'), 82 | help='the delay selection policy') 83 | udo_parser.add_argument('-rl_max_delay_time', type=int, default=5, 84 | help='the delay selection policy') 85 | 86 | # tuning configuration 87 | udo_parser.add_argument('-sample_rate', type=int, default="1", 88 | help='sampled rate from workload') 89 | udo_parser.add_argument('-default_query_time_out', type=int, default="5", 90 | help='default timeout in seconds for each query') 91 | udo_parser.add_argument('-time_out_ratio', type=float, default="1.1", 92 | help='timeout ratio respect to default time') 93 | 94 | # load json file 95 | udo_parser.add_argument('--load_json', 96 | help='Load settings from file in json format. Command line options override values in file.') 97 | 98 | args = udo_parser.parse_args() 99 | 100 | # json file the higher priority 101 | if args.load_json: 102 | with open(args.load_json, 'rt') as f: 103 | t_args = argparse.Namespace() 104 | t_args.__dict__.update(json.load(f)) 105 | args = udo_parser.parse_args(namespace=t_args) 106 | 107 | # change 108 | args = vars(args) 109 | 110 | # init queries 111 | if args['queries']: 112 | queries = dict() 113 | for file_name in sorted(os.listdir(args['queries'])): 114 | if file_name.endswith(".sql"): 115 | with open(os.path.join(args['queries'], file_name)) as f: 116 | content = f.read() 117 | queries[file_name] = content 118 | else: 119 | print("please specify queries") 120 | exit() 121 | # print(constants.QUERIES) 122 | 123 | # init indices with external files 124 | if args['indices']: 125 | # analyze queries to extract indexes 126 | with open(args['indices']) as f: 127 | content = f.readlines() 128 | indices = [x.strip() for x in content] 129 | candidate_indices = [] 130 | for index_str in indices: 131 | index_information = index_str.split(";") 132 | candidate_indices.append((index_information[0], index_information[1], index_information[2])) 133 | else: 134 | print("please specify candidate indices") 135 | exit() 136 | 137 | # check the validate of configurations file 138 | if not args['system'] or not args['db']: 139 | print("Please specific a database") 140 | exit() 141 | 142 | if not args['duration'] or not args['horizon']: 143 | print("Wrong parameters. Please check the input parameters") 144 | exit() 145 | 146 | if not args['sys_params']: 147 | print("Please specify the system parameters") 148 | exit() 149 | 150 | dbms_conf = { 151 | "host": "127.0.0.1", 152 | "db": args['db'], 153 | "user": args['username'], 154 | "passwd": args['password'], 155 | } 156 | 157 | with open(args['sys_params'], 'rt') as f: 158 | sys_params = json.load(f) 159 | 160 | # create a dbms driver 161 | driver = None 162 | if args['system'] == "mysql": 163 | driver = MysqlDriver(dbms_conf, sys_params) 164 | elif args['system'] == "postgres": 165 | driver = PostgresDriver(dbms_conf, sys_params) 166 | 167 | # obtain index cardinality information 168 | driver.connect() 169 | cardinality_info = driver.cardinalities() 170 | 171 | # obtain index applicable queries 172 | for i in range(len(candidate_indices)): 173 | contain_query = [] 174 | for query_id, query_str in queries.items(): 175 | # print(candidate_indices[i][2]) 176 | if "where" in query_str: 177 | where_clause = query_str[query_str.index("where"):].lower() 178 | else: 179 | where_clause = query_str.lower() 180 | contain_columns = candidate_indices[i][2].lower().split(",") 181 | if all(contain_column in where_clause for contain_column in contain_columns): 182 | # if any(contain_column in where_clause for contain_column in contain_columns): 183 | contain_query.append(query_id) 184 | index_cardinality = cardinality_info[candidate_indices[i][1].lower()] 185 | candidate_indices[i] += (contain_query, index_cardinality,) 186 | 187 | # filter indices which contains at least has one appliable query 188 | candidate_indices = [candidate_index for candidate_index in candidate_indices if len(candidate_index[3]) > 0] 189 | 190 | # print(len(index.candidate_indices)) 191 | 192 | tuning_config = { 193 | "duration": args['duration'], 194 | "horizon": args['horizon'], 195 | "sample_rate": args['sample_rate'], 196 | "default_query_time_out": args['default_query_time_out'], 197 | "time_out_ratio": args['time_out_ratio'] 198 | } 199 | 200 | # if the agent is udo 201 | if args['agent'] == 'udo': 202 | if not args['heavy_horizon']: 203 | print("Please specific the step of heavy configurations") 204 | exit() 205 | # run udo 206 | print(tuning_config) 207 | horizon = args['horizon'] 208 | tuning_config['heavy_horizon'] = args['heavy_horizon'] 209 | tuning_config['light_horizon'] = horizon - int(args['heavy_horizon']) 210 | tuning_config['rl_max_delay_time'] = args['rl_max_delay_time'] 211 | tuning_config['rl_update'] = args['rl_update'] 212 | tuning_config['rl_select'] = args['rl_select'] 213 | run_udo_agent(driver=driver, queries=queries, candidate_indices=candidate_indices, tuning_config=tuning_config) 214 | elif args['agent'] == 'udo-s': 215 | # run simplified udo 216 | run_simplifed_udo_agent(driver=driver, queries=queries, candidate_indices=candidate_indices, 217 | tuning_config=tuning_config) 218 | elif args['agent'] == 'ddpg': 219 | # run ddpg deep rl 220 | run_ddpg_agent(driver=driver, queries=queries, candidate_indices=candidate_indices, tuning_config=tuning_config) 221 | elif args['agent'] == 'sarsa': 222 | # run sarsa deep rl 223 | run_sarsa_agent(driver=driver, queries=queries, candidate_indices=candidate_indices, 224 | tuning_config=tuning_config) 225 | -------------------------------------------------------------------------------- /udo/mcts/uct_node.py: -------------------------------------------------------------------------------- 1 | # ----------------------------------------------------------------------- 2 | # Copyright (c) 2021 Cornell Database Group 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining 5 | # a copy of this software and associated documentation files (the 6 | # "Software"), to deal in the Software without restriction, including 7 | # without limitation the rights to use, copy, modify, merge, publish, 8 | # distribute, sublicense, and/or sell copies of the Software, and to 9 | # permit persons to whom the Software is furnished to do so, subject to 10 | # the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be 13 | # included in all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 17 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT 18 | # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 19 | # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 20 | # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 21 | # OTHER DEALINGS IN THE SOFTWARE. 22 | # ----------------------------------------------------------------------- 23 | 24 | import logging 25 | import math 26 | import random 27 | 28 | from .mcts_node import * 29 | 30 | 31 | class uct_node(mcts_node): 32 | def __init__(self, round, tree_level, tree_height, state, env, space_type=SpaceType.All, 33 | selection_policy=SelectionPolicy.UCBV, update_policy=UpdatePolicy.RAVE, terminate_action=None, 34 | ucb1_er=1.2, ucbv_bound=0.1, ucbv_const=0.1): 35 | self.create_in = round 36 | # construct the transaction space, index space and parameter space 37 | self.env = env 38 | self.state = state 39 | if space_type is SpaceType.Light: 40 | self.actions = env.retrieve_light_actions(state) 41 | elif space_type is SpaceType.Heavy: 42 | self.actions = env.retrieve_heavy_actions(state) 43 | else: 44 | self.actions = env.retrieve_actions(state) 45 | self.nr_action = len(self.actions) 46 | self.tree_level = tree_level 47 | # for the children node 48 | self.children = [None] * self.nr_action 49 | # for the number of tries of children node 50 | self.nr_tries = [0] * self.nr_action 51 | self.first_moment = [0] * self.nr_action 52 | self.second_moment = [0] * self.nr_action 53 | self.total_visit = 0 54 | self.priority_actions = self.actions.copy() 55 | self.tree_height = tree_height 56 | self.terminate_action = terminate_action 57 | # update policy 58 | self.selection_policy = selection_policy 59 | self.update_policy = update_policy 60 | self.space_type = space_type 61 | # tuning parameters for UCB1 62 | self.explore_rate = ucb1_er 63 | # tuning parameters for UCBV 64 | self.bound = ucbv_bound 65 | self.const = ucbv_const 66 | 67 | def select_action(self): 68 | """select action according to rl policy""" 69 | if len(self.priority_actions) > 0: 70 | selected_action = random.choice(self.priority_actions) 71 | self.priority_actions.remove(selected_action) 72 | selected_action_idx = self.actions.index(selected_action) 73 | return selected_action, selected_action_idx 74 | else: 75 | # select actions according to ucb1 76 | best_action_idx = -1 77 | best_ucb_score = -1 78 | for action_idx in range(self.nr_action): 79 | if self.nr_tries[action_idx] > 0: 80 | mean = self.first_moment[action_idx] / self.nr_tries[action_idx] 81 | variance = max((self.second_moment[action_idx] / self.nr_tries[action_idx] - (mean * mean)), 0) 82 | if self.selection_policy == SelectionPolicy.UCB1: 83 | ucb_score = mean + self.explore_rate * math.sqrt( 84 | math.log(self.total_visit) / self.nr_tries[action_idx]) 85 | elif self.selection_policy == SelectionPolicy.UCBV: 86 | ucb_score = mean + math.sqrt( 87 | self.const * variance * math.log(self.total_visit) / self.nr_tries[action_idx]) + ( 88 | 3 * self.bound * math.log(self.total_visit) / self.nr_tries[action_idx]) 89 | else: 90 | raise ValueError('Selection policy is unavailable.') 91 | if ucb_score > best_ucb_score: 92 | best_ucb_score = ucb_score 93 | best_action_idx = action_idx 94 | return self.actions[best_action_idx], best_action_idx 95 | 96 | def playout(self, current_level_action): 97 | """randomly select other actions until reaching to terminate state""" 98 | current_level_state = self.state 99 | selected_action_path = [] 100 | for current_level in range(self.tree_level + 1, self.tree_height): 101 | current_level_state, current_state_id = self.env.transition(current_level_state, 102 | current_level_action) 103 | if self.space_type is SpaceType.Light: 104 | current_candidate_actions = self.env.retrieve_light_actions(current_level_state) 105 | elif self.space_type is SpaceType.Heavy: 106 | current_candidate_actions = self.env.retrieve_heavy_actions(current_level_state) 107 | elif self.space_type == SpaceType.All: 108 | current_candidate_actions = self.env.retrieve_actions(current_level_state) 109 | # choose one random action from the action space 110 | if self.terminate_action is not None: 111 | current_candidate_actions.append(self.terminate_action) 112 | current_level_action = random.choice(current_candidate_actions) 113 | selected_action_path.append(current_level_action) 114 | return selected_action_path 115 | 116 | def sample(self, round): 117 | """sample actions from search tree""" 118 | if self.nr_action == 0: 119 | return [] 120 | else: 121 | # inner node 122 | selected_action, selected_action_idx = self.select_action() 123 | if self.terminate_action is not None and selected_action == self.terminate_action: 124 | # for the terminate action, stop expansion 125 | return [selected_action] 126 | can_expand = (self.create_in != round) and (self.tree_level < self.tree_height) 127 | if can_expand and not self.children[selected_action_idx]: 128 | state, state_idx = self.env.transition(self.state, selected_action) 129 | self.children[selected_action_idx] = uct_node(round=0, tree_level=self.tree_level + 1, 130 | tree_height=self.tree_height, state=state, env=self.env, 131 | space_type=self.space_type, 132 | selection_policy=self.selection_policy, 133 | update_policy=self.update_policy, 134 | terminate_action=self.terminate_action, 135 | ucb1_er=self.explore_rate, 136 | ucbv_bound=self.bound, 137 | ucbv_const=self.const) 138 | child = self.children[selected_action_idx] 139 | # recursively sample the tree 140 | if child: 141 | return [selected_action] + child.sample(round) 142 | else: 143 | return [selected_action] + self.playout(selected_action) 144 | 145 | def update_statistics_with_mcts_reward(self, reward, selected_actions): 146 | """update mcts statistics using reward information""" 147 | rewards = {selected_action: reward for selected_action in selected_actions} 148 | self.update_statistics_with_delta_reward(rewards, selected_actions) 149 | 150 | def update_statistics_with_delta_reward(self, rewards, selected_actions): 151 | """update mcts statistics using intermediate reward information""" 152 | if self.update_policy is UpdatePolicy.RAVE: 153 | for action_idx in range(self.nr_action): 154 | current_action = self.actions[action_idx] 155 | # we use the RAVE update policy 156 | if current_action in selected_actions: 157 | reward = rewards[current_action] 158 | self.total_visit += 1 159 | self.nr_tries[action_idx] += 1 160 | self.first_moment[action_idx] += reward 161 | self.second_moment[action_idx] += reward * reward 162 | if current_action in self.priority_actions: 163 | self.priority_actions.remove(current_action) 164 | # update the reward for subtree 165 | if self.children[action_idx] is not None: 166 | next_selected_actions = list(filter(lambda x: x != current_action, selected_actions)) 167 | self.children[action_idx].update_statistics_with_delta_reward(rewards, next_selected_actions) 168 | else: 169 | current_action = selected_actions[self.tree_level] 170 | for action_idx in range(self.nr_action): 171 | if self.actions[action_idx] == current_action: 172 | reward = rewards[current_action] 173 | self.total_visit += 1 174 | self.nr_tries[action_idx] += 1 175 | self.first_moment[action_idx] += reward 176 | self.second_moment[action_idx] += reward * reward 177 | # self.mean_reward[action_idx] = self.first_moment[action_idx] / self.nr_tries[action_idx] 178 | if self.children[action_idx] is not None: 179 | self.children[action_idx].update_statistics_with_delta_reward(rewards, selected_actions) 180 | 181 | def update_batch(self, update_infos): 182 | for (rewards, select_actions) in update_infos: 183 | rewards = {select_actions[i]: rewards[i] for i in range(len(rewards))} 184 | self.update_statistics_with_delta_reward(rewards, select_actions) 185 | 186 | def print_reward_info(self): 187 | """print the reward information""" 188 | mean_reward = [0] * self.nr_action 189 | for i in range(self.nr_action): 190 | if self.nr_tries[i] > 0: 191 | mean_reward[i] = self.first_moment[i] / self.nr_tries[i] 192 | else: 193 | mean_reward[i] = 0 194 | logging.debug(f"first layer avg reward {mean_reward}") 195 | return mean_reward 196 | 197 | def best_actions(self): 198 | """best actions from the search tree""" 199 | best_mean = 0 200 | best_action_idx = -1 201 | for action_idx in range(self.nr_action): 202 | if self.nr_tries[action_idx] > 0: 203 | mean = self.first_moment[action_idx] / self.nr_tries[action_idx] 204 | if mean > best_mean: 205 | best_mean = mean 206 | best_action_idx = action_idx 207 | if self.children[best_action_idx] is not None: 208 | return [self.actions[best_action_idx]] + self.children[best_action_idx].best_actions() 209 | else: 210 | return [self.actions[best_action_idx]] 211 | -------------------------------------------------------------------------------- /udo/agent/ddpg_agent.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import random 3 | import time 4 | 5 | import gym 6 | import udo_optimization 7 | 8 | import numpy as np 9 | import tensorflow as tf 10 | from tensorflow.keras import layers 11 | 12 | 13 | def run_ddpg_agent(driver, queries, candidate_indices, tuning_config): 14 | """Run DDPG agent for universal optimization""" 15 | env = gym.make("udo_optimization-v0", driver=driver, queries=queries, candidate_indices=candidate_indices, 16 | config=tuning_config) 17 | 18 | num_states = env.observation_space.shape[0] 19 | num_actions = env.nA_index 20 | logging.debug(f"Size of State Space -> {num_states}") 21 | logging.debug(f"Size of Action Space -> {num_actions}") 22 | 23 | upper_bound = 1 24 | lower_bound = -1 25 | 26 | logging.debug(f"Max Value of Action -> {upper_bound}") 27 | logging.debug(f"Min Value of Action -> {lower_bound}") 28 | 29 | # the following code is from Keras RL example, https://github.com/keras-team/keras-io/blob/master/examples/rl/ddpg_pendulum.py 30 | class OUActionNoise: 31 | def __init__(self, mean, std_deviation, theta=0.15, dt=1e-2, x_initial=None): 32 | self.theta = theta 33 | self.mean = mean 34 | self.std_dev = std_deviation 35 | self.dt = dt 36 | self.x_initial = x_initial 37 | self.reset() 38 | 39 | def __call__(self): 40 | # Formula taken from https://www.wikipedia.org/wiki/Ornstein-Uhlenbeck_process. 41 | x = ( 42 | self.x_prev 43 | + self.theta * (self.mean - self.x_prev) * self.dt 44 | + self.std_dev * np.sqrt(self.dt) * np.random.normal(size=self.mean.shape) 45 | ) 46 | # Store x into x_prev 47 | # Makes next noise dependent on current one 48 | self.x_prev = x 49 | return x 50 | 51 | def reset(self): 52 | if self.x_initial is not None: 53 | self.x_prev = self.x_initial 54 | else: 55 | self.x_prev = np.zeros_like(self.mean) 56 | 57 | class Buffer: 58 | def __init__(self, buffer_capacity=100000, batch_size=64): 59 | # Number of "experiences" to store at max 60 | self.buffer_capacity = buffer_capacity 61 | # Num of tuples to train on. 62 | self.batch_size = batch_size 63 | 64 | # Its tells us num of times record() was called. 65 | self.buffer_counter = 0 66 | 67 | # Instead of list of tuples as the exp.replay concept go 68 | # We use different np.arrays for each tuple element 69 | self.state_buffer = np.zeros((self.buffer_capacity, num_states)) 70 | self.action_buffer = np.zeros((self.buffer_capacity, num_actions)) 71 | self.reward_buffer = np.zeros((self.buffer_capacity, 1)) 72 | self.next_state_buffer = np.zeros((self.buffer_capacity, num_states)) 73 | 74 | # Takes (s,a,r,s') obervation tuple as input 75 | def record(self, obs_tuple): 76 | # Set index to zero if buffer_capacity is exceeded, 77 | # replacing old records 78 | index = self.buffer_counter % self.buffer_capacity 79 | 80 | self.state_buffer[index] = obs_tuple[0] 81 | self.action_buffer[index] = obs_tuple[1] 82 | self.reward_buffer[index] = obs_tuple[2] 83 | self.next_state_buffer[index] = obs_tuple[3] 84 | 85 | self.buffer_counter += 1 86 | 87 | # Eager execution is turned on by default in TensorFlow 2. Decorating with tf.function allows 88 | # TensorFlow to build a static graph out of the logic and computations in our function. 89 | # This provides a large speed up for blocks of code that contain many small TensorFlow operations such as this one. 90 | @tf.function 91 | def update( 92 | self, state_batch, action_batch, reward_batch, next_state_batch, 93 | ): 94 | # Training and updating Actor & Critic networks. 95 | # See Pseudo Code. 96 | with tf.GradientTape() as tape: 97 | target_actions = target_actor(next_state_batch, training=True) 98 | y = reward_batch + gamma * target_critic( 99 | [next_state_batch, target_actions], training=True 100 | ) 101 | critic_value = critic_model([state_batch, action_batch], training=True) 102 | critic_loss = tf.math.reduce_mean(tf.math.square(y - critic_value)) 103 | 104 | critic_grad = tape.gradient(critic_loss, critic_model.trainable_variables) 105 | critic_optimizer.apply_gradients( 106 | zip(critic_grad, critic_model.trainable_variables) 107 | ) 108 | 109 | with tf.GradientTape() as tape: 110 | actions = actor_model(state_batch, training=True) 111 | critic_value = critic_model([state_batch, actions], training=True) 112 | # Used `-value` as we want to maximize the value given 113 | # by the critic for our actions 114 | actor_loss = -tf.math.reduce_mean(critic_value) 115 | 116 | actor_grad = tape.gradient(actor_loss, actor_model.trainable_variables) 117 | actor_optimizer.apply_gradients( 118 | zip(actor_grad, actor_model.trainable_variables) 119 | ) 120 | 121 | # We compute the loss and update parameters 122 | def learn(self): 123 | # Get sampling range 124 | record_range = min(self.buffer_counter, self.buffer_capacity) 125 | # Randomly sample indices 126 | batch_indices = np.random.choice(record_range, self.batch_size) 127 | 128 | # Convert to tensors 129 | state_batch = tf.convert_to_tensor(self.state_buffer[batch_indices]) 130 | action_batch = tf.convert_to_tensor(self.action_buffer[batch_indices]) 131 | reward_batch = tf.convert_to_tensor(self.reward_buffer[batch_indices]) 132 | reward_batch = tf.cast(reward_batch, dtype=tf.float32) 133 | next_state_batch = tf.convert_to_tensor(self.next_state_buffer[batch_indices]) 134 | 135 | self.update(state_batch, action_batch, reward_batch, next_state_batch) 136 | 137 | # This update target parameters slowly 138 | # Based on rate `tau`, which is much less than one. 139 | @tf.function 140 | def update_target(target_weights, weights, tau): 141 | for (a, b) in zip(target_weights, weights): 142 | a.assign(b * tau + a * (1 - tau)) 143 | 144 | def get_actor(): 145 | # Initialize weights between -3e-3 and 3-e3 146 | last_init = tf.random_uniform_initializer(minval=-0.003, maxval=0.003) 147 | 148 | inputs = layers.Input(shape=(num_states,)) 149 | out = layers.Dense(256, activation="relu")(inputs) 150 | out = layers.Dense(256, activation="relu")(out) 151 | outputs = layers.Dense(num_actions, activation="tanh", kernel_initializer=last_init)(out) 152 | 153 | # Our upper bound is 2.0 for Pendulum. 154 | outputs = outputs * upper_bound 155 | model = tf.keras.Model(inputs, outputs) 156 | return model 157 | 158 | def get_critic(): 159 | # State as input 160 | state_input = layers.Input(shape=(num_states)) 161 | state_out = layers.Dense(16, activation="relu")(state_input) 162 | state_out = layers.Dense(32, activation="relu")(state_out) 163 | 164 | # Action as input 165 | action_input = layers.Input(shape=(num_actions)) 166 | action_out = layers.Dense(32, activation="relu")(action_input) 167 | 168 | # Both are passed through seperate layer before concatenating 169 | concat = layers.Concatenate()([state_out, action_out]) 170 | 171 | out = layers.Dense(256, activation="relu")(concat) 172 | out = layers.Dense(256, activation="relu")(out) 173 | outputs = layers.Dense(1)(out) 174 | 175 | # Outputs single value for give state-action 176 | model = tf.keras.Model([state_input, action_input], outputs) 177 | 178 | return model 179 | 180 | def policy(state, noise_object, previous_actions): 181 | sampled_actions = tf.squeeze(actor_model(state)) 182 | noise = noise_object() 183 | # Adding noise to action 184 | sampled_actions = sampled_actions.numpy() + noise 185 | logging.debug(f"sampled_actions: {sampled_actions}") 186 | 187 | # We make sure action is within bounds 188 | legal_action = np.clip(sampled_actions, lower_bound, upper_bound) 189 | 190 | action_weights = [np.squeeze(legal_action)] 191 | print(action_weights) 192 | max_action_weight = -1000 193 | nr_action = len(action_weights[0]) 194 | start_action = random.randrange(0, nr_action) 195 | action = 0 196 | logging.debug(f"start_action {start_action}") 197 | # pick up the action which has the highest weight 198 | for action_idx in range(0, nr_action): 199 | current_action = (start_action + action_idx) % nr_action 200 | if (current_action not in previous_actions) and (action_weights[0][current_action] > max_action_weight): 201 | max_action_weight = action_weights[0][current_action] 202 | action = current_action 203 | 204 | # action = np.random.choice(np.flatnonzero(action_weights == max(action_weights))) 205 | # action = np.argmax(action_weights) 206 | return action 207 | 208 | std_dev = 0.2 209 | ou_noise = OUActionNoise(mean=np.zeros(1), std_deviation=float(std_dev) * np.ones(1)) 210 | 211 | actor_model = get_actor() 212 | critic_model = get_critic() 213 | 214 | target_actor = get_actor() 215 | target_critic = get_critic() 216 | 217 | # Making the weights equal initially 218 | target_actor.set_weights(actor_model.get_weights()) 219 | target_critic.set_weights(critic_model.get_weights()) 220 | 221 | # Learning rate for actor-critic models 222 | critic_lr = 0.002 223 | actor_lr = 0.001 224 | 225 | critic_optimizer = tf.keras.optimizers.Adam(critic_lr) 226 | actor_optimizer = tf.keras.optimizers.Adam(actor_lr) 227 | 228 | # Discount factor for future rewards 229 | gamma = 0.99 230 | # Used to update target networks 231 | tau = 0.005 232 | 233 | buffer = Buffer(50000, 64) 234 | 235 | # set the horizon 236 | env.horizon = tuning_config['horizon'] 237 | duration_in_seconds = tuning_config['duration'] * 3600 238 | 239 | # To store reward history of each episode 240 | ep_reward_list = [] 241 | # To store average reward history of last few episodes 242 | avg_reward_list = [] 243 | 244 | start_time = time.time() 245 | current_time = time.time() 246 | 247 | ep = 0 248 | while (current_time - start_time) < duration_in_seconds: 249 | # start a new episode 250 | prev_state = env.reset() 251 | episodic_reward = 0 252 | previous_actions = [] 253 | while True: 254 | # Uncomment this to see the Actor in action 255 | tf_prev_state = tf.expand_dims(tf.convert_to_tensor(prev_state), 0) 256 | # select action according to policy 257 | action = policy(tf_prev_state, ou_noise, previous_actions) 258 | previous_actions.append(action) 259 | # receive state and reward from environment. 260 | state, reward, done, info = env.step(action) 261 | 262 | buffer.record((prev_state, action, reward, state)) 263 | episodic_reward += reward 264 | 265 | ep += 1 266 | 267 | buffer.learn() 268 | update_target(target_actor.variables, actor_model.variables, tau) 269 | update_target(target_critic.variables, critic_model.variables, tau) 270 | 271 | # End this episode when `done` is True 272 | if done: 273 | current_time = time.time() 274 | break 275 | 276 | prev_state = state 277 | 278 | ep_reward_list.append(episodic_reward) 279 | 280 | # Mean of last 40 episodes 281 | avg_reward = np.mean(ep_reward_list[-40:]) 282 | logging.info(f"Episode * {ep} * Avg Reward is ==> {avg_reward}") 283 | avg_reward_list.append(avg_reward) 284 | logging.info(f"episode: {ep}") 285 | logging.info(f"evaluate duration: {(current_time - start_time)}") 286 | env.print_state_summary(env.best_state) 287 | -------------------------------------------------------------------------------- /udo/agent/udo_agent.py: -------------------------------------------------------------------------------- 1 | # ----------------------------------------------------------------------- 2 | # Copyright (c) 2021 Cornell Database Group 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining 5 | # a copy of this software and associated documentation files (the 6 | # "Software"), to deal in the Software without restriction, including 7 | # without limitation the rights to use, copy, modify, merge, publish, 8 | # distribute, sublicense, and/or sell copies of the Software, and to 9 | # permit persons to whom the Software is furnished to do so, subject to 10 | # the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be 13 | # included in all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 17 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT 18 | # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 19 | # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 20 | # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 21 | # OTHER DEALINGS IN THE SOFTWARE. 22 | # ----------------------------------------------------------------------- 23 | 24 | import logging 25 | import math 26 | import random 27 | import time 28 | 29 | import gym 30 | 31 | from udo.mcts.mcts_node import SpaceType 32 | from udo.mcts.uct_node import uct_node 33 | from udo.optimizer import order_optimizer 34 | 35 | 36 | def run_udo_agent(driver, queries, candidate_indices, tuning_config): 37 | duration_in_seconds = tuning_config['duration'] * 3600 38 | 39 | env = gym.make('udo_optimization-v0', driver=driver, queries=queries, candidate_indices=candidate_indices, 40 | config=tuning_config) 41 | 42 | nr_query = len(queries) # number of queries 43 | index_card_info = list(map(lambda x: x[4], candidate_indices)) 44 | index_to_applicable_queries = env.index_to_applicable_queries 45 | 46 | start_tune_time = time.time() 47 | logging.debug(f"start time: {start_tune_time}") 48 | 49 | optimizer = order_optimizer.OrderOptimizer(index_card_info) 50 | 51 | # number of indices equal heavy_tree_height + 1 52 | heavy_tree_height = tuning_config['heavy_horizon'] 53 | light_tree_height = tuning_config['light_horizon'] 54 | max_delay_time = tuning_config['rl_max_delay_time'] 55 | 56 | init_state = env.state_decoder(0) 57 | micro_episode = 5 58 | 59 | terminate_action = env.index_candidate_num 60 | light_tree_cache = dict() 61 | query_evaluate_sample_rate = tuning_config['sample_rate'] 62 | # reset the environment 63 | env.reset() 64 | default_runtime = env.default_runtime 65 | 66 | heavy_root = uct_node(round=0, tree_level=0, tree_height=heavy_tree_height, state=init_state, env=env, 67 | space_type=SpaceType.Heavy) 68 | idx_build_time = 0 69 | best_simulation_time = sum(default_runtime) 70 | best_configs = {} 71 | t1 = 1 72 | start_time = time.time() 73 | end_episode_time = time.time() 74 | while (end_episode_time - start_time) < duration_in_seconds: 75 | selected_heavy_action_batch = [] 76 | configuration_to_evaluate = [] 77 | # remove_terminate_action_batch = [] 78 | logging.debug(f"delay_time: {max_delay_time}") 79 | for d in range(max_delay_time): 80 | selected_heavy_actions = heavy_root.sample(t1 + d) 81 | selected_heavy_action_batch.append(selected_heavy_actions) 82 | # add all intermediate steps 83 | remove_terminate_heavy_actions = [heavy_action for heavy_action in selected_heavy_actions if 84 | heavy_action != terminate_action] 85 | for i in range(len(remove_terminate_heavy_actions)): 86 | current_actions = frozenset(remove_terminate_heavy_actions[:i + 1]) 87 | if current_actions not in configuration_to_evaluate: 88 | configuration_to_evaluate.append(current_actions) 89 | # show those actions 90 | logging.debug(f"selected heavy actions: {selected_heavy_action_batch}") 91 | logging.debug(f"evaluate configurations: {configuration_to_evaluate}") 92 | # evaluated_order = range(0, len(remove_terminate_action_batch)) 93 | evaluated_order = optimizer.greedy_min_cost_order(configuration_to_evaluate) 94 | logging.debug(f"evaluated order: {evaluated_order}") 95 | macro_performance_info = dict() 96 | for current_order_idx in range(len(evaluated_order)): 97 | current_order = evaluated_order[current_order_idx] 98 | selected_heavy_action_frozen = frozenset(configuration_to_evaluate[current_order]) 99 | # generate add action 100 | add_action = set(selected_heavy_action_frozen) 101 | drop_action = set() 102 | if current_order_idx > 0: 103 | # take the previous created indices 104 | previous_set = set(configuration_to_evaluate[evaluated_order[current_order_idx - 1]]) 105 | if t1 > 1 or current_order_idx > 0: 106 | add_action = add_action - previous_set 107 | drop_action = previous_set - set(selected_heavy_action_frozen) 108 | logging.debug(f"invoke action: {add_action}") 109 | logging.debug(f"drop action: {drop_action}") 110 | # build the indices 111 | time_start = time.time() 112 | env.index_step(add_action, drop_action) 113 | time_end = time.time() 114 | idx_build_time += (time_end - time_start) 115 | 116 | logging.debug(f"selected_heavy_action_frozen: {selected_heavy_action_frozen}") 117 | if selected_heavy_action_frozen in light_tree_cache: 118 | light_root = light_tree_cache[selected_heavy_action_frozen] 119 | else: 120 | light_root = uct_node(round=0, tree_level=0, tree_height=light_tree_height, state=init_state, env=env, 121 | space_type=SpaceType.Light) 122 | light_tree_cache[selected_heavy_action_frozen] = light_root 123 | # for the light tree 124 | best_reward = 0 125 | # query to consider for the current select heavy configuration 126 | query_to_consider = set( 127 | [applicable_query for applicable_queries in 128 | list(map(lambda x: index_to_applicable_queries[x], selected_heavy_action_frozen)) for applicable_query 129 | in 130 | applicable_queries]) 131 | logging.debug(f"query to consider: {query_to_consider}") 132 | best_micro_performance = dict() 133 | for t2 in range(1, micro_episode + 1): 134 | # for the micro episode 135 | env.reset() 136 | selected_light_actions = light_root.sample(t1 * micro_episode + t2) 137 | # evaluate the light actions 138 | for selected_light_action in selected_light_actions: 139 | # move to next state 140 | state = env.step_without_evaluation(selected_light_action) 141 | # obtain sample number 142 | sample_num = math.ceil(query_evaluate_sample_rate * len(query_to_consider)) 143 | # generate sample queries 144 | sampled_query_list = random.sample(list(query_to_consider), k=sample_num) 145 | logging.debug(f"sampled_query_list: {sampled_query_list}") 146 | # obtain run time info by running queries within timeout 147 | run_time = env.evaluate(sampled_query_list) 148 | # the total time of sampled queries 149 | total_run_time = sum(run_time) 150 | # the default time of sampled queries 151 | default_runtime_of_sampled_queries = [default_runtime[select_query] for select_query in 152 | sampled_query_list] 153 | sum_default_time_of_sampled_queries = sum(default_runtime_of_sampled_queries) 154 | # the relative ration of the improvement, the less of total_run_time, the better 155 | light_reward = sum_default_time_of_sampled_queries / total_run_time 156 | logging.debug(f"default_runtime_of_sampled_queries {default_runtime_of_sampled_queries}") 157 | logging.debug(f"light_action: {selected_light_actions}") 158 | logging.debug(f"light_reward: {light_reward}") 159 | 160 | other_default_time = sum(default_runtime[select_query] for select_query in range(nr_query) if 161 | select_query not in sampled_query_list) 162 | estimate_workload_time = (other_default_time + total_run_time) 163 | logging.debug(f"estimate whole workload time: {estimate_workload_time}") 164 | if estimate_workload_time < best_simulation_time: 165 | best_simulation_time = estimate_workload_time 166 | best_configs = {"heavy": selected_heavy_action_frozen, "light": selected_light_actions} 167 | 168 | current_time = time.time() 169 | logging.debug(f"current global time: {(current_time - start_time)}") 170 | logging.debug(f"global time for indices: {idx_build_time}") 171 | 172 | light_root.update_statistics_with_mcts_reward(light_reward, selected_light_actions) 173 | # update the best gain for each query 174 | for sample_query_id in range(len(sampled_query_list)): 175 | sample_query = sampled_query_list[sample_query_id] 176 | run_time_of_sampled_query = run_time[sample_query_id] 177 | if sample_query in best_micro_performance: 178 | # if we get better improvement, then set it to better time 179 | if run_time_of_sampled_query < best_micro_performance[sample_query]: 180 | best_micro_performance[sample_query] = run_time_of_sampled_query 181 | else: 182 | best_micro_performance[sample_query] = run_time[sample_query_id] 183 | if sum(run_time) < best_reward: 184 | best_reward = sum(run_time) 185 | # save the performance of current selected heavy action 186 | macro_performance_info[selected_heavy_action_frozen] = best_micro_performance 187 | 188 | logging.debug(f"macro_performance_info: {macro_performance_info}") 189 | # sys.stdout.flush() 190 | # ### obtain the best macro performances info 191 | best_slot_performance = dict() 192 | for selected_heavy_action_frozen, performance in macro_performance_info.items(): 193 | for query, query_run_time in performance.items(): 194 | if query in best_slot_performance: 195 | if query_run_time < best_slot_performance[query]: 196 | best_slot_performance[query] = query_run_time 197 | else: 198 | best_slot_performance[query] = query_run_time 199 | logging.debug(f"best_slot_performance {best_slot_performance}") 200 | # after testing the performance of each configuration 201 | # generate the update information based on delta improvement 202 | update_info_slot = [] 203 | for selected_heavy_actions in selected_heavy_action_batch: 204 | # generate intermediate result 205 | update_reward = [] 206 | previous_performance = dict() 207 | for i in range(len(selected_heavy_actions)): 208 | if selected_heavy_actions[i] != terminate_action: 209 | selected_heavy_action_frozen = frozenset(selected_heavy_actions[:i + 1]) 210 | logging.debug(f"current selected_heavy_action_frozen: {selected_heavy_action_frozen}") 211 | applicable_query_performance = macro_performance_info[selected_heavy_action_frozen] 212 | # generate reward based on the difference between previous performance and current performance 213 | # the query for current indices 214 | query_to_consider_new = index_to_applicable_queries[selected_heavy_actions[i]] 215 | # get the query to consider 216 | delta_improvement = 0 217 | for query in query_to_consider_new: 218 | previous_runtime = default_runtime[query] 219 | current_runtime = default_runtime[query] 220 | if query in applicable_query_performance: 221 | current_runtime = applicable_query_performance[query] 222 | if query in previous_performance: 223 | previous_runtime = previous_performance[query] 224 | delta_improvement += (previous_runtime - current_runtime) 225 | previous_performance = applicable_query_performance 226 | logging.debug(f"applicable_query_performance: {applicable_query_performance}") 227 | logging.debug("previous_performance: {previous_performance}") 228 | delta_improvement = max(delta_improvement, 0) 229 | update_reward.append(delta_improvement) 230 | # update the tree based on the simulation results 231 | update_info_slot.append((update_reward, [heavy_action for heavy_action in selected_heavy_actions if 232 | heavy_action != terminate_action])) 233 | logging.debug(f"update_info_slot: {update_info_slot}") 234 | heavy_root.update_batch(update_info_slot) 235 | previous_set = set(configuration_to_evaluate[evaluated_order[-1]]) 236 | heavy_root.print_reward_info() 237 | end_episode_time = time.time() 238 | logging.debug(f"current time: {(end_episode_time - start_time)}") 239 | logging.debug(f"time for indices: {idx_build_time}") 240 | env.print_action_summary(heavy_root.best_actions()) 241 | t1 += max_delay_time 242 | 243 | best_heavy_actions = heavy_root.best_actions() 244 | best_heavy_action_alternatives = best_heavy_actions 245 | if len(best_heavy_action_alternatives) < heavy_tree_height: 246 | first_layer_reward_info = heavy_root.print_reward_info() 247 | first_layer_reward_sorted_info = sorted(first_layer_reward_info, reverse=True) 248 | # topK_rewards = first_layer_reward_sorted_info[:heavy_tree_height] 249 | # topK_actions = [first_layer_reward_info.index(reward) for reward in topK_rewards] 250 | # may exist a issue index performance is tie, but very unlikely happened 251 | indexed_columns = set() 252 | for top_action in best_heavy_actions: 253 | top_index = candidate_indices[top_action] 254 | current_index_columns = top_index[2].split(",") 255 | indexed_columns.update(current_index_columns) 256 | # for top_action in topK_actions: 257 | while len(best_heavy_action_alternatives) < heavy_tree_height: 258 | # test whether this index appear before 259 | top_reward = first_layer_reward_sorted_info[0] 260 | top_action = first_layer_reward_info.index(top_reward) 261 | top_index = candidate_indices[top_action] 262 | current_index_columns = top_index[2].split(",") 263 | if not any(current_index_column in indexed_columns for current_index_column in current_index_columns): 264 | best_heavy_action_alternatives.append(top_action) 265 | indexed_columns.update(current_index_columns) 266 | first_layer_reward_sorted_info.pop(0) 267 | 268 | best_frozen_heavy_configs = frozenset(best_heavy_action_alternatives) 269 | if best_frozen_heavy_configs in light_tree_cache: 270 | light_root = light_tree_cache[best_frozen_heavy_configs] 271 | else: 272 | light_root = uct_node(round=0, tree_level=0, tree_height=light_tree_height, state=init_state, env=env, 273 | space_type=SpaceType.Light) 274 | 275 | add_action = set(best_heavy_actions) - previous_set 276 | drop_action = previous_set - set(best_heavy_actions) 277 | 278 | # best frozen heavy configs 279 | # best_frozen_heavy_configs = frozenset(best_heavy_actions) 280 | # if best_frozen_heavy_configs in light_tree_cache: 281 | # light_root = light_tree_cache[best_frozen_heavy_configs] 282 | # else: 283 | # light_root = uct_node(round=0, tree_level=0, tree_height=light_tree_height, state=init_state, env=env, 284 | # space_type=SpaceType.Light) 285 | 286 | # additional step, tuning the final index configuration 287 | # build the indices 288 | # add_action = set(best_heavy_actions) - previous_set 289 | # drop_action = previous_set - set(best_heavy_actions) 290 | 291 | # really build index and evaluate the performance 292 | env.index_step(add_action, drop_action) 293 | micro_episode_final_tune = 50 294 | default_total_time = sum(default_runtime) 295 | best_light_runtime = sum(default_runtime) 296 | best_light_config_simulation = [] 297 | for t2 in range(1, micro_episode_final_tune): 298 | # for the micro episode 299 | env.reset() 300 | selected_light_actions = light_root.sample(t1 * micro_episode + t2) 301 | # evaluate the light actions 302 | for selected_light_action in selected_light_actions: 303 | # move to next state 304 | state = env.step_without_evaluation(selected_light_action) 305 | # obtain run time info by running queries within timeout 306 | run_time = env.evaluate([query for query in range(nr_query)]) 307 | # the total time of total queries 308 | total_run_time = sum(run_time) 309 | if total_run_time < best_light_runtime: 310 | best_light_runtime = total_run_time 311 | best_light_config_simulation = selected_light_actions 312 | # the relative ration of the improvement, the less of total_run_time, the better 313 | light_reward = default_total_time / total_run_time 314 | logging.debug(f"light_action: {selected_light_actions}") 315 | logging.debug(f"light_reward: {light_reward}") 316 | light_root.update_statistics_with_mcts_reward(light_reward, selected_light_actions) 317 | best_light_actions = light_root.best_actions() 318 | # best heavy action from tree search 319 | logging.info(f"best configurations during the simulation {best_configs}") 320 | logging.info(f"best heavy action from tree search {heavy_root.best_actions()}") 321 | logging.info(f"best_light_actions {best_light_actions}") 322 | 323 | final_tune_time = time.time() 324 | logging.info(f"end: {final_tune_time}") 325 | 326 | logging.info(f"Summary: Total Tuning Time {(final_tune_time - start_tune_time) / 3600} hours") 327 | logging.info(f"Index Recommendation from MCTS according to metric 1") 328 | env.print_action_summary(best_heavy_actions) 329 | logging.info(f"Index Recommendation from MCTS according to metric 2") 330 | env.print_action_summary(best_heavy_action_alternatives) 331 | logging.info(f"System Parameter Recommendation from MCTS according to metric 1") 332 | env.print_action_summary(best_light_actions) 333 | logging.info(f"System Parameter Recommendation form MCTS according to metric 2") 334 | env.print_action_summary(best_light_config_simulation) 335 | -------------------------------------------------------------------------------- /udo/udo-optimization/udo_optimization/envs/udo_env.py: -------------------------------------------------------------------------------- 1 | # ----------------------------------------------------------------------- 2 | # Copyright (c) 2021 Cornell Database Group 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining 5 | # a copy of this software and associated documentation files (the 6 | # "Software"), to deal in the Software without restriction, including 7 | # without limitation the rights to use, copy, modify, merge, publish, 8 | # distribute, sublicense, and/or sell copies of the Software, and to 9 | # permit persons to whom the Software is furnished to do so, subject to 10 | # the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be 13 | # included in all copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 17 | # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT 18 | # IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR 19 | # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, 20 | # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 21 | # OTHER DEALINGS IN THE SOFTWARE. 22 | # ----------------------------------------------------------------------- 23 | 24 | import logging 25 | import math 26 | import random 27 | import time 28 | 29 | import gym 30 | import numpy as np 31 | from gym import spaces 32 | 33 | 34 | class UDOEnv(gym.Env): 35 | """the environment of DBMS optimizer""" 36 | 37 | def __init__(self, driver, queries, candidate_indices, config): 38 | """ init method 39 | Args: 40 | driver: DBMS connector 41 | queries: queries to optimize 42 | candidate_indices: candidate indices to select 43 | config: the configuration for tuning 44 | 45 | Return: the environment 46 | """ 47 | super(UDOEnv, self).__init__() 48 | 49 | self.driver = driver 50 | 51 | # initial index action tuning space 52 | self.candidate_indices = candidate_indices 53 | self.index_candidate_num = len(self.candidate_indices) 54 | logging.debug(f"the total number of index candidates is {self.index_candidate_num}") 55 | 56 | # initial system parameter space 57 | self.parameter_candidate = self.driver.get_system_parameter_space() 58 | self.parameter_candidate_num = len(self.parameter_candidate) 59 | logging.debug(f"the total number of tuning system parameters is {self.parameter_candidate_num}") 60 | 61 | # combine the actions from 2 sub actions 62 | # action space 63 | self.nA_index = self.index_candidate_num 64 | self.nA_parameter = sum(self.parameter_candidate) 65 | self.nA = int(self.nA_index + self.nA_parameter) 66 | self.action_space = spaces.Discrete(self.nA) 67 | logging.debug(f"action space {self.nA}") 68 | 69 | # state space 70 | self.nS_index = int(math.pow(2, self.index_candidate_num)) 71 | self.nS_parameter = np.prod(self.parameter_candidate) 72 | # self.nS = self.nS_index * self.nS_parameter 73 | logging.debug(f"index state space {self.nS_index}") 74 | logging.debug(f"parameter state space {self.nS_parameter}") 75 | 76 | # our transition matrix is a deterministic matrix 77 | # the observation space 78 | observation_space_array = np.concatenate([np.full(self.index_candidate_num, 1), self.parameter_candidate]) 79 | self.observation_space = spaces.MultiDiscrete(observation_space_array) 80 | self.current_state = np.concatenate( 81 | [np.zeros(self.index_candidate_num, dtype=int), np.zeros(self.parameter_candidate_num, dtype=int)]) 82 | 83 | # the MDP init setting 84 | # index build time 85 | 86 | # horizon 87 | self.horizon = config['horizon'] 88 | self._step = 0 89 | # current step 90 | # for runtime statistics 91 | self.accumulated_index_time = 0 92 | self.start_time = time.time() 93 | 94 | # get all queries in the given benchmark 95 | self.queries = queries 96 | self.query_ids = list(self.queries.keys()) 97 | self.nr_query = len(self.queries) 98 | self.query_sqls = [self.queries[query_id] for query_id in self.query_ids] 99 | query_to_id = {self.query_ids[idx]: idx for idx in range(self.nr_query)} 100 | self.index_to_applicable_queries = list( 101 | map(lambda x: list(map(lambda y: query_to_id[y], x[3])), self.candidate_indices)) 102 | 103 | # the default run time 104 | default_time_out_per_query = config['default_query_time_out'] 105 | time_out_ratio = config['time_out_ratio'] 106 | input_runtime_out = [default_time_out_per_query] * self.nr_query 107 | self.default_runtime = self.driver.run_queries_with_timeout(self.query_sqls, input_runtime_out) 108 | self.runtime_out = [ 109 | time_out_ratio * query_runtime if query_runtime < default_time_out_per_query 110 | else default_time_out_per_query for query_runtime in self.default_runtime] 111 | 112 | self.sample_rate = config['sample_rate'] 113 | self.best_state = None 114 | self.best_run_performance = sum(self.default_runtime) 115 | logging.info("start to tuning your database") 116 | logging.info(f"timeout for queries {self.runtime_out}") 117 | 118 | def state_decoder(self, num): 119 | """decode a vector to a state number""" 120 | index_pos = int(num % self.nS_index) 121 | index_state_string = np.binary_repr(int(index_pos), width=self.index_candidate_num)[::-1] 122 | # index stata represented in string 123 | index_state = list(map(int, index_state_string)) 124 | parameter_value = int(num / self.nS_index) 125 | parameter_pos = [0] * self.parameter_candidate_num 126 | for i in range(self.parameter_candidate_num): 127 | parameter_pos[i] = (parameter_value % self.parameter_candidate[i]) 128 | parameter_value = int(parameter_value / self.parameter_candidate[i]) 129 | return index_state + parameter_pos 130 | 131 | def state_encoder(self, state): 132 | """encode a state number to a vector""" 133 | index_state = state[:self.index_candidate_num] 134 | parameter_state = state[self.index_candidate_num:] 135 | index_pos = int("".join([str(int(a)) for a in reversed(index_state)]), 2) 136 | parameter_pos = 0 137 | parameter_base = 1 138 | for i in range(len(self.parameter_candidate)): 139 | parameter_pos = parameter_pos + parameter_state[i] * parameter_base 140 | parameter_base = parameter_base * self.parameter_candidate[i] 141 | pos = index_pos + parameter_pos * self.nS_index 142 | return int(pos) 143 | 144 | def retrieve_heavy_actions(self, state): 145 | """obtain available heavy actions given a state""" 146 | index_state = state[:self.index_candidate_num] 147 | # check which indices are available or not 148 | candidate_index_action = [i for i in range(len(index_state)) if index_state[i] == 0] 149 | return candidate_index_action 150 | 151 | def retrieve_light_actions(self, state): 152 | """obtain available light actions given a state""" 153 | parameter_state = state[self.index_candidate_num:] 154 | candidate_parameter_action = [] 155 | parameter_sum = 0 156 | for i in range(len(parameter_state)): 157 | # if the current parameter state is a default state, we switch its value 158 | if parameter_state[i] == 0: 159 | for j in range(1, self.parameter_candidate[i]): 160 | candidate_parameter_action.append(self.nA_index + parameter_sum + j) 161 | parameter_sum += self.parameter_candidate[i] 162 | all_light_actions = candidate_parameter_action 163 | return all_light_actions 164 | 165 | def retrieve_actions(self, state): 166 | """retrieve available actions for a state including both light actions and heavy actions""" 167 | heavy_action = self.retrieve_heavy_actions(state) 168 | light_action = self.retrieve_light_actions(state) 169 | return heavy_action + light_action 170 | 171 | def transition(self, state, action): 172 | """transition from a state and an action to a next state""" 173 | assert action < self.nA 174 | index_state = state[:self.index_candidate_num] 175 | parameter_state = state[self.index_candidate_num:] 176 | if (action < self.nA_index): 177 | # action is the index action 178 | index_action = action 179 | index_state[index_action] = 1 180 | else: 181 | # action is the parameter action 182 | parameter_action = action - self.nA_index 183 | parameter_value = 0 184 | parameter_type = 0 185 | while parameter_type < len(self.parameter_candidate): 186 | parameter_range = self.parameter_candidate[parameter_type] 187 | # test whether cover this range 188 | if parameter_action < (parameter_value + parameter_range): 189 | parameter_value = parameter_action - parameter_value 190 | break 191 | parameter_value = parameter_value + parameter_range 192 | parameter_type += 1 193 | parameter_state[parameter_type] = parameter_value 194 | next_state = index_state + parameter_state 195 | return next_state, self.state_encoder(next_state) 196 | 197 | def step_without_evaluation(self, action): 198 | """move to the next state given the action and current state""" 199 | state = self.current_state 200 | # parameter state and action 201 | index_current_state = state[:self.index_candidate_num] 202 | parameter_current_state = state[self.index_candidate_num:] 203 | parameter_action = action - self.nA_index 204 | parameter_value = 0 205 | for parameter_type in range(len(self.parameter_candidate)): 206 | parameter_range = self.parameter_candidate[parameter_type] 207 | if parameter_action < (parameter_value + parameter_range): 208 | parameter_value = parameter_action - parameter_value 209 | break 210 | parameter_value = parameter_value + parameter_range 211 | parameter_current_state[parameter_type] = parameter_value 212 | next_state = np.concatenate([index_current_state, parameter_current_state]) 213 | self.current_state = next_state 214 | return next_state 215 | 216 | def retrieve_light_action_command(self, action): 217 | parameter_action = action - self.nA_index 218 | parameter_value = 0 219 | parameter_type = 0 220 | for parameter_type in range(len(self.parameter_candidate)): 221 | parameter_range = self.parameter_candidate[parameter_type] 222 | if parameter_action < (parameter_value + parameter_range): 223 | parameter_value = parameter_action - parameter_value 224 | break 225 | parameter_value = parameter_value + parameter_range 226 | return self.driver.get_system_parameter_command(parameter_type, parameter_value) 227 | 228 | def evaluate(self, sampled_queries): 229 | """evaluate current state using the sampled queries""" 230 | state = self.current_state 231 | index_current_state = state[:self.index_candidate_num] 232 | parameter_current_state = state[self.index_candidate_num:] 233 | # change system parameter 234 | self.driver.change_system_parameter(parameter_current_state) 235 | # invoke queries 236 | run_time = self.driver.run_queries_with_timeout( 237 | [self.query_sqls[sampled_query] for sampled_query in sampled_queries], 238 | [self.runtime_out[sampled_query] for sampled_query in sampled_queries]) 239 | logging.debug(f"evaluate time {sum(run_time)}") 240 | return run_time 241 | 242 | def evaluate_query_cost(self, sampled_queries): 243 | cost = self.driver.analyze_queries_cost( 244 | [self.query_sqls[sampled_query] for sampled_query in sampled_queries]) 245 | return cost 246 | 247 | def index_step(self, add_actions, remove_actions): 248 | """index step""" 249 | for add_action in add_actions: 250 | index_to_create = self.candidate_indices[add_action] 251 | # build index action 252 | logging.debug(f"create index {index_to_create}") 253 | self.driver.build_index(index_to_create) 254 | for remove_action in remove_actions: 255 | index_to_drop = self.candidate_indices[remove_action] 256 | # drop index action 257 | logging.debug(f"drop index {index_to_drop}") 258 | self.driver.drop_index(index_to_drop) 259 | 260 | def index_add_step(self, add_action): 261 | """create indexes""" 262 | index_to_create = self.candidate_indices[add_action] 263 | # an index build action 264 | logging.debug(f"create index {index_to_create}") 265 | self.driver.build_index(index_to_create) 266 | 267 | def index_drop_step(self, remove_actions): 268 | """drop indexes""" 269 | for remove_action in remove_actions: 270 | index_to_drop = self.candidate_indices[remove_action] 271 | # drop index action 272 | logging.debug(f"drop index {index_to_drop}") 273 | self.driver.drop_index(index_to_drop) 274 | 275 | def step(self, action): 276 | """invoke an action and move to the next step""" 277 | state = self.current_state 278 | # parameter state and action 279 | index_current_state = state[:self.index_candidate_num] 280 | parameter_current_state = state[self.index_candidate_num:] 281 | if action < self.nA_index: 282 | # index action, create indices 283 | if index_current_state[action] == 0: 284 | index_start_time = time.time() 285 | self.index_add_step(action) 286 | index_end_time = time.time() 287 | index_time = (index_end_time - index_start_time) 288 | self.accumulated_index_time = self.accumulated_index_time + index_time 289 | index_current_state[action] = 1 290 | else: 291 | # parameter action, switch to different parameters 292 | parameter_action = action - self.nA_index 293 | parameter_value = 0 294 | for parameter_type in range(len(self.parameter_candidate)): 295 | parameter_range = self.parameter_candidate[parameter_type] 296 | if parameter_action < (parameter_value + parameter_range): 297 | parameter_value = parameter_action - parameter_value 298 | break 299 | parameter_value = parameter_value + parameter_range 300 | parameter_current_state[parameter_type] = parameter_value 301 | self.driver.change_system_parameter(parameter_current_state) 302 | 303 | # heavy actions, and we only consider queries applicable to selected indices 304 | active_indices = [index_pos for index_pos in range(self.nA_index) if index_current_state[index_pos] == 1] 305 | query_to_consider = set( 306 | [applicable_query for applicable_queries in 307 | list(map(lambda x: self.index_to_applicable_queries[x], active_indices)) for applicable_query in 308 | applicable_queries]) 309 | 310 | if len(query_to_consider) == 0: 311 | query_to_consider = range(self.nr_query) 312 | 313 | # obtain sample number 314 | sample_num = math.ceil(self.sample_rate * len(query_to_consider)) 315 | # generate sample queries 316 | sample_queries = random.sample(list(query_to_consider), k=sample_num) 317 | 318 | # invoke queries 319 | run_time = self.driver.run_queries_with_timeout( 320 | [self.query_sqls[sample_query] for sample_query in sample_queries], 321 | [self.runtime_out[sampled_query] for sampled_query in sample_queries]) 322 | 323 | next_state = np.concatenate([index_current_state, parameter_current_state]) 324 | self.current_state = next_state 325 | 326 | # scale the reward 327 | reward = sum(self.default_runtime) / sum(run_time) 328 | current_time = time.time() 329 | estimate_workload_time = sum(run_time) + sum( 330 | [self.default_runtime[query] for query in range(self.nr_query) if query not in sample_queries]) 331 | 332 | logging.debug(f"run time: {run_time}") 333 | logging.debug(f"index time: {self.accumulated_index_time}") 334 | logging.debug(f"evaluate time: {sum(run_time)}") 335 | logging.debug(f"next state: {next_state}") 336 | logging.debug(f"tuning duration: {(current_time - self.start_time)}", ) 337 | logging.debug(f"estimate whole workload time: {estimate_workload_time}") 338 | 339 | if estimate_workload_time < self.best_run_performance: 340 | self.best_run_performance = estimate_workload_time 341 | self.best_state = self.current_state 342 | 343 | self._step += 1 344 | if self._step == self.horizon: 345 | self._step = 0 346 | return next_state, reward, True, {} 347 | else: 348 | return next_state, reward, False, {} 349 | 350 | def reset(self): 351 | """reset the state""" 352 | if self.current_state is not None: 353 | state = self.current_state 354 | index_current_state = state[:self.index_candidate_num] 355 | parameter_current_state = state[self.index_candidate_num:] 356 | # drop all indices at the current state 357 | for i in range(self.index_candidate_num): 358 | if index_current_state[i] == 1: 359 | index_drop_sql = self.candidate_indices[i] 360 | self.driver.drop_index(index_drop_sql) 361 | # set the parameter to default value, the first value 362 | # self.driver.change_system_parameter(np.zeros(self.parameter_candidate_num, dtype=int)) 363 | self.current_state = np.concatenate( 364 | [np.zeros(self.index_candidate_num, dtype=int), np.zeros(self.parameter_candidate_num, dtype=int)]) 365 | return self.current_state 366 | 367 | def print_state_summary(self, state): 368 | index_state = state[:self.index_candidate_num] 369 | parameter_state = state[self.index_candidate_num:] 370 | logging.info(f"Best index configurations:") 371 | for i in range(len(index_state)): 372 | if index_state[i] == 1: 373 | index_to_create = self.candidate_indices[i] 374 | self.driver.build_index_command(index_to_create) 375 | logging.info(f"Best system parameters:") 376 | for i in range(len(parameter_state)): 377 | logging.info(self.driver.get_system_parameter_command(i, parameter_state[i])) 378 | 379 | def print_action_summary(self, actions): 380 | """print action summary given actions""" 381 | best_indices = [self.candidate_indices[action_idx] for action_idx in actions if action_idx < self.nA_index] 382 | best_sys_actions = [action_idx for action_idx in actions if action_idx >= self.nA_index] 383 | if best_indices is not None and len(best_indices) > 0: 384 | logging.info(f"Best index configurations:") 385 | for best_index in best_indices: 386 | logging.info(self.driver.build_index_command(best_index)) 387 | if best_sys_actions is not None and len(best_sys_actions) > 0: 388 | logging.info(f"Best system parameters:") 389 | for sys_action in best_sys_actions: 390 | logging.info(self.retrieve_light_action_command(sys_action)) 391 | --------------------------------------------------------------------------------