From 57390eb26be8d39021244248c5d4bd95652215b4 Mon Sep 17 00:00:00 2001
From: tobitege <10787084+tobitege@users.noreply.github.com>
Date: Sat, 14 Sep 2024 06:26:42 +0200
Subject: [PATCH] (enh) docker pull (if not found locally) with progress info
 (#3682)

---
 openhands/core/main.py                   |  2 +-
 openhands/runtime/builder/docker.py      | 47 +++++++++++++++++++++---
 openhands/runtime/utils/runtime_build.py | 17 +++++----
 3 files changed, 52 insertions(+), 14 deletions(-)

diff --git a/openhands/core/main.py b/openhands/core/main.py
index fc59ba89cd..c25ba9a0d8 100644
--- a/openhands/core/main.py
+++ b/openhands/core/main.py
@@ -71,7 +71,7 @@ def create_runtime(
 
     # runtime and tools
     runtime_cls = get_runtime_cls(config.runtime)
-    logger.info(f'Initializing runtime: {runtime_cls}')
+    logger.info(f'Initializing runtime: {runtime_cls.__name__}')
     runtime: Runtime = runtime_cls(
         config=config,
         event_stream=event_stream,
diff --git a/openhands/runtime/builder/docker.py b/openhands/runtime/builder/docker.py
index 0b9f49b0d1..441a253fdd 100644
--- a/openhands/runtime/builder/docker.py
+++ b/openhands/runtime/builder/docker.py
@@ -69,21 +69,56 @@ class DockerRuntimeBuilder(RuntimeBuilder):
             bool: Whether the Docker image exists in the registry or in the local store
         """
         try:
-            logger.info(f'Checking, if image {image_name} exists locally.')
+            logger.info(f'Checking, if image exists locally:\n{image_name}')
             self.docker_client.images.get(image_name)
-            logger.info(f'Image {image_name} found locally.')
+            logger.info('Image found locally.')
             return True
         except docker.errors.ImageNotFound:
             try:
                 logger.info(
                     'Image not found locally. Trying to pull it, please wait...'
                 )
-                self.docker_client.images.pull(image_name)
-                logger.info(f'Image {image_name} pulled successfully.')
+
+                layers = {}
+                for line in self.docker_client.api.pull(
+                    image_name, stream=True, decode=True
+                ):
+                    if 'id' in line and 'progressDetail' in line:
+                        layer_id = line['id']
+                        if layer_id not in layers:
+                            layers[layer_id] = {
+                                'last_logged': -10
+                            }  # Initialize last logged at -10%
+
+                        if (
+                            'total' in line['progressDetail']
+                            and 'current' in line['progressDetail']
+                        ):
+                            total = line['progressDetail']['total']
+                            current = line['progressDetail']['current']
+                            percentage = (current / total) * 100
+
+                            # Log if percentage is at least 10% higher than last logged
+                            if percentage - layers[layer_id]['last_logged'] >= 10:
+                                logger.info(
+                                    f'Layer {layer_id}: {percentage:.0f}% downloaded'
+                                )
+                                layers[layer_id]['last_logged'] = percentage
+
+                    elif 'status' in line:
+                        logger.info(line['status'])
+
+                logger.info('Image pulled')
                 return True
             except docker.errors.ImageNotFound:
                 logger.info('Could not find image locally or in registry.')
                 return False
-            except Exception:
-                logger.info('Could not pull image directly.')
+            except Exception as e:
+                msg = 'Image could not be pulled: '
+                ex_msg = str(e)
+                if 'Not Found' in ex_msg:
+                    msg += 'image not found in registry.'
+                else:
+                    msg += f'{ex_msg}'
+                logger.warning(msg)
                 return False
diff --git a/openhands/runtime/utils/runtime_build.py b/openhands/runtime/utils/runtime_build.py
index 44467c7202..2baa6732d0 100644
--- a/openhands/runtime/utils/runtime_build.py
+++ b/openhands/runtime/utils/runtime_build.py
@@ -39,8 +39,11 @@ def _put_source_code_to_dir(temp_dir: str):
     Parameters:
     - temp_dir (str): The directory to put the source code in
     """
+    if not os.path.isdir(temp_dir):
+        raise RuntimeError(f'Temp directory {temp_dir} does not exist')
+
     project_root = os.path.dirname(os.path.dirname(os.path.abspath(openhands.__file__)))
-    logger.info(f'Using project root: {project_root}')
+    logger.info(f'Building source distribution using project root: {project_root}')
 
     # Fetch the correct version from pyproject.toml
     package_version = _get_package_version()
@@ -63,12 +66,12 @@ def _put_source_code_to_dir(temp_dir: str):
         logger.error(err_logs)
 
     if result.returncode != 0:
-        logger.error(f'Build failed: {result}')
-        raise Exception(f'Build failed: {result}')
+        logger.error(f'Image build failed:\n{result}')
+        raise RuntimeError(f'Image build failed:\n{result}')
 
     if not os.path.exists(tarball_path):
         logger.error(f'Source distribution not found at {tarball_path}')
-        raise Exception(f'Source distribution not found at {tarball_path}')
+        raise RuntimeError(f'Source distribution not found at {tarball_path}')
     logger.info(f'Source distribution created at {tarball_path}')
 
     # Unzip the tarball
@@ -150,14 +153,14 @@ def prep_docker_build_folder(
         file.write(dockerfile_content)
 
     # Get the MD5 hash of the dir_path directory
-    hash = dirhash(dir_path, 'md5')
+    dist_hash = dirhash(dir_path, 'md5')
     logger.info(
         f'Input base image: {base_image}\n'
         f'Skip init: {skip_init}\n'
         f'Extra deps: {extra_deps}\n'
-        f'Hash for docker build directory [{dir_path}] (contents: {os.listdir(dir_path)}): {hash}\n'
+        f'Hash for docker build directory [{dir_path}] (contents: {os.listdir(dir_path)}): {dist_hash}\n'
     )
-    return hash
+    return dist_hash
 
 
 def get_runtime_image_repo_and_tag(base_image: str) -> tuple[str, str]: