[train] RayTrainReportCallback should only save a checkpoint on rank 0 for xgboost/lightgbm

[justinvyu]

Why are these changes needed?

This PR adds a condition to only save and report a checkpoint on the rank 0 worker for xgboost and lightgbm. This prevents unnecessary checkpoints being created, since all data parallel workers have the same model states. Note: this also accounts for usage in Tune, where ray.train.get_context().get_world_rank() returns None.

This also includes a drive-by fix for checkpoint_at_end in the xgboost callback. Now, we no longer do a separate checkpoint_at_end if the checkpoint frequency happens to line up with the last iteration. For example, if saving every 5 iterations: [iter] 0 1 2 3 4 (checkpoint) 5 6 7 8 9 (checkpoint) (checkpoint), we no longer have this "duplicate" checkpoint at the end after this fix.

Related issue number

Reported on slack

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[justinvyu]

Is this test unnecessarily complicated? Any better ways to do this?

[hongpeng-guo]

Just went through some MagicMock docs. I think this unit-testing function looks good too me :)
For simplicity, probably we can just test the class variable RayTrainReportCallback._last_checkpoint_iteration after the function callback.after_iteration() is called. In this way, we can reduce the mock of ray.train.report. Then the code from line 154 ~ 169 could be simplified as:

with mock.patch("ray.train.get_context") as mock_get_context:
        mock_context = mock.MagicMock()
        mock_context.get_world_rank.return_value = rank
        mock_get_context.return_value = mock_context

        booster.num_boosted_rounds.return_value = 2
        callback.after_iteration(booster, epoch=1, evals_log={})
        assertEqual(callback._last_checkpoint_iteration, 1) # same as epoch number

However, this test assumes ray.train.report and RayTrainReportCallback._get_checkpoint() always works. It's also a little indirect. Correct me if I am wrong :)

[justinvyu]

You're right, the test was a little too indirect and I was able to find a similar way to simplify it. PTAL

[hongpeng-guo]

Nice simplification! Learnt a lot from the code 👍

[hongpeng-guo]

Just went through some MagicMock docs. I think this unit-testing function looks good too me :)
For simplicity, probably we can just test the class variable RayTrainReportCallback._last_checkpoint_iteration after the function callback.after_iteration() is called. In this way, we can reduce the mock of ray.train.report. Then the code from line 154 ~ 169 could be simplified as:

with mock.patch("ray.train.get_context") as mock_get_context:
        mock_context = mock.MagicMock()
        mock_context.get_world_rank.return_value = rank
        mock_get_context.return_value = mock_context

        booster.num_boosted_rounds.return_value = 2
        callback.after_iteration(booster, epoch=1, evals_log={})
        assertEqual(callback._last_checkpoint_iteration, 1) # same as epoch number

However, this test assumes ray.train.report and RayTrainReportCallback._get_checkpoint() always works. It's also a little indirect. Correct me if I am wrong :)

[matthewdeng]

Clean!

[hongpeng-guo]

LGTM!