neural_lam.custom_loggers

Custom logging utilities (e.g., MLFlow wrappers) used in Neural-LAM.

Classes

CustomMLFlowLogger

Custom MLFlow logger that adds the log_image() functionality not

Module Contents

class neural_lam.custom_loggers.CustomMLFlowLogger(experiment_name: str, tracking_uri: str, run_name: str, save_dir: str)

Bases: pytorch_lightning.loggers.MLFlowLogger

Custom MLFlow logger that adds the log_image() functionality not present in the default implementation from pytorch-lightning as of version 2.0.3 at least.

Initialize the logger, ensure save_dir exists, and start the MLflow run.

Parameters:

• experiment_name (str) – Target MLflow experiment.

• tracking_uri (str) – MLflow tracking server URI.

• run_name (str) – Human-readable run name stored as mlflow.runName.

• save_dir (str) – Directory where log_image writes temporary figure files. Created eagerly with exist_ok=True.

Notes

Starts the MLflow run with log_system_metrics=True and also records run_id as an MLflow param.

log_image(key: str, images: list[matplotlib.pyplot.Figure], step: int | None = None) → None

Log one or more matplotlib figures as images to MLFlow.

When images contains more than one figure, each is logged under a key suffixed with its position in the list (key_0, key_1, …). A single-figure list is logged under the bare key.

Parameters:

• key (str) – Key to log the image under. If step is given, _{step} is appended before any per-figure index suffix.

• images (list of matplotlib.figure.Figure) – Figures to log.

• step (int or None, optional) – Step to associate with the log entry. None logs without a step suffix.

Raises:

botocore.exceptions.NoCredentialsError – If AWS credentials for the MLflow artifact store are missing; re-raised for the caller to handle.

property save_dir: str

Returns the directory where the MLFlow artifacts are saved. Used to define the path to save output when using the logger.

Returns:

Path to the directory where the artifacts are saved.

Return type:

str