neural_lam.utils.logging
========================

.. py:module:: neural_lam.utils.logging

.. autoapi-nested-parse::

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



Functions
---------

.. autoapisummary::

   neural_lam.utils.logging.init_training_logger_metrics
   neural_lam.utils.logging.log_on_rank_zero
   neural_lam.utils.logging.setup_training_logger


Module Contents
---------------

.. py:function:: init_training_logger_metrics(training_logger: Any, val_steps: list[int]) -> None

   Configure validation metric aggregation for the active training logger.

   :param training_logger: Logger instance used during training.
   :type training_logger: Any
   :param val_steps: Autoregressive rollout lengths to log as separate metrics.
   :type val_steps: list of int


.. py:function:: log_on_rank_zero(msg: str, level: str = 'info', *args: Any, **kwargs: Any) -> None

   Log a message only on rank zero using loguru logger.

   :param msg: The message to log.
   :type msg: str
   :param level: The logging level (e.g. "info", "warning", "error").
   :type level: str, default "info"
   :param \*args: Positional arguments passed to the logger.
   :type \*args: Any
   :param \*\*kwargs: Keyword arguments passed to the logger.
   :type \*\*kwargs: Any


.. py:function:: setup_training_logger(datastore: neural_lam.datastore.base.BaseDatastore, args: argparse.Namespace, run_name: str, run_dir: str) -> pytorch_lightning.loggers.Logger

   Set up the training logger (WandB or MLFlow).

   :param datastore: Datastore providing metadata for logging configuration.
   :type datastore: BaseDatastore
   :param args: Parsed training arguments controlling the logger backend.
   :type args: argparse.Namespace
   :param run_name: Name of the run.
   :type run_name: str
   :param run_dir: Directory under which all artifacts for this run are written
                   (logger ``save_dir``, checkpoints, Lightning ``default_root_dir``).
                   Typically ``runs/<run_name>``.
   :type run_dir: str

   :returns: The initialized logger object.
   :rtype: pl.loggers.Logger

   :raises ValueError: If ``args.logger`` is not ``'wandb'`` or ``'mlflow'``.

   .. rubric:: 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.