# zenml integration install mlflow
# zenml experiment-tracker register mlflow_tracker -f mlflow ...
from zenml import pipeline, step
from zenml.integrations.mlflow.experiment_trackers import MLFlowExperimentTracker
from zenml.client import Client
@step(experiment_tracker="mlflow_tracker")
def train_model(X_train, y_train, X_test, y_test):
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
model = RandomForestClassifier()
model.fit(X_train, y_train)
# ZenML automatically logs parameters and the model
y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
# Only need to log custom metrics
return model, {"accuracy": accuracy}
@pipeline
def ml_pipeline():
X, y = load_data()
X_train, X_test, y_train, y_test = preprocess_data(X, y)
model, metrics = train_model(X_train, y_train, X_test, y_test)
ml_pipeline()