Graphviz

zenml.integrations.graphviz special

Initialization of the Graphviz integration.

GraphvizIntegration (Integration)

Definition of Graphviz integration for ZenML.

Source code in zenml/integrations/graphviz/__init__.py

class GraphvizIntegration(Integration):
    """Definition of Graphviz integration for ZenML."""

    NAME = GRAPHVIZ
    REQUIREMENTS = ["graphviz>=0.17"]

visualizers special

Initialization of Graphviz visualizers.

pipeline_run_dag_visualizer

Implementation of the Graphviz pipeline run DAG visualizer.

PipelineRunDagVisualizer (BaseVisualizer)

Visualize the lineage of runs in a pipeline.

Source code in zenml/integrations/graphviz/visualizers/pipeline_run_dag_visualizer.py

class PipelineRunDagVisualizer(BaseVisualizer):
    """Visualize the lineage of runs in a pipeline."""

    ARTIFACT_DEFAULT_COLOR = "blue"
    ARTIFACT_CACHED_COLOR = "green"
    ARTIFACT_SHAPE = "box"
    ARTIFACT_PREFIX = "artifact_"
    STEP_COLOR = "#431D93"
    STEP_SHAPE = "ellipse"
    STEP_PREFIX = "step_"
    FONT = "Roboto"

    @abstractmethod
    def visualize(
        self, object: PipelineRunView, *args: Any, **kwargs: Any
    ) -> graphviz.Digraph:
        """Creates a pipeline lineage diagram using graphviz.

        Args:
            object: The pipeline run view to visualize.
            *args: Additional arguments to pass to the visualization.
            **kwargs: Additional keyword arguments to pass to the visualization.

        Returns:
            A graphviz digraph object.

        Raises:
            RuntimeError: If the `dot` package isn't installed.
        """
        logger.warning(
            "This integration is not completed yet. Results might be unexpected."
        )
        if not shutil.which("dot"):
            raise RuntimeError(
                "Unable to find the `dot` package on your system. Please "
                "install it (https://graphviz.org/download/) and try again."
            )

        dot = graphviz.Digraph(comment=object.name)

        # link the steps together
        for step in object.steps:
            # add each step as a node
            dot.node(
                self.STEP_PREFIX + str(step.id),
                step.entrypoint_name,
                shape=self.STEP_SHAPE,
            )
            # for each parent of a step, add an edge

            for artifact_name, artifact in step.outputs.items():
                dot.node(
                    self.ARTIFACT_PREFIX + str(artifact.id),
                    f"{artifact_name} \n" f"({artifact.data_type})",
                    shape=self.ARTIFACT_SHAPE,
                )
                dot.edge(
                    self.STEP_PREFIX + str(step.id),
                    self.ARTIFACT_PREFIX + str(artifact.id),
                )

            for artifact_name, artifact in step.inputs.items():
                dot.edge(
                    self.ARTIFACT_PREFIX + str(artifact.id),
                    self.STEP_PREFIX + str(step.id),
                )

        with tempfile.NamedTemporaryFile(delete=False, suffix=".html") as f:
            dot.render(filename=f.name, format="png", view=True, cleanup=True)
        return dot

visualize(self, object, *args, **kwargs)

Creates a pipeline lineage diagram using graphviz.

Parameters:

Name	Type	Description	Default
object	PipelineRunView	The pipeline run view to visualize.	required
*args	Any	Additional arguments to pass to the visualization.	()
**kwargs	Any	Additional keyword arguments to pass to the visualization.	{}

Returns:

Type	Description
Digraph	A graphviz digraph object.

Exceptions:

Type	Description
RuntimeError	If the dot package isn't installed.

Source code in zenml/integrations/graphviz/visualizers/pipeline_run_dag_visualizer.py

@abstractmethod
def visualize(
    self, object: PipelineRunView, *args: Any, **kwargs: Any
) -> graphviz.Digraph:
    """Creates a pipeline lineage diagram using graphviz.

    Args:
        object: The pipeline run view to visualize.
        *args: Additional arguments to pass to the visualization.
        **kwargs: Additional keyword arguments to pass to the visualization.

    Returns:
        A graphviz digraph object.

    Raises:
        RuntimeError: If the `dot` package isn't installed.
    """
    logger.warning(
        "This integration is not completed yet. Results might be unexpected."
    )
    if not shutil.which("dot"):
        raise RuntimeError(
            "Unable to find the `dot` package on your system. Please "
            "install it (https://graphviz.org/download/) and try again."
        )

    dot = graphviz.Digraph(comment=object.name)

    # link the steps together
    for step in object.steps:
        # add each step as a node
        dot.node(
            self.STEP_PREFIX + str(step.id),
            step.entrypoint_name,
            shape=self.STEP_SHAPE,
        )
        # for each parent of a step, add an edge

        for artifact_name, artifact in step.outputs.items():
            dot.node(
                self.ARTIFACT_PREFIX + str(artifact.id),
                f"{artifact_name} \n" f"({artifact.data_type})",
                shape=self.ARTIFACT_SHAPE,
            )
            dot.edge(
                self.STEP_PREFIX + str(step.id),
                self.ARTIFACT_PREFIX + str(artifact.id),
            )

        for artifact_name, artifact in step.inputs.items():
            dot.edge(
                self.ARTIFACT_PREFIX + str(artifact.id),
                self.STEP_PREFIX + str(step.id),
            )

    with tempfile.NamedTemporaryFile(delete=False, suffix=".html") as f:
        dot.render(filename=f.name, format="png", view=True, cleanup=True)
    return dot