config S3 chunk size with internal env var

Netflix · May 27, 2024 · a024c05 · a024c05
1 parent da2ad85
commit a024c05
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/metaflow/plugins/datastores/s3_storage.py b/metaflow/plugins/datastores/s3_storage.py
@@ -13,6 +13,8 @@
 from metaflow.datastore.datastore_storage import CloseAfterUse, DataStoreStorage
 
 
+S3_CHUNK_SIZE = int(os.environ.get("_S3_CHUNK_SIZE", 1024**2))
+
 try:
     # python2
     from urlparse import urlparse
@@ -151,7 +153,7 @@ def iter_results():
 
         return CloseAfterUse(iter_results(), closer=s3)
 
-    def stream_bytes(self, paths, chunk_size=2**1024):
+    def stream_bytes(self, paths, chunk_size=S3_CHUNK_SIZE):
         if len(paths) == 0:
             return CloseAfterUse(iter([]))