Visualizations

Data visualization is a powerful tool for understanding your ML pipeline outputs. ZenML provides built-in capabilities to visualize artifacts, helping you gain insights into your data, model performance, and pipeline execution.

Accessing Visualizations

ZenML automatically generates visualizations for many common data types, making it easy to inspect your artifacts without additional code.

Dashboard Visualizations

The ZenML dashboard displays visualizations for artifacts produced by your pipeline runs:

To view visualizations in the dashboard:

1. Navigate to the Runs tab

2. Select a specific pipeline run

3. Click on any step to view its outputs

4. Select an artifact to view its visualizations

Notebook Visualizations

You can also display artifact visualizations in Jupyter notebooks using the visualize() method:

from zenml.client import Client

# Get an artifact from a previous pipeline run
run = Client().get_pipeline_run("<PIPELINE_RUN_ID>")
artifact = run.steps["<STEP_NAME>"].outputs[<OUTPUT_NAME>][0]

# Display the visualization
artifact.visualize()

Supported Visualization Types

ZenML supports visualizations for many common data types out of the box:

• A statistical representation of a Pandas Dataframe represented as a png image.

• Drift detection reports by Evidently, Great Expectations, and whylogs.

• A Hugging Face datasets viewer embedded as a HTML iframe.

Creating Custom Visualizations

It is simple to associate a custom visualization with an artifact in ZenML, if the visualization is one of the supported visualization types. Currently, the following visualization types are supported:

• HTML: Embedded HTML visualizations such as data validation reports,

• Image: Visualizations of image data such as Pillow images (e.g. PIL.Image) or certain numeric numpy arrays,

• CSV: Tables, such as the pandas DataFrame .describe() output,

• Markdown: Markdown strings or pages.

• JSON: JSON strings or objects.

There are three ways how you can add custom visualizations to the dashboard:

• If you are already handling HTML, Markdown, CSV or JSON data in one of your steps, you can have them visualized in just a few lines of code by casting them to a special class inside your step.

• If you want to automatically extract visualizations for all artifacts of a certain data type, you can define type-specific visualization logic by building a custom materializer.

Visualization via Special Return Types

If you already have HTML, Markdown, CSV or JSON data available as a string inside your step, you can simply cast them to one of the following types and return them from your step:

• zenml.types.HTMLString for strings in HTML format, e.g., "<h1>Header</h1>Some text",

• zenml.types.MarkdownString for strings in Markdown format, e.g., "# Header\nSome text",

• zenml.types.CSVString for strings in CSV format, e.g., "a,b,c\n1,2,3".

• zenml.types.JSONString for strings in JSON format, e.g., {"key": "value"}.

Example:

from zenml import step
from zenml.types import CSVString

@step
def my_step() -> CSVString:
    some_csv = "a,b,c\n1,2,3"
    return CSVString(some_csv)

This would create the following visualization in the dashboard:

Another example is visualizing a matplotlib plot by embedding the image in an HTML string:

import matplotlib.pyplot as plt
import base64
import io

from zenml.types import HTMLString
from zenml import step, pipeline

@step
def create_matplotlib_visualization() -> HTMLString:
    """Creates a matplotlib visualization and returns it as embedded HTML."""
    # Create plot
    fig, ax = plt.subplots()
    ax.plot([1, 2, 3, 4], [1, 4, 2, 3])
    ax.set_title('Sample Plot')
    
    # Convert plot to base64 string
    buf = io.BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight', dpi=300)
    plt.close(fig)  # Clean up
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    
    # Create HTML with embedded image
    html = f'''
    <div style="text-align: center;">
        <img src="data:image/png;base64,{image_base64}" 
             style="max-width: 100%; height: auto;">
    </div>
    '''
    
    return HTMLString(html)

@pipeline
def visualization_pipeline():
    create_matplotlib_visualization()

if __name__ == "__main__":
    visualization_pipeline()

Visualization via Materializers

If you want to automatically extract visualizations for all artifacts of a certain data type, you can do so by overriding the save_visualizations() method of the corresponding materializer. Let's look at an example of how to visualize matplotlib figures in your ZenML dashboard:

Example: Matplotlib Figure Visualization

1. Custom Class First, we create a custom class to hold our matplotlib figure:

from typing import Any
from pydantic import BaseModel

class MatplotlibVisualization(BaseModel):
    """Custom class to hold matplotlib figures."""
    figure: Any  # This will hold the matplotlib figure

2. Materializer Next, we create a custom materializer that handles this class and implements the visualization logic:

import os
from typing import Dict
from zenml.materializers.base_materializer import BaseMaterializer
from zenml.enums import VisualizationType
from zenml.io import fileio

class MatplotlibMaterializer(BaseMaterializer):
    """Materializer that handles matplotlib figures."""
    ASSOCIATED_TYPES = (MatplotlibVisualization,)

    def save_visualizations(
        self, data: MatplotlibVisualization
    ) -> Dict[str, VisualizationType]:
        """Create and save visualizations for the matplotlib figure."""
        visualization_path = os.path.join(self.uri, "visualization.png")
        with fileio.open(visualization_path, 'wb') as f:
            data.figure.savefig(f, format='png', bbox_inches='tight')
        return {visualization_path: VisualizationType.IMAGE}

3. Step Finally, we create a step that returns our custom type:

import matplotlib.pyplot as plt
from zenml import step

@step
def create_matplotlib_visualization() -> MatplotlibVisualization:
    """Creates a matplotlib visualization."""
    fig, ax = plt.subplots()
    ax.plot([1, 2, 3, 4], [1, 4, 2, 3])
    ax.set_title('Sample Plot')
    return MatplotlibVisualization(figure=fig)

For another example, see our Hugging Face datasets materializer which visualizes datasets by embedding their preview viewer.

Controlling Visualizations

Access to Visualizations

In order for the visualizations to show up on the dashboard, the following must be true:

Configuring a Service Connector

Visualizations are usually stored alongside the artifact, in the artifact store. Therefore, if a user would like to see the visualization displayed on the ZenML dashboard, they must give access to the server to connect to the artifact store.

The service connector documentation goes deeper into the concept of service connectors and how they can be configured to give the server permission to access the artifact store. For a concrete example, see the AWS S3 artifact store documentation.

Configuring Artifact Stores

If all visualizations of a certain pipeline run are not showing up in the dashboard, it might be that your ZenML server does not have the required dependencies or permissions to access that artifact store. See the custom artifact store docs page for more information.

Enabling/Disabling Visualizations

You can control whether visualizations are generated at the pipeline or step level:

# Disable visualizations for a pipeline
@pipeline(enable_artifact_visualization=False)
def my_pipeline():
    ...

# Disable visualizations for a step
@step(enable_artifact_visualization=False)
def my_step():
    ...

You can also configure this in YAML:

enable_artifact_visualization: False

steps:
  my_step:
    enable_artifact_visualization: True

Conclusion

Visualizing artifacts is a powerful way to gain insights from your ML pipelines. ZenML's built-in visualization capabilities make it easy to understand your data and model outputs, identify issues, and communicate results.

By leveraging these visualization tools, you can better understand your ML workflows, debug problems more effectively, and make more informed decisions about your models.