go

mosaicml · Mar 6, 2024 · e41dca6 · e41dca6
1 parent ed1cd04
commit e41dca6
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 7 deletions.
diff --git a/composer/loggers/mlflow_logger.py b/composer/loggers/mlflow_logger.py
@@ -437,6 +437,7 @@ def post_close(self):
             import mlflow
 
             assert isinstance(self._run_id, str)
+            mlflow.flush_async_logging()
             self._mlflow_client.set_terminated(self._run_id)
             mlflow.end_run()
 

diff --git a/setup.py b/setup.py
@@ -219,7 +219,7 @@ def package_files(prefix: str, directory: str, extension: str):
 ]
 
 extra_deps['mlflow'] = [
-    'mlflow>=2.9.2,<3.0',
+    'mlflow>=2.11.0,<3.0',
 ]
 
 extra_deps['pandas'] = ['pandas>=2.0.0,<3.0']

diff --git a/tests/loggers/test_mlflow_logger.py b/tests/loggers/test_mlflow_logger.py
@@ -587,8 +587,8 @@ def test_mlflow_logging_works(tmp_path, device):
                       eval_interval=eval_interval,
                       device=device)
     trainer.fit()
-    # Allow async logging to finish.
-    time.sleep(3)
+    # Allow system metrics to be collected.
+    time.sleep(2)
     test_mlflow_logger.post_close()
 
     run = _get_latest_mlflow_run(
@@ -707,8 +707,6 @@ def test_mlflow_ignore_metrics(tmp_path, device):
                       eval_interval=eval_interval,
                       device=device)
     trainer.fit()
-    # Allow async logging to finish.
-    time.sleep(3)
     test_mlflow_logger.post_close()
 
     run = _get_latest_mlflow_run(
@@ -751,8 +749,6 @@ def test_mlflow_ignore_hyperparameters(tmp_path):
                                       ignore_hyperparameters=['num*', 'mlflow_run_id', 'nothing'])
 
     Trainer(model=SimpleConvModel(), loggers=test_mlflow_logger, max_duration=f'4ba')
-    # Allow async logging to finish.
-    time.sleep(3)
     test_mlflow_logger.post_close()
 
     run = _get_latest_mlflow_run(