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

Example 46 with ClusterService

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

the class NativeMemoryEntryContextTests method testTrainingDataEntryContext_getIndicesService.

public void testTrainingDataEntryContext_getIndicesService() {
    ClusterService clusterService = mock(ClusterService.class);
    NativeMemoryEntryContext.TrainingDataEntryContext trainingDataEntryContext = new NativeMemoryEntryContext.TrainingDataEntryContext(0, "test", "test", null, clusterService, 0, 0);
    assertEquals(clusterService, trainingDataEntryContext.getClusterService());
}

Example 47 with ClusterService

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

the class ModelCacheTests method testGet_modelDoesNotFitInCache.

public void testGet_modelDoesNotFitInCache() throws ExecutionException, InterruptedException {
    String modelId = "test-model-id";
    int dimension = 2;
    String cacheSize = "1kb";
    Model mockModel = new Model(new ModelMetadata(KNNEngine.DEFAULT, SpaceType.DEFAULT, dimension, ModelState.CREATED, ZonedDateTime.now(ZoneOffset.UTC).toString(), "", ""), new byte[BYTES_PER_KILOBYTES + 1], modelId);
    ModelDao modelDao = mock(ModelDao.class);
    when(modelDao.get(modelId)).thenReturn(mockModel);
    Settings settings = Settings.builder().put(MODEL_CACHE_SIZE_LIMIT_SETTING.getKey(), cacheSize).build();
    ClusterSettings clusterSettings = new ClusterSettings(settings, ImmutableSet.of(MODEL_CACHE_SIZE_LIMIT_SETTING));
    ClusterService clusterService = mock(ClusterService.class);
    when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
    when(clusterService.getSettings()).thenReturn(settings);
    ModelCache.initialize(modelDao, clusterService);
    ModelCache modelCache = new ModelCache();
    assertEquals(mockModel, modelCache.get(modelId));
    assertFalse(modelCache.contains(modelId));
}

Example 48 with ClusterService

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

the class ModelCacheTests method testRebuild_afterSettingUpdate.

public void testRebuild_afterSettingUpdate() throws ExecutionException, InterruptedException {
    String modelId = "test-model-id";
    int dimension = 2;
    int modelSize = 2 * BYTES_PER_KILOBYTES;
    Model mockModel = new Model(new ModelMetadata(KNNEngine.DEFAULT, SpaceType.DEFAULT, dimension, ModelState.CREATED, ZonedDateTime.now(ZoneOffset.UTC).toString(), "", ""), new byte[modelSize], modelId);
    String cacheSize1 = "1kb";
    String cacheSize2 = "4kb";
    ModelDao modelDao = mock(ModelDao.class);
    when(modelDao.get(modelId)).thenReturn(mockModel);
    Settings settings = Settings.builder().put(MODEL_CACHE_SIZE_LIMIT_SETTING.getKey(), cacheSize1).build();
    ClusterSettings clusterSettings = new ClusterSettings(settings, ImmutableSet.of(MODEL_CACHE_SIZE_LIMIT_SETTING));
    ClusterService clusterService = mock(ClusterService.class);
    when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
    when(clusterService.getSettings()).thenReturn(settings);
    ModelCache.initialize(modelDao, clusterService);
    ModelCache modelCache = new ModelCache();
    // Add element to cache - element should not remain in cache
    modelCache.get(modelId);
    assertEquals(0, modelCache.getTotalWeightInKB());
    // Rebuild and make sure cache is empty
    Settings newSettings = Settings.builder().put(MODEL_CACHE_SIZE_LIMIT_SETTING.getKey(), cacheSize2).build();
    clusterService.getClusterSettings().applySettings(newSettings);
    assertEquals(0, modelCache.getTotalWeightInKB());
    // Add element again - element should remain in cache
    modelCache.get(modelId);
    assertEquals((modelSize / BYTES_PER_KILOBYTES) + 1, modelCache.getTotalWeightInKB());
}

Example 49 with ClusterService

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

the class ModelCacheTests method testContains.

public void testContains() throws ExecutionException, InterruptedException {
    String modelId1 = "test-model-id-1";
    int dimension = 2;
    int modelSize1 = 100;
    Model mockModel1 = new Model(new ModelMetadata(KNNEngine.DEFAULT, SpaceType.DEFAULT, dimension, ModelState.CREATED, ZonedDateTime.now(ZoneOffset.UTC).toString(), "", ""), new byte[modelSize1], modelId1);
    String modelId2 = "test-model-id-2";
    String cacheSize = "10%";
    ModelDao modelDao = mock(ModelDao.class);
    when(modelDao.get(modelId1)).thenReturn(mockModel1);
    Settings settings = Settings.builder().put(MODEL_CACHE_SIZE_LIMIT_SETTING.getKey(), cacheSize).build();
    ClusterSettings clusterSettings = new ClusterSettings(settings, ImmutableSet.of(MODEL_CACHE_SIZE_LIMIT_SETTING));
    ClusterService clusterService = mock(ClusterService.class);
    when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
    when(clusterService.getSettings()).thenReturn(settings);
    ModelCache.initialize(modelDao, clusterService);
    ModelCache modelCache = new ModelCache();
    assertFalse(modelCache.contains(modelId1));
    modelCache.get(modelId1);
    assertTrue(modelCache.contains(modelId1));
    assertFalse(modelCache.contains(modelId2));
}

Example 50 with ClusterService

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

the class ModelCacheTests method testGetTotalWeight.

public void testGetTotalWeight() throws ExecutionException, InterruptedException {
    String modelId1 = "test-model-id-1";
    String modelId2 = "test-model-id-2";
    int dimension = 2;
    String cacheSize = "10%";
    int size1 = BYTES_PER_KILOBYTES;
    Model mockModel1 = new Model(new ModelMetadata(KNNEngine.DEFAULT, SpaceType.DEFAULT, dimension, ModelState.CREATED, ZonedDateTime.now(ZoneOffset.UTC).toString(), "", ""), new byte[size1], modelId1);
    int size2 = BYTES_PER_KILOBYTES * 3;
    Model mockModel2 = new Model(new ModelMetadata(KNNEngine.DEFAULT, SpaceType.DEFAULT, dimension, ModelState.CREATED, ZonedDateTime.now(ZoneOffset.UTC).toString(), "", ""), new byte[size2], modelId2);
    ModelDao modelDao = mock(ModelDao.class);
    when(modelDao.get(modelId1)).thenReturn(mockModel1);
    when(modelDao.get(modelId2)).thenReturn(mockModel2);
    Settings settings = Settings.builder().put(MODEL_CACHE_SIZE_LIMIT_SETTING.getKey(), cacheSize).build();
    ClusterSettings clusterSettings = new ClusterSettings(settings, ImmutableSet.of(MODEL_CACHE_SIZE_LIMIT_SETTING));
    ClusterService clusterService = mock(ClusterService.class);
    when(clusterService.getClusterSettings()).thenReturn(clusterSettings);
    when(clusterService.getSettings()).thenReturn(settings);
    ModelCache.initialize(modelDao, clusterService);
    ModelCache modelCache = new ModelCache();
    modelCache.get(modelId1);
    modelCache.get(modelId2);
    modelCache.get(modelId1);
    modelCache.get(modelId2);
    assertEquals((size1 + size2) / BYTES_PER_KILOBYTES + 2, modelCache.getTotalWeightInKB());
}