Monitoring: Accurately capture predict failures (#4525) (#4555)

* fix: correctly capture ml failure stats on predict failures

* refactor: remove redundant braces

* refactor: dont track any 400 errors and add positive test case

* fix: apply spotless

---------

(cherry picked from commit a28d344)

Signed-off-by: Pavan Yekbote <pybot@amazon.com>
Co-authored-by: Pavan Yekbote <pybot@amazon.com>

Author: opensearch-trigger-bot[bot]

plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java

@@ -461,7 +461,7 @@ private void runPredict(
                                 mlModelManager.trackPredictDuration(modelId, startTime);
                                 internalListener.onResponse(output);
                             }
-                        }, e -> handlePredictFailure(mlTask, internalListener, e, false, modelId, actionName));
+                        }, e -> handlePredictFailure(mlTask, internalListener, e, shouldTrackRemoteFailure(e), modelId, actionName));
                         predictor.asyncPredict(mlInput, trackPredictDurationListener);
                     } else {
                         MLOutput output = mlModelManager.trackPredictDuration(modelId, () -> predictor.predict(mlInput));
@@ -478,7 +478,7 @@ private void runPredict(
                     return;
                 } catch (Exception e) {
                     log.error("Failed to predict model " + modelId, e);
-                    handlePredictFailure(mlTask, internalListener, e, false, modelId, actionName);
+                    handlePredictFailure(mlTask, internalListener, e, shouldTrackRemoteFailure(e), modelId, actionName);
                     return;
                 }
             } else if (FunctionName.needDeployFirst(algorithm)) {
@@ -604,4 +604,22 @@ public void validateOutputSchema(String modelId, ModelTensorOutput output) {
             }
         }
     }
+
+    boolean shouldTrackRemoteFailure(Exception e) {
+        // Don't track failures for user configuration issues
+        if (e instanceof IllegalArgumentException) {
+            return false;
+        }
+
+        // Don't track any 4xx client errors (user/configuration issues)
+        if (e instanceof OpenSearchStatusException) {
+            RestStatus status = ((OpenSearchStatusException) e).status();
+            if (status.getStatus() >= 400 && status.getStatus() < 500) {
+                return false;
+            }
+        }
+
+        // Track failures for infrastructure/service issues
+        return true;
+    }
 }

plugin/src/test/java/org/opensearch/ml/task/MLPredictTaskRunnerTests.java

@@ -51,6 +51,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.common.transport.TransportAddress;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.index.get.GetResult;
@@ -721,4 +722,25 @@ private void setupMocks(boolean runOnLocalNode, boolean failedToParseQueryInput,
             }).when(client).get(any(), any());
         }
     }
+
+    public void testShouldTrackRemoteFailure() {
+        // Test IllegalArgumentException - should not track
+        assertFalse(taskRunner.shouldTrackRemoteFailure(new IllegalArgumentException("Invalid argument")));
+
+        // Test OpenSearchStatusException with 4xx status codes - should not track
+        assertFalse(taskRunner.shouldTrackRemoteFailure(new OpenSearchStatusException("Bad request", RestStatus.BAD_REQUEST)));
+        assertFalse(taskRunner.shouldTrackRemoteFailure(new OpenSearchStatusException("Unauthorized", RestStatus.UNAUTHORIZED)));
+        assertFalse(taskRunner.shouldTrackRemoteFailure(new OpenSearchStatusException("Forbidden", RestStatus.FORBIDDEN)));
+        assertFalse(taskRunner.shouldTrackRemoteFailure(new OpenSearchStatusException("Not found", RestStatus.NOT_FOUND)));
+
+        // Test OpenSearchStatusException with 5xx status codes - should track
+        assertTrue(taskRunner.shouldTrackRemoteFailure(new OpenSearchStatusException("Server error", RestStatus.INTERNAL_SERVER_ERROR)));
+        assertTrue(
+            taskRunner.shouldTrackRemoteFailure(new OpenSearchStatusException("Service unavailable", RestStatus.SERVICE_UNAVAILABLE))
+        );
+
+        // Test other exceptions - should track
+        assertTrue(taskRunner.shouldTrackRemoteFailure(new RuntimeException("Runtime error")));
+        assertTrue(taskRunner.shouldTrackRemoteFailure(new IOException("IO error")));
+    }
 }