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

Viewed 271

I want to use an Azure Machine Learning compute cluster as a compute target to run a Kedro pipeline integrated with Mlflow.

Here's the code snippet (hooks.py) that integrates experiment tracking using Mlflow and Azure ML as backend/artifact stores.

"""Project hooks."""
from typing import Any, Dict, Iterable, Optional
import git
import os
import mlflow
import mlflow.sklearn
from kedro.config import ConfigLoader
from kedro.framework.hooks import hook_impl
from kedro.io import DataCatalog
from kedro.pipeline.node import Node
from kedro.versioning import Journal
from azureml.core import Workspace
from azureml.core.experiment import Experiment

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()

This works well on a compute instance that I created in my Azure ML workspace, simply by doing the following :

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

That's ok but what I really want is to use compute clusters to deal with hyperparameter tuning and other heavy workloads... I Just want to mention here that I still want to git clone to the compute instance and submit the run to the compute cluster from within the compute instance (but if anyone has a better approach, please feel free to share).

I know of two ways (listed below) to specify a compute cluster as a compute target in Azure ML but both require to pass a script parameter.

  1. Pure Azure ML ScriptRunConfig() method to submit experiments by specifying script and compute_target parameters. See Submit remote run with Azure Ml
  2. Mlflow integration with Azure ML : that requires to add an MLproject file to the project folder. See Submit an mlflow project run.

I tried for quite some time now to figure out how to do that within the Kedro structure but without success. So my question here, what's the best way to push experiment runs in a Kedro Pipeline to Azure ML compute clusters?

Thank you in advance for your help !

0 Answers
Related