9595import org .opensearch .ml .common .MLTask ;
9696import org .opensearch .ml .common .connector .Connector ;
9797import org .opensearch .ml .common .exception .MLException ;
98+ import org .opensearch .ml .common .exception .MLLimitExceededException ;
9899import org .opensearch .ml .common .exception .MLResourceNotFoundException ;
99100import org .opensearch .ml .common .exception .MLValidationException ;
100101import org .opensearch .ml .common .model .MLModelState ;
@@ -318,9 +319,10 @@ private void uploadMLModelMeta(MLRegisterModelMetaInput mlRegisterModelMetaInput
318319 * @param mlTask ML task
319320 */
320321 public void registerMLModel (MLRegisterModelInput registerModelInput , MLTask mlTask ) {
321- mlStats . getStat ( MLNodeLevelStat . ML_NODE_TOTAL_REQUEST_COUNT ). increment ();
322+
322323 checkAndAddRunningTask (mlTask , maxRegisterTasksPerNode );
323324 try {
325+ mlStats .getStat (MLNodeLevelStat .ML_NODE_TOTAL_REQUEST_COUNT ).increment ();
324326 mlStats .getStat (MLNodeLevelStat .ML_NODE_EXECUTING_TASK_COUNT ).increment ();
325327 mlStats .createCounterStatIfAbsent (mlTask .getFunctionName (), REGISTER , ML_ACTION_REQUEST_COUNT ).increment ();
326328
@@ -380,7 +382,6 @@ public void registerMLModel(MLRegisterModelInput registerModelInput, MLTask mlTa
380382 handleException (registerModelInput .getFunctionName (), mlTask .getTaskId (), e );
381383 }
382384 } catch (Exception e ) {
383- mlStats .createCounterStatIfAbsent (mlTask .getFunctionName (), REGISTER , MLActionLevelStat .ML_ACTION_FAILURE_COUNT ).increment ();
384385 handleException (registerModelInput .getFunctionName (), mlTask .getTaskId (), e );
385386 } finally {
386387 mlStats .getStat (MLNodeLevelStat .ML_NODE_EXECUTING_TASK_COUNT ).increment ();
@@ -392,9 +393,9 @@ private void indexRemoteModel(MLRegisterModelInput registerModelInput, MLTask ml
392393 FunctionName functionName = mlTask .getFunctionName ();
393394 try (ThreadContext .StoredContext context = client .threadPool ().getThreadContext ().stashContext ()) {
394395 mlStats .getStat (MLNodeLevelStat .ML_NODE_TOTAL_REQUEST_COUNT ).increment ();
395-
396396 mlStats .createCounterStatIfAbsent (functionName , REGISTER , ML_ACTION_REQUEST_COUNT ).increment ();
397397 mlStats .getStat (MLNodeLevelStat .ML_NODE_EXECUTING_TASK_COUNT ).increment ();
398+
398399 String modelName = registerModelInput .getModelName ();
399400 String version = modelVersion == null ? registerModelInput .getVersion () : modelVersion ;
400401 Instant now = Instant .now ();
@@ -462,7 +463,6 @@ private void registerModelFromUrl(MLRegisterModelInput registerModelInput, MLTas
462463 FunctionName functionName = mlTask .getFunctionName ();
463464 try (ThreadContext .StoredContext context = client .threadPool ().getThreadContext ().stashContext ()) {
464465 mlStats .getStat (MLNodeLevelStat .ML_NODE_TOTAL_REQUEST_COUNT ).increment ();
465-
466466 mlStats .createCounterStatIfAbsent (functionName , REGISTER , ML_ACTION_REQUEST_COUNT ).increment ();
467467 mlStats .getStat (MLNodeLevelStat .ML_NODE_EXECUTING_TASK_COUNT ).increment ();
468468 String modelName = registerModelInput .getModelName ();
@@ -689,7 +689,12 @@ private void deleteModel(String modelId) {
689689 }
690690
691691 private void handleException (FunctionName functionName , String taskId , Exception e ) {
692- mlStats .createCounterStatIfAbsent (functionName , REGISTER , MLActionLevelStat .ML_ACTION_FAILURE_COUNT ).increment ();
692+ if (!(e instanceof MLLimitExceededException )
693+ && !(e instanceof MLResourceNotFoundException )
694+ && !(e instanceof IllegalArgumentException )) {
695+ mlStats .createCounterStatIfAbsent (functionName , REGISTER , MLActionLevelStat .ML_ACTION_FAILURE_COUNT ).increment ();
696+ mlStats .getStat (MLNodeLevelStat .ML_NODE_TOTAL_FAILURE_COUNT ).increment ();
697+ }
693698 Map <String , Object > updated = ImmutableMap .of (ERROR_FIELD , MLExceptionUtils .getRootCauseMessage (e ), STATE_FIELD , FAILED );
694699 mlTaskManager .updateMLTask (taskId , updated , TIMEOUT_IN_MILLIS , true );
695700 }
@@ -713,6 +718,7 @@ public void deployModel(
713718 ActionListener <String > listener
714719 ) {
715720 mlStats .createCounterStatIfAbsent (functionName , ActionName .DEPLOY , ML_ACTION_REQUEST_COUNT ).increment ();
721+ mlStats .getStat (MLNodeLevelStat .ML_NODE_TOTAL_REQUEST_COUNT ).increment ();
716722 List <String > workerNodes = mlTask .getWorkerNodes ();
717723 if (modelCacheHelper .isModelDeployed (modelId )) {
718724 if (workerNodes != null && workerNodes .size () > 0 ) {
@@ -835,7 +841,13 @@ public void deployModel(
835841 }
836842
837843 private void handleDeployModelException (String modelId , FunctionName functionName , ActionListener <String > listener , Exception e ) {
838- mlStats .createCounterStatIfAbsent (functionName , ActionName .DEPLOY , MLActionLevelStat .ML_ACTION_FAILURE_COUNT ).increment ();
844+
845+ if (!(e instanceof MLLimitExceededException )
846+ && !(e instanceof MLResourceNotFoundException )
847+ && !(e instanceof IllegalArgumentException )) {
848+ mlStats .createCounterStatIfAbsent (functionName , ActionName .DEPLOY , MLActionLevelStat .ML_ACTION_FAILURE_COUNT ).increment ();
849+ mlStats .getStat (MLNodeLevelStat .ML_NODE_TOTAL_FAILURE_COUNT ).increment ();
850+ }
839851 removeModel (modelId );
840852 listener .onFailure (e );
841853 }
@@ -858,7 +870,7 @@ public void getModel(String modelId, ActionListener<MLModel> listener) {
858870 }
859871
860872 /**
861- * Get model from model index with includes/exludes filter.
873+ * Get model from model index with includes/excludes filter.
862874 *
863875 * @param modelId model id
864876 * @param includes fields included
@@ -1045,6 +1057,7 @@ public synchronized Map<String, String> undeployModel(String[] modelIds) {
10451057 if (modelCacheHelper .isModelDeployed (modelId )) {
10461058 modelUndeployStatus .put (modelId , UNDEPLOYED );
10471059 mlStats .getStat (MLNodeLevelStat .ML_NODE_TOTAL_MODEL_COUNT ).decrement ();
1060+ mlStats .getStat (MLNodeLevelStat .ML_NODE_TOTAL_REQUEST_COUNT ).increment ();
10481061 mlStats
10491062 .createCounterStatIfAbsent (getModelFunctionName (modelId ), ActionName .UNDEPLOY , ML_ACTION_REQUEST_COUNT )
10501063 .increment ();
0 commit comments