Sync MLflow tracking in Kedro and Pytorch Lightning

[hamiddimyati]

Hi team,

I have model_training pipeline which contain training deep neural networks using Pytorch Lightning. I want to log all metrics during training and saved those as MlflowMetricHistoryDataset within Kedro catalog. How to make it possible?

I tried using MLflow tracking as logger for Lightning trainer. The metrics are logged in the same MLflow directory but in different run name (not model_training). Any helps for my issue are really appreciated. Thanks

[UPDATE]
Now, I can link the MLflow as Lightning logger with Kedro pipeline run. However, still wondering how to save the metrics and also the trained model in Kedro catalog, as MlflowMetricsHistoryDataset and MlflowArtifactDataset respectively.