From 82bbdf46250387fbb9bdffb8ac161a42de1a768c Mon Sep 17 00:00:00 2001 From: Henry Lindeman Date: Fri, 22 Mar 2024 10:25:14 -0700 Subject: [PATCH] move memory CB check into try block to catch exception and hand to listener for register remote model Signed-off-by: Henry Lindeman --- .../src/main/java/org/opensearch/ml/model/MLModelManager.java | 2 +- 1 file changed, 1 insertion(+), 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 index bdcb78d6ad..a109ce040e 100644 --- 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 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();