From e9aee74eb8e29eceddfb7e449eb49874f972adb7 Mon Sep 17 00:00:00 2001 From: bigning Date: Fri, 9 Aug 2024 10:01:29 -0700 Subject: [PATCH] restore google cloud object store test (#3538) --- tests/utils/object_store/test_gs_object_store.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/utils/object_store/test_gs_object_store.py b/tests/utils/object_store/test_gs_object_store.py index b52d3cf977..93b8a05ccd 100644 --- a/tests/utils/object_store/test_gs_object_store.py +++ b/tests/utils/object_store/test_gs_object_store.py @@ -9,7 +9,6 @@ from botocore.exceptions import ClientError from torch.utils.data import DataLoader -from composer.loggers import RemoteUploaderDownloader from composer.optim import DecoupledSGDW from composer.trainer import Trainer from composer.utils import GCSObjectStore @@ -17,15 +16,15 @@ def get_gcs_os_from_trainer(trainer: Trainer) -> GCSObjectStore: - rud = [dest for dest in trainer.logger.destinations if isinstance(dest, RemoteUploaderDownloader)][0] - gcs_os = rud.remote_backend + assert trainer._checkpoint_saver is not None + assert trainer._checkpoint_saver.remote_uploader is not None + gcs_os = trainer._checkpoint_saver.remote_uploader.remote_backend assert isinstance(gcs_os, GCSObjectStore) return gcs_os @pytest.mark.gpu # json auth is hard to set up on github actions / CPU tests @pytest.mark.remote -@pytest.mark.skip(reason='Waiting for new GCP key to be approved') def test_gs_object_store_integration_hmac_auth(expected_use_gcs_sdk_val=False, client_should_be_none=True): model = SimpleModel() train_dataset = RandomClassificationDataset() @@ -35,7 +34,7 @@ def test_gs_object_store_integration_hmac_auth(expected_use_gcs_sdk_val=False, c model=model, optimizers=optimizer, train_dataloader=train_dataloader, - save_folder='gs://mosaicml-internal-integration-testing/checkpoints/{run_name}', + save_folder='gs://mosaicml-runtime-internal-integration-testing/checkpoints/{run_name}', save_filename='test-model.pt', max_duration='1ba', precision='amp_bf16', @@ -54,7 +53,7 @@ def test_gs_object_store_integration_hmac_auth(expected_use_gcs_sdk_val=False, c model=model, optimizers=optimizer, train_dataloader=train_dataloader, - load_path=f'gs://mosaicml-internal-integration-testing/checkpoints/{run_name}/test-model.pt', + load_path=f'gs://mosaicml-runtime-internal-integration-testing/checkpoints/{run_name}/test-model.pt', max_duration='2ba', precision='amp_bf16', ) @@ -64,7 +63,6 @@ def test_gs_object_store_integration_hmac_auth(expected_use_gcs_sdk_val=False, c @pytest.mark.gpu @pytest.mark.remote -@pytest.mark.skip(reason='Waiting for new GCP key to be approved') def test_gs_object_store_integration_json_auth(): with mock.patch.dict(os.environ): if 'GCS_KEY' in os.environ: