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 whensave
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.