make retry time window configurable in ml processor config

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/MLProcessorConfig.java

@@ -21,6 +21,7 @@
 import org.opensearch.dataprepper.plugins.ml_inference.processor.configuration.AwsAuthenticationOptions;
 import org.opensearch.dataprepper.plugins.ml_inference.processor.configuration.ServiceName;
 
+import java.time.Duration;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -31,6 +32,7 @@
         "It supports both synchronous and asynchronous invocations based on your use case.")
 public class MLProcessorConfig {
     private static final int DEFAULT_MAX_BATCH_SIZE = 100;
+    public static final Duration DEFAULT_RETRY_WINDOW = Duration.ofMinutes(10);
 
     @JsonProperty("aws")
     @NotNull
@@ -82,6 +84,11 @@ public class MLProcessorConfig {
     @JsonProperty("max_batch_size")
     private int maxBatchSize = DEFAULT_MAX_BATCH_SIZE;
 
+    @JsonPropertyDescription("The time duration for which the ml_inference processor retains events for retry attempts."
+            + "Supports ISO_8601 notation Strings (\"PT20.345S\", \"PT15M\", etc.) as well as simple notation Strings for seconds (\"60s\") and milliseconds (\"1500ms\")")
+    @JsonProperty("retry_time_window")
+    private Duration retryTimeWindow = DEFAULT_RETRY_WINDOW;
+
     @JsonProperty("dlq")
     private PluginModel dlq;
 

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

@@ -31,7 +31,6 @@ 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;
@@ -42,6 +41,7 @@ public abstract class AbstractBatchJobCreator implements MLBatchJobCreator {
     protected final List<String> tagsOnFailure;
     protected final MlCommonRequester mlCommonRequester;
     protected DlqPushHandler dlqPushHandler = null;
+    protected final long maxRetryTimeWindow;
     private static final Aws4Signer signer;
     static {
         signer = Aws4Signer.create();
@@ -61,6 +61,7 @@ public AbstractBatchJobCreator(MLProcessorConfig mlProcessorConfig,
         this.tagsOnFailure = mlProcessorConfig.getTagsOnFailure();
         this.mlCommonRequester = new MlCommonRequester(signer, mlProcessorConfig, awsCredentialsSupplier);
         this.dlqPushHandler = dlqPushHandler;
+        this.maxRetryTimeWindow = mlProcessorConfig.getRetryTimeWindow().toMillis();
     }
 
     // Add common logic here that both subclasses can share
@@ -133,7 +134,7 @@ class RetryRecord {
         }
 
         boolean isExpired() {
-            return System.currentTimeMillis() - createdTime > MAX_RETRY_WINDOW_MS;
+            return System.currentTimeMillis() - createdTime > maxRetryTimeWindow;
         }
 
         void incrementRetryCount() {

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

@@ -268,7 +268,7 @@ private void handleExpiredRecords(List<RetryRecord> expiredRecords, List<Record<
             String errorMessage = String.format(
                     "Record expired after %d retries over %d minutes",
                     expiredRecord.getRetryCount(),
-                    MAX_RETRY_WINDOW_MS / 60000
+                    maxRetryTimeWindow / 60000
             );
 
             LOG.error(NOISY, "Record expired from throttle queue: {}", errorMessage);

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

@@ -51,6 +51,7 @@ public class SageMakerBatchJobCreator extends AbstractBatchJobCreator {
     // Added batch processing fields
     @Getter
     private final ConcurrentLinkedQueue<Record<Event>> batch_records = new ConcurrentLinkedQueue<>();
+    @Getter
     private final ConcurrentLinkedQueue<Record<Event>> processedBatchRecords = new ConcurrentLinkedQueue<>();
     @Getter
     private final ConcurrentLinkedQueue<RetryRecord> retryQueue = new ConcurrentLinkedQueue<>();
@@ -287,7 +288,7 @@ private void processRetryQueue() {
                 expiredRecords.add(retryRecord.getRecord());
                 LOG.debug("Record expired after {} attempts over {} ms",
                         retryRecord.getRetryCount(),
-                        MAX_RETRY_WINDOW_MS);
+                        maxRetryTimeWindow);
                 return true;
             }
             return false;
@@ -312,10 +313,10 @@ private void processRetryQueue() {
 
     private void handleExpiredRecords(List<Record<Event>> expiredRecords) {
         LOG.warn("{} records expired from retry queue after {} ms timeout",
-                expiredRecords.size(), MAX_RETRY_WINDOW_MS);
+                expiredRecords.size(), maxRetryTimeWindow);
         handleFailure(expiredRecords, processedBatchRecords,
                 new MLBatchJobException(HttpURLConnection.HTTP_BAD_REQUEST,
-                        "Records expired after " + MAX_RETRY_WINDOW_MS/(60*1000) + " minute retry window"),
+                        "Records expired after " + maxRetryTimeWindow/(60*1000) + " minute retry window"),
                 HttpURLConnection.HTTP_BAD_REQUEST);
     }
 

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

@@ -36,6 +36,7 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.dataprepper.plugins.ml_inference.processor.MLProcessorConfig.DEFAULT_RETRY_WINDOW;
 import static org.opensearch.dataprepper.plugins.ml_inference.processor.common.AbstractBatchJobCreator.NUMBER_OF_FAILED_BATCH_JOBS_CREATION;
 import static org.opensearch.dataprepper.plugins.ml_inference.processor.common.AbstractBatchJobCreator.NUMBER_OF_RECORDS_FAILED_IN_BATCH_JOB;
 import static org.opensearch.dataprepper.plugins.ml_inference.processor.common.AbstractBatchJobCreator.NUMBER_OF_RECORDS_SUCCEEDED_IN_BATCH_JOB;
@@ -66,6 +67,7 @@ public class BedrockBatchJobCreatorTest {
     void setUp() {
         MockitoAnnotations.openMocks(this);
         when(mlProcessorConfig.getOutputPath()).thenReturn("s3://offlinebatch/output");
+        when(mlProcessorConfig.getRetryTimeWindow()).thenReturn(DEFAULT_RETRY_WINDOW);
         counter = new Counter() {
             @Override
             public void increment(double v) {}