Skip to content

Mlflow

Initialization for the ZenML MLflow integration.

The MLflow integrations currently enables you to use MLflow tracking as a convenient way to visualize your experiment runs within the MLflow UI.

MlflowIntegration

Bases: Integration

Definition of MLflow integration for ZenML.

Source code in src/zenml/integrations/mlflow/__init__.py 
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
class MlflowIntegration(Integration):
    """Definition of MLflow integration for ZenML."""

    NAME = MLFLOW

    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
        "python-rapidjson",
        "pydantic",
        "numpy",
        "pandas",
    ]

    @classmethod
    def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
        """Method to get the requirements for the integration.

        Args:
            target_os: The target operating system to get the requirements for.

        Returns:
            A list of requirements.
        """
        from zenml.integrations.numpy import NumpyIntegration
        from zenml.integrations.pandas import PandasIntegration

        reqs = [
            "mlflow>=2.1.1,<3",
            # TODO: remove this requirement once rapidjson is fixed
            "python-rapidjson<1.15",
            # When you do:
            #       pip install zenml
            # You get all our required dependencies. However, if you follow it
            # with:
            #       zenml integration install mlflow
            # This downgrades pydantic to v1 even though mlflow does not have
            # any issues with v2. This is why we have to pin it here so a
            # downgrade will not happen.
            "pydantic>=2.8.0,<2.9.0",
        ]
        if sys.version_info.minor >= 12:
            logger.debug(
                "The MLflow integration on Python 3.12 and above is not yet "
                "fully supported: The extra dependencies 'mlserver' and "
                "'mlserver-mlflow' will be skipped."
            )
        else:
            reqs.extend([
                "mlserver>=1.3.3",
                "mlserver-mlflow>=1.3.3",
            ])

        reqs.extend(NumpyIntegration.get_requirements(target_os=target_os))
        reqs.extend(PandasIntegration.get_requirements(target_os=target_os))
        return reqs

    @classmethod
    def activate(cls) -> None:
        """Activate the MLflow integration."""
        from zenml.integrations.mlflow import services  # noqa

    @classmethod
    def flavors(cls) -> List[Type[Flavor]]:
        """Declare the stack component flavors for the MLflow integration.

        Returns:
            List of stack component flavors for this integration.
        """
        from zenml.integrations.mlflow.flavors import (
            MLFlowExperimentTrackerFlavor,
            MLFlowModelDeployerFlavor,
            MLFlowModelRegistryFlavor,
        )

        return [
            MLFlowModelDeployerFlavor,
            MLFlowExperimentTrackerFlavor,
            MLFlowModelRegistryFlavor,
        ]

activate() classmethod

Activate the MLflow integration.

Source code in src/zenml/integrations/mlflow/__init__.py 
90
91
92
93
@classmethod
def activate(cls) -> None:
    """Activate the MLflow integration."""
    from zenml.integrations.mlflow import services  # noqa

flavors() classmethod

Declare the stack component flavors for the MLflow integration.

Returns:

Type Description
List[Type[Flavor]]

List of stack component flavors for this integration.
Source code in src/zenml/integrations/mlflow/__init__.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
@classmethod
def flavors(cls) -> List[Type[Flavor]]:
    """Declare the stack component flavors for the MLflow integration.

    Returns:
        List of stack component flavors for this integration.
    """
    from zenml.integrations.mlflow.flavors import (
        MLFlowExperimentTrackerFlavor,
        MLFlowModelDeployerFlavor,
        MLFlowModelRegistryFlavor,
    )

    return [
        MLFlowModelDeployerFlavor,
        MLFlowExperimentTrackerFlavor,
        MLFlowModelRegistryFlavor,
    ]

get_requirements(target_os=None) classmethod

Method to get the requirements for the integration.

Parameters:

Name Type Description Default
target_os Optional[str]

The target operating system to get the requirements for.

 None

Returns:

Type Description
List[str]

A list of requirements.
Source code in src/zenml/integrations/mlflow/__init__.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
@classmethod
def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
    """Method to get the requirements for the integration.

    Args:
        target_os: The target operating system to get the requirements for.

    Returns:
        A list of requirements.
    """
    from zenml.integrations.numpy import NumpyIntegration
    from zenml.integrations.pandas import PandasIntegration

    reqs = [
        "mlflow>=2.1.1,<3",
        # TODO: remove this requirement once rapidjson is fixed
        "python-rapidjson<1.15",
        # When you do:
        #       pip install zenml
        # You get all our required dependencies. However, if you follow it
        # with:
        #       zenml integration install mlflow
        # This downgrades pydantic to v1 even though mlflow does not have
        # any issues with v2. This is why we have to pin it here so a
        # downgrade will not happen.
        "pydantic>=2.8.0,<2.9.0",
    ]
    if sys.version_info.minor >= 12:
        logger.debug(
            "The MLflow integration on Python 3.12 and above is not yet "
            "fully supported: The extra dependencies 'mlserver' and "
            "'mlserver-mlflow' will be skipped."
        )
    else:
        reqs.extend([
            "mlserver>=1.3.3",
            "mlserver-mlflow>=1.3.3",
        ])

    reqs.extend(NumpyIntegration.get_requirements(target_os=target_os))
    reqs.extend(PandasIntegration.get_requirements(target_os=target_os))
    return reqs