Add Rank-Zero Printing and Improve Wandb Initialization

[sadamov]

Description: This PR introduces rank_zero_print and init_wandb utility functions to enhance printing behavior in multi-process environments and streamline Weights and Biases (wandb) initialization for logging and experiment tracking.

Changes:

• Added rank_zero_print function:

  • Ensures printing only occurs from rank 0 process in multi-process setup.
  • Wraps print function for selective printing based on process rank.

• Introduced init_wandb function:

  • Initializes wandb based on provided arguments.
  • Handles new runs and resuming existing runs.
  • Generates meaningful run names with relevant information.
  • Creates WandbLogger for PyTorch Lightning integration.
  • Saves constants.py file to wandb for reference.

• Added necessary imports for time, PyTorch Lightning, wandb, and rank-zero utilities.

Benefits:

• Prevents duplicate or conflicting output in multi-process environments.
• Simplifies wandb setup and configuration for experiment tracking and logging.
• Improves organization and identification of experiments with informative run names.
• Enhances reproducibility by saving constants.py file to wandb.

[sadamov]

Okay I requested another review, and from my side this PR should be ready for merge.

[sadamov]

@joeloskarsson the issue was that wandb is only initialized once the trainer.fit() started. So a nice way to solve this issue is: removing the wandb.init() as you suggested and then saving the configs in the ar_model:

    def on_run_end(self):
        if self.trainer.is_global_zero:
            wandb.save("neural_lam/data_config.yaml")

[sadamov]

@leifdenby can I add this PR to the roadmap for v.0.2.0? and then add it as a feature to the changelog as well?