Fix Optuna tests

kubeflow · Aug 17, 2021 · 6f36d29 · 6f36d29
1 parent 5316709
commit 6f36d29
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/pkg/suggestion/v1beta1/goptuna/converter.go b/pkg/suggestion/v1beta1/goptuna/converter.go
@@ -75,13 +75,13 @@ func toGoptunaSampler(algorithm *api_v1_beta1.AlgorithmSpec) (goptuna.Sampler, g
  return nil, nil, err
  }
  opts = append(opts, tpe.SamplerOptionSeed(int64(seed)))
- } else if s.Name == "startup_trials" {
+ } else if s.Name == "n_startup_trials" {
  n, err := strconv.Atoi(s.Value)
  if err != nil {
  return nil, nil, err
  }
  opts = append(opts, tpe.SamplerOptionNumberOfStartupTrials(n))
- } else if s.Name == "ei_candidates" {
+ } else if s.Name == "n_ei_candidates" {
  n, err := strconv.Atoi(s.Value)
  if err != nil {
  return nil, nil, err

diff --git a/test/suggestion/v1beta1/test_optuna_service.py b/test/suggestion/v1beta1/test_optuna_service.py
@@ -33,10 +33,10 @@ def setup_method(self):
  servicers, grpc_testing.strict_real_time())
 
  @pytest.mark.parametrize(
- ["algorithm_name", "algorithm_settings"], 
+ ["algorithm_name", "algorithm_settings"],
  [
- ["tpe", {"startup_trials": "20", "ei_candidates": "10", "random_state": "71"}], 
- ["multivariate-tpe", {"startup_trials": "20", "ei_candidates": "10", "random_state": "71"}], 
+ ["tpe", {"n_startup_trials": "20", "n_ei_candidates": "10", "random_state": "71"}],
+ ["multivariate-tpe", {"n_startup_trials": "20", "n_ei_candidates": "10", "random_state": "71"}],
  ["cmaes", {"restart_strategy": "ipop", "sigma": "2", "random_state": "71"}],
  ["random", {"random_state": "71"}],
  ],
@@ -47,7 +47,7 @@ def test_get_suggestion(self, algorithm_name, algorithm_settings):
  spec=api_pb2.ExperimentSpec(
  algorithm=api_pb2.AlgorithmSpec(
  algorithm_name=algorithm_name,
- algorithm_settings = [
+ algorithm_settings=[
  api_pb2.AlgorithmSetting(
  name=name,
  value=value