diff --git a/src/main/java/org/opensearch/knn/index/KNNSettings.java b/src/main/java/org/opensearch/knn/index/KNNSettings.java index 01806ad6a..3a790a78f 100644 --- a/src/main/java/org/opensearch/knn/index/KNNSettings.java +++ b/src/main/java/org/opensearch/knn/index/KNNSettings.java @@ -825,6 +825,14 @@ public static double getCircuitBreakerUnsetPercentage() { return KNNSettings.state().getSettingValue(KNNSettings.KNN_CIRCUIT_BREAKER_UNSET_PERCENTAGE); } + /** + * Gets the remote build service endpoint. + * @return String representation of the remote build service endpoint URL + */ + public static String getRemoteBuildServiceEndpoint() { + return KNNSettings.state().getSettingValue(KNNSettings.KNN_REMOTE_BUILD_SERVICE_ENDPOINT); + } + public static boolean isFaissAVX2Disabled() { try { return KNNSettings.state().getSettingValue(KNNSettings.KNN_FAISS_AVX2_DISABLED); diff --git a/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java b/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java index e001fbd07..3e7fc101f 100644 --- a/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java +++ b/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java @@ -17,6 +17,7 @@ import org.opensearch.knn.index.remote.RemoteBuildResponse; import org.opensearch.knn.index.remote.RemoteIndexClient; import org.opensearch.knn.index.remote.RemoteIndexClientFactory; +import org.opensearch.knn.index.remote.RemoteStatusResponse; import org.opensearch.repositories.RepositoriesService; import org.opensearch.repositories.Repository; import org.opensearch.repositories.RepositoryMissingException; diff --git a/src/main/java/org/opensearch/knn/index/engine/KNNEngine.java b/src/main/java/org/opensearch/knn/index/engine/KNNEngine.java index b1f4338af..4355d5221 100644 --- a/src/main/java/org/opensearch/knn/index/engine/KNNEngine.java +++ b/src/main/java/org/opensearch/knn/index/engine/KNNEngine.java @@ -11,6 +11,7 @@ import org.opensearch.knn.index.engine.faiss.Faiss; import org.opensearch.knn.index.engine.lucene.Lucene; import org.opensearch.knn.index.engine.nmslib.Nmslib; +import org.opensearch.knn.index.remote.RemoteIndexParameters; import java.util.List; import java.util.Map; @@ -177,11 +178,6 @@ public KNNLibraryIndexingContext getKNNLibraryIndexingContext( return knnLibrary.getKNNLibraryIndexingContext(knnMethodContext, knnMethodConfigContext); } - @Override - public Map getRemoteIndexingParameters(Map indexInfoParameters) { - return knnLibrary.getRemoteIndexingParameters(indexInfoParameters); - } - @Override public KNNLibrarySearchContext getKNNLibrarySearchContext(String methodName) { return knnLibrary.getKNNLibrarySearchContext(methodName); @@ -221,4 +217,9 @@ public ResolvedMethodContext resolveMethod( public boolean supportsRemoteIndexBuild() { return knnLibrary.supportsRemoteIndexBuild(); } + + @Override + public RemoteIndexParameters createRemoteIndexingParameters(Map indexInfoParameters) { + return knnLibrary.createRemoteIndexingParameters(indexInfoParameters); + } } diff --git a/src/main/java/org/opensearch/knn/index/engine/KNNLibrary.java b/src/main/java/org/opensearch/knn/index/engine/KNNLibrary.java index 29eb36dd2..bd62f1fd4 100644 --- a/src/main/java/org/opensearch/knn/index/engine/KNNLibrary.java +++ b/src/main/java/org/opensearch/knn/index/engine/KNNLibrary.java @@ -7,6 +7,7 @@ import org.opensearch.common.ValidationException; import org.opensearch.knn.index.SpaceType; +import org.opensearch.knn.index.remote.RemoteIndexParameters; import java.util.Collections; import java.util.List; @@ -149,11 +150,7 @@ default boolean supportsRemoteIndexBuild() { return false; } - /** - * Get the remote build supported index parameter mapping to be sent to the remote build service. - * @param indexInfoParameters the index parameters from BuildIndexParams - */ - default Map getRemoteIndexingParameters(Map indexInfoParameters) { - throw new UnsupportedOperationException("This method must be implemented by the implementing class"); + default RemoteIndexParameters createRemoteIndexingParameters(Map indexInfoParameters) { + throw new UnsupportedOperationException("Remote build service does not support this engine"); } } diff --git a/src/main/java/org/opensearch/knn/index/engine/faiss/Faiss.java b/src/main/java/org/opensearch/knn/index/engine/faiss/Faiss.java index bdbf46a0b..8a0409dcd 100644 --- a/src/main/java/org/opensearch/knn/index/engine/faiss/Faiss.java +++ b/src/main/java/org/opensearch/knn/index/engine/faiss/Faiss.java @@ -14,30 +14,14 @@ import org.opensearch.knn.index.engine.MethodResolver; import org.opensearch.knn.index.engine.NativeLibrary; import org.opensearch.knn.index.engine.ResolvedMethodContext; +import org.opensearch.knn.index.remote.RemoteIndexParameters; -import java.util.HashMap; import java.util.Map; import java.util.function.Function; -import static org.opensearch.knn.common.KNNConstants.INDEX_DESCRIPTION_PARAMETER; import static org.opensearch.knn.common.KNNConstants.METHOD_HNSW; import static org.opensearch.knn.common.KNNConstants.METHOD_IVF; -import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION; -import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_EF_SEARCH; -import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_M; -import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_NLIST; -import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_NLIST_DEFAULT; -import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_NPROBES; -import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_NPROBES_DEFAULT; -import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_SPACE_TYPE; import static org.opensearch.knn.common.KNNConstants.NAME; -import static org.opensearch.knn.common.KNNConstants.PARAMETERS; -import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE; -import static org.opensearch.knn.index.remote.KNNRemoteConstants.ALGORITHM; -import static org.opensearch.knn.index.remote.KNNRemoteConstants.ALGORITHM_PARAMETERS; -import static org.opensearch.knn.index.KNNSettings.INDEX_KNN_DEFAULT_ALGO_PARAM_EF_CONSTRUCTION; -import static org.opensearch.knn.index.KNNSettings.INDEX_KNN_DEFAULT_ALGO_PARAM_EF_SEARCH; -import static org.opensearch.knn.index.KNNSettings.INDEX_KNN_DEFAULT_SPACE_TYPE; /** * Implements NativeLibrary for the faiss native library @@ -127,67 +111,6 @@ public Float scoreToRadialThreshold(Float score, SpaceType spaceType) { return spaceType.scoreToDistanceTranslation(score); } - // TODO refactor to make the index parameter fetching more intelligent and less cumbersome - /** - * Get the parameters that need to be passed to the remote build service for training - * - * @param indexInfoParameters result of indexInfo.getParameters() to parse - * @return Map of parameters to be used as "index_parameters" - */ - @Override - public Map getRemoteIndexingParameters(Map indexInfoParameters) { - Map indexParameters = new HashMap<>(); - String methodName = (String) indexInfoParameters.get(NAME); - indexParameters.put(ALGORITHM, methodName); - indexParameters.put(METHOD_PARAMETER_SPACE_TYPE, indexInfoParameters.getOrDefault(SPACE_TYPE, INDEX_KNN_DEFAULT_SPACE_TYPE)); - - assert (indexInfoParameters.containsKey(PARAMETERS)); - Object innerParams = indexInfoParameters.get(PARAMETERS); - assert (innerParams instanceof Map); - { - Map algorithmParams = new HashMap<>(); - Map innerMap = (Map) innerParams; - switch (methodName) { - case METHOD_HNSW -> { - algorithmParams.put( - METHOD_PARAMETER_EF_CONSTRUCTION, - innerMap.getOrDefault(METHOD_PARAMETER_EF_CONSTRUCTION, INDEX_KNN_DEFAULT_ALGO_PARAM_EF_CONSTRUCTION) - ); - algorithmParams.put( - METHOD_PARAMETER_EF_SEARCH, - innerMap.getOrDefault(METHOD_PARAMETER_EF_SEARCH, INDEX_KNN_DEFAULT_ALGO_PARAM_EF_SEARCH) - ); - Object indexDescription = indexInfoParameters.get(INDEX_DESCRIPTION_PARAMETER); - assert indexDescription instanceof String; - algorithmParams.put(METHOD_PARAMETER_M, getMFromIndexDescription((String) indexDescription)); - } - case METHOD_IVF -> { - algorithmParams.put( - METHOD_PARAMETER_NLIST, - innerMap.getOrDefault(METHOD_PARAMETER_NLIST, METHOD_PARAMETER_NLIST_DEFAULT) - ); - algorithmParams.put( - METHOD_PARAMETER_NPROBES, - innerMap.getOrDefault(METHOD_PARAMETER_NPROBES, METHOD_PARAMETER_NPROBES_DEFAULT) - ); - } - } - indexParameters.put(ALGORITHM_PARAMETERS, algorithmParams); - } - return indexParameters; - } - - public static int getMFromIndexDescription(String indexDescription) { - int commaIndex = indexDescription.indexOf(","); - if (commaIndex == -1) { - throw new IllegalArgumentException("Invalid index description: " + indexDescription); - } - String hnswPart = indexDescription.substring(0, commaIndex); - int m = Integer.parseInt(hnswPart.substring(4)); - assert (m > 1 && m < 100); - return m; - } - @Override public ResolvedMethodContext resolveMethod( KNNMethodContext knnMethodContext, @@ -202,4 +125,12 @@ public ResolvedMethodContext resolveMethod( public boolean supportsRemoteIndexBuild() { return true; } + + @Override + public RemoteIndexParameters createRemoteIndexingParameters(Map indexInfoParameters) { + if (METHOD_HNSW.equals(indexInfoParameters.get(NAME))) { + return FaissHNSWMethod.getRemoteIndexingParameters(indexInfoParameters); + } + throw new IllegalArgumentException("Unsupported method for remote indexing"); + } } diff --git a/src/main/java/org/opensearch/knn/index/engine/faiss/FaissHNSWMethod.java b/src/main/java/org/opensearch/knn/index/engine/faiss/FaissHNSWMethod.java index b1b0f0063..b801da89e 100644 --- a/src/main/java/org/opensearch/knn/index/engine/faiss/FaissHNSWMethod.java +++ b/src/main/java/org/opensearch/knn/index/engine/faiss/FaissHNSWMethod.java @@ -13,12 +13,14 @@ import org.opensearch.knn.index.engine.AbstractKNNMethod; import org.opensearch.knn.index.engine.DefaultHnswSearchContext; import org.opensearch.knn.index.engine.Encoder; +import org.opensearch.knn.index.engine.KNNMethodContext; import org.opensearch.knn.index.engine.MethodComponent; import org.opensearch.knn.index.engine.MethodComponentContext; import org.opensearch.knn.index.engine.Parameter; -import org.opensearch.knn.index.engine.KNNMethodContext; import org.opensearch.knn.index.engine.TrainingConfigValidationInput; import org.opensearch.knn.index.engine.TrainingConfigValidationOutput; +import org.opensearch.knn.index.remote.RemoteFaissHNSWIndexParameters; +import org.opensearch.knn.index.remote.RemoteIndexParameters; import java.util.Arrays; import java.util.Collections; @@ -29,11 +31,14 @@ import java.util.stream.Collectors; import static org.opensearch.knn.common.KNNConstants.FAISS_HNSW_DESCRIPTION; +import static org.opensearch.knn.common.KNNConstants.INDEX_DESCRIPTION_PARAMETER; import static org.opensearch.knn.common.KNNConstants.METHOD_ENCODER_PARAMETER; import static org.opensearch.knn.common.KNNConstants.METHOD_HNSW; import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION; import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_EF_SEARCH; import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_M; +import static org.opensearch.knn.common.KNNConstants.PARAMETERS; +import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE; /** * Faiss HNSW method implementation @@ -147,4 +152,41 @@ protected Function indexInfoParameters) { + RemoteFaissHNSWIndexParameters.RemoteFaissHNSWIndexParametersBuilder builder = RemoteFaissHNSWIndexParameters.builder(); + assert (indexInfoParameters.get(SPACE_TYPE) instanceof String); + String spaceType = (String) indexInfoParameters.get(SPACE_TYPE); + builder.algorithm(METHOD_HNSW).spaceType(spaceType); + + Object innerParams = indexInfoParameters.get(PARAMETERS); + assert (innerParams instanceof Map); + Map innerMap = (Map) innerParams; + assert (innerMap.get(METHOD_PARAMETER_EF_CONSTRUCTION) instanceof Integer); + builder.efConstruction((Integer) innerMap.get(METHOD_PARAMETER_EF_CONSTRUCTION)); + assert (innerMap.get(METHOD_PARAMETER_EF_SEARCH) instanceof Integer); + builder.efSearch((Integer) innerMap.get(METHOD_PARAMETER_EF_SEARCH)); + Object indexDescription = indexInfoParameters.get(INDEX_DESCRIPTION_PARAMETER); + assert indexDescription instanceof String; + builder.m(getMFromIndexDescription((String) indexDescription)); + + return builder.build(); + } + + public static int getMFromIndexDescription(String indexDescription) { + int commaIndex = indexDescription.indexOf(","); + if (commaIndex == -1) { + throw new IllegalArgumentException("Invalid index description: " + indexDescription); + } + String hnswPart = indexDescription.substring(0, commaIndex); + int m = Integer.parseInt(hnswPart.substring(4)); + return m; + } } diff --git a/src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequest.java b/src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequest.java index 42896fcc6..d28e3dfe5 100644 --- a/src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequest.java +++ b/src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequest.java @@ -16,7 +16,6 @@ import org.opensearch.knn.index.vectorvalues.KNNVectorValues; import java.io.IOException; -import java.util.Map; import static org.opensearch.knn.common.KNNConstants.DIMENSION; import static org.opensearch.knn.common.KNNConstants.KNN_ENGINE; @@ -49,7 +48,7 @@ public class RemoteBuildRequest implements ToXContentObject { protected int docCount; protected String vectorDataType; protected String engine; - protected Map indexParameters; + protected RemoteIndexParameters indexParameters; /** * Constructor for RemoteBuildRequest. @@ -80,8 +79,6 @@ public RemoteBuildRequest( KNNCodecUtil.initializeVectorValues(vectorValues); assert (vectorValues.dimension() > 0); - Map indexParameters = indexInfo.getKnnEngine().getRemoteIndexingParameters(indexInfo.getParameters()); - this.repositoryType = repositoryType; this.containerName = containerName; this.vectorPath = blobName + VECTOR_BLOB_FILE_EXTENSION; @@ -91,7 +88,7 @@ public RemoteBuildRequest( this.docCount = indexInfo.getTotalLiveDocs(); this.vectorDataType = vectorDataType; this.engine = indexInfo.getKnnEngine().getName(); - this.indexParameters = indexParameters; + this.indexParameters = indexInfo.getKnnEngine().createRemoteIndexingParameters(indexInfo.getParameters()); } @Override diff --git a/src/main/java/org/opensearch/knn/index/remote/RemoteBuildResponse.java b/src/main/java/org/opensearch/knn/index/remote/RemoteBuildResponse.java index e28889d5d..88bcdce3a 100644 --- a/src/main/java/org/opensearch/knn/index/remote/RemoteBuildResponse.java +++ b/src/main/java/org/opensearch/knn/index/remote/RemoteBuildResponse.java @@ -42,7 +42,7 @@ public static RemoteBuildResponse fromXContent(XContentParser parser) throws IOE } } if (StringUtils.isBlank(builder.jobId)) { - throw new IOException("Invalid response format, missing + " + JOB_ID_FIELD); + throw new IOException("Invalid response format, missing " + JOB_ID_FIELD); } return builder.build(); } diff --git a/src/main/java/org/opensearch/knn/index/remote/RemoteFaissHNSWIndexParameters.java b/src/main/java/org/opensearch/knn/index/remote/RemoteFaissHNSWIndexParameters.java new file mode 100644 index 000000000..aca33f298 --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/remote/RemoteFaissHNSWIndexParameters.java @@ -0,0 +1,30 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.remote; + +import lombok.experimental.SuperBuilder; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.knn.common.KNNConstants; + +import java.io.IOException; + +import static org.opensearch.knn.index.remote.KNNRemoteConstants.ALGORITHM_PARAMETERS; + +@SuperBuilder +public class RemoteFaissHNSWIndexParameters extends RemoteIndexParameters { + int m; + int efConstruction; + int efSearch; + + @Override + void addAlgorithmParameters(XContentBuilder builder) throws IOException { + builder.startObject(ALGORITHM_PARAMETERS); + builder.field(KNNConstants.METHOD_PARAMETER_M, m); + builder.field(KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION, efConstruction); + builder.field(KNNConstants.METHOD_PARAMETER_EF_SEARCH, efSearch); + builder.endObject(); + } +} diff --git a/src/main/java/org/opensearch/knn/index/remote/RemoteIndexClient.java b/src/main/java/org/opensearch/knn/index/remote/RemoteIndexClient.java index 8040a3271..dfdea6bf7 100644 --- a/src/main/java/org/opensearch/knn/index/remote/RemoteIndexClient.java +++ b/src/main/java/org/opensearch/knn/index/remote/RemoteIndexClient.java @@ -5,8 +5,6 @@ package org.opensearch.knn.index.remote; -import org.opensearch.knn.index.codec.nativeindex.remote.RemoteStatusResponse; - import java.io.IOException; /** diff --git a/src/main/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClient.java b/src/main/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClient.java index cda7909d1..91447f565 100644 --- a/src/main/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClient.java +++ b/src/main/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClient.java @@ -25,7 +25,9 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.knn.index.KNNSettings; -import org.opensearch.knn.index.codec.nativeindex.remote.RemoteStatusResponse; +import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams; +import org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy; +import org.opensearch.knn.plugin.KNNPlugin; import java.io.Closeable; import java.io.IOException; @@ -37,7 +39,6 @@ import static org.apache.hc.core5.http.HttpStatus.SC_OK; import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_CLIENT_PASSWORD_SETTING; import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_CLIENT_USERNAME_SETTING; -import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING; import static org.opensearch.knn.index.remote.KNNRemoteConstants.BUILD_ENDPOINT; /** @@ -68,8 +69,12 @@ protected static CloseableHttpClient getHttpClient() { return HttpClientHolder.httpClient; } + /** + * Creates the client, setting the endpoint per-instance so the same endpoint is used per-build operation + * (per call to {@link RemoteIndexBuildStrategy#buildAndWriteIndex(BuildIndexParams)}) + */ public RemoteIndexHTTPClient() { - String endpoint = KNNSettings.state().getSettingValue(KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING.getKey()); + String endpoint = KNNSettings.getRemoteBuildServiceEndpoint(); if (endpoint == null || endpoint.isEmpty()) { throw new IllegalArgumentException("No endpoint set for RemoteIndexClient"); } @@ -143,7 +148,8 @@ private String toJson(ToXContentObject object) throws IOException { } /** - * Set the global auth header to use the refreshed secure settings + * Set the global auth header to use the refreshed secure settings. + * Called by {@link KNNPlugin#reload(Settings)} when the nodes reload API is called. * @param settings Settings to use to get the credentials */ public static void reloadAuthHeader(Settings settings) { diff --git a/src/main/java/org/opensearch/knn/index/remote/RemoteIndexParameters.java b/src/main/java/org/opensearch/knn/index/remote/RemoteIndexParameters.java new file mode 100644 index 000000000..1b7831f3c --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/remote/RemoteIndexParameters.java @@ -0,0 +1,34 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.remote; + +import lombok.experimental.SuperBuilder; +import org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; + +import java.io.IOException; + +import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_SPACE_TYPE; +import static org.opensearch.knn.index.remote.KNNRemoteConstants.ALGORITHM; + +@SuperBuilder +public abstract class RemoteIndexParameters implements ToXContentObject { + String spaceType; + String algorithm; + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.field(METHOD_PARAMETER_SPACE_TYPE, spaceType); + builder.field(ALGORITHM, algorithm); + addAlgorithmParameters(builder); + builder.endObject(); + return builder; + } + + abstract void addAlgorithmParameters(XContentBuilder builder) throws IOException; + +} diff --git a/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteStatusResponse.java b/src/main/java/org/opensearch/knn/index/remote/RemoteStatusResponse.java similarity index 80% rename from src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteStatusResponse.java rename to src/main/java/org/opensearch/knn/index/remote/RemoteStatusResponse.java index a9a87722a..2751eee34 100644 --- a/src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteStatusResponse.java +++ b/src/main/java/org/opensearch/knn/index/remote/RemoteStatusResponse.java @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.knn.index.codec.nativeindex.remote; +package org.opensearch.knn.index.remote; import lombok.AllArgsConstructor; import lombok.Getter; diff --git a/src/test/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClientTests.java b/src/test/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClientTests.java index d4d46b343..f11bee7ab 100644 --- a/src/test/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClientTests.java +++ b/src/test/java/org/opensearch/knn/index/remote/RemoteIndexHTTPClientTests.java @@ -49,6 +49,7 @@ import java.util.Set; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -64,10 +65,11 @@ import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD; import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_CLIENT_PASSWORD_SETTING; import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_CLIENT_USERNAME_SETTING; +import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_SERVICE_ENDPOINT; import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING; import static org.opensearch.knn.index.SpaceType.L2; import static org.opensearch.knn.index.VectorDataType.FLOAT; -import static org.opensearch.knn.index.engine.faiss.Faiss.getMFromIndexDescription; +import static org.opensearch.knn.index.engine.faiss.FaissHNSWMethod.getMFromIndexDescription; import static org.opensearch.knn.index.remote.KNNRemoteConstants.BUCKET; import static org.opensearch.knn.index.remote.KNNRemoteConstants.BUILD_ENDPOINT; import static org.opensearch.knn.index.remote.KNNRemoteConstants.DOC_ID_FILE_EXTENSION; @@ -122,7 +124,7 @@ public void testBuildRequest() { try (MockedStatic knnSettingsStaticMock = Mockito.mockStatic(KNNSettings.class)) { knnSettingsStaticMock.when(KNNSettings::state).thenReturn(knnSettingsMock); - when(knnSettingsMock.getSettingValue(KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING.getKey())).thenReturn(MOCK_ENDPOINT); + when(KNNSettings.getRemoteBuildServiceEndpoint()).thenReturn(MOCK_ENDPOINT); KNNSettings.state().setClusterService(clusterService); BuildIndexParams indexInfo = createTestBuildIndexParams(); @@ -226,7 +228,7 @@ public void testSubmitVectorBuild() throws IOException, URISyntaxException { try (MockedStatic knnSettingsStaticMock = Mockito.mockStatic(KNNSettings.class)) { knnSettingsStaticMock.when(KNNSettings::state).thenReturn(knnSettingsMock); - when(knnSettingsMock.getSettingValue(KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING.getKey())).thenReturn(MOCK_ENDPOINT); + when(KNNSettings.getRemoteBuildServiceEndpoint()).thenReturn(MOCK_ENDPOINT); KNNSettings.state().setClusterService(clusterService); BuildIndexParams buildIndexParams = createTestBuildIndexParams(); @@ -255,7 +257,10 @@ public void testSecureSettingsReloadAndException() throws IOException { final MockSecureSettings secureSettings = new MockSecureSettings(); secureSettings.setString(KNN_REMOTE_BUILD_CLIENT_USERNAME_SETTING.getKey(), USERNAME); secureSettings.setString(KNN_REMOTE_BUILD_CLIENT_PASSWORD_SETTING.getKey(), PASSWORD); - final Settings settings = Settings.builder().setSecureSettings(secureSettings).build(); + final Settings settings = Settings.builder() + .setSecureSettings(secureSettings) + .put(KNN_REMOTE_BUILD_SERVICE_ENDPOINT, MOCK_ENDPOINT) + .build(); CloseableHttpClient mockHttpClient = mock(CloseableHttpClient.class); try (MockedStatic clientStaticMock = Mockito.mockStatic(RemoteIndexHTTPClient.class)) { @@ -264,7 +269,7 @@ public void testSecureSettingsReloadAndException() throws IOException { try (MockedStatic knnSettingsStaticMock = Mockito.mockStatic(KNNSettings.class)) { KNNSettings knnSettingsMock = mock(KNNSettings.class); knnSettingsStaticMock.when(KNNSettings::state).thenReturn(knnSettingsMock); - when(knnSettingsMock.getSettingValue(KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING.getKey())).thenReturn(MOCK_ENDPOINT); + when(KNNSettings.getRemoteBuildServiceEndpoint()).thenReturn(MOCK_ENDPOINT); KNNSettings.state().setClusterService(clusterService); when(mockHttpClient.execute(any(HttpPost.class), any(HttpClientResponseHandler.class))).thenAnswer( @@ -349,9 +354,12 @@ private IndexSettings createTestIndexSettings() { } private void setupTestClusterSettings() { - ClusterSettings clusterSettings = mock(ClusterSettings.class); - when(clusterSettings.get(KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING)).thenReturn(MOCK_ENDPOINT); - when(clusterService.getClusterSettings()).thenReturn(clusterSettings); + Settings settings = Settings.builder().put(KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING.getKey(), MOCK_ENDPOINT).build(); + Set> settingsSet = new HashSet<>(); + settingsSet.add(KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING); + ClusterSettings clusterSettings = new ClusterSettings(settings, settingsSet); + doReturn(clusterSettings).when(clusterService).getClusterSettings(); + doReturn(settings).when(clusterService).getSettings(); KNNSettings.state().setClusterService(clusterService); }