nemo_automodel.components.loggers.mlflow_utils#

Module Contents#

Classes#

MLflowLogger

MLflow logger for experiment tracking and model management.

Functions#

build_mlflow

Build MLflow logger from configuration.

Data#

logger

API#

nemo_automodel.components.loggers.mlflow_utils.logger#

‘getLogger(…)’

class nemo_automodel.components.loggers.mlflow_utils.MLflowLogger(
experiment_name: str,
run_name: Optional[str] = None,
tracking_uri: Optional[str] = None,
tags: Optional[Dict[str, str]] = None,
artifact_location: Optional[str] = None,
**kwargs,
)#

MLflow logger for experiment tracking and model management.

Initialization

Initialize MLflow logger.

Parameters:

  • experiment_name – Name of the MLflow experiment

  • run_name – Name of the current run (optional)

  • tracking_uri – MLflow tracking server URI (optional)

  • tags – Dictionary of tags to add to the run

  • artifact_location – Location to store artifacts (optional)

  • **kwargs – Additional arguments passed to mlflow.start_run()

log_params(params: Dict[str, Any]) → None#

Log parameters to MLflow.

Parameters:

params – Dictionary of parameters to log

log_metrics(
metrics: Dict[str, float],
step: Optional[int] = None,
) → None#

Log metrics to MLflow.

Parameters:

  • metrics – Dictionary of metrics to log

  • step – Step number for the metrics (optional)

log_artifacts(
local_dir: str,
artifact_path: Optional[str] = None,
) → None#

Log artifacts to MLflow.

Parameters:

  • local_dir – Local directory containing artifacts

  • artifact_path – Path within the run’s artifact directory (optional)

log_artifact(
local_path: str,
artifact_path: Optional[str] = None,
) → None#

Log a single artifact to MLflow.

Parameters:

  • local_path – Local path to the artifact

  • artifact_path – Path within the run’s artifact directory (optional)

log_model(
model: torch.nn.Module,
artifact_path: str = 'model',
registered_model_name: Optional[str] = None,
**kwargs,
) → None#

Log a PyTorch model to MLflow.

Parameters:

  • model – PyTorch model to log

  • artifact_path – Path within the run’s artifact directory

  • registered_model_name – Name for model registry (optional)

  • **kwargs – Additional arguments for mlflow.pytorch.log_model()

__enter__()#
__exit__(exc_type, exc_val, exc_tb)#
nemo_automodel.components.loggers.mlflow_utils.build_mlflow(
cfg,
) → nemo_automodel.components.loggers.mlflow_utils.MLflowLogger#

Build MLflow logger from configuration.

Parameters:

cfg – Configuration object containing MLflow settings

Returns:

MLflowLogger instance