ray-project--ray/python/ray/air/tests/test_integration_mlflow.py

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__]))