feat: add monitoring

jina-ai · Apr 5, 2022 · 10e8ac4 · 10e8ac4
1 parent 25e1f98
commit 10e8ac4
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 24 deletions.
diff --git a/server/clip_server/executors/clip_torch.py b/server/clip_server/executors/clip_torch.py
@@ -1,10 +1,12 @@
+import contextlib
 import io
 import os
 from typing import TYPE_CHECKING, Optional, List
 
 import torch
 from PIL import Image
 from jina import Executor, requests
+from prometheus_client import Summary
 
 from clip_server.model import clip
 
@@ -27,36 +29,55 @@ def __init__(
             self._device = device
         self._model, self._preprocess = clip.load(name, device=self._device, jit=jit)
 
+        self.summary_text = self.get_summary('text_preproc_second', 'Time to preprocess text')
+        self.summary_image = self.get_summary('image_preproc_second', 'Time to preprocess image')
+        self.summary_encode = self.get_summary('encode_second', 'Time to encode')
+
+    def get_summary(self, title, details):
+        return (
+            Summary(
+                title,
+                details,
+                registry=self.runtime_args.metrics_registry,
+            ).time()
+            if self.runtime_args.metrics_registry
+            else contextlib.nullcontext
+        )
+
     def _preproc_image(self, d: 'Document'):
-        d.tensor = self._preprocess(Image.open(io.BytesIO(d.blob))).to(self._device)
-        return d
+        with self.summary_image:
+            d.tensor = self._preprocess(Image.open(io.BytesIO(d.blob))).to(self._device)
+            return d
 
     def _preproc_text(self, da: 'DocumentArray') -> List[str]:
-        texts = da.texts
-        da.tensors = clip.tokenize(da.texts).to(self._device)
-        da[:, 'mime_type'] = 'text'
-        return texts
+        with self.summary_text:
+            texts = da.texts
+            da.tensors = clip.tokenize(da.texts).to(self._device)
+            da[:, 'mime_type'] = 'text'
+            return texts
 
     @requests
     async def encode(self, docs: 'DocumentArray', **kwargs):
         _img_da = docs.find({'blob': {'$exists': True}})
         _txt_da = docs.find({'text': {'$exists': True}})
 
-        with torch.inference_mode():
-            # for image
-            if _img_da:
-                _img_da.apply(self._preproc_image)
-                _img_da.embeddings = (
-                    self._model.encode_image(_img_da.tensors).cpu().numpy()
-                )
-
-            # for text
-            if _txt_da:
-                texts = self._preproc_text(_txt_da)
-                _txt_da.embeddings = (
-                    self._model.encode_text(_txt_da.tensors).cpu().numpy()
-                )
-                _txt_da.texts = texts
-
-        # drop tensors
-        docs.tensors = None
+        with self.summary_encode:
+
+            with torch.inference_mode():
+                # for image
+                if _img_da:
+                    _img_da.apply(self._preproc_image)
+                    _img_da.embeddings = (
+                        self._model.encode_image(_img_da.tensors).cpu().numpy()
+                    )
+
+                # for text
+                if _txt_da:
+                    texts = self._preproc_text(_txt_da)
+                    _txt_da.embeddings = (
+                        self._model.encode_text(_txt_da.tensors).cpu().numpy()
+                    )
+                    _txt_da.texts = texts
+
+            # drop tensors
+            docs.tensors = None
diff --git a/server/clip_server/torch-flow.yml b/server/clip_server/torch-flow.yml
@@ -2,10 +2,14 @@ jtype: Flow
 version: '1'
 with:
   port: 51000
+  monitoring: true
+  port_monitoring: 8000
 executors:
   - name: clip_t
     uses:
       jtype: CLIPEncoder
       metas:
         py_modules:
           - executors/clip_torch.py
+    monitoring: true
+    port_monitoring: 9000