blazefl 2.0.0.dev4__tar.gz → 2.0.0.dev6__tar.gz

blazefl-2.0.0.dev6/.python-version

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

{blazefl-2.0.0.dev4 → blazefl-2.0.0.dev6}/PKG-INFO

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

{blazefl-2.0.0.dev4 → blazefl-2.0.0.dev6}/examples/quickstart-fedavg/config/config.yaml

@@ -14,5 +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
-parallel: true
+execution_mode: multi-process
 ipc_mode: storage

{blazefl-2.0.0.dev4 → blazefl-2.0.0.dev6}/examples/quickstart-fedavg/main.py

@@ -9,6 +9,7 @@ from blazefl.contrib import (
     FedAvgBaseClientTrainer,
     FedAvgBaseServerHandler,
     FedAvgProcessPoolClientTrainer,
+    FedAvgThreadPoolClientTrainer,
 )
 from blazefl.utils import seed_everything
 from hydra.core import hydra_config
@@ -23,7 +24,9 @@ class FedAvgPipeline:
     def __init__(
         self,
         handler: FedAvgBaseServerHandler,
-        trainer: FedAvgBaseClientTrainer | FedAvgProcessPoolClientTrainer,
+        trainer: FedAvgBaseClientTrainer
+        | FedAvgProcessPoolClientTrainer
+        | FedAvgThreadPoolClientTrainer,
         writer: SummaryWriter,
     ) -> None:
         self.handler = handler
@@ -97,41 +100,60 @@ def main(cfg: DictConfig):
         sample_ratio=cfg.sample_ratio,
         batch_size=cfg.batch_size,
     )
-    trainer: FedAvgBaseClientTrainer | FedAvgProcessPoolClientTrainer | None = None
-    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 = FedAvgBaseClientTrainer(
-            model_selector=model_selector,
-            model_name=cfg.model_name,
-            dataset=dataset,
-            device=device,
-            num_clients=cfg.num_clients,
-            epochs=cfg.epochs,
-            lr=cfg.lr,
-            batch_size=cfg.batch_size,
-        )
+    trainer: (
+        FedAvgBaseClientTrainer
+        | FedAvgProcessPoolClientTrainer
+        | FedAvgThreadPoolClientTrainer
+        | None
+    ) = None
+    match cfg.execution_mode:
+        case "multi-process":
+            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,
+            )
+        case "single-thread":
+            trainer = FedAvgBaseClientTrainer(
+                model_selector=model_selector,
+                model_name=cfg.model_name,
+                dataset=dataset,
+                device=device,
+                num_clients=cfg.num_clients,
+                epochs=cfg.epochs,
+                lr=cfg.lr,
+                batch_size=cfg.batch_size,
+            )
+        case "multi-thread":
+            trainer = FedAvgThreadPoolClientTrainer(
+                model_selector=model_selector,
+                model_name=cfg.model_name,
+                dataset=dataset,
+                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,
+            )
+        case _:
+            raise ValueError(f"Invalid execution mode: {cfg.execution_mode}")
     pipeline = FedAvgPipeline(handler=handler, trainer=trainer, writer=writer)
     try:
         pipeline.main()
     except KeyboardInterrupt:
-        logging.info("KeyboardInterrupt: Stopping the pipeline.")
-    except Exception as e:
-        logging.exception(f"An error occurred: {e}")
+        logging.info("KeyboardInterrupt")
 
 
 if __name__ == "__main__":

{blazefl-2.0.0.dev4 → blazefl-2.0.0.dev6}/examples/step-by-step-dsfl/main.py

@@ -120,9 +120,7 @@ def main(
     try:
         pipeline.main()
     except KeyboardInterrupt:
-        logging.info("KeyboardInterrupt: Stopping the pipeline.")
-    except Exception as e:
-        logging.exception(f"An error occurred: {e}")
+        logging.info("KeyboardInterrupt")
 
 
 if __name__ == "__main__":

{blazefl-2.0.0.dev4 → blazefl-2.0.0.dev6}/pyproject.toml

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

{blazefl-2.0.0.dev4 → blazefl-2.0.0.dev6}/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.dev4 → blazefl-2.0.0.dev6}/src/blazefl/contrib/fedavg.py

@@ -1,10 +1,12 @@
 import random
+import threading
 from copy import deepcopy
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Literal
 
 import torch
+import torch.multiprocessing as mp
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
@@ -14,6 +16,7 @@ from blazefl.core import (
     ModelSelector,
     PartitionedDataset,
     ProcessPoolClientTrainer,
+    ThreadPoolClientTrainer,
 )
 from blazefl.utils import (
     RandomState,
@@ -545,12 +548,15 @@ class FedAvgProcessPoolClientTrainer(
         self.num_clients = num_clients
         self.seed = seed
         self.ipc_mode = ipc_mode
+        self.manager = mp.Manager()
+        self.stop_event = self.manager.Event()
 
     @staticmethod
     def worker(
         config: FedAvgClientConfig | Path,
         payload: FedAvgDownlinkPackage | Path,
         device: str,
+        stop_event: threading.Event,
     ) -> FedAvgUplinkPackage | Path:
         """
         Process a single client's local training and evaluation.
@@ -578,6 +584,7 @@ class FedAvgProcessPoolClientTrainer(
             config_path: Path,
             payload_path: Path,
             device: str,
+            stop_event: threading.Event,
         ) -> Path:
             config = torch.load(config_path, weights_only=False)
             assert isinstance(config, FedAvgClientConfig)
@@ -587,6 +594,7 @@ class FedAvgProcessPoolClientTrainer(
                 config=config,
                 payload=payload,
                 device=device,
+                stop_event=stop_event,
             )
             torch.save(package, config_path)
             return config_path
@@ -595,6 +603,7 @@ class FedAvgProcessPoolClientTrainer(
             config: FedAvgClientConfig,
             payload: FedAvgDownlinkPackage,
             device: str,
+            stop_event: threading.Event,
         ) -> FedAvgUplinkPackage:
             if config.state_path.exists():
                 state = torch.load(config.state_path, weights_only=False)
@@ -616,16 +625,17 @@ class FedAvgProcessPoolClientTrainer(
                 device=device,
                 epochs=config.epochs,
                 lr=config.lr,
+                stop_event=stop_event,
             )
             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)
+            return _storage_worker(config, payload, device, stop_event)
         elif isinstance(config, FedAvgClientConfig) and isinstance(
             payload, FedAvgDownlinkPackage
         ):
-            return _shared_memory_worker(config, payload, device)
+            return _shared_memory_worker(config, payload, device, stop_event)
         else:
             raise TypeError(
                 "Invalid types for config and payload."
@@ -640,6 +650,7 @@ class FedAvgProcessPoolClientTrainer(
         device: str,
         epochs: int,
         lr: float,
+        stop_event: threading.Event,
     ) -> FedAvgUplinkPackage:
         """
         Train the model with the given training data loader.
@@ -664,6 +675,8 @@ class FedAvgProcessPoolClientTrainer(
 
         data_size = 0
         for _ in range(epochs):
+            if stop_event.is_set():
+                break
             for data, target in train_loader:
                 data = data.to(device)
                 target = target.to(device)
@@ -714,3 +727,120 @@ 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 = threading.Event()
+
+    def worker(
+        self,
+        cid: int,
+        device: str,
+        payload: FedAvgDownlinkPackage,
+        stop_event: threading.Event,
+    ) -> 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 = self.train(
+            model=model,
+            model_parameters=payload.model_parameters,
+            train_loader=train_loader,
+            device=device,
+            epochs=self.epochs,
+            lr=self.lr,
+            stop_event=stop_event,
+        )
+        return package
+
+    def train(
+        self,
+        model: torch.nn.Module,
+        model_parameters: torch.Tensor,
+        train_loader: DataLoader,
+        device: str,
+        epochs: int,
+        lr: float,
+        stop_event: threading.Event,
+    ) -> FedAvgUplinkPackage:
+        """
+        Train the model with the given training data loader.
+
+        Args:
+            model (torch.nn.Module): The model to train.
+            model_parameters (torch.Tensor): Initial global model parameters.
+            train_loader (DataLoader): DataLoader for the training data.
+            device (str): Device to run the training on.
+            epochs (int): Number of local training epochs.
+            lr (float): Learning rate for the optimizer.
+
+        Returns:
+            FedAvgUplinkPackage: Uplink package containing updated model parameters
+            and data size.
+        """
+        model.to(device)
+        deserialize_model(model, model_parameters)
+        model.train()
+        optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+        criterion = torch.nn.CrossEntropyLoss()
+
+        data_size = 0
+        for _ in range(epochs):
+            if stop_event.is_set():
+                break
+            for data, target in train_loader:
+                data = data.to(device)
+                target = target.to(device)
+
+                output = model(data)
+                loss = criterion(output, target)
+
+                data_size += len(target)
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+        model_parameters = serialize_model(model)
+
+        return FedAvgUplinkPackage(model_parameters, data_size)
+
+    def uplink_package(self) -> list[FedAvgUplinkPackage]:
+        package = deepcopy(self.cache)
+        self.cache = []
+        return package

{blazefl-2.0.0.dev4 → blazefl-2.0.0.dev6}/src/blazefl/core/client_trainer.py

@@ -1,11 +1,12 @@
-import multiprocessing as mp
 import signal
+import threading
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from multiprocessing.pool import ApplyResult
 from pathlib import Path
 from typing import Literal, Protocol, TypeVar
 
 import torch
+import torch.multiprocessing as mp
 from tqdm import tqdm
 
 from blazefl.utils import move_tensor_to_shared_memory
@@ -82,6 +83,7 @@ class ProcessPoolClientTrainer(
     device_count: int
     cache: list[UplinkPackage]
     ipc_mode: Literal["storage", "shared_memory"] = "storage"
+    stop_event: threading.Event
 
     def get_client_config(self, cid: int) -> ClientConfig:
         """
@@ -111,7 +113,10 @@ class ProcessPoolClientTrainer(
 
     @staticmethod
     def worker(
-        config: ClientConfig | Path, payload: DownlinkPackage | Path, device: str
+        config: ClientConfig | Path,
+        payload: DownlinkPackage | Path,
+        device: str,
+        stop_event: threading.Event,
     ) -> UplinkPackage | Path:
         """
         Process a single client's training task.
@@ -157,11 +162,13 @@ class ProcessPoolClientTrainer(
         else:  # shared_memory
             move_tensor_to_shared_memory(payload)
 
-        with mp.Pool(
+        self.stop_event.clear()
+        pool = mp.Pool(
             processes=self.num_parallels,
             initializer=signal.signal,
             initargs=(signal.SIGINT, signal.SIG_IGN),
-        ) as pool:
+        )
+        try:
             jobs: list[ApplyResult] = []
             for cid in cid_list:
                 config = self.get_client_config(cid)
@@ -171,12 +178,15 @@ class ProcessPoolClientTrainer(
                     torch.save(config, config_path)
                     jobs.append(
                         pool.apply_async(
-                            self.worker, (config_path, payload_path, device)
+                            self.worker,
+                            (config_path, payload_path, device, self.stop_event),
                         )
                     )
                 else:  # shared_memory
                     jobs.append(
-                        pool.apply_async(self.worker, (config, payload, device))
+                        pool.apply_async(
+                            self.worker, (config, payload, device, self.stop_event)
+                        )
                     )
 
             for job in tqdm(jobs, desc="Client", leave=False):
@@ -187,6 +197,10 @@ class ProcessPoolClientTrainer(
                 else:  # shared_memory
                     package = result
                 self.cache.append(package)
+        finally:
+            self.stop_event.set()
+            pool.close()
+            pool.join()
 
 
 class ThreadPoolClientTrainer(
@@ -197,12 +211,14 @@ class ThreadPoolClientTrainer(
     device: str
     device_count: int
     cache: list[UplinkPackage]
+    stop_event: threading.Event
 
     def worker(
         self,
         cid: int,
         device: str,
         payload: DownlinkPackage,
+        stop_event: threading.Event,
     ) -> UplinkPackage:
         """
         Process a single client's training task in a thread.
@@ -211,6 +227,7 @@ class ThreadPoolClientTrainer(
             cid (int): The client ID.
             device (str): The device to use for processing this client.
             payload (DownlinkPackage): The data package received from the server.
+            stop_event (threading.Event): Event to signal stopping the worker.
 
         Returns:
             UplinkPackage: The uplink package containing the client's results.
@@ -223,7 +240,9 @@ class ThreadPoolClientTrainer(
         return self.device
 
     def local_process(self, payload: DownlinkPackage, cid_list: list[int]) -> None:
-        with ThreadPoolExecutor(max_workers=self.num_parallels) as executor:
+        self.stop_event.clear()
+        executor = ThreadPoolExecutor(max_workers=self.num_parallels)
+        try:
             futures = []
             for cid in cid_list:
                 device = self.get_client_device(cid)
@@ -232,11 +251,18 @@ class ThreadPoolClientTrainer(
                     cid,
                     device,
                     payload,
+                    self.stop_event,
                 )
                 futures.append(future)
 
             for future in tqdm(
-                as_completed(futures), total=len(futures), desc="Client", leave=False
+                as_completed(futures),
+                total=len(futures),
+                desc="Client",
+                leave=False,
             ):
                 result = future.result()
                 self.cache.append(result)
+        finally:
+            self.stop_event.set()
+            executor.shutdown(wait=True, cancel_futures=True)

{blazefl-2.0.0.dev4 → blazefl-2.0.0.dev6}/src/blazefl/core/client_trainer.pyi

@@ -1,3 +1,4 @@
+import threading
 from blazefl.utils import move_tensor_to_shared_memory as move_tensor_to_shared_memory
 from multiprocessing.pool import ApplyResult as ApplyResult
 from pathlib import Path
@@ -18,10 +19,11 @@ class ProcessPoolClientTrainer(BaseClientTrainer[UplinkPackage, DownlinkPackage]
     device_count: int
     cache: list[UplinkPackage]
     ipc_mode: Literal['storage', 'shared_memory']
+    stop_event: threading.Event
     def get_client_config(self, cid: int) -> ClientConfig: ...
     def get_client_device(self, cid: int) -> str: ...
     @staticmethod
-    def worker(config: ClientConfig | Path, payload: DownlinkPackage | Path, device: str) -> UplinkPackage | Path: ...
+    def worker(config: ClientConfig | Path, payload: DownlinkPackage | Path, device: str, stop_event: threading.Event) -> UplinkPackage | Path: ...
     def local_process(self, payload: DownlinkPackage, cid_list: list[int]) -> None: ...
 
 class ThreadPoolClientTrainer(BaseClientTrainer[UplinkPackage, DownlinkPackage], Protocol[UplinkPackage, DownlinkPackage]):
@@ -29,6 +31,7 @@ class ThreadPoolClientTrainer(BaseClientTrainer[UplinkPackage, DownlinkPackage],
     device: str
     device_count: int
     cache: list[UplinkPackage]
-    def worker(self, cid: int, device: str, payload: DownlinkPackage) -> UplinkPackage: ...
+    stop_event: threading.Event
+    def worker(self, cid: int, device: str, payload: DownlinkPackage, stop_event: threading.Event) -> UplinkPackage: ...
     def get_client_device(self, cid: int) -> str: ...
     def local_process(self, payload: DownlinkPackage, cid_list: list[int]) -> None: ...