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

BroadcastTask

Broadcast an empty task and wait for every targeted client response.

ModelEval

FLIP model-collection evaluation: evaluate a set of server-provided models on every client.

InitEvaluation

InitTraining

ScatterAndGather

FLIP's FedAvg controller — a thin subclass of NVFLARE's stock ScatterAndGather.

ScatterAndGatherLDM

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.Controller

Broadcast 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.CrossSiteModelEval

FLIP model-collection evaluation: evaluate a set of server-provided models on every client.

Subclasses NVFLARE’s CrossSiteModelEval to 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_dir attributes 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_model call 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.ScatterAndGather

FLIP’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_rounds allocator-aware cleanup (which bounds server RSS on large-model jobs), the per-round aggregator.reset, allow_empty_global_weights, and any future upstream fixes. Only the four FLIP-specific hooks are overridden:

  • __init__() — carries the FLIP model_id and a FLIP client, 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 churn memory_gc_rounds exists to bound. FLIP never snapshotted, so 0 also 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_DIFF head update into full WEIGHTS before aggregation, since FLIP’s aggregator expects WEIGHTS (FLIP#684), and relays each genuinely accepted result to the hub as a CLIENT_RESULT_RECEIVED fact.

  • handle_event() — relays FLIP metrics on FlipEvents.SEND_RESULT.

  • _check_abort_signal() — fires FlipEvents.ABORTED so 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_DIFF update into full WEIGHTS for the WEIGHTS aggregator.

FedAvg here aggregates WEIGHTS while clients return a WEIGHT_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 aggregates WEIGHT_DIFF directly) 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; None when 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_round is 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.ScatterAndGather

Two-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_name train_ae then train_dm) and NVFLARE runs them sequentially — so the round loop itself is stock’s, inherited verbatim (with its per-round aggregator.reset, failed-client tracking, memory_gc_rounds allocator 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-side model_locator for the phase-1 → phase-2 weight handoff. Delegates the shared FedAvg wiring/validation to the base, passing num_rounds = ae + dm so the inherited persist/state code sees a valid total before a phase starts.

  • start_controller() — resolves the model_locator and applies the app config.json per-phase round-count overrides, then delegates component validation and initial-model loading to stock.

  • control_flow() — points self._num_rounds at this phase’s round count (stock’s loop, NUM_ROUNDS prop/header and last-round persist all key off it); for train_dm first 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