Quick code clean up

experiments/03-python-overhead/run_cpp.py

@@ -12,7 +12,7 @@ def main():
     parser.add_argument("--iters", type=int, default=50)
     args = parser.parse_args()
 
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     db = DBType.from_str(args.dbname)
     cstr = config.get_odbc_connection_string(db)
 

experiments/14-analytical-perf-cost/ana_runner.py

@@ -56,7 +56,7 @@ def noop(_signal, _frame):
     else:
         out_dir = pathlib.Path(".")
 
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     directory = Directory(config)
     asyncio.run(directory.refresh())
     ec = EngineConnections.connect_sync(
@@ -120,7 +120,7 @@ def noop(_signal, _frame):
 
 
 def run_warmup(args, query_bank: List[str], queries: List[int], engine: Engine):
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     directory = Directory(config)
     asyncio.run(directory.refresh())
     ec = EngineConnections.connect_sync(

src/brad/admin/bootstrap_schema.py

@@ -48,7 +48,7 @@ def register_admin_action(subparser) -> None:
 # This method is called by `brad.exec.admin.main`.
 def bootstrap_schema(args):
     # 1. Load the config.
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
 
     # 2. Load and validate the user-provided schema.
     user = UserProvidedBlueprint.load_from_yaml_file(args.schema_file)

src/brad/admin/bulk_load.py

@@ -328,7 +328,7 @@ async def _truncate_aurora_tables(
 
 
 async def bulk_load_impl(args, manifest: Dict[str, Any]) -> None:
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     assets = AssetManager(config)
     blueprint_mgr = BlueprintManager(config, assets, manifest["schema_name"])
     await blueprint_mgr.load()

src/brad/admin/control.py

@@ -40,7 +40,7 @@ def register_admin_action(subparser) -> None:
 
 async def control_impl(args) -> None:
     # 1. Load the config, blueprint, and provisioning.
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     assets = AssetManager(config)
 
     blueprint_mgr = BlueprintManager(config, assets, args.schema_name)

src/brad/admin/drop_schema.py

@@ -33,7 +33,7 @@ def register_admin_action(subparser) -> None:
 # This method is called by `brad.exec.admin.main`.
 def drop_schema(args):
     # 1. Load the config and blueprint.
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
 
     # 2. Connect to the underlying engines without an explicit database.
     directory = Directory(config)

src/brad/admin/modify_blueprint.py

@@ -186,7 +186,7 @@ async def run_transition(
 # This method is called by `brad.exec.admin.main`.
 def modify_blueprint(args):
     # 1. Load the config.
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
 
     # 2. Load the existing blueprint.
     assets = AssetManager(config)

src/brad/admin/run_on.py

@@ -40,7 +40,7 @@ def register_admin_action(subparser) -> None:
 def run_on(args):
     # 1. Load the config and blueprint.
     engine = Engine.from_str(args.engine)
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     assets = AssetManager(config)
     blueprint_mgr = BlueprintManager(config, assets, args.schema_name)
     blueprint_mgr.load_sync()

src/brad/admin/run_planner.py

@@ -127,7 +127,7 @@ def run_planner(args) -> None:
     independently of the rest of BRAD.
     """
     # 1. Load the config.
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
 
     # 2. Load the planner config.
     planner_config = PlannerConfig(args.planner_config_file)

src/brad/admin/train_router.py

@@ -111,7 +111,7 @@ async def set_up_estimator(
 # This method is called by `brad.exec.admin.main`.
 def train_router(args):
     schema_name = extract_schema_name(args.schema_file)
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     policy = RoutingPolicy.from_str(args.policy)
 
     if args.std_dataset_path is not None:

src/brad/admin/workload_logs.py

@@ -59,7 +59,7 @@ def workload_logs(args) -> None:
 
 
 def inspect_logs(args) -> None:
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
 
     timestamp_format = "%Y-%m-%d %H:%M:%S"
     window_start = datetime.strptime(args.window_start, timestamp_format)

src/brad/blueprint/provisioning.py

@@ -1,20 +1,14 @@
 class Provisioning:
-    def __init__(
-        self, instance_type: str, num_nodes: int, is_paused: bool = False
-    ) -> None:
+    def __init__(self, instance_type: str, num_nodes: int) -> None:
         self._instance_type = instance_type
         self._num_nodes = num_nodes
-        self._is_paused = is_paused
 
     def instance_type(self) -> str:
         return self._instance_type
 
     def num_nodes(self) -> int:
         return self._num_nodes
 
-    def is_paused(self) -> bool:
-        return self._is_paused
-
     def clone(self) -> "Provisioning":
         return Provisioning(self._instance_type, self._num_nodes)
 
@@ -27,8 +21,6 @@ def __repr__(self) -> str:
                 self._instance_type,
                 "(",
                 str(self._num_nodes),
-                ", ",
-                str(self._is_paused),
                 ")",
             ]
         )

src/brad/calibration/measure_load.py

@@ -142,7 +142,7 @@ def main() -> None:
         return
 
     schema_name = os.environ[args.schema_name_var]
-    config = ConfigFile(os.environ[args.config_file_var])
+    config = ConfigFile.load(os.environ[args.config_file_var])
     queries = load_queries(args.query_file)
 
     directory = Directory(config)

src/brad/config/file.py

@@ -3,7 +3,7 @@
 import pathlib
 import re
 import yaml
-from typing import Optional, Dict
+from typing import Optional, Dict, Any
 from datetime import timedelta
 
 from brad.config.engine import Engine
@@ -13,10 +13,14 @@
 
 
 class ConfigFile:
-    def __init__(self, path: str):
-        self._raw_path = path
-        with open(path, "r", encoding="UTF-8") as file:
-            self._raw = yaml.load(file, Loader=yaml.Loader)
+    @classmethod
+    def load(cls, file_path: str) -> "ConfigFile":
+        with open(file_path, "r", encoding="UTF-8") as file:
+            raw = yaml.load(file, Loader=yaml.Loader)
+        return cls(raw)
+
+    def __init__(self, raw_parsed: Dict[str, Any]):
+        self._raw = raw_parsed
 
     def get_cluster_ids(self) -> Dict[Engine, str]:
         return {
@@ -25,10 +29,6 @@ def get_cluster_ids(self) -> Dict[Engine, str]:
             Engine.Athena: "brad-db0",  # TODO(Amadou): I don't want to break existing configs. Coordinate with Geoff on this.
         }
 
-    @property
-    def raw_path(self) -> str:
-        return self._raw_path
-
     @property
     def daemon_log_path(self) -> Optional[pathlib.Path]:
         return self._extract_log_path("daemon_log_file")

src/brad/daemon/daemon.py

@@ -217,7 +217,7 @@ async def _run_setup(self) -> None:
                 target=start_front_end,
                 args=(
                     fe_index,
-                    self._config.raw_path,
+                    self._config,
                     self._schema_name,
                     self._path_to_planner_config,
                     self._debug_mode,

src/brad/exec/daemon.py

@@ -71,7 +71,7 @@ def main(args):
     # descriptors!).
     mp.set_start_method("spawn")
 
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     temp_config = (
         TempConfig.load_from_file(args.temp_config_file)
         if args.temp_config_file is not None

src/brad/front_end/start_front_end.py

@@ -12,7 +12,7 @@
 
 def start_front_end(
     fe_index: int,
-    config_path: str,
+    config: ConfigFile,
     schema_name: str,
     path_to_planner_config: str,
     debug_mode: bool,
@@ -23,7 +23,6 @@ def start_front_end(
     Schedule this method to run in a child process to launch a BRAD front
     end server.
     """
-    config = ConfigFile(config_path)
     set_up_logging(filename=config.front_end_log_file(fe_index), debug_mode=debug_mode)
 
     event_loop = asyncio.new_event_loop()

src/brad/provisioning/physical.py

@@ -89,11 +89,11 @@ def should_trigger_replan(self, overrides=None) -> bool:
     # Update physical provisioning.
     def update_blueprint(self, new_blueprint: Blueprint):
         aurora_instance_type = new_blueprint.aurora_provisioning().instance_type()
-        aurora_paused = new_blueprint.aurora_provisioning().is_paused()
         aurora_instance_count = new_blueprint.aurora_provisioning().num_nodes()
+        aurora_paused = aurora_instance_count == 0
         redshift_instance_type = new_blueprint.redshift_provisioning().instance_type()
         redshift_instance_count = new_blueprint.redshift_provisioning().num_nodes()
-        redshift_paused = new_blueprint.redshift_provisioning().is_paused()
+        redshift_paused = redshift_instance_count == 0
         print("Rescaling Aurora...")
         self._rds_provisioning.rescale(
             immediate=True,

tests/test_neighborhood.py

@@ -44,7 +44,7 @@ def test_neighborhood_change():
     initial = bootstrap_blueprint(user)
     print()  # Flush stdout.
     print("Running test")
-    config = ConfigFile(
+    config = ConfigFile.load(
         "config.yml"
     )  # TODO: Support configs in tests. This will not work.
     assets = AssetManager(config)

tests/test_workload.py

@@ -3,7 +3,7 @@
 
 
 def test_read_from_s3():
-    config = ConfigFile("./config/config.yml")
+    config = ConfigFile.load("./config/config.yml")
 
     workload = workload_from_s3_logs(config, 3)
 

tools/calibration/table_sizes.py

@@ -30,7 +30,7 @@ def delete_s3_object(client, bucket: str, key: str) -> None:
 
 
 async def main_impl(args) -> None:
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     assets = AssetManager(config)
     mgr = BlueprintManager(config, assets, args.schema_name)
     await mgr.load()

tools/calibration/transactions/imdb_extended/retrieve_metrics.py

@@ -98,7 +98,7 @@ def main():
             instance_identifier=args.instance_id
         )
     elif args.config_file is not None:
-        config = ConfigFile(args.config_file)
+        config = ConfigFile.load(args.config_file)
         directory = Directory(config)
         asyncio.run(directory.refresh())
         client = PerfInsightsClient(resource_id=directory.aurora_writer().resource_id())

tools/data_collection/collect_athena_data.py

@@ -205,7 +205,7 @@ def main():
             query_start_offset,
         )
 
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     client = boto3.client(
         "athena",
         aws_access_key_id=config.aws_access_key,

tools/one_off/rename_indexes.py

@@ -13,7 +13,7 @@ def main():
     args = parser.parse_args()
 
     bp = UserProvidedBlueprint.load_from_yaml_file(args.schema_file)
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     ecs = EngineConnections.connect_sync(
         config, bp.schema_name, autocommit=False, specific_engines={Engine.Aurora}
     )

tools/query_dataset/adjust_queries.py

@@ -469,7 +469,7 @@ def main():
 
     table_cols, table_indexed_cols = load_schema(args.schema_file)
 
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     directory = Directory(config)
     asyncio.run(directory.refresh())
     conn = ConnectionFactory.connect_to_sync(

tools/query_dataset/extract_query_stats.py

@@ -95,7 +95,7 @@ def main():
     parser.add_argument("--queries-file", type=str, required=True)
     args = parser.parse_args()
 
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     conn = asyncio.run(ConnectionFactory.connect_to_sidecar(args.schema_name, config))
     cursor = conn.cursor_sync()
 

tools/query_dataset/run_inference.py

@@ -35,7 +35,7 @@ def main():
 
     engine = Engine.from_str(args.engine)
     queries = load_queries(args.queries_file)
-    config = ConfigFile(args.config_file)
+    config = ConfigFile.load(args.config_file)
     conn = asyncio.run(ConnectionFactory.connect_to_sidecar(args.schema_name, config))
 
     model = TrainedModel.load(

workloads/IMDB_extended/run_transactions.py

@@ -244,7 +244,7 @@ def main():
     if args.brad_direct:
         assert args.config_file is not None
         assert args.schema_name is not None
-        config = ConfigFile(args.config_file)
+        config = ConfigFile.load(args.config_file)
         directory = Directory(config)
         asyncio.run(directory.refresh())
     else:

workloads/IMDB_extended/workload_utils/connect.py

@@ -28,7 +28,7 @@ def connect_to_db(
         assert args.schema_name is not None
         assert args.config_file is not None
 
-        config = ConfigFile(args.config_file)
+        config = ConfigFile.load(args.config_file)
         if directory is None:
             directory_to_use = Directory(config)
             asyncio.run(directory_to_use.refresh())