Repository: jina-ai/serve
Pull request: feat: allow different evaluation aggregations

Problem statement:
change MeanEvaluationCallback to EvaluationCallback supporting multiple operations in optimizer
**Describe the feature**
MeanEvaluationCallback is the class used to collect evaluations from FlowRunner runs.

However, it only accepts a single way to collect results into a final metric. We can easily support multiple times the same way `SimpleAggregateRanker` did

Hints from issue discussion:

Commit message: Merge branch 'master' of https://github.com/jina-ai/jina into feat-evaluation-callback

Patch to review:
diff --git a/jina/optimizers/__init__.py b/jina/optimizers/__init__.py
index 62311baaf57c8..8315f2ca6d183 100644
--- a/jina/optimizers/__init__.py
+++ b/jina/optimizers/__init__.py
@@ -1,6 +1,7 @@
 import os
 from collections import defaultdict
 from typing import Optional
+import numpy as np
 
 import yaml
 
@@ -16,6 +17,7 @@
 from ..importer import ImportExtensions
 from ..jaml import JAMLCompatible, JAML
 from ..logging.predefined import default_logger as logger
+from ..types.request import Response
 
 if False:
     from .flow_runner import FlowRunner
@@ -48,7 +50,7 @@ def get_final_evaluation(self) -> float:
         """
         raise NotImplementedError
 
-    def __call__(self, response):
+    def __call__(self, response: 'Response'):
         """
         Collects the results of evaluators in the response object for aggregation.
 
@@ -58,31 +60,44 @@ def __call__(self, response):
         raise NotImplementedError
 
 
-class MeanEvaluationCallback(OptimizerCallback):
+class EvaluationCallback(OptimizerCallback):
     """
-    Calculates the mean of all evaluations during a single :py:class:`FlowRunner`
-    execution from the :py:class:`FlowOptimizer`.
+    Calculates an aggrefation of all evaluations during a single :py:class:`FlowRunner`
+    execution from the :py:class:`FlowOptimizer` with a given operator.
+
+    Valid operators are: `['min', 'max', 'mean', 'median', 'sum', 'prod']`
     """
 
-    def __init__(self, eval_name: Optional[str] = None):
+    AGGREGATE_FUNCTIONS = ['min', 'max', 'mean', 'median', 'sum', 'prod']
+
+    def __init__(self, eval_name: Optional[str] = None, operation: str = 'mean'):
         """
         :param eval_name: evaluation name as required by the evaluator. Not needed if only 1 evaluator is used
+        :param operation: name of the operation to be performed when getting the final evaluation
         """
+        self._operation = operation
+        if operation in self.AGGREGATE_FUNCTIONS:
+            self.np_aggregate_function = getattr(np, operation)
+        else:
+            raise ValueError(
+                f'The operation "{operation}" is not among the supported operators "{self.AGGREGATE_FUNCTIONS}".'
+            )
+
         self._eval_name = eval_name
-        self._evaluation_values = defaultdict(float)
+        self._evaluation_values = defaultdict(list)
         self._n_docs = 0
 
     def get_empty_copy(self):
         """
-        Return an empty copy of the :class:`OptimizerCallback`.
+        Return an empty copy of the :class:`EvaluationCallback`.
 
         :return: Evaluation values
         """
-        return MeanEvaluationCallback(self._eval_name)
+        return EvaluationCallback(self._eval_name, self._operation)
 
     def get_final_evaluation(self):
         """
-        Calculates and returns mean evaluation value on the metric defined in the :method:`__init__`.
+        Calculates and returns evaluation value on the metric defined in the :method:`__init__`.
 
         :return:: The aggregation of all evaluation collected via :method:`__call__`
         """
@@ -95,9 +110,9 @@ def get_final_evaluation(self):
                     f'More than one evaluation metric found. Please define the right eval_name. Currently {evaluation_name} is used'
                 )
 
-        return self._evaluation_values[evaluation_name] / self._n_docs
+        return self.np_aggregate_function(self._evaluation_values[evaluation_name])
 
-    def __call__(self, response):
+    def __call__(self, response: 'Response'):
         """
         Store the evaluation values
 
@@ -107,7 +122,7 @@ def __call__(self, response):
         logger.info(f'Num of docs evaluated: {self._n_docs}')
         for doc in response.data.docs:
             for evaluation in doc.evaluations:
-                self._evaluation_values[evaluation.op_name] += evaluation.value
+                self._evaluation_values[evaluation.op_name].append(evaluation.value)
 
 
 class ResultProcessor(JAMLCompatible):
diff --git a/tests/integration/graph/test_graph_executor.py b/tests/integration/graph/test_graph_executor.py
index f116ee11a3e9c..c8b1ed25f9843 100644
--- a/tests/integration/graph/test_graph_executor.py
+++ b/tests/integration/graph/test_graph_executor.py
@@ -60,6 +60,7 @@ def graph():
 
 def test_flow_with_graph_executor(graph, mocker):
     def validate_resp(resp):
+        print(f' type {type(resp)} => {resp}')
         assert len(resp.data.docs) == 1
         for doc in resp.data.docs:
             graph = GraphDocument(doc)
diff --git a/tests/integration/optimizers/optimizer_conf.yml b/tests/integration/optimizers/optimizer_conf.yml
index 2e93c9c21e709..2e3f0103485ba 100644
--- a/tests/integration/optimizers/optimizer_conf.yml
+++ b/tests/integration/optimizers/optimizer_conf.yml
@@ -8,7 +8,7 @@
        documents: ${{JINA_OPTIMIZER_DATA_FILE}}
        request_size: 1
        execution_endpoint: 'search'
-   evaluation_callback: !MeanEvaluationCallback {}
+   evaluation_callback: !EvaluationCallback {}
    parameter_yaml: ${{JINA_OPTIMIZER_PARAMETER_FILE}}
    workspace_base_dir: ${{JINA_OPTIMIZER_WORKSPACE_DIR}}
    n_trials: 5
diff --git a/tests/integration/optimizers/test_optimizer.py b/tests/integration/optimizers/test_optimizer.py
index 4439fdbe6ee10..21af0fd3e6a31 100644
--- a/tests/integration/optimizers/test_optimizer.py
+++ b/tests/integration/optimizers/test_optimizer.py
@@ -7,7 +7,7 @@
 
 from jina import Document
 from jina.jaml import JAML
-from jina.optimizers import FlowOptimizer, MeanEvaluationCallback
+from jina.optimizers import FlowOptimizer, EvaluationCallback
 from jina.optimizers import run_optimizer_cli
 from jina.optimizers.flow_runner import SingleFlowRunner, MultiFlowRunner
 from jina.parsers.optimizer import set_optimizer_parser
@@ -62,7 +62,7 @@ def test_optimizer_single_flow(tmpdir, config, sampler):
     opt = FlowOptimizer(
         flow_runner=eval_flow_runner,
         parameter_yaml=os.path.join(cur_dir, 'parameter.yml'),
-        evaluation_callback=MeanEvaluationCallback(),
+        evaluation_callback=EvaluationCallback(),
         workspace_base_dir=str(tmpdir),
         n_trials=5,
         sampler=sampler,
@@ -94,7 +94,7 @@ def test_optimizer_multi_flow(tmpdir, config):
     opt = FlowOptimizer(
         flow_runner=multi_flow_runner,
         parameter_yaml=os.path.join(cur_dir, 'parameter.yml'),
-        evaluation_callback=MeanEvaluationCallback(),
+        evaluation_callback=EvaluationCallback(),
         workspace_base_dir=str(tmpdir),
         n_trials=5,
     )
@@ -123,7 +123,7 @@ def test_yaml_multi_flow(tmpdir, config):
             documents: {jsonlines_file}
             request_size: 1
             execution_endpoint: 'search'
-  evaluation_callback: !MeanEvaluationCallback {{}}
+  evaluation_callback: !EvaluationCallback {{}}
   parameter_yaml: '{os.path.join(cur_dir, 'parameter.yml')}'
   workspace_base_dir: {tmpdir}
   n_trials: 5
@@ -160,7 +160,7 @@ def test_yaml_single_flow(tmpdir, config):
       documents: {jsonlines_file}
       request_size: 1
       execution_endpoint: 'search'
-  evaluation_callback: !MeanEvaluationCallback {{}}
+  evaluation_callback: !EvaluationCallback {{}}
   parameter_yaml: '{os.path.join(cur_dir, 'parameter.yml')}'
   workspace_base_dir: {tmpdir}
   n_trials: 5
diff --git a/tests/integration/optimizers/test_optimizer_pod_choice.py b/tests/integration/optimizers/test_optimizer_pod_choice.py
index b5b8273ed826d..033ab2fa62d4f 100644
--- a/tests/integration/optimizers/test_optimizer_pod_choice.py
+++ b/tests/integration/optimizers/test_optimizer_pod_choice.py
@@ -3,7 +3,7 @@
 import pytest
 
 from jina import Document
-from jina.optimizers import FlowOptimizer, MeanEvaluationCallback
+from jina.optimizers import FlowOptimizer, EvaluationCallback
 from jina.optimizers.flow_runner import SingleFlowRunner
 
 cur_dir = os.path.dirname(os.path.abspath(__file__))
@@ -44,7 +44,7 @@ def test_optimizer_single_flow_option1(tmpdir, config):
     opt = FlowOptimizer(
         flow_runner=eval_flow_runner,
         parameter_yaml=os.path.join(cur_dir, 'parameter_pod_choice.yml'),
-        evaluation_callback=MeanEvaluationCallback(),
+        evaluation_callback=EvaluationCallback(),
         workspace_base_dir=str(tmpdir),
         n_trials=10,
     )
@@ -67,7 +67,7 @@ def test_optimizer_single_flow_option2(tmpdir, config):
     opt = FlowOptimizer(
         flow_runner=eval_flow_runner,
         parameter_yaml=os.path.join(cur_dir, 'parameter_pod_choice.yml'),
-        evaluation_callback=MeanEvaluationCallback(),
+        evaluation_callback=EvaluationCallback(),
         workspace_base_dir=str(tmpdir),
         n_trials=20,
     )
diff --git a/tests/unit/optimizers/test_evaluation_callback.py b/tests/unit/optimizers/test_evaluation_callback.py
new file mode 100644
index 0000000000000..57eb1e9cca211
--- /dev/null
+++ b/tests/unit/optimizers/test_evaluation_callback.py
@@ -0,0 +1,69 @@
+import pytest
+
+from jina.optimizers import EvaluationCallback
+from jina.types.request import Request
+from jina import Document
+
+EVAL_NAME = 'evaluation'
+NUM_REQUESTS = 3
+NUM_DOCS_PER_REQUEST = 4
+
+
+@pytest.fixture()
+def operator(request):
+    return request.param
+
+
+@pytest.fixture()
+def callback(operator):
+    return EvaluationCallback(EVAL_NAME, operator)
+
+
+@pytest.fixture()
+def responses():
+    def create_documents(multiplier):
+        d1 = Document()
+        score1 = d1.evaluations.add()
+        score1.op_name = 'evaluation'
+        score1.value = multiplier * 1.0
+        d2 = Document()
+        score2 = d2.evaluations.add()
+        score2.op_name = 'evaluation'
+        score2.value = multiplier * 1.0
+        d3 = Document()
+        score3 = d3.evaluations.add()
+        score3.op_name = 'evaluation'
+        score3.value = multiplier * 2.0
+        d4 = Document()
+        score4 = d4.evaluations.add()
+        score4.op_name = 'evaluation'
+        score4.value = multiplier * 2.0
+        return [d1, d2, d3, d4]
+
+    def create_request(multiplier):
+        req = Request()
+        req.request_type = 'data'
+        for doc in create_documents(multiplier):
+            req.docs.append(doc)
+        return req
+
+    return [create_request(multiplier=i + 1) for i in range(NUM_REQUESTS)]
+
+
+@pytest.mark.parametrize(
+    'operator, expected',
+    [
+        ('mean', 3.0),
+        ('max', 6.0),
+        ('min', 1.0),
+        ('sum', 36.0),
+        ('prod', 82944.0),
+        ('median', 2.5),
+    ],
+    indirect=['operator'],
+)
+def test_evaluation_callback(callback, responses, operator, expected):
+    for resp in responses:
+        callback(resp)
+
+    assert callback.get_final_evaluation() == expected


Code review comments:
diff_hunk: @@ -58,31 +60,44 @@ def __call__(self, response):
         raise NotImplementedError
 
 
-class MeanEvaluationCallback(OptimizerCallback):
+class EvaluationCallback(OptimizerCallback):
     """
-    Calculates the mean of all evaluations during a single :py:class:`FlowRunner`
-    execution from the :py:class:`FlowOptimizer`.
+    Calculates an aggrefation of all evaluations during a single :py:class:`FlowRunner`
path: jina/optimizers/__init__.py
text: @author:
```suggestion
    Calculates an aggregation of all evaluations during a single :py:class:`FlowRunner`
```

diff_hunk: @@ -60,6 +60,7 @@ def graph():
 
 def test_flow_with_graph_executor(graph, mocker):
     def validate_resp(resp):
+        print(f' type {type(resp)} => {resp}')
path: tests/integration/graph/test_graph_executor.py
text: @user1:
remove print

@author:
ohh

diff_hunk: @@ -58,31 +60,44 @@ def __call__(self, response):
         raise NotImplementedError
 
 
-class MeanEvaluationCallback(OptimizerCallback):
+class EvaluationCallback(OptimizerCallback):
     """
-    Calculates the mean of all evaluations during a single :py:class:`FlowRunner`
-    execution from the :py:class:`FlowOptimizer`.
+    Calculates an aggrefation of all evaluations during a single :py:class:`FlowRunner`
path: jina/optimizers/__init__.py
text: @user1:
small typo: aggregation

diff_hunk: @@ -1,6 +1,7 @@
 import os
 from collections import defaultdict
 from typing import Optional
+import numpy as np
path: jina/optimizers/__init__.py
text: @user1:
move this down to `import yaml`?