Merge branch 'opensearch-project:main' into main

Author: owaiskazi19

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

@@ -464,7 +464,7 @@ public void registerMLRemoteModel(
                             mlRegisterModelInput.getTenantId(),
                             new MLResourceNotFoundException("Failed to get model group due to index missing")
                         );
-                        listener.onFailure(e);
+                        listener.onFailure(new OpenSearchStatusException("Model group not found", RestStatus.NOT_FOUND));
                     } else {
                         log.error("Failed to get model group", e);
                         handleException(mlRegisterModelInput.getFunctionName(), mlTask.getTaskId(), mlRegisterModelInput.getTenantId(), e);
@@ -2077,40 +2077,44 @@ public void getConnector(String connectorId, String tenantId, ActionListener<Con
             .tenantId(tenantId)
             .build();
 
-        sdkClient.getDataObjectAsync(getDataObjectRequest).whenComplete((r, throwable) -> {
-            log.debug("Completed Get Connector Request, id:{}", connectorId);
-            if (throwable != null) {
-                Exception cause = SdkClientUtils.unwrapAndConvertToException(throwable);
-                if (ExceptionsHelper.unwrap(cause, IndexNotFoundException.class) != null) {
-                    log.error("Failed to get connector index", cause);
-                    listener.onFailure(new OpenSearchStatusException("Failed to find connector", RestStatus.NOT_FOUND));
+        try (ThreadContext.StoredContext ctx = client.threadPool().getThreadContext().stashContext()) {
+            sdkClient.getDataObjectAsync(getDataObjectRequest).whenComplete((r, throwable) -> {
+                log.debug("Completed Get Connector Request, id:{}", connectorId);
+                ctx.restore();
+                if (throwable != null) {
+                    Exception cause = SdkClientUtils.unwrapAndConvertToException(throwable);
+                    if (ExceptionsHelper.unwrap(cause, IndexNotFoundException.class) != null) {
+                        log.error("Failed to get connector index", cause);
+                        listener.onFailure(new OpenSearchStatusException("Failed to find connector", RestStatus.NOT_FOUND));
+                    } else {
+                        log.error("Failed to get ML connector {}", connectorId, cause);
+                        listener.onFailure(cause);
+                    }
                 } else {
-                    log.error("Failed to get ML connector {}", connectorId, cause);
-                    listener.onFailure(cause);
-                }
-            } else {
-                try {
-                    GetResponse gr = r.parser() == null ? null : GetResponse.fromXContent(r.parser());
-                    if (gr != null && gr.isExists()) {
-                        try (
-                            XContentParser parser = MLNodeUtils
-                                .createXContentParserFromRegistry(NamedXContentRegistry.EMPTY, gr.getSourceAsBytesRef())
-                        ) {
-                            ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
-                            Connector connector = Connector.createConnector(parser);
-                            listener.onResponse(connector);
-                        } catch (Exception e) {
-                            log.error("Failed to parse connector:{}", connectorId);
-                            listener.onFailure(e);
+                    try {
+                        GetResponse gr = r.parser() == null ? null : GetResponse.fromXContent(r.parser());
+                        if (gr != null && gr.isExists()) {
+                            try (
+                                XContentParser parser = MLNodeUtils
+                                    .createXContentParserFromRegistry(NamedXContentRegistry.EMPTY, gr.getSourceAsBytesRef())
+                            ) {
+                                ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
+                                Connector connector = Connector.createConnector(parser);
+                                listener.onResponse(connector);
+                            } catch (Exception e) {
+                                log.error("Failed to parse connector:{}", connectorId);
+                                listener.onFailure(e);
+                            }
+                        } else {
+                            listener
+                                .onFailure(new OpenSearchStatusException("Failed to find connector:" + connectorId, RestStatus.NOT_FOUND));
                         }
-                    } else {
-                        listener.onFailure(new OpenSearchStatusException("Failed to find connector:" + connectorId, RestStatus.NOT_FOUND));
+                    } catch (Exception e) {
+                        listener.onFailure(e);
                     }
-                } catch (Exception e) {
-                    listener.onFailure(e);
                 }
-            }
-        });
+            });
+        }
     }
 
     /**

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

@@ -75,6 +75,7 @@
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.get.GetRequest;
 import org.opensearch.action.get.GetResponse;
 import org.opensearch.action.index.IndexResponse;
@@ -92,9 +93,11 @@
 import org.opensearch.core.common.breaker.CircuitBreakingException;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.index.shard.ShardId;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
+import org.opensearch.index.IndexNotFoundException;
 import org.opensearch.index.get.GetResult;
 import org.opensearch.ml.breaker.MLCircuitBreakerService;
 import org.opensearch.ml.breaker.ThresholdCircuitBreaker;
@@ -492,6 +495,46 @@ public void testRegisterMLRemoteModel() throws PrivilegedActionException, IOExce
         verify(mlTaskManager).updateMLTask(anyString(), any(), anyMap(), anyLong(), anyBoolean());
     }
 
+    @Test
+    public void testRegisterMLRemoteModelModelGroupNotFoundException() throws PrivilegedActionException, IOException {
+        // Create listener and capture the failure
+        ArgumentCaptor<Exception> exceptionCaptor = ArgumentCaptor.forClass(Exception.class);
+        ActionListener<MLRegisterModelResponse> listener = mock(ActionListener.class);
+
+        // Setup mocks
+        doNothing().when(mlTaskManager).checkLimitAndAddRunningTask(any(), any());
+        when(mlCircuitBreakerService.checkOpenCB()).thenReturn(null);
+        when(threadPool.executor(REGISTER_THREAD_POOL)).thenReturn(taskExecutorService);
+        when(modelHelper.downloadPrebuiltModelMetaList(any(), any())).thenReturn(Collections.singletonList("demo"));
+        when(modelHelper.isModelAllowed(any(), any())).thenReturn(true);
+
+        // Create test inputs
+        MLRegisterModelInput pretrainedInput = mockRemoteModelInput(true);
+        MLTask pretrainedTask = MLTask.builder().taskId("pretrained").modelId("pretrained").functionName(FunctionName.REMOTE).build();
+
+        // Mock index handler
+        mock_MLIndicesHandler_initModelIndex(mlIndicesHandler, true);
+
+        // Mock client.get() to throw IndexNotFoundException
+        doAnswer(invocation -> {
+            ActionListener<GetResponse> getModelGroupListener = invocation.getArgument(1);
+            getModelGroupListener.onFailure(new IndexNotFoundException("Test", "test"));
+            return null;
+        }).when(client).get(any(), any());
+
+        // Execute method under test
+        modelManager.registerMLRemoteModel(sdkClient, pretrainedInput, pretrainedTask, listener);
+
+        // Verify the listener's onFailure was called with correct exception
+        verify(listener).onFailure(exceptionCaptor.capture());
+        Exception exception = exceptionCaptor.getValue();
+
+        // Verify exception type and message
+        assertTrue(exception instanceof OpenSearchStatusException);
+        assertEquals("Model group not found", exception.getMessage());
+        assertEquals(RestStatus.NOT_FOUND, ((OpenSearchStatusException) exception).status());
+    }
+
     public void testRegisterMLRemoteModel_SkipMemoryCBOpen() throws IOException {
         ActionListener<MLRegisterModelResponse> listener = mock(ActionListener.class);
         doNothing().when(mlTaskManager).checkLimitAndAddRunningTask(any(), any());

plugin/src/test/java/org/opensearch/ml/rest/RestMLRemoteInferenceIT.java

@@ -740,120 +740,6 @@ public void testCohereGenerateTextModel() throws IOException, InterruptedExcepti
         assertFalse(((String) responseMap.get("text")).isEmpty());
     }
 
-    public void testCohereClassifyModel() throws IOException, InterruptedException {
-        // Skip test if key is null
-        if (COHERE_KEY == null) {
-            return;
-        }
-        String entity = "{\n"
-            + "  \"name\": \"Cohere classify model Connector\",\n"
-            + "  \"description\": \"The connector to public Cohere classify model service\",\n"
-            + "  \"version\": 1,\n"
-            + "  \"client_config\": {\n"
-            + "    \"max_connection\": 20,\n"
-            + "    \"connection_timeout\": 50000,\n"
-            + "    \"read_timeout\": 50000\n"
-            + "  },\n"
-            + "  \"protocol\": \"http\",\n"
-            + "  \"parameters\": {\n"
-            + "      \"endpoint\": \"api.cohere.ai\",\n"
-            + "      \"auth\": \"API_Key\",\n"
-            + "      \"content_type\": \"application/json\",\n"
-            + "      \"max_tokens\": \"20\"\n"
-            + "  },\n"
-            + "  \"credential\": {\n"
-            + "      \"cohere_key\": \""
-            + COHERE_KEY
-            + "\"\n"
-            + "  },\n"
-            + "  \"actions\": [\n"
-            + "      {\n"
-            + "      \"action_type\": \"predict\",\n"
-            + "          \"method\": \"POST\",\n"
-            + "          \"url\": \"https://${parameters.endpoint}/v1/classify\",\n"
-            + "          \"headers\": { \n"
-            + "          \"Authorization\": \"Bearer ${credential.cohere_key}\"\n"
-            + "          },\n"
-            + "          \"request_body\": \"{ \\\"inputs\\\": ${parameters.inputs}, \\\"examples\\\": ${parameters.examples}, \\\"truncate\\\": \\\"END\\\" }\"\n"
-            + "      }\n"
-            + "  ]\n"
-            + "}";
-        Response response = createConnector(entity);
-        Map responseMap = parseResponseToMap(response);
-        String connectorId = (String) responseMap.get("connector_id");
-        response = registerRemoteModel("cohere classify model", connectorId);
-        responseMap = parseResponseToMap(response);
-        String taskId = (String) responseMap.get("task_id");
-        waitForTask(taskId, MLTaskState.COMPLETED);
-        response = getTask(taskId);
-        responseMap = parseResponseToMap(response);
-        String modelId = (String) responseMap.get("model_id");
-        response = deployRemoteModel(modelId);
-        responseMap = parseResponseToMap(response);
-        taskId = (String) responseMap.get("task_id");
-        waitForTask(taskId, MLTaskState.COMPLETED);
-        String predictInput = "{\n"
-            + "  \"parameters\": {\n"
-            + "      \"inputs\": [\n"
-            + "            \"Confirm your email address\",\n"
-            + "            \"hey i need u to send some $\"\n"
-            + "        ],\n"
-            + "        \"examples\": [\n"
-            + "            {\n"
-            + "                \"text\": \"Dermatologists don't like her!\",\n"
-            + "                \"label\": \"Spam\"\n"
-            + "            },\n"
-            + "            {\n"
-            + "                \"text\": \"Hello, open to this?\",\n"
-            + "                \"label\": \"Spam\"\n"
-            + "            },\n"
-            + "            {\n"
-            + "                \"text\": \"I need help please wire me $1000 right now\",\n"
-            + "                \"label\": \"Spam\"\n"
-            + "            },\n"
-            + "            {\n"
-            + "                \"text\": \"Nice to know you ;)\",\n"
-            + "                \"label\": \"Spam\"\n"
-            + "            },\n"
-            + "            {\n"
-            + "                \"text\": \"Please help me?\",\n"
-            + "                \"label\": \"Spam\"\n"
-            + "            },\n"
-            + "            {\n"
-            + "                \"text\": \"Your parcel will be delivered today\",\n"
-            + "                \"label\": \"Not spam\"\n"
-            + "            },\n"
-            + "            {\n"
-            + "                \"text\": \"Review changes to our Terms and Conditions\",\n"
-            + "                \"label\": \"Not spam\"\n"
-            + "            },\n"
-            + "            {\n"
-            + "                \"text\": \"Weekly sync notes\",\n"
-            + "                \"label\": \"Not spam\"\n"
-            + "            },\n"
-            + "            {\n"
-            + "                \"text\": \"Re: Follow up from todays meeting\",\n"
-            + "                \"label\": \"Not spam\"\n"
-            + "            },\n"
-            + "            {\n"
-            + "                \"text\": \"Pre-read for tomorrow\",\n"
-            + "                \"label\": \"Not spam\"\n"
-            + "            }\n"
-            + "        ]\n"
-            + "  }\n"
-            + "}";
-
-        response = predictRemoteModel(modelId, predictInput);
-        responseMap = parseResponseToMap(response);
-        List responseList = (List) responseMap.get("inference_results");
-        responseMap = (Map) responseList.get(0);
-        responseList = (List) responseMap.get("output");
-        responseMap = (Map) responseList.get(0);
-        responseMap = (Map) responseMap.get("dataAsMap");
-        responseList = (List) responseMap.get("classifications");
-        assertFalse(responseList.isEmpty());
-    }
-
     public static Response createConnector(String input) throws IOException {
         try {
             return TestHelper.makeRequest(client(), "POST", "/_plugins/_ml/connectors/_create", null, TestHelper.toHttpEntity(input), null);