Skip to content

matteovhaxt/capstone

BranchesTags

Repository files navigation

Capstone

Table of Contents

Overview

This repository contains a federated learning environment to train the CIFAR-10 dataset. The training runs are configurable via the Makefile. The main objective of this project is conducting experiments to assess whether an adaptive approach to client weight quantization can achieve better outcomes that non- or uniformly quantized runs, measured by time and transmitted bytes.

Get started

This repo uses the Makefile for setup and runs (uv is required).

Install dependencies and create virtual environment:

make setup

Your first run with simulated local clients:

make dev

Set the necessary environment variables:

cp .env.example .env

Your first end-to-end production training run:

make start

Other useful scripts:

  • make start-detached: run the production training w/o streaming
  • make build: build the server app docker image
  • make format: run ruff code formatting and fixes
  • make test: run the testing suite using pytest

System architecture

The application entry point is the server_app.py that is located on your machine, or the remote repository when using the run-experiment.yml GitHub action. Upon starting a run, the ServerApp is sent to the Superlink that is running on the production server. The main component of the training run is the BaseFedAvg strategy. It starts the client instances via the RunService, that communicates to the dstack container running on the VPS. After the specified minimum of clients have been started, the training run begins. During the run, metrics are captured by various trackers and communicated to Weights & Biases via the MonitoringService.

graph LR
    subgraph "Developer Machine"
        SERVERAPP[Server App<br/>server_app.py]
    end

    subgraph "VPS Host"
        subgraph "Docker Compose Stack"
            SUPERLINK[Flower SuperLink]
            DSTACK[dstack Server]
        end
    end

    subgraph "RunPod Cloud"
        subgraph "GPU Instance"
            SN[SuperNode<br/>RTX4090]
            CLIENT[Client App<br/>client_app.py]
        end
    end

    subgraph "External Services"
        WANDB[Weights & Biases<br/>Monitoring]
    end

    SERVERAPP -->|connects via<br/>ServerAppIo API| SUPERLINK
    SERVERAPP -->|submits tasks| DSTACK
    SERVERAPP -->|logs metrics| WANDB

    DSTACK -->|provisions & launches| SN

    SUPERLINK <-->|Fleet API| SN
    SN -->|executes| CLIENT
    CLIENT -->|model updates<br/>quantized gradients| SUPERLINK

    CLIENT -.->|training metrics| WANDB

    classDef vps fill:#e1f5ff,stroke:#0288d1,stroke-width:2px,color:#000
    classDef runpod fill:#fff3e0,stroke:#f57c00,stroke-width:2px,color:#000
    classDef external fill:#f3e5f5,stroke:#7b1fa2,stroke-width:2px,color:#000
    classDef dev fill:#e8f5e9,stroke:#388e3c,stroke-width:2px,color:#000

    class SUPERLINK,DSTACK vps
    class SN,CLIENT runpod
    class WANDB external
    class SERVERAPP dev
Loading

Repository structure

  • src/server: server app entry points, strategies, and services
  • src/client: client app entry points and helpers
  • src/core: shared data, model, training, and quantization logic
  • tests: unit tests for the client, server, and core layers

Main components

Server

  • src/server/server_app.py: main entry point that builds the strategy and runs training
  • src/server/services/monitoring_service.py: initializes the W&B client and exposes a method to log metrics
  • src/server/services/run_service.py: starts, polls status of, and stops RunPod GPU instances using dstack
  • src/server/services/quantization_service.py: assigns quantization factors to clients based on predefined split
  • src/server/strategies/base.py: the strategy class that contains the services and trackers, and hooks into the flwr strategy methods to extend the training functionality
  • src/server/strategies/factory.py: helper method to construct the BaseFedAvg class according to the defined configuration
  • src/server/strategies/utils.py: utility functions and small tracker classes used in the strategy

The diagram below shows the data flow of the BaseFedAvg strategy.

graph LR
    subgraph "Setup"
        START[start<br/>Reset trackers]
    end
    
    subgraph "Per Round"
        ROUND{Round<br/>loop}
        
        subgraph "Configure Phase"
            CONFIG[configure_train<br/>Build messages]
            QUANT{Quantization?}
            ASSIGN[assign_bits]
            ATTACH[Attach to<br/>message configs]
        end
        
        subgraph "Client Training"
            CLIENTS[Clients execute<br/>local training]
        end
        
        subgraph "Aggregate Phase"
            AGG[aggregate_train<br/>Collect replies]
            MEASURE[measure_bytes<br/>simulate_latency]
            DEQUANT[maybe_dequantize<br/>_arrays]
            FEDAVG[FedAvg<br/>weighted average]
        end
        
        EVAL[aggregate_evaluate]
    end
    
    subgraph "Services"
        RUNSVC[RunService<br/>start/stop runs]
        QUANTSVC[QuantizationService]
        MONITOR[MonitoringService<br/>log_metrics]
    end
    
    subgraph "Cleanup"
        END[Stop runs]
    end
    
    START --> ROUND
    START -.-> RUNSVC
    
    ROUND --> CONFIG
    CONFIG --> QUANT
    QUANT -->|Yes| ASSIGN
    QUANT -->|No| CLIENTS
    ASSIGN -.-> QUANTSVC
    ASSIGN --> ATTACH
    ATTACH --> CLIENTS
    
    CLIENTS --> AGG
    AGG --> MEASURE
    MEASURE --> DEQUANT
    DEQUANT --> FEDAVG
    FEDAVG --> MONITOR
    
    FEDAVG -.-> EVAL
    EVAL -.-> MONITOR
    
    MONITOR --> ROUND
    ROUND -->|Complete| END
    END -.-> RUNSVC
    
    classDef setup fill:#e8f5e9,stroke:#388e3c,stroke-width:2px,color:#000
    classDef configure fill:#e1f5ff,stroke:#0288d1,stroke-width:2px,color:#000
    classDef aggregate fill:#fff3e0,stroke:#f57c00,stroke-width:2px,color:#000
    classDef services fill:#f3e5f5,stroke:#7b1fa2,stroke-width:2px,color:#000
    classDef client fill:#ffebee,stroke:#c62828,stroke-width:2px,color:#000
    
    class START,END setup
    class CONFIG,QUANT,ASSIGN,ATTACH configure
    class AGG,MEASURE,DEQUANT,FEDAVG,EVAL aggregate
    class RUNSVC,QUANTSVC,MONITOR services
    class CLIENTS client
Loading

Client

  • src/client/client_app.py: train and evaluate methods that run on the client each round, parses run and training config, loads quantization configuration, runs training epochs, packages model updates to send to server
  • src/client/utils.py: utility functions that modularizes the behavior used in the client rounds

The diagram below shows the client round data flow and processes.

graph LR
    subgraph CT["Client Training"]
        subgraph CA["Client App"]
            CLIENTS[Clients execute<br/>local training]
            
            TRAIN_MSG[train<br/>handle Message]
            PREP_TRAIN[prepare_model_and_loaders]
            RESOLVE[resolve_quantization_bits]
            TRAIN_FN[train_fn<br/>local epochs]
            PACKAGE[package_model_update]
            TRAIN_REPLY[Message reply<br/>arrays + metrics + metadata]
            
            EVAL_MSG[evaluate<br/>handle Message]
            PREP_EVAL[prepare_model_and_loaders]
            TEST_FN[test_fn]
            EVAL_REPLY[Message reply<br/>metrics]
        end
    end
    
    CLIENTS -.-> TRAIN_MSG
    CLIENTS -.-> EVAL_MSG
    
    TRAIN_MSG --> PREP_TRAIN
    PREP_TRAIN --> RESOLVE
    RESOLVE --> TRAIN_FN
    TRAIN_FN --> PACKAGE
    PACKAGE --> TRAIN_REPLY
    
    EVAL_MSG --> PREP_EVAL
    PREP_EVAL --> TEST_FN
    TEST_FN --> EVAL_REPLY
    
    classDef allnodes fill:#ffccbc,stroke:#d84315,stroke-width:2px,color:#000
    
    class CLIENTS,TRAIN_MSG,PREP_TRAIN,RESOLVE,TRAIN_FN,PACKAGE,TRAIN_REPLY,EVAL_MSG,PREP_EVAL,TEST_FN,EVAL_REPLY allnodes
Loading

Core

  • src/core/data.py: dataset loading, partitioning and transforms
  • src/core/model.py: CNN model definition and load function
  • src/core/training.py: train and test functions
  • src/core/quantization.py: the core (de-)quantization implementation
  • src/core/config.py: RunConfig and TrainConfig data classes

Deployment

The deployment pipeline is centered around the deploy-server.yml GitHub action. It runs tests, builds the image, pushes it to GitHub container registry, uses SSH inside of the VPS to fetch the latest images and start the compose stack.
The run-experiment.yml action can be used to start a training run remotely.

graph LR
    subgraph DEV["Development"]
        DEVELOPER[Developer]
        SOURCE[Source Code<br/>Python/Config]
        LOCAL[Local Tests<br/>pytest]
    end
    
    subgraph VC["Version Control"]
        GIT[Git Repository<br/>GitHub]
    end
    
    subgraph CI_DEPLOY["CI Pipeline - Deploy Server"]
        RUN_TESTS_D[Run Tests<br/>pytest]
        BUILD[Build Container<br/>Docker]
        PUSH[Push Image<br/>GhCR]
        SSH_DEPLOY[SSH Deploy<br/>docker-compose up]
    end
    
    subgraph REGISTRY["Container Registry"]
        GHCR[GitHub Container<br/>Registry<br/>ghcr.io/matteosmalmassari/federated<br/>-quantization-FL]
    end
    
    subgraph CI_EXP["CI Pipeline - Run Experiment"]
        RUN_TESTS_E[Run Tests<br/>pytest]
        EXECUTE[Execute<br/>make start]
    end
    
    subgraph PROD["Production VPS"]
        COMPOSE[docker-compose.yml]
        DSTACK_CONTAINER[dstack Server<br/>Container]
        SUPERLINK[Flower SuperLink<br/>Container]
    end
    
    subgraph RUNTIME["Runtime"]
        SERVER_APP[Server App<br/>Connects to SuperLink]
        SUPERNODES[SuperNodes<br/>on RunPod]
    end
    
    DEVELOPER -->|writes| SOURCE
    DEVELOPER -->|runs| LOCAL
    SOURCE -->|git push| GIT
    
    GIT -->|on push to main| RUN_TESTS_D
    RUN_TESTS_D -->|if pass| BUILD
    BUILD -->|build| PUSH
    PUSH -->|SSH + deploy| SSH_DEPLOY
    PUSH -->|push| GHCR
    
    SSH_DEPLOY -->|updates| COMPOSE
    COMPOSE -->|starts| DSTACK_CONTAINER
    COMPOSE -->|starts| SUPERLINK
    
    GIT -->|manual trigger| RUN_TESTS_E
    RUN_TESTS_E -->|if pass| EXECUTE
    
    EXECUTE -->|launches| SERVER_APP
    SERVER_APP -->|submits tasks via| DSTACK_CONTAINER
    SERVER_APP -->|connects to| SUPERLINK
    DSTACK_CONTAINER -->|provisions| SUPERNODES
    SUPERNODES -->|communicate with| SUPERLINK
    
    classDef dev fill:#e8f5e9,stroke:#388e3c,stroke-width:2px,color:#000
    classDef vc fill:#fff9c4,stroke:#f57f17,stroke-width:2px,color:#000
    classDef ci fill:#ffccbc,stroke:#d84315,stroke-width:2px,color:#000
    classDef registry fill:#f3e5f5,stroke:#7b1fa2,stroke-width:2px,color:#000
    classDef prod fill:#e1f5ff,stroke:#0288d1,stroke-width:2px,color:#000
    classDef runtime fill:#fff3e0,stroke:#f57c00,stroke-width:2px,color:#000
    
    class DEVELOPER,SOURCE,LOCAL dev
    class GIT vc
    class RUN_TESTS_D,BUILD,PUSH,SSH_DEPLOY,RUN_TESTS_E,EXECUTE ci
    class GHCR registry
    class COMPOSE,DSTACK_CONTAINER,SUPERLINK prod
    class SERVER_APP,SUPERNODES runtime
Loading

Contributing

The project uses the conventional commits specification for commit messages.

About

My bachelors thesis project about communication optimization for federated learning.

Resources

Readme
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

No releases published

Packages

 
 
 

Contributors

Languages