Define a TrainStepRetryableException and catch it in train_loop.py

[kiukchung]

🚀 Feature

Currently train_loop catches RuntimeError from train_step and retries but Exception does not retry. This implies that RuntimeErrors are considered retryable versus Exception is not. Clearly define Retryable vs NonRetryable exceptions for train_step and use that to decide whether the train_step should be rolled back and retried.

Motivation

Better, clearer API and exception handling and retry logic

Pitch

See description

Alternatives

N/A

Additional context

Link to code: https://github.com/pytorch/elastic/blob/master/torchelastic/train_loop.py#L123

[kiukchung]

no longer relevant in pet v0.2.0