From eb0d3f049a18d006c3508f5eaadc92e22eb5645a Mon Sep 17 00:00:00 2001 From: Prince Kumar Date: Mon, 3 Aug 2026 06:14:24 +0000 Subject: [PATCH 1/6] feat(adaptive-readahead): introducing new cache adaptive-readahead --- fsspec/caching.py | 91 ++++++ fsspec/prefetch.py | 564 ++++++++++++++++++++++++++++++++++ fsspec/tests/test_caches.py | 30 ++ fsspec/tests/test_prefetch.py | 135 ++++++++ 4 files changed, 820 insertions(+) create mode 100644 fsspec/prefetch.py create mode 100644 fsspec/tests/test_prefetch.py diff --git a/fsspec/caching.py b/fsspec/caching.py index 3cfb343a6..678fd1c84 100644 --- a/fsspec/caching.py +++ b/fsspec/caching.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import collections import functools import logging @@ -282,6 +283,95 @@ def _fetch(self, start: int | None, end: int | None) -> bytes: return part + self.cache[:l] +class AdaptiveReadaheadCache(BaseCache): + """Adaptive readahead cache using the generic prefetch engine. + + This cache remains API-compatible with existing fsspec cache wiring and + falls back to classic ``ReadAheadCache`` when async prefetching is not + available. + """ + + name = "adaptive_readahead" + + def __init__( + self, + blocksize: int, + fetcher: Fetcher, + size: int, + concurrency: int = 4, + max_prefetch_size: int | None = None, + ) -> None: + super().__init__(blocksize, fetcher, size) + self._fallback = ReadAheadCache(blocksize, fetcher, size) + self._prefetcher = None + + async def _default_fetcher_async( + start_offset: int, + total_size: int, + split_factor: int = 1, + ) -> bytes: + del split_factor + return await asyncio.to_thread( + self.fetcher, start_offset, start_offset + total_size + ) + + try: + import fsspec.asyn + + from .prefetch import BackgroundPrefetcher + + self._prefetcher = BackgroundPrefetcher( + fetcher=_default_fetcher_async, + size=size, + concurrency=concurrency, + max_prefetch_size=max_prefetch_size, + loop=fsspec.asyn.get_loop(), + ) + except Exception as e: + logger.debug( + "AdaptiveReadaheadCache falling back to ReadAheadCache: %s", + e, + exc_info=True, + ) + self._prefetcher = None + + def _fetch(self, start: int | None, end: int | None) -> bytes: + if self._prefetcher is None: + out = self._fallback._fetch(start, end) + self.hit_count = self._fallback.hit_count + self.miss_count = self._fallback.miss_count + self.total_requested_bytes = self._fallback.total_requested_bytes + return out + + out = self._prefetcher.fetch(start, end) + self.miss_count += 1 + self.total_requested_bytes += len(out) + return out + + def close(self) -> None: + if self._prefetcher is not None: + self._prefetcher.close() + self._prefetcher = None + + def __getstate__(self) -> dict[str, Any]: + # The prefetcher owns asyncio primitives that are not picklable. + self.close() + state = self.__dict__.copy() + state["_prefetcher"] = None + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + self.__dict__.update(state) + self._prefetcher = None + + def __del__(self): + try: + self.close() + except Exception: + # Best-effort cleanup during GC. + pass + + class FirstChunkCache(BaseCache): """Caches the first block of a file only @@ -1016,6 +1106,7 @@ def register_cache(cls: type[BaseCache], clobber: bool = False) -> None: MMapCache, BytesCache, ReadAheadCache, + AdaptiveReadaheadCache, BlockCache, FirstChunkCache, AllBytes, diff --git a/fsspec/prefetch.py b/fsspec/prefetch.py new file mode 100644 index 000000000..944490375 --- /dev/null +++ b/fsspec/prefetch.py @@ -0,0 +1,564 @@ +from __future__ import annotations + +import asyncio +import logging +import weakref +from collections import deque + +import fsspec.asyn + +logger = logging.getLogger(__name__) + +# Kept for compatibility with existing tests/consumers. +HAS_CPYTHON_API = False + + +def _fast_slice(src_bytes: bytes, offset: int, read_size: int) -> bytes: + if read_size == 0: + return b"" + if offset < 0 or offset + read_size > len(src_bytes): + raise ValueError("Slice indices out of bounds") + return src_bytes[offset : offset + read_size] + + +class RunningAverageTracker: + """Tracks a running average of values over a sliding window.""" + + def __init__(self, maxlen: int = 10): + self._history: deque[int] = deque(maxlen=maxlen) + self._sum = 0 + + def add(self, value: int) -> None: + if value <= 0: + raise ValueError( + "Internal error, RunningAverageTracker tried inserting negative value" + ) + if len(self._history) == self._history.maxlen: + self._sum -= self._history[0] + + self._history.append(value) + self._sum += value + + @property + def average(self) -> int: + count = len(self._history) + if count == 0: + return 1024 * 1024 + return self._sum // count + + @property + def is_variable(self) -> bool: + count = len(self._history) + if count < 2: + return False + return len(set(self._history)) > 1 + + @property + def last_value(self) -> int: + if not self._history: + raise RuntimeError("No entry found in history") + return self._history[-1] + + def clear(self) -> None: + self._history.clear() + self._sum = 0 + + +class PrefetchProducer: + """Background worker that fetches sequential blocks of data.""" + + MIN_CHUNK_SIZE = 5 * 1024 * 1024 + MIN_PREFETCH_SIZE = 128 * 1024 * 1024 + MIN_STREAKS_FOR_PREFETCHING = 3 + VARIABLE_IO_THRESHOLD = 64 * 1024 * 1024 + + def __init__( + self, + fetcher, + size: int, + concurrency: int, + queue: asyncio.Queue, + wakeup_event: asyncio.Event, + consumer: "PrefetchConsumer", + tracker: RunningAverageTracker, + orchestrator: "BackgroundPrefetcher", + user_max_prefetch_size=None, + ): + self.fetcher = fetcher + self.size = size + self.concurrency = concurrency + self.queue = queue + self.wakeup_event = wakeup_event + self.consumer = consumer + self.tracker = tracker + self.orchestrator = weakref.proxy(orchestrator) + self._user_max_prefetch_size = user_max_prefetch_size + + self.current_offset = 0 + self.is_stopped = False + self._active_tasks = set() + self._producer_task = None + + @property + def max_prefetch_size(self) -> int: + if self._user_max_prefetch_size is not None: + return min( + self._user_max_prefetch_size, + max(2 * self.tracker.average, self.MIN_PREFETCH_SIZE), + ) + return max(2 * self.tracker.average, self.MIN_PREFETCH_SIZE) + + def start(self) -> None: + self.is_stopped = False + self.wakeup_event.clear() + self._producer_task = asyncio.create_task(self._loop()) + + async def stop(self) -> None: + self.is_stopped = True + self.wakeup_event.set() + + tasks_to_wait = [] + if self._producer_task and not self._producer_task.done(): + self._producer_task.cancel() + tasks_to_wait.append(self._producer_task) + + for task in list(self._active_tasks): + if not task.done(): + tasks_to_wait.append(task) + + self._active_tasks.clear() + + while not self.queue.empty(): + try: + item = self.queue.get_nowait() + if ( + isinstance(item, asyncio.Task) + and item.done() + and not item.cancelled() + ): + item.exception() + except asyncio.QueueEmpty: + break + + if tasks_to_wait: + await asyncio.gather(*tasks_to_wait, return_exceptions=True) + + self.wakeup_event.clear() + + async def restart(self, new_offset: int) -> None: + await self.stop() + self.current_offset = new_offset + self.start() + + async def _loop(self) -> None: + try: + while not self.is_stopped: + await self.wakeup_event.wait() + self.wakeup_event.clear() + + if self.is_stopped: + break + + await self._process_prefetch_cycle() + + except asyncio.CancelledError: + pass + except Exception as e: + self.is_stopped = True + self.orchestrator.set_error(e) + await self.queue.put(e) + + def _calculate_prefetch_params(self) -> tuple[int, int, int]: + avg_io_size = self.tracker.average + streak = self.consumer.sequential_streak + is_variable = self.tracker.is_variable + last_read_size = self.tracker.last_value + + exceeds_user_max = ( + self._user_max_prefetch_size is not None + and avg_io_size > self._user_max_prefetch_size + ) + + if ( + is_variable and avg_io_size > self.VARIABLE_IO_THRESHOLD + ) or exceeds_user_max: + prefetch_multiplier = 1 + elif streak < self.MIN_STREAKS_FOR_PREFETCHING: + prefetch_multiplier = 1 + else: + prefetch_multiplier = streak - self.MIN_STREAKS_FOR_PREFETCHING + 1 + + if self.queue.empty() or prefetch_multiplier == 1: + io_size = last_read_size + else: + io_size = avg_io_size + + prefetch_size = min(prefetch_multiplier * io_size, self.max_prefetch_size) + if self.consumer.offset + prefetch_size < self.consumer.target_offset: + prefetch_size = self.consumer.target_offset - self.consumer.offset + + if is_variable: + effective_prefetch_size = prefetch_size + else: + effective_prefetch_size = (prefetch_size // io_size) * io_size + if effective_prefetch_size == 0: + effective_prefetch_size = prefetch_size + + return prefetch_size, io_size, effective_prefetch_size + + async def _process_prefetch_cycle(self) -> None: + prefetch_size, io_size, effective_prefetch_size = ( + self._calculate_prefetch_params() + ) + + while ( + not self.is_stopped + and (self.current_offset - self.consumer.offset) < prefetch_size + and self.current_offset < self.size + ): + user_offset = self.consumer.offset + space_remaining = self.size - self.current_offset + prefetch_space_available = prefetch_size - ( + self.current_offset - user_offset + ) + + if prefetch_size >= self.MIN_CHUNK_SIZE: + if prefetch_space_available >= self.MIN_CHUNK_SIZE: + actual_size = min(max(self.MIN_CHUNK_SIZE, io_size), space_remaining) + else: + break + else: + actual_size = min(io_size, space_remaining) + + if prefetch_space_available < actual_size: + if self.tracker.is_variable or prefetch_space_available == prefetch_size: + actual_size = prefetch_space_available + else: + break + + streak = self.consumer.sequential_streak + if streak < self.MIN_STREAKS_FOR_PREFETCHING: + sfactor = self.concurrency + else: + sfactor = min( + self.concurrency, + max(1, actual_size * self.concurrency // effective_prefetch_size), + ) + + download_task = asyncio.create_task( + self.fetcher(self.current_offset, actual_size, split_factor=sfactor) + ) + self._active_tasks.add(download_task) + download_task.add_done_callback(self._active_tasks.discard) + + await self.queue.put(download_task) + self.current_offset += actual_size + + if self.current_offset >= self.size: + self.is_stopped = True + + +class PrefetchConsumer: + """Consumes prefetched chunks and manages byte slicing.""" + + def __init__( + self, + queue: asyncio.Queue, + wakeup_event: asyncio.Event, + tracker: RunningAverageTracker, + orchestrator: "BackgroundPrefetcher", + ): + self.queue = queue + self.wakeup_event = wakeup_event + self.tracker = tracker + self.orchestrator = weakref.proxy(orchestrator) + self.sequential_streak = 0 + self.offset = 0 + self.target_offset = 0 + self._current_block = b"" + self._current_block_idx = 0 + + def seek(self, new_offset: int) -> None: + self.offset = new_offset + self.target_offset = new_offset + self.sequential_streak = 0 + self._current_block = b"" + self._current_block_idx = 0 + + def clear_buffer(self) -> None: + self._current_block = b"" + self._current_block_idx = 0 + + async def _advance(self, size: int, save_data: bool) -> list[bytes]: + if size <= 0: + return [] + + chunks = [] + processed = 0 + self.target_offset = self.offset + size + + while processed < size: + available = len(self._current_block) - self._current_block_idx + trigger_wakeup = False + + if not available: + is_producer_stopped = ( + self.orchestrator.producer is None + or self.orchestrator.producer.is_stopped + ) + if is_producer_stopped and self.queue.empty(): + break + + if self.queue.empty(): + self.wakeup_event.set() + + task = await self.queue.get() + + if isinstance(task, Exception): + self.orchestrator.set_error(task) + raise task + + try: + block = await task + self.sequential_streak += 1 + if self.sequential_streak >= PrefetchProducer.MIN_STREAKS_FOR_PREFETCHING: + exceeds_user_max = ( + self.orchestrator.max_prefetch_size is not None + and self.tracker.average > self.orchestrator.max_prefetch_size + ) + is_massive_variable = ( + self.tracker.is_variable + and self.tracker.average > PrefetchProducer.VARIABLE_IO_THRESHOLD + ) + if not (is_massive_variable or exceeds_user_max): + trigger_wakeup = True + + self._current_block = block + self._current_block_idx = 0 + available = len(self._current_block) + except asyncio.CancelledError: + raise + except Exception as e: + self.orchestrator.set_error(e) + raise e + + if not self._current_block: + break + + needed = size - processed + take = min(needed, available) + + if save_data: + if take == len(self._current_block) and self._current_block_idx == 0: + chunk = self._current_block + else: + chunk = await asyncio.to_thread( + _fast_slice, + self._current_block, + self._current_block_idx, + take, + ) + chunks.append(chunk) + + self._current_block_idx += take + processed += take + self.offset += take + if trigger_wakeup: + self.wakeup_event.set() + + return chunks + + async def consume(self, size: int) -> bytes: + if size <= 0: + return b"" + + chunks = await self._advance(size, save_data=True) + + if not chunks: + return b"" + + if len(chunks) == 1: + return chunks[0] + + return await asyncio.to_thread(b"".join, chunks) + + async def skip(self, size: int) -> None: + await self._advance(size, save_data=False) + + +class BackgroundPrefetcher: + """Coordinates read behavior and background prefetch work.""" + + producer = None + + def __init__( + self, + fetcher, + size: int, + concurrency: int, + max_prefetch_size=None, + loop=None, + ): + self.size = size + self.concurrency = concurrency + self.max_prefetch_size = max_prefetch_size + + if max_prefetch_size is not None and max_prefetch_size <= 0: + raise ValueError( + "max_prefetch_size should be a positive integer to use adaptive prefetching!" + ) + + self.loop = loop + self._error = None + self.is_stopped = False + self.user_offset = 0 + self.read_tracker = RunningAverageTracker(maxlen=10) + + self.queue = None + self.wakeup_event = None + self._async_lock = None + self.consumer = None + self.producer = None + + def _start(): + self.queue = asyncio.Queue() + self.wakeup_event = asyncio.Event() + self._async_lock = asyncio.Lock() + + self.consumer = PrefetchConsumer( + queue=self.queue, + wakeup_event=self.wakeup_event, + tracker=self.read_tracker, + orchestrator=self, + ) + + self.producer = PrefetchProducer( + fetcher=fetcher, + size=self.size, + concurrency=self.concurrency, + queue=self.queue, + wakeup_event=self.wakeup_event, + consumer=self.consumer, + tracker=self.read_tracker, + orchestrator=self, + user_max_prefetch_size=max_prefetch_size, + ) + self.producer.start() + + try: + current_loop = asyncio.get_running_loop() + except RuntimeError: + current_loop = None + + if current_loop is self.loop and self.loop is not None: + _start() + elif self.loop is not None: + + async def _start_wrapper(): + _start() + + fsspec.asyn.sync(self.loop, _start_wrapper) + elif current_loop is not None: + self.loop = current_loop + _start() + else: + raise RuntimeError("No event loop found") + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self.aclose() + + def set_error(self, e: Exception) -> None: + self._error = e + + async def _restart_producer(self, new_offset: int) -> None: + self._error = None + await self.producer.restart(new_offset) + self.consumer.seek(new_offset) + self.read_tracker.clear() + + async def _async_fetch(self, start: int, end: int) -> bytes: + async with self._async_lock: + try: + if self.is_stopped: + raise RuntimeError("The file instance has been closed.") + + if self._error: + self.user_offset = start + await self._restart_producer(start) + elif start != self.user_offset: + block_offset = self.consumer.offset - self.consumer._current_block_idx + if self.user_offset < start <= self.producer.current_offset: + skip_amount = start - self.user_offset + await self.consumer.skip(skip_amount) + self.user_offset = start + elif block_offset <= start < self.consumer.offset: + self.consumer._current_block_idx = start - block_offset + self.consumer.offset = start + self.consumer.target_offset = start + self.user_offset = start + else: + self.user_offset = start + await self._restart_producer(start) + + requested_size = end - start + self.read_tracker.add(requested_size) + + chunk = await self.consumer.consume(requested_size) + self.user_offset += len(chunk) + return chunk + except asyncio.CancelledError as e: + self._error = e + raise + except Exception as e: + self._error = e + if self.producer and not self.producer.is_stopped: + await self.producer.stop() + raise + + async def _async_close(self) -> None: + async with self._async_lock: + if self.is_stopped: + return + + self.is_stopped = True + if self.producer: + await self.producer.stop() + + self.consumer.clear_buffer() + + async def afetch(self, start: int | None, end: int | None) -> bytes: + if start is None: + start = 0 + if end is None: + end = self.size + + end = min(end, self.size) + + if start >= self.size or start >= end: + return b"" + + if self.is_stopped: + raise RuntimeError( + "The file instance has been closed. This can occur if a close operation " + "is executed concurrently while a read operation is still in progress." + ) + + return await self._async_fetch(start, end) + + def fetch(self, start: int | None, end: int | None) -> bytes: + return fsspec.asyn.sync(self.loop, self.afetch, start, end) + + async def aclose(self) -> None: + await self._async_close() + + def close(self) -> None: + fsspec.asyn.sync(self.loop, self._async_close) diff --git a/fsspec/tests/test_caches.py b/fsspec/tests/test_caches.py index a1977f307..e35356427 100644 --- a/fsspec/tests/test_caches.py +++ b/fsspec/tests/test_caches.py @@ -1,3 +1,4 @@ +import asyncio import pickle import string @@ -359,3 +360,32 @@ def test_cache_kwargs(mocker): # It is a random location that cannot be predicted. # The important thing is the 'overwrite' kwarg fs.fs.put.assert_called_with(fs.fs.put.call_args[0][0], ["/test"], overwrite=True) + + +def test_adaptive_readahead_cache_with_async_fetcher(): + data = string.ascii_letters.encode() + + cache = caches["adaptive_readahead"]( + 8, + letters_fetcher, + len(data), + # Keep this below the default prefetch minimum so behavior is deterministic + # and bounded for unit testing. + concurrency=2, + max_prefetch_size=64, + ) + try: + assert cache._fetch(0, 0) == b"" + assert cache._fetch(0, 5) == data[0:5] + assert cache._fetch(5, 12) == data[5:12] + assert cache._fetch(12, 20) == data[12:20] + finally: + cache.close() + + +def test_adaptive_readahead_cache_fallback_without_loop(): + data = string.ascii_letters.encode() + cache = caches["adaptive_readahead"](8, letters_fetcher, len(data)) + + assert cache._fetch(0, 10) == data[0:10] + assert cache._fetch(10, 17) == data[10:17] diff --git a/fsspec/tests/test_prefetch.py b/fsspec/tests/test_prefetch.py new file mode 100644 index 000000000..5a6b179c0 --- /dev/null +++ b/fsspec/tests/test_prefetch.py @@ -0,0 +1,135 @@ +import asyncio +from unittest import mock + +import fsspec.asyn +import pytest + +from fsspec.prefetch import BackgroundPrefetcher, RunningAverageTracker, _fast_slice + + +@pytest.fixture +def prefetcher_factory(): + prefetchers = [] + + def _make_prefetcher(**kwargs): + if "loop" not in kwargs: + kwargs["loop"] = fsspec.asyn.get_loop() + + bp = BackgroundPrefetcher(**kwargs) + prefetchers.append(bp) + return bp + + yield _make_prefetcher + + for bp in prefetchers: + bp.is_stopped = False + bp.close() + + +class MockFetcher: + def __init__(self, data, fail_at_call=None): + self.data = data + self.calls = [] + self.fail_at_call = fail_at_call + self.call_count = 0 + + async def __call__(self, start, size, split_factor=1): + self.call_count += 1 + self.calls.append({"start": start, "size": size, "split_factor": split_factor}) + await asyncio.sleep(0.001) + + if self.fail_at_call is not None and self.call_count >= self.fail_at_call: + raise OSError("Simulated Network Timeout") + + return self.data[start : start + size] + + +def test_fast_slice_direct(): + src = b"0123456789" + assert _fast_slice(src, 2, 4) == b"2345" + assert _fast_slice(src, 5, 0) == b"" + assert _fast_slice(src, 0, 10) == b"0123456789" + + +def test_running_average_tracker(): + tracker = RunningAverageTracker(maxlen=3) + assert tracker.average == 1024 * 1024 + + tracker.add(512) + tracker.add(512) + assert tracker.average == 512 + + tracker.add(2048) + assert tracker.average == 1024 + + +def test_fetch_default_args_and_out_of_bounds(prefetcher_factory): + fetcher = MockFetcher(b"12345") + bp = prefetcher_factory(fetcher=fetcher, size=5, concurrency=4) + + assert bp.fetch(None, None) == b"12345" + assert bp.fetch(None, 2) == b"12" + assert bp.fetch(5, 10) == b"" + assert bp.fetch(10, 20) == b"" + assert bp.fetch(2, 2) == b"" + assert bp.fetch(4, 2) == b"" + + +def test_seek_logic(prefetcher_factory): + data = b"0123456789" * 10 + fetcher = MockFetcher(data) + bp = prefetcher_factory(fetcher=fetcher, size=100, concurrency=4) + + assert bp.fetch(0, 10) == data[0:10] + assert bp.fetch(10, 20) == data[10:20] + assert bp.user_offset == 20 + assert bp.fetch(50, 60) == data[50:60] + assert bp.user_offset == 60 + assert bp.fetch(10, 20) == data[10:20] + assert bp.user_offset == 20 + + +def test_producer_error_propagation_and_recovery(prefetcher_factory): + fetcher = MockFetcher(b"A" * 2000, fail_at_call=3) + bp = prefetcher_factory(fetcher=fetcher, size=2000, concurrency=4) + + for i in range(2): + bp.fetch(i * 100, (i + 1) * 100) + + with pytest.raises(OSError, match="Simulated Network Timeout"): + bp.fetch(400, 500) + + assert isinstance(bp._error, OSError) + + fetcher.fail_at_call = None + data = bp.fetch(400, 500) + assert data == b"A" * 100 + assert bp._error is None + + +def test_read_after_close(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + bp.close() + + assert bp.is_stopped is True + with pytest.raises(RuntimeError, match="The file instance has been closed"): + bp.fetch(0, 10) + + +def test_init_invalid_max_prefetch_size(): + with pytest.raises(ValueError, match=r"max_prefetch_size should be a positive integer"): + BackgroundPrefetcher( + fetcher=MockFetcher(b""), + size=1000, + concurrency=4, + max_prefetch_size=0, + loop=fsspec.asyn.get_loop(), + ) + + +def test_fast_slice_pypy_fallback(): + src = b"0123456789_pypy_fallback_test" + with mock.patch("fsspec.prefetch.HAS_CPYTHON_API", False): + assert _fast_slice(src, 11, 13) == b"pypy_fallback" + assert _fast_slice(src, 5, 0) == b"" + assert _fast_slice(src, 0, len(src)) == src From db257af8a2b3e9af4cb6d769873a77e2d27b2c51 Mon Sep 17 00:00:00 2001 From: Prince Kumar Date: Thu, 6 Aug 2026 13:43:17 +0000 Subject: [PATCH 2/6] rebase with master --- fsspec/caching.py | 39 ++++++++++++++++++++++++++++--------- fsspec/tests/test_caches.py | 8 ++++---- 2 files changed, 34 insertions(+), 13 deletions(-) diff --git a/fsspec/caching.py b/fsspec/caching.py index 678fd1c84..37368302c 100644 --- a/fsspec/caching.py +++ b/fsspec/caching.py @@ -14,6 +14,8 @@ from operator import itemgetter from typing import TYPE_CHECKING, Any, ClassVar, Generic, NamedTuple, TypeVar +import fsspec.asyn + if TYPE_CHECKING: import mmap @@ -284,14 +286,28 @@ def _fetch(self, start: int | None, end: int | None) -> bytes: class AdaptiveReadaheadCache(BaseCache): - """Adaptive readahead cache using the generic prefetch engine. + """Cache with adaptive asynchronous prefetching. + + Optimized for sequential and near-sequential reads by dynamically + adjusting the amount of data prefetched in the background. The cache + uses the generic prefetch engine when async loop support is available, + and falls back to ``ReadAheadCache`` when it is not. - This cache remains API-compatible with existing fsspec cache wiring and - falls back to classic ``ReadAheadCache`` when async prefetching is not - available. + Parameters + ---------- + blocksize: int + Nominal read size used by callers. + fetcher: Fetcher + Function of the form ``f(start, end)`` that returns bytes. + size: int + Total size of the file. + concurrency: int + Maximum number of concurrent background fetch tasks. + max_prefetch_size: int | None + Optional upper bound for adaptive prefetch size in bytes. """ - name = "adaptive_readahead" + name = "adaptive" def __init__( self, @@ -316,8 +332,6 @@ async def _default_fetcher_async( ) try: - import fsspec.asyn - from .prefetch import BackgroundPrefetcher self._prefetcher = BackgroundPrefetcher( @@ -327,9 +341,16 @@ async def _default_fetcher_async( max_prefetch_size=max_prefetch_size, loop=fsspec.asyn.get_loop(), ) + logger.info( + "AdaptiveReadaheadCache enabled (blocksize=%d, size=%d, concurrency=%d, max_prefetch_size=%s)", + blocksize, + size, + concurrency, + max_prefetch_size, + ) except Exception as e: - logger.debug( - "AdaptiveReadaheadCache falling back to ReadAheadCache: %s", + logger.info( + "AdaptiveReadaheadCache fallback to ReadAheadCache: %s", e, exc_info=True, ) diff --git a/fsspec/tests/test_caches.py b/fsspec/tests/test_caches.py index e35356427..410a8a545 100644 --- a/fsspec/tests/test_caches.py +++ b/fsspec/tests/test_caches.py @@ -362,10 +362,10 @@ def test_cache_kwargs(mocker): fs.fs.put.assert_called_with(fs.fs.put.call_args[0][0], ["/test"], overwrite=True) -def test_adaptive_readahead_cache_with_async_fetcher(): +def test_adaptive_cache_with_async_fetcher(): data = string.ascii_letters.encode() - cache = caches["adaptive_readahead"]( + cache = caches["adaptive"]( 8, letters_fetcher, len(data), @@ -383,9 +383,9 @@ def test_adaptive_readahead_cache_with_async_fetcher(): cache.close() -def test_adaptive_readahead_cache_fallback_without_loop(): +def test_adaptive_cache_fallback_without_loop(): data = string.ascii_letters.encode() - cache = caches["adaptive_readahead"](8, letters_fetcher, len(data)) + cache = caches["adaptive"](8, letters_fetcher, len(data)) assert cache._fetch(0, 10) == data[0:10] assert cache._fetch(10, 17) == data[10:17] From 2a8031c8a486641e2aa79b5a614e7382788480f6 Mon Sep 17 00:00:00 2001 From: Prince Kumar Date: Fri, 7 Aug 2026 04:02:15 +0000 Subject: [PATCH 3/6] matching the exact implementation --- fsspec/caching.py | 11 +- fsspec/implementations/local.py | 2 +- fsspec/prefetch.py | 421 ++++++++++++++++++++++++++++---- 3 files changed, 379 insertions(+), 55 deletions(-) diff --git a/fsspec/caching.py b/fsspec/caching.py index 37368302c..51e905540 100644 --- a/fsspec/caching.py +++ b/fsspec/caching.py @@ -14,8 +14,6 @@ from operator import itemgetter from typing import TYPE_CHECKING, Any, ClassVar, Generic, NamedTuple, TypeVar -import fsspec.asyn - if TYPE_CHECKING: import mmap @@ -331,6 +329,13 @@ async def _default_fetcher_async( self.fetcher, start_offset, start_offset + total_size ) + try: + from . import asyn as fsspec_asyn + except ImportError as e: + raise ImportError( + "AdaptiveReadaheadCache requires fsspec.asyn to be available" + ) from e + try: from .prefetch import BackgroundPrefetcher @@ -339,7 +344,7 @@ async def _default_fetcher_async( size=size, concurrency=concurrency, max_prefetch_size=max_prefetch_size, - loop=fsspec.asyn.get_loop(), + loop=fsspec_asyn.get_loop(), ) logger.info( "AdaptiveReadaheadCache enabled (blocksize=%d, size=%d, concurrency=%d, max_prefetch_size=%s)", diff --git a/fsspec/implementations/local.py b/fsspec/implementations/local.py index fd7531e7e..5ea89f8fd 100644 --- a/fsspec/implementations/local.py +++ b/fsspec/implementations/local.py @@ -8,7 +8,7 @@ import tempfile from functools import lru_cache -from fsspec import AbstractFileSystem +from ..spec import AbstractFileSystem from fsspec.compression import compr from fsspec.core import get_compression from fsspec.utils import isfilelike, stringify_path diff --git a/fsspec/prefetch.py b/fsspec/prefetch.py index 944490375..ed32d4a8c 100644 --- a/fsspec/prefetch.py +++ b/fsspec/prefetch.py @@ -1,34 +1,71 @@ -from __future__ import annotations - import asyncio +import ctypes import logging import weakref from collections import deque -import fsspec.asyn +from . import asyn as fsspec_asyn logger = logging.getLogger(__name__) -# Kept for compatibility with existing tests/consumers. -HAS_CPYTHON_API = False +try: + PyBytes_FromStringAndSize = ctypes.pythonapi.PyBytes_FromStringAndSize + PyBytes_FromStringAndSize.argtypes = (ctypes.c_void_p, ctypes.c_ssize_t) + PyBytes_FromStringAndSize.restype = ctypes.py_object + + PyBytes_AsString = ctypes.pythonapi.PyBytes_AsString + PyBytes_AsString.argtypes = (ctypes.py_object,) + PyBytes_AsString.restype = ctypes.c_void_p + HAS_CPYTHON_API = True +except Exception: + PyBytes_FromStringAndSize = None + PyBytes_AsString = None + HAS_CPYTHON_API = False -def _fast_slice(src_bytes: bytes, offset: int, read_size: int) -> bytes: +# Please refer to following discussion to understand why this is required at this point +# Discussion = https://github.com/fsspec/gcsfs/pull/795#discussion_r3032749881 +def _fast_slice(src_bytes, offset, read_size): if read_size == 0: return b"" if offset < 0 or offset + read_size > len(src_bytes): raise ValueError("Slice indices out of bounds") - return src_bytes[offset : offset + read_size] + + if HAS_CPYTHON_API: + dest_bytes = PyBytes_FromStringAndSize(None, read_size) + src_ptr = PyBytes_AsString(src_bytes) + dest_ptr = PyBytes_AsString(dest_bytes) + # Releases the GIL + ctypes.memmove(dest_ptr, src_ptr + offset, read_size) + return dest_bytes + else: + # Standard fallback for PyPy/non-CPython + return src_bytes[offset : offset + read_size] class RunningAverageTracker: - """Tracks a running average of values over a sliding window.""" + """Tracks a running average of values over a sliding window. + + This is used to monitor read sizes and adaptively scale the + prefetching strategy based on recent user behavior. + """ - def __init__(self, maxlen: int = 10): - self._history: deque[int] = deque(maxlen=maxlen) + def __init__(self, maxlen=10): + """Initializes the tracker with a specific window size. + + Args: + maxlen (int): The maximum number of historical values to keep. + """ + logger.debug("Initializing RunningAverageTracker with maxlen: %d", maxlen) + self._history = deque(maxlen=maxlen) self._sum = 0 - def add(self, value: int) -> None: + def add(self, value: int): + """Adds a new value to the sliding window and updates the rolling sum. + + Args: + value (int): The integer value to add to the history. + """ if value <= 0: raise ValueError( "Internal error, RunningAverageTracker tried inserting negative value" @@ -38,38 +75,78 @@ def add(self, value: int) -> None: self._history.append(value) self._sum += value + logger.debug( + "RunningAverageTracker added value: %d, new sum: %d", value, self._sum + ) @property def average(self) -> int: + """Calculates and returns the current running average. + + Returns: + int: The integer average of the current history. + """ count = len(self._history) if count == 0: - return 1024 * 1024 + return 1024 * 1024 # 1MB return self._sum // count @property def is_variable(self) -> bool: + """Determines if the history contains distinct chunk sizes.""" count = len(self._history) if count < 2: return False + return len(set(self._history)) > 1 @property def last_value(self) -> int: + """Returns the most recent entry in the history.""" if not self._history: raise RuntimeError("No entry found in history") + return self._history[-1] - def clear(self) -> None: + def clear(self): + """Clears the history and resets the sum to zero.""" + logger.debug("Clearing RunningAverageTracker history.") self._history.clear() self._sum = 0 class PrefetchProducer: - """Background worker that fetches sequential blocks of data.""" + """Background worker that fetches sequential blocks of data. + This class handles the network requests. It spawns asynchronous tasks + to fetch data ahead of the user's current reading position and + places those task promises into a queue for the consumer. + """ + + # If the request is too small, and prefetch window is expanded till 5MB + # we then make request in 5MB blocks. MIN_CHUNK_SIZE = 5 * 1024 * 1024 + + # If user doesn't specify any max_prefetch_size, the prefetcher defaults + # to maximum of 2 * io_size and 128MB MIN_PREFETCH_SIZE = 128 * 1024 * 1024 + + # The prefetching starts on the third read. MIN_STREAKS_FOR_PREFETCHING = 3 + + # Threshold for disabling proactive prefetching on large, variable reads. + # + # If the average read size exceeds this value and patterns are variable, + # prefetching shifts from an I/O bottleneck to a memory(CPU) bottleneck. When a user + # requests random massive sizes (e.g., jumping between 64MB and INF), the + # producer still fetches chunks based on the rolling average. The consumer + # then has to pick up multiple chunks and stitch them together to match the + # exact requested size. + # + # For small average read sizes, this byte assembly is fast and the bottleneck + # remains the network I/O. However, for massive reads (>= 64MB), the extra + # step of copying and assembling huge byte strings in memory severely slows + # down the operation. VARIABLE_IO_THRESHOLD = 64 * 1024 * 1024 def __init__( @@ -84,11 +161,31 @@ def __init__( orchestrator: "BackgroundPrefetcher", user_max_prefetch_size=None, ): + """Initializes the background producer. + + Args: + fetcher (Callable): A coroutine function to fetch bytes from a remote source. + size (int): Total size of the file being fetched. + concurrency (int): Maximum number of concurrent fetch tasks. + queue (asyncio.Queue): The shared queue to push download tasks into. + wakeup_event (asyncio.Event): Event used to wake the producer from an idle state. + consumer (PrefetchConsumer): The consumer reading the prefetched chunks. + tracker (RunningAverageTracker): Tracker for history of read sizes. + orchestrator (BackgroundPrefetcher): The parent object managing the operation. + user_max_prefetch_size (int, optional): A hard limit for prefetch size overrides. + """ + logger.debug( + "Initializing PrefetchProducer: size=%d, concurrency=%d, user_max_prefetch_size=%s", + size, + concurrency, + user_max_prefetch_size, + ) self.fetcher = fetcher self.size = size self.concurrency = concurrency self.queue = queue self.wakeup_event = wakeup_event + self.consumer = consumer self.tracker = tracker self.orchestrator = weakref.proxy(orchestrator) @@ -101,6 +198,11 @@ def __init__( @property def max_prefetch_size(self) -> int: + """Calculates the maximum prefetch size based on user intent or io size. + + Returns: + int: The maximum number of bytes to prefetch ahead. + """ if self._user_max_prefetch_size is not None: return min( self._user_max_prefetch_size, @@ -108,12 +210,25 @@ def max_prefetch_size(self) -> int: ) return max(2 * self.tracker.average, self.MIN_PREFETCH_SIZE) - def start(self) -> None: + def start(self): + """Starts the background producer loop. + + This clears any previous wakeup events and spawns the main loop task. + """ + logger.debug("Starting PrefetchProducer loop.") self.is_stopped = False self.wakeup_event.clear() self._producer_task = asyncio.create_task(self._loop()) - async def stop(self) -> None: + async def stop(self): + """Cancels all active fetch tasks and shuts down the producer loop. + + This method ensures the queue is flushed and waits for cancelled + tasks to finish cleaning up. + """ + logger.debug( + "Stopping PrefetchProducer. Active fetch tasks: %d", len(self._active_tasks) + ) self.is_stopped = True self.wakeup_event.set() @@ -126,8 +241,12 @@ async def stop(self) -> None: if not task.done(): tasks_to_wait.append(task) + # We do not cancel the network task, instead we wait on them. + # This is intentionally done to avoid MRD stream disruption. self._active_tasks.clear() + # Clear out any leftover items in the queue + cleared_items = 0 while not self.queue.empty(): try: item = self.queue.get_nowait() @@ -137,20 +256,38 @@ async def stop(self) -> None: and not item.cancelled() ): item.exception() + cleared_items += 1 except asyncio.QueueEmpty: break + if cleared_items > 0: + logger.debug( + "Cleared %d leftover items from the queue during stop.", cleared_items + ) + if tasks_to_wait: + logger.debug( + "Waiting for %d cancelled tasks to finish their teardown.", + len(tasks_to_wait), + ) await asyncio.gather(*tasks_to_wait, return_exceptions=True) self.wakeup_event.clear() - async def restart(self, new_offset: int) -> None: + async def restart(self, new_offset: int): + """Stops current tasks and restarts the background loop at a new byte offset. + + Args: + new_offset (int): The new byte position to start prefetching from. + """ + logger.debug("Restarting PrefetchProducer at new offset: %d", new_offset) await self.stop() self.current_offset = new_offset self.start() - async def _loop(self) -> None: + async def _loop(self): + """The main background loop that delegates calculations and spawns tasks.""" + logger.debug("PrefetchProducer internal loop is now running.") try: while not self.is_stopped: await self.wakeup_event.wait() @@ -162,13 +299,24 @@ async def _loop(self) -> None: await self._process_prefetch_cycle() except asyncio.CancelledError: - pass + logger.debug("PrefetchProducer loop was cancelled.") except Exception as e: + logger.error( + "PrefetchProducer loop encountered an unexpected error: %s", + e, + exc_info=True, + ) self.is_stopped = True self.orchestrator.set_error(e) await self.queue.put(e) def _calculate_prefetch_params(self) -> tuple[int, int, int]: + """ + Evaluates current trackers and state to determine sizes. + + Returns: + tuple: (prefetch_size, io_size, effective_prefetch_size) + """ avg_io_size = self.tracker.average streak = self.consumer.sequential_streak is_variable = self.tracker.is_variable @@ -179,9 +327,13 @@ def _calculate_prefetch_params(self) -> tuple[int, int, int]: and avg_io_size > self._user_max_prefetch_size ) + # Disable prefetching ahead if variable AND average > 64MB, or if it exceeds user max if ( is_variable and avg_io_size > self.VARIABLE_IO_THRESHOLD ) or exceeds_user_max: + logger.debug( + "Large IO detected (variable > 64MB or > user max). Disabling background prefetching." + ) prefetch_multiplier = 1 elif streak < self.MIN_STREAKS_FOR_PREFETCHING: prefetch_multiplier = 1 @@ -206,11 +358,19 @@ def _calculate_prefetch_params(self) -> tuple[int, int, int]: return prefetch_size, io_size, effective_prefetch_size - async def _process_prefetch_cycle(self) -> None: + async def _process_prefetch_cycle(self): + """Executes a single cycle of enqueuing fetch tasks.""" prefetch_size, io_size, effective_prefetch_size = ( self._calculate_prefetch_params() ) + logger.debug( + "Producer awake. Current offset: %d, User offset: %d, Prefetch size: %d", + self.current_offset, + self.consumer.offset, + prefetch_size, + ) + while ( not self.is_stopped and (self.current_offset - self.consumer.offset) < prefetch_size @@ -224,14 +384,19 @@ async def _process_prefetch_cycle(self) -> None: if prefetch_size >= self.MIN_CHUNK_SIZE: if prefetch_space_available >= self.MIN_CHUNK_SIZE: - actual_size = min(max(self.MIN_CHUNK_SIZE, io_size), space_remaining) + actual_size = min( + max(self.MIN_CHUNK_SIZE, io_size), space_remaining + ) else: break else: actual_size = min(io_size, space_remaining) if prefetch_space_available < actual_size: - if self.tracker.is_variable or prefetch_space_available == prefetch_size: + if ( + self.tracker.is_variable + or prefetch_space_available == prefetch_size + ): actual_size = prefetch_space_available else: break @@ -242,9 +407,19 @@ async def _process_prefetch_cycle(self) -> None: else: sfactor = min( self.concurrency, - max(1, actual_size * self.concurrency // effective_prefetch_size), + max( + 1, + actual_size * self.concurrency // effective_prefetch_size, + ), ) + logger.debug( + "Spawning fetch task. Offset: %d, Size: %d, Split Factor: %d", + self.current_offset, + actual_size, + sfactor, + ) + download_task = asyncio.create_task( self.fetcher(self.current_offset, actual_size, split_factor=sfactor) ) @@ -255,11 +430,16 @@ async def _process_prefetch_cycle(self) -> None: self.current_offset += actual_size if self.current_offset >= self.size: + logger.debug("Producer reached EOF. Exiting background loop.") self.is_stopped = True class PrefetchConsumer: - """Consumes prefetched chunks and manages byte slicing.""" + """Consumes prefetched chunks from the queue and manages byte slicing. + + This class pulls data out of the shared queue and slices it into the + exact byte sizes requested by the user. It also manages the local block buffer. + """ def __init__( self, @@ -268,6 +448,15 @@ def __init__( tracker: RunningAverageTracker, orchestrator: "BackgroundPrefetcher", ): + """Initializes the consumer. + + Args: + queue (asyncio.Queue): The shared queue containing fetch tasks. + wakeup_event (asyncio.Event): Event used to wake the producer when more data is needed. + tracker (RunningAverageTracker): Tracker for history of read sizes. + orchestrator (BackgroundPrefetcher): The parent object managing the operation. + """ + logger.debug("Initializing PrefetchConsumer.") self.queue = queue self.wakeup_event = wakeup_event self.tracker = tracker @@ -278,18 +467,33 @@ def __init__( self._current_block = b"" self._current_block_idx = 0 - def seek(self, new_offset: int) -> None: + def seek(self, new_offset: int): + """Clears the buffer and resets the internal offset for a hard seek. + + Args: + new_offset (int): The byte position the consumer is jumping to. + """ + logger.debug( + "Consumer executing hard seek to offset %d. Clearing internal buffer.", + new_offset, + ) self.offset = new_offset self.target_offset = new_offset self.sequential_streak = 0 self._current_block = b"" self._current_block_idx = 0 - def clear_buffer(self) -> None: + def clear_buffer(self): + """Discards the local byte buffer. Useful during shutdown or resets.""" + logger.debug("Consumer local block buffer cleared.") self._current_block = b"" self._current_block_idx = 0 async def _advance(self, size: int, save_data: bool) -> list[bytes]: + """Internal method to advance the offset and optionally extract data. + + Handles queue exhaustion, producer wakeups, and streak tracking. + """ if size <= 0: return [] @@ -307,31 +511,48 @@ async def _advance(self, size: int, save_data: bool) -> list[bytes]: or self.orchestrator.producer.is_stopped ) if is_producer_stopped and self.queue.empty(): + logger.debug("Consumer reached EOF.") break if self.queue.empty(): + logger.debug("Queue is empty. Waking up producer.") self.wakeup_event.set() task = await self.queue.get() if isinstance(task, Exception): + logger.error("Consumer retrieved an exception: %s", task) self.orchestrator.set_error(task) raise task try: block = await task + self.sequential_streak += 1 - if self.sequential_streak >= PrefetchProducer.MIN_STREAKS_FOR_PREFETCHING: + if ( + self.sequential_streak + >= PrefetchProducer.MIN_STREAKS_FOR_PREFETCHING + ): exceeds_user_max = ( self.orchestrator.max_prefetch_size is not None - and self.tracker.average > self.orchestrator.max_prefetch_size + and self.tracker.average + > self.orchestrator.max_prefetch_size ) is_massive_variable = ( self.tracker.is_variable - and self.tracker.average > PrefetchProducer.VARIABLE_IO_THRESHOLD + and self.tracker.average + > PrefetchProducer.VARIABLE_IO_THRESHOLD ) + + # Suppress proactive wakeups to prevent large CPU assembly + # on erratic large reads or exceeding max if not (is_massive_variable or exceeds_user_max): trigger_wakeup = True + else: + logger.debug( + "Suppressing proactive producer wakeup due to massive variable" + " workload or exceeding user max prefetch." + ) self._current_block = block self._current_block_idx = 0 @@ -339,6 +560,7 @@ async def _advance(self, size: int, save_data: bool) -> list[bytes]: except asyncio.CancelledError: raise except Exception as e: + logger.error("Consumer caught an error: %s", e, exc_info=True) self.orchestrator.set_error(e) raise e @@ -352,11 +574,9 @@ async def _advance(self, size: int, save_data: bool) -> list[bytes]: if take == len(self._current_block) and self._current_block_idx == 0: chunk = self._current_block else: + # Native Python slicing was GIL bound in my experiments. chunk = await asyncio.to_thread( - _fast_slice, - self._current_block, - self._current_block_idx, - take, + _fast_slice, self._current_block, self._current_block_idx, take ) chunks.append(chunk) @@ -369,6 +589,20 @@ async def _advance(self, size: int, save_data: bool) -> list[bytes]: return chunks async def consume(self, size: int) -> bytes: + """Pulls exactly 'size' bytes from the local block or the task queue. + + If the local block is exhausted, this will wait on the queue for the next + available chunk of data. + + Args: + size (int): The exact number of bytes to retrieve. + + Returns: + bytes: The requested bytes. This may be shorter than 'size' if EOF is reached. + + Raises: + Exception: Re-raises any exceptions encountered by the producer fetch tasks. + """ if size <= 0: return b"" @@ -383,27 +617,50 @@ async def consume(self, size: int) -> bytes: return await asyncio.to_thread(b"".join, chunks) async def skip(self, size: int) -> None: + """Advances the consumer offset without allocating memory.""" await self._advance(size, save_data=False) class BackgroundPrefetcher: - """Coordinates read behavior and background prefetch work.""" + """Orchestrator that manages reading behavior and coordinates background work. + + This acts as the main public interface for the file reader. It tracks the + user's reading history, routes seek operations, and links the producer's + network tasks with the consumer's data slicing logic. + """ producer = None def __init__( - self, - fetcher, - size: int, - concurrency: int, - max_prefetch_size=None, - loop=None, + self, fetcher, size: int, concurrency: int, max_prefetch_size=None, loop=None ): + """Initializes the background prefetcher. + + Args: + fetcher (Callable): A coroutine of the form `f(start, end)` which gets bytes from the remote. + size (int): Total byte size of the file being read. + concurrency (int): Number of concurrent network requests to use for large chunks. + max_prefetch_size (int, optional): Maximum bytes to prefetch ahead of the current user offset. + loop (asyncio.AbstractEventLoop, optional): The event loop to attach the prefetcher to. + If executing synchronously, this should be the fsspec background loop. If executing + asynchronously (asynchronous=True), this should be None so it can automatically + inherit the user's currently running event loop. + + Raises: + ValueError: If max_prefetch_size is provided but is not a positive integer. + """ + logger.debug( + "Starting BackgroundPrefetcher. Size: %d, Concurrency: %d, Max Prefetch: %s", + size, + concurrency, + max_prefetch_size, + ) self.size = size self.concurrency = concurrency self.max_prefetch_size = max_prefetch_size if max_prefetch_size is not None and max_prefetch_size <= 0: + logger.error("Invalid max_prefetch_size provided: %s", max_prefetch_size) raise ValueError( "max_prefetch_size should be a positive integer to use adaptive prefetching!" ) @@ -421,6 +678,7 @@ def __init__( self.producer = None def _start(): + # Ensures all primitives bind directly to `self.loop` self.queue = asyncio.Queue() self.wakeup_event = asyncio.Event() self._async_lock = asyncio.Lock() @@ -451,61 +709,102 @@ def _start(): current_loop = None if current_loop is self.loop and self.loop is not None: + # We are already safely running inside the fsspec background loop _start() elif self.loop is not None: - + # We are on the main thread; schedule setup on the fsspec background loop async def _start_wrapper(): _start() - fsspec.asyn.sync(self.loop, _start_wrapper) + fsspec_asyn.sync(self.loop, _start_wrapper) elif current_loop is not None: + # asynchronous=True: use the user's active event loop self.loop = current_loop _start() else: + # asynchronous=True but called completely outside of an async context raise RuntimeError("No event loop found") + logger.debug("BackgroundPrefetcher initialization complete.") + def __enter__(self): + """Context manager entry point.""" return self def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit point. Ensures the prefetcher is cleanly closed.""" self.close() async def __aenter__(self): + """Async context manager entry point.""" return self async def __aexit__(self, exc_type, exc_val, exc_tb): + """Async context manager exit point. Ensures the prefetcher is cleanly closed.""" await self.aclose() - def set_error(self, e: Exception) -> None: + def set_error(self, e: Exception): + logger.error("Global error state set in BackgroundPrefetcher: %s", e) self._error = e - async def _restart_producer(self, new_offset: int) -> None: + async def _restart_producer(self, new_offset: int): + logger.debug( + "Handling seek request. Restarting producer at offset: %d", new_offset + ) self._error = None await self.producer.restart(new_offset) self.consumer.seek(new_offset) self.read_tracker.clear() - async def _async_fetch(self, start: int, end: int) -> bytes: + async def _async_fetch(self, start, end): + """Core internal async fetching logic, protected safely by the async lock.""" async with self._async_lock: try: if self.is_stopped: raise RuntimeError("The file instance has been closed.") + logger.debug("Executing _async_fetch for range %d - %d.", start, end) + + # If the prefetcher is in error state, let's do a hard seek to start offset. if self._error: + logger.info( + "Recovering from error state. Restarting producer at offset: %d", + start, + ) self.user_offset = start await self._restart_producer(start) elif start != self.user_offset: - block_offset = self.consumer.offset - self.consumer._current_block_idx + block_offset = ( + self.consumer.offset - self.consumer._current_block_idx + ) if self.user_offset < start <= self.producer.current_offset: + logger.debug( + "Soft seek detected. Skipping ahead from %d to %d.", + self.user_offset, + start, + ) skip_amount = start - self.user_offset await self.consumer.skip(skip_amount) self.user_offset = start elif block_offset <= start < self.consumer.offset: + logger.debug( + "Local seek performed. User offset moved from %d to %d. " + "Adjusting buffer index from %d to %d.", + self.user_offset, + start, + self.consumer._current_block_idx, + start - block_offset, + ) self.consumer._current_block_idx = start - block_offset self.consumer.offset = start self.consumer.target_offset = start self.user_offset = start else: + logger.debug( + "Hard seek detected. Moving user offset from %d to %d.", + self.user_offset, + start, + ) self.user_offset = start await self._restart_producer(start) @@ -514,39 +813,55 @@ async def _async_fetch(self, start: int, end: int) -> bytes: chunk = await self.consumer.consume(requested_size) self.user_offset += len(chunk) + + logger.debug("Completed _async_fetch. Returned %d bytes.", len(chunk)) return chunk except asyncio.CancelledError as e: self._error = e raise except Exception as e: + logger.error( + "Exception raised during asynchronous fetch: %s", e, exc_info=True + ) self._error = e if self.producer and not self.producer.is_stopped: await self.producer.stop() raise - async def _async_close(self) -> None: + async def _async_close(self): + """Asynchronous teardown logic protected by the async lock.""" async with self._async_lock: if self.is_stopped: return self.is_stopped = True + logger.debug("Acquired async lock. Tearing down producer and buffers.") + if self.producer: await self.producer.stop() self.consumer.clear_buffer() + logger.debug("BackgroundPrefetcher closed successfully.") async def afetch(self, start: int | None, end: int | None) -> bytes: + """Asynchronous API counterpart to `_fetch`.""" if start is None: start = 0 if end is None: end = self.size end = min(end, self.size) + logger.debug( + "Asynchronous afetch called for bounds start=%s, end=%s.", start, end + ) if start >= self.size or start >= end: return b"" if self.is_stopped: + logger.error( + "Cannot fetch data: BackgroundPrefetcher is stopped or closed." + ) raise RuntimeError( "The file instance has been closed. This can occur if a close operation " "is executed concurrently while a read operation is still in progress." @@ -555,10 +870,14 @@ async def afetch(self, start: int | None, end: int | None) -> bytes: return await self._async_fetch(start, end) def fetch(self, start: int | None, end: int | None) -> bytes: - return fsspec.asyn.sync(self.loop, self.afetch, start, end) + """Synchronous API wrapper delegating to `afetch`.""" + # Delegates all boundaries, checking, and fetching to the async event loop perfectly + return fsspec_asyn.sync(self.loop, self.afetch, start, end) - async def aclose(self) -> None: + async def aclose(self): + """Safely shuts down the prefetcher from an asynchronous context.""" await self._async_close() - def close(self) -> None: - fsspec.asyn.sync(self.loop, self._async_close) + def close(self): + """Safely shuts down the prefetcher from a synchronous context.""" + fsspec_asyn.sync(self.loop, self._async_close) \ No newline at end of file From bddaf64442607adeb6f846ea96de1a55133e619c Mon Sep 17 00:00:00 2001 From: Prince Kumar Date: Fri, 7 Aug 2026 04:53:18 +0000 Subject: [PATCH 4/6] complete unit test --- fsspec/caching.py | 2 +- fsspec/{prefetch.py => prefetcher.py} | 0 fsspec/tests/test_prefetch.py | 135 ----- fsspec/tests/test_prefetcher.py | 726 ++++++++++++++++++++++++++ 4 files changed, 727 insertions(+), 136 deletions(-) rename fsspec/{prefetch.py => prefetcher.py} (100%) delete mode 100644 fsspec/tests/test_prefetch.py create mode 100644 fsspec/tests/test_prefetcher.py diff --git a/fsspec/caching.py b/fsspec/caching.py index 51e905540..8b150e818 100644 --- a/fsspec/caching.py +++ b/fsspec/caching.py @@ -337,7 +337,7 @@ async def _default_fetcher_async( ) from e try: - from .prefetch import BackgroundPrefetcher + from .prefetcher import BackgroundPrefetcher self._prefetcher = BackgroundPrefetcher( fetcher=_default_fetcher_async, diff --git a/fsspec/prefetch.py b/fsspec/prefetcher.py similarity index 100% rename from fsspec/prefetch.py rename to fsspec/prefetcher.py diff --git a/fsspec/tests/test_prefetch.py b/fsspec/tests/test_prefetch.py deleted file mode 100644 index 5a6b179c0..000000000 --- a/fsspec/tests/test_prefetch.py +++ /dev/null @@ -1,135 +0,0 @@ -import asyncio -from unittest import mock - -import fsspec.asyn -import pytest - -from fsspec.prefetch import BackgroundPrefetcher, RunningAverageTracker, _fast_slice - - -@pytest.fixture -def prefetcher_factory(): - prefetchers = [] - - def _make_prefetcher(**kwargs): - if "loop" not in kwargs: - kwargs["loop"] = fsspec.asyn.get_loop() - - bp = BackgroundPrefetcher(**kwargs) - prefetchers.append(bp) - return bp - - yield _make_prefetcher - - for bp in prefetchers: - bp.is_stopped = False - bp.close() - - -class MockFetcher: - def __init__(self, data, fail_at_call=None): - self.data = data - self.calls = [] - self.fail_at_call = fail_at_call - self.call_count = 0 - - async def __call__(self, start, size, split_factor=1): - self.call_count += 1 - self.calls.append({"start": start, "size": size, "split_factor": split_factor}) - await asyncio.sleep(0.001) - - if self.fail_at_call is not None and self.call_count >= self.fail_at_call: - raise OSError("Simulated Network Timeout") - - return self.data[start : start + size] - - -def test_fast_slice_direct(): - src = b"0123456789" - assert _fast_slice(src, 2, 4) == b"2345" - assert _fast_slice(src, 5, 0) == b"" - assert _fast_slice(src, 0, 10) == b"0123456789" - - -def test_running_average_tracker(): - tracker = RunningAverageTracker(maxlen=3) - assert tracker.average == 1024 * 1024 - - tracker.add(512) - tracker.add(512) - assert tracker.average == 512 - - tracker.add(2048) - assert tracker.average == 1024 - - -def test_fetch_default_args_and_out_of_bounds(prefetcher_factory): - fetcher = MockFetcher(b"12345") - bp = prefetcher_factory(fetcher=fetcher, size=5, concurrency=4) - - assert bp.fetch(None, None) == b"12345" - assert bp.fetch(None, 2) == b"12" - assert bp.fetch(5, 10) == b"" - assert bp.fetch(10, 20) == b"" - assert bp.fetch(2, 2) == b"" - assert bp.fetch(4, 2) == b"" - - -def test_seek_logic(prefetcher_factory): - data = b"0123456789" * 10 - fetcher = MockFetcher(data) - bp = prefetcher_factory(fetcher=fetcher, size=100, concurrency=4) - - assert bp.fetch(0, 10) == data[0:10] - assert bp.fetch(10, 20) == data[10:20] - assert bp.user_offset == 20 - assert bp.fetch(50, 60) == data[50:60] - assert bp.user_offset == 60 - assert bp.fetch(10, 20) == data[10:20] - assert bp.user_offset == 20 - - -def test_producer_error_propagation_and_recovery(prefetcher_factory): - fetcher = MockFetcher(b"A" * 2000, fail_at_call=3) - bp = prefetcher_factory(fetcher=fetcher, size=2000, concurrency=4) - - for i in range(2): - bp.fetch(i * 100, (i + 1) * 100) - - with pytest.raises(OSError, match="Simulated Network Timeout"): - bp.fetch(400, 500) - - assert isinstance(bp._error, OSError) - - fetcher.fail_at_call = None - data = bp.fetch(400, 500) - assert data == b"A" * 100 - assert bp._error is None - - -def test_read_after_close(prefetcher_factory): - bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) - bp.close() - - assert bp.is_stopped is True - with pytest.raises(RuntimeError, match="The file instance has been closed"): - bp.fetch(0, 10) - - -def test_init_invalid_max_prefetch_size(): - with pytest.raises(ValueError, match=r"max_prefetch_size should be a positive integer"): - BackgroundPrefetcher( - fetcher=MockFetcher(b""), - size=1000, - concurrency=4, - max_prefetch_size=0, - loop=fsspec.asyn.get_loop(), - ) - - -def test_fast_slice_pypy_fallback(): - src = b"0123456789_pypy_fallback_test" - with mock.patch("fsspec.prefetch.HAS_CPYTHON_API", False): - assert _fast_slice(src, 11, 13) == b"pypy_fallback" - assert _fast_slice(src, 5, 0) == b"" - assert _fast_slice(src, 0, len(src)) == src diff --git a/fsspec/tests/test_prefetcher.py b/fsspec/tests/test_prefetcher.py new file mode 100644 index 000000000..26f3d69bc --- /dev/null +++ b/fsspec/tests/test_prefetcher.py @@ -0,0 +1,726 @@ +import asyncio +from unittest import mock + +import fsspec.asyn +import pytest + +from fsspec.prefetcher import BackgroundPrefetcher, RunningAverageTracker, _fast_slice + + +@pytest.fixture +def prefetcher_factory(): + prefetchers = [] + + def _make_prefetcher(**kwargs): + if "loop" not in kwargs: + kwargs["loop"] = fsspec.asyn.get_loop() + + bp = BackgroundPrefetcher(**kwargs) + prefetchers.append(bp) + return bp + + yield _make_prefetcher + + for bp in prefetchers: + bp.is_stopped = False + bp.close() + + +class MockFetcher: + def __init__(self, data, fail_at_call=None, hang_at_call=None): + self.data = data + self.calls = [] + self.fail_at_call = fail_at_call + self.hang_at_call = hang_at_call + self.call_count = 0 + + async def __call__(self, start, size, split_factor=1): + self.call_count += 1 + self.calls.append({"start": start, "size": size, "split_factor": split_factor}) + + await asyncio.sleep(0.001) + + if self.hang_at_call is not None and self.call_count >= self.hang_at_call: + await asyncio.sleep(1000) + + if self.fail_at_call is not None and self.call_count >= self.fail_at_call: + raise OSError("Simulated Network Timeout") + + return self.data[start : start + size] + + +def test_fast_slice_direct(): + src = b"0123456789" + assert _fast_slice(src, 2, 4) == b"2345" + assert _fast_slice(src, 5, 0) == b"" + assert _fast_slice(src, 0, 10) == b"0123456789" + + +def test_running_average_tracker(): + tracker = RunningAverageTracker(maxlen=3) + assert tracker.average == 1024 * 1024 # Default 1MB fallback + + tracker.add(512) + tracker.add(512) + assert tracker.average == 512 + + tracker.add(2048) + assert tracker.average == 1024 # (512 + 512 + 2048) // 3 + + tracker.clear() + assert tracker.average == 1024 * 1024 + + +def test_max_prefetch_size_property(prefetcher_factory): + bp1 = prefetcher_factory(fetcher=MockFetcher(b""), size=10000, concurrency=4) + assert bp1.producer.max_prefetch_size == bp1.producer.MIN_PREFETCH_SIZE + + bp2 = prefetcher_factory(fetcher=MockFetcher(b""), size=1000000000, concurrency=4) + # Give it a history so it calculates 2x the io_size + bp2.read_tracker.add(100 * 1024 * 1024) + assert bp2.producer.max_prefetch_size == 200 * 1024 * 1024 + + +def test_sequential_read_spanning_blocks(prefetcher_factory): + data = b"A" * 100 + b"B" * 100 + b"C" * 100 + fetcher = MockFetcher(data) + bp = prefetcher_factory(fetcher=fetcher, size=300, concurrency=4) + bp.read_tracker.add(100) # Seed the adaptive tracker + + assert bp.fetch(0, 100) == b"A" * 100 + assert bp.fetch(100, 150) == b"B" * 50 + assert bp.consumer._current_block_idx == 50 + assert bp.fetch(150, 250) == b"B" * 50 + b"C" * 50 + assert bp.fetch(250, 300) == b"C" * 50 + assert bp.fetch(300, 310) == b"" + + +def test_fetch_default_args_and_out_of_bounds(prefetcher_factory): + fetcher = MockFetcher(b"12345") + bp = prefetcher_factory(fetcher=fetcher, size=5, concurrency=4) + + assert bp.fetch(None, None) == b"12345" + assert bp.fetch(None, 2) == b"12" + assert bp.fetch(5, 10) == b"" + assert bp.fetch(10, 20) == b"" + assert bp.fetch(2, 2) == b"" + assert bp.fetch(4, 2) == b"" + + +def test_seek_logic(prefetcher_factory): + data = b"0123456789" * 10 + fetcher = MockFetcher(data) + bp = prefetcher_factory(fetcher=fetcher, size=100, concurrency=4) + + assert bp.fetch(0, 10) == data[0:10] + assert bp.fetch(10, 20) == data[10:20] + assert bp.user_offset == 20 + assert bp.fetch(50, 60) == data[50:60] + assert bp.user_offset == 60 + assert bp.fetch(10, 20) == data[10:20] + assert bp.user_offset == 20 + + +def test_exception_placed_in_queue(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + async def inject_error(): + await bp.queue.put(ValueError("Injected Producer Error")) + + fsspec.asyn.sync(bp.loop, inject_error) + + with pytest.raises(ValueError, match="Injected Producer Error"): + bp.fetch(0, 50) + + assert isinstance(bp._error, ValueError) + + +def test_producer_concurrency_streak_and_min_chunk(prefetcher_factory): + data = b"X" * 1000 + fetcher = MockFetcher(data) + + bp = prefetcher_factory(fetcher=fetcher, size=1000, concurrency=4) + bp.read_tracker.add(50) + + # Temporarily lower chunk limit for test + original_min_chunk = bp.producer.MIN_CHUNK_SIZE + bp.producer.MIN_CHUNK_SIZE = 10 + + # Do 6 reads to push the streak well past the MIN_STREAKS threshold + # Update these values as BackgroundPrefetcher constant changes. + target_streak = bp.producer.MIN_STREAKS_FOR_PREFETCHING + 3 + for i in range(target_streak): + bp.fetch(i * 50, (i + 1) * 50) + + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + + split_factors = [call["split_factor"] for call in fetcher.calls] + assert split_factors[0] == 4 + assert max(split_factors) > 1 + assert max(split_factors) <= 4 + + bp.producer.MIN_CHUNK_SIZE = original_min_chunk + + +def test_producer_loop_space_constraints(prefetcher_factory): + data = b"Y" * 100 + fetcher = MockFetcher(data) + + bp = prefetcher_factory(fetcher=fetcher, size=100, concurrency=4) + bp.read_tracker.add(60) + + original_min_chunk = bp.producer.MIN_CHUNK_SIZE + bp.producer.MIN_CHUNK_SIZE = 200 + + assert bp.fetch(0, 10) == b"Y" * 10 + + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + sizes = [call["size"] for call in fetcher.calls] + assert all(s <= 100 for s in sizes) + + bp.producer.MIN_CHUNK_SIZE = original_min_chunk + + +def test_producer_error_propagation_and_recovery(prefetcher_factory): + fetcher = MockFetcher(b"A" * 2000, fail_at_call=3) + bp = prefetcher_factory(fetcher=fetcher, size=2000, concurrency=4) + + for i in range(2): + bp.fetch(i * 100, (i + 1) * 100) + + # 3rd read triggers the network timeout + with pytest.raises(OSError, match="Simulated Network Timeout"): + bp.fetch(400, 500) + + # The prefetcher is now in an error state + assert isinstance(bp._error, OSError) + + # Disable the mock failure so it can succeed on retry + fetcher.fail_at_call = None + + # The next fetch should seamlessly recover, wiping the error and returning data + data = bp.fetch(400, 500) + assert data == b"A" * 100 + assert bp._error is None + + +def test_read_after_close(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + bp.close() + + assert bp.is_stopped is True + with pytest.raises(RuntimeError, match="The file instance has been closed"): + bp.fetch(0, 10) + + +def test_read_recovers_after_error(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + # Simulate an error state from a previous failed read + bp._error = ValueError("Pre-existing error") + + # The new error-recovery logic allows a subsequent read to clear the error and succeed + assert bp.fetch(0, 10) == b"X" * 10 + assert bp._error is None + + +def test_empty_queue_when_stopped(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 500), size=500, concurrency=4) + bp.is_stopped = True + + with pytest.raises(RuntimeError, match="The file instance has been closed"): + bp.fetch(0, 100) + + +def test_cancel_all_tasks_cleans_queue_with_exceptions(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + async def inject_task(): + async def dummy_exception_task(): + raise ValueError("Hidden error") + + task = asyncio.create_task(dummy_exception_task()) + await bp.queue.put(task) + await asyncio.sleep(0.05) + + fsspec.asyn.sync(bp.loop, inject_task) + bp.close() + assert bp.queue.empty() + + +def test_cleanup_cancels_active_tasks(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"Z" * 1000), size=1000, concurrency=4) + + async def inject_task(): + async def dummy_task(): + await asyncio.sleep(3) + + task = asyncio.create_task(dummy_task()) + bp.producer._active_tasks.add(task) + + fsspec.asyn.sync(bp.loop, inject_task) + + assert len(bp.producer._active_tasks) > 0 + assert bp.is_stopped is False + + bp.close() + + assert bp.is_stopped is True + assert len(bp.producer._active_tasks) == 0 + + +def test_read_task_cancellation(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 1000), size=1000, concurrency=4) + + async def inject_and_read(): + bp.is_stopped = True + while not bp.queue.empty(): + bp.queue.get_nowait() + + cancel_task = asyncio.create_task(asyncio.sleep(10)) + cancel_task.cancel() + await bp.queue.put(cancel_task) + + with pytest.raises(asyncio.CancelledError): + await bp.consumer.consume(10) + + fsspec.asyn.sync(bp.loop, inject_and_read) + + +def test_async_fetch_exception_trapping(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + async def bad_consume(*args, **kwargs): + raise RuntimeError("Simulated async crash") + + bp.consumer.consume = bad_consume + + with pytest.raises(RuntimeError, match="Simulated async crash"): + bp.fetch(0, 10) + + # Orchestrator should capture the error internally and halt producer processing correctly + assert isinstance(bp._error, RuntimeError) + assert bp.producer.is_stopped is True + + +def test_read_past_eof_internal(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 50), size=50, concurrency=4) + bp.user_offset = 50 + res = bp.fetch(50, 60) + assert res == b"" + + +def test_fetch_with_exact_block_matches(prefetcher_factory): + data = b"X" * 100 + bp = prefetcher_factory(fetcher=MockFetcher(data), size=100, concurrency=4) + bp.read_tracker.add(50) + + assert bp.fetch(0, 50) == b"X" * 50 + assert bp.consumer._current_block_idx == 50 + assert bp.fetch(50, 100) == b"X" * 50 + + +def test_queue_empty_race_condition(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + async def inject(): + bp.queue.put_nowait(asyncio.create_task(asyncio.sleep(0))) + with mock.patch.object(bp.queue, "get_nowait", side_effect=asyncio.QueueEmpty): + await bp.producer.stop() + + fsspec.asyn.sync(bp.loop, inject) + + +def test_producer_space_remaining_break(prefetcher_factory): + bp = prefetcher_factory( + fetcher=MockFetcher(b"X" * 1000), + size=1000, + concurrency=4, + max_prefetch_size=150, + ) + bp.fetch(0, 10) + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + + +def test_producer_min_chunk_logic(prefetcher_factory): + bp1 = prefetcher_factory( + fetcher=MockFetcher(b"X" * 1000), + size=1000, + concurrency=4, + max_prefetch_size=300, + ) + bp1.producer.MIN_CHUNK_SIZE = 100 + + fsspec.asyn.sync(bp1.loop, asyncio.sleep, 0.1) + + bp2 = prefetcher_factory( + fetcher=MockFetcher(b"X" * 1000), + size=1000, + concurrency=4, + max_prefetch_size=150, + ) + bp2.producer.MIN_CHUNK_SIZE = 100 + fsspec.asyn.sync(bp2.loop, asyncio.sleep, 0.1) + + +def test_producer_loop_exception(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"A" * 100), size=100, concurrency=4) + error_object = ValueError("Producer crash") + + with mock.patch( + "fsspec.prefetcher.RunningAverageTracker.average", new_callable=mock.PropertyMock + ) as mocked_avg: + mocked_avg.side_effect = error_object + with pytest.raises(ValueError, match="Producer crash"): + bp.fetch(0, 10) + + assert bp.is_stopped is False + assert bp._error == error_object + + +def test_seek_same_offset(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b""), size=100, concurrency=4) + bp.fetch(0, 10) + + +def test_read_history_maxlen(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 2000), size=2000, concurrency=4) + for i in range(12): + bp.fetch(i * 10, (i + 1) * 10) + assert len(bp.read_tracker._history) == 10 + + +def test_fast_slice_branch(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 200), size=200, concurrency=4) + assert bp.fetch(0, 10) == b"X" * 10 + assert bp.fetch(10, 20) == b"X" * 10 + + +def test_async_fetch_not_block_break(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b""), size=100, concurrency=4) + + async def fake_consume(size): + return b"" + + bp.consumer.consume = fake_consume + bp.user_offset = 0 + + res = bp.fetch(0, 50) + assert res == b"" + + +def test_fetch_stopped_before_execution(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + bp.is_stopped = True + bp._error = None + + with pytest.raises(RuntimeError, match="The file instance has been closed"): + bp.fetch(0, 10) + + +def test_async_fetch_zero_copy_remainder(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X"), size=100, concurrency=4) + bp.consumer._current_block = b"ABCDE" + bp.consumer._current_block_idx = 0 + bp.user_offset = 0 + res = bp.fetch(0, 5) + assert res == b"ABCDE" + assert bp.consumer._current_block_idx == 5 + + +def test_read_runtime_error_on_stopped_empty(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X"), size=100, concurrency=4) + bp.is_stopped = True + bp.producer.is_stopped = True + + while not bp.queue.empty(): + bp.queue.get_nowait() + + res = fsspec.asyn.sync(bp.loop, bp.consumer.consume, 10) + assert res == b"" + + +def test_init_invalid_max_prefetch_size(): + with pytest.raises( + ValueError, + match=r"max_prefetch_size should be a positive integer", + ): + BackgroundPrefetcher( + fetcher=MockFetcher(b""), size=1000, concurrency=4, max_prefetch_size=0 + ) + + +def test_init_valid_max_prefetch_size_edge_case(prefetcher_factory): + bp = prefetcher_factory( + fetcher=MockFetcher(b""), size=1000, concurrency=4, max_prefetch_size=100 + ) + assert bp.producer._user_max_prefetch_size == 100 + + +def test_consumer_zero_size_checks(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + # 1. Test consume size <= 0 + res_consume_zero = fsspec.asyn.sync(bp.loop, bp.consumer.consume, 0) + assert res_consume_zero == b"" + res_consume_neg = fsspec.asyn.sync(bp.loop, bp.consumer.consume, -5) + assert res_consume_neg == b"" + + # 2. Test _advance size <= 0 directly + # (consume catches it early, so we call _advance directly to hit its internal check) + res_advance_zero = fsspec.asyn.sync( + bp.loop, bp.consumer._advance, 0, save_data=True + ) + assert res_advance_zero == [] + res_advance_neg = fsspec.asyn.sync( + bp.loop, bp.consumer._advance, -10, save_data=False + ) + assert res_advance_neg == [] + + +def test_producer_min_chunk_inner_break(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + bp = prefetcher_factory( + fetcher=fetcher, size=1000, concurrency=4, max_prefetch_size=400 + ) + + bp.read_tracker.add(100) + + original_min_chunk = bp.producer.MIN_CHUNK_SIZE + bp.producer.MIN_CHUNK_SIZE = 200 + + async def trigger_loop(): + bp.producer.current_offset = 250 + bp.consumer.offset = 0 + bp.consumer.target_offset = 0 + # streak=6 makes prefetch_multiplier = 4 (6 - 3 + 1) + # prefetch_size = 4 * 100 = 400 + bp.consumer.sequential_streak = 6 + bp.wakeup_event.set() + await asyncio.sleep(0.05) + + fsspec.asyn.sync(bp.loop, trigger_loop) + + assert fetcher.call_count == 0 + + bp.producer.MIN_CHUNK_SIZE = original_min_chunk + + +def test_producer_loop_break_on_stopped_after_wakeup(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + bp = prefetcher_factory(fetcher=fetcher, size=1000, concurrency=4) + + async def trigger_stop_and_wake(): + bp.producer.is_stopped = True + bp.wakeup_event.set() + await asyncio.sleep(0.05) + + fsspec.asyn.sync(bp.loop, trigger_stop_and_wake) + + # Verify the producer gracefully exited without doing work + assert fetcher.call_count == 0 + + +def test_massive_read_disables_proactive_prefetching(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + + # max_prefetch_size = 40 + bp = prefetcher_factory( + fetcher=fetcher, size=1000, concurrency=4, max_prefetch_size=40 + ) + + # Do enough reads to build a sequential streak and trigger large averages + # Reading 60 bytes at a time. Average = 60. Threshold = 50. + for i in range(4): + bp.fetch(i * 60, (i + 1) * 60) + + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + + # Because average (60) > threshold (40), prefetch_multiplier is pinned to 1. + # The producer should only fetch what the user specifically read (4 * 60 = 240) + # and should NOT have pre-fetched any additional data ahead into the queue. + assert bp.producer.current_offset == 240 + + +def test_normal_read_allows_proactive_prefetching(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + + # max_prefetch_size = 200 makes dynamic threshold = 100 + bp = prefetcher_factory( + fetcher=fetcher, size=1000, concurrency=4, max_prefetch_size=200 + ) + + # Reading 60 bytes at a time. Average = 60. Threshold = 100. + for i in range(4): + bp.fetch(i * 60, (i + 1) * 60) + + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + + # Because average (60) <= threshold (100), the producer allows prefetching. + # It calculates a normal prefetch_multiplier > 1 and pre-fetches data ahead. + assert bp.producer.current_offset > 240 + + +def test_target_offset_expands_prefetch(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + bp = prefetcher_factory(fetcher=fetcher, size=1000, concurrency=4) + + # Seed tracker to keep the default `max_prefetch_size` calculation small + bp.read_tracker.add(10) + + # The consumer requests a massive chunk (500 bytes), far exceeding normal prefetch windows + bp.fetch(0, 500) + + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + + # The new target_offset logic should explicitly tell the producer to expand its + # boundary to cover the requested 500 bytes, overriding the tiny multiplier logic. + assert bp.consumer.target_offset == 500 + assert bp.producer.current_offset >= 500 + + +def test_producer_min_chunk_inner_empty_queue_shrink(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + bp = prefetcher_factory( + fetcher=fetcher, size=1000, concurrency=4, max_prefetch_size=400 + ) + + bp.read_tracker.add(100) + + original_min_chunk = bp.producer.MIN_CHUNK_SIZE + bp.producer.MIN_CHUNK_SIZE = 200 + + async def trigger_loop(): + # Setup conditions where the queue is empty and the user is waiting + # This makes prefetch_space_available exactly equal to prefetch_size + bp.producer.current_offset = 0 + bp.consumer.offset = 0 + bp.consumer.target_offset = 0 + bp.consumer.sequential_streak = 6 + bp.wakeup_event.set() + await asyncio.sleep(0.05) + + fsspec.asyn.sync(bp.loop, trigger_loop) + + # Because space_available == prefetch_size, it triggers the shrink condition + # instead of breaking, ensuring the blocked consumer gets its data. + assert fetcher.call_count > 0 + + bp.producer.MIN_CHUNK_SIZE = original_min_chunk + + +def test_async_context_manager_and_afetch(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + async def run_async(): + async with bp as ctx: + res = await ctx.afetch(0, 10) + assert res == b"X" * 10 + # Test default bounds -> (0, 100). Validates a backwards hard seek internally. + res_all = await ctx.afetch(None, None) + assert len(res_all) == 100 + + assert bp.is_stopped is True + assert bp.consumer._current_block == b"" # Buffer cleanly cleared on async exit + + fsspec.asyn.sync(bp.loop, run_async) + + +def test_init_with_explicit_loop(prefetcher_factory): + """Verify that passing an explicit loop assigns it correctly.""" + loop = fsspec.asyn.get_loop() + bp = prefetcher_factory( + fetcher=MockFetcher(b"X"), size=100, concurrency=1, loop=loop + ) + assert bp.loop is loop + + +@pytest.mark.asyncio +async def test_init_with_running_loop(): + """Verify asynchronous=True behavior where it inherits the user's running loop.""" + current = asyncio.get_running_loop() + # explicitly passing loop=None simulates asynchronous=True behavior + bp = BackgroundPrefetcher( + fetcher=MockFetcher(b"X"), size=100, concurrency=1, loop=None + ) + assert bp.loop is current + await bp.aclose() + + +def test_init_within_fsspec_loop(prefetcher_factory): + """Verify the edge case where the prefetcher is initialized while already running inside the target loop.""" + loop = fsspec.asyn.get_loop() + + async def init_inside_loop(): + bp = BackgroundPrefetcher( + fetcher=MockFetcher(b"X"), size=100, concurrency=1, loop=loop + ) + assert bp.loop is loop + await bp.aclose() + + # Submit the initialization directly onto the background thread + fsspec.asyn.sync(loop, init_inside_loop) + + +def test_init_no_loop_raises_error(): + """Verify synchronous execution strictly fails if no explicit loop or active loop is found.""" + with pytest.raises(RuntimeError, match="No event loop found"): + BackgroundPrefetcher( + fetcher=MockFetcher(b"X"), size=100, concurrency=1, loop=None + ) + + +def test_local_seek_optimization(prefetcher_factory): + data = b"0123456789" * 10 + fetcher = MockFetcher(data) + bp = prefetcher_factory(fetcher=fetcher, size=100, concurrency=4) + + # First fetch: read first 50 bytes (0-50) + # This will trigger a fetch. The block size fetched will be 50 bytes (covering [0, 50]). + assert bp.fetch(0, 50) == data[0:50] + assert len(bp.consumer._current_block) == 50 + initial_fetch_calls = fetcher.call_count + assert initial_fetch_calls > 0 + + # 1. Perform a backward seek *within* the currently buffered block [0, 50]. + # Seek back to 5 and read 10 bytes (5 to 15). + # This should be a zero-cost local seek and NOT increment the fetcher call count. + assert bp.fetch(5, 15) == data[5:15] + assert fetcher.call_count == initial_fetch_calls + assert bp.user_offset == 15 + + # 2. Perform a forward seek *within* the currently buffered block [0, 50]. + # Seek forward to 30 and read 10 bytes (30 to 40). + # This should also be resolved locally and NOT increment the fetcher call count. + assert bp.fetch(30, 40) == data[30:40] + assert fetcher.call_count == initial_fetch_calls + assert bp.user_offset == 40 + + # 3. Read further to trigger a hard seek and load the next block [60, 100]. + assert bp.fetch(60, 80) == data[60:80] + calls_after_next = fetcher.call_count + assert calls_after_next > initial_fetch_calls + + # The currently buffered block in the consumer is now [60, 100]. + # Seek backward to 10 (which is outside [60, 100]). + # This should trigger a hard seek and increment the fetcher call count. + assert bp.fetch(10, 20) == data[10:20] + assert fetcher.call_count > calls_after_next + + +@mock.patch("fsspec.prefetcher.HAS_CPYTHON_API", False) +def test_fast_slice_pypy_fallback(): + """ + Tests that when HAS_CPYTHON_API is False (e.g., on PyPy), _fast_slice + correctly falls back to standard Python slicing and respects boundaries. + """ + src = b"0123456789_pypy_fallback_test" + + # 1. Verify standard extraction works perfectly + assert _fast_slice(src, 11, 13) == b"pypy_fallback" + + # 2. Verify zero-length read + assert _fast_slice(src, 5, 0) == b"" + + # 3. Verify exact bounds + assert _fast_slice(src, 0, len(src)) == src \ No newline at end of file From 33543a0114e051107c4985daf1f614ffae0ab481 Mon Sep 17 00:00:00 2001 From: Prince Kumar Date: Fri, 7 Aug 2026 11:18:47 +0000 Subject: [PATCH 5/6] adding more test for adpative readahead cache --- fsspec/tests/test_caches.py | 39 +++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/fsspec/tests/test_caches.py b/fsspec/tests/test_caches.py index 410a8a545..b290d8249 100644 --- a/fsspec/tests/test_caches.py +++ b/fsspec/tests/test_caches.py @@ -2,6 +2,7 @@ import pickle import string +import fsspec import pytest from fsspec.caching import ( @@ -383,6 +384,44 @@ def test_adaptive_cache_with_async_fetcher(): cache.close() +def test_adaptive_cache_registered(): + assert "adaptive" in caches + + +def test_adaptive_cache_fallback_when_prefetcher_init_fails(monkeypatch): + import fsspec.prefetcher as prefetcher_mod + + class FailingPrefetcher: + def __init__(self, *args, **kwargs): + raise RuntimeError("prefetch init failed") + + monkeypatch.setattr(prefetcher_mod, "BackgroundPrefetcher", FailingPrefetcher) + + data = string.ascii_letters.encode() + cache = caches["adaptive"](8, letters_fetcher, len(data)) + + # On prefetch setup failure, adaptive must still serve reads through readahead. + assert cache._fetch(0, 10) == data[0:10] + assert cache._fetch(10, 17) == data[10:17] + assert cache._prefetcher is None + + +def test_adaptive_cache_selected_in_open_flow(): + from fsspec.spec import AbstractBufferedFile + + data = string.ascii_letters.encode() + + class TestFile(AbstractBufferedFile): + DEFAULT_BLOCK_SIZE = 8 + + def _fetch_range(self, start, end): + return data[start:end] + + with TestFile(None, "afile", mode="rb", cache_type="adaptive", size=len(data)) as f: + assert f.cache.name == "adaptive" + assert f.read(12) == data[:12] + + def test_adaptive_cache_fallback_without_loop(): data = string.ascii_letters.encode() cache = caches["adaptive"](8, letters_fetcher, len(data)) From e7296232392c9c9e3f1a6da123bbdfa6186f6585 Mon Sep 17 00:00:00 2001 From: Prince Kumar Date: Fri, 7 Aug 2026 16:50:28 +0000 Subject: [PATCH 6/6] linting issue --- fsspec/implementations/local.py | 3 ++- fsspec/prefetcher.py | 15 +++++---------- fsspec/tests/test_caches.py | 2 -- fsspec/tests/test_prefetcher.py | 7 ++++--- 4 files changed, 11 insertions(+), 16 deletions(-) diff --git a/fsspec/implementations/local.py b/fsspec/implementations/local.py index 5ea89f8fd..4623df4d9 100644 --- a/fsspec/implementations/local.py +++ b/fsspec/implementations/local.py @@ -8,11 +8,12 @@ import tempfile from functools import lru_cache -from ..spec import AbstractFileSystem from fsspec.compression import compr from fsspec.core import get_compression from fsspec.utils import isfilelike, stringify_path +from ..spec import AbstractFileSystem + logger = logging.getLogger("fsspec.local") diff --git a/fsspec/prefetcher.py b/fsspec/prefetcher.py index ed32d4a8c..e73b40e64 100644 --- a/fsspec/prefetcher.py +++ b/fsspec/prefetcher.py @@ -237,9 +237,7 @@ async def stop(self): self._producer_task.cancel() tasks_to_wait.append(self._producer_task) - for task in list(self._active_tasks): - if not task.done(): - tasks_to_wait.append(task) + tasks_to_wait.extend(task for task in self._active_tasks if not task.done()) # We do not cancel the network task, instead we wait on them. # This is intentionally done to avoid MRD stream disruption. @@ -301,10 +299,9 @@ async def _loop(self): except asyncio.CancelledError: logger.debug("PrefetchProducer loop was cancelled.") except Exception as e: - logger.error( + logger.exception( "PrefetchProducer loop encountered an unexpected error: %s", e, - exc_info=True, ) self.is_stopped = True self.orchestrator.set_error(e) @@ -560,7 +557,7 @@ async def _advance(self, size: int, save_data: bool) -> list[bytes]: except asyncio.CancelledError: raise except Exception as e: - logger.error("Consumer caught an error: %s", e, exc_info=True) + logger.exception("Consumer caught an error: %s", e) self.orchestrator.set_error(e) raise e @@ -820,9 +817,7 @@ async def _async_fetch(self, start, end): self._error = e raise except Exception as e: - logger.error( - "Exception raised during asynchronous fetch: %s", e, exc_info=True - ) + logger.exception("Exception raised during asynchronous fetch: %s", e) self._error = e if self.producer and not self.producer.is_stopped: await self.producer.stop() @@ -880,4 +875,4 @@ async def aclose(self): def close(self): """Safely shuts down the prefetcher from a synchronous context.""" - fsspec_asyn.sync(self.loop, self._async_close) \ No newline at end of file + fsspec_asyn.sync(self.loop, self._async_close) diff --git a/fsspec/tests/test_caches.py b/fsspec/tests/test_caches.py index b290d8249..f486f9a76 100644 --- a/fsspec/tests/test_caches.py +++ b/fsspec/tests/test_caches.py @@ -1,8 +1,6 @@ -import asyncio import pickle import string -import fsspec import pytest from fsspec.caching import ( diff --git a/fsspec/tests/test_prefetcher.py b/fsspec/tests/test_prefetcher.py index 26f3d69bc..6a03ad83c 100644 --- a/fsspec/tests/test_prefetcher.py +++ b/fsspec/tests/test_prefetcher.py @@ -1,9 +1,9 @@ import asyncio from unittest import mock -import fsspec.asyn import pytest +import fsspec.asyn from fsspec.prefetcher import BackgroundPrefetcher, RunningAverageTracker, _fast_slice @@ -368,7 +368,8 @@ def test_producer_loop_exception(prefetcher_factory): error_object = ValueError("Producer crash") with mock.patch( - "fsspec.prefetcher.RunningAverageTracker.average", new_callable=mock.PropertyMock + "fsspec.prefetcher.RunningAverageTracker.average", + new_callable=mock.PropertyMock, ) as mocked_avg: mocked_avg.side_effect = error_object with pytest.raises(ValueError, match="Producer crash"): @@ -723,4 +724,4 @@ def test_fast_slice_pypy_fallback(): assert _fast_slice(src, 5, 0) == b"" # 3. Verify exact bounds - assert _fast_slice(src, 0, len(src)) == src \ No newline at end of file + assert _fast_slice(src, 0, len(src)) == src