Add background task support and rewrite rescan-images command.

Author: achimnol

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

@@ -1,18 +1,21 @@
+import json
 import sys
+
 import click
 from tabulate import tabulate
 
 from . import admin
+from ...compat import asyncio_run
 from ...session import Session
 from ..pretty import print_done, print_warn, print_fail, print_error
 
 
 @admin.command()
 @click.option('--operation', is_flag=True, help='Get operational images only')
-def images(operation):
-    '''
+def images(operation: bool) -> None:
+    """
     Show the list of registered images in this cluster.
-    '''
+    """
     fields = [
         ('Name', 'name'),
         ('Registry', 'registry'),
@@ -40,25 +43,37 @@ def images(operation):
 @click.option('-r', '--registry', type=str, default=None,
               help='The name (usually hostname or "lablup") '
                    'of the Docker registry configured.')
-def rescan_images(registry):
-    '''Update the kernel image metadata from all configured docker registries.'''
-    with Session() as session:
-        try:
-            result = session.Image.rescan_images(registry)
-        except Exception as e:
-            print_error(e)
-            sys.exit(1)
-        if result['ok']:
-            print_done("Updated the image metadata from the configured registries.")
-        else:
-            print_fail(f"Rescanning has failed: {result['msg']}")
+def rescan_images(registry: str) -> None:
+    """
+    Update the kernel image metadata from all configured docker registries.
+    """
+
+    async def rescan_images_impl(registry: str) -> None:
+        async with Session() as session:
+            try:
+                result = await session.Image.rescan_images(registry)
+            except Exception as e:
+                print_error(e)
+                sys.exit(1)
+            if not result['ok']:
+                print_fail(f"Rescanning has failed: {result['msg']}")
+                sys.exit(1)
+            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.")
+
+    asyncio_run(rescan_images_impl(registry))
 
 
 @admin.command()
 @click.argument('alias', type=str)
 @click.argument('target', type=str)
 def alias_image(alias, target):
-    '''Add an image alias.'''
+    """Add an image alias."""
     with Session() as session:
         try:
             result = session.Image.alias_image(alias, target)
@@ -74,7 +89,7 @@ def alias_image(alias, target):
 @admin.command()
 @click.argument('alias', type=str)
 def dealias_image(alias):
-    '''Remove an image alias.'''
+    """Remove an image alias."""
     with Session() as session:
         try:
             result = session.Image.dealias_image(alias)

src/ai/backend/client/cli/run.py

@@ -1145,9 +1145,9 @@ def events(name, owner_access_key):
     async def _run_events():
         async with AsyncSession() as session:
             compute_session = session.ComputeSession(name, owner_access_key)
-            async with compute_session.stream_events() as sse_response:
-                async for ev in sse_response.fetch_events():
-                    print(click.style(ev['event'], fg='cyan', bold=True), json.loads(ev['data']))
+            async with compute_session.listen_events() as response:
+                async for ev in response:
+                    print(click.style(ev.event, fg='cyan', bold=True), json.loads(ev.data))
 
     try:
         asyncio_run(_run_events())

src/ai/backend/client/compat.py

@@ -1,9 +1,16 @@
-'''
+"""
 A compatibility module for backported codes from Python 3.6+ standard library.
-'''
+"""
 
 import asyncio
 
+__all__ = (
+    'current_loop',
+    'all_tasks',
+    'asyncio_run',
+    'asyncio_run_forever',
+)
+
 
 if hasattr(asyncio, 'get_running_loop'):  # Python 3.7+
     current_loop = asyncio.get_running_loop
@@ -60,11 +67,11 @@ def _asyncio_run(coro, *, debug=False):
 
 
 def asyncio_run_forever(server_context, *, debug=False):
-    '''
+    """
     A proposed-but-not-implemented asyncio.run_forever() API based on
     @vxgmichel's idea.
     See discussions on https://github.com/python/asyncio/pull/465
-    '''
+    """
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
     loop.set_debug(debug)

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

@@ -0,0 +1,38 @@
+from typing import Union
+from uuid import UUID
+
+from ..request import (
+    Request,
+    SSEContextManager,
+)
+
+
+class BackgroundTask:
+    """
+    Provides server-sent events streaming functions.
+    """
+
+    session = None
+    '''The client session instance that this function class is bound to.'''
+
+    task_id: UUID
+
+    def __init__(self, task_id: Union[UUID, str]) -> None:
+        self.task_id = task_id if isinstance(task_id, UUID) else UUID(task_id)
+
+    # only supported in AsyncAPISession
+    def listen_events(self) -> SSEContextManager:
+        """
+        Opens an event stream of the background task updates.
+
+        :returns: a context manager that returns an :class:`SSEResponse` object.
+        """
+        params = {
+            'task_id': str(self.task_id),
+        }
+        request = Request(
+            self.session,
+            'GET', '/events/background-task',
+            params=params,
+        )
+        return request.connect_events()