From daeff3dfafb94e534090e3474b9cb5c775d46d02 Mon Sep 17 00:00:00 2001
From: tobitege <10787084+tobitege@users.noreply.github.com>
Date: Thu, 29 Aug 2024 00:17:58 +0200
Subject: [PATCH] startup handling and logging of docker images tweaked (#3645)

---
 openhands/runtime/builder/docker.py      | 33 ++++++++++++++----------
 openhands/runtime/utils/runtime_build.py | 18 +++++++------
 2 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/openhands/runtime/builder/docker.py b/openhands/runtime/builder/docker.py
index b1f76d854..0b9f49b0d 100644
--- a/openhands/runtime/builder/docker.py
+++ b/openhands/runtime/builder/docker.py
@@ -61,22 +61,29 @@ class DockerRuntimeBuilder(RuntimeBuilder):
         return target_image_hash_name
 
     def image_exists(self, image_name: str) -> bool:
-        """Check if the image exists in the registry (try to pull it first) AND in the local store.
+        """Check if the image exists in the registry (try to pull it first) or in the local store.
 
         Args:
             image_name (str): The Docker image to check (<image repo>:<image tag>)
         Returns:
-            bool: Whether the Docker image exists in the registry and in the local store
+            bool: Whether the Docker image exists in the registry or in the local store
         """
-        # Try to pull the Docker image from the registry
         try:
-            self.docker_client.images.pull(image_name)
-        except Exception:
-            logger.info(f'Cannot pull image {image_name} directly')
-
-        images = self.docker_client.images.list()
-        if images:
-            for image in images:
-                if image_name in image.tags:
-                    return True
-        return False
+            logger.info(f'Checking, if image {image_name} exists locally.')
+            self.docker_client.images.get(image_name)
+            logger.info(f'Image {image_name} 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.')
+                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.')
+                return False
diff --git a/openhands/runtime/utils/runtime_build.py b/openhands/runtime/utils/runtime_build.py
index 8114c4f18..6f6fa5ba3 100644
--- a/openhands/runtime/utils/runtime_build.py
+++ b/openhands/runtime/utils/runtime_build.py
@@ -76,23 +76,25 @@ def _put_source_code_to_dir(temp_dir: str):
     Parameters:
     - temp_dir (str): The directory to put the source code in
     """
+    project_tar = 'project.tar.gz'
+    project_path = os.path.join(temp_dir, project_tar)
+    logger.info('Building source distribution...')
+
     # Build the project source tarball
     tarball_path = _create_project_source_dist()
     filename = os.path.basename(tarball_path)
     filename = filename.removesuffix('.tar.gz')
 
     # Move the project tarball to temp_dir
-    _res = shutil.copy(tarball_path, os.path.join(temp_dir, 'project.tar.gz'))
+    _res = shutil.copy(tarball_path, project_path)
     if _res:
         os.remove(tarball_path)
-    logger.info(
-        f'Source distribution moved to {os.path.join(temp_dir, "project.tar.gz")}'
-    )
+    logger.info('Source distribution moved to ' + project_path)
 
     # Unzip the tarball
-    shutil.unpack_archive(os.path.join(temp_dir, 'project.tar.gz'), temp_dir)
+    shutil.unpack_archive(project_path, temp_dir)
     # Remove the tarball
-    os.remove(os.path.join(temp_dir, 'project.tar.gz'))
+    os.remove(project_path)
     # Rename the directory containing the code to 'code'
     os.rename(os.path.join(temp_dir, filename), os.path.join(temp_dir, 'code'))
     logger.info(f'Unpacked source code directory: {os.path.join(temp_dir, "code")}')
@@ -260,9 +262,9 @@ def build_runtime_image(
     # Scenario 2: If a Docker image with the exact hash is not found, we will FIRST try to re-build it
     # by leveraging the `generic_runtime_image_name` to save some time
     # from re-building the dependencies (e.g., poetry install, apt install)
-    elif runtime_builder.image_exists(generic_runtime_image_name) and not force_rebuild:
+    if not force_rebuild and runtime_builder.image_exists(generic_runtime_image_name):
         logger.info(
-            f'Cannot find docker Image [{hash_runtime_image_name}]\n'
+            f'Could not find docker image [{hash_runtime_image_name}]\n'
             f'Will try to re-build it from latest [{generic_runtime_image_name}] image to potentially save '
             f'time for dependencies installation.\n'
         )