ZenML: Build Portable, Production-Ready ML and LLM Pipelines
Hey friends, it's me again, Ruby Abdullah. This time I want to talk to you about one of the tools that I honestly think you absolutely must master if you are serious about working in machine learning or LLMs professionally, not just messing around in a notebook. That tool is called ZenML.
Let me be real with you. Back when I first started learning ML, I would stack the entire process into a single notebook file. Load the data in the top cell, training in the middle, evaluation at the bottom, and when I wanted to deploy I would just copy and paste things all over the place. It felt fast, but the moment the project grew and started heading into production, everything turned into a mess. Hard to reproduce, hard to track, and when something broke I had no idea where to restart from. This is exactly where ZenML comes in to save our lives.
In this article I will explain ZenML from scratch until you can build your own end-to-end pipeline. We will cover installation, the core concepts like @step and @pipeline, the concept of stacks which is the heart of ZenML, artifacts and caching, integrations with tools like MLflow, and how to deploy. Ready? Let's get started.
Introduction
Before we dive into the code, I want you to first understand what ZenML actually is and why it matters.
ZenML is an open-source MLOps framework whose job is to make your machine learning workflows structured, reproducible, and easy to move from your laptop to the cloud without having to rewrite your code. So picture this, friends. You have model training code. It runs locally on your laptop. Then your boss says, "Hey, move this to the cloud, use a Kubernetes orchestrator, and track it with MLflow." Without ZenML, you would have to tinker with a ton of code. But with ZenML, you just change the stack configuration, and your core code stays exactly the same. Pretty cool, right?
The two main concepts in ZenML are the pipeline and the step. A step is a single unit of work, for example "load data" or "train model". A pipeline is a series of steps that are connected together. Now, ZenML separates the ML logic (which you write) from the infrastructure (where that code runs). This separation is what makes your code so portable.
Why does this matter for your career? Because in the industry, what companies look for is not just someone who can build an accurate model in a notebook. What they want is someone who can bring that model into production, maintain it, track experiments, and collaborate with a team. ZenML teaches you the right MLOps mindset from the very beginning. So in my opinion, learning ZenML is a really great investment for your future.
One more thing I love about ZenML: it does not force you to use any specific tool. Whether you want to use scikit-learn, PyTorch, TensorFlow, or even LLM frameworks like LangChain, all of them work. ZenML acts like glue that binds all of your components into one neat flow.
Installation
Alright, now let's get into the practical part. Installing ZenML is super easy, just one command.
pip install zenml
If you want the full version with a visual dashboard (and I highly recommend using this), install the server version:
pip install "zenml[server]"
I always recommend that you use a virtual environment so you don't create a mess with dependencies from other projects. Here is how I usually set it up:
python -m venv zenml-env
source zenml-env/bin/activate # on Windows: zenml-env\Scripts\activate
pip install "zenml[server]" scikit-learn
After the install finishes, we need to initialize ZenML inside our project folder. This is important because ZenML will create a hidden .zen folder that stores all the configuration and metadata.
zenml init
You will see a success message saying the ZenML repository has been created. Now, to check if everything is correct, try running this command:
zenml status
If you want to see that nice visual dashboard, run:
zenml login --local
This command will launch a local server and open a browser to the ZenML dashboard. Here you can see all the pipelines you have run, their steps, the artifacts they produced, and much more. I personally love this dashboard because it makes debugging so much easier. We can visually see which step failed, how long each step took, and what data flowed between the steps.
Oh, and one tip from me: if you are using a slightly older version of ZenML, the command might be zenml up instead of zenml login --local. So if you hit an error, first check your ZenML version with zenml version. I wrote this article based on a fairly recent version of ZenML.
Basic Usage
Now for the most exciting part, we start writing code. I will explain the two most important decorators in ZenML: @step and @pipeline.
Building Your First Step
A step is like a regular Python function, but we decorate it with @step. This decorator tells ZenML that this function is a unit of work that can be tracked, cached, and run on specific infrastructure. Take a look at this simple example:
from zenml import step
@step
def loaddata() -> dict:
"""Load simple data for the example."""
data = {
"features": [[1, 2], [3, 4], [5, 6], [7, 8]],
"labels": [0, 0, 1, 1],
}
return data
Easy right? This function just returns a dictionary. But because of the @step decorator, ZenML will automatically record its output as an artifact and store it in the artifact store. We will talk about the artifact store later.
One important thing: ZenML really values type hints. Notice I wrote -> dict in that function. These type hints are not just for show, ZenML uses this information to decide how to store and retrieve the artifact. So get into the habit of always adding type hints to your steps, friends.
Building a Pipeline
Once you have steps, we chain them into a pipeline using the @pipeline decorator. A pipeline is a function that calls steps in sequence. Here is an example:
from zenml import pipeline
@pipeline
def simplepipeline():
data = loaddata()
return data
Then to run it, just call the pipeline function:
if name == "main":
simplepipeline()
Save all the code above into one file, for example run.py, then run it with python run.py. ZenML will execute your pipeline, save all the artifacts, and record the metadata. If you open the dashboard, you will see this pipeline appear there complete with a visualization of its flow.
Understanding Data Flow Between Steps
What makes ZenML powerful is how it manages data flow between steps. When you write data = loaddata() and then data is used in another step, ZenML automatically understands that there is a dependency between those two steps. It will run them in the correct order and automatically pass the artifact from one step to the next. You don't need to manually handle how data moves, ZenML takes care of all of it behind the scenes.
This is really different from writing a regular script. In a regular script, if one part fails, you have to rerun from the beginning. In ZenML, because each step's output is saved as an artifact, you can resume from just the failed step without rerunning from zero. This is what makes experiment iteration so much faster.
Advanced Usage
Now we level up. In this section I will cover the concepts that make ZenML truly production-ready: stacks, artifacts, caching, and integrations. Then at the end I will give you a complete end-to-end pipeline example.
The Stack Concept
This is the most important concept in ZenML, so please pay close attention, friends. A stack is the collection of infrastructure and tools where your pipeline runs. The analogy is like this: your pipeline code is the recipe, and the stack is the kitchen. The same recipe can be cooked in a home kitchen or a restaurant kitchen, the result is the same but the scale and capacity differ.
A stack consists of several components. The three most important ones are:
The Orchestrator determines where and how your steps run. By default ZenML uses a local orchestrator, so everything runs on your machine. But you can switch to Kubernetes, Airflow, Kubeflow, or a cloud orchestrator like Vertex AI and SageMaker. What is cool is that your code does not change at all, you just swap the orchestrator in the stack. The Artifact Store is where all the outputs from your steps are stored. By default it is in a local folder, but for production you can use cloud storage like Amazon S3, Google Cloud Storage, or Azure Blob Storage. All artifacts are stored neatly and can be accessed again anytime. The Experiment Tracker is optional but very useful. It records the metrics, parameters, and results of your experiments. The most popular one is MLflow, but there is also Weights & Biases and Neptune.To see the currently active stack, run:
zenml stack describe
And to see all the stacks you have:
zenml stack list
When you first use ZenML, you automatically get a default stack containing a local orchestrator and a local artifact store. This is enough for learning and development. Later, when you want to move to production, you just create a new stack.
Artifacts and Caching
Every output from a step in ZenML is stored as an artifact. These artifacts are versioned automatically, so you have a complete history of all the data and models that were ever produced. This is really important for reproducibility. If someone one day asks "which data was yesterday's model version trained on?", you can answer with certainty because everything is recorded.
Now what I love most about this artifact system is caching. This is the feature that makes ZenML super efficient. The concept is this: if you rerun a pipeline and there is a step whose input and code have not changed at all, ZenML will not rerun that step. It immediately grabs the result from the cache. Imagine you have a data preprocessing step that takes 10 minutes. If you only change the training step, you don't need to wait another 10 minutes for preprocessing, ZenML instantly uses the cached result. It saves a ton of time.
Caching is enabled by default. But if you want to turn off caching for a specific step, for example a step that fetches real-time data that is always changing, you can set it like this:
from zenml import step
@step(enablecache=False)
def fetchrealtimedata() -> dict:
"""This step always reruns, it does not use the cache."""
# fetch the latest data from an API
return {"data": "always fresh"}
You can also disable caching at the pipeline level with @pipeline(enablecache=False). Super flexible, right?
Integration with MLflow and Cloud
ZenML has dozens of built-in integrations. To install an integration, you use the zenml integration install command. For example for MLflow:
zenml integration install mlflow -y
After that, you register the MLflow experiment tracker with ZenML and add it to a stack:
zenml experiment-tracker register mlflowtracker --flavor=mlflow
zenml stack register mlflowstack -o default -a default -e mlflowtracker --set
The command above creates a new stack called mlflowstack that uses the default orchestrator, the default artifact store, and the MLflow experiment tracker. The --set flag immediately activates this stack. Now after that, inside your training step you just activate the experiment tracker:
from zenml import step
from zenml.client import Client
experimenttracker = Client().activestack.experimenttracker
@step(experimenttracker=experimenttracker.name)
def trainwithtracking(data: dict) -> float:
import mlflow
mlflow.sklearn.autolog()
# your training code goes here
accuracy = 0.95
mlflow.logmetric("accuracy", accuracy)
return accuracy
For the cloud, the concept is the same. For example if you want to use AWS, install the s3 and aws integrations, register an artifact store that points to your S3 bucket, and register a cloud orchestrator. Once again, your pipeline code does not change. This is ZenML's main strength that I keep emphasizing: separate the logic from the infrastructure.
End-to-End Pipeline Example
Alright friends, now let's combine everything we have learned into one complete pipeline: load data, train, evaluate. I use the Iris dataset from scikit-learn so everyone can try it right away. Here is the complete code, just copy and run it:
from zenml import step, pipeline
from typing import Tuple
from typingextensions import Annotated
import pandas as pd
from sklearn.datasets import loadiris
from sklearn.modelselection import traintestsplit
from sklearn.ensemble import RandomForestClassifier
from sklearn.base import ClassifierMixin
@step
def loaddata() -> Tuple[
Annotated[pd.DataFrame, "Xtrain"],
Annotated[pd.DataFrame, "Xtest"],
Annotated[pd.Series, "ytrain"],
Annotated[pd.Series, "ytest"],
]:
"""Load the Iris dataset and split into train and test."""
iris = loadiris(asframe=True)
X = iris.data
y = iris.target
Xtrain, Xtest, ytrain, ytest = traintestsplit(
X, y, testsize=0.2, randomstate=42
)
return Xtrain, Xtest, ytrain, ytest
@step
def trainmodel(
Xtrain: pd.DataFrame, ytrain: pd.Series
) -> ClassifierMixin:
"""Train a Random Forest model."""
model = RandomForestClassifier(nestimators=100, randomstate=42)
model.fit(Xtrain, ytrain)
return model
@step
def evaluatemodel(
model: ClassifierMixin, Xtest: pd.DataFrame, ytest: pd.Series
) -> float:
"""Evaluate the model accuracy on the test data."""
accuracy = model.score(Xtest, ytest)
print(f"Model accuracy: {accuracy:.4f}")
return accuracy
@pipeline
def trainingpipeline():
Xtrain, Xtest, ytrain, ytest = loaddata()
model = trainmodel(Xtrain, ytrain)
accuracy = evaluatemodel(model, Xtest, ytest)
return accuracy
if name == "main":
trainingpipeline()
Take a close look at the code above. There are three steps: loaddata, trainmodel, and evaluatemodel. In loaddata I use Annotated to give a name to each output artifact, so in the dashboard later it is clear which one is Xtrain, which is ytest, and so on. This is a best practice that makes your pipeline much easier to read.
In the pipeline, I just chain those three steps. ZenML automatically understands that trainmodel needs the output of loaddata, and evaluatemodel needs the output of trainmodel and loaddata. It figures out the order by itself. Run python run.py, then open the dashboard, and you will see a nice DAG (Directed Acyclic Graph) showing the flow of your pipeline.
Try running it twice. On the second run, if you did not change anything, notice that ZenML will say the steps are "cached". That is the caching we discussed. Now try changing nestimators in trainmodel to 50, then run it again. This time the loaddata step stays cached (because it did not change), but trainmodel and evaluatemodel will rerun. Super efficient, right?
Deploy and Serving
After your model is ready, the next step is to deploy it so it can be used for predictions. ZenML has several ways to do this. The most common is using a model deployer integration like MLflow, Seldon, or BentoML.
The concept is that you add a deployment step to your pipeline. This step takes the trained model and deploys it as an API endpoint that can receive prediction requests. As a simple example using the MLflow deployer, you first install the integration:
zenml integration install mlflow -y
zenml model-deployer register mlflowdeployer --flavor=mlflow
Then in the pipeline, you can add a step to deploy a model that passes a certain accuracy threshold. So a bad model will not get deployed. This is called a continuous deployment pipeline, one of the most powerful MLOps patterns. Your model becomes an endpoint that is ready to receive requests, and you can send new data to get predictions in real time.
For real production, you would usually run this pipeline on a schedule using an orchestrator like Airflow or Kubeflow, then deploy to a Kubernetes cluster. Once again, because of ZenML's architecture that separates logic from infrastructure, the transition from local to production is really smooth.
Best Practices
After using ZenML for quite a while across various projects, there are some best practices I want to share with you so you don't repeat the mistakes I made in the past.
First, always use clear type hints. I said this at the start but I am repeating it because it is important. Type hints are not just to keep your code neat, ZenML uses this information to serialize artifacts. If your type hints are wrong or missing, you can get confusing errors. So get into the habit of writing correct type hints on every step input and output. Second, build small and focused steps. Don't build one giant step that does everything. Split it into small steps that each have a single responsibility. This makes your caching more effective (because unchanged steps can be cached), and makes debugging easier (if there is an error, you know exactly which step is the problem). Third, use caching wisely. Caching is your friend, but be careful with steps that should always be fresh, like steps that fetch data from an external source that keeps changing. For steps like that, turn off caching withenablecache=False. Conversely, for heavy and deterministic steps, leave caching on so your iterations are fast.
Fourth, separate configuration from code. Don't hardcode parameters like learning rate or number of epochs directly in the code. ZenML supports YAML configuration files where you can put all your parameters. So you can change experiments without touching the code at all. This makes team collaboration much cleaner.
Fifth, start from a local stack, then move to the cloud. Don't immediately set up complex cloud infrastructure from the start. Develop and test on a local stack first until your pipeline runs smoothly. Only after you are confident, move to the production stack. Because your code is portable, this transition will not be a hassle.
Sixth, consistently use an experiment tracker. From the start of the project, get into the habit of using an experiment tracker like MLflow. Don't wait until the project is huge to start thinking about tracking. Recording all your experiments from the beginning will save you when you need to compare dozens of models to find the best one.
Seventh, give artifacts descriptive names. Use Annotated to name your artifacts like I showed earlier. In the dashboard, artifacts with clear names are much easier to trace than ones that are just "output0", "output1".
Conclusion
Alright friends, we have reached the end of the article. We have covered a whole lot about ZenML. We started from installation and initialization, then moved into the core concepts of @step and @pipeline. We continued to deeper concepts like the stack with its three main components (orchestrator, artifact store, experiment tracker), the super efficient artifact and caching system, integrations with MLflow and the cloud, and a complete end-to-end pipeline example from load data, train, to evaluate. We also discussed how to deploy a model and a set of best practices from my own experience.
If I may summarize the single most important lesson from ZenML, it is this: separate the ML logic from the infrastructure. This is the principle that makes your code portable, reproducible, and production-ready. With ZenML, you can focus on the part that truly matters, which is building a good model, without having to worry about how that code runs on different infrastructures.
For those of you who are serious about a career in machine learning and LLMs, I highly recommend diving deep into ZenML. MLOps skills like this are what will set you apart from most people who can only play around in a notebook. Start with small projects, try building a simple pipeline like the one I showed earlier, then slowly increase the complexity. Later you will feel for yourself how comfortable it is to work with a structured workflow.
That is all from me. I hope this article is helpful and gets you even more excited to learn. If you have any questions, don't hesitate to reach out to me. Happy coding, and see you in the next article. Keep up the spirit, friends.