feat: add monitoring

server/clip_server/executors/clip_onnx.py

@@ -13,7 +13,7 @@
 )
 from clip_server.model import clip
 from clip_server.model.clip_onnx import CLIPOnnxModel
-from jina import Executor, requests, DocumentArray
+from jina import Executor, requests, DocumentArray, monitor
 
 
 class CLIPEncoder(Executor):
@@ -78,6 +78,16 @@ def __init__(
 
         self._model.start_sessions(sess_options=sess_options, providers=providers)
 
+    @monitor()
+    def _preproc_images(self, docs: 'DocumentArray'):
+        return preproc_image(
+            docs, preprocess_fn=self._preprocess_tensor, return_np=True
+        )
+
+    @monitor()
+    def _preproc_texts(self, docs: 'DocumentArray'):
+        return preproc_text(docs, return_np=True)
+
     @requests(on='/rank')
     async def rank(self, docs: 'DocumentArray', parameters: Dict, **kwargs):
         await self.encode(docs['@r,m'])
@@ -94,9 +104,7 @@ async def encode(self, docs: 'DocumentArray', **kwargs):
         # for image
         if _img_da:
             for minibatch, _contents in _img_da.map_batch(
-                partial(
-                    preproc_image, preprocess_fn=self._preprocess_tensor, return_np=True
-                ),
+                self._preproc_images,
                 batch_size=self._minibatch_size,
                 pool=self._pool,
             ):
@@ -112,7 +120,7 @@ async def encode(self, docs: 'DocumentArray', **kwargs):
                     # for text
         if _txt_da:
             for minibatch, _contents in _txt_da.map_batch(
-                partial(preproc_text, return_np=True),
+                self._preproc_texts,
                 batch_size=self._minibatch_size,
                 pool=self._pool,
             ):

server/clip_server/executors/clip_tensorrt.py

@@ -11,7 +11,7 @@
 )
 from clip_server.model import clip
 from clip_server.model.clip_trt import CLIPTensorRTModel
-from jina import Executor, requests, DocumentArray
+from jina import Executor, requests, DocumentArray, monitor
 
 
 class CLIPEncoder(Executor):
@@ -46,6 +46,19 @@ def __init__(
 
         self._model.start_engines()
 
+    @monitor()
+    def _preproc_images(self, docs: 'DocumentArray'):
+        return preproc_image(
+            docs,
+            preprocess_fn=self._preprocess_tensor,
+            device=self._device,
+            return_np=False,
+        )
+
+    @monitor()
+    def _preproc_texts(self, docs: 'DocumentArray'):
+        return preproc_text(docs, device=self._device, return_np=False)
+
     @requests(on='/rank')
     async def rank(self, docs: 'DocumentArray', parameters: Dict, **kwargs):
         await self.encode(docs['@r,m'])
@@ -62,12 +75,7 @@ async def encode(self, docs: 'DocumentArray', **kwargs):
         # for image
         if _img_da:
             for minibatch, _contents in _img_da.map_batch(
-                partial(
-                    preproc_image,
-                    preprocess_fn=self._preprocess_tensor,
-                    device=self._device,
-                    return_np=False,
-                ),
+                self._preproc_images,
                 batch_size=self._minibatch_size,
                 pool=self._pool,
             ):
@@ -89,7 +97,7 @@ async def encode(self, docs: 'DocumentArray', **kwargs):
                     # for text
         if _txt_da:
             for minibatch, _contents in _txt_da.map_batch(
-                partial(preproc_text, device=self._device, return_np=False),
+                self._preproc_texts,
                 batch_size=self._minibatch_size,
                 pool=self._pool,
             ):

server/clip_server/executors/clip_torch.py

@@ -13,7 +13,7 @@
     set_rank,
 )
 from clip_server.model import clip
-from jina import Executor, requests, DocumentArray
+from jina import Executor, requests, DocumentArray, monitor
 
 
 class CLIPEncoder(Executor):
@@ -59,9 +59,21 @@ def __init__(
 
         self._pool = ThreadPool(processes=num_worker_preprocess)
 
+    @monitor()
+    def _preproc_images(self, docs: 'DocumentArray'):
+        return preproc_image(
+            docs,
+            preprocess_fn=self._preprocess_tensor,
+            device=self._device,
+            return_np=False,
+        )
+
+    @monitor()
+    def _preproc_texts(self, docs: 'DocumentArray'):
+        return preproc_text(docs, device=self._device, return_np=False)
+
     @requests(on='/rank')
     async def rank(self, docs: 'DocumentArray', parameters: Dict, **kwargs):
-
         await self.encode(docs['@r,m'])
 
         set_rank(docs)
@@ -77,12 +89,7 @@ async def encode(self, docs: 'DocumentArray', **kwargs):
             # for image
             if _img_da:
                 for minibatch, _contents in _img_da.map_batch(
-                    partial(
-                        preproc_image,
-                        preprocess_fn=self._preprocess_tensor,
-                        device=self._device,
-                        return_np=False,
-                    ),
+                    self._preproc_images,
                     batch_size=self._minibatch_size,
                     pool=self._pool,
                 ):
@@ -104,7 +111,7 @@ async def encode(self, docs: 'DocumentArray', **kwargs):
             # for text
             if _txt_da:
                 for minibatch, _contents in _txt_da.map_batch(
-                    partial(preproc_text, device=self._device, return_np=False),
+                    self._preproc_texts,
                     batch_size=self._minibatch_size,
                     pool=self._pool,
                 ):

server/clip_server/onnx-flow.yml

@@ -2,11 +2,14 @@ jtype: Flow
 version: '1'
 with:
   port: 51000
+  monitoring: True
+  port_monitoring: 9090
 executors:
   - name: clip_o
     uses:
       jtype: CLIPEncoder
       metas:
         py_modules:
           - executors/clip_onnx.py
-
+    monitoring: true
+    port_monitoring: 9091

server/clip_server/tensorrt-flow.yml

@@ -2,11 +2,14 @@ jtype: Flow
 version: '1'
 with:
   port: 51000
+  monitoring: True
+  port_monitoring: 9090
 executors:
   - name: clip_r
     uses:
       jtype: CLIPEncoder
       metas:
         py_modules:
           - executors/clip_tensorrt.py
-
+    monitoring: true
+    port_monitoring: 9091

server/clip_server/torch-flow.yml

@@ -2,10 +2,14 @@ jtype: Flow
 version: '1'
 with:
   port: 51000
+  monitoring: True
+  port_monitoring: 9090
 executors:
   - name: clip_t
     uses:
       jtype: CLIPEncoder
       metas:
         py_modules:
           - executors/clip_torch.py
+    monitoring: true
+    port_monitoring: 9091

server/setup.py

@@ -42,13 +42,15 @@
     zip_safe=False,
     setup_requires=['setuptools>=18.0', 'wheel'],
     install_requires=[
+        'protobuf>=3.13.0,<=3.20.1',
         'ftfy',
         'torch',
         'regex',
         'torchvision',
-        'jina>=3.3.21',
+        'jina>=3.4.0',
         'docarray>=0.11.0',
         'packaging',
+        'prometheus-client',
     ],
     extras_require={
         'onnx': [