refactored how workspace works in gymlib tests

Author: wangpatrick57

env/tests/gymlib_integtest_util.py

@@ -35,22 +35,20 @@ class GymlibIntegtestManager:
     BENCHMARK = "tpch"
     SCALE_FACTOR = 0.01
     DBGYM_CONFIG_PATH = Path("env/tests/gymlib_integtest_dbgym_config.yaml")
-
-    # This is set at most once by set_up_workspace().
-    DBGYM_WORKSPACE: Optional[DBGymWorkspace] = None
+    WORKSPACE_PATH: Optional[Path] = None
 
     @staticmethod
     def set_up_workspace() -> None:
         """
         Set up the workspace if it has not already been set up.
         None of the integtest_*.py files will delete the workspace so that future tests run faster.
         """
-        workspace_path = get_workspace_path_from_config(
+        GymlibIntegtestManager.WORKSPACE_PATH = get_workspace_path_from_config(
             GymlibIntegtestManager.DBGYM_CONFIG_PATH
         )
 
         # This if statement prevents us from setting up the workspace twice, which saves time.
-        if not workspace_path.exists():
+        if not GymlibIntegtestManager.WORKSPACE_PATH.exists():
             subprocess.run(
                 ["./env/tests/_set_up_gymlib_integtest_workspace.sh"],
                 env={
@@ -64,23 +62,13 @@ def set_up_workspace() -> None:
                 check=True,
             )
 
-        # Once we get here, we have an invariant that the workspace exists. We need this
-        # invariant to be true in order to create the DBGymWorkspace.
-        #
-        # However, it also can't be created more than once so we need to check `is None`.
-        if GymlibIntegtestManager.DBGYM_WORKSPACE is None:
-            # Reset this in case it had been created by a test *not* using GymlibIntegtestManager.set_up_workspace().
-            DBGymWorkspace._num_times_created_this_run = 0
-            GymlibIntegtestManager.DBGYM_WORKSPACE = DBGymWorkspace(workspace_path)
-
     @staticmethod
-    def get_dbgym_workspace() -> DBGymWorkspace:
-        assert GymlibIntegtestManager.DBGYM_WORKSPACE is not None
-        return GymlibIntegtestManager.DBGYM_WORKSPACE
+    def get_workspace_path() -> Path:
+        assert GymlibIntegtestManager.WORKSPACE_PATH is not None
+        return GymlibIntegtestManager.WORKSPACE_PATH
 
     @staticmethod
     def get_default_metadata() -> TuningMetadata:
-        dbgym_workspace = GymlibIntegtestManager.get_dbgym_workspace()
         assert GymlibIntegtestManager.BENCHMARK == "tpch"
         suffix = get_workload_suffix(
             GymlibIntegtestManager.BENCHMARK,
@@ -91,23 +79,25 @@ def get_default_metadata() -> TuningMetadata:
         return TuningMetadata(
             workload_path=fully_resolve_path(
                 get_workload_symlink_path(
-                    dbgym_workspace.dbgym_workspace_path,
+                    GymlibIntegtestManager.get_workspace_path(),
                     GymlibIntegtestManager.BENCHMARK,
                     GymlibIntegtestManager.SCALE_FACTOR,
                     suffix,
                 ),
             ),
             pristine_dbdata_snapshot_path=fully_resolve_path(
                 get_dbdata_tgz_symlink_path(
-                    dbgym_workspace.dbgym_workspace_path,
+                    GymlibIntegtestManager.get_workspace_path(),
                     GymlibIntegtestManager.BENCHMARK,
                     GymlibIntegtestManager.SCALE_FACTOR,
                 ),
             ),
             dbdata_parent_path=fully_resolve_path(
-                get_tmp_path_from_workspace_path(dbgym_workspace.dbgym_workspace_path),
+                get_tmp_path_from_workspace_path(
+                    GymlibIntegtestManager.get_workspace_path()
+                ),
             ),
             pgbin_path=fully_resolve_path(
-                get_pgbin_symlink_path(dbgym_workspace.dbgym_workspace_path),
+                get_pgbin_symlink_path(GymlibIntegtestManager.get_workspace_path()),
             ),
         )

env/tests/integtest_pg_conn.py

@@ -10,12 +10,18 @@
     get_is_postgres_running,
     get_running_postgres_ports,
 )
+from util.workspace import DBGymWorkspace
 
 
 class PostgresConnTests(unittest.TestCase):
+    workspace: DBGymWorkspace
+
     @staticmethod
     def setUpClass() -> None:
         GymlibIntegtestManager.set_up_workspace()
+        PostgresConnTests.workspace = DBGymWorkspace(
+            GymlibIntegtestManager.get_workspace_path()
+        )
 
     def setUp(self) -> None:
         self.assertFalse(
@@ -38,7 +44,7 @@ def tearDown(self) -> None:
 
     def create_pg_conn(self, pgport: int = DEFAULT_POSTGRES_PORT) -> PostgresConn:
         return PostgresConn(
-            GymlibIntegtestManager.get_dbgym_workspace(),
+            PostgresConnTests.workspace,
             pgport,
             self.metadata.pristine_dbdata_snapshot_path,
             self.metadata.dbdata_parent_path,

env/tests/integtest_replay.py

@@ -10,16 +10,22 @@
     SysKnobsDelta,
     TuningArtifactsWriter,
 )
+from util.workspace import DBGymWorkspace
 
 
 class ReplayTests(unittest.TestCase):
+    workspace: DBGymWorkspace
+
     @staticmethod
     def setUpClass() -> None:
         GymlibIntegtestManager.set_up_workspace()
+        ReplayTests.workspace = DBGymWorkspace(
+            GymlibIntegtestManager.get_workspace_path()
+        )
 
     def test_replay(self) -> None:
         writer = TuningArtifactsWriter(
-            GymlibIntegtestManager.get_dbgym_workspace(),
+            ReplayTests.workspace,
             GymlibIntegtestManager.get_default_metadata(),
         )
         writer.write_step(
@@ -41,7 +47,7 @@ def test_replay(self) -> None:
             )
         )
         replay_data = replay(
-            GymlibIntegtestManager.get_dbgym_workspace(),
+            ReplayTests.workspace,
             writer.tuning_artifacts_path,
         )
 

env/tests/integtest_tuning_artifacts.py

@@ -9,13 +9,19 @@
     TuningArtifactsReader,
     TuningArtifactsWriter,
 )
+from util.workspace import DBGymWorkspace
 
 
 class PostgresConnTests(unittest.TestCase):
     @staticmethod
     def setUpClass() -> None:
         GymlibIntegtestManager.set_up_workspace()
 
+    def setUp(self) -> None:
+        # We re-create a workspace for each test because each test will create its own TuningArtifactsWriter.
+        DBGymWorkspace._num_times_created_this_run = 0
+        self.workspace = DBGymWorkspace(GymlibIntegtestManager.get_workspace_path())
+
     @staticmethod
     def make_config(letter: str) -> DBMSConfigDelta:
         return DBMSConfigDelta(
@@ -26,7 +32,7 @@ def make_config(letter: str) -> DBMSConfigDelta:
 
     def test_get_delta_at_step(self) -> None:
         writer = TuningArtifactsWriter(
-            GymlibIntegtestManager.get_dbgym_workspace(),
+            self.workspace,
             GymlibIntegtestManager.get_default_metadata(),
         )
 
@@ -51,7 +57,7 @@ def test_get_delta_at_step(self) -> None:
 
     def test_get_all_deltas_in_order(self) -> None:
         writer = TuningArtifactsWriter(
-            GymlibIntegtestManager.get_dbgym_workspace(),
+            self.workspace,
             GymlibIntegtestManager.get_default_metadata(),
         )
 
@@ -72,7 +78,7 @@ def test_get_all_deltas_in_order(self) -> None:
 
     def test_get_metadata(self) -> None:
         writer = TuningArtifactsWriter(
-            GymlibIntegtestManager.get_dbgym_workspace(),
+            self.workspace,
             GymlibIntegtestManager.get_default_metadata(),
         )
         reader = TuningArtifactsReader(writer.tuning_artifacts_path)

env/tests/integtest_workload.py

@@ -3,17 +3,22 @@
 from benchmark.tpch.constants import DEFAULT_TPCH_SEED, NUM_TPCH_QUERIES
 from env.tests.gymlib_integtest_util import GymlibIntegtestManager
 from env.workload import Workload
+from util.workspace import DBGymWorkspace
 
 
 class WorkloadTests(unittest.TestCase):
+    workspace: DBGymWorkspace
+
     @staticmethod
     def setUpClass() -> None:
         GymlibIntegtestManager.set_up_workspace()
+        WorkloadTests.workspace = DBGymWorkspace(
+            GymlibIntegtestManager.get_workspace_path()
+        )
 
     def test_workload(self) -> None:
         workload_path = GymlibIntegtestManager.get_default_metadata().workload_path
-
-        workload = Workload(GymlibIntegtestManager.get_dbgym_workspace(), workload_path)
+        workload = Workload(WorkloadTests.workspace, workload_path)
 
         # Check the order of query IDs.
         self.assertEqual(

env/tuning_artifacts.py

@@ -1,7 +1,7 @@
 import json
 from dataclasses import asdict, dataclass
 from pathlib import Path
-from typing import Any, NewType, TypedDict
+from typing import Any, NewType
 
 from util.workspace import DBGymWorkspace, is_fully_resolved
 
@@ -84,6 +84,7 @@ def __init__(
         self.tuning_artifacts_path = (
             self.dbgym_workspace.dbgym_this_run_path / "tuning_artifacts"
         )
+        # exist_ok is False because you should only create one TuningArtifactsWriter per run.
         self.tuning_artifacts_path.mkdir(parents=False, exist_ok=False)
         assert is_fully_resolved(self.tuning_artifacts_path)
         self.next_step_num = 0