[Bug] Fix KLSM masking issue (#30)

* [Bug] Fix masking for KLSM LTuner model
* [Minor] Change table creation on mlos experiment

endure/ltune/model/kaplsm_tuner.py

@@ -127,12 +127,14 @@ def _forward_impl(self, x: Tensor, temp=1e-3, hard=False) -> Tensor:
 
         mask = nn.functional.one_hot(max_levels, num_classes=self.num_kap)
         cum_sum = torch.cumsum(mask, dim=1)
-        mask = mask - cum_sum + cum_sum[-1:None]  # Reverse cumulative sum
-        k = mask.unsqueeze(-1) * k
+        mask = 1 - cum_sum + mask  # Sets all values BEFORE max_level to 1
         default = torch.zeros(self.capacity_range)
         default[0] = 1
         default = default.to(k.device).to(torch.long)
-        k[k.sum(-1) == 0] += default
+        # Set K values outside actual level to 1
+        k = mask.unsqueeze(-1) * k
+        k[mask == 0] += default
+
         k = torch.flatten(k, start_dim=1)
 
         out = torch.concat([bits, t, k], dim=-1)

jobs/ltune_train.py

@@ -181,8 +181,7 @@ def run(self) -> Optional[Trainer]:
 
         return trainer
 
-
-if __name__ == "__main__":
+def main():
     from endure.data.io import Reader
 
     config = Reader.read_config("endure.toml")
@@ -194,3 +193,7 @@ def run(self) -> Optional[Trainer]:
 
     job = LTuneTrainJob(config)
     job.run()
+
+
+if __name__ == "__main__":
+    main()

jobs/mlos_exp_runs.py

@@ -10,7 +10,7 @@
 from mlos_core.optimizers import SmacOptimizer
 
 NUM_ROUNDS = 100
-NUM_TRIALS = 10
+NUM_TRIALS = 100
 
 
 class ExperimentMLOS:
@@ -41,12 +41,12 @@ def _create_parameter_space(self, system: System) -> CS.ConfigurationSpace:
             CS.UniformIntegerHyperparameter(
                 name="size_ratio",
                 lower=self.bounds.size_ratio_range[0],
-                upper=self.bounds.size_ratio_range[1],
+                upper=self.bounds.size_ratio_range[1] - 1,  # ConfigSpace is inclusive
             ),
         ]
         kap_params = [
             CS.UniformIntegerHyperparameter(
-                name=f"kap_{i}", lower=1, upper=self.bounds.size_ratio_range[1] - 1
+                name=f"kap_{i}", lower=1, upper=self.bounds.size_ratio_range[1] - 2
             )
             for i in range(self.bounds.max_considered_levels)
         ]
@@ -59,7 +59,9 @@ def _create_parameter_space(self, system: System) -> CS.ConfigurationSpace:
 
     def _create_optimizer(self, parameter_space: CS.ConfigurationSpace):
         return SmacOptimizer(
-            parameter_space=parameter_space, optimization_targets=["cost"]
+            parameter_space=parameter_space,
+            optimization_targets=["cost"],
+            n_random_init=1,
         )
 
     def _train_model(
@@ -89,6 +91,7 @@ def _train_model(
 
     def run(self) -> None:
         system = System()
+        self.db.create_tables()
         for rep_wl in self.config["workloads"]:
             workload = Workload(
                 z0=rep_wl["z0"],
@@ -108,10 +111,13 @@ def run(self) -> None:
 
 class MLOSDatabase:
     def __init__(self, config: dict, db_path: str = "mlos_exp.db") -> None:
-        self.config = config
         self.log: logging.Logger = logging.getLogger(config["log"]["name"])
-        connector = sqlite3.connect(db_path)
-        cursor = connector.cursor()
+        self.connector = sqlite3.connect(db_path)
+        self.db_path = db_path
+        self.config = config
+
+    def create_tables(self) -> None:
+        cursor = self.connector.cursor()
         cursor.execute(
             """
             CREATE TABLE IF NOT EXISTS environments (
@@ -147,10 +153,10 @@ def __init__(self, config: dict, db_path: str = "mlos_exp.db") -> None:
             );
             """
         )
-        connector.commit()
+        self.connector.commit()
         cursor.close()
-        self.log.info(f"Created database: {db_path}")
-        self.connector = connector
+
+        return
 
     def log_workload(self, workload: Workload, system: System) -> int:
         cursor = self.connector.cursor()