[Remote Vector Index Build] Introduce RemoteIndexClient skeleton and Build Request construction (#2560)

Add RemoteIndexClient initial implementation, its accompanying dependencies, and Build Request, Retry Strategy, and test files

Signed-off-by: owenhalpert <ohalpert@gmail.com>

Author: owenhalpert

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * [Remote Vector Index Build] Introduce Remote Native Index Build feature flag, settings, and initial skeleton [#2525](https://github.com/opensearch-project/k-NN/pull/2525)
 * [Remote Vector Index Build] Implement vector data upload and vector data size threshold setting [#2550](https://github.com/opensearch-project/k-NN/pull/2550)
 * [Remote Vector Index Build] Implement data download and IndexOutput write functionality [#2554](https://github.com/opensearch-project/k-NN/pull/2554)
+* [Remote Vector Index Build] Introduce Client Skeleton + basic Build Request implementation [#2560](https://github.com/opensearch-project/k-NN/pull/2560)
 * Add concurrency optimizations with native memory graph loading and force eviction (#2265) [https://github.com/opensearch-project/k-NN/pull/2345]
 ### Enhancements
 * Introduce node level circuit breakers for k-NN [#2509](https://github.com/opensearch-project/k-NN/pull/2509)

build.gradle

@@ -321,7 +321,9 @@ dependencies {
     api "net.java.dev.jna:jna-platform:5.13.0"
     // OpenSearch core is using slf4j 1.7.36. Therefore, we cannot change the version here.
     implementation 'org.slf4j:slf4j-api:1.7.36'
-
+    api "org.apache.httpcomponents.client5:httpclient5:${versions.httpclient5}"
+    api "org.apache.httpcomponents.core5:httpcore5:${versions.httpcore5}"
+    api "org.apache.httpcomponents.core5:httpcore5-h2:${versions.httpcore5}"
     zipArchive group: 'org.opensearch.plugin', name:'opensearch-security', version: "${opensearch_build}"
 }
 

src/main/java/org/opensearch/knn/index/KNNSettings.java

@@ -15,10 +15,12 @@
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.Booleans;
+import org.opensearch.common.settings.SecureSetting;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.common.settings.SecureString;
 import org.opensearch.core.common.unit.ByteSizeUnit;
 import org.opensearch.core.common.unit.ByteSizeValue;
 import org.opensearch.index.IndexModule;
@@ -100,6 +102,11 @@ public class KNNSettings {
     public static final String KNN_INDEX_REMOTE_VECTOR_BUILD = "index.knn.remote_index_build.enabled";
     public static final String KNN_REMOTE_VECTOR_REPO = "knn.remote_index_build.vector_repo";
     public static final String KNN_INDEX_REMOTE_VECTOR_BUILD_THRESHOLD = "index.knn.remote_index_build.size_threshold";
+    public static final String KNN_REMOTE_BUILD_SERVICE_ENDPOINT = "knn.remote_index_build.client.endpoint";
+    public static final String KNN_REMOTE_BUILD_CLIENT_POLL_INTERVAL = "knn.remote_index_build.client.poll_interval";
+    public static final String KNN_REMOTE_BUILD_CLIENT_TIMEOUT = "knn.remote_index_build.client.timeout";
+    public static final String KNN_REMOTE_BUILD_CLIENT_USERNAME = "knn.remote_index_build.client.username";
+    public static final String KNN_REMOTE_BUILD_CLIENT_PASSWORD = "knn.remote_index_build.client.password";
 
     /**
      * Default setting values
@@ -133,6 +140,10 @@ public class KNNSettings {
     // TODO: Tune this default value based on benchmarking
     public static final ByteSizeValue KNN_INDEX_REMOTE_VECTOR_BUILD_THRESHOLD_DEFAULT_VALUE = new ByteSizeValue(50, ByteSizeUnit.MB);
 
+    // TODO: Tune these default values based on benchmarking
+    public static final Integer KNN_DEFAULT_REMOTE_BUILD_CLIENT_TIMEOUT_MINUTES = 60;
+    public static final Integer KNN_DEFAULT_REMOTE_BUILD_CLIENT_POLL_INTERVAL_SECONDS = 30;
+
     /**
      * Settings Definition
      */
@@ -409,6 +420,47 @@ public class KNNSettings {
         Dynamic,
         IndexScope
     );
+    /**
+     * Remote build service endpoint to be used for remote index build.
+     */
+    public static final Setting<String> KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING = Setting.simpleString(
+        KNN_REMOTE_BUILD_SERVICE_ENDPOINT,
+        NodeScope,
+        Dynamic
+    );
+
+    /**
+     * Time the remote build service client will wait before falling back to CPU index build.
+     */
+    public static final Setting<TimeValue> KNN_REMOTE_BUILD_CLIENT_TIMEOUT_SETTING = Setting.timeSetting(
+        KNN_REMOTE_BUILD_CLIENT_TIMEOUT,
+        TimeValue.timeValueMinutes(KNN_DEFAULT_REMOTE_BUILD_CLIENT_TIMEOUT_MINUTES),
+        NodeScope,
+        Dynamic
+    );
+
+    /**
+     * Setting to control how often the remote build service client polls the build service for the status of the job.
+     */
+    public static final Setting<TimeValue> KNN_REMOTE_BUILD_CLIENT_POLL_INTERVAL_SETTING = Setting.timeSetting(
+        KNN_REMOTE_BUILD_CLIENT_POLL_INTERVAL,
+        TimeValue.timeValueSeconds(KNN_DEFAULT_REMOTE_BUILD_CLIENT_POLL_INTERVAL_SECONDS),
+        NodeScope,
+        Dynamic
+    );
+
+    /**
+     * Keystore settings for build service HTTP authorization
+     */
+    public static final Setting<SecureString> KNN_REMOTE_BUILD_CLIENT_USERNAME_SETTING = SecureSetting.secureString(
+        KNN_REMOTE_BUILD_CLIENT_USERNAME,
+        null
+    );
+    public static final Setting<SecureString> KNN_REMOTE_BUILD_CLIENT_PASSWORD_SETTING = SecureSetting.secureString(
+        KNN_REMOTE_BUILD_CLIENT_PASSWORD,
+        null
+    );
+
     /**
      * Dynamic settings
      */
@@ -600,6 +652,26 @@ private Setting<?> getSetting(String key) {
             return KNN_INDEX_REMOTE_VECTOR_BUILD_THRESHOLD_SETTING;
         }
 
+        if (KNN_REMOTE_BUILD_SERVICE_ENDPOINT.equals(key)) {
+            return KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING;
+        }
+
+        if (KNN_REMOTE_BUILD_CLIENT_TIMEOUT.equals(key)) {
+            return KNN_REMOTE_BUILD_CLIENT_TIMEOUT_SETTING;
+        }
+
+        if (KNN_REMOTE_BUILD_CLIENT_POLL_INTERVAL.equals(key)) {
+            return KNN_REMOTE_BUILD_CLIENT_POLL_INTERVAL_SETTING;
+        }
+
+        if (KNN_REMOTE_BUILD_CLIENT_USERNAME.equals(key)) {
+            return KNN_REMOTE_BUILD_CLIENT_USERNAME_SETTING;
+        }
+
+        if (KNN_REMOTE_BUILD_CLIENT_PASSWORD.equals(key)) {
+            return KNN_REMOTE_BUILD_CLIENT_PASSWORD_SETTING;
+        }
+
         throw new IllegalArgumentException("Cannot find setting by key [" + key + "]");
     }
 
@@ -628,7 +700,12 @@ public List<Setting<?>> getSettings() {
             KNN_DERIVED_SOURCE_ENABLED_SETTING,
             KNN_INDEX_REMOTE_VECTOR_BUILD_SETTING,
             KNN_REMOTE_VECTOR_REPO_SETTING,
-            KNN_INDEX_REMOTE_VECTOR_BUILD_THRESHOLD_SETTING
+            KNN_INDEX_REMOTE_VECTOR_BUILD_THRESHOLD_SETTING,
+            KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING,
+            KNN_REMOTE_BUILD_CLIENT_TIMEOUT_SETTING,
+            KNN_REMOTE_BUILD_CLIENT_POLL_INTERVAL_SETTING,
+            KNN_REMOTE_BUILD_CLIENT_USERNAME_SETTING,
+            KNN_REMOTE_BUILD_CLIENT_PASSWORD_SETTING
         );
         return Stream.concat(settings.stream(), Stream.concat(getFeatureFlags().stream(), dynamicCacheSettings.values().stream()))
             .collect(Collectors.toList());
@@ -748,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);

src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/DefaultVectorRepositoryAccessor.java

@@ -33,10 +33,10 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
-import static org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy.DOC_ID_FILE_EXTENSION;
-import static org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy.VECTORS_PATH;
-import static org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy.VECTOR_BLOB_FILE_EXTENSION;
 import static org.opensearch.knn.index.codec.util.KNNCodecUtil.initializeVectorValues;
+import static org.opensearch.knn.index.remote.KNNRemoteConstants.DOC_ID_FILE_EXTENSION;
+import static org.opensearch.knn.index.remote.KNNRemoteConstants.VECTORS_PATH;
+import static org.opensearch.knn.index.remote.KNNRemoteConstants.VECTOR_BLOB_FILE_EXTENSION;
 
 @Log4j2
 @AllArgsConstructor

src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java

@@ -6,14 +6,18 @@
 package org.opensearch.knn.index.codec.nativeindex.remote;
 
 import lombok.extern.log4j.Log4j2;
-import org.apache.commons.lang.NotImplementedException;
 import org.opensearch.common.StopWatch;
 import org.opensearch.common.UUIDs;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.IndexSettings;
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategy;
 import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
+import org.opensearch.knn.index.remote.RemoteBuildRequest;
+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;
@@ -38,10 +42,6 @@ public class RemoteIndexBuildStrategy implements NativeIndexBuildStrategy {
     private final NativeIndexBuildStrategy fallbackStrategy;
     private final IndexSettings indexSettings;
 
-    static final String VECTOR_BLOB_FILE_EXTENSION = ".knnvec";
-    static final String DOC_ID_FILE_EXTENSION = ".knndid";
-    static final String VECTORS_PATH = "_vectors";
-
     /**
      * Public constructor, intended to be called by {@link org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategyFactory} based in
      * part on the return value from {@link RemoteIndexBuildStrategy#shouldBuildIndexRemotely}
@@ -125,18 +125,20 @@ public void buildAndWriteIndex(BuildIndexParams indexInfo) throws IOException {
             time_in_millis = stopWatch.stop().totalTime().millis();
             log.debug("Repository write took {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName());
 
+            RemoteIndexClient client = RemoteIndexClientFactory.getRemoteIndexClient();
+            RemoteBuildRequest request = new RemoteBuildRequest(indexSettings, indexInfo, getRepository().getMetadata(), blobName);
             stopWatch = new StopWatch().start();
-            submitVectorBuild();
+            RemoteBuildResponse remoteBuildResponse = client.submitVectorBuild(request);
             time_in_millis = stopWatch.stop().totalTime().millis();
             log.debug("Submit vector build took {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName());
 
             stopWatch = new StopWatch().start();
-            String downloadPath = awaitVectorBuild();
+            RemoteStatusResponse remoteStatusResponse = client.awaitVectorBuild(remoteBuildResponse);
             time_in_millis = stopWatch.stop().totalTime().millis();
             log.debug("Await vector build took {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName());
 
             stopWatch = new StopWatch().start();
-            vectorRepositoryAccessor.readFromRepository(downloadPath, indexInfo.getIndexOutputWithBuffer());
+            vectorRepositoryAccessor.readFromRepository(remoteStatusResponse.getIndexPath(), indexInfo.getIndexOutputWithBuffer());
             time_in_millis = stopWatch.stop().totalTime().millis();
             log.debug("Repository read took {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName());
         } catch (Exception e) {
@@ -163,20 +165,4 @@ private BlobStoreRepository getRepository() throws RepositoryMissingException {
         assert repository instanceof BlobStoreRepository : "Repository should be instance of BlobStoreRepository";
         return (BlobStoreRepository) repository;
     }
-
-    /**
-     * Submit vector build request to remote vector build service
-     *
-     */
-    private void submitVectorBuild() {
-        throw new NotImplementedException();
-    }
-
-    /**
-     * Wait on remote vector build to complete
-     * @return String     The path from which we should perform download, delimited by "/"
-     */
-    private String awaitVectorBuild() throws NotImplementedException {
-        throw new NotImplementedException();
-    }
 }

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;
@@ -216,4 +217,9 @@ public ResolvedMethodContext resolveMethod(
     public boolean supportsRemoteIndexBuild() {
         return knnLibrary.supportsRemoteIndexBuild();
     }
+
+    @Override
+    public RemoteIndexParameters createRemoteIndexingParameters(Map<String, Object> indexInfoParameters) {
+        return knnLibrary.createRemoteIndexingParameters(indexInfoParameters);
+    }
 }

src/main/java/org/opensearch/knn/index/engine/KNNLibrary.java

@@ -7,9 +7,11 @@
 
 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;
+import java.util.Map;
 
 /**
  * KNNLibrary is an interface that helps the plugin communicate with k-NN libraries
@@ -147,4 +149,8 @@ default List<String> mmapFileExtensions() {
     default boolean supportsRemoteIndexBuild() {
         return false;
     }
+
+    default RemoteIndexParameters createRemoteIndexingParameters(Map<String, Object> indexInfoParameters) {
+        throw new UnsupportedOperationException("Remote build service does not support this engine");
+    }
 }

src/main/java/org/opensearch/knn/index/engine/faiss/Faiss.java

@@ -14,12 +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.Map;
 import java.util.function.Function;
 
 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.NAME;
 
 /**
  * Implements NativeLibrary for the faiss native library
@@ -123,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.createRemoteIndexingParameters(indexInfoParameters);
+        }
+        throw new IllegalArgumentException("Unsupported method for remote indexing");
+    }
 }

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,39 @@ protected Function<TrainingConfigValidationInput, TrainingConfigValidationOutput
             return encoder.validateEncoderConfig(trainingConfigValidationInput);
         };
     }
+
+    /**
+     * 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 createRemoteIndexingParameters(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();
+    }
+
+    private 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);
+        return Integer.parseInt(hnswPart.substring(4));
+    }
 }