feat(framework) Use types in template configs (#3875)

src/py/flwr/cli/new/templates/app/code/client.hf.py.tpl

@@ -17,10 +17,11 @@ from $import_name.task import (
 
 # Flower client
 class FlowerClient(NumPyClient):
-    def __init__(self, net, trainloader, testloader):
+    def __init__(self, net, trainloader, testloader, local_epochs):
         self.net = net
         self.trainloader = trainloader
         self.testloader = testloader
+        self.local_epochs = local_epochs
 
     def get_parameters(self, config):
         return get_weights(self.net)
@@ -33,7 +34,7 @@ class FlowerClient(NumPyClient):
         train(
             self.net,
             self.trainloader,
-            epochs=int(self.context.run_config["local-epochs"]),
+            epochs=self.local_epochs,
         )
         return self.get_parameters(config={}), len(self.trainloader), {}
 
@@ -52,9 +53,10 @@ def client_fn(context: Context):
     partition_id = int(context.node_config["partition-id"])
     num_partitions = int(context.node_config["num-partitions"])
     trainloader, valloader = load_data(partition_id, num_partitions)
+    local_epochs = context.run_config["local-epochs"]
 
     # Return Client instance
-    return FlowerClient(net, trainloader, valloader).to_client()
+    return FlowerClient(net, trainloader, valloader, local_epochs).to_client()
 
 
 # Flower ClientApp

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

@@ -19,13 +19,22 @@ from $import_name.task import (
 
 # Define Flower Client and client_fn
 class FlowerClient(NumPyClient):
-    def __init__(self, data):
-        num_layers = int(self.context.run_config["num-layers"])
-        hidden_dim = int(self.context.run_config["hidden-dim"])
-        num_classes = 10
-        batch_size = int(self.context.run_config["batch-size"])
-        learning_rate = float(self.context.run_config["lr"])
-        num_epochs = int(self.context.run_config["local-epochs"])
+    def __init__(
+        self,
+        data,
+        num_layers,
+        hidden_dim,
+        num_classes,
+        batch_size,
+        learning_rate,
+        num_epochs,
+    ):
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.num_classes = num_classes
+        self.batch_size = batch_size
+        self.learning_rate = learning_rate
+        self.num_epochs = num_epochs
 
         self.train_images, self.train_labels, self.test_images, self.test_labels = data
         self.model = MLP(
@@ -65,8 +74,17 @@ def client_fn(context: Context):
     num_partitions = int(context.node_config["num-partitions"])
     data = load_data(partition_id, num_partitions)
 
+    num_layers = context.run_config["num-layers"]
+    hidden_dim = context.run_config["hidden-dim"]
+    num_classes = 10
+    batch_size = context.run_config["batch-size"]
+    learning_rate = context.run_config["lr"]
+    num_epochs = context.run_config["local-epochs"]
+
     # Return Client instance
-    return FlowerClient(data).to_client()
+    return FlowerClient(
+        data, num_layers, hidden_dim, num_classes, batch_size, learning_rate, num_epochs
+    ).to_client()
 
 
 # Flower ClientApp

src/py/flwr/cli/new/templates/app/code/client.pytorch.py.tpl

@@ -16,18 +16,19 @@ from $import_name.task import (
 
 # Define Flower Client and client_fn
 class FlowerClient(NumPyClient):
-    def __init__(self, net, trainloader, valloader):
+    def __init__(self, net, trainloader, valloader, local_epochs):
         self.net = net
         self.trainloader = trainloader
         self.valloader = valloader
+        self.local_epochs = local_epochs
 
     def fit(self, parameters, config):
         set_weights(self.net, parameters)
         results = train(
             self.net,
             self.trainloader,
             self.valloader,
-            int(self.context.run_config["local-epochs"]),
+            self.local_epochs,
             DEVICE,
         )
         return get_weights(self.net), len(self.trainloader.dataset), results
@@ -44,9 +45,10 @@ def client_fn(context: Context):
     partition_id = int(context.node_config["partition-id"])
     num_partitions = int(context.node_config["num-partitions"])
     trainloader, valloader = load_data(partition_id, num_partitions)
+    local_epochs = context.run_config["local-epochs"]
 
     # Return Client instance
-    return FlowerClient(net, trainloader, valloader).to_client()
+    return FlowerClient(net, trainloader, valloader, local_epochs).to_client()
 
 
 # Flower ClientApp

src/py/flwr/cli/new/templates/app/code/client.tensorflow.py.tpl

@@ -8,12 +8,17 @@ from $import_name.task import load_data, load_model
 
 # Define Flower Client and client_fn
 class FlowerClient(NumPyClient):
-    def __init__(self, model, x_train, y_train, x_test, y_test):
+    def __init__(
+        self, model, x_train, y_train, x_test, y_test, epochs, batch_size, verbose
+    ):
         self.model = model
         self.x_train = x_train
         self.y_train = y_train
         self.x_test = x_test
         self.y_test = y_test
+        self.epochs = epochs
+        self.batch_size = batch_size
+        self.verbose = verbose
 
     def get_parameters(self, config):
         return self.model.get_weights()
@@ -23,9 +28,9 @@ class FlowerClient(NumPyClient):
         self.model.fit(
             self.x_train,
             self.y_train,
-            epochs=int(self.context.run_config["local-epochs"]),
-            batch_size=int(self.context.run_config["batch-size"]),
-            verbose=bool(self.context.run_config.get("verbose")),
+            epochs=self.epochs,
+            batch_size=self.batch_size,
+            verbose=self.verbose,
         )
         return self.model.get_weights(), len(self.x_train), {}
 
@@ -42,9 +47,14 @@ def client_fn(context: Context):
     partition_id = int(context.node_config["partition-id"])
     num_partitions = int(context.node_config["num-partitions"])
     x_train, y_train, x_test, y_test = load_data(partition_id, num_partitions)
+    epochs = context.run_config["local-epochs"]
+    batch_size = context.run_config["batch-size"]
+    verbose = context.run_config.get("verbose")
 
     # Return Client instance
-    return FlowerClient(net, x_train, y_train, x_test, y_test).to_client()
+    return FlowerClient(
+        net, x_train, y_train, x_test, y_test, epochs, batch_size, verbose
+    ).to_client()
 
 
 # Flower ClientApp

src/py/flwr/cli/new/templates/app/code/server.hf.py.tpl

@@ -7,7 +7,7 @@ from flwr.server import ServerApp, ServerAppComponents, ServerConfig
 
 def server_fn(context: Context):
     # Read from config
-    num_rounds = int(context.run_config["num-server-rounds"])
+    num_rounds = context.run_config["num-server-rounds"]
 
     # Define strategy
     strategy = FedAvg(
@@ -18,5 +18,6 @@ def server_fn(context: Context):
 
     return ServerAppComponents(strategy=strategy, config=config)
 
+
 # Create ServerApp
 app = ServerApp(server_fn=server_fn)

src/py/flwr/cli/new/templates/app/code/server.jax.py.tpl

@@ -7,13 +7,14 @@ from flwr.server import ServerApp, ServerAppComponents, ServerConfig
 
 def server_fn(context: Context):
     # Read from config
-    num_rounds = int(context.run_config["num-server-rounds"])
+    num_rounds = context.run_config["num-server-rounds"]
 
     # Define strategy
     strategy = FedAvg()
     config = ServerConfig(num_rounds=num_rounds)
 
     return ServerAppComponents(strategy=strategy, config=config)
 
+
 # Create ServerApp
 app = ServerApp(server_fn=server_fn)

src/py/flwr/cli/new/templates/app/code/server.mlx.py.tpl

@@ -7,13 +7,14 @@ from flwr.server.strategy import FedAvg
 
 def server_fn(context: Context):
     # Read from config
-    num_rounds = int(context.run_config["num-server-rounds"])
+    num_rounds = context.run_config["num-server-rounds"]
 
     # Define strategy
     strategy = FedAvg()
     config = ServerConfig(num_rounds=num_rounds)
 
     return ServerAppComponents(strategy=strategy, config=config)
 
+
 # Create ServerApp
 app = ServerApp(server_fn=server_fn)

src/py/flwr/cli/new/templates/app/code/server.numpy.py.tpl

@@ -7,13 +7,14 @@ from flwr.server.strategy import FedAvg
 
 def server_fn(context: Context):
     # Read from config
-    num_rounds = int(context.run_config["num-server-rounds"])
+    num_rounds = context.run_config["num-server-rounds"]
 
     # Define strategy
     strategy = FedAvg()
     config = ServerConfig(num_rounds=num_rounds)
 
     return ServerAppComponents(strategy=strategy, config=config)
 
+
 # Create ServerApp
 app = ServerApp(server_fn=server_fn)

src/py/flwr/cli/new/templates/app/code/server.pytorch.py.tpl

@@ -13,7 +13,7 @@ parameters = ndarrays_to_parameters(ndarrays)
 
 def server_fn(context: Context):
     # Read from config
-    num_rounds = int(context.run_config["num-server-rounds"])
+    num_rounds = context.run_config["num-server-rounds"]
 
     # Define strategy
     strategy = FedAvg(

src/py/flwr/cli/new/templates/app/code/server.sklearn.py.tpl

@@ -7,7 +7,7 @@ from flwr.server.strategy import FedAvg
 
 def server_fn(context: Context):
     # Read from config
-    num_rounds = int(context.run_config["num-server-rounds"])
+    num_rounds = context.run_config["num-server-rounds"]
 
     # Define strategy
     strategy = FedAvg(
@@ -19,5 +19,6 @@ def server_fn(context: Context):
 
     return ServerAppComponents(strategy=strategy, config=config)
 
+
 # Create ServerApp
 app = ServerApp(server_fn=server_fn)

src/py/flwr/cli/new/templates/app/code/server.tensorflow.py.tpl

@@ -13,7 +13,7 @@ parameters = ndarrays_to_parameters(load_model().get_weights())
 
 def server_fn(context: Context):
     # Read from config
-    num_rounds = int(context.run_config["num-server-rounds"])
+    num_rounds = context.run_config["num-server-rounds"]
 
     # Define strategy
     strategy = strategy = FedAvg(

src/py/flwr/cli/new/templates/app/pyproject.flowertune.toml.tpl

@@ -30,7 +30,7 @@ serverapp = "$import_name.app:server"
 clientapp = "$import_name.app:client"
 
 [tool.flwr.app.config]
-num-server-rounds = "3"
+num-server-rounds = 3
 
 [tool.flwr.federations]
 default = "localhost"

src/py/flwr/cli/new/templates/app/pyproject.hf.toml.tpl

@@ -28,8 +28,8 @@ serverapp = "$import_name.server_app:app"
 clientapp = "$import_name.client_app:app"
 
 [tool.flwr.app.config]
-num-server-rounds = "3"
-local-epochs = "1"
+num-server-rounds = 3
+local-epochs = 1
 
 [tool.flwr.federations]
 default = "localhost"

src/py/flwr/cli/new/templates/app/pyproject.jax.toml.tpl

@@ -25,7 +25,7 @@ serverapp = "$import_name.server_app:app"
 clientapp = "$import_name.client_app:app"
 
 [tool.flwr.app.config]
-num-server-rounds = "3"
+num-server-rounds = 3
 
 [tool.flwr.federations]
 default = "localhost"

src/py/flwr/cli/new/templates/app/pyproject.mlx.toml.tpl

@@ -25,12 +25,12 @@ serverapp = "$import_name.server_app:app"
 clientapp = "$import_name.client_app:app"
 
 [tool.flwr.app.config]
-num-server-rounds = "3"
-local-epochs = "1"
-num-layers = "2"
-hidden-dim = "32"
-batch-size = "256"
-lr = "0.1"
+num-server-rounds = 3
+local-epochs = 1
+num-layers = 2
+hidden-dim = 32
+batch-size = 256
+lr = 0.1
 
 [tool.flwr.federations]
 default = "localhost"

src/py/flwr/cli/new/templates/app/pyproject.numpy.toml.tpl

@@ -23,7 +23,7 @@ serverapp = "$import_name.server_app:app"
 clientapp = "$import_name.client_app:app"
 
 [tool.flwr.app.config]
-num-server-rounds = "3"
+num-server-rounds = 3
 
 [tool.flwr.federations]
 default = "localhost"

src/py/flwr/cli/new/templates/app/pyproject.pytorch.toml.tpl

@@ -25,8 +25,8 @@ serverapp = "$import_name.server_app:app"
 clientapp = "$import_name.client_app:app"
 
 [tool.flwr.app.config]
-num-server-rounds = "3"
-local-epochs = "1"
+num-server-rounds = 3
+local-epochs = 1
 
 [tool.flwr.federations]
 default = "localhost"

src/py/flwr/cli/new/templates/app/pyproject.sklearn.toml.tpl

@@ -24,7 +24,7 @@ serverapp = "$import_name.server_app:app"
 clientapp = "$import_name.client_app:app"
 
 [tool.flwr.app.config]
-num-server-rounds = "3"
+num-server-rounds = 3
 
 [tool.flwr.federations]
 default = "localhost"

src/py/flwr/cli/new/templates/app/pyproject.tensorflow.toml.tpl

@@ -24,10 +24,10 @@ serverapp = "$import_name.server_app:app"
 clientapp = "$import_name.client_app:app"
 
 [tool.flwr.app.config]
-num-server-rounds = "3"
-local-epochs = "1"
-batch-size = "32"
-verbose = ""  # Empty string means False
+num-server-rounds = 3
+local-epochs = 1
+batch-size = 32
+verbose = false
 
 [tool.flwr.federations]
 default = "localhost"