diff --git a/python/ray/air/tests/test_integration_mlflow.py b/python/ray/air/tests/test_integration_mlflow.py index 143321bdf68a4..b0f4b400dab1b 100644 --- a/python/ray/air/tests/test_integration_mlflow.py +++ b/python/ray/air/tests/test_integration_mlflow.py @@ -42,14 +42,16 @@ def clear_env_vars(): class MLflowTest(unittest.TestCase): def setUp(self): - self.tracking_uri = tempfile.mkdtemp() - self.registry_uri = tempfile.mkdtemp() + 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") - assert client.get_experiment_by_name("existing_experiment").experiment_id == "0" + # 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: # Remove tune session if initialized to clean up for next test @@ -70,10 +72,10 @@ def testMlFlowLoggerCallbackConfig(self): logger.mlflow_util._mlflow.get_registry_uri(), self.registry_uri ) self.assertListEqual( - [e.name for e in logger.mlflow_util._mlflow.list_experiments()], - ["existing_experiment", "test_exp"], + [e.name for e in logger.mlflow_util._mlflow.search_experiments()], + ["test_exp", "existing_experiment", "Default"], ) - self.assertEqual(logger.mlflow_util.experiment_id, "1") + self.assertEqual(logger.mlflow_util.experiment_id, "2") # Check if client recognizes already existing experiment. logger = MLflowLoggerCallback( @@ -82,7 +84,7 @@ def testMlFlowLoggerCallbackConfig(self): registry_uri=self.registry_uri, ) logger.setup() - self.assertEqual(logger.mlflow_util.experiment_id, "0") + self.assertEqual(logger.mlflow_util.experiment_id, "1") # Pass in experiment name as env var. clear_env_vars() @@ -91,7 +93,7 @@ def testMlFlowLoggerCallbackConfig(self): tracking_uri=self.tracking_uri, registry_uri=self.registry_uri ) logger.setup() - self.assertEqual(logger.mlflow_util.experiment_id, "1") + self.assertEqual(logger.mlflow_util.experiment_id, "2") # Pass in existing experiment name as env var. clear_env_vars() @@ -100,16 +102,16 @@ def testMlFlowLoggerCallbackConfig(self): tracking_uri=self.tracking_uri, registry_uri=self.registry_uri ) logger.setup() - self.assertEqual(logger.mlflow_util.experiment_id, "0") + self.assertEqual(logger.mlflow_util.experiment_id, "1") # Pass in existing experiment id as env var. clear_env_vars() - os.environ["MLFLOW_EXPERIMENT_ID"] = "0" + 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, "0") + self.assertEqual(logger.mlflow_util.experiment_id, "1") # Pass in non existing experiment id as env var. # This should create a new experiment. @@ -124,18 +126,18 @@ def testMlFlowLoggerCallbackConfig(self): # 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"] = "0" + 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, "0") + 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"] = "0" + os.environ["MLFLOW_EXPERIMENT_ID"] = "1" logger = MLflowLoggerCallback( tracking_uri=self.tracking_uri, registry_uri=self.registry_uri, tags=tags ) @@ -159,7 +161,7 @@ def testMlFlowLoggerLogging(self): # 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=["1"]) + 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") @@ -174,7 +176,7 @@ def testMlFlowLoggerLogging(self): # 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=["1"]) + all_runs = logger.mlflow_util._mlflow.search_runs(experiment_ids=["2"]) self.assertEqual(len(all_runs), 1) # Check metrics are logged properly. @@ -278,7 +280,7 @@ def setUp(self): self.dirpath = tempfile.mkdtemp() import mlflow - mlflow.set_tracking_uri(self.dirpath) + mlflow.set_tracking_uri("sqlite:///" + self.dirpath + "/mlflow.sqlite") mlflow.create_experiment(name="existing_experiment") self.mlflow_util = _MLflowLoggerUtil() @@ -301,7 +303,7 @@ 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 == "0" + assert self.mlflow_util.experiment_id == "1" def test_run_started_with_correct_experiment(self): experiment_name = "my_experiment_name" @@ -322,7 +324,7 @@ def test_run_started_with_correct_experiment(self): 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 == "0" + assert self.mlflow_util.experiment_id == "1" del os.environ["MLFLOW_EXPERIMENT_NAME"] def test_id_precedence(self): @@ -337,7 +339,7 @@ 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 == "1" + assert self.mlflow_util.experiment_id == "2" def test_setup_fail(self): with self.assertRaises(ValueError): diff --git a/python/requirements/ml/requirements_train.txt b/python/requirements/ml/requirements_train.txt index c633ddbecf610..4a6ac751c6212 100644 --- a/python/requirements/ml/requirements_train.txt +++ b/python/requirements/ml/requirements_train.txt @@ -3,8 +3,7 @@ -r requirements_dl.txt mlflow==1.30.0; python_version <= '3.7' -# Todo (krfricke): Re-upgrade to 2.x -mlflow==1.30.0; python_version > '3.7' +mlflow==2.4.1; python_version > '3.7' tensorboardX==2.4.1 # Dependencies for Hugging Face examples & tests: diff --git a/python/requirements/ml/requirements_tune.txt b/python/requirements/ml/requirements_tune.txt index 04f2afb2a10f6..c1ab3924f489a 100644 --- a/python/requirements/ml/requirements_tune.txt +++ b/python/requirements/ml/requirements_tune.txt @@ -26,8 +26,7 @@ jupyterlab==3.6.1 lightgbm==3.3.5 matplotlib!=3.4.3 mlflow==1.30.0; python_version <= '3.7' -# Todo (krfricke): Re-upgrade to 2.x -mlflow==1.30.0; python_version > '3.7' +mlflow==2.4.1; python_version > '3.7' # Unavailable for arm64 in more recent versions mxnet==1.8.0.post0; platform_machine != "arm64" nevergrad==0.4.3.post7