Repository: jina-ai/serve
Pull request: feat: allow optimizer to work on executor alternatives

Problem statement:
unittests for Optimizer
**Describe the feature**
There are no unittests for Optimizers



allow optimizer to have executor class as parameter
**Describe the feature**
Right now, FlowOptimizer only receives a set of parameters that can be updated for a given Flow template. This Flow template is static and does not consider the potential switching between specific Pod instances or Executor types



optimizer GridSampler not tested
**Describe the bug**
optimizer GridSampler not tested

Hints from issue discussion:

Commit message: test: add grid sampler tests

Patch to review:
diff --git a/jina/jaml/__init__.py b/jina/jaml/__init__.py
index a0b02c59935b8..3bcbc1cdc61b4 100644
--- a/jina/jaml/__init__.py
+++ b/jina/jaml/__init__.py
@@ -445,7 +445,7 @@ def load_config(
         *,
         allow_py_modules: bool = True,
         substitute: bool = True,
-        context: Dict[str, Any] = None,
+        context: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> 'JAMLCompatible':
         """A high-level interface for loading configuration with features
diff --git a/jina/optimizers/__init__.py b/jina/optimizers/__init__.py
index 6d77ff8b39b76..59dd3b65dc386 100644
--- a/jina/optimizers/__init__.py
+++ b/jina/optimizers/__init__.py
@@ -11,7 +11,7 @@
     CategoricalParameter,
     DiscreteUniformParameter,
 )
-from .parameters import load_optimization_parameters
+from .parameters import load_optimization_parameters, PodAlternativeParameter
 from ..helper import colored
 from ..importer import ImportExtensions
 from ..jaml import JAMLCompatible, JAML
@@ -187,12 +187,12 @@ def __init__(
         self._sampler = sampler
         self._direction = direction
         self._seed = seed
+        self.parameters = load_optimization_parameters(self._parameter_yaml)
 
     def _trial_parameter_sampler(self, trial: 'Trial'):
         trial_parameters = {}
-        parameters = load_optimization_parameters(self._parameter_yaml)
-        for param in parameters:
-            trial_parameters[param.jaml_variable] = param.suggest(trial)
+        for param in self.parameters:
+            param.update_trial_params(trial, trial_parameters)
 
         trial.workspace = (
             self._workspace_base_dir
diff --git a/jina/optimizers/flow_runner.py b/jina/optimizers/flow_runner.py
index 4c42d39bcb570..a84f8909451e6 100644
--- a/jina/optimizers/flow_runner.py
+++ b/jina/optimizers/flow_runner.py
@@ -1,13 +1,13 @@
 import os
 import shutil
 from collections.abc import Iterable
-from typing import Optional, Union, List
+from typing import Union, List, Optional, Callable
 
 from ..types.document.generators import from_lines
-from ..flow import Flow
 from ..helper import colored
 from ..logging.predefined import default_logger as logger
 from ..jaml import JAMLCompatible
+from ..flow import Flow
 
 
 class FlowRunner(JAMLCompatible):
@@ -17,7 +17,7 @@ def run(
         self,
         trial_parameters: dict,
         workspace: str = 'workspace',
-        callback=None,
+        callback: Optional[Callable[..., None]] = None,
         **kwargs,
     ):
         """
@@ -95,13 +95,13 @@ def run(
         self,
         trial_parameters: dict,
         workspace: str = 'workspace',
-        callback=None,
+        callback: Optional[Callable[..., None]] = None,
         **kwargs,
     ):
         """
         Running method of :class:`SingleFlowRunner`.
 
-        :param trial_parameters: parameters need to be tried
+        :param trial_parameters: parameters that need to be tried
         :param workspace: path of workspace
         :param callback: callback function
         :param kwargs: keyword argument
@@ -134,7 +134,7 @@ def run(
         self,
         trial_parameters: dict,
         workspace: str = 'workspace',
-        callback=None,
+        callback: Optional[Callable[..., None]] = None,
         **kwargs,
     ):
         """
diff --git a/jina/optimizers/parameters.py b/jina/optimizers/parameters.py
index eca8c8676ba33..8a972404b00b8 100644
--- a/jina/optimizers/parameters.py
+++ b/jina/optimizers/parameters.py
@@ -1,5 +1,4 @@
-from abc import abstractmethod
-from typing import Optional, Sequence, Union
+from typing import Optional, Sequence, List, Dict, Union
 
 from ..jaml import JAML, JAMLCompatible
 
@@ -23,7 +22,7 @@ def __init__(
         else:
             self.jaml_variable = jaml_variable
 
-    def suggest(self, trial: 'Trial'):
+    def _suggest(self, trial: 'Trial'):
         """
         Suggest an instance of the parameter for a given trial.
 
@@ -32,6 +31,16 @@ def suggest(self, trial: 'Trial'):
         """
         raise NotImplementedError
 
+    def update_trial_params(self, trial: 'Trial', trial_params: Dict):
+        """
+        Update the trial parameters given the variables and the trials
+
+        :param trial: An instance of an Optuna Trial object
+        :param trial_params: the parameters to update
+            # noqa: DAR201
+        """
+        trial_params[self.jaml_variable] = self._suggest(trial)
+
 
 class IntegerParameter(OptimizationParameter):
     """
@@ -61,7 +70,7 @@ def __init__(
         self.step_size = step_size
         self.log = log
 
-    def suggest(self, trial: 'Trial'):
+    def _suggest(self, trial: 'Trial'):
         """
         Suggest an instance of the parameter for a given trial.
 
@@ -89,7 +98,7 @@ def __init__(self, low: float, high: float, *args, **kwargs):
         self.low = low
         self.high = high
 
-    def suggest(self, trial: 'Trial'):
+    def _suggest(self, trial: 'Trial'):
         """
         Suggest an instance of the parameter for a given trial.
 
@@ -115,7 +124,7 @@ def __init__(self, low: float, high: float, *args, **kwargs):
         self.low = low
         self.high = high
 
-    def suggest(self, trial: 'Trial'):
+    def _suggest(self, trial: 'Trial'):
         """
         Suggest an instance of the parameter for a given trial.
 
@@ -137,12 +146,12 @@ class CategoricalParameter(OptimizationParameter):
     """
 
     def __init__(
-        self, choices: Sequence[Union[None, bool, int, float, str]], *args, **kwargs
+        self, choices: Sequence[Optional[Union[bool, int, float, str]]], *args, **kwargs
     ):
         super().__init__(*args, **kwargs)
         self.choices = choices
 
-    def suggest(self, trial: 'Trial'):
+    def _suggest(self, trial: 'Trial'):
         """
         Suggest an instance of the parameter for a given trial.
 
@@ -165,7 +174,7 @@ def __init__(self, low: float, high: float, q: float, *args, **kwargs):
         self.high = high
         self.q = q
 
-    def suggest(self, trial: 'Trial'):
+    def _suggest(self, trial: 'Trial'):
         """
         Suggest an instance of the parameter for a given trial.
 
@@ -180,6 +189,34 @@ def suggest(self, trial: 'Trial'):
         )
 
 
+class PodAlternativeParameter(CategoricalParameter):
+    """Base class for all optimization parameters."""
+
+    def __init__(
+        self,
+        inner_parameters: Dict[str, List[OptimizationParameter]],
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        assert len(self.choices) == len(
+            inner_parameters.keys()
+        ), 'Every pod alternative needs to have its set of parameters'
+        self.inner_parameters = inner_parameters
+
+    def update_trial_params(self, trial: 'Trial', trial_params: Dict):
+        """
+        Update the trial parameters given the variables and the trials
+
+        :param trial: An instance of an Optuna Trial object
+        :param trial_params: the parameters to update
+            # noqa: DAR201
+        """
+        super().update_trial_params(trial, trial_params)
+        for param in self.inner_parameters[trial_params[self.jaml_variable]]:
+            param.update_trial_params(trial, trial_params)
+
+
 def load_optimization_parameters(filepath: str):
     """
     Loads optimization parameters from a `.yml` file and parses it with the JAML parser.
diff --git a/tests/integration/optimizers/flow_pod_choice.yml b/tests/integration/optimizers/flow_pod_choice.yml
new file mode 100644
index 0000000000000..289d77a26649f
--- /dev/null
+++ b/tests/integration/optimizers/flow_pod_choice.yml
@@ -0,0 +1,13 @@
+!Flow
+version: '1'
+env:
+  JINA_DUMMYCRAFTER_PARAM1_VAR: ${{JINA_DUMMYCRAFTER_PARAM1}}
+  JINA_DUMMYCRAFTER_PARAM2_VAR: ${{JINA_DUMMYCRAFTER_PARAM2}}
+  JINA_DUMMYCRAFTER_PARAM3_VAR: ${{JINA_DUMMYCRAFTER_PARAM3}}
+  JINA_DUMMYCRAFTER_PARAM4_VAR: ${{JINA_DUMMYCRAFTER_PARAM4}}
+  JINA_DUMMYCRAFTER_PARAM5_VAR: ${{JINA_DUMMYCRAFTER_PARAM5}}
+  JINA_DUMMYCRAFTER_PARAM6_VAR: ${{JINA_DUMMYCRAFTER_PARAM6}}
+  JINA_DUMMYCRAFTER_CHOICE_VAR: ${{JINA_DUMMYCRAFTER_CHOICE}}
+pods:
+  - uses: ${{root.env.JINA_DUMMYCRAFTER_CHOICE_VAR}}
+  - uses: pods/evaluate.yml
diff --git a/tests/integration/optimizers/parameter_pod_choice.yml b/tests/integration/optimizers/parameter_pod_choice.yml
new file mode 100644
index 0000000000000..cf7f76ef72389
--- /dev/null
+++ b/tests/integration/optimizers/parameter_pod_choice.yml
@@ -0,0 +1,45 @@
+- !PodAlternativeParameter
+  parameter_name: pod_choice
+  jaml_variable: JINA_DUMMYCRAFTER_CHOICE
+  choices:
+    - pods/craft_option1.yml
+    - pods/craft_option2.yml
+  inner_parameters:
+    pods/craft_option1.yml:
+      - !IntegerParameter
+        jaml_variable: JINA_DUMMYCRAFTER_PARAM1
+        high: 1
+        low: 0
+        step_size: 1
+        parameter_name: param1
+      - !IntegerParameter
+        jaml_variable: JINA_DUMMYCRAFTER_PARAM2
+        high: 1
+        low: 1
+        step_size: 1
+        parameter_name: param2
+      - !IntegerParameter
+        jaml_variable: JINA_DUMMYCRAFTER_PARAM3
+        high: 2
+        low: 1
+        step_size: 1
+        parameter_name: param3
+    pods/craft_option2.yml:
+      - !IntegerParameter
+        jaml_variable: JINA_DUMMYCRAFTER_PARAM4
+        high: 1
+        low: 0
+        step_size: 1
+        parameter_name: param4
+      - !IntegerParameter
+        jaml_variable: JINA_DUMMYCRAFTER_PARAM5
+        high: 1
+        low: 1
+        step_size: 1
+        parameter_name: param5
+      - !IntegerParameter
+        jaml_variable: JINA_DUMMYCRAFTER_PARAM6
+        high: 2
+        low: 1
+        step_size: 1
+        parameter_name: param6
diff --git a/tests/integration/optimizers/pods/craft_option1.yml b/tests/integration/optimizers/pods/craft_option1.yml
new file mode 100644
index 0000000000000..753aed0c856eb
--- /dev/null
+++ b/tests/integration/optimizers/pods/craft_option1.yml
@@ -0,0 +1,7 @@
+!DummyCrafterOption1
+with:
+  param1: ${{JINA_DUMMYCRAFTER_PARAM1_VAR}}
+  param2: ${{JINA_DUMMYCRAFTER_PARAM2_VAR}}
+  param3: ${{JINA_DUMMYCRAFTER_PARAM3_VAR}}
+metas:
+  py_modules: 'dummy_crafter_pod_choice.py'
diff --git a/tests/integration/optimizers/pods/craft_option2.yml b/tests/integration/optimizers/pods/craft_option2.yml
new file mode 100644
index 0000000000000..5471a16ed9058
--- /dev/null
+++ b/tests/integration/optimizers/pods/craft_option2.yml
@@ -0,0 +1,7 @@
+!DummyCrafterOption2
+with:
+  param4: ${{JINA_DUMMYCRAFTER_PARAM4_VAR}}
+  param5: ${{JINA_DUMMYCRAFTER_PARAM5_VAR}}
+  param6: ${{JINA_DUMMYCRAFTER_PARAM6_VAR}}
+metas:
+  py_modules: 'dummy_crafter_pod_choice.py'
diff --git a/tests/integration/optimizers/pods/dummy_crafter.py b/tests/integration/optimizers/pods/dummy_crafter.py
index 330c2a3888fac..a2cea24364d8c 100644
--- a/tests/integration/optimizers/pods/dummy_crafter.py
+++ b/tests/integration/optimizers/pods/dummy_crafter.py
@@ -1,5 +1,3 @@
-from typing import Dict
-
 from jina import Executor, requests
 
 from jina.optimizers.parameters import IntegerParameter
@@ -49,7 +47,7 @@ def good_params(self):
         )
 
     @requests
-    def craft(self, docs, *args, **kwargs) -> Dict:
+    def craft(self, docs, *args, **kwargs):
         for doc in docs:
             if not self.good_params:
                 doc.text = ''
diff --git a/tests/integration/optimizers/pods/dummy_crafter_pod_choice.py b/tests/integration/optimizers/pods/dummy_crafter_pod_choice.py
new file mode 100644
index 0000000000000..eeb2fbcef0372
--- /dev/null
+++ b/tests/integration/optimizers/pods/dummy_crafter_pod_choice.py
@@ -0,0 +1,110 @@
+from jina import Executor, requests
+
+from jina.optimizers.parameters import IntegerParameter
+from jina.logging.logger import JinaLogger
+
+
+class DummyCrafterOption1(Executor):
+    DEFAULT_OPTIMIZATION_PARAMETER = [
+        IntegerParameter(
+            executor_name='DummyCrafterOption1',
+            parameter_name='param1',
+            low=0,
+            high=1,
+            step_size=1,
+        ),
+        IntegerParameter(
+            executor_name='DummyCrafterOption1',
+            parameter_name='param2',
+            low=0,
+            high=1,
+            step_size=1,
+        ),
+        IntegerParameter(
+            executor_name='DummyCrafterOption1',
+            parameter_name='param3',
+            low=0,
+            high=1,
+            step_size=1,
+        ),
+    ]
+
+    GOOD_PARAM_1 = 0
+    GOOD_PARAM_2 = 1
+    GOOD_PARAM_3 = 1
+
+    def __init__(self, param1: int, param2: int, param3: int, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.param1 = param1
+        self.param2 = param2
+        self.param3 = param3
+        self.logger = JinaLogger(self.__class__.__name__)
+
+    @property
+    def good_params(self):
+        return (
+            self.param1 == DummyCrafterOption1.GOOD_PARAM_1
+            and self.param2 == DummyCrafterOption1.GOOD_PARAM_2
+            and self.param3 == DummyCrafterOption1.GOOD_PARAM_3
+        )
+
+    @requests
+    def craft(self, docs, *args, **kwargs):
+        for doc in docs:
+            if not self.good_params:
+                doc.text = ''
+            if self.good_params and doc.text == 'DummyCrafterOption1':
+                doc.text = 'hello'
+
+
+class DummyCrafterOption2(Executor):
+    DEFAULT_OPTIMIZATION_PARAMETER = [
+        IntegerParameter(
+            executor_name='DummyCrafterOption2',
+            parameter_name='param4',
+            low=0,
+            high=1,
+            step_size=1,
+        ),
+        IntegerParameter(
+            executor_name='DummyCrafterOption2',
+            parameter_name='param5',
+            low=0,
+            high=1,
+            step_size=1,
+        ),
+        IntegerParameter(
+            executor_name='DummyCrafterOption2',
+            parameter_name='param6',
+            low=0,
+            high=1,
+            step_size=1,
+        ),
+    ]
+
+    GOOD_PARAM_4 = 0
+    GOOD_PARAM_5 = 1
+    GOOD_PARAM_6 = 1
+
+    def __init__(self, param4: int, param5: int, param6: int, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.param4 = param4
+        self.param5 = param5
+        self.param6 = param6
+        self.logger = JinaLogger(self.__class__.__name__)
+
+    @property
+    def good_params(self):
+        return (
+            self.param4 == DummyCrafterOption2.GOOD_PARAM_4
+            and self.param5 == DummyCrafterOption2.GOOD_PARAM_5
+            and self.param6 == DummyCrafterOption2.GOOD_PARAM_6
+        )
+
+    @requests
+    def craft(self, docs, *args, **kwargs):
+        for doc in docs:
+            if not self.good_params:
+                doc.text = ''
+            if self.good_params and doc.text == 'DummyCrafterOption2':
+                doc.text = 'hello'
diff --git a/tests/integration/optimizers/pods/dummy_evaluate.py b/tests/integration/optimizers/pods/dummy_evaluate.py
index 69f4512809d13..44c6fdac7be05 100644
--- a/tests/integration/optimizers/pods/dummy_evaluate.py
+++ b/tests/integration/optimizers/pods/dummy_evaluate.py
@@ -1,10 +1,11 @@
 from jina import Executor, requests, DocumentArray
+from jina.logging.logger import JinaLogger
 
 
 class DummyTextEvaluator(Executor):
-    @property
-    def metric(self) -> str:
-        return 'DummyTextEvaluator'
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.logger = JinaLogger(self.__class__.__name__)
 
     @requests
     def evaluate(self, docs: 'DocumentArray', groundtruths: 'DocumentArray', **kwargs):
diff --git a/tests/integration/optimizers/test_optimizer.py b/tests/integration/optimizers/test_optimizer.py
index f31209aeaeb9c..4439fdbe6ee10 100644
--- a/tests/integration/optimizers/test_optimizer.py
+++ b/tests/integration/optimizers/test_optimizer.py
@@ -46,21 +46,31 @@ def document_generator(num_doc):
         yield doc, groundtruth_doc
 
 
-def test_optimizer_single_flow(tmpdir, config):
+@pytest.mark.parametrize('sampler', ['TPESampler', 'GridSampler', 'RandomSampler'])
+def test_optimizer_single_flow(tmpdir, config, sampler):
     eval_flow_runner = SingleFlowRunner(
         flow_yaml=os.path.join(cur_dir, 'flow.yml'),
         documents=document_generator(10),
         request_size=1,
         execution_endpoint='search',
     )
+    grid_sampler_search_space = {
+        'JINA_DUMMYCRAFTER_PARAM1': [0, 1],
+        'JINA_DUMMYCRAFTER_PARAM2': [0, 1, 2],
+        'JINA_DUMMYCRAFTER_PARAM3': [1],
+    }
     opt = FlowOptimizer(
         flow_runner=eval_flow_runner,
         parameter_yaml=os.path.join(cur_dir, 'parameter.yml'),
         evaluation_callback=MeanEvaluationCallback(),
         workspace_base_dir=str(tmpdir),
         n_trials=5,
+        sampler=sampler,
     )
-    result = opt.optimize_flow()
+    if sampler == 'GridSampler':
+        result = opt.optimize_flow(search_space=grid_sampler_search_space)
+    else:
+        result = opt.optimize_flow()
     validate_result(result, tmpdir)
 
 
diff --git a/tests/integration/optimizers/test_optimizer_pod_choice.py b/tests/integration/optimizers/test_optimizer_pod_choice.py
new file mode 100644
index 0000000000000..4b07c618bbd65
--- /dev/null
+++ b/tests/integration/optimizers/test_optimizer_pod_choice.py
@@ -0,0 +1,86 @@
+import json
+import os
+
+import pytest
+import yaml
+from google.protobuf.json_format import MessageToJson
+
+from jina import Document
+from jina.jaml import JAML
+from jina.optimizers import FlowOptimizer, MeanEvaluationCallback
+from jina.optimizers import run_optimizer_cli
+from jina.optimizers.flow_runner import SingleFlowRunner, MultiFlowRunner
+from jina.parsers.optimizer import set_optimizer_parser
+
+cur_dir = os.path.dirname(os.path.abspath(__file__))
+
+
+@pytest.fixture
+def config(tmpdir):
+    os.environ['JINA_OPTIMIZER_WORKSPACE_DIR'] = str(tmpdir)
+    os.environ['JINA_OPTIMIZER_PARAMETER_FILE'] = os.path.join(cur_dir, 'parameter.yml')
+    os.environ['JINA_OPTIMIZER_DATA_FILE'] = os.path.join(cur_dir, 'data.jsonlines')
+    yield
+    del os.environ['JINA_OPTIMIZER_WORKSPACE_DIR']
+    del os.environ['JINA_OPTIMIZER_PARAMETER_FILE']
+    del os.environ['JINA_OPTIMIZER_DATA_FILE']
+
+
+def document_generator_option1(num_doc):
+    for _ in range(num_doc):
+        doc = Document(content='DummyCrafterOption1')
+        groundtruth_doc = Document(content='hello')
+        yield doc, groundtruth_doc
+
+
+def document_generator_option2(num_doc):
+    for _ in range(num_doc):
+        doc = Document(content='DummyCrafterOption2')
+        groundtruth_doc = Document(content='hello')
+        yield doc, groundtruth_doc
+
+
+def test_optimizer_single_flow_option1(tmpdir, config):
+    eval_flow_runner = SingleFlowRunner(
+        flow_yaml=os.path.join(cur_dir, 'flow_pod_choice.yml'),
+        documents=document_generator_option1(10),
+        request_size=1,
+        execution_endpoint='search',
+    )
+    opt = FlowOptimizer(
+        flow_runner=eval_flow_runner,
+        parameter_yaml=os.path.join(cur_dir, 'parameter_pod_choice.yml'),
+        evaluation_callback=MeanEvaluationCallback(),
+        workspace_base_dir=str(tmpdir),
+        n_trials=10,
+    )
+    result = opt.optimize_flow()
+    assert (
+        result.best_parameters['JINA_DUMMYCRAFTER_CHOICE'] == 'pods/craft_option1.yml'
+    )
+    assert result.best_parameters['JINA_DUMMYCRAFTER_PARAM1'] == 0
+    assert result.best_parameters['JINA_DUMMYCRAFTER_PARAM2'] == 1
+    assert result.best_parameters['JINA_DUMMYCRAFTER_PARAM3'] == 1
+
+
+def test_optimizer_single_flow_option2(tmpdir, config):
+    eval_flow_runner = SingleFlowRunner(
+        flow_yaml=os.path.join(cur_dir, 'flow_pod_choice.yml'),
+        documents=document_generator_option2(10),
+        request_size=1,
+        execution_endpoint='search',
+    )
+    opt = FlowOptimizer(
+        flow_runner=eval_flow_runner,
+        parameter_yaml=os.path.join(cur_dir, 'parameter_pod_choice.yml'),
+        evaluation_callback=MeanEvaluationCallback(),
+        workspace_base_dir=str(tmpdir),
+        n_trials=20,
+    )
+    result = opt.optimize_flow()
+    assert (
+        result.best_parameters['JINA_DUMMYCRAFTER_CHOICE'] == 'pods/craft_option2.yml'
+    )
+    assert result.best_parameters['JINA_DUMMYCRAFTER_PARAM4'] == 0
+    assert result.best_parameters['JINA_DUMMYCRAFTER_PARAM5'] == 1
+    assert result.best_parameters['JINA_DUMMYCRAFTER_PARAM6'] == 1
diff --git a/tests/integration/sparse_pipeline/__init__.py b/tests/integration/sparse_pipeline/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/tests/unit/optimizers/__init__.py b/tests/unit/optimizers/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/tests/unit/optimizers/test_parameters.py b/tests/unit/optimizers/test_parameters.py
new file mode 100644
index 0000000000000..dd964888792db
--- /dev/null
+++ b/tests/unit/optimizers/test_parameters.py
@@ -0,0 +1,168 @@
+import pytest
+import optuna
+
+from jina.optimizers.parameters import (
+    IntegerParameter,
+    UniformParameter,
+    LogUniformParameter,
+    CategoricalParameter,
+    DiscreteUniformParameter,
+    PodAlternativeParameter,
+)
+
+
+@pytest.fixture()
+def optuna_sampler():
+    return optuna.samplers.TPESampler()
+
+
+def test_integer_parameter(optuna_sampler):
+    parameter = IntegerParameter(
+        jaml_variable='JINA_DUMMY',
+        high=10,
+        low=0,
+        step_size=1,
+        parameter_name='integer',
+    )
+
+    def objective(trial):
+        trial_parameters = {}
+        parameter.update_trial_params(trial, trial_parameters)
+        assert 'JINA_DUMMY' in trial_parameters.keys()
+        assert 0 <= trial_parameters['JINA_DUMMY'] <= 10
+        return 0.0
+
+    study = optuna.create_study(direction='maximize', sampler=optuna_sampler)
+    study.optimize(objective, n_trials=5)
+
+
+def test_uniform_parameter(optuna_sampler):
+    parameter = UniformParameter(
+        jaml_variable='JINA_DUMMY', high=10, low=0, parameter_name='uniform'
+    )
+
+    def objective(trial):
+        trial_parameters = {}
+        parameter.update_trial_params(trial, trial_parameters)
+        assert 'JINA_DUMMY' in trial_parameters.keys()
+        assert 0 <= trial_parameters['JINA_DUMMY'] <= 10
+        return 0.0
+
+    study = optuna.create_study(direction='maximize', sampler=optuna_sampler)
+    study.optimize(objective, n_trials=5)
+
+
+def test_log_uniform_parameter(optuna_sampler):
+    parameter = LogUniformParameter(
+        jaml_variable='JINA_DUMMY', high=10, low=1, parameter_name='loguniform'
+    )
+
+    def objective(trial):
+        trial_parameters = {}
+        parameter.update_trial_params(trial, trial_parameters)
+        assert 'JINA_DUMMY' in trial_parameters.keys()
+        assert 0 <= trial_parameters['JINA_DUMMY'] <= 10
+        return 0.0
+
+    study = optuna.create_study(direction='maximize', sampler=optuna_sampler)
+    study.optimize(objective, n_trials=5)
+
+
+def test_categorical_parameter(optuna_sampler):
+    parameter = CategoricalParameter(
+        jaml_variable='JINA_DUMMY',
+        choices=[f'choice-{i}' for i in range(10)],
+        parameter_name='categorical',
+    )
+
+    def objective(trial):
+        trial_parameters = {}
+        parameter.update_trial_params(trial, trial_parameters)
+        assert 'JINA_DUMMY' in trial_parameters.keys()
+        assert trial_parameters['JINA_DUMMY'] in [f'choice-{i}' for i in range(10)]
+        return 0.0
+
+    study = optuna.create_study(direction='maximize', sampler=optuna_sampler)
+    study.optimize(objective, n_trials=5)
+
+
+def test_discrete_uniform_parameter(optuna_sampler):
+    parameter = DiscreteUniformParameter(
+        jaml_variable='JINA_DUMMY',
+        high=10,
+        low=0,
+        q=1,
+        parameter_name='discreteuniform',
+    )
+
+    def objective(trial):
+        trial_parameters = {}
+        parameter.update_trial_params(trial, trial_parameters)
+        assert 'JINA_DUMMY' in trial_parameters.keys()
+        assert 0 <= trial_parameters['JINA_DUMMY'] <= 10
+        return 0.0
+
+    study = optuna.create_study(direction='maximize', sampler=optuna_sampler)
+    study.optimize(objective, n_trials=5)
+
+
+def test_pod_alternative_parameter(optuna_sampler):
+    inner_parameters = {
+        'pod1': [
+            IntegerParameter(
+                jaml_variable='JINA_INTEGER_DUMMY_POD1',
+                high=10,
+                low=0,
+                parameter_name='integerparam_pod1',
+            ),
+            CategoricalParameter(
+                jaml_variable='JINA_CAT_DUMMY_POD1',
+                choices=[f'choice-{i}' for i in range(10)],
+                parameter_name='categorical_pod1',
+            ),
+        ],
+        'pod2': [
+            IntegerParameter(
+                jaml_variable='JINA_INTEGER_DUMMY_POD2',
+                high=10,
+                low=0,
+                parameter_name='integerparam_pod2',
+            ),
+            CategoricalParameter(
+                jaml_variable='JINA_CAT_DUMMY_POD2',
+                choices=[f'choice-{i}' for i in range(10)],
+                parameter_name='categorical_pod2',
+            ),
+        ],
+    }
+    parameter = PodAlternativeParameter(
+        jaml_variable='JINA_DUMMY',
+        choices=['pod1', 'pod2'],
+        inner_parameters=inner_parameters,
+        parameter_name='podalternative',
+    )
+
+    def objective(trial):
+        trial_parameters = {}
+        parameter.update_trial_params(trial, trial_parameters)
+        assert 'JINA_DUMMY' in trial_parameters.keys()
+        assert trial_parameters['JINA_DUMMY'] in ['pod1', 'pod2']
+        if trial_parameters['JINA_DUMMY'] == 'pod1':
+            assert 'JINA_INTEGER_DUMMY_POD1' in trial_parameters.keys()
+            assert 0 <= trial_parameters['JINA_INTEGER_DUMMY_POD1'] <= 10
+            assert 'JINA_CAT_DUMMY_POD1' in trial_parameters.keys()
+            assert trial_parameters['JINA_CAT_DUMMY_POD1'] in [
+                f'choice-{i}' for i in range(10)
+            ]
+        if trial_parameters['JINA_DUMMY'] == 'pod2':
+            assert 'JINA_INTEGER_DUMMY_POD2' in trial_parameters.keys()
+            assert 0 <= trial_parameters['JINA_INTEGER_DUMMY_POD2'] <= 10
+            assert 'JINA_CAT_DUMMY_POD2' in trial_parameters.keys()
+            assert trial_parameters['JINA_CAT_DUMMY_POD2'] in [
+                f'choice-{i}' for i in range(10)
+            ]
+
+        return 0.0
+
+    study = optuna.create_study(direction='maximize', sampler=optuna_sampler)
+    study.optimize(objective, n_trials=5)


Code review comments:
diff_hunk: @@ -0,0 +1,86 @@
+import json
+import os
+
+import pytest
+import yaml
+from google.protobuf.json_format import MessageToJson
+
+from jina import Document
+from jina.jaml import JAML
+from jina.optimizers import FlowOptimizer, MeanEvaluationCallback
+from jina.optimizers import run_optimizer_cli
+from jina.optimizers.flow_runner import SingleFlowRunner, MultiFlowRunner
+from jina.parsers.optimizer import set_optimizer_parser
path: tests/integration/optimizers/test_optimizer_pod_choice.py
text: @user1:
some unused imports