BranchKey Python Client

Official Python client for the BranchKey federated learning and analytics platform. Upload model weights, compute and upload federated analytics, download aggregated results, and track training runs — without raw data leaving your site.

Full documentation: app.branchkey.com/docs. This page is the short version.

3.0.0 is a breaking release. The AMQP/RabbitMQ transport is gone; WebSocket is the only transport and needs no configuration. Drop any rabbitmq_config or use_websocket argument — both now raise TypeError. Nothing else changed. pip install "branchkey<3" pins the previous behaviour.

Installation

pip install branchkey

Requires Python 3.10 or higher (tested on 3.10 – 3.14).

Quick start

Create a leaf entity on the platform to obtain its credentials — see Getting started → Access and setup.

from branchkey import Client, Credentials

client = Client(
    Credentials(
        id="your-leaf-uuid",
        name="my-client",
        session_token="your-session-token-uuid",
        owner_id="your-user-uuid",
        tree_id="your-tree-uuid",
        branch_id="your-branch-uuid",
    )
)

# 1. Upload this round's model weights.
#    `weighting` is this site's influence during aggregation — usually the
#    number of training samples. `parameters` is a list of numpy arrays.
file_path = client.save_weights("model_weights", weighting=1000, parameters=parameters)
file_id = client.file_upload(file_path)

# 2. Wait for the aggregated result and download it.
aggregation_id = client.queue.get(block=True)
client.file_download(aggregation_id)  # -> {client.output_dir}/{aggregation_id}.npz

client.queue.get(block=False) polls instead of blocking, and client.get_latest_aggregation_id() returns the most recent notification or None.

The downloaded archive contains the aggregated layers only, under the keys layer_0, layer_1, …, with no weighting field:

import numpy as np

npz = np.load(f"{client.output_dir}/{aggregation_id}.npz")
parameters = [npz[k] for k in sorted(npz.files) if k.startswith("layer_")]

Federated analytics

Federated analytics answers a question about the data itself — "what is the mean age across the federation?" — with no model trained. Hand save_analytics a raw column and the SDK reduces it to six combinable numbers before anything is written to disk or sent over the wire:

file_path = client.save_analytics(
    {
        "age": patients["age"].to_numpy(),
        "tumour_volume": patients["volume"].to_numpy(),
    }
)
file_id = client.file_upload(file_path)

The archive holds six entries per column, and nothing else:

age -> age_n, age_sum, age_sumsq, age_min, age_max, age_nan

n is the count, sum is Σx, sumsq is Σx², min/max are the extremes, and nan is how many values were missing and left out. Not a single patient's age is in the file. That is a property of the library, not a request made of you: there is no argument that turns it off.

Sums are sent rather than statistics because sums combine across sites and statistics do not — two sites each reporting a variance of zero can pool to a variance of 1600. From the bundle the platform derives count, sum, min, max, range, mean, variance and standard deviation, all exact. Columns must be 1-D and numeric; NaN is dropped by default and counted in _nan.

Already computed the number yourself — a preprocessing planner's output, a label histogram? Send it with save_fields instead, and choose its combining operation in the branch configuration.

Where to go next

Support

• Documentation: app.branchkey.com/docs
• Website: branchkey.com
• Email: info@branchkey.com

Licence

GPL-3.0 — the full text ships with the package as LICENSE.

BranchKey — Federated Learning Platform