[Remote Vector Index Build] Introduce RemoteIndexClient skeleton and Build Request construction

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)
 ### Enhancements
 * Introduce node level circuit breakers for k-NN [#2509](https://github.com/opensearch-project/k-NN/pull/2509)
 ### Bug Fixes

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/common/KNNConstants.java

@@ -165,4 +165,25 @@ public class KNNConstants {
     public static final String DERIVED_VECTOR_FIELD_ATTRIBUTE_KEY = "knn-derived-source-enabled";
     public static final String DERIVED_VECTOR_FIELD_ATTRIBUTE_TRUE_VALUE = "true";
     public static final String DERIVED_VECTOR_FIELD_ATTRIBUTE_FALSE_VALUE = "false";
+
+    // Remote build constants
+    public static final String BUILD_ENDPOINT = "/_build";
+    public static final String STATUS_ENDPOINT = "/_status";
+    public static final String S3 = "s3";
+    public static final String BUCKET = "bucket";
+    // Build request keys
+    public static final String ALGORITHM = "algorithm";
+    public static final String ALGORITHM_PARAMETERS = "algorithm_parameters";
+    public static final String INDEX_PARAMETERS = "index_parameters";
+    public static final String DOC_COUNT = "doc_count";
+    public static final String TENANT_ID = "tenant_id";
+    public static final String DOC_ID_PATH = "doc_id_path";
+    public static final String VECTOR_PATH = "vector_path";
+    public static final String CONTAINER_NAME = "container_name";
+    public static final String REPOSITORY_TYPE = "repository_type";
+    // Server responses
+    public static final String JOB_ID = "job_id";
+    public static final String TASK_STATUS = "task_status";
+    public static final String INDEX_PATH = "index_path";
+    public static final String ERROR_MESSAGE = "error_message";
 }

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());

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

@@ -6,14 +6,19 @@
 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.HTTPRemoteBuildRequest;
+import org.opensearch.knn.index.remote.RemoteBuildRequest;
+import org.opensearch.knn.index.remote.RemoteBuildRequestBuilder;
+import org.opensearch.knn.index.remote.RemoteBuildResponse;
+import org.opensearch.knn.index.remote.RemoteIndexClient;
+import org.opensearch.knn.index.remote.RemoteIndexClientFactory;
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.repositories.Repository;
 import org.opensearch.repositories.RepositoryMissingException;
@@ -38,8 +43,8 @@ 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";
+    public static final String VECTOR_BLOB_FILE_EXTENSION = ".knnvec";
+    public static final String DOC_ID_FILE_EXTENSION = ".knndid";
     static final String VECTORS_PATH = "_vectors";
 
     /**
@@ -125,18 +130,26 @@ 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());
 
+            // TODO future implementations will set the following two params depending on some setting to denote the protocol
+            RemoteIndexClient client = RemoteIndexClientFactory.getRemoteIndexClient(RemoteIndexClientFactory.TYPE_HTTP);
+            RemoteBuildRequest request = RemoteBuildRequestBuilder.builder(HTTPRemoteBuildRequest.class)
+                .indexSettings(indexSettings)
+                .indexInfo(indexInfo)
+                .repositoryMetadata(getRepository().getMetadata())
+                .blobName(blobName)
+                .build();
             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 +176,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/codec/nativeindex/remote/RemoteStatusResponse.java

@@ -0,0 +1,15 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.codec.nativeindex.remote;
+
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+
+@Getter
+@AllArgsConstructor
+public class RemoteStatusResponse {
+    private String indexPath;
+}

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

@@ -177,6 +177,11 @@ 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);

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

@@ -10,6 +10,7 @@
 
 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 +148,12 @@ default List<String> mmapFileExtensions() {
     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");
+    }
 }

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

@@ -15,11 +15,29 @@
 import org.opensearch.knn.index.engine.NativeLibrary;
 import org.opensearch.knn.index.engine.ResolvedMethodContext;
 
+import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Function;
 
+import static org.opensearch.knn.common.KNNConstants.ALGORITHM;
+import static org.opensearch.knn.common.KNNConstants.ALGORITHM_PARAMETERS;
+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.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
@@ -109,6 +127,67 @@ 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,