move to conftest

mosaicml · Nov 19, 2024 · 902e4e0 · 902e4e0
1 parent 2ed8a4d
commit 902e4e0
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 25 deletions.
diff --git a/tests/callbacks/test_callbacks.py b/tests/callbacks/test_callbacks.py
@@ -20,14 +20,6 @@
 from tests.common import EventCounterCallback
 
 
-@pytest.fixture(autouse=True)
-def setup_mlflow_tracking(monkeypatch, tmp_path):
-    mlflow = pytest.importorskip('mlflow')
-    tracking_uri = str(tmp_path / 'mlruns')
-    monkeypatch.setenv(mlflow.environment_variables.MLFLOW_TRACKING_URI.name, tracking_uri)
-    os.makedirs(tracking_uri, exist_ok=True)
-
-
 def test_callbacks_map_to_events():
     # callback methods must be 1:1 mapping with events
     # exception for private methods

diff --git a/tests/callbacks/test_loggers_across_callbacks.py b/tests/callbacks/test_loggers_across_callbacks.py
@@ -17,14 +17,6 @@
 )
 
 
-@pytest.fixture(autouse=True)
-def setup_mlflow_tracking(monkeypatch, tmp_path):
-    mlflow = pytest.importorskip('mlflow')
-    tracking_uri = str(tmp_path / 'mlruns')
-    monkeypatch.setenv(mlflow.environment_variables.MLFLOW_TRACKING_URI.name, tracking_uri)
-    os.makedirs(tracking_uri, exist_ok=True)
-
-
 @pytest.mark.parametrize('logger_cls', get_cbs_and_marks(loggers=True))
 @pytest.mark.parametrize('callback_cls', get_cbs_and_marks(callbacks=True))
 @pytest.mark.filterwarnings('ignore::UserWarning')

diff --git a/tests/fixtures/autouse_fixtures.py b/tests/fixtures/autouse_fixtures.py
@@ -148,3 +148,11 @@ def remove_run_name_env_var():
         os.environ['COMPOSER_RUN_NAME'] = composer_run_name
     if run_name is not None:
         os.environ['RUN_NAME'] = run_name
+
+
+@pytest.fixture(autouse=True)
+def setup_mlflow_tracking(monkeypatch, tmp_path):
+    mlflow = pytest.importorskip('mlflow')
+    tracking_uri = str(tmp_path / 'mlruns')
+    monkeypatch.setenv(mlflow.environment_variables.MLFLOW_TRACKING_URI.name, tracking_uri)
+    os.makedirs(tracking_uri, exist_ok=True)
diff --git a/tests/loggers/test_mlflow_logger.py b/tests/loggers/test_mlflow_logger.py
@@ -25,14 +25,6 @@
 )
 
 
-@pytest.fixture(autouse=True)
-def setup_mlflow_tracking(monkeypatch, tmp_path):
-    mlflow = pytest.importorskip('mlflow')
-    tracking_uri = str(tmp_path / 'mlruns')
-    monkeypatch.setenv(mlflow.environment_variables.MLFLOW_TRACKING_URI.name, tracking_uri)
-    os.makedirs(tracking_uri, exist_ok=True)
-
-
 def _get_latest_mlflow_run(experiment_name, tracking_uri=None):
     pytest.importorskip('mlflow')
     from mlflow import MlflowClient

diff --git a/tests/utils/test_dist.py b/tests/utils/test_dist.py
@@ -71,6 +71,6 @@ def test_busy_wait_for_local_rank_zero(tmp_path):
     end_time = time.time()
     total_time = end_time - start_time
     gathered_times = dist.all_gather_object(total_time)
-    assert os.listdir(gathered_tmp_path) == []
+    assert os.listdir(gathered_tmp_path) == ['mlruns']
     assert len(gathered_times) == 2
     assert abs(gathered_times[0] - gathered_times[1]) < 0.1