simulate a DB tuning environment to test DDPG performance

server/analysis/ddpg/__init__.py

@@ -3,8 +3,3 @@
 #
 # Copyright (c) 2017-18, Carnegie Mellon University Database Group
 #
-
-
-from analysis.ddpg.ddpg import DDPG
-
-__all__ = ["DDPG"]

server/analysis/simulation.py

@@ -0,0 +1,151 @@
+#
+# OtterTune - simulation.py
+#
+# Copyright (c) 2017-18, Carnegie Mellon University Database Group
+#
+
+import random
+import os
+import sys
+try:
+    import matplotlib.pyplot as plt
+except (ModuleNotFoundError, ImportError):
+    plt = None
+import numpy as np
+import torch
+sys.path.append("../")
+from analysis.util import get_analysis_logger  # noqa
+from analysis.ddpg.ddpg import DDPG  # noqa
+
+LOG = get_analysis_logger(__name__)
+
+
+class Environment(object):
+    def __init__(self, n_knob, n_metric, mode=0):
+        self.knob_dim = n_knob
+        self.metric_dim = n_metric
+        self.mode = mode
+
+    def identity_sqrt(self, knob_data):
+        n1 = self.knob_dim // 4
+        n2 = self.knob_dim // 4
+        part1 = np.sum(knob_data[0: n1])
+        part2 = np.sum(np.sqrt(knob_data[n1: n1 + n2]))
+        reward = np.array([part1 + part2]) / (self.knob_dim // 2)
+        metric_data = np.zeros(self.metric_dim)
+        return reward, metric_data
+
+    def borehole(self, knob_data):
+        # ref: http://www.sfu.ca/~ssurjano/borehole.html
+        # pylint: disable=invalid-name
+        rw = knob_data[0] * (0.15 - 0.05) + 0.05
+        r = knob_data[1] * (50000 - 100) + 100
+        Tu = knob_data[2] * (115600 - 63070) + 63070
+        Hu = knob_data[3] * (1110 - 990) + 990
+        Tl = knob_data[4] * (116 - 63.1) + 63.1
+        Hl = knob_data[5] * (820 - 700) + 700
+        L = knob_data[6] * (1680 - 1120) + 1120
+        Kw = knob_data[7] * (12045 - 9855) + 9855
+
+        frac = 2 * L * Tu / (np.log(r / rw) * rw ** 2 * Kw)
+        reward = 2 * np.pi * Tu * (Hu - Hl) / (np.log(r / rw) * (1 + frac + Tu / Tl))
+        return np.array([reward]), np.zeros(self.metric_dim)
+
+    def threshold(self, knob_data):
+        n1 = self.knob_dim // 4
+        n2 = self.knob_dim // 4
+        part1 = np.sum(knob_data[0: n1] > 0.9)
+        part2 = np.sum(knob_data[n1: n1 + n2] < 0.1)
+        reward = np.array([part1 + part2])
+        metric_data = np.zeros(self.metric_dim)
+        return reward, metric_data
+
+    def simulate(self, knob_data):
+        if self.mode == 0:
+            return self.identity_sqrt(knob_data)
+        elif self.mode == 1:
+            return self.threshold(knob_data)
+        elif self.mode == 2:
+            return self.borehole(knob_data)
+
+
+def train_ddpg(env, gamma=0.99, tau=0.002, lr=0.01, batch_size=32, n_loops=1000):
+    results = []
+    x_axis = []
+    ddpg = DDPG(n_actions=env.knob_dim, n_states=env.metric_dim, gamma=gamma, tau=tau,
+                clr=lr, alr=lr, batch_size=batch_size)
+    knob_data = np.random.rand(env.knob_dim)
+    prev_metric_data = np.zeros(env.metric_dim)
+    for i in range(n_loops):
+        reward, metric_data = env.simulate(knob_data)
+        ddpg.add_sample(prev_metric_data, knob_data, reward, metric_data, False)
+        ddpg.update()
+        if i % 20 == 0:
+            results.append(run_ddpg(env, ddpg))
+            x_axis.append(i)
+        prev_metric_data = metric_data
+        knob_data = ddpg.choose_action(prev_metric_data)
+    return np.array(results), np.array(x_axis)
+
+
+def run_ddpg(env, ddpg):
+    total_reward = 0.0
+    n_samples = 100
+    prev_metric_data = np.zeros(env.metric_dim)
+    for _ in range(n_samples):
+        knob_data = ddpg.choose_action(prev_metric_data)
+        reward, prev_metric_data = env.simulate(knob_data)
+        total_reward += reward
+    return total_reward / n_samples
+
+
+def plotlines(x_axis, data1, data2, label1, label2, title, path):
+    if plt:
+        plt.plot(x_axis, data1, color='red', label=label1)
+        plt.plot(x_axis, data2, color='blue', label=label2)
+        plt.legend()
+        plt.xlabel("loops")
+        plt.ylabel("rewards")
+        plt.title(title)
+        plt.savefig(path)
+        plt.clf()
+
+
+def main(knob_dim=192, metric_dim=60, lr=0.001, mode=0, n_loops=1000):
+    if not plt:
+        LOG.info("Cannot import matplotlib. Will write results to files instead of figures.")
+    random.seed(0)
+    np.random.seed(0)
+    torch.manual_seed(0)
+    env = Environment(knob_dim, metric_dim, mode=mode)
+
+    n_repeats = 5
+    for i in range(n_repeats):
+        if i == 0:
+            results1, x_axis = train_ddpg(env, gamma=0, lr=lr, n_loops=n_loops)
+        else:
+            results1 += train_ddpg(env, gamma=0, lr=lr, n_loops=n_loops)[0]
+    for i in range(n_repeats):
+        if i == 0:
+            results2, x_axis = train_ddpg(env, gamma=0.99, lr=lr, n_loops=n_loops)
+        else:
+            results2 += train_ddpg(env, gamma=0.99, lr=lr, n_loops=n_loops)[0]
+    results1 /= n_repeats
+    results2 /= n_repeats
+    title = "knob_{}_lr_{}".format(knob_dim, lr)
+    if plt:
+        if not os.path.exists("figures"):
+            os.mkdir("figures")
+        filename = "figures/{}.pdf".format(title)
+        plotlines(x_axis, results1, results2, "gamma=0", "gamma=0.99", title, filename)
+    else:
+        with open(title + '_1.csv', 'w') as f1:
+            for i, result in zip(x_axis, results1):
+                f1.write(str(i) + ',' + str(result[0]) + '\n')
+        with open(title + '_2.csv', 'w') as f2:
+            for i, result in zip(x_axis, results2):
+                f2.write(str(i) + ',' + str(result[0]) + '\n')
+
+
+if __name__ == '__main__':
+    main()