chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,442 @@
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from collections import namedtuple
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from mlflow.tracking import MlflowClient
|
||||
|
||||
import ray
|
||||
from ray._private.dict import flatten_dict
|
||||
from ray.air._internal.mlflow import _MLflowLoggerUtil
|
||||
from ray.air.integrations.mlflow import MLflowLoggerCallback, _NoopModule, setup_mlflow
|
||||
from ray.train.torch import TorchTrainer
|
||||
from ray.tune import Tuner
|
||||
|
||||
|
||||
class MockTrial(
|
||||
namedtuple("MockTrial", ["config", "trial_name", "trial_id", "local_path"])
|
||||
):
|
||||
def __hash__(self):
|
||||
return hash(self.trial_id)
|
||||
|
||||
def __str__(self):
|
||||
return self.trial_name
|
||||
|
||||
|
||||
class Mock_MLflowLoggerUtil(_MLflowLoggerUtil):
|
||||
def save_artifacts(self, dir, run_id):
|
||||
self.artifact_saved = True
|
||||
self.artifact_info = {"dir": dir, "run_id": run_id}
|
||||
|
||||
|
||||
def clear_env_vars():
|
||||
os.environ.pop("MLFLOW_EXPERIMENT_NAME", None)
|
||||
os.environ.pop("MLFLOW_EXPERIMENT_ID", None)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_start_4_cpus():
|
||||
"""Automatically start and stop Ray for each test."""
|
||||
ray.init(num_cpus=4)
|
||||
yield
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
def test_setup_mlflow_in_train_worker(ray_start_4_cpus):
|
||||
"""Test that setup_mlflow works in a Train worker."""
|
||||
|
||||
def train_func(config):
|
||||
setup_mlflow(
|
||||
experiment_name="test_exp",
|
||||
create_experiment_if_not_exists=True,
|
||||
)
|
||||
|
||||
trainer = TorchTrainer(train_func)
|
||||
trainer.fit()
|
||||
|
||||
|
||||
def test_setup_mlflow_in_tune_trial(ray_start_4_cpus):
|
||||
"""Test that setup_mlflow works in a Tune trial."""
|
||||
|
||||
def train_func(config):
|
||||
setup_mlflow(
|
||||
experiment_name="test_exp",
|
||||
create_experiment_if_not_exists=True,
|
||||
)
|
||||
|
||||
tuner = Tuner(train_func)
|
||||
result_grid = tuner.fit()
|
||||
|
||||
assert all(res.error is None for res in result_grid)
|
||||
|
||||
|
||||
class MLflowTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tracking_uri = "sqlite:///" + tempfile.mkdtemp() + "/mlflow.sqlite"
|
||||
self.registry_uri = "sqlite:///" + tempfile.mkdtemp() + "/mlflow.sqlite"
|
||||
|
||||
client = MlflowClient(
|
||||
tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
|
||||
)
|
||||
client.create_experiment(name="existing_experiment")
|
||||
# Mlflow > 2 creates a "Default" experiment which has ID 0, so we start our
|
||||
# test with ID 1.
|
||||
assert client.get_experiment_by_name("existing_experiment").experiment_id == "1"
|
||||
|
||||
def tearDown(self) -> None:
|
||||
pass
|
||||
|
||||
def testMlFlowLoggerCallbackConfig(self):
|
||||
# Explicitly pass in all args.
|
||||
logger = MLflowLoggerCallback(
|
||||
tracking_uri=self.tracking_uri,
|
||||
registry_uri=self.registry_uri,
|
||||
experiment_name="test_exp",
|
||||
)
|
||||
logger.setup()
|
||||
self.assertEqual(
|
||||
logger.mlflow_util._mlflow.get_tracking_uri(), self.tracking_uri
|
||||
)
|
||||
self.assertEqual(
|
||||
logger.mlflow_util._mlflow.get_registry_uri(), self.registry_uri
|
||||
)
|
||||
self.assertListEqual(
|
||||
[e.name for e in logger.mlflow_util._mlflow.search_experiments()],
|
||||
["test_exp", "existing_experiment", "Default"],
|
||||
)
|
||||
self.assertEqual(logger.mlflow_util.experiment_id, "2")
|
||||
|
||||
# Check if client recognizes already existing experiment.
|
||||
logger = MLflowLoggerCallback(
|
||||
experiment_name="existing_experiment",
|
||||
tracking_uri=self.tracking_uri,
|
||||
registry_uri=self.registry_uri,
|
||||
)
|
||||
logger.setup()
|
||||
self.assertEqual(logger.mlflow_util.experiment_id, "1")
|
||||
|
||||
# Pass in experiment name as env var.
|
||||
clear_env_vars()
|
||||
os.environ["MLFLOW_EXPERIMENT_NAME"] = "test_exp"
|
||||
logger = MLflowLoggerCallback(
|
||||
tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
|
||||
)
|
||||
logger.setup()
|
||||
self.assertEqual(logger.mlflow_util.experiment_id, "2")
|
||||
|
||||
# Pass in existing experiment name as env var.
|
||||
clear_env_vars()
|
||||
os.environ["MLFLOW_EXPERIMENT_NAME"] = "existing_experiment"
|
||||
logger = MLflowLoggerCallback(
|
||||
tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
|
||||
)
|
||||
logger.setup()
|
||||
self.assertEqual(logger.mlflow_util.experiment_id, "1")
|
||||
|
||||
# Pass in existing experiment id as env var.
|
||||
clear_env_vars()
|
||||
os.environ["MLFLOW_EXPERIMENT_ID"] = "1"
|
||||
logger = MLflowLoggerCallback(
|
||||
tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
|
||||
)
|
||||
logger.setup()
|
||||
self.assertEqual(logger.mlflow_util.experiment_id, "1")
|
||||
|
||||
# Pass in non existing experiment id as env var.
|
||||
# This should create a new experiment.
|
||||
clear_env_vars()
|
||||
os.environ["MLFLOW_EXPERIMENT_ID"] = "500"
|
||||
with self.assertRaises(ValueError):
|
||||
logger = MLflowLoggerCallback(
|
||||
tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
|
||||
)
|
||||
logger.setup()
|
||||
|
||||
# Experiment id env var should take precedence over name env var.
|
||||
clear_env_vars()
|
||||
os.environ["MLFLOW_EXPERIMENT_NAME"] = "test_exp"
|
||||
os.environ["MLFLOW_EXPERIMENT_ID"] = "1"
|
||||
logger = MLflowLoggerCallback(
|
||||
tracking_uri=self.tracking_uri, registry_uri=self.registry_uri
|
||||
)
|
||||
logger.setup()
|
||||
self.assertEqual(logger.mlflow_util.experiment_id, "1")
|
||||
|
||||
# Using tags
|
||||
tags = {"user_name": "John", "git_commit_hash": "abc123"}
|
||||
clear_env_vars()
|
||||
os.environ["MLFLOW_EXPERIMENT_NAME"] = "test_tags"
|
||||
os.environ["MLFLOW_EXPERIMENT_ID"] = "1"
|
||||
logger = MLflowLoggerCallback(
|
||||
tracking_uri=self.tracking_uri, registry_uri=self.registry_uri, tags=tags
|
||||
)
|
||||
logger.setup()
|
||||
self.assertEqual(logger.tags, tags)
|
||||
|
||||
@patch("ray.air.integrations.mlflow._MLflowLoggerUtil", Mock_MLflowLoggerUtil)
|
||||
def testMlFlowLoggerLogging(self):
|
||||
clear_env_vars()
|
||||
trial_config = {"par1": "a", "par2": "b"}
|
||||
trial = MockTrial(trial_config, "trial1", 0, "artifact")
|
||||
|
||||
logger = MLflowLoggerCallback(
|
||||
tracking_uri=self.tracking_uri,
|
||||
registry_uri=self.registry_uri,
|
||||
experiment_name="test1",
|
||||
save_artifact=True,
|
||||
tags={"hello": "world"},
|
||||
)
|
||||
logger.setup()
|
||||
|
||||
# Check if run is created with proper tags.
|
||||
logger.on_trial_start(iteration=0, trials=[], trial=trial)
|
||||
all_runs = logger.mlflow_util._mlflow.search_runs(experiment_ids=["2"])
|
||||
self.assertEqual(len(all_runs), 1)
|
||||
# all_runs is a pandas dataframe.
|
||||
all_runs = all_runs.to_dict(orient="records")
|
||||
run = logger.mlflow_util._mlflow.get_run(all_runs[0]["run_id"])
|
||||
self.assertDictEqual(
|
||||
run.data.tags,
|
||||
{"hello": "world", "trial_name": "trial1", "mlflow.runName": "trial1"},
|
||||
)
|
||||
self.assertEqual(logger._trial_runs[trial], run.info.run_id)
|
||||
# Params should be logged.
|
||||
self.assertDictEqual(run.data.params, trial_config)
|
||||
|
||||
# When same trial is started again, new run should not be created.
|
||||
logger.on_trial_start(iteration=0, trials=[], trial=trial)
|
||||
all_runs = logger.mlflow_util._mlflow.search_runs(experiment_ids=["2"])
|
||||
self.assertEqual(len(all_runs), 1)
|
||||
|
||||
# Check metrics are logged properly.
|
||||
result = {
|
||||
"metric1": 0.8,
|
||||
"metric2": 1,
|
||||
"metric3": None,
|
||||
"training_iteration": 0,
|
||||
}
|
||||
logger.on_trial_result(0, [], trial, result)
|
||||
run = logger.mlflow_util._mlflow.get_run(run_id=run.info.run_id)
|
||||
# metric3 is not logged since it cannot be converted to float.
|
||||
self.assertDictEqual(
|
||||
run.data.metrics, {"metric1": 0.8, "metric2": 1.0, "training_iteration": 0}
|
||||
)
|
||||
|
||||
# Check that artifact is logged on termination.
|
||||
logger.on_trial_complete(0, [], trial)
|
||||
self.assertTrue(logger.mlflow_util.artifact_saved)
|
||||
self.assertDictEqual(
|
||||
logger.mlflow_util.artifact_info,
|
||||
{"dir": "artifact", "run_id": run.info.run_id},
|
||||
)
|
||||
|
||||
# Check if params are logged at the end.
|
||||
run = logger.mlflow_util._mlflow.get_run(run_id=run.info.run_id)
|
||||
self.assertDictEqual(run.data.params, trial_config)
|
||||
|
||||
@patch("ray.air.integrations.mlflow._MLflowLoggerUtil", Mock_MLflowLoggerUtil)
|
||||
def testMlFlowLoggerLogging_logAtEnd(self):
|
||||
clear_env_vars()
|
||||
trial_config = {"par1": "a", "par2": "b"}
|
||||
trial = MockTrial(trial_config, "trial1", 0, "artifact")
|
||||
|
||||
logger = MLflowLoggerCallback(
|
||||
tracking_uri=self.tracking_uri,
|
||||
registry_uri=self.registry_uri,
|
||||
experiment_name="test_log_at_end",
|
||||
tags={"hello": "world"},
|
||||
log_params_on_trial_end=True,
|
||||
)
|
||||
logger.setup()
|
||||
exp_id = logger.mlflow_util.experiment_id
|
||||
|
||||
logger.on_trial_start(iteration=0, trials=[], trial=trial)
|
||||
all_runs = logger.mlflow_util._mlflow.search_runs(experiment_ids=[exp_id])
|
||||
self.assertEqual(len(all_runs), 1)
|
||||
# all_runs is a pandas dataframe.
|
||||
all_runs = all_runs.to_dict(orient="records")
|
||||
run = logger.mlflow_util._mlflow.get_run(all_runs[0]["run_id"])
|
||||
|
||||
# Params should NOT be logged at start.
|
||||
self.assertDictEqual(run.data.params, {})
|
||||
|
||||
# Check that params are logged at the end.
|
||||
logger.on_trial_complete(0, [], trial)
|
||||
run = logger.mlflow_util._mlflow.get_run(run_id=run.info.run_id)
|
||||
self.assertDictEqual(run.data.params, trial_config)
|
||||
|
||||
def testMlFlowSetupExplicit(self):
|
||||
clear_env_vars()
|
||||
trial_config = {"par1": 4, "par2": 9.0}
|
||||
|
||||
# No MLflow config passed in.
|
||||
with self.assertRaises(ValueError):
|
||||
setup_mlflow(trial_config)
|
||||
|
||||
# Invalid experiment-id
|
||||
with self.assertRaises(ValueError):
|
||||
setup_mlflow(trial_config, experiment_id="500")
|
||||
|
||||
# Set to experiment that does not already exist.
|
||||
with self.assertRaises(ValueError):
|
||||
setup_mlflow(
|
||||
trial_config,
|
||||
experiment_id="500",
|
||||
experiment_name="new_experiment",
|
||||
tracking_uri=self.tracking_uri,
|
||||
)
|
||||
|
||||
mlflow = setup_mlflow(
|
||||
trial_config,
|
||||
experiment_id="500",
|
||||
experiment_name="existing_experiment",
|
||||
tracking_uri=self.tracking_uri,
|
||||
)
|
||||
mlflow.end_run()
|
||||
|
||||
@patch("ray.train.get_context")
|
||||
def testMlFlowSetupRankNonRankZero(self, mock_get_context):
|
||||
"""Assert that non-rank-0 workers get a noop module"""
|
||||
mock_context = MagicMock()
|
||||
mock_context.get_world_rank.return_value = 1
|
||||
|
||||
mock_get_context.return_value = mock_context
|
||||
|
||||
mlflow = setup_mlflow({})
|
||||
assert isinstance(mlflow, _NoopModule)
|
||||
|
||||
mlflow.log_metrics()
|
||||
mlflow.sklearn.save_model(None, "model_directory")
|
||||
|
||||
|
||||
class MLflowUtilTest(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.dirpath = tempfile.mkdtemp()
|
||||
import mlflow
|
||||
|
||||
mlflow.set_tracking_uri("sqlite:///" + self.dirpath + "/mlflow.sqlite")
|
||||
mlflow.create_experiment(name="existing_experiment")
|
||||
|
||||
self.mlflow_util = _MLflowLoggerUtil()
|
||||
self.tracking_uri = mlflow.get_tracking_uri()
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.dirpath)
|
||||
|
||||
def test_experiment_id(self):
|
||||
self.mlflow_util.setup_mlflow(tracking_uri=self.tracking_uri, experiment_id="0")
|
||||
assert self.mlflow_util.experiment_id == "0"
|
||||
|
||||
def test_experiment_id_env_var(self):
|
||||
os.environ["MLFLOW_EXPERIMENT_ID"] = "0"
|
||||
self.mlflow_util.setup_mlflow(tracking_uri=self.tracking_uri)
|
||||
assert self.mlflow_util.experiment_id == "0"
|
||||
del os.environ["MLFLOW_EXPERIMENT_ID"]
|
||||
|
||||
def test_experiment_name(self):
|
||||
self.mlflow_util.setup_mlflow(
|
||||
tracking_uri=self.tracking_uri, experiment_name="existing_experiment"
|
||||
)
|
||||
assert self.mlflow_util.experiment_id == "1"
|
||||
|
||||
def test_run_started_with_correct_experiment(self):
|
||||
experiment_name = "my_experiment_name"
|
||||
# Make sure run is started under the correct experiment.
|
||||
self.mlflow_util.setup_mlflow(
|
||||
tracking_uri=self.tracking_uri, experiment_name=experiment_name
|
||||
)
|
||||
run = self.mlflow_util.start_run(set_active=True)
|
||||
assert (
|
||||
run.info.experiment_id
|
||||
== self.mlflow_util._mlflow.get_experiment_by_name(
|
||||
experiment_name
|
||||
).experiment_id
|
||||
)
|
||||
|
||||
self.mlflow_util.end_run()
|
||||
|
||||
def test_experiment_name_env_var(self):
|
||||
os.environ["MLFLOW_EXPERIMENT_NAME"] = "existing_experiment"
|
||||
self.mlflow_util.setup_mlflow(tracking_uri=self.tracking_uri)
|
||||
assert self.mlflow_util.experiment_id == "1"
|
||||
del os.environ["MLFLOW_EXPERIMENT_NAME"]
|
||||
|
||||
def test_id_precedence(self):
|
||||
os.environ["MLFLOW_EXPERIMENT_ID"] = "0"
|
||||
self.mlflow_util.setup_mlflow(
|
||||
tracking_uri=self.tracking_uri, experiment_name="new_experiment"
|
||||
)
|
||||
assert self.mlflow_util.experiment_id == "0"
|
||||
del os.environ["MLFLOW_EXPERIMENT_ID"]
|
||||
|
||||
def test_new_experiment(self):
|
||||
self.mlflow_util.setup_mlflow(
|
||||
tracking_uri=self.tracking_uri, experiment_name="new_experiment"
|
||||
)
|
||||
assert self.mlflow_util.experiment_id == "2"
|
||||
|
||||
def test_setup_fail(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.mlflow_util.setup_mlflow(
|
||||
tracking_uri=self.tracking_uri,
|
||||
experiment_name="new_experiment2",
|
||||
create_experiment_if_not_exists=False,
|
||||
)
|
||||
|
||||
def test_log_params(self):
|
||||
params = {"a": "a", "x": {"y": "z"}}
|
||||
self.mlflow_util.setup_mlflow(
|
||||
tracking_uri=self.tracking_uri, experiment_name="new_experiment"
|
||||
)
|
||||
run = self.mlflow_util.start_run()
|
||||
run_id = run.info.run_id
|
||||
self.mlflow_util.log_params(params_to_log=params, run_id=run_id)
|
||||
|
||||
run = self.mlflow_util._mlflow.get_run(run_id=run_id)
|
||||
assert run.data.params == flatten_dict(params)
|
||||
|
||||
params2 = {"b": "b"}
|
||||
self.mlflow_util.start_run(set_active=True)
|
||||
self.mlflow_util.log_params(params_to_log=params2, run_id=run_id)
|
||||
run = self.mlflow_util._mlflow.get_run(run_id=run_id)
|
||||
assert run.data.params == flatten_dict(
|
||||
{
|
||||
**params,
|
||||
**params2,
|
||||
}
|
||||
)
|
||||
|
||||
self.mlflow_util.end_run()
|
||||
|
||||
def test_log_metrics(self):
|
||||
metrics = {"a": 1.0, "x": {"y": 2.0}}
|
||||
self.mlflow_util.setup_mlflow(
|
||||
tracking_uri=self.tracking_uri, experiment_name="new_experiment"
|
||||
)
|
||||
run = self.mlflow_util.start_run()
|
||||
run_id = run.info.run_id
|
||||
self.mlflow_util.log_metrics(metrics_to_log=metrics, run_id=run_id, step=0)
|
||||
|
||||
run = self.mlflow_util._mlflow.get_run(run_id=run_id)
|
||||
assert run.data.metrics == flatten_dict(metrics)
|
||||
|
||||
metrics2 = {"b": 1.0}
|
||||
self.mlflow_util.start_run(set_active=True)
|
||||
self.mlflow_util.log_metrics(metrics_to_log=metrics2, run_id=run_id, step=0)
|
||||
assert self.mlflow_util._mlflow.get_run(
|
||||
run_id=run_id
|
||||
).data.metrics == flatten_dict(
|
||||
{
|
||||
**metrics,
|
||||
**metrics2,
|
||||
}
|
||||
)
|
||||
self.mlflow_util.end_run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Reference in New Issue
Block a user