Examples with MODEL_INDEX_NAME org.opensearch.knn.common.KNNConstants.MODEL_INDEX_NAME used on opensource projects

Example 1 with MODEL_INDEX_NAME

use of org.opensearch.knn.common.KNNConstants.MODEL_INDEX_NAME in project k-NN by opensearch-project.

the class UpdateModelMetadataTransportActionTests method testMasterOperation.

public void testMasterOperation() throws InterruptedException {
    // Setup the Model system index
    createIndex(MODEL_INDEX_NAME);
    // Setup the model
    String modelId = "test-model";
    ModelMetadata modelMetadata = new ModelMetadata(KNNEngine.DEFAULT, SpaceType.L2, 128, ModelState.CREATED, ZonedDateTime.now(ZoneOffset.UTC).toString(), "", "");
    // Get update  transport action
    UpdateModelMetadataTransportAction updateModelMetadataTransportAction = node().injector().getInstance(UpdateModelMetadataTransportAction.class);
    // Generate update request
    UpdateModelMetadataRequest updateModelMetadataRequest = new UpdateModelMetadataRequest(modelId, false, modelMetadata);
    // Get cluster state, update metadata, check cluster state - all asynchronously
    final CountDownLatch inProgressLatch1 = new CountDownLatch(1);
    client().admin().cluster().prepareState().execute(ActionListener.wrap(stateResponse1 -> {
        ClusterState clusterState1 = stateResponse1.getState();
        updateModelMetadataTransportAction.masterOperation(updateModelMetadataRequest, clusterState1, ActionListener.wrap(acknowledgedResponse -> {
            assertTrue(acknowledgedResponse.isAcknowledged());
            client().admin().cluster().prepareState().execute(ActionListener.wrap(stateResponse2 -> {
                ClusterState updatedClusterState = stateResponse2.getState();
                IndexMetadata indexMetadata = updatedClusterState.metadata().index(MODEL_INDEX_NAME);
                assertNotNull(indexMetadata);
                Map<String, String> modelMetadataMap = indexMetadata.getCustomData(MODEL_METADATA_FIELD);
                assertNotNull(modelMetadataMap);
                String modelAsString = modelMetadataMap.get(modelId);
                assertNotNull(modelAsString);
                ModelMetadata modelMetadataCopy = ModelMetadata.fromString(modelAsString);
                assertEquals(modelMetadata, modelMetadataCopy);
                inProgressLatch1.countDown();
            }, e -> fail("Update failed:" + e)));
        }, e -> fail("Update failed: " + e)));
    }, e -> fail("Update failed: " + e)));
    assertTrue(inProgressLatch1.await(60, TimeUnit.SECONDS));
    // Generate remove request
    UpdateModelMetadataRequest removeModelMetadataRequest = new UpdateModelMetadataRequest(modelId, true, modelMetadata);
    final CountDownLatch inProgressLatch2 = new CountDownLatch(1);
    client().admin().cluster().prepareState().execute(ActionListener.wrap(stateResponse1 -> {
        ClusterState clusterState1 = stateResponse1.getState();
        updateModelMetadataTransportAction.masterOperation(removeModelMetadataRequest, clusterState1, ActionListener.wrap(acknowledgedResponse -> {
            assertTrue(acknowledgedResponse.isAcknowledged());
            client().admin().cluster().prepareState().execute(ActionListener.wrap(stateResponse2 -> {
                ClusterState updatedClusterState = stateResponse2.getState();
                IndexMetadata indexMetadata = updatedClusterState.metadata().index(MODEL_INDEX_NAME);
                assertNotNull(indexMetadata);
                Map<String, String> modelMetadataMap = indexMetadata.getCustomData(MODEL_METADATA_FIELD);
                assertNotNull(modelMetadataMap);
                String modelAsString = modelMetadataMap.get(modelId);
                assertNull(modelAsString);
                inProgressLatch2.countDown();
            }, e -> fail("Update failed")));
        }, e -> fail("Update failed")));
    }, e -> fail("Update failed")));
    assertTrue(inProgressLatch2.await(60, TimeUnit.SECONDS));
}