blazefl 2.0.0.dev2__tar.gz → 2.0.0.dev4__tar.gz

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

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

{blazefl-2.0.0.dev2 → blazefl-2.0.0.dev4}/examples/experimental-freethreaded/main.py

@@ -12,11 +12,10 @@ from blazefl.contrib import (
     FedAvgServerHandler,
 )
 from blazefl.contrib.fedavg import FedAvgDownlinkPackage, FedAvgUplinkPackage
-from blazefl.core import ModelSelector, PartitionedDataset
+from blazefl.core import ModelSelector, MultiThreadClientTrainer, PartitionedDataset
 from blazefl.utils import seed_everything
 from omegaconf import DictConfig, OmegaConf
 
-from core.client_trainer import MultiThreadClientTrainer
 from dataset import PartitionedCIFAR10
 from models import FedAvgModelSelector
 

{blazefl-2.0.0.dev2 → blazefl-2.0.0.dev4}/examples/experimental-freethreaded/pyproject.toml

@@ -5,7 +5,7 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
-    "blazefl>=2.0.0dev1",
+    "blazefl>=2.0.0dev2",
     "hydra-core>=1.3.2",
     "torch>=2.7.1",
     "torchvision>=0.22.1",

{blazefl-2.0.0.dev2 → blazefl-2.0.0.dev4}/examples/experimental-freethreaded/uv.lock

@@ -10,16 +10,16 @@ sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d
 
 [[package]]
 name = "blazefl"
-version = "2.0.0.dev1"
+version = "2.0.0.dev2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "numpy" },
     { name = "torch" },
     { name = "tqdm" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/9d/4a/b82e557a19b3dc957e43fc29c73bb3b9ddbb78575de20e2a98a7db954c51/blazefl-2.0.0.dev1.tar.gz", hash = "sha256:aeef9d1a6835ae240b9772144a1ca8ee2691a9252a980173b1c81d1c7c453596", size = 591232, upload_time = "2025-06-10T08:28:25.786Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/9b/92/3f668d7259a81d13d5d2e6089d8022629a09058ddbcae33f5873d3871656/blazefl-2.0.0.dev2.tar.gz", hash = "sha256:28a3e09d6f6cec8d8e8bb4da90adaa0c795980c74f9d4ce65024ba762726b595", size = 613563, upload_time = "2025-06-10T09:24:08.981Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/99/c0/3a4e12feee7a3a9e52a3eba30816a7ce7d4a11e3b11989b72c01419e1116/blazefl-2.0.0.dev1-py3-none-any.whl", hash = "sha256:1a158f673ac8f644b0a273d6cdb6d32b0ff17a1f479dffd930a06873f023042a", size = 23714, upload_time = "2025-06-10T08:28:24.297Z" },
+    { url = "https://files.pythonhosted.org/packages/e3/7c/841c8b9d22caaefc1dcc167076292fd1479309b3e11eb6fc1b24141792a8/blazefl-2.0.0.dev2-py3-none-any.whl", hash = "sha256:88639018166dafa79fde59aa5575bfb7277bf9e0e4497f22a59e9c7ec34efa4b", size = 23999, upload_time = "2025-06-10T09:24:07.553Z" },
 ]
 
 [[package]]
@@ -50,7 +50,7 @@ dev = [
 
 [package.metadata]
 requires-dist = [
-    { name = "blazefl", specifier = ">=2.0.0.dev1" },
+    { name = "blazefl", specifier = ">=2.0.0.dev2" },
     { name = "hydra-core", specifier = ">=1.3.2" },
     { name = "torch", specifier = ">=2.7.1" },
     { name = "torchvision", specifier = ">=0.22.1" },

{blazefl-2.0.0.dev2 → blazefl-2.0.0.dev4}/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.dev2 → blazefl-2.0.0.dev4}/examples/quickstart-fedavg/main.py

@@ -6,9 +6,9 @@ import hydra
 import torch
 import torch.multiprocessing as mp
 from blazefl.contrib import (
-    FedAvgParallelClientTrainer,
-    FedAvgSerialClientTrainer,
-    FedAvgServerHandler,
+    FedAvgBaseClientTrainer,
+    FedAvgBaseServerHandler,
+    FedAvgProcessPoolClientTrainer,
 )
 from blazefl.utils import seed_everything
 from hydra.core import hydra_config
@@ -22,8 +22,8 @@ from models import FedAvgModelSelector
 class FedAvgPipeline:
     def __init__(
         self,
-        handler: FedAvgServerHandler,
-        trainer: FedAvgSerialClientTrainer | FedAvgParallelClientTrainer,
+        handler: FedAvgBaseServerHandler,
+        trainer: FedAvgBaseClientTrainer | FedAvgProcessPoolClientTrainer,
         writer: SummaryWriter,
     ) -> None:
         self.handler = handler
@@ -87,7 +87,7 @@ def main(cfg: DictConfig):
     )
     model_selector = FedAvgModelSelector(num_classes=10)
 
-    handler = FedAvgServerHandler(
+    handler = FedAvgBaseServerHandler(
         model_selector=model_selector,
         model_name=cfg.model_name,
         dataset=dataset,
@@ -97,32 +97,33 @@ def main(cfg: DictConfig):
         sample_ratio=cfg.sample_ratio,
         batch_size=cfg.batch_size,
     )
-    trainer: FedAvgSerialClientTrainer | FedAvgParallelClientTrainer | None = None
-    if cfg.serial:
-        trainer = FedAvgSerialClientTrainer(
+    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 = FedAvgParallelClientTrainer(
+        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.dev4/examples/step-by-step-dsfl/algorithm/__init__.py

@@ -0,0 +1,3 @@
+from algorithm.dsfl import DSFLBaseServerHandler, DSFLProcessPoolClientTrainer
+
+__all__ = ["DSFLBaseServerHandler", "DSFLProcessPoolClientTrainer"]

{blazefl-2.0.0.dev2 → blazefl-2.0.0.dev4}/examples/step-by-step-dsfl/algorithm/dsfl.py

@@ -7,8 +7,8 @@ from pathlib import Path
 import torch
 import torch.nn.functional as F
 from blazefl.core import (
-    ParallelClientTrainer,
-    ServerHandler,
+    BaseServerHandler,
+    ProcessPoolClientTrainer,
 )
 from blazefl.utils import (
     FilteredDataset,
@@ -35,7 +35,7 @@ class DSFLDownlinkPackage:
     next_indices: torch.Tensor
 
 
-class DSFLServerHandler(ServerHandler[DSFLUplinkPackage, DSFLDownlinkPackage]):
+class DSFLBaseServerHandler(BaseServerHandler[DSFLUplinkPackage, DSFLDownlinkPackage]):
     def __init__(
         self,
         model_selector: DSFLModelSelector,
@@ -116,7 +116,7 @@ class DSFLServerHandler(ServerHandler[DSFLUplinkPackage, DSFLDownlinkPackage]):
             era_soft_labels = F.softmax(mean_soft_labels / self.era_temperature, dim=0)
             global_soft_labels.append(era_soft_labels)
 
-        DSFLServerHandler.distill(
+        DSFLBaseServerHandler.distill(
             self.model,
             self.kd_optimizer,
             self.dataset,
@@ -205,7 +205,7 @@ class DSFLServerHandler(ServerHandler[DSFLUplinkPackage, DSFLDownlinkPackage]):
         return avg_loss, avg_acc
 
     def get_summary(self) -> dict[str, float]:
-        server_loss, server_acc = DSFLServerHandler.evaulate(
+        server_loss, server_acc = DSFLBaseServerHandler.evaulate(
             self.model,
             self.dataset.get_dataloader(
                 type_="test",
@@ -233,7 +233,7 @@ class DSFLServerHandler(ServerHandler[DSFLUplinkPackage, DSFLDownlinkPackage]):
 
 
 @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
 
 
@@ -257,8 +256,8 @@ class DSFLClientState:
     kd_optimizer: dict[str, torch.Tensor] | None
 
 
-class DSFLParallelClientTrainer(
-    ParallelClientTrainer[DSFLUplinkPackage, DSFLDownlinkPackage, DSFLDiskSharedData]
+class DSFLProcessPoolClientTrainer(
+    ProcessPoolClientTrainer[DSFLUplinkPackage, DSFLDownlinkPackage, DSFLClientConfig]
 ):
     def __init__(
         self,
@@ -300,82 +299,90 @@ class DSFLParallelClientTrainer(
         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)
-            DSFLServerHandler.distill(
+                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,
         )
-        DSFLParallelClientTrainer.train(
+        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 = DSFLParallelClientTrainer.predict(
+        soft_labels = DSFLProcessPoolClientTrainer.predict(
             model=model,
             open_loader=open_loader,
             device=device,
         )
 
         # 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 = DSFLServerHandler.evaulate(
+        loss, acc = DSFLBaseServerHandler.evaulate(
             model=model,
             test_loader=test_loader,
             device=device,
@@ -383,19 +390,19 @@ class DSFLParallelClientTrainer(
 
         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 DSFLParallelClientTrainer(
         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 DSFLParallelClientTrainer(
             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.dev2 → blazefl-2.0.0.dev4}/examples/step-by-step-dsfl/main.py

@@ -10,7 +10,7 @@ from hydra.core import hydra_config
 from omegaconf import DictConfig, OmegaConf
 from torch.utils.tensorboard.writer import SummaryWriter
 
-from algorithm import DSFLParallelClientTrainer, DSFLServerHandler
+from algorithm import DSFLBaseServerHandler, DSFLProcessPoolClientTrainer
 from dataset import DSFLPartitionedDataset
 from models import DSFLModelSelector
 
@@ -18,8 +18,8 @@ from models import DSFLModelSelector
 class DSFLPipeline:
     def __init__(
         self,
-        handler: DSFLServerHandler,
-        trainer: DSFLParallelClientTrainer,
+        handler: DSFLBaseServerHandler,
+        trainer: DSFLProcessPoolClientTrainer,
         writer: SummaryWriter,
     ) -> None:
         self.handler = handler
@@ -82,7 +82,7 @@ def main(
 
     match cfg.algorithm.name:
         case "dsfl":
-            handler = DSFLServerHandler(
+            handler = DSFLBaseServerHandler(
                 model_selector=model_selector,
                 model_name=cfg.model_name,
                 dataset=dataset,
@@ -96,7 +96,7 @@ def main(
                 device=device,
                 sample_ratio=cfg.sample_ratio,
             )
-            trainer = DSFLParallelClientTrainer(
+            trainer = DSFLProcessPoolClientTrainer(
                 model_selector=model_selector,
                 model_name=cfg.model_name,
                 dataset=dataset,

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

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

{blazefl-2.0.0.dev2 → blazefl-2.0.0.dev4}/src/blazefl/contrib/__init__.py

@@ -6,13 +6,13 @@ extending the core functionalities of BlazeFL.
 """
 
 from blazefl.contrib.fedavg import (
-    FedAvgParallelClientTrainer,
-    FedAvgSerialClientTrainer,
-    FedAvgServerHandler,
+    FedAvgBaseClientTrainer,
+    FedAvgBaseServerHandler,
+    FedAvgProcessPoolClientTrainer,
 )
 
 __all__ = [
-    "FedAvgServerHandler",
-    "FedAvgParallelClientTrainer",
-    "FedAvgSerialClientTrainer",
+    "FedAvgBaseServerHandler",
+    "FedAvgProcessPoolClientTrainer",
+    "FedAvgBaseClientTrainer",
 ]