Moving monkey patching into mlos_core

microsoft · Aug 20, 2024 · e2e2626 · e2e2626
1 parent 6a6c28a
commit e2e2626
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 12 deletions.
diff --git a/mlos_bench/mlos_bench/optimizers/convert_configspace.py b/mlos_bench/mlos_bench/optimizers/convert_configspace.py
@@ -26,10 +26,7 @@
 from mlos_bench.tunables.tunable import Tunable, TunableValue
 from mlos_bench.tunables.tunable_groups import TunableGroups
 from mlos_bench.util import try_parse_val
-from mlos_core.spaces.converters.util import (
-    monkey_patch_hp_quantization,
-    QUANTIZATION_BINS_META_KEY,
-)
+from mlos_core.spaces.converters.util import QUANTIZATION_BINS_META_KEY
 
 _LOG = logging.getLogger(__name__)
 
@@ -80,6 +77,10 @@ def _tunable_to_configspace(
     meta: Dict[Hashable, TunableValue] = {"cost": cost}
     if group_name is not None:
         meta["group"] = group_name
+    if tunable.is_numerical and tunable.quantization_bins:
+        # Temporary workaround to dropped quantization support in ConfigSpace 1.0
+        # See Also: https://github.com/automl/ConfigSpace/issues/390
+        meta[QUANTIZATION_BINS_META_KEY] = tunable.quantization_bins
 
     if tunable.type == "categorical":
         return ConfigurationSpace(
@@ -140,14 +141,6 @@ def _tunable_to_configspace(
     else:
         raise TypeError(f"Invalid Parameter Type: {tunable.type}")
 
-    if tunable.quantization_bins:
-        # Temporary workaround to dropped quantization support in ConfigSpace 1.0
-        # See Also: https://github.com/automl/ConfigSpace/issues/390
-        new_meta = dict(range_hp.meta or {})
-        new_meta[QUANTIZATION_BINS_META_KEY] = tunable.quantization_bins
-        range_hp.meta = new_meta
-        monkey_patch_hp_quantization(range_hp)
-
     if not tunable.special:
         return ConfigurationSpace({tunable.name: range_hp})
 

diff --git a/mlos_core/mlos_core/optimizers/optimizer.py b/mlos_core/mlos_core/optimizers/optimizer.py
@@ -14,6 +14,7 @@
 import pandas as pd
 
 from mlos_core.spaces.adapters.adapter import BaseSpaceAdapter
+from mlos_core.spaces.converters.util import monkey_patch_cs_quantization
 from mlos_core.util import config_to_dataframe
 
 
@@ -44,6 +45,9 @@ def __init__(
         space_adapter : BaseSpaceAdapter
             The space adapter class to employ for parameter space transformations.
         """
+        # Temporary workaround to dropped quantization support in ConfigSpace 1.0
+        # See Also: https://github.com/automl/ConfigSpace/issues/390
+        monkey_patch_cs_quantization(parameter_space)
         self.parameter_space: ConfigSpace.ConfigurationSpace = parameter_space
         self.optimizer_parameter_space: ConfigSpace.ConfigurationSpace = (
             parameter_space if space_adapter is None else space_adapter.target_parameter_space