Repository: jina-ai/serve
Pull request: feat: custom monitoring with context manager
# Custom monitoring with context manager
## Context

having only the possibility to monitor custom part of Executor with a decorator is a limitation. We need to have a context manager as well. See full issue here : https://github.com/jina-ai/jina/issues/4876

## What this PR do

This pr allow developer of Executor to use custom metrics with a context manager

example:

```python
from jina import requests, Executor
from docarray import DocumentArray

def preprocessing(self, docs: DocumentArray):
    ...

def model_inference(self, tensor):
    ...

class MyExecutor(Executor):

    @requests
    def encode(self, docs: DocumentArray, **kwargs):
        with self.monitor('preprocessing_seconds', 'Time preprocessing the requests'):
            docs.tensors = preprocessing(docs)
        with self.monitor('model_inference_seconds', 'Time doing inference the requests'):
            docs.embedding = model_inference(docs.tensors)
```



Problem statement:
Make the monitor decorator a context manager as well
If you always force user to create new functions, how do you do this
```python
def foo():
   shared_var
   some_code_block_to_measure1
   more_shared_var
   some_code_block_to_measure2
   more_shared_var
   some_code_block_to_measure3
```
adding decorator to the three blocks means pulling them out with nontrivial refactoring, and as a developer, I really do not appreciate it.
Hence, it is better to use with context manager
```python
def foo():
   shared_var
   with monitor(...):
     some_code_block_to_measure1
   more_shared_var
   with monitor(...):
     some_code_block_to_measure2
   more_shared_var
   with monitor(...):
     some_code_block_to_measure3
```
As the developer of CAS, I would rather not have monitor feature at all, then seeing monitor feature breaks my project structure. The reason is simple, developers do not trust the pattern the @monitor imposed on the over their own pattern. If developers have to refactor their function structure very awkwardly to add monitoring, they'd rather not have monitoring.

Hints from issue discussion:

Commit message: fix(docs): update example

Patch to review:
diff --git a/docs/fundamentals/executor/monitoring-executor.md b/docs/fundamentals/executor/monitoring-executor.md
index d3524206c713b..5360f26408f88 100644
--- a/docs/fundamentals/executor/monitoring-executor.md
+++ b/docs/fundamentals/executor/monitoring-executor.md
@@ -41,13 +41,17 @@ Adding the custom monitoring on a method is as straightforward as decorating the
 ````
 
 ```python
-from jina import Executor, monitor
+from jina import Executor, monitor, requests
 
 
 class MyExecutor(Executor):
     @monitor()
     def my_method(self):
         ...
+
+    @requests
+    def foo(self, **kwargs):
+        self.my_method()
 ```
 
 This will create a [Prometheus summary](https://prometheus.io/docs/concepts/metric_types/#summary)
@@ -64,6 +68,25 @@ def method(self):
     ...
 ```
 
+You can as well monitor internal part of your executor with a Context Manager:
+
+```python
+from jina import Executor, requests
+from docarray import DocumentArray
+
+
+def process():
+    ...
+
+
+class MyExecutor(Executor):
+    @requests
+    def foo(self, docs, **kwargs):
+        with self.monitor('processing_seconds', 'Time processing my document'):
+            docs = process(docs)
+```
+
+
 ````{admonition} respect Prometheus naming
 :class: caution
 You should respect Prometheus naming [conventions](https://prometheus.io/docs/practices/naming/#metric-names). 
@@ -99,11 +122,13 @@ The encode function is composed of two sub-functions.
 
 By default, only the `encode` function will be monitored. 
 
+````{tab} Decorator
+
 ```{code-block} python
 ---
 emphasize-lines: 6, 10
 ---
-from jina import requests,monitor,Executor
+from jina import requests,monitor, Executor
 from docarray import DocumentArray
 
 class MyExecutor(Executor):
@@ -121,6 +146,32 @@ class MyExecutor(Executor):
         docs.tensors = self.preprocessing(docs)
         docs.embedding = self.model_inference(docs.tensors)
 ```
+````
+
+````{tab} Context Manager
+```{code-block} python
+---
+emphasize-lines: 14, 16
+---
+from jina import requests, Executor
+from docarray import DocumentArray
+
+def preprocessing(self, docs: DocumentArray):
+    ...
+
+def model_inference(self, tensor):
+    ...
+
+class MyExecutor(Executor):
+
+    @requests
+    def encode(self, docs: DocumentArray, **kwargs):
+        with self.monitor('preprocessing_seconds', 'Time preprocessing the requests'):
+            docs.tensors = preprocessing(docs)
+        with self.monitor('model_inference_seconds', 'Time doing inference the requests'):
+            docs.embedding = model_inference(docs.tensors)
+```
+````
 
 
 ### Defining custom metrics directly with the Prometheus client
diff --git a/jina/serve/executors/__init__.py b/jina/serve/executors/__init__.py
index ae5fe515a1b9c..51bd81fc6b5b6 100644
--- a/jina/serve/executors/__init__.py
+++ b/jina/serve/executors/__init__.py
@@ -490,7 +490,7 @@ def to_docker_compose_yaml(
             == BaseExecutor.StandaloneExecutorType.EXTERNAL,
         )
 
-    def get_metrics(
+    def monitor(
         self, name: Optional[str] = None, documentation: Optional[str] = None
     ) -> Optional['Summary']:
         """
@@ -512,6 +512,6 @@ def get_metrics(
                     namespace='jina',
                     labelnames=('runtime_name',),
                 ).labels(self.runtime_args.name)
-            return self._metrics_buffer[name]
+            return self._metrics_buffer[name].time()
         else:
-            return None
+            return contextlib.nullcontext()
diff --git a/jina/serve/executors/decorators.py b/jina/serve/executors/decorators.py
index 88e3e972813f2..1bb7c2aade933 100644
--- a/jina/serve/executors/decorators.py
+++ b/jina/serve/executors/decorators.py
@@ -229,12 +229,7 @@ def _decorator(func: Callable):
 
         @functools.wraps(func)
         def _f(self, *args, **kwargs):
-            metric = self.get_metrics(name_, documentation_)
-
-            if metric:
-                with metric.time():
-                    return func(self, *args, **kwargs)
-            else:
+            with self.monitor(name_, documentation_):
                 return func(self, *args, **kwargs)
 
         return _f
diff --git a/tests/integration/monitoring/test_executor.py b/tests/integration/monitoring/test_executor.py
index 0caf658d6d2c4..563e57c923785 100644
--- a/tests/integration/monitoring/test_executor.py
+++ b/tests/integration/monitoring/test_executor.py
@@ -59,3 +59,41 @@ def proces_2(self, docs):
             f'jina_proces_2_seconds_count{{runtime_name="executor0/rep-0"}} 1.0'
             in str(resp.content)
         )
+
+
+def test_context_manager_interface(port_generator):
+    class DummyExecutor(Executor):
+        @requests(on='/foo')
+        def foo(self, docs, **kwargs):
+
+            with self.monitor(
+                name='proces_seconds', documentation='process time in seconds '
+            ):
+                self._proces(docs)
+
+            with self.monitor(
+                name='proces_2_seconds', documentation='process 2 time in seconds '
+            ):
+                self.proces_2(docs)
+
+        def _proces(self, docs):
+            ...
+
+        def proces_2(self, docs):
+            ...
+
+    port = port_generator()
+    with Flow(monitoring=True, port_monitoring=port_generator()).add(
+        uses=DummyExecutor, monitoring=True, port_monitoring=port
+    ) as f:
+        f.post('/foo', inputs=DocumentArray.empty(4))
+
+        resp = req.get(f'http://localhost:{port}/')
+        assert (
+            f'jina_proces_seconds_count{{runtime_name="executor0/rep-0"}} 1.0'
+            in str(resp.content)
+        )
+        assert (
+            f'jina_proces_2_seconds_count{{runtime_name="executor0/rep-0"}} 1.0'
+            in str(resp.content)
+        )
diff --git a/tests/unit/serve/runtimes/worker/test_worker_runtime.py b/tests/unit/serve/runtimes/worker/test_worker_runtime.py
index 3a612a94178a3..ad321e20f72ec 100644
--- a/tests/unit/serve/runtimes/worker/test_worker_runtime.py
+++ b/tests/unit/serve/runtimes/worker/test_worker_runtime.py
@@ -444,3 +444,70 @@ def start_runtime(args, cancel_event):
     runtime_thread.join()
 
     assert not AsyncNewLoopRuntime.is_ready(f'{args.host}:{args.port}')
+
+
+@pytest.mark.asyncio
+@pytest.mark.slow
+@pytest.mark.timeout(5)
+async def test_decorator_monitoring(port_generator):
+    class DummyExecutor(Executor):
+        @requests
+        def foo(self, docs, **kwargs):
+
+            with self.monitor(
+                name='proces_seconds', documentation='process time in seconds '
+            ):
+                self._proces(docs)
+
+            with self.monitor(
+                name='proces_2_seconds', documentation='process 2 time in seconds '
+            ):
+                self.proces_2(docs)
+
+        def _proces(self, docs):
+            ...
+
+        def proces_2(self, docs):
+            ...
+
+    port = port_generator()
+    args = set_pod_parser().parse_args(
+        ['--monitoring', '--port-monitoring', str(port), '--uses', 'DummyExecutor']
+    )
+
+    cancel_event = multiprocessing.Event()
+
+    def start_runtime(args, cancel_event):
+        with WorkerRuntime(args, cancel_event=cancel_event) as runtime:
+            runtime.run_forever()
+
+    runtime_thread = Process(
+        target=start_runtime,
+        args=(args, cancel_event),
+        daemon=True,
+    )
+    runtime_thread.start()
+
+    assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
+        timeout=5.0,
+        ctrl_address=f'{args.host}:{args.port}',
+        ready_or_shutdown_event=Event(),
+    )
+
+    assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
+        timeout=5.0,
+        ctrl_address=f'{args.host}:{args.port}',
+        ready_or_shutdown_event=Event(),
+    )
+
+    await GrpcConnectionPool.send_request_async(
+        _create_test_data_message(), f'{args.host}:{args.port}', timeout=1.0
+    )
+
+    resp = req.get(f'http://localhost:{port}/')
+    assert f'jina_proces_seconds_count{{runtime_name="None"}} 1.0' in str(resp.content)
+
+    cancel_event.set()
+    runtime_thread.join()
+
+    assert not AsyncNewLoopRuntime.is_ready(f'{args.host}:{args.port}')


Code review comments:
diff_hunk: @@ -59,3 +59,41 @@ def proces_2(self, docs):
             f'jina_proces_2_seconds_count{{runtime_name="executor0/rep-0"}} 1.0'
             in str(resp.content)
         )
+
+
+def test_context_manager_interface(port_generator):
+    class DummyExecutor(Executor):
+        @requests(on='/foo')
+        def foo(self, docs, **kwargs):
+
+            with self.monitor(
+                name='proces_seconds', documentation='process time in seconds '
+            ):
+                self._proces(docs)
+
+            with self.monitor(
+                name='proces_2_seconds', documentation='process 2 time in seconds '
+            ):
+                self.proces_2(docs)
+
+        def _proces(self, docs):
path: tests/integration/monitoring/test_executor.py
text: @user1:
```suggestion
        def _process(self, docs):
```