Repository: microsoft/MLOS

Issue:
Allow user to specify custom probabilities for special values of the tunables.
We also need customizable probabilities for picking special vs. regular values by the optimizer.
These values are currently hardcoded in `convert_configspace.py`.


Discussion:

Resolving patch:
diff --git a/mlos_bench/mlos_bench/optimizers/convert_configspace.py b/mlos_bench/mlos_bench/optimizers/convert_configspace.py
--- a/mlos_bench/mlos_bench/optimizers/convert_configspace.py
+++ b/mlos_bench/mlos_bench/optimizers/convert_configspace.py
@@ -8,7 +8,7 @@
 
 import logging
 
-from typing import Dict, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 from ConfigSpace import (
     CategoricalHyperparameter,
@@ -35,6 +35,14 @@ class TunableValueKind:
     RANGE = "range"
 
 
+def _normalize_weights(weights: List[float]) -> List[float]:
+    """
+    Helper function for normalizing weights to probabilities.
+    """
+    total = sum(weights)
+    return [w / total for w in weights]
+
+
 def _tunable_to_configspace(
         tunable: Tunable, group_name: Optional[str] = None, cost: int = 0) -> ConfigurationSpace:
     """
@@ -61,8 +69,11 @@ def _tunable_to_configspace(
     if tunable.type == "categorical":
         return ConfigurationSpace({
             tunable.name: CategoricalHyperparameter(
-                name=tunable.name, choices=tunable.categories,
-                default_value=tunable.default, meta=meta)
+                name=tunable.name,
+                choices=tunable.categories,
+                weights=_normalize_weights(tunable.weights) if tunable.weights else None,
+                default_value=tunable.default,
+                meta=meta)
         })
 
     if tunable.type == "int":
@@ -75,28 +86,44 @@ def _tunable_to_configspace(
     if not tunable.special:
         return ConfigurationSpace({
             tunable.name: hp_type(
-                name=tunable.name, lower=tunable.range[0], upper=tunable.range[1],
+                name=tunable.name,
+                lower=tunable.range[0],
+                upper=tunable.range[1],
                 default_value=tunable.default if tunable.in_range(tunable.default) else None,
                 meta=meta)
         })
 
+    # Compute the probabilities of switching between regular and special values.
+    special_weights: Optional[List[float]] = None
+    switch_weights = [0.5, 0.5]  # FLAML requires uniform weights.
+    if tunable.weights and tunable.range_weight is not None:
+        special_weights = _normalize_weights(tunable.weights)
+        switch_weights = _normalize_weights([sum(tunable.weights), tunable.range_weight])
+
     # Create three hyperparameters: one for regular values,
     # one for special values, and one to choose between the two.
     (special_name, type_name) = special_param_names(tunable.name)
     conf_space = ConfigurationSpace({
         tunable.name: hp_type(
-            name=tunable.name, lower=tunable.range[0], upper=tunable.range[1],
+            name=tunable.name,
+            lower=tunable.range[0],
+            upper=tunable.range[1],
             default_value=tunable.default if tunable.in_range(tunable.default) else None,
-            meta=meta),
+            meta=meta
+        ),
         special_name: CategoricalHyperparameter(
-            name=special_name, choices=tunable.special,
+            name=special_name,
+            choices=tunable.special,
+            weights=special_weights,
             default_value=tunable.default if tunable.default in tunable.special else None,
-            meta=meta),
+            meta=meta
+        ),
         type_name: CategoricalHyperparameter(
             name=type_name,
             choices=[TunableValueKind.SPECIAL, TunableValueKind.RANGE],
+            weights=switch_weights,
             default_value=TunableValueKind.SPECIAL,
-            weights=[0.5, 0.5]),  # TODO: Make weights configurable; FLAML requires uniform weights.
+        ),
     })
     conf_space.add_condition(EqualsCondition(
         conf_space[special_name], conf_space[type_name], TunableValueKind.SPECIAL))
diff --git a/mlos_bench/mlos_bench/tunables/tunable.py b/mlos_bench/mlos_bench/tunables/tunable.py
--- a/mlos_bench/mlos_bench/tunables/tunable.py
+++ b/mlos_bench/mlos_bench/tunables/tunable.py
@@ -33,6 +33,9 @@ class TunableDict(TypedDict, total=False):
     values: Optional[List[Optional[str]]]
     range: Optional[Union[Sequence[int], Sequence[float]]]
     special: Optional[Union[List[int], List[float]]]
+    values_weights: Optional[List[float]]
+    special_weights: Optional[List[float]]
+    range_weight: Optional[float]
     meta: Dict[str, Any]
 
 
@@ -79,6 +82,10 @@ def __init__(self, name: str, config: TunableDict):
             config_range = (config_range[0], config_range[1])
             self._range = config_range
         self._special: Union[List[int], List[float]] = config.get("special") or []
+        self._weights: List[float] = (
+            config.get("values_weights") or config.get("special_weights") or []
+        )
+        self._range_weight: Optional[float] = config.get("range_weight")
         self._current_value = None
         self._sanity_check()
         self.value = self._default
@@ -96,11 +103,27 @@ def _sanity_check(self) -> None:
                 raise ValueError(f"Values must be unique for the categorical type tunable {self}")
             if self._special:
                 raise ValueError(f"Categorical tunable cannot have special values: {self}")
+            if self._range_weight is not None:
+                raise ValueError(f"Categorical tunable cannot have range_weight: {self}")
+            if self._weights:
+                if len(self._weights) != len(self._values):
+                    raise ValueError(f"Must specify weights for all values: {self}")
+                if any(w < 0 for w in self._weights):
+                    raise ValueError(f"All weights must be non-negative: {self}")
         elif self.is_numerical:
             if self._values is not None:
                 raise ValueError(f"Values must be None for the numerical type tunable {self}")
             if not self._range or len(self._range) != 2 or self._range[0] >= self._range[1]:
                 raise ValueError(f"Invalid range for tunable {self}: {self._range}")
+            if self._weights:
+                if self._range_weight is None:
+                    raise ValueError(f"Must specify weight for the range: {self}")
+                if len(self._weights) != len(self._special):
+                    raise ValueError("Must specify weights for all special values {self}")
+                if any(w < 0 for w in self._weights + [self._range_weight]):
+                    raise ValueError(f"All weights must be non-negative: {self}")
+            elif self._range_weight is not None:
+                raise ValueError(f"Must specify both weights and range_weight or none: {self}")
         else:
             raise ValueError(f"Invalid parameter type for tunable {self}: {self._type}")
         if not self.is_valid(self.default):
@@ -369,6 +392,32 @@ def is_special(self) -> bool:
         """
         return self.value in self._special
 
+    @property
+    def weights(self) -> Optional[List[float]]:
+        """
+        Get the weights of the categories or special values of the tunable.
+        Return None if there are none.
+
+        Returns
+        -------
+        weights : [float]
+            A list of weights or None.
+        """
+        return self._weights
+
+    @property
+    def range_weight(self) -> Optional[float]:
+        """
+        Get weight of the range of the numeric tunable.
+        Return None if there are no weights or a tunable is categorical.
+
+        Returns
+        -------
+        weight : float
+            Weight of the range or None.
+        """
+        return self._range_weight
+
     @property
     def type(self) -> str:
         """


Test patch:
diff --git a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/good/full/full-tunable-params-test.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/good/full/full-tunable-params-test.jsonc
--- a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/good/full/full-tunable-params-test.jsonc
+++ b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/good/full/full-tunable-params-test.jsonc
@@ -9,7 +9,9 @@
                 "default": 10,
                 "range": [1, 500],
                 "meta": {"suffix": "MB"},
-                "special": [-1]
+                "special": [-1],
+                "special_weights": [0.1],
+                "range_weight": 0.9
             },
             "float": {
                 "description": "Float",
@@ -44,7 +46,6 @@
                 "values": [true, 1, "true"],
                 "default": true
             }
-
         }
     },
     "covariant_group_name-2": {
@@ -75,7 +76,8 @@
                 "type": "categorical",
                 "default": "yes",
                 "meta": {"quote": true},
-                "values": ["yes", "no"]
+                "values": ["yes", "no"],
+                "values_weights": [50, 50]
             }
         }
     }
diff --git a/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py b/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py
--- a/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py
+++ b/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py
@@ -38,7 +38,8 @@
                 "description": "Idling method",
                 "type": "categorical",
                 "default": "halt",
-                "values": ["halt", "mwait", "noidle"]
+                "values": ["halt", "mwait", "noidle"],
+                "values_weights": [33, 33, 33]  // FLAML requires uniform weights
             }
         }
     },
@@ -50,7 +51,11 @@
                 "type": "int",
                 "default": -1,
                 "range": [0, 500000],
-                "special": [-1, 0]
+                "special": [-1, 0],
+                // FLAML requires uniform weights, separately for
+                // specials and switching between specials and range.
+                "special_weights": [0.25, 0.25],
+                "range_weight": 0.5
             },
             "kernel_sched_latency_ns": {
                 "description": "Initial value for the scheduler period",
diff --git a/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py b/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py
--- a/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py
+++ b/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py
@@ -28,7 +28,58 @@ def test_categorical_required_params() -> None:
     json_config = """
     {
         "type": "categorical",
-        "values_missing": ["foo", "bar", "foo"],
+        "values_missing": ["foo", "bar", "baz"],
+        "default": "foo"
+    }
+    """
+    config = json.loads(json_config)
+    with pytest.raises(ValueError):
+        Tunable(name='test', config=config)
+
+
+def test_categorical_weights() -> None:
+    """
+    Instantiate a categorical tunable with weights.
+    """
+    json_config = """
+    {
+        "type": "categorical",
+        "values": ["foo", "bar", "baz"],
+        "values_weights": [25, 25, 50],
+        "default": "foo"
+    }
+    """
+    config = json.loads(json_config)
+    tunable = Tunable(name='test', config=config)
+    assert tunable.weights == [25, 25, 50]
+
+
+def test_categorical_weights_wrong_count() -> None:
+    """
+    Try to instantiate a categorical tunable with incorrect number of weights.
+    """
+    json_config = """
+    {
+        "type": "categorical",
+        "values": ["foo", "bar", "baz"],
+        "values_weights": [50, 50],
+        "default": "foo"
+    }
+    """
+    config = json.loads(json_config)
+    with pytest.raises(ValueError):
+        Tunable(name='test', config=config)
+
+
+def test_categorical_weights_wrong_values() -> None:
+    """
+    Try to instantiate a categorical tunable with invalid weights.
+    """
+    json_config = """
+    {
+        "type": "categorical",
+        "values": ["foo", "bar", "baz"],
+        "values_weights": [-1, 50, 50],
         "default": "foo"
     }
     """
@@ -174,6 +225,165 @@ def test_numerical_tunable_reversed_range(tunable_type: str) -> None:
         Tunable(name=f'test_{tunable_type}', config=config)
 
 
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_weights(tunable_type: str) -> None:
+    """
+    Instantiate a numerical tunable with weighted special values.
+    """
+    json_config = f"""
+    {{
+        "type": "{tunable_type}",
+        "range": [0, 100],
+        "special": [0],
+        "special_weights": [0.1],
+        "range_weight": 0.9,
+        "default": 0
+    }}
+    """
+    config = json.loads(json_config)
+    tunable = Tunable(name='test', config=config)
+    assert tunable.special == [0]
+    assert tunable.weights == [0.1]
+    assert tunable.range_weight == 0.9
+
+
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_weights_no_specials(tunable_type: str) -> None:
+    """
+    Raise an error if special_weights are specified but no special values.
+    """
+    json_config = f"""
+    {{
+        "type": "{tunable_type}",
+        "range": [0, 100],
+        "special_weights": [0.1, 0.9],
+        "default": 0
+    }}
+    """
+    config = json.loads(json_config)
+    with pytest.raises(ValueError):
+        Tunable(name='test', config=config)
+
+
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_weights_non_normalized(tunable_type: str) -> None:
+    """
+    Instantiate a numerical tunable with non-normalized weights
+    of the special values.
+    """
+    json_config = f"""
+    {{
+        "type": "{tunable_type}",
+        "range": [0, 100],
+        "special": [-1, 0],
+        "special_weights": [0, 10],
+        "range_weight": 90,
+        "default": 0
+    }}
+    """
+    config = json.loads(json_config)
+    tunable = Tunable(name='test', config=config)
+    assert tunable.special == [-1, 0]
+    assert tunable.weights == [0, 10]  # Zero weights are ok
+    assert tunable.range_weight == 90
+
+
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_weights_wrong_count(tunable_type: str) -> None:
+    """
+    Try to instantiate a numerical tunable with incorrect number of weights.
+    """
+    json_config = f"""
+    {{
+        "type": "{tunable_type}",
+        "range": [0, 100],
+        "special": [0],
+        "special_weights": [0.1, 0.1, 0.8],
+        "range_weight": 0.1,
+        "default": 0
+    }}
+    """
+    config = json.loads(json_config)
+    with pytest.raises(ValueError):
+        Tunable(name='test', config=config)
+
+
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_weights_no_range_weight(tunable_type: str) -> None:
+    """
+    Try to instantiate a numerical tunable with weights but no range_weight.
+    """
+    json_config = f"""
+    {{
+        "type": "{tunable_type}",
+        "range": [0, 100],
+        "special": [0, -1],
+        "special_weights": [0.1, 0.2],
+        "default": 0
+    }}
+    """
+    config = json.loads(json_config)
+    with pytest.raises(ValueError):
+        Tunable(name='test', config=config)
+
+
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_range_weight_no_weights(tunable_type: str) -> None:
+    """
+    Try to instantiate a numerical tunable with specials but no range_weight.
+    """
+    json_config = f"""
+    {{
+        "type": "{tunable_type}",
+        "range": [0, 100],
+        "special": [0, -1],
+        "range_weight": 0.3,
+        "default": 0
+    }}
+    """
+    config = json.loads(json_config)
+    with pytest.raises(ValueError):
+        Tunable(name='test', config=config)
+
+
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_range_weight_no_specials(tunable_type: str) -> None:
+    """
+    Try to instantiate a numerical tunable with specials but no range_weight.
+    """
+    json_config = f"""
+    {{
+        "type": "{tunable_type}",
+        "range": [0, 100],
+        "range_weight": 0.3,
+        "default": 0
+    }}
+    """
+    config = json.loads(json_config)
+    with pytest.raises(ValueError):
+        Tunable(name='test', config=config)
+
+
+@pytest.mark.parametrize("tunable_type", ["int", "float"])
+def test_numerical_weights_wrong_values(tunable_type: str) -> None:
+    """
+    Try to instantiate a numerical tunable with incorrect number of weights.
+    """
+    json_config = f"""
+    {{
+        "type": "{tunable_type}",
+        "range": [0, 100],
+        "special": [0],
+        "special_weights": [-1],
+        "range_weight": 10,
+        "default": 0
+    }}
+    """
+    config = json.loads(json_config)
+    with pytest.raises(ValueError):
+        Tunable(name='test', config=config)
+
+
 def test_bad_type() -> None:
     """
     Disallow bad types.
diff --git a/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_test.py b/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_test.py
--- a/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_test.py
+++ b/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_test.py
@@ -51,12 +51,14 @@ def configuration_space() -> ConfigurationSpace:
         "kernel_sched_latency_ns": (0, 1000000000),
     })
 
+    # NOTE: FLAML requires distribution to be uniform
     spaces["vmSize"].default_value = "Standard_B4ms"
     spaces["idle"].default_value = "halt"
     spaces["kernel_sched_migration_cost_ns"].default_value = 250000
     spaces[kernel_sched_migration_cost_ns_special].default_value = -1
+    spaces[kernel_sched_migration_cost_ns_special].probabilities = (0.5, 0.5)
     spaces[kernel_sched_migration_cost_ns_type].default_value = TunableValueKind.SPECIAL
-    spaces[kernel_sched_migration_cost_ns_type].probabilities = (0.5, 0.5)  # FLAML requires distribution to be uniform
+    spaces[kernel_sched_migration_cost_ns_type].probabilities = (0.5, 0.5)
     spaces["kernel_sched_latency_ns"].default_value = 2000000
 
     spaces.add_condition(EqualsCondition(