Rehaul index parameter logic, minor edits, settings change

Signed-off-by: owenhalpert <ohalpert@gmail.com>
owenhalpert · Mar 4, 2025 · aef03d0 · aef03d0
1 parent 82a50cb
commit aef03d0
Show file tree

Hide file tree

Showing 14 changed files with 164 additions and 111 deletions.
diff --git 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/...main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java b/...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
@@ -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<String, Object> getRemoteIndexingParameters(Map<String, Object> 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<String, Object> 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
@@ -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<String, Object> getRemoteIndexingParameters(Map<String, Object> indexInfoParameters) {
-        throw new UnsupportedOperationException("This method must be implemented by the implementing class");
+    default RemoteIndexParameters createRemoteIndexingParameters(Map<String, Object> 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
@@ -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<String, Object> getRemoteIndexingParameters(Map<String, Object> indexInfoParameters) {
-        Map<String, Object> 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<String, Object> algorithmParams = new HashMap<>();
-            Map<String, Object> innerMap = (Map<String, Object>) 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<String, Object> 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
@@ -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<TrainingConfigValidationInput, TrainingConfigValidationOutput
             return encoder.validateEncoderConfig(trainingConfigValidationInput);
         };
     }
+
+    // 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"
+     */
+    public static RemoteIndexParameters getRemoteIndexingParameters(Map<String, Object> 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<String, Object> innerMap = (Map<String, Object>) 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
@@ -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<String, Object> indexParameters;
+    protected RemoteIndexParameters indexParameters;
 
     /**
      * Constructor for RemoteBuildRequest.
@@ -80,8 +79,6 @@ public RemoteBuildRequest(
         KNNCodecUtil.initializeVectorValues(vectorValues);
         assert (vectorValues.dimension() > 0);
 
-        Map<String, Object> 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
@@ -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
@@ -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
@@ -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
@@ -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) {
-Original file line number
+Diff line change
@@ Expand Up / @@ -5,8 +5,6 @@ @@
     package org.opensearch.knn.index.remote;
-    import org.opensearch.knn.index.codec.nativeindex.remote.RemoteStatusResponse;
     import java.io.IOException;
     /**
@@ Expand Down @@