code hangs when creating a dir

[nickKyr]

Hello,

I have written a small class derived from Callback in order to log some outputs during training. Initially, a directory is created, for the corresponding files/outputs to be stored:

class epochCallback(pl.callbacks.Callback):

    @pl.utilities.rank_zero_only
    def on_train_start(self, trainer, pl_module):
        if not restart:
            print("create output directory (rank {})".format(self.model.global_rank))
                if not isdir(aiOutputDir +"outputs_train"):
                    os.mkdir(aiOutputDir +"outputs_train")

My intention is to create the directory only once, for the main rank/GPU, that's why I use the rank_zero_only decorator.
When I try it on multiple GPUs (ddp), the code hangs. If I create an empty directory in advance, then it runs! It seems to me like a barrier is needed before creating the dir... If this is the case, how do I add this in lightning? I have also noticed that the print statement is printed for all GPUs, so I am not sure if the decorator works as it should... Should I use os.environ() instead of the decorator, and specify explicitly the global ID rank which will create the directory?

Any ideas? What is the standard practice here?

Thanks in advance,
Nikos

[awaelchli]

Can you try to remove the decorator and instead do this?

def on_train_start(self, trainer, pl_module):
        print("global_rank is", trainer.global_rank)
        if trainer.is_global_zero:
            print("create output directory (rank {})".format(self.model.global_rank))
                if not isdir(aiOutputDir +"outputs_train"):
                    os.mkdir(aiOutputDir +"outputs_train")

[nickKyr]

Thanks Adrian, your approach seems to be working. I have tried it multi-GPUs and in multi-GPU nodes...