From 149d4ef86ab5b42e7373066ba84ca1e985555011 Mon Sep 17 00:00:00 2001
From: Tully Foote <tullyfoote@intrinsic.ai>
Date: Mon, 3 Feb 2025 14:46:35 -0800
Subject: [PATCH] Improve the robustness of image cleanup (#313)

* improve robustness of image cleanup

* Don't automatically clean up a named image
That makes the name useless

* Remove unused output_callback

As flagged in the review
---
 src/rocker/cli.py  |  4 ++--
 src/rocker/core.py | 18 +++++++++++++++---
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/rocker/cli.py b/src/rocker/cli.py
index e74791b7..70d2cfbc 100644
--- a/src/rocker/cli.py
+++ b/src/rocker/cli.py
@@ -69,13 +69,13 @@ def main():
     exit_code = dig.build(**vars(args))
     if exit_code != 0:
         print("Build failed exiting")
-        if not args_dict['persist_image']:
+        if not (args_dict['persist_image'] or args_dict.get('image_name')):
             dig.clear_image()
         return exit_code
     # Convert command into string
     args.command = ' '.join(args.command)
     result = dig.run(**args_dict)
-    if not args_dict['persist_image']:
+    if not (args_dict['persist_image'] or args_dict.get('image_name')):
         print(f'Clearing Image: {dig.image_id}s\nTo not clean up use --persist-images')
         dig.clear_image()
     return result
diff --git a/src/rocker/core.py b/src/rocker/core.py
index a6a85738..0117613a 100644
--- a/src/rocker/core.py
+++ b/src/rocker/core.py
@@ -252,12 +252,23 @@ def docker_build(docker_client = None, output_callback = None, **kwargs):
         print("no more output and success not detected")
         return None
 
-def docker_remove_image(image_id, docker_client = None, output_callback = None, **kwargs):
+def docker_remove_image(
+        image_id,
+        docker_client = None,
+        fail_on_error = False,
+        force = False,
+        **kwargs):
 
     if not docker_client:
         docker_client = get_docker_client()
 
-    docker_client.remove_image(image_id)
+    try:
+        docker_client.remove_image(image_id, force=force)
+    except docker.errors.APIError as ex:
+        ## removing the image can fail if there's child images
+        if fail_on_error:
+            return False
+    return True
 
 class SIGWINCHPassthrough(object):
     def __init__ (self, process):
@@ -426,7 +437,8 @@ def run(self, command='', **kwargs):
 
     def clear_image(self):
         if self.image_id:
-            docker_remove_image(self.image_id)
+            if not docker_remove_image(self.image_id):
+                print(f'Failed to clear image {self.image_id} it likely has child images.')
             self.image_id = None
             self.built = False