Examples with MLStatsNodesRequest org.opensearch.ml.action.stats.MLStatsNodesRequest used on opensource projects

Example 1 with MLStatsNodesRequest

use of org.opensearch.ml.action.stats.MLStatsNodesRequest in project ml-commons by opensearch-project.

the class RestStatsMLActionTests method testGetRequestSpecifyStats.

public void testGetRequestSpecifyStats() {
    Map<String, String> param = ImmutableMap.<String, String>builder().put("nodeId", "111,222").put("stat", StatNames.ML_EXECUTING_TASK_COUNT).build();
    FakeRestRequest fakeRestRequest = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.GET).withPath(MachineLearningPlugin.ML_BASE_URI + "/{nodeId}/stats/{stat}").withParams(param).build();
    MLStatsNodesRequest request = restAction.getRequest(fakeRestRequest);
    Assert.assertEquals(request.getStatsToBeRetrieved().size(), 1);
    Assert.assertTrue(request.getStatsToBeRetrieved().contains(StatNames.ML_EXECUTING_TASK_COUNT));
}

Example 2 with MLStatsNodesRequest

use of org.opensearch.ml.action.stats.MLStatsNodesRequest in project ml-commons by opensearch-project.

the class RestStatsMLActionTests method testGetRequestAllStats.

public void testGetRequestAllStats() {
    Map<String, String> param = ImmutableMap.<String, String>builder().put("nodeId", "111,222").put("stat", MLStatsNodesRequest.ALL_STATS_KEY).build();
    FakeRestRequest fakeRestRequest = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.GET).withPath(MachineLearningPlugin.ML_BASE_URI + "/{nodeId}/stats/{stat}").withParams(param).build();
    MLStatsNodesRequest request = restAction.getRequest(fakeRestRequest);
    Assert.assertEquals(0, request.getStatsToBeRetrieved().size());
    Assert.assertTrue(request.isRetrieveAllStats());
}

Example 3 with MLStatsNodesRequest

use of org.opensearch.ml.action.stats.MLStatsNodesRequest in project ml-commons by opensearch-project.

the class MLTaskDispatcher method dispatchTask.

/**
 * Select least loaded node based on ML_EXECUTING_TASK_COUNT and JVM_HEAP_USAGE
 * @param listener Action listener
 */
public void dispatchTask(ActionListener<DiscoveryNode> listener) {
    // todo: add ML node type setting check
    // DiscoveryNode[] mlNodes = getEligibleMLNodes();
    DiscoveryNode[] mlNodes = getEligibleDataNodes();
    MLStatsNodesRequest MLStatsNodesRequest = new MLStatsNodesRequest(mlNodes);
    MLStatsNodesRequest.addAll(ImmutableSet.of(ML_EXECUTING_TASK_COUNT, JVM_HEAP_USAGE.getName()));
    client.execute(MLStatsNodesAction.INSTANCE, MLStatsNodesRequest, ActionListener.wrap(mlStatsResponse -> {
        // Check JVM pressure
        List<MLStatsNodeResponse> candidateNodeResponse = mlStatsResponse.getNodes().stream().filter(stat -> (long) stat.getStatsMap().get(JVM_HEAP_USAGE.getName()) < DEFAULT_JVM_HEAP_USAGE_THRESHOLD).collect(Collectors.toList());
        if (candidateNodeResponse.size() == 0) {
            String errorMessage = "All nodes' memory usage exceeds limitation " + DEFAULT_JVM_HEAP_USAGE_THRESHOLD + ". No eligible node available to run ml jobs ";
            log.warn(errorMessage);
            listener.onFailure(new LimitExceededException(errorMessage));
            return;
        }
        // Check # of executing ML task
        candidateNodeResponse = candidateNodeResponse.stream().filter(stat -> (Long) stat.getStatsMap().get(ML_EXECUTING_TASK_COUNT) < maxMLBatchTaskPerNode).collect(Collectors.toList());
        if (candidateNodeResponse.size() == 0) {
            String errorMessage = "All nodes' executing ML task count reach limitation.";
            log.warn(errorMessage);
            listener.onFailure(new LimitExceededException(errorMessage));
            return;
        }
        // sort nodes by JVM usage percentage and # of executing ML task
        Optional<MLStatsNodeResponse> targetNode = candidateNodeResponse.stream().sorted((MLStatsNodeResponse r1, MLStatsNodeResponse r2) -> {
            int result = ((Long) r1.getStatsMap().get(ML_EXECUTING_TASK_COUNT)).compareTo((Long) r2.getStatsMap().get(ML_EXECUTING_TASK_COUNT));
            if (result == 0) {
                // JVM heap usage.
                return ((Long) r1.getStatsMap().get(JVM_HEAP_USAGE.getName())).compareTo((Long) r2.getStatsMap().get(JVM_HEAP_USAGE.getName()));
            }
            return result;
        }).findFirst();
        listener.onResponse(targetNode.get().getNode());
    }, exception -> {
        log.error("Failed to get node's task stats", exception);
        listener.onFailure(exception);
    }));
}

Example 4 with MLStatsNodesRequest

use of org.opensearch.ml.action.stats.MLStatsNodesRequest in project ml-commons by opensearch-project.

the class RestStatsMLAction method getRequest.

/**
 * Creates a MLStatsNodesRequest from a RestRequest
 *
 * @param request RestRequest
 * @return MLStatsNodesRequest
 */
@VisibleForTesting
MLStatsNodesRequest getRequest(RestRequest request) {
    // todo: add logic to triage request based on node type(ML node or data node)
    MLStatsNodesRequest mlStatsRequest = new MLStatsNodesRequest(splitCommaSeparatedParam(request, "nodeId").orElse(null));
    mlStatsRequest.timeout(request.param("timeout"));
    List<String> requestedStats = splitCommaSeparatedParam(request, "stat").map(Arrays::asList).orElseGet(Collections::emptyList);
    Set<String> validStats = mlStats.getStats().keySet();
    if (isAllStatsRequested(requestedStats)) {
        mlStatsRequest.setRetrieveAllStats(true);
    } else {
        mlStatsRequest.addAll(getStatsToBeRetrieved(request, validStats, requestedStats));
    }
    return mlStatsRequest;
}

Example 5 with MLStatsNodesRequest

use of org.opensearch.ml.action.stats.MLStatsNodesRequest in project ml-commons by opensearch-project.

the class RestStatsMLActionTests method testGetRequestEmptyStats.

public void testGetRequestEmptyStats() {
    Map<String, String> param = ImmutableMap.<String, String>builder().put("nodeId", "111,222").build();
    FakeRestRequest fakeRestRequest = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.GET).withPath(MachineLearningPlugin.ML_BASE_URI + "/{nodeId}/stats/").withParams(param).build();
    MLStatsNodesRequest request = restAction.getRequest(fakeRestRequest);
    Assert.assertEquals(0, request.getStatsToBeRetrieved().size());
    Assert.assertTrue(request.isRetrieveAllStats());
}