TuningArtifactsWriter (cmu-db#59)

**Summary**: Changed `TuningAgent` interface to `TuningArtifactsWriter`.

**Details**:
* This makes it much more flexible because you don't need to enclose all
of the tuning logic into a class.

env/integtest_replay.py

@@ -3,7 +3,7 @@
 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 (
+from env.tuning_artifacts import (
     DBMSConfigDelta,
     IndexesDelta,
     QueryKnobsDelta,
@@ -36,7 +36,7 @@ def test_replay(self) -> None:
         )
         agent.step()
         replay_data = replay(
-            IntegtestWorkspace.get_dbgym_cfg(), agent.tuning_agent_artifacts_dpath
+            IntegtestWorkspace.get_dbgym_cfg(), agent.tuning_artifacts_dpath
         )
 
         # We do some very simple sanity checks here due to the inherent randomness of executing a workload.

env/integtest_tuning_agent.py → env/integtest_tuning_artifacts.py

@@ -1,7 +1,7 @@
 import unittest
 
 from env.integtest_util import IntegtestWorkspace, MockTuningAgent
-from env.tuning_agent import (
+from env.tuning_artifacts import (
     DBMSConfigDelta,
     IndexesDelta,
     QueryKnobsDelta,
@@ -33,7 +33,7 @@ def test_get_delta_at_step(self) -> None:
         agent.delta_to_return = PostgresConnTests.make_config("c")
         agent.step()
 
-        reader = TuningAgentArtifactsReader(agent.tuning_agent_artifacts_dpath)
+        reader = TuningAgentArtifactsReader(agent.tuning_artifacts_dpath)
 
         self.assertEqual(
             reader.get_delta_at_step(1), PostgresConnTests.make_config("b")
@@ -58,7 +58,7 @@ def test_get_all_deltas_in_order(self) -> None:
         agent.delta_to_return = PostgresConnTests.make_config("c")
         agent.step()
 
-        reader = TuningAgentArtifactsReader(agent.tuning_agent_artifacts_dpath)
+        reader = TuningAgentArtifactsReader(agent.tuning_artifacts_dpath)
 
         self.assertEqual(
             reader.get_all_deltas_in_order(),
@@ -71,7 +71,7 @@ def test_get_all_deltas_in_order(self) -> None:
 
     def test_get_metadata(self) -> None:
         agent = MockTuningAgent(IntegtestWorkspace.get_dbgym_cfg())
-        reader = TuningAgentArtifactsReader(agent.tuning_agent_artifacts_dpath)
+        reader = TuningAgentArtifactsReader(agent.tuning_artifacts_dpath)
         metadata = reader.get_metadata()
         expected_metadata = IntegtestWorkspace.get_default_metadata()
         self.assertEqual(metadata, expected_metadata)

env/integtest_util.py

@@ -4,7 +4,7 @@
 
 import yaml
 
-from env.tuning_agent import DBMSConfigDelta, TuningAgent, TuningAgentMetadata
+from env.tuning_artifacts import DBMSConfigDelta, TuningAgent, TuningMetadata
 from util.workspace import (
     DBGymConfig,
     fully_resolve_path,
@@ -57,12 +57,12 @@ def get_workspace_path() -> Path:
             return Path(yaml.safe_load(f)["dbgym_workspace_path"])
 
     @staticmethod
-    def get_default_metadata() -> TuningAgentMetadata:
+    def get_default_metadata() -> TuningMetadata:
         dbgym_cfg = IntegtestWorkspace.get_dbgym_cfg()
         workspace_path = fully_resolve_path(
             dbgym_cfg, IntegtestWorkspace.get_workspace_path()
         )
-        return TuningAgentMetadata(
+        return TuningMetadata(
             workload_path=fully_resolve_path(
                 dbgym_cfg,
                 get_default_workload_path(
@@ -94,7 +94,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self.delta_to_return: Optional[DBMSConfigDelta] = None
 
-    def _get_metadata(self) -> TuningAgentMetadata:
+    def _get_metadata(self) -> TuningMetadata:
         return IntegtestWorkspace.get_default_metadata()
 
     def _step(self) -> DBMSConfigDelta:

env/replay.py

@@ -2,14 +2,14 @@
 from pathlib import Path
 
 from env.pg_conn import PostgresConn
-from env.tuning_agent import TuningAgentArtifactsReader
+from env.tuning_artifacts import TuningAgentArtifactsReader
 from env.workload import Workload
 from util.pg import DEFAULT_POSTGRES_PORT
 from util.workspace import DBGymConfig
 
 
 def replay(
-    dbgym_cfg: DBGymConfig, tuning_agent_artifacts_dpath: Path
+    dbgym_cfg: DBGymConfig, tuning_artifacts_dpath: Path
 ) -> list[tuple[float, int]]:
     """
     Returns the total runtime and the number of timed out queries for each step.
@@ -18,7 +18,7 @@ def replay(
     """
     replay_data: list[tuple[float, int]] = []
 
-    reader = TuningAgentArtifactsReader(tuning_agent_artifacts_dpath)
+    reader = TuningAgentArtifactsReader(tuning_artifacts_dpath)
     pg_conn = PostgresConn(
         dbgym_cfg,
         DEFAULT_POSTGRES_PORT,

env/tuning_agent.py → env/tuning_artifacts.py

@@ -5,17 +5,17 @@
 
 from util.workspace import DBGymConfig, is_fully_resolved
 
-# PostgresConn doesn't use these types because PostgresConn is used internally by tuning agents.
-# These types are only given as the outputs of tuning agents.
+# PostgresConn doesn't use these types because PostgresConn is used internally by tuning agents
+# while these types are only used in the interface between the orchestrator and the 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]])
 
 
 @dataclass
-class TuningAgentMetadata:
-    """Metadata for tuning agent"""
+class TuningMetadata:
+    """Metadata for the tuning process."""
 
     workload_path: Path
     pristine_dbdata_snapshot_path: Path
@@ -68,26 +68,26 @@ class DBMSConfigDelta:
     qknobs: QueryKnobsDelta
 
 
-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"
+def get_delta_at_step_fpath(tuning_artifacts_dpath: Path, step_num: int) -> Path:
+    return tuning_artifacts_dpath / f"step{step_num}_delta.json"
 
 
-def get_metadata_fpath(tuning_agent_artifacts_dpath: Path) -> Path:
-    return tuning_agent_artifacts_dpath / "metadata.json"
+def get_metadata_fpath(tuning_artifacts_dpath: Path) -> Path:
+    return tuning_artifacts_dpath / "metadata.json"
 
 
 class TuningAgent:
     def __init__(self, dbgym_cfg: DBGymConfig) -> None:
         self.dbgym_cfg = dbgym_cfg
-        self.tuning_agent_artifacts_dpath = self.dbgym_cfg.cur_task_runs_artifacts_path(
-            "tuning_agent_artifacts", mkdir=True
+        self.tuning_artifacts_dpath = self.dbgym_cfg.cur_task_runs_artifacts_path(
+            "tuning_artifacts", mkdir=True
         )
-        assert is_fully_resolved(self.tuning_agent_artifacts_dpath)
+        assert is_fully_resolved(self.tuning_artifacts_dpath)
         self.next_step_num = 0
 
         # Write metadata file
         metadata = self._get_metadata()
-        with get_metadata_fpath(self.tuning_agent_artifacts_dpath).open("w") as f:
+        with get_metadata_fpath(self.tuning_artifacts_dpath).open("w") as f:
             json.dump(metadata.asdict(), f)
 
     def step(self) -> None:
@@ -97,12 +97,12 @@ def step(self) -> None:
         curr_step_num = self.next_step_num
         self.next_step_num += 1
         dbms_cfg_delta = self._step()
-        with get_delta_at_step_fpath(
-            self.tuning_agent_artifacts_dpath, curr_step_num
-        ).open("w") as f:
+        with get_delta_at_step_fpath(self.tuning_artifacts_dpath, curr_step_num).open(
+            "w"
+        ) as f:
             json.dump(asdict(dbms_cfg_delta), f)
 
-    def _get_metadata(self) -> TuningAgentMetadata:
+    def _get_metadata(self) -> TuningMetadata:
         """
         This should be overridden by subclasses.
 
@@ -120,20 +120,18 @@ def _step(self) -> DBMSConfigDelta:
 
 
 class TuningAgentArtifactsReader:
-    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)
+    def __init__(self, tuning_artifacts_dpath: Path) -> None:
+        self.tuning_artifacts_dpath = tuning_artifacts_dpath
+        assert is_fully_resolved(self.tuning_artifacts_dpath)
         num_steps = 0
-        while get_delta_at_step_fpath(
-            self.tuning_agent_artifacts_dpath, num_steps
-        ).exists():
+        while get_delta_at_step_fpath(self.tuning_artifacts_dpath, num_steps).exists():
             num_steps += 1
         self.num_steps = num_steps
 
-    def get_metadata(self) -> TuningAgentMetadata:
-        with get_metadata_fpath(self.tuning_agent_artifacts_dpath).open("r") as f:
+    def get_metadata(self) -> TuningMetadata:
+        with get_metadata_fpath(self.tuning_artifacts_dpath).open("r") as f:
             data = json.load(f)
-            return TuningAgentMetadata(
+            return TuningMetadata(
                 workload_path=Path(data["workload_path"]),
                 pristine_dbdata_snapshot_path=Path(
                     data["pristine_dbdata_snapshot_path"]
@@ -144,7 +142,7 @@ def get_metadata(self) -> TuningAgentMetadata:
 
     def get_delta_at_step(self, step_num: int) -> DBMSConfigDelta:
         assert step_num >= 0 and step_num < self.num_steps
-        with get_delta_at_step_fpath(self.tuning_agent_artifacts_dpath, step_num).open(
+        with get_delta_at_step_fpath(self.tuning_artifacts_dpath, step_num).open(
             "r"
         ) as f:
             data = json.load(f)