From 1308b8d7494ce3c8765cf2090322d6dd92b4ae29 Mon Sep 17 00:00:00 2001 From: Robert Steiner Date: Wed, 13 Nov 2024 20:09:37 +0100 Subject: [PATCH] refactor(framework:skip) Refactor docker build matrix (#4224) Signed-off-by: Robert Steiner --- .github/workflows/docker-build-main.yml | 33 +- .github/workflows/framework-release.yml | 8 +- .github/workflows/release-nightly.yml | 40 +- dev/build-docker-image-matrix.py | 466 ++++++++++++++++-------- 4 files changed, 357 insertions(+), 190 deletions(-) diff --git a/.github/workflows/docker-build-main.yml b/.github/workflows/docker-build-main.yml index e54257048245..38a9cd56942b 100644 --- a/.github/workflows/docker-build-main.yml +++ b/.github/workflows/docker-build-main.yml @@ -14,7 +14,7 @@ jobs: outputs: 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 @@ -25,21 +25,26 @@ jobs: run: | 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" + FLWR_VERSION_REF="git+${{ github.server_url }}/${{ github.repository }}.git@${{ github.sha }}" + python dev/build-docker-image-matrix.py --flwr-version "${FLWR_VERSION_REF}" --matrix unstable > 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 + strategy: + fail-fast: false + matrix: ${{ fromJson(needs.parameters.outputs.matrix).base }} with: - namespace-repository: flwr/base - file-dir: src/docker/base/ubuntu + namespace-repository: ${{ matrix.images.namespace_repository }} + file-dir: ${{ matrix.images.file_dir }} build-args: | PIP_VERSION=${{ needs.parameters.outputs.pip-version }} SETUPTOOLS_VERSION=${{ needs.parameters.outputs.setuptools-version }} - FLWR_VERSION_REF=${{ needs.parameters.outputs.flwr-version-ref }} - tags: unstable + ${{ matrix.images.build_args_encoded }} + tags: ${{ matrix.images.tags_encoded }} secrets: dockerhub-user: ${{ secrets.DOCKERHUB_USERNAME }} dockerhub-token: ${{ secrets.DOCKERHUB_TOKEN }} @@ -48,21 +53,15 @@ 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/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 + build-args: BASE_IMAGE=${{ matrix.images.base_image }} + tags: ${{ matrix.images.tags_encoded }} secrets: dockerhub-user: ${{ secrets.DOCKERHUB_USERNAME }} dockerhub-token: ${{ secrets.DOCKERHUB_TOKEN }} diff --git a/.github/workflows/framework-release.yml b/.github/workflows/framework-release.yml index 6e9225c4ad16..6af0c281882b 100644 --- a/.github/workflows/framework-release.yml +++ b/.github/workflows/framework-release.yml @@ -71,7 +71,7 @@ jobs: - id: matrix run: | - python dev/build-docker-image-matrix.py --flwr-version "${{ needs.publish.outputs.flwr-version }}" > matrix.json + python dev/build-docker-image-matrix.py --flwr-version "${{ needs.publish.outputs.flwr-version }}" --matrix stable > matrix.json echo "matrix=$(cat matrix.json)" >> $GITHUB_OUTPUT build-base-images: @@ -88,8 +88,8 @@ jobs: build-args: | PIP_VERSION=${{ needs.parameters.outputs.pip-version }} SETUPTOOLS_VERSION=${{ needs.parameters.outputs.setuptools-version }} - ${{ matrix.images.build_args }} - tags: ${{ matrix.images.tag }} + ${{ matrix.images.build_args_encoded }} + tags: ${{ matrix.images.tags_encoded }} secrets: dockerhub-user: ${{ secrets.DOCKERHUB_USERNAME }} dockerhub-token: ${{ secrets.DOCKERHUB_TOKEN }} @@ -106,7 +106,7 @@ jobs: namespace-repository: ${{ matrix.images.namespace_repository }} file-dir: ${{ matrix.images.file_dir }} build-args: BASE_IMAGE=${{ matrix.images.base_image }} - tags: ${{ matrix.images.tags }} + tags: ${{ matrix.images.tags_encoded }} secrets: dockerhub-user: ${{ secrets.DOCKERHUB_USERNAME }} dockerhub-token: ${{ secrets.DOCKERHUB_TOKEN }} diff --git a/.github/workflows/release-nightly.yml b/.github/workflows/release-nightly.yml index 32f76cc86c5b..d1de7bed531e 100644 --- a/.github/workflows/release-nightly.yml +++ b/.github/workflows/release-nightly.yml @@ -13,11 +13,10 @@ jobs: name: Relase nightly on PyPI if: github.repository == 'adap/flower' outputs: - name: ${{ steps.release.outputs.name }} - version: ${{ steps.release.outputs.version }} 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 @@ -33,27 +32,30 @@ jobs: echo "skip=true" >> $GITHUB_OUTPUT fi - echo "name=$(poetry version | awk {'print $1'})" >> $GITHUB_OUTPUT - 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" + NAME=$(poetry version | awk {'print $1'}) + VERSION=$(poetry version -s) + python dev/build-docker-image-matrix.py --flwr-version "${VERSION}" --matrix nightly --flwr-package "${NAME}" > 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 + strategy: + fail-fast: false + matrix: ${{ fromJson(needs.release-nightly.outputs.matrix).base }} with: - namespace-repository: flwr/base - file-dir: src/docker/base/ubuntu + namespace-repository: ${{ matrix.images.namespace_repository }} + file-dir: ${{ matrix.images.file_dir }} build-args: | PIP_VERSION=${{ needs.release-nightly.outputs.pip-version }} SETUPTOOLS_VERSION=${{ needs.release-nightly.outputs.setuptools-version }} - FLWR_VERSION=${{ needs.release-nightly.outputs.version }} - FLWR_PACKAGE=${{ needs.release-nightly.outputs.name }} - tags: | - ${{ needs.release-nightly.outputs.version }} - nightly + ${{ matrix.images.build_args_encoded }} + tags: ${{ matrix.images.tags_encoded }} secrets: dockerhub-user: ${{ secrets.DOCKERHUB_USERNAME }} dockerhub-token: ${{ secrets.DOCKERHUB_TOKEN }} @@ -65,20 +67,12 @@ 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/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: | - ${{ needs.release-nightly.outputs.version }} - nightly + build-args: BASE_IMAGE=${{ matrix.images.base_image }} + tags: ${{ matrix.images.tags_encoded }} secrets: dockerhub-user: ${{ secrets.DOCKERHUB_USERNAME }} dockerhub-token: ${{ secrets.DOCKERHUB_TOKEN }} diff --git a/dev/build-docker-image-matrix.py b/dev/build-docker-image-matrix.py index 6368166b7e51..9d255ac5471f 100644 --- a/dev/build-docker-image-matrix.py +++ b/dev/build-docker-image-matrix.py @@ -18,12 +18,19 @@ - **Unstable**: Triggered on main branch commits. Builds simplified matrix (latest Python, Ubuntu only). """ +import sys import argparse import json -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field from enum import Enum from typing import Any, Callable, Dict, List, Optional +# when we switch to Python 3.11 in the ci, we need to change the DistroName to: +# class DistroName(StrEnum): +# ALPINE = "alpine" +# UBUNTU = "ubuntu" +assert sys.version_info < (3, 11), "Script requires Python 3.9 or lower." + class DistroName(str, Enum): ALPINE = "alpine" @@ -49,146 +56,89 @@ class Distro: @dataclass class Variant: distro: Distro - file_dir_fn: Callable[[Distro, dict], str] - tag_fn: Callable[[Distro, str, str, dict], str] - build_args_fn: Callable[[Distro, str, str, dict], str] extras: Optional[Any] = None - def file_dir(self) -> str: - return self.file_dir_fn(self.distro, self.extras) - - def tag(self, flwr_version: str, python_version: str) -> str: - return self.tag_fn(self.distro, flwr_version, python_version, self.extras) - - def build_args(self, flwr_version: str, python_version: str) -> str: - return self.build_args_fn( - self.distro, flwr_version, python_version, self.extras - ) - @dataclass class CpuVariant: pass -CPU_BUILD_ARGS = """PYTHON_VERSION={python_version} -FLWR_VERSION={flwr_version} -DISTRO={distro_name} -DISTRO_VERSION={distro_version} -""" - - @dataclass class CudaVariant: version: str -CUDA_BUILD_ARGS = ( - CPU_BUILD_ARGS - + """CUDA_VERSION={cuda_version} -""" -) - - -LATEST_SUPPORTED_CUDA_VERSIONS = "12.4.1" CUDA_VERSIONS_CONFIG = [ ("11.2.2", "20.04"), ("11.8.0", "22.04"), ("12.1.0", "22.04"), ("12.3.2", "22.04"), - (LATEST_SUPPORTED_CUDA_VERSIONS, "22.04"), ] +LATEST_SUPPORTED_CUDA_VERSION = Variant( + Distro(DistroName.UBUNTU, "22.04"), + CudaVariant(version="12.4.1"), +) +# ubuntu base image +UBUNTU_VARIANT = Variant( + Distro(DistroName.UBUNTU, "24.04"), + CpuVariant(), +) -def remove_patch_version(version: str) -> str: - return ".".join(version.split(".")[0:2]) - - -# ubuntu base images for each supported python version -UBUNTU_VARIANTS = [ - Variant( - Distro(DistroName.UBUNTU, "24.04"), - lambda distro, _: f"{DOCKERFILE_ROOT}/base/{distro.name.value}", - lambda distro, flwr_version, python_version, _: f"{flwr_version}-py{python_version}-{distro.name.value}{distro.version}", - lambda distro, flwr_version, python_version, _: CPU_BUILD_ARGS.format( - python_version=python_version, - flwr_version=flwr_version, - distro_name=distro.name, - distro_version=distro.version, - ), - CpuVariant(), - ) -] +# alpine base image +ALPINE_VARIANT = Variant( + Distro(DistroName.ALPINE, "3.19"), + CpuVariant(), +) -# alpine base images for the latest supported python version -ALPINE_VARIANTS = [ - Variant( - Distro(DistroName.ALPINE, "3.19"), - lambda distro, _: f"{DOCKERFILE_ROOT}/base/{distro.name.value}", - lambda distro, flwr_version, python_version, _: f"{flwr_version}-py{python_version}-{distro.name.value}{distro.version}", - lambda distro, flwr_version, python_version, _: CPU_BUILD_ARGS.format( - python_version=python_version, - flwr_version=flwr_version, - distro_name=distro.name, - distro_version=distro.version, - ), - CpuVariant(), - ) -] -# ubuntu cuda base images for each supported python and cuda version +# ubuntu cuda base images CUDA_VARIANTS = [ Variant( Distro(DistroName.UBUNTU, ubuntu_version), - lambda distro, _: f"{DOCKERFILE_ROOT}/base/{distro.name.value}-cuda", - lambda distro, flwr_version, python_version, extras: f"{flwr_version}-py{python_version}-cu{remove_patch_version(extras.version)}-{distro.name.value}{distro.version}", - lambda distro, flwr_version, python_version, extras: CUDA_BUILD_ARGS.format( - python_version=python_version, - flwr_version=flwr_version, - distro_name=distro.name, - distro_version=distro.version, - cuda_version=extras.version, - ), CudaVariant(version=cuda_version), ) for (cuda_version, ubuntu_version) in CUDA_VERSIONS_CONFIG -] +] + [LATEST_SUPPORTED_CUDA_VERSION] + + +def remove_patch_version(version: str) -> str: + return ".".join(version.split(".")[0:2]) @dataclass -class BaseImage: - variant: Variant - python_version: str - namespace_repository: str - file_dir: str - tag: str - flwr_version: str - build_args: str +class BaseImageBuilder: + file_dir_fn: Callable[[Any], str] + tags_fn: Callable[[Any], list[str]] + build_args_fn: Callable[[Any], str] + build_args: Any + tags: list[str] = field(init=False) + file_dir: str = field(init=False) + tags_encoded: str = field(init=False) + build_args_encoded: str = field(init=False) -def new_base_image( - flwr_version: str, python_version: str, variant: Variant -) -> Dict[str, Any]: - return BaseImage( - variant, - python_version, - "flwr/base", - variant.file_dir(), - variant.tag(flwr_version, python_version), - flwr_version, - variant.build_args(flwr_version, python_version), - ) +@dataclass +class BaseImage(BaseImageBuilder): + namespace_repository: str = "flwr/base" + @property + def file_dir(self) -> str: + return self.file_dir_fn(self.build_args) -def generate_base_images( - flwr_version: str, python_versions: List[str], variants: List[Variant] -) -> List[Dict[str, Any]]: - return [ - new_base_image(flwr_version, python_version, variant) - for variant in variants - for python_version in python_versions - ] + @property + def tags(self) -> str: + return self.tags_fn(self.build_args) + + @property + def tags_encoded(self) -> str: + return "\n".join(self.tags) + + @property + def build_args_encoded(self) -> str: + return self.build_args_fn(self.build_args) @dataclass @@ -196,7 +146,7 @@ class BinaryImage: namespace_repository: str file_dir: str base_image: str - tags: List[str] + tags_encoded: str def new_binary_image( @@ -211,7 +161,7 @@ def new_binary_image( return BinaryImage( f"flwr/{name}", f"{DOCKERFILE_ROOT}/{name}", - base_image.tag, + base_image.tags[0], "\n".join(tags), ) @@ -231,53 +181,95 @@ def generate_binary_images( def tag_latest_alpine_with_flwr_version(image: BaseImage) -> List[str]: if ( - image.variant.distro.name == DistroName.ALPINE - and image.python_version == LATEST_SUPPORTED_PYTHON_VERSION + image.build_args.variant.distro.name == DistroName.ALPINE + and image.build_args.python_version == LATEST_SUPPORTED_PYTHON_VERSION ): - return [image.tag, image.flwr_version] + return image.tags + [image.build_args.flwr_version] else: - return [image.tag] + return image.tags def tag_latest_ubuntu_with_flwr_version(image: BaseImage) -> List[str]: if ( - image.variant.distro.name == DistroName.UBUNTU - and image.python_version == LATEST_SUPPORTED_PYTHON_VERSION - and isinstance(image.variant.extras, CpuVariant) + image.build_args.variant.distro.name == DistroName.UBUNTU + and image.build_args.python_version == LATEST_SUPPORTED_PYTHON_VERSION + and isinstance(image.build_args.variant.extras, CpuVariant) ): - return [image.tag, image.flwr_version] + return image.tags + [image.build_args.flwr_version] else: - return [image.tag] + return image.tags -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() +# +# Build matrix for stable releases +# +def build_stable_matrix(flwr_version: str) -> List[BaseImage]: + @dataclass + class StableBaseImageBuildArgs: + variant: Variant + python_version: str + flwr_version: str - flwr_version = args.flwr_version + cpu_build_args = """PYTHON_VERSION={python_version} +FLWR_VERSION={flwr_version} +DISTRO={distro_name} +DISTRO_VERSION={distro_version} +""" - ubuntu_base_images = generate_base_images( - flwr_version, - SUPPORTED_PYTHON_VERSIONS, - UBUNTU_VARIANTS, - ) + cpu_build_args_variants = [ + StableBaseImageBuildArgs(UBUNTU_VARIANT, python_version, flwr_version) + for python_version in SUPPORTED_PYTHON_VERSIONS + ] + [ + StableBaseImageBuildArgs( + ALPINE_VARIANT, LATEST_SUPPORTED_PYTHON_VERSION, flwr_version + ) + ] - alpine_base_images = generate_base_images( - flwr_version, - [LATEST_SUPPORTED_PYTHON_VERSION], - ALPINE_VARIANTS, - ) + cpu_base_images = [ + BaseImage( + file_dir_fn=lambda args: f"{DOCKERFILE_ROOT}/base/{args.variant.distro.name.value}", + tags_fn=lambda args: [ + f"{args.flwr_version}-py{args.python_version}-{args.variant.distro.name.value}{args.variant.distro.version}" + ], + build_args_fn=lambda args: cpu_build_args.format( + python_version=args.python_version, + flwr_version=args.flwr_version, + distro_name=args.variant.distro.name, + distro_version=args.variant.distro.version, + ), + build_args=build_args_variant, + ) + for build_args_variant in cpu_build_args_variants + ] - cuda_base_images = generate_base_images( - flwr_version, - SUPPORTED_PYTHON_VERSIONS, - CUDA_VARIANTS, - ) + cuda_build_args_variants = [ + StableBaseImageBuildArgs(variant, python_version, flwr_version) + for variant in CUDA_VARIANTS + for python_version in SUPPORTED_PYTHON_VERSIONS + ] + + cuda_build_args = cpu_build_args + """CUDA_VERSION={cuda_version}""" + + cuda_base_image = [ + BaseImage( + file_dir_fn=lambda args: f"{DOCKERFILE_ROOT}/base/{args.variant.distro.name.value}-cuda", + tags_fn=lambda args: [ + f"{args.flwr_version}-py{args.python_version}-cu{remove_patch_version(args.variant.extras.version)}-{args.variant.distro.name.value}{args.variant.distro.version}", + ], + build_args_fn=lambda args: cuda_build_args.format( + python_version=args.python_version, + flwr_version=args.flwr_version, + distro_name=args.variant.distro.name, + distro_version=args.variant.distro.version, + cuda_version=args.variant.extras.version, + ), + build_args=build_args_variant, + ) + for build_args_variant in cuda_build_args_variants + ] - base_images = ubuntu_base_images + alpine_base_images + # base_images = cpu_base_images + cuda_base_image + base_images = cpu_base_images binary_images = ( # ubuntu and alpine images for the latest supported python version @@ -285,8 +277,9 @@ def tag_latest_ubuntu_with_flwr_version(image: BaseImage) -> List[str]: "superlink", base_images, tag_latest_alpine_with_flwr_version, - lambda image: image.python_version == LATEST_SUPPORTED_PYTHON_VERSION - and isinstance(image.variant.extras, CpuVariant), + lambda image: image.build_args.python_version + == LATEST_SUPPORTED_PYTHON_VERSION + and isinstance(image.build_args.variant.extras, CpuVariant), ) # ubuntu images for each supported python version + generate_binary_images( @@ -294,12 +287,12 @@ def tag_latest_ubuntu_with_flwr_version(image: BaseImage) -> List[str]: base_images, tag_latest_alpine_with_flwr_version, lambda image: ( - image.variant.distro.name == DistroName.UBUNTU - and isinstance(image.variant.extras, CpuVariant), + image.build_args.variant.distro.name == DistroName.UBUNTU + and isinstance(image.build_args.variant.extras, CpuVariant) ) or ( - image.variant.distro.name == DistroName.ALPINE - and image.python_version == LATEST_SUPPORTED_PYTHON_VERSION + image.build_args.variant.distro.name == DistroName.ALPINE + and image.build_args.python_version == LATEST_SUPPORTED_PYTHON_VERSION ), ) # ubuntu images for each supported python version @@ -307,17 +300,198 @@ def tag_latest_ubuntu_with_flwr_version(image: BaseImage) -> List[str]: "serverapp", base_images, tag_latest_ubuntu_with_flwr_version, - lambda image: image.variant.distro.name == DistroName.UBUNTU, + lambda image: image.build_args.variant.distro.name == DistroName.UBUNTU, ) # ubuntu images for each supported python version + generate_binary_images( "clientapp", base_images, tag_latest_ubuntu_with_flwr_version, - lambda image: image.variant.distro.name == DistroName.UBUNTU, + lambda image: image.build_args.variant.distro.name == DistroName.UBUNTU, ) ) + return base_images, binary_images + + +# +# Build matrix for unstable releases +# +def build_unstable_matrix(flwr_version_ref: str) -> List[BaseImage]: + @dataclass + class UnstableBaseImageBuildArgs: + variant: Variant + python_version: str + flwr_version_ref: str + + cpu_ubuntu_build_args_variant = UnstableBaseImageBuildArgs( + UBUNTU_VARIANT, LATEST_SUPPORTED_PYTHON_VERSION, flwr_version_ref + ) + + cpu_build_args = """PYTHON_VERSION={python_version} +FLWR_VERSION_REF={flwr_version_ref} +DISTRO={distro_name} +DISTRO_VERSION={distro_version} +""" + + cpu_base_image = BaseImage( + file_dir_fn=lambda args: f"{DOCKERFILE_ROOT}/base/{args.variant.distro.name.value}", + tags_fn=lambda _: ["unstable"], + build_args_fn=lambda args: cpu_build_args.format( + python_version=args.python_version, + flwr_version_ref=args.flwr_version_ref, + distro_name=args.variant.distro.name, + distro_version=args.variant.distro.version, + ), + build_args=cpu_ubuntu_build_args_variant, + ) + + cuda_build_args_variant = UnstableBaseImageBuildArgs( + LATEST_SUPPORTED_CUDA_VERSION, LATEST_SUPPORTED_PYTHON_VERSION, flwr_version_ref + ) + + cuda_build_args = cpu_build_args + """CUDA_VERSION={cuda_version}""" + + cuda_base_image = BaseImage( + file_dir_fn=lambda args: f"{DOCKERFILE_ROOT}/base/{args.variant.distro.name.value}-cuda", + tags_fn=lambda _: ["unstable-cuda"], + build_args_fn=lambda args: cuda_build_args.format( + python_version=args.python_version, + flwr_version_ref=args.flwr_version_ref, + distro_name=args.variant.distro.name, + distro_version=args.variant.distro.version, + cuda_version=args.variant.extras.version, + ), + build_args=cuda_build_args_variant, + ) + + # base_images = [cpu_base_image, cuda_base_image] + base_images = [cpu_base_image] + + binary_images = ( + generate_binary_images( + "superlink", + base_images, + lambda image: image.tags, + lambda image: isinstance(image.build_args.variant.extras, CpuVariant), + ) + + generate_binary_images( + "supernode", + base_images, + lambda image: image.tags, + lambda image: isinstance(image.build_args.variant.extras, CpuVariant), + ) + + generate_binary_images("serverapp", base_images, lambda image: image.tags) + + generate_binary_images("clientapp", base_images, lambda image: image.tags) + ) + + return base_images, binary_images + + +# +# Build matrix for nightly releases +# +def build_nightly_matrix(flwr_version: str, flwr_package: str) -> List[BaseImage]: + @dataclass + class NightlyBaseImageBuildArgs: + variant: Variant + python_version: str + flwr_version: str + flwr_package: str + + cpu_ubuntu_build_args_variant = NightlyBaseImageBuildArgs( + UBUNTU_VARIANT, LATEST_SUPPORTED_PYTHON_VERSION, flwr_version, flwr_package + ) + + cpu_build_args = """PYTHON_VERSION={python_version} +FLWR_VERSION={flwr_version} +FLWR_PACKAGE={flwr_package} +DISTRO={distro_name} +DISTRO_VERSION={distro_version} +""" + + cpu_base_image = BaseImage( + file_dir_fn=lambda args: f"{DOCKERFILE_ROOT}/base/{args.variant.distro.name.value}", + tags_fn=lambda args: [args.flwr_version, "nightly"], + build_args_fn=lambda args: cpu_build_args.format( + python_version=args.python_version, + flwr_version=args.flwr_version, + flwr_package=args.flwr_package, + distro_name=args.variant.distro.name, + distro_version=args.variant.distro.version, + ), + build_args=cpu_ubuntu_build_args_variant, + ) + + cuda_build_args_variant = NightlyBaseImageBuildArgs( + LATEST_SUPPORTED_CUDA_VERSION, + LATEST_SUPPORTED_PYTHON_VERSION, + flwr_version, + flwr_package, + ) + + cuda_build_args = cpu_build_args + """CUDA_VERSION={cuda_version}""" + + cuda_base_image = BaseImage( + file_dir_fn=lambda args: f"{DOCKERFILE_ROOT}/base/{args.variant.distro.name.value}-cuda", + tags_fn=lambda args: [f"{args.flwr_version}-cuda", "nightly-cuda"], + build_args_fn=lambda args: cuda_build_args.format( + python_version=args.python_version, + flwr_version=args.flwr_version, + flwr_package=args.flwr_package, + distro_name=args.variant.distro.name, + distro_version=args.variant.distro.version, + cuda_version=args.variant.extras.version, + ), + build_args=cuda_build_args_variant, + ) + + # base_images = [cpu_base_image, cuda_base_image] + base_images = [cpu_base_image] + + binary_images = ( + generate_binary_images( + "superlink", + base_images, + lambda image: image.tags, + lambda image: isinstance(image.build_args.variant.extras, CpuVariant), + ) + + generate_binary_images( + "supernode", + base_images, + lambda image: image.tags, + lambda image: isinstance(image.build_args.variant.extras, CpuVariant), + ) + + generate_binary_images("serverapp", base_images, lambda image: image.tags) + + generate_binary_images("clientapp", base_images, lambda image: image.tags) + ) + + 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("--flwr-package", type=str, default="flwr") + arg_parser.add_argument( + "--matrix", choices=["stable", "nightly", "unstable"], default="stable" + ) + + args = arg_parser.parse_args() + + flwr_version = args.flwr_version + flwr_package = args.flwr_package + matrix = args.matrix + + if matrix == "stable": + base_images, binary_images = build_stable_matrix(flwr_version) + elif matrix == "nightly": + base_images, binary_images = build_nightly_matrix(flwr_version, flwr_package) + else: + base_images, binary_images = build_unstable_matrix(flwr_version) + print( json.dumps( {