Provide more granular methods to configure the embedded S3 client. (#325

)
predibase · Mar 13, 2024 · e7ccc63 · e7ccc63
1 parent e637d89
commit e7ccc63
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/server/lorax_server/utils/sources/s3.py b/server/lorax_server/utils/sources/s3.py
@@ -59,13 +59,21 @@ def _get_bucket_resource(bucket_name: str) -> "Bucket":
         )
     )
 
+    S3_ENDPOINT_URL = os.environ.get("S3_ENDPOINT_URL", None)
     R2_ACCOUNT_ID = os.environ.get("R2_ACCOUNT_ID", None)
+
     if R2_ACCOUNT_ID:
         s3 = boto3.resource('s3', 
             endpoint_url = f'https://{R2_ACCOUNT_ID}.r2.cloudflarestorage.com',
             config=config
            )
         return s3.Bucket(bucket_name)
+    elif S3_ENDPOINT_URL:
+        s3 = boto3.resource('s3', 
+            endpoint_url = f'{S3_ENDPOINT_URL}',
+            config=config
+           )
+        return s3.Bucket(bucket_name)
     else:
         s3 = boto3.resource('s3', config=config)
         return s3.Bucket(bucket_name)