From 2e4cfa26a0768cc15bfce47aaad379f7bb840695 Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 16 Aug 2024 07:59:26 -0700
Subject: [PATCH 1/4] Use number of bins instead of quantization interval in
 mlos_bench tunables (#835)

Closes #803

> Future PR will rename the config schema in order to reduce confusion
on the change in semantics, but also keep this PR smaller.

---------

Co-authored-by: Brian Kroth <bpkroth@users.noreply.github.com>
Co-authored-by: Brian Kroth <bpkroth@microsoft.com>
---
 .../tunables/tunable-params-schema.json       | 11 ++-
 .../optimizers/convert_configspace.py         | 56 +++++++++++-----
 .../optimizers/grid_search_optimizer.py       |  6 +-
 ...le-params-int-bad-float-quantization.jsonc |  2 +-
 .../good/full/full-tunable-params-test.jsonc  |  6 +-
 .../optimizers/grid_search_optimizer_test.py  |  7 +-
 .../tests/tunable_groups_fixtures.py          |  1 +
 .../tunables/test_tunables_size_props.py      | 31 ++++++---
 .../tests/tunables/tunable_definition_test.py |  6 +-
 .../tunable_to_configspace_quant_test.py      | 67 +++++++++++++++++++
 .../tunables/tunable_to_configspace_test.py   |  7 +-
 mlos_bench/mlos_bench/tunables/tunable.py     | 58 +++++++---------
 12 files changed, 179 insertions(+), 79 deletions(-)
 create mode 100644 mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_quant_test.py

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 2380ad62fe..fd6c61b91d 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
@@ -125,7 +125,8 @@
         },
         "quantization": {
             "description": "The number of buckets to quantize the range into.",
-            "$comment": "type left unspecified here"
+            "type": "integer",
+            "exclusiveMinimum": 1
         },
         "log_scale": {
             "description": "Whether to use log instead of linear scale for the range search.",
@@ -217,9 +218,7 @@
                     "$ref": "#/$defs/tunable_param_distribution"
                 },
                 "quantization": {
-                    "$ref": "#/$defs/quantization",
-                    "type": "integer",
-                    "exclusiveMinimum": 1
+                    "$ref": "#/$defs/quantization"
                 },
                 "log": {
                     "$ref": "#/$defs/log_scale"
@@ -267,9 +266,7 @@
                     "$ref": "#/$defs/tunable_param_distribution"
                 },
                 "quantization": {
-                    "$ref": "#/$defs/quantization",
-                    "type": "number",
-                    "exclusiveMinimum": 0
+                    "$ref": "#/$defs/quantization"
                 },
                 "log": {
                     "$ref": "#/$defs/log_scale"
diff --git a/mlos_bench/mlos_bench/optimizers/convert_configspace.py b/mlos_bench/mlos_bench/optimizers/convert_configspace.py
index 6244bba7b3..ad7968cf69 100644
--- a/mlos_bench/mlos_bench/optimizers/convert_configspace.py
+++ b/mlos_bench/mlos_bench/optimizers/convert_configspace.py
@@ -11,8 +11,6 @@
 
 from ConfigSpace import (
     Beta,
-    BetaFloatHyperparameter,
-    BetaIntegerHyperparameter,
     CategoricalHyperparameter,
     Configuration,
     ConfigurationSpace,
@@ -20,12 +18,10 @@
     Float,
     Integer,
     Normal,
-    NormalFloatHyperparameter,
-    NormalIntegerHyperparameter,
     Uniform,
-    UniformFloatHyperparameter,
-    UniformIntegerHyperparameter,
 )
+from ConfigSpace.functional import quantize
+from ConfigSpace.hyperparameters import NumericalHyperparameter
 from ConfigSpace.types import NotSet
 
 from mlos_bench.tunables.tunable import Tunable, TunableValue
@@ -53,6 +49,37 @@ def _normalize_weights(weights: List[float]) -> List[float]:
     return [w / total for w in weights]
 
 
+def _monkey_patch_quantization(hp: NumericalHyperparameter, quantization_bins: int) -> 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 quantization_bins <= 1:
+        raise ValueError(f"{quantization_bins=} :: must be greater than 1.")
+
+    # Temporary workaround to dropped quantization support in ConfigSpace 1.0
+    # See Also: https://github.com/automl/ConfigSpace/issues/390
+    if not hasattr(hp, "sample_value_mlos_orig"):
+        setattr(hp, "sample_value_mlos_orig", hp.sample_value)
+
+    assert hasattr(hp, "sample_value_mlos_orig")
+    setattr(
+        hp,
+        "sample_value",
+        lambda size=None, **kwargs: quantize(
+            hp.sample_value_mlos_orig(size, **kwargs),
+            bounds=(hp.lower, hp.upper),
+            bins=quantization_bins,
+        ).astype(type(hp.default_value)),
+    )
+
+
 def _tunable_to_configspace(
     tunable: Tunable,
     group_name: Optional[str] = None,
@@ -77,6 +104,7 @@ def _tunable_to_configspace(
     cs : ConfigurationSpace
         A ConfigurationSpace object that corresponds to the Tunable.
     """
+    # pylint: disable=too-complex
     meta: Dict[Hashable, TunableValue] = {"cost": cost}
     if group_name is not None:
         meta["group"] = group_name
@@ -110,20 +138,12 @@ def _tunable_to_configspace(
     elif tunable.distribution is not None:
         raise TypeError(f"Invalid Distribution Type: {tunable.distribution}")
 
-    range_hp: Union[
-        BetaFloatHyperparameter,
-        BetaIntegerHyperparameter,
-        NormalFloatHyperparameter,
-        NormalIntegerHyperparameter,
-        UniformFloatHyperparameter,
-        UniformIntegerHyperparameter,
-    ]
+    range_hp: NumericalHyperparameter
     if tunable.type == "int":
         range_hp = Integer(
             name=tunable.name,
             bounds=(int(tunable.range[0]), int(tunable.range[1])),
             log=bool(tunable.is_log),
-            # TODO: Restore quantization support (#803).
             distribution=distribution,
             default=(
                 int(tunable.default)
@@ -137,7 +157,6 @@ def _tunable_to_configspace(
             name=tunable.name,
             bounds=tunable.range,
             log=bool(tunable.is_log),
-            # TODO: Restore quantization support (#803).
             distribution=distribution,
             default=(
                 float(tunable.default)
@@ -149,6 +168,11 @@ def _tunable_to_configspace(
     else:
         raise TypeError(f"Invalid Parameter Type: {tunable.type}")
 
+    if tunable.quantization:
+        # 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)
+
     if not tunable.special:
         return ConfigurationSpace({tunable.name: range_hp})
 
diff --git a/mlos_bench/mlos_bench/optimizers/grid_search_optimizer.py b/mlos_bench/mlos_bench/optimizers/grid_search_optimizer.py
index 8bcd090415..0fc9209619 100644
--- a/mlos_bench/mlos_bench/optimizers/grid_search_optimizer.py
+++ b/mlos_bench/mlos_bench/optimizers/grid_search_optimizer.py
@@ -47,7 +47,7 @@ def __init__(
         self._suggested_configs: Set[Tuple[TunableValue, ...]] = set()
 
     def _sanity_check(self) -> None:
-        size = np.prod([tunable.cardinality for (tunable, _group) in self._tunables])
+        size = np.prod([tunable.cardinality or np.inf for (tunable, _group) in self._tunables])
         if size == np.inf:
             raise ValueError(
                 f"Unquantized tunables are not supported for grid search: {self._tunables}"
@@ -79,9 +79,9 @@ def _get_grid(self) -> Tuple[Tuple[str, ...], Dict[Tuple[TunableValue, ...], Non
             for config in generate_grid(
                 self.config_space,
                 {
-                    tunable.name: int(tunable.cardinality)
+                    tunable.name: tunable.cardinality or 0  # mypy wants an int
                     for (tunable, _group) in self._tunables
-                    if tunable.quantization or tunable.type == "int"
+                    if tunable.is_numerical and tunable.cardinality
                 },
             )
         ]
diff --git a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-float-quantization.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-float-quantization.jsonc
index 194682b859..cff352e7b8 100644
--- a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-float-quantization.jsonc
+++ b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-float-quantization.jsonc
@@ -6,7 +6,7 @@
                 "type": "float",
                 "default": 10,
                 "range": [1, 500],
-                "quantization": 0   // <-- should be greater than 0
+                "quantization": 1   // <-- should be greater than 1
             }
         }
     }
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 ae7291e5fc..6d83b248af 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
@@ -7,7 +7,7 @@
                 "description": "Int",
                 "type": "int",
                 "default": 10,
-                "range": [1, 500],
+                "range": [0, 500],
                 "meta": {"suffix": "MB"},
                 "special": [-1],
                 "special_weights": [0.1],
@@ -26,7 +26,7 @@
                 "description": "Int",
                 "type": "int",
                 "default": 10,
-                "range": [1, 500],
+                "range": [0, 500],
                 "meta": {"suffix": "MB"},
                 "special": [-1],
                 "special_weights": [0.1],
@@ -48,7 +48,7 @@
                 "meta": {"scale": 1000, "prefix": "/proc/var/random/", "base": 2.71828},
                 "range": [1.1, 111.1],
                 "special": [-1.1],
-                "quantization": 10,
+                "quantization": 11,
                 "distribution": {
                     "type": "uniform"
                 },
diff --git a/mlos_bench/mlos_bench/tests/optimizers/grid_search_optimizer_test.py b/mlos_bench/mlos_bench/tests/optimizers/grid_search_optimizer_test.py
index 769bf8859d..bf49cc82b5 100644
--- a/mlos_bench/mlos_bench/tests/optimizers/grid_search_optimizer_test.py
+++ b/mlos_bench/mlos_bench/tests/optimizers/grid_search_optimizer_test.py
@@ -9,6 +9,7 @@
 import random
 from typing import Dict, List
 
+import numpy as np
 import pytest
 
 from mlos_bench.environments.status import Status
@@ -40,7 +41,7 @@ def grid_search_tunables_config() -> dict:
                     "type": "float",
                     "range": [0, 1],
                     "default": 0.5,
-                    "quantization": 0.25,
+                    "quantization": 5,
                 },
             },
         },
@@ -99,7 +100,9 @@ def test_grid_search_grid(
 ) -> None:
     """Make sure that grid search optimizer initializes and works correctly."""
     # Check the size.
-    expected_grid_size = math.prod(tunable.cardinality for tunable, _group in grid_search_tunables)
+    expected_grid_size = math.prod(
+        tunable.cardinality or np.inf for tunable, _group in grid_search_tunables
+    )
     assert expected_grid_size > len(grid_search_tunables)
     assert len(grid_search_tunables_grid) == expected_grid_size
     # Check for specific example configs inclusion.
diff --git a/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py b/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py
index 2c10e2ba75..e606496a49 100644
--- a/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py
+++ b/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py
@@ -62,6 +62,7 @@
                 "type": "int",
                 "default": 2000000,
                 "range": [0, 1000000000],
+                "quantization": 11,
                 "log": false
             }
         }
diff --git a/mlos_bench/mlos_bench/tests/tunables/test_tunables_size_props.py b/mlos_bench/mlos_bench/tests/tunables/test_tunables_size_props.py
index fcbca29ed9..cfc5b43bac 100644
--- a/mlos_bench/mlos_bench/tests/tunables/test_tunables_size_props.py
+++ b/mlos_bench/mlos_bench/tests/tunables/test_tunables_size_props.py
@@ -4,7 +4,6 @@
 #
 """Unit tests for checking tunable size properties."""
 
-import numpy as np
 import pytest
 
 from mlos_bench.tunables.tunable import Tunable
@@ -23,9 +22,9 @@ def test_tunable_int_size_props() -> None:
             "default": 3,
         },
     )
-    assert tunable.span == 4
-    assert tunable.cardinality == 5
     expected = [1, 2, 3, 4, 5]
+    assert tunable.span == 4
+    assert tunable.cardinality == len(expected)
     assert list(tunable.quantized_values or []) == expected
     assert list(tunable.values or []) == expected
 
@@ -41,7 +40,7 @@ def test_tunable_float_size_props() -> None:
         },
     )
     assert tunable.span == 3.5
-    assert tunable.cardinality == np.inf
+    assert tunable.cardinality is None
     assert tunable.quantized_values is None
     assert tunable.values is None
 
@@ -68,11 +67,17 @@ def test_tunable_quantized_int_size_props() -> None:
     """Test quantized tunable int size properties."""
     tunable = Tunable(
         name="test",
-        config={"type": "int", "range": [100, 1000], "default": 100, "quantization": 100},
+        config={
+            "type": "int",
+            "range": [100, 1000],
+            "default": 100,
+            "quantization": 10,
+        },
     )
-    assert tunable.span == 900
-    assert tunable.cardinality == 10
     expected = [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000]
+    assert tunable.span == 900
+    assert tunable.cardinality == len(expected)
+    assert tunable.quantization == len(expected)
     assert list(tunable.quantized_values or []) == expected
     assert list(tunable.values or []) == expected
 
@@ -81,10 +86,16 @@ def test_tunable_quantized_float_size_props() -> None:
     """Test quantized tunable float size properties."""
     tunable = Tunable(
         name="test",
-        config={"type": "float", "range": [0, 1], "default": 0, "quantization": 0.1},
+        config={
+            "type": "float",
+            "range": [0, 1],
+            "default": 0,
+            "quantization": 11,
+        },
     )
-    assert tunable.span == 1
-    assert tunable.cardinality == 11
     expected = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
+    assert tunable.span == 1
+    assert tunable.cardinality == len(expected)
+    assert tunable.quantization == len(expected)
     assert pytest.approx(list(tunable.quantized_values or []), 0.0001) == expected
     assert pytest.approx(list(tunable.values or []), 0.0001) == expected
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 7403841f8d..0ad08eefb6 100644
--- a/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py
+++ b/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py
@@ -234,13 +234,15 @@ def test_numerical_quantization(tunable_type: TunableValueTypeName) -> None:
     {{
         "type": "{tunable_type}",
         "range": [0, 100],
-        "quantization": 10,
+        "quantization": 11,
         "default": 0
     }}
     """
     config = json.loads(json_config)
     tunable = Tunable(name="test", config=config)
-    assert tunable.quantization == 10
+    expected = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
+    assert tunable.quantization == len(expected)
+    assert pytest.approx(list(tunable.quantized_values or []), 1e-8) == expected
     assert not tunable.is_log
 
 
diff --git a/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_quant_test.py b/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_quant_test.py
new file mode 100644
index 0000000000..606fcd9b77
--- /dev/null
+++ b/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_quant_test.py
@@ -0,0 +1,67 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""Unit tests for ConfigSpace quantization monkey patching."""
+
+import numpy as np
+from ConfigSpace import UniformFloatHyperparameter, UniformIntegerHyperparameter
+from numpy.random import RandomState
+
+from mlos_bench.optimizers.convert_configspace import _monkey_patch_quantization
+from mlos_bench.tests import SEED
+
+
+def test_configspace_quant_int() -> None:
+    """Check the quantization of an integer hyperparameter."""
+    quantized_values = set(range(0, 101, 10))
+    hp = UniformIntegerHyperparameter("hp", lower=0, upper=100, log=False)
+
+    # 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)
+    # 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)
+
+    # 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)
+    # 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."""
+    quantized_values = set(range(0, 101, 10))
+    hp = UniformIntegerHyperparameter("hp", lower=0, upper=100, log=False)
+
+    # 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)
+    # 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)
+    # 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)
+    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)))
+    assert len(samples_set - quantized_values_new) < len(samples_set)
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
index 8e472a7ea5..2b52000225 100644
--- a/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_test.py
+++ b/mlos_bench/mlos_bench/tests/tunables/tunable_to_configspace_test.py
@@ -13,9 +13,11 @@
     UniformFloatHyperparameter,
     UniformIntegerHyperparameter,
 )
+from ConfigSpace.hyperparameters import NumericalHyperparameter
 
 from mlos_bench.optimizers.convert_configspace import (
     TunableValueKind,
+    _monkey_patch_quantization,
     _tunable_to_configspace,
     special_param_names,
     tunable_groups_to_configspace,
@@ -41,8 +43,6 @@ def configuration_space() -> ConfigurationSpace:
         special_param_names("kernel_sched_migration_cost_ns")
     )
 
-    # TODO: Add quantization support tests (#803).
-
     # NOTE: FLAML requires distribution to be uniform
     spaces = ConfigurationSpace(
         {
@@ -101,6 +101,9 @@ def configuration_space() -> ConfigurationSpace:
             TunableValueKind.RANGE,
         )
     )
+    hp = spaces["kernel_sched_latency_ns"]
+    assert isinstance(hp, NumericalHyperparameter)
+    _monkey_patch_quantization(hp, quantization_bins=10)
     return spaces
 
 
diff --git a/mlos_bench/mlos_bench/tunables/tunable.py b/mlos_bench/mlos_bench/tunables/tunable.py
index 8f9bb48bff..b4a58e0a50 100644
--- a/mlos_bench/mlos_bench/tunables/tunable.py
+++ b/mlos_bench/mlos_bench/tunables/tunable.py
@@ -64,7 +64,7 @@ class TunableDict(TypedDict, total=False):
     default: TunableValue
     values: Optional[List[Optional[str]]]
     range: Optional[Union[Sequence[int], Sequence[float]]]
-    quantization: Optional[Union[int, float]]
+    quantization: Optional[int]
     log: Optional[bool]
     distribution: Optional[DistributionDict]
     special: Optional[Union[List[int], List[float]]]
@@ -109,7 +109,7 @@ def __init__(self, name: str, config: TunableDict):
             self._values = [str(v) if v is not None else v for v in self._values]
         self._meta: Dict[str, Any] = config.get("meta", {})
         self._range: Optional[Union[Tuple[int, int], Tuple[float, float]]] = None
-        self._quantization: Optional[Union[int, float]] = config.get("quantization")
+        self._quantization: Optional[int] = config.get("quantization")
         self._log: Optional[bool] = config.get("log")
         self._distribution: Optional[DistributionName] = None
         self._distribution_params: Dict[str, float] = {}
@@ -182,19 +182,8 @@ def _sanity_check_numerical(self) -> 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._quantization is not None:
-            if self.dtype == int:
-                if not isinstance(self._quantization, int):
-                    raise ValueError(f"Quantization of a int param should be an int: {self}")
-                if self._quantization <= 1:
-                    raise ValueError(f"Number of quantization points is <= 1: {self}")
-            if self.dtype == float:
-                if not isinstance(self._quantization, (float, int)):
-                    raise ValueError(
-                        f"Quantization of a float param should be a float or int: {self}"
-                    )
-                if self._quantization <= 0:
-                    raise ValueError(f"Number of quantization points is <= 0: {self}")
+        if self._quantization is not None and self._quantization <= 1:
+            raise ValueError(f"Number of quantization bins is <= 1: {self}")
         if self._distribution is not None and self._distribution not in {
             "uniform",
             "normal",
@@ -391,7 +380,6 @@ def is_valid(self, value: TunableValue) -> bool:
         is_valid : bool
             True if the value is valid, False otherwise.
         """
-        # FIXME: quantization check?
         if self.is_categorical and self._values:
             return value in self._values
         elif self.is_numerical and self._range:
@@ -592,14 +580,14 @@ def span(self) -> Union[int, float]:
         return num_range[1] - num_range[0]
 
     @property
-    def quantization(self) -> Optional[Union[int, float]]:
+    def quantization(self) -> Optional[int]:
         """
-        Get the quantization factor, if specified.
+        Get the number of quantization bins, if specified.
 
         Returns
         -------
-        quantization : int, float, None
-            The quantization factor, or None.
+        quantization : int | None
+            Number of quantization bins, or None.
         """
         if self.is_categorical:
             return None
@@ -618,41 +606,45 @@ def quantized_values(self) -> Optional[Union[Iterable[int], Iterable[float]]]:
         """
         num_range = self.range
         if self.type == "float":
-            if not self._quantization:
+            if not self.quantization:
                 return None
             # Be sure to return python types instead of numpy types.
-            cardinality = self.cardinality
-            assert isinstance(cardinality, int)
             return (
                 float(x)
                 for x in np.linspace(
                     start=num_range[0],
                     stop=num_range[1],
-                    num=cardinality,
+                    num=self.quantization,
                     endpoint=True,
                 )
             )
         assert self.type == "int", f"Unhandled tunable type: {self}"
-        return range(int(num_range[0]), int(num_range[1]) + 1, int(self._quantization or 1))
+        return range(
+            int(num_range[0]),
+            int(num_range[1]) + 1,
+            int(self.span / (self.quantization - 1)) if self.quantization else 1,
+        )
 
     @property
-    def cardinality(self) -> Union[int, float]:
+    def cardinality(self) -> Optional[int]:
         """
-        Gets the cardinality of elements in this tunable, or else infinity.
+        Gets the cardinality of elements in this tunable, or else None. (i.e., when the
+        tunable is continuous float and not quantized).
 
         If the tunable has quantization set, this
 
         Returns
         -------
-        cardinality : int, float
-            Either the number of points in the tunable or else infinity.
+        cardinality : int
+            Either the number of points in the tunable or else None.
         """
         if self.is_categorical:
             return len(self.categories)
-        if not self.quantization and self.type == "float":
-            return np.inf
-        q_factor = self.quantization or 1
-        return int(self.span / q_factor) + 1
+        if self.quantization:
+            return self.quantization
+        if self.type == "int":
+            return int(self.span) + 1
+        return None
 
     @property
     def is_log(self) -> Optional[bool]:

From fadfacb9bd1b4acb848f04e5929749935b7edf44 Mon Sep 17 00:00:00 2001
From: Sergiy Matusevych <sergiym@microsoft.com>
Date: Fri, 16 Aug 2024 10:01:49 -0700
Subject: [PATCH 2/4] Rename `quantization` -> `quantization_bins` (#844)

Merge after (or instead of) #835

diff from #835 :: https://github.com/motus/MLOS/pull/15/files

Closes #803

---------

Co-authored-by: Brian Kroth <bpkroth@users.noreply.github.com>
Co-authored-by: Brian Kroth <bpkroth@microsoft.com>
---
 .../tunables/tunable-params-schema.json       | 12 +++++-----
 .../optimizers/convert_configspace.py         |  4 ++--
 ...e-params-float-bad-quantization-type.jsonc |  2 +-
 ...le-params-int-bad-float-quantization.jsonc |  2 +-
 ...able-params-int-bad-int-quantization.jsonc |  2 +-
 ...e-params-int-wrong-quantization-type.jsonc |  2 +-
 .../good/full/full-tunable-params-test.jsonc  |  6 ++---
 .../optimizers/grid_search_optimizer_test.py  |  2 +-
 .../tests/tunable_groups_fixtures.py          |  2 +-
 .../tunables/test_tunables_size_props.py      |  8 +++----
 .../tests/tunables/tunable_definition_test.py |  6 ++---
 mlos_bench/mlos_bench/tunables/tunable.py     | 24 +++++++++----------
 12 files changed, 36 insertions(+), 36 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 fd6c61b91d..e278c797a4 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
@@ -123,7 +123,7 @@
             "maxItems": 2,
             "uniqueItems": true
         },
-        "quantization": {
+        "quantization_bins": {
             "description": "The number of buckets to quantize the range into.",
             "type": "integer",
             "exclusiveMinimum": 1
@@ -187,7 +187,7 @@
             },
             "required": ["type", "default", "values"],
             "not": {
-                "required": ["range", "special", "special_weights", "range_weight", "log", "quantization", "distribution"]
+                "required": ["range", "special", "special_weights", "range_weight", "log", "quantization_bins", "distribution"]
             },
             "$comment": "TODO: add check that default is in values",
             "unevaluatedProperties": false
@@ -217,8 +217,8 @@
                 "distribution": {
                     "$ref": "#/$defs/tunable_param_distribution"
                 },
-                "quantization": {
-                    "$ref": "#/$defs/quantization"
+                "quantization_bins": {
+                    "$ref": "#/$defs/quantization_bins"
                 },
                 "log": {
                     "$ref": "#/$defs/log_scale"
@@ -265,8 +265,8 @@
                 "distribution": {
                     "$ref": "#/$defs/tunable_param_distribution"
                 },
-                "quantization": {
-                    "$ref": "#/$defs/quantization"
+                "quantization_bins": {
+                    "$ref": "#/$defs/quantization_bins"
                 },
                 "log": {
                     "$ref": "#/$defs/log_scale"
diff --git a/mlos_bench/mlos_bench/optimizers/convert_configspace.py b/mlos_bench/mlos_bench/optimizers/convert_configspace.py
index ad7968cf69..52623cee33 100644
--- a/mlos_bench/mlos_bench/optimizers/convert_configspace.py
+++ b/mlos_bench/mlos_bench/optimizers/convert_configspace.py
@@ -168,10 +168,10 @@ def _tunable_to_configspace(
     else:
         raise TypeError(f"Invalid Parameter Type: {tunable.type}")
 
-    if tunable.quantization:
+    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)
+        _monkey_patch_quantization(range_hp, tunable.quantization_bins)
 
     if not tunable.special:
         return ConfigurationSpace({tunable.name: range_hp})
diff --git a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-float-bad-quantization-type.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-float-bad-quantization-type.jsonc
index 0ed8884457..379356b556 100644
--- a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-float-bad-quantization-type.jsonc
+++ b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-float-bad-quantization-type.jsonc
@@ -6,7 +6,7 @@
                 "type": "float",
                 "default": 10,
                 "range": [0, 10],
-                "quantization": true  // <-- this is invalid
+                "quantization_bins": true  // <-- this is invalid
             }
         }
     }
diff --git a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-float-quantization.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-float-quantization.jsonc
index cff352e7b8..27a3986df4 100644
--- a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-float-quantization.jsonc
+++ b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-float-quantization.jsonc
@@ -6,7 +6,7 @@
                 "type": "float",
                 "default": 10,
                 "range": [1, 500],
-                "quantization": 1   // <-- should be greater than 1
+                "quantization_bins": 1   // <-- should be greater than 1
             }
         }
     }
diff --git a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-int-quantization.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-int-quantization.jsonc
index 199cf681ca..df648be385 100644
--- a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-int-quantization.jsonc
+++ b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-bad-int-quantization.jsonc
@@ -6,7 +6,7 @@
                 "type": "int",
                 "default": 10,
                 "range": [1, 500],
-                "quantization": 1   // <-- should be greater than 1
+                "quantization_bins": 1   // <-- should be greater than 1
             }
         }
     }
diff --git a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-wrong-quantization-type.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-wrong-quantization-type.jsonc
index 1b7af4ffcd..7ed215b363 100644
--- a/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-wrong-quantization-type.jsonc
+++ b/mlos_bench/mlos_bench/tests/config/schemas/tunable-params/test-cases/bad/invalid/tunable-params-int-wrong-quantization-type.jsonc
@@ -6,7 +6,7 @@
                 "type": "int",
                 "default": 10,
                 "range": [1, 500],
-                "quantization": "yes"  // <-- this is invalid
+                "quantization_bins": "yes"  // <-- this is invalid
             }
         }
     }
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 6d83b248af..5e045e4171 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
@@ -12,7 +12,7 @@
                 "special": [-1],
                 "special_weights": [0.1],
                 "range_weight": 0.9,
-                "quantization": 50,
+                "quantization_bins": 50,
                 "distribution": {
                     "type": "beta",
                     "params": {
@@ -31,7 +31,7 @@
                 "special": [-1],
                 "special_weights": [0.1],
                 "range_weight": 0.9,
-                "quantization": 50,
+                "quantization_bins": 50,
                 "distribution": {
                     "type": "normal",
                     "params": {
@@ -48,7 +48,7 @@
                 "meta": {"scale": 1000, "prefix": "/proc/var/random/", "base": 2.71828},
                 "range": [1.1, 111.1],
                 "special": [-1.1],
-                "quantization": 11,
+                "quantization_bins": 11,
                 "distribution": {
                     "type": "uniform"
                 },
diff --git a/mlos_bench/mlos_bench/tests/optimizers/grid_search_optimizer_test.py b/mlos_bench/mlos_bench/tests/optimizers/grid_search_optimizer_test.py
index bf49cc82b5..703381e488 100644
--- a/mlos_bench/mlos_bench/tests/optimizers/grid_search_optimizer_test.py
+++ b/mlos_bench/mlos_bench/tests/optimizers/grid_search_optimizer_test.py
@@ -41,7 +41,7 @@ def grid_search_tunables_config() -> dict:
                     "type": "float",
                     "range": [0, 1],
                     "default": 0.5,
-                    "quantization": 5,
+                    "quantization_bins": 5,
                 },
             },
         },
diff --git a/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py b/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py
index e606496a49..59082fac9e 100644
--- a/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py
+++ b/mlos_bench/mlos_bench/tests/tunable_groups_fixtures.py
@@ -62,7 +62,7 @@
                 "type": "int",
                 "default": 2000000,
                 "range": [0, 1000000000],
-                "quantization": 11,
+                "quantization_bins": 11,
                 "log": false
             }
         }
diff --git a/mlos_bench/mlos_bench/tests/tunables/test_tunables_size_props.py b/mlos_bench/mlos_bench/tests/tunables/test_tunables_size_props.py
index cfc5b43bac..c3344f6988 100644
--- a/mlos_bench/mlos_bench/tests/tunables/test_tunables_size_props.py
+++ b/mlos_bench/mlos_bench/tests/tunables/test_tunables_size_props.py
@@ -71,13 +71,13 @@ def test_tunable_quantized_int_size_props() -> None:
             "type": "int",
             "range": [100, 1000],
             "default": 100,
-            "quantization": 10,
+            "quantization_bins": 10,
         },
     )
     expected = [100, 200, 300, 400, 500, 600, 700, 800, 900, 1000]
     assert tunable.span == 900
     assert tunable.cardinality == len(expected)
-    assert tunable.quantization == len(expected)
+    assert tunable.quantization_bins == len(expected)
     assert list(tunable.quantized_values or []) == expected
     assert list(tunable.values or []) == expected
 
@@ -90,12 +90,12 @@ def test_tunable_quantized_float_size_props() -> None:
             "type": "float",
             "range": [0, 1],
             "default": 0,
-            "quantization": 11,
+            "quantization_bins": 11,
         },
     )
     expected = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
     assert tunable.span == 1
     assert tunable.cardinality == len(expected)
-    assert tunable.quantization == len(expected)
+    assert tunable.quantization_bins == len(expected)
     assert pytest.approx(list(tunable.quantized_values or []), 0.0001) == expected
     assert pytest.approx(list(tunable.values or []), 0.0001) == expected
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 0ad08eefb6..a7b12c2692 100644
--- a/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py
+++ b/mlos_bench/mlos_bench/tests/tunables/tunable_definition_test.py
@@ -234,14 +234,14 @@ def test_numerical_quantization(tunable_type: TunableValueTypeName) -> None:
     {{
         "type": "{tunable_type}",
         "range": [0, 100],
-        "quantization": 11,
+        "quantization_bins": 11,
         "default": 0
     }}
     """
     config = json.loads(json_config)
     tunable = Tunable(name="test", config=config)
     expected = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
-    assert tunable.quantization == len(expected)
+    assert tunable.quantization_bins == len(expected)
     assert pytest.approx(list(tunable.quantized_values or []), 1e-8) == expected
     assert not tunable.is_log
 
@@ -393,7 +393,7 @@ def test_numerical_quantization_wrong(tunable_type: TunableValueTypeName) -> Non
     {{
         "type": "{tunable_type}",
         "range": [0, 100],
-        "quantization": 0,
+        "quantization_bins": 0,
         "default": 0
     }}
     """
diff --git a/mlos_bench/mlos_bench/tunables/tunable.py b/mlos_bench/mlos_bench/tunables/tunable.py
index b4a58e0a50..4d2781ad11 100644
--- a/mlos_bench/mlos_bench/tunables/tunable.py
+++ b/mlos_bench/mlos_bench/tunables/tunable.py
@@ -64,7 +64,7 @@ class TunableDict(TypedDict, total=False):
     default: TunableValue
     values: Optional[List[Optional[str]]]
     range: Optional[Union[Sequence[int], Sequence[float]]]
-    quantization: Optional[int]
+    quantization_bins: Optional[int]
     log: Optional[bool]
     distribution: Optional[DistributionDict]
     special: Optional[Union[List[int], List[float]]]
@@ -109,7 +109,7 @@ def __init__(self, name: str, config: TunableDict):
             self._values = [str(v) if v is not None else v for v in self._values]
         self._meta: Dict[str, Any] = config.get("meta", {})
         self._range: Optional[Union[Tuple[int, int], Tuple[float, float]]] = None
-        self._quantization: Optional[int] = config.get("quantization")
+        self._quantization_bins: Optional[int] = config.get("quantization_bins")
         self._log: Optional[bool] = config.get("log")
         self._distribution: Optional[DistributionName] = None
         self._distribution_params: Dict[str, float] = {}
@@ -162,7 +162,7 @@ def _sanity_check_categorical(self) -> None:
             raise ValueError(f"Categorical tunable cannot have range_weight: {self}")
         if self._log is not None:
             raise ValueError(f"Categorical tunable cannot have log parameter: {self}")
-        if self._quantization is not None:
+        if self._quantization_bins is not None:
             raise ValueError(f"Categorical tunable cannot have quantization parameter: {self}")
         if self._distribution is not None:
             raise ValueError(f"Categorical parameters do not support `distribution`: {self}")
@@ -182,7 +182,7 @@ def _sanity_check_numerical(self) -> 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._quantization is not None and self._quantization <= 1:
+        if self._quantization_bins is not None and self._quantization_bins <= 1:
             raise ValueError(f"Number of quantization bins is <= 1: {self}")
         if self._distribution is not None and self._distribution not in {
             "uniform",
@@ -580,18 +580,18 @@ def span(self) -> Union[int, float]:
         return num_range[1] - num_range[0]
 
     @property
-    def quantization(self) -> Optional[int]:
+    def quantization_bins(self) -> Optional[int]:
         """
         Get the number of quantization bins, if specified.
 
         Returns
         -------
-        quantization : int | None
+        quantization_bins : int | None
             Number of quantization bins, or None.
         """
         if self.is_categorical:
             return None
-        return self._quantization
+        return self._quantization_bins
 
     @property
     def quantized_values(self) -> Optional[Union[Iterable[int], Iterable[float]]]:
@@ -606,7 +606,7 @@ def quantized_values(self) -> Optional[Union[Iterable[int], Iterable[float]]]:
         """
         num_range = self.range
         if self.type == "float":
-            if not self.quantization:
+            if not self.quantization_bins:
                 return None
             # Be sure to return python types instead of numpy types.
             return (
@@ -614,7 +614,7 @@ def quantized_values(self) -> Optional[Union[Iterable[int], Iterable[float]]]:
                 for x in np.linspace(
                     start=num_range[0],
                     stop=num_range[1],
-                    num=self.quantization,
+                    num=self.quantization_bins,
                     endpoint=True,
                 )
             )
@@ -622,7 +622,7 @@ def quantized_values(self) -> Optional[Union[Iterable[int], Iterable[float]]]:
         return range(
             int(num_range[0]),
             int(num_range[1]) + 1,
-            int(self.span / (self.quantization - 1)) if self.quantization else 1,
+            int(self.span / (self.quantization_bins - 1)) if self.quantization_bins else 1,
         )
 
     @property
@@ -640,8 +640,8 @@ def cardinality(self) -> Optional[int]:
         """
         if self.is_categorical:
             return len(self.categories)
-        if self.quantization:
-            return self.quantization
+        if self.quantization_bins:
+            return self.quantization_bins
         if self.type == "int":
             return int(self.span) + 1
         return None

From 6134b146129d2244feaeea1a03b553916a449f3f Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@users.noreply.github.com>
Date: Fri, 16 Aug 2024 12:58:38 -0500
Subject: [PATCH 3/4] =?UTF-8?q?Bump=20version:=200.6.0=20=E2=86=92=200.6.1?=
 =?UTF-8?q?=20(#845)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Bump version after #844 which represents a breaking config schema
change.
---
 .bumpversion.cfg                 | 2 +-
 doc/source/version.py            | 2 +-
 mlos_bench/mlos_bench/version.py | 2 +-
 mlos_core/mlos_core/version.py   | 2 +-
 mlos_viz/mlos_viz/version.py     | 2 +-
 5 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/.bumpversion.cfg b/.bumpversion.cfg
index 44468a477e..f5326c2c01 100644
--- a/.bumpversion.cfg
+++ b/.bumpversion.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.6.0
+current_version = 0.6.1
 commit = True
 tag = True
 
diff --git a/doc/source/version.py b/doc/source/version.py
index 53c4de37f3..77a706ac1e 100644
--- a/doc/source/version.py
+++ b/doc/source/version.py
@@ -7,7 +7,7 @@
 """
 
 # NOTE: This should be managed by bumpversion.
-VERSION = '0.6.0'
+VERSION = '0.6.1'
 
 if __name__ == "__main__":
     print(VERSION)
diff --git a/mlos_bench/mlos_bench/version.py b/mlos_bench/mlos_bench/version.py
index 58b3db26e0..9bb7876ac4 100644
--- a/mlos_bench/mlos_bench/version.py
+++ b/mlos_bench/mlos_bench/version.py
@@ -5,7 +5,7 @@
 """Version number for the mlos_bench package."""
 
 # NOTE: This should be managed by bumpversion.
-VERSION = "0.6.0"
+VERSION = "0.6.1"
 
 if __name__ == "__main__":
     print(VERSION)
diff --git a/mlos_core/mlos_core/version.py b/mlos_core/mlos_core/version.py
index 8990c83170..64b1c52022 100644
--- a/mlos_core/mlos_core/version.py
+++ b/mlos_core/mlos_core/version.py
@@ -5,7 +5,7 @@
 """Version number for the mlos_core package."""
 
 # NOTE: This should be managed by bumpversion.
-VERSION = "0.6.0"
+VERSION = "0.6.1"
 
 if __name__ == "__main__":
     print(VERSION)
diff --git a/mlos_viz/mlos_viz/version.py b/mlos_viz/mlos_viz/version.py
index 3daffe21ed..a38eddb9e9 100644
--- a/mlos_viz/mlos_viz/version.py
+++ b/mlos_viz/mlos_viz/version.py
@@ -5,7 +5,7 @@
 """Version number for the mlos_viz package."""
 
 # NOTE: This should be managed by bumpversion.
-VERSION = "0.6.0"
+VERSION = "0.6.1"
 
 if __name__ == "__main__":
     print(VERSION)

From 8afe9c37c2fba1ba11e1bb2fa8f445c59eb0214a Mon Sep 17 00:00:00 2001
From: Brian Kroth <bpkroth@users.noreply.github.com>
Date: Fri, 16 Aug 2024 13:17:27 -0500
Subject: [PATCH 4/4] Devcontainer tweaks (#846)

Try to add github.com ssh keys to the ssh known_hosts file to avoid
prompting in scripted operations in downstream consumers.
---
 .devcontainer/Dockerfile           | 7 +++++++
 .github/workflows/devcontainer.yml | 4 ++++
 2 files changed, 11 insertions(+)

diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
index cf4b37aca9..8702d314cb 100644
--- a/.devcontainer/Dockerfile
+++ b/.devcontainer/Dockerfile
@@ -92,3 +92,10 @@ RUN umask 0002 \
     && mkdir -p /opt/conda/pkgs/cache/ && chown -R vscode:conda /opt/conda/pkgs/cache/
 RUN mkdir -p /home/vscode/.conda/envs \
     && ln -s /opt/conda/envs/mlos /home/vscode/.conda/envs/mlos
+
+# Try and prime the devcontainer's ssh known_hosts keys with the github one for scripted calls.
+RUN mkdir -p /home/vscode/.ssh \
+    && ( \
+        grep -q ^github.com /home/vscode/.ssh/known_hosts \
+        || ssh-keyscan github.com | tee -a /home/vscode/.ssh/known_hosts \
+    )
diff --git a/.github/workflows/devcontainer.yml b/.github/workflows/devcontainer.yml
index c07a902fc6..b57257a907 100644
--- a/.github/workflows/devcontainer.yml
+++ b/.github/workflows/devcontainer.yml
@@ -128,6 +128,10 @@ jobs:
       run: |
         docker exec --user vscode --env USER=vscode mlos-devcontainer printenv
 
+    - name: Check that github.com is in the ssh known_hosts file
+      run: |
+        docker exec --user vscode --env USER=vscode mlos-devcontainer grep ^github.com /home/vscode/.ssh/known_hosts
+
     - name: Update the conda env in the devcontainer
       timeout-minutes: 10
       run: |