Read every page of S3 model versions

great_ai/large_file/large_file/large_file_s3.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 
 import boto3
 
@@ -70,22 +70,23 @@ class LargeFileS3(LargeFileBase):
     def _find_remote_instances(self) -> List[DataInstance]:
         logger.debug(f"Fetching S3 versions of {self._name}")
 
-        found_objects = self._client.list_objects_v2(
-            Bucket=self.bucket_name, Prefix=self._name
-        )
-        return (
-            [
-                DataInstance(
-                    name=o["Key"].split(S3_NAME_VERSION_SEPARATOR)[0],
-                    version=int(o["Key"].split(S3_NAME_VERSION_SEPARATOR)[-1]),
-                    remote_path=o["Key"],
-                )
-                for o in found_objects["Contents"]
-                if o["Key"].split(S3_NAME_VERSION_SEPARATOR)[0] == self._name
-            ]
-            if "Contents" in found_objects
-            else []
-        )
+        prefix = f"{self._name}{S3_NAME_VERSION_SEPARATOR}"
+        parameters: Dict[str, Any] = {"Bucket": self.bucket_name, "Prefix": prefix}
+        instances = []
+        while True:
+            page = self._client.list_objects_v2(**parameters)
+            for item in page.get("Contents", []):
+                key = item["Key"]
+                suffix = key[len(prefix) :] if key.startswith(prefix) else ""
+                if suffix.isdigit():
+                    instances.append(
+                        DataInstance(
+                            name=self._name, version=int(suffix), remote_path=key
+                        )
+                    )
+            if not page.get("IsTruncated"):
+                return instances
+            parameters["ContinuationToken"] = page["NextContinuationToken"]
 
     def _download(
         self, remote_path: Any, local_path: Path, hide_progress: bool

tests/large_file/test_large_file.py

@@ -93,7 +93,7 @@ def test_initialized_with_dict(client: Any) -> None:
     )
 
     s3.list_objects_v2.assert_called_once_with(
-        Bucket=credentials["large_files_bucket_name"], Prefix="test-file"
+        Bucket=credentials["large_files_bucket_name"], Prefix="test-file/"
     )
 
     assert lf._version == 2
@@ -137,7 +137,7 @@ def test_initialized_with_file(client: Any) -> None:
 
     assert s3.list_objects_v2.called
     s3.list_objects_v2.assert_called_once_with(
-        Bucket=credentials["large_files_bucket_name"], Prefix="test-file"
+        Bucket=credentials["large_files_bucket_name"], Prefix="test-file/"
     )
 
     assert lf._version == 2

tests/large_file/test_storage_safety.py

@@ -121,3 +121,29 @@ def test_gridfs_matches_only_exact_model_names() -> None:
         reader = LargeFileMongo("model")
     assert reader.version == 0
     assert [instance.name for instance in reader._instances] == ["model"]
+
+
+def test_s3_follows_every_listing_page() -> None:
+    backend = MagicMock()
+    backend.list_objects_v2.side_effect = [
+        {
+            "Contents": [{"Key": "model/0"}, {"Key": "model-other/99"}],
+            "IsTruncated": True,
+            "NextContinuationToken": "next",
+        },
+        {
+            "Contents": [
+                {"Key": "model/9"},
+                {"Key": "model/backup"},
+                {"Key": "model/9/extra"},
+            ]
+        },
+    ]
+    with patch.object(LargeFileS3, "_client", backend):
+        reader = LargeFileS3("model")
+    assert reader.version == 9
+    assert [instance.version for instance in reader._instances] == [0, 9]
+    assert backend.list_objects_v2.call_args_list[0].kwargs["Prefix"] == "model/"
+    assert (
+        backend.list_objects_v2.call_args_list[1].kwargs["ContinuationToken"] == "next"
+    )