Completed migrating clients to the new strategy API. (#395)

* Migrated examples/customized/custom_client.py to the new composable client API.

* ruff format & ruff check --fix.

* Migrated some examples to use the latest client API using strategies.

* Expanded the test coverage for registry-style components across models, datasources, and samplers.

* Migrated all remaining examples to the new client API.

Author: baochunli

.github/workflows/delete_workflow_runs.yml

@@ -1,8 +1,7 @@
 # This workflow will delete old workflow runs, on a schedule or manually triggered.
 name: Delete old workflow runs
 permissions:
-  contents: read
-  pull-requests: read
+  contents: read/write
 
 on:
   workflow_dispatch:

examples/basic/basic.py

@@ -14,8 +14,8 @@
 from torchvision.transforms import ToTensor
 
 from plato.clients import simple
-from plato.datasources import base
 from plato.config import Config
+from plato.datasources import base
 from plato.servers import fedavg
 from plato.trainers.composable import ComposableTrainer
 from plato.trainers.strategies.base import (
@@ -38,8 +38,12 @@ def __init__(self):
         Config()
         base_path = Path(Config.params.get("base_path", "./runtime"))
         data_dir = Path(Config.params.get("data_path", base_path / "data"))
-        self.trainset = MNIST(str(data_dir), train=True, download=True, transform=ToTensor())
-        self.testset = MNIST(str(data_dir), train=False, download=True, transform=ToTensor())
+        self.trainset = MNIST(
+            str(data_dir), train=True, download=True, transform=ToTensor()
+        )
+        self.testset = MNIST(
+            str(data_dir), train=False, download=True, transform=ToTensor()
+        )
 
 
 class MNISTTrainingStepStrategy(TrainingStepStrategy):

examples/basic/server_strategies.py

@@ -62,8 +62,12 @@ def __init__(self):
         Config()
         base_path = Path(Config.params.get("base_path", "./runtime"))
         data_dir = Path(Config.params.get("data_path", base_path / "data"))
-        self.trainset = MNIST(str(data_dir), train=True, download=True, transform=ToTensor())
-        self.testset = MNIST(str(data_dir), train=False, download=True, transform=ToTensor())
+        self.trainset = MNIST(
+            str(data_dir), train=True, download=True, transform=ToTensor()
+        )
+        self.testset = MNIST(
+            str(data_dir), train=False, download=True, transform=ToTensor()
+        )
 
 
 class MNISTTrainingStepStrategy(TrainingStepStrategy):

examples/client_selection/afl/afl.py

@@ -16,7 +16,7 @@
 
 def main():
     """A Plato federated learning training session using the AFL algorithm."""
-    client = afl_client.Client()
+    client = afl_client.create_client()
     server = afl_server.Server()
     server.run(client)
 

examples/client_selection/afl/afl_client.py

@@ -12,8 +12,7 @@
 
 import logging
 import math
-from types import SimpleNamespace
-from typing import Iterable, Optional
+from typing import Iterable, List, Optional
 
 import torch
 
@@ -164,38 +163,51 @@ def _get_pre_training_loss(context: ClientContext) -> Optional[float]:
         return 0.0
 
 
-class Client(simple.Client):
-    """A federated learning client for AFL."""
-
-    def __init__(
-        self,
-        model=None,
-        datasource=None,
-        algorithm=None,
-        trainer=None,
-        callbacks=None,
-        trainer_callbacks: Optional[Iterable] = None,
+def _ensure_pretraining_callback(
+    trainer_callbacks: Optional[Iterable],
+) -> List:
+    """Ensure AFL's pre-training loss callback is present once."""
+    callbacks_list = list(trainer_callbacks) if trainer_callbacks else []
+    if not any(
+        cb == AFLPreTrainingLossCallback
+        or getattr(cb, "__class__", None) == AFLPreTrainingLossCallback
+        for cb in callbacks_list
     ):
-        callbacks_list = list(trainer_callbacks) if trainer_callbacks else []
-        if not any(
-            cb == AFLPreTrainingLossCallback
-            or getattr(cb, "__class__", None) == AFLPreTrainingLossCallback
-            for cb in callbacks_list
-        ):
-            callbacks_list.append(AFLPreTrainingLossCallback)
-
-        super().__init__(
-            model=model,
-            datasource=datasource,
-            algorithm=algorithm,
-            trainer=trainer,
-            callbacks=callbacks,
-            trainer_callbacks=callbacks_list,
-        )
-        self._configure_composable(
-            lifecycle_strategy=self.lifecycle_strategy,
-            payload_strategy=self.payload_strategy,
-            training_strategy=self.training_strategy,
-            reporting_strategy=AFLReportingStrategy(),
-            communication_strategy=self.communication_strategy,
-        )
+        callbacks_list.append(AFLPreTrainingLossCallback)
+    return callbacks_list
+
+
+def create_client(
+    *,
+    model=None,
+    datasource=None,
+    algorithm=None,
+    trainer=None,
+    callbacks=None,
+    trainer_callbacks: Optional[Iterable] = None,
+):
+    """Build an AFL client configured with valuation hooks."""
+    callbacks_list = _ensure_pretraining_callback(trainer_callbacks)
+
+    client = simple.Client(
+        model=model,
+        datasource=datasource,
+        algorithm=algorithm,
+        trainer=trainer,
+        callbacks=callbacks,
+        trainer_callbacks=callbacks_list,
+    )
+
+    client._configure_composable(
+        lifecycle_strategy=client.lifecycle_strategy,
+        payload_strategy=client.payload_strategy,
+        training_strategy=client.training_strategy,
+        reporting_strategy=AFLReportingStrategy(),
+        communication_strategy=client.communication_strategy,
+    )
+
+    return client
+
+
+# Maintain compatibility for previous imports that expected a Client callable.
+Client = create_client

examples/client_selection/afl/afl_strategy.py

@@ -16,7 +16,7 @@
 
 def main():
     """A Plato federated learning training session using AFL strategy."""
-    client = afl_client.Client()
+    client = afl_client.create_client()
     server = afl_server_strategy.Server()
     server.run(client)
 

examples/client_selection/oort/oort.py

@@ -18,7 +18,7 @@
 def main():
     """A Plato federated learning training session using Oort strategy."""
     trainer = oort_trainer.Trainer
-    client = oort_client.Client(trainer=trainer)
+    client = oort_client.create_client(trainer=trainer)
     server = oort_server.Server(trainer=trainer)
     server.run(client)
 

examples/client_selection/oort/oort_client.py

@@ -39,31 +39,35 @@ def build_report(self, context: ClientContext, report):
         return report
 
 
-class Client(simple.Client):
-    """A federated learning client that calculates its statistical utility."""
-
-    def __init__(
-        self,
-        model=None,
-        datasource=None,
-        algorithm=None,
-        trainer=None,
-        callbacks=None,
-        trainer_callbacks=None,
-    ):
-        super().__init__(
-            model=model,
-            datasource=datasource,
-            algorithm=algorithm,
-            trainer=trainer,
-            callbacks=callbacks,
-            trainer_callbacks=trainer_callbacks,
-        )
+def create_client(
+    *,
+    model=None,
+    datasource=None,
+    algorithm=None,
+    trainer=None,
+    callbacks=None,
+    trainer_callbacks=None,
+):
+    """Build an Oort client configured with statistical utility reporting."""
+    client = simple.Client(
+        model=model,
+        datasource=datasource,
+        algorithm=algorithm,
+        trainer=trainer,
+        callbacks=callbacks,
+        trainer_callbacks=trainer_callbacks,
+    )
 
-        self._configure_composable(
-            lifecycle_strategy=self.lifecycle_strategy,
-            payload_strategy=self.payload_strategy,
-            training_strategy=self.training_strategy,
-            reporting_strategy=OortReportingStrategy(),
-            communication_strategy=self.communication_strategy,
-        )
+    client._configure_composable(
+        lifecycle_strategy=client.lifecycle_strategy,
+        payload_strategy=client.payload_strategy,
+        training_strategy=client.training_strategy,
+        reporting_strategy=OortReportingStrategy(),
+        communication_strategy=client.communication_strategy,
+    )
+
+    return client
+
+
+# Maintain compatibility for imports expecting a Client callable.
+Client = create_client

examples/client_selection/pisces/pisces.py

@@ -17,7 +17,7 @@
 def main():
     """Pisces: an asynchronous client selection and server aggregation algorithm."""
     trainer = pisces_trainer.Trainer
-    client = pisces_client.Client(trainer=trainer)
+    client = pisces_client.create_client(trainer=trainer)
     server = pisces_server.Server(trainer=trainer)
 
     server.run(client)

examples/client_selection/pisces/pisces_client.py

@@ -59,35 +59,38 @@ def build_report(self, context: ClientContext, report):
         return report
 
 
-class Client(simple.Client):
+def create_client(
+    *,
+    model=None,
+    datasource=None,
+    algorithm=None,
+    trainer=None,
+    callbacks=None,
+    trainer_callbacks=None,
+    loss_decay: float = 1e-2,
+):
     """
-    A Pisces federated learning client who sends weight updates and client statistical utility.
+    Build a Pisces client that reports statistical utility with an EMA of squared loss.
     """
-
-    def __init__(
-        self,
-        model=None,
-        datasource=None,
-        algorithm=None,
-        trainer=None,
-        callbacks=None,
-        trainer_callbacks=None,
-        *,
-        loss_decay: float = 1e-2,
-    ):
-        super().__init__(
-            model=model,
-            datasource=datasource,
-            algorithm=algorithm,
-            trainer=trainer,
-            callbacks=callbacks,
-            trainer_callbacks=trainer_callbacks,
-        )
-
-        self._configure_composable(
-            lifecycle_strategy=self.lifecycle_strategy,
-            payload_strategy=self.payload_strategy,
-            training_strategy=self.training_strategy,
-            reporting_strategy=PiscesReportingStrategy(loss_decay),
-            communication_strategy=self.communication_strategy,
-        )
+    client = simple.Client(
+        model=model,
+        datasource=datasource,
+        algorithm=algorithm,
+        trainer=trainer,
+        callbacks=callbacks,
+        trainer_callbacks=trainer_callbacks,
+    )
+
+    client._configure_composable(
+        lifecycle_strategy=client.lifecycle_strategy,
+        payload_strategy=client.payload_strategy,
+        training_strategy=client.training_strategy,
+        reporting_strategy=PiscesReportingStrategy(loss_decay),
+        communication_strategy=client.communication_strategy,
+    )
+
+    return client
+
+
+# Maintain compatibility for imports expecting a Client callable.
+Client = create_client