Merge pull request #59 from shanejbrown/main

Fix bug with MP images artifacts and enable parallel MP builds

buildrunner/docker/multiplatform_image_builder.py

@@ -7,7 +7,7 @@
 """
 import logging
 import os
-from multiprocessing import Process
+from multiprocessing import Manager, Process
 from platform import machine, system
 from typing import List
 
@@ -19,14 +19,24 @@
 
 class ImageInfo:
     """Image information repo with associated tags"""
-    def __init__(self, repo: str, tags: List[str] = None):
+    def __init__(self, repo: str, tags: List[str] = None, platform: str = None, digest: str = None,):
+        """
+        Args:
+            repo (str): The repo name for the image.
+            tags (List[str], optional): The tags for the image. Defaults to None.
+            platform (str, optional): The platform for the image. Defaults to None.
+            digest (str, optional): The digest id for the image. Defaults to None.
+        """
         self._repo = repo
 
         if tags is None:
             self._tags = ["latest"]
         else:
             self._tags = tags
 
+        self._platform = platform
+        self._digest = digest
+
     @property
     def repo(self) -> str:
         """The repo name for the image."""
@@ -37,6 +47,22 @@ def tags(self) -> List[str]:
         """The tags for the image."""
         return self._tags
 
+    @property
+    def digest(self) -> str:
+        """The digest id for the image."""
+        return self._digest
+
+    def trunc_digest(self) -> str:
+        """The truncated digest id for the image."""
+        if self._digest is None:
+            return "Digest not available"
+        return self._digest.replace("sha256:", "")[:12]
+
+    @property
+    def platform(self) -> str:
+        """The platform for the image."""
+        return self._platform
+
     def formatted_list(self) -> List[str]:
         """Returns a list of formatted image names"""
         return [f"{self._repo}:{tag}" for tag in self._tags]
@@ -191,7 +217,8 @@ def _build_single_image(self,
                             path: str = ".",
                             file: str = "Dockerfile",
                             tags: List[str] = None,
-                            docker_registry: str = None,) -> None:
+                            docker_registry: str = None,
+                            built_images: list = None) -> None:
         """
         Builds a single image for the given platform
 
@@ -202,9 +229,13 @@ def _build_single_image(self,
             path (str, optional): The path to the Dockerfile. Defaults to ".".
             file (str, optional): The path/name of the Dockerfile (ie. <path>/Dockerfile). Defaults to "Dockerfile".
             tags (List[str], optional): The tags to apply to the image. Defaults to None.
+            docker_registry (str, optional): The docker registry to push the image to. Defaults to None.
+            built_images (list, optional): A list of built images. Defaults to None.
         """
         if tags is None:
             tags = ["latest"]
+        if built_images is None:
+            built_images = []
 
         assert os.path.isdir(path) and os.path.exists(f"{file}"), \
             f"Either path {path}({os.path.isdir(path)}) or file " \
@@ -223,14 +254,21 @@ def _build_single_image(self,
 
         # Check that the images were built and in the registry
         # Docker search is not currently implemented in python-on-wheels
+        image_id = None
         for tag_name in tagged_names:
             try:
                 images = docker.image.pull(tag_name)
                 assert images, f"Failed to build {tag_name}"
+                image_id = docker.image.inspect(tag_name).id
                 docker.image.remove(images, force=True)
             except python_on_whales.exceptions.DockerException as err:
                 LOGGER.error(f"Failed to build {tag_name}: {err}")
                 raise err
+
+        built_images.append(ImageInfo(repo=name,
+                                      tags=tags,
+                                      platform=platform,
+                                      digest=image_id,))
         return image
 
     # pylint: disable=too-many-locals
@@ -241,7 +279,7 @@ def build_multiple_images(self,
                               name: str = None,
                               tags: List[str] = None,
                               push=True,
-                              do_multiprocessing: bool = False,
+                              do_multiprocessing: bool = True,
                               docker_registry: str = None,) -> List[ImageInfo]:
         """
         Builds multiple images for the given platforms. One image will be built for each platform.
@@ -253,7 +291,7 @@ def build_multiple_images(self,
             name (str, optional): The name of the image. Defaults to None.
             tags (List[str], optional): The tags to apply to the image. Defaults to None.
             push (bool, optional): Whether to push the image to the registry. Defaults to True.
-            do_multiprocessing (bool, optional): Whether to use multiprocessing to build the images. Defaults to False.
+            do_multiprocessing (bool, optional): Whether to use multiprocessing to build the images. Defaults to True.
 
         Returns:
             List[ImageInfo]: The list of intermediate built images, these images are ephemeral
@@ -274,18 +312,31 @@ def build_multiple_images(self,
         base_image_name = f"{self._registry_info.ip_addr}:{self._registry_info.port}/{santized_name}"
 
         # Keeps track of the built images {name: [ImageInfo(image_names)]]}
-        self._intermediate_built_images[name] = []
-
+        manager = Manager()
+        self._intermediate_built_images[name] = manager.list()
         processes = []
         for platform in platforms:
             curr_name = f"{base_image_name}-{platform.replace('/', '-')}"
             LOGGER.debug(f"Building {curr_name} for {platform}")
             if do_multiprocessing:
                 processes.append(Process(target=self._build_single_image,
-                                         args=(curr_name, platform, push, path, file, tags, docker_registry)))
+                                         args=(curr_name,
+                                               platform,
+                                               push,
+                                               path,
+                                               file,
+                                               tags,
+                                               docker_registry,
+                                               self._intermediate_built_images[name])))
             else:
-                self._build_single_image(curr_name, platform, push, path, file, tags, docker_registry)
-            self._intermediate_built_images[name].append(ImageInfo(curr_name, tags))
+                self._build_single_image(curr_name,
+                                         platform,
+                                         push,
+                                         path,
+                                         file,
+                                         tags,
+                                         docker_registry,
+                                         self._intermediate_built_images[name])
 
         for proc in processes:
             proc.start()
@@ -425,3 +476,15 @@ def tag_single_platform(self, name: str, tags: List[str] = None,  dest_name: str
             except python_on_whales.exceptions.DockerException as err:
                 LOGGER.error(f"Failed while tagging {dest_name}: {err}")
                 raise err
+
+    def get_built_images(self, name: str) -> List[str]:
+        """
+        Returns the list of tagged images for the given name
+
+        Args:
+            name (str): The name of the image to find
+
+        Returns:
+            List[str]: The list of built images for the given name
+        """
+        return self._intermediate_built_images[name]

buildrunner/steprunner/tasks/push.py

@@ -98,6 +98,22 @@ def run(self, context):  # pylint: disable=too-many-branches
                     tags=repo.tags,
                     dest_name=repo.repository)
 
+                # add image as artifact
+                if not self._commit_only:
+                    images = self.step_runner.multi_platform.get_built_images(self.get_unique_build_name())
+                    image_ids = ",".join([image.trunc_digest() for image in images])
+                    platforms = [image.platform for image in images]
+                    self.step_runner.build_runner.add_artifact(
+                        os.path.join(self.step_runner.name, image_ids.replace(",", "/")),
+                        {
+                            'type': 'docker-image',
+                            'docker:image': image_ids,
+                            'docker:repository': repo.repository,
+                            'docker:tags': repo.tags,
+                            'docker:platforms': platforms,
+                        },
+                    )
+
         # Tag single platform images
         else:
             # first see if a run task produced an image (via a post-build config)

tests/test_multiplatform.py

@@ -1,9 +1,9 @@
 import os
 from typing import List
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
-from python_on_whales import docker
+from python_on_whales import Image, docker
 from python_on_whales.exceptions import DockerException
 
 from buildrunner.docker.multiplatform_image_builder import (
@@ -384,24 +384,35 @@ def test_push_with_dest_names():
      ['test-build-image-2001-linux-amd64', 'test-build-image-2001-linux-arm64']
     )
 ])
-def test_build(name, platforms, expected_image_names):
-    with patch('buildrunner.docker.multiplatform_image_builder.MultiplatformImageBuilder._build_single_image'):
-        test_path = f'{TEST_DIR}/test-files/multiplatform'
-        with MultiplatformImageBuilder() as mp:
-            built_images = mp.build_multiple_images(name=name,
-                                    platforms=platforms,
-                                    path=test_path,
-                                    file=f'{test_path}/Dockerfile',
-                                    do_multiprocessing=False)
+@patch('buildrunner.docker.multiplatform_image_builder.docker.image.remove')
+@patch('buildrunner.docker.multiplatform_image_builder.docker.image.inspect')
+@patch('buildrunner.docker.multiplatform_image_builder.docker.image.pull')
+@patch('buildrunner.docker.multiplatform_image_builder.docker.buildx.build')
+def test_build(mock_build, mock_pull, mock_inspect, mock_remove, name, platforms, expected_image_names):
+    mock_inspect.return_value = MagicMock()
+    mock_inspect.return_value.id = 'myfakeimageid'
+    test_path = f'{TEST_DIR}/test-files/multiplatform'
+    with MultiplatformImageBuilder() as mp:
+        built_images = mp.build_multiple_images(name=name,
+                                platforms=platforms,
+                                path=test_path,
+                                file=f'{test_path}/Dockerfile',
+                                do_multiprocessing=False)
 
-            assert len(built_images) ==  len(platforms)
-            assert len(built_images) ==  len(expected_image_names)
+        assert len(built_images) ==  len(platforms)
+        assert len(built_images) ==  len(expected_image_names)
 
-            missing_images = actual_images_match_expected(built_images, expected_image_names)
-            assert missing_images == [], f'Failed to find {missing_images} in {[image.repo for image in built_images]}'
+        missing_images = actual_images_match_expected(built_images, expected_image_names)
+        assert missing_images == [], f'Failed to find {missing_images} in {[image.repo for image in built_images]}'
 
 
-def test_build_multiple_builds():
+@patch('buildrunner.docker.multiplatform_image_builder.docker.image.remove')
+@patch('buildrunner.docker.multiplatform_image_builder.docker.image.inspect')
+@patch('buildrunner.docker.multiplatform_image_builder.docker.image.pull')
+@patch('buildrunner.docker.multiplatform_image_builder.docker.buildx.build')
+def test_build_multiple_builds(mock_build, mock_pull, mock_inspect, mock_remove):
+    mock_inspect.return_value = MagicMock()
+    mock_inspect.return_value.id = 'myfakeimageid'
     name1 = 'test-build-multi-image-2001'
     platforms1 = ['linux/amd64', 'linux/arm64']
     expected_image_names1 = ['test-build-multi-image-2001-linux-amd64', 'test-build-multi-image-2001-linux-arm64']
@@ -411,33 +422,32 @@ def test_build_multiple_builds():
     expected_image_names2 = ['test-build-multi-image-2002-linux-amd64', 'test-build-multi-image-2002-linux-arm64']
 
     test_path = f'{TEST_DIR}/test-files/multiplatform'
-    with patch('buildrunner.docker.multiplatform_image_builder.MultiplatformImageBuilder._build_single_image'):
-        with MultiplatformImageBuilder() as mp:
-            # Build set 1
-            built_images1 = mp.build_multiple_images(name=name1,
-                                    platforms=platforms1,
-                                    path=test_path,
-                                    file=f'{test_path}/Dockerfile',
-                                    do_multiprocessing=False)
+    with MultiplatformImageBuilder() as mp:
+        # Build set 1
+        built_images1 = mp.build_multiple_images(name=name1,
+                                platforms=platforms1,
+                                path=test_path,
+                                file=f'{test_path}/Dockerfile',
+                                do_multiprocessing=False)
 
-            # Build set 2
-            built_images2 = mp.build_multiple_images(name=name2,
-                                    platforms=platforms2,
-                                    path=test_path,
-                                    file=f'{test_path}/Dockerfile',
-                                    do_multiprocessing=False)
+        # Build set 2
+        built_images2 = mp.build_multiple_images(name=name2,
+                                platforms=platforms2,
+                                path=test_path,
+                                file=f'{test_path}/Dockerfile',
+                                do_multiprocessing=False)
 
-            # Check set 1
-            assert len(built_images1) ==  len(platforms1)
-            assert len(built_images1) ==  len(expected_image_names1)
-            missing_images = actual_images_match_expected(built_images1, expected_image_names1)
-            assert missing_images == [], f'Failed to find {missing_images} in {[image.repo for image in built_images1]}'
-
-            # Check set 2
-            assert len(built_images2) ==  len(platforms2)
-            assert len(built_images2) ==  len(expected_image_names2)
-            missing_images = actual_images_match_expected(built_images2, expected_image_names2)
-            assert missing_images == [], f'Failed to find {missing_images} in {[image.repo for image in built_images2]}'
+        # Check set 1
+        assert len(built_images1) ==  len(platforms1)
+        assert len(built_images1) ==  len(expected_image_names1)
+        missing_images = actual_images_match_expected(built_images1, expected_image_names1)
+        assert missing_images == [], f'Failed to find {missing_images} in {[image.repo for image in built_images1]}'
+
+        # Check set 2
+        assert len(built_images2) ==  len(platforms2)
+        assert len(built_images2) ==  len(expected_image_names2)
+        missing_images = actual_images_match_expected(built_images2, expected_image_names2)
+        assert missing_images == [], f'Failed to find {missing_images} in {[image.repo for image in built_images2]}'
 
 
 @pytest.mark.parametrize("name, tags, platforms, expected_image_names",[