Merge branch 'main' into tenant-aware-integ

Author: dhrubo-os

ml-algorithms/build.gradle

@@ -85,8 +85,7 @@ dependencies {
     implementation('net.minidev:json-smart:2.5.2')
     implementation group: 'org.json', name: 'json', version: '20231013'
     implementation group: 'software.amazon.awssdk', name: 'netty-nio-client', version: "2.30.18"
-    // TODO: Replace the custom jar with dependency implementation 'io.modelcontextprotocol.sdk:mcp:0.10.0' after its release
-    api files('libs/mcp-0.10.0-SNAPSHOT.jar')
+    api('io.modelcontextprotocol.sdk:mcp:0.9.0')
     testImplementation("com.fasterxml.jackson.core:jackson-annotations:${versions.jackson}")
     testImplementation("com.fasterxml.jackson.core:jackson-databind:${versions.jackson_databind}")
     testImplementation group: 'com.networknt' , name: 'json-schema-validator', version: '1.4.0'

ml-algorithms/libs/mcp-0.10.0-SNAPSHOT.jar

@@ -82,8 +82,6 @@ public List<MLToolSpec> getMcpToolSpecs() {
             Duration readTimeout = Duration.ofSeconds(super.getConnectorClientConfig().getReadTimeout());
 
             Consumer<HttpRequest.Builder> headerConfig = builder -> {
-                builder.header("Content-Type", "application/json");
-
                 for (Map.Entry<String, String> entry : connector.getDecryptedHeaders().entrySet()) {
                     builder.header(entry.getKey(), entry.getValue());
                 }

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/McpConnectorExecutor.java

@@ -467,6 +467,7 @@ configurations.all {
     resolutionStrategy.force "jakarta.json:jakarta.json-api:2.1.3"
     resolutionStrategy.force "org.opensearch:opensearch:${opensearch_version}"
     resolutionStrategy.force "org.bouncycastle:bcprov-jdk18on:1.78.1"
+    resolutionStrategy.force 'io.projectreactor:reactor-core:3.7.0'
 }
 
 apply plugin: 'com.netflix.nebula.ospackage'

plugin/build.gradle

@@ -139,6 +139,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLDepl
         MLDeployModelRequest deployModelRequest = MLDeployModelRequest.fromActionRequest(request);
         String modelId = deployModelRequest.getModelId();
         String tenantId = deployModelRequest.getTenantId();
+        log.debug("Received deploy request for modelId: {}", modelId);
         if (!TenantAwareHelper.validateTenantId(mlFeatureEnabledSetting, tenantId, listener)) {
             return;
         }
@@ -215,6 +216,7 @@ private void deployModel(
     ) {
         String[] targetNodeIds = deployModelRequest.getModelNodeIds();
         boolean deployToAllNodes = targetNodeIds == null || targetNodeIds.length == 0;
+        log.info("Starting model deployment for model: {}", modelId);
         if (!allowCustomDeploymentPlan && !deployToAllNodes) {
             throw new IllegalArgumentException("Don't allow custom deployment plan");
         }
@@ -283,6 +285,7 @@ private void deployModel(
             .build();
         mlTaskManager.createMLTask(mlTask, ActionListener.wrap(response -> {
             String taskId = response.getId();
+            log.debug("ML deploy task {} created for modelId: {}", taskId, modelId);
             mlTask.setTaskId(taskId);
             if (algorithm == FunctionName.REMOTE) {
                 mlTaskManager.add(mlTask, eligibleNodeIds);
@@ -396,6 +399,7 @@ private ActionListener<MLDeployModelNodesResponse> deployModelNodesResponseListe
             if (mlTaskManager.contains(taskId)) {
                 mlTaskManager.updateMLTask(taskId, tenantId, Map.of(STATE_FIELD, MLTaskState.RUNNING), TASK_SEMAPHORE_TIMEOUT, false);
             }
+            log.debug("Model deployment successful for model: {}", modelId);
             listener.onResponse(new MLDeployModelResponse(taskId, MLTaskType.DEPLOY_MODEL, MLTaskState.COMPLETED.name()));
         }, e -> {
             log.error("Failed to deploy model {}", modelId, e);
@@ -423,6 +427,7 @@ void updateModelDeployStatusAndTriggerOnNodesAction(
         List<DiscoveryNode> eligibleNodes,
         boolean deployToAllNodes
     ) {
+        log.debug("Triggering deploy on nodes for modelId: {}", modelId);
         MLDeployModelInput deployModelInput = new MLDeployModelInput(
             modelId,
             taskId,
@@ -438,6 +443,7 @@ void updateModelDeployStatusAndTriggerOnNodesAction(
             deployModelInput
         );
         ActionListener<MLDeployModelNodesResponse> actionListener = ActionListener.wrap(r -> {
+            log.debug("Successfully triggered model deployment on nodes for model: {}", modelId);
             if (mlTaskManager.contains(taskId)) {
                 mlTaskManager
                     .updateMLTask(taskId, mlModel.getTenantId(), Map.of(STATE_FIELD, MLTaskState.RUNNING), TASK_SEMAPHORE_TIMEOUT, false);
@@ -462,6 +468,7 @@ void updateModelDeployStatusAndTriggerOnNodesAction(
         });
 
         List<String> workerNodes = eligibleNodes.stream().map(DiscoveryNode::getId).toList();
+        log.debug("Updating model state to DEPLOYING for modelId: {}", modelId);
         mlModelManager
             .updateModel(
                 modelId,

plugin/src/main/java/org/opensearch/ml/action/deploy/TransportDeployModelAction.java

@@ -140,6 +140,8 @@ private MLDeployModelNodeResponse createDeployModelNodeResponse(MLDeployModelNod
 
         String localNodeId = clusterService.localNode().getId();
 
+        log.debug("Starting model deployment for model: {} on node {}", modelId, localNodeId);
+
         ActionListener<MLForwardResponse> taskDoneListener = ActionListener
             .wrap(res -> { log.info("deploy model task done {}", taskId); }, ex -> {
                 logException("Deploy model task failed: " + taskId, ex, log);
@@ -155,6 +157,12 @@ private MLDeployModelNodeResponse createDeployModelNodeResponse(MLDeployModelNod
             deployToAllNodes,
             mlTask,
             ActionListener.wrap(r -> {
+                log
+                    .debug(
+                        "Model deployed successfully on local node: {}. Sending DONE message to coordinating node: {}",
+                        localNodeId,
+                        coordinatingNodeId
+                    );
                 MLForwardInput mlForwardInput = MLForwardInput
                     .builder()
                     .requestType(MLForwardRequestType.DEPLOY_MODEL_DONE)
@@ -175,6 +183,13 @@ private MLDeployModelNodeResponse createDeployModelNodeResponse(MLDeployModelNod
                         );
                 }
             }, e -> {
+                log
+                    .warn(
+                        "Model deployment failed on local node: {}. Sending FAILED message to coordinating node {} with error: {}",
+                        localNodeId,
+                        coordinatingNodeId,
+                        e.getMessage()
+                    );
                 MLForwardInput mlForwardInput = MLForwardInput
                     .builder()
                     .requestType(MLForwardRequestType.DEPLOY_MODEL_DONE)

plugin/src/main/java/org/opensearch/ml/action/deploy/TransportDeployModelOnNodeAction.java

@@ -100,6 +100,7 @@ public TransportUndeployModelAction(
 
     @Override
     protected void doExecute(Task task, MLUndeployModelNodesRequest request, ActionListener<MLUndeployModelNodesResponse> listener) {
+        log.info("Executing undeploy for models: {}", Arrays.toString(request.getModelIds()));
         ActionListener<MLUndeployModelNodesResponse> wrappedListener = ActionListener.wrap(undeployModelNodesResponse -> {
             processUndeployModelResponseAndUpdate(request.getTenantId(), undeployModelNodesResponse, listener);
         }, listener::onFailure);
@@ -112,6 +113,7 @@ void processUndeployModelResponseAndUpdate(
         ActionListener<MLUndeployModelNodesResponse> listener
     ) {
         List<MLUndeployModelNodeResponse> responses = undeployModelNodesResponse.getNodes();
+        log.debug("Processing undeploy model responses from nodes");
         if (responses == null || responses.isEmpty()) {
             listener.onResponse(undeployModelNodesResponse);
             return;
@@ -135,9 +137,10 @@ void processUndeployModelResponseAndUpdate(
 
             Map<String, String> modelUndeployStatus = r.getModelUndeployStatus();
             for (Map.Entry<String, String> entry : modelUndeployStatus.entrySet()) {
+                String modelId = entry.getKey();
                 String status = entry.getValue();
+                log.debug("Model status of model {} on node {}: {}", modelId, r.getNode().getId(), status);
                 if (UNDEPLOYED.equals(status)) {
-                    String modelId = entry.getKey();
                     if (!actualRemovedNodesMap.containsKey(modelId)) {
                         actualRemovedNodesMap.put(modelId, new ArrayList<>());
                     }
@@ -154,6 +157,7 @@ void processUndeployModelResponseAndUpdate(
         MLSyncUpNodesRequest syncUpRequest = new MLSyncUpNodesRequest(nodeFilter.getAllNodes(), syncUpInput);
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
             if (!actualRemovedNodesMap.isEmpty()) {
+                log.debug("Models undeployed from nodes: {}", actualRemovedNodesMap);
                 BulkDataObjectRequest bulkRequest = BulkDataObjectRequest.builder().globalIndex(ML_MODEL_INDEX).build();
                 Map<String, Boolean> deployToAllNodes = new HashMap<>();
                 for (String modelId : actualRemovedNodesMap.keySet()) {
@@ -166,8 +170,10 @@ void processUndeployModelResponseAndUpdate(
                      *  we need to update both planning worker nodes (count) and current worker nodes (count)
                      *  and deployToAllNodes value in model index.
                      */
+                    log.debug("Updating metadata for model {}: removedNodes={}", modelId, removedNodes);
                     Map<String, Object> updateDocument = new HashMap<>();
-                    if (modelWorkNodesBeforeRemoval.get(modelId).length == removedNodeCount) { // undeploy all nodes.
+                    if (modelWorkNodesBeforeRemoval.get(modelId).length == removedNodeCount) {
+                        log.debug("All nodes removed for model {}. Marking as undeployed.", modelId);// undeploy all nodes.
                         updateDocument.put(MLModel.PLANNING_WORKER_NODES_FIELD, ImmutableList.of());
                         updateDocument.put(MLModel.PLANNING_WORKER_NODE_COUNT_FIELD, 0);
                         updateDocument.put(MLModel.CURRENT_WORKER_NODE_COUNT_FIELD, 0);
@@ -180,6 +186,12 @@ void processUndeployModelResponseAndUpdate(
                             .stream(modelWorkNodesBeforeRemoval.get(modelId))
                             .filter(x -> !removedNodes.contains(x))
                             .collect(Collectors.toList());
+                        log
+                            .debug(
+                                "Partially undeployed for model {} with remaining planning worker nodes: {}",
+                                modelId,
+                                newPlanningWorkerNodes
+                            );
                         updateDocument.put(MLModel.PLANNING_WORKER_NODES_FIELD, newPlanningWorkerNodes);
                         updateDocument.put(MLModel.PLANNING_WORKER_NODE_COUNT_FIELD, newPlanningWorkerNodes.size());
                         updateDocument.put(MLModel.CURRENT_WORKER_NODE_COUNT_FIELD, newPlanningWorkerNodes.size());
@@ -195,6 +207,7 @@ void processUndeployModelResponseAndUpdate(
                     bulkRequest.add(updateRequest).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
                 }
                 syncUpInput.setDeployToAllNodes(deployToAllNodes);
+                log.debug("Sending bulk metadata update request for undeploy");
                 ActionListener<BulkResponse> actionListener = ActionListener.wrap(r -> {
                     log
                         .debug(
@@ -203,6 +216,7 @@ void processUndeployModelResponseAndUpdate(
                         );
                 }, e -> { log.error("Failed to update model state as undeployed", e); });
                 ActionListener<BulkResponse> wrappedListener = ActionListener.runAfter(actionListener, () -> {
+                    log.debug("Triggering sync-up after bulk update for undeploy");
                     syncUpUndeployedModels(syncUpRequest);
                     listener.onResponse(undeployModelNodesResponse);
                 });
@@ -288,11 +302,14 @@ private MLUndeployModelNodeResponse createUndeployModelNodeResponse(MLUndeployMo
 
         boolean specifiedModelIds = modelIds != null && modelIds.length > 0;
         String[] removedModelIds = specifiedModelIds ? modelIds : mlModelManager.getAllModelIds();
+
+        log.debug("Models to undeploy: {}", Arrays.toString(removedModelIds));
         if (removedModelIds != null) {
             for (String modelId : removedModelIds) {
                 FunctionName functionName = mlModelManager.getModelFunctionName(modelId);
                 String[] workerNodes = mlModelManager.getWorkerNodes(modelId, functionName);
                 modelWorkerNodesMap.put(modelId, workerNodes);
+                log.debug("Retrieved worker nodes for model {}: {}", modelId, Arrays.toString(workerNodes));
             }
         }
 

plugin/src/main/java/org/opensearch/ml/action/undeploy/TransportUndeployModelAction.java

@@ -129,12 +129,14 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLUnde
         String[] modelIds = undeployModelsRequest.getModelIds();
         String tenantId = undeployModelsRequest.getTenantId();
         String[] targetNodeIds = undeployModelsRequest.getNodeIds();
+        log.info("Executing undeploy model action for modelIds: {}", Arrays.toString(modelIds));
 
         if (!TenantAwareHelper.validateTenantId(mlFeatureEnabledSetting, tenantId, listener)) {
             return;
         }
 
         if (modelIds == null) {
+            log.error("No modelIds provided in undeploy.");
             listener.onFailure(new IllegalArgumentException("Must set specific model ids to undeploy"));
             return;
         }
@@ -186,10 +188,12 @@ private void undeployModels(
         String tenantId,
         ActionListener<MLUndeployModelsResponse> listener
     ) {
+        log.debug("Initiating undeploy on nodes: {}, for modelIds: {}", Arrays.toString(targetNodeIds), Arrays.toString(modelIds));
         MLUndeployModelNodesRequest mlUndeployModelNodesRequest = new MLUndeployModelNodesRequest(targetNodeIds, modelIds);
         mlUndeployModelNodesRequest.setTenantId(tenantId);
 
         client.execute(MLUndeployModelAction.INSTANCE, mlUndeployModelNodesRequest, ActionListener.wrap(response -> {
+            log.info("Undeploy response received from nodes");
             /*
              * The method TransportUndeployModelsAction.processUndeployModelResponseAndUpdate(...) performs
              * undeploy action of models by removing the models from the nodes cache and updating the index when it's able to find it.
@@ -213,9 +217,11 @@ private void undeployModels(
                 return modelCacheMissForModelIds;
             });
             if (response.getNodes().isEmpty() || modelNotFoundInNodesCache) {
+                log.warn("No node found running the model(s): {}", Arrays.toString(modelIds));
                 bulkSetModelIndexToUndeploy(modelIds, listener, response);
                 return;
             }
+            log.info("Successfully undeployed model(s) from nodes: {}", Arrays.toString(modelIds));
             listener.onResponse(new MLUndeployModelsResponse(response));
         }, listener::onFailure));
     }

plugin/src/main/java/org/opensearch/ml/action/undeploy/TransportUndeployModelsAction.java

@@ -102,6 +102,7 @@ public void run() {
         initMLConfig();
         if (!clusterService.state().metadata().indices().containsKey(ML_MODEL_INDEX)) {
             // no need to run sync up job if no model index
+            log.info("Skipping sync up job - ML model index not found");
             return;
         }
         log.debug("ML sync job starts");
@@ -111,6 +112,7 @@ public void run() {
 
         // gather running model/tasks on nodes
         client.execute(MLSyncUpAction.INSTANCE, gatherInfoRequest, ActionListener.wrap(r -> {
+            log.debug("Received sync up responses from nodes");
             List<MLSyncUpNodeResponse> responses = r.getNodes();
             if (r.failures() != null && !r.failures().isEmpty()) {
                 log
@@ -130,6 +132,7 @@ public void run() {
             Map<String, Set<String>> expiredModelToNodes = new HashMap<>();
             for (MLSyncUpNodeResponse response : responses) {
                 String nodeId = response.getNode().getId();
+                log.debug("Processing sync response from node: {}", nodeId);
                 String[] expiredModelIds = response.getExpiredModelIds();
                 if (expiredModelIds != null && expiredModelIds.length > 0) {
                     Arrays
@@ -215,6 +218,7 @@ private void undeployExpiredModels(
         Map<String, Set<String>> deployingModels
     ) {
         String[] targetNodeIds = getAllNodes(clusterService);
+        log.debug("Sending requests to undeploy expired models: {}", expiredModels);
         MLUndeployModelsRequest mlUndeployModelsRequest = new MLUndeployModelsRequest(
             expiredModels.toArray(new String[expiredModels.size()]),
             targetNodeIds,
@@ -279,6 +283,7 @@ void initMLConfig() {
     @VisibleForTesting
     void refreshModelState(Map<String, Set<String>> modelWorkerNodes, Map<String, Set<String>> deployingModels) {
         if (!updateModelStateSemaphore.tryAcquire()) {
+            log.debug("Model state refresh already in progress. Skipping this cycle.");
             return;
         }
         try {