Checkpoint_conversion: import GCS via gcloud_stub for decoupled mode#4003
Merged
copybara-service[bot] merged 1 commit intoMay 28, 2026
Conversation
The top-level `from google.cloud.storage import Client, transfer_manager` in checkpoint_conversion/utils/utils.py broke pytest collection for the ROCm decoupled tests (DECOUPLE_GCLOUD=TRUE), since the package isn't installed in that environment. - gcloud_stub.gcs_storage(): also import and attach the transfer_manager submodule (it isn't auto-imported by `from google.cloud import storage`); extend _gcs_stubs() with a no-op transfer_manager stub. - checkpoint_conversion/utils/utils.py: drop the direct google.cloud import and bind Client/transfer_manager via gcs_storage(), matching the existing pattern in src/maxtext/utils/gcs_utils.py. (cherry picked from commit 698624b)
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
shralex
approved these changes
May 28, 2026
darisoy
approved these changes
May 28, 2026
e664b16
into
AI-Hypercomputer:main
49 of 51 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
src/maxtext/checkpoint_conversion/utils/utils.pydoes a top-levelfrom google.cloud.storage import Client, transfer_manager. When tests run withDECOUPLE_GCLOUD=TRUE(andgoogle-cloud-storageis not installed in that environment), pytest collection fails on this file withModuleNotFoundErrorbeforetests/conftest.pycan apply thedecoupledmarker filter to deselect the test, so the whole collection aborts.This PR routes the import through the existing
gcloud_stub.gcs_storage()helper (same pattern already used bysrc/maxtext/utils/gcs_utils.py), so the module imports cleanly in decoupled mode while behaving identically in normal mode.gcloud_stub.gcs_storage(): also imports and attaches thetransfer_managersubmodule (it isn't auto-imported byfrom google.cloud import storage)._gcs_stubs()is extended with a matching no-optransfer_managerstub exposingupload_many_from_filenames.checkpoint_conversion/utils/utils.py: drops the directfrom google.cloud.storage import ...and bindsClient/transfer_managerat module top viagcs_storage(). No call sites change, the public symbol surface (Client,transfer_manager) is preserved.(cherry picked from commit 698624b)
Tests
google-cloud-storage:python3 -c 'from maxtext.checkpoint_conversion.utils import utils; print(utils.Client, utils.transfer_manager)'.pytest --collect-only -m 'not cpu_only and not tpu_only and not post_training and decoupled and not scheduled_only' --ignore=tests/post_trainingcollects 719/1064 tests with exit 0 (was 1 collection error before the fix).google-cloud-storageis installed,gcs_storage()returns the real module with both.Clientand.transfer_managerattributes.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.