MLflow Flavors

If you’re a data scientist or ML engineer, undoubtedly, you’ve encountered this scenario countless times:
Working in an organization where various teams using different ML libraries and frameworks based on project requirements/preferences.

Looks flexible?
But it for sure gave you some headache some time especially when you should deal with managing and deploying models across these diverse environments.

What I exactly mean is that suppose Team A is deeply into TensorFlow for its deep learning capabilities (we all love that😅), while Team B relies heavily on Spark ML for its distributed computing ability (that is also lovable😎).

Now, there may be a separate team specifically responsible for deploying all ML models (a.k.a the deployment team).

Having been part of such a team previously, I can confidently say that you might encounter a decent amount of frustration when dealing with discrepancies in how different models are saved, logged, and retrieved across various machine learning libraries and frameworks😅!

It looks like each library has its own approach to serialization and deserialization.

In fact this deployment team might be dreaming of having some kind of “wrappers” for specific ML libraries that can provide a “unified” interface for saving, loading, and handling models, regardless of their foundational framework.

And that is what MLflow calls flavors!
So, whether it’s Tensorflow, PyTorch or Spark ML, MLflow ensures that models can be “encapsulated” within their respective “flavors”.

So, when each team saves their TensorFlow model or Spark ML model using MLflow, they are neatly packaged under the TensorFlow or Spark flavors!

In fact, I personally loved this “abstraction” notion, especially when I knew that each flavor defines a standardized 𝚙𝚛𝚎𝚍𝚒𝚌𝚝() method regardless of their underlying library!

For 𝚜𝚔𝚕𝚎𝚊𝚛𝚗 flavor for example, you can have standard methods for all types of its models.
You can save and load model with 𝚜𝚊𝚟𝚎_𝚖𝚘𝚍𝚎𝚕() and 𝚕𝚘𝚊𝚍_𝚖𝚘𝚍𝚎𝚕() and can predict with 𝚙𝚛𝚎𝚍𝚒𝚌𝚝().

A Caveat:
While that is cool and seems universal, it’s not always feasible to accommodate every unique model scenario generated by a specific library.
The complexities of certain models may pose challenges in fitting them into this standardized framework.

A Solution:
However, there’s a nice workaround.
That is MLflow offers the flexibility to craft a “custom flavor” by extending the foundational python model base class, which underlies all named flavors’ variants.
With the correct implementation, you can embed any code or model from any library within a “custom” class, all while enjoying the uniformity benefits, we mentioned with a named flavor.
So, fortunately we can tailor solutions to our specific needs.
In MLflow jargon that is 𝚙𝚢𝚝𝚑𝚘𝚗_𝚏𝚞𝚗𝚌𝚝𝚒𝚘𝚗 or 𝚖𝚕𝚏𝚕𝚘𝚠.𝚙𝚢𝚏𝚞𝚗𝚌().

If you’re wrestled with this kind of issues you might need to give a it a try 🙂.