move memory CB check into try block to catch exception and hand to li…

…stener for register remote model Signed-off-by: Henry Lindeman <[email protected]>
HenryL27 · Mar 22, 2024 · 82bbdf4 · 82bbdf4
1 parent 624dfba
commit 82bbdf4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java
@@ -346,8 +346,8 @@ public void registerMLRemoteModel(
         MLTask mlTask,
         ActionListener<MLRegisterModelResponse> listener
     ) {
-        checkAndAddRunningTask(mlTask, maxRegisterTasksPerNode);
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
+            checkAndAddRunningTask(mlTask, maxRegisterTasksPerNode);
             mlStats.getStat(MLNodeLevelStat.ML_REQUEST_COUNT).increment();
             mlStats.createCounterStatIfAbsent(mlTask.getFunctionName(), REGISTER, ML_ACTION_REQUEST_COUNT).increment();
             mlStats.getStat(MLNodeLevelStat.ML_EXECUTING_TASK_COUNT).increment();