gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-07-27 23:47:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add progress status for docker env (#215)

Browse Source 
* add progress status for docker env
This commit is contained in:
XianBW
2024-08-19 19:10:30 +08:00
committed by GitHub
parent 44fd1563a2
commit aa05ed87bf
1 changed files with 45 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rdagent/utils/env.py
+45 -2
View File
@@ -6,6 +6,7 @@ Tries to create uniform environment for the agent to run;
"""
# 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
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 @@ class DockerEnv(Env[DockerConf]):
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}")