From 538d4ef2e52de795b90d3f75b2e1e877ab85c18d Mon Sep 17 00:00:00 2001
From: XianBW <36835909+XianBW@users.noreply.github.com>
Date: Mon, 19 Aug 2024 19:10:30 +0800
Subject: [PATCH] feat: add progress status for docker env (#215)

* add progress status for docker env
---
 rdagent/utils/env.py | 47 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 45 insertions(+), 2 deletions(-)

diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py
index c1bc7ba4..196620f3 100644
--- a/rdagent/utils/env.py
+++ b/rdagent/utils/env.py
@@ -6,6 +6,7 @@
 """
 # TODO: move the scenario specific docker env into other folders.
 
+import json
 import os
 import subprocess
 import sys
@@ -19,6 +20,7 @@
 import docker.models.containers
 from pydantic import BaseModel
 from pydantic_settings import BaseSettings
+from rich.progress import Progress, TextColumn
 
 from rdagent.log import rdagent_logger as logger
 
@@ -189,14 +191,55 @@ def prepare(self):
         client = docker.from_env()
         if self.conf.build_from_dockerfile and self.conf.dockerfile_folder_path.exists():
             logger.info(f"Building the image from dockerfile: {self.conf.dockerfile_folder_path}")
-            image, logs = client.images.build(
+            resp_stream = client.api.build(
                 path=str(self.conf.dockerfile_folder_path), tag=self.conf.image, network_mode=self.conf.network
             )
+            if isinstance(resp_stream, str):
+                logger.info(resp_stream)
+            with Progress(TextColumn("{task.description}")) as p:
+                task = p.add_task("[cyan]Building image...")
+                for part in resp_stream:
+                    status_dict = json.loads(part)
+                    if "error" in status_dict:
+                        p.update(task, description=f"[red]error: {status_dict['error']}")
+                        raise docker.errors.BuildError(status_dict["error"])
+                    if "stream" in status_dict:
+                        p.update(task, description=status_dict["stream"])
             logger.info(f"Finished building the image from dockerfile: {self.conf.dockerfile_folder_path}")
         try:
             client.images.get(self.conf.image)
         except docker.errors.ImageNotFound:
-            client.images.pull(self.conf.image)
+            image_pull = client.api.pull(self.conf.image, stream=True, decode=True)
+            current_status = ""
+            layer_set = set()
+            completed_layers = 0
+            with Progress(TextColumn("{task.description}"), TextColumn("{task.fields[progress]}")) as sp:
+                main_task = sp.add_task("[cyan]Pulling image...", progress="")
+                status_task = sp.add_task("[bright_magenta]layer status", progress="")
+                for line in image_pull:
+                    if "error" in line:
+                        sp.update(status_task, description=f"[red]error", progress=line["error"])
+                        raise docker.errors.APIError(line["error"])
+
+                    layer_id = line["id"]
+                    status = line["status"]
+                    p_text = line.get("progress", None)
+
+                    if layer_id not in layer_set:
+                        layer_set.add(layer_id)
+
+                    if p_text:
+                        current_status = p_text
+
+                    if status == "Pull complete" or status == "Already exists":
+                        completed_layers += 1
+
+                    sp.update(main_task, progress=f"[green]{completed_layers}[white]/{len(layer_set)} layers completed")
+                    sp.update(
+                        status_task,
+                        description=f"[bright_magenta]layer {layer_id} [yellow]{status}",
+                        progress=current_status,
+                    )
         except docker.errors.APIError as e:
             raise RuntimeError(f"Error while pulling the image: {e}")