Enable MlflowArtifactDataSet to load from a run_id

[Galileo-Galilei]

Description

I sometimes want to retrieve artifacts from an existing mlflow run, but there is no built-in functionality in kedro-mlflow to do so. I have to do it manually (create a mlflow client, feed in credentials and configuration, download the artifact, feed in the new path to the undelying dataset). This is error prone and often breaks the kedro principles by performing such IO operations inside a node.

Context

It is a common workflow to retrieve artifacts from a mlflow runs to feed a new run (e.g. retrieve embeddings trained in a ml pipeline and stored in mlflow in a previous run to feed another training pipeline for a NLP task like classification).

Possible Implementation

(Optional) Suggest an idea for implementing the addition or change.

The idea would be to add a load_args argument at the MlflowArtifactDataSet level to handle this.

catalog.yml

my_artifact_dataset:
    type: kedro_mlflow.io.MlflowArtifactDataSet
    load_args:
        run_id: 123456789
        credentials: mlflow_connect
    data_set:
        type: pandas.CSVDataSet
        load_args:
            sep: ";"

kedro_mlflow/io/mlflow_dataset.py

Add a _load method which creates a mlflow client, download the file to self._filepath and load the dataset of a run_id is provided.

Some important notes:

• the existing run_id attribute must be deprecated and moved to a save_args argument
• credentials (i.e. environment variables) must be properly managed before and after downloading to avoid messing up with the global session.
• We should likely add a logger to have informationon where the dataset is retrieved from (mlflow of filepath)

Possible Alternatives

It may be more user friendly to create to different datasets (MlflowArtifactLoggerDataSet and MlflowArtifactLoaderDataSet to separate these 2 behaviours, but I think it is still understandable with one dataset which performs both operations.