BranchKey Python Client

Official Python client for the BranchKey federated learning and analytics platform. This library provides a simple interface to upload model weights, compute and upload federated analytics, download aggregated results, and track training runs.

Installation

pip install branchkey

Requirements: Python 3.9 or higher

Quick Start

1. Get Credentials

Create a leaf entity through the BranchKey platform to obtain credentials via the /v2/entities API endpoint.

2. Initialise Client

from branchkey import (
    Client,
    Credentials,
    APIConfig,
    RabbitMQConfig,
    WebSocketConfig,
    RunConfig,
    RetryConfig,
)

# Create credentials
credentials = 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",
)

# Initialise client with default settings
client = Client(credentials)

# Or with custom configuration
client = Client(
    credentials=credentials,
    api_config=APIConfig(
        host="https://app.branchkey.com",
        ssl=True,
    ),
    rabbitmq_config=RabbitMQConfig(
        port=5671,
        ssl=True,
    ),
    run_config=RunConfig(
        wait_for_run=False,
        check_interval_s=30,
    ),
)

3. Upload Model Weights

import numpy as np

# Prepare model weights
weighting = 1000  # Weight for aggregation (typically number of samples)
parameters = [layer1_weights, layer2_weights, ...]

# Save and upload
file_path = client.save_weights("model_weights", weighting, parameters)
file_id = client.file_upload(file_path)
print(f"Uploaded: {file_id}")

4. Download Aggregated Results

# Wait for aggregation notification
aggregation_id = client.queue.get(block=True)  # Blocks until aggregation ready
client.file_download(aggregation_id)
print(f"Downloaded to: {client.output_dir}/{aggregation_id}.npz")

# Or check without blocking
if not client.queue.empty():
    aggregation_id = client.queue.get(block=False)
    client.file_download(aggregation_id)

Configuration

All configuration uses immutable dataclasses for type safety and clarity.

Credentials

from branchkey import Credentials

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

# Or from a dictionary
credentials = Credentials.from_dict(creds_dict)

API Configuration

from branchkey import APIConfig

api_config = APIConfig(
    host="https://app.branchkey.com",  # API endpoint (default)
    ssl=True,                           # Verify SSL certificates (default)
    proxies=None,                       # Optional proxy dict
)

Transport: WebSocket vs AMQP (RabbitMQ)

The client supports two transport mechanisms for receiving aggregation notifications. WebSocket is the default. The AMQP/RabbitMQ transport is deprecated and will be removed in branchkey 3.0.0; selecting it emits a DeprecationWarning.

WebSocket (default)

from branchkey import Client, Credentials, WebSocketConfig

client = Client(
    credentials=credentials,
    websocket_config=WebSocketConfig(
        max_reconnect_attempts=0,       # 0 = infinite retry (default)
        reconnect_backoff_factor=2.0,   # Exponential backoff multiplier
        reconnect_max_delay=60,         # Max delay in seconds
    ),
    use_websocket=True,  # Default; may be omitted
)

# Receive aggregations via polling
aggregation_id = client.get_latest_aggregation_id()
if aggregation_id:
    client.file_download(aggregation_id)

AMQP/RabbitMQ (deprecated)

from branchkey import Client, Credentials, RabbitMQConfig

client = Client(
    credentials=credentials,
    rabbitmq_config=RabbitMQConfig(
        host=None,                      # Auto-derived from API host
        port=5671,                      # TLS port (default)
        ssl=True,                       # Use TLS (default)
        max_reconnect_attempts=0,       # 0 = infinite retry (default)
        reconnect_backoff_factor=2.0,   # Exponential backoff multiplier
        reconnect_max_delay=60,         # Max delay in seconds
    ),
    use_websocket=False,  # Required - the default is now WebSocket
)

# Receive aggregations via queue
aggregation_id = client.queue.get(block=True)

Passing rabbitmq_config without use_websocket=False warns, because the config is ignored: the client connects over WebSocket.

Run Configuration

from branchkey import RunConfig

run_config = RunConfig(
    wait_for_run=False,     # Wait if run is paused before uploading
    check_interval_s=30,    # Run status check interval in seconds
)

HTTP Retry Configuration

The client automatically retries failed HTTP requests with exponential backoff:

from branchkey import RetryConfig

retry_config = RetryConfig(
    max_retries=3,                                   # Maximum retry attempts
    backoff_factor=1.0,                              # Backoff multiplier (seconds)
    total_timeout=30,                                # Request timeout in seconds
    status_forcelist=(408, 429, 500, 502, 503, 504), # HTTP codes to retry
    allowed_methods=("GET", "POST", "PUT"),          # Methods that support retry
)

client = Client(credentials, retry_config=retry_config)

Retry Behaviour:

• Retries on: 408, 429, 5xx errors, connection timeouts
• Does NOT retry: Other 4xx client errors (400, 401, 403, 404)
• Backoff delays: Exponential (1s, 2s, 4s, ...)

Configuration Examples:

# Production: More retries, longer timeout
production_retry = RetryConfig(max_retries=5, backoff_factor=2.0, total_timeout=60)

# Development: Faster failure
dev_retry = RetryConfig(max_retries=1, backoff_factor=0.5, total_timeout=10)

Complete Configuration Example

from branchkey import (
    Client,
    Credentials,
    APIConfig,
    RabbitMQConfig,
    WebSocketConfig,
    RunConfig,
    RetryConfig,
)

client = Client(
    credentials=Credentials(
        id="leaf-uuid",
        name="my-leaf",
        session_token="token",
        tree_id="tree-uuid",
        branch_id="branch-uuid",
        owner_id="user-uuid",
    ),
    api_config=APIConfig(
        host="https://app.branchkey.com",
        ssl=True,
    ),
    rabbitmq_config=RabbitMQConfig(
        port=5671,
        ssl=True,
        max_reconnect_attempts=10,
    ),
    websocket_config=WebSocketConfig(
        max_reconnect_attempts=10,
    ),
    run_config=RunConfig(
        wait_for_run=True,
        check_interval_s=15,
    ),
    retry_config=RetryConfig(
        max_retries=5,
        backoff_factor=2.0,
    ),
    use_websocket=False,  # False for AMQP, True for WebSocket
    output_dir="./aggregated_output",  # Directory for downloaded files
)

Model Weight Format

Model weights are stored in compressed NPZ format.

Structure

# Format: (weighting, [list_of_parameter_arrays])
weighting = 1000  # Weight for aggregation (see below)
parameters = [layer1, layer2, ...]  # List of numpy arrays

Weighting Options

The weighting parameter controls how much influence this update has during aggregation:

1. By Sample Count (Most Common)

weighting = len(train_dataset)  # e.g., 1000 samples
# Client with 1000 samples has 2x influence of client with 500 samples

2. Equal Weighting

weighting = 1  # All clients have equal influence

3. Quality-Based Weighting

validation_accuracy = 0.85
weighting = len(train_dataset) * validation_accuracy  # Weight by quality

PyTorch Example

import numpy as np

# Using client helper
weighting = len(train_dataset)
parameters = []
for name, param in model.named_parameters():
    parameters.append(param.data.cpu().detach().numpy())

file_path = client.save_weights("model_weights", weighting, parameters)
file_id = client.file_upload(file_path)

# Using convert_pytorch_numpy
weighting, parameters = client.convert_pytorch_numpy(
    model.named_parameters(),
    weighting=len(train_dataset)
)
file_path = client.save_weights("model_weights", weighting, parameters)
file_id = client.file_upload(file_path)

TensorFlow/Keras Example

import numpy as np

weighting = len(train_dataset)
parameters = [layer.numpy() for layer in model.trainable_weights]

file_path = client.save_weights("model_weights", weighting, parameters)
file_id = client.file_upload(file_path)

Loading Aggregated Weights

import numpy as np

# Load aggregated weights from NPZ file
npz_data = np.load(f"{client.output_dir}/aggregation_id.npz")

# Note: Aggregated results only contain layers (no weighting)
layer_keys = sorted([k for k in npz_data.files if k.startswith('layer_')])
parameters = [npz_data[k] for k in layer_keys]

# Apply to PyTorch model
import torch
for i, param in enumerate(model.parameters()):
    param.data = torch.from_numpy(parameters[i])

Federated Analytics

Federated learning sends model weights. Federated analytics answers a question about the data itself — "what is the mean age across the whole federation?" — without any model being trained. There are two methods, and the difference between them is who does the maths.

save_analytics is the paved road. Reach for save_fields when your own code produces the number and there is no raw column to reduce.

save_analytics — the six-value bundle

Hand over the raw column. The SDK reduces it before anything is written to disk or sent over the wire:

import numpy as np

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 named 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.

Why sums and not the statistics themselves

This is the question everyone asks, so: sums combine across sites and statistics do not.

• Site A holds [10, 10, 10] — mean 10, variance 0
• Site B holds [90, 90, 90] — mean 90, variance 0
• Pooled, that is [10, 10, 10, 90, 90, 90] — mean 50, variance 1600

Two variances of zero pool to 1600, because the pooled figure depends on how far apart the site means are, and per-site variances cannot express that. Send the bundle and it comes out exactly right:

n = 6,  Σx = 300,  Σx² = 24600
mean     = Σx / n           = 50
variance = Σx²/n − mean²    = 4100 − 2500 = 1600

So Σx makes the mean combinable, and Σx² makes the variance combinable.

How the platform combines them

Nothing is configured — the arithmetic decides. From those, the platform derives count, sum, min, max, range, mean, variance and std, all exact, with no approximation anywhere.

Medians and percentiles are not available from a fixed-size bundle. They are genuinely not decomposable, and no bundle of any size gives them.

What is accepted, and what is rejected

Columns must be 1-D and numeric — integer, float or boolean. Strings, objects, datetimes and complex numbers are rejected by name, because Σx and Σx² mean nothing for them; send those with save_fields instead. A multi-dimensional value such as a spacing vector is also rejected, for the same reason: use save_fields.

NaN is dropped by default (nan_policy="omit") and a warning tells you how many values went from which column. n then counts only the values that contributed, so the mean and variance stay exact over the data that was actually present, and <column>_nan carries how many were dropped — so n + nan is the number of records the site held, and a reader can see that the federation dropped 12 of 300 values rather than being shown a smaller n and left to assume a smaller cohort. Pass nan_policy="raise" if a missing value means your export is wrong rather than the record is incomplete. Under neither policy can a NaN reach Σx — one missing value at one site would otherwise turn the whole federation's sum, mean, variance and std into NaN, with nothing to say which site caused it. Infinities are always rejected. A column that is empty, or entirely NaN, is rejected rather than sent as a placeholder that would corrupt the federation's min and max.

<column>_nan is written on every payload, including when it is zero and including under nan_policy="raise" where it can only be zero. The field set a site sends must be a property of the code, not of that site's data or arguments: if it varied, two participants in the same round would disagree about which fields they send, and the branch would be terminally rejected whichever way it was configured.

min and max are the only two entries that are real individual records — the minimum age is one actual patient's age. They are kept because range and normalisation need them; it is stated here rather than left to be discovered.

The bundle is computed and stored in float64. Σx² loses precision for very large magnitudes, since the variance then comes out as a small difference between two big numbers; for ages, spacings, intensities and case counts float64 is comfortably sufficient.

save_fields — values you computed yourself

When your own code produces the number, send it as it is. Each entry keeps its name through the round trip, and you choose the combining operation per field in the branch configuration (min, max, sum, mean, must_match, …).

import numpy as np

file_path = client.save_fields(
    {
        "target_spacing": np.array([1.0, 0.8, 0.8]),  # nnU-Net planner output
        "num_channels": 4,                            # must_match across sites
        "n_cases": 312,                               # sum across sites
    },
    kind="federated_analytics",
)
file_id = client.file_upload(file_path)

save_fields sends exactly what you give it. If a value is a raw record, that record leaves the site — only save_analytics guarantees otherwise.

Because the operation cannot be inferred for a field like target_spacing — min, max and mean are all plausible and each produces a different preprocessing plan — the platform asks you to assign one before the first aggregation runs. Bundle entries from save_analytics are recognised by their suffix and pre-filled.

Payload kinds

Every archive declares its kind, and the declaration is sent to the platform with the upload. The kind does not describe the payload's content — it selects how the platform combines it:

An archive written by save_weights carries no explicit declaration and is treated as federated_learning — so files written by earlier releases of this SDK behave exactly as before, including still being rejected if they contain no layer arrays.

kind="federated_learning" on save_fields is for named layer tensors (e.g. FedBN, where aggregation runs over named parameters rather than positional ones); it requires a weighting.

Field names and values

For save_fields, values may be numpy arrays, scalars, or nested sequences; a bare scalar becomes a one-element entry. Dtypes must be numeric, boolean or string — anything the platform cannot combine (object, datetime, structured) is rejected by name.

Two classes of name are rejected, before anything is written:

• weighting and __bk_payload_kind__ — reserved by the platform.
• layer_0, layer_1, … — the positional model-parameter namespace. Other names containing "layer" (bn1.weight, layer_norm) are fine.

Field shapes do not have to agree with each other — num_channels is a scalar and target_spacing a 3-vector. What must agree is the same field across sites, which the platform checks and rejects by name.

Performance Metrics

Submit training or testing metrics:

import json

metrics = {"accuracy": 0.95, "loss": 0.12}
client.send_performance_metrics(
    aggregation_id="aggregation-uuid",
    data=json.dumps(metrics),
    mode="test"  # "test", "train", or "non-federated"
)

Client Properties

client.run_status        # Current run status: "start", "stop", or "pause"
client.run_number        # Current run iteration
client.leaf_id           # Your leaf UUID
client.branch_id         # Parent branch UUID
client.tree_id           # Tree UUID
client.is_initialized    # Initialisation status
client.use_websocket     # True if using WebSocket transport
client.output_dir        # Directory for downloaded aggregated files

Branch Configuration

Fetch branch configuration including model-specific settings:

config = client.get_branch_config()
model_config = config.get("model_config", {})
sklearn_params = model_config.get("sklearn_params", {})

Advanced Features

Proxy Support

from branchkey import Client, Credentials, APIConfig

proxies = {
    'http': 'http://user:password@proxy.example.com:8080',
    'https': 'http://user:password@proxy.example.com:8080',
}

client = Client(
    credentials=credentials,
    api_config=APIConfig(proxies=proxies),
)

Context Manager

Use the client as a context manager for automatic cleanup:

from branchkey import Client, Credentials

with Client(credentials) as client:
    # Upload model
    file_path = client.save_weights("model", 1000, parameters)
    file_id = client.file_upload(file_path)

    # Download aggregation
    if not client.queue.empty():
        aggregation_id = client.queue.get(block=False)
        client.file_download(aggregation_id)
# Connections automatically closed

Error Handling

try:
    file_id = client.file_upload(file_path)
except Exception as e:
    print(f"Upload failed: {e}")
    # Logs include:
    # - HTTP status codes
    # - Response content preview
    # - Retry attempt information

Public API

from branchkey import (
    # Main client
    Client,

    # Configuration (frozen dataclasses)
    Credentials,
    APIConfig,
    RabbitMQConfig,
    WebSocketConfig,
    RunConfig,
    RetryConfig,

    # Utilities
    get_metadata,

    # Payload kinds
    PAYLOAD_KIND_FEDERATED_LEARNING,
    PAYLOAD_KIND_FEDERATED_ANALYTICS,

    # Analytics bundle
    ANALYTICS_BUNDLE_SUFFIXES,   # ("n", "sum", "sumsq", "min", "max", "nan")
    NAN_POLICY_OMIT,
    NAN_POLICY_RAISE,
)

Support

• Website: https://branchkey.com
• Repository: https://gitlab.com/branchkey/client_application
• Email: info@branchkey.com

BranchKey - Federated Learning Platform