# naira-fl-sdk — instructions for AI coding agents

You are building a **federated-learning model profile** with the `naira-fl-sdk`
Python package. A model profile is a small directory (a PyTorch model, a trainer,
and a `model.yaml` manifest) that the SDK validates, simulates, and packages into a
`.zip` for upload to a Naira FL platform. Your job is done when `naira-fl-sdk
validate` passes, `naira-fl-sdk test` trains cleanly, and you have produced the
`.zip`.

Install:

    pip install naira-fl-sdk[torch]

The CLI is `naira-fl-sdk`. It has a built-in self-correction loop (`validate` and
`test`) — use it. Read every error message and fix the code until the command passes.

---

## STEPS (do them in order)

1. **Scaffold.** Run:

       naira-fl-sdk init <name>

   `<name>` must match `^[a-zA-Z0-9][a-zA-Z0-9_-]*$` (letters, digits, `_`, `-`;
   must start alphanumeric). This creates `<name>/` with starter `model.py`,
   `trainer.py`, `model.yaml`, and `requirements.txt`.

2. **Implement** the three files per the SPEC below:
   - `model.py` — your `torch.nn.Module` subclass.
   - `trainer.py` — a single `FLTrainer` subclass implementing the required methods.
   - `model.yaml` — the manifest: model class, metrics, data spec, hyperparameters,
     and any extra pip deps under `requirements:`.

3. **Validate.** Run:

       naira-fl-sdk validate ./<name>

   If you see `FAIL:` or `ERROR:` lines, read them, fix the files, and re-run until
   you get `PASS: model profile is valid`. `WARN:` lines are non-blocking.

4. **Test (local FL simulation, no NVFlare needed).** Run:

       naira-fl-sdk test ./<name> --rounds 2 --data <data-dir>

   This runs `validate_data()`, then `setup()`, then N rounds of `train()`/
   `validate()`, then `test()`. Fix any runtime error and re-run until it trains and
   prints metrics cleanly.

   IMPORTANT: the tooling does NOT check that your returned metric keys match the
   manifest, nor that `num_samples` is present. A mismatch passes locally but fails
   on the platform. After it runs, **manually confirm** that the keys in each printed
   dict exactly equal the `name`s declared under the matching `metrics.<phase>` in
   `model.yaml`.

5. **Package.** Run:

       naira-fl-sdk package ./<name>

   This writes `<name>.zip`.

6. **Done.** Upload `<name>.zip` to your Naira platform admin. (The SDK's job ends
   here; pairing the profile with a strategy and running real FL jobs happens on the
   platform.)

---

## SPEC

### Imports

    from naira_fl_sdk import FLTrainer, ValidationResult

Both are public. Don't import from submodules.

### FLTrainer interface (implement in `trainer.py`)

Define exactly **one** subclass of `FLTrainer`. The simulator finds it by scanning
for an `FLTrainer` subclass, so the class name does not matter. Methods:

- `setup(self, data_path: str, hyperparams: dict) -> None` — **required.** Called
  once at client start. Load data from `data_path`, build the model, optimizer,
  loss, and dataloaders. `hyperparams` comes from your manifest defaults (plus
  platform overrides). Cast values to their declared types — overrides may arrive as
  strings (e.g. `int(hyperparams["batch_size"])`).
- `get_model(self)` -> your `nn.Module` — **required.** Return the model instance.
- `train(self, model) -> dict` — **required.** Train one round. The `model` arrives
  with the latest global weights already loaded — do NOT re-initialize them. Return a
  dict whose keys match `metrics.train`. **Always include `"num_samples": <int>`**
  (the number of local training samples) — weighted aggregation across clients
  depends on it.
- `validate(self, model) -> dict` — **required.** Evaluate on local validation data.
  Return a dict whose keys match `metrics.validate`. May return `{}` if there is no
  validation data.
- `test(self, model) -> dict` — **optional** (default returns `{}`). Final evaluation
  after all rounds. Return a dict whose keys match `metrics.test`. Handle the
  optional test file being absent gracefully.
- `validate_data(self, data_path: str, hyperparams: dict) -> ValidationResult` —
  **required.** Runs BEFORE any training and **gates** it. Check required files
  exist, schema/columns/shapes are right, etc. Return
  `ValidationResult(valid=True)` when ready, or
  `ValidationResult(valid=False, errors=[...], warnings=[...])` on problems. Treat
  missing *optional* files as warnings, not errors.

`model.py` is imported by `trainer.py` with a **flat import** (`from model import
YourModel`) — the SDK puts the profile directory on `sys.path`. Don't use a package-
relative import.

### model.yaml schema

Required: `name`, `description`, `model` (with `file` and `class`), `trainer`
(with `file`), `metrics`.
Optional: `framework`, `requirements`, `data`, `hyperparameters`.

- `name` — must match the name regex above.
- `framework` — `pytorch` (default) or `tensorflow`.
- `model.file` — Python file (default `model.py`). `model.class` — the model class
  name. NOTE the YAML key is literally **`class`**, not `class_name`.
- `trainer.file` — Python file (default `trainer.py`).
- `requirements` — list of extra pip requirement strings (e.g. `scikit-learn>=1.3`).
- `data` — `description`, `required_files: [{name, description}]`,
  `optional_files: [{name, description}]`. Describe what client data must look like.
- `hyperparameters` — map of `name -> {default, type, description}`.
- `metrics` — `train`, `validate`, `test` lists of metric defs. You must define at
  least one `train` or `validate` metric. Each metric def:
  - `name` (required) — the dict key your trainer returns for it.
  - `display_name` — human label (optional).
  - `type` — `float` (default) or `int`.
  - `format` — `number` (default), `percentage`, or `duration`.
  - `aggregate` — how to combine across clients: `weighted_avg` (default), `sum`,
    `mean`, `harmonic_mean`, `max`, or `min`. (Use `sum` for `num_samples`.)
  - `group` — optional UI grouping label.
  - `primary` — `true` to surface it in overview dashboards.

---

## RULES (common mistakes — do not violate)

- Returned metric dict keys must EXACTLY equal the `metrics.<phase>` `name`s in
  `model.yaml`. Not checked by the CLI — verify by hand.
- `train()` must include `"num_samples": <int>`. Not checked by the CLI.
- Do not re-initialize model weights in `train`/`validate`/`test` — they arrive
  pre-loaded with the global weights.
- Cast hyperparameters to their declared types inside the trainer.
- `validate_data()` gates training; return `valid=False` with clear `errors` on real
  problems, and use `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
  to stderr on EVERY command. It is benign — ignore it. The real failure signals are
  `FAIL:`/`ERROR:` lines, a non-zero exit code, or a runtime traceback.

---

## COMPLETE WORKED EXAMPLE (verified — passes `validate` and `test`)

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

### model.yaml

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

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

```python
"""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,
        }
```
