ZenML

Overview

ZenML is an open-source MLOps framework that enables developers to build and manage reproducible machine learning pipelines. Founded in 2022, the platform provides a Pythonic interface for defining end-to-end ML workflows, from data ingestion and preprocessing to model training, evaluation, and deployment. Its core objective is to bridge the gap between local ML experimentation and production-grade MLOps, allowing teams to maintain consistency and traceability across their development lifecycle.

The framework centers on the concept of 'stacks,' which are modular configurations of MLOps tools for different stages of the ML pipeline. These stacks can include orchestrators (like Kubeflow Pipelines or Apache Airflow), artifact stores, data validators, model deployers, and experiment trackers. This modularity allows developers to choose and combine their preferred tools without being locked into a specific vendor ecosystem. For example, a user could configure a ZenML stack to use MLflow for experiment tracking and MinIO for artifact storage, orchestrated by Kubeflow Pipelines, all defined and managed through ZenML's Python API.

ZenML supports a range of use cases, from individual researchers needing robust experiment tracking to enterprise teams implementing complex, distributed ML systems. It focuses on facilitating reproducibility by versioning pipeline code, configurations, data artifacts, and models. This ensures that any experiment or production run can be replicated, a critical requirement for debugging, auditing, and regulatory compliance. The platform offers both an open-source version (ZenML OSS) and a managed cloud offering (ZenML Cloud), catering to different operational needs and team sizes. ZenML Cloud adds features like centralized experiment dashboards, collaboration tools, and enhanced security controls, including SOC 2 Type II compliance and GDPR adherence.

Its design principles prioritize developer experience, providing a familiar Python interface that integrates directly with popular ML frameworks such as TensorFlow, PyTorch, and Scikit-learn. The platform's emphasis on flexibility and integration makes it suitable for organizations looking to standardize their MLOps practices while retaining control over their underlying infrastructure and tool choices.

Key features

• ML Pipeline Orchestration: Define, run, and manage end-to-end machine learning pipelines using Python, integrating with orchestrators like Kubeflow Pipelines, Airflow, and Apache Spark. These pipelines specify the sequence of steps, data flow, and resource requirements.
• Reproducibility: Automatically track and version pipeline runs, code versions, configurations, and artifacts, ensuring experiments and production deployments are reproducible. This includes capturing environment details and dependencies for consistent results.
• Experiment Tracking and Versioning: Integrate with popular experiment trackers such as MLflow and Weights & Biases to log metrics, parameters, and models for comparative analysis and version control. This allows for detailed inspection of past runs.
• MLOps Stack Abstraction: Abstract various MLOps tools into modular components (stacks), enabling users to mix and match orchestrators, artifact stores, model deployers, and experiment trackers according to their needs. This provides flexibility in infrastructure choices.
• Local to Production Workflows: Facilitate the transition of ML code from local development and experimentation environments to scalable production deployments without extensive refactoring. This is achieved through consistent pipeline definitions across environments.
• Artifact Management: Provide structured storage and versioning for all pipeline artifacts, including datasets, trained models, and intermediate outputs, ensuring data lineage and accessibility. Supports various artifact stores like S3, GCS, and MinIO.
• Data Validation Integration: Connect with data validation tools like Great Expectations to ensure data quality and integrity at different stages of the pipeline, preventing data-related issues from impacting model performance.
• Model Deployment: Enable seamless deployment of trained models to production environments with integrations for tools like Seldon Core and BentoML, supporting various serving patterns and endpoint management.
• Collaboration Features (ZenML Cloud): Offer team collaboration tools, centralized dashboards, and access control for managing ML workflows across multiple users in a secure, shared environment.

Pricing

ZenML offers both an open-source version and a cloud-hosted platform with a free tier and paid subscription plans. The detailed pricing is available on the ZenML pricing page.

Common integrations

• Orchestrators:
  • Kubeflow Pipelines: Kubeflow Orchestrator integration details
  • Apache Airflow: Airflow Orchestrator integration guide
  • Apache Spark: Spark Orchestrator documentation
• Experiment Trackers:
  • MLflow: MLflow Experiment Tracker integration
  • Weights & Biases: Weights & Biases integration guide
• Artifact Stores:
  • AWS S3: S3 Artifact Store configuration
  • Google Cloud Storage (GCS): GCS Artifact Store setup
  • MinIO: MinIO Artifact Store documentation
• Container Registries:
  • Docker Hub, GCR, ECR: Container Registry configurations
• Data Validators:
  • Great Expectations: Great Expectations integration
• Model Deployers:
  • BentoML: BentoML Model Deployer guide
  • Seldon Core: Seldon Core Model Deployer integration

Alternatives

• MLflow: An open-source platform for managing the ML lifecycle, including experiment tracking, reproducible runs, and model deployment.
• Kubeflow Pipelines: A component of Kubeflow that provides a platform for building and deploying portable, scalable ML workflows based on Docker containers.
• Weights & Biases: A development platform for machine learning that helps teams track experiments, manage datasets, and collaborate on model development.
• AWS Step Functions: A serverless workflow service that allows developers to orchestrate distributed applications and microservices, including ML workflows.
• Google Cloud Vertex AI Pipelines: A fully managed service for building, deploying, and managing ML pipelines on Google Cloud, often used for production ML systems.

Getting started

To begin using ZenML, you typically install the Python package, initialize a ZenML repository, and define your first pipeline. The following example demonstrates a basic pipeline that simulates data loading, model training, and evaluation steps.

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

from zenml import pipeline, step
from zenml.client import Client

# Define a ZenML step for data loading
@step
def data_loader() -> pd.DataFrame:
    """Loads a simple dummy dataset."""
    # In a real scenario, this would load from a database or storage service
    data = {
        'feature1': [10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
        'feature2': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
        'target': [0, 0, 0, 0, 1, 1, 1, 1, 1, 1]
    }
    df = pd.DataFrame(data)
    print("Data loaded successfully.")
    return df

# Define a ZenML step for data splitting
@step
def data_splitter(data: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame, pd.Series, pd.Series]:
    """Splits the data into training and testing sets."""
    X = data[['feature1', 'feature2']]
    y = data['target']
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    print("Data split into training and testing sets.")
    return X_train, X_test, y_train, y_test

# Define a ZenML step for model training
@step
def model_trainer(
    X_train: pd.DataFrame, y_train: pd.Series
) -> LogisticRegression:
    """Trains a Logistic Regression model."""
    model = LogisticRegression()
    model.fit(X_train, y_train)
    print("Model training complete.")
    return model

# Define a ZenML step for model evaluation
@step
def model_evaluator(
    model: LogisticRegression, X_test: pd.DataFrame, y_test: pd.Series
) -> float:
    """Evaluates the trained model."""
    predictions = model.predict(X_test)
    accuracy = accuracy_score(y_test, predictions)
    print(f"Model accuracy: {accuracy}")
    return accuracy

# Define the ZenML pipeline orchestrating the steps
@pipeline
def simple_ml_pipeline():
    """A simple ML pipeline using ZenML."""
    df = data_loader()
    X_train, X_test, y_train, y_test = data_splitter(data=df)
    model = model_trainer(X_train=X_train, y_train=y_train)
    accuracy = model_evaluator(model=model, X_test=X_test, y_test=y_test)

# To run this pipeline:
# 1. Install ZenML: pip install zenml
# 2. Initialize a ZenML repository: zenml init
# 3. Run the pipeline from your Python script:
if __name__ == "__main__":
    # Ensure ZenML is initialized in your project directory
    # client = Client()
    # client.initialize()

    print("Running simple_ml_pipeline...")
    simple_ml_pipeline()
    print("Pipeline finished.")

This example demonstrates how to define individual machine learning operations as @step functions and then compose them into an end-to-end @pipeline. ZenML automatically handles the execution order, data passing between steps, and tracking of artifacts. After defining this pipeline, you would typically run zenml init in your project directory to set up the ZenML repository and then execute the Python script. ZenML will log the pipeline run, its steps, and artifacts, which can be viewed in the ZenML CLI or dashboard.