Move quantization_bins specification into CS.NumericalHyperparameter …

…meta field

mlos_bench/mlos_bench/optimizers/convert_configspace.py

@@ -26,7 +26,10 @@
 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_quantization
+from mlos_core.spaces.converters.util import (
+    monkey_patch_hp_quantization,
+    QUANTIZATION_BINS_META_KEY,
+)
 
 _LOG = logging.getLogger(__name__)
 
@@ -140,7 +143,10 @@ def _tunable_to_configspace(
     if tunable.quantization_bins:
         # Temporary workaround to dropped quantization support in ConfigSpace 1.0
         # See Also: https://github.com/automl/ConfigSpace/issues/390
-        monkey_patch_quantization(range_hp, tunable.quantization_bins)
+        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})

mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_test.py

@@ -23,7 +23,10 @@
 )
 from mlos_bench.tunables.tunable import Tunable
 from mlos_bench.tunables.tunable_groups import TunableGroups
-from mlos_core.spaces.converters.util import monkey_patch_quantization
+from mlos_core.spaces.converters.util import (
+    monkey_patch_cs_quantization,
+    QUANTIZATION_BINS_META_KEY,
+)
 
 # pylint: disable=redefined-outer-name
 
@@ -63,7 +66,11 @@ def configuration_space() -> ConfigurationSpace:
                 bounds=(0, 1000000000),
                 log=False,
                 default=2000000,
-                meta={"group": "kernel", "cost": 0},
+                meta={
+                    "group": "kernel",
+                    "cost": 0,
+                    QUANTIZATION_BINS_META_KEY: 11,
+                },
             ),
             "kernel_sched_migration_cost_ns": Integer(
                 name="kernel_sched_migration_cost_ns",
@@ -101,9 +108,7 @@ def configuration_space() -> ConfigurationSpace:
             TunableValueKind.RANGE,
         )
     )
-    hp = spaces["kernel_sched_latency_ns"]
-    assert isinstance(hp, NumericalHyperparameter)
-    monkey_patch_quantization(hp, quantization_bins=11)
+    monkey_patch_cs_quantization(spaces)
     return spaces
 
 

mlos_core/mlos_core/spaces/converters/util.py

@@ -4,21 +4,35 @@
 #
 """Helper functions for config space converters."""
 
+from ConfigSpace import ConfigurationSpace
 from ConfigSpace.functional import quantize
-from ConfigSpace.hyperparameters import NumericalHyperparameter
+from ConfigSpace.hyperparameters import Hyperparameter, NumericalHyperparameter
 
 
-def monkey_patch_quantization(hp: NumericalHyperparameter, quantization_bins: int) -> None:
+QUANTIZATION_BINS_META_KEY = "quantization_bins"
+
+
+def monkey_patch_hp_quantization(hp: Hyperparameter) -> None:
     """
     Monkey-patch quantization into the Hyperparameter.
 
     Parameters
     ----------
     hp : NumericalHyperparameter
         ConfigSpace hyperparameter to patch.
-    quantization_bins : int
-        Number of bins to quantize the hyperparameter into.
     """
+    if not isinstance(hp, NumericalHyperparameter):
+        return
+    assert isinstance(hp, NumericalHyperparameter)
+    quantization_bins = (hp.meta or {}).get(QUANTIZATION_BINS_META_KEY)
+    if quantization_bins is None:
+        return  # no quantization requested
+
+    try:
+        quantization_bins = int(quantization_bins)
+    except ValueError as ex:
+        raise ValueError(f"{quantization_bins=} :: must be an integer.") from ex
+
     if quantization_bins <= 1:
         raise ValueError(f"{quantization_bins=} :: must be greater than 1.")
 
@@ -37,3 +51,16 @@ def monkey_patch_quantization(hp: NumericalHyperparameter, quantization_bins: in
             bins=quantization_bins,
         ).astype(type(hp.default_value)),
     )
+
+
+def monkey_patch_cs_quantization(cs: ConfigurationSpace) -> None:
+    """
+    Monkey-patch quantization into the Hyperparameters of a ConfigSpace.
+
+    Parameters
+    ----------
+    cs : ConfigurationSpace
+        ConfigSpace to patch.
+    """
+    for hp in cs.values():
+        monkey_patch_hp_quantization(hp)

mlos_core/mlos_core/tests/spaces/adapters/llamatune_test.py

@@ -13,7 +13,10 @@
 import pytest
 
 from mlos_core.spaces.adapters import LlamaTuneAdapter
-from mlos_core.spaces.converters.util import monkey_patch_quantization
+from mlos_core.spaces.converters.util import (
+    monkey_patch_cs_quantization,
+    QUANTIZATION_BINS_META_KEY,
+)
 
 # Explicitly test quantized values with llamatune space adapter.
 # TODO: Add log scale sampling tests as well.
@@ -32,24 +35,47 @@ def construct_parameter_space(  # pylint: disable=too-many-arguments
     input_space = CS.ConfigurationSpace(seed=seed)
 
     for idx in range(n_continuous_params):
-        input_space.add(CS.UniformFloatHyperparameter(name=f"cont_{idx}", lower=0, upper=64))
+        input_space.add(
+            CS.UniformFloatHyperparameter(
+                name=f"cont_{idx}",
+                lower=0,
+                upper=64,
+            )
+        )
     for idx in range(n_quantized_continuous_params):
-        param_int = CS.UniformFloatHyperparameter(name=f"cont_{idx}", lower=0, upper=64)
-        monkey_patch_quantization(param_int, 6)
-        input_space.add(param_int)
+        input_space.add(
+            CS.UniformFloatHyperparameter(
+                name=f"cont_{idx}",
+                lower=0,
+                upper=64,
+                meta={QUANTIZATION_BINS_META_KEY: 6},
+            )
+        )
     for idx in range(n_integer_params):
-        input_space.add(CS.UniformIntegerHyperparameter(name=f"int_{idx}", lower=-1, upper=256))
+        input_space.add(
+            CS.UniformIntegerHyperparameter(
+                name=f"int_{idx}",
+                lower=-1,
+                upper=256,
+            )
+        )
     for idx in range(n_quantized_integer_params):
-        param_float = CS.UniformIntegerHyperparameter(name=f"int_{idx}", lower=0, upper=256)
-        monkey_patch_quantization(param_float, 17)
-        input_space.add(param_float)
+        input_space.add(
+            CS.UniformIntegerHyperparameter(
+                name=f"int_{idx}",
+                lower=0,
+                upper=256,
+                meta={QUANTIZATION_BINS_META_KEY: 17},
+            )
+        )
     for idx in range(n_categorical_params):
         input_space.add(
             CS.CategoricalHyperparameter(
                 name=f"str_{idx}", choices=[f"option_{idx}" for idx in range(5)]
             )
         )
 
+    monkey_patch_cs_quantization(input_space)
     return input_space
 
 
@@ -322,15 +348,15 @@ def test_special_parameter_values_biasing() -> None:  # pylint: disable=too-comp
             special_values_instances["int_2"][100] += 1
 
     assert (1 - eps) * int(num_configs * bias_percentage) <= special_values_instances["int_1"][0]
-    assert (1 - eps) * int(num_configs * bias_percentage / 2) <= special_values_instances["int_1"][
-        1
-    ]
-    assert (1 - eps) * int(num_configs * bias_percentage / 2) <= special_values_instances["int_2"][
-        2
-    ]
-    assert (1 - eps) * int(num_configs * bias_percentage * 1.5) <= special_values_instances[
-        "int_2"
-    ][100]
+    assert (1 - eps) * int(num_configs * bias_percentage / 2) <= (
+        special_values_instances["int_1"][1]
+    )
+    assert (1 - eps) * int(num_configs * bias_percentage / 2) <= (
+        special_values_instances["int_2"][2]
+    )
+    assert (1 - eps) * int(num_configs * bias_percentage * 1.5) <= (
+        special_values_instances["int_2"][100]
+    )
 
 
 def test_max_unique_values_per_param() -> None:

mlos_core/mlos_core/tests/spaces/monkey_patch_quantization_test.py

@@ -5,62 +5,99 @@
 """Unit tests for ConfigSpace quantization monkey patching."""
 
 import numpy as np
-from ConfigSpace import UniformFloatHyperparameter, UniformIntegerHyperparameter
+from ConfigSpace import (
+    ConfigurationSpace,
+    UniformFloatHyperparameter,
+    UniformIntegerHyperparameter,
+)
 from numpy.random import RandomState
 
-from mlos_core.spaces.converters.util import monkey_patch_quantization
+from mlos_core.spaces.converters.util import (
+    monkey_patch_cs_quantization,
+    QUANTIZATION_BINS_META_KEY,
+)
 from mlos_core.tests import SEED
 
 
 def test_configspace_quant_int() -> None:
     """Check the quantization of an integer hyperparameter."""
+    quantization_bins = 11
     quantized_values = set(range(0, 101, 10))
-    hp = UniformIntegerHyperparameter("hp", lower=0, upper=100, log=False)
+    hp = UniformIntegerHyperparameter(
+        "hp",
+        lower=0,
+        upper=100,
+        log=False,
+        meta={QUANTIZATION_BINS_META_KEY: quantization_bins},
+    )
+    cs = ConfigurationSpace()
+    cs.add(hp)
 
     # Before patching: expect that at least one value is not quantized.
     assert not set(hp.sample_value(100)).issubset(quantized_values)
 
-    monkey_patch_quantization(hp, 11)
+    monkey_patch_cs_quantization(cs)
     # After patching: *all* values must belong to the set of quantized values.
     assert hp.sample_value() in quantized_values  # check scalar type
     assert set(hp.sample_value(100)).issubset(quantized_values)  # batch version
 
 
 def test_configspace_quant_float() -> None:
     """Check the quantization of a float hyperparameter."""
-    quantized_values = set(np.linspace(0, 1, num=5, endpoint=True))
-    hp = UniformFloatHyperparameter("hp", lower=0, upper=1, log=False)
+    # 5 is a nice number of bins to avoid floating point errors.
+    quantization_bins = 5
+    quantized_values = set(np.linspace(0, 1, num=quantization_bins, endpoint=True))
+    hp = UniformFloatHyperparameter(
+        "hp",
+        lower=0,
+        upper=1,
+        log=False,
+        meta={QUANTIZATION_BINS_META_KEY: quantization_bins},
+    )
+    cs = ConfigurationSpace()
+    cs.add(hp)
 
     # Before patching: expect that at least one value is not quantized.
     assert not set(hp.sample_value(100)).issubset(quantized_values)
 
-    # 5 is a nice number of bins to avoid floating point errors.
-    monkey_patch_quantization(hp, 5)
+    monkey_patch_cs_quantization(cs)
     # After patching: *all* values must belong to the set of quantized values.
     assert hp.sample_value() in quantized_values  # check scalar type
     assert set(hp.sample_value(100)).issubset(quantized_values)  # batch version
 
 
 def test_configspace_quant_repatch() -> None:
     """Repatch the same hyperparameter with different number of bins."""
+    quantization_bins = 11
     quantized_values = set(range(0, 101, 10))
-    hp = UniformIntegerHyperparameter("hp", lower=0, upper=100, log=False)
+    hp = UniformIntegerHyperparameter(
+        "hp",
+        lower=0,
+        upper=100,
+        log=False,
+        meta={QUANTIZATION_BINS_META_KEY: quantization_bins},
+    )
+    cs = ConfigurationSpace()
+    cs.add(hp)
 
     # Before patching: expect that at least one value is not quantized.
     assert not set(hp.sample_value(100)).issubset(quantized_values)
 
-    monkey_patch_quantization(hp, 11)
+    monkey_patch_cs_quantization(cs)
     # After patching: *all* values must belong to the set of quantized values.
     samples = hp.sample_value(100, seed=RandomState(SEED))
     assert set(samples).issubset(quantized_values)
 
     # Patch the same hyperparameter again and check that the results are the same.
-    monkey_patch_quantization(hp, 11)
+    monkey_patch_cs_quantization(cs)
     # After patching: *all* values must belong to the set of quantized values.
     assert all(samples == hp.sample_value(100, seed=RandomState(SEED)))
 
     # Repatch with the higher number of bins and make sure we get new values.
-    monkey_patch_quantization(hp, 21)
+    new_meta = dict(hp.meta or {})
+    new_meta[QUANTIZATION_BINS_META_KEY] = 21
+    hp.meta = new_meta
+    monkey_patch_cs_quantization(cs)
     samples_set = set(hp.sample_value(100, seed=RandomState(SEED)))
     quantized_values_new = set(range(5, 96, 10))
     assert samples_set.issubset(set(range(0, 101, 5)))