Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Backport 1.2] [BUG FIX] Add space type default and ef search parameter in warmup #285

Merged
merged 1 commit into from
Feb 10, 2022
Merged

[Backport 1.2] [BUG FIX] Add space type default and ef search parameter in warmup #285

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions src/main/java/org/opensearch/knn/index/IndexUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,17 +11,23 @@

package org.opensearch.knn.index;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import org.opensearch.cluster.metadata.IndexMetadata;
import org.opensearch.cluster.metadata.MappingMetadata;
import org.opensearch.common.ValidationException;
import org.opensearch.knn.common.KNNConstants;
import org.opensearch.knn.index.util.KNNEngine;
import org.opensearch.knn.indices.ModelDao;
import org.opensearch.knn.indices.ModelMetadata;

import java.io.File;
import java.util.Collections;
import java.util.Map;

import static org.opensearch.knn.common.KNNConstants.BYTES_PER_KILOBYTES;
import static org.opensearch.knn.common.KNNConstants.HNSW_ALGO_EF_SEARCH;
import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE;

public class IndexUtil {

Expand Down Expand Up @@ -156,4 +162,26 @@ public static ValidationException validateKnnField(IndexMetadata indexMetadata,

return null;
}

/**
* Gets the load time parameters for a given engine.
*
* @param spaceType Space for this particular segment
* @param knnEngine Engine used for the native library indices being loaded in
* @param indexName Name of OpenSearch index that the segment files belong to
* @return load parameters that will be passed to the JNI.
*/
public static Map<String, Object> getParametersAtLoading(SpaceType spaceType, KNNEngine knnEngine, String indexName) {
Map<String, Object> loadParameters = Maps.newHashMap(ImmutableMap.of(
SPACE_TYPE, spaceType.getValue()
));

// For nmslib, we need to add the dynamic ef_search parameter that needs to be passed in when the
// hnsw graphs are loaded into memory
if (KNNEngine.NMSLIB.equals(knnEngine)) {
loadParameters.put(HNSW_ALGO_EF_SEARCH, KNNSettings.getEfSearchParam(indexName));
}

return Collections.unmodifiableMap(loadParameters);
}
}
9 changes: 6 additions & 3 deletions src/main/java/org/opensearch/knn/index/KNNIndexShard.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

package org.opensearch.knn.index;

import com.google.common.collect.ImmutableMap;
import org.apache.lucene.index.FieldInfo;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
Expand All @@ -31,6 +30,7 @@
import java.util.stream.Collectors;

import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE;
import static org.opensearch.knn.index.IndexUtil.getParametersAtLoading;
import static org.opensearch.knn.index.codec.KNNCodecUtil.buildEngineFileName;

/**
Expand Down Expand Up @@ -86,7 +86,7 @@ public void warmup() throws IOException {
new NativeMemoryEntryContext.IndexEntryContext(
key,
NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(),
ImmutableMap.of(SPACE_TYPE, value.getValue()),
getParametersAtLoading(value, KNNEngine.getEngineNameFromPath(key), getIndexName()),
getIndexName()
), true);
} catch (ExecutionException ex) {
Expand Down Expand Up @@ -122,7 +122,10 @@ private Map<String, SpaceType> getEnginePaths(IndexReader indexReader, KNNEngine

for (FieldInfo fieldInfo : reader.getFieldInfos()) {
if (fieldInfo.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) {
SpaceType spaceType = SpaceType.getSpace(fieldInfo.attributes().get(SPACE_TYPE));
// Space Type will not be present on ES versions 7.1 and 7.4 because the only available space type
// was L2. So, if Space Type is not present, just fall back to L2
String spaceTypeName = fieldInfo.attributes().getOrDefault(SPACE_TYPE, SpaceType.L2.getValue());
SpaceType spaceType = SpaceType.getSpace(spaceTypeName);
String engineFileName = buildEngineFileName(reader.getSegmentInfo().info.name,
knnEngine.getLatestBuildVersion(), fieldInfo.name, fileExtension);

Expand Down
16 changes: 3 additions & 13 deletions src/main/java/org/opensearch/knn/index/KNNWeight.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
import static org.opensearch.knn.common.KNNConstants.KNN_ENGINE;
import static org.opensearch.knn.common.KNNConstants.MODEL_ID;
import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE;
import static org.opensearch.knn.index.IndexUtil.getParametersAtLoading;
import static org.opensearch.knn.plugin.stats.KNNCounter.GRAPH_QUERY_ERRORS;

/**
Expand Down Expand Up @@ -134,24 +135,13 @@ public Scorer scorer(LeafReaderContext context) throws IOException {
KNNCounter.GRAPH_QUERY_REQUESTS.increment();

// We need to first get index allocation
NativeMemoryAllocation indexAllocation = null;

Map<String, Object> loadParameters = new HashMap<String, Object>() {{
put(SPACE_TYPE, spaceType.getValue());
}};

// For nmslib, we set efSearch from an index setting. This has the advantage of being able to dynamically
// update this value, which we cannot do at the moment for mapping parameters.
if (knnEngine.equals(KNNEngine.NMSLIB)) {
loadParameters.put(KNNConstants.HNSW_ALGO_EF_SEARCH, KNNSettings.getEfSearchParam(knnQuery.getIndexName()));
}

NativeMemoryAllocation indexAllocation;
try {
indexAllocation = nativeMemoryCacheManager.get(
new NativeMemoryEntryContext.IndexEntryContext(
indexPath.toString(),
NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(),
loadParameters,
getParametersAtLoading(spaceType, knnEngine, knnQuery.getIndexName()),
knnQuery.getIndexName()
), true);
} catch (ExecutionException e) {
Expand Down
72 changes: 72 additions & 0 deletions src/test/java/org/opensearch/knn/index/IndexUtilTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*
* Modifications Copyright OpenSearch Contributors. See
* GitHub history for details.
*/

package org.opensearch.knn.index;

import com.google.common.collect.ImmutableMap;
import org.opensearch.cluster.ClusterState;
import org.opensearch.cluster.metadata.IndexMetadata;
import org.opensearch.cluster.metadata.Metadata;
import org.opensearch.cluster.service.ClusterService;
import org.opensearch.common.settings.Settings;
import org.opensearch.knn.KNNTestCase;
import org.opensearch.knn.index.util.KNNEngine;

import java.util.Map;

import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.opensearch.knn.common.KNNConstants.HNSW_ALGO_EF_SEARCH;
import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE;
import static org.opensearch.knn.index.IndexUtil.getParametersAtLoading;
import static org.opensearch.knn.index.KNNSettings.KNN_ALGO_PARAM_EF_SEARCH;

public class IndexUtilTests extends KNNTestCase {
public void testGetLoadParameters() {
// Test faiss to ensure that space type gets set properly
SpaceType spaceType1 = SpaceType.COSINESIMIL;
KNNEngine knnEngine1 = KNNEngine.FAISS;
String indexName = "my-test-index";

Map<String, Object> loadParameters = getParametersAtLoading(spaceType1, knnEngine1, indexName);
assertEquals(1, loadParameters.size());
assertEquals(spaceType1.getValue(), loadParameters.get(SPACE_TYPE));

// Test nmslib to ensure both space type and ef search are properly set
SpaceType spaceType2 = SpaceType.L1;
KNNEngine knnEngine2 = KNNEngine.NMSLIB;
int efSearchValue = 413;

// We use the constant for the setting here as opposed to the identifier of efSearch in nmslib jni
Map<String, Object> indexSettings = ImmutableMap.of(
KNN_ALGO_PARAM_EF_SEARCH, efSearchValue
);

// Because ef search comes from an index setting, we need to mock the long line of calls to get those
// index settings
Settings settings = Settings.builder().loadFromMap(indexSettings).build();
IndexMetadata indexMetadata = mock(IndexMetadata.class);
when(indexMetadata.getSettings()).thenReturn(settings);
Metadata metadata = mock(Metadata.class);
when(metadata.index(anyString())).thenReturn(indexMetadata);
ClusterState clusterState = mock(ClusterState.class);
when(clusterState.getMetadata()).thenReturn(metadata);
ClusterService clusterService = mock(ClusterService.class);
when(clusterService.state()).thenReturn(clusterState);
KNNSettings.state().setClusterService(clusterService);

loadParameters = getParametersAtLoading(spaceType2, knnEngine2, indexName);
assertEquals(2, loadParameters.size());
assertEquals(spaceType2.getValue(), loadParameters.get(SPACE_TYPE));
assertEquals(efSearchValue, loadParameters.get(HNSW_ALGO_EF_SEARCH));
}
}