Use an Azure ML compute cluster to run Kedro + Mlflow pipeline

[Downfor-u]

How to use an Azure Machine Learning compute cluster to run a Kedro pipeline that integrates with Mlflow and Azure ML ?

For now, my code does not reference any remote compute. I just get my code on an Azure ML compute instance and run inside it i.e. doing the following :

1. git clone the source code into the compute instance
2. Do a kedro run in the compute instance Terminal

class ProjectHooks:
    @hook_impl
    def register_config_loader(
        self,
        conf_paths: Iterable[str],
        env: str,
        extra_params: Dict[str, Any],
    ) -> ConfigLoader:
        return ConfigLoader(conf_paths)

    @hook_impl
    def register_catalog(
        self,
        catalog: Optional[Dict[str, Dict[str, Any]]],
        credentials: Dict[str, Dict[str, Any]],
        load_versions: Dict[str, str],
        save_version: str,
        journal: Journal,
    ) -> DataCatalog:
        return DataCatalog.from_config(
            catalog, credentials, load_versions, save_version, journal
        )


class ModelTrackingHooks:
    """Namespace for grouping all model-tracking hooks with MLflow together."""

    @hook_impl
    def before_pipeline_run(self, run_params: Dict[str, Any]) -> None:
        """Hook implementation to start an MLflow run
        with the same run_id as the Kedro pipeline run.
        """

        # Get Azure workspace
        ws = Workspace.get(name=workspace_name,
                           subscription_id=subscription_id,
                           resource_group=resource_group)

        # Set tracking uri
        mlflow.set_tracking_uri(ws.get_mlflow_tracking_uri())

        # Create an Azure ML experiment in the workspace
        experiment = Experiment(workspace=ws, name='kedro-mlflow-experiment')
        mlflow.set_experiment(experiment.name)

        mlflow.start_run(run_name=run_params["run_id"])
        mlflow.log_params(run_params)

    @hook_impl
    def after_node_run(
        self, node: Node, outputs: Dict[str, Any], inputs: Dict[str, Any]
    ) -> None:
        """Hook implementation to add model tracking after some node runs.
        In this example, we will:
        * Log the parameters after the data splitting node runs.
        * Log the model after the model training node runs.
        * Log the model's metrics after the model evaluating node runs.
        """
        if node._func_name == "function_name":
            mlflow.log_metrics(...)

    @hook_impl
    def after_pipeline_run(self) -> None:
        """Hook implementation to end the MLflow run
        after the Kedro pipeline finishes.
        """
        mlflow.end_run()

What I want to do next is to update the code to specify a compute cluster as a compute target while still running the code from the compute instance. How can I do this ?

Thank you for your help !

[datajoely]

Ho @Downfor-u keen to help here - but I think we need a little more info, this feels like more of the questions are on Azure side. Kedro doesn't concern itself with the infrastructure side.

Does the run work at all without the MLFlow elements? Or does it not run at all?

[Downfor-u]

Thank you for your reply @datajoely !

Actually, the integration between Kedro, Mlflow and Azure ML works well. What I tried to achieve was to extend the code to specify a compute cluster as a compute target (and not rely only on the compute instance).

I already succeded in doing this with the ScriptRunConfig() class of Azure ML (see Submit remote run with Azure Ml) but I can't figure out how to add that within my Kedro's code structure.

Does that make sense ?

[datajoely]

I have two ideas, but not sure if they will work perfectly:

1. Create a script file that loads the context and performs a context.run() via from kedro.framework.context import load_context you then call this via the ScriptRunConfig operator.
2. You try packaging and installing the Kedro project and seeing if you can get the ScriptRunConfig or another operator to run that.

[Downfor-u]

Thanks for the ideas ! I'll try them out and go back to you for an update.

[madamak]

@Downfor-u
I am trying to do something similar and would be very interested if you have an update on this!

[michalbrys]

We recently released a plugin for Kedro AzureML Pipelines integration.
I'll leave a link for the future reference: https://github.com/getindata/kedro-azureml