fix: async progress bar does not display (#779)

jina-ai · Jul 22, 2022 · 7c8285b · 7c8285b
1 parent fa62d8e
commit 7c8285b
Showing 1 changed file with 32 additions and 24 deletions.
diff --git a/client/clip_client/client.py b/client/clip_client/client.py
@@ -284,19 +284,22 @@ async def aencode(self, content, **kwargs):
         )
 
         results = DocumentArray()
-        async for da in self._async_client.post(
-            **self._get_post_payload(content, kwargs)
-        ):
-            if not results:
-                self._pbar.start_task(self._r_task)
-            results.extend(da)
-            self._pbar.update(
-                self._r_task,
-                advance=len(da),
-                total_size=str(
-                    filesize.decimal(int(os.environ.get('JINA_GRPC_RECV_BYTES', '0')))
-                ),
-            )
+        with self._pbar:
+            async for da in self._async_client.post(
+                **self._get_post_payload(content, kwargs)
+            ):
+                if not results:
+                    self._pbar.start_task(self._r_task)
+                results.extend(da)
+                self._pbar.update(
+                    self._r_task,
+                    advance=len(da),
+                    total_size=str(
+                        filesize.decimal(
+                            int(os.environ.get('JINA_GRPC_RECV_BYTES', '0'))
+                        )
+                    ),
+                )
 
         return self._unboxed_result(results)
 
@@ -415,16 +418,21 @@ async def arank(self, docs: Iterable['Document'], **kwargs) -> 'DocumentArray':
             total=len(docs),
         )
         results = DocumentArray()
-        async for da in self._async_client.post(**self._get_rank_payload(docs, kwargs)):
-            if not results:
-                self._pbar.start_task(self._r_task)
-            results.extend(da)
-            self._pbar.update(
-                self._r_task,
-                advance=len(da),
-                total_size=str(
-                    filesize.decimal(int(os.environ.get('JINA_GRPC_RECV_BYTES', '0')))
-                ),
-            )
+        with self._pbar:
+            async for da in self._async_client.post(
+                **self._get_rank_payload(docs, kwargs)
+            ):
+                if not results:
+                    self._pbar.start_task(self._r_task)
+                results.extend(da)
+                self._pbar.update(
+                    self._r_task,
+                    advance=len(da),
+                    total_size=str(
+                        filesize.decimal(
+                            int(os.environ.get('JINA_GRPC_RECV_BYTES', '0'))
+                        )
+                    ),
+                )
 
         return results