syft-flwr 0.4.3__py3-none-any.whl

syft_flwr/__init__.py

@@ -0,0 +1,15 @@
+__version__ = "0.4.3"
+
+from syft_flwr.bootstrap import bootstrap
+from syft_flwr.run import run
+
+__all__ = ["bootstrap", "run"]
+
+
+# Register the mount provider for syft_rds when syft_flwr is initializedAdd commentMore actions
+from syft_rds.syft_runtime.mounts import register_mount_provider
+
+from .mounts import SyftFlwrMountProvider
+
+# Register the mount provider
+register_mount_provider("syft_flwr", SyftFlwrMountProvider())

syft_flwr/bootstrap.py

@@ -0,0 +1,102 @@
+import time
+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
+
+    pyproject_conf["tool"]["syft_flwr"] = {}
+
+    # configure unique app name for each syft_flwr run
+    base_app_name = pyproject_conf["project"]["name"]
+    pyproject_conf["tool"]["syft_flwr"]["app_name"] = (
+        f"{aggregator}_{base_app_name}_{int(time.time())}"
+    )
+
+    # always override the datasites and aggregator
+    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")
+
+    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/consts.py

@@ -0,0 +1,2 @@
+# Environment variable to control encryption
+SYFT_FLWR_ENCRYPTION_ENABLED = "SYFT_FLWR_ENCRYPTION_ENABLED"

syft_flwr/flower_client.py

@@ -0,0 +1,199 @@
+import base64
+import sys
+import traceback
+
+from flwr.client import ClientApp
+from flwr.common import Context
+from flwr.common.constant import ErrorCode, MessageType
+from flwr.common.message import Error, Message
+from flwr.common.record import RecordDict
+from loguru import logger
+from syft_event import SyftEvents
+from syft_event.types import Request
+from typing_extensions import Optional, Union
+
+from syft_flwr.serde import bytes_to_flower_message, flower_message_to_bytes
+from syft_flwr.utils import create_flwr_message, setup_client
+
+
+class MessageHandler:
+    """Handles message processing for Flower client."""
+
+    def __init__(
+        self, client_app: ClientApp, context: Context, encryption_enabled: bool
+    ):
+        self.client_app = client_app
+        self.context = context
+        self.encryption_enabled = encryption_enabled
+
+    def prepare_reply(self, data: bytes) -> Union[str, bytes]:
+        """Prepare reply data based on encryption setting."""
+        if self.encryption_enabled:
+            logger.info(f"🔒 Preparing ENCRYPTED reply, size: {len(data)/2**20:.2f} MB")
+            return base64.b64encode(data).decode("utf-8")
+        else:
+            logger.info(f"📤 Preparing PLAINTEXT reply, size: {len(data)/2**20:.2f} MB")
+            return data
+
+    def process_message(self, message: Message) -> Union[str, bytes]:
+        """Process normal Flower message and return reply."""
+        logger.info(f"Processing message with metadata: {message.metadata}")
+        reply_message = self.client_app(message=message, context=self.context)
+        reply_bytes = flower_message_to_bytes(reply_message)
+        return self.prepare_reply(reply_bytes)
+
+    def create_error_reply(
+        self, message: Optional[Message], error: Error
+    ) -> Union[str, bytes]:
+        """Create error reply message."""
+        error_reply = create_flwr_message(
+            content=RecordDict(),
+            reply_to=message,
+            message_type=message.metadata.message_type
+            if message
+            else MessageType.SYSTEM,
+            dst_node_id=message.metadata.src_node_id if message else 0,
+            group_id=message.metadata.group_id if message else "",
+            error=error,
+        )
+        error_bytes = flower_message_to_bytes(error_reply)
+        logger.info(f"Error reply size: {len(error_bytes)/2**20:.2f} MB")
+        return self.prepare_reply(error_bytes)
+
+
+class RequestProcessor:
+    """Processes incoming requests and handles encryption/decryption."""
+
+    def __init__(
+        self, message_handler: MessageHandler, box: SyftEvents, client_email: str
+    ):
+        self.message_handler = message_handler
+        self.box = box
+        self.client_email = client_email
+
+    def decode_request_body(self, request_body: Union[bytes, str]) -> bytes:
+        """Decode request body, handling base64 if encrypted."""
+        if not self.message_handler.encryption_enabled:
+            return request_body
+
+        try:
+            # Convert to string if bytes
+            if isinstance(request_body, bytes):
+                request_body_str = request_body.decode("utf-8")
+            else:
+                request_body_str = request_body
+            # Decode base64
+            decoded = base64.b64decode(request_body_str)
+            logger.debug("🔓 Decoded base64 message")
+            return decoded
+        except Exception:
+            # Not base64 or decoding failed, use as-is
+            return request_body
+
+    def is_stop_signal(self, message: Message) -> bool:
+        """Check if message is a stop signal."""
+        if message.metadata.message_type != MessageType.SYSTEM:
+            return False
+
+        # Check for stop action in config
+        if "config" in message.content and "action" in message.content["config"]:
+            return message.content["config"]["action"] == "stop"
+
+        # Alternative stop signal format
+        return message.metadata.group_id == "final"
+
+    def process(self, request: Request) -> Optional[Union[str, bytes]]:
+        """Process incoming request and return response."""
+        original_sender = request.headers.get("X-Syft-Original-Sender", "unknown")
+        encryption_status = (
+            "🔐 ENCRYPTED"
+            if self.message_handler.encryption_enabled
+            else "📥 PLAINTEXT"
+        )
+
+        logger.info(
+            f"{encryption_status} Received request from {original_sender}, "
+            f"id: {request.id}, size: {len(request.body) / 1024 / 1024:.2f} MB"
+        )
+
+        # Parse message
+        try:
+            request_body = self.decode_request_body(request.body)
+            message = bytes_to_flower_message(request_body)
+
+            if self.message_handler.encryption_enabled:
+                logger.debug(
+                    f"🔓 Successfully decrypted message from {original_sender}"
+                )
+        except Exception as e:
+            logger.error(
+                f"❌ Failed to deserialize message from {original_sender}: {e}"
+            )
+            logger.debug(
+                f"Request body preview (first 200 bytes): {str(request.body[:200])}"
+            )
+
+            # Can't create error reply without valid message - skip response
+            logger.warning(
+                "Skipping error reply (cannot create without valid parsed message)"
+            )
+            return None
+
+        # Handle message
+        try:
+            # Check for stop signal
+            if self.is_stop_signal(message):
+                logger.info("Received stop signal")
+                self.box._stop_event.set()
+                return None
+
+            # Process normal message
+            return self.message_handler.process_message(message)
+
+        except Exception as e:
+            error_message = f"Client: '{self.client_email}'. Error: {str(e)}. Traceback: {traceback.format_exc()}"
+            logger.error(error_message)
+
+            error = Error(
+                code=ErrorCode.CLIENT_APP_RAISED_EXCEPTION, reason=error_message
+            )
+            return self.message_handler.create_error_reply(message, error)
+
+
+def syftbox_flwr_client(client_app: ClientApp, context: Context, app_name: str):
+    """Run the Flower ClientApp with SyftBox."""
+    # Setup
+    client, encryption_enabled, syft_flwr_app_name = setup_client(app_name)
+    box = SyftEvents(
+        app_name=syft_flwr_app_name,
+        client=client,
+        cleanup_expiry="1d",  # Keep request/response files for 1 days
+        cleanup_interval="1d",  # Run cleanup daily
+    )
+
+    logger.info(f"Started SyftBox Flower Client on: {box.client.email}")
+    logger.info(f"syft_flwr app name: {syft_flwr_app_name}")
+
+    # Check if cleanup is running
+    if box.is_cleanup_running():
+        logger.info("Cleanup service is active")
+
+    # Create handlers
+    message_handler = MessageHandler(client_app, context, encryption_enabled)
+    processor = RequestProcessor(message_handler, box, box.client.email)
+
+    # Register message handler
+    @box.on_request(
+        "/messages", auto_decrypt=encryption_enabled, encrypt_reply=encryption_enabled
+    )
+    def handle_messages(request: Request) -> Optional[Union[str, bytes]]:
+        return processor.process(request)
+
+    # Run
+    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,50 @@
+import traceback
+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
+from syft_flwr.utils import setup_client
+
+
+def syftbox_flwr_server(
+    server_app: ServerApp,
+    context: Context,
+    datasites: list[str],
+    app_name: str,
+) -> Context:
+    """Run the Flower ServerApp with SyftBox."""
+    client, _, syft_flwr_app_name = setup_client(app_name)
+
+    # Construct the SyftGrid
+    syft_grid = SyftGrid(
+        app_name=syft_flwr_app_name, datasites=datasites, client=client
+    )
+
+    # Set the run id (random for now)
+    run_id = randint(0, 1000)
+    syft_grid.set_run(run_id)
+
+    logger.info(f"Started SyftBox Flower Server on: {syft_grid._client.email}")
+    logger.info(f"syft_flwr app name: {syft_flwr_app_name}")
+
+    try:
+        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}")
+    except Exception as e:
+        logger.error(f"Server encountered an error: {str(e)}")
+        logger.error(f"Traceback: {traceback.format_exc()}")
+        updated_context = context
+    finally:
+        syft_grid.send_stop_signal(group_id="final", reason="Server stopped")
+        logger.info("Sending stop signals to the clients")
+
+    return updated_context