fix test

Signed-off-by: Andrew Xue <andewzxue@gmail.com>
ray-project · Nov 8, 2023 · dd4f333 · dd4f333
1 parent b2fab9d
commit dd4f333
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/python/ray/data/_internal/stats.py b/python/ray/data/_internal/stats.py
@@ -419,13 +419,12 @@ def _run_iter_update_loop():
             self._iter_update_thread.start()
 
     def clear_stats_actor_iter_metrics(self, tags: Dict[str, str]):
-        if self._iter_update_thread is not None:
-            # Stop background update thread.
-            self._collect_iter_stats = False
-            with self._iter_update_condition:
-                self._iter_update_condition.notify()
-            self._iter_update_thread.join()
-            self._iter_update_thread = None
+        # Stop background update thread.
+        self._collect_iter_stats = False
+        with self._iter_update_condition:
+            self._iter_update_condition.notify()
+        self._iter_update_thread.join()
+        self._iter_update_thread = None
 
         # Clear metrics from _StatsActor.
         global _stats_actor

diff --git a/python/ray/data/tests/test_stats.py b/python/ray/data/tests/test_stats.py
@@ -156,7 +156,9 @@ def patch_update_stats_actor():
 def patch_update_stats_actor_iter():
     with patch(
         "ray.data._internal.stats.StatsActorManager.update_stats_actor_iter_metrics"
-    ) as update_fn:
+    ) as update_fn, patch(
+        "ray.data._internal.stats.StatsActorManager.clear_stats_actor_iter_metrics"
+    ):
         yield update_fn