How to save a copy of the running script into the log folder?

[ItamarKanter]

Hi,
I want to save a copy of the running script of each experiment.
I have tried to apply the following callback:

import os
from pathlib import Path

class MyCopyingCallback(pl.Callback):

    def on_init_end(self, trainer):
        log_dir = trainer.logger.log_dir
        Path(log_dir).mkdir(parents=True, exist_ok=True)
        shutil.copy2(os.path.realpath(__file__), os.path.join(log_dir, os.path.basename(os.path.realpath(__file__))))

The problem is that I ended up with two copies of my running script under two different folders (I'm using the ddp backend)

Any ideas on how to solve this issue?

[awaelchli]

Hey!
You can use the rank id to determine in which process you are, like so:

if trainer.global_rank == 0:
    # do something only on first process

In your example, this could make sense:

import os
from pathlib import Path

class MyCopyingCallback(pl.Callback):

    def on_init_end(self, trainer):
        if trainer.global_rank != 0:
            return 
        log_dir = trainer.logger.log_dir
        Path(log_dir).mkdir(parents=True, exist_ok=True)
        shutil.copy2(os.path.realpath(__file__), os.path.join(log_dir, os.path.basename(os.path.realpath(__file__))))

[awaelchli]

I made the same callback also a while ago. Have a look here: https://github.com/awaelchli/pytorch-lightning-snippets/blob/master/checkpoint/code_snapshot.py
It makes a zip file of all files, so maybe not exactly what you want.

And there is also a rank_zero_only decorator, for an alternative way.