excluding circuit breaker for Agent

Signed-off-by: Dhrubo Saha <[email protected]>

Author: dhrubo-os

plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java

@@ -975,7 +975,10 @@ private <T> ThreadedActionListener<T> threadedActionListener(String threadPoolNa
      * @param runningTaskLimit limit
      */
     public void checkAndAddRunningTask(MLTask mlTask, Integer runningTaskLimit) {
-        if (Objects.nonNull(mlTask) && mlTask.getFunctionName() != FunctionName.REMOTE) {
+
+        // for agent and remote model prediction we don't need to check circuit breaker
+        if (Objects.nonNull(mlTask) && mlTask.getFunctionName() != FunctionName.REMOTE
+                && mlTask.getFunctionName() != FunctionName.AGENT) {
             checkOpenCircuitBreaker(mlCircuitBreakerService, mlStats);
         }
         mlTaskManager.checkLimitAndAddRunningTask(mlTask, runningTaskLimit);

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

@@ -112,8 +112,7 @@ public void dispatchTask(
             if (clusterService.localNode().getId().equals(nodeId)) {
                 // Execute ML task locally
                 log.debug("Execute ML request {} locally on node {}", request.getRequestID(), nodeId);
-                checkOpenCircuitBreaker(mlCircuitBreakerService, mlStats);
-                executeTask(request, listener);
+                checkCBAndExecute(functionName, request, listener);
             } else {
                 // Execute ML task remotely
                 log.debug("Execute ML request {} remotely on node {}", request.getRequestID(), nodeId);
@@ -130,7 +129,8 @@ public void dispatchTask(
     protected abstract void executeTask(Request request, ActionListener<Response> listener);
 
     protected void checkCBAndExecute(FunctionName functionName, Request request, ActionListener<Response> listener) {
-        if (functionName != FunctionName.REMOTE) {
+        // for agent and remote model prediction we don't need to check circuit breaker
+        if (functionName != FunctionName.REMOTE && functionName != FunctionName.AGENT) {
             checkOpenCircuitBreaker(mlCircuitBreakerService, mlStats);
         }
         executeTask(request, listener);

plugin/src/test/java/org/opensearch/ml/model/MLModelManagerTests.java

@@ -361,6 +361,18 @@ public void testRegisterMLModel_CircuitBreakerOpen() {
         verify(mlTaskManager).updateMLTask(anyString(), any(), anyMap(), anyLong(), anyBoolean());
     }
 
+    public void testRegisterMLModel_CircuitBreakerNotOpenForAgent() {
+        registerModelInput.setFunctionName(FunctionName.AGENT);
+        doNothing().when(mlTaskManager).checkLimitAndAddRunningTask(any(), any());
+        when(mlCircuitBreakerService.checkOpenCB()).thenReturn(thresholdCircuitBreaker);
+        when(thresholdCircuitBreaker.getName()).thenReturn("Disk Circuit Breaker");
+        when(thresholdCircuitBreaker.getThreshold()).thenReturn(87);
+        expectedEx.expect(CircuitBreakingException.class);
+        expectedEx.expectMessage("Disk Circuit Breaker is open, please check your resources!");
+        modelManager.registerMLModel(registerModelInput, mlTask);
+        verify(mlTaskManager).updateMLTask(anyString(), any(), anyMap(), anyLong(), anyBoolean());
+    }
+
     public void testRegisterMLModel_InitModelIndexFailure() {
         doNothing().when(mlTaskManager).checkLimitAndAddRunningTask(any(), any());
         when(mlCircuitBreakerService.checkOpenCB()).thenReturn(null);

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

@@ -155,4 +155,16 @@ public void testRun_NoCircuitbreakerforRemote() {
         Long value = (Long) mlStats.getStat(MLNodeLevelStat.ML_CIRCUIT_BREAKER_TRIGGER_COUNT).getValue();
         assertEquals(0L, value.longValue());
     }
+
+    public void testRun_NoCircuitbreakerforAgent() {
+        when(mlCircuitBreakerService.checkOpenCB()).thenReturn(thresholdCircuitBreaker);
+        when(thresholdCircuitBreaker.getName()).thenReturn("Memory Circuit Breaker");
+        when(thresholdCircuitBreaker.getThreshold()).thenReturn(87);
+        TransportService transportService = mock(TransportService.class);
+        ActionListener listener = mock(ActionListener.class);
+        MLTaskRequest request = new MLTaskRequest(false);
+        mlTaskRunner.run(FunctionName.AGENT, request, transportService, listener);
+        Long value = (Long) mlStats.getStat(MLNodeLevelStat.ML_CIRCUIT_BREAKER_TRIGGER_COUNT).getValue();
+        assertEquals(0L, value.longValue());
+    }
 }