-
Notifications
You must be signed in to change notification settings - Fork 143
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Backport 2.x] support batch task management by periodically polling …
…the remote task via a cron job (#3458) * support batch task management by periodically bolling the remote task via a cron job (#3421) * support batch task management by periocially bolling the remote task via a cron job Signed-off-by: Bhavana Goud Ramaram <[email protected]> * address comments and resolve dependencies to avoid conflicts Signed-off-by: Bhavana Goud Ramaram <[email protected]> * add unit tests Signed-off-by: Bhavana Goud Ramaram <[email protected]> * renamed files and added more tests Signed-off-by: Bhavana Goud Ramaram <[email protected]> --------- Signed-off-by: Bhavana Goud Ramaram <[email protected]> (cherry picked from commit 161d789) * fix failing BWC tests Signed-off-by: Bhavana Goud Ramaram <[email protected]> * fix missing path in failing BWC tests Signed-off-by: Bhavana Goud Ramaram <[email protected]> * fix failing BWC tests Signed-off-by: Bhavana Goud Ramaram <[email protected]> * add missing braces Signed-off-by: Bhavana Goud Ramaram <[email protected]> * add missing braces Signed-off-by: Bhavana Goud Ramaram <[email protected]> * add missing braces Signed-off-by: Bhavana Goud Ramaram <[email protected]> * add missing braces Signed-off-by: Bhavana Goud Ramaram <[email protected]> * add missing braces Signed-off-by: Bhavana Goud Ramaram <[email protected]> * add to yml file Signed-off-by: Bhavana Goud Ramaram <[email protected]> * add to yml file Signed-off-by: Bhavana Goud Ramaram <[email protected]> * add to yml file Signed-off-by: Bhavana Goud Ramaram <[email protected]> * refactored code Signed-off-by: Bhavana Goud Ramaram <[email protected]> --------- Signed-off-by: Bhavana Goud Ramaram <[email protected]> Co-authored-by: Bhavana Goud Ramaram <[email protected]> (cherry picked from commit f083b7e)
- Loading branch information
1 parent
7552412
commit 600990c
Showing
27 changed files
with
1,544 additions
and
101 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -30,5 +30,6 @@ public enum MLTaskState { | |
| CANCELLED, | ||
| COMPLETED_WITH_ERROR, | ||
| CANCELLING, | ||
| EXPIRED | ||
| EXPIRED, | ||
| UNREACHABLE | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
86 changes: 86 additions & 0 deletions
86
ml-algorithms/src/main/java/org/opensearch/ml/engine/utils/S3Utils.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,86 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.ml.engine.utils; | ||
|
|
||
| import java.security.AccessController; | ||
| import java.security.PrivilegedActionException; | ||
| import java.security.PrivilegedExceptionAction; | ||
|
|
||
| import com.google.common.annotations.VisibleForTesting; | ||
|
|
||
| import lombok.extern.log4j.Log4j2; | ||
| import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; | ||
| import software.amazon.awssdk.auth.credentials.AwsCredentials; | ||
| import software.amazon.awssdk.auth.credentials.AwsSessionCredentials; | ||
| import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; | ||
| import software.amazon.awssdk.core.sync.RequestBody; | ||
| import software.amazon.awssdk.regions.Region; | ||
| import software.amazon.awssdk.services.s3.S3Client; | ||
| import software.amazon.awssdk.services.s3.model.PutObjectRequest; | ||
|
|
||
| @Log4j2 | ||
| public class S3Utils { | ||
| @VisibleForTesting | ||
| public static S3Client initS3Client(String accessKey, String secretKey, String sessionToken, String region) { | ||
| AwsCredentials credentials = sessionToken == null | ||
| ? AwsBasicCredentials.create(accessKey, secretKey) | ||
| : AwsSessionCredentials.create(accessKey, secretKey, sessionToken); | ||
|
|
||
| try { | ||
| S3Client s3 = AccessController | ||
| .doPrivileged( | ||
| (PrivilegedExceptionAction<S3Client>) () -> S3Client | ||
| .builder() | ||
| .region(Region.of(region)) // Specify the region here | ||
| .credentialsProvider(StaticCredentialsProvider.create(credentials)) | ||
| .build() | ||
| ); | ||
| return s3; | ||
| } catch (PrivilegedActionException e) { | ||
| throw new RuntimeException("Can't load credentials", e); | ||
| } | ||
| } | ||
|
|
||
| public static void putObject(S3Client s3Client, String bucketName, String key, String content) { | ||
| try { | ||
| AccessController.doPrivileged((PrivilegedExceptionAction<Void>) () -> { | ||
| PutObjectRequest request = PutObjectRequest.builder().bucket(bucketName).key(key).build(); | ||
|
|
||
| s3Client.putObject(request, RequestBody.fromString(content)); | ||
| log.debug("Successfully uploaded file to S3: s3://{}/{}", bucketName, key); | ||
| return null; // Void return type for doPrivileged | ||
| }); | ||
| } catch (PrivilegedActionException e) { | ||
| throw new RuntimeException("Failed to upload file to S3: s3://" + bucketName + "/" + key, e); | ||
| } | ||
| } | ||
|
|
||
| public static String getS3BucketName(String s3Uri) { | ||
| // Remove the "s3://" prefix | ||
| String uriWithoutPrefix = s3Uri.substring(5); | ||
| // Find the first slash after the bucket name | ||
| int slashIndex = uriWithoutPrefix.indexOf('/'); | ||
| // If there is no slash, the entire remaining string is the bucket name | ||
| if (slashIndex == -1) { | ||
| return uriWithoutPrefix; | ||
| } | ||
| // Otherwise, the bucket name is the substring up to the first slash | ||
| return uriWithoutPrefix.substring(0, slashIndex); | ||
| } | ||
|
|
||
| public static String getS3KeyName(String s3Uri) { | ||
| String uriWithoutPrefix = s3Uri.substring(5); | ||
| // Find the first slash after the bucket name | ||
| int slashIndex = uriWithoutPrefix.indexOf('/'); | ||
| // If there is no slash, it means there is no key, return an empty string or handle as needed | ||
| if (slashIndex == -1) { | ||
| return ""; | ||
| } | ||
| // The key name is the substring after the first slash | ||
| return uriWithoutPrefix.substring(slashIndex + 1); | ||
| } | ||
|
|
||
| } |
Oops, something went wrong.