diff --git a/buildrunner/docker/multiplatform_images.py b/buildrunner/docker/multiplatform_images.py index 509b988f..56019591 100644 --- a/buildrunner/docker/multiplatform_images.py +++ b/buildrunner/docker/multiplatform_images.py @@ -1,9 +1,35 @@ +from ast import In from multiprocessing import Process import os import platform from python_on_whales import docker +class ImageInfo: + def __init__(self, repo:str, tags:list[str]=['latest']): + self._repo = repo + self._tags = tags + + @property + def repo(self) -> str: + return self._repo + + @property + def tags(self) -> list[str]: + return self._tags + + def formatted_list(self) -> list[str]: + return [f"{self._repo}:{tag}" for tag in self._tags] + + def __str__(self): + if len(self._tags) == 1: + return f"{self._repo}:{self._tags[0]}" + return f"{self._repo} tags={self._tags}" + + def __repr__(self): + return self.__str__() + + class MultiplatformImages: """A collection of images that can be built and pushed together.""" @@ -24,6 +50,15 @@ def __del__(self): if self._use_local_registry: self._stop_local_registry() + if self._built_images != {}: + for name, images in self._built_images.items(): + for image in images: + print(image.tags) + for tag in image.tags: + print(f"Removing image {image.repo}:{tag} for {name}") + docker.image.remove(f'{image.repo}:{tag}', force=True) + + @property def registry_ip(self) -> int: return self._reg_ip @@ -63,37 +98,45 @@ def _stop_local_registry(self): docker.stop(self._reg_name) docker.remove(self._reg_name, volumes=True, force=True) - def build_image(self, name:str, platform:str, push:bool, path:str='.', file:str='Dockerfile'): + def build_image(self, name:str, platform:str, push:bool, path:str='.', file:str='Dockerfile', tags:list[str]=['latest'],): push = True print(f'path: {path}, file: {file}') assert os.path.isdir(path) and os.path.isfile(f'{path}/{file}') - image = docker.buildx.build(path, tags=[name], platforms=[platform], push=push, file=file) + tagged_names = [f'{name}:{tag}' for tag in tags] + print(f"Building {tagged_names} for {platform}") + docker.buildx.build(path, tags=tagged_names, platforms=[platform], push=push, file=file) + # TODO also need to do tags other than latest def build(self, platforms:list[str], path:str='.', file:str='Dockerfile', name:str='test', + tags:list[str]=['latest'], do_multiprocessing:bool=False) -> None: """ Builds the images for the given platforms """ - - print(f"Building {name} for platforms {platforms} from {file}") + # TODO handle tags and multiple names. Tags are a list of strings + print(f"Building {name}:{tags} for platforms {platforms} from {file}") push=True base_image_name = f"{self._reg_ip}:{self._reg_port}/{name}" - image_names=[] + + # Keeps track of the built images {name: [ImageInfo(image_names)]]} + self._built_images[name]=[] p = [] for platform in platforms: curr_name = f"{base_image_name}-{platform.replace('/','-')}" print(f"Building {curr_name} for {platform}") if do_multiprocessing: - p.append(Process(target=self.build_image,args=(curr_name, platform, push, path, file))) + p.append(Process(target=self.build_image,args=(curr_name, platform, push, path, file, tags))) else: - self.build_image(name, platform, push, path, file) - image_names.append(curr_name) + # Print args for build_image + print(f"curr_name: {curr_name}, platform: {platform}, push: {push}, path: {path}, file: {file}, tags: {tags}") + self.build_image(curr_name, platform, push, path, file, tags) + self._built_images[name].append(ImageInfo(curr_name, tags)) for proc in p: proc.start() @@ -101,12 +144,42 @@ def build(self, for proc in p: proc.join() - # Keeps track of the images built and store them in a dict or something in self {name: [image_names]]} - self._built_images[name] = image_names - return image_names + return self._built_images[name] + + def push(self, src_key:str, dest_names:str) -> None: + # TODO change dest name to be list of strings + print(f"Creating manifest list {dest_names}") + + initial_timeout_seconds = 60 + timeout_step_seconds = 60 + timeout_max_seconds = 600 + retries = 5 + + src_names = self._built_images[src_key] + + timeout_seconds = initial_timeout_seconds + # TODO handle tags ':' + while retries > 0: + retries -= 1 + print(f"Creating manifest list {dest_names} with timeout {timeout_seconds} seconds") + p = Process(target=docker.buildx.imagetools.create, kwargs={"sources": src_names, "tags": [dest_names],}) + p.start() + p.join(timeout_seconds) + if p.is_alive(): + print(f"Timeout after {timeout_seconds} seconds, killing process") + p.kill() + print(f"Process killed and retries remaining {retries}") + if retries == 0: + raise Exception(f"Timeout after {retries} retries and {timeout_seconds} seconds each try") + else: + # Process finished within timeout + print(f"Process finished") + break + timeout_seconds += timeout_step_seconds + if timeout_seconds > timeout_max_seconds: + timeout_seconds = timeout_max_seconds - def push(self) -> None: - pass + print(f"Finished creating manifest list {dest_names}") def _find_native_platform(self, name:str) -> str: """ @@ -119,13 +192,13 @@ def _find_native_platform(self, name:str) -> str: if name not in self._built_images.keys() or self._built_images[name] == []: return None - match_platform = [image for image in self._built_images[name] if image.endswith(pattern)] + match_platform = [image for image in self._built_images[name] if image.repo.endswith(pattern)] # No matches found, change os if match_platform == []: if host_os == 'Darwin': pattern = f'linux-{host_arch}' - match_platform = [image for image in self._built_images[name] if image.endswith(pattern)] + match_platform = [image for image in self._built_images[name] if image.repo.endswith(pattern)] assert len(match_platform) <= 1, f"Found more than one match for {name} and {pattern}: {match_platform}" @@ -136,12 +209,17 @@ def _find_native_platform(self, name:str) -> str: return match_platform[0] - def load_single_platform(self, name:str) -> None: + def tag_single_platform(self, name:str, tags:list[str]=['latest']) -> None: """ Loads a single platform image into the host registry. If a matching image to the native platform is not found, then the first image is loaded. """ - self._find_native_platform(name) + source_image = self._find_native_platform(name) + + for image in source_image.formatted_list(): + docker.pull(image) + for tag in tags: + docker.tag(image, f'{name}:{tag}') # def tag(self, src_image, dest_tag: str) -> None: diff --git a/tests/test_multiplatform.py b/tests/test_multiplatform.py index 22c54c49..2f30d8de 100644 --- a/tests/test_multiplatform.py +++ b/tests/test_multiplatform.py @@ -6,7 +6,7 @@ import pytest -from buildrunner.docker.multiplatform_images import MultiplatformImages +from buildrunner.docker.multiplatform_images import ImageInfo, MultiplatformImages TEST_DIR = os.path.basename(os.path.dirname(__file__)) @@ -50,31 +50,35 @@ def test_use_local_registry(): 'test-images-2000', 'linux', 'arm64', - {'test-images-2000': ['localhost:32828/test-images-2000-linux-amd64', 'localhost:32828/test-images-2000-linux-arm64']}, - 'localhost:32828/test-images-2000-linux-arm64' + {'test-images-2000': [ImageInfo('localhost:32828/test-images-2000-linux-amd64', ['latest']), + ImageInfo('localhost:32828/test-images-2000-linux-arm64', ['latest'])]}, + 'localhost:32828/test-images-2000-linux-arm64:latest' ), # OS does not match for Darwin change to linux ( 'test-images-2000', 'Darwin', 'arm64', - {'test-images-2000': ['localhost:32828/test-images-2000-linux-amd64', 'localhost:32828/test-images-2000-linux-arm64']}, - 'localhost:32828/test-images-2000-linux-arm64' + {'test-images-2000': [ImageInfo('localhost:32828/test-images-2000-linux-amd64', ['latest']), + ImageInfo('localhost:32828/test-images-2000-linux-arm64', ['latest'])]}, + 'localhost:32828/test-images-2000-linux-arm64:latest' ), # No match found, get the first image ( 'test-images-2000', 'linux', 'arm', - {'test-images-2000': ['localhost:32828/test-images-2000-linux-amd64', 'localhost:32828/test-images-2000-linux-arm64']}, - 'localhost:32828/test-images-2000-linux-amd64' + {'test-images-2000': [ImageInfo('localhost:32828/test-images-2000-linux-amd64', ['0.1.0']), + ImageInfo('localhost:32828/test-images-2000-linux-arm64', ['0.2.0'])]}, + 'localhost:32828/test-images-2000-linux-amd64:0.1.0' ), # Built_images for name does not exist in dictionary ( 'test-images-2001', 'linux', 'arm64', - {'test-images-2000': ['localhost:32828/test-images-2000-linux-amd64', 'localhost:32828/test-images-2000-linux-arm64']}, + {'test-images-2000': [ImageInfo('localhost:32828/test-images-2000-linux-amd64', ['latest']), + ImageInfo('localhost:32828/test-images-2000-linux-arm64', ['latest'])]}, None ), # Built_images for name is empty @@ -96,15 +100,43 @@ def test_find_native_platform(name, mp = MultiplatformImages(use_local_registry=False) mp._built_images = built_images found_platform = mp._find_native_platform(name) - assert found_platform == expected_image + assert str(found_platform) == str(expected_image) -def test_load_single_platform(): - pass +@pytest.mark.parametrize("name, platforms, expected_image_names",[ + ('test-image-2000', + ['linux/arm64'], + ['test-image-2000-linux-arm64'] + )]) +def test_tag_single_platform(name, platforms, expected_image_names): + tag='latest' + mp = MultiplatformImages() + built_images = mp.build(name=name, + platforms=platforms, + path=f'{TEST_DIR}/test-files/multiplatform', + file='Dockerfile', + do_multiprocessing=False) + print(built_images[0]) + mp.tag_single_platform(name) + found_image = docker.image.list(filters={'reference': f'{name}*'}) + assert len(found_image) == 1 + assert f'{name}:{tag}' in found_image[0].repo_tags def test_push(): - pass + registry = 'shanejbrown' + dest_name = {f'{registry}/test-image-2001': ['latest', '0.1.0']} + + build_name = 'test-image-2001' + platforms = ['linux/arm64','linux/amd64'] + mp = MultiplatformImages() + built_images = mp.build(name=build_name, + platforms=platforms, + path=f'{TEST_DIR}/test-files/multiplatform', + file='Dockerfile', + do_multiprocessing=False) + + mp.push(src_key=build_name, dest_names=dest_name) @pytest.mark.parametrize("name, platforms, expected_image_names",[ @@ -131,7 +163,7 @@ def test_build(name, platforms, expected_image_names): for image in expected_image_names: found = False for built_image in built_images: - if built_image.endswith(image): + if built_image.repo.endswith(image): found = True break assert found == True, f"Could not find {image} in {built_images}" @@ -167,7 +199,7 @@ def test_build_multiple_builds(): for image in expected_image_names1: found = False for built_image in built_images1: - if built_image.endswith(image): + if built_image.repo.endswith(image): found = True break assert found == True, f"Could not find {image} in {built_images1}" @@ -178,7 +210,39 @@ def test_build_multiple_builds(): for image in expected_image_names2: found = False for built_image in built_images2: - if built_image.endswith(image): + if built_image.repo.endswith(image): found = True break assert found == True, f"Could not find {image} in {built_images2}" + + +@pytest.mark.parametrize("name, tags, platforms, expected_image_names",[ + ('test-image-2000', + ['latest', '0.1.0'], + ['linux/arm64'], + ['test-image-2000-linux-arm64'] + ), + ('test-image-2001', + ['latest', '0.2.0'], + ['linux/amd64', 'linux/arm64'], + ['test-image-2001-linux-amd64', 'test-image-2001-linux-arm64'] + ) +]) +def test_build_with_tags(name, tags, platforms, expected_image_names): + mp = MultiplatformImages() + built_images = mp.build(name=name, + platforms=platforms, + path=f'{TEST_DIR}/test-files/multiplatform', + file='Dockerfile', + tags=tags, + do_multiprocessing=False) + + assert len(built_images) == len(platforms) + assert len(built_images) == len(expected_image_names) + for image in expected_image_names: + found = False + for built_image in built_images: + if built_image.repo.endswith(image): + found = True + break + assert found == True, f"Could not find {image} in {built_images}" \ No newline at end of file