allow retry sagemaker batch job creation for longer time window

Signed-off-by: Xun Zhang <xunzh@amazon.com>

Author: Zhangxunmt

data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/AbstractBatchJobCreator.java

@@ -31,6 +31,8 @@ public abstract class AbstractBatchJobCreator implements MLBatchJobCreator {
     public static final String NUMBER_OF_RECORDS_FAILED_IN_BATCH_JOB = "recordsFailedInBatchJobCreation";
     public static final String NUMBER_OF_RECORDS_SUCCEEDED_IN_BATCH_JOB = "recordsSucceededInBatchJobCreation";
     protected static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+    protected static final long MAX_RETRY_WINDOW_MS = 600_000; // 10 minutes
+    protected static final int TOO_MANY_REQUESTS = 429;
     protected final MLProcessorConfig mlProcessorConfig;
     protected final AwsCredentialsSupplier awsCredentialsSupplier;
     protected final Counter numberOfBatchJobsSuccessCounter;
@@ -40,7 +42,6 @@ public abstract class AbstractBatchJobCreator implements MLBatchJobCreator {
     protected final List<String> tagsOnFailure;
     protected final MlCommonRequester mlCommonRequester;
     protected DlqPushHandler dlqPushHandler = null;
-
     private static final Aws4Signer signer;
     static {
         signer = Aws4Signer.create();
@@ -119,4 +120,32 @@ protected DlqObject createDlqObjectFromEvent(final Event event,
                 .withPluginId(dlqPushHandler.getDlqPluginSetting().getName())
                 .build();
     }
+
+    class RetryRecord {
+        private final Record<Event> record;
+        private final long createdTime;
+        private int retryCount;
+
+        RetryRecord(Record<Event> record) {
+            this.record = record;
+            this.createdTime = System.currentTimeMillis();
+            this.retryCount = 0;
+        }
+
+        boolean isExpired() {
+            return System.currentTimeMillis() - createdTime > MAX_RETRY_WINDOW_MS;
+        }
+
+        void incrementRetryCount() {
+            retryCount++;
+        }
+
+        Record<Event> getRecord() {
+            return record;
+        }
+
+        int getRetryCount() {
+            return retryCount;
+        }
+    }
 }

data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/BedrockBatchJobCreator.java

@@ -35,9 +35,8 @@
 
 public class BedrockBatchJobCreator extends AbstractBatchJobCreator {
     private final AwsCredentialsSupplier awsCredentialsSupplier;
-    private static final long MAX_RETRY_WINDOW_MS = 300_000; // 5 minutes
     @Getter
-    private final ConcurrentLinkedQueue<ThrottledRecord> throttledRecords = new ConcurrentLinkedQueue<>();
+    private final ConcurrentLinkedQueue<RetryRecord> throttledRecords = new ConcurrentLinkedQueue<>();
     private final Lock processingLock;
 
     private static final String BEDROCK_PAYLOAD_TEMPLATE = "{\"parameters\": {\"inputDataConfig\": {\"s3InputDataConfig\": {\"s3Uri\": \"s3://\"}}," +
@@ -55,13 +54,13 @@ public void createMLBatchJob(List<Record<Event>> inputRecords, List<Record<Event
     }
 
     private void processRecords(List<Record<Event>> records, List<Record<Event>> resultRecords,
-                                List<ThrottledRecord> throttledRecords) {
+                                List<RetryRecord> throttledRecords) {
         List<Record<Event>> failedRecords = new ArrayList<>();
         List<DlqObject> dlqObjects = new ArrayList<>();
 
         for (int i = 0; i < records.size(); i++) {
             Record<Event> record = records.get(i);
-            ThrottledRecord throttledRecord = throttledRecords != null ? throttledRecords.get(i) : null;
+            RetryRecord throttledRecord = throttledRecords != null ? throttledRecords.get(i) : null;
 
             processRecord(record, resultRecords, failedRecords, dlqObjects, throttledRecord);
 
@@ -89,7 +88,7 @@ private void processRecords(List<Record<Event>> records, List<Record<Event>> res
 
     private void processRecord(Record<Event> record, List<Record<Event>> resultRecords,
                                List<Record<Event>> failedRecords, List<DlqObject> dlqObjects,
-                               ThrottledRecord throttledRecord) {
+                               RetryRecord throttledRecord) {
         try {
             String s3Uri = generateS3Uri(record);
             String payload = createPayloadBedrock(s3Uri, mlProcessorConfig);
@@ -122,9 +121,9 @@ private void processRecord(Record<Event> record, List<Record<Event>> resultRecor
                 if (e instanceof MLBatchJobException) {
                     MLBatchJobException mlException = (MLBatchJobException) e;
                     statusCode = mlException.getStatusCode();
-                    if (statusCode == 429) {
-                        ThrottledRecord newThrottledRecord = throttledRecord != null ?
-                                throttledRecord : new ThrottledRecord(record);
+                    if (shouldRetry(statusCode, mlException.getMessage())) {
+                        RetryRecord newThrottledRecord = throttledRecord != null ?
+                                throttledRecord : new RetryRecord(record);
                         throttledRecords.offer(newThrottledRecord);
                         LOG.info("Request {} throttled{}, added to retry queue: {}",
                                 throttledRecord != null ? "still" : "",
@@ -158,11 +157,29 @@ public void addProcessedBatchRecordsToResults(List<Record<Event>> resultRecords)
 
         try {
             processThrottledRecords(resultRecords);
+        } catch (Exception e) {
+            LOG.error("Error in batch processing throttled records. Error: {}", e.getMessage());
         } finally {
             processingLock.unlock();
         }
     }
 
+    private boolean shouldRetry(int statusCode, String errorMessage) {
+        if (statusCode == TOO_MANY_REQUESTS) {
+            return true;
+        }
+
+        if (errorMessage == null) {
+            return false;
+        }
+
+        // Check for quota-related messages
+        return (statusCode == HttpURLConnection.HTTP_BAD_REQUEST) &&
+                (errorMessage.contains("quota for number of concurrent invoke-model jobs") ||
+                        errorMessage.contains("throttling") ||
+                        errorMessage.contains("request was denied due to remote server throttling"));
+    }
+
     private void handleFailure(Record<Event> record,
                                List<Record<Event>> resultRecords,
                                List<Record<Event>> failedRecords,
@@ -203,11 +220,11 @@ private void pushToDlq(List<DlqObject> dlqObjects) {
     }
 
     private void processThrottledRecords(List<Record<Event>> resultRecords) {
-        List<ThrottledRecord> expiredRecords = new ArrayList<>();
-        List<ThrottledRecord> recordsToRetry = new ArrayList<>();
+        List<RetryRecord> expiredRecords = new ArrayList<>();
+        List<RetryRecord> recordsToRetry = new ArrayList<>();
 
         // Process throttled records
-        ThrottledRecord throttledRecord;
+        RetryRecord throttledRecord;
         while ((throttledRecord = throttledRecords.poll()) != null) {
             if (throttledRecord.isExpired()) {
                 expiredRecords.add(throttledRecord);
@@ -224,30 +241,30 @@ private void processThrottledRecords(List<Record<Event>> resultRecords) {
         retryThrottledRecords(recordsToRetry, resultRecords);
     }
 
-    private void retryThrottledRecords(List<ThrottledRecord> recordsToRetry, List<Record<Event>> resultRecords) {
+    private void retryThrottledRecords(List<RetryRecord> recordsToRetry, List<Record<Event>> resultRecords) {
         if (recordsToRetry.isEmpty()) {
             return;
         }
 
         LOG.info("Retrying {} throttled records", recordsToRetry.size());
         processRecords(
                 recordsToRetry.stream()
-                        .map(ThrottledRecord::getRecord)
+                        .map(RetryRecord::getRecord)
                         .collect(Collectors.toCollection(ArrayList::new)),
                 resultRecords,
                 recordsToRetry
         );
     }
 
-    private void handleExpiredRecords(List<ThrottledRecord> expiredRecords, List<Record<Event>> resultRecords) {
+    private void handleExpiredRecords(List<RetryRecord> expiredRecords, List<Record<Event>> resultRecords) {
         if (expiredRecords.isEmpty()) {
             return;
         }
 
         List<Record<Event>> failedRecords = new ArrayList<>();
         List<DlqObject> dlqObjects = new ArrayList<>();
 
-        for (ThrottledRecord expiredRecord : expiredRecords) {
+        for (RetryRecord expiredRecord : expiredRecords) {
             String errorMessage = String.format(
                     "Record expired after %d retries over %d minutes",
                     expiredRecord.getRetryCount(),
@@ -302,32 +319,4 @@ private String createPayloadBedrock(String S3Uri, MLProcessorConfig mlProcessorC
             throw new RuntimeException("Failed to create payload for BedRock batch job", e);
         }
     }
-
-    class ThrottledRecord {
-        private final Record<Event> record;
-        private final long createdTime;
-        private int retryCount;
-
-        ThrottledRecord(Record<Event> record) {
-            this.record = record;
-            this.createdTime = System.currentTimeMillis();
-            this.retryCount = 0;
-        }
-
-        boolean isExpired() {
-            return System.currentTimeMillis() - createdTime > MAX_RETRY_WINDOW_MS;
-        }
-
-        void incrementRetryCount() {
-            retryCount++;
-        }
-
-        Record<Event> getRecord() {
-            return record;
-        }
-
-        int getRetryCount() {
-            return retryCount;
-        }
-    }
 }