Skip to content

480 ODELIA call finalize_training also in swarm training - #502

Draft
oleschwen wants to merge 28 commits into
mainfrom
480-swarm-finalize_training-after-the-while-flareis_running-loop-never-runs-launcher-sigterms-the-subprocess
Draft

480 ODELIA call finalize_training also in swarm training#502
oleschwen wants to merge 28 commits into
mainfrom
480-swarm-finalize_training-after-the-while-flareis_running-loop-never-runs-launcher-sigterms-the-subprocess

Conversation

@oleschwen

@oleschwen oleschwen commented Jul 29, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Implemented call to finalize_training via callback so that the code is actually executed after training also in swarm training mode (which was, so far, implemented and intended to run, but was not executed in our swarm training setup).

Additionally, slightly extended test for output of 3DCNN local training (for consistency).

Builds on #491, should be reviewed and merged only afterwards.

oleschwen added 12 commits July 28, 2026 11:20
@oleschwen
check that output of finalize_training is present in local training run
519b336
@oleschwen
call finalize_training at end of trainer.fit also in swarm training i…
c8a1a77 
…n minimal example

implemented via callback as the code in main after the while flare.is_running() loop is not executed in swarm training in our setup
@oleschwen
for test, require only two clients
6eaa58f
@oleschwen
include 2 epochs of 3DCNN training in CI (timeout 3 minutes), refacto…
28a4b17 
…red check for expected output
@oleschwen
check for expected output files after 3DCNN local training
e194d92
@oleschwen
call finalize_training at end of trainer.fit via callback in 3DCNN lo…
e11a6eb 
…cal training

test whether expected output is printed
(test for swarm training will follow)
@oleschwen
removed unnecessary output
589b030
@oleschwen
only one client can be aggregator, thus check output of both
f160da4
@oleschwen
Merge branch 'main' into 480-swarm-finalize_training-after-the-while-…
f084335 
…flareis_running-loop-never-runs-launcher-sigterms-the-subprocess
@oleschwen
try running one epoch of 3DCNN training (two clients, synthetic datas…
aff3cac 
…et) as part of CI
@oleschwen
test that expected output from finalize_training is present
70af2aa
@oleschwen
Merge branch 'main' into 480-swarm-finalize_training-after-the-while-…
ef070d7 
…flareis_running-loop-never-runs-launcher-sigterms-the-subprocess
@oleschwen oleschwen linked an issue Jul 29, 2026 that may be closed by this pull request
Swarm: finalize_training after the while flare.is_running() loop never runs (launcher SIGTERMs the subprocess) #480
Closed
3 tasks
@oleschwen oleschwen self-assigned this Jul 29, 2026
@oleschwen
Merge branch 'main' into 480-swarm-finalize_training-after-the-while-…
305fbb1 
…flareis_running-loop-never-runs-launcher-sigterms-the-subprocess
oleschwen added a commit that referenced this pull request Jul 29, 2026
@oleschwen
Continue based on #502 (not yet merged)
8eedd68 
Merge branch '480-swarm-finalize_training-after-the-while-flareis_running-loop-never-runs-launcher-sigterms-the-subprocess' into 372-odelia-include-evaluation-on-test-data-in-localswarm-training
Ultimate-Storm pushed a commit that referenced this pull request Jul 29, 2026
@gfidlab-ci-user @claude
revert: drop the #480 main.py change — oleschwen's #502 owns it
d3633d3 
PR #502 ("WIP 480 ODELIA call finalize_training also in swarm training") is
actively being worked and takes a different, broader approach: it removes the
post-loop finalize entirely and reworks finalization in threedcnn_ptl.py, also
covering minimal_training_pytorch_cnn and the CI harness. Landing a competing
in-loop edit to the same shared main.py would duplicate that work and conflict.

This PR keeps only the deploy-test fixes, which no open PR touches.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This was referenced Jul 29, 2026
Closed
Merged
@oleschwen
Merge branch 'main' into 480-swarm-finalize_training-after-the-while-…
937d021 
…flareis_running-loop-never-runs-launcher-sigterms-the-subprocess
@oleschwen oleschwen changed the title WIP 480 ODELIA call finalize_training also in swarm training 480 ODELIA call finalize_training also in swarm training Jul 30, 2026
@oleschwen
oleschwen marked this pull request as ready for review July 30, 2026 10:58
@oleschwen
oleschwen requested a review from Ultimate-Storm July 30, 2026 10:58
Ultimate-Storm
Ultimate-Storm approved these changes Jul 31, 2026
View reviewed changes

@Ultimate-Storm Ultimate-Storm left a comment •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approved

Two non-blocking nits

  • Naming/coupling: GT_PredProb_Output_Callback now also performs checkpoint finalization and carries logger/model/checkpointing/path_run_dir/env_vars for that purpose. A future reader looking for "where does finalize happen" is unlikely to look in a class named for CSV export. A separate small FinalizeCheckpointsCallback would keep the responsibility obvious — happy for that to be a follow-up, not a blocker.
  • Per-round cost: finalize_training now runs every round rather than once. It's idempotent and mostly copies already-written checkpoints, but it does a shutil.copy of last.ckpt each round. Fine at current model sizes; worth remembering if checkpoints grow.

Context: I'd briefly opened a competing in-loop fix in #505 and reverted it once I saw this PR — #505 is now deploy-test-only, so there's no overlap.

@oleschwen

oleschwen commented Jul 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Good points, let me think about this some more before merging.
The name is indeed misleading, I was thinking ahead to #372 where finalize_training should also output PredProb (for the test data), but it does more, that's easy to fix or split in separate callbacks.
Calling finalize_training one per round is something I specifically wanted to avoid (also in the interest of not running inference on the test data multiple times). Apparently the callback is not the right mechanism for this purpose, and a test running only one round does not capture that it is called multiple times. See #480 (comment).

@oleschwen
oleschwen marked this pull request as draft July 31, 2026 11:50
oleschwen added 10 commits August 13, 2026 11:13
@oleschwen
Calling finalize_training via callback does not work as intended: it …
ab150c8 
…is called in every round, not just at the end
@oleschwen
test that finalize_training is called (currently fails)
0eb48b3
@oleschwen
setting shutdown_timeout for minimal_example to allow finalize_traini…
4722db1 
…ng to be executed after training has finished
@oleschwen
removed callback that does not work as intended (executed every round…
cf80283 
… rather than just at the end)
@oleschwen
calling finalize_training only at the end of training
f4dce5e
@oleschwen
run two rounds so that we can check whether finalize_training is exec…
94d09c2 
…uted only once
@oleschwen
check if finalize_training is called exactly once
6a6b39a
@oleschwen
commented failing test, unclear why output appears twice in console log
f5be086
@oleschwen
fixed comment
2c3a923
@oleschwen
removed unnecessary include
99fd90f
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Ultimate-Storm Ultimate-Storm Ultimate-Storm approved these changes

Assignees

@oleschwen oleschwen

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Swarm: finalize_training after the while flare.is_running() loop never runs (launcher SIGTERMs the subprocess)

2 participants

@oleschwen @Ultimate-Storm