mlflow_tracker.py

from hamilton import driver
from hamilton.io.materialization import to
from hamilton.plugins.h_mlflow import MLFlowTracker

dr = (
    driver.Builder()
    .with_modules(model_training_2)
    .with_adapters(
        MLFlowTracker(
            experiment_name="my_experiment",
            run_name="hamilton_dag",
        )
    )
    .build()
)