Honor parameter distribution in Skopt and Goptuna suggestion services

pkg/suggestion/v1beta1/goptuna/converter.go

@@ -120,22 +120,37 @@ func toGoptunaSearchSpace(parameters []*api_v1_beta1.ParameterSpec) (map[string]
 				return nil, err
 			}
 
+			distribution := p.GetFeasibleSpace().GetDistribution()
 			stepstr := p.GetFeasibleSpace().GetStep()
-			if stepstr == "" {
-				searchSpace[p.Name] = goptuna.UniformDistribution{
-					High: high,
-					Low:  low,
-				}
-			} else {
-				step, err := strconv.ParseFloat(stepstr, 64)
-				if err != nil {
-					return nil, err
+			switch distribution {
+			case api_v1_beta1.Distribution_UNIFORM,
+				api_v1_beta1.Distribution_DISTRIBUTION_UNSPECIFIED:
+				if stepstr == "" {
+					searchSpace[p.Name] = goptuna.UniformDistribution{
+						High: high,
+						Low:  low,
+					}
+				} else {
+					step, err := strconv.ParseFloat(stepstr, 64)
+					if err != nil {
+						return nil, err
+					}
+					searchSpace[p.Name] = goptuna.DiscreteUniformDistribution{
+						High: high,
+						Low:  low,
+						Q:    step,
+					}
 				}
-				searchSpace[p.Name] = goptuna.DiscreteUniformDistribution{
+			case api_v1_beta1.Distribution_LOG_UNIFORM:
+				searchSpace[p.Name] = goptuna.LogUniformDistribution{
 					High: high,
 					Low:  low,
-					Q:    step,
 				}
+			default:
+				return nil, fmt.Errorf(
+					"unsupported distribution %v for parameter %s; the goptuna "+
+						"suggestion service supports only UNIFORM and LOG_UNIFORM "+
+						"distributions", distribution, p.Name)
 			}
 		} else if p.ParameterType == api_v1_beta1.ParameterType_INT {
 			high, err := strconv.Atoi(p.GetFeasibleSpace().GetMax())
@@ -146,6 +161,14 @@ func toGoptunaSearchSpace(parameters []*api_v1_beta1.ParameterSpec) (map[string]
 			if err != nil {
 				return nil, err
 			}
+			distribution := p.GetFeasibleSpace().GetDistribution()
+			if distribution != api_v1_beta1.Distribution_UNIFORM &&
+				distribution != api_v1_beta1.Distribution_DISTRIBUTION_UNSPECIFIED {
+				return nil, fmt.Errorf(
+					"unsupported distribution %v for int parameter %s; the goptuna "+
+						"suggestion service supports only UNIFORM for int parameters",
+					distribution, p.Name)
+			}
 			stepstr := p.GetFeasibleSpace().GetStep()
 			if stepstr == "" {
 				searchSpace[p.Name] = goptuna.IntUniformDistribution{

pkg/suggestion/v1beta1/goptuna/converter_test.go

@@ -162,6 +162,53 @@ func Test_toGoptunaSearchSpace(t *testing.T) {
 				},
 			},
 		},
+		"Double parameter type with log-uniform distribution": {
+			parameters: []*api_v1_beta1.ParameterSpec{
+				{
+					Name:          "param-double",
+					ParameterType: api_v1_beta1.ParameterType_DOUBLE,
+					FeasibleSpace: &api_v1_beta1.FeasibleSpace{
+						Max:          "5.5",
+						Min:          "1.5",
+						Distribution: api_v1_beta1.Distribution_LOG_UNIFORM,
+					},
+				},
+			},
+			wantSearchSpace: map[string]interface{}{
+				"param-double": goptuna.LogUniformDistribution{
+					High: 5.5,
+					Low:  1.5,
+				},
+			},
+		},
+		"Double parameter type with unsupported normal distribution": {
+			parameters: []*api_v1_beta1.ParameterSpec{
+				{
+					Name:          "param-double",
+					ParameterType: api_v1_beta1.ParameterType_DOUBLE,
+					FeasibleSpace: &api_v1_beta1.FeasibleSpace{
+						Max:          "5.5",
+						Min:          "1.5",
+						Distribution: api_v1_beta1.Distribution_NORMAL,
+					},
+				},
+			},
+			wantError: cmpopts.AnyError,
+		},
+		"Int parameter type with unsupported log-uniform distribution": {
+			parameters: []*api_v1_beta1.ParameterSpec{
+				{
+					Name:          "param-int",
+					ParameterType: api_v1_beta1.ParameterType_INT,
+					FeasibleSpace: &api_v1_beta1.FeasibleSpace{
+						Max:          "5",
+						Min:          "1",
+						Distribution: api_v1_beta1.Distribution_LOG_UNIFORM,
+					},
+				},
+			},
+			wantError: cmpopts.AnyError,
+		},
 	}
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {

pkg/suggestion/v1beta1/skopt/base_service.py

@@ -17,6 +17,7 @@
 
 import skopt
 
+from pkg.apis.manager.v1beta1.python import api_pb2
 from pkg.suggestion.v1beta1.internal.constant import (
     CATEGORICAL,
     DISCRETE,
@@ -59,19 +60,30 @@ def create_optimizer(self):
         skopt_search_space = []
 
         for param in self.search_space.params:
-            if param.type == INTEGER:
-                skopt_search_space.append(
-                    skopt.space.Integer(int(param.min), int(param.max), name=param.name)
-                )
-            elif param.type == DOUBLE:
-                skopt_search_space.append(
-                    skopt.space.Real(
-                        float(param.min),
-                        float(param.max),
-                        "log-uniform",
-                        name=param.name,
+            if param.type == INTEGER or param.type == DOUBLE:
+                if param.distribution in [api_pb2.UNIFORM, None]:
+                    prior = "uniform"
+                elif param.distribution == api_pb2.LOG_UNIFORM:
+                    prior = "log-uniform"
+                else:
+                    raise ValueError(
+                        f"Unsupported distribution "
+                        f"{api_pb2.Distribution.Name(param.distribution)} for parameter "
+                        f"{param.name}. The Skopt suggestion service supports only "
+                        f"UNIFORM and LOG_UNIFORM distributions."
+                    )
+                if param.type == INTEGER:
+                    skopt_search_space.append(
+                        skopt.space.Integer(
+                            int(param.min), int(param.max), prior, name=param.name
+                        )
+                    )
+                else:
+                    skopt_search_space.append(
+                        skopt.space.Real(
+                            float(param.min), float(param.max), prior, name=param.name
+                        )
                     )
-                )
             elif param.type == CATEGORICAL or param.type == DISCRETE:
                 skopt_search_space.append(
                     skopt.space.Categorical(param.list, name=param.name)

test/unit/v1beta1/suggestion/test_skopt_service.py

@@ -16,9 +16,16 @@
 
 import grpc
 import grpc_testing
+import pytest
 import utils
 
 from pkg.apis.manager.v1beta1.python import api_pb2
+from pkg.suggestion.v1beta1.internal import constant
+from pkg.suggestion.v1beta1.internal.search_space import (
+    HyperParameter,
+    HyperParameterSearchSpace,
+)
+from pkg.suggestion.v1beta1.skopt.base_service import BaseSkoptService
 from pkg.suggestion.v1beta1.skopt.service import SkoptService
 
 
@@ -310,5 +317,55 @@ def test_validate_algorithm_settings(self):
         self.assertEqual(details, "{name} should be great or equal than zero".format(name=wrong_algorithm_setting.name))
 
 
+class TestSkoptDistribution:
+    @pytest.mark.parametrize(
+        ["param_type", "distribution", "want_prior"],
+        [
+            ["double", api_pb2.UNIFORM, "uniform"],
+            ["double", api_pb2.LOG_UNIFORM, "log-uniform"],
+            ["int", api_pb2.UNIFORM, "uniform"],
+            ["int", api_pb2.LOG_UNIFORM, "log-uniform"],
+        ],
+    )
+    def test_distribution_maps_to_prior(self, param_type, distribution, want_prior):
+        # The declared distribution used to be discarded: double parameters were always
+        # built with a hardcoded "log-uniform" prior and int parameters were always
+        # uniform (#2688). The prior must now follow the requested distribution.
+        search_space = HyperParameterSearchSpace()
+        search_space.goal = constant.MAX_GOAL
+        if param_type == "double":
+            param = HyperParameter.double("x", "0.1", "0.9", "", distribution)
+        else:
+            param = HyperParameter.int("x", "1", "9", "", distribution)
+        search_space.params = [param]
+
+        service = BaseSkoptService(search_space=search_space)
+        dimension = service.skopt_optimizer.space.dimensions[0]
+        assert dimension.prior == want_prior
+
+    @pytest.mark.parametrize(
+        ["param_type", "distribution"],
+        [
+            ["double", api_pb2.NORMAL],
+            ["double", api_pb2.LOG_NORMAL],
+            ["int", api_pb2.NORMAL],
+            ["int", api_pb2.LOG_NORMAL],
+        ],
+    )
+    def test_unsupported_distribution_raises(self, param_type, distribution):
+        # NORMAL and LOG_NORMAL have no skopt prior, so they must fail loudly instead of
+        # being silently substituted with a different distribution.
+        search_space = HyperParameterSearchSpace()
+        search_space.goal = constant.MAX_GOAL
+        if param_type == "double":
+            param = HyperParameter.double("x", "0.1", "0.9", "", distribution)
+        else:
+            param = HyperParameter.int("x", "1", "9", "1", distribution)
+        search_space.params = [param]
+
+        with pytest.raises(ValueError, match="Unsupported distribution"):
+            BaseSkoptService(search_space=search_space)
+
+
 if __name__ == "__main__":
     unittest.main()