API Reference¶

class AppConfig(BaseImmutableConfig):
    """
    The master configuration schema for the entire simulation application.
    Aggregates all sub-configurations and validates cross-domain logical consistency.
    """

    comm_strategy: CommStrategyConfig = Field(default_factory=AsyncStrategy)
    mem_strategy: MemStrategyConfig = Field(default_factory=MemStrategyConfig)

    data: DataConfig = Field(default_factory=DataConfig)
    model: ModelConfig = Field(default_factory=ModelConfig)
    simulation: SimulationConfig = Field(default_factory=SimulationConfig)
    evaluation: EvaluationConfig = Field(default_factory=EvaluationConfig)
    optimization: OptimizationConfig = Field(default_factory=OptimizationConfig)
    checkpoints: CheckpointConfig = Field(default_factory=CheckpointConfig)
    visualization: VisualizationConfig = Field(default_factory=VisualizationConfig)

    @model_validator(mode="after")
    def check_logical_consistency(self) -> "AppConfig":
        """
        Validates that the communication strategy does not request more clients
        than are available in the simulation pool.

        Returns:
            AppConfig: The validated configuration object.

        Raises:
            ValueError: If synchronous sample size exceeds total clients.
        """
        if (
            self.comm_strategy.type == "sync"
            and self.comm_strategy.sample_size > self.simulation.num_clients
        ):
            raise ValueError(
                f"Config Error: Sample size ({self.comm_strategy.sample_size}) cannot "
                f"exceed total clients ({self.simulation.num_clients})."
            )
        return self

    @model_validator(mode="after")
    def check_model_compatibility(self) -> "AppConfig":
        """
        Validates that the selected model architecture can accept the number
        of input channels provided by the selected dataset.

        Returns:
            AppConfig: The validated configuration object.

        Raises:
            ValueError: If a strict channel mismatch occurs.
        """
        dataset = self.data.dataset
        model = self.model.model_name

        if model.required_channels and dataset.num_channels != model.required_channels:
            raise ValueError(
                f"Config Error: '{model}' requires {model.required_channels} channel(s), "
                f"but '{dataset}' has {dataset.num_channels}. Choose a different model."
            )

        return self

    @model_validator(mode="after")
    def check_checkpoint_interval(self) -> "AppConfig":
        """
        Checks if the checkpoint interval exceeds or equals the simulation timeout,
        warning the user that intermediate checkpoints will not be saved.

        Returns:
            AppConfig: The validated configuration object.
        """
        interval = self.checkpoints.interval_seconds
        timeout = self.simulation.timeout_seconds

        if interval >= timeout:  # pragma: no branch
            logger.warning(
                f"Config Warning: Checkpoint interval '{interval}' is equal or greater than the "
                f"simulation timeout '{timeout}'. The simulation will save a final checkpoint "
                f"upon termination or interruption, and no intermediate checkpoints."
            )

        return self

    @model_validator(mode="after")
    def sanitize_visualization_config(self) -> "AppConfig":
        """
        Disables visualizations automatically if the client count exceeds a readability
        threshold (150 clients) to prevent resource exhaustion and unreadable plots.

        Returns:
            AppConfig: The sanitized configuration object with updated visualization flags.
        """
        limit = 150  # Threshold for readable plots

        disable_split = (
            self.visualization.visualize_data_split
            and self.simulation.num_clients > limit
        )
        disable_arrivals = (
            self.visualization.visualize_client_arrivals
            and self.simulation.num_clients > limit
        )

        if not (disable_split or disable_arrivals):
            return self

        if disable_split:  # pragma: no branch
            logger.warning(
                f"Config Warning: Too many clients ({self.simulation.num_clients}) for "
                "data split visualization. Disabling to prevent unreadable plot."
            )
        if disable_arrivals:  # pragma: no branch
            logger.warning(
                f"Config Warning: Too many clients ({self.simulation.num_clients}) for "
                "arrival visualization. Disabling to prevent unreadable plot."
            )

        new_viz_config = self.visualization.model_copy(
            update={
                "visualize_data_split": False
                if disable_split
                else self.visualization.visualize_data_split,
                "visualize_client_arrivals": False
                if disable_arrivals
                else self.visualization.visualize_client_arrivals,
            }
        )

        object.__setattr__(self, "visualization", new_viz_config)

        return self

    @model_validator(mode="after")
    def check_batch_size_validity(self) -> "AppConfig":
        """
        Validates that training and evaluation batch sizes do not exceed
        their respective total dataset sizes.

        Returns:
            AppConfig: The validated configuration object.

        Raises:
            ValueError: If a configured batch size exceeds the available dataset size.
        """
        train_size = self.data.dataset.train_size
        test_size = self.data.dataset.test_size
        batch_size = self.optimization.batch_size
        batch_size_eval = self.evaluation.batch_size

        if batch_size >= train_size:
            raise ValueError(
                f"Config Error: Batch size ({batch_size}) cannot be equal to or exceed "
                f"dataset size ({train_size}) for {self.data.dataset.name}."
            )

        if batch_size_eval >= test_size:
            raise ValueError(
                f"Config Error: Evaluation batch size ({batch_size_eval}) cannot be equal to or exceed "
                f"test dataset size ({test_size}) for {self.data.dataset.name}."
            )
        return self

class AsyncStrategy(BaseImmutableConfig):
    """Configuration for asynchronous federated learning strategies."""

    type: Literal["async"] = "async"
    buffer_size: int = Field(
        default=3,
        gt=0,
        description="Number of client updates that triggers a global model update.",
    )

    @property
    def agg_target(self) -> int:
        """
        Retrieves the target number of client updates required for a global aggregation.

        Returns:
            int: The configured buffer size limit.
        """
        return self.buffer_size

class CheckpointConfig(BaseImmutableConfig):
    """
    Configuration for managing state serialization and disk I/O.
    Controls both interval-based heavy checkpoints and best-model artifacts.
    """

    interval_seconds: float = Field(
        default=400.0,
        gt=0,
        description="The interval (in wall-clock seconds) at which the simulator saves a resumable checkpoint.",
    )
    keep_best: bool = Field(
        default=False,
        description="If set to `True`, the simulator continuously saves a separate copy of the global model that achieved the highest accuracy on the test set.",
    )

class DataConfig(BaseImmutableConfig):
    """Configuration for dataset selection and distributed partitioning."""

    dataset: DatasetType = Field(
        default=DatasetType.MNIST,
        description="The target dataset for the simulation.",
    )
    dirichlet_alpha: float = Field(
        default=0.1,
        gt=0.0,
        description="Dirichlet distribution parameter.",
    )
    split_seed: int = Field(
        default=42,
        ge=0,
        description="The random seed ensuring reproducibility during the dataset partitioning process.",
    )

class DatasetType(StrEnum):
    """
    Enumeration of supported federated learning datasets.

    Attributes:
        MNIST (str): The MNIST dataset of handwritten digits.
        FASHION_MNIST (str): The Fashion-MNIST dataset of clothing articles.
        CIFAR10 (str): The CIFAR-10 dataset of 10 object classes.
        CIFAR100 (str): The CIFAR-100 dataset of 100 object classes.
    """

    MNIST = "mnist"
    FASHION_MNIST = "fashion_mnist"
    CIFAR10 = "cifar10"
    CIFAR100 = "cifar100"

    @property
    def source(self) -> str:
        """
        Retrieves the source library mapping for the dataset.

        Returns:
            str: The name of the upstream library (e.g., "torchvision").
        """
        match self:
            case (
                DatasetType.MNIST
                | DatasetType.FASHION_MNIST
                | DatasetType.CIFAR10
                | DatasetType.CIFAR100
            ):
                return "torchvision"

    @property
    def source_name(self) -> str:
        """
        Retrieves the exact dataset class name used by the source library.

        Returns:
            str: The string identifier for the source dataset class.
        """
        match self:
            case DatasetType.MNIST:
                return "MNIST"
            case DatasetType.FASHION_MNIST:
                return "FashionMNIST"
            case DatasetType.CIFAR10:
                return "CIFAR10"
            case DatasetType.CIFAR100:
                return "CIFAR100"

    @property
    def train_size(self) -> int:
        """
        Retrieves the total number of samples in the raw training split.

        Returns:
            int: The training dataset size.
        """
        match self:
            case DatasetType.MNIST | DatasetType.FASHION_MNIST:
                return 60000
            case DatasetType.CIFAR10 | DatasetType.CIFAR100:
                return 50000

    @property
    def test_size(self) -> int:
        """
        Retrieves the total number of samples in the raw evaluation split.

        Returns:
            int: The evaluation dataset size.
        """
        return 10000

    @property
    def num_classes(self) -> int:
        """
        Retrieves the total number of target classes/labels in the dataset.

        Returns:
            int: The number of classes.
        """
        match self:
            case DatasetType.CIFAR100:
                return 100
            case _:
                return 10

    @property
    def num_channels(self) -> int:
        """
        Retrieves the number of color channels in the dataset images.

        Returns:
            int: 1 for grayscale, 3 for RGB.
        """
        match self:
            case DatasetType.MNIST | DatasetType.FASHION_MNIST:
                return 1
            case _:
                return 3

    @property
    def image_size(self) -> int:
        """
        Retrieves the pixel height and width of the images (assumes square aspect ratio).

        Returns:
            int: The image dimension in pixels.
        """
        match self:
            case DatasetType.MNIST | DatasetType.FASHION_MNIST:
                return 28
            case _:
                return 32

    @property
    def mean(self) -> tuple[float] | tuple[float, float, float]:
        """
        Retrieves the channel-wise mean values for dataset normalization.

        Returns:
            tuple[float] | tuple[float, float, float]: A tuple of means for each channel.
        """
        match self:
            case DatasetType.MNIST:
                return (0.1307,)
            case DatasetType.FASHION_MNIST:
                return (0.2860,)
            case DatasetType.CIFAR10:
                return (0.4914, 0.4822, 0.4465)
            case DatasetType.CIFAR100:
                return (0.5071, 0.4865, 0.4409)

    @property
    def std(self) -> tuple[float] | tuple[float, float, float]:
        """
        Retrieves the channel-wise standard deviation values for dataset normalization.

        Returns:
            tuple[float] | tuple[float, float, float]: A tuple of standard deviations for each channel.
        """
        match self:
            case DatasetType.MNIST:
                return (0.3081,)
            case DatasetType.FASHION_MNIST:
                return (0.3530,)
            case DatasetType.CIFAR10:
                return (0.2470, 0.2435, 0.2616)
            case DatasetType.CIFAR100:
                return (0.2673, 0.2564, 0.2762)

    @property
    def apply_crop_transform(self) -> bool:
        """Determines whether random cropping should be applied during training.

        Returns:
            bool: True if random cropping is enabled, False otherwise.
        """
        match self:
            case DatasetType.CIFAR10 | DatasetType.CIFAR100:
                return True
            case _:
                return False

    @property
    def apply_horizontal_flip_transform(self) -> bool:
        """
        Determines whether random horizontal flipping should be applied during training.

        Returns:
            bool: True if horizontal flipping is enabled, False otherwise.
        """
        match self:
            case DatasetType.CIFAR10 | DatasetType.CIFAR100:
                return True
            case _:
                return False

class DefaultDirs(StrEnum):
    """
    Enumerates the standardized directories for storing simulation artifacts.

    Attributes:
        DATA (str): The default directory for storing generated simulation input data such as raw datasets, data splits and simualtion clocks. Defaults to `"data"`.
        OUTPUTS (str): The default directory for storing execution logs, metrics logs, and runtime metadata. Defaults to `"outputs"`.
        CHECKPOINTS (str): The default directory for storing resumable simulation and best model checkpoints. Defaults to `"checkpoints"`.
    """

    DATA = "data"
    OUTPUTS = "outputs"
    CHECKPOINTS = "checkpoints"

class DeviceType(StrEnum):
    """
    Enumeration of supported hardware accelerator backends.

    Attributes:
        CPU (str): Central Processing Unit backend.
        MPS (str): Apple Metal Performance Shaders backend.
        CUDA (str): NVIDIA CUDA backend.
        AUTO (str): Automatically selects the best available backend.
    """

    CPU = "cpu"
    MPS = "mps"
    CUDA = "cuda"
    AUTO = "auto"

class EvaluationConfig(BaseImmutableConfig):
    """Configuration for the server-side global evaluation process."""

    batch_size: int = Field(
        default=32,
        gt=0,
        description="The number of test dataset samples processed per batch during global model evaluation (for metric generation).",
    )
    num_workers: int = Field(
        default=0,
        ge=0,
        description="The number of subprocesses used for data loading, corresponding to the PyTorch DataLoader parameter.",
    )

class MemStrategyConfig(BaseImmutableConfig):
    """Configuration defining the memory tracking behavior of clients."""

    type: MemoryType = Field(
        default=MemoryType.DISABLED, description="Type of client memory augmentation."
    )

class MemoryType(StrEnum):
    """
    Enumeration of client-side memory tracking strategies.

    Attributes:
        DISABLED (str): Strategy indicating no memory tracking.
        MODELS (str): Strategy for tracking historical model weights.
        GRADS (str): Strategy for tracking historical gradients.
    """

    DISABLED = "disabled"
    MODELS = "models"
    GRADS = "gradients"

    @property
    def requires_buffer_reset(self) -> bool:
        """
        Determines if the server buffer must be flushed after a global update.

        Returns:
            bool: True if the server buffer requires resetting, False if updates
                are accumulated continuously (e.g., gradient memory).
        """
        match self:
            case MemoryType.GRADS:
                return False
            case _:
                return True

    @property
    def has_memory(self) -> bool:
        """
        Determines if the selected strategy necessitates tracking client-side memory states.

        Returns:
            bool: True if memory is actively used, False if disabled.
        """
        match self:
            case MemoryType.DISABLED:
                return False
            case _:
                return True

class ModelConfig(BaseImmutableConfig):
    """Configuration detailing the target neural network architecture."""

    model_name: ModelType = Field(
        default=ModelType.CNN, description="Model architecture to use."
    )

class ModelType(StrEnum):
    """
    Enumeration of supported neural network architectures.

    Attributes:
        LOG_REG (str): Logistic regression architecture.
        CNN (str): Simple Convolutional Neural Network architecture.
        RESNET18 (str): ResNet-18 architecture.
    """

    LOG_REG = "logreg"
    CNN = "cnn"
    RESNET18 = "resnet18"

    @property
    def required_channels(self) -> int | None:
        """
        Determines the strict number of input channels required by the architecture.

        Returns:
            int | None: The required integer channel count, or None if the model
                dynamically adapts to any input shape.
        """
        match self:
            case ModelType.RESNET18:
                return 3
            case _:
                return None

class OptimizationConfig(BaseImmutableConfig):
    """Configuration for the local client-side optimization process."""

    learning_rate: float = Field(
        default=0.1,
        gt=0.0,
        description="The step size applied during local client training.",
    )
    weight_decay: float = Field(
        default=0.0,
        ge=0.0,
        description="The L2 penalty (weight decay) applied by the PyTorch optimizer to prevent overfitting.",
    )
    num_local_steps: int = Field(
        default=100,
        gt=0,
        description="The exact number of local SGD steps (batches) a client performs before communicating with the server.",
    )
    batch_size: int = Field(
        default=32,
        gt=0,
        description="The number of samples processed per local training step.",
    )

class SimulationConfig(BaseImmutableConfig):
    """Configuration for the top-level simulation environment and hardware settings."""

    device: DeviceType = Field(
        default=DeviceType.AUTO,
        description="The hardware accelerator used for the simulation.",
    )
    num_clients: int = Field(default=10, gt=1, description="Total number of clients.")
    timeout_seconds: float = Field(
        default=300.0,
        gt=0,
        description="Simulation duration in wall-clock seconds.",
    )
    client_rate_std: float = Field(
        default=1.0,
        ge=0.0,
        description="Standard deviation of client latency.",
    )
    rate_seed: int = Field(
        default=42,
        ge=0,
        description="The random seed used to generate client arrival times and latency distributions.",
    )
    torch_seed: int = Field(
        default=42, ge=0, description="The random seed for all PyTorch operations."
    )

class SyncStrategy(BaseImmutableConfig):
    """Configuration for synchronous federated learning strategies."""

    type: Literal["sync"] = "sync"
    sample_size: int = Field(
        default=3,
        gt=0,
        description="Number of clients sampled by the server at each round.",
    )

    @property
    def agg_target(self) -> int:
        """
        Retrieves the target number of client updates required for a global aggregation.

        Returns:
            int: The sample size per synchronous round.
        """
        return self.sample_size

class VisualizationConfig(BaseImmutableConfig):
    """
    Configuration for creating and saving data split and client arrival visualizations.

    Note: Requires `matplotlib` to be installed if enabled.
    """

    visualize_data_split: bool = Field(
        default=False,
        description="Generates and saves a chart in .png format illustrating the distribution of dataset samples across the clients.",
    )

    visualize_client_arrivals: bool = Field(
        default=False,
        description="Generates and saves a timeline plot in .png format depicting the simulated arrival times and latencies of the clients.",
    )