neural_lam.utils.logging

neural_lam.utils.logging#

Setup and configuration of training loggers (WandB / MLFlow).

Functions#

init_training_logger_metrics(→ None)

Configure validation metric aggregation for the active training logger.

log_on_rank_zero(→ None)

Log a message only on rank zero using loguru logger.

setup_training_logger(→ pytorch_lightning.loggers.Logger)

Set up the training logger (WandB or MLFlow).

Module Contents#

neural_lam.utils.logging.init_training_logger_metrics(training_logger: Any, val_steps: list[int]) → None[source]#

Configure validation metric aggregation for the active training logger.

Parameters:

  • training_logger (Any) – Logger instance used during training.

  • val_steps (list of int) – Autoregressive rollout lengths to log as separate metrics.

neural_lam.utils.logging.log_on_rank_zero(msg: str, level: str = 'info', *args: Any, **kwargs: Any) → None[source]#

Log a message only on rank zero using loguru logger.

Parameters:

  • msg (str) – The message to log.

  • level (str, default "info") – The logging level (e.g. “info”, “warning”, “error”).

  • *args (Any) – Positional arguments passed to the logger.

  • **kwargs (Any) – Keyword arguments passed to the logger.

neural_lam.utils.logging.setup_training_logger(datastore: neural_lam.datastore.base.BaseDatastore, args: argparse.Namespace, run_name: str, run_dir: str) → pytorch_lightning.loggers.Logger[source]#

Set up the training logger (WandB or MLFlow).

Parameters:

  • datastore (BaseDatastore) – Datastore providing metadata for logging configuration.

  • args (argparse.Namespace) – Parsed training arguments controlling the logger backend.

  • run_name (str) – Name of the run.

  • run_dir (str) – Directory under which all artifacts for this run are written (logger save_dir, checkpoints, Lightning default_root_dir). Typically runs/<run_name>.

Returns:

The initialized logger object.

Return type:

pl.loggers.Logger

Raises:

ValueError – If args.logger is not 'wandb' or 'mlflow'.

Notes

When --wandb_id is given, resume="allow" is set automatically: W&B resumes the run if it exists, or creates it with that ID otherwise. This allows the same job script to be safely resubmitted on HPC systems. The run name is set to None when resuming to preserve the existing name.