TFX with MLflow?

Hi, I have been using Tfx pipeline for my project. I wanted to integrate it with MLflow for tracking and storing the files, artifacts etc ?

I want to store outputs of components ( such as statistics_gen, schema_gen ,evaluator etc.) to Mlflow

Is it possible, if yes how ?

Hi @Aditya_Soni ,

Yes, it’s achievable. I conducted some research (though I’m not an expert and haven’t directly worked on this) and discovered valuable resources that can assist you in beginning this integration.

One resource is the MLflow TFX integration plugin: MLflow Plugins — MLflow 2.7.1 documentation. This plugin provides a number of components that can be used to track and store TFX pipeline artifacts in MLflow.

Another resource is the MLflow Python API: Python API — MLflow 2.7.1 documentation. This API allows you to track and store any type of artifact in MLflow, including TFX pipeline artifacts.

MLflow TFX integration plugin documentation: mlflow.tensorflow — MLflow 2.7.1 documentation

I hope this helps you!

Thanks.