Repository navigation
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request extends ExtendedGcsFileSystem to support Rapid Cache Ultra (RCU) buckets via the high-performance gRPC data path, replacing the BucketType enum with a more detailed BucketLayout dataclass. It also introduces fallback handling to HTTP uploads if RCU is disabled mid-write, updates dependencies, and adds comprehensive unit tests. The reviewer suggested simplifying the _uses_grpc_data_path helper by directly returning the BucketLayout property to avoid duplication, and recommended checking fresh_layout.uses_grpc_data_path instead of CachingTier.RAPID_CACHE_ULTRA in _init_aaow for more robust error propagation.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1138 +/- ##
==========================================
+ Coverage 92.24% 92.46% +0.22%
==========================================
Files 20 20
Lines 4047 4141 +94
==========================================
+ Hits 3733 3829 +96
+ Misses 314 312 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
50b22ef to
d32c65f
Compare
zhixiangli
left a comment
zhixiangli
left a comment
There was a problem hiding this comment.
Thanks for this! I exercised the mid-run disable flow against real GCS (scratch bucket with rapid-caches create --cache-type=rapid-cache-ultra, then rapid-caches disable while a filesystem still had RCU cached), and the core design holds up:
- GetStorageLayout stops reporting RCU within a few seconds of the disable.
- The cached-RCU write open fails with
InvalidArgument 400 "Project … may not create storageClass RAPID buckets with locationConstraint …". pipe_filefalls back to HTTP and keepsmetadata.- A stale
ZonalFileraisesRapidCacheUltraDisabledErrorand still closes; new opens getGCSFile.
A few comments inline, in priority order:
ZonalFile._ensure_aaow:close()re-opens the writer after any open failure other than the RCU error. That doubles the open and the new layout re-check, and can create a 0-byte object after a failed write (reproduced on real GCS).simple_upload/initiate_uploadhave no HTTP fallback, unlike_put_file/_pipe_file._init_aaowre-checks the layout on everyGoogleAPICallError, including "object already finalized".- Nit: look up the layout once in
_put_file/_pipe_file.
I have patches and tests for 1 and 2 locally, checked against d32c65f: the new tests fail on this branch and pass with the fixes. Happy to share them if useful.
d32c65f to
8fdea6c
Compare
zhixiangli
left a comment
zhixiangli
left a comment
There was a problem hiding this comment.
Thanks for the stacked change. The re-check in _init_aaow and the HTTP retry hold up on real buckets: on a regional bucket cached as RCU, pipe_file (small and 6 MiB), put_file, core.simple_upload and core.initiate_upload all fell back to HTTP, kept metadata, and refreshed the cached layout.
One blocking issue: a GCSFile opened before invalidate_storage_layout_cache() can crash on close and leave a 0-byte object. Details inline. The remaining comments are minor.
| def invalidate_storage_layout_cache(self, bucket=None): | ||
| """Forget cached bucket layouts so the next operation looks them up again. | ||
|
|
||
| The layout of each bucket (zonal, hierarchical namespace, Rapid Cache | ||
| Ultra) is looked up the first time this filesystem instance uses the | ||
| bucket. It stays cached until it is invalidated here, or until an | ||
| appendable write to a Rapid Cache Ultra bucket fails to open and the | ||
| re-checked layout no longer reports Rapid Cache Ultra. Call this method | ||
| after enabling or disabling Rapid Cache Ultra on a bucket that this | ||
| instance already used, so that later reads and writes pick the | ||
| matching data path. It only clears in-memory state and sends no | ||
| request. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| bucket: str, optional | ||
| A bucket name, or any path inside the bucket such as | ||
| ``"gs://bucket/key"``. When omitted, the cached layouts of all | ||
| buckets are dropped. A bucket without a cached layout is ignored. | ||
| """ |
There was a problem hiding this comment.
Blocking. The docstring tells callers to invalidate after enabling RCU. A GCSFile that was already open when the cache was invalidated then breaks on close():
ExtendedGcsFileSystem._openchooses the file class once, when the file is opened.GCSFile._initiate_uploadcallscore.initiate_upload, which looks up the layout again when the file flushes. If the layout now says gRPC, it hands the call toextended_gcsfs.initiate_upload.- That call opens an
AsyncAppendableObjectWriter, whichGCSFilestores asself.locationand then passes to_call("POST", ...).
Reproduced on a real zonal bucket by opening with STANDARD cached, then invalidating, writing and closing:
AttributeError: 'AsyncAppendableObjectWriter' object has no attribute 'startswith' (core.py:622)
object after close: 0 bytes, unfinalized; writer stream not closed
This happens for every write size, because flush(force=True) always calls _initiate_upload. (GCSFile._simple_upload is never called.) The same path already existed when a layout lookup failed with UNKNOWN at open and recovered before close; this API makes it a documented case. On zonal and RCU buckets, opening the writer replaces the object immediately, so an overwrite also loses the previous contents.
Suggestion: keep GCSFile on HTTP. Move the HTTP body of core.initiate_upload into _http_initiate_upload(...) and call it from GCSFile._initiate_upload. Optionally, call it from the RapidCacheUltraDisabledError fallback in extended_gcsfs.initiate_upload too, which skips a redundant layout lookup. I checked this change by patching it in at runtime: small and 6 MiB writes succeed on a regional bucket. On a zonal bucket, the UNKNOWN-at-open case now fails with a clear HttpError 400 instead of leaving a corrupt object.
If that is out of scope here, please at least say in the docstring that file handles which are already open keep their data path and should be closed before invalidating. A regression test would also help: open with STANDARD cached, invalidate, have the lookup return STANDARD_RCU, write and close, then assert an HTTP POST and that init_aaow was not awaited.
| if await fs._uses_grpc_data_path(bucket): | ||
| finalize_on_close = fs._resolve_finalize_on_close( | ||
| await fs._lookup_bucket_layout(bucket), kwargs.get("finalize_on_close") | ||
| ) |
There was a problem hiding this comment.
Nit: _put_file and _pipe_file now look up the layout once and pass it to _resolve_finalize_on_close. For consistency, consider the same here (and in initiate_upload):
layout = await fs._lookup_bucket_layout(bucket)
if layout.uses_grpc_data_path:
finalize_on_close = fs._resolve_finalize_on_close(layout, kwargs.get("finalize_on_close"))The two *_delegates_to_core_for_non_zonal tests in test_zonal_unit.py patch _uses_grpc_data_path. With this change they would need to patch _lookup_bucket_layout instead.
| # Once RCU is disabled, new objects fail with InvalidArgument ("Project ... | ||
| # may not create storageClass RAPID buckets with locationConstraint ...") and | ||
| # appends with FailedPrecondition ("Cannot write rapid object to zone ..."). | ||
| except api_exceptions.GoogleAPICallError as e: |
There was a problem hiding this comment.
Nit: catching the base class looks right to me:
- The SDK already retries transient errors inside
writer.open()and raisesRetryErroronce retries run out.RetryErroris not aGoogleAPICallError, so transient errors never reach this handler. - The RCU error is raised only after a fresh lookup confirms that RCU is off, and the original error is kept as
__cause__. - I also saw
MethodNotImplemented(501) on an append to a bucket without RCU, so the server can return more than the two codes listed.
Could the comment say that the broad catch is intentional? As written, it reads as if only InvalidArgument and FailedPrecondition are expected.
| @pytest.mark.parametrize("func", ["simple_upload", "initiate_upload"]) | ||
| async def test_module_upload_helpers_retry_over_http_when_rcu_disabled( | ||
| rcu_fs, open_error, caplog, func | ||
| ): | ||
| """simple_upload and initiate_upload retry over HTTP after RCU was disabled.""" | ||
| fs, init_aaow = rcu_fs | ||
| fs._storage_layout_cache[BUCKET] = layouts.STANDARD_RCU | ||
| init_aaow.side_effect = open_error | ||
| fs._get_bucket_layout.return_value = layouts.STANDARD | ||
|
|
||
| with ( | ||
| mock.patch(f"gcsfs.core.{func}", new_callable=mock.AsyncMock) as http_func, | ||
| caplog.at_level(logging.WARNING, logger="gcsfs"), | ||
| ): | ||
| if func == "simple_upload": | ||
| await simple_upload(fs, BUCKET, "key", b"data", {"k": "v"}) | ||
| else: | ||
| await initiate_upload(fs, BUCKET, "key", metadata={"k": "v"}) | ||
|
|
||
| http_func.assert_awaited_once() | ||
| assert {"k": "v"} in http_func.await_args.args | ||
| assert any("Retrying the write over HTTP" in r.getMessage() for r in caplog.records) | ||
| assert not any("will be ignored" in r.getMessage() for r in caplog.records) | ||
|
|
There was a problem hiding this comment.
Nit: mocking gcsfs.core.simple_upload / initiate_upload skips the hand-off between core and extended_gcsfs. That hand-off only ends because _init_aaow refreshes the cached layout first; if that refresh regressed, core would keep handing the call back and recurse without end, and these tests would still pass. Suggest one test that patches only fs._call and zb_hns_utils.init_aaow, calls core.simple_upload, and asserts one multipart POST plus fs._storage_layout_cache[BUCKET] == layouts.STANDARD.
| @pytest.fixture | ||
| def open_error(): | ||
| """The error GCS returns when a new object's writer opens after RCU was disabled.""" | ||
| return api_exceptions.InvalidArgument( | ||
| "Project 123 may not create storageClass RAPID buckets with " | ||
| "locationConstraint US-CENTRAL1." | ||
| ) |
There was a problem hiding this comment.
Nit: the comment in _init_aaow says appends fail with FailedPrecondition ("Cannot write rapid object to zone ..."), but only InvalidArgument is tested. Consider parametrizing the fixture over both.
| @pytest.mark.parametrize( | ||
| "error, fresh", | ||
| [ | ||
| ("open_error", layouts.STANDARD_RCU), | ||
| ("open_error", BucketLayout.UNKNOWN), | ||
| (FileExistsError(PATH), None), | ||
| ], | ||
| ) | ||
| def test_zonal_file_failed_open_is_not_retried_on_close( | ||
| rcu_fs, open_error, error, fresh | ||
| ): | ||
| """close() after a failed writer open neither reopens nor re-checks the layout.""" | ||
| fs, init_aaow = rcu_fs | ||
| fs._storage_layout_cache[BUCKET] = layouts.STANDARD_RCU | ||
| init_aaow.side_effect = open_error if error == "open_error" else error | ||
| fs._get_bucket_layout.return_value = fresh | ||
|
|
||
| with pytest.raises(type(init_aaow.side_effect)): | ||
| with fs.open(PATH, "wb") as f: | ||
| f.write(b"x") | ||
|
|
||
| assert init_aaow.await_count == 1 | ||
| assert fs._get_bucket_layout.await_count == (1 if error == "open_error" else 0) |
There was a problem hiding this comment.
Nit: consider adding the empty-file case: with fs.open(PATH, "wb"): pass after RCU is disabled. Closing opens the writer, which fails. Expected: RapidCacheUltraDisabledError, f.closed is True, a second close() does nothing, and init_aaow.await_count == 1. A real run showed exactly this, so the test only locks it in.
…yout tests - Add a skip_if_grpc_read fixture that skips when TEST_BUCKET is a Zonal or RCU bucket (both read over gRPC), and use it in the test_core.py tests that assert HTTP read internals: test_cat_file_routing_and_thresholds, test_cat_file_explicit_concurrency and test_user_agent_includes_cache_type_and_source_in_read (replacing its inline zonal-only check). - Cover the HIERARCHICAL_RAPID_CACHE and UNKNOWN layouts in test_get_threshold_for_disk_reads_http_buckets_delegates_to_super. - Drop the unused _is_zonal_bucket patch from the test_fetch_range_split_concurrent_* tests; _fetch_range_split does not check the bucket type.
- zb_hns_utils.close_aaow re-raises a failed finalize, so put_file, pipe_file, simple_upload, upload_chunk and ZonalFile.close() no longer report an unfinalized write as successful. Closes that do not finalize stay best-effort. - ZonalFile._close_impl still closes the writer, without finalizing, when the closing flush fails, and then raises the flush error. - ZonalFile._ensure_aaow marks the file force-flushed when the writer fails to open, so close() does not open it again. - Document that on zonal and RCU buckets an overwrite replaces the existing object as soon as the writer opens, so a failed write leaves the partially written object in its place, unfinalized. - Tests: a failed finalize in close_aaow, a failed closing flush, a failed "xb" writer open, one ignored upload parameter per warning case, the finalize_on_close precedence matrix, and the commit and empty-file tests on the RCU layouts.
On Zonal and RCU buckets the appendable writer replaces the existing object as soon as it opens, so a write that fails after that point has already changed the bucket. - _put_file opens the local file before the writer, so a missing or unreadable lpath raises without touching the object, as on the HTTP path. - Once the writer is open, _put_file and _pipe_file update the dircache in a finally, so a failed append or finalize no longer leaves the old listing cached. - ZonalFile._close_impl invalidates the path and its parent when the closing flush fails, since fsspec's close() then skips that step. Tests: - _put_file with a missing lpath raises FileNotFoundError without opening the writer. - The put/pipe append-failure tests also check the dircache update, and a failed ZonalFile closing flush drops the cached parent listing. - With the real close_aaow, a failed finalize propagates from _put_file, _pipe_file and simple_upload, the writer is closed once, and put/pipe still update the dircache. - zonal_write_mocks gives every open an open writer that carries its bucket and object name, so later opens in a test are closed too. The manual _is_stream_open resets are gone, and test_zonal_file_overwrite checks that both opens finalize their writer.
…t when an RCU write fails to open
…pload pipe_file, put_file and simple_upload resolve finalize_on_close before the writer opens, so invalidating the layout cache while an upload runs no longer closes an object written with storage_class RAPID without finalizing it.
…heck errors The module-level simple_upload and initiate_upload now catch RapidCacheUltraDisabledError from _init_aaow, log the same warning as _put_file and _pipe_file, and fall back to the core HTTP helpers. The ignored-parameters warning is logged only after the gRPC writer opens. _put_file and _pipe_file look the bucket layout up once for routing and finalize_on_close. A comment in _init_aaow records the two errors a write gets once RCU is disabled (InvalidArgument for new objects, FailedPrecondition for appends), and the open_error fixture uses the server message seen for new objects. A new test checks that close() after a failed ZonalFile writer open neither reopens the writer nor re-checks the layout.
GCSFile._initiate_upload now calls the new core._http_initiate_upload, which holds the HTTP body of core.initiate_upload. A file opened while the bucket used the HTTP data path therefore keeps its resumable upload when the layout later resolves to RCU or zonal, for example after invalidate_storage_layout_cache(), instead of getting an AsyncAppendableObjectWriter as its session URL and failing on close. core.initiate_upload keeps its gRPC dispatch for other callers. The module-level simple_upload and initiate_upload look the bucket layout up once, and the comment above the GoogleAPICallError catch in _init_aaow says the broad catch is intentional. New tests cover a GCSFile opened before RCU was enabled, the real core -> extended_gcsfs -> core hand-off in simple_upload after RCU was disabled, FailedPrecondition append errors (the open_error fixture is now parametrized), and closing an empty ZonalFile after RCU was disabled.
8fdea6c to
d52b6ed
Compare
Summary
Adds
ExtendedGcsFileSystem.invalidate_storage_layout_cache(bucket=None)to drop cached bucket layouts, and handles Rapid Cache Ultra (RCU) being disabled on a bucket while a process is running. When an appendable write to a bucket cached as RCU fails to open with aGoogleAPICallError,_init_aaowlooks the layout up once more. If the bucket no longer reports RCU, it caches the new layout and raisesRapidCacheUltraDisabledError(anOSErrorsubclass).pipe_file,put_fileand the module-levelsimple_uploadandinitiate_uploadthen retry the upload over HTTP, and openZonalFilehandles raise the error and can only be closed.Depends on #1135 (stacked on #1134). Only the last four commits are new here.
Changes
gcsfs/extended_gcsfs.py):invalidate_storage_layout_cache(bucket=None)drops the cached layout of one bucket, given as a bucket name or any path inside it, or of every bucket whenbucketis omitted. It only clears in-memory state and is separate frominvalidate_cache(path), so writes, deletes and moves do not triggerGetStorageLayoutcalls.gcsfs/extended_gcsfs.py,gcsfs/zonal_file.py):RapidCacheUltraDisabledError(OSError)._init_aaowraises it without opening a writer when the cached layout is known and no longer uses the gRPC data path, for example in aZonalFileopened before another write refreshed the cache. AnUNKNOWNlayout still attempts the open, withoutstorage_class.google.api_core.exceptions.GoogleAPICallError,_init_aaowreuses a non-RCU layout that another write already cached, or calls_get_bucket_layoutonce. Precondition errors such asFileExistsErrorskip the re-check. If the lookup fails (UNKNOWN) or still reports RCU, the cached layout is kept and the original error is re-raised. Otherwise the new layout is cached andRapidCacheUltraDisabledErroris raised from the original error. The re-check is not limited to one status code, because once RCU is disabled new objects fail withInvalidArgumentand appends withFailedPrecondition._put_file,_pipe_fileand the module-levelsimple_uploadandinitiate_uploadcatchRapidCacheUltraDisabledErrorfrom_init_aaow, log a warning and retry over HTTP (super()._put_file/super()._pipe_file,core.simple_upload/core.initiate_upload). The ignored-parameters warning is logged only after the gRPC writer opens, and errors duringappendare not retried.ZonalFile.commit()raisesValueErrorinstead of reopening when the writer failed to open (feat(rcu): route Rapid Cache Ultra writes to the gRPC appendable writer #1135 marks such a handle force-flushed, soclose()and leaving awithblock do not open it again). OpenZonalFilehandles do not switch toGCSFile, which would turn"ab"into an overwrite.gcsfs/tests/test_rcu_unit.py: the_init_aaowre-check on flat and HNS RCU buckets (RCU disabled or still enabled, a failed lookup, zonal buckets, precondition errors, a layout refreshed by another write), the HTTP retry and warnings inpipe_file,put_file,simple_uploadandinitiate_upload, no retry after a failedappend,finalize_on_closestaying as resolved when the layout is invalidated mid-upload (pipe_file,put_file,simple_upload), andZonalFilewrites, stale handles,commit()andclose()after RCU is disabled or the writer fails to open (no second open or layout re-check).gcsfs/tests/test_hns_unit.py,gcsfs/tests/test_hns.py:invalidate_storage_layout_cache(one bucket, a path inside a bucket, an uncached bucket, all buckets, a new lookup after invalidation). Tests callinvalidate_storage_layout_cache()instead of clearing_storage_layout_cachedirectly.