flwr 1.21.0__py3-none-any.whl → 1.22.0__py3-none-any.whl

flwr/cli/app.py

@@ -25,6 +25,7 @@ from .log import log
 from .login import login
 from .ls import ls
 from .new import new
+from .pull import pull
 from .run import run
 from .stop import stop
 
@@ -46,6 +47,7 @@ app.command()(log)
 app.command()(ls)
 app.command()(stop)
 app.command()(login)
+app.command()(pull)
 
 typer_click_object = get_command(app)
 

flwr/cli/new/new.py

@@ -35,15 +35,16 @@ class MlFramework(str, Enum):
     """Available frameworks."""
 
     PYTORCH = "PyTorch"
-    PYTORCH_MSG_API = "PyTorch (Message API)"
     TENSORFLOW = "TensorFlow"
     SKLEARN = "sklearn"
     HUGGINGFACE = "HuggingFace"
     JAX = "JAX"
     MLX = "MLX"
     NUMPY = "NumPy"
+    XGBOOST = "XGBoost"
     FLOWERTUNE = "FlowerTune"
     BASELINE = "Flower Baseline"
+    PYTORCH_LEGACY_API = "PyTorch (Legacy API, deprecated)"
 
 
 class LlmChallengeName(str, Enum):
@@ -155,8 +156,8 @@ def new(
     if framework_str == MlFramework.BASELINE:
         framework_str = "baseline"
 
-    if framework_str == MlFramework.PYTORCH_MSG_API:
-        framework_str = "pytorch_msg_api"
+    if framework_str == MlFramework.PYTORCH_LEGACY_API:
+        framework_str = "pytorch_legacy_api"
 
     print(
         typer.style(
@@ -201,7 +202,7 @@ def new(
         }
 
         # Challenge specific context
-        fraction_fit = "0.2" if llm_challenge_str == "code" else "0.1"
+        fraction_train = "0.2" if llm_challenge_str == "code" else "0.1"
         if llm_challenge_str == "generalnlp":
             challenge_name = "General NLP"
             num_clients = "20"
@@ -220,7 +221,7 @@ def new(
             dataset_name = "flwrlabs/code-alpaca-20k"
 
         context["llm_challenge_str"] = llm_challenge_str
-        context["fraction_fit"] = fraction_fit
+        context["fraction_train"] = fraction_train
         context["challenge_name"] = challenge_name
         context["num_clients"] = num_clients
         context["dataset_name"] = dataset_name
@@ -247,14 +248,15 @@ def new(
             MlFramework.TENSORFLOW.value,
             MlFramework.SKLEARN.value,
             MlFramework.NUMPY.value,
-            "pytorch_msg_api",
+            MlFramework.XGBOOST.value,
+            "pytorch_legacy_api",
         ]
         if framework_str in frameworks_with_tasks:
             files[f"{import_name}/task.py"] = {
                 "template": f"app/code/task.{template_name}.py.tpl"
             }
 
-        if framework_str == "pytorch_msg_api":
+        if framework_str == "pytorch_legacy_api":
             # Use custom __init__ that better captures name of framework
             files[f"{import_name}/__init__.py"] = {
                 "template": f"app/code/__init__.{framework_str}.py.tpl"

flwr/cli/new/templates/app/README.flowertune.md.tpl

@@ -26,7 +26,7 @@ pip install -e .
 ## Experimental setup
 
 The dataset is divided into $num_clients partitions in an IID fashion, a partition is assigned to each ClientApp.
-We randomly sample a fraction ($fraction_fit) of the total nodes to participate in each round, for a total of `200` rounds.
+We randomly sample a fraction ($fraction_train) of the total nodes to participate in each round, for a total of `200` rounds.
 All settings are defined in `pyproject.toml`.
 
 > [!IMPORTANT]

flwr/cli/new/templates/app/code/client.baseline.py.tpl

@@ -1,58 +1,75 @@
 """$project_name: A Flower Baseline."""
 
 import torch
-from flwr.client import ClientApp, NumPyClient
-from flwr.common import Context
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
 
 from $import_name.dataset import load_data
-from $import_name.model import Net, get_weights, set_weights, test, train
-
-
-class FlowerClient(NumPyClient):
-    """A class defining the client."""
-
-    def __init__(self, net, trainloader, valloader, local_epochs):
-        self.net = net
-        self.trainloader = trainloader
-        self.valloader = valloader
-        self.local_epochs = local_epochs
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self.net.to(self.device)
-
-    def fit(self, parameters, config):
-        """Traim model using this client's data."""
-        set_weights(self.net, parameters)
-        train_loss = train(
-            self.net,
-            self.trainloader,
-            self.local_epochs,
-            self.device,
-        )
-        return (
-            get_weights(self.net),
-            len(self.trainloader.dataset),
-            {"train_loss": train_loss},
-        )
-
-    def evaluate(self, parameters, config):
-        """Evaluate model using this client's data."""
-        set_weights(self.net, parameters)
-        loss, accuracy = test(self.net, self.valloader, self.device)
-        return loss, len(self.valloader.dataset), {"accuracy": accuracy}
-
-
-def client_fn(context: Context):
-    """Construct a Client that will be run in a ClientApp."""
-    # Load model and data
-    net = Net()
+from $import_name.model import Net
+from $import_name.model import test as test_fn
+from $import_name.model import train as train_fn
+
+# Flower ClientApp
+app = ClientApp()
+
+
+@app.train()
+def train(msg: Message, context: Context):
+    """Train the model on local data."""
+
+    # Load the model and initialize it with the received weights
+    model = Net()
+    model.load_state_dict(msg.content["arrays"].to_torch_state_dict())
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # Load the data
     partition_id = int(context.node_config["partition-id"])
     num_partitions = int(context.node_config["num-partitions"])
-    trainloader, valloader = load_data(partition_id, num_partitions)
+    trainloader, _ = load_data(partition_id, num_partitions)
     local_epochs = context.run_config["local-epochs"]
 
-    # Return Client instance
-    return FlowerClient(net, trainloader, valloader, local_epochs).to_client()
+    # Call the training function
+    train_loss = train_fn(
+        model,
+        trainloader,
+        local_epochs,
+        device,
+    )
 
+    # Construct and return reply Message
+    model_record = ArrayRecord(model.state_dict())
+    metrics = {
+        "train_loss": train_loss,
+        "num-examples": len(trainloader.dataset),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)
 
-# Flower ClientApp
-app = ClientApp(client_fn)
+
+@app.evaluate()
+def evaluate(msg: Message, context: Context):
+    """Evaluate the model on local data."""
+
+    # Load the model and initialize it with the received weights
+    model = Net()
+    model.load_state_dict(msg.content["arrays"].to_torch_state_dict())
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # Load the data
+    partition_id = int(context.node_config["partition-id"])
+    num_partitions = int(context.node_config["num-partitions"])
+    _, valloader = load_data(partition_id, num_partitions)
+
+    # Call the evaluation function
+    eval_loss, eval_acc = test_fn(model, valloader, device)
+
+    # Construct and return reply Message
+    metrics = {
+        "eval_loss": eval_loss,
+        "eval_acc": eval_acc,
+        "num-examples": len(valloader.dataset),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"metrics": metric_record})
+    return Message(content=content, reply_to=msg)

flwr/cli/new/templates/app/code/client.huggingface.py.tpl

@@ -1,41 +1,67 @@
 """$project_name: A Flower / $framework_str app."""
 
 import torch
-from flwr.client import ClientApp, NumPyClient
-from flwr.common import Context
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
 from transformers import AutoModelForSequenceClassification
 
-from $import_name.task import get_weights, load_data, set_weights, test, train
+from $import_name.task import load_data
+from $import_name.task import test as test_fn
+from $import_name.task import train as train_fn
 
+# Flower ClientApp
+app = ClientApp()
+
+
+@app.train()
+def train(msg: Message, context: Context):
+    """Train the model on local data."""
+
+    # Get this client's dataset partition
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    model_name = context.run_config["model-name"]
+    trainloader, _ = load_data(partition_id, num_partitions, model_name)
+
+    # Load model
+    num_labels = context.run_config["num-labels"]
+    net = AutoModelForSequenceClassification.from_pretrained(
+        model_name, num_labels=num_labels
+    )
 
-# Flower client
-class FlowerClient(NumPyClient):
-    def __init__(self, net, trainloader, testloader, local_epochs):
-        self.net = net
-        self.trainloader = trainloader
-        self.testloader = testloader
-        self.local_epochs = local_epochs
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self.net.to(self.device)
+    # Initialize it with the received weights
+    net.load_state_dict(msg.content["arrays"].to_torch_state_dict())
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    net.to(device)
 
-    def fit(self, parameters, config):
-        set_weights(self.net, parameters)
-        train(self.net, self.trainloader, epochs=self.local_epochs, device=self.device)
-        return get_weights(self.net), len(self.trainloader), {}
+    # Train the model on local data
+    train_loss = train_fn(
+        net,
+        trainloader,
+        context.run_config["local-steps"],
+        device,
+    )
 
-    def evaluate(self, parameters, config):
-        set_weights(self.net, parameters)
-        loss, accuracy = test(self.net, self.testloader, self.device)
-        return float(loss), len(self.testloader), {"accuracy": accuracy}
+    # Construct and return reply Message
+    model_record = ArrayRecord(net.state_dict())
+    metrics = {
+        "train_loss": train_loss,
+        "num-examples": len(trainloader.dataset),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)
 
 
-def client_fn(context: Context):
+@app.evaluate()
+def evaluate(msg: Message, context: Context):
+    """Evaluate the model on local data."""
 
     # Get this client's dataset partition
     partition_id = context.node_config["partition-id"]
     num_partitions = context.node_config["num-partitions"]
     model_name = context.run_config["model-name"]
-    trainloader, valloader = load_data(partition_id, num_partitions, model_name)
+    _, valloader = load_data(partition_id, num_partitions, model_name)
 
     # Load model
     num_labels = context.run_config["num-labels"]
@@ -43,13 +69,25 @@ def client_fn(context: Context):
         model_name, num_labels=num_labels
     )
 
-    local_epochs = context.run_config["local-epochs"]
-
-    # Return Client instance
-    return FlowerClient(net, trainloader, valloader, local_epochs).to_client()
+    # Initialize it with the received weights
+    net.load_state_dict(msg.content["arrays"].to_torch_state_dict())
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    net.to(device)
 
+    # Evaluate the model on local data
+    val_loss, val_accuracy = test_fn(
+        net,
+        valloader,
+        device,
+    )
 
-# Flower ClientApp
-app = ClientApp(
-    client_fn,
-)
+    # Construct and return reply Message
+    model_record = ArrayRecord(net.state_dict())
+    metrics = {
+        "val_loss": val_loss,
+        "val_accuracy": val_accuracy,
+        "num-examples": len(valloader.dataset),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)

flwr/cli/new/templates/app/code/client.jax.py.tpl

@@ -1,50 +1,71 @@
 """$project_name: A Flower / $framework_str app."""
 
 import jax
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
 
-from flwr.client import ClientApp, NumPyClient
-from flwr.common import Context
-from $import_name.task import (
-    evaluation,
-    get_params,
-    load_data,
-    load_model,
-    loss_fn,
-    set_params,
-    train,
-)
-
-
-# Define Flower Client and client_fn
-class FlowerClient(NumPyClient):
-    def __init__(self, input_dim):
-        self.train_x, self.train_y, self.test_x, self.test_y = load_data()
-        self.grad_fn = jax.grad(loss_fn)
-        self.params = load_model((input_dim,))
-
-    def fit(self, parameters, config):
-        set_params(self.params, parameters)
-        self.params, loss, num_examples = train(
-            self.params, self.grad_fn, self.train_x, self.train_y
-        )
-        return get_params(self.params), num_examples, {"loss": float(loss)}
-
-    def evaluate(self, parameters, config):
-        set_params(self.params, parameters)
-        loss, num_examples = evaluation(
-            self.params, self.grad_fn, self.test_x, self.test_y
-        )
-        return float(loss), num_examples, {"loss": float(loss)}
-
-
-def client_fn(context: Context):
+from $import_name.task import evaluation as evaluation_fn
+from $import_name.task import get_params, load_data, load_model, loss_fn, set_params
+from $import_name.task import train as train_fn
+
+# Flower ClientApp
+app = ClientApp()
+
+
+@app.train()
+def train(msg: Message, context: Context):
+    """Train the model on local data."""
+
+    # Read from config
     input_dim = context.run_config["input-dim"]
 
-    # Return Client instance
-    return FlowerClient(input_dim).to_client()
+    # Load data and model
+    train_x, train_y, _, _ = load_data()
+    model = load_model((input_dim,))
+    grad_fn = jax.grad(loss_fn)
 
+    # Set model parameters
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
+    set_params(model, ndarrays)
 
-# Flower ClientApp
-app = ClientApp(
-    client_fn,
-)
+    # Train the model on local data
+    model, loss, num_examples = train_fn(model, grad_fn, train_x, train_y)
+
+    # Construct and return reply Message
+    model_record = ArrayRecord(get_params(model))
+    metrics = {
+        "train_loss": float(loss),
+        "num-examples": num_examples,
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)
+
+
+@app.evaluate()
+def evaluate(msg: Message, context: Context):
+    """Evaluate the model on local data."""
+
+    # Read from config
+    input_dim = context.run_config["input-dim"]
+
+    # Load data and model
+    _, _, test_x, test_y = load_data()
+    model = load_model((input_dim,))
+    grad_fn = jax.grad(loss_fn)
+
+    # Set model parameters
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
+    set_params(model, ndarrays)
+
+    # Evaluate the model on local data
+    loss, num_examples = evaluation_fn(model, grad_fn, test_x, test_y)
+
+    # Construct and return reply Message
+    metrics = {
+        "test_loss": float(loss),
+        "num-examples": num_examples,
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"metrics": metric_record})
+    return Message(content=content, reply_to=msg)

flwr/cli/new/templates/app/code/client.mlx.py.tpl

@@ -3,10 +3,9 @@
 import mlx.core as mx
 import mlx.nn as nn
 import mlx.optimizers as optim
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
 
-from flwr.client import ClientApp, NumPyClient
-from flwr.common import Context
-from flwr.common.config import UserConfig
 from $import_name.task import (
     MLP,
     batch_iterate,
@@ -17,57 +16,87 @@ from $import_name.task import (
     set_params,
 )
 
+# Flower ClientApp
+app = ClientApp()
+
+
+@app.train()
+def train(msg: Message, context: Context):
+    """Train the model on local data."""
+
+    # Read config
+    num_layers = context.run_config["num-layers"]
+    input_dim = context.run_config["input-dim"]
+    hidden_dim = context.run_config["hidden-dim"]
+    batch_size = context.run_config["batch-size"]
+    learning_rate = context.run_config["lr"]
+    num_epochs = context.run_config["local-epochs"]
 
-# Define Flower Client and client_fn
-class FlowerClient(NumPyClient):
-    def __init__(
-        self,
-        data,
-        run_config: UserConfig,
-        num_classes,
-    ):
-        num_layers = run_config["num-layers"]
-        hidden_dim = run_config["hidden-dim"]
-        input_dim = run_config["input-dim"]
-        batch_size = run_config["batch-size"]
-        learning_rate = run_config["lr"]
-        self.num_epochs = run_config["local-epochs"]
-
-        self.train_images, self.train_labels, self.test_images, self.test_labels = data
-        self.model = MLP(num_layers, input_dim, hidden_dim, num_classes)
-        self.optimizer = optim.SGD(learning_rate=learning_rate)
-        self.loss_and_grad_fn = nn.value_and_grad(self.model, loss_fn)
-        self.batch_size = batch_size
-
-    def fit(self, parameters, config):
-        set_params(self.model, parameters)
-        for _ in range(self.num_epochs):
-            for X, y in batch_iterate(
-                self.batch_size, self.train_images, self.train_labels
-            ):
-                _, grads = self.loss_and_grad_fn(self.model, X, y)
-                self.optimizer.update(self.model, grads)
-                mx.eval(self.model.parameters(), self.optimizer.state)
-        return get_params(self.model), len(self.train_images), {}
-
-    def evaluate(self, parameters, config):
-        set_params(self.model, parameters)
-        accuracy = eval_fn(self.model, self.test_images, self.test_labels)
-        loss = loss_fn(self.model, self.test_images, self.test_labels)
-        return loss.item(), len(self.test_images), {"accuracy": accuracy.item()}
-
-
-def client_fn(context: Context):
+    # Instantiate model and apply global parameters
+    model = MLP(num_layers, input_dim, hidden_dim, output_dim=10)
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
+    set_params(model, ndarrays)
+
+    # Define optimizer and loss function
+    optimizer = optim.SGD(learning_rate=learning_rate)
+    loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
+
+    # Load data
     partition_id = context.node_config["partition-id"]
     num_partitions = context.node_config["num-partitions"]
-    data = load_data(partition_id, num_partitions)
-    num_classes = 10
+    train_images, train_labels, _, _ = load_data(partition_id, num_partitions)
 
-    # Return Client instance
-    return FlowerClient(data, context.run_config, num_classes).to_client()
+    # Train the model on local data
+    for _ in range(num_epochs):
+        for X, y in batch_iterate(batch_size, train_images, train_labels):
+            _, grads = loss_and_grad_fn(model, X, y)
+            optimizer.update(model, grads)
+            mx.eval(model.parameters(), optimizer.state)
 
+    # Compute train accuracy and loss
+    accuracy = eval_fn(model, train_images, train_labels)
+    loss = loss_fn(model, train_images, train_labels)
+    # Construct and return reply Message
+    model_record = ArrayRecord(get_params(model))
+    metrics = {
+        "num-examples": len(train_images),
+        "accuracy": float(accuracy.item()),
+        "loss": float(loss.item()),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)
 
-# Flower ClientApp
-app = ClientApp(
-    client_fn,
-)
+
+@app.evaluate()
+def evaluate(msg: Message, context: Context):
+    """Evaluate the model on local data."""
+
+    # Read config
+    num_layers = context.run_config["num-layers"]
+    input_dim = context.run_config["input-dim"]
+    hidden_dim = context.run_config["hidden-dim"]
+
+    # Instantiate model and apply global parameters
+    model = MLP(num_layers, input_dim, hidden_dim, output_dim=10)
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
+    set_params(model, ndarrays)
+
+    # Load data
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    _, _, test_images, test_labels = load_data(partition_id, num_partitions)
+
+    # Evaluate the model on local data
+    accuracy = eval_fn(model, test_images, test_labels)
+    loss = loss_fn(model, test_images, test_labels)
+
+    # Construct and return reply Message
+    metrics = {
+        "num-examples": len(test_images),
+        "accuracy": float(accuracy.item()),
+        "loss": float(loss.item()),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"metrics": metric_record})
+    return Message(content=content, reply_to=msg)

flwr/cli/new/templates/app/code/client.numpy.py.tpl

@@ -1,23 +1,46 @@
 """$project_name: A Flower / $framework_str app."""
 
-from flwr.client import ClientApp, NumPyClient
-from flwr.common import Context
-from $import_name.task import get_dummy_model
+import numpy as np
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
 
+# Flower ClientApp
+app = ClientApp()
 
-class FlowerClient(NumPyClient):
 
-    def fit(self, parameters, config):
-        model = get_dummy_model()
-        return [model], 1, {}
+@app.train()
+def train(msg: Message, context: Context):
+    """Train the model on local data."""
 
-    def evaluate(self, parameters, config):
-        return float(0.0), 1, {"accuracy": float(1.0)}
+    # The model is the global arrays
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
 
+    # Simulate local training (here we just add random noise to model parameters)
+    model = [m + np.random.rand(*m.shape) for m in ndarrays]
 
-def client_fn(context: Context):
-    return FlowerClient().to_client()
+    # Construct and return reply Message
+    model_record = ArrayRecord(model)
+    metrics = {
+        "random_metric": np.random.rand(),
+        "num-examples": 1,
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)
 
 
-# Flower ClientApp
-app = ClientApp(client_fn=client_fn)
+@app.evaluate()
+def evaluate(msg: Message, context: Context):
+    """Evaluate the model on local data."""
+
+    # The model is the global arrays
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
+
+    # Return reply Message
+    metrics = {
+        "random_metric": np.random.rand(3).tolist(),
+        "num-examples": 1,
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"metrics": metric_record})
+    return Message(content=content, reply_to=msg)