From b44ac96a00156b736315134646510cdfffb8d9a3 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Tue, 28 Jul 2026 04:05:52 +0000 Subject: [PATCH] checkpoint_conversion: import GCS via gcloud_stub in load_dynamic `load_dynamic.py` imports `google.cloud.storage` at module scope, which breaks pytest collection of `tests/unit/checkpointing_test.py` in decoupled mode (`DECOUPLE_GCLOUD=TRUE`), where `google-cloud-storage` is intentionally not installed. Same failure and same fix as #4003 for `utils/utils.py`, following the existing `gcloud_stub` pattern. - `load_dynamic.py`: bind `storage = gcs_storage()` instead of importing `google.cloud.storage` directly. - `checkpointing_test.py`: three tests patched `google.cloud.storage.Client`, which cannot be resolved when the package is absent. Patch the module-level `load_dynamic.storage` binding instead, matching the existing `gcs_storage` test in the same file. Non-decoupled behavior is unchanged: `gcs_storage()` returns the real `google.cloud.storage` module when it is installed. --- src/maxtext/checkpoint_conversion/utils/load_dynamic.py | 6 +++++- tests/unit/checkpointing_test.py | 6 +++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py index 098bd403f1..079c35c00d 100644 --- a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py +++ b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py @@ -73,12 +73,12 @@ from flax import nnx import flax.traverse_util -from google.cloud import storage import huggingface_hub import jax from maxtext.checkpoint_conversion.utils import hf_model_configs from maxtext.checkpoint_conversion.utils import param_mapping from maxtext.checkpoint_conversion.utils import tensor_handling +from maxtext.common.gcloud_stub import gcs_storage from maxtext.utils import gcs_utils from maxtext.utils import globals as maxtext_globals from maxtext.utils import max_logging @@ -86,6 +86,10 @@ from orbax.checkpoint._src.arrays import sharding as sharding_utils +# Route GCS through the decoupling helper so this module imports cleanly in +# decoupled environments where google-cloud-storage is intentionally absent. +storage = gcs_storage() + HF_MODEL_CONFIGS = hf_model_configs.HF_MODEL_CONFIGS get_hf_loading_function = tensor_handling.get_hf_loading_function diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py index 421a6d150f..b4b67a6eaf 100644 --- a/tests/unit/checkpointing_test.py +++ b/tests/unit/checkpointing_test.py @@ -149,7 +149,7 @@ class LoadDynamicTest(parameterized.TestCase): """Tests for cache downloads and dynamic loading of safetensors.""" @mock.patch("huggingface_hub.HfFileSystem") - @mock.patch("google.cloud.storage.Client") + @mock.patch.object(load_dynamic.storage, "Client") def test_build_gcs_cache_worker_cache_hit(self, mock_storage_client, mock_hf_fs): mock_client_instance = mock_storage_client.return_value mock_bucket = mock_client_instance.bucket.return_value @@ -161,7 +161,7 @@ def test_build_gcs_cache_worker_cache_hit(self, mock_storage_client, mock_hf_fs) mock_blob.upload_from_file.assert_not_called() @mock.patch("huggingface_hub.HfFileSystem") - @mock.patch("google.cloud.storage.Client") + @mock.patch.object(load_dynamic.storage, "Client") def test_build_gcs_cache_worker_cache_miss_success(self, mock_storage_client, mock_hf_fs): mock_fs_instance = mock_hf_fs.return_value mock_remote_file = mock.MagicMock() @@ -177,7 +177,7 @@ def test_build_gcs_cache_worker_cache_miss_success(self, mock_storage_client, mo mock_blob.upload_from_file.assert_called_once_with(mock_remote_file, client=mock_client_instance) @mock.patch("huggingface_hub.HfFileSystem") - @mock.patch("google.cloud.storage.Client") + @mock.patch.object(load_dynamic.storage, "Client") def test_build_gcs_cache_worker_retry_and_fail(self, mock_storage_client, mock_hf_fs): mock_fs_instance = mock_hf_fs.return_value mock_fs_instance.open.side_effect = Exception("Download failed")