Use mlflow.flush_async_loggin

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
mlflow · Feb 27, 2024 · 4b9a2ab · 4b9a2ab
1 parent e02987f
commit 4b9a2ab
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/transformers/test_transformers_autolog.py b/tests/transformers/test_transformers_autolog.py
@@ -1,5 +1,4 @@
 import random
-import time
 
 import numpy as np
 import optuna
@@ -504,7 +503,7 @@ def test_trainer_hyperparameter_tuning_does_not_log_sklearn_model(
     exp = mlflow.set_experiment(experiment_name="hyperparam_trainer")
 
     transformers_hyperparameter_trainer.train()
-    time.sleep(5)  # wait for async logging to finish
+    mlflow.flush_async_logging()
 
     last_run = mlflow.last_active_run()
     assert last_run.data.metrics["epoch"] == 3.0
@@ -531,7 +530,7 @@ def test_trainer_hyperparameter_tuning_functional_does_not_log_sklearn_model(
     exp = mlflow.set_experiment(experiment_name="hyperparam_trainer_functional")
 
     transformers_hyperparameter_functional.train()
-    time.sleep(5)  # wait for async logging to finish
+    mlflow.flush_async_logging()
 
     last_run = mlflow.last_active_run()
     assert last_run.data.metrics["epoch"] == 1.0