Nnunet test

[scarere]

PR Type

[Feature | Fix | Documentation | Other ]

Short Description

This PR improves the functionality of nnUNetClient

• Added a shutdown method to nnUNetClient to ensure processes spawned by dataloaders are terminated
• Added nnunet_plans to properties returned by get_properties for nnUNetClient
  • If the plans are not provided in the server config, then the client generates them and sends them back to the server
  • This allows nnunet plans to be an optional parameter as the server can request a random client to generate them
• Added the FlServerWithInitializer server class.
  • Ensures a config is sent to the client when requesting initial parameters
    • nnUNetclients no longer require a starting checkpoint because of this
    • Sever config should be passed to the client when retrieving the initial parameters adap/flower#3770
  • Adds an initialize hook method to allow further server initialization before training
• Added an NnUNetServer subclass of FlServerWithInitializer
  • uses the initialize hook method to request nnunet plans from a random client if the plans were not provided in the config
• Added load_msd_data to fl4heath.utils.load_data to dynamically download and extract any MSD dataset
• Added an nnunet example which requires no set up
  • Environment variables are dynamically set
  • Works with any MSD dataset and will download and prepare the dataset on the fly
• fixed the startup messages in run_smoke_test

Tests Added

• Added two smoke tests with the nnunet example. One for a 2d config and one for a 3d_fullres config
  • Uses the Task04_Hippocampus dataset which is the smallest of the MSD datasets and is downloaded on the fly
  • Does not check outputs, just ensures that the code runs without errors

[scarere]

The nnunet smoke tests work locally for me, but none of the other tests do. I get tolerance errors on at least two of them but probably all of them. I tried checking out main and I still get the errors so I think this must be related to my setup somehow.

[scarere]

Not sure why nnunet smoke tests are failing. Here is the error found in the logs

2024-07-30T19:28:01.6031749Z   File "/home/runner/work/FL4Health/FL4Health/.venv/lib/python3.9/site-packages/torch/utils/cpp_extension.py", line 28, in <module>
2024-07-30T19:28:01.6032383Z     from pkg_resources import packaging  # type: ignore[attr-defined]
2024-07-30T19:28:01.6032813Z torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
2024-07-30T19:28:01.6033547Z ImportError: cannot import name 'packaging' from 'pkg_resources' (/home/runner/work/FL4Health/FL4Health/.venv/lib/python3.9/site-packages/pkg_resources/__init__.py)

[jewelltaylor]

Looks great! I will take one last pass once you have the chance to address a few of the small comments I had. But this is pretty much good to go from my POV

[jewelltaylor]

LGTM!