diff --git a/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/MLProcessorConfig.java b/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/MLProcessorConfig.java index 7456182b50..cb4a7ed5e6 100644 --- a/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/MLProcessorConfig.java +++ b/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; diff --git a/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/AbstractBatchJobCreator.java b/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/AbstractBatchJobCreator.java index 8786e67fd5..ab070cffb6 100644 --- a/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/AbstractBatchJobCreator.java +++ b/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/AbstractBatchJobCreator.java @@ -31,6 +31,7 @@ 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 int TOO_MANY_REQUESTS = 429; protected final MLProcessorConfig mlProcessorConfig; protected final AwsCredentialsSupplier awsCredentialsSupplier; protected final Counter numberOfBatchJobsSuccessCounter; @@ -40,7 +41,7 @@ public abstract class AbstractBatchJobCreator implements MLBatchJobCreator { protected final List tagsOnFailure; protected final MlCommonRequester mlCommonRequester; protected DlqPushHandler dlqPushHandler = null; - + protected final long maxRetryTimeWindow; private static final Aws4Signer signer; static { signer = Aws4Signer.create(); @@ -60,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 @@ -119,4 +121,32 @@ protected DlqObject createDlqObjectFromEvent(final Event event, .withPluginId(dlqPushHandler.getDlqPluginSetting().getName()) .build(); } + + class RetryRecord { + private final Record record; + private final long createdTime; + private int retryCount; + + RetryRecord(Record record) { + this.record = record; + this.createdTime = System.currentTimeMillis(); + this.retryCount = 0; + } + + boolean isExpired() { + return System.currentTimeMillis() - createdTime > maxRetryTimeWindow; + } + + void incrementRetryCount() { + retryCount++; + } + + Record getRecord() { + return record; + } + + int getRetryCount() { + return retryCount; + } + } } diff --git a/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/BedrockBatchJobCreator.java b/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/BedrockBatchJobCreator.java index 2a67093092..5cb0b8786d 100644 --- a/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/BedrockBatchJobCreator.java +++ b/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 throttledRecords = new ConcurrentLinkedQueue<>(); + private final ConcurrentLinkedQueue 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> inputRecords, List> records, List> resultRecords, - List throttledRecords) { + List throttledRecords) { List> failedRecords = new ArrayList<>(); List dlqObjects = new ArrayList<>(); for (int i = 0; i < records.size(); i++) { Record 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> records, List> res private void processRecord(Record record, List> resultRecords, List> failedRecords, List dlqObjects, - ThrottledRecord throttledRecord) { + RetryRecord throttledRecord) { try { String s3Uri = generateS3Uri(record); String payload = createPayloadBedrock(s3Uri, mlProcessorConfig); @@ -122,9 +121,9 @@ private void processRecord(Record record, List> 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> 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 record, List> resultRecords, List> failedRecords, @@ -203,11 +220,11 @@ private void pushToDlq(List dlqObjects) { } private void processThrottledRecords(List> resultRecords) { - List expiredRecords = new ArrayList<>(); - List recordsToRetry = new ArrayList<>(); + List expiredRecords = new ArrayList<>(); + List recordsToRetry = new ArrayList<>(); // Process throttled records - ThrottledRecord throttledRecord; + RetryRecord throttledRecord; while ((throttledRecord = throttledRecords.poll()) != null) { if (throttledRecord.isExpired()) { expiredRecords.add(throttledRecord); @@ -224,7 +241,7 @@ private void processThrottledRecords(List> resultRecords) { retryThrottledRecords(recordsToRetry, resultRecords); } - private void retryThrottledRecords(List recordsToRetry, List> resultRecords) { + private void retryThrottledRecords(List recordsToRetry, List> resultRecords) { if (recordsToRetry.isEmpty()) { return; } @@ -232,14 +249,14 @@ private void retryThrottledRecords(List recordsToRetry, List expiredRecords, List> resultRecords) { + private void handleExpiredRecords(List expiredRecords, List> resultRecords) { if (expiredRecords.isEmpty()) { return; } @@ -247,11 +264,11 @@ private void handleExpiredRecords(List expiredRecords, List> failedRecords = new ArrayList<>(); List 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(), - MAX_RETRY_WINDOW_MS / 60000 + maxRetryTimeWindow / 60000 ); LOG.error(NOISY, "Record expired from throttle queue: {}", errorMessage); @@ -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 record; - private final long createdTime; - private int retryCount; - - ThrottledRecord(Record 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 getRecord() { - return record; - } - - int getRetryCount() { - return retryCount; - } - } } diff --git a/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/SageMakerBatchJobCreator.java b/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/SageMakerBatchJobCreator.java index cf18051ed5..440c17c33c 100644 --- a/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/SageMakerBatchJobCreator.java +++ b/data-prepper-plugins/ml-inference-processor/src/main/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/SageMakerBatchJobCreator.java @@ -29,7 +29,9 @@ import java.net.HttpURLConnection; import java.util.ArrayList; import java.util.Collection; +import java.util.HashSet; import java.util.List; +import java.util.Set; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.Lock; @@ -42,16 +44,21 @@ public class SageMakerBatchJobCreator extends AbstractBatchJobCreator { private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + private final AwsCredentialsSupplier awsCredentialsSupplier; private final S3Client s3Client; private final Lock batchProcessingLock; // Added batch processing fields @Getter private final ConcurrentLinkedQueue> batch_records = new ConcurrentLinkedQueue<>(); + @Getter private final ConcurrentLinkedQueue> processedBatchRecords = new ConcurrentLinkedQueue<>(); + @Getter + private final ConcurrentLinkedQueue retryQueue = new ConcurrentLinkedQueue<>(); private final int maxBatchSize; private final AtomicLong lastUpdateTimestamp = new AtomicLong(-1); private static final long INACTIVITY_TIMEOUT_MS = 60000; // 1 minute in milliseconds + private volatile boolean retryRecordsAddedToBatch = false; private static final String SAGEMAKER_PAYLOAD_TEMPLATE = "{\"parameters\":{\"TransformInput\":{\"ContentType\":\"application/json\"," + "\"DataSource\":{\"S3DataSource\":{\"S3DataType\":\"ManifestFile\",\"S3Uri\":\"\"}}," @@ -111,6 +118,11 @@ public void checkAndProcessBatch() { } try { + // First, process any records in the retry queue if retry records haven't been added to the current batch + if (!retryRecordsAddedToBatch) { + processRetryQueue(); + } + if (batch_records.isEmpty()) { return; } @@ -134,6 +146,8 @@ public void checkAndProcessBatch() { List> currentBatch = new ArrayList<>(batch_records); batch_records.clear(); lastUpdateTimestamp.set(-1); + // Reset the flag when we process the batch + retryRecordsAddedToBatch = false; processCurrentBatch(currentBatch); } @@ -163,28 +177,9 @@ private void processCurrentBatch(List> currentBatch) { ); if (result.isSuccess()) { - LOG.info("Successfully created SageMaker batch job for manifest URL: {}", manifestUrl); - processedBatchRecords.addAll(currentBatch); - incrementSuccessCounter(); - numberOfRecordsSuccessCounter.increment(currentBatch.size()); + handleSuccess(currentBatch, manifestUrl); } else { - Exception lastException = result.getLastException(); - int statusCode; - String errorMessage; - - if (lastException instanceof MLBatchJobException) { - MLBatchJobException mlException = (MLBatchJobException) lastException; - statusCode = mlException.getStatusCode(); - errorMessage = String.format("Failed to Create SageMaker batch job after %d attempts: %s", - result.getAttemptsMade(), mlException.getMessage()); - } else { - statusCode = HttpURLConnection.HTTP_INTERNAL_ERROR; - errorMessage = String.format("Failed to Create SageMaker batch job after %d attempts: %s", - result.getAttemptsMade(), lastException.getMessage()); - } - - handleFailure(currentBatch, processedBatchRecords, new MLBatchJobException(statusCode, errorMessage), statusCode); - LOG.error("SageMaker batch job failed for manifest URL: {}. Status: {}, Error: {}", manifestUrl, statusCode, errorMessage); + handleRetryOrFailure(currentBatch, result, manifestUrl); } } catch (IllegalArgumentException e) { LOG.error(NOISY, "Invalid arguments for SageMaker batch job. Error: {}", e.getMessage()); @@ -198,11 +193,75 @@ private void processCurrentBatch(List> currentBatch) { } } + private void handleSuccess(List> currentBatch, String manifestUrl) { + LOG.info("Successfully created SageMaker batch job for manifest URL: {}", manifestUrl); + removeCurrentBatchFromRetryQueue(currentBatch); + // Add to processed records and update counters + processedBatchRecords.addAll(currentBatch); + incrementSuccessCounter(); + numberOfRecordsSuccessCounter.increment(currentBatch.size()); + } + + private void handleRetryOrFailure(List> currentBatch, + RetryUtil.RetryResult result, String manifestUrl) { + Exception lastException = result.getLastException(); + + if (lastException instanceof MLBatchJobException) { + MLBatchJobException mlException = (MLBatchJobException) lastException; + int statusCode = mlException.getStatusCode(); + + if (statusCode == TOO_MANY_REQUESTS || + (statusCode == HttpURLConnection.HTTP_BAD_REQUEST && isThrottlingError(mlException.getMessage()))) { + handleThrottling(currentBatch); + return; + } + + String errorMessage = String.format("Failed to Create SageMaker batch job after %d attempts: %s", + result.getAttemptsMade(), mlException.getMessage()); + handleFailure(currentBatch, processedBatchRecords, + new MLBatchJobException(statusCode, errorMessage), statusCode); + LOG.error("SageMaker batch job failed for manifest URL: {}. Status: {}, Error: {}", manifestUrl, statusCode, errorMessage); + } else { + handleFailure(currentBatch, processedBatchRecords, lastException, + HttpURLConnection.HTTP_INTERNAL_ERROR); + LOG.error("SageMaker batch job failed for manifest URL: {}. Status: {}, Error: {}", manifestUrl, HttpURLConnection.HTTP_INTERNAL_ERROR, lastException.getMessage()); + } + } + + private boolean isThrottlingError(String errorMessage) { + if (errorMessage == null) { + return false; + } + + return errorMessage.toLowerCase().contains("throttling") || + errorMessage.toLowerCase().contains("request was denied due to remote server throttling"); + } + + private void handleThrottling(List> currentBatch) { + LOG.warn("Rate limited (429). Adding {} records to retry queue", currentBatch.size()); + + // Create a set of records already in retry queue + if (!retryQueue.isEmpty()) { + Set> existingRetryRecords = retryQueue.stream() + .map(RetryRecord::getRecord) + .collect(Collectors.toSet()); + + currentBatch.forEach(record -> { + if (!existingRetryRecords.contains(record)) { + retryQueue.add(new RetryRecord(record)); + } + }); + } else { + currentBatch.forEach(record -> {retryQueue.add(new RetryRecord(record));}); + } + } + private void handleFailure(List> failedRecords, ConcurrentLinkedQueue> resultRecords, Throwable throwable, int statusCode) { if (failedRecords.isEmpty()) { incrementFailureCounter(); return; } + removeCurrentBatchFromRetryQueue(failedRecords); resultRecords.addAll(addFailureTags(failedRecords)); incrementFailureCounter(); numberOfRecordsFailedCounter.increment(failedRecords.size()); @@ -222,6 +281,61 @@ private void handleFailure(List> failedRecords, ConcurrentLinkedQu } } + private void processRetryQueue() { + List> expiredRecords = new ArrayList<>(); + retryQueue.removeIf(retryRecord -> { + if (retryRecord.isExpired()) { + expiredRecords.add(retryRecord.getRecord()); + LOG.debug("Record expired after {} attempts over {} ms", + retryRecord.getRetryCount(), + maxRetryTimeWindow); + return true; + } + return false; + }); + + if (!expiredRecords.isEmpty()) { + handleExpiredRecords(expiredRecords); + } + + // Add non-expired records to batch_records + retryQueue.forEach(retryRecord -> { + batch_records.add(retryRecord.getRecord()); + retryRecord.incrementRetryCount(); + }); + retryRecordsAddedToBatch = true; // Set the flag + if (!retryQueue.isEmpty()) { + lastUpdateTimestamp.set(System.currentTimeMillis()); + LOG.info("Added {} records to batch for retry, retry queue size: {}", + retryQueue.size(), retryQueue.size()); + } + } + + private void handleExpiredRecords(List> expiredRecords) { + LOG.warn("{} records expired from retry queue after {} ms timeout", + expiredRecords.size(), maxRetryTimeWindow); + handleFailure(expiredRecords, processedBatchRecords, + new MLBatchJobException(HttpURLConnection.HTTP_BAD_REQUEST, + "Records expired after " + maxRetryTimeWindow/(60*1000) + " minute retry window"), + HttpURLConnection.HTTP_BAD_REQUEST); + } + + private void removeCurrentBatchFromRetryQueue(List> currentBatch) { + // Remove processed records from retry queue + int initialRetryQueueSize = retryQueue.size(); + if (initialRetryQueueSize > 0) { + Set> processedRecords = new HashSet<>(currentBatch); + retryQueue.removeIf(retryRecord -> + processedRecords.contains(retryRecord.getRecord())); + + int removedCount = initialRetryQueueSize - retryQueue.size(); + if (removedCount > 0) { + LOG.info("Removed {} processed records from retry queue. Remaining: {}", + removedCount, retryQueue.size()); + } + } + } + // Shutdown methods @Override public void prepareForShutdown() { diff --git a/data-prepper-plugins/ml-inference-processor/src/test/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/BedrockBatchJobCreatorTest.java b/data-prepper-plugins/ml-inference-processor/src/test/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/BedrockBatchJobCreatorTest.java index 43179938d0..19d5971d58 100644 --- a/data-prepper-plugins/ml-inference-processor/src/test/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/BedrockBatchJobCreatorTest.java +++ b/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) {} @@ -205,7 +207,7 @@ void testCreateMLBatchJob_Throttled() { // Verify throttled record was processed assertEquals(1, bedrockBatchJobCreator.getThrottledRecords().size()); - BedrockBatchJobCreator.ThrottledRecord throttledRecord = bedrockBatchJobCreator.getThrottledRecords().peek(); + BedrockBatchJobCreator.RetryRecord throttledRecord = bedrockBatchJobCreator.getThrottledRecords().peek(); assertNotNull(throttledRecord); assertEquals(1, throttledRecord.getRetryCount()); } @@ -235,7 +237,7 @@ void testCreateMLBatchJob_ThrottledMultipleTimes() { bedrockBatchJobCreator.addProcessedBatchRecordsToResults(resultRecords); // Verify retry count increased - BedrockBatchJobCreator.ThrottledRecord throttledRecord = bedrockBatchJobCreator.getThrottledRecords().peek(); + BedrockBatchJobCreator.RetryRecord throttledRecord = bedrockBatchJobCreator.getThrottledRecords().peek(); assertNotNull(throttledRecord); assertEquals(1, throttledRecord.getRetryCount()); diff --git a/data-prepper-plugins/ml-inference-processor/src/test/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/SageMakerBatchJobCreatorTest.java b/data-prepper-plugins/ml-inference-processor/src/test/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/SageMakerBatchJobCreatorTest.java index d6de6f5b4a..797e0cbc22 100644 --- a/data-prepper-plugins/ml-inference-processor/src/test/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/SageMakerBatchJobCreatorTest.java +++ b/data-prepper-plugins/ml-inference-processor/src/test/java/org/opensearch/dataprepper/plugins/ml_inference/processor/common/SageMakerBatchJobCreatorTest.java @@ -36,18 +36,21 @@ import software.amazon.awssdk.services.s3.S3Client; import software.amazon.awssdk.services.s3.model.PutObjectRequest; +import java.lang.reflect.Constructor; import java.lang.reflect.Field; import java.lang.reflect.Method; import java.net.HttpURLConnection; import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Set; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.lessThanOrEqualTo; @@ -64,11 +67,11 @@ 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; import static org.opensearch.dataprepper.plugins.ml_inference.processor.common.AbstractBatchJobCreator.NUMBER_OF_SUCCESSFUL_BATCH_JOBS_CREATION; -import static org.opensearch.dataprepper.plugins.ml_inference.processor.common.AbstractBatchJobCreator.OBJECT_MAPPER; public class SageMakerBatchJobCreatorTest { @Mock @@ -95,6 +98,12 @@ public class SageMakerBatchJobCreatorTest { @Mock private PluginSetting pluginSetting; + @Mock + private Counter recordsSuccessCounter; + + @Mock + private Counter recordsFailedCounter; + private SageMakerBatchJobCreator sageMakerBatchJobCreator; private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); @@ -112,13 +121,14 @@ void setUp() { when(mlProcessorConfig.getAwsAuthenticationOptions()).thenReturn(awsAuthenticationOptions); final EventKey sourceKey = eventKeyFactory.createEventKey("key"); when(mlProcessorConfig.getInputKey()).thenReturn(sourceKey); + when(mlProcessorConfig.getRetryTimeWindow()).thenReturn(DEFAULT_RETRY_WINDOW); when(awsAuthenticationOptions.getAwsRegion()).thenReturn(Region.US_EAST_1); when(awsCredentialsSupplier.getProvider(any())).thenReturn(awsCredentialsProvider); counter = mock(Counter.class); when(pluginMetrics.counter(NUMBER_OF_SUCCESSFUL_BATCH_JOBS_CREATION)).thenReturn(counter); when(pluginMetrics.counter(NUMBER_OF_FAILED_BATCH_JOBS_CREATION)).thenReturn(counter); - when(pluginMetrics.counter(NUMBER_OF_RECORDS_FAILED_IN_BATCH_JOB)).thenReturn(counter); - when(pluginMetrics.counter(NUMBER_OF_RECORDS_SUCCEEDED_IN_BATCH_JOB)).thenReturn(counter); + when(pluginMetrics.counter(NUMBER_OF_RECORDS_FAILED_IN_BATCH_JOB)).thenReturn(recordsFailedCounter); + when(pluginMetrics.counter(NUMBER_OF_RECORDS_SUCCEEDED_IN_BATCH_JOB)).thenReturn(recordsSuccessCounter); when(pluginSetting.getPipelineName()).thenReturn("pipeline_sagemaker"); when(pluginSetting.getName()).thenReturn("pipeline_sagemaker"); @@ -543,6 +553,298 @@ void testConcurrentCheckAndProcessBatch() throws InterruptedException { assertTrue(sageMakerBatchJobCreator.getBatch_records().isEmpty(), "Batch should be empty after processing"); } + @Test + void testRetryLogic_ThrottlingFailure() throws Exception { + // Arrange + Event event = createMockEvent("test-bucket", "input.jsonl"); + Record record = new Record<>(event); + batch_records.add(record); + + try (MockedStatic mockedRetryUtil = mockStatic(RetryUtil.class)) { + // Mock RetryUtil to simulate throttling failure + mockedRetryUtil.when(() -> RetryUtil.retryWithBackoffWithResult(any(), any())) + .thenReturn(new RetryUtil.RetryResult(false, new MLBatchJobException(429, "Rate limited"), 1)); + + // Act + sageMakerBatchJobCreator.checkAndProcessBatch(); + + // Assert + ConcurrentLinkedQueue retryQueue = sageMakerBatchJobCreator.getRetryQueue(); + assertEquals(1, retryQueue.size()); + assertTrue(batch_records.isEmpty()); + assertTrue(processedBatchRecords.isEmpty()); + verify(counter, times(0)).increment(); + } + } + + @Test + void testRetryLogic_SuccessfulRetry() throws Exception { + // Arrange + Event event = createMockEvent("test-bucket", "input.jsonl"); + Record record = new Record<>(event); + batch_records.add(record); + + try (MockedStatic mockedRetryUtil = mockStatic(RetryUtil.class)) { + // RetryUtil returns failure with 429 status + mockedRetryUtil.when(() -> RetryUtil.retryWithBackoffWithResult(any(), any())) + .thenReturn(new RetryUtil.RetryResult(false, new MLBatchJobException(429, "Rate limited"), 1)) + .thenReturn(new RetryUtil.RetryResult(true, null, 1)); + // Act + sageMakerBatchJobCreator.checkAndProcessBatch(); + + // Assert + ConcurrentLinkedQueue retryQueue = sageMakerBatchJobCreator.getRetryQueue(); + assertEquals(1, retryQueue.size()); + assertTrue(batch_records.isEmpty()); + + // Act - Second attempt + sageMakerBatchJobCreator.checkAndProcessBatch(); + // Assert + assertTrue(retryQueue.isEmpty()); + assertEquals(1, processedBatchRecords.size()); + verify(counter, times(1)).increment(); + } + } + + @Test + void testRetryLogic_FailedRetry() throws Exception { + // Arrange + Event event = createMockEvent("test-bucket", "input.jsonl"); + Record record = new Record<>(event); + batch_records.add(record); + + try (MockedStatic mockedRetryUtil = mockStatic(RetryUtil.class)) { + // RetryUtil returns failure with 429 status + mockedRetryUtil.when(() -> RetryUtil.retryWithBackoffWithResult(any(), any())) + .thenReturn(new RetryUtil.RetryResult(false, new MLBatchJobException(429, "Rate limited"), 1)) + .thenReturn(new RetryUtil.RetryResult(false, new MLBatchJobException(404, "timeout"), 1)); + // Act + sageMakerBatchJobCreator.checkAndProcessBatch(); + + // Assert + ConcurrentLinkedQueue retryQueue = sageMakerBatchJobCreator.getRetryQueue(); + assertEquals(1, retryQueue.size()); + assertTrue(batch_records.isEmpty()); + + // Act - Second attempt + when(dlqPushHandler.getDlqPluginSetting()).thenReturn(pluginSetting); + when(event.toJsonString()).thenReturn("event"); // Add this + sageMakerBatchJobCreator.checkAndProcessBatch(); + // Assert + assertTrue(retryQueue.isEmpty()); + assertEquals(1, processedBatchRecords.size()); + verify(dlqPushHandler).perform(any()); + verify(counter, times(1)).increment(); + } + } + + @Test + void testRetryLogic_ExpiredRetry() throws Exception { + // Arrange + Event event = createMockEvent("test-bucket", "input.jsonl"); + Record record = new Record<>(event); + batch_records.add(record); + + try (MockedStatic mockedRetryUtil = mockStatic(RetryUtil.class)) { + // Mock initial failure + mockedRetryUtil.when(() -> RetryUtil.retryWithBackoffWithResult(any(), any())) + .thenReturn(new RetryUtil.RetryResult(false, new MLBatchJobException(429, "Rate limited"), 1)); + + // First attempt - adds to retry queue + sageMakerBatchJobCreator.checkAndProcessBatch(); + + // Get retry queue and simulate record expiration + ConcurrentLinkedQueue retryQueue = sageMakerBatchJobCreator.getRetryQueue(); + AbstractBatchJobCreator.RetryRecord retryRecord = retryQueue.peek(); + setPrivateField(retryRecord, "createdTime", System.currentTimeMillis() - (mlProcessorConfig.getRetryTimeWindow().toMillis() + 1)); + + // Act - Process expired record + when(dlqPushHandler.getDlqPluginSetting()).thenReturn(pluginSetting); + when(event.toJsonString()).thenReturn("event"); // Add this + sageMakerBatchJobCreator.checkAndProcessBatch(); + + // Assert + assertTrue(retryQueue.isEmpty()); + assertEquals(1, processedBatchRecords.size()); + verify(dlqPushHandler).perform(any()); + verify(counter, times(1)).increment(); + } + } + + @Test + void testProcessBothRetryAndNewRecords() throws Exception { + // Create record A for retry queue + Event eventA = createMockEvent("test-bucket", "inputA.jsonl"); + Record recordA = new Record<>(eventA); + + // Create RetryRecord using reflection + Object retryRecordA = createRetryRecord(recordA); + sageMakerBatchJobCreator.getRetryQueue().add((AbstractBatchJobCreator.RetryRecord) retryRecordA); + + // Create record B for batch_records + Event eventB = createMockEvent("test-bucket", "inputB.jsonl"); + Record recordB = new Record<>(eventB); + List> newRecords = Collections.singletonList(recordB); + List> resultRecords = new ArrayList<>(); + + try (MockedStatic mockedRetryUtil = mockStatic(RetryUtil.class)) { + // Mock successful processing + mockedRetryUtil.when(() -> RetryUtil.retryWithBackoffWithResult(any(), any())) + .thenReturn(new RetryUtil.RetryResult(true, null, 3)); + + // Add record B to batch_records + sageMakerBatchJobCreator.createMLBatchJob(newRecords, resultRecords); + // Verify initial state + assertEquals(1, sageMakerBatchJobCreator.getRetryQueue().size(), "Retry queue should have record A"); + assertEquals(1, sageMakerBatchJobCreator.getBatch_records().size(), "Batch records should have record B"); + + sageMakerBatchJobCreator.checkAndProcessBatch(); + // Verify final state + + // Verify retry queue is empty + assertTrue(sageMakerBatchJobCreator.getRetryQueue().isEmpty(), + "Retry queue should be empty after processing"); + + // Verify batch_records is empty + assertTrue(sageMakerBatchJobCreator.getBatch_records().isEmpty(), + "Batch records should be empty after processing"); + + // Verify success counters + verify(counter).increment(); + verify(recordsSuccessCounter).increment(2); // Both records A and B + } + } + + @Test + void testProcessBothRetryAndNewRecords_WithThrottledAgain() throws Exception { + // Create record A for retry queue + Event eventA = createMockEvent("test-bucket", "inputA.jsonl"); + Record recordA = new Record<>(eventA); + Object retryRecordA = createRetryRecord(recordA); + sageMakerBatchJobCreator.getRetryQueue().add((AbstractBatchJobCreator.RetryRecord) retryRecordA); + + // Create record B for batch_records + Event eventB = createMockEvent("test-bucket", "inputB.jsonl"); + Record recordB = new Record<>(eventB); + List> newRecords = Collections.singletonList(recordB); + List> resultRecords = new ArrayList<>(); + + try (MockedStatic mockedRetryUtil = mockStatic(RetryUtil.class)) { + // Mock failure with throttling + mockedRetryUtil.when(() -> RetryUtil.retryWithBackoffWithResult(any(), any())) + .thenReturn(new RetryUtil.RetryResult(false, new MLBatchJobException(429, "Rate limited"), 1)); + + sageMakerBatchJobCreator.createMLBatchJob(newRecords, resultRecords); + + // Verify initial state + assertEquals(1, sageMakerBatchJobCreator.getRetryQueue().size(), + "Retry queue should have record A"); + assertEquals(1, sageMakerBatchJobCreator.getBatch_records().size(), + "Batch records should have record B"); + // Get retry count of record A before processing + AbstractBatchJobCreator.RetryRecord retryRecordBefore = + sageMakerBatchJobCreator.getRetryQueue().peek(); + int initialRetryCount = getRetryCount(retryRecordBefore); + + // Process both records - should be throttled + sageMakerBatchJobCreator.checkAndProcessBatch(); + + // Verify retry queue contains both records + assertEquals(2, sageMakerBatchJobCreator.getRetryQueue().size(), + "Retry queue should contain both records after throttling"); + assertEquals(0, sageMakerBatchJobCreator.getProcessedBatchRecords().size(), + "Retry queue should contain both records after throttling"); + // Verify batch_records is empty + assertTrue(sageMakerBatchJobCreator.getBatch_records().isEmpty(), + "Batch records should be empty after processing"); + // Verify record A's retry count was incremented + AbstractBatchJobCreator.RetryRecord updatedRetryRecordA = + sageMakerBatchJobCreator.getRetryQueue().stream() + .filter(r -> getRecordFromRetryRecord(r).getData().getJsonNode().get("key").asText().equals("inputA.jsonl")) + .findFirst() + .orElse(null); + + assertNotNull(updatedRetryRecordA, "Record A should still be in retry queue"); + assertEquals(initialRetryCount + 1, getRetryCount(updatedRetryRecordA), + "Retry count for record A should be incremented"); + + // Verify record B was added to retry queue + boolean recordBInRetryQueue = sageMakerBatchJobCreator.getRetryQueue().stream() + .anyMatch(r -> getRecordFromRetryRecord(r).getData().getJsonNode().get("key").asText().equals("inputB.jsonl")); + assertTrue(recordBInRetryQueue, "Record B should be in retry queue"); + } + } + + @Test + void testProcessBothRetryAndNewRecords_WithFailure() throws Exception { + // Create record A for retry queue + Event eventA = createMockEvent("test-bucket", "inputA.jsonl"); + Record recordA = new Record<>(eventA); + Object retryRecordA = createRetryRecord(recordA); + sageMakerBatchJobCreator.getRetryQueue().add((AbstractBatchJobCreator.RetryRecord) retryRecordA); + when(eventA.toJsonString()).thenReturn("eventA"); + + // Create record B for batch_records + Event eventB = createMockEvent("test-bucket", "inputB.jsonl"); + Record recordB = new Record<>(eventB); + List> newRecords = Collections.singletonList(recordB); + List> resultRecords = new ArrayList<>(); + when(eventB.toJsonString()).thenReturn("eventB"); + + when(dlqPushHandler.getDlqPluginSetting()).thenReturn(pluginSetting); + // Capture DLQ objects + ArgumentCaptor> dlqObjectsCaptor = ArgumentCaptor.forClass(List.class); + + try (MockedStatic mockedRetryUtil = mockStatic(RetryUtil.class)) { + // Mock 404 failure response + mockedRetryUtil.when(() -> RetryUtil.retryWithBackoffWithResult(any(), any())) + .thenReturn(new RetryUtil.RetryResult(false, new MLBatchJobException(404, "Not Found"), 3)); + // Add record B to batch_records + sageMakerBatchJobCreator.createMLBatchJob(newRecords, resultRecords); + + // Verify initial state + assertEquals(1, sageMakerBatchJobCreator.getRetryQueue().size(), + "Retry queue should have record A"); + assertEquals(1, sageMakerBatchJobCreator.getBatch_records().size(), + "Batch records should have record B"); + + // Process both records - should fail with 404 + sageMakerBatchJobCreator.checkAndProcessBatch(); + + // Verify final state + // Both records should be removed from their respective queues + assertTrue(sageMakerBatchJobCreator.getRetryQueue().isEmpty(), + "Retry queue should be empty after failure"); + assertTrue(sageMakerBatchJobCreator.getBatch_records().isEmpty(), + "Batch records should be empty after failure"); + + // Verify both records were added to processedBatchRecords with failure tags + assertEquals(2, sageMakerBatchJobCreator.getProcessedBatchRecords().size(), + "Both records should be in result records"); + + // Verify DLQ handling + verify(dlqPushHandler).perform(dlqObjectsCaptor.capture()); + List dlqObjects = dlqObjectsCaptor.getValue(); + assertEquals(2, dlqObjects.size(), "Both records should be sent to DLQ"); + + // Verify DLQ objects content + Set dlqKeys = dlqObjects.stream() + .map(dlqObject -> ((MLBatchJobFailedDlqData) dlqObject.getFailedData()).getS3Key()) + .collect(Collectors.toSet()); + assertTrue(dlqKeys.contains("inputA.jsonl"), "DLQ should contain record A"); + assertTrue(dlqKeys.contains("inputB.jsonl"), "DLQ should contain record B"); + // Verify all DLQ objects have correct status code + dlqObjects.forEach(dlqObject -> { + MLBatchJobFailedDlqData failedData = (MLBatchJobFailedDlqData) dlqObject.getFailedData(); + assertEquals(404, failedData.getStatus(), "DLQ object should have 404 status code"); + assertEquals("Failed to Create SageMaker batch job after 3 attempts: Not Found", failedData.getMessage(), "DLQ object should have correct error message"); + }); + // Verify counters + verify(counter).increment(); // Failure counter + verify(recordsFailedCounter).increment(2); // Both records failed + } + } + @Test void generateJobName_ReturnsValidJobName() { // when @@ -571,4 +873,48 @@ private void callPrivateProcessRemainingBatch() throws Exception { method.setAccessible(true); method.invoke(sageMakerBatchJobCreator); } + + private void setPrivateField(Object object, String fieldName, Object value) { + try { + Field field = object.getClass().getDeclaredField(fieldName); + field.setAccessible(true); + field.set(object, value); + } catch (Exception e) { + throw new RuntimeException("Failed to set private field: " + fieldName, e); + } + } + + private Object createRetryRecord(Record record) throws Exception { + // Get RetryRecord class + Class retryRecordClass = Class.forName( + "org.opensearch.dataprepper.plugins.ml_inference.processor.common.AbstractBatchJobCreator$RetryRecord"); + + // Get constructor + Constructor constructor = retryRecordClass.getDeclaredConstructor( + AbstractBatchJobCreator.class, Record.class); + constructor.setAccessible(true); + + // Create instance + return constructor.newInstance(sageMakerBatchJobCreator, record); + } + + private int getRetryCount(Object retryRecord) { + try { + Field retryCountField = retryRecord.getClass().getDeclaredField("retryCount"); + retryCountField.setAccessible(true); + return (int) retryCountField.get(retryRecord); + } catch (Exception e) { + throw new RuntimeException("Failed to get retry count", e); + } + } + + private Record getRecordFromRetryRecord(Object retryRecord) { + try { + Field recordField = retryRecord.getClass().getDeclaredField("record"); + recordField.setAccessible(true); + return (Record) recordField.get(retryRecord); + } catch (Exception e) { + throw new RuntimeException("Failed to get record from RetryRecord", e); + } + } }