[CI] Ensure that mlflow callback cleans up background-saving threads on trainer teardown.

[justinxzhao]

#3670 de-duplicates logged metrics to MLFlow by first checking if the metrics for a given training step have already been logged, and skipping if so.

This stateful mechanic ensures that if training completes at the same step as evaluation (default configuration), the on_eval_end callback logs metrics and the on_trainer_teardown callback skips self.save_fn().

However, the save_fn block in on_trainer_teardown also takes care of two cleanup steps:

1. self.save_thread.join() to join the background-saving thread.
2. self.save_fn((..., False)) to break the infinite background loop, _log_mlflow_loop, which is started in on_trainer_train_setup if background saving is enabled.

These were improperly skipped by the new stateful mechanic, causing pytest to hang and reach timeout.

In on_trainer_teardown, the fix is to 1) always join the save_thread and 2) always call self.save_fn((None, None, None, False)) to ensure that the background saving loop breaks.

This PR also fixes the recent integration_test_e timeouts.

Credit to @jeffkinnison who pinpointed the culprit and issue, and worked with me on this solution.

CC: @Infernaught @geoffreyangus

[jeffkinnison]

LGTM for a quick fix, just a couple of suggestions for how we could potentially iterate on this fix.

[github-actions]

Unit Test Results

  6 files  ±0    6 suites  ±0   40m 50s ⏱️ - 1m 25s
31 tests ±0  26 ✔️ ±0    5 💤 ±0  0 ❌ ±0 
82 runs  ±0  66 ✔️ ±0  16 💤 ±0  0 ❌ ±0 

Results for commit 364c39f. ± Comparison against base commit 3a0e2b2.

♻️ This comment has been updated with latest results.

[justinxzhao]

test_mlflow.py is working with the timeout removal.

[jeffkinnison]

test_mlflow.py is working with the timeout removal.

Oh, interesting. Do we know why the timeout was causing problems?