def test_kedro_mlflow_config_setup_tracking_priority(kedro_project_with_mlflow_conf):
"""Test if the mlflow_tracking uri set is the one of mlflow.yml
if it also exist in credentials.
"""
# create a ".kedro.yml" file to identify "tmp_path" as the root of a kedro project
(kedro_project_with_mlflow_conf / "conf/base/credentials.yml").write_text(
yaml.dump(dict(my_mlflow_creds=dict(mlflow_tracking_uri="mlruns2")))
)
config = KedroMlflowConfig(
server=dict(
mlflow_tracking_uri="mlruns1",
credentials="my_mlflow_creds",
),
)
bootstrap_project(kedro_project_with_mlflow_conf)
with KedroSession.create(project_path=kedro_project_with_mlflow_conf) as session:
context = session.load_context()
config.setup(context)
assert (
mlflow.get_tracking_uri()
== (kedro_project_with_mlflow_conf / "mlruns1").as_uri()
)
# reset folder to avoid interference with other tests
(kedro_project_with_mlflow_conf / "conf/base/credentials.yml").write_text("")
def test_kedro_mlflow_config_setup_set_experiment_globally(
kedro_project_with_mlflow_conf,
):
mlflow_tracking_uri = (kedro_project_with_mlflow_conf / "mlruns").as_uri()
# the config must restore properly the experiment
config = KedroMlflowConfig(
server=dict(mlflow_tracking_uri="mlruns"),
tracking=dict(experiment=dict(name="incredible_exp")),
)
bootstrap_project(kedro_project_with_mlflow_conf)
with KedroSession.create(project_path=kedro_project_with_mlflow_conf) as session:
context = session.load_context() # setup config
config.setup(context)
mlflow_client = MlflowClient(mlflow_tracking_uri)
runs_list_before_interactive_run = mlflow_client.list_run_infos(
config.tracking.experiment._experiment.experiment_id
)
with mlflow.start_run():
mlflow.log_param("a", 1)
my_run_id = mlflow.active_run().info.run_id
runs_list_after_interactive_run = mlflow_client.list_run_infos(
config.tracking.experiment._experiment.experiment_id
)
assert (
len(runs_list_after_interactive_run) - len(runs_list_before_interactive_run)
== 1
)
assert runs_list_after_interactive_run[0].run_id == my_run_id
def test_ui_open_http_uri(monkeypatch, mocker, tmp_path):
config = {
"output_dir": tmp_path,
"kedro_version": kedro_version,
"project_name": "This is a fake project",
"repo_name": "fake-project-with-http-uri",
"python_package": "fake_project_with_http_uri",
"include_example": True,
}
cookiecutter(
str(TEMPLATE_PATH),
output_dir=config["output_dir"],
no_input=True,
extra_context=config,
)
project_path = tmp_path / config["repo_name"]
shutil.rmtree(project_path / "src" /
"tests") # avoid conflicts with pytest
mlflow_config = KedroMlflowConfig(server=dict(
mlflow_tracking_uri="http://google.com"))
with open((project_path / "conf" / "local" / "mlflow.yml").as_posix(),
"w") as fhandler:
yaml.dump(
mlflow_config.dict(),
fhandler,
default_flow_style=False,
)
monkeypatch.chdir(project_path.as_posix())
cli_runner = CliRunner()
# This does not test anything : the goal is to check whether it raises an error
# context_mocker = mocker.patch(
# "kedro.framework.session.session.KedroSession.load_context"
# )
mocker.patch(
"kedro_mlflow.config.kedro_mlflow_config.KedroMlflowConfig.setup")
open_mocker = mocker.patch(
"webbrowser.open") # make the test succeed, but no a real test
cli_runner.invoke(
cli_ui
) # result=cli_runner.invoke(cli_ui); print(result.exception) to debug
open_mocker.assert_called_once()
def test_kedro_mlflow_config_setup_set_tracking_uri(kedro_project_with_mlflow_conf):
mlflow_tracking_uri = (kedro_project_with_mlflow_conf / "awesome_tracking").as_uri()
config = KedroMlflowConfig(
server=dict(mlflow_tracking_uri="awesome_tracking"),
tracking=dict(experiment=dict(name="exp1")),
)
bootstrap_project(kedro_project_with_mlflow_conf)
with KedroSession.create(project_path=kedro_project_with_mlflow_conf) as session:
context = session.load_context() # setup config
config.setup(context)
assert mlflow.get_tracking_uri() == mlflow_tracking_uri
def after_context_created(
self,
context: KedroContext,
) -> None:
"""Hooks to be invoked after a `KedroContext` is created. This is the earliest
hook triggered within a Kedro run. The `KedroContext` stores useful information
such as `credentials`, `config_loader` and `env`.
Args:
context: The context that was created.
"""
try:
conf_mlflow_yml = context.config_loader.get(
"mlflow*", "mlflow*/**")
except MissingConfigException:
raise KedroMlflowConfigError(
"No 'mlflow.yml' config file found in environment. Use ``kedro mlflow init`` command in CLI to create a default config file."
)
mlflow_config = KedroMlflowConfig.parse_obj(conf_mlflow_yml)
mlflow_config.setup(context) # setup global mlflow configuration
# store in context for interactive use
# we use __setattr__ instead of context.mlflow because
# the class will become frozen in kedro>=0.19
context.__setattr__("mlflow", mlflow_config)
self.mlflow_config = mlflow_config # store for further reuse
def test_kedro_mlflow_config_new_experiment_does_not_exists(
kedro_project_with_mlflow_conf,
):
config = KedroMlflowConfig(
server=dict(mlflow_tracking_uri="mlruns"),
tracking=dict(experiment=dict(name="exp1")),
)
bootstrap_project(kedro_project_with_mlflow_conf)
with KedroSession.create(project_path=kedro_project_with_mlflow_conf) as session:
context = session.load_context() # setup config
config.setup(context)
assert "exp1" in [
exp.name for exp in config.server._mlflow_client.list_experiments()
]
def test_kedro_mlflow_config_setup_export_credentials(kedro_project_with_mlflow_conf):
(kedro_project_with_mlflow_conf / "conf/base/credentials.yml").write_text(
yaml.dump(dict(my_mlflow_creds=dict(fake_mlflow_cred="my_fake_cred")))
)
# the config must restore properly the experiment
config = KedroMlflowConfig(
server=dict(credentials="my_mlflow_creds"),
)
bootstrap_project(kedro_project_with_mlflow_conf)
with KedroSession.create(project_path=kedro_project_with_mlflow_conf) as session:
context = session.load_context() # setup config
config.setup(context)
assert os.environ["fake_mlflow_cred"] == "my_fake_cred"
def test_kedro_mlflow_config_experiment_exists(kedro_project_with_mlflow_conf):
# create an experiment with the same name
mlflow_tracking_uri = (
kedro_project_with_mlflow_conf / "conf" / "local" / "mlruns"
).as_uri()
MlflowClient(mlflow_tracking_uri).create_experiment("exp1")
config = KedroMlflowConfig(
server=dict(mlflow_tracking_uri="mlruns"),
tracking=dict(experiment=dict(name="exp1")),
)
bootstrap_project(kedro_project_with_mlflow_conf)
with KedroSession.create(project_path=kedro_project_with_mlflow_conf) as session:
context = session.load_context() # setup config
config.setup(context)
assert "exp1" in [
exp.name for exp in config.server._mlflow_client.list_experiments()
]
def test_kedro_mlflow_config_experiment_was_deleted(kedro_project_with_mlflow_conf):
# create an experiment with the same name and then delete it
mlflow_tracking_uri = (kedro_project_with_mlflow_conf / "mlruns").as_uri()
mlflow_client = MlflowClient(mlflow_tracking_uri)
mlflow_client.create_experiment("exp1")
mlflow_client.delete_experiment(
mlflow_client.get_experiment_by_name("exp1").experiment_id
)
# the config must restore properly the experiment
config = KedroMlflowConfig(
server=dict(mlflow_tracking_uri="mlruns"),
tracking=dict(experiment=dict(name="exp1")),
)
bootstrap_project(kedro_project_with_mlflow_conf)
with KedroSession.create(project_path=kedro_project_with_mlflow_conf) as session:
context = session.load_context() # setup config
config.setup(context)
assert "exp1" in [
exp.name for exp in config.server._mlflow_client.list_experiments()
]
def test_kedro_mlflow_config_init():
# kedro_project_with_mlflow_conf is a global fixture in conftest
config = KedroMlflowConfig()
assert config.dict() == dict(
server=dict(
mlflow_tracking_uri=None, # not setup, not modified yet
credentials=None,
),
tracking=dict(
disable_tracking=dict(pipelines=[]),
experiment=dict(name="Default", restore_if_deleted=True),
run=dict(id=None, name=None, nested=True),
params=dict(
dict_params=dict(
flatten=False,
recursive=True,
sep=".",
),
long_params_strategy="fail",
),
),
ui=dict(port="5000", host="127.0.0.1"),
)
def test_mlflow_yml_rendering(template_mlflowyml):
# the mlflow yml file must be consistent with the default in KedroMlflowConfig for readibility
with open(template_mlflowyml, "r") as file_handler:
mlflow_config = yaml.safe_load(file_handler)
# note: Using Pydantic model Construct method skip all validations
# and here we do not want to check the path
expected_config = KedroMlflowConfig.construct(
project_path="fake/path",
tracking=dict(
disable_tracking=dict(pipelines=[]),
experiment=dict(name="fake_project", restore_if_deleted=True),
params=dict(
dict_params=dict(flatten=False, recursive=True, sep="."),
long_params_strategy="fail",
),
run=dict(id=None, name=None, nested=True),
), # check for proper rendering
)
assert mlflow_config == expected_config.dict(exclude={"project_path"})
def test_from_dict_to_dict_idempotent(kedro_project_with_mlflow_conf):
config = KedroMlflowConfig()
original_config_dict = config.dict()
# modify config
reloaded_config = KedroMlflowConfig.parse_obj(original_config_dict)
assert config == reloaded_config