blazefl 2.0.0.dev4__tar.gz → 2.0.0.dev5__tar.gz

blazefl-2.0.0.dev5/.python-version

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blazefl
-Version: 2.0.0.dev4
+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.dev4 → blazefl-2.0.0.dev5}/pyproject.toml

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

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

@@ -1,4 +1,5 @@
 import random
+import threading
 from copy import deepcopy
 from dataclasses import dataclass
 from pathlib import Path
@@ -14,6 +15,7 @@ from blazefl.core import (
     ModelSelector,
     PartitionedDataset,
     ProcessPoolClientTrainer,
+    ThreadPoolClientTrainer,
 )
 from blazefl.utils import (
     RandomState,
@@ -640,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.
@@ -664,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)
@@ -714,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

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

@@ -1,5 +1,7 @@
+import logging
 import multiprocessing as mp
 import signal
+import threading
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from multiprocessing.pool import ApplyResult
 from pathlib import Path
@@ -197,6 +199,7 @@ class ThreadPoolClientTrainer(
     device: str
     device_count: int
     cache: list[UplinkPackage]
+    stop_event: threading.Event | None
 
     def worker(
         self,
@@ -223,20 +226,31 @@ 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:
-            futures = []
-            for cid in cid_list:
-                device = self.get_client_device(cid)
-                future = executor.submit(
-                    self.worker,
-                    cid,
-                    device,
-                    payload,
-                )
-                futures.append(future)
-
-            for future in tqdm(
-                as_completed(futures), total=len(futures), desc="Client", leave=False
-            ):
-                result = future.result()
-                self.cache.append(result)
+        if self.stop_event is None:
+            self.stop_event = threading.Event()
+        self.stop_event.clear()
+        try:
+            with ThreadPoolExecutor(max_workers=self.num_parallels) as executor:
+                futures = []
+                for cid in cid_list:
+                    device = self.get_client_device(cid)
+                    future = executor.submit(
+                        self.worker,
+                        cid,
+                        device,
+                        payload,
+                    )
+                    futures.append(future)
+
+                for future in tqdm(
+                    as_completed(futures),
+                    total=len(futures),
+                    desc="Client",
+                    leave=False,
+                ):
+                    result = future.result()
+                    self.cache.append(result)
+        except KeyboardInterrupt:
+            logging.warning("Training interrupted by user.")
+            self.stop_event.set()
+        return

{blazefl-2.0.0.dev4 → blazefl-2.0.0.dev5}/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
@@ -29,6 +30,7 @@ class ThreadPoolClientTrainer(BaseClientTrainer[UplinkPackage, DownlinkPackage],
     device: str
     device_count: int
     cache: list[UplinkPackage]
+    stop_event: threading.Event | None
     def worker(self, cid: int, device: str, payload: DownlinkPackage) -> UplinkPackage: ...
     def get_client_device(self, cid: int) -> str: ...
     def local_process(self, payload: DownlinkPackage, cid_list: list[int]) -> None: ...

{blazefl-2.0.0.dev4 → blazefl-2.0.0.dev5}/tests/test_contrib/test_fedavg.py

@@ -13,6 +13,7 @@ from src.blazefl.contrib.fedavg import (
     FedAvgBaseClientTrainer,
     FedAvgBaseServerHandler,
     FedAvgProcessPoolClientTrainer,
+    FedAvgThreadPoolClientTrainer,
 )
 from src.blazefl.core import ModelSelector, PartitionedDataset
 
@@ -86,7 +87,9 @@ def tmp_state_dir(tmp_path):
     return state_dir
 
 
-def test_server_and_base_integration(model_selector, partitioned_dataset, device):
+def test_base_server_and_base_trainer_integration(
+    model_selector, partitioned_dataset, device
+):
     model_name = "dummy"
     global_round = 1
     num_clients = 3
@@ -134,7 +137,7 @@ def test_server_and_base_integration(model_selector, partitioned_dataset, device
 
 
 @pytest.mark.parametrize("ipc_mode", ["storage", "shared_memory"])
-def test_server_and_process_pool_integration(
+def test_base_handler_and_process_pool_trainer_integration(
     model_selector, partitioned_dataset, device, tmp_share_dir, tmp_state_dir, ipc_mode
 ):
     model_name = "dummy"
@@ -195,7 +198,7 @@ def run_local_process(trainer, downlink, cids):
         trainer.local_process(downlink, cids)
 
 
-def test_server_and_process_pool_integration_keyboard_interrupt(
+def test_base_handler_and_process_pool_trainer_integration_keyboard_interrupt(
     model_selector, partitioned_dataset, device, tmp_share_dir, tmp_state_dir
 ):
     model_name = "dummy"
@@ -252,7 +255,9 @@ def test_server_and_process_pool_integration_keyboard_interrupt(
         if len(spawned_pids) == num_parallels:
             break
         if time.time() - start_time > timeout:
-            raise AssertionError("Timeout reached while waiting for spawned processes.")
+            pytest.fail(
+                f"Process did not spawn {len(spawned_pids)} processes within {timeout}s"
+            )
     assert proc.is_alive()
 
     os.kill(proc.pid, signal.SIGINT)
@@ -265,3 +270,118 @@ def test_server_and_process_pool_integration_keyboard_interrupt(
         if psutil.pid_exists(pid):
             orphan_pids.append(pid)
     assert len(orphan_pids) == 0
+
+
+def test_base_handler_and_thread_pool_trainer_integration(
+    model_selector, partitioned_dataset, device
+):
+    model_name = "dummy"
+    global_round = 2
+    num_clients = 3
+    sample_ratio = 1.0
+    epochs = 1
+    batch_size = 2
+    lr = 0.01
+    seed = 42
+    num_parallels = 2
+
+    server = FedAvgBaseServerHandler(
+        model_selector=model_selector,
+        model_name=model_name,
+        dataset=partitioned_dataset,
+        global_round=global_round,
+        num_clients=num_clients,
+        sample_ratio=sample_ratio,
+        device=device,
+        batch_size=batch_size,
+    )
+
+    trainer = FedAvgThreadPoolClientTrainer(
+        model_selector=model_selector,
+        model_name=model_name,
+        dataset=partitioned_dataset,
+        device=device,
+        num_clients=num_clients,
+        epochs=epochs,
+        batch_size=batch_size,
+        lr=lr,
+        seed=seed,
+        num_parallels=num_parallels,
+    )
+
+    for round_ in range(1, global_round + 1):
+        cids = server.sample_clients()
+        downlink = server.downlink_package()
+        trainer.local_process(downlink, cids)
+        uplinks = trainer.uplink_package()
+        assert len(uplinks) == num_clients
+
+        done = False
+        for pkg in uplinks:
+            done = server.load(pkg)
+        assert done is True
+        assert server.round == round_
+
+    assert server.if_stop() is True
+
+
+def test_base_handler_and_thread_pool_trainer_integration_keyboard_interrupt(
+    model_selector,
+    partitioned_dataset,
+    device,
+):
+    model_name = "dummy"
+    global_round = 1
+    num_clients = 10
+    sample_ratio = 1.0
+    epochs = 10**5
+    batch_size = 2
+    lr = 0.01
+    seed = 42
+    num_parallels = 10
+
+    server = FedAvgBaseServerHandler(
+        model_selector=model_selector,
+        model_name=model_name,
+        dataset=partitioned_dataset,
+        global_round=global_round,
+        num_clients=num_clients,
+        sample_ratio=sample_ratio,
+        device=device,
+        batch_size=batch_size,
+    )
+
+    trainer = FedAvgThreadPoolClientTrainer(
+        model_selector=model_selector,
+        model_name=model_name,
+        dataset=partitioned_dataset,
+        device=device,
+        num_clients=num_clients,
+        epochs=epochs,
+        batch_size=batch_size,
+        lr=lr,
+        seed=seed,
+        num_parallels=num_parallels,
+    )
+
+    cids = server.sample_clients()
+    downlink = server.downlink_package()
+
+    proc = Process(target=run_local_process, args=(trainer, downlink, cids))
+    proc.start()
+    assert proc.pid is not None
+
+    p = psutil.Process(proc.pid)
+    timeout = 5
+    while p.num_threads() < num_parallels + 1:
+        if not p.is_running() or time.time() - p.create_time() > timeout:
+            pytest.fail(
+                f"Process did not spawn {num_parallels} threads within {timeout}s"
+            )
+        time.sleep(0.1)
+    assert proc.is_alive()
+
+    os.kill(proc.pid, signal.SIGINT)
+
+    proc.join(timeout=5)
+    assert not proc.is_alive()

{blazefl-2.0.0.dev4 → blazefl-2.0.0.dev5}/uv.lock

@@ -82,7 +82,7 @@ wheels = [
 
 [[package]]
 name = "blazefl"
-version = "2.0.0.dev4"
+version = "2.0.0.dev5"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },

blazefl-2.0.0.dev4/examples/step-by-step-dsfl/.python-version

@@ -1 +0,0 @@
-3.12