Metadata-Version: 2.4
Name: branchkey
Version: 2.9.3
Summary: Client application to interface with the BranchKey system
Home-page: https://branchkey.com
Author: BranchKey
Author-email: info@branchkey.com
Project-URL: Homepage, https://branchkey.com
Project-URL: Repository, https://gitlab.com/branchkey/client_application
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
Classifier: Operating System :: OS Independent
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: requests==2.32.3
Requires-Dist: numpy==1.26.4
Requires-Dist: pika==1.3.2
Requires-Dist: pysocks==1.7.1
Requires-Dist: websockets>=12.0
Requires-Dist: aiohttp>=3.9.0
Dynamic: author
Dynamic: author-email
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: license-file
Dynamic: project-url
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# BranchKey Python Client

![BK_logo](https://branchkey.com/branding/bk-logo-medium.png)

[![PyPI version](https://badge.fury.io/py/branchkey.svg)](https://badge.fury.io/py/branchkey)
[![Python](https://img.shields.io/pypi/pyversions/branchkey.svg)](https://pypi.org/project/branchkey/)
[![License: GPL v3](https://img.shields.io/badge/License-GPLv3-blue.svg)](https://www.gnu.org/licenses/gpl-3.0)

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

```bash
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

```python
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

```python
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

```python
# 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

```python
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

```python
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)

```python
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)

```python
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

```python
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:

```python
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:**

```python
# 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

```python
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

```python
# 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)**

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

**2. Equal Weighting**

```python
weighting = 1  # All clients have equal influence
```

**3. Quality-Based Weighting**

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

### PyTorch Example

```python
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)
```

```python
# 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

```python
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

```python
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` | `save_fields` |
|---|---|---|
| You hand it | a **raw column** of your own records | values **you have already computed** |
| The SDK does | reduces each column to six combinable statistics | writes your values as given |
| Raw values leave the site | **never — enforced by the library** | whatever you pass is what is sent |
| Combining operation | fixed by the arithmetic, pre-filled for you | you choose it per field in the branch config |
| Use it for | age, volume, intensity, any measured column | nnU-Net planner output, label histograms, channel counts |

`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**:

```python
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

| Entry | Combiner |
|---|---|
| `_n`, `_sum`, `_sumsq`, `_nan` | summed across sites |
| `_min` | minimum across sites |
| `_max` | maximum across sites |

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`, …).

```python
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:

| Kind | Written by | Layout | Combined by | `weighting` |
|---|---|---|---|---|
| `federated_learning` | `save_weights` (and `save_fields(kind="federated_learning")`) | `weighting` + `layer_0..layer_n` | Weighted average of model parameters | Required |
| `federated_analytics` | `save_analytics`, `save_fields` (default) | `weighting` + named fields | One operation per field, across sites | Not used |

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:

```python
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

```python
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:

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

## Advanced Features

### Proxy Support

```python
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:

```python
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

```python
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

```python
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](https://branchkey.com)
- **Repository**: [https://gitlab.com/branchkey/client_application](https://gitlab.com/branchkey/client_application)
- **Email**: info@branchkey.com

---

**BranchKey** - Federated Learning Platform
