xla

06ba6e92 - Support async checkpointing through CheckpointManager (#5697)

Support async checkpointing through CheckpointManager (#5697) * Support async checkpointing through CheckpointManager * Allow threads to exit when CheckpointManager is freed * Use rank from tracked process group * Add TODO