Hi, you guys created a really cool framework and I am hoping to refactor our pipelines with it! I was wondering if there is a way to add a custom MLflow flavor to the

to.mlflow()

materializer since our models are not in their named flavors?

Hi Jernej! Happy that you're liking Hamilton so far I worked on building the MLFlow plugin and I'll gladly help you integrate custom flavors! I don't have much experience with custom MLFlow flavors, so having a code example of your custom flavor and how you use it to save models (without Hamilton) would be useful. Looking at the MLFlow plugin code, I have an idea of what needs to be changed, but this would help me ground development and testing. Otherwise, I'm available today to discuss your requirements (Slack or call)!

Hi Thierry, great! I am also new to custom flavors; so far I was following the MLflow example on their website and taking their python file flavor.py as template. Effectively, I had to create a custom module

custom_flavor.py

that implemented

save_model,log_model,load_model,_CustomModelWrapper

. Then I just import it and use it in the MLflow context manager:

import mlflow
import custom_flavor

with mlflow.start_run():
        custom_flavor.log_model(
            custom_model=test_model,
            artifact_path=ARTIFACT_PATH,
        )

Looking at your mlflow_extension.py, seems to me that this part of the code trips me up: https://github.com/DAGWorks-Inc/hamilton/blob/main/hamilton/plugins/mlflow_extensions.py#L50C8-L54C75

I implemented a solution that should work. # Current features Currently, the MLFlow plugin would allow for 2 main patterns.

to.mlflow(...)  # infer the flavor

to.mlflow(..., flavor="sklearn")  # mlflow supported flavor

^Note that this should support the

pyfunc

flavor, which is a convenient way to create custom estimators (MLFlow reference) # Solution The most straight solution is to allow

flavor

to receive the Python module defining the custom flavor.

import custom_flavor
to.mlflow(..., flavor=custom_flavor)

# consequently, this would also work
import mlflow
to.mlflow(..., flavor=mlflow.sklearn)

I will produce a release candidate of Hamilton that you'll be able to install and try to make sure everything is smooth before we add it to the official release 😁 Will ping you when it's available

Awesome, that's exactly what I need. Happy to test it and get back to you 🙂

It should be available through

pip install --upgrade sf-hamilton==1.71.0rc1

Hi Thierry, sorry it took me so long: https://github.com/DAGWorks-Inc/hamilton/pull/1068

@Jernej Frank I think @Thierry Jean’s tests exercised loading

Awesome, thanks for the heads up!

Yep I’ll get to this, this afternoon my time