lit-mlflow

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 uv:

uv 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()
    ]
)

Docs

uv run mkdocs build -f ./mkdocs.yml -d ./_build/

Update template

copier update --trust -A --vcs-ref=HEAD

Credits

This project was generated with