flwr-nightly 1.10.0.dev20240705__py3-none-any.whl → 1.10.0.dev20240708__py3-none-any.whl

flwr/server/driver/grpc_driver.py

@@ -16,8 +16,8 @@
 
 import time
 import warnings
-from logging import DEBUG, ERROR, WARNING
-from typing import Iterable, List, Optional, Tuple, cast
+from logging import DEBUG, WARNING
+from typing import Iterable, List, Optional, cast
 
 import grpc
 
@@ -27,8 +27,6 @@ from flwr.common.logger import log
 from flwr.common.serde import message_from_taskres, message_to_taskins
 from flwr.common.typing import Run
 from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
-    CreateRunRequest,
-    CreateRunResponse,
     GetNodesRequest,
     GetNodesResponse,
     PullTaskResRequest,
@@ -53,167 +51,103 @@ Call `connect()` on the `GrpcDriverStub` instance before calling any of the othe
 """
 
 
-class GrpcDriverStub:
-    """`GrpcDriverStub` provides access to the gRPC Driver API/service.
+class GrpcDriver(Driver):
+    """`GrpcDriver` provides an interface to the Driver API.
 
     Parameters
     ----------
-    driver_service_address : Optional[str]
-        The IPv4 or IPv6 address of the Driver API server.
-        Defaults to `"[::]:9091"`.
+    run_id : int
+        The identifier of the run.
+    driver_service_address : str (default: "[::]:9091")
+        The address (URL, IPv6, IPv4) of the SuperLink Driver API service.
     root_certificates : Optional[bytes] (default: None)
         The PEM-encoded root certificates as a byte string.
         If provided, a secure connection using the certificates will be
         established to an SSL-enabled Flower server.
     """
 
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
+        run_id: int,
         driver_service_address: str = DEFAULT_SERVER_ADDRESS_DRIVER,
         root_certificates: Optional[bytes] = None,
     ) -> None:
-        self.driver_service_address = driver_service_address
-        self.root_certificates = root_certificates
-        self.channel: Optional[grpc.Channel] = None
-        self.stub: Optional[DriverStub] = None
+        self._run_id = run_id
+        self._addr = driver_service_address
+        self._cert = root_certificates
+        self._run: Optional[Run] = None
+        self._grpc_stub: Optional[DriverStub] = None
+        self._channel: Optional[grpc.Channel] = None
+        self.node = Node(node_id=0, anonymous=True)
+
+    @property
+    def _is_connected(self) -> bool:
+        """Check if connected to the Driver API server."""
+        return self._channel is not None
 
-    def is_connected(self) -> bool:
-        """Return True if connected to the Driver API server, otherwise False."""
-        return self.channel is not None
+    def _connect(self) -> None:
+        """Connect to the Driver API.
 
-    def connect(self) -> None:
-        """Connect to the Driver API."""
+        This will not call GetRun.
+        """
         event(EventType.DRIVER_CONNECT)
-        if self.channel is not None or self.stub is not None:
+        if self._is_connected:
             log(WARNING, "Already connected")
             return
-        self.channel = create_channel(
-            server_address=self.driver_service_address,
-            insecure=(self.root_certificates is None),
-            root_certificates=self.root_certificates,
+        self._channel = create_channel(
+            server_address=self._addr,
+            insecure=(self._cert is None),
+            root_certificates=self._cert,
         )
-        self.stub = DriverStub(self.channel)
-        log(DEBUG, "[Driver] Connected to %s", self.driver_service_address)
+        self._grpc_stub = DriverStub(self._channel)
+        log(DEBUG, "[Driver] Connected to %s", self._addr)
 
-    def disconnect(self) -> None:
+    def _disconnect(self) -> None:
         """Disconnect from the Driver API."""
         event(EventType.DRIVER_DISCONNECT)
-        if self.channel is None or self.stub is None:
+        if not self._is_connected:
             log(DEBUG, "Already disconnected")
             return
-        channel = self.channel
-        self.channel = None
-        self.stub = None
+        channel: grpc.Channel = self._channel
+        self._channel = None
+        self._grpc_stub = None
         channel.close()
         log(DEBUG, "[Driver] Disconnected")
 
-    def create_run(self, req: CreateRunRequest) -> CreateRunResponse:
-        """Request for run ID."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverStub` instance not connected")
-
-        # Call Driver API
-        res: CreateRunResponse = self.stub.CreateRun(request=req)
-        return res
-
-    def get_run(self, req: GetRunRequest) -> GetRunResponse:
-        """Get run information."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverStub` instance not connected")
-
-        # Call gRPC Driver API
-        res: GetRunResponse = self.stub.GetRun(request=req)
-        return res
-
-    def get_nodes(self, req: GetNodesRequest) -> GetNodesResponse:
-        """Get client IDs."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverStub` instance not connected")
-
-        # Call gRPC Driver API
-        res: GetNodesResponse = self.stub.GetNodes(request=req)
-        return res
-
-    def push_task_ins(self, req: PushTaskInsRequest) -> PushTaskInsResponse:
-        """Schedule tasks."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverStub` instance not connected")
-
-        # Call gRPC Driver API
-        res: PushTaskInsResponse = self.stub.PushTaskIns(request=req)
-        return res
-
-    def pull_task_res(self, req: PullTaskResRequest) -> PullTaskResResponse:
-        """Get task results."""
-        # Check if channel is open
-        if self.stub is None:
-            log(ERROR, ERROR_MESSAGE_DRIVER_NOT_CONNECTED)
-            raise ConnectionError("`GrpcDriverStub` instance not connected")
-
-        # Call Driver API
-        res: PullTaskResResponse = self.stub.PullTaskRes(request=req)
-        return res
-
-
-class GrpcDriver(Driver):
-    """`Driver` class provides an interface to the Driver API.
-
-    Parameters
-    ----------
-    run_id : int
-        The identifier of the run.
-    stub : Optional[GrpcDriverStub] (default: None)
-        The ``GrpcDriverStub`` instance used to communicate with the SuperLink.
-        If None, an instance connected to "[::]:9091" will be created.
-    """
-
-    def __init__(  # pylint: disable=too-many-arguments
-        self,
-        run_id: int,
-        stub: Optional[GrpcDriverStub] = None,
-    ) -> None:
-        self._run_id = run_id
-        self._run: Optional[Run] = None
-        self.stub = stub if stub is not None else GrpcDriverStub()
-        self.node = Node(node_id=0, anonymous=True)
+    def _init_run(self) -> None:
+        # Check if is initialized
+        if self._run is not None:
+            return
+        # Get the run info
+        req = GetRunRequest(run_id=self._run_id)
+        res: GetRunResponse = self._stub.GetRun(req)
+        if not res.HasField("run"):
+            raise RuntimeError(f"Cannot find the run with ID: {self._run_id}")
+        self._run = Run(
+            run_id=res.run.run_id,
+            fab_id=res.run.fab_id,
+            fab_version=res.run.fab_version,
+            override_config=dict(res.run.override_config.items()),
+        )
 
     @property
     def run(self) -> Run:
         """Run information."""
-        self._get_stub_and_run_id()
-        return Run(**vars(cast(Run, self._run)))
+        self._init_run()
+        return Run(**vars(self._run))
 
-    def _get_stub_and_run_id(self) -> Tuple[GrpcDriverStub, int]:
-        # Check if is initialized
-        if self._run is None:
-            # Connect
-            if not self.stub.is_connected():
-                self.stub.connect()
-            # Get the run info
-            req = GetRunRequest(run_id=self._run_id)
-            res = self.stub.get_run(req)
-            if not res.HasField("run"):
-                raise RuntimeError(f"Cannot find the run with ID: {self._run_id}")
-            self._run = Run(
-                run_id=res.run.run_id,
-                fab_id=res.run.fab_id,
-                fab_version=res.run.fab_version,
-            )
-
-        return self.stub, self._run.run_id
+    @property
+    def _stub(self) -> DriverStub:
+        """Driver stub."""
+        if not self._is_connected:
+            self._connect()
+        return cast(DriverStub, self._grpc_stub)
 
     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
+            # Assume self._run being initialized
+            message.metadata.run_id == self._run_id
             and message.metadata.src_node_id == self.node.node_id
             and message.metadata.message_id == ""
             and message.metadata.reply_to_message == ""
@@ -234,7 +168,7 @@ class GrpcDriver(Driver):
         This method constructs a new `Message` with given content and metadata.
         The `run_id` and `src_node_id` will be set automatically.
         """
-        _, run_id = self._get_stub_and_run_id()
+        self._init_run()
         if ttl:
             warnings.warn(
                 "A custom TTL was set, but note that the SuperLink does not enforce "
@@ -245,7 +179,7 @@ class GrpcDriver(Driver):
 
         ttl_ = DEFAULT_TTL if ttl is None else ttl
         metadata = Metadata(
-            run_id=run_id,
+            run_id=self._run_id,
             message_id="",  # Will be set by the server
             src_node_id=self.node.node_id,
             dst_node_id=dst_node_id,
@@ -258,9 +192,11 @@ class GrpcDriver(Driver):
 
     def get_node_ids(self) -> List[int]:
         """Get node IDs."""
-        stub, run_id = self._get_stub_and_run_id()
+        self._init_run()
         # Call GrpcDriverStub method
-        res = stub.get_nodes(GetNodesRequest(run_id=run_id))
+        res: GetNodesResponse = self._stub.GetNodes(
+            GetNodesRequest(run_id=self._run_id)
+        )
         return [node.node_id for node in res.nodes]
 
     def push_messages(self, messages: Iterable[Message]) -> Iterable[str]:
@@ -269,7 +205,7 @@ class GrpcDriver(Driver):
         This method takes an iterable of messages and sends each message
         to the node specified in `dst_node_id`.
         """
-        stub, _ = self._get_stub_and_run_id()
+        self._init_run()
         # Construct TaskIns
         task_ins_list: List[TaskIns] = []
         for msg in messages:
@@ -280,7 +216,9 @@ class GrpcDriver(Driver):
             # Add to list
             task_ins_list.append(taskins)
         # Call GrpcDriverStub method
-        res = stub.push_task_ins(PushTaskInsRequest(task_ins_list=task_ins_list))
+        res: PushTaskInsResponse = self._stub.PushTaskIns(
+            PushTaskInsRequest(task_ins_list=task_ins_list)
+        )
         return list(res.task_ids)
 
     def pull_messages(self, message_ids: Iterable[str]) -> Iterable[Message]:
@@ -289,9 +227,9 @@ class GrpcDriver(Driver):
         This method is used to collect messages from the SuperLink that correspond to a
         set of given message IDs.
         """
-        stub, _ = self._get_stub_and_run_id()
+        self._init_run()
         # Pull TaskRes
-        res = stub.pull_task_res(
+        res: PullTaskResResponse = self._stub.PullTaskRes(
             PullTaskResRequest(node=self.node, task_ids=message_ids)
         )
         # Convert TaskRes to Message
@@ -331,7 +269,7 @@ class GrpcDriver(Driver):
     def close(self) -> None:
         """Disconnect from the SuperLink if connected."""
         # Check if `connect` was called before
-        if not self.stub.is_connected():
+        if not self._is_connected:
             return
         # Disconnect
-        self.stub.disconnect()
+        self._disconnect()

flwr/server/run_serverapp.py

@@ -25,10 +25,13 @@ from flwr.common import Context, EventType, RecordSet, event
 from flwr.common.config import get_flwr_dir, get_project_config, get_project_dir
 from flwr.common.logger import log, update_console_handler, warn_deprecated_feature
 from flwr.common.object_ref import load_app
-from flwr.proto.driver_pb2 import CreateRunRequest  # pylint: disable=E0611
+from flwr.proto.driver_pb2 import (  # pylint: disable=E0611
+    CreateRunRequest,
+    CreateRunResponse,
+)
 
 from .driver import Driver
-from .driver.grpc_driver import GrpcDriver, GrpcDriverStub
+from .driver.grpc_driver import GrpcDriver
 from .server_app import LoadServerAppError, ServerApp
 
 ADDRESS_DRIVER_API = "0.0.0.0:9091"
@@ -144,22 +147,27 @@ def run_server_app() -> None:  # pylint: disable=too-many-branches
             "For more details, use: ``flower-server-app -h``"
         )
 
-    stub = GrpcDriverStub(
-        driver_service_address=args.superlink, root_certificates=root_certificates
-    )
+    # Initialize GrpcDriver
     if args.run_id is not None:
         # User provided `--run-id`, but not `server-app`
-        run_id = args.run_id
+        driver = GrpcDriver(
+            run_id=args.run_id,
+            driver_service_address=args.superlink,
+            root_certificates=root_certificates,
+        )
     else:
         # User provided `server-app`, but not `--run-id`
         # Create run if run_id is not provided
-        stub.connect()
+        driver = GrpcDriver(
+            run_id=0,  # Will be overwritten
+            driver_service_address=args.superlink,
+            root_certificates=root_certificates,
+        )
+        # Create run
         req = CreateRunRequest(fab_id=args.fab_id, fab_version=args.fab_version)
-        res = stub.create_run(req)
-        run_id = res.run_id
-
-    # Initialize GrpcDriver
-    driver = GrpcDriver(run_id=run_id, stub=stub)
+        res: CreateRunResponse = driver._stub.CreateRun(req)  # pylint: disable=W0212
+        # Overwrite driver._run_id
+        driver._run_id = res.run_id  # pylint: disable=W0212
 
     # Dynamically obtain ServerApp path based on run_id
     if args.run_id is not None:

flwr/server/superlink/driver/driver_servicer.py

@@ -69,7 +69,11 @@ class DriverServicer(driver_pb2_grpc.DriverServicer):
         """Create run ID."""
         log(DEBUG, "DriverServicer.CreateRun")
         state: State = self.state_factory.state()
-        run_id = state.create_run(request.fab_id, request.fab_version)
+        run_id = state.create_run(
+            request.fab_id,
+            request.fab_version,
+            dict(request.override_config.items()),
+        )
         return CreateRunResponse(run_id=run_id)
 
     def PushTaskIns(

flwr/server/superlink/state/in_memory_state.py

@@ -275,7 +275,12 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
         """Retrieve stored `node_id` filtered by `client_public_keys`."""
         return self.public_key_to_node_id.get(client_public_key)
 
-    def create_run(self, fab_id: str, fab_version: str) -> int:
+    def create_run(
+        self,
+        fab_id: str,
+        fab_version: str,
+        override_config: Dict[str, str],
+    ) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id
         with self.lock:
@@ -283,7 +288,10 @@ class InMemoryState(State):  # pylint: disable=R0902,R0904
 
             if run_id not in self.run_ids:
                 self.run_ids[run_id] = Run(
-                    run_id=run_id, fab_id=fab_id, fab_version=fab_version
+                    run_id=run_id,
+                    fab_id=fab_id,
+                    fab_version=fab_version,
+                    override_config=override_config,
                 )
                 return run_id
         log(ERROR, "Unexpected run creation failure.")

flwr/server/superlink/state/sqlite_state.py

@@ -15,6 +15,7 @@
 """SQLite based implemenation of server state."""
 
 
+import json
 import re
 import sqlite3
 import time
@@ -61,9 +62,10 @@ CREATE INDEX IF NOT EXISTS idx_online_until ON node (online_until);
 
 SQL_CREATE_TABLE_RUN = """
 CREATE TABLE IF NOT EXISTS run(
-    run_id          INTEGER UNIQUE,
-    fab_id          TEXT,
-    fab_version     TEXT
+    run_id                INTEGER UNIQUE,
+    fab_id                TEXT,
+    fab_version           TEXT,
+    override_config       TEXT
 );
 """
 
@@ -613,7 +615,12 @@ class SqliteState(State):  # pylint: disable=R0904
             return node_id
         return None
 
-    def create_run(self, fab_id: str, fab_version: str) -> int:
+    def create_run(
+        self,
+        fab_id: str,
+        fab_version: str,
+        override_config: Dict[str, str],
+    ) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
         # Sample a random int64 as run_id
         run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
@@ -622,8 +629,13 @@ class SqliteState(State):  # pylint: disable=R0904
         query = "SELECT COUNT(*) FROM run WHERE run_id = ?;"
         # If run_id does not exist
         if self.query(query, (run_id,))[0]["COUNT(*)"] == 0:
-            query = "INSERT INTO run (run_id, fab_id, fab_version) VALUES (?, ?, ?);"
-            self.query(query, (run_id, fab_id, fab_version))
+            query = (
+                "INSERT INTO run (run_id, fab_id, fab_version, override_config)"
+                "VALUES (?, ?, ?, ?);"
+            )
+            self.query(
+                query, (run_id, fab_id, fab_version, json.dumps(override_config))
+            )
             return run_id
         log(ERROR, "Unexpected run creation failure.")
         return 0
@@ -687,7 +699,10 @@ class SqliteState(State):  # pylint: disable=R0904
         try:
             row = self.query(query, (run_id,))[0]
             return Run(
-                run_id=run_id, fab_id=row["fab_id"], fab_version=row["fab_version"]
+                run_id=run_id,
+                fab_id=row["fab_id"],
+                fab_version=row["fab_version"],
+                override_config=json.loads(row["override_config"]),
             )
         except sqlite3.IntegrityError:
             log(ERROR, "`run_id` does not exist.")

flwr/server/superlink/state/state.py

@@ -16,7 +16,7 @@
 
 
 import abc
-from typing import List, Optional, Set
+from typing import Dict, List, Optional, Set
 from uuid import UUID
 
 from flwr.common.typing import Run
@@ -157,7 +157,12 @@ class State(abc.ABC):  # pylint: disable=R0904
         """Retrieve stored `node_id` filtered by `client_public_keys`."""
 
     @abc.abstractmethod
-    def create_run(self, fab_id: str, fab_version: str) -> int:
+    def create_run(
+        self,
+        fab_id: str,
+        fab_version: str,
+        override_config: Dict[str, str],
+    ) -> int:
         """Create a new run for the specified `fab_id` and `fab_version`."""
 
     @abc.abstractmethod

flwr/simulation/app.py

@@ -29,12 +29,14 @@ from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 
 from flwr.client import ClientFnExt
 from flwr.common import EventType, event
-from flwr.common.logger import log, set_logger_propagation
+from flwr.common.constant import NODE_ID_NUM_BYTES
+from flwr.common.logger import log, set_logger_propagation, warn_unsupported_feature
 from flwr.server.client_manager import ClientManager
 from flwr.server.history import History
 from flwr.server.server import Server, init_defaults, run_fl
 from flwr.server.server_config import ServerConfig
 from flwr.server.strategy import Strategy
+from flwr.server.superlink.state.utils import generate_rand_int_from_bytes
 from flwr.simulation.ray_transport.ray_actor import (
     ClientAppActor,
     VirtualClientEngineActor,
@@ -51,7 +53,7 @@ Invalid Arguments in method:
 `start_simulation(
     *,
     client_fn: ClientFn,
-    num_clients: Optional[int] = None,
+    num_clients: int,
     clients_ids: Optional[List[str]] = None,
     client_resources: Optional[Dict[str, float]] = None,
     server: Optional[Server] = None,
@@ -70,13 +72,29 @@ REASON:
 
 """
 
+NodeToPartitionMapping = Dict[int, int]
+
+
+def _create_node_id_to_partition_mapping(
+    num_clients: int,
+) -> NodeToPartitionMapping:
+    """Generate a node_id:partition_id mapping."""
+    nodes_mapping: NodeToPartitionMapping = {}  # {node-id; partition-id}
+    for i in range(num_clients):
+        while True:
+            node_id = generate_rand_int_from_bytes(NODE_ID_NUM_BYTES)
+            if node_id not in nodes_mapping:
+                break
+        nodes_mapping[node_id] = i
+    return nodes_mapping
+
 
 # pylint: disable=too-many-arguments,too-many-statements,too-many-branches
 def start_simulation(
     *,
     client_fn: ClientFnExt,
-    num_clients: Optional[int] = None,
-    clients_ids: Optional[List[str]] = None,
+    num_clients: int,
+    clients_ids: Optional[List[str]] = None,  # UNSUPPORTED, WILL BE REMOVED
     client_resources: Optional[Dict[str, float]] = None,
     server: Optional[Server] = None,
     config: Optional[ServerConfig] = None,
@@ -102,13 +120,14 @@ def start_simulation(
         (model, dataset, hyperparameters, ...) should be (re-)created in either the
         call to `client_fn` or the call to any of the client methods (e.g., load
         evaluation data in the `evaluate` method itself).
-    num_clients : Optional[int]
-        The total number of clients in this simulation. This must be set if
-        `clients_ids` is not set and vice-versa.
+    num_clients : int
+        The total number of clients in this simulation.
     clients_ids : Optional[List[str]]
+        UNSUPPORTED, WILL BE REMOVED. USE `num_clients` INSTEAD.
         List `client_id`s for each client. This is only required if
         `num_clients` is not set. Setting both `num_clients` and `clients_ids`
         with `len(clients_ids)` not equal to `num_clients` generates an error.
+        Using this argument will raise an error.
     client_resources : Optional[Dict[str, float]] (default: `{"num_cpus": 1, "num_gpus": 0.0}`)
         CPU and GPU resources for a single client. Supported keys
         are `num_cpus` and `num_gpus`. To understand the GPU utilization caused by
@@ -158,7 +177,6 @@ def start_simulation(
         is an advanced feature. For all details, please refer to the Ray documentation:
         https://docs.ray.io/en/latest/ray-core/scheduling/index.html
 
-
     Returns
     -------
     hist : flwr.server.history.History
@@ -170,6 +188,14 @@ def start_simulation(
         {"num_clients": len(clients_ids) if clients_ids is not None else num_clients},
     )
 
+    if clients_ids is not None:
+        warn_unsupported_feature(
+            "Passing `clients_ids` to `start_simulation` is deprecated and not longer "
+            "used by `start_simulation`. Use `num_clients` exclusively instead."
+        )
+        log(ERROR, "`clients_ids` argument used.")
+        sys.exit()
+
     # Set logger propagation
     loop: Optional[asyncio.AbstractEventLoop] = None
     try:
@@ -196,20 +222,8 @@ def start_simulation(
         initialized_config,
     )
 
-    # clients_ids takes precedence
-    cids: List[str]
-    if clients_ids is not None:
-        if (num_clients is not None) and (len(clients_ids) != num_clients):
-            log(ERROR, INVALID_ARGUMENTS_START_SIMULATION)
-            sys.exit()
-        else:
-            cids = clients_ids
-    else:
-        if num_clients is None:
-            log(ERROR, INVALID_ARGUMENTS_START_SIMULATION)
-            sys.exit()
-        else:
-            cids = [str(x) for x in range(num_clients)]
+    # Create node-id to partition-id mapping
+    nodes_mapping = _create_node_id_to_partition_mapping(num_clients)
 
     # Default arguments for Ray initialization
     if not ray_init_args:
@@ -308,10 +322,11 @@ def start_simulation(
     )
 
     # Register one RayClientProxy object for each client with the ClientManager
-    for cid in cids:
+    for node_id, partition_id in nodes_mapping.items():
         client_proxy = RayActorClientProxy(
             client_fn=client_fn,
-            cid=cid,
+            node_id=node_id,
+            partition_id=partition_id,
             actor_pool=pool,
         )
         initialized_server.client_manager().register(client=client_proxy)