syft-flwr 0.1.0__py3-none-any.whl

README.md

@@ -0,0 +1,6 @@
+# syft_flwr
+
+`syft_flwr` is an open source framework that facilitate federated learning projects using [Flower](https://github.com/adap/flower) over the [SyftBox](https://github.com/OpenMined/syftbox) protocol
+
+## Installation
+`pip install syft_flwr`

pyproject.toml

@@ -0,0 +1,52 @@
+[project]
+name = "syft-flwr"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.9.2"
+dependencies = [
+    "flwr[simulation]==1.17.0",
+    "flwr-datasets[vision]>=0.5.0",
+    "loguru>=0.7.3",
+    "safetensors>=0.5.0",
+    "typing-extensions>=4.13.0",
+    "tomli>=2.2.1",
+    "tomli-w>=1.2.0",
+    "syft-core==0.2.3",
+    "syft-event==0.2.0",
+    "syft-rpc==0.2.0",
+    "syft-rds==0.1.0",
+]
+
+[project.scripts]
+syft_flwr = "syft_flwr.cli:main"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.uv]
+dev-dependencies = [
+    "ipykernel>=6.29.5",
+    "ipywidgets>=8.1.6",
+    "pytest>=8.3.4",
+    "pre-commit>=4.0.1",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/syft_flwr"]
+only-include = ["src", "pyproject.toml", "/README.md"]
+exclude = ["src/**/__pycache__"]
+
+[tool.hatch.build.targets.sdist]
+only-include = ["src", "pyproject.toml", "/README.md"]
+exclude = ["src/**/__pycache__", "examples", "notebooks", "justfile"]
+
+[tool.ruff]
+exclude = [".archive"]
+
+[tool.ruff.lint]
+extend-select = ["I"]
+
+[tool.ruff.lint.per-file-ignores]
+"**/__init__.py" = ["F401"]

syft_flwr/__init__.py

@@ -0,0 +1,6 @@
+__version__ = "0.1.0"
+
+from syft_flwr.bootstrap import bootstrap
+from syft_flwr.run import run
+
+__all__ = ["bootstrap", "run"]

syft_flwr/bootstrap.py

@@ -0,0 +1,97 @@
+from pathlib import Path
+
+from loguru import logger
+from typing_extensions import List, Union
+
+from syft_flwr import __version__
+from syft_flwr.config import load_flwr_pyproject, write_toml
+from syft_flwr.utils import is_valid_datasite
+
+__all__ = ["bootstrap"]
+
+
+MAIN_TEMPLATE_PATH = Path(__file__).parent / "templates" / "main.py.tpl"
+MAIN_TEMPLATE_CONTENT = MAIN_TEMPLATE_PATH.read_text()
+assert MAIN_TEMPLATE_CONTENT
+
+
+def __copy_main_py(flwr_project_dir: Path) -> None:
+    """Copy the content below to `main.py` file to the syft-flwr project"""
+
+    main_py_path = flwr_project_dir / "main.py"
+
+    if main_py_path.exists():
+        raise Exception(f"The file '{main_py_path}' already exists")
+
+    main_py_path.write_text(MAIN_TEMPLATE_CONTENT)
+
+
+def __update_pyproject_toml(
+    flwr_project_dir: Union[str, Path],
+    aggregator: str,
+    datasites: List[str],
+) -> None:
+    """Update the `pyproject.toml` file to the syft-flwr project"""
+    flwr_project_dir = Path(flwr_project_dir)
+    flwr_pyproject = Path(flwr_project_dir, "pyproject.toml")
+    pyproject_conf = load_flwr_pyproject(flwr_pyproject, check_module=False)
+
+    # TODO: remove this after we find out how to pass the right context to the clients
+    pyproject_conf["tool"]["flwr"]["app"]["config"]["partition-id"] = 0
+    pyproject_conf["tool"]["flwr"]["app"]["config"]["num-partitions"] = 1
+    # TODO end
+
+    # add syft_flwr as a dependency
+    if "dependencies" not in pyproject_conf["project"]:
+        pyproject_conf["project"]["dependencies"] = []
+
+    deps: list = pyproject_conf["project"]["dependencies"]
+    deps = [dep for dep in deps if not dep.startswith("syft_flwr")]
+    deps.append(f"syft_flwr=={__version__}")
+    pyproject_conf["project"]["dependencies"] = deps
+
+    # always override the datasites and aggregator
+    pyproject_conf["tool"]["syft_flwr"] = {}
+    pyproject_conf["tool"]["syft_flwr"]["datasites"] = datasites
+    pyproject_conf["tool"]["syft_flwr"]["aggregator"] = aggregator
+
+    write_toml(flwr_pyproject, pyproject_conf)
+
+
+def __validate_flwr_project_dir(flwr_project_dir: Union[str, Path]) -> Path:
+    flwr_pyproject = flwr_project_dir / "pyproject.toml"
+    flwr_main_py = flwr_project_dir / "main.py"
+
+    if flwr_main_py.exists():
+        raise FileExistsError(f"File '{flwr_main_py}' already exists")
+
+    if not flwr_project_dir.exists():
+        raise FileNotFoundError(f"Directory '{flwr_project_dir}' not found")
+
+    if not flwr_pyproject.exists():
+        raise FileNotFoundError(f"File '{flwr_pyproject}' not found")
+
+
+def bootstrap(
+    flwr_project_dir: Union[str, Path], aggregator: str, datasites: List[str]
+) -> None:
+    """Bootstrap a new syft-flwr project from the flwr project at the given path"""
+    flwr_project_dir = Path(flwr_project_dir)
+
+    if not is_valid_datasite(aggregator):
+        raise ValueError(f"'{aggregator}' is not a valid datasite")
+
+    if len(datasites) < 2:
+        raise ValueError("You must provide at least two datasites as Flower clients")
+
+    for ds in datasites:
+        if not is_valid_datasite(ds):
+            raise ValueError(f"{ds} is not a valid datasite")
+
+    __validate_flwr_project_dir(flwr_project_dir)
+    __update_pyproject_toml(flwr_project_dir, aggregator, datasites)
+    __copy_main_py(flwr_project_dir)
+
+    logger.info(
+        f"Successfully bootstrapped syft-flwr project at {flwr_project_dir} with datasites {datasites} and aggregator {aggregator} ✅"
+    )

syft_flwr/cli.py

@@ -0,0 +1,116 @@
+from pathlib import Path
+
+import typer
+from rich import print as rprint
+from typing_extensions import Annotated, List, Tuple
+
+app = typer.Typer(
+    name="syft_flwr",
+    no_args_is_help=True,
+    pretty_exceptions_enable=False,
+    context_settings={"help_option_names": ["-h", "--help"]},
+)
+
+
+@app.command()
+def version() -> None:
+    """Print syft_flwr version"""
+    from syft_flwr import __version__
+
+    print(__version__)
+
+
+PROJECT_DIR_OPTS = typer.Argument(help="Path to a Flower project")
+AGGREGATOR_OPTS = typer.Option(
+    "-a",
+    "--aggregator",
+    "-s",
+    "--server",
+    help="Datasite email of the Flower Server",
+)
+DATASITES_OPTS = typer.Option(
+    "-d",
+    "--datasites",
+    help="Datasites addresses",
+)
+MOCK_DATASET_PATHS_OPTS = typer.Option(
+    "-m",
+    "--mock-dataset-paths",
+    help="Mock dataset paths",
+)
+
+
+def prompt_for_missing_args(
+    aggregator: str, datasites: List[str]
+) -> Tuple[Path, str, List[str]]:
+    if not aggregator:
+        aggregator = typer.prompt(
+            "Enter the datasite email of the Aggregator (Flower Server)"
+        )
+    if not datasites:
+        datasites = typer.prompt(
+            "Enter a comma-separated email of datasites of the Flower Clients"
+        )
+        datasites = datasites.split(",")
+
+    return aggregator, datasites
+
+
+def prompt_for_missing_mock_paths(mock_dataset_paths: List[str]) -> List[str]:
+    if not mock_dataset_paths:
+        mock_paths = typer.prompt("Enter comma-separated paths to mock datasets")
+        mock_dataset_paths = mock_paths.split(",")
+
+    return mock_dataset_paths
+
+
+@app.command()
+def bootstrap(
+    project_dir: Annotated[Path, PROJECT_DIR_OPTS],
+    aggregator: Annotated[str, AGGREGATOR_OPTS] = None,
+    datasites: Annotated[List[str], DATASITES_OPTS] = None,
+) -> None:
+    """Bootstrap a new syft_flwr project from a flwr project"""
+    from syft_flwr.bootstrap import bootstrap
+
+    aggregator, datasites = prompt_for_missing_args(aggregator, datasites)
+
+    try:
+        project_dir = project_dir.absolute()
+        rprint(f"[cyan]Bootstrapping project at '{project_dir}'[/cyan]")
+        rprint(f"[cyan]Aggregator: {aggregator}[/cyan]")
+        rprint(f"[cyan]Datasites: {datasites}[/cyan]")
+        bootstrap(project_dir, aggregator, datasites)
+        rprint(f"[green]Bootstrapped project at '{project_dir}'[/green]")
+    except Exception as e:
+        rprint(f"[red]Error[/red]: {e}")
+        raise typer.Exit(1)
+
+
+@app.command()
+def run(
+    project_dir: Annotated[Path, PROJECT_DIR_OPTS],
+    mock_dataset_paths: Annotated[List[str], MOCK_DATASET_PATHS_OPTS] = None,
+) -> None:
+    """Run a syft_flwr project in simulation mode over mock data"""
+    from syft_flwr.run_simulation import run
+
+    try:
+        mock_dataset_paths: list[str] = prompt_for_missing_mock_paths(
+            mock_dataset_paths
+        )
+        project_dir = Path(project_dir).expanduser().resolve()
+        rprint(f"[cyan]Running syft_flwr project at '{project_dir}'[/cyan]")
+        rprint(f"[cyan]Mock dataset paths: {mock_dataset_paths}[/cyan]")
+        run(project_dir, mock_dataset_paths)
+    except Exception as e:
+        rprint(f"[red]Error[/red]: {e}")
+        raise typer.Exit(1)
+
+
+def main() -> None:
+    app()
+
+
+if __name__ == "__main__":
+    main()

syft_flwr/config.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import tomli
+import tomli_w
+from flwr.common.config import validate_config
+from loguru import logger
+
+
+def load_toml(path: str):
+    with open(path, "rb") as fp:
+        return tomli.load(fp)
+
+
+def write_toml(path: str, val: dict):
+    with open(path, "wb") as fp:
+        tomli_w.dump(val, fp)
+
+
+def load_flwr_pyproject(path: Path, check_module: bool = True) -> dict:
+    """Load the flower's pyproject.toml file and validate it."""
+
+    if path.name != "pyproject.toml":
+        path = path / "pyproject.toml"
+
+    pyproject = load_toml(path)
+    is_valid, errors, warnings = validate_config(pyproject, check_module, path.parent)
+
+    if not is_valid:
+        raise Exception(errors)
+
+    if warnings:
+        logger.warning(warnings)
+
+    return pyproject

syft_flwr/flower_client.py

@@ -0,0 +1,65 @@
+import sys
+import traceback
+
+from flwr.client import ClientApp
+from flwr.common import Context
+from flwr.common.constant import ErrorCode
+from flwr.common.message import Error, Message
+from loguru import logger
+from syft_event import SyftEvents
+from syft_event.types import Request
+
+from syft_flwr.flwr_compatibility import RecordDict, create_flwr_message
+from syft_flwr.serde import bytes_to_flower_message, flower_message_to_bytes
+
+
+def syftbox_flwr_client(client_app: ClientApp, context: Context):
+    """Run the Flower ClientApp with SyftBox."""
+
+    box = SyftEvents("flwr")
+    client_email = box.client.email
+    logger.info(f"Started SyftBox Flower Client on: {client_email}")
+
+    @box.on_request("/messages")
+    def handle_messages(request: Request) -> None:
+        logger.info(
+            f"Received request id: {request.id}, size: {len(request.body) / 1024 / 1024} (MB)"
+        )
+        message: Message = bytes_to_flower_message(request.body)
+
+        try:
+            reply_message: Message = client_app(message=message, context=context)
+            res_bytes: bytes = flower_message_to_bytes(reply_message)
+            logger.info(f"Reply message size: {len(res_bytes)/2**20} MB")
+            return res_bytes
+
+        except Exception as e:
+            error_traceback = traceback.format_exc()
+            error_message = f"Client: '{client_email}'. Error: {str(e)}. Traceback: {error_traceback}"
+            logger.error(error_message)
+
+            error = Error(
+                code=ErrorCode.CLIENT_APP_RAISED_EXCEPTION, reason=f"{error_message}"
+            )
+
+            error_reply: Message = create_flwr_message(
+                content=RecordDict(),
+                reply_to=message,
+                message_type=message.metadata.message_type,
+                src_node_id=message.metadata.dst_node_id,
+                dst_node_id=message.metadata.src_node_id,
+                group_id=message.metadata.group_id,
+                run_id=message.metadata.run_id,
+                error=error,
+            )
+            error_bytes: bytes = flower_message_to_bytes(error_reply)
+            logger.info(f"Error reply message size: {len(error_bytes)/2**20} MB")
+            return error_bytes
+
+    try:
+        box.run_forever()
+    except Exception as e:
+        logger.error(
+            f"Fatal error in syftbox_flwr_client: {str(e)}\n{traceback.format_exc()}"
+        )
+        sys.exit(1)

syft_flwr/flower_server.py

@@ -0,0 +1,25 @@
+from random import randint
+
+from flwr.common import Context
+from flwr.server import ServerApp
+from flwr.server.run_serverapp import run as run_server
+from loguru import logger
+
+from syft_flwr.grid import SyftGrid
+
+
+def syftbox_flwr_server(server_app: ServerApp, context: Context, datasites: list[str]):
+    """Run the Flower ServerApp with SyftBox."""
+    syft_grid = SyftGrid(datasites=datasites)
+    run_id = randint(0, 1000)
+    syft_grid.set_run(run_id)
+    logger.info(f"Started SyftBox Flower Server on: {syft_grid._client.email}")
+
+    updated_context = run_server(
+        syft_grid,
+        context=context,
+        loaded_server_app=server_app,
+        server_app_dir="",
+    )
+    logger.info(f"Server completed with context: {updated_context}")
+    return updated_context

syft_flwr/flwr_compatibility.py

@@ -0,0 +1,121 @@
+import flwr
+from flwr.common import Metadata
+from flwr.common.message import Error, Message
+from packaging.version import Version
+from typing_extensions import Optional
+
+
+def flwr_later_than_1_17():
+    return Version(flwr.__version__) >= Version("1.17.0")
+
+
+# Version-dependent imports
+if flwr_later_than_1_17():
+    from flwr.common.record import RecordDict
+    from flwr.server.grid import Grid
+else:
+    from flwr.common.record import RecordSet as RecordDict
+    from flwr.server.driver import Driver as Grid
+
+
+__all__ = ["Grid", "RecordDict"]
+
+
+def check_reply_to_field(metadata: Metadata) -> bool:
+    """Check if reply_to field is empty based on Flower version."""
+    if flwr_later_than_1_17():
+        return metadata.reply_to_message_id == ""
+    else:
+        return metadata.reply_to_message == ""
+
+
+def create_flwr_message(
+    content: RecordDict,
+    message_type: str,
+    src_node_id: int,
+    dst_node_id: int,
+    group_id: str,
+    run_id: int,
+    ttl: Optional[float] = None,
+    error: Optional[Error] = None,
+    reply_to: Optional[Message] = None,
+) -> Message:
+    """Create a Flower message with version-compatible parameters."""
+    if flwr_later_than_1_17():
+        return _create_message_v1_17_plus(
+            content,
+            message_type,
+            dst_node_id,
+            group_id,
+            ttl,
+            error,
+            reply_to,
+        )
+    else:
+        return _create_message_pre_v1_17(
+            content,
+            message_type,
+            src_node_id,
+            dst_node_id,
+            group_id,
+            run_id,
+            ttl,
+            error,
+        )
+
+
+def _create_message_v1_17_plus(
+    content: RecordDict,
+    message_type: str,
+    dst_node_id: int,
+    group_id: str,
+    ttl: Optional[float],
+    error: Optional[Error],
+    reply_to: Optional[Message],
+) -> Message:
+    """Create message for Flower version 1.17+."""
+    if reply_to is not None:
+        if error is not None:
+            return Message(reply_to=reply_to, error=error)
+        return Message(content=content, reply_to=reply_to)
+    else:
+        if error is not None:
+            raise ValueError("Error and reply_to cannot both be None")
+        return Message(
+            content=content,
+            dst_node_id=dst_node_id,
+            message_type=message_type,
+            ttl=ttl,
+            group_id=group_id,
+        )
+
+
+def _create_message_pre_v1_17(
+    content: RecordDict,
+    message_type: str,
+    src_node_id: int,
+    dst_node_id: int,
+    group_id: str,
+    run_id: int,
+    ttl: Optional[float],
+    error: Optional[Error],
+) -> Message:
+    """Create message for Flower versions before 1.17."""
+    from flwr.common import DEFAULT_TTL
+
+    ttl_ = DEFAULT_TTL if ttl is None else ttl
+    metadata = Metadata(
+        run_id=run_id,
+        message_id="",  # Will be set when saving to file
+        src_node_id=src_node_id,
+        dst_node_id=dst_node_id,
+        reply_to_message="",
+        group_id=group_id,
+        ttl=ttl_,
+        message_type=message_type,
+    )
+
+    if error is not None:
+        return Message(metadata=metadata, error=error)
+    else:
+        return Message(metadata=metadata, content=content)