fix: Progress display of 'admin image rescan' command (#182)

* It's a regression due to renaming of bgtask events.
lablup · Oct 28, 2021 · 22d0231 · 22d0231
1 parent f1e697f
commit 22d0231
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/changes/182.fix b/changes/182.fix
@@ -0,0 +1 @@
+Fix a regression of progress display in the `admin image rescan` command
diff --git a/src/ai/backend/client/cli/admin/image.py b/src/ai/backend/client/cli/admin/image.py
@@ -60,23 +60,23 @@ async def rescan_images_impl(registry: str) -> None:
                 print_fail(f"Failed to begin registry scanning: {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)
+            bgtask_id = result['task_id']
+            bgtask = session.BackgroundTask(bgtask_id)
             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':
+                            if ev.event == 'bgtask_updated':
                                 pbar.total = data['total_progress']
                                 pbar.write(data['message'])
                                 pbar.update(data['current_progress'] - pbar.n)
-                            elif ev.event == 'task_failed':
+                            elif ev.event == 'bgtask_failed':
                                 error_msg = data['message']
                                 completion_msg_func = \
                                     lambda: print_fail(f"Error occurred: {error_msg}")
-                            elif ev.event == 'task_cancelled':
+                            elif ev.event == 'bgtask_cancelled':
                                 completion_msg_func = \
                                     lambda: print_warn("Registry scanning has been "
                                                        "cancelled in the middle.")