blazefl 2.0.0.dev6__tar.gz → 2.0.0.dev7__tar.gz

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

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

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/examples/experimental-freethreaded/config/config.yaml

@@ -15,3 +15,4 @@ dataset_split_dir: /tmp/experimental-freethreaded/split
 share_dir: /tmp/experimental-freethreaded/share
 state_dir: /tmp/experimental-freethreaded/state
 execution_mode: multi-thread
+ipc_mode: storage

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/examples/experimental-freethreaded/main.py

@@ -1,5 +1,4 @@
 import logging
-from copy import deepcopy
 from datetime import datetime
 from pathlib import Path
 
@@ -7,12 +6,13 @@ import hydra
 import torch
 import torch.multiprocessing as mp
 from blazefl.contrib import (
-    FedAvgParallelClientTrainer,
-    FedAvgSerialClientTrainer,
-    FedAvgServerHandler,
+    FedAvgBaseClientTrainer,
+    FedAvgBaseServerHandler,
+    FedAvgProcessPoolClientTrainer,
+)
+from blazefl.contrib.fedavg import (
+    FedAvgThreadPoolClientTrainer,
 )
-from blazefl.contrib.fedavg import FedAvgDownlinkPackage, FedAvgUplinkPackage
-from blazefl.core import ModelSelector, MultiThreadClientTrainer, PartitionedDataset
 from blazefl.utils import seed_everything
 from omegaconf import DictConfig, OmegaConf
 
@@ -20,75 +20,13 @@ from dataset import PartitionedCIFAR10
 from models import FedAvgModelSelector
 
 
-class FedAvgMultiThreadClientTrainer(
-    MultiThreadClientTrainer[
-        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
-
-    def process_client(
-        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 = FedAvgParallelClientTrainer.train(
-            model=model,
-            model_parameters=payload.model_parameters,
-            train_loader=train_loader,
-            device=device,
-            epochs=self.epochs,
-            lr=self.lr,
-        )
-        return package
-
-    def uplink_package(self) -> list[FedAvgUplinkPackage]:
-        package = deepcopy(self.cache)
-        self.cache = []
-        return package
-
-
 class FedAvgPipeline:
     def __init__(
         self,
-        handler: FedAvgServerHandler,
-        trainer: FedAvgSerialClientTrainer
-        | FedAvgParallelClientTrainer
-        | FedAvgMultiThreadClientTrainer,
+        handler: FedAvgBaseServerHandler,
+        trainer: FedAvgBaseClientTrainer
+        | FedAvgProcessPoolClientTrainer
+        | FedAvgThreadPoolClientTrainer,
     ) -> None:
         self.handler = handler
         self.trainer = trainer
@@ -145,7 +83,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,
@@ -156,14 +94,14 @@ def main(cfg: DictConfig):
         batch_size=cfg.batch_size,
     )
     trainer: (
-        FedAvgSerialClientTrainer
-        | FedAvgParallelClientTrainer
-        | FedAvgMultiThreadClientTrainer
+        FedAvgBaseClientTrainer
+        | FedAvgProcessPoolClientTrainer
+        | FedAvgThreadPoolClientTrainer
         | None
     ) = None
     match cfg.execution_mode:
-        case "serial":
-            trainer = FedAvgSerialClientTrainer(
+        case "single-thread":
+            trainer = FedAvgBaseClientTrainer(
                 model_selector=model_selector,
                 model_name=cfg.model_name,
                 dataset=dataset,
@@ -174,7 +112,7 @@ def main(cfg: DictConfig):
                 batch_size=cfg.batch_size,
             )
         case "multi-process":
-            trainer = FedAvgParallelClientTrainer(
+            trainer = FedAvgProcessPoolClientTrainer(
                 model_selector=model_selector,
                 model_name=cfg.model_name,
                 dataset=dataset,
@@ -187,9 +125,10 @@ def main(cfg: DictConfig):
                 lr=cfg.lr,
                 batch_size=cfg.batch_size,
                 num_parallels=cfg.num_parallels,
+                ipc_mode=cfg.ipc_mode,
             )
         case "multi-thread":
-            trainer = FedAvgMultiThreadClientTrainer(
+            trainer = FedAvgThreadPoolClientTrainer(
                 model_selector=model_selector,
                 model_name=cfg.model_name,
                 dataset=dataset,
@@ -207,9 +146,7 @@ def main(cfg: DictConfig):
     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.dev6 → blazefl-2.0.0.dev7}/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.0dev2",
+    "blazefl>=2.0.0dev6",
     "hydra-core>=1.3.2",
     "torch>=2.7.1",
     "torchvision>=0.22.1",

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/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.dev2"
+version = "2.0.0.dev6"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "numpy" },
     { name = "torch" },
     { name = "tqdm" },
 ]
-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" }
+sdist = { url = "https://files.pythonhosted.org/packages/d7/76/d247942d051447dc1c6dbc25b2d55b4cdf146cc0ebd3594c0c33628f465c/blazefl-2.0.0.dev6.tar.gz", hash = "sha256:d4875f03872917dfd99b3a70ac2c6efc3483ee0e42e484d7328e0810fe5141b9", size = 615380, upload_time = "2025-06-16T16:26:50.929Z" }
 wheels = [
-    { 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" },
+    { url = "https://files.pythonhosted.org/packages/c8/30/eba32caaed89d04c925f676d48fb50bad7bed3ce2c86c6ebcaf2eb92aff1/blazefl-2.0.0.dev6-py3-none-any.whl", hash = "sha256:e38c3739be9ef2bc7bb8ba213da3319e4be8f9ea9f27ecd661cbc70c0cedd655", size = 26215, upload_time = "2025-06-16T16:26:49.512Z" },
 ]
 
 [[package]]
@@ -50,7 +50,7 @@ dev = [
 
 [package.metadata]
 requires-dist = [
-    { name = "blazefl", specifier = ">=2.0.0.dev2" },
+    { name = "blazefl", specifier = ">=2.0.0.dev6" },
     { name = "hydra-core", specifier = ">=1.3.2" },
     { name = "torch", specifier = ">=2.7.1" },
     { name = "torchvision", specifier = ">=0.22.1" },

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/examples/quickstart-fedavg/dataset/dataset.py

@@ -3,6 +3,7 @@ from pathlib import Path
 
 import torch
 import torchvision
+from blazefl.contrib import FedAvgPartitionType
 from blazefl.core import PartitionedDataset
 from blazefl.utils import FilteredDataset
 from torch.utils.data import DataLoader, Dataset
@@ -15,7 +16,7 @@ from dataset.functional import (
 )
 
 
-class PartitionedCIFAR10(PartitionedDataset):
+class PartitionedCIFAR10(PartitionedDataset[FedAvgPartitionType]):
     def __init__(
         self,
         root: Path,
@@ -107,24 +108,22 @@ class PartitionedCIFAR10(PartitionedDataset):
             self.path.joinpath("test.pkl"),
         )
 
-    def get_dataset(self, type_: str, cid: int | None) -> Dataset:
+    def get_dataset(self, type_: FedAvgPartitionType, cid: int | None) -> Dataset:
         match type_:
-            case "train":
+            case FedAvgPartitionType.TRAIN:
                 dataset = torch.load(
                     self.path.joinpath(type_, f"{cid}.pkl"),
                     weights_only=False,
                 )
-            case "test":
+            case FedAvgPartitionType.TEST:
                 dataset = torch.load(
                     self.path.joinpath(f"{type_}.pkl"), weights_only=False
                 )
-            case _:
-                raise ValueError(f"Invalid dataset type: {type_}")
         assert isinstance(dataset, Dataset)
         return dataset
 
     def get_dataloader(
-        self, type_: str, cid: int | None, batch_size: int | None = None
+        self, type_: FedAvgPartitionType, cid: int | None, batch_size: int | None = None
     ) -> DataLoader:
         dataset = self.get_dataset(type_, cid)
         assert isinstance(dataset, Sized)

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/examples/quickstart-fedavg/pyproject.toml

@@ -32,4 +32,5 @@ blazefl = { workspace = true }
 [dependency-groups]
 dev = [
     "mypy>=1.13.0",
+    "pre-commit>=4.2.0",
 ]

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/examples/step-by-step-dsfl/algorithm/dsfl.py

@@ -1,10 +1,12 @@
 import random
+import threading
 from collections import defaultdict
 from copy import deepcopy
 from dataclasses import dataclass
 from pathlib import Path
 
 import torch
+import torch.multiprocessing as mp
 import torch.nn.functional as F
 from blazefl.core import (
     BaseServerHandler,
@@ -17,7 +19,7 @@ from blazefl.utils import (
 )
 from torch.utils.data import DataLoader, Subset
 
-from dataset import DSFLPartitionedDataset
+from dataset import DSFLPartitionedDataset, DSFLPartitionType
 from models import DSFLModelSelector
 
 
@@ -125,6 +127,7 @@ class DSFLBaseServerHandler(BaseServerHandler[DSFLUplinkPackage, DSFLDownlinkPac
             self.kd_epochs,
             self.kd_batch_size,
             self.device,
+            stop_event=None,
         )
 
         self.global_soft_labels = torch.stack(global_soft_labels)
@@ -140,10 +143,11 @@ class DSFLBaseServerHandler(BaseServerHandler[DSFLUplinkPackage, DSFLDownlinkPac
         kd_epochs: int,
         kd_batch_size: int,
         device: str,
+        stop_event: threading.Event | None,
     ) -> None:
         model.to(device)
         model.train()
-        open_dataset = dataset.get_dataset(type_="open", cid=None)
+        open_dataset = dataset.get_dataset(type_=DSFLPartitionType.OPEN, cid=None)
         open_loader = DataLoader(
             Subset(open_dataset, global_indices),
             batch_size=kd_batch_size,
@@ -156,6 +160,8 @@ class DSFLBaseServerHandler(BaseServerHandler[DSFLUplinkPackage, DSFLDownlinkPac
             batch_size=kd_batch_size,
         )
         for _ in range(kd_epochs):
+            if stop_event is not None and stop_event.is_set():
+                break
             for data, soft_label in zip(
                 open_loader, global_soft_label_loader, strict=True
             ):
@@ -208,7 +214,7 @@ class DSFLBaseServerHandler(BaseServerHandler[DSFLUplinkPackage, DSFLDownlinkPac
         server_loss, server_acc = DSFLBaseServerHandler.evaulate(
             self.model,
             self.dataset.get_dataloader(
-                type_="test",
+                type_=DSFLPartitionType.TEST,
                 cid=None,
                 batch_size=self.kd_batch_size,
             ),
@@ -300,6 +306,8 @@ class DSFLProcessPoolClientTrainer(
         self.num_clients = num_clients
         self.seed = seed
         self.ipc_mode = "storage"
+        self.manager = mp.Manager()
+        self.stop_event = self.manager.Event()
 
         if self.device == "cuda":
             self.device_count = torch.cuda.device_count()
@@ -309,6 +317,7 @@ class DSFLProcessPoolClientTrainer(
         config: DSFLClientConfig | Path,
         payload: DSFLDownlinkPackage | Path,
         device: str,
+        stop_event: threading.Event,
     ) -> Path:
         assert isinstance(config, Path) and isinstance(payload, Path)
         config_path, payload_path = config, payload
@@ -334,7 +343,7 @@ class DSFLProcessPoolClientTrainer(
             seed_everything(c.seed, device=device)
 
         # Distill
-        open_dataset = c.dataset.get_dataset(type_="open", cid=None)
+        open_dataset = c.dataset.get_dataset(type_=DSFLPartitionType.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()
@@ -349,11 +358,12 @@ class DSFLProcessPoolClientTrainer(
                 kd_epochs=c.kd_epochs,
                 kd_batch_size=c.kd_batch_size,
                 device=device,
+                stop_event=stop_event,
             )
 
         # Train
         train_loader = c.dataset.get_dataloader(
-            type_="train",
+            type_=DSFLPartitionType.TRAIN,
             cid=c.cid,
             batch_size=c.batch_size,
         )
@@ -363,6 +373,7 @@ class DSFLProcessPoolClientTrainer(
             train_loader=train_loader,
             device=device,
             epochs=c.epochs,
+            stop_event=stop_event,
         )
 
         # Predict
@@ -378,7 +389,7 @@ class DSFLProcessPoolClientTrainer(
 
         # Evaluate
         test_loader = c.dataset.get_dataloader(
-            type_="test",
+            type_=DSFLPartitionType.TEST,
             cid=c.cid,
             batch_size=c.batch_size,
         )
@@ -411,12 +422,15 @@ class DSFLProcessPoolClientTrainer(
         train_loader: DataLoader,
         device: str,
         epochs: int,
+        stop_event: threading.Event,
     ) -> None:
         model.to(device)
         model.train()
         criterion = torch.nn.CrossEntropyLoss()
 
         for _ in range(epochs):
+            if stop_event.is_set():
+                break
             for data, target in train_loader:
                 data = data.to(device)
                 target = target.to(device)
@@ -431,7 +445,9 @@ class DSFLProcessPoolClientTrainer(
 
     @staticmethod
     def predict(
-        model: torch.nn.Module, open_loader: DataLoader, device: str
+        model: torch.nn.Module,
+        open_loader: DataLoader,
+        device: str,
     ) -> torch.Tensor:
         model.to(device)
         model.eval()

blazefl-2.0.0.dev7/examples/step-by-step-dsfl/dataset/__init__.py

@@ -0,0 +1,3 @@
+from dataset.dataset import DSFLPartitionedDataset, DSFLPartitionType
+
+__all__ = ["DSFLPartitionedDataset", "DSFLPartitionType"]

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/examples/step-by-step-dsfl/dataset/dataset.py

@@ -1,4 +1,5 @@
 from collections.abc import Sized
+from enum import StrEnum
 from pathlib import Path
 
 import numpy as np
@@ -15,7 +16,13 @@ from dataset.functional import (
 )
 
 
-class DSFLPartitionedDataset(PartitionedDataset):
+class DSFLPartitionType(StrEnum):
+    TRAIN = "train"
+    OPEN = "open"
+    TEST = "test"
+
+
+class DSFLPartitionedDataset(PartitionedDataset[DSFLPartitionType]):
     def __init__(
         self,
         root: Path,
@@ -68,7 +75,7 @@ class DSFLPartitionedDataset(PartitionedDataset):
             train=False,
             download=True,
         )
-        for type_ in ["train", "open", "test"]:
+        for type_ in [ds.value for ds in DSFLPartitionType]:
             self.path.joinpath(type_).mkdir(parents=True)
 
         match self.partition:
@@ -141,19 +148,19 @@ class DSFLPartitionedDataset(PartitionedDataset):
             self.path.joinpath("test", "default.pkl"),
         )
 
-    def get_dataset(self, type_: str, cid: int | None) -> Dataset:
+    def get_dataset(self, type_: DSFLPartitionType, cid: int | None) -> Dataset:
         match type_:
-            case "train":
+            case DSFLPartitionType.TRAIN:
                 dataset = torch.load(
                     self.path.joinpath(type_, f"{cid}.pkl"),
                     weights_only=False,
                 )
-            case "open":
+            case DSFLPartitionType.OPEN:
                 dataset = torch.load(
                     self.path.joinpath(f"{type_}.pkl"),
                     weights_only=False,
                 )
-            case "test":
+            case DSFLPartitionType.TEST:
                 if cid is not None:
                     dataset = torch.load(
                         self.path.joinpath(type_, f"{cid}.pkl"),
@@ -163,13 +170,11 @@ class DSFLPartitionedDataset(PartitionedDataset):
                     dataset = torch.load(
                         self.path.joinpath(type_, "default.pkl"), weights_only=False
                     )
-            case _:
-                raise ValueError(f"Invalid dataset type: {type_}")
         assert isinstance(dataset, Dataset)
         return dataset
 
     def get_dataloader(
-        self, type_: str, cid: int | None, batch_size: int | None = None
+        self, type_: DSFLPartitionType, cid: int | None, batch_size: int | None = None
     ) -> DataLoader:
         dataset = self.get_dataset(type_, cid)
         assert isinstance(dataset, Sized)

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/examples/step-by-step-dsfl/pyproject.toml

@@ -36,4 +36,5 @@ blazefl = { workspace = true }
 [dependency-groups]
 dev = [
     "mypy>=1.13.0",
+    "pre-commit>=4.2.0",
 ]

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

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

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/src/blazefl/contrib/__init__.py

@@ -9,6 +9,8 @@ from blazefl.contrib.fedavg import (
     FedAvgBaseClientTrainer,
     FedAvgBaseServerHandler,
     FedAvgDownlinkPackage,
+    FedAvgPartitionedDataset,
+    FedAvgPartitionType,
     FedAvgProcessPoolClientTrainer,
     FedAvgThreadPoolClientTrainer,
     FedAvgUplinkPackage,
@@ -21,4 +23,6 @@ __all__ = [
     "FedAvgThreadPoolClientTrainer",
     "FedAvgUplinkPackage",
     "FedAvgDownlinkPackage",
+    "FedAvgPartitionType",
+    "FedAvgPartitionedDataset",
 ]

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/src/blazefl/contrib/fedavg.py

@@ -2,6 +2,7 @@ import random
 import threading
 from copy import deepcopy
 from dataclasses import dataclass
+from enum import StrEnum
 from pathlib import Path
 from typing import Literal
 
@@ -57,8 +58,16 @@ class FedAvgDownlinkPackage:
     model_parameters: torch.Tensor
 
 
+class FedAvgPartitionType(StrEnum):
+    TRAIN = "train"
+    TEST = "test"
+
+
+FedAvgPartitionedDataset = PartitionedDataset[FedAvgPartitionType]
+
+
 class FedAvgBaseServerHandler(
-    BaseServerHandler[FedAvgUplinkPackage, FedAvgDownlinkPackage]
+    BaseServerHandler[FedAvgUplinkPackage, FedAvgDownlinkPackage],
 ):
     """
     Server-side handler for the Federated Averaging (FedAvg) algorithm.
@@ -83,7 +92,7 @@ class FedAvgBaseServerHandler(
         self,
         model_selector: ModelSelector,
         model_name: str,
-        dataset: PartitionedDataset,
+        dataset: FedAvgPartitionedDataset,
         global_round: int,
         num_clients: int,
         sample_ratio: float,
@@ -241,7 +250,7 @@ class FedAvgBaseServerHandler(
         server_loss, server_acc = FedAvgBaseServerHandler.evaluate(
             self.model,
             self.dataset.get_dataloader(
-                type_="test",
+                type_=FedAvgPartitionType.TEST,
                 cid=None,
                 batch_size=self.batch_size,
             ),
@@ -289,7 +298,7 @@ class FedAvgBaseClientTrainer(
         self,
         model_selector: ModelSelector,
         model_name: str,
-        dataset: PartitionedDataset,
+        dataset: FedAvgPartitionedDataset,
         device: str,
         num_clients: int,
         epochs: int,
@@ -339,14 +348,9 @@ class FedAvgBaseClientTrainer(
         model_parameters = payload.model_parameters
         for cid in tqdm(cid_list, desc="Client", leave=False):
             data_loader = self.dataset.get_dataloader(
-                type_="train", cid=cid, batch_size=self.batch_size
+                type_=FedAvgPartitionType.TRAIN, cid=cid, batch_size=self.batch_size
             )
             pack = self.train(model_parameters, data_loader)
-            val_loader = self.dataset.get_dataloader(
-                type_="val", cid=cid, batch_size=self.batch_size
-            )
-            loss, acc = self.evaluate(val_loader)
-            pack.metadata = {"loss": loss, "acc": acc}
             self.cache.append(pack)
 
     def train(
@@ -457,7 +461,7 @@ class FedAvgClientConfig:
 
     model_selector: ModelSelector
     model_name: str
-    dataset: PartitionedDataset
+    dataset: FedAvgPartitionedDataset
     epochs: int
     batch_size: int
     lr: float
@@ -501,7 +505,7 @@ class FedAvgProcessPoolClientTrainer(
         model_name: str,
         share_dir: Path,
         state_dir: Path,
-        dataset: PartitionedDataset,
+        dataset: FedAvgPartitionedDataset,
         device: str,
         num_clients: int,
         epochs: int,
@@ -614,7 +618,7 @@ class FedAvgProcessPoolClientTrainer(
 
             model = config.model_selector.select_model(config.model_name)
             train_loader = config.dataset.get_dataloader(
-                type_="train",
+                type_=FedAvgPartitionType.TRAIN,
                 cid=config.cid,
                 batch_size=config.batch_size,
             )
@@ -739,7 +743,7 @@ class FedAvgThreadPoolClientTrainer(
         self,
         model_selector: ModelSelector,
         model_name: str,
-        dataset: PartitionedDataset,
+        dataset: FedAvgPartitionedDataset,
         device: str,
         num_clients: int,
         epochs: int,
@@ -773,7 +777,7 @@ class FedAvgThreadPoolClientTrainer(
     ) -> FedAvgUplinkPackage:
         model = self.model_selector.select_model(self.model_name)
         train_loader = self.dataset.get_dataloader(
-            type_="train",
+            type_=FedAvgPartitionType.TRAIN,
             cid=cid,
             batch_size=self.batch_size,
         )

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/src/blazefl/core/partitioned_dataset.py

@@ -1,9 +1,12 @@
-from typing import Protocol
+from enum import StrEnum
+from typing import Protocol, TypeVar
 
 from torch.utils.data import DataLoader, Dataset
 
+PartitionType = TypeVar("PartitionType", bound=StrEnum, contravariant=True)
 
-class PartitionedDataset(Protocol):
+
+class PartitionedDataset(Protocol[PartitionType]):
     """
     Abstract base class for partitioned datasets in federated learning.
 
@@ -14,7 +17,7 @@ class PartitionedDataset(Protocol):
         NotImplementedError: If the methods are not implemented in a subclass.
     """
 
-    def get_dataset(self, type_: str, cid: int | None) -> Dataset:
+    def get_dataset(self, type_: PartitionType, cid: int | None) -> Dataset:
         """
         Retrieve a dataset for a specific type and client ID.
 
@@ -28,7 +31,7 @@ class PartitionedDataset(Protocol):
         ...
 
     def get_dataloader(
-        self, type_: str, cid: int | None, batch_size: int | None
+        self, type_: PartitionType, cid: int | None, batch_size: int | None
     ) -> DataLoader:
         """
         Retrieve a DataLoader for a specific type, client ID, and batch size.

blazefl-2.0.0.dev7/src/blazefl/core/partitioned_dataset.pyi

@@ -0,0 +1,9 @@
+from enum import StrEnum
+from torch.utils.data import DataLoader, Dataset
+from typing import Protocol, TypeVar
+
+PartitionType = TypeVar('PartitionType', bound=StrEnum, contravariant=True)
+
+class PartitionedDataset(Protocol[PartitionType]):
+    def get_dataset(self, type_: PartitionType, cid: int | None) -> Dataset: ...
+    def get_dataloader(self, type_: PartitionType, cid: int | None, batch_size: int | None) -> DataLoader: ...

{blazefl-2.0.0.dev6 → blazefl-2.0.0.dev7}/uv.lock

@@ -82,7 +82,7 @@ wheels = [
 
 [[package]]
 name = "blazefl"
-version = "2.0.0.dev6"
+version = "2.0.0.dev7"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },
@@ -828,6 +828,7 @@ dependencies = [
 [package.dev-dependencies]
 dev = [
     { name = "mypy" },
+    { name = "pre-commit" },
 ]
 
 [package.metadata]
@@ -839,7 +840,10 @@ requires-dist = [
 ]
 
 [package.metadata.requires-dev]
-dev = [{ name = "mypy", specifier = ">=1.13.0" }]
+dev = [
+    { name = "mypy", specifier = ">=1.13.0" },
+    { name = "pre-commit", specifier = ">=4.2.0" },
+]
 
 [[package]]
 name = "requests"
@@ -1063,6 +1067,7 @@ dependencies = [
 [package.dev-dependencies]
 dev = [
     { name = "mypy" },
+    { name = "pre-commit" },
 ]
 
 [package.metadata]
@@ -1074,7 +1079,10 @@ requires-dist = [
 ]
 
 [package.metadata.requires-dev]
-dev = [{ name = "mypy", specifier = ">=1.13.0" }]
+dev = [
+    { name = "mypy", specifier = ">=1.13.0" },
+    { name = "pre-commit", specifier = ">=4.2.0" },
+]
 
 [[package]]
 name = "sympy"

blazefl-2.0.0.dev6/examples/step-by-step-dsfl/dataset/__init__.py

@@ -1,3 +0,0 @@
-from dataset.dataset import DSFLPartitionedDataset
-
-__all__ = ["DSFLPartitionedDataset"]

blazefl-2.0.0.dev6/src/blazefl/core/partitioned_dataset.pyi

@@ -1,6 +0,0 @@
-from torch.utils.data import DataLoader, Dataset
-from typing import Protocol
-
-class PartitionedDataset(Protocol):
-    def get_dataset(self, type_: str, cid: int | None) -> Dataset: ...
-    def get_dataloader(self, type_: str, cid: int | None, batch_size: int | None) -> DataLoader: ...