blazefl 2.0.0.dev3__tar.gz → 2.0.0.dev5__tar.gz

blazefl-2.0.0.dev5/.python-version

@@ -0,0 +1 @@
+3.13

{blazefl-2.0.0.dev3 → blazefl-2.0.0.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blazefl
-Version: 2.0.0.dev3
+Version: 2.0.0.dev5
 Summary: A blazing-fast and lightweight simulation framework for Federated Learning.
 Author-email: kitsuyaazuma <kitsuyaazuma@gmail.com>
 License-File: LICENSE

{blazefl-2.0.0.dev3 → blazefl-2.0.0.dev5}/examples/quickstart-fedavg/config/config.yaml

@@ -14,4 +14,5 @@ dataset_root_dir: /tmp/quickstart-fedavg/dataset
 dataset_split_dir: /tmp/quickstart-fedavg/split
 share_dir: /tmp/quickstart-fedavg/share
 state_dir: /tmp/quickstart-fedavg/state
-serial: false
+parallel: true
+ipc_mode: storage

{blazefl-2.0.0.dev3 → blazefl-2.0.0.dev5}/examples/quickstart-fedavg/main.py

@@ -98,31 +98,32 @@ def main(cfg: DictConfig):
         batch_size=cfg.batch_size,
     )
     trainer: FedAvgBaseClientTrainer | FedAvgProcessPoolClientTrainer | None = None
-    if cfg.serial:
-        trainer = FedAvgBaseClientTrainer(
+    if cfg.parallel:
+        trainer = FedAvgProcessPoolClientTrainer(
             model_selector=model_selector,
             model_name=cfg.model_name,
             dataset=dataset,
+            share_dir=share_dir,
+            state_dir=state_dir,
+            seed=cfg.seed,
             device=device,
             num_clients=cfg.num_clients,
             epochs=cfg.epochs,
             lr=cfg.lr,
             batch_size=cfg.batch_size,
+            num_parallels=cfg.num_parallels,
+            ipc_mode=cfg.ipc_mode,
         )
     else:
-        trainer = FedAvgProcessPoolClientTrainer(
+        trainer = FedAvgBaseClientTrainer(
             model_selector=model_selector,
             model_name=cfg.model_name,
             dataset=dataset,
-            share_dir=share_dir,
-            state_dir=state_dir,
-            seed=cfg.seed,
             device=device,
             num_clients=cfg.num_clients,
             epochs=cfg.epochs,
             lr=cfg.lr,
             batch_size=cfg.batch_size,
-            num_parallels=cfg.num_parallels,
         )
     pipeline = FedAvgPipeline(handler=handler, trainer=trainer, writer=writer)
     try:

{blazefl-2.0.0.dev3 → blazefl-2.0.0.dev5}/examples/step-by-step-dsfl/algorithm/dsfl.py

@@ -233,7 +233,7 @@ class DSFLBaseServerHandler(BaseServerHandler[DSFLUplinkPackage, DSFLDownlinkPac
 
 
 @dataclass
-class DSFLDiskSharedData:
+class DSFLClientConfig:
     model_selector: DSFLModelSelector
     model_name: str
     dataset: DSFLPartitionedDataset
@@ -245,7 +245,6 @@ class DSFLDiskSharedData:
     kd_lr: float
     cid: int
     seed: int
-    payload: DSFLDownlinkPackage
     state_path: Path
 
 
@@ -258,7 +257,7 @@ class DSFLClientState:
 
 
 class DSFLProcessPoolClientTrainer(
-    ProcessPoolClientTrainer[DSFLUplinkPackage, DSFLDownlinkPackage, DSFLDiskSharedData]
+    ProcessPoolClientTrainer[DSFLUplinkPackage, DSFLDownlinkPackage, DSFLClientConfig]
 ):
     def __init__(
         self,
@@ -300,68 +299,76 @@ class DSFLProcessPoolClientTrainer(
         self.device = device
         self.num_clients = num_clients
         self.seed = seed
+        self.ipc_mode = "storage"
 
         if self.device == "cuda":
             self.device_count = torch.cuda.device_count()
 
     @staticmethod
-    def process_client(path: Path, device: str) -> Path:
-        data = torch.load(path, weights_only=False)
-        assert isinstance(data, DSFLDiskSharedData)
-
-        model = data.model_selector.select_model(data.model_name)
-        optimizer = torch.optim.SGD(model.parameters(), lr=data.lr)
+    def worker(
+        config: DSFLClientConfig | Path,
+        payload: DSFLDownlinkPackage | Path,
+        device: str,
+    ) -> Path:
+        assert isinstance(config, Path) and isinstance(payload, Path)
+        config_path, payload_path = config, payload
+        c = torch.load(config_path, weights_only=False)
+        p = torch.load(payload_path, weights_only=False)
+        assert isinstance(c, DSFLClientConfig) and isinstance(p, DSFLDownlinkPackage)
+
+        model = c.model_selector.select_model(c.model_name)
+        optimizer = torch.optim.SGD(model.parameters(), lr=c.lr)
         kd_optimizer: torch.optim.SGD | None = None
 
         state: DSFLClientState | None = None
-        if data.state_path.exists():
-            state = torch.load(data.state_path, weights_only=False)
+        if c.state_path.exists():
+            state = torch.load(c.state_path, weights_only=False)
             assert isinstance(state, DSFLClientState)
             RandomState.set_random_state(state.random)
             model.load_state_dict(state.model)
             optimizer.load_state_dict(state.optimizer)
             if state.kd_optimizer is not None:
-                kd_optimizer = torch.optim.SGD(model.parameters(), lr=data.kd_lr)
+                kd_optimizer = torch.optim.SGD(model.parameters(), lr=c.kd_lr)
                 kd_optimizer.load_state_dict(state.kd_optimizer)
         else:
-            seed_everything(data.seed, device=device)
+            seed_everything(c.seed, device=device)
 
         # Distill
-        open_dataset = data.dataset.get_dataset(type_="open", cid=None)
-        if data.payload.indices is not None and data.payload.soft_labels is not None:
-            global_soft_labels = list(torch.unbind(data.payload.soft_labels, dim=0))
-            global_indices = data.payload.indices.tolist()
+        open_dataset = c.dataset.get_dataset(type_="open", cid=None)
+        if p.indices is not None and p.soft_labels is not None:
+            global_soft_labels = list(torch.unbind(p.soft_labels, dim=0))
+            global_indices = p.indices.tolist()
             if kd_optimizer is None:
-                kd_optimizer = torch.optim.SGD(model.parameters(), lr=data.kd_lr)
+                kd_optimizer = torch.optim.SGD(model.parameters(), lr=c.kd_lr)
             DSFLBaseServerHandler.distill(
                 model=model,
                 optimizer=kd_optimizer,
-                dataset=data.dataset,
+                dataset=c.dataset,
                 global_soft_labels=global_soft_labels,
                 global_indices=global_indices,
-                kd_epochs=data.kd_epochs,
-                kd_batch_size=data.kd_batch_size,
+                kd_epochs=c.kd_epochs,
+                kd_batch_size=c.kd_batch_size,
                 device=device,
             )
 
         # Train
-        train_loader = data.dataset.get_dataloader(
+        train_loader = c.dataset.get_dataloader(
             type_="train",
-            cid=data.cid,
-            batch_size=data.batch_size,
+            cid=c.cid,
+            batch_size=c.batch_size,
         )
         DSFLProcessPoolClientTrainer.train(
             model=model,
             optimizer=optimizer,
             train_loader=train_loader,
             device=device,
-            epochs=data.epochs,
+            epochs=c.epochs,
         )
 
         # Predict
         open_loader = DataLoader(
-            Subset(open_dataset, data.payload.next_indices.tolist()),
-            batch_size=data.batch_size,
+            Subset(open_dataset, p.next_indices.tolist()),
+            batch_size=c.batch_size,
         )
         soft_labels = DSFLProcessPoolClientTrainer.predict(
             model=model,
@@ -370,10 +377,10 @@ class DSFLProcessPoolClientTrainer(
         )
 
         # Evaluate
-        test_loader = data.dataset.get_dataloader(
+        test_loader = c.dataset.get_dataloader(
             type_="test",
-            cid=data.cid,
-            batch_size=data.batch_size,
+            cid=c.cid,
+            batch_size=c.batch_size,
         )
         loss, acc = DSFLBaseServerHandler.evaulate(
             model=model,
@@ -383,19 +390,19 @@ class DSFLProcessPoolClientTrainer(
 
         package = DSFLUplinkPackage(
             soft_labels=soft_labels,
-            indices=data.payload.next_indices,
+            indices=p.next_indices,
             metadata={"loss": loss, "acc": acc},
         )
 
-        torch.save(package, path)
+        torch.save(package, config_path)
         state = DSFLClientState(
             random=RandomState.get_random_state(device=device),
             model=model.state_dict(),
             optimizer=optimizer.state_dict(),
             kd_optimizer=kd_optimizer.state_dict() if kd_optimizer else None,
         )
-        torch.save(state, data.state_path)
-        return path
+        torch.save(state, c.state_path)
+        return config_path
 
     @staticmethod
     def train(
@@ -442,10 +449,8 @@ class DSFLProcessPoolClientTrainer(
         soft_labels = torch.cat(soft_labels_list, dim=0)
         return soft_labels.cpu()
 
-    def get_shared_data(
-        self, cid: int, payload: DSFLDownlinkPackage
-    ) -> DSFLDiskSharedData:
-        data = DSFLDiskSharedData(
+    def get_client_config(self, cid: int) -> DSFLClientConfig:
+        data = DSFLClientConfig(
             model_selector=self.model_selector,
             model_name=self.model_name,
             dataset=self.dataset,
@@ -457,7 +462,6 @@ class DSFLProcessPoolClientTrainer(
             kd_lr=self.kd_lr,
             cid=cid,
             seed=self.seed,
-            payload=payload,
             state_path=self.state_dir.joinpath(f"{cid}.pt"),
         )
         return data

{blazefl-2.0.0.dev3 → blazefl-2.0.0.dev5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "blazefl"
-version = "2.0.0.dev3"
+version = "2.0.0.dev5"
 description = "A blazing-fast and lightweight simulation framework for Federated Learning."
 readme = "README.md"
 authors = [

{blazefl-2.0.0.dev3 → blazefl-2.0.0.dev5}/src/blazefl/contrib/__init__.py

@@ -8,11 +8,17 @@ extending the core functionalities of BlazeFL.
 from blazefl.contrib.fedavg import (
     FedAvgBaseClientTrainer,
     FedAvgBaseServerHandler,
+    FedAvgDownlinkPackage,
     FedAvgProcessPoolClientTrainer,
+    FedAvgThreadPoolClientTrainer,
+    FedAvgUplinkPackage,
 )
 
 __all__ = [
     "FedAvgBaseServerHandler",
     "FedAvgProcessPoolClientTrainer",
     "FedAvgBaseClientTrainer",
+    "FedAvgThreadPoolClientTrainer",
+    "FedAvgUplinkPackage",
+    "FedAvgDownlinkPackage",
 ]

{blazefl-2.0.0.dev3 → blazefl-2.0.0.dev5}/src/blazefl/contrib/fedavg.py

@@ -1,7 +1,9 @@
 import random
+import threading
 from copy import deepcopy
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Literal
 
 import torch
 from torch.utils.data import DataLoader
@@ -13,6 +15,7 @@ from blazefl.core import (
     ModelSelector,
     PartitionedDataset,
     ProcessPoolClientTrainer,
+    ThreadPoolClientTrainer,
 )
 from blazefl.utils import (
     RandomState,
@@ -431,7 +434,7 @@ class FedAvgBaseClientTrainer(
 
 
 @dataclass
-class FedAvgDiskSharedData:
+class FedAvgClientConfig:
     """
     Data structure representing shared data for parallel client training
     in the Federated Averaging (FedAvg) algorithm.
@@ -448,7 +451,6 @@ class FedAvgDiskSharedData:
         lr (float): Learning rate for the optimizer.
         cid (int): Client ID.
         seed (int): Seed for reproducibility.
-        payload (FedAvgDownlinkPackage): Downlink package with global model parameters.
         state_path (Path): Path to save the client's random state.
     """
 
@@ -460,13 +462,12 @@ class FedAvgDiskSharedData:
     lr: float
     cid: int
     seed: int
-    payload: FedAvgDownlinkPackage
     state_path: Path
 
 
 class FedAvgProcessPoolClientTrainer(
     ProcessPoolClientTrainer[
-        FedAvgUplinkPackage, FedAvgDownlinkPackage, FedAvgDiskSharedData
+        FedAvgUplinkPackage, FedAvgDownlinkPackage, FedAvgClientConfig
     ]
 ):
     """
@@ -488,6 +489,8 @@ class FedAvgProcessPoolClientTrainer(
         lr (float): Learning rate for the optimizer.
         seed (int): Seed for reproducibility.
         num_parallels (int): Number of parallel processes for training.
+        ipc_mode (Literal["storage", "shared_memory"]):
+            Inter-process communication mode.
         device_count (int | None): Number of CUDA devices available (if using GPU).
     """
 
@@ -505,6 +508,7 @@ class FedAvgProcessPoolClientTrainer(
         lr: float,
         seed: int,
         num_parallels: int,
+        ipc_mode: Literal["storage", "shared_memory"],
     ) -> None:
         """
         Initialize the FedAvgParalleClientTrainer.
@@ -542,50 +546,93 @@ class FedAvgProcessPoolClientTrainer(
         self.device = device
         self.num_clients = num_clients
         self.seed = seed
+        self.ipc_mode = ipc_mode
 
     @staticmethod
-    def process_client(path: Path, device: str) -> Path:
+    def worker(
+        config: FedAvgClientConfig | Path,
+        payload: FedAvgDownlinkPackage | Path,
+        device: str,
+    ) -> FedAvgUplinkPackage | Path:
         """
         Process a single client's local training and evaluation.
 
-        This method is executed by a parallel process and handles data loading,
-        training, evaluation, and saving results to a shared file.
+        This method is executed by a worker process and handles loading client
+        configuration and payload, performing the client-specific training,
+        and returning the result.
 
         Args:
-            path (Path): Path to the shared data file containing client-specific
-            information.
-            device (str): Device to use for processing.
+            config (FedAvgClientConfig | Path):
+                The client's configuration data, or a path to a file containing
+                the configuration if `ipc_mode` is "storage".
+            payload (FedAvgDownlinkPackage | Path):
+                The downlink payload from the server, or a path to a file
+                containing the payload if `ipc_mode` is "storage".
+            device (str): Device to use for processing (e.g., "cpu", "cuda:0").
 
         Returns:
-            Path: Path to the file with the processed results.
-        """
-        data = torch.load(path, weights_only=False)
-        assert isinstance(data, FedAvgDiskSharedData)
-
-        if data.state_path.exists():
-            state = torch.load(data.state_path, weights_only=False)
-            assert isinstance(state, RandomState)
-            RandomState.set_random_state(state)
+            FedAvgUplinkPackage | Path:
+                The uplink package containing the client's results, or a path to
+                a file containing the package if `ipc_mode` is "storage".
+        """
+
+        def _storage_worker(
+            config_path: Path,
+            payload_path: Path,
+            device: str,
+        ) -> Path:
+            config = torch.load(config_path, weights_only=False)
+            assert isinstance(config, FedAvgClientConfig)
+            payload = torch.load(payload_path, weights_only=False)
+            assert isinstance(payload, FedAvgDownlinkPackage)
+            package = _shared_memory_worker(
+                config=config,
+                payload=payload,
+                device=device,
+            )
+            torch.save(package, config_path)
+            return config_path
+
+        def _shared_memory_worker(
+            config: FedAvgClientConfig,
+            payload: FedAvgDownlinkPackage,
+            device: str,
+        ) -> FedAvgUplinkPackage:
+            if config.state_path.exists():
+                state = torch.load(config.state_path, weights_only=False)
+                assert isinstance(state, RandomState)
+                RandomState.set_random_state(state)
+            else:
+                seed_everything(config.seed, device=device)
+
+            model = config.model_selector.select_model(config.model_name)
+            train_loader = config.dataset.get_dataloader(
+                type_="train",
+                cid=config.cid,
+                batch_size=config.batch_size,
+            )
+            package = FedAvgProcessPoolClientTrainer.train(
+                model=model,
+                model_parameters=payload.model_parameters,
+                train_loader=train_loader,
+                device=device,
+                epochs=config.epochs,
+                lr=config.lr,
+            )
+            torch.save(RandomState.get_random_state(device=device), config.state_path)
+            return package
+
+        if isinstance(config, Path) and isinstance(payload, Path):
+            return _storage_worker(config, payload, device)
+        elif isinstance(config, FedAvgClientConfig) and isinstance(
+            payload, FedAvgDownlinkPackage
+        ):
+            return _shared_memory_worker(config, payload, device)
         else:
-            seed_everything(data.seed, device=device)
-
-        model = data.model_selector.select_model(data.model_name)
-        train_loader = data.dataset.get_dataloader(
-            type_="train",
-            cid=data.cid,
-            batch_size=data.batch_size,
-        )
-        package = FedAvgProcessPoolClientTrainer.train(
-            model=model,
-            model_parameters=data.payload.model_parameters,
-            train_loader=train_loader,
-            device=device,
-            epochs=data.epochs,
-            lr=data.lr,
-        )
-        torch.save(package, path)
-        torch.save(RandomState.get_random_state(device=device), data.state_path)
-        return path
+            raise TypeError(
+                "Invalid types for config and payload."
+                " Expected FedAvgClientConfig and FedAvgDownlinkPackage or Path."
+            )
 
     @staticmethod
     def train(
@@ -595,6 +642,7 @@ class FedAvgProcessPoolClientTrainer(
         device: str,
         epochs: int,
         lr: float,
+        stop_event: threading.Event | None = None,
     ) -> FedAvgUplinkPackage:
         """
         Train the model with the given training data loader.
@@ -619,6 +667,8 @@ class FedAvgProcessPoolClientTrainer(
 
         data_size = 0
         for _ in range(epochs):
+            if stop_event is not None and stop_event.is_set():
+                break
             for data, target in train_loader:
                 data = data.to(device)
                 target = target.to(device)
@@ -636,21 +686,17 @@ class FedAvgProcessPoolClientTrainer(
 
         return FedAvgUplinkPackage(model_parameters, data_size)
 
-    def get_shared_data(
-        self, cid: int, payload: FedAvgDownlinkPackage
-    ) -> FedAvgDiskSharedData:
+    def get_client_config(self, cid: int) -> FedAvgClientConfig:
         """
-        Generate the shared data for a specific client.
+        Generate the client configuration for a specific client.
 
         Args:
             cid (int): Client ID.
-            payload (FedAvgDownlinkPackage): Downlink package with global model
-            parameters.
 
         Returns:
-            FedAvgDiskSharedData: Shared data structure for the client.
+            FedAvgClientConfig: Client configuration data structure.
         """
-        data = FedAvgDiskSharedData(
+        data = FedAvgClientConfig(
             model_selector=self.model_selector,
             model_name=self.model_name,
             dataset=self.dataset,
@@ -659,7 +705,6 @@ class FedAvgProcessPoolClientTrainer(
             lr=self.lr,
             cid=cid,
             seed=self.seed,
-            payload=payload,
             state_path=self.state_dir.joinpath(f"{cid}.pt"),
         )
         return data
@@ -674,3 +719,67 @@ class FedAvgProcessPoolClientTrainer(
         package = deepcopy(self.cache)
         self.cache = []
         return package
+
+
+class FedAvgThreadPoolClientTrainer(
+    ThreadPoolClientTrainer[
+        FedAvgUplinkPackage,
+        FedAvgDownlinkPackage,
+    ]
+):
+    def __init__(
+        self,
+        model_selector: ModelSelector,
+        model_name: str,
+        dataset: PartitionedDataset,
+        device: str,
+        num_clients: int,
+        epochs: int,
+        batch_size: int,
+        lr: float,
+        seed: int,
+        num_parallels: int,
+    ) -> None:
+        self.num_parallels = num_parallels
+        self.device = device
+        if self.device == "cuda":
+            self.device_count = torch.cuda.device_count()
+        self.cache: list[FedAvgUplinkPackage] = []
+
+        self.model_selector = model_selector
+        self.model_name = model_name
+        self.dataset = dataset
+        self.epochs = epochs
+        self.batch_size = batch_size
+        self.lr = lr
+        self.num_clients = num_clients
+        self.seed = seed
+        self.stop_event = None
+
+    def worker(
+        self,
+        cid: int,
+        device: str,
+        payload: FedAvgDownlinkPackage,
+    ) -> FedAvgUplinkPackage:
+        model = self.model_selector.select_model(self.model_name)
+        train_loader = self.dataset.get_dataloader(
+            type_="train",
+            cid=cid,
+            batch_size=self.batch_size,
+        )
+        package = FedAvgProcessPoolClientTrainer.train(
+            model=model,
+            model_parameters=payload.model_parameters,
+            train_loader=train_loader,
+            device=device,
+            epochs=self.epochs,
+            lr=self.lr,
+            stop_event=self.stop_event,
+        )
+        return package
+
+    def uplink_package(self) -> list[FedAvgUplinkPackage]:
+        package = deepcopy(self.cache)
+        self.cache = []
+        return package