flwr-nightly 1.10.0.dev20240712__tar.gz → 1.10.0.dev20240714__tar.gz

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.10.0.dev20240712
+Version: 1.10.0.dev20240714
 Summary: Flower: A Friendly Federated Learning Framework
 Home-page: https://flower.ai
 License: Apache-2.0

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "flwr-nightly"
-version = "1.10.0.dev20240712"
+version = "1.10.0.dev20240714"
 description = "Flower: A Friendly Federated Learning Framework"
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/cli/new/templates/app/code/client.hf.py.tpl

@@ -1,6 +1,7 @@
 """$project_name: A Flower / HuggingFace Transformers app."""
 
 from flwr.client import ClientApp, NumPyClient
+from flwr.common import Context
 from transformers import AutoModelForSequenceClassification
 
 from $import_name.task import (
@@ -38,12 +39,15 @@ class FlowerClient(NumPyClient):
         return float(loss), len(self.testloader), {"accuracy": accuracy}
 
 
-def client_fn(cid):
+def client_fn(context: Context):
     # Load model and data
     net = AutoModelForSequenceClassification.from_pretrained(
         CHECKPOINT, num_labels=2
     ).to(DEVICE)
-    trainloader, valloader = load_data(int(cid), 2)
+
+    partition_id = int(context.node_config['partition-id'])
+    num_partitions = int(context.node_config['num-partitions])
+    trainloader, valloader = load_data(partition_id, num_partitions)
 
     # Return Client instance
     return FlowerClient(net, trainloader, valloader).to_client()

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/cli/new/templates/app/code/client.jax.py.tpl

@@ -2,6 +2,7 @@
 
 import jax
 from flwr.client import NumPyClient, ClientApp
+from flwr.common import Context
 
 from $import_name.task import (
     evaluation,
@@ -44,7 +45,7 @@ class FlowerClient(NumPyClient):
         )
         return float(loss), num_examples, {"loss": float(loss)}
 
-def client_fn(cid):
+def client_fn(context: Context):
     # Return Client instance
     return FlowerClient().to_client()
 

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/cli/new/templates/app/code/client.mlx.py.tpl

@@ -4,6 +4,7 @@ import mlx.core as mx
 import mlx.nn as nn
 import mlx.optimizers as optim
 from flwr.client import NumPyClient, ClientApp
+from flwr.common import Context
 
 from $import_name.task import (
     batch_iterate,
@@ -57,8 +58,10 @@ class FlowerClient(NumPyClient):
         return loss.item(), len(self.test_images), {"accuracy": accuracy.item()}
 
 
-def client_fn(cid):
-    data = load_data(int(cid), 2)
+def client_fn(context: Context):
+    partition_id = int(context.node_config["partition-id"])
+    num_partitions = int(context.node_config["num-partitions"])
+    data = load_data(partition_id, num_partitions)
 
     # Return Client instance
     return FlowerClient(data).to_client()

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/cli/new/templates/app/code/client.numpy.py.tpl

@@ -1,6 +1,7 @@
 """$project_name: A Flower / NumPy app."""
 
 from flwr.client import NumPyClient, ClientApp
+from flwr.common import Context
 import numpy as np
 
 
@@ -15,7 +16,7 @@ class FlowerClient(NumPyClient):
         return float(0.0), 1, {"accuracy": float(1.0)}
 
 
-def client_fn(cid: str):
+def client_fn(context: Context):
     return FlowerClient().to_client()
 
 

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/cli/new/templates/app/code/client.pytorch.py.tpl

@@ -1,6 +1,7 @@
 """$project_name: A Flower / PyTorch app."""
 
 from flwr.client import NumPyClient, ClientApp
+from flwr.common import Context
 
 from $import_name.task import (
     Net,
@@ -31,10 +32,12 @@ class FlowerClient(NumPyClient):
         return loss, len(self.valloader.dataset), {"accuracy": accuracy}
 
 
-def client_fn(cid):
+def client_fn(context: Context):
     # Load model and data
     net = Net().to(DEVICE)
-    trainloader, valloader = load_data(int(cid), 2)
+    partition_id = int(context.node_config["partition-id"])
+    num_partitions = int(context.node_config["num-partitions"])
+    trainloader, valloader = load_data(partition_id, num_partitions)
 
     # Return Client instance
     return FlowerClient(net, trainloader, valloader).to_client()

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/cli/new/templates/app/code/client.sklearn.py.tpl

@@ -4,6 +4,7 @@ import warnings
 
 import numpy as np
 from flwr.client import NumPyClient, ClientApp
+from flwr.common import Context
 from flwr_datasets import FederatedDataset
 from sklearn.linear_model import LogisticRegression
 from sklearn.metrics import log_loss
@@ -68,8 +69,9 @@ class FlowerClient(NumPyClient):
 
 fds = FederatedDataset(dataset="mnist", partitioners={"train": 2})
 
-def client_fn(cid: str):
-    dataset = fds.load_partition(int(cid), "train").with_format("numpy")
+def client_fn(context: Context):
+    partition_id = int(context.node_config["partition-id"])
+    dataset = fds.load_partition(partition_id, "train").with_format("numpy")
 
     X, y = dataset["image"].reshape((len(dataset), -1)), dataset["label"]
 

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/cli/new/templates/app/code/client.tensorflow.py.tpl

@@ -1,6 +1,7 @@
 """$project_name: A Flower / TensorFlow app."""
 
 from flwr.client import NumPyClient, ClientApp
+from flwr.common import Context
 
 from $import_name.task import load_data, load_model
 
@@ -28,10 +29,12 @@ class FlowerClient(NumPyClient):
         return loss, len(self.x_test), {"accuracy": accuracy}
 
 
-def client_fn(cid):
+def client_fn(context: Context):
     # Load model and data
     net = load_model()
-    x_train, y_train, x_test, y_test = load_data(int(cid), 2)
+
+    partition_id = int(context.node_config["partition-id"])
+    x_train, y_train, x_test, y_test = load_data(partition_id, 2)
 
     # Return Client instance
     return FlowerClient(net, x_train, y_train, x_test, y_test).to_client()

flwr_nightly-1.10.0.dev20240714/src/py/flwr/cli/new/templates/app/code/server.hf.py.tpl

@@ -0,0 +1,19 @@
+"""$project_name: A Flower / HuggingFace Transformers app."""
+
+from flwr.common import Context
+from flwr.server.strategy import FedAvg
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+
+
+def server_fn(context: Context):
+    # Define strategy
+    strategy = FedAvg(
+        fraction_fit=1.0,
+        fraction_evaluate=1.0,
+    )
+    config = ServerConfig(num_rounds=3)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

flwr_nightly-1.10.0.dev20240714/src/py/flwr/cli/new/templates/app/code/server.jax.py.tpl

@@ -0,0 +1,16 @@
+"""$project_name: A Flower / JAX app."""
+
+from flwr.common import Context
+from flwr.server.strategy import FedAvg
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+
+
+def server_fn(context: Context):
+    # Define strategy
+    strategy = FedAvg()
+    config = ServerConfig(num_rounds=3)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

flwr_nightly-1.10.0.dev20240714/src/py/flwr/cli/new/templates/app/code/server.mlx.py.tpl

@@ -0,0 +1,16 @@
+"""$project_name: A Flower / MLX app."""
+
+from flwr.common import Context
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from flwr.server.strategy import FedAvg
+
+
+def server_fn(context: Context):
+    # Define strategy
+    strategy = FedAvg()
+    config = ServerConfig(num_rounds=3)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

flwr_nightly-1.10.0.dev20240714/src/py/flwr/cli/new/templates/app/code/server.numpy.py.tpl

@@ -0,0 +1,16 @@
+"""$project_name: A Flower / NumPy app."""
+
+from flwr.common import Context
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from flwr.server.strategy import FedAvg
+
+
+def server_fn(context: Context):
+    # Define strategy
+    strategy = FedAvg()
+    config = ServerConfig(num_rounds=3)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

flwr_nightly-1.10.0.dev20240714/src/py/flwr/cli/new/templates/app/code/server.pytorch.py.tpl

@@ -0,0 +1,27 @@
+"""$project_name: A Flower / PyTorch app."""
+
+from flwr.common import Context, ndarrays_to_parameters
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from flwr.server.strategy import FedAvg
+
+from $import_name.task import Net, get_weights
+
+
+# Initialize model parameters
+ndarrays = get_weights(Net())
+parameters = ndarrays_to_parameters(ndarrays)
+
+def server_fn(context: Context):
+    # Define strategy
+    strategy = FedAvg(
+        fraction_fit=1.0,
+        fraction_evaluate=1.0,
+        min_available_clients=2,
+        initial_parameters=parameters,
+    )
+    config = ServerConfig(num_rounds=3)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

flwr_nightly-1.10.0.dev20240714/src/py/flwr/cli/new/templates/app/code/server.sklearn.py.tpl

@@ -0,0 +1,20 @@
+"""$project_name: A Flower / Scikit-Learn app."""
+
+from flwr.common import Context
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from flwr.server.strategy import FedAvg
+
+
+def server_fn(context: Context):
+    # Define strategy
+    strategy = FedAvg(
+        fraction_fit=1.0,
+        fraction_evaluate=1.0,
+        min_available_clients=2,
+    )
+    config = ServerConfig(num_rounds=3)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

flwr_nightly-1.10.0.dev20240714/src/py/flwr/cli/new/templates/app/code/server.tensorflow.py.tpl

@@ -0,0 +1,27 @@
+"""$project_name: A Flower / TensorFlow app."""
+
+from flwr.common import Context, ndarrays_to_parameters
+from flwr.server import ServerApp, ServerAppComponents, ServerConfig
+from flwr.server.strategy import FedAvg
+
+from $import_name.task import load_model
+
+# Define config
+config = ServerConfig(num_rounds=3)
+
+parameters = ndarrays_to_parameters(load_model().get_weights())
+
+def server_fn(context: Context):
+    # Define strategy
+    strategy = strategy = FedAvg(
+        fraction_fit=1.0,
+        fraction_evaluate=1.0,
+        min_available_clients=2,
+        initial_parameters=parameters,
+    )
+    config = ServerConfig(num_rounds=3)
+
+    return ServerAppComponents(strategy=strategy, config=config)
+
+# Create ServerApp
+app = ServerApp(server_fn=server_fn)

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/cli/new/templates/app/code/task.hf.py.tpl

@@ -16,9 +16,9 @@ DEVICE = torch.device("cpu")
 CHECKPOINT = "distilbert-base-uncased"  # transformer model checkpoint
 
 
-def load_data(partition_id, num_clients):
+def load_data(partition_id: int, num_partitions: int):
     """Load IMDB data (training and eval)"""
-    fds = FederatedDataset(dataset="imdb", partitioners={"train": num_clients})
+    fds = FederatedDataset(dataset="imdb", partitioners={"train": num_partitions})
     partition = fds.load_partition(partition_id)
     # Divide data: 80% train, 20% test
     partition_train_test = partition.train_test_split(test_size=0.2, seed=42)

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/cli/new/templates/app/code/task.mlx.py.tpl

@@ -43,8 +43,8 @@ def batch_iterate(batch_size, X, y):
         yield X[ids], y[ids]
 
 
-def load_data(partition_id, num_clients):
-    fds = FederatedDataset(dataset="mnist", partitioners={"train": num_clients})
+def load_data(partition_id: int, num_partitions: int):
+    fds = FederatedDataset(dataset="mnist", partitioners={"train": num_partitions})
     partition = fds.load_partition(partition_id)
     partition_splits = partition.train_test_split(test_size=0.2, seed=42)
 

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/cli/new/templates/app/code/task.pytorch.py.tpl

@@ -34,7 +34,7 @@ class Net(nn.Module):
         return self.fc3(x)
 
 
-def load_data(partition_id, num_partitions):
+def load_data(partition_id: int, num_partitions: int):
     """Load partition CIFAR10 data."""
     fds = FederatedDataset(dataset="cifar10", partitioners={"train": num_partitions})
     partition = fds.load_partition(partition_id)

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/cli/run/run.py

@@ -62,7 +62,7 @@ def run(
     config_overrides: Annotated[
         Optional[str],
         typer.Option(
-            "--config",
+            "--run-config",
             "-c",
             help="Override configuration key-value pairs",
         ),

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/client/app.py

@@ -28,7 +28,7 @@ from grpc import RpcError
 from flwr.client.client import Client
 from flwr.client.client_app import ClientApp, LoadClientAppError
 from flwr.client.typing import ClientFnExt
-from flwr.common import GRPC_MAX_MESSAGE_LENGTH, EventType, Message, event
+from flwr.common import GRPC_MAX_MESSAGE_LENGTH, Context, EventType, Message, event
 from flwr.common.address import parse_address
 from flwr.common.constant import (
     MISSING_EXTRA_REST,
@@ -138,8 +138,8 @@ def start_client(
 
     Starting an SSL-enabled gRPC client using system certificates:
 
-    >>> def client_fn(node_id: int, partition_id: Optional[int]):
-    >>>     return FlowerClient()
+    >>> def client_fn(context: Context):
+    >>>     return FlowerClient().to_client()
     >>>
     >>> start_client(
     >>>     server_address=localhost:8080,
@@ -253,8 +253,7 @@ def _start_client_internal(
         if client_fn is None:
             # Wrap `Client` instance in `client_fn`
             def single_client_factory(
-                node_id: int,  # pylint: disable=unused-argument
-                partition_id: Optional[int],  # pylint: disable=unused-argument
+                context: Context,  # pylint: disable=unused-argument
             ) -> Client:
                 if client is None:  # Added this to keep mypy happy
                     raise ValueError(
@@ -349,7 +348,6 @@ def _start_client_internal(
                     node_state = NodeState(
                         node_id=-1,
                         node_config={},
-                        partition_id=None,
                     )
                 else:
                     # Call create_node fn to register node
@@ -361,7 +359,6 @@ def _start_client_internal(
                     node_state = NodeState(
                         node_id=node_id,
                         node_config=node_config,
-                        partition_id=None,
                     )
 
             app_state_tracker.register_signal_handler()

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/client/client_app.py

@@ -30,21 +30,41 @@ from flwr.common.logger import warn_deprecated_feature, warn_preview_feature
 from .typing import ClientAppCallable
 
 
+def _alert_erroneous_client_fn() -> None:
+    raise ValueError(
+        "A `ClientApp` cannot make use of a `client_fn` that does "
+        "not have a signature in the form: `def client_fn(context: "
+        "Context)`. You can import the `Context` like this: "
+        "`from flwr.common import Context`"
+    )
+
+
 def _inspect_maybe_adapt_client_fn_signature(client_fn: ClientFnExt) -> ClientFnExt:
     client_fn_args = inspect.signature(client_fn).parameters
+    first_arg = list(client_fn_args.keys())[0]
+
+    if len(client_fn_args) != 1:
+        _alert_erroneous_client_fn()
+
+    first_arg_type = client_fn_args[first_arg].annotation
 
-    if not all(key in client_fn_args for key in ["node_id", "partition_id"]):
+    if first_arg_type is str or first_arg == "cid":
+        # Warn previous signature for `client_fn` seems to be used
         warn_deprecated_feature(
-            "`client_fn` now expects a signature `def client_fn(node_id: int, "
-            "partition_id: Optional[int])`.\nYou provided `client_fn` with signature: "
-            f"{dict(client_fn_args.items())}"
+            "`client_fn` now expects a signature `def client_fn(context: Context)`."
+            "The provided `client_fn` has signature: "
+            f"{dict(client_fn_args.items())}. You can import the `Context` like this:"
+            " `from flwr.common import Context`"
         )
 
         # Wrap depcreated client_fn inside a function with the expected signature
         def adaptor_fn(
-            node_id: int, partition_id: Optional[int]  # pylint: disable=unused-argument
-        ) -> Client:
-            return client_fn(str(partition_id))  # type: ignore
+            context: Context,
+        ) -> Client:  # pylint: disable=unused-argument
+            # if patition-id is defined, pass it. Else pass node_id that should
+            # always be defined during Context init.
+            cid = context.node_config.get("partition-id", context.node_id)
+            return client_fn(str(cid))  # type: ignore
 
         return adaptor_fn
 
@@ -71,7 +91,7 @@ class ClientApp:
     >>> class FlowerClient(NumPyClient):
     >>>     # ...
     >>>
-    >>> def client_fn(node_id: int, partition_id: Optional[int]):
+    >>> def client_fn(context: Context):
     >>>    return FlowerClient().to_client()
     >>>
     >>> app = ClientApp(client_fn)

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/client/message_handler/message_handler.py

@@ -92,7 +92,7 @@ def handle_legacy_message_from_msgtype(
     client_fn: ClientFnExt, message: Message, context: Context
 ) -> Message:
     """Handle legacy message in the inner most mod."""
-    client = client_fn(message.metadata.dst_node_id, context.partition_id)
+    client = client_fn(context)
 
     # Check if NumPyClient is returend
     if isinstance(client, NumPyClient):

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/client/node_state.py

@@ -36,12 +36,13 @@ class NodeState:
     """State of a node where client nodes execute runs."""
 
     def __init__(
-        self, node_id: int, node_config: Dict[str, str], partition_id: Optional[int]
+        self,
+        node_id: int,
+        node_config: Dict[str, str],
     ) -> None:
         self.node_id = node_id
         self.node_config = node_config
         self.run_infos: Dict[int, RunInfo] = {}
-        self._partition_id = partition_id
 
     def register_context(
         self,
@@ -59,7 +60,6 @@ class NodeState:
                     node_config=self.node_config,
                     state=RecordSet(),
                     run_config=initial_run_config.copy(),
-                    partition_id=self._partition_id,
                 ),
             )
 

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/client/node_state_tests.py

@@ -41,7 +41,7 @@ def test_multirun_in_node_state() -> None:
     expected_values = {0: "1", 1: "1" * 3, 2: "1" * 2, 3: "1", 5: "1"}
 
     # NodeState
-    node_state = NodeState(node_id=0, node_config={}, partition_id=None)
+    node_state = NodeState(node_id=0, node_config={})
 
     for task in tasks:
         run_id = task.run_id

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/client/typing.py

@@ -15,7 +15,7 @@
 """Custom types for Flower clients."""
 
 
-from typing import Callable, Optional
+from typing import Callable
 
 from flwr.common import Context, Message
 
@@ -23,7 +23,7 @@ from .client import Client as Client
 
 # Compatibility
 ClientFn = Callable[[str], Client]
-ClientFnExt = Callable[[int, Optional[int]], Client]
+ClientFnExt = Callable[[Context], Client]
 
 ClientAppCallable = Callable[[Message, Context], Message]
 Mod = Callable[[Message, Context, ClientAppCallable], Message]

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/common/constant.py

@@ -57,6 +57,9 @@ APP_DIR = "apps"
 FAB_CONFIG_FILE = "pyproject.toml"
 FLWR_HOME = "FLWR_HOME"
 
+# Constants entries in Node config for Simulation
+PARTITION_ID_KEY = "partition-id"
+NUM_PARTITIONS_KEY = "num-partitions"
 
 GRPC_ADAPTER_METADATA_FLOWER_VERSION_KEY = "flower-version"
 GRPC_ADAPTER_METADATA_SHOULD_EXIT_KEY = "should-exit"

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/common/context.py

@@ -16,7 +16,7 @@
 
 
 from dataclasses import dataclass
-from typing import Dict, Optional
+from typing import Dict
 
 from .record import RecordSet
 
@@ -43,17 +43,12 @@ class Context:
         A config (key/value mapping) held by the entity in a given run and that will
         stay local. It can be used at any point during the lifecycle of this entity
         (e.g. across multiple rounds)
-    partition_id : Optional[int] (default: None)
-        An index that specifies the data partition that the ClientApp using this Context
-        object should make use of. Setting this attribute is better suited for
-        simulation or proto typing setups.
     """
 
     node_id: int
     node_config: Dict[str, str]
     state: RecordSet
     run_config: Dict[str, str]
-    partition_id: Optional[int]
 
     def __init__(  # pylint: disable=too-many-arguments
         self,
@@ -61,10 +56,8 @@ class Context:
         node_config: Dict[str, str],
         state: RecordSet,
         run_config: Dict[str, str],
-        partition_id: Optional[int] = None,
     ) -> None:
         self.node_id = node_id
         self.node_config = node_config
         self.state = state
         self.run_config = run_config
-        self.partition_id = partition_id

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/server/superlink/fleet/vce/backend/raybackend.py

@@ -21,6 +21,7 @@ from typing import Callable, Dict, List, Tuple, Union
 import ray
 
 from flwr.client.client_app import ClientApp
+from flwr.common.constant import PARTITION_ID_KEY
 from flwr.common.context import Context
 from flwr.common.logger import log
 from flwr.common.message import Message
@@ -168,7 +169,7 @@ class RayBackend(Backend):
 
         Return output message and updated context.
         """
-        partition_id = context.partition_id
+        partition_id = context.node_config[PARTITION_ID_KEY]
 
         try:
             # Submit a task to the pool

{flwr_nightly-1.10.0.dev20240712 → flwr_nightly-1.10.0.dev20240714}/src/py/flwr/server/superlink/fleet/vce/vce_api.py

@@ -29,7 +29,12 @@ from typing import Callable, Dict, Optional
 
 from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
 from flwr.client.node_state import NodeState
-from flwr.common.constant import PING_MAX_INTERVAL, ErrorCode
+from flwr.common.constant import (
+    NUM_PARTITIONS_KEY,
+    PARTITION_ID_KEY,
+    PING_MAX_INTERVAL,
+    ErrorCode,
+)
 from flwr.common.logger import log
 from flwr.common.message import Error
 from flwr.common.object_ref import load_app
@@ -73,7 +78,7 @@ def worker(
             task_ins: TaskIns = taskins_queue.get(timeout=1.0)
             node_id = task_ins.task.consumer.node_id
 
-            # Register and retrieve runstate
+            # Register and retrieve context
             node_states[node_id].register_context(run_id=task_ins.run_id)
             context = node_states[node_id].retrieve_context(run_id=task_ins.run_id)
 
@@ -283,9 +288,15 @@ def start_vce(
 
     # Construct mapping of NodeStates
     node_states: Dict[int, NodeState] = {}
+    # Number of unique partitions
+    num_partitions = len(set(nodes_mapping.values()))
     for node_id, partition_id in nodes_mapping.items():
         node_states[node_id] = NodeState(
-            node_id=node_id, node_config={}, partition_id=partition_id
+            node_id=node_id,
+            node_config={
+                PARTITION_ID_KEY: str(partition_id),
+                NUM_PARTITIONS_KEY: str(num_partitions),
+            },
         )
 
     # Load backend config