MLflow log python function arguments in experiments

How to Log python function arguments/parameters and return value in MLflow.

MLflow is widely used for Record and query experiments: code, data, config, and results. Data scientists or Machine Learning engineers are familiar with MLflow, they use MLflow to log metrics while training, testing, and artifacts like models and versions. In some cases, we would like to dynamically log the user-defined function arguments and the function return type. Sadly this feature is not available out of the box by MLflow. So in this blog, I will show how I created a python decorator that will help us log the user-defined function's arguments passed and also the return value dynamically with a decorator call.

• Step 1: Create a user-defined function that accepts some arguments and returns a value

  def add(x,y):
    return x+y
  

  I kept the function very simple, its a function that accepts 2 integers and performs addition, and returns the result

• Step 2: Create a python decorator that takes a function as input and logs the arguments passed to that function with the return value

def log_function_param(func):
    def wrapper(*args,**kwargs):
        if args:
            log_param(f"__call_id__{str(uuid4())[:8]} . {str(func.__name__)} function args ", str(','.join([str(x) for x in args])))
        if kwargs:
            log_param(f"__call_id__{str(uuid4())[:8]} . {str(func.__name__)} function kwargs",str(kwargs.items()))
        start_time = time.time()
        val = func(*args,**kwargs)
        end_time = time.time()
        log_param(f"__call_id__{str(uuid4())[:8]} . {str(func.__name__)} function execution time ", str(end_time - start_time)+" secs")
        log_param(f"__call_id__{str(uuid4())[:8]} . {str(func.__name__)} function return value ", str(val))
        return val
    return wrapper

Let us now understand the snippet. The above snippet is a python decorator implementation which upon calling executes the function and passes the arguments to the functions call, in this decorator I have added couple of log_param which is imported from mlflow. the args stores all the arguments passed as a tuple, so if any user sends bunch of inputs are arguments then it will log all the arguments For example, if I call the function add(4,5,6,7,6,5,34,45,56,67) then the value of args will be (4,5,6,7,6,5,34,45,56,67). Similarly if I call the function add(a=5,b=10,c=2) with key and value then kwargs will be a dictionary with key as dict key and the values as dict value, {"a":5,"b":10,"c":2}

- Step 3: Execute the code and view the result in MLflow experiments

%% main.py
import mlflow
import time
from mlflow import log_metric, log_param, log_artifacts
from uuid import uuid4

def log_function_param(func):
    def wrapper(*args,**kwargs):
        if args:
            log_param(f"__call_id__{str(uuid4())[:8]} . {str(func.__name__)} function args ", str(','.join([str(x) for x in args])))
        if kwargs:
            log_param(f"__call_id__{str(uuid4())[:8]} . {str(func.__name__)} function kwargs",str(kwargs.items()))
        start_time = time.time()
        val = func(*args,**kwargs)
        end_time = time.time()
        log_param(f"__call_id__{str(uuid4())[:8]} . {str(func.__name__)} function execution time ", str(end_time - start_time)+" secs")
        log_param(f"__call_id__{str(uuid4())[:8]} . {str(func.__name__)} function return value ", str(val))
        return val
    return wrapper

@log_function_param
def add(x,y):
    return x+y

if __name__ == "__main__":
    mlflow.set_experiment("Test Exp 2")
    r = add(x=5,y=7)

execute the py file

python main.py

Then check the MLflow experiment with name Test Exp 2 by running the following command in cmd

mlflow ui

executing the same script, but now the add function sends integer as arguments

if __name__ == "__main__":
    mlflow.set_experiment("Test Exp 2")
    r = add(55,17)

This python decorator will log any user-defined function any number or arguments and return value by adding @log_function_param decorator call above the function you want to log.