How to use a callback to run a periodical process inside the training loop?
I want to modify a certain function during the training loop (let's save every 10 000 global training step). I am using multiple GPUs. Currently, I have implemented it inside the training_step.  While the update is happening I want to make sure other DDP processes works. In the following code, I make sure if the global_rank is zero do the update.
def training_step(self, batch, batch_idx) -> Dict:

      if (not batch_idx==0 and batch_idx%10000==0) and (self.trainer.global_rank==0):

        #do the update

Sometimes this process crashes because, while one GPU is updating the function, other trying to use it.
So is there a call back that I can use to execute my command, where I can make sure the process waits until this step is finish. More like on_end_of_global_training_step?
One more question: how can I get the global_step count inside the training loop?