Add save and export to training servicer

[thodkatz]

It builds upon #227 .

Adding the options to save and export a model for the training service.

• Save: pytorch state dict, including the optimizer
• Export: save it as a bioimageio package (.zip)

When we save, the training is paused, and then we resume. For the export, we just pause.

TODO:

• The preprocessing and the postprocessing of the pytorch-3dunet should match the ones provided by bioimageio (opened an issue pre/post-processing compatibility of pytorch-3dunet and bioimageio-core #231)
• Investigate kwargs of the model, so bioimageio can instantiate it

[codecov]

Codecov Report

Attention: Patch coverage is 28.98551% with 98 lines in your changes missing coverage. Please review.

Project coverage is 62.57%. Comparing base (90b4ce0) to head (6e4ec16).
Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
tiktorch/trainer.py	21.73%	54 Missing ⚠️
tiktorch/proto/training_pb2.py	5.88%	16 Missing ⚠️
tiktorch/server/session/backend/supervisor.py	23.07%	10 Missing ⚠️
tiktorch/server/session/backend/base.py	33.33%	6 Missing ⚠️
tiktorch/server/session/backend/commands.py	53.84%	6 Missing ⚠️
tiktorch/proto/training_pb2_grpc.py	50.00%	4 Missing ⚠️
tiktorch/server/session/process.py	66.66%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #228      +/-   ##
==========================================
- Coverage   63.63%   62.57%   -1.07%     
==========================================
  Files          47       47              
  Lines        2742     2840      +98     
==========================================
+ Hits         1745     1777      +32     
- Misses        997     1063      +66     

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

[k-dominik]

is there a reason for the +1? Maybe it deserves a comment. Pytorch themselves seem to save just the epoch: https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-a-general-checkpoint-for-inference-and-or-resuming-training

[k-dominik]

The wording to me here is a bit unclear. As far as I remember pytorch-3dunet will keep "latest" and "best" checkpoints. This sounds a bit as if a history were retained (or is this actually the case)?