feat: add endpoint for s3 (#14)

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

Hide file tree

Showing 3 changed files with 17 additions and 12 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
@@ -16,6 +16,7 @@
 
 package org.pytorch.serve.plugins.dragonfly.objectstorage;
 
+import java.net.URI;
 import java.net.URL;
 import java.time.Duration;
 import org.pytorch.serve.plugins.dragonfly.config.ObjectStorageConfig;
@@ -27,20 +28,24 @@
 import software.amazon.awssdk.services.s3.presigner.model.GetObjectPresignRequest;
 
 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();
+
+    S3Presigner.Builder builder = S3Presigner.builder();
+    builder.region(region);
+    builder.credentialsProvider(StaticCredentialsProvider.create(awsCreds));
+
+    if (endpoint == null || endpoint.isEmpty()) {
+      builder.endpointOverride(URI.create(endpoint));
+    }
+
+    presigner = builder.build();
   }
 
   @Override