From f219178958a272f9940dfac22908f8367961efd0 Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 2 Feb 2024 15:32:11 -0800
Subject: [PATCH] use values_weights instead of just weights for categoricals

---
 .../config/schemas/tunables/tunable-params-schema.json      | 6 +++---
 .../test-cases/good/full/full-tunable-params-test.jsonc     | 2 +-
 mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py      | 2 +-
 .../mlos_bench/tests/tunables/tunable_definition_test.py    | 6 +++---
 mlos_bench/mlos_bench/tunables/tunable.py                   | 4 ++--
 5 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/mlos_bench/mlos_bench/config/schemas/tunables/tunable-params-schema.json b/mlos_bench/mlos_bench/config/schemas/tunables/tunable-params-schema.json
index d7a1ad69248..e00118fe762 100644
--- a/mlos_bench/mlos_bench/config/schemas/tunables/tunable-params-schema.json
+++ b/mlos_bench/mlos_bench/config/schemas/tunables/tunable-params-schema.json
@@ -36,7 +36,7 @@
                     "minItems": 1,
                     "uniqueItems": true
                 },
-                "weights": {
+                "values_weights": {
                     "type": "array",
                     "items": {
                         "type": "number"
@@ -92,7 +92,7 @@
             },
             "required": ["type", "default", "range"],
             "not": {
-                "required": ["values", "weights"]
+                "required": ["values", "values_weights"]
             },
             "$comment": "TODO: add check that default is in range",
             "unevaluatedProperties": false
@@ -139,7 +139,7 @@
             },
             "required": ["type", "default", "range"],
             "not": {
-                "required": ["values", "weights"]
+                "required": ["values", "values_weights"]
             },
             "$comment": "TODO: add check that default is in range",
             "unevaluatedProperties": false
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
index 4c271fce15e..733385bb969 100644
--- 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
@@ -77,7 +77,7 @@
                 "default": "yes",
                 "meta": {"quote": true},
                 "values": ["yes", "no"],
-                "weights": [50, 50]
+                "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
index dd6a65e041a..920005b7258 100644
--- a/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py
+++ b/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py
@@ -39,7 +39,7 @@
                 "type": "categorical",
                 "default": "halt",
                 "values": ["halt", "mwait", "noidle"],
-                "weights": [33, 33, 33]  // FLAML requires uniform weights
+                "values_weights": [33, 33, 33]  // FLAML requires uniform weights
             }
         }
     },
diff --git a/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py b/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py
index 3a3897a9e52..d4817d3cfbe 100644
--- a/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py
+++ b/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py
@@ -45,7 +45,7 @@ def test_categorical_weights() -> None:
     {
         "type": "categorical",
         "values": ["foo", "bar", "baz"],
-        "weights": [25, 25, 50],
+        "values_weights": [25, 25, 50],
         "default": "foo"
     }
     """
@@ -62,7 +62,7 @@ def test_categorical_weights_wrong_count() -> None:
     {
         "type": "categorical",
         "values": ["foo", "bar", "baz"],
-        "weights": [50, 50],
+        "values_weights": [50, 50],
         "default": "foo"
     }
     """
@@ -79,7 +79,7 @@ def test_categorical_weights_wrong_values() -> None:
     {
         "type": "categorical",
         "values": ["foo", "bar", "baz"],
-        "weights": [-1, 50, 50],
+        "values_weights": [-1, 50, 50],
         "default": "foo"
     }
     """
diff --git a/mlos_bench/mlos_bench/tunables/tunable.py b/mlos_bench/mlos_bench/tunables/tunable.py
index ef4e8df51d6..10192a70c46 100644
--- a/mlos_bench/mlos_bench/tunables/tunable.py
+++ b/mlos_bench/mlos_bench/tunables/tunable.py
@@ -33,7 +33,7 @@ class TunableDict(TypedDict, total=False):
     values: Optional[List[Optional[str]]]
     range: Optional[Union[Sequence[int], Sequence[float]]]
     special: Optional[Union[List[int], List[float]]]
-    weights: Optional[List[float]]
+    values_weights: Optional[List[float]]
     special_weights: Optional[List[float]]
     range_weight: Optional[float]
     meta: Dict[str, Any]
@@ -83,7 +83,7 @@ def __init__(self, name: str, config: TunableDict):
             self._range = config_range
         self._special: Union[List[int], List[float]] = config.get("special") or []
         self._weights: List[float] = (
-            config.get("weights") or config.get("special_weights") or []
+            config.get("values_weights") or config.get("special_weights") or []
         )
         self._range_weight: Optional[float] = config.get("range_weight")
         self._current_value = None