[no_early_kickoff] [Train] Recommend Trainer.restore on errors raised by trainer.fit()

[justinvyu]

Why are these changes needed?

Currently, TuneErrors that recommend Tuner.restore get propagated to users even when they're using Trainer.fit. This PR fixes this by wrapping raised errors with a TrainingFailedError that includes a message on how to restore the run, and also how to configure a new run to retry on training failures. This PR also separates AIR "error propagation" tests into a new file.

Notes on the TODOs

I added a few TODOs in the code regarding "driver error propagation" in Tune. Currently, errors that happen in the Tune driver (ex: within user defined callback hooks) get surfaced to the user in different ways.

For example:

• An error within on_trial_result will surface as a TuneError(TuneError(OriginalError))
• An error within on_checkpoint will not surface at all. It'll just log a warning.
• An error within on_step_begin will surface as OriginalError.

This needs to be fixed in 2.5 (all hooks should handle errors in the same way), tracker issue here: TODO.

For the purposes of 2.4, wrapping errors raised by trainer.fit() TrainingFailedError error (which already existed) is good enough to fix the issue of "Tune concepts showing up when you're not even using Tune."

Related issue number

Closes #33566

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 :(

[woshiyyya]

Can you explain the difference between result.error and the errors captured in tuner.fit()? Should we have a unified way to handle them?

[justinvyu]

result.error: Error that happens in the trainable (ex: runtime error in the training loop per worker -> the trainer actor will raise -> Tune will record this error associated with the trial -> can be accessed by result.error)
tuner.fit(): Error that happens in the Tune driver itself (aka in the execution loop happening on the driver node).

These should be handled separately because driver error should crash the entire experiment execution, whereas individual trial errors should not crash other ongoing trials, unless otherwise configured.

In the trainer.fit case, there's only 1 trial so the full experiment will crash on trial failure.

[woshiyyya]

lgtm, just a few minor comments.

[Yard1]

Looks good to me, thanks!