Skip to content

Checkpoint_conversion: import GCS via gcloud_stub for decoupled mode#4003

Merged
copybara-service[bot] merged 1 commit into
AI-Hypercomputer:mainfrom
ROCm:upstream-fix-decoupled-gcs-collection
May 28, 2026
Merged

Checkpoint_conversion: import GCS via gcloud_stub for decoupled mode#4003
copybara-service[bot] merged 1 commit into
AI-Hypercomputer:mainfrom
ROCm:upstream-fix-decoupled-gcs-collection

Conversation

@gulsumgudukbay
Copy link
Copy Markdown
Collaborator

@gulsumgudukbay gulsumgudukbay commented May 28, 2026

Description

src/maxtext/checkpoint_conversion/utils/utils.py does a top-level from google.cloud.storage import Client, transfer_manager. When tests run with DECOUPLE_GCLOUD=TRUE (and google-cloud-storage is not installed in that environment), pytest collection fails on this file with ModuleNotFoundError before tests/conftest.py can apply the decoupled marker 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 by src/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 the transfer_manager submodule (it isn't auto-imported by from google.cloud import storage). _gcs_stubs() is extended with a matching no-op transfer_manager stub exposing upload_many_from_filenames.
  • checkpoint_conversion/utils/utils.py: drops the direct from google.cloud.storage import ... and binds Client / transfer_manager at module top via gcs_storage(). No call sites change, the public symbol surface (Client, transfer_manager) is preserved.

(cherry picked from commit 698624b)

Tests

  • Module imports cleanly in a venv without 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_training collects 719/1064 tests with exit 0 (was 1 collection error before the fix).
  • In a venv where google-cloud-storage is installed, gcs_storage() returns the real module with both .Client and .transfer_manager attributes.
  • CI green on this branch.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

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
Copy link
Copy Markdown

codecov Bot commented May 28, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@copybara-service copybara-service Bot merged commit e664b16 into AI-Hypercomputer:main May 28, 2026
49 of 51 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants