refactor(framework) Refactor docker build matrix

Signed-off-by: Robert Steiner <[email protected]>
adap · Sep 24, 2024 · 32e117a · 32e117a
1 parent fc6962d
commit 32e117a
Show file tree

Hide file tree

Showing 5 changed files with 116 additions and 49 deletions.
diff --git a/.github/actions/format-docker-build-args/action.yml b/.github/actions/format-docker-build-args/action.yml
@@ -0,0 +1,34 @@
+name: "format-docker-build-args"
+description: "Format the build args to work with the wretry.action"
+inputs:
+  build-args:
+    description: "List of build arguments."
+outputs:
+  build-args:
+    description: "Build args formatted to work with the wretry.action"
+    value: ${{ steps.build-args.outputs.build-args }}
+runs:
+  using: "composite"
+  steps:
+    - id: build-args
+      shell: python
+      run: |
+        import os
+
+        # Adds two spaces to the line breaks to ensure proper indentation
+        # when passing the multi-line string to the wretry.action.
+        # Without it, the multi-line string is passed like this:
+        #
+        # build-args: |
+        #   ARG1=
+        # ARG2=
+        # ARG3=
+        #
+        # This causes the Docker action to interpret ARG2 and ARG3 as keys instead
+        # of values of the multi-line string.
+        build_args = '''${{ inputs.build-args }}'''.replace("\n", "\n  ")
+
+        with open(os.environ['GITHUB_OUTPUT'], 'a') as fh:
+            print("build-args<<EOF", file=fh)
+            print(build_args, file=fh)
+            print("EOF", file=fh)
diff --git a/.github/workflows/_docker-build.yml b/.github/workflows/_docker-build.yml
@@ -58,24 +58,17 @@ jobs:
           hash = hashlib.sha256('''${{ inputs.namespace-repository }}
           ${{ inputs.file-dir }}
           ${{ inputs.build-args }}'''.encode())
-          # Adds two spaces to the line breaks to ensure proper indentation
-          # when passing the multi-line string to the wretry.action.
-          # Without it, the multi-line string is passed like this:
-          #
-          # build-args: |
-          #   ARG1=
-          # ARG2=
-          # ARG3=
-          #
-          # This causes the Docker action to interpret ARG2 and ARG3 as keys instead
-          # of values of the multi-line string.
-          build_args = '''${{ inputs.build-args }}'''.replace("\n", "\n  ")
 
           with open(os.environ['GITHUB_OUTPUT'], 'a') as fh:
               print(f"id={hash.hexdigest()}", file=fh)
-              print("build-args<<EOF", file=fh)
-              print(build_args, file=fh)
-              print("EOF", file=fh)
+
+      - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
+
+      - name: Format build args
+        uses: ./.github/actions/format-docker-build-args
+        id: build-args
+        with:
+          build-args: ${{ inputs.build-args }}
 
       - name: Extract metadata (tags, labels) for Docker
         id: meta
@@ -105,7 +98,7 @@ jobs:
             context: "{{defaultContext}}:${{ inputs.file-dir }}"
             outputs: type=image,name=${{ inputs.namespace-repository }},push-by-digest=true,name-canonical=true,push=true
             build-args: |
-              ${{ steps.build-id.outputs.build-args }}
+              ${{ steps.build-args.outputs.build-args }}
 
       - name: Export digest
         run: |

diff --git a/.github/workflows/docker-build-main.yml b/.github/workflows/docker-build-main.yml
@@ -15,6 +15,7 @@ jobs:
       pip-version: ${{ steps.versions.outputs.pip-version }}
       setuptools-version: ${{ steps.versions.outputs.setuptools-version }}
       flwr-version-ref: ${{ steps.versions.outputs.flwr-version-ref }}
+      matrix: ${{ steps.versions.outputs.matrix }}
     steps:
       - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
 
@@ -26,15 +27,17 @@ jobs:
           echo "pip-version=${{ steps.bootstrap.outputs.pip-version }}" >> "$GITHUB_OUTPUT"
           echo "setuptools-version=${{ steps.bootstrap.outputs.setuptools-version }}" >> "$GITHUB_OUTPUT"
           echo "flwr-version-ref=git+${{ github.server_url }}/${{ github.repository }}.git@${{ github.sha }}" >> "$GITHUB_OUTPUT"
+          python dev/build-docker-image-matrix.py --flwr-version unstable --simple > matrix.json
+          echo "matrix=$(cat matrix.json)" >> $GITHUB_OUTPUT
 
   build-docker-base-images:
     name: Build base images
     if: github.repository == 'adap/flower'
     uses: ./.github/workflows/_docker-build.yml
     needs: parameters
     with:
-      namespace-repository: flwr/base
-      file-dir: src/docker/base/ubuntu
+      namespace-repository: ${{ fromJson(needs.parameters.outputs.matrix).base.images[0].namespace_repository }}
+      file-dir: ${{ fromJson(needs.parameters.outputs.matrix).base.images[0].file_dir }}
       build-args: |
         PIP_VERSION=${{ needs.parameters.outputs.pip-version }}
         SETUPTOOLS_VERSION=${{ needs.parameters.outputs.setuptools-version }}
@@ -48,19 +51,12 @@ jobs:
     name: Build binary images
     if: github.repository == 'adap/flower'
     uses: ./.github/workflows/_docker-build.yml
-    needs: build-docker-base-images
+    needs: [parameters, build-docker-base-images]
     strategy:
       fail-fast: false
-      matrix:
-        images: [
-          { repository: "flwr/superlink", file_dir: "src/docker/superlink" },
-          { repository: "flwr/supernode", file_dir: "src/docker/supernode" },
-          { repository: "flwr/serverapp", file_dir: "src/docker/serverapp" },
-          { repository: "flwr/superexec", file_dir: "src/docker/superexec" },
-          { repository: "flwr/clientapp", file_dir: "src/docker/clientapp" }
-        ]
+      matrix: ${{ fromJson(needs.parameters.outputs.matrix).binary }}
     with:
-      namespace-repository: ${{ matrix.images.repository }}
+      namespace-repository: ${{ matrix.images.namespace_repository }}
       file-dir: ${{ matrix.images.file_dir }}
       build-args: BASE_IMAGE=unstable
       tags: unstable

diff --git a/.github/workflows/release-nightly.yml b/.github/workflows/release-nightly.yml
@@ -18,6 +18,7 @@ jobs:
       skip: ${{ steps.release.outputs.skip }}
       pip-version: ${{ steps.release.outputs.pip-version }}
       setuptools-version: ${{ steps.release.outputs.setuptools-version }}
+      matrix: ${{ steps.release.outputs.matrix }}
     steps:
       - uses: actions/checkout@v4
       - name: Bootstrap
@@ -37,15 +38,17 @@ jobs:
           echo "version=$(poetry version -s)" >> $GITHUB_OUTPUT
           echo "pip-version=${{ steps.bootstrap.outputs.pip-version }}" >> "$GITHUB_OUTPUT"
           echo "setuptools-version=${{ steps.bootstrap.outputs.setuptools-version }}" >> "$GITHUB_OUTPUT"
+          python dev/build-docker-image-matrix.py --flwr-version nightly --simple > matrix.json
+          echo "matrix=$(cat matrix.json)" >> $GITHUB_OUTPUT
 
   build-docker-base-images:
     name: Build nightly base images
     if: github.repository == 'adap/flower' && needs.release-nightly.outputs.skip != 'true'
     uses: ./.github/workflows/_docker-build.yml
     needs: release-nightly
     with:
-      namespace-repository: flwr/base
-      file-dir: src/docker/base/ubuntu
+      namespace-repository: ${{ fromJson(needs.release-nightly.outputs.matrix).base.images[0].namespace_repository }}
+      file-dir: ${{ fromJson(needs.release-nightly.outputs.matrix).base.images[0].file_dir }}
       build-args: |
         PIP_VERSION=${{ needs.release-nightly.outputs.pip-version }}
         SETUPTOOLS_VERSION=${{ needs.release-nightly.outputs.setuptools-version }}
@@ -65,16 +68,9 @@ jobs:
     needs: [release-nightly, build-docker-base-images]
     strategy:
       fail-fast: false
-      matrix:
-        images: [
-          { repository: "flwr/superlink", file_dir: "src/docker/superlink" },
-          { repository: "flwr/supernode", file_dir: "src/docker/supernode" },
-          { repository: "flwr/serverapp", file_dir: "src/docker/serverapp" },
-          { repository: "flwr/superexec", file_dir: "src/docker/superexec" },
-          { repository: "flwr/clientapp", file_dir: "src/docker/clientapp" }
-        ]
+      matrix: ${{ fromJson(needs.release-nightly.outputs.matrix).binary }}
     with:
-      namespace-repository: ${{ matrix.images.repository }}
+      namespace-repository: ${{ matrix.images.namespace_repository }}
       file-dir: ${{ matrix.images.file_dir }}
       build-args: BASE_IMAGE=${{ needs.release-nightly.outputs.version }}
       tags: |

diff --git a/dev/build-docker-image-matrix.py b/dev/build-docker-image-matrix.py
@@ -29,6 +29,9 @@ class Distro:
 
 DOCKERFILE_ROOT = "src/docker"
 
+UBUNTU_22_04 = Distro(DistroName.UBUNTU, "22.04")
+ALPINE_3_19 = Distro(DistroName.ALPINE, "3.19")
+
 
 @dataclass
 class BaseImage:
@@ -121,26 +124,23 @@ def tag_latest_ubuntu_with_flwr_version(image: BaseImage) -> List[str]:
         return [image.tag]
 
 
-if __name__ == "__main__":
-    arg_parser = argparse.ArgumentParser(
-        description="Generate Github Docker workflow matrix"
-    )
-    arg_parser.add_argument("--flwr-version", type=str, required=True)
-    args = arg_parser.parse_args()
-
-    flwr_version = args.flwr_version
+def generate_complete_matrix(flwr_version: str):
+    """Generates a matrix comprising Ubuntu and Alpine images. For Alpine, the matrix
+    includes only the latest supported Python version, whereas for Ubuntu, it includes all
+    supported Python versions.
+    """
 
     # ubuntu base images for each supported python version
     ubuntu_base_images = generate_base_images(
         flwr_version,
         SUPPORTED_PYTHON_VERSIONS,
-        [Distro(DistroName.UBUNTU, "22.04")],
+        [UBUNTU_22_04],
     )
     # alpine base images for the latest supported python version
     alpine_base_images = generate_base_images(
         flwr_version,
         [LATEST_SUPPORTED_PYTHON_VERSION],
-        [Distro(DistroName.ALPINE, "3.19")],
+        [ALPINE_3_19],
     )
 
     base_images = ubuntu_base_images + alpine_base_images
@@ -154,6 +154,7 @@ def tag_latest_ubuntu_with_flwr_version(image: BaseImage) -> List[str]:
             lambda image: image.python_version == LATEST_SUPPORTED_PYTHON_VERSION,
         )
         # ubuntu images for each supported python version
+        # and alpine image for the latest supported python version
         + generate_binary_images(
             "supernode",
             base_images,
@@ -187,6 +188,53 @@ def tag_latest_ubuntu_with_flwr_version(image: BaseImage) -> List[str]:
         )
     )
 
+    return base_images, binary_images
+
+
+def generate_ubuntu_python_latest_matrix(flwr_version: str):
+    """Generates a matrix comprising Ubuntu images. It includes only the latest supported Python
+    version.
+    """
+
+    # ubuntu base image for the latest supported python version
+    base_images = generate_base_images(
+        flwr_version,
+        [LATEST_SUPPORTED_PYTHON_VERSION],
+        [UBUNTU_22_04],
+    )
+
+    binary_images = (
+        # ubuntu and alpine images for the latest supported python version
+        generate_binary_images("superlink", base_images)
+        # ubuntu images for the latest supported python version
+        + generate_binary_images("supernode", base_images)
+        # ubuntu images for the latest supported python version
+        + generate_binary_images("serverapp", base_images)
+        # ubuntu images for the latest supported python version
+        + generate_binary_images("superexec", base_images)
+        # ubuntu images for the latest supported python version
+        + generate_binary_images("clientapp", base_images)
+    )
+
+    return base_images, binary_images
+
+
+if __name__ == "__main__":
+    arg_parser = argparse.ArgumentParser(
+        description="Generate Github Docker workflow matrix"
+    )
+    arg_parser.add_argument("--flwr-version", type=str, required=True)
+    arg_parser.add_argument("--simple", action="store_true")
+    args = arg_parser.parse_args()
+
+    flwr_version = args.flwr_version
+    simple = args.simple
+
+    if simple:
+        base_images, binary_images = generate_ubuntu_python_latest_matrix(flwr_version)
+    else:
+        base_images, binary_images = generate_complete_matrix(flwr_version)
+
     print(
         json.dumps(
             {