naira-fl-sdk 0.3.2

Python SDK for building federated learning model profiles

naira-fl-sdk

Python SDK for building federated learning model profiles. Define your model, implement the trainer interface, and the SDK handles NVFlare integration automatically.

Install

pip install naira-fl-sdk

With PyTorch:

pip install naira-fl-sdk[torch]

Quick Start

1. Scaffold a new model profile

naira-fl-sdk init my-model

This creates a directory with starter files:

my-model/
├── model.yaml       # Model manifest (config, metrics, data spec)
├── model.py         # PyTorch model definition
├── trainer.py       # FLTrainer implementation
└── requirements.txt # Extra pip dependencies

2. Implement your trainer

Subclass FLTrainer and implement the required methods:

from naira_fl_sdk import FLTrainer, ValidationResult

class MyTrainer(FLTrainer):

    def setup(self, data_path: str, hyperparams: dict) -> None:
        """Load data and initialize model, optimizer, loss."""
        ...

    def get_model(self):
        """Return your nn.Module instance."""
        return self.model

    def train(self, model) -> dict:
        """Train one round. Return metrics matching model.yaml."""
        return {"train_loss": avg_loss, "num_samples": n}

    def validate(self, model) -> dict:
        """Validate one round. Return metrics matching model.yaml."""
        return {"val_loss": avg_loss, "accuracy": acc}

    def validate_data(self, data_path: str, hyperparams: dict) -> ValidationResult:
        """Check that client data is suitable for training."""
        return ValidationResult(valid=True)

3. Configure model.yaml

name: my-model
description: "My federated learning model"
framework: pytorch

model:
  file: model.py
  class: MyModel

trainer:
  file: trainer.py

requirements:
  - scikit-learn>=1.0

data:
  description: "Expects preprocessed tensors"
  required_files:
    - name: "train_data.pt"
      description: "Training data tensor"
    - name: "train_labels.pt"
      description: "Training labels tensor"
  optional_files:
    - name: "val_data.pt"
      description: "Validation data tensor"

hyperparameters:
  batch_size:
    default: 32
    type: int
    description: "Training batch size"
  learning_rate:
    default: 0.001
    type: float
    description: "Learning rate"

metrics:
  train:
    - name: train_loss
      display_name: "Training Loss"
      primary: true
    - name: num_samples
      type: int
      aggregate: sum
  validate:
    - name: val_loss
      display_name: "Validation Loss"
      primary: true
    - name: accuracy
      display_name: "Accuracy"
      format: percentage
      primary: true
  test: []

4. Validate and test

# Check model profile structure
naira-fl-sdk validate ./my-model

# Run local FL simulation (no NVFlare needed)
naira-fl-sdk test ./my-model --rounds 3 --data ./sample-data

5. Package and upload

naira-fl-sdk package ./my-model -o my-model.zip

Upload the zip through the FL platform admin UI.

CLI Reference

Command	Description
naira-fl-sdk init <name>	Scaffold a new model profile
naira-fl-sdk validate <dir>	Validate model profile structure
naira-fl-sdk test <dir>	Run local FL simulation
naira-fl-sdk package <dir>	Package into a zip for upload

FLTrainer API

Method	Required	Description
setup(data_path, hyperparams)	Yes	Initialize model, data, optimizer
get_model()	Yes	Return the nn.Module instance
train(model)	Yes	Train one round, return metrics dict
validate(model)	Yes	Validate one round, return metrics dict
test(model)	No	Final evaluation after training
validate_data(data_path, hyperparams)	Yes	Pre-training data validation

Metrics

Metrics are defined per-phase in model.yaml under metrics.train, metrics.validate, and metrics.test.

Each metric supports:

Field	Default	Description
name	(required)	Key returned by trainer methods
display_name	None	Human-readable label
type	float	float or int
format	number	number, percentage, or duration
aggregate	weighted_avg	How to aggregate across clients
primary	false	Show in overview dashboards

Aggregation strategies: weighted_avg, sum, mean, harmonic_mean, max, min.

Data Validation

The validate_data() method runs before training starts, giving clients early feedback about data issues. Return a ValidationResult:

from naira_fl_sdk import ValidationResult

def validate_data(self, data_path, hyperparams):
    errors, warnings = [], []

    if not Path(data_path, "train_data.pt").exists():
        errors.append("Missing train_data.pt")

    return ValidationResult(
        valid=len(errors) == 0,
        errors=errors,
        warnings=warnings,
    )

Requirements

• Python >= 3.9
• Dependencies: pyyaml, click, pydantic
• Optional: torch >= 2.0 (for model training)

License

MIT