Hydra: [Bug] MLFlow folder not created with Hydra context

[nathansomavarapu]

🐛 Bug

This might not be a bug it might have more to do with my misunderstanding, but I can't find much info on it so I thought I would ask. How can MLflow be used in conjunction with hydra? When I try and use the MLflow context manager inside a main function with the Hydra decorator the MLflow logging directory is never created.

To reproduce

import hydra
import mlflow

# Commenting the below line make the script work as expected
@hydra.main(config_path="conf", config_name="test")
def main(cfg):
    params = [('test', 'test')]

    with mlflow.start_run() as run:
        for k,v in params:
            mlflow.log_param(k, str(v))

if __name__ == "__main__":
    main()

* Stack trace/error message *

// N/A

Expected Behavior

Expect a folder named mlruns to be created with a single run with the parameter test and value test.

System information

• Hydra Version : 1.0
• Python version : 3.7
• Virtual environment type and version : Conda
• Operating system : Ubuntu 16.04

Additional context

Add any other context about the problem here.

[jieru-hu]

@nathansomavarapu thx for filing the issue. Could you update your issue's title with more info?

The following thread might help:
https://stackoverflow.com/questions/62296590/mlflow-and-hydra-causing-crash-when-used-together

We are working on a new type of plugin that will support MLFlow more naturally.
see https://github.com/facebookresearch/hydra/issues/385

[nathansomavarapu]

I'll take a look at that thread! Sorry for the title.

[nathansomavarapu]

Ok that link was pretty helpful I was actually able to get my expected behavior by adding the following line before the MLflow context manager

mlflow.set_tracking_uri('file://' + utils.get_original_cwd() + '/mlruns')

So in total a working example for anyone who might run across this is,

import hydra
import mlflow
from hydra import utils

@hydra.main(config_path="conf", config_name="test")
def main(cfg):
    params = [('test', 'test')]

    mlflow.set_tracking_uri('file://' + utils.get_original_cwd() + '/mlruns')

    with mlflow.start_run() as run:
        for k,v in params:
            mlflow.log_param(k, str(v))

if __name__ == "__main__":
    main()