Add context.restore() to avoid missing info

Signed-off-by: zane-neo <[email protected]>
opensearch-project · Oct 10, 2023 · 257959d · 257959d
1 parent 2c6f577
commit 257959d
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java b/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java
@@ -92,6 +92,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLTask
             ActionListener<MLModel> modelActionListener = new ActionListener<>() {
                 @Override
                 public void onResponse(MLModel mlModel) {
+                    context.restore();
                     modelCacheHelper.setModelInfo(modelId, mlModel);
                     FunctionName functionName = mlModel.getAlgorithm();
                     mlPredictionTaskRequest.getMlInput().setAlgorithm(functionName);