Skip to main content

Production-grade Python Client SDK for room-based Federated Learning

Project description

DistFL

Production-Grade Federated Learning Client SDK

Python PyPI License Tests


Bring your own model (PyTorch or Scikit-Learn), connect to a DistFL server, train locally on private data, and let the server aggregate updates — all via compressed WebSocket communication. No raw data ever leaves the client.

┌──────────────┐     ┌──────────────┐     ┌──────────────┐
│   Client A   │     │   Client B   │     │   Client C   │
│  (Hospital)  │     │  (Bank)      │     │  (Lab)       │
│  Local Data  │     │  Local Data  │     │  Local Data  │
└──────┬───────┘     └──────┬───────┘     └──────┬───────┘
       │   model updates    │   (gzip+WS)        │
       └────────────────────┼────────────────────┘
                            │
                    ┌───────▼────────┐
                    │  DistFL Server │
                    │  (Go Backend)  │
                    │  FedAvg Agg.   │
                    └───────┬────────┘
                            │
                    aggregated global model
                    broadcast to all clients

✨ Features

Category Details
BYOM Use any PyTorch nn.Module or Scikit-Learn estimator with partial_fit
Simple Lifecycle initialize()validate()start() — 3 calls to go from zero to training
Room-Based FL Create rooms, share invite codes, configure training params per room
Compressed WebSocket GZIP-compressed binary messages over persistent WebSocket connections
Auto Reconnect Exponential backoff with configurable delays and heartbeat pings
Crash Recovery SQLite-backed state persistence — no duplicate round submissions after restart
Live Dashboard Built-in web UI with real-time loss curves, ΔW tracking, and training logs
Prediction Extract globally-aggregated weights and run inference locally
CLI distfl run, distfl create-room, distfl join-room, distfl ui, distfl status

📦 Installation

pip install distfl-client

From source:

git clone https://github.com/AbhaySingh002/new-repo-code.git
cd new-repo-code/DistFL
pip install -e ".[dev]"

🚀 Quick Start

1. Room Creator

The creator relies on the FLClient to initialize the global model architecture, create a new room on the server, and wait for other participants to join before starting.

[!NOTE] Why partial_fit?
Federated Learning requires all clients to share the exact same weight matrix architecture. For Scikit-Learn models like SGDClassifier, the shape of the weights (coef_ and intercept_) isn't initialized until it sees training data. We run a single dummy partial_fit on the creator side to establish this shape before sending it to the server.

from sklearn.linear_model import SGDClassifier
from fl_client import FLClient
import pandas as pd
import numpy as np

# Prepare model (scikit-learn requires partial_fit to initialize weights)
model = SGDClassifier(loss="log_loss", penalty="l2", max_iter=1,
                      learning_rate="constant", eta0=0.01)
df = pd.read_csv("./data.csv")
X = df.drop(columns=["label"]).values[:10].astype(np.float64)
y = df["label"].values[:10].astype(np.int64)
model.partial_fit(X, y, classes=[0, 1])

# Create room
client = FLClient(server_url="wss://fedlearn-server.onrender.com")
room = client.create_room(
    model=model,
    data_path="./data.csv",
    target="label",
    training_config={"local_epochs": 1, "batch_size": 32, "learning_rate": 0.01},
    room_name="Phishing Detection",
)

room_id = room["id"]
print(f"✅ Room created: {room_id}")
print(f"   Invite code: {room['invite_code']}")

# Wait for participants, then start
client.wait_for_clients(min_clients=2, timeout=120)
client.start_training()

2. Room Joiner

Each participant joins an existing room, validates their local dataset, and starts training. This follows the 3-Step Lifecycle:

  1. initialize() — Connects to the server, fetches the room's data schema and model configuration, and injects the latest global model weights into your local model.
  2. validate() — Checks your local dataset (data.csv) against the room's expected schema (e.g., ensuring it has the correct target column and feature count) and performs a dummy forward pass to catch shape errors early.
  3. start() — Signals readiness to the server and blocks while entering the federated training loop.
from sklearn.linear_model import SGDClassifier
from fl_client import FLClient

model = SGDClassifier(loss="log_loss", penalty="l2", max_iter=1,
                      learning_rate="constant", eta0=0.01)
# ... partial_fit to initialize shape (same architecture as creator)

client = FLClient(server_url="wss://fedlearn-server.onrender.com")
client.join(room_id, invite_code="abc123", model=model)
client.validate("./data.csv")
client.ready()
client.start(max_rounds=5)  # Blocks until training completes

print("✅ Training complete!")

3. PyTorch Models

import torch.nn as nn
from fl_client import FLClient

class PhishingMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(30, 64), nn.ReLU(),
            nn.Linear(64, 32), nn.ReLU(),
            nn.Linear(32, 2),
        )
    def forward(self, x):
        return self.net(x)

client = FLClient(server_url="wss://fedlearn-server.onrender.com")
room = client.create_room(
    model=PhishingMLP(),
    data_path="./data.csv",
    target="label",
    training_config={"local_epochs": 2, "batch_size": 64, "learning_rate": 0.001},
    room_name="PyTorch FL Room",
)

4. Prediction After Training

Because clients can disconnect, crash, or experience network drops, the DistFL SDK maintains a local SQLite State Database.

[!TIP] Why connect to the DB?
The server does not hold your data. Your final, fully-trained aggregated model weights are saved to your local fl_client_state.db at the end of training. By loading the state for your specific client_id (e.g. worker-1), you can extract these weights and run predictions locally without ever needing to communicate with the server again.

from fl_client.storage.db import StateDB
from fl_client.model.wrapper import wrap_model

db = StateDB("fl_client_state.db")
state = db.load_state("worker-1")

wrapper = wrap_model(model)
wrapper.set_weights(state.last_weights)

predictions = model.predict(X_test)
accuracy = (predictions == y_test).mean()
print(f"✅ Accuracy: {accuracy * 100:.2f}%")

💻 CLI Reference

# Full lifecycle from a YAML config
distfl run --config config.yaml

# Create a room
distfl create-room --server-url wss://fedlearn-server.onrender.com --room-name "My Room"

# Join a room and train
distfl join-room ROOM_ID --data ./data.csv --server-url wss://fedlearn-server.onrender.com

# Launch the real-time web dashboard
distfl ui --port 5050

# Inspect persisted client state
distfl status --client-id worker-1 --db-path fl_client_state.db

# Clear persisted state
distfl clear --client-id worker-1 --db-path fl_client_state.db

⚙️ Configuration

All options can be set via YAML file, CLI flags, or environment variables (FL_ prefix):

# Server connection
server_url: "wss://fedlearn-server.onrender.com"
room_id: ""                          # Leave empty to create a new room
client_id: ""                        # Auto-generated if omitted

# Dataset
data_path: "./data.csv"
label_column: "label"

# Training hyperparameters
batch_size: 32
local_epochs: 2
learning_rate: 0.001

# State persistence
db_path: "fl_client_state.db"        # SQLite for crash recovery

# Networking
reconnect_max_delay: 60.0            # Max backoff delay (seconds)
reconnect_base_delay: 1.0            # Initial reconnect delay
heartbeat_interval: 30.0             # WebSocket ping interval

# Dashboard
dashboard_port: 5050                 # Real-time metrics UI (0 = disabled)

# Logging
log_level: "INFO"                    # DEBUG, INFO, WARNING, ERROR

🧪 Supported Frameworks

Framework Requirements Weight Extraction
PyTorch Any nn.Module state_dict() → 3D float32 lists
Scikit-Learn Estimator with partial_fit (e.g. SGDClassifier, SGDRegressor) coef_ + intercept_ → 3D float32 lists

🧪 Testing

# Install dev dependencies
pip install -e ".[dev]"

# Run all 52 unit tests
python -m pytest tests/ -v

Test Coverage

Module Tests What's Covered
test_compressor.py 7 Compress/decompress round-trip, empty data, large payloads
test_connection.py 8 WS URL construction, connect/disconnect, send/receive
test_serializer.py 9 Serialize/deserialize, shape preservation, JSON round-trip
test_storage.py 7 SQLite save/load, upsert, clear, round logging
test_trainer.py 4 Train results, finite loss, accuracy metrics, multi-epoch
test_validation.py 17 NaN/Inf/shape/range checks, loss validation, weight shapes

🔐 Privacy & Security

  • Data never leaves the client — only model weight updates are transmitted
  • GZIP compression — reduces bandwidth and adds a layer of obfuscation
  • Server-side validation — NaN, Inf, out-of-range, shape mismatch, L2 norm, and duplicate submission checks
  • Invite codes — rooms can be access-controlled via invite codes
  • Crash recovery — SQLite persistence prevents duplicate round submissions

📄 License

MIT


Built with ❤️ for privacy-preserving machine learning

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

distfl_client-1.0.3.tar.gz (47.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

distfl_client-1.0.3-py3-none-any.whl (55.1 kB view details)

Uploaded Python 3

File details

Details for the file distfl_client-1.0.3.tar.gz.

File metadata

  • Download URL: distfl_client-1.0.3.tar.gz
  • Upload date:
  • Size: 47.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for distfl_client-1.0.3.tar.gz
Algorithm Hash digest
SHA256 5378749b01a4b71466682363af0131dc8b66683b9f14b42e95f0f0b40b535e01
MD5 62a3dddb3f437b2e5deac64109dc35ad
BLAKE2b-256 ade6d19c667b4ee1cd5bc418e65dc978ad66e389d12c2dd346f84abaab1a4922

See more details on using hashes here.

Provenance

The following attestation bundles were made for distfl_client-1.0.3.tar.gz:

Publisher: publish.yml on DistFl/DistFl-Client

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file distfl_client-1.0.3-py3-none-any.whl.

File metadata

  • Download URL: distfl_client-1.0.3-py3-none-any.whl
  • Upload date:
  • Size: 55.1 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for distfl_client-1.0.3-py3-none-any.whl
Algorithm Hash digest
SHA256 f74173c36663d5f3eaed0ee79f423a9599526ce05d8d3fc09e6943f5cd8e26f3
MD5 1ea8e05525d54a1b568b2041fa36dd0a
BLAKE2b-256 24623748d1032e87eb1a1640b3d096babc6f38806e73cac039afd258f3c5e124

See more details on using hashes here.

Provenance

The following attestation bundles were made for distfl_client-1.0.3-py3-none-any.whl:

Publisher: publish.yml on DistFl/DistFl-Client

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page