lint

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

Author: iamjustinhsu

python/ray/data/_internal/block_batching/iter_batches.py

@@ -1,4 +1,5 @@
 import collections
+import time
 from contextlib import contextmanager, nullcontext
 from typing import Any, Callable, Dict, Iterator, Optional
 
@@ -15,7 +16,7 @@
 )
 from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
 from ray.data._internal.memory_tracing import trace_deallocation
-from ray.data._internal.stats import DatasetStats, StatsManager
+from ray.data._internal.stats import DatasetStats, _StatsManager
 from ray.data._internal.util import make_async_gen
 from ray.data.block import Block, DataBatch
 from ray.data.context import DataContext
@@ -92,6 +93,8 @@ class BatchIterator:
             formatting to be overlapped with the UDF. Defaults to 1.
     """
 
+    UPDATE_METRICS_INTERVAL_S: float = 5.0
+
     def __init__(
         self,
         ref_bundles: Iterator[RefBundle],
@@ -136,6 +139,7 @@ def __init__(
             else WaitBlockPrefetcher()
         )
         self._yielded_first_batch = False
+        self._metrics_last_updated: float = 0.0
 
     def _prefetch_blocks(
         self, ref_bundles: Iterator[RefBundle]
@@ -219,34 +223,29 @@ def _iter_batches(self) -> Iterator[DataBatch]:
             preserve_ordering=False,
         )
 
-        with self._epoch_context():
-            while True:
-                with self.get_next_batch_context():
-                    try:
-                        batch = next(async_batch_iter)
-                    except StopIteration:
-                        break
-                with self.yield_batch_context(batch):
-                    yield batch.data
+        self.before_epoch_start()
+
+        while True:
+            with self.get_next_batch_context():
+                try:
+                    batch = next(async_batch_iter)
+                except StopIteration:
+                    break
+            with self.yield_batch_context(batch):
+                yield batch.data
+
+        self.after_epoch_end()
 
     def __iter__(self) -> Iterator[DataBatch]:
         return self._iter_batches()
 
-    @contextmanager
-    def _epoch_context(self):
-        """Context manager for epoch lifecycle: setup and cleanup.
-
-        Ensures proper initialization before iteration and cleanup after,
-        even if the iteration is interrupted or fails.
-        """
-        # Setup: Initialize epoch state
+    def before_epoch_start(self):
         self._yielded_first_batch = False
-        StatsManager.register_dataset_tag(self._dataset_tag)
-        try:
-            yield
-        finally:
-            # Cleanup: Clear iteration metrics after epoch completes or fails
-            StatsManager.clear_iteration_metrics(self._dataset_tag)
+
+    def after_epoch_end(self):
+        _StatsManager.update_iteration_metrics(
+            self._stats, self._dataset_tag, force_update=True
+        )
 
     @contextmanager
     def get_next_batch_context(self):
@@ -271,7 +270,10 @@ def get_next_batch_context(self):
     def yield_batch_context(self, batch: Batch):
         with self._stats.iter_user_s.timer() if self._stats else nullcontext():
             yield
-        StatsManager.update_iteration_metrics(self._stats, self._dataset_tag)
+        now = time.time()
+        if (now - self._metrics_last_updated) > self.UPDATE_METRICS_INTERVAL_S:
+            _StatsManager.update_iteration_metrics(self._stats, self._dataset_tag)
+            self._metrics_last_updated = now
 
 
 def _format_in_threadpool(

python/ray/data/_internal/execution/streaming_executor.py

@@ -44,7 +44,7 @@
 )
 from ray.data._internal.metadata_exporter import Topology as TopologyMetadata
 from ray.data._internal.progress_bar import ProgressBar
-from ray.data._internal.stats import DatasetStats, StatsManager, Timer
+from ray.data._internal.stats import DatasetStats, Timer, _StatsManager
 from ray.data.context import OK_PREFIX, WARN_PREFIX, DataContext
 from ray.util.debug import log_once
 from ray.util.metrics import Gauge
@@ -71,6 +71,8 @@ class StreamingExecutor(Executor, threading.Thread):
     a way that maximizes throughput under resource constraints.
     """
 
+    UPDATE_METRICS_INTERVAL_S: float = 5.0
+
     def __init__(
         self,
         data_context: DataContext,
@@ -115,6 +117,8 @@ def __init__(
             register_dataset_logger(self._dataset_id)
         )
 
+        self._metrics_last_updated: float = 0.0
+
         self._sched_loop_duration_s = Gauge(
             "data_sched_loop_duration_s",
             description="Duration of the scheduling loop in seconds",
@@ -223,7 +227,7 @@ def execute(
         op_to_id = {
             op: self._get_operator_id(op, i) for i, op in enumerate(self._topology)
         }
-        StatsManager.register_dataset_to_stats_actor(
+        _StatsManager.register_dataset_to_stats_actor(
             self._dataset_id,
             self._get_operator_tags(),
             TopologyMetadata.create_topology_metadata(dag, op_to_id),
@@ -271,9 +275,6 @@ def shutdown(self, force: bool, exception: Optional[Exception] = None):
                 else DatasetState.FAILED.name,
                 force_update=True,
             )
-            # Once Dataset execution completes, mark it as complete
-            # and remove last cached execution stats.
-            StatsManager.clear_last_execution_stats(self._dataset_id)
             # Freeze the stats and save it.
             self._final_stats = self._generate_stats()
             stats_summary_string = self._final_stats.to_summary().to_string(
@@ -660,13 +661,19 @@ def _get_state_dict(self, state):
         }
 
     def _update_stats_metrics(self, state: str, force_update: bool = False):
-        StatsManager.update_execution_metrics(
-            self._dataset_id,
-            [op.metrics for op in self._topology],
-            self._get_operator_tags(),
-            self._get_state_dict(state=state),
-            force_update=force_update,
-        )
+        now = time.time()
+        if (
+            force_update
+            or (now - self._metrics_last_updated) > self.UPDATE_METRICS_INTERVAL_S
+        ):
+            _StatsManager.update_execution_metrics(
+                self._dataset_id,
+                [op.metrics for op in self._topology],
+                self._get_operator_tags(),
+                self._get_state_dict(state=state),
+                force_update=force_update,
+            )
+            self._metrics_last_updated = now
 
     def _use_rich_progress(self):
         return self._data_context.enable_rich_progress_bars

python/ray/data/_internal/iterator/stream_split_iterator.py

@@ -81,7 +81,6 @@ def gen_blocks() -> Iterator[RefBundle]:
             future: ObjectRef[
                 Optional[ObjectRef[Block]]
             ] = self._coord_actor.get.remote(cur_epoch, self._output_split_idx)
-
             while True:
                 block_ref_and_md: Optional[RefBundle] = ray.get(future)
                 if not block_ref_and_md: