From 2458f0851d817309e75db1bc5e7b9bedb4732c09 Mon Sep 17 00:00:00 2001 From: Grain Team Date: Thu, 7 May 2026 16:43:16 -0700 Subject: [PATCH] Internal PiperOrigin-RevId: 912221544 --- grain/_src/python/dataset/transformations/prefetch.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/grain/_src/python/dataset/transformations/prefetch.py b/grain/_src/python/dataset/transformations/prefetch.py index 76449dd38..3f930fb7a 100644 --- a/grain/_src/python/dataset/transformations/prefetch.py +++ b/grain/_src/python/dataset/transformations/prefetch.py @@ -188,6 +188,8 @@ def __init__( self._buffer = collections.deque() self._lock = threading.Lock() self._executor_wrapper = None + self._next_depth_measure_step = 0 + self._depth_measure_freq = 100 assert isinstance(read_options.num_threads, int) assert isinstance(read_options.prefetch_buffer_size, int) @@ -257,7 +259,9 @@ def __next__(self) -> T: if not self._buffer: # Fill the buffer on the first iteration. self._fill_buffer() - self._measure_prefetch_depth() + if self._next_returned_index >= self._next_depth_measure_step: + self._measure_prefetch_depth() + self._next_depth_measure_step += self._depth_measure_freq element = self._buffer.popleft() # Prefetch elements until the buffer is full again. self._fill_buffer()