Remove rank restrictions on loggers

[edward-io]

🚀 Feature

Allow loggers to log on non rank 0 processes.

Motivation

We want to log system data across all ranks (e.g. using a custom logger to log GPU memory, latency stats, etc). This is important if we want to detect stragglers or anomalous conditions in a rank.

Pitch

The trainer should not check that the process is on global rank 0 before calling the logger.

We already have rank_zero_experiment (https://github.com/PyTorchLightning/pytorch-lightning/blob/a64cc373946a755ce6c3aef57c1be607dfe29a0c/pytorch_lightning/loggers/base.py#L33) and @rank_zero_only decorating all loggers. We can simply remove the global zero checks in logger_connector.py here (https://github.com/PyTorchLightning/pytorch-lightning/blob/6b47cbe3ca8aa3fd82211bc9fa32e734753a6950/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py#L105-L109), and here (https://github.com/PyTorchLightning/pytorch-lightning/blob/6b47cbe3ca8aa3fd82211bc9fa32e734753a6950/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py#L180).

Alternatives

Additional context

Part of #7740

[ananthsub]

@awaelchli @carmocca this is similar to last year when we lifted rank_zero_only restrictions for saving checkpoints

[awaelchli]

@edward-io thanks for this suggestion. The first occurrence you pointed out can probably be without the check for global zero but I'm not sure what the second one has to do with the loggers since it prints to stdout. What would you like to do there?

[edward-io]

@awaelchli good point, it might not be useful to print it out. I don't have a strong opinion, what about @ananthsub?

[ananthsub]

@awaelchli good point, it might not be useful to print it out. I don't have a strong opinion, what about @ananthsub?

it's fine to print out only on rank 0. the first one is the one which can be broadened, as #8608 does

[tchaton]

Hey @ananthsub @awaelchli,

Could we imagine this re-design inspired from Metrics.

Motivations:

• Provide an option for users to log on all ranks if needed
• Remove logging control within the LoggerConnector as each Logger will be responsible for it now.

from pytorch_lightning.utilities import rank_zero_only as rank_zero_only_fn

class BaseLogger

    def __init__(self, rank_zero_only: bool = True):
        self.rank_zero_only = rank_zero_only

        self.log_metrics = self._rank_zero_only_wrapper(self.log_metrics, rank_zero_only_fn)
        self.experiment = self._rank_zero_only_wrapper(self.experiment, rank_zero_experiment)


    def _rank_zero_only_wrapper(self, fn: Callable, decorator: Callback) -> Callable:

        @functools.wraps(fn)
        def wrapped_func(*args: Any, **kwargs: Any) -> Optional[Any]:
            return fn(*args, **kwargs)

        return decorator(wrapped_func) if self.rank_zero_only else wrapped_func