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]

For AI Agents / LLMs

Building a model profile with a coding agent (Claude Code, Cursor, etc.)? Copy the whole section below into the agent. It is a self-contained, copy-paste brief: the workflow, the full API + manifest spec, the rules agents most often get wrong, and a complete worked example verified to pass validate and test.

Your job as the agent is done when naira-fl-sdk validate passes, naira-fl-sdk test trains cleanly, and you have produced the .zip.

Steps (in order)

1. Scaffold: naira-fl-sdk init <name> — <name> must match ^[a-zA-Z0-9][a-zA-Z0-9_-]*$. Creates <name>/ with starter model.py, trainer.py, model.yaml, requirements.txt.
2. Implement model.py (a torch.nn.Module), trainer.py (one FLTrainer subclass), and model.yaml per the spec below; add extra pip deps to requirements:.
3. Validate: naira-fl-sdk validate ./<name> — read any FAIL:/ERROR: lines, fix, re-run until PASS. WARN: is non-blocking.
4. Test (local FL sim, no NVFlare): naira-fl-sdk test ./<name> --rounds 2 --data <data-dir> — runs validate_data() → setup() → N×(train,validate) → test. Fix runtime errors, re-run until it trains cleanly. The CLI does NOT verify metric keys or num_samples — after it runs, manually confirm each printed dict's keys exactly match the matching metrics.<phase> names in model.yaml.
5. Package: naira-fl-sdk package ./<name> — writes <name>.zip.
6. Done: upload <name>.zip to your Naira platform admin. (SDK scope ends here.)

Spec

Imports (both are public — don't reach into submodules):

from naira_fl_sdk import FLTrainer, ValidationResult

FLTrainer interface (define exactly one subclass in trainer.py; the simulator finds it by scanning for an FLTrainer subclass, so the class name doesn't matter):

model.py is imported from trainer.py with a flat import (from model import YourModel) — the SDK puts the profile dir on sys.path.

model.yaml schema. Required: name, description, model (file + class), trainer (file), metrics. Optional: framework (pytorch default / tensorflow), requirements, data, hyperparameters.

• The model class key in YAML is literally class (not class_name).
• data: description, required_files: [{name, description}], optional_files: [{name, description}].
• hyperparameters: map of name -> {default, type, description}.
• metrics: train/validate/test lists (at least one train or validate metric). Each metric def: name (the dict key your trainer returns), display_name, type (float|int), format (number|percentage|duration), aggregate (weighted_avg|sum|mean|harmonic_mean|max|min — use sum for num_samples), group, primary.

Rules (common mistakes — do not violate)

• Returned metric dict keys must exactly equal the metrics.<phase> names. Not checked by the CLI — verify by hand.
• train() must include "num_samples": <int>. Not checked by the CLI.
• Don't re-initialize weights in train/validate/test — they arrive pre-loaded.
• Cast hyperparameters to their declared types (overrides may arrive as strings).
• validate_data() gates training: valid=False + clear errors on real problems; warnings for non-blocking issues (e.g. a missing optional file).
• Import the model with a flat from model import ..., not a relative import.
• A pydantic UserWarning ("Field name 'validate' … shadows an attribute") prints on stderr on every command — it is benign, ignore it. Real failure signals are FAIL:/ERROR: lines, a non-zero exit, or a runtime traceback.

Complete worked example (verified — passes validate and test)

A binary diabetes-risk classifier over tabular CSV, exercising all three metric phases, extra requirements, hyperparameter casting, num_samples, optional files, and a thorough validate_data().

model.yaml:

name: diabetes-risk-prediction
description: "Predict Type 2 diabetes onset from clinical diagnostic measurements (Pima Indians dataset). Binary classifier using an MLP with batch normalization and dropout."
framework: pytorch

model:
  file: model.py
  class: DiabetesNet

trainer:
  file: trainer.py

requirements:
  - scikit-learn>=1.3
  - pandas>=2.0

data:
  description: "CSV file with 8 clinical features and a binary Outcome column (0 = no diabetes, 1 = diabetes). Features: Pregnancies, Glucose, BloodPressure, SkinThickness, Insulin, BMI, DiabetesPedigreeFunction, Age."
  required_files:
    - name: "train.csv"
      description: "Training data CSV with header row. Must contain all 8 feature columns and Outcome column."
  optional_files:
    - name: "val.csv"
      description: "Validation data CSV (same schema as train.csv). If absent, 20% of training data is used for validation."
    - name: "test.csv"
      description: "Held-out test data CSV (same schema). Used for final evaluation after training completes."

hyperparameters:
  learning_rate:
    default: 0.001
    type: float
    description: "Adam optimizer learning rate"
  batch_size:
    default: 32
    type: int
    description: "Mini-batch size for training and evaluation"
  epochs_per_round:
    default: 5
    type: int
    description: "Number of local training epochs per FL round"
  hidden_dim:
    default: 64
    type: int
    description: "Hidden layer dimension in the MLP"
  dropout_rate:
    default: 0.3
    type: float
    description: "Dropout probability for regularization"

metrics:
  train:
    - name: train_loss
      display_name: "Training Loss"
      type: float
      format: number
      aggregate: weighted_avg
      primary: true
    - name: train_accuracy
      display_name: "Training Accuracy"
      type: float
      format: percentage
      aggregate: weighted_avg
      group: "accuracy"
    - name: num_samples
      display_name: "Samples Processed"
      type: int
      format: number
      aggregate: sum
  validate:
    - name: val_loss
      display_name: "Validation Loss"
      type: float
      format: number
      aggregate: weighted_avg
      primary: true
    - name: accuracy
      display_name: "Accuracy"
      type: float
      format: percentage
      aggregate: weighted_avg
      group: "classification"
      primary: true
    - name: precision
      display_name: "Precision"
      type: float
      format: percentage
      aggregate: weighted_avg
      group: "classification"
    - name: recall
      display_name: "Recall (Sensitivity)"
      type: float
      format: percentage
      aggregate: weighted_avg
      group: "classification"
    - name: f1_score
      display_name: "F1 Score"
      type: float
      format: percentage
      aggregate: harmonic_mean
      group: "classification"
    - name: auc_roc
      display_name: "AUC-ROC"
      type: float
      format: percentage
      aggregate: weighted_avg
      group: "classification"
  test:
    - name: test_loss
      display_name: "Test Loss"
      type: float
      format: number
      aggregate: weighted_avg
    - name: test_accuracy
      display_name: "Test Accuracy"
      type: float
      format: percentage
      aggregate: weighted_avg
      group: "final"
      primary: true
    - name: test_precision
      display_name: "Test Precision"
      type: float
      format: percentage
      aggregate: weighted_avg
      group: "final"
    - name: test_recall
      display_name: "Test Recall"
      type: float
      format: percentage
      aggregate: weighted_avg
      group: "final"
    - name: test_f1_score
      display_name: "Test F1 Score"
      type: float
      format: percentage
      aggregate: harmonic_mean
      group: "final"
    - name: test_auc_roc
      display_name: "Test AUC-ROC"
      type: float
      format: percentage
      aggregate: weighted_avg
      group: "final"

model.py:

"""DiabetesNet - MLP binary classifier for Type 2 diabetes prediction."""

import torch.nn as nn


class DiabetesNet(nn.Module):
    """Multi-layer perceptron with batch normalization and dropout.

    Architecture: 8 -> hidden_dim -> hidden_dim//2 -> 2
    """

    def __init__(self, input_dim=8, hidden_dim=64, dropout_rate=0.3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.BatchNorm1d(hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            nn.Linear(hidden_dim // 2, 2),
        )

    def forward(self, x):
        return self.net(x)

trainer.py:

"""DiabetesTrainer - FLTrainer implementation for diabetes risk prediction."""

import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from pathlib import Path
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
)
from torch.utils.data import DataLoader, TensorDataset

from naira_fl_sdk import FLTrainer, ValidationResult
from model import DiabetesNet

FEATURE_COLUMNS = [
    "Pregnancies",
    "Glucose",
    "BloodPressure",
    "SkinThickness",
    "Insulin",
    "BMI",
    "DiabetesPedigreeFunction",
    "Age",
]
TARGET_COLUMN = "Outcome"
ALL_COLUMNS = FEATURE_COLUMNS + [TARGET_COLUMN]


def _load_csv(path: Path, mean=None, std=None):
    """Load a CSV, standardize features, return tensors and stats."""
    df = pd.read_csv(path)
    X = df[FEATURE_COLUMNS].values.astype("float32")
    y = df[TARGET_COLUMN].values.astype("int64")

    X = torch.tensor(X)
    y = torch.tensor(y)

    if mean is None:
        mean = X.mean(dim=0)
        std = X.std(dim=0)
        std[std == 0] = 1.0  # avoid division by zero

    X = (X - mean) / std
    return X, y, mean, std


class DiabetesTrainer(FLTrainer):
    """Federated trainer for diabetes risk prediction.

    Implements all FLTrainer methods including test() for final evaluation
    and thorough validate_data() with errors and warnings.
    """

    def validate_data(self, data_path: str, hyperparams: dict) -> ValidationResult:
        errors = []
        warnings = []
        dp = Path(data_path)

        # --- Required: train.csv ---
        train_csv = dp / "train.csv"
        if not train_csv.exists():
            errors.append("Missing required file: train.csv")
        else:
            try:
                df = pd.read_csv(train_csv)
            except Exception as e:
                errors.append(f"Cannot parse train.csv: {e}")
                return ValidationResult(valid=False, errors=errors, warnings=warnings)

            # Check columns
            missing_cols = set(ALL_COLUMNS) - set(df.columns)
            if missing_cols:
                errors.append(f"train.csv missing columns: {sorted(missing_cols)}")
            else:
                # Check data types - features should be numeric
                for col in FEATURE_COLUMNS:
                    if not pd.api.types.is_numeric_dtype(df[col]):
                        errors.append(f"Column '{col}' must be numeric, got {df[col].dtype}")

                # Check target is binary
                unique_targets = set(df[TARGET_COLUMN].unique())
                if not unique_targets.issubset({0, 1}):
                    errors.append(f"Outcome must be 0 or 1, got values: {unique_targets}")

                # Check for NaN/missing values
                nan_counts = df[ALL_COLUMNS].isna().sum()
                cols_with_nan = nan_counts[nan_counts > 0]
                if len(cols_with_nan) > 0:
                    for col, count in cols_with_nan.items():
                        warnings.append(f"train.csv has {count} NaN values in '{col}' — rows will be dropped")

                # Check minimum sample count
                if len(df) < 10:
                    errors.append(f"train.csv has only {len(df)} rows — need at least 10 for training")
                elif len(df) < 50:
                    warnings.append(f"train.csv has only {len(df)} rows — may produce poor model quality")

                # Check class balance
                if len(df) > 0 and unique_targets.issubset({0, 1}):
                    positive_rate = df[TARGET_COLUMN].mean()
                    if positive_rate < 0.05 or positive_rate > 0.95:
                        warnings.append(
                            f"Severe class imbalance: {positive_rate:.1%} positive — "
                            f"model may be biased toward majority class"
                        )

                # Check for zero-value placeholders (common in this dataset)
                zero_check_cols = ["Glucose", "BloodPressure", "BMI"]
                for col in zero_check_cols:
                    if col in df.columns:
                        zero_count = (df[col] == 0).sum()
                        if zero_count > 0:
                            warnings.append(
                                f"'{col}' has {zero_count} zero values — "
                                f"these may be missing data coded as 0"
                            )

        # --- Optional: val.csv ---
        val_csv = dp / "val.csv"
        if not val_csv.exists():
            warnings.append("No val.csv found — 20% of training data will be used for validation")
        else:
            try:
                vdf = pd.read_csv(val_csv)
                missing = set(ALL_COLUMNS) - set(vdf.columns)
                if missing:
                    errors.append(f"val.csv missing columns: {sorted(missing)}")
            except Exception as e:
                errors.append(f"Cannot parse val.csv: {e}")

        # --- Optional: test.csv ---
        test_csv = dp / "test.csv"
        if not test_csv.exists():
            warnings.append("No test.csv found — final test evaluation will be skipped")
        else:
            try:
                tdf = pd.read_csv(test_csv)
                missing = set(ALL_COLUMNS) - set(tdf.columns)
                if missing:
                    errors.append(f"test.csv missing columns: {sorted(missing)}")
            except Exception as e:
                errors.append(f"Cannot parse test.csv: {e}")

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

    def setup(self, data_path: str, hyperparams: dict) -> None:
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.hyperparams = hyperparams
        dp = Path(data_path)
        batch_size = int(hyperparams.get("batch_size", 32))

        # Load training data
        X_train, y_train, self.mean, self.std = _load_csv(dp / "train.csv")

        self.train_loader = DataLoader(
            TensorDataset(X_train, y_train),
            batch_size=batch_size,
            shuffle=True,
        )

        # Load or split validation data
        val_csv = dp / "val.csv"
        if val_csv.exists():
            X_val, y_val, _, _ = _load_csv(val_csv, self.mean, self.std)
        else:
            # Split 80/20 from training data
            n = len(X_train)
            perm = torch.randperm(n)
            split = int(n * 0.8)
            train_idx, val_idx = perm[:split], perm[split:]
            X_val, y_val = X_train[val_idx], y_train[val_idx]
            X_train, y_train = X_train[train_idx], y_train[train_idx]
            self.train_loader = DataLoader(
                TensorDataset(X_train, y_train),
                batch_size=batch_size,
                shuffle=True,
            )

        self.val_loader = DataLoader(
            TensorDataset(X_val, y_val),
            batch_size=batch_size,
        )

        # Load test data if available
        test_csv = dp / "test.csv"
        if test_csv.exists():
            X_test, y_test, _, _ = _load_csv(test_csv, self.mean, self.std)
            self.test_loader = DataLoader(
                TensorDataset(X_test, y_test),
                batch_size=batch_size,
            )
        else:
            self.test_loader = None

        # Initialize model
        hidden_dim = int(hyperparams.get("hidden_dim", 64))
        dropout_rate = float(hyperparams.get("dropout_rate", 0.3))
        self.model = DiabetesNet(
            input_dim=8,
            hidden_dim=hidden_dim,
            dropout_rate=dropout_rate,
        ).to(self.device)

        # Optimizer and loss
        self.optimizer = optim.Adam(
            self.model.parameters(),
            lr=float(hyperparams.get("learning_rate", 0.001)),
        )
        # Weight positive class to handle imbalance
        pos_count = (y_train == 1).sum().float()
        neg_count = (y_train == 0).sum().float()
        pos_weight = neg_count / max(pos_count, 1.0)
        self.criterion = nn.CrossEntropyLoss(
            weight=torch.tensor([1.0, pos_weight]).to(self.device)
        )

    def get_model(self):
        return self.model

    def train(self, model) -> dict:
        model.train()
        epochs = int(self.hyperparams.get("epochs_per_round", 5))
        total_loss, total_correct, total_samples = 0.0, 0, 0

        for _ in range(epochs):
            for X_batch, y_batch in self.train_loader:
                X_batch, y_batch = X_batch.to(self.device), y_batch.to(self.device)
                self.optimizer.zero_grad()
                logits = model(X_batch)
                loss = self.criterion(logits, y_batch)
                loss.backward()
                self.optimizer.step()

                total_loss += loss.item() * X_batch.size(0)
                total_correct += (logits.argmax(1) == y_batch).sum().item()
                total_samples += X_batch.size(0)

        return {
            "train_loss": total_loss / total_samples,
            "train_accuracy": total_correct / total_samples,
            "num_samples": total_samples,
        }

    def validate(self, model) -> dict:
        return self._run_evaluation(model, self.val_loader, prefix="")

    def test(self, model) -> dict:
        if self.test_loader is None:
            return {}
        return self._run_evaluation(model, self.test_loader, prefix="test_")

    def _run_evaluation(self, model, loader, prefix: str) -> dict:
        model.eval()
        all_labels, all_preds, all_probs = [], [], []
        total_loss, total = 0.0, 0

        with torch.no_grad():
            for X_batch, y_batch in loader:
                X_batch, y_batch = X_batch.to(self.device), y_batch.to(self.device)
                logits = model(X_batch)
                loss = self.criterion(logits, y_batch)
                total_loss += loss.item() * X_batch.size(0)
                total += X_batch.size(0)

                probs = torch.softmax(logits, dim=1)[:, 1]
                preds = logits.argmax(1)

                all_labels.extend(y_batch.cpu().tolist())
                all_preds.extend(preds.cpu().tolist())
                all_probs.extend(probs.cpu().tolist())

        # Compute sklearn metrics
        acc = accuracy_score(all_labels, all_preds)
        prec = precision_score(all_labels, all_preds, zero_division=0)
        rec = recall_score(all_labels, all_preds, zero_division=0)
        f1 = f1_score(all_labels, all_preds, zero_division=0)
        try:
            auc = roc_auc_score(all_labels, all_probs)
        except ValueError:
            auc = 0.0  # single class in batch

        # Validation metrics use unprefixed names, test uses test_ prefix
        if prefix:
            return {
                f"{prefix}loss": total_loss / total,
                f"{prefix}accuracy": acc,
                f"{prefix}precision": prec,
                f"{prefix}recall": rec,
                f"{prefix}f1_score": f1,
                f"{prefix}auc_roc": auc,
            }
        return {
            "val_loss": total_loss / total,
            "accuracy": acc,
            "precision": prec,
            "recall": rec,
            "f1_score": f1,
            "auc_roc": auc,
        }

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

FLTrainer API

Metrics

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

Each metric supports:

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