Skip to content

Commit

Permalink
feat(framework:skip) Add GPU base image to build matrix
Browse files Browse the repository at this point in the history 
Signed-off-by: Robert Steiner <[email protected]>
  • Loading branch information
@Robert-Steiner
Robert-Steiner committed Nov 4, 2024
1 parent ba690fd commit e11b3aa
Show file tree
Hide file tree
Showing 2 changed files with 159 additions and 27 deletions.
5 changes: 1 addition & 4 deletions .github/workflows/framework-release.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}
Expand Down
181 changes: 158 additions & 23 deletions dev/build-docker-image-matrix.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
"""
Usage: python dev/build-docker-image-matrix.py --flwr-version <flower version e.g. 1.9.0>
Usage: python dev/build-docker-image-matrix.py --flwr-version <flower version e.g. 1.13.0>
"""

import argparse
Expand Down Expand Up @@ -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
]

Expand Down Expand Up @@ -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]
Expand All @@ -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:
Expand All @@ -130,37 +241,46 @@ 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
generate_binary_images(
"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
),
)
Expand All @@ -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))
},
Expand Down

0 comments on commit e11b3aa

Please sign in to comment.