Skip to content

Commit

Permalink
Add clear_image to ImageGenerator to enable non-persistence and not t…
Browse files Browse the repository at this point in the history
…aking up all disk space. Especially for the tests.

Also add option to top level for cleaning up the image on the CLI
  • Loading branch information
tfoote committed Nov 9, 2024
1 parent 759bbf7 commit d1b3553
Show file tree
Hide file tree
Showing 6 changed files with 31 additions and 1 deletion.
8 changes: 7 additions & 1 deletion src/rocker/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down Expand Up @@ -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():
Expand Down
11 changes: 11 additions & 0 deletions src/rocker/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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 = {}
Expand Down
8 changes: 8 additions & 0 deletions test/test_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,13 +68,15 @@ 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):
dig = DockerImageGenerator([], {}, 'ubuntu:bionic')
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):
Expand All @@ -85,33 +87,38 @@ 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):
dig = DockerImageGenerator([], {}, 'ubuntu:bionic')
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):
dig = DockerImageGenerator([], {}, 'ubuntu:bionic')
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):
dig = DockerImageGenerator([], {}, 'ubuntu:bionic')
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):
Expand All @@ -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):
Expand Down
2 changes: 2 additions & 0 deletions test/test_extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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):
Expand Down
2 changes: 2 additions & 0 deletions test/test_nvidia.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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()
Expand Down
1 change: 1 addition & 0 deletions test/test_rmw_extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

0 comments on commit d1b3553

Please sign in to comment.