mlflow_materializers.py

from hamilton import driver
from hamilton.io.materialization import to

model_saver = to.mlflow(
    id="trained_model_mlflow",  # name given to the saver
    dependencies=["trained_model"],  # node returning the model
    register_as="my_predictor",  # name of the model in the MLFlow registry
)

dr = (
    driver.Builder()
    .with_modules(model_training)
    .with_materializers(model_saver)
    .build()
)
dr.execute(["trained_model_mlflow"])