Preemption/Retry logic

[dlwh]

Summarizing/reprising a conversation elsewhere:

In Levanter, we have our own version of ray_tpu. Overall I like your version better and we'd prefer to use it.

There are a few features that we have that aren't in this repo. I'll make separate issues for each

One we have is a messy logic for diagnosing failures as being due to preemption or retries, and then retrying the job automatically based on whether we detect failures or preemptions. Ray doesn't seem to be great at handling this itself, at least not in the fairly old version we use.

That logic is here-ish: https://github.com/stanford-crfm/levanter/blob/94afdc17f6091249e70e90cb27b4378d0553ff56/src/levanter/infra/ray_tpu.py#279

What we currently do is catch exceptions and inspect them as being either a node failure or other preemption, or an actual error. We then have counters for both types of failures.