flip.nvflare.controllers ======================== .. py:module:: flip.nvflare.controllers .. autoapi-nested-parse:: 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 ---------- .. toctree:: :maxdepth: 1 /autoapi/flip/nvflare/controllers/broadcast_task/index /autoapi/flip/nvflare/controllers/fed_evaluation/index /autoapi/flip/nvflare/controllers/init_evaluation/index /autoapi/flip/nvflare/controllers/init_training/index /autoapi/flip/nvflare/controllers/scatter_and_gather/index /autoapi/flip/nvflare/controllers/scatter_and_gather_ldm/index Classes ------- .. autoapisummary:: flip.nvflare.controllers.BroadcastTask flip.nvflare.controllers.ModelEval flip.nvflare.controllers.InitEvaluation flip.nvflare.controllers.InitTraining flip.nvflare.controllers.ScatterAndGather flip.nvflare.controllers.ScatterAndGatherLDM Package Contents ---------------- .. py:class:: BroadcastTask(task_name: str, timeout: int = 600, participating_clients: list[str] | None = None) Bases: :py:obj:`nvflare.apis.impl.controller.Controller` Broadcast an empty task and wait for every targeted client response. .. py:attribute:: task_name .. py:attribute:: timeout :value: 600 .. py:attribute:: participating_clients :value: None .. py:method:: start_controller(fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: control_flow(abort_signal: nvflare.apis.signal.Signal, fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: stop_controller(fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: 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 .. py:class:: 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: :py:obj:`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 :class:`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. .. py:attribute:: _evaluation_task_name :value: 'evaluation' .. py:attribute:: _eval_results_dir :value: 'evaluation_results' .. py:attribute:: _eval_results :type: dict .. py:attribute:: _all_models_dxo :value: None .. py:method:: start_controller(fl_ctx: nvflare.apis.fl_context.FLContext) .. py:method:: control_flow(abort_signal: nvflare.apis.signal.Signal, fl_ctx: nvflare.apis.fl_context.FLContext) .. py:method:: _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. .. py:method:: _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. .. py:method:: _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. .. py:method:: handle_event(event_type: str, fl_ctx: nvflare.apis.fl_context.FLContext) .. py:class:: InitEvaluation(model_id: str = '', min_clients: int = FlipConstants.MIN_CLIENTS, flip: flip.FLIP = FLIP(), cleanup_timeout: int = 600) Bases: :py:obj:`nvflare.apis.impl.controller.Controller` .. py:attribute:: _model_id_fallback :value: '' .. py:attribute:: _model_id :type: str | None :value: None .. py:attribute:: _min_clients .. py:attribute:: flip .. py:attribute:: _cleanup_timeout :value: 600 .. py:method:: _resolve_model_id(fl_ctx: nvflare.apis.fl_context.FLContext) -> str .. py:method:: start_controller(fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: control_flow(abort_signal: nvflare.apis.signal.Signal, fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: stop_controller(fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: 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 .. py:method:: _set_init_evaluation_status(fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: _check_abort_signal(fl_ctx: nvflare.apis.fl_context.FLContext, abort_signal: nvflare.apis.signal.Signal) -> bool .. py:method:: _process_cleanup_result(client_task: nvflare.apis.controller_spec.ClientTask, fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: _accept_cleanup_result(client_name: str, result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext) -> bool | None .. py:class:: InitTraining(model_id: str = '', min_clients: int = FlipConstants.MIN_CLIENTS, flip: flip.FLIP = FLIP(), cleanup_timeout: int = 600) Bases: :py:obj:`nvflare.apis.impl.controller.Controller` .. py:attribute:: _model_id_fallback :value: '' .. py:attribute:: _model_id :type: str | None :value: None .. py:attribute:: _min_clients .. py:attribute:: flip .. py:attribute:: _cleanup_timeout :value: 600 .. py:method:: _resolve_model_id(fl_ctx: nvflare.apis.fl_context.FLContext) -> str .. py:method:: start_controller(fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: control_flow(abort_signal: nvflare.apis.signal.Signal, fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: stop_controller(fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: 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 .. py:method:: _set_init_training_status(fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: _check_abort_signal(fl_ctx: nvflare.apis.fl_context.FLContext, abort_signal: nvflare.apis.signal.Signal) -> bool .. py:method:: _process_cleanup_result(client_task: nvflare.apis.controller_spec.ClientTask, fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: _accept_cleanup_result(client_name: str, result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext) -> bool | None .. py:class:: ScatterAndGather(*args, model_id: str = '', **kwargs) Bases: :py:obj:`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: * :meth:`__init__` — carries the FLIP ``model_id`` and a :class:`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. * :meth:`_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. * :meth:`handle_event` — relays FLIP metrics on ``FlipEvents.SEND_RESULT``. * :meth:`_check_abort_signal` — fires ``FlipEvents.ABORTED`` so downstream components (e.g. ``PersistToS3AndCleanup``) can persist results on an aborted run. .. py:attribute:: _model_id_fallback :value: '' .. py:attribute:: _model_id :type: str | None :value: None .. py:attribute:: flip .. py:attribute:: _round_acceptances :type: dict[int, set[str]] .. py:method:: _resolve_model_id(fl_ctx: nvflare.apis.fl_context.FLContext) -> str .. py:method:: _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. .. py:method:: _accept_train_result(client_name: str, result: nvflare.apis.shareable.Shareable, fl_ctx: nvflare.apis.fl_context.FLContext, is_unknown_task: bool = False) -> bool .. py:method:: _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. .. py:method:: _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. .. py:method:: handle_event(event_type: str, fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: _check_abort_signal(fl_ctx: nvflare.apis.fl_context.FLContext, abort_signal: nvflare.apis.signal.Signal) -> bool .. py:class:: 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: :py:obj:`flip.nvflare.controllers.scatter_and_gather.ScatterAndGather` Two-phase (autoencoder → diffusion model) FedAvg controller for latent diffusion. Subclasses FLIP's :class:`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: * :meth:`__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. * :meth:`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. * :meth:`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. * :meth:`stop_controller` — cancels outstanding tasks (stock only flips the phase). .. py:attribute:: _num_rounds_ae :value: 5 .. py:attribute:: _num_rounds_dm :value: 5 .. py:attribute:: _fatal_error_delay :value: 5 .. py:attribute:: model_locator_id :value: '' .. py:attribute:: model_locator :value: None .. py:method:: start_controller(fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: control_flow(abort_signal: nvflare.apis.signal.Signal, fl_ctx: nvflare.apis.fl_context.FLContext) -> None .. py:method:: stop_controller(fl_ctx: nvflare.apis.fl_context.FLContext) -> None