[Refactor] Remove config for some LTune modules (#8)

* [Refactor] Remove config for LTuneDataSet

* [Refactor] Remove LTune generator and model builder reliance on config

Author: ephoris

endure.toml

@@ -312,8 +312,6 @@ norm_layer = "Batch"
 
 categorical_mode = "reinmax"
 
-k_clip = true
-
 # kwargs specific to LTune models during forward pass
 [ltune.model.train_kwargs]
 temp = 1

endure/ltune/data/dataset.py

@@ -1,44 +1,33 @@
 import glob
-import logging
 import numpy as np
 import os
 import pandas as pd
 import pyarrow.parquet as pa
 import torch
 import torch.utils.data
-from typing import Any
 
+from endure.ltune.data.input_features import kINPUT_FEATS
 
-class LTuneIterableDataSet(torch.utils.data.IterableDataset):
+
+class LTuneDataSet(torch.utils.data.IterableDataset):
     def __init__(
         self,
-        config: dict[str, Any],
         folder: str,
         format: str = "parquet",
         shuffle: bool = False,
-    ):
-        self.log = logging.getLogger(config["log"]["name"])
-        self._config = config
+    ) -> None:
         self._format = format
         self._fnames = glob.glob(os.path.join(folder, "*." + format))
         self._shuffle = shuffle
 
     def _get_input_cols(self):
-        return self._config["ltune"]["input_features"]
+        return kINPUT_FEATS
 
     def _load_data(self, fname):
         if self._format == "parquet":
             df = pa.read_table(fname).to_pandas()
         else:
             df = pd.read_csv(fname)
-        if self._config["ltune"]["data"]["normalize_inputs"]:
-            df = self._normalize_df(df)
-
-        return df
-
-    def _normalize_df(self, df):
-        df[["z0", "z1", "q", "w"]] -= [0.5, 0.5, 0.5, 0.5]
-        df[["z0", "z1", "q", "w"]] /= [0.3, 0.3, 0.3, 0.3]
 
         return df
 

endure/ltune/data/generator.py

@@ -1,22 +1,26 @@
-from typing import Any, Union
-import logging
+from typing import List, Tuple, Union
 
 import numpy as np
 
-from endure.lsm.types import LSMDesign, System, Policy
+from endure.lsm.types import System
+from endure.ltune.data.input_features import kSYSTEM_HEADER, kWORKLOAD_HEADER
 
 
-class LTuneGenerator:
+class LTuneDataGenerator:
     def __init__(
         self,
-        config: dict[str, Any],
-        format: str = "parquet",
+        page_sizes: List[int] = [4, 8, 16],
+        entry_sizes: List[int] = [1024, 2048, 4096, 8192],
+        memory_budget_range: Tuple[float, float] = (5.0, 20.0),
+        selectivity_range: Tuple[float, float] = (1e-7, 1e-9),
+        elements_range: Tuple[int, int] = (100000000, 1000000000),
         precision: int = 3,
     ) -> None:
-        self.log = logging.getLogger(config["log"]["name"])
-        self._config = config
-        self._header = self._gen_workload_header() + self._gen_system_header()
-        self.format = format
+        self.entry_sizes = entry_sizes
+        self.memory_budget_range = memory_budget_range
+        self.page_sizes = page_sizes
+        self.selectivity_range = selectivity_range
+        self.elements_range = elements_range
         self.precision = precision
 
     def _sample_workload(self, dimensions: int) -> list:
@@ -31,25 +35,25 @@ def _sample_workload(self, dimensions: int) -> list:
     # TODO: Will want to configure environment to simulate larger ranges over
     # potential system values
     def _sample_entry_per_page(self, entry_size: int = 8192) -> int:
+        # Potential page sizes are 4KB, 8KB, 16KB
         KB_TO_BITS = 8 * 1024
-        page_sizes = np.array(self._config["generator"]["page_sizes"])
+        page_sizes = np.array(self.page_sizes)
         entries_per_page = (page_sizes * KB_TO_BITS) / entry_size
         return np.random.choice(entries_per_page)
 
     def _sample_selectivity(self) -> float:
-        low, high = self._config["generator"]["selectivity_range"]
+        low, high = self.selectivity_range
         return (high - low) * np.random.rand() + low
 
     def _sample_entry_size(self) -> int:
-        choices = self._config["generator"]["entry_sizes"]
-        return np.random.choice(choices)
+        return np.random.choice(self.entry_sizes)
 
     def _sample_memory_budget(self) -> float:
-        low, high = self._config["generator"]["memory_budget"]
+        low, high = self.memory_budget_range
         return (high - low) * np.random.rand() + low
 
     def _sample_total_elements(self) -> int:
-        low, high = self._config["generator"]["elements_range"]
+        low, high = self.elements_range
         return np.random.randint(low=low, high=high)
 
     def _sample_system(self) -> System:
@@ -63,10 +67,10 @@ def _sample_system(self) -> System:
         return system
 
     def _gen_system_header(self) -> list:
-        return ["B", "s", "E", "H", "N"]
+        return kSYSTEM_HEADER
 
     def _gen_workload_header(self) -> list:
-        return ["z0", "z1", "q", "w"]
+        return kWORKLOAD_HEADER
 
     def generate_header(self) -> list:
         return self._gen_workload_header() + self._gen_system_header()
@@ -89,22 +93,11 @@ def generate_row_csv(self) -> list:
 
         return line
 
-    def generate_row_parquet(self) -> dict[str, Union[int, float]]:
+    def generate_row(self) -> dict[str, Union[int, float]]:
         header = self.generate_header()
         row = self.generate_row_csv()
         line = {}
         for key, val in zip(header, row):
             line[key] = val
 
         return line
-
-    def generate_row(
-        self,
-        row_type: str = "parquet"
-    ) -> Union[list, dict[str, Union[int, float]]]:
-        if row_type == "parquet":
-            row = self.generate_row_parquet()
-        else:  # format == 'csv'
-            row = self.generate_row_csv()
-
-        return row

endure/ltune/data/input_features.py

@@ -0,0 +1,16 @@
+kSYSTEM_HEADER = [
+    "entry_p_page",
+    "selec",
+    "entry_size",
+    "max_h",
+    "num_elem"
+]
+
+kWORKLOAD_HEADER = [
+    "z0",
+    "z1",
+    "q",
+    "w",
+]
+
+kINPUT_FEATS = kSYSTEM_HEADER + kWORKLOAD_HEADER

endure/ltune/loss.py

@@ -6,6 +6,7 @@
 import toml
 
 from endure.lcm.model.builder import LearnedCostModelBuilder
+from endure.lsm.types import STR_POLICY_DICT
 
 
 class LearnedCostModelLoss(torch.nn.Module):
@@ -24,7 +25,9 @@ def __init__(self, config: dict[str, Any], model_path: str):
             max_levels=lcm_cfg["lsm"]["max_levels"],
             **lcm_cfg["lcm"]["model"],
         )
-        lcm_model = lcm_cfg["job"]["LCMTrain"]["model"]
+        lcm_model = STR_POLICY_DICT.get(lcm_cfg["lsm"]["design"], None)
+        if lcm_model is None:
+            raise TypeError(f"Illegal LCM model choice: {lcm_model=}")
         self.model = self.lcm_builder.build_model(lcm_model)
 
         data = torch.load(

endure/ltune/model/builder.py

@@ -1,58 +1,65 @@
 import torch
 import logging
-from typing import Any, Optional
+from typing import Any, Optional, Tuple, Type
 from torch import nn
-from reinmax import reinmax
+from endure.lsm.types import Policy
 
 from endure.ltune.model import ClassicTuner, QLSMTuner, KapLSMTuner
+from endure.ltune.data.input_features import kINPUT_FEATS
 
 
 class LTuneModelBuilder:
-    def __init__(self, config: dict[str, Any]):
-        self._config = config
-        self.log = logging.getLogger(self._config["log"]["name"])
+    def __init__(
+        self,
+        hidden_length: int = 1,
+        hidden_width: int = 64,
+        norm_layer: str = "Batch",
+        dropout: float = 0.0,
+        categorical_mode: str = "gumbel",
+        size_ratio_range: Tuple[int, int] = (2, 31),
+        max_levels: int = 16,
+    ) -> None:
+        self.hidden_length = hidden_length
+        self.hidden_width = hidden_width
+        self.dropout = dropout
+        self.categorical_mode = categorical_mode
+        self.max_levels = max_levels
+        self.size_ratio_min, self.size_ratio_max = size_ratio_range
+        self.capacity_range = self.size_ratio_max - self.size_ratio_min + 1
+
+        self.norm_layer = nn.BatchNorm1d
+        if norm_layer == "Layer":
+            self.norm_layer = nn.LayerNorm
+
         self._models = {
-            # "Tier": ClassicTuner,
-            # "Level": ClassicTuner,
-            "KLSM": KapLSMTuner,
-            "Classic": ClassicTuner,
-            "QLSM": QLSMTuner,
+            Policy.Classic: ClassicTuner,
+            Policy.QFixed: QLSMTuner,
+            Policy.KHybrid: KapLSMTuner,
         }
 
     def get_choices(self):
         return self._models.keys()
 
-    def build_model(self, choice: Optional[str] = None) -> torch.nn.Module:
-        lsm_design: str = self._config["lsm"]["design"]
-        if choice is None:
-            choice = lsm_design
-
-        model_params = self._config["ltune"]["model"]
-        capacity_range = (
-            self._config["lsm"]["size_ratio"]["max"] -
-            self._config["lsm"]["size_ratio"]["min"] + 1
-        )
-        args = {
-            'num_feats': len(self._config["ltune"]["input_features"]),
-            'capacity_range': capacity_range,
-            'hidden_length': model_params["hidden_length"],
-            'hidden_width': model_params["hidden_width"],
-            'dropout_percentage': model_params["dropout"],
-        }
+    def build_model(self, policy: Policy) -> torch.nn.Module:
+        feat_list = kINPUT_FEATS
 
-        if model_params["norm_layer"] == "Batch":
-            args['norm_layer'] = nn.BatchNorm1d
-        elif model_params["norm_layer"] == "Layer":
-            args['norm_layer'] = nn.LayerNorm
+        kwargs = {
+            "num_feats": len(feat_list),
+            "capacity_range": self.capacity_range,
+            "hidden_length": self.hidden_length,
+            "hidden_width": self.hidden_width,
+            "dropout_percentage": self.dropout,
+            "norm_layer": self.norm_layer,
+        }
 
-        model_class = self._models.get(choice, None)
+        model_class = self._models.get(policy, None)
         if model_class is None:
-            raise NotImplementedError(f"Model for LSM Design not implemented yet")
+            raise NotImplementedError(f"Tuner for LSM Design not implemented.")
 
         if model_class is KapLSMTuner:
-            args['num_kap'] = self._config['lsm']['max_levels']
-            args['categorical_mode'] = model_params.get('categorical_mode', 'gumbel')
+            kwargs["num_kap"] = self.max_levels
+            kwargs["categorical_mode"] = self.categorical_mode
 
-        model = model_class(**args)
+        model = model_class(**kwargs)
 
         return model

endure/ltune/util/ltune_eval.py

@@ -7,7 +7,7 @@
 from endure.lcm.util import eval_lcm_impl
 from endure.lsm.cost import EndureCost
 from endure.lsm.types import LSMDesign, System, Policy, STR_POLICY_DICT
-from endure.ltune.data.generator import LTuneGenerator
+from endure.ltune.data.generator import LTuneDataGenerator
 from endure.ltune.loss import LearnedCostModelLoss
 import endure.lsm.solver as Solver
 
@@ -20,7 +20,7 @@ def __init__(
         design_type: str = "Level",
     ) -> None:
         self.policy = STR_POLICY_DICT.get(design_type, Policy.KHybrid)
-        self.gen = LTuneGenerator(config)
+        self.gen = LTuneDataGenerator()
         self.loss = LearnedCostModelLoss(
             config,
             config["job"]["LTuneTrain"]["loss_fn_path"]

jobs/ltune_data_gen.py

@@ -9,7 +9,7 @@
 import pyarrow.parquet as pq
 
 from endure.data.io import Reader
-from endure.ltune.data.generator import LTuneGenerator
+from endure.ltune.data.generator import LTuneDataGenerator
 
 
 class LTuneDataGenJob:
@@ -23,7 +23,7 @@ def __init__(self, config):
         )
 
     def _choose_generator(self):
-        return LTuneGenerator(self.config)
+        return LTuneDataGenerator()
 
     def generate_csv_file(self, generator, idx: int, pos: int) -> int:
         fname_prefix = self.setting["file_prefix"]
@@ -52,7 +52,7 @@ def generate_csv_file(self, generator, idx: int, pos: int) -> int:
         return idx
 
     def generate_parquet_file(
-        self, generator: LTuneGenerator, idx: int, pos: int
+        self, generator: LTuneDataGenerator, idx: int, pos: int
     ) -> int:
         fname_prefix = self.setting["file_prefix"]
         fname = f"{fname_prefix}-{idx:04}.parquet"
@@ -71,7 +71,7 @@ def generate_parquet_file(
             ncols=80,
             disable=self.config["log"]["disable_tqdm"],
         ):
-            table.append(generator.generate_row_parquet())
+            table.append(generator.generate_row())
         table = pa.Table.from_pylist(table)
         pq.write_table(table, fpath)