Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(framework) Add FAB hash to load_client_app_fn #4305

Merged
merged 47 commits into from
Oct 10, 2024
Merged
Show file tree
Hide file tree
Changes from 46 commits
Commits
Show all changes
47 commits
Select commit Hold shift + click to select a range
54d76eb
feat(framework) Add FAB hash to build
chongshenng Oct 8, 2024
574c915
feat(framework) Add FAB constants
chongshenng Oct 8, 2024
36a4a79
Add Union
chongshenng Oct 8, 2024
e1f2010
feat(framework) Add FAB hash to install
chongshenng Oct 8, 2024
7f86734
Merge branch 'add-fab-constants' into add-fab-hash-install
chongshenng Oct 8, 2024
3073985
Merge branch 'add-fab-hash-build' into add-fab-hash-to-client-fn
chongshenng Oct 8, 2024
3501a29
Merge branch 'add-fab-hash-install' into add-fab-hash-to-client-fn
chongshenng Oct 8, 2024
358a2e2
feat(framework) Add FAB hash to client functions
chongshenng Oct 8, 2024
5a3aa54
Add FAB hash to get_project_dir
chongshenng Oct 8, 2024
62c703b
Add missing fab_hash argument
chongshenng Oct 8, 2024
3a3c432
Use rglob
chongshenng Oct 8, 2024
0741204
Add util function
chongshenng Oct 8, 2024
6f51376
Add allowed extensions
chongshenng Oct 8, 2024
d2b73ca
Remove unused import
chongshenng Oct 8, 2024
22d94a0
Merge branch 'add-fab-constants' into add-fab-hash-install
chongshenng Oct 8, 2024
44c3711
First change
chongshenng Oct 8, 2024
31b2ccd
Small refactor of tests
chongshenng Oct 8, 2024
70b8631
Disable too-many-locals
chongshenng Oct 8, 2024
b727dce
Merge branch 'main' into add-fab-hash-build
chongshenng Oct 8, 2024
b314f14
Update function name
chongshenng Oct 8, 2024
a5615aa
Merge branch 'main' into add-fab-hash-install
chongshenng Oct 8, 2024
1df7ddc
Update function name
chongshenng Oct 8, 2024
a914218
Merge branch 'main' into add-fab-hash-to-client-fn
chongshenng Oct 8, 2024
97d6029
Merge branch 'add-fab-constants' into add-fab-hash-to-client-fn
chongshenng Oct 8, 2024
983905c
Merge branch 'add-fab-hash-build' into add-fab-hash-to-client-fn
chongshenng Oct 8, 2024
685bd81
Update PR
chongshenng Oct 8, 2024
49a5fa0
Bugfix
chongshenng Oct 8, 2024
95ed8d5
Merge branch 'add-fab-hash-build' into add-fab-hash-to-client-fn
chongshenng Oct 8, 2024
b78b66b
Merge branch 'main' into add-fab-hash-install
chongshenng Oct 9, 2024
354ddde
Merge branch 'add-fab-hash-build' into add-fab-hash-install
chongshenng Oct 9, 2024
13d796d
Merge branch 'main' into add-fab-hash-to-client-fn
chongshenng Oct 9, 2024
9552846
Merge branch 'add-fab-hash-install' into add-fab-hash-to-client-fn
chongshenng Oct 9, 2024
b5c9d53
Address comments
chongshenng Oct 9, 2024
015079a
Move FAB validation to install.py
chongshenng Oct 9, 2024
baee2c8
Sort imports
chongshenng Oct 9, 2024
97ce98e
Fix
chongshenng Oct 9, 2024
cd3ebb3
Isort
chongshenng Oct 9, 2024
7ae9b60
Disable too-many-locals
chongshenng Oct 9, 2024
96921ed
Merge branch 'main' into add-fab-hash-install
danieljanes Oct 10, 2024
d8f1334
Merge branch 'main' into add-fab-hash-to-client-fn
chongshenng Oct 10, 2024
a773a4d
Merge branch 'add-fab-hash-install' into add-fab-hash-to-client-fn
chongshenng Oct 10, 2024
6108c2e
Bugfix
chongshenng Oct 10, 2024
f3b9ef3
Merge branch 'add-fab-hash-install' into add-fab-hash-to-client-fn
chongshenng Oct 10, 2024
9da8d65
Revert to typer.secho
chongshenng Oct 10, 2024
8479186
Merge branch 'add-fab-hash-install' into add-fab-hash-to-client-fn
chongshenng Oct 10, 2024
f5bab06
Add realistic FAB hash
chongshenng Oct 10, 2024
3b6b235
Merge branch 'main' into add-fab-hash-to-client-fn
danieljanes Oct 10, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 60 additions & 29 deletions src/py/flwr/cli/build.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,26 +14,50 @@
# ==============================================================================
"""Flower command line interface `build` command."""

import hashlib
import os
import shutil
import tempfile
import zipfile
from pathlib import Path
from typing import Annotated, Optional
from typing import Annotated, Any, Optional, Union

import pathspec
import tomli_w
import typer

from flwr.common.constant import FAB_ALLOWED_EXTENSIONS, FAB_DATE, FAB_HASH_TRUNCATION

from .config_utils import load_and_validate
from .utils import get_sha256_hash, is_valid_project_name
from .utils import is_valid_project_name


def write_to_zip(
zipfile_obj: zipfile.ZipFile, filename: str, contents: Union[bytes, str]
) -> zipfile.ZipFile:
"""Set a fixed date and write contents to a zip file."""
zip_info = zipfile.ZipInfo(filename)
zip_info.date_time = FAB_DATE
zipfile_obj.writestr(zip_info, contents)
return zipfile_obj


def get_fab_filename(conf: dict[str, Any], fab_hash: str) -> str:
"""Get the FAB filename based on the given config and FAB hash."""
publisher = conf["tool"]["flwr"]["app"]["publisher"]
name = conf["project"]["name"]
version = conf["project"]["version"].replace(".", "-")
fab_hash_truncated = fab_hash[:FAB_HASH_TRUNCATION]
return f"{publisher}.{name}.{version}.{fab_hash_truncated}.fab"

# pylint: disable=too-many-locals

# pylint: disable=too-many-locals, too-many-statements
def build(
app: Annotated[
Optional[Path],
typer.Option(help="Path of the Flower App to bundle into a FAB"),
] = None,
) -> str:
) -> tuple[str, str]:
"""Build a Flower App into a Flower App Bundle (FAB).

You can run ``flwr build`` without any arguments to bundle the app located in the
Expand Down Expand Up @@ -85,16 +109,8 @@ def build(
# Load .gitignore rules if present
ignore_spec = _load_gitignore(app)

# Set the name of the zip file
fab_filename = (
f"{conf['tool']['flwr']['app']['publisher']}"
f".{conf['project']['name']}"
f".{conf['project']['version'].replace('.', '-')}.fab"
)
list_file_content = ""

allowed_extensions = {".py", ".toml", ".md"}

# Remove the 'federations' field from 'tool.flwr' if it exists
if (
"tool" in conf
Expand All @@ -105,38 +121,53 @@ def build(

toml_contents = tomli_w.dumps(conf)

with zipfile.ZipFile(fab_filename, "w", zipfile.ZIP_DEFLATED) as fab_file:
fab_file.writestr("pyproject.toml", toml_contents)
with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as temp_file:
temp_filename = temp_file.name

with zipfile.ZipFile(temp_filename, "w", zipfile.ZIP_DEFLATED) as fab_file:
write_to_zip(fab_file, "pyproject.toml", toml_contents)

# Continue with adding other files
for root, _, files in os.walk(app, topdown=True):
files = [
# Continue with adding other files
all_files = [
f
for f in files
if not ignore_spec.match_file(Path(root) / f)
and f != fab_filename
and Path(f).suffix in allowed_extensions
and f != "pyproject.toml" # Exclude the original pyproject.toml
for f in app.rglob("*")
if not ignore_spec.match_file(f)
and f.name != temp_filename
and f.suffix in FAB_ALLOWED_EXTENSIONS
and f.name != "pyproject.toml" # Exclude the original pyproject.toml
]

for file in files:
file_path = Path(root) / file
for file_path in all_files:
# Read the file content manually
with open(file_path, "rb") as f:
file_contents = f.read()

archive_path = file_path.relative_to(app)
fab_file.write(file_path, archive_path)
write_to_zip(fab_file, str(archive_path), file_contents)

# Calculate file info
sha256_hash = get_sha256_hash(file_path)
sha256_hash = hashlib.sha256(file_contents).hexdigest()
file_size_bits = os.path.getsize(file_path) * 8 # size in bits
list_file_content += f"{archive_path},{sha256_hash},{file_size_bits}\n"

# Add CONTENT and CONTENT.jwt to the zip file
fab_file.writestr(".info/CONTENT", list_file_content)
# Add CONTENT and CONTENT.jwt to the zip file
write_to_zip(fab_file, ".info/CONTENT", list_file_content)

# Get hash of FAB file
content = Path(temp_filename).read_bytes()
fab_hash = hashlib.sha256(content).hexdigest()

# Set the name of the zip file
fab_filename = get_fab_filename(conf, fab_hash)

# Once the temporary zip file is created, rename it to the final filename
shutil.move(temp_filename, fab_filename)

typer.secho(
f"🎊 Successfully built {fab_filename}", fg=typer.colors.GREEN, bold=True
)

return fab_filename
return fab_filename, fab_hash


def _load_gitignore(app: Path) -> pathspec.PathSpec:
Expand Down
80 changes: 60 additions & 20 deletions src/py/flwr/cli/install.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
# ==============================================================================
"""Flower command line interface `install` command."""


import hashlib
import shutil
import subprocess
import tempfile
Expand All @@ -25,7 +25,8 @@

import typer

from flwr.common.config import get_flwr_dir
from flwr.common.config import get_flwr_dir, get_metadata_from_config
from flwr.common.constant import FAB_HASH_TRUNCATION

from .config_utils import load_and_validate
from .utils import get_sha256_hash
Expand Down Expand Up @@ -91,9 +92,11 @@ def install_from_fab(
fab_name: Optional[str]
if isinstance(fab_file, bytes):
fab_file_archive = BytesIO(fab_file)
fab_hash = hashlib.sha256(fab_file).hexdigest()
fab_name = None
elif isinstance(fab_file, Path):
fab_file_archive = fab_file
fab_hash = hashlib.sha256(fab_file.read_bytes()).hexdigest()
fab_name = fab_file.stem
else:
raise ValueError("fab_file must be either a Path or bytes")
Expand Down Expand Up @@ -126,14 +129,16 @@ def install_from_fab(
shutil.rmtree(info_dir)

installed_path = validate_and_install(
tmpdir_path, fab_name, flwr_dir, skip_prompt
tmpdir_path, fab_hash, fab_name, flwr_dir, skip_prompt
)

return installed_path


# pylint: disable=too-many-locals
def validate_and_install(
project_dir: Path,
fab_hash: str,
fab_name: Optional[str],
flwr_dir: Optional[Path],
skip_prompt: bool = False,
Expand All @@ -149,28 +154,17 @@ def validate_and_install(
)
raise typer.Exit(code=1)

publisher = config["tool"]["flwr"]["app"]["publisher"]
project_name = config["project"]["name"]
version = config["project"]["version"]
version, fab_id = get_metadata_from_config(config)
publisher, project_name = fab_id.split("/")
config_metadata = (publisher, project_name, version, fab_hash)

if (
fab_name
and fab_name != f"{publisher}.{project_name}.{version.replace('.', '-')}"
):
typer.secho(
"❌ FAB file has incorrect name. The file name must follow the format "
"`<publisher>.<project_name>.<version>.fab`.",
fg=typer.colors.RED,
bold=True,
)
raise typer.Exit(code=1)
if fab_name:
_validate_fab_and_config_metadata(fab_name, config_metadata)

install_dir: Path = (
(get_flwr_dir() if not flwr_dir else flwr_dir)
/ "apps"
/ publisher
/ project_name
/ version
/ f"{publisher}.{project_name}.{version}.{fab_hash[:FAB_HASH_TRUNCATION]}"
)
if install_dir.exists():
if skip_prompt:
Expand Down Expand Up @@ -226,3 +220,49 @@ def _verify_hashes(list_content: str, tmpdir: Path) -> bool:
if not file_path.exists() or get_sha256_hash(file_path) != hash_expected:
return False
return True


def _validate_fab_and_config_metadata(
fab_name: str, config_metadata: tuple[str, str, str, str]
) -> None:
"""Validate metadata from the FAB filename and config."""
publisher, project_name, version, fab_hash = config_metadata

fab_name = fab_name.removesuffix(".fab")

fab_publisher, fab_project_name, fab_version, fab_shorthash = fab_name.split(".")
fab_version = fab_version.replace("-", ".")

# Check FAB filename format
if (
f"{fab_publisher}.{fab_project_name}.{fab_version}"
!= f"{publisher}.{project_name}.{version}"
or len(fab_shorthash) != FAB_HASH_TRUNCATION # Verify hash length
):
typer.secho(
"❌ FAB file has incorrect name. The file name must follow the format "
"`<publisher>.<project_name>.<version>.<8hexchars>.fab`.",
fg=typer.colors.RED,
bold=True,
)
raise typer.Exit(code=1)

# Verify hash is a valid hexadecimal
try:
_ = int(fab_shorthash, 16)
except Exception as e:
typer.secho(
f"❌ FAB file has an invalid hexadecimal string `{fab_shorthash}`.",
fg=typer.colors.RED,
bold=True,
)
raise typer.Exit(code=1) from e

# Verify shorthash matches
if fab_shorthash != fab_hash[:FAB_HASH_TRUNCATION]:
typer.secho(
"❌ The hash in the FAB file name does not match the hash of the FAB.",
fg=typer.colors.RED,
bold=True,
)
raise typer.Exit(code=1)
10 changes: 5 additions & 5 deletions src/py/flwr/cli/run/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
# ==============================================================================
"""Flower command line interface `run` command."""

import hashlib
import json
import subprocess
import sys
Expand Down Expand Up @@ -134,6 +133,7 @@ def run(
_run_without_superexec(app, federation_config, config_overrides, federation)


# pylint: disable=too-many-locals
def _run_with_superexec(
app: Path,
federation_config: dict[str, Any],
Expand Down Expand Up @@ -179,9 +179,9 @@ def _run_with_superexec(
channel.subscribe(on_channel_state_change)
stub = ExecStub(channel)

fab_path = Path(build(app))
content = fab_path.read_bytes()
fab = Fab(hashlib.sha256(content).hexdigest(), content)
fab_path, fab_hash = build(app)
content = Path(fab_path).read_bytes()
fab = Fab(fab_hash, content)

req = StartRunRequest(
fab=fab_to_proto(fab),
Expand All @@ -193,7 +193,7 @@ def _run_with_superexec(
res = stub.StartRun(req)

# Delete FAB file once it has been sent to the SuperExec
fab_path.unlink()
Path(fab_path).unlink()
typer.secho(f"🎊 Successfully started run {res.run_id}", fg=typer.colors.GREEN)

if stream:
Expand Down
6 changes: 3 additions & 3 deletions src/py/flwr/client/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -202,7 +202,7 @@ def start_client_internal(
*,
server_address: str,
node_config: UserConfig,
load_client_app_fn: Optional[Callable[[str, str], ClientApp]] = None,
load_client_app_fn: Optional[Callable[[str, str, str], ClientApp]] = None,
client_fn: Optional[ClientFnExt] = None,
client: Optional[Client] = None,
grpc_max_message_length: int = GRPC_MAX_MESSAGE_LENGTH,
Expand Down Expand Up @@ -298,7 +298,7 @@ def single_client_factory(

client_fn = single_client_factory

def _load_client_app(_1: str, _2: str) -> ClientApp:
def _load_client_app(_1: str, _2: str, _3: str) -> ClientApp:
return ClientApp(client_fn=client_fn)

load_client_app_fn = _load_client_app
Expand Down Expand Up @@ -529,7 +529,7 @@ def _on_backoff(retry_state: RetryState) -> None:
else:
# Load ClientApp instance
client_app: ClientApp = load_client_app_fn(
fab_id, fab_version
fab_id, fab_version, run.fab_hash
)

# Execute ClientApp
Expand Down
7 changes: 5 additions & 2 deletions src/py/flwr/client/clientapp/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -132,8 +132,11 @@ def run_clientapp( # pylint: disable=R0914
)

try:
# Load ClientApp
client_app: ClientApp = load_client_app_fn(run.fab_id, run.fab_version)
if fab:
# Load ClientApp
client_app: ClientApp = load_client_app_fn(
run.fab_id, run.fab_version, fab.hash_str
)

# Execute ClientApp
reply_message = client_app(message=message, context=context)
Expand Down
Loading
Loading