[Q&A] How to save the global model at the client side after the last round?

[vcvandersluis]

Python version (python3 -V)

3.9.18

NVFlare version (python3 -m pip list | grep "nvflare")

2.3.8

NVFlare branch (if running examples, please use the branch that corresponds to the NVFlare version, git branch)

No response

Operating system

Ubuntu

Have you successfully run any of the following examples?

• hello-numpy-sag with simulator
• hello-pt with simulator
• hello-numpy-sag with POC
• hello-pt with POC

Please describe your question

I want to save the global model at the client side after the last training round.

I achieved this by adding the following code

            data_shareable.set_header("ONLY_SAVE_GLOBAL_MODEL", True)
            save_model_task = Task(
                name="train",
                data=data_shareable,
                props={},
                timeout=self._train_timeout,
                before_task_sent_cb=self._prepare_train_task_data,
                result_received_cb=self._process_train_result,
            )

            self.broadcast_and_wait(
                task=save_model_task,
                min_responses=self._min_clients,
                wait_time_after_min_received=self._wait_time_after_min_received,
                fl_ctx=fl_ctx,
                abort_signal=abort_signal,
            )

to the control_flow function of the Scatter and Gather class: (outside the while training loop)
At this location:

NVFlare/nvflare/app_common/workflows/scatter_and_gather.py

Line 286 in 2874b12

.

And in the client train function I put this code (at the start of the function):

        if shareable.get_header("ONLY_SAVE_GLOBAL_MODEL") == True:
            self.ONLY_SAVE_GLOBAL_MODEL = shareable.get_header("ONLY_SAVE_GLOBAL_MODEL")

        print("self.ONLY_SAVE_GLOBAL_MODEL", self.ONLY_SAVE_GLOBAL_MODEL)
        if self.ONLY_SAVE_GLOBAL_MODEL == True:
            self.save_global_model()
            self.system_panic(
                    f"Global model is saved stop the run",
                    fl_ctx=fl_ctx,
                )

This will result in the global model being saved after the last round, but the run will then be aborted.
(which makes sense because I call the system_panic function)

Is it possible to save the global model at the client side after the last round without raising an error?

I know I can download the global model using the admin console, but I prefer to do it automatically at the end of the training.

I already tried to change the workflow in the server/client config, but I did not succeed. (Adapting the LearnerExecutor & adding a second Scatter and Gather)

I hope someone can help me with this issue.

LOG:
self.ONLY_SAVE_GLOBAL_MODEL True
2024-03-22 09:54:50,484 - ClientRunner - ERROR - [identity=site-1, run=simulate_job, peer=simulator_server, peer_run=simulate_job, task_name=train, task_id=3fb274f6-8c8a-47f2-a71d-b83260b23df5]: Aborting current RUN due to FATAL_SYSTEM_ERROR received: Global model is saved stop the run
2024-03-22 09:54:50,486 - ClientRunner - INFO - [identity=site-1, run=simulate_job]: ABORT (RUN) command received

[SYangster]

Hi thanks for the question. You are on the right track, however we don’t need to raise an error when saving the global model.

In the controller when you call self.broadcast_and_wait() with the ONLY_SAVE_GLOBAL_MODEL shareable header, it is expecting self._min_clients number of responses. After saving the global model on the client, rather than doing system_panic() you can instead respond with a shareable via return make_reply(ReturnCode.OK).

[vcvandersluis]

Thanks for your answer, when I use return make_reply(ReturnCode.OK).
I get this error:

2024-03-26 11:41:40,804 - ServerRunner - INFO - [identity=simulator_server, run=simulate_job, wf=scatter_gather_ctl, peer=site-1, peer_run=simulate_job]: got result from client site-1 for task: name=save_global_model, id=3b39249c-a6ea-4ee0-aa3d-bda000ded498
2024-03-26 11:41:40,807 - IntimeModelSelector - ERROR - [identity=simulator_server, run=simulate_job, wf=scatter_gather_ctl, peer=site-1, peer_run=simulate_job, peer_rc=OK, task_name=save_global_model, task_id=3b39249c-a6ea-4ee0-aa3d-bda000ded498]: shareable data is not a valid DXO. Received Exception: ValueError: the shareable is not a valid DXO - expect content_type DXO but got None
2024-03-26 11:41:40,822 - IntimeModelSelector - ERROR - Traceback (most recent call last):
  File "/miniconda3/envs/fed13_nvidia_flare_2_3/lib/python3.9/site-packages/nvflare/app_common/widgets/intime_model_selector.py", line 81, in _before_accept
    dxo = from_shareable(shareable)
  File "/miniconda3/envs/fed13_nvidia_flare_2_3/lib/python3.9/site-packages/nvflare/apis/dxo.py", line 167, in from_shareable
    raise ValueError("the shareable is not a valid DXO - expect content_type DXO but got {}".format(content_type))
ValueError: the shareable is not a valid DXO - expect content_type DXO but got None

2024-03-26 11:41:40,827 - InTimeAccumulateWeightedAggregator - ERROR - [identity=simulator_server, run=simulate_job, wf=scatter_gather_ctl, peer=site-1, peer_run=simulate_job, peer_rc=OK, task_name=save_global_model, task_id=3b39249c-a6ea-4ee0-aa3d-bda000ded498]: shareable data is not a valid DXO
2024-03-26 11:41:40,836 - InTimeAccumulateWeightedAggregator - ERROR - Traceback (most recent call last):
  File "/miniconda3/envs/fed13_nvidia_flare_2_3/lib/python3.9/site-packages/nvflare/app_common/aggregators/intime_accumulate_model_aggregator.py", line 177, in accept
    dxo = from_shareable(shareable)
  File "/miniconda3/envs/fed13_nvidia_flare_2_3/lib/python3.9/site-packages/nvflare/apis/dxo.py", line 167, in from_shareable
    raise ValueError("the shareable is not a valid DXO - expect content_type DXO but got {}".format(content_type))
ValueError: the shareable is not a valid DXO - expect content_type DXO but got None

2024-03-26 11:41:40,843 - ScatterAndGather - INFO - [identity=simulator_server, run=simulate_job, wf=scatter_gather_ctl, peer=site-1, peer_run=simulate_job, peer_rc=OK, task_name=save_global_model, task_id=3b39249c-a6ea-4ee0-aa3d-bda000ded498]: Contribution from site-1 REJECTED by the aggregator at round 5.

I also tried this:

dxo2 = DXO(data_kind=DataKind.FL_MODEL, data=global_weights)
return dxo2.to_shareable()

This gave:

2024-03-26 11:53:35,419 - ServerRunner - INFO - [identity=simulator_server, run=simulate_job, wf=scatter_gather_ctl, peer=site-1, peer_run=simulate_job]: got result from client site-1 for task: name=save_global_model, id=ccc22929-31e1-43ff-87a0-a0a5a94ad20f
2024-03-26 11:53:35,422 - InTimeAccumulateWeightedAggregator - ERROR - [identity=simulator_server, run=simulate_job, wf=scatter_gather_ctl, peer=site-1, peer_run=simulate_job, peer_rc=OK, task_name=save_global_model, task_id=ccc22929-31e1-43ff-87a0-a0a5a94ad20f]: cannot handle data kind FL_MODEL, expecting DataKind.WEIGHT_DIFF, DataKind.WEIGHTS, or DataKind.COLLECTION.
2024-03-26 11:53:35,432 - ScatterAndGather - INFO - [identity=simulator_server, run=simulate_job, wf=scatter_gather_ctl, peer=site-1, peer_run=simulate_job, peer_rc=OK, task_name=save_global_model, task_id=ccc22929-31e1-43ff-87a0-a0a5a94ad20f]: Contribution from site-1 REJECTED by the aggregator at round 1.

It seems that the InTimeAccumulateWeightedAggregator class expects DataKind.WEIGHT_DIFF.
How can I prevent this error?

(I mentioned above that I am using version 2.3.8. This is wrong I use version 2.4.0)

[SYangster]

I see you changed the task name from “train” to “save_global_model” which is good. However since this is no longer a “train” task we should remove the self._prepare_train_task_data and self._process_train_result callbacks from the save_model_task (the callbacks can be None), so we don’t unnecessarily invoke the InTimeAccumulateWeightedAggregator or IntimeModelSelector during our new "save_global_model" task.

[YuanTingHsieh]

@vcvandersluis
^ so @SYangster mean:

            data_shareable.set_header("ONLY_SAVE_GLOBAL_MODEL", True)
            save_model_task = Task(
                name="train",
                data=data_shareable,
                props={},
                timeout=self._train_timeout,
            )

            self.broadcast_and_wait(
                task=save_model_task,
                min_responses=self._min_clients,
                wait_time_after_min_received=self._wait_time_after_min_received,
                fl_ctx=fl_ctx,
                abort_signal=abort_signal,
            )

[vcvandersluis]

It works now!

Thanks! @SYangster and @YuanTingHsieh