syft-flwr 0.1.0__py3-none-any.whl

syft_flwr/grid.py

@@ -0,0 +1,171 @@
+import time
+from typing import Iterable, cast
+
+from flwr.common.message import Message
+from flwr.common.typing import Run
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
+from loguru import logger
+from syft_core import Client
+from syft_rpc import rpc, rpc_db
+from typing_extensions import Optional
+
+from syft_flwr.flwr_compatibility import (
+    Grid,
+    RecordDict,
+    check_reply_to_field,
+    create_flwr_message,
+)
+from syft_flwr.serde import bytes_to_flower_message, flower_message_to_bytes
+from syft_flwr.utils import str_to_int
+
+# this is what superlink super node do
+AGGREGATOR_NODE_ID = 1
+
+
+class SyftGrid(Grid):
+    def __init__(
+        self,
+        datasites: list[str] = [],
+        client: Client = None,
+    ) -> None:
+        self._client = Client.load() if client is None else client
+        self._run: Optional[Run] = None
+        self.node = Node(node_id=AGGREGATOR_NODE_ID)
+        self.datasites = datasites
+        self.client_map = {str_to_int(ds): ds for ds in self.datasites}
+        logger.debug(
+            f"Initialize SyftGrid for '{self._client.email}' with datasites: {self.datasites}"
+        )
+
+    def set_run(self, run_id: int) -> None:
+        # TODO: In Grpc Grid case, the superlink is the one which sets up the run id,
+        # do we need to do the same here, where the run id is set from an external context.
+
+        # Convert to Flower Run object
+        self._run = Run.create_empty(run_id)
+
+    @property
+    def run(self) -> Run:
+        """Run ID."""
+        return Run(**vars(cast(Run, self._run)))
+
+    def _check_message(self, message: Message) -> None:
+        # Check if the message is valid
+        if not (
+            message.metadata.run_id == cast(Run, self._run).run_id
+            and message.metadata.src_node_id == self.node.node_id
+            and message.metadata.message_id == ""
+            and check_reply_to_field(message.metadata)
+            and message.metadata.ttl > 0
+        ):
+            logger.debug(f"Invalid message with metadata: {message.metadata}")
+            raise ValueError(f"Invalid message: {message}")
+
+    def create_message(
+        self,
+        content: RecordDict,
+        message_type: str,
+        dst_node_id: int,
+        group_id: str,
+        ttl: Optional[float] = None,
+    ) -> Message:
+        """Create a new message with specified parameters."""
+        return create_flwr_message(
+            content=content,
+            message_type=message_type,
+            dst_node_id=dst_node_id,
+            group_id=group_id,
+            ttl=ttl,
+            run_id=cast(Run, self._run).run_id,
+            src_node_id=self.node.node_id,
+        )
+
+    def get_node_ids(self) -> list[int]:
+        """Get node IDs of all connected nodes."""
+        # it is map from datasites to node id
+        return list(self.client_map.keys())
+
+    def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
+        """Push messages to specified node IDs."""
+        # Construct Messages
+        run_id = cast(Run, self._run).run_id
+        message_ids = []
+        for msg in messages:
+            # Set metadata
+            msg.metadata.__dict__["_run_id"] = run_id
+            msg.metadata.__dict__["_src_node_id"] = self.node.node_id
+            # RPC URL
+            dest_datasite = self.client_map[msg.metadata.dst_node_id]
+            url = rpc.make_url(dest_datasite, app_name="flwr", endpoint="messages")
+            # Check message
+            self._check_message(msg)
+            # Serialize message
+            msg_bytes = flower_message_to_bytes(msg)
+            # Send message
+            future = rpc.send(url=url, body=msg_bytes, client=self._client)
+            logger.debug(
+                f"Pushed message to {url} with metadata {msg.metadata}; size {len(msg_bytes) / 1024 / 1024} (Mb)"
+            )
+            # Save future
+            rpc_db.save_future(future=future, namespace="flwr", client=self._client)
+            message_ids.append(future.id)
+
+        return message_ids
+
+    def pull_messages(self, message_ids):
+        """Pull messages based on message IDs."""
+        messages = {}
+
+        for msg_id in message_ids:
+            future = rpc_db.get_future(future_id=msg_id, client=self._client)
+            response = future.resolve()
+            if response is None:
+                continue
+
+            response.raise_for_status()
+
+            if not response.body:
+                raise ValueError(f"Empty response: {response}")
+
+            message: Message = bytes_to_flower_message(response.body)
+            if message.has_error():
+                error = message.error
+                logger.error(
+                    f"Message {msg_id} error with code={error.code}, reason={error.reason}"
+                )
+                continue
+
+            logger.debug(
+                f"Pulled message from {response.url} with metadata: {message.metadata}, size: {len(response.body) / 1024 / 1024} (Mb)"
+            )
+            messages[msg_id] = message
+            rpc_db.delete_future(future_id=msg_id, client=self._client)
+
+        return messages
+
+    def send_and_receive(
+        self,
+        messages: Iterable[Message],
+        *,
+        timeout: Optional[float] = None,
+    ) -> Iterable[Message]:
+        """Push messages to specified node IDs and pull the reply messages.
+
+        This method sends a list of messages to their destination node IDs and then
+        waits for the replies. It continues to pull replies until either all replies are
+        received or the specified timeout duration is exceeded.
+        """
+        # Push messages
+        msg_ids = set(self.push_messages(messages))
+
+        # Pull messages
+        end_time = time.time() + (timeout if timeout is not None else 0.0)
+        ret = {}
+        while timeout is None or time.time() < end_time:
+            res_msgs = self.pull_messages(msg_ids)
+            ret.update(res_msgs)
+            msg_ids.difference_update(res_msgs.keys())
+            if len(msg_ids) == 0:
+                break
+            time.sleep(3)
+        return ret.values()

syft_flwr/run.py

@@ -0,0 +1,61 @@
+import warnings
+from pathlib import Path
+from uuid import uuid4
+
+from flwr.client.client_app import LoadClientAppError
+from flwr.common import Context
+from flwr.common.object_ref import load_app
+from flwr.server.server_app import LoadServerAppError
+
+from syft_flwr.config import load_flwr_pyproject
+from syft_flwr.flower_client import syftbox_flwr_client
+from syft_flwr.flower_server import syftbox_flwr_server
+from syft_flwr.flwr_compatibility import RecordDict
+from syft_flwr.run_simulation import run
+
+__all__ = ["syftbox_run_flwr_client", "syftbox_run_flwr_server", "run"]
+
+
+# Suppress Pydantic deprecation warnings
+warnings.filterwarnings("ignore", category=DeprecationWarning, module="pydantic")
+
+
+def syftbox_run_flwr_client(flower_project_dir: Path) -> None:
+    pyproject_conf = load_flwr_pyproject(flower_project_dir)
+    client_ref = pyproject_conf["tool"]["flwr"]["app"]["components"]["clientapp"]
+
+    context = Context(
+        run_id=uuid4().int,
+        node_id=uuid4().int,
+        node_config=pyproject_conf["tool"]["flwr"]["app"]["config"],
+        state=RecordDict(),
+        run_config=pyproject_conf["tool"]["flwr"]["app"]["config"],
+    )
+    client_app = load_app(
+        client_ref,
+        LoadClientAppError,
+        flower_project_dir,
+    )
+
+    syftbox_flwr_client(client_app, context)
+
+
+def syftbox_run_flwr_server(flower_project_dir: Path) -> None:
+    pyproject_conf = load_flwr_pyproject(flower_project_dir)
+    datasites = pyproject_conf["tool"]["syft_flwr"]["datasites"]
+    server_ref = pyproject_conf["tool"]["flwr"]["app"]["components"]["serverapp"]
+
+    context = Context(
+        run_id=uuid4().int,
+        node_id=uuid4().int,
+        node_config=pyproject_conf["tool"]["flwr"]["app"]["config"],
+        state=RecordDict(),
+        run_config=pyproject_conf["tool"]["flwr"]["app"]["config"],
+    )
+    server_app = load_app(
+        server_ref,
+        LoadServerAppError,
+        flower_project_dir,
+    )
+
+    syftbox_flwr_server(server_app, context, datasites)

syft_flwr/run_simulation.py

@@ -0,0 +1,204 @@
+import asyncio
+import os
+import uuid
+from pathlib import Path
+
+from loguru import logger
+from syft_rds.client.rds_client import RDSClient
+from syft_rds.orchestra import remove_rds_stack_dir, setup_rds_server
+from typing_extensions import Union
+
+from syft_flwr.config import load_flwr_pyproject
+
+
+def _setup_mock_rds_clients(
+    project_dir: Path, aggregator: str, datasites: list[str]
+) -> tuple[str, list[RDSClient], RDSClient]:
+    """Setup mock RDS clients for the given project directory"""
+    key = project_dir.name + "_" + str(uuid.uuid4())
+    remove_rds_stack_dir(key)
+
+    ds_stack = setup_rds_server(email=aggregator, key=key)
+    ds_client = ds_stack.init_session(host=aggregator)
+
+    do_clients = []
+    for datasite in datasites:
+        do_stack = setup_rds_server(email=datasite, key=key)
+        do_client = do_stack.init_session(host=datasite)
+        do_clients.append(do_client)
+
+    return key, do_clients, ds_client
+
+
+async def _run_main_py(
+    main_py_path: Path,
+    config_path: Path,
+    client_email: str,
+    log_dir: Path,
+    dataset_path: Union[str, Path] | None = None,
+) -> int:
+    """Run the `main.py` file for a given client"""
+    log_file_path = log_dir / f"{client_email}.log"
+
+    # setting up env variables
+    env = os.environ.copy()
+    env["SYFTBOX_CLIENT_CONFIG_PATH"] = str(config_path)
+    env["DATA_DIR"] = str(dataset_path)
+
+    # running the main.py file asynchronously in a subprocess
+    try:
+        with open(log_file_path, "w") as f:
+            process = await asyncio.create_subprocess_exec(
+                "python",
+                str(main_py_path),
+                "-s",
+                stdout=f,
+                stderr=f,
+                env=env,
+            )
+            return_code = await process.wait()
+            logger.debug(
+                f"`{client_email}` returns code {return_code} for running `{main_py_path}`"
+            )
+            return return_code
+    except Exception as e:
+        logger.error(f"Error running `{main_py_path}` for `{client_email}`: {e}")
+        return 1
+
+
+async def _run_simulated_flwr_project(
+    project_dir: Path,
+    do_clients: list[RDSClient],
+    ds_client: RDSClient,
+    mock_dataset_paths: list[Union[str, Path]],
+) -> bool:
+    """Run all clients and server concurrently"""
+    run_success = True
+
+    log_dir = project_dir / "simulation_logs"
+    log_dir.mkdir(parents=True, exist_ok=True)
+    logger.info(f"📝 Log directory: {log_dir}")
+
+    main_py_path = project_dir / "main.py"
+
+    logger.info(
+        f"Running DS client '{ds_client.email}' with config path {ds_client._syftbox_client.config_path}"
+    )
+    ds_task: asyncio.Task = asyncio.create_task(
+        _run_main_py(
+            main_py_path,
+            ds_client._syftbox_client.config_path,
+            ds_client.email,
+            log_dir,
+        )
+    )
+
+    client_tasks: list[asyncio.Task] = []
+    for client, mock_dataset_path in zip(do_clients, mock_dataset_paths):
+        # check if the client has a mock dataset path
+        logger.info(
+            f"Running DO client '{client.email}' with config path {client._syftbox_client.config_path} on mock dataset {mock_dataset_path}"
+        )
+        client_tasks.append(
+            asyncio.create_task(
+                _run_main_py(
+                    main_py_path,
+                    client._syftbox_client.config_path,
+                    client.email,
+                    log_dir,
+                    mock_dataset_path,
+                )
+            )
+        )
+
+    ds_return_code = await ds_task
+    if ds_return_code != 0:
+        run_success = False
+
+    # log out ds client logs
+    with open(log_dir / f"{ds_client.email}.log", "r") as log_file:
+        log_content = log_file.read().strip()
+        logger.info(f"DS client '{ds_client.email}' logs:\n{log_content}")
+
+    # cancel all client tasks if DS client returns
+    logger.debug("Cancelling DO client tasks as DS client returned")
+    for task in client_tasks:
+        if not task.done():
+            task.cancel()
+
+    await asyncio.gather(*client_tasks, return_exceptions=True)
+
+    return run_success
+
+
+def _validate_bootstraped_project(project_dir: Path) -> None:
+    """Validate a bootstraped `syft_flwr` project directory"""
+    if not project_dir.exists():
+        raise FileNotFoundError(f"Project directory {project_dir} does not exist")
+
+    if not project_dir.is_dir():
+        raise NotADirectoryError(f"Project directory {project_dir} is not a directory")
+
+    if not (project_dir / "main.py").exists():
+        raise FileNotFoundError(f"main.py not found at {project_dir}")
+
+    if not (project_dir / "pyproject.toml").exists():
+        raise FileNotFoundError(f"pyproject.toml not found at {project_dir}")
+
+
+def _validate_mock_dataset_paths(mock_dataset_paths: list[str]) -> list[Path]:
+    """Validate the mock dataset paths"""
+    resolved_paths = []
+    for path in mock_dataset_paths:
+        path = Path(path).expanduser().resolve()
+        if not path.exists():
+            raise ValueError(f"Mock dataset path {path} does not exist")
+        resolved_paths.append(path)
+    return resolved_paths
+
+
+def run(
+    project_dir: Union[str, Path], mock_dataset_paths: list[Union[str, Path]]
+) -> None:
+    """Run a syft_flwr project in simulation mode over mock data"""
+
+    project_dir = Path(project_dir).expanduser().resolve()
+    _validate_bootstraped_project(project_dir)
+    mock_dataset_paths = _validate_mock_dataset_paths(mock_dataset_paths)
+
+    pyproject_conf = load_flwr_pyproject(project_dir)
+    datasites = pyproject_conf["tool"]["syft_flwr"]["datasites"]
+    aggregator = pyproject_conf["tool"]["syft_flwr"]["aggregator"]
+
+    key, do_clients, ds_client = _setup_mock_rds_clients(
+        project_dir, aggregator, datasites
+    )
+
+    async def main():
+        try:
+            run_success = await _run_simulated_flwr_project(
+                project_dir, do_clients, ds_client, mock_dataset_paths
+            )
+            if run_success:
+                logger.success("Simulation completed successfully ✅")
+            else:
+                logger.error("Simulation failed ❌")
+        except Exception as e:
+            logger.error(f"Simulation failed ❌: {e}")
+        finally:
+            # Clean up the RDS stack
+            remove_rds_stack_dir(key)
+            logger.debug(f"Removed RDS stack: {key}")
+
+    try:
+        loop = asyncio.get_running_loop()
+        logger.debug(f"Running in an environment with an existing event loop {loop}")
+        # We are in an environment with an existing event loop (like Jupyter)
+        asyncio.create_task(main())
+    except RuntimeError:
+        logger.debug("No existing event loop, creating and running one")
+        # No existing event loop, create and run one (for scripts)
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        loop.run_until_complete(main())
+        loop.close()

syft_flwr/serde.py

@@ -0,0 +1,15 @@
+from flwr.common.message import Message
+from flwr.common.serde import message_from_proto, message_to_proto
+from flwr.proto.message_pb2 import Message as ProtoMessage
+
+
+def bytes_to_flower_message(data: bytes) -> Message:
+    message_pb = ProtoMessage()
+    message_pb.ParseFromString(data)
+    message = message_from_proto(message_pb)
+    return message
+
+
+def flower_message_to_bytes(message: Message) -> bytes:
+    msg_proto = message_to_proto(message)
+    return msg_proto.SerializeToString()

syft_flwr/strategy/__init__.py

@@ -0,0 +1,3 @@
+from syft_flwr.strategy.fedavg import FedAvgWithModelSaving
+
+__all__ = ["FedAvgWithModelSaving"]

syft_flwr/strategy/fedavg.py

@@ -0,0 +1,33 @@
+from pathlib import Path
+
+from flwr.common import parameters_to_ndarrays
+from flwr.server.strategy import FedAvg
+from loguru import logger
+from safetensors.numpy import save_file
+
+
+class FedAvgWithModelSaving(FedAvg):
+    """This is a custom strategy that behaves exactly like
+    FedAvg with the difference of storing of the state of
+    the global model to disk after each round.
+    Ref: https://discuss.flower.ai/t/how-do-i-save-the-global-model-after-training/71/2
+    """
+
+    def __init__(self, save_path: str, *args, **kwargs):
+        self.save_path = Path(save_path)
+        self.save_path.mkdir(exist_ok=True, parents=True)
+        super().__init__(*args, **kwargs)
+
+    def _save_global_model(self, server_round: int, parameters):
+        """A new method to save the parameters to disk."""
+        ndarrays = parameters_to_ndarrays(parameters)
+        tensor_dict = {f"layer_{i}": array for i, array in enumerate(ndarrays)}
+        filename = self.save_path / f"parameters_round_{server_round}.safetensors"
+        save_file(tensor_dict, str(filename))
+
+        logger.info(f"Checkpoint saved to: {filename}")
+
+    def evaluate(self, server_round: int, parameters):
+        """Evaluate model parameters using an evaluation function."""
+        self._save_global_model(server_round, parameters)
+        return super().evaluate(server_round, parameters)

syft_flwr/templates/main.py.tpl

@@ -0,0 +1,27 @@
+import os
+from pathlib import Path
+
+from syft_core import Client
+
+from syft_flwr.config import load_flwr_pyproject
+from syft_flwr.run import syftbox_run_flwr_client, syftbox_run_flwr_server
+
+DATA_DIR = os.getenv("DATA_DIR")
+OUTPUT_DIR = os.getenv("OUTPUT_DIR")
+
+
+flower_project_dir = Path(__file__).parent.absolute()
+client = Client.load()
+config = load_flwr_pyproject(flower_project_dir)
+
+is_client = client.email in config["tool"]["syft_flwr"]["datasites"]
+is_server = client.email in config["tool"]["syft_flwr"]["aggregator"]
+
+if is_client:
+    # run by each DO
+    syftbox_run_flwr_client(flower_project_dir)
+elif is_server:
+    # run by the DS
+    syftbox_run_flwr_server(flower_project_dir)
+else:
+    raise ValueError(f"{client.email} is not in config.datasites or config.aggregator")

syft_flwr/utils.py

@@ -0,0 +1,36 @@
+import os
+import re
+import zlib
+from pathlib import Path
+
+EMAIL_REGEX = r"^[^@]+@[^@]+\.[^@]+$"
+
+
+def is_valid_datasite(datasite: str) -> bool:
+    return re.match(EMAIL_REGEX, datasite)
+
+
+def str_to_int(input_string: str) -> int:
+    """Convert a string to an int32"""
+    return zlib.crc32(input_string.encode())
+
+
+def get_syftbox_dataset_path() -> Path:
+    """Get the path to the syftbox dataset from the environment variable"""
+    data_dir = Path(os.getenv("DATA_DIR", ".data/"))
+    if not data_dir.exists():
+        raise FileNotFoundError(
+            f"Path {data_dir} does not exist (must be a valid file or directory)"
+        )
+    return data_dir
+
+
+def run_syft_flwr() -> bool:
+    """Util function to check if we are running with syft_flwr or plain flwr
+    Currently only checks the `DATA_DIR` environment variable.
+    """
+    try:
+        get_syftbox_dataset_path()
+        return True
+    except FileNotFoundError:
+        return False

syft_flwr-0.1.0.dist-info/METADATA

@@ -0,0 +1,25 @@
+Metadata-Version: 2.4
+Name: syft-flwr
+Version: 0.1.0
+Summary: Add your description here
+License-File: LICENSE
+Requires-Python: >=3.9.2
+Requires-Dist: flwr-datasets[vision]>=0.5.0
+Requires-Dist: flwr[simulation]==1.17.0
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: safetensors>=0.5.0
+Requires-Dist: syft-core==0.2.3
+Requires-Dist: syft-event==0.2.0
+Requires-Dist: syft-rds==0.1.0
+Requires-Dist: syft-rpc==0.2.0
+Requires-Dist: tomli-w>=1.2.0
+Requires-Dist: tomli>=2.2.1
+Requires-Dist: typing-extensions>=4.13.0
+Description-Content-Type: text/markdown
+
+# 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`

syft_flwr-0.1.0.dist-info/RECORD

@@ -0,0 +1,22 @@
+README.md,sha256=rxa8da_OJlI6RVDAXa3tfmZC_dGEN_VFPJXJNMfXl-Q,250
+pyproject.toml,sha256=k7335xXFv_VDmvxmIk_wy6LriaTc92j0th88WDjJ-F0,1148
+syft_flwr/__init__.py,sha256=Tm2mxtuCAXctyyi89c9psk0PuGdv1TOauCg7-gFPyRk,127
+syft_flwr/bootstrap.py,sha256=P6bDLSnCQqrbxQ6-QBc0sR75L5ods8oLKJqItZzNlao,3446
+syft_flwr/cli.py,sha256=imctwdQMxQeGQZaiKSX1Mo2nU_-RmA-cGB3H4huuUeA,3274
+syft_flwr/config.py,sha256=4hwkovGtFOLNULjJwoGYcA0uT4y3vZSrxndXqYXquMY,821
+syft_flwr/flower_client.py,sha256=up15MSEBExiZPgcXHo1EIflWaVN0nO6y08gXrzpumSM,2418
+syft_flwr/flower_server.py,sha256=dz7dvEgFm76s8G266RfN46TohnMPGAk8K6DPWb23LsE,784
+syft_flwr/flwr_compatibility.py,sha256=vURf9rfsZ1uPm04szw6RpGYxtlG3BE4tW3YijptiGyk,3197
+syft_flwr/grid.py,sha256=76Wc1O1W8iGMr9QS97o9wTaU_87UB4fqYA5HnlIHy_U,6087
+syft_flwr/run.py,sha256=gLm_zjzCuWVNKVK2ocJzx_izXDJAW411nl6IXDyrkFM,2008
+syft_flwr/run_simulation.py,sha256=t3shhfzAWDUlf6iQmwf5sS9APZQE9mkaZ9MLCYs9Ng0,6922
+syft_flwr/serde.py,sha256=5fCI-cRUOh5wE7cXQd4J6jex1grRGnyD1Jx-VlEDOXM,495
+syft_flwr/utils.py,sha256=3dDYEB7btq9hxZ9UsfQWh3i44OerAhGXc5XaX5wO3-o,955
+syft_flwr/strategy/__init__.py,sha256=mpUmExjjFkqU8gg41XsOBKfO3aqCBe7XPJSU-_P7smU,97
+syft_flwr/strategy/fedavg.py,sha256=UCgXL_0woRjs7iHaqN16UYNJ7J6ogY3sO6620g1W3sU,1364
+syft_flwr/templates/main.py.tpl,sha256=p0uK97jvLGk3LJdy1_HF1R5BQgIjaTGkYnr-csfh39M,791
+syft_flwr-0.1.0.dist-info/METADATA,sha256=wnKYSmLIbSoKAd7A1Wk7PatwD_g5TuByahxzJYJitcc,799
+syft_flwr-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+syft_flwr-0.1.0.dist-info/entry_points.txt,sha256=o7oT0dCoHn-3WyIwdDw1lBh2q-GvhB_8s0hWeJU4myc,49
+syft_flwr-0.1.0.dist-info/licenses/LICENSE,sha256=0msOUar8uPZTqkAOTBp4rCzd7Jl9eRhfKiNufwrsg7k,11361
+syft_flwr-0.1.0.dist-info/RECORD,,

syft_flwr-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

syft_flwr-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+syft_flwr = syft_flwr.cli:main