Replay changing config (#56)

**Summary**: Replay now applies config deltas.

**Demo**:
Modified test that passes.
<img width="626" alt="Screenshot 2024-12-24 at 20 18 09"
src="https://github.com/user-attachments/assets/bcbf5e9c-f68c-452c-8a4e-47b1f226c63f"
/>


**Details**:
* Undoing query knobs is not yet handled. This depends on whether the
interface takes in config deltas or full configs for query knobs which
depends on how the agents work.

Author: wangpatrick57

env/integtest_replay.py

@@ -1,7 +1,14 @@
 import unittest
 
+from benchmark.tpch.constants import DEFAULT_TPCH_SEED
 from env.integtest_util import IntegtestWorkspace, MockTuningAgent
 from env.replay import replay
+from env.tuning_agent import (
+    DBMSConfigDelta,
+    IndexesDelta,
+    QueryKnobsDelta,
+    SysKnobsDelta,
+)
 
 
 class ReplayTests(unittest.TestCase):
@@ -11,7 +18,35 @@ def setUpClass() -> None:
 
     def test_replay(self) -> None:
         agent = MockTuningAgent(IntegtestWorkspace.get_dbgym_cfg())
-        replay(IntegtestWorkspace.get_dbgym_cfg(), agent.tuning_agent_artifacts_dpath)
+        agent.delta_to_return = DBMSConfigDelta(
+            indexes=IndexesDelta(
+                ["CREATE INDEX idx_orders_custkey ON orders(o_custkey)"]
+            ),
+            sysknobs=SysKnobsDelta(
+                {"shared_buffers": "2GB"},
+            ),
+            qknobs=QueryKnobsDelta(
+                {
+                    f"S{DEFAULT_TPCH_SEED}-Q1": [
+                        "set enable_hashagg = off",
+                        "set enable_sort = on",
+                    ],
+                }
+            ),
+        )
+        agent.step()
+        replay_data = replay(
+            IntegtestWorkspace.get_dbgym_cfg(), agent.tuning_agent_artifacts_dpath
+        )
+
+        # We do some very simple sanity checks here due to the inherent randomness of executing a workload.
+        # We check that there is one data point for the initial config and one for the config change.
+        self.assertEqual(len(replay_data), 2)
+        # We check that the second step is faster.
+        self.assertLess(replay_data[1][0], replay_data[0][0])
+        # We check that no queries timed out in either step.
+        self.assertEqual(replay_data[0][1], 0)
+        self.assertEqual(replay_data[1][1], 0)
 
 
 if __name__ == "__main__":

env/integtest_tuning_agent.py

@@ -23,37 +23,45 @@ def make_config(letter: str) -> DBMSConfigDelta:
             qknobs=QueryKnobsDelta({letter: [letter]}),
         )
 
-    def test_get_step_delta(self) -> None:
+    def test_get_delta_at_step(self) -> None:
         agent = MockTuningAgent(IntegtestWorkspace.get_dbgym_cfg())
 
-        agent.config_to_return = PostgresConnTests.make_config("a")
+        agent.delta_to_return = PostgresConnTests.make_config("a")
         agent.step()
-        agent.config_to_return = PostgresConnTests.make_config("b")
+        agent.delta_to_return = PostgresConnTests.make_config("b")
         agent.step()
-        agent.config_to_return = PostgresConnTests.make_config("c")
+        agent.delta_to_return = PostgresConnTests.make_config("c")
         agent.step()
 
         reader = TuningAgentArtifactsReader(agent.tuning_agent_artifacts_dpath)
 
-        self.assertEqual(reader.get_step_delta(1), PostgresConnTests.make_config("b"))
-        self.assertEqual(reader.get_step_delta(0), PostgresConnTests.make_config("a"))
-        self.assertEqual(reader.get_step_delta(1), PostgresConnTests.make_config("b"))
-        self.assertEqual(reader.get_step_delta(2), PostgresConnTests.make_config("c"))
+        self.assertEqual(
+            reader.get_delta_at_step(1), PostgresConnTests.make_config("b")
+        )
+        self.assertEqual(
+            reader.get_delta_at_step(0), PostgresConnTests.make_config("a")
+        )
+        self.assertEqual(
+            reader.get_delta_at_step(1), PostgresConnTests.make_config("b")
+        )
+        self.assertEqual(
+            reader.get_delta_at_step(2), PostgresConnTests.make_config("c")
+        )
 
-    def test_get_all_deltas(self) -> None:
+    def test_get_all_deltas_in_order(self) -> None:
         agent = MockTuningAgent(IntegtestWorkspace.get_dbgym_cfg())
 
-        agent.config_to_return = PostgresConnTests.make_config("a")
+        agent.delta_to_return = PostgresConnTests.make_config("a")
         agent.step()
-        agent.config_to_return = PostgresConnTests.make_config("b")
+        agent.delta_to_return = PostgresConnTests.make_config("b")
         agent.step()
-        agent.config_to_return = PostgresConnTests.make_config("c")
+        agent.delta_to_return = PostgresConnTests.make_config("c")
         agent.step()
 
         reader = TuningAgentArtifactsReader(agent.tuning_agent_artifacts_dpath)
 
         self.assertEqual(
-            reader.get_all_deltas(),
+            reader.get_all_deltas_in_order(),
             [
                 PostgresConnTests.make_config("a"),
                 PostgresConnTests.make_config("b"),

env/integtest_util.py

@@ -92,14 +92,14 @@ def get_default_metadata() -> TuningAgentMetadata:
 class MockTuningAgent(TuningAgent):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
-        self.config_to_return: Optional[DBMSConfigDelta] = None
+        self.delta_to_return: Optional[DBMSConfigDelta] = None
 
     def _get_metadata(self) -> TuningAgentMetadata:
         return IntegtestWorkspace.get_default_metadata()
 
     def _step(self) -> DBMSConfigDelta:
-        assert self.config_to_return is not None
-        ret = self.config_to_return
-        # Setting this ensures you must set self.config_to_return every time.
-        self.config_to_return = None
+        assert self.delta_to_return is not None
+        ret = self.delta_to_return
+        # Setting this ensures you must set self.delta_to_return every time.
+        self.delta_to_return = None
         return ret

env/replay.py

@@ -1,18 +1,24 @@
-import logging
+from collections import defaultdict
 from pathlib import Path
 
 from env.pg_conn import PostgresConn
 from env.tuning_agent import TuningAgentArtifactsReader
 from env.workload import Workload
-from util.log import DBGYM_OUTPUT_LOGGER_NAME
 from util.pg import DEFAULT_POSTGRES_PORT
 from util.workspace import DBGymConfig
 
 
-# TODO: make it return the full replay data.
-def replay(dbgym_cfg: DBGymConfig, tuning_agent_artifacts_dpath: Path) -> None:
-    reader = TuningAgentArtifactsReader(tuning_agent_artifacts_dpath)
+def replay(
+    dbgym_cfg: DBGymConfig, tuning_agent_artifacts_dpath: Path
+) -> list[tuple[float, int]]:
+    """
+    Returns the total runtime and the number of timed out queries for each step.
 
+    The first step will use no configuration changes.
+    """
+    replay_data: list[tuple[float, int]] = []
+
+    reader = TuningAgentArtifactsReader(tuning_agent_artifacts_dpath)
     pg_conn = PostgresConn(
         dbgym_cfg,
         DEFAULT_POSTGRES_PORT,
@@ -28,21 +34,40 @@ def replay(dbgym_cfg: DBGymConfig, tuning_agent_artifacts_dpath: Path) -> None:
 
     pg_conn.restore_pristine_snapshot()
     pg_conn.restart_postgres()
-    total_runtime, num_timed_out_queries = time_workload(pg_conn, workload)
-    print(f"Total runtime: {total_runtime / 1e6} seconds")
-    print(f"Number of timed out queries: {num_timed_out_queries}")
+    qknobs: defaultdict[str, list[str]] = defaultdict(list)
+    replay_data.append(time_workload(pg_conn, workload, qknobs))
+
+    for delta in reader.get_all_deltas_in_order():
+        pg_conn.restart_with_changes(delta.sysknobs)
+
+        for index in delta.indexes:
+            pg_conn.psql(index)
+
+        for query, knobs in delta.qknobs.items():
+            # TODO: account for deleting a knob if we are representing knobs as deltas.
+            qknobs[query].extend(knobs)
+
+        replay_data.append(time_workload(pg_conn, workload, qknobs))
+
     pg_conn.shutdown_postgres()
+    return replay_data
 
 
-def time_workload(pg_conn: PostgresConn, workload: Workload) -> tuple[float, int]:
+def time_workload(
+    pg_conn: PostgresConn, workload: Workload, qknobs: dict[str, list[str]]
+) -> tuple[float, int]:
     """
-    It returns the total runtime and the number of timed out queries.
+    Returns the total runtime and the number of timed out queries.
     """
     total_runtime: float = 0
     num_timed_out_queries: int = 0
 
-    for query in workload.get_queries_in_order():
-        runtime, did_time_out, _ = pg_conn.time_query(query)
+    for qid in workload.get_query_order():
+        query = workload.get_query(qid)
+        this_query_knobs = qknobs[qid]
+        runtime, did_time_out, _ = pg_conn.time_query(
+            query, query_knobs=this_query_knobs
+        )
         total_runtime += runtime
         if did_time_out:
             num_timed_out_queries += 1

env/tuning_agent.py

@@ -9,6 +9,7 @@
 # These types are only given as the outputs of tuning agents.
 IndexesDelta = NewType("IndexesDelta", list[str])
 SysKnobsDelta = NewType("SysKnobsDelta", dict[str, str])
+# TODO: I'm not decided whether these should be deltas or full configs. I'm going to figure this out once I integrate Proto-X and UDO.
 QueryKnobsDelta = NewType("QueryKnobsDelta", dict[str, list[str]])
 
 
@@ -67,7 +68,7 @@ class DBMSConfigDelta:
     qknobs: QueryKnobsDelta
 
 
-def get_step_delta_fpath(tuning_agent_artifacts_dpath: Path, step_num: int) -> Path:
+def get_delta_at_step_fpath(tuning_agent_artifacts_dpath: Path, step_num: int) -> Path:
     return tuning_agent_artifacts_dpath / f"step{step_num}_delta.json"
 
 
@@ -96,7 +97,7 @@ def step(self) -> None:
         curr_step_num = self.next_step_num
         self.next_step_num += 1
         dbms_cfg_delta = self._step()
-        with get_step_delta_fpath(
+        with get_delta_at_step_fpath(
             self.tuning_agent_artifacts_dpath, curr_step_num
         ).open("w") as f:
             json.dump(asdict(dbms_cfg_delta), f)
@@ -123,7 +124,7 @@ def __init__(self, tuning_agent_artifacts_dpath: Path) -> None:
         self.tuning_agent_artifacts_dpath = tuning_agent_artifacts_dpath
         assert is_fully_resolved(self.tuning_agent_artifacts_dpath)
         num_steps = 0
-        while get_step_delta_fpath(
+        while get_delta_at_step_fpath(
             self.tuning_agent_artifacts_dpath, num_steps
         ).exists():
             num_steps += 1
@@ -141,9 +142,9 @@ def get_metadata(self) -> TuningAgentMetadata:
                 pgbin_path=Path(data["pgbin_path"]),
             )
 
-    def get_step_delta(self, step_num: int) -> DBMSConfigDelta:
+    def get_delta_at_step(self, step_num: int) -> DBMSConfigDelta:
         assert step_num >= 0 and step_num < self.num_steps
-        with get_step_delta_fpath(self.tuning_agent_artifacts_dpath, step_num).open(
+        with get_delta_at_step_fpath(self.tuning_agent_artifacts_dpath, step_num).open(
             "r"
         ) as f:
             data = json.load(f)
@@ -153,5 +154,5 @@ def get_step_delta(self, step_num: int) -> DBMSConfigDelta:
                 qknobs=data["qknobs"],
             )
 
-    def get_all_deltas(self) -> list[DBMSConfigDelta]:
-        return [self.get_step_delta(step_num) for step_num in range(self.num_steps)]
+    def get_all_deltas_in_order(self) -> list[DBMSConfigDelta]:
+        return [self.get_delta_at_step(step_num) for step_num in range(self.num_steps)]