Skip to content

lit-mlflow: An improved Lightning mlflow logger

Welcome to lit-mlflow's documentation!

Build Documentation Docs with MkDocs Poetry linting: ruff pre-commit Checked with pyright security: bandit Semantic Versions Copier License

An improved Lightning mlflow logger. Works seamlessly with PyTorch Lightning on Databricks and offers more control compared to the mlflow.pytorch.autolog function.

Features

  • Makes mlflow logging work with lightning and Databricks

Installation

With pip: 

python -m pip install lit-mlflow

With poetry: 

poetry add lit-mlflow

How to use it

Replace mlflow.autolog() with the MlFlowAutoCallback:

from lit_mlflow import MlFlowAutoCallback
import lightning.pytorch as pl

trainer = pl.Trainer(
    callbacks=[
        MlFlowAutoCallback()
    ]
)

To support Databricks mlflow, use the DbxMLFlowLogger instead of the MlFlowLogger:

from lit_mlflow import DbxMLFlowLogger
import lightning.pytorch as pl

trainer = pl.Trainer(
    logger=[
        DbxMLFlowLogger()
    ]
)

See a more complete example in the notebooks folder.

API

Check the API reference for more details.