Custom training classes · Docs

Overview

The SDK ships sensible defaults for every moving piece, so most users only supply a model. This guide covers the three classes you can override and how source extraction works.

Source Extraction — One Mechanism, Three Use Cases

For each class you pass (model=, executor=, persistor=), the SDK:

Calls inspect.getsource(cls) to grab the class body.
If that fails (happens in some notebooks where torch patches inspect), falls back to scanning IPython cell history.
Also grabs every import statement from the same cell so those symbols are available on the worker.
Writes the result to scripts/<filename>.py in your bucket.

warn

The key rule: keep imports in the same cell as the class.

This fails:

Python

# Cell 1
import torch.nn as nn

# Cell 2
class MyModel(nn.Module):   # nn will be missing on the worker
    ...

This works:

Python

# Cell 1
import torch.nn as nn

class MyModel(nn.Module):
    ...

Model — the Only Required Class

Python

import torch.nn as nn
from torchvision import models

class MyResNet(nn.Module):
    def __init__(self, num_classes: int):
        super().__init__()
        self.backbone = models.resnet18(weights=None)
        self.backbone.fc = nn.Linear(self.backbone.fc.in_features, num_classes)

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

job = sdk.rt_submit(model=MyResNet, name="my-job", shards_dir="./shards")

Constraints

Constructor takes num_classes: int.
forward(self, x) accepts (B, 3, img_size, img_size) and returns logits (B, num_classes).

Extra constructor kwargs (e.g. pretrained=True, dropout=0.1) go through ModelConfig.model_args:

Python

from resontech import ModelConfig

sdk.rt_submit(
    model=MyResNet,
    name="my-job",
    shards_dir="./shards",
    model_config=ModelConfig(model_args={"pretrained": True, "dropout": 0.1}),
)

note

num_classes is injected separately from TrainingConfig.num_classes.

Custom Executor (optional)

By default the SDK ships a generic executor that reads TrainingConfig args and calls rt_train_model from the stitched adapter. Supply your own only if you need behaviour the default can't express (custom data loaders, non-standard metrics, multi-task heads, …).

Python

# Subclass the FL runtime's Executor base class:
from nvflare.apis.executor import Executor
from nvflare.apis.fl_context import FLContext
from nvflare.apis.shareable import Shareable

class MyExecutor(Executor):
    def __init__(self, local_epochs: int, batch_size: int,
                 learning_rate: float, **kwargs):
        super().__init__()
        self.local_epochs = local_epochs
        self.batch_size = batch_size
        self.learning_rate = learning_rate

    def execute(self, task_name, shareable: Shareable,
                fl_ctx: FLContext, abort_signal):
        ...

job = sdk.rt_submit(
    model=MyResNet,
    name="custom-executor-job",
    shards_dir="./shards",
    executor=MyExecutor,
)

The SDK:

Extracts MyExecutor's source.
Uses its class name in config_fed_client.json → executor.path.
Writes it to scripts/custom_client_executor.py in your bucket.

You get every TrainingConfig field (fixed + extra) as kwargs.

Custom Persistor (optional)

The default custom_persistor.py wraps your model class with the standard ModelPersistor interface: loads initial weights, saves aggregated ones per round, emits the final checkpoint to model_out/. Override only if you need bespoke serialisation (quantised weights, partial-model updates, custom filename schemes).

Python

from nvflare.app_common.abstract.model_persistor import ModelPersistor
from nvflare.app_common.abstract.model import ModelLearnable

class MyPersistor(ModelPersistor):
    def load_model(self, fl_ctx) -> ModelLearnable: ...
    def save_model(self, learnable: ModelLearnable, fl_ctx): ...

job = sdk.rt_submit(
    model=MyResNet,
    name="custom-persistor-job",
    shards_dir="./shards",
    persistor=MyPersistor,
)

Preview Everything Before Submitting

Python

from resontech.rt.builder import RTJobBuilder

builder = RTJobBuilder(
    training_config=training,
    federation_config=federation,
    model_config=model_config,
)

scripts = builder.render_scripts(MyResNet, executor=MyExecutor, persistor=MyPersistor)
print(scripts["model_def.py"])
print("---")
print(scripts["custom_client_executor.py"])
print("---")
print(scripts["custom_persistor.py"])

Troubleshooting Source Extraction

"Could not extract source of 'X'"

You're in an environment where neither inspect.getsource nor IPython history gives usable source. Workaround: save the class to a .py file, import it normally, and pass its full dotted path via ModelConfig(model_class="my_module.MyResNet"). The SDK will skip source extraction and trust the import path.

"NameError: name 'nn' is not defined" on the worker

Imports live in a different cell from the class. Consolidate them into the same cell and resubmit. Applies to model=, executor=, persistor=.

Custom executor is silently ignored

Make sure you pass the class, not an instance:

sdk.rt_submit(..., executor=MyExecutor)       # ✅
sdk.rt_submit(..., executor=MyExecutor())     # ❌

More failure modes in the troubleshooting reference.

PreviousTraining hyperparameters Next How FL works on ResonTech