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 :
git clonethe source code into the Azure ML compute instance- Do a
kedro runin 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.
- Pure Azure ML
ScriptRunConfig()method to submit experiments by specifyingscriptandcompute_targetparameters. See Submit remote run with Azure Ml - 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 !