feat: Implement progress bar for rescanning images using bgtask API

Author: achimnol

src/ai/backend/client/cli/admin/images.py

@@ -3,10 +3,11 @@
 
 import click
 from tabulate import tabulate
+from tqdm import tqdm
 
 from . import admin
 from ...compat import asyncio_run
-from ...session import Session
+from ...session import Session, AsyncSession
 from ..pretty import print_done, print_warn, print_fail, print_error
 
 
@@ -49,7 +50,7 @@ def rescan_images(registry: str) -> None:
     """
 
     async def rescan_images_impl(registry: str) -> None:
-        async with Session() as session:
+        async with AsyncSession() as session:
             try:
                 result = await session.Image.rescan_images(registry)
             except Exception as e:
@@ -61,10 +62,27 @@ async def rescan_images_impl(registry: str) -> None:
             print_done("Started updating the image metadata from the configured registries.")
             task_id = result['task_id']
             bgtask = session.BackgroundTask(task_id)
-            async with bgtask.listen_events() as response:
-                async for ev in response:
-                    print(click.style(ev.event, fg='cyan', bold=True), json.loads(ev.data))
-            print_done("Finished registry scanning.")
+            try:
+                completion_msg_func = lambda: print_done("Finished registry scanning.")
+                with tqdm(unit='image') as pbar:
+                    async with bgtask.listen_events() as response:
+                        async for ev in response:
+                            data = json.loads(ev.data)
+                            if ev.event == 'task_updated':
+                                pbar.total = data['total_progress']
+                                pbar.write(data['message'])
+                                pbar.update(data['current_progress'] - pbar.n)
+                            elif ev.event == 'task_failed':
+                                error_msg = data['message']
+                                completion_msg_func = \
+                                    lambda: print_fail(f"Error occurred: {error_msg}")
+                            elif ev.event == 'task_cancelled':
+                                pbar.write(ev.data)
+                                completion_msg_func = \
+                                    lambda: print_warn("Registry scanning has been "
+                                                       "cancelled in the middle.")
+            finally:
+                completion_msg_func()
 
     asyncio_run(rescan_images_impl(registry))
 

src/ai/backend/client/func/image.py

@@ -9,23 +9,23 @@
 
 
 class Image:
-    '''
+    """
     Provides a shortcut of :func:`Admin.query()
     <ai.backend.client.admin.Admin.query>` that fetches the information about
     available images.
-    '''
+    """
 
     session = None
-    '''The client session instance that this function class is bound to.'''
+    """The client session instance that this function class is bound to."""
 
     @api_function
     @classmethod
     async def list(cls,
                    operation: bool = False,
                    fields: Iterable[str] = None) -> Sequence[dict]:
-        '''
+        """
         Fetches the list of registered images in this cluster.
-        '''
+        """
 
         if fields is None:
             fields = (
@@ -56,7 +56,7 @@ async def list(cls,
     async def rescan_images(cls, registry: str):
         q = 'mutation($registry: String) {' \
             '  rescan_images(registry:$registry) {' \
-            '   ok msg' \
+            '   ok msg task_id' \
             '  }' \
             '}'
         variables = {

src/ai/backend/client/request.py

@@ -651,7 +651,6 @@ async def fetch_events(self) -> AsyncIterator[SSEMessage]:
                 if len(msg_lines) == 0:
                     continue
                 event_type = 'message'
-                event_data = ''
                 event_id = None
                 event_retry = None
                 data_lines = []
@@ -678,6 +677,8 @@ async def fetch_events(self) -> AsyncIterator[SSEMessage]:
                     id=event_id,
                     retry=event_retry,
                 )
+                if event_type == 'server_close':
+                    break
             else:
                 msg_lines.append(line.decode('utf-8'))