diff --git a/.github/workflows/framework-release.yml b/.github/workflows/framework-release.yml index e608329872de..6e9225c4ad16 100644 --- a/.github/workflows/framework-release.yml +++ b/.github/workflows/framework-release.yml @@ -86,12 +86,9 @@ jobs: namespace-repository: ${{ matrix.images.namespace_repository }} file-dir: ${{ matrix.images.file_dir }} build-args: | - PYTHON_VERSION=${{ matrix.images.python_version }} PIP_VERSION=${{ needs.parameters.outputs.pip-version }} SETUPTOOLS_VERSION=${{ needs.parameters.outputs.setuptools-version }} - DISTRO=${{ matrix.images.distro.name }} - DISTRO_VERSION=${{ matrix.images.distro.version }} - FLWR_VERSION=${{ matrix.images.flwr_version }} + ${{ matrix.images.build_args }} tags: ${{ matrix.images.tag }} secrets: dockerhub-user: ${{ secrets.DOCKERHUB_USERNAME }} diff --git a/dev/build-docker-image-matrix.py b/dev/build-docker-image-matrix.py index 52c96e3cca7a..0cc44ac057cf 100644 --- a/dev/build-docker-image-matrix.py +++ b/dev/build-docker-image-matrix.py @@ -1,5 +1,5 @@ """ -Usage: python dev/build-docker-image-matrix.py --flwr-version +Usage: python dev/build-docker-image-matrix.py --flwr-version """ import argparse @@ -31,34 +31,144 @@ class Distro: @dataclass -class BaseImage: +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"), +] + + +def remove_path_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 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 +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_path_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 +] + + +@dataclass +class BaseImage: + variant: Variant python_version: str namespace_repository: str file_dir: str tag: str flwr_version: str + build_args: str def new_base_image( - flwr_version: str, python_version: str, distro: Distro + flwr_version: str, python_version: str, variant: Variant ) -> Dict[str, Any]: return BaseImage( - distro, + variant, python_version, "flwr/base", - f"{DOCKERFILE_ROOT}/base/{distro.name.value}", - f"{flwr_version}-py{python_version}-{distro.name.value}{distro.version}", + variant.file_dir(), + variant.tag(flwr_version, python_version), flwr_version, + variant.build_args(flwr_version, python_version), ) def generate_base_images( - flwr_version: str, python_versions: List[str], distros: List[Dict[str, str]] + flwr_version: str, python_versions: List[str], variants: List[Variant] ) -> List[Dict[str, Any]]: return [ - new_base_image(flwr_version, python_version, distro) - for distro in distros + new_base_image(flwr_version, python_version, variant) + for variant in variants for python_version in python_versions ] @@ -103,7 +213,7 @@ def generate_binary_images( def tag_latest_alpine_with_flwr_version(image: BaseImage) -> List[str]: if ( - image.distro.name == DistroName.ALPINE + image.variant.distro.name == DistroName.ALPINE and image.python_version == LATEST_SUPPORTED_PYTHON_VERSION ): return [image.tag, image.flwr_version] @@ -113,8 +223,9 @@ def tag_latest_alpine_with_flwr_version(image: BaseImage) -> List[str]: def tag_latest_ubuntu_with_flwr_version(image: BaseImage) -> List[str]: if ( - image.distro.name == DistroName.UBUNTU + image.variant.distro.name == DistroName.UBUNTU and image.python_version == LATEST_SUPPORTED_PYTHON_VERSION + and isinstance(image.variant.extras, CpuVariant) ): return [image.tag, image.flwr_version] else: @@ -130,20 +241,25 @@ def tag_latest_ubuntu_with_flwr_version(image: BaseImage) -> List[str]: flwr_version = args.flwr_version - # ubuntu base images for each supported python version ubuntu_base_images = generate_base_images( flwr_version, SUPPORTED_PYTHON_VERSIONS, - [Distro(DistroName.UBUNTU, "24.04")], + UBUNTU_VARIANTS, ) - # 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_VARIANTS, ) - base_images = ubuntu_base_images + alpine_base_images + cuda_base_images = generate_base_images( + flwr_version, + SUPPORTED_PYTHON_VERSIONS, + CUDA_VARIANTS, + ) + + base_images = ubuntu_base_images + alpine_base_images # + cuda_base_images binary_images = ( # ubuntu and alpine images for the latest supported python version @@ -151,16 +267,20 @@ 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, + lambda image: image.python_version == LATEST_SUPPORTED_PYTHON_VERSION + and isinstance(image.variant.extras, CpuVariant), ) # ubuntu images for each supported python version + generate_binary_images( "supernode", base_images, tag_latest_alpine_with_flwr_version, - lambda image: image.distro.name == DistroName.UBUNTU + lambda image: ( + image.variant.distro.name == DistroName.UBUNTU + and isinstance(image.variant.extras, CpuVariant), + ) or ( - image.distro.name == DistroName.ALPINE + image.variant.distro.name == DistroName.ALPINE and image.python_version == LATEST_SUPPORTED_PYTHON_VERSION ), ) @@ -169,28 +289,43 @@ def tag_latest_ubuntu_with_flwr_version(image: BaseImage) -> List[str]: "serverapp", base_images, tag_latest_ubuntu_with_flwr_version, - lambda image: image.distro.name == DistroName.UBUNTU, + lambda image: image.variant.distro.name == DistroName.UBUNTU, ) # ubuntu images for each supported python version + generate_binary_images( "superexec", base_images, tag_latest_ubuntu_with_flwr_version, - lambda image: image.distro.name == DistroName.UBUNTU, + lambda image: image.variant.distro.name == DistroName.UBUNTU + and isinstance(image.variant.extras, CpuVariant), ) # ubuntu images for each supported python version + generate_binary_images( "clientapp", base_images, tag_latest_ubuntu_with_flwr_version, - lambda image: image.distro.name == DistroName.UBUNTU, + lambda image: image.variant.distro.name == DistroName.UBUNTU, ) ) print( json.dumps( { - "base": {"images": list(map(lambda image: asdict(image), base_images))}, + "base": { + "images": list( + map( + lambda image: asdict( + image, + dict_factory=lambda x: { + k: v + for (k, v) in x + if v is not None and callable(v) is False + }, + ), + base_images, + ) + ) + }, "binary": { "images": list(map(lambda image: asdict(image), binary_images)) },