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.

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

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

class kedro_mlflow.io.metrics.mlflow_metrics_dataset.MlflowMetricsHistoryDataset(run_id: str = None, prefix: str | None = None)

Bases: AbstractDataset

This class represent MLflow metrics dataset.

__init__(run_id: str = None, prefix: str | None = None)

Initialise MlflowMetricsHistoryDataset.

Parameters:

  • prefix (Optional[str]) – Prefix for metrics logged in MLflow.

  • run_id (str) – ID of MLflow run.

property run_id

Get run id.

If active run is not found, tries to find last experiment.

Raise DatasetError exception if run id can’t be found.

Returns:

String contains run_id.

Return type:

str

Models DataSet

class kedro_mlflow.io.models.mlflow_abstract_model_dataset.MlflowModelRegistryDataset(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

Absract 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 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:

  • 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: MlflowModelRegistryDataset

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.

property model_uri
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: MlflowModelRegistryDataset

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.

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

Bases: MlflowModelRegistryDataset

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

__init__(model_name: str, stage_or_version: str | int = 'latest', 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.