From d1b3553c0138c78ff2f0669d31a4b239eeefcaa7 Mon Sep 17 00:00:00 2001 From: Tully Foote Date: Sat, 9 Nov 2024 00:37:09 -0800 Subject: [PATCH] Add clear_image to ImageGenerator to enable non-persistence and not taking up all disk space. Especially for the tests. Also add option to top level for cleaning up the image on the CLI --- src/rocker/cli.py | 8 +++++++- src/rocker/core.py | 11 +++++++++++ test/test_core.py | 8 ++++++++ test/test_extension.py | 2 ++ test/test_nvidia.py | 2 ++ test/test_rmw_extension.py | 1 + 6 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/rocker/cli.py b/src/rocker/cli.py index 2840c8b..9e6743b 100644 --- a/src/rocker/cli.py +++ b/src/rocker/cli.py @@ -35,6 +35,7 @@ def main(): parser.add_argument('--noexecute', action='store_true', help='Deprecated') parser.add_argument('--nocache', action='store_true') parser.add_argument('--nocleanup', action='store_true', help='do not remove the docker container when stopped') + parser.add_argument('--persist-image', action='store_true', help='do not remove the docker image when stopped', default=False) #TODO(tfoote) Add a name to it if persisting parser.add_argument('--pull', action='store_true') parser.add_argument('--version', action='version', version='%(prog)s ' + get_rocker_version()) @@ -68,10 +69,15 @@ def main(): exit_code = dig.build(**vars(args)) if exit_code != 0: print("Build failed exiting") + if not args_dict['persist_image']: + dig.clear_image() return exit_code # Convert command into string args.command = ' '.join(args.command) - return dig.run(**args_dict) + result = dig.run(**args_dict) + if not args_dict['persist_image']: + dig.clear_image() + return result def detect_image_os(): diff --git a/src/rocker/core.py b/src/rocker/core.py index b520f9d..03af51c 100644 --- a/src/rocker/core.py +++ b/src/rocker/core.py @@ -251,6 +251,12 @@ 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): + + if not docker_client: + docker_client = get_docker_client() + + docker_client.remove_image(image_id) class SIGWINCHPassthrough(object): def __init__ (self, process): @@ -413,6 +419,11 @@ def run(self, command='', **kwargs): print("Docker run failed\n", ex) return ex.returncode + def clear_image(self): + if self.image_id: + docker_remove_image(self.image_id) + self.image_id = None + self.built = False def write_files(extensions, args_dict, target_directory): all_files = {} diff --git a/test/test_core.py b/test/test_core.py index a258a5f..85da664 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -68,6 +68,7 @@ def test_run_before_build(self): self.assertEqual(dig.run('true'), 1) self.assertEqual(dig.build(), 0) self.assertEqual(dig.run('true'), 0) + dig.clear_image() @pytest.mark.docker def test_return_code_no_extensions(self): @@ -75,6 +76,7 @@ def test_return_code_no_extensions(self): self.assertEqual(dig.build(), 0) self.assertEqual(dig.run('true'), 0) self.assertEqual(dig.run('false'), 1) + dig.clear_image() @pytest.mark.docker def test_return_code_multiple_extensions(self): @@ -85,12 +87,14 @@ def test_return_code_multiple_extensions(self): self.assertEqual(dig.build(), 0) self.assertEqual(dig.run('true'), 0) self.assertEqual(dig.run('false'), 1) + dig.clear_image() @pytest.mark.docker def test_noexecute(self): dig = DockerImageGenerator([], {}, 'ubuntu:bionic') self.assertEqual(dig.build(), 0) self.assertEqual(dig.run('true', noexecute=True), 0) + dig.clear_image() @pytest.mark.docker def test_dry_run(self): @@ -98,6 +102,7 @@ def test_dry_run(self): self.assertEqual(dig.build(), 0) self.assertEqual(dig.run('true', mode='dry-run'), 0) self.assertEqual(dig.run('false', mode='dry-run'), 0) + dig.clear_image() @pytest.mark.docker def test_non_interactive(self): @@ -105,6 +110,7 @@ def test_non_interactive(self): self.assertEqual(dig.build(), 0) self.assertEqual(dig.run('true', mode='non-interactive'), 0) self.assertEqual(dig.run('false', mode='non-interactive'), 1) + dig.clear_image() @pytest.mark.docker def test_device(self): @@ -112,6 +118,7 @@ def test_device(self): self.assertEqual(dig.build(), 0) self.assertEqual(dig.run('true', devices=['/dev/random']), 0) self.assertEqual(dig.run('true', devices=['/dev/does_not_exist']), 0) + dig.clear_image() @pytest.mark.docker def test_network(self): @@ -120,6 +127,7 @@ def test_network(self): networks = ['bridge', 'host', 'none'] for n in networks: self.assertEqual(dig.run('true', network=n), 0) + dig.clear_image() @pytest.mark.docker def test_extension_manager(self): diff --git a/test/test_extension.py b/test/test_extension.py index 7b94670..d264c35 100644 --- a/test/test_extension.py +++ b/test/test_extension.py @@ -454,6 +454,7 @@ def test_user_collisions(self): self.assertTrue(exit_code == 0, f"Build failed with exit code {exit_code}") run_exit_code = dig.run(**build_args) self.assertTrue(run_exit_code == 0, f"Run failed with exit code {run_exit_code}") + dig.clear_image() # Test colliding UID and name @@ -464,6 +465,7 @@ def test_user_collisions(self): self.assertTrue(exit_code == 0, f"Build failed with exit code {exit_code}") run_exit_code = dig.run(**build_args) self.assertTrue(run_exit_code == 0, f"Run failed with exit code {run_exit_code}") + dig.clear_image() class PulseExtensionTest(unittest.TestCase): diff --git a/test/test_nvidia.py b/test/test_nvidia.py index e37b7dc..6fb7e63 100644 --- a/test/test_nvidia.py +++ b/test/test_nvidia.py @@ -292,6 +292,7 @@ def test_no_cuda(self): dig = DockerImageGenerator([], {}, tag) self.assertEqual(dig.build(), 0) self.assertNotEqual(dig.run(), 0) + dig.clear_image() @pytest.mark.docker def test_cuda_install(self): @@ -302,6 +303,7 @@ def test_cuda_install(self): dig = DockerImageGenerator(active_extensions, {}, tag) self.assertEqual(dig.build(), 0) self.assertEqual(dig.run(), 0) + dig.clear_image() def test_cuda_env_subs(self): plugins = list_plugins() diff --git a/test/test_rmw_extension.py b/test/test_rmw_extension.py index 47bd1bd..5c40a45 100644 --- a/test/test_rmw_extension.py +++ b/test/test_rmw_extension.py @@ -92,3 +92,4 @@ def test_rmw_extension(self): self.assertEqual(dig.build(), 0) self.assertEqual(dig.run(command='dpkg -l ros-rolling-rmw-cyclonedds-cpp'), 0) self.assertIn('-e RMW_IMPLEMENTATION=rmw_cyclonedds_cpp', dig.generate_docker_cmd('', mode='dry-run')) + dig.clear_image()