Clean printing to terminal alongside progress bar

[juneskiafc]

🚀 Feature

A way to print to terminal without breaking up the progress bar.

Motivation

A lot of people print stuff to terminal while training/validating/testing, and currently a simple call to print() will break the progress bar. A way to get around this is to set up a custom progress bar, with methods for calling tqdm.write, and passing that as a callback to the trainer. However, this feels like a lot of effort for just getting clean terminal output alongside the progress bar.

Pitch

Ability to get the ProgressBar - or the current active tqdm instance (main_progress_bar, val_progres_bar, etc) - through the Trainer or the LightningModule.
Something like pbar = trainer.pbar would feel intuitive.
Then, the user can call pbar.write(), and get clean printing like tqdm.write().

[github-actions]

Hi! thanks for your contribution!, great first issue!

[awaelchli]

hey! that's a neat idea, I really like it. How about we integrate it with the already existing LightningModule.print?
You can currently access the progress bar like so:

self.trainer.progress_bar_callback.main_progress_bar.write("hello")

[juneskiafc]

That sounds great! I actually didn't know you could do trainer.progres_bar_callback, maybe I missed it in the docs.
I think that way of accessing the tqdm.write() method is okay, maybe that should be added to the docs under the Trainer properties section for now?
I do think there is value forcing LightningModule.print to print without interfering with the progress bar, though.

[asnorkin]

I'll happily work on this :)

I propose to change current LightningModule.print() method:

class LightningModule:
    ...
    def print(self, *args, **kwargs) -> None:
        if self.trainer.is_global_zero:
            print(*args, **kwargs)

to this one:

class LightningModule:
    ...
    def print(self, *args, **kwargs) -> None:
        if self.trainer.is_global_zero:
            progress_bar = self.trainer.progress_bar_callback
            if progress_bar.is_enabled:
                progress_bar.print(*args, **kwargs)
            else:
                print(*args, **kwargs)

And add progress_bar.print(*args, **kwargs) method to the ProgressBar class.

[awaelchli]

sounds reasonable! in your snippet above, make sure to also include a None check for progress bar callback on trainer.
cheers