Datasets

Artifact DataSet

class kedro_mlflow.io.artifacts.mlflow_artifact_dataset.MlflowArtifactDataset(dataset: str | Dict, run_id: str = None, artifact_path: str = None, credentials: Dict[str, Any] = None)

Bases: AbstractVersionedDataset

This class is a wrapper for any kedro AbstractDataset. It decorates their save method to log the dataset in mlflow when save is called.

load() → Any

MlflowArtifactDataset is a factory for DataSet and consequently does not implements abtracts methods

save(data: Any) → None

MlflowArtifactDataset is a factory for DataSet and consequently does not implements abtracts methods

Metrics DataSet

class kedro_mlflow.io.metrics.mlflow_metric_dataset.MlflowMetricDataset(key: str = None, run_id: str = None, load_args: Dict[str, Any] = None, save_args: Dict[str, Any] = None)

Bases: MlflowAbstractMetricDataset

DEFAULT_SAVE_MODE = 'overwrite'
SUPPORTED_SAVE_MODES = {'append', 'overwrite'}
__init__(key: str = None, run_id: str = None, load_args: Dict[str, Any] = None, save_args: Dict[str, Any] = None)

Initialise MlflowMetricDataset. :param run_id: The ID of the mlflow run where the metric should be logged :type run_id: str

load() → None

Loads data by delegation to the provided load method.

Returns:

Data returned by the provided load method.

Raises:

DatasetError – When underlying load method raises error.

save(data: float) → None

Saves data by delegation to the provided save method.

Parameters:

data – the value to be saved by provided save method.

Raises:

  • DatasetError – when underlying save method raises error.

  • FileNotFoundError – when save method got file instead of dir, on Windows.

  • NotADirectoryError – when save method got file instead of dir, on Unix.

class kedro_mlflow.io.metrics.mlflow_metric_history_dataset.MlflowMetricHistoryDataset(key: str = None, run_id: str = None, load_args: Dict[str, Any] = None, save_args: Dict[str, Any] = None)

Bases: MlflowAbstractMetricDataset

__init__(key: str = None, run_id: str = None, load_args: Dict[str, Any] = None, save_args: Dict[str, Any] = None)

Initialise MlflowMetricDataset. :param run_id: The ID of the mlflow run where the metric should be logged :type run_id: str

load() → None

Loads data by delegation to the provided load method.

Returns:

Data returned by the provided load method.

Raises:

DatasetError – When underlying load method raises error.

save(data: List[int] | Dict[int, float] | List[Dict[str, float | str]]) → None

Saves data by delegation to the provided save method.

Parameters:

data – the value to be saved by provided save method.

Raises:

  • DatasetError – when underlying save method raises error.

  • FileNotFoundError – when save method got file instead of dir, on Windows.

  • NotADirectoryError – when save method got file instead of dir, on Unix.

Models DataSet

class kedro_mlflow.io.models.mlflow_abstract_model_dataset.MlflowAbstractModelDataSet(filepath: str, flavor: str, pyfunc_workflow: str | None = None, load_args: Dict[str, Any] = None, save_args: Dict[str, Any] = None, version: Version = None)

Bases: AbstractVersionedDataset

Abstract mother class for model datasets.

__init__(filepath: str, flavor: str, pyfunc_workflow: str | None = None, load_args: Dict[str, Any] = None, save_args: Dict[str, Any] = None, version: Version = None) → None

Initialize the Kedro MlflowAbstractModelDataSet.

Parameters are passed from the Data Catalog.

During save, the model is first logged to MLflow. During load, the model is pulled from MLflow run with run_id.

Parameters:

  • filepath (str) – Path to store the dataset locally.

  • flavor (str) – Built-in or custom MLflow model flavor module. Must be Python-importable.

  • pyfunc_workflow (str, optional) – Either python_model or loader_module. See https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#workflows.

  • load_args (Dict[str, Any], optional) – Arguments to load_model function from specified flavor. Defaults to {}.

  • save_args (Dict[str, Any], optional) – Arguments to log_model function from specified flavor. Defaults to {}.

  • version (Version, optional) – Specific version to load.

Raises:

DatasetError – When passed flavor does not exist.

class kedro_mlflow.io.models.mlflow_model_tracking_dataset.MlflowModelTrackingDataset(flavor: str, run_id: str | None = None, artifact_path: str | None = 'model', pyfunc_workflow: str | None = None, load_args: Dict[str, Any] | None = None, save_args: Dict[str, Any] | None = None)

Bases: MlflowAbstractModelDataSet

Wrapper for saving, logging and loading for all MLflow model flavor.

__init__(flavor: str, run_id: str | None = None, artifact_path: str | None = 'model', pyfunc_workflow: str | None = None, load_args: Dict[str, Any] | None = None, save_args: Dict[str, Any] | None = None) → None

Initialize the Kedro MlflowModelDataSet.

Parameters are passed from the Data Catalog.

During save, the model is first logged to MLflow. During load, the model is pulled from MLflow run with run_id.

Parameters:

  • flavor (str) – Built-in or custom MLflow model flavor module. Must be Python-importable.

  • run_id (Optional[str], optional) – MLflow run ID to use to load the model from or save the model to. Defaults to None.

  • artifact_path (str, optional) – the run relative path to the model.

  • pyfunc_workflow (str, optional) – Either python_model or loader_module. See https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#workflows.

  • load_args (Dict[str, Any], optional) – Arguments to load_model function from specified flavor. Defaults to None.

  • save_args (Dict[str, Any], optional) – Arguments to log_model function from specified flavor. Defaults to None.

Raises:

DatasetError – When passed flavor does not exist.

load() → Any

Loads an MLflow model from local path or from MLflow run.

Returns:

Deserialized model.

Return type:

Any

property model_uri
save(model: Any) → None

Save a model to local path and then logs it to MLflow.

Parameters:

model (Any) – A model object supported by the given MLflow flavor.

class kedro_mlflow.io.models.mlflow_model_local_filesystem_dataset.MlflowModelLocalFileSystemDataset(filepath: str, flavor: str, pyfunc_workflow: str | None = None, load_args: Dict[str, Any] = None, save_args: Dict[str, Any] = None, log_args: Dict[str, Any] = None, version: Version = None)

Bases: MlflowAbstractModelDataSet

Wrapper for saving, logging and loading for all MLflow model flavor.

__init__(filepath: str, flavor: str, pyfunc_workflow: str | None = None, load_args: Dict[str, Any] = None, save_args: Dict[str, Any] = None, log_args: Dict[str, Any] = None, version: Version = None) → None

Initialize the Kedro MlflowModelDataSet.

Parameters are passed from the Data Catalog.

During save, the model is saved locally at filepath During load, the model is loaded from the local filepath.

Parameters:

  • flavor (str) – Built-in or custom MLflow model flavor module. Must be Python-importable.

  • filepath (str) – Path to store the dataset locally.

  • pyfunc_workflow (str, optional) – Either python_model or loader_module. See https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#workflows.

  • load_args (Dict[str, Any], optional) – Arguments to load_model function from specified flavor. Defaults to None.

  • save_args (Dict[str, Any], optional) – Arguments to save_model function from specified flavor. Defaults to None.

  • version (Version, optional) – Kedro version to use. Defaults to None.

Raises:

DatasetError – When passed flavor does not exist.

load() → Any

Loads an MLflow model from local path or from MLflow run.

Returns:

Deserialized model.

Return type:

Any

save(model: Any) → None

Save a model to local path and then logs it to MLflow.

Parameters:

model (Any) – A model object supported by the given MLflow flavor.

class kedro_mlflow.io.models.mlflow_model_registry_dataset.MlflowModelRegistryDataset(model_name: str, stage_or_version: str | int | None = None, alias: str | None = None, flavor: str | None = 'mlflow.pyfunc', pyfunc_workflow: str | None = 'python_model', load_args: Dict[str, Any] | None = None)

Bases: MlflowAbstractModelDataSet

Wrapper for saving, logging and loading for all MLflow model flavor.

__init__(model_name: str, stage_or_version: str | int | None = None, alias: str | None = None, flavor: str | None = 'mlflow.pyfunc', pyfunc_workflow: str | None = 'python_model', load_args: Dict[str, Any] | None = None) → None

Initialize the Kedro MlflowModelRegistryDataset.

Parameters are passed from the Data Catalog.

During “load”, the model is pulled from MLflow model registry by its name. “save” is not supported.

Parameters:

  • model_name (str) – The name of the registered model is the mlflow registry

  • stage_or_version (str) – A valid stage (either “staging” or “production”) or version number for the registred model. Default to “latest” which fetch the last version and the higher “stage” available.

  • flavor (str) – Built-in or custom MLflow model flavor module. Must be Python-importable.

  • pyfunc_workflow (str, optional) – Either python_model or loader_module. See https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#workflows.

  • load_args (Dict[str, Any], optional) – Arguments to load_model function from specified flavor. Defaults to None.

Raises:

DatasetError – When passed flavor does not exist.

load() → Any

Loads an MLflow model from local path or from MLflow run.

Returns:

Deserialized model.

Return type:

Any

save(model: Any) → None

Saves data by delegation to the provided save method.

Parameters:

data – the value to be saved by provided save method.

Raises:

  • DatasetError – when underlying save method raises error.

  • FileNotFoundError – when save method got file instead of dir, on Windows.

  • NotADirectoryError – when save method got file instead of dir, on Unix.