Add push and tags support

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,50 +98,88 @@ 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()
 
         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:

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}"