feat: add endpoint for s3

Signed-off-by: Gaius <[email protected]>
dragonflyoss · Nov 20, 2023 · 9a79802 · 9a79802
1 parent 4f13967
commit 9a79802
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 11 deletions.
diff --git a/src/main/java/org/pytorch/serve/plugins/dragonfly/config/ObjectStorageConfig.java b/src/main/java/org/pytorch/serve/plugins/dragonfly/config/ObjectStorageConfig.java
@@ -49,7 +49,7 @@ public class ObjectStorageConfig {
 
   // Endpoint of OSS.
   @SerializedName("endpoint")
-  private String endPoint;
+  private String endpoint;
 
   // Access key ID of OSS.
   @SerializedName("access_key_id")
@@ -123,8 +123,8 @@ public String getContainerName() {
     return containerName;
   }
 
-  public String getEendPoint() {
-    return endPoint;
+  public String getEndpoint() {
+    return endpoint;
   }
 
   public String getAccessKeyId() {

diff --git a/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/OSS.java b/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/OSS.java
@@ -27,11 +27,11 @@ public class OSS implements ObjectStorage {
 
   public OSS(ObjectStorageConfig objectStorageConfig) {
 
-    String endPoint = objectStorageConfig.getEendPoint();
+    String endpoint = objectStorageConfig.getEndpoint();
     String accessKeyId = objectStorageConfig.getAccessKeyId();
     String accessKeySecret = objectStorageConfig.getAccessKeySecret();
 
-    ossClient = new OSSClientBuilder().build(endPoint, accessKeyId, accessKeySecret);
+    ossClient = new OSSClientBuilder().build(endpoint, accessKeyId, accessKeySecret);
   }
 
   @Override

diff --git a/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/S3.java b/src/main/java/org/pytorch/serve/plugins/dragonfly/objectstorage/S3.java
@@ -18,6 +18,7 @@
 
 import java.net.URL;
 import java.time.Duration;
+import org.apache.commons.lang3.StringUtils;
 import org.pytorch.serve.plugins.dragonfly.config.ObjectStorageConfig;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
@@ -31,16 +32,21 @@ public class S3 implements ObjectStorage {
   private S3Presigner presigner;
 
   public S3(ObjectStorageConfig objectStorageConfig) {
+    String endpoint = objectStorageConfig.getEndpoint();
     String accessKey = objectStorageConfig.getAccessKey();
     String secretKey = objectStorageConfig.getSecretKey();
     Region region = Region.of(objectStorageConfig.getRegion());
-
     AwsBasicCredentials awsCreds = AwsBasicCredentials.create(accessKey, secretKey);
-    presigner =
-        S3Presigner.builder()
-            .region(region)
-            .credentialsProvider(StaticCredentialsProvider.create(awsCreds))
-            .build();
+
+    val builder = S3Presigner.builder();
+    builder.region(region);
+    builder.credentialsProvider(StaticCredentialsProvider.create(awsCreds));
+
+    if (StringUtils.isNotBlank(endpoint)) {
+      builder.endpointOverride(new URI(endpoint));
+    }
+
+    presigner = builder.build();
   }
 
   @Override