Examples with ClusterService org.opensearch.cluster.service.ClusterService used on opensource projects

Example 11 with ClusterService

use of org.opensearch.cluster.service.ClusterService in project k-NN by opensearch-project.

the class TrainingJobRouterTransportActionTests method testMultiNode_withCapacity_withPreferredAvailable.

public void testMultiNode_withCapacity_withPreferredAvailable() {
    // Mock datanodes in the cluster through mocking the cluster service
    List<String> nodeIds = ImmutableList.of("node-1", "node-2", "node-3");
    String preferredNode = nodeIds.get(2);
    ImmutableOpenMap<String, DiscoveryNode> discoveryNodesMap = generateDiscoveryNodes(nodeIds);
    ClusterService clusterService = generateMockedClusterService(discoveryNodesMap);
    // Create a response to be returned with job route decision info
    List<TrainingJobRouteDecisionInfoNodeResponse> responseList = new ArrayList<>();
    // First node has capacity
    responseList.add(new TrainingJobRouteDecisionInfoNodeResponse(discoveryNodesMap.get(nodeIds.get(0)), 0));
    // Second node has capacity
    responseList.add(new TrainingJobRouteDecisionInfoNodeResponse(discoveryNodesMap.get(nodeIds.get(1)), 0));
    // Third node with capacity
    responseList.add(new TrainingJobRouteDecisionInfoNodeResponse(discoveryNodesMap.get(nodeIds.get(2)), 0));
    TrainingJobRouteDecisionInfoResponse infoResponse = new TrainingJobRouteDecisionInfoResponse(ClusterName.DEFAULT, responseList, Collections.emptyList());
    TransportService transportService = mock(TransportService.class);
    Client client = mock(Client.class);
    // Setup the action
    TrainingJobRouterTransportAction transportAction = new TrainingJobRouterTransportAction(transportService, new ActionFilters(Collections.emptySet()), clusterService, client);
    // Select the node
    DiscoveryNode selectedNode = transportAction.selectNode(preferredNode, infoResponse);
    assertEquals(preferredNode, selectedNode.getId());
}

Example 12 with ClusterService

use of org.opensearch.cluster.service.ClusterService in project k-NN by opensearch-project.

the class TrainingJobRouterTransportActionTests method testMultiNode_withoutCapacity.

public void testMultiNode_withoutCapacity() {
    // Mock datanodes in the cluster through mocking the cluster service
    List<String> nodeIds = ImmutableList.of("node-1", "node-2", "node-3");
    ImmutableOpenMap<String, DiscoveryNode> discoveryNodesMap = generateDiscoveryNodes(nodeIds);
    ClusterService clusterService = generateMockedClusterService(discoveryNodesMap);
    // Create a response to be returned with job route decision info
    List<TrainingJobRouteDecisionInfoNodeResponse> responseList = new ArrayList<>();
    // First node has no capacity
    responseList.add(new TrainingJobRouteDecisionInfoNodeResponse(discoveryNodesMap.get(nodeIds.get(0)), 1));
    // Second node has no capacity
    responseList.add(new TrainingJobRouteDecisionInfoNodeResponse(discoveryNodesMap.get(nodeIds.get(1)), 1));
    // Third node has no capacity
    responseList.add(new TrainingJobRouteDecisionInfoNodeResponse(discoveryNodesMap.get(nodeIds.get(1)), 1));
    TrainingJobRouteDecisionInfoResponse infoResponse = new TrainingJobRouteDecisionInfoResponse(ClusterName.DEFAULT, responseList, Collections.emptyList());
    TransportService transportService = mock(TransportService.class);
    Client client = mock(Client.class);
    // Setup the action
    TrainingJobRouterTransportAction transportAction = new TrainingJobRouterTransportAction(transportService, new ActionFilters(Collections.emptySet()), clusterService, client);
    // Select the node
    DiscoveryNode selectedNode = transportAction.selectNode(null, infoResponse);
    assertNull(selectedNode);
}

Example 13 with ClusterService

use of org.opensearch.cluster.service.ClusterService in project k-NN by opensearch-project.

the class TrainingJobRouterTransportActionTests method testTrainingIndexSize.

@SuppressWarnings("unchecked")
public void testTrainingIndexSize() {
    String trainingIndexName = "training-index";
    int dimension = 133;
    int vectorCount = 1000000;
    // 519,531.25 KB ~= 520 MB
    int expectedSize = dimension * vectorCount * Float.BYTES / BYTES_PER_KILOBYTES + 1;
    // Setup the request
    TrainingModelRequest trainingModelRequest = new TrainingModelRequest(null, KNNMethodContext.getDefault(), dimension, trainingIndexName, "training-field", null, "description");
    // Mock client to return the right number of docs
    TotalHits totalHits = new TotalHits(vectorCount, TotalHits.Relation.EQUAL_TO);
    SearchHits searchHits = new SearchHits(new SearchHit[2], totalHits, 1.0f);
    SearchResponse searchResponse = mock(SearchResponse.class);
    when(searchResponse.getHits()).thenReturn(searchHits);
    Client client = mock(Client.class);
    doAnswer(invocationOnMock -> {
        ((ActionListener<SearchResponse>) invocationOnMock.getArguments()[1]).onResponse(searchResponse);
        return null;
    }).when(client).search(any(), any());
    // Setup the action
    ClusterService clusterService = mock(ClusterService.class);
    TransportService transportService = mock(TransportService.class);
    TrainingJobRouterTransportAction transportAction = new TrainingJobRouterTransportAction(transportService, new ActionFilters(Collections.emptySet()), clusterService, client);
    ActionListener<Integer> listener = ActionListener.wrap(size -> assertEquals(expectedSize, size.intValue()), e -> fail(e.getMessage()));
    transportAction.getTrainingIndexSizeInKB(trainingModelRequest, listener);
}

Example 14 with ClusterService

use of org.opensearch.cluster.service.ClusterService in project k-NN by opensearch-project.

the class TrainingJobRouterTransportActionTests method generateMockedClusterService.

private ClusterService generateMockedClusterService(ImmutableOpenMap<String, DiscoveryNode> discoveryNodeMap) {
    DiscoveryNodes discoveryNodes = mock(DiscoveryNodes.class);
    when(discoveryNodes.getDataNodes()).thenReturn(discoveryNodeMap);
    ClusterState clusterState = mock(ClusterState.class);
    when(clusterState.nodes()).thenReturn(discoveryNodes);
    ClusterService clusterService = mock(ClusterService.class);
    when(clusterService.state()).thenReturn(clusterState);
    return clusterService;
}

Example 15 with ClusterService

use of org.opensearch.cluster.service.ClusterService in project k-NN by opensearch-project.

the class TrainingModelRequestTests method testValidation_invalid_trainingIndexDoesNotExist.

public void testValidation_invalid_trainingIndexDoesNotExist() {
    // Check that validation produces exception when the training index doesnt exist
    // Setup the training request
    String modelId = "test-model-id";
    KNNMethodContext knnMethodContext = mock(KNNMethodContext.class);
    when(knnMethodContext.validate()).thenReturn(null);
    when(knnMethodContext.isTrainingRequired()).thenReturn(true);
    int dimension = 10;
    String trainingIndex = "test-training-index";
    String trainingField = "test-training-field";
    TrainingModelRequest trainingModelRequest = new TrainingModelRequest(modelId, knnMethodContext, dimension, trainingIndex, trainingField, null, null);
    // Mock the model dao to return null so that no exception is produced
    ModelDao modelDao = mock(ModelDao.class);
    when(modelDao.getMetadata(modelId)).thenReturn(null);
    Metadata metadata = mock(Metadata.class);
    when(metadata.index(trainingIndex)).thenReturn(null);
    ClusterState clusterState = mock(ClusterState.class);
    when(clusterState.metadata()).thenReturn(metadata);
    ClusterService clusterService = mock(ClusterService.class);
    when(clusterService.state()).thenReturn(clusterState);
    // Initialize static components with the mocks
    TrainingModelRequest.initialize(modelDao, clusterService);
    // Test that validation produces model already exists error message
    ActionRequestValidationException exception = trainingModelRequest.validate();
    assertNotNull(exception);
    List<String> validationErrors = exception.validationErrors();
    assertEquals(1, validationErrors.size());
    assertTrue(validationErrors.get(0).contains("does not exist"));
}