Refactor Checkpointer

[fegin]

Stack from ghstack (oldest at bottom):

• Integrate TorchFT #834
• -> Refactor Checkpointer #871

Several bugs fixes, refactors, and feature improvement for the next PR (integration with TorchFT)

1. Code refactor for better readability
2. Remove the time based checkpoint condition, this is not used and can cause deadlocks when integrating with TorchFT. This will also make code simpler.
3. Fixes a async_with_pinned_memory bug.
4. The original keep_last_k implementation can cause exceptions in certain case and is also slow. Fixes the bugs and use a thread to purge checkpoints.

[fduwjj]

Can this test memory leak as well?

[fegin]

No, unfortunately. Memory leakage requires some more thorough test. I'm not sure if there is an easy way to test with unittest.

[fduwjj]

For my learning purpose, why do we only do the purge for dist.get_rank() == 0?

[fegin]

We assume that all ranks can access the same files. That's the assumption of DCP. If we let all ranks to purge, then the file systems will complain.

[tianyu-l]

I'm not entitled to review every detail of this PR -- so will leave it to others.
But it looks quite good to me.

[tianyu-l]

Suggested change

	# We probably want to create a Trainer objecta.
	# We probably want to create a Trainer object.

[tianyu-l]

Hmm I still didn't get why we need a Trainer (yet).

[fegin]

@tianyu-l , a Trainer class is not necessary but cleaner.

class Trainer(Stateful):
      def __init__(self, job_config)-> None:
              move_all_init_code in train.py here
              
      def train(self) -> None:
             training_loop
             self.checkpoint.save(state={"trainer": self})
      
      def state_dict(self) -> Dict[str, Any]:
            return {"step": self.step, "log_steps": self.log_steps, ....}
            
       def load_state_dict(self, sd) -> None:
            self.step = sd["step"]
            ...

While we can simplify the original train() by splitting it to two separate functions and keeps TrainerState, this one single trainer class is more natural, as we keep the states and methods in one class/object.

[tianyu-l]

I see, makes sense to me!

[fduwjj]

Nit: IIUC, force is only for unit test right? Might also mention this in the comment?

[fegin]

It's not, it is also used if the training is finished.

[fduwjj]

LGTM