Skip to content

Instantly share code, notes, and snippets.

@ywolff
Created February 15, 2022 01:18
Show Gist options
  • Save ywolff/57658f8178a3b296ccb0167cf00afe7d to your computer and use it in GitHub Desktop.
Save ywolff/57658f8178a3b296ccb0167cf00afe7d to your computer and use it in GitHub Desktop.
from functools import wraps
import mlflow
from src.constants import PROJECT_EXPERIMENT_NAME
def mlflow_run(wrapped_function):
@wraps(wrapped_function)
def wrapper(*args, **kwargs):
mlflow.set_experiment(PROJECT_EXPERIMENT_NAME)
with mlflow.start_run(): # recover parent run thanks to MLFLOW_RUN_ID env variable
with mlflow.start_run(run_name=wrapped_function.__name__, nested=True): # start child run
return wrapped_function(*args, **kwargs)
return wrapper
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment