Race condition in acquiring TPU heads? #17
Description
I'm pretty sure that the two-step logic to acquire TPU slices has a race condition. It may be too low probability to care about, but I think it's there.
The reserve call gets and then releases a tpu head to get its name, then it sleeps, then it allocates using that TPU head resource. If, during that sleep, another driver process manages to request a TPU, it can wind up with the same TPU head, in which case one of the two driver processes will block on the other.
I think this can be solved by creating the second pg (the one with tpu.name) inside of reserve and not releasing the outer pg until the inner one is ready... Alternatively, in the Levanter version of this logic, we don't release the TPU-head resource which also solves this.