diff --git a/docs/source/prefetcher.rst b/docs/source/prefetcher.rst index a641d8102..110ef0cc7 100644 --- a/docs/source/prefetcher.rst +++ b/docs/source/prefetcher.rst @@ -2,8 +2,7 @@ GCSFS Adaptive Concurrent Prefetching: Architecture & Usage Guide ================================================================= -Prefetcher is not enabled by default. To enable, you need to pass the environment variable -`USE_EXPERIMENTAL_ADAPTIVE_PREFETCHING='true'` and `DEFAULT_GCSFS_CONCURRENCY`=4. As currently written, this implementation is +Prefetcher is enabled by default when cache_type is not set explicitly with `DEFAULT_GCSFS_CONCURRENCY=4`. To disable, you can pass the environment variable `USE_EXPERIMENTAL_ADAPTIVE_PREFETCHING='false'` or pass `use_experimental_adaptive_prefetching=False` when opening a file. As currently written, this implementation is separate from the fsspec-style caching layer, but the intent is to eventually make this available to all asynchronous filesystems using the standard `cache_type=` argument. How it interacts with the existing cache types ("readahead", "first", etc.) remains to be decided, and in the meantime, use at your own risk. @@ -71,17 +70,30 @@ Interaction with GCSFile The prefetcher is integrated into the ``GCSFile`` and replaces the standard sequential fetching mechanism when enabled. -Enabling the Feature --------------------- +Feature Configuration & Disabling +--------------------------------- -To use this architecture, set the following environment variables: +Adaptive prefetching is enabled by default when ``cache_type`` is not explicitly set by the user, using ``DEFAULT_GCSFS_CONCURRENCY=4``. + +Prefetching can be disabled in three ways: + +1. Explicitly specify a ``cache_type`` when opening a file (e.g., ``cache_type="readahead"`` or ``cache_type="none"`` or any other cache_type): + +.. code-block:: python + + gcs.open("bucket/file.txt", "rb", cache_type="readahead") + +2. Set the environment variable: .. code-block:: bash - export DEFAULT_GCSFS_CONCURRENCY=4 - export USE_EXPERIMENTAL_ADAPTIVE_PREFETCHING='true' + export USE_EXPERIMENTAL_ADAPTIVE_PREFETCHING='false' + +3. Pass ``use_experimental_adaptive_prefetching=False`` directly when opening a file: + +.. code-block:: python -We recommend setting ``cache_type="none"`` for optimal results. The engine avoids prefetching for random workloads, and other cache types create unnecessary memory copies that degrade performance. + gcs.open("bucket/file.txt", "rb", use_experimental_adaptive_prefetching=False) Under the Hood Lifecycle ------------------------ diff --git a/gcsfs/core.py b/gcsfs/core.py index 1bae97e60..a99b82f17 100644 --- a/gcsfs/core.py +++ b/gcsfs/core.py @@ -2305,6 +2305,34 @@ def sign(self, path, expiration=100, **kwargs): GoogleCredentials.load_tokens() +def _get_prefetcher_and_cache_config(cache_type, kwargs): + """ + Resolves effective cache_type and whether prefetch reader should be enabled. + + Rules: + - If user explicitly sets cache_type (cache_type is not None), prefetcher is disabled and cache_type is used. + - If cache_type is None and prefetcher is enabled (default), cache_type is "none" and prefetcher is active. + - If cache_type is None and prefetcher is disabled, fallback to default_cache_type. + """ + if cache_type is not None: + use_prefetch_reader = False + else: + if "use_experimental_adaptive_prefetching" in kwargs: + val = kwargs["use_experimental_adaptive_prefetching"] + use_prefetch_reader = ( + val.lower() in ("true", "1") if isinstance(val, str) else bool(val) + ) + else: + use_prefetch_reader = os.environ.get( + "USE_EXPERIMENTAL_ADAPTIVE_PREFETCHING", "true" + ).lower() in ( + "true", + "1", + ) + cache_type = "none" if use_prefetch_reader else "readahead" + return cache_type, use_prefetch_reader + + class GCSFile(fsspec.spec.AbstractBufferedFile): def __init__( self, @@ -2313,7 +2341,7 @@ def __init__( mode="rb", block_size=DEFAULT_BLOCK_SIZE, autocommit=True, - cache_type="readahead", + cache_type=None, cache_options=None, acl=None, consistency="md5", @@ -2376,6 +2404,10 @@ def __init__( raise OSError("Attempt to open a bucket") self.generation = _coalesce_generation(generation, path_generation) self.concurrency = kwargs.get("concurrency", DEFAULT_CONCURRENCY) + cache_type, use_prefetch_reader = _get_prefetcher_and_cache_config( + cache_type, kwargs + ) + super().__init__( gcsfs, path, @@ -2394,34 +2426,6 @@ def __init__( self.consistency = consistency self.checker = get_consistency_checker(consistency) - # Ideally, all of these fields should be part of `cache_options`. Because current - # `fsspec` caches do not accept arbitrary `*args` and `**kwargs`, passing them - # there currently causes instantiation errors. We are holding off on introducing - # them as explicit keyword arguments to ensure existing user workloads are not - # disrupted. This will be refactored once the upstream `fsspec` changes are merged. - use_prefetch_reader = kwargs.get( - "use_experimental_adaptive_prefetching", False - ) or os.environ.get( - "USE_EXPERIMENTAL_ADAPTIVE_PREFETCHING", "false" - ).lower() in ( - "true", - "1", - ) - - if "r" in mode and use_prefetch_reader: - max_prefetch_size = kwargs.get("max_prefetch_size", MAX_PREFETCH_SIZE) - from .prefetcher import BackgroundPrefetcher - - self._prefetch_engine = BackgroundPrefetcher( - self._async_fetch_range, - self.size, - max_prefetch_size=max_prefetch_size, - concurrency=self.concurrency, - loop=self.gcsfs.loop, - ) - else: - self._prefetch_engine = None - # _supports_append is an internal argument not meant to be used directly. # If True, allows opening file in append mode. This is generally not supported # by GCS, but may be supported by subclasses (e.g. ZonalFile). This flag should @@ -2453,6 +2457,20 @@ def __init__( self.blocksize = GCS_MIN_BLOCK_SIZE self.location = None + if "r" in mode and use_prefetch_reader: + max_prefetch_size = kwargs.get("max_prefetch_size", MAX_PREFETCH_SIZE) + from .prefetcher import BackgroundPrefetcher + + self._prefetch_engine = BackgroundPrefetcher( + self._async_fetch_range, + self.size, + max_prefetch_size=max_prefetch_size, + concurrency=self.concurrency, + loop=self.gcsfs.loop, + ) + else: + self._prefetch_engine = None + @property def details(self): if self._details is None: diff --git a/gcsfs/tests/test_core.py b/gcsfs/tests/test_core.py index ae42c8d66..035c073ac 100644 --- a/gcsfs/tests/test_core.py +++ b/gcsfs/tests/test_core.py @@ -1598,7 +1598,7 @@ def test_errors(gcs): def test_read_small(gcs): fn = TEST_BUCKET + "/2014-01-01.csv" - with gcs.open(fn, "rb", block_size=10) as f: + with gcs.open(fn, "rb", block_size=10, cache_type="readahead") as f: out = [] while True: data = f.read(3) @@ -1725,7 +1725,7 @@ def test_readline_from_cache(gcs): with gcs.open(a, "wb") as f: f.write(data) - with gcs.open(a, "rb") as f: + with gcs.open(a, "rb", cache_type="readahead") as f: result = f.readline() assert result == b"a,b\n" assert f.loc == 4 @@ -2813,13 +2813,41 @@ async def mock_fail_seq(path, start, end, **kwargs): ) -def test_gcsfile_prefetch_disabled_fallback(gcs): - """Verify that omitting the flag entirely skips the prefetcher initialization.""" - fn = f"{TEST_BUCKET}/no_prefetch.txt" +def test_gcsfile_prefetch_and_cache_type_rules(gcs): + """Verify that prefetcher is only used when cache_type is not set by user, and default cache_type is 'none'.""" + fn = f"{TEST_BUCKET}/cache_rules.txt" gcs.pipe(fn, b"HelloWorld") + # 1. Default: cache_type is not set -> prefetcher active, cache_type is "none" + with gcs.open(fn, "rb") as f: + assert getattr(f, "_prefetch_engine", None) is not None + assert f.cache_type == "none" + assert f.read() == b"HelloWorld" + + # 2. Prefetcher disabled, no cache_type set -> no prefetcher, cache_type falls back to "readahead" with gcs.open(fn, "rb", use_experimental_adaptive_prefetching=False) as f: assert getattr(f, "_prefetch_engine", None) is None + assert f.cache_type == "readahead" + assert f.read() == b"HelloWorld" + + # 3. User sets cache_type="readahead" -> prefetcher NOT used, cache_type is "readahead" + with gcs.open(fn, "rb", cache_type="readahead") as f: + assert getattr(f, "_prefetch_engine", None) is None + assert f.cache_type == "readahead" + assert f.read() == b"HelloWorld" + + # 4. User sets cache_type="readahead" even with prefetcher=True -> prefetcher NOT used + with gcs.open( + fn, "rb", cache_type="readahead", use_experimental_adaptive_prefetching=True + ) as f: + assert getattr(f, "_prefetch_engine", None) is None + assert f.cache_type == "readahead" + assert f.read() == b"HelloWorld" + + # 5. User explicitly sets cache_type="none" -> prefetcher NOT used + with gcs.open(fn, "rb", cache_type="none") as f: + assert getattr(f, "_prefetch_engine", None) is None + assert f.cache_type == "none" assert f.read() == b"HelloWorld" diff --git a/gcsfs/tests/test_extended_gcsfs.py b/gcsfs/tests/test_extended_gcsfs.py index 52b7d3760..c042acbdc 100644 --- a/gcsfs/tests/test_extended_gcsfs.py +++ b/gcsfs/tests/test_extended_gcsfs.py @@ -161,7 +161,9 @@ def test_read_small_zb(extended_gcsfs, gcs_bucket_mocks): with gcs_bucket_mocks( csv_data, bucket_type_val=BucketType.ZONAL_HIERARCHICAL ) as mocks: - with extended_gcsfs.open(csv_file_path, "rb", block_size=10) as f: + with extended_gcsfs.open( + csv_file_path, "rb", block_size=10, cache_type="readahead_chunked" + ) as f: out = [] i = 1 while True: @@ -194,7 +196,7 @@ def test_readline_zb(extended_gcsfs, gcs_bucket_mocks): def test_readline_from_cache_zb(extended_gcsfs, gcs_bucket_mocks): data = text_files["zonal/test/a"] with gcs_bucket_mocks(data, bucket_type_val=BucketType.ZONAL_HIERARCHICAL): - with extended_gcsfs.open(a, "rb") as f: + with extended_gcsfs.open(a, "rb", cache_type="readahead_chunked") as f: result = f.readline() assert result == b"a,b\n" assert f.loc == 4 @@ -376,11 +378,27 @@ def test_multithreaded_read_overlapping_ranges_zb( assert mocks["pool"].close.call_count == len(read_tasks) -def test_default_cache_is_readahead_chunked(extended_gcsfs, gcs_bucket_mocks): +def test_default_cache_is_none_with_prefetcher(extended_gcsfs, gcs_bucket_mocks): data = text_files["zonal/test/b"] with gcs_bucket_mocks(data, bucket_type_val=BucketType.ZONAL_HIERARCHICAL): + # 1. Default: cache_type not set -> prefetcher enabled, cache is BaseCache ("none") with extended_gcsfs.open(b, "rb") as f: + assert isinstance(f.cache, caching.BaseCache) + assert f._prefetch_engine is not None + + # 2. Prefetcher disabled, cache_type not set -> no prefetcher, cache falls back to ReadAhead + with extended_gcsfs.open( + b, "rb", use_experimental_adaptive_prefetching=False + ) as f: + import fsspec + + assert isinstance(f.cache, fsspec.caching.ReadAheadCache) + assert f._prefetch_engine is None + + # 3. Explicit cache_type="readahead_chunked" -> no prefetcher, cache is ReadAheadChunked + with extended_gcsfs.open(b, "rb", cache_type="readahead_chunked") as f: assert isinstance(f.cache, caching.ReadAheadChunked) + assert f._prefetch_engine is None def test_multithreaded_read_chunk_boundary_zb( @@ -932,7 +950,7 @@ def test_get_file_from_zonal_bucket(extended_gcsfs, gcs_bucket_mocks): assert f.read() == json_data if mocks: mocks["downloader"].download_ranges.assert_awaited() - mocks["downloader"].close.assert_awaited_once() + mocks["downloader"].close.assert_awaited() async def create_mrd_side_effect(client, bucket, object_name, generation): @@ -985,7 +1003,7 @@ def test_get_list_from_zonal_bucket(extended_gcsfs): with open(l2, "rb") as f: assert f.read() == files[file2] - assert mock_create_mrd.call_count == 2 + assert mock_create_mrd.call_count == 4 def test_get_directory_from_zonal_bucket(extended_gcsfs): @@ -1033,7 +1051,7 @@ def test_get_directory_from_zonal_bucket(extended_gcsfs): with open(os.path.join(local_dir, "accounts.2.json"), "rb") as f: assert f.read() == files[file2] - assert mock_create_mrd.call_count == 2 + assert mock_create_mrd.call_count == 4 @pytest.mark.asyncio @@ -1102,10 +1120,6 @@ async def mock_is_zonal(bucket): def test_read_block_zb(extended_gcsfs, gcs_bucket_mocks, subtests): - file_size = len( - json_data - ) # We need the file size to predict if readahead will trigger - for param in read_block_params: with subtests.test(id=param.id): offset, length, delimiter, expected_data = param.values @@ -1139,15 +1153,7 @@ def test_read_block_zb(extended_gcsfs, gcs_bucket_mocks, subtests): # for delimiters. We just assert that it requested ranges. assert len(actual_ranges) >= 1 else: - req_end = offset + length - if req_end >= file_size: - expected_chunks = 1 - else: - expected_chunks = 2 - - assert ( - len(actual_ranges) == expected_chunks - ), f"Expected {expected_chunks} chunks (Request + Readahead), got {len(actual_ranges)}" + assert len(actual_ranges) >= 1 actual_offsets = sorted( range_[0] for range_ in actual_ranges ) diff --git a/gcsfs/tests/test_zonal_file.py b/gcsfs/tests/test_zonal_file.py index c9d7962fc..2164f1b1d 100644 --- a/gcsfs/tests/test_zonal_file.py +++ b/gcsfs/tests/test_zonal_file.py @@ -635,7 +635,7 @@ def fake_sync(loop, func, *args, **kwargs): assert result == [b"split_data"] mock_gcsfs._fetch_range_split.assert_awaited_once_with( zf.path, - concurrency=1, + concurrency=4, start=10, chunk_lengths=[5], size=zf.size, @@ -711,7 +711,7 @@ def test_zonal_file_pool_size_initialization(mock_sync, mock_gcsfs): mode="rb", use_experimental_adaptive_prefetching=True, ) - assert zf2.pool_size == 1 + assert zf2.pool_size == 4 assert zf2._prefetch_engine is not None zf2.close() @@ -721,11 +721,67 @@ def test_zonal_file_pool_size_initialization(mock_sync, mock_gcsfs): mode="rb", use_experimental_adaptive_prefetching=False, ) - assert zf3.pool_size == 1 + assert zf3.pool_size == 4 assert zf3._prefetch_engine is None zf3.close() +@mock.patch("gcsfs.zonal_file.asyn.sync") +def test_zonal_file_cache_type_default_resolution(mock_sync, mock_gcsfs): + """Tests dynamic cache_type resolution for ZonalFile.""" + # 1. Default prefetcher enabled -> cache_type="none" + zf_default = ZonalFile( + gcsfs=mock_gcsfs, path="gs://test-bucket/test-key", mode="rb" + ) + assert zf_default.cache_type == "none" + assert zf_default._prefetch_engine is not None + zf_default.close() + + # 2. Prefetcher disabled (opt-out), no cache_type set -> cache_type="readahead", no prefetcher + zf_no_prefetch = ZonalFile( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="rb", + use_experimental_adaptive_prefetching=False, + ) + assert zf_no_prefetch.cache_type == "readahead" + assert zf_no_prefetch._prefetch_engine is None + zf_no_prefetch.close() + + # 3. Explicit cache_type="readahead_chunked" -> cache_type="readahead_chunked", no prefetcher + zf_readahead = ZonalFile( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="rb", + cache_type="readahead_chunked", + ) + assert zf_readahead.cache_type == "readahead_chunked" + assert zf_readahead._prefetch_engine is None + zf_readahead.close() + + # 4. Explicit cache_type="none" -> cache_type="none", no prefetcher + zf_explicit_none = ZonalFile( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="rb", + cache_type="none", + ) + assert zf_explicit_none.cache_type == "none" + assert zf_explicit_none._prefetch_engine is None + zf_explicit_none.close() + + # 5. Explicit cache_type="bytes" -> cache_type="bytes", no prefetcher + zf_bytes = ZonalFile( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="rb", + cache_type="bytes", + ) + assert zf_bytes.cache_type == "bytes" + assert zf_bytes._prefetch_engine is None + zf_bytes.close() + + @mock.patch("gcsfs.zonal_file.asyn.sync") def test_zonal_file_fetch_range_mutually_exclusive(mock_sync, mock_gcsfs): """Tests that providing both end and chunk_lengths raises a ValueError.""" diff --git a/gcsfs/zb_hns_utils.py b/gcsfs/zb_hns_utils.py index 3ced93a51..a884e2b6a 100644 --- a/gcsfs/zb_hns_utils.py +++ b/gcsfs/zb_hns_utils.py @@ -19,7 +19,10 @@ ) MRD_MAX_RANGES = 1000 # MRD supports up to 1000 ranges per request -DEFAULT_CONCURRENCY = int(os.environ.get("DEFAULT_GCSFS_CONCURRENCY", "1")) +try: + DEFAULT_CONCURRENCY = int(os.environ.get("DEFAULT_GCSFS_CONCURRENCY", "4")) +except ValueError: + DEFAULT_CONCURRENCY = 4 MAX_PREFETCH_SIZE = 256 * 1024 * 1024 logger = logging.getLogger("gcsfs") diff --git a/gcsfs/zonal_file.py b/gcsfs/zonal_file.py index 908f82b35..a2d08d527 100644 --- a/gcsfs/zonal_file.py +++ b/gcsfs/zonal_file.py @@ -28,7 +28,7 @@ def __init__( mode="rb", block_size=DEFAULT_BLOCK_SIZE, autocommit=True, - cache_type="readahead_chunked", + cache_type=None, cache_options=None, acl=None, consistency="md5",