diff --git a/docs/source/prefetcher.rst b/docs/source/prefetcher.rst index a641d810..cbf5505c 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 with `DEFAULT_GCSFS_CONCURRENCY=4` and `cache_type="none"`. 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,22 @@ 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 with ``DEFAULT_GCSFS_CONCURRENCY=4`` and ``cache_type="none"``. + +To disable prefetching and revert to the legacy readahead cache, set the environment variable: .. code-block:: bash - export DEFAULT_GCSFS_CONCURRENCY=4 - export USE_EXPERIMENTAL_ADAPTIVE_PREFETCHING='true' + export USE_EXPERIMENTAL_ADAPTIVE_PREFETCHING='false' + +or 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 1bae97e6..7b7cfb4b 100644 --- a/gcsfs/core.py +++ b/gcsfs/core.py @@ -2313,7 +2313,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 +2376,27 @@ def __init__( raise OSError("Attempt to open a bucket") self.generation = _coalesce_generation(generation, path_generation) self.concurrency = kwargs.get("concurrency", DEFAULT_CONCURRENCY) + # 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. + 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", + ) + + if cache_type is None: + cache_type = "none" if use_prefetch_reader else "readahead" + super().__init__( gcsfs, path, @@ -2394,20 +2415,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 diff --git a/gcsfs/prefetcher.py b/gcsfs/prefetcher.py index 9f3b52b9..f0e551ba 100644 --- a/gcsfs/prefetcher.py +++ b/gcsfs/prefetcher.py @@ -872,4 +872,16 @@ async def aclose(self): def close(self): """Safely shuts down the prefetcher from a synchronous context.""" - fsspec.asyn.sync(self.loop, self._async_close) + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + + if loop and loop.is_running(): + loop.create_task(self._async_close()) + else: + try: + fsspec.asyn.sync(self.loop, self._async_close) + except NotImplementedError: + if self.loop and self.loop.is_running(): + self.loop.create_task(self._async_close()) diff --git a/gcsfs/tests/test_core.py b/gcsfs/tests/test_core.py index ae42c8d6..1a6caa09 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 @@ -2814,12 +2814,22 @@ 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.""" + """Verify that disabling prefetcher defaults to readahead cache unless cache_type="none" is explicitly specified.""" fn = f"{TEST_BUCKET}/no_prefetch.txt" gcs.pipe(fn, b"HelloWorld") + # When prefetcher disabled and no cache_type specified, defaults 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" + + # When prefetcher disabled and cache_type="none" explicitly specified, remains "none" + with gcs.open( + fn, "rb", cache_type="none", use_experimental_adaptive_prefetching=False + ) 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 52b7d376..336b2511 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,10 +378,15 @@ 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): with extended_gcsfs.open(b, "rb") as f: + assert isinstance(f.cache, caching.BaseCache) + assert f._prefetch_engine is not None + with extended_gcsfs.open( + b, "rb", use_experimental_adaptive_prefetching=False + ) as f: assert isinstance(f.cache, caching.ReadAheadChunked) @@ -932,7 +939,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 +992,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 +1040,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 +1109,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 +1142,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 c9d7962f..6499cbb9 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,56 @@ 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) -> cache_type="readahead_chunked" + 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_chunked" + assert zf_no_prefetch._prefetch_engine is None + zf_no_prefetch.close() + + # 3. Explicit cache_type="none" with prefetcher disabled -> cache_type="none" and no prefetcher + zf_explicit_none = ZonalFile( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="rb", + cache_type="none", + use_experimental_adaptive_prefetching=False, + ) + assert zf_explicit_none.cache_type == "none" + assert zf_explicit_none._prefetch_engine is None + zf_explicit_none.close() + + # 4. Explicit cache_type="bytes" -> cache_type="bytes" + zf_bytes = ZonalFile( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="rb", + cache_type="bytes", + ) + assert zf_bytes.cache_type == "bytes" + 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 3ced93a5..a884e2b6 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 908f82b3..9e3d5805 100644 --- a/gcsfs/zonal_file.py +++ b/gcsfs/zonal_file.py @@ -1,4 +1,5 @@ import logging +import os from fsspec import asyn from google.cloud.storage.asyncio.async_appendable_object_writer import ( @@ -28,7 +29,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", @@ -94,6 +95,18 @@ def __init__( "Only read, write and append operations are currently supported for Zonal buckets." ) + if cache_type is None: + val = kwargs.get("use_experimental_adaptive_prefetching") + if val is not None: + use_prefetch = ( + val.lower() in ("true", "1") if isinstance(val, str) else bool(val) + ) + else: + use_prefetch = os.environ.get( + "USE_EXPERIMENTAL_ADAPTIVE_PREFETCHING", "true" + ).lower() in ("true", "1") + cache_type = "none" if use_prefetch else "readahead_chunked" + super().__init__( gcsfs, path,