Add s3fs bucket assertion

awslabs · Mar 12, 2024 · c369130 · c369130
1 parent 38a5343
commit c369130
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py b/s3torchconnector/tst/e2e/test_e2e_s3_lightning_checkpoint.py
@@ -3,6 +3,8 @@
 import lightning
 import pytest
 import random
+
+import s3fs
 import torch
 
 from pathlib import Path
@@ -73,6 +75,8 @@ def test_delete_checkpoint(checkpoint_directory):
 
 
 def test_load_trained_checkpoint(checkpoint_directory):
+    # Force s3fs to acknowledge the bucket existence
+    assert s3fs.S3FileSystem().exists(checkpoint_directory.bucket)
     nonce = random.randrange(2**64)
     dataset = WikiText2(data_dir=Path(f"/tmp/data/{nonce}"))
     dataloader = DataLoader(dataset, num_workers=3)