State checkpointer refactor

[fatemetkl]

PR Type

Fix

Short Description

Clickup Ticket(s): https://app.clickup.com/t/868d4w5rt

Unifying our PerRoundStateCheckpointer and the Snapshotting functionality that supports early stopping.
This PR adds checkpointing.state_checkpointer modules with classes that can be used to checkpoint the state of the server and clients.

Tests Added

• tests/checkpointing/test_state_checkpointer.py
• tests/utils/test_early_stopper.py

I also manually tested our nnunet_example to make sure NnUnetServerCheckpointAndStateModule works as expected. We have a smoke test (load_from_checkpoint_example) for ServerPerRoundStateCheckpointer.

[codecov]

Codecov Report

Attention: Patch coverage is 94.31818% with 15 lines in your changes missing coverage. Please review.

Project coverage is 77.63%. Comparing base (2444580) to head (24581fb).
Report is 25 commits behind head on main.

Files with missing lines	Patch %	Lines
fl4health/checkpointing/state_checkpointer.py	93.22%	12 Missing ⚠️
fl4health/checkpointing/client_module.py	90.00%	1 Missing ⚠️
fl4health/checkpointing/server_module.py	94.11%	1 Missing ⚠️
fl4health/servers/nnunet_server.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #394      +/-   ##
==========================================
+ Coverage   76.83%   77.63%   +0.79%     
==========================================
  Files         156      157       +1     
  Lines        9472     9553      +81     
==========================================
+ Hits         7278     7416     +138     
+ Misses       2194     2137      -57     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[fatemetkl]

We don't need to call _load_server_state() here because it will be called in FlServer's fit_with_per_round_checkpointing()

[emersodb]

Nearly ready to go I think a couple of small questions and suggested refactors, but looks great!

[emersodb]

Good to go!