flip.nvflare.controllers
FLIP Controllers module containing NVFLARE workflow controllers.
Controllers orchestrate federated learning workflows.
- Exports:
InitTraining: Initialization controller for training setup
ScatterAndGather: Main training loop controller with FedAvg aggregation
ScatterAndGatherLDM: Dual-phase training controller for LDM (autoencoder + diffusion model)
BroadcastTask: Broadcast a lifecycle task and wait for client responses
InitEvaluation: Initialization controller for evaluation setup
ModelEval: Main evaluation loop controller
Submodules
Classes
Broadcast an empty task and wait for every targeted client response. |
|
FLIP model-collection evaluation: evaluate a set of server-provided models on every client. |
|
FLIP's FedAvg controller — a thin subclass of NVFLARE's stock |
|
Two-phase (autoencoder → diffusion model) FedAvg controller for latent diffusion. |
Package Contents
- class flip.nvflare.controllers.BroadcastTask(task_name: str, timeout: int = 600, participating_clients: list[str] | None = None)
Bases:
nvflare.apis.impl.controller.ControllerBroadcast an empty task and wait for every targeted client response.
- task_name
- timeout = 600
- participating_clients = None
- start_controller(fl_ctx: nvflare.apis.fl_context.FLContext) None
- control_flow(abort_signal: nvflare.apis.signal.Signal, fl_ctx: nvflare.apis.fl_context.FLContext) None
- stop_controller(fl_ctx: nvflare.apis.fl_context.FLContext) None
- process_result_of_unknown_task(client: nvflare.apis.client.Client, task_name: str, client_task_id: str, result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext) None
- class flip.nvflare.controllers.ModelEval(task_check_period=0.5, submit_model_timeout=600, validation_timeout: int = 6000, model_locator_id='', formatter_id='', submit_model_task_name=AppConstants.TASK_SUBMIT_MODEL, evaluation_task_name=PTConstants.EvalTaskName, cleanup_models=False, participating_clients=None, wait_for_clients_timeout=300)
Bases:
nvflare.app_common.workflows.cross_site_model_eval.CrossSiteModelEvalFLIP model-collection evaluation: evaluate a set of server-provided models on every client.
Subclasses NVFLARE’s
CrossSiteModelEvalto inherit its constructor validation, unknown-task routing, and stop logic. Unlike cross-site validation, where each client’s submitted model is validated by every other client, evaluation loads a collection of models server-side and bundles them into a single evaluation task broadcast to all clients.Sharing the base also defines the
_validation_task_name/_cross_val_dirattributes the previous standalone fork referenced but never assigned.- _evaluation_task_name = 'evaluation'
- _eval_results_dir = 'evaluation_results'
- _eval_results: dict
- _all_models_dxo = None
- start_controller(fl_ctx: nvflare.apis.fl_context.FLContext)
- control_flow(abort_signal: nvflare.apis.signal.Signal, fl_ctx: nvflare.apis.fl_context.FLContext)
- _locate_server_models(fl_ctx: nvflare.apis.fl_context.FLContext) bool
Load the whole collection of models to evaluate from the model locator.
Unlike stock (one DXO located per model name), evaluation loads every model in a single
locate_modelcall and keeps the collection DXO to broadcast to clients.
- _before_send_validate_task_cb(client_task: nvflare.apis.controller_spec.ClientTask, fl_ctx: nvflare.apis.fl_context.FLContext)
Bundle the whole model collection into the task data sent to each client.
- _accept_val_result(client_name: str, result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext)
Record each client’s evaluation result, keyed flatly by client.
- handle_event(event_type: str, fl_ctx: nvflare.apis.fl_context.FLContext)
- class flip.nvflare.controllers.InitEvaluation(model_id: str = '', min_clients: int = FlipConstants.MIN_CLIENTS, flip: flip.FLIP = FLIP(), cleanup_timeout: int = 600)
Bases:
nvflare.apis.impl.controller.Controller- _model_id_fallback = ''
- _model_id: str | None = None
- _min_clients
- flip
- _cleanup_timeout = 600
- _resolve_model_id(fl_ctx: nvflare.apis.fl_context.FLContext) str
- start_controller(fl_ctx: nvflare.apis.fl_context.FLContext) None
- control_flow(abort_signal: nvflare.apis.signal.Signal, fl_ctx: nvflare.apis.fl_context.FLContext) None
- stop_controller(fl_ctx: nvflare.apis.fl_context.FLContext) None
- process_result_of_unknown_task(client: nvflare.apis.client.Client, task_name: str, client_task_id: str, result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext) None
- _set_init_evaluation_status(fl_ctx: nvflare.apis.fl_context.FLContext) None
- _check_abort_signal(fl_ctx: nvflare.apis.fl_context.FLContext, abort_signal: nvflare.apis.signal.Signal) bool
- _process_cleanup_result(client_task: nvflare.apis.controller_spec.ClientTask, fl_ctx: nvflare.apis.fl_context.FLContext) None
- _accept_cleanup_result(client_name: str, result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext) bool | None
- class flip.nvflare.controllers.InitTraining(model_id: str = '', min_clients: int = FlipConstants.MIN_CLIENTS, flip: flip.FLIP = FLIP(), cleanup_timeout: int = 600)
Bases:
nvflare.apis.impl.controller.Controller- _model_id_fallback = ''
- _model_id: str | None = None
- _min_clients
- flip
- _cleanup_timeout = 600
- _resolve_model_id(fl_ctx: nvflare.apis.fl_context.FLContext) str
- start_controller(fl_ctx: nvflare.apis.fl_context.FLContext) None
- control_flow(abort_signal: nvflare.apis.signal.Signal, fl_ctx: nvflare.apis.fl_context.FLContext) None
- stop_controller(fl_ctx: nvflare.apis.fl_context.FLContext) None
- process_result_of_unknown_task(client: nvflare.apis.client.Client, task_name: str, client_task_id: str, result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext) None
- _set_init_training_status(fl_ctx: nvflare.apis.fl_context.FLContext) None
- _check_abort_signal(fl_ctx: nvflare.apis.fl_context.FLContext, abort_signal: nvflare.apis.signal.Signal) bool
- _process_cleanup_result(client_task: nvflare.apis.controller_spec.ClientTask, fl_ctx: nvflare.apis.fl_context.FLContext) None
- _accept_cleanup_result(client_name: str, result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext) bool | None
- class flip.nvflare.controllers.ScatterAndGather(*args, model_id: str = '', **kwargs)
Bases:
nvflare.app_common.workflows.scatter_and_gather.ScatterAndGatherFLIP’s FedAvg controller — a thin subclass of NVFLARE’s stock
ScatterAndGather.FLIP previously vendored a full copy of the stock controller, which drifted out of date. This subclass inherits stock’s round loop verbatim — and with it
memory_gc_roundsallocator-aware cleanup (which bounds server RSS on large-model jobs), the per-roundaggregator.reset,allow_empty_global_weights, and any future upstream fixes. Only the four FLIP-specific hooks are overridden:__init__()— carries the FLIPmodel_idand aFLIPclient, and disables stock’s per-round component snapshot by default (snapshot_every_n_rounds=0). Snapshotting re-serialises the full global model every round; for a ~759 MiB model that re-introduces the very per-round memory churnmemory_gc_roundsexists to bound. FLIP never snapshotted, so0also preserves prior behaviour. It stays configurable for callers that want resilience._accept_train_result()— reports a client-side execution exception to the hub, converts a (possibly partial, frozen-backbone)WEIGHT_DIFFhead update into fullWEIGHTSbefore aggregation, since FLIP’s aggregator expectsWEIGHTS(FLIP#684), and relays each genuinely accepted result to the hub as aCLIENT_RESULT_RECEIVEDfact.handle_event()— relays FLIP metrics onFlipEvents.SEND_RESULT._check_abort_signal()— firesFlipEvents.ABORTEDso downstream components (e.g.PersistToS3AndCleanup) can persist results on an aborted run.
- _model_id_fallback = ''
- _model_id: str | None = None
- flip
- _round_acceptances: dict[int, set[str]]
- _resolve_model_id(fl_ctx: nvflare.apis.fl_context.FLContext) str
- _diff_to_weights(result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext) nvflare.apis.shareable.Shareable
Convert a client
WEIGHT_DIFFupdate into fullWEIGHTSfor the WEIGHTS aggregator.FedAvg here aggregates
WEIGHTSwhile clients return aWEIGHT_DIFF; rebuild the full weights by adding the diff onto the current global model. Partial-safe (FLIP#684): a frozen-backbone fine-tune sends only its trainable head, and keys absent from the diff keep their global value, so a head-only update reconstructs correctly. FedOpt (which aggregatesWEIGHT_DIFFdirectly) is left untouched. On any error the original result is returned so the base class can apply its own handling.
- _accept_train_result(client_name: str, result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext, is_unknown_task: bool = False) bool
- _client_update_size(result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext) int | None
Best-effort byte size of the client’s update;
Nonewhen it cannot be sized.
- _report_client_result(client_name: str, size_bytes: int | None, fl_ctx: nvflare.apis.fl_context.FLContext) None
Emit a CLIENT_RESULT_RECEIVED fact and refresh the per-round acceptance counts.
Called only after the base class (and its aggregator) accepted the result — the event’s contract is “per accepted result”. Best-effort telemetry: any failure here is logged and must never block result acceptance. The counts are shared as sticky fl_ctx props so ServerEventHandler can attach them to the ROUND_AGGREGATED event on ROUND_DONE. NVFLARE’s
_current_roundis 0-based; the wire contract is 1-based.
- handle_event(event_type: str, fl_ctx: nvflare.apis.fl_context.FLContext) None
- _check_abort_signal(fl_ctx: nvflare.apis.fl_context.FLContext, abort_signal: nvflare.apis.signal.Signal) bool
- class flip.nvflare.controllers.ScatterAndGatherLDM(model_id: str = '', min_clients: int = 1, num_rounds_ae: int = 5, num_rounds_dm: int = 5, start_round: int = 0, model_locator_id: str = '', wait_time_after_min_received: int = 10, aggregator_id=AppConstants.DEFAULT_AGGREGATOR_ID, persistor_id=AppConstants.DEFAULT_PERSISTOR_ID, shareable_generator_id=AppConstants.DEFAULT_SHAREABLE_GENERATOR_ID, train_task_name=AppConstants.TASK_TRAIN, train_timeout: int = 0, ignore_result_error: bool = True, fatal_error_delay: int = 5, task_check_period: float = 0.5, persist_every_n_rounds: int = 1, memory_gc_rounds: int = 1)
Bases:
flip.nvflare.controllers.scatter_and_gather.ScatterAndGatherTwo-phase (autoencoder → diffusion model) FedAvg controller for latent diffusion.
Subclasses FLIP’s
ScatterAndGather(itself a thin subclass of stock NVFLARE ScatterAndGather). Each instance runs a single phase — the job wires one controller per phase (train_task_nametrain_aethentrain_dm) and NVFLARE runs them sequentially — so the round loop itself is stock’s, inherited verbatim (with its per-roundaggregator.reset, failed-client tracking,memory_gc_roundsallocator cleanup, and any future upstream fixes), alongside all the shared FLIP hooks (_accept_train_result, the metrics relay,FlipEvents.ABORTED). Only what is genuinely LDM-specific is overridden:__init__()— two per-phase round counts (num_rounds_ae/num_rounds_dm) and a server-sidemodel_locatorfor the phase-1 → phase-2 weight handoff. Delegates the shared FedAvg wiring/validation to the base, passingnum_rounds = ae + dmso the inherited persist/state code sees a valid total before a phase starts.start_controller()— resolves themodel_locatorand applies the appconfig.jsonper-phase round-count overrides, then delegates component validation and initial-model loading to stock.control_flow()— pointsself._num_roundsat this phase’s round count (stock’s loop,NUM_ROUNDSprop/header and last-round persist all key off it); fortrain_dmfirst loads the phase-1 autoencoder weights via the persistor/model-locator; then delegates the round loop to stock.stop_controller()— cancels outstanding tasks (stock only flips the phase).
- _num_rounds_ae = 5
- _num_rounds_dm = 5
- _fatal_error_delay = 5
- model_locator_id = ''
- model_locator = None
- start_controller(fl_ctx: nvflare.apis.fl_context.FLContext) None
- control_flow(abort_signal: nvflare.apis.signal.Signal, fl_ctx: nvflare.apis.fl_context.FLContext) None
- stop_controller(fl_ctx: nvflare.apis.fl_context.FLContext) None