Add metric collection to the remote index build.

Signed-off-by: owenhalpert <[email protected]>

Author: owenhalpert

CHANGELOG.md

@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Features
 * [Remote Vector Index Build] Client polling mechanism, encoder check, method parameter retrieval [#2576](https://github.com/opensearch-project/k-NN/pull/2576)
 * [Remote Vector Index Build] Move client to separate module [#2603](https://github.com/opensearch-project/k-NN/pull/2603)
+* [Remote Vector Index Build] Add metric collection for remote build process [#2615](https://github.com/opensearch-project/k-NN/pull/2615)
 Add filter function to KNNQueryBuilder with unit tests and integration tests [#2599](https://github.com/opensearch-project/k-NN/pull/2599)
 ### Enhancements
 ### Bug Fixes

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

@@ -16,7 +16,6 @@
 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.codec.util.KNNCodecUtil;
 import org.opensearch.knn.index.engine.KNNMethodContext;
 import org.opensearch.knn.index.remote.RemoteIndexWaiter;
 import org.opensearch.knn.index.remote.RemoteIndexWaiterFactory;
@@ -43,6 +42,21 @@
 import static org.opensearch.knn.index.KNNSettings.KNN_INDEX_REMOTE_VECTOR_BUILD_SETTING;
 import static org.opensearch.knn.index.KNNSettings.KNN_INDEX_REMOTE_VECTOR_BUILD_THRESHOLD_SETTING;
 import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_VECTOR_REPO_SETTING;
+import static org.opensearch.knn.index.codec.util.KNNCodecUtil.initializeVectorValues;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.BUILD_REQUEST_FAILURE_COUNT;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.BUILD_REQUEST_SUCCESS_COUNT;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.INDEX_BUILD_FAILURE_COUNT;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.INDEX_BUILD_SUCCESS_COUNT;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.READ_FAILURE_COUNT;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.READ_SUCCESS_COUNT;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.READ_TIME;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.REMOTE_INDEX_BUILD_CURRENT_OPERATIONS;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.REMOTE_INDEX_BUILD_CURRENT_SIZE;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.REMOTE_INDEX_BUILD_TIME;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.WAITING_TIME;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.WRITE_FAILURE_COUNT;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.WRITE_SUCCESS_COUNT;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.WRITE_TIME;
 
 /**
  * This class orchestrates building vector indices. It handles uploading data to a repository, submitting a remote
@@ -120,63 +134,154 @@ public static boolean shouldBuildIndexRemotely(IndexSettings indexSettings, long
      *      2. Triggers index build
      *      3. Awaits on vector build to complete
      *      4. Downloads index file and writes to indexOutput
-     *
-     * @param indexInfo
-     * @throws IOException
      */
     @Override
     public void buildAndWriteIndex(BuildIndexParams indexInfo) throws IOException {
-        StopWatch stopWatch;
+        StopWatch remoteBuildStopWatch = new StopWatch();
+        KNNVectorValues<?> knnVectorValues = indexInfo.getKnnVectorValuesSupplier().get();
+        initializeVectorValues(knnVectorValues);
+        startRemoteIndexBuildStats((long) indexInfo.getTotalLiveDocs() * knnVectorValues.bytesPerVector(), remoteBuildStopWatch);
+
+        try {
+            RepositoryContext repositoryContext = getRepositoryContext(indexInfo);
+
+            // 1. Write required data to repository
+            writeToRepository(repositoryContext, indexInfo);
+
+            // 2. Trigger remote index build
+            RemoteIndexClient client = RemoteIndexClientFactory.getRemoteIndexClient(KNNSettings.getRemoteBuildServiceEndpoint());
+            RemoteBuildResponse remoteBuildResponse = submitBuild(repositoryContext, indexInfo, client);
+
+            // 3. Await vector build completion
+            RemoteBuildStatusResponse remoteBuildStatusResponse = awaitIndexBuild(remoteBuildResponse, indexInfo, client);
+
+            // 4. Download index file and write to indexOutput
+            readFromRepository(indexInfo, repositoryContext, remoteBuildStatusResponse);
+
+            endRemoteIndexBuildStats((long) indexInfo.getTotalLiveDocs() * knnVectorValues.bytesPerVector(), remoteBuildStopWatch);
+            INDEX_BUILD_SUCCESS_COUNT.increment();
+        } catch (Exception e) {
+            log.warn("Failed to build index remotely", e);
+            handleFailure(indexInfo, knnVectorValues.bytesPerVector(), remoteBuildStopWatch);
+        }
+    }
+
+    /**
+     * Writes the required vector and doc ID data to the repository
+     */
+    private void writeToRepository(RepositoryContext repositoryContext, BuildIndexParams indexInfo) throws IOException,
+        InterruptedException {
+        BlobStoreRepository repository = repositoryContext.blobStoreRepository;
+        BlobPath blobPath = repositoryContext.blobPath;
+        VectorRepositoryAccessor vectorRepositoryAccessor = new DefaultVectorRepositoryAccessor(
+            repository.blobStore().blobContainer(blobPath)
+        );
+
+        StopWatch stopWatch = new StopWatch().start();
         long time_in_millis;
         try {
-            BlobStoreRepository repository = getRepository();
-            BlobPath blobPath = repository.basePath().add(indexSettings.getUUID() + VECTORS_PATH);
-            VectorRepositoryAccessor vectorRepositoryAccessor = new DefaultVectorRepositoryAccessor(
-                repository.blobStore().blobContainer(blobPath)
-            );
-            stopWatch = new StopWatch().start();
-            // We create a new time based UUID per file in order to avoid conflicts across shards. It is also very difficult to get the
-            // shard id in this context.
-            String blobName = UUIDs.base64UUID() + "_" + indexInfo.getFieldName() + "_" + indexInfo.getSegmentWriteState().segmentInfo.name;
             vectorRepositoryAccessor.writeToRepository(
-                blobName,
+                repositoryContext.blobName,
                 indexInfo.getTotalLiveDocs(),
                 indexInfo.getVectorDataType(),
                 indexInfo.getKnnVectorValuesSupplier()
             );
             time_in_millis = stopWatch.stop().totalTime().millis();
+            WRITE_SUCCESS_COUNT.increment();
+            WRITE_TIME.incrementBy(time_in_millis);
             log.debug("Repository write took {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName());
+        } catch (Exception e) {
+            time_in_millis = stopWatch.stop().totalTime().millis();
+            WRITE_FAILURE_COUNT.increment();
+            log.error("Repository write failed after {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName(), e);
+            throw e;
+        }
+    }
 
-            final RemoteIndexClient client = RemoteIndexClientFactory.getRemoteIndexClient(KNNSettings.getRemoteBuildServiceEndpoint());
+    /**
+     * Submits a remote build request to the remote index build service
+     * @return RemoteBuildResponse containing the response from the remote service
+     */
+    private RemoteBuildResponse submitBuild(RepositoryContext repositoryContext, BuildIndexParams indexInfo, RemoteIndexClient client)
+        throws IOException {
+        final RemoteBuildResponse remoteBuildResponse;
+        StopWatch stopWatch = new StopWatch().start();
+        long time_in_millis;
+        try {
             final RemoteBuildRequest buildRequest = buildRemoteBuildRequest(
                 indexSettings,
                 indexInfo,
-                repository.getMetadata(),
-                blobPath.buildAsString() + blobName,
+                repositoryContext.blobStoreRepository.getMetadata(),
+                repositoryContext.blobPath.buildAsString() + repositoryContext.blobName,
                 knnMethodContext
             );
-            stopWatch = new StopWatch().start();
-            final RemoteBuildResponse remoteBuildResponse = client.submitVectorBuild(buildRequest);
+            remoteBuildResponse = client.submitVectorBuild(buildRequest);
+            BUILD_REQUEST_SUCCESS_COUNT.increment();
             time_in_millis = stopWatch.stop().totalTime().millis();
             log.debug("Submit vector build took {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName());
+            return remoteBuildResponse;
+        } catch (IOException e) {
+            BUILD_REQUEST_FAILURE_COUNT.increment();
+            time_in_millis = stopWatch.stop().totalTime().millis();
+            log.error("Submit vector build failed after {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName(), e);
+            throw e;
+        }
+    }
 
+    /**
+     * Awaits the vector build to complete
+     * @return RemoteBuildStatusResponse containing the completed status response from the remote service.
+     * This will only be returned with a COMPLETED_INDEX_BUILD status, otherwise the method will throw an exception.
+     */
+    private RemoteBuildStatusResponse awaitIndexBuild(
+        RemoteBuildResponse remoteBuildResponse,
+        BuildIndexParams indexInfo,
+        RemoteIndexClient client
+    ) throws IOException, InterruptedException {
+        RemoteBuildStatusResponse remoteBuildStatusResponse;
+        StopWatch stopWatch = new StopWatch().start();
+        long time_in_millis;
+        try {
             final RemoteBuildStatusRequest remoteBuildStatusRequest = RemoteBuildStatusRequest.builder()
                 .jobId(remoteBuildResponse.getJobId())
                 .build();
             RemoteIndexWaiter waiter = RemoteIndexWaiterFactory.getRemoteIndexWaiter(client);
-            stopWatch = new StopWatch().start();
-            RemoteBuildStatusResponse remoteBuildStatusResponse = waiter.awaitVectorBuild(remoteBuildStatusRequest);
+            remoteBuildStatusResponse = waiter.awaitVectorBuild(remoteBuildStatusRequest);
             time_in_millis = stopWatch.stop().totalTime().millis();
+            WAITING_TIME.incrementBy(time_in_millis);
             log.debug("Await vector build took {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName());
+            return remoteBuildStatusResponse;
+        } catch (InterruptedException | IOException e) {
+            time_in_millis = stopWatch.stop().totalTime().millis();
+            log.error("Await vector build failed after {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName(), e);
+            throw e;
+        }
+    }
 
-            stopWatch = new StopWatch().start();
-            vectorRepositoryAccessor.readFromRepository(remoteBuildStatusResponse.getFileName(), indexInfo.getIndexOutputWithBuffer());
+    /**
+     * Downloads the index file from the repository and writes to the indexOutput
+     */
+    private void readFromRepository(
+        BuildIndexParams indexInfo,
+        RepositoryContext repositoryContext,
+        RemoteBuildStatusResponse remoteBuildStatusResponse
+    ) throws IOException {
+        StopWatch stopWatch = new StopWatch().start();
+        long time_in_millis;
+        try {
+            repositoryContext.vectorRepositoryAccessor.readFromRepository(
+                remoteBuildStatusResponse.getFileName(),
+                indexInfo.getIndexOutputWithBuffer()
+            );
             time_in_millis = stopWatch.stop().totalTime().millis();
+            READ_SUCCESS_COUNT.increment();
+            READ_TIME.incrementBy(time_in_millis);
             log.debug("Repository read took {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName());
         } catch (Exception e) {
-            // TODO: This needs more robust failure handling
-            log.warn("Failed to build index remotely", e);
-            fallbackStrategy.buildAndWriteIndex(indexInfo);
+            time_in_millis = stopWatch.stop().totalTime().millis();
+            READ_FAILURE_COUNT.increment();
+            log.error("Repository read failed after {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName(), e);
+            throw e;
         }
     }
 
@@ -223,7 +328,7 @@ static RemoteBuildRequest buildRemoteBuildRequest(
         String vectorDataType = indexInfo.getVectorDataType().getValue();
 
         KNNVectorValues<?> vectorValues = indexInfo.getKnnVectorValuesSupplier().get();
-        KNNCodecUtil.initializeVectorValues(vectorValues);
+        initializeVectorValues(vectorValues);
         assert (vectorValues.dimension() > 0);
 
         return RemoteBuildRequest.builder()
@@ -240,4 +345,51 @@ static RemoteBuildRequest buildRemoteBuildRequest(
             .build();
     }
 
+    /**
+     * Helper method to collect remote index build metrics on start
+     */
+    private void startRemoteIndexBuildStats(long size, StopWatch stopWatch) {
+        stopWatch.start();
+        REMOTE_INDEX_BUILD_CURRENT_OPERATIONS.increment();
+        REMOTE_INDEX_BUILD_CURRENT_SIZE.incrementBy(size);
+    }
+
+    /**
+     * Helper method to collect remote index build metrics on success
+     */
+    private void endRemoteIndexBuildStats(long size, StopWatch stopWatch) {
+        long time_in_millis = stopWatch.stop().totalTime().millis();
+        REMOTE_INDEX_BUILD_CURRENT_OPERATIONS.decrement();
+        REMOTE_INDEX_BUILD_CURRENT_SIZE.decrementBy(size);
+        REMOTE_INDEX_BUILD_TIME.incrementBy(time_in_millis);
+    }
+
+    /**
+     * Helper method to collect remote index build metrics on failure and invoke fallback strategy
+     */
+    private void handleFailure(BuildIndexParams indexParams, long bytesPerVector, StopWatch stopWatch) throws IOException {
+        endRemoteIndexBuildStats(indexParams.getTotalLiveDocs() * bytesPerVector, stopWatch);
+        INDEX_BUILD_FAILURE_COUNT.increment();
+        fallbackStrategy.buildAndWriteIndex(indexParams);
+    }
+
+    /**
+     * Record to hold various repository related objects
+     */
+    private record RepositoryContext(BlobStoreRepository blobStoreRepository, BlobPath blobPath,
+        VectorRepositoryAccessor vectorRepositoryAccessor, String blobName) {
+    }
+
+    /**
+     * Helper method to get repository context
+     */
+    private RepositoryContext getRepositoryContext(BuildIndexParams indexInfo) {
+        BlobStoreRepository repository = getRepository();
+        BlobPath blobPath = repository.basePath().add(indexSettings.getUUID() + VECTORS_PATH);
+        String blobName = UUIDs.base64UUID() + "_" + indexInfo.getFieldName() + "_" + indexInfo.getSegmentWriteState().segmentInfo.name;
+        VectorRepositoryAccessor vectorRepositoryAccessor = new DefaultVectorRepositoryAccessor(
+            repository.blobStore().blobContainer(blobPath)
+        );
+        return new RepositoryContext(repository, blobPath, vectorRepositoryAccessor, blobName);
+    }
 }

src/main/java/org/opensearch/knn/index/remote/RemoteIndexPoller.java

@@ -17,6 +17,8 @@
 
 import static org.opensearch.knn.index.KNNSettings.getRemoteBuildClientPollInterval;
 import static org.opensearch.knn.index.KNNSettings.getRemoteBuildClientTimeout;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.STATUS_REQUEST_FAILURE_COUNT;
+import static org.opensearch.knn.plugin.stats.KNNRemoteIndexBuildValue.STATUS_REQUEST_SUCCESS_COUNT;
 
 /**
  * Implementation of a {@link RemoteIndexWaiter} that awaits the vector build by polling.
@@ -63,7 +65,14 @@ public RemoteBuildStatusResponse awaitVectorBuild(RemoteBuildStatusRequest remot
         sleepWithJitter(pollInterval * INITIAL_DELAY_FACTOR);
 
         while (System.nanoTime() - startTime < timeout) {
-            RemoteBuildStatusResponse remoteBuildStatusResponse = client.getBuildStatus(remoteBuildStatusRequest);
+            RemoteBuildStatusResponse remoteBuildStatusResponse;
+            try {
+                remoteBuildStatusResponse = client.getBuildStatus(remoteBuildStatusRequest);
+            } catch (IOException e) {
+                STATUS_REQUEST_FAILURE_COUNT.increment();
+                throw e;
+            }
+            STATUS_REQUEST_SUCCESS_COUNT.increment();
             String taskStatus = remoteBuildStatusResponse.getTaskStatus();
             if (StringUtils.isBlank(taskStatus)) {
                 throw new IOException(String.format("Invalid response format, missing %s", TASK_STATUS));

src/main/java/org/opensearch/knn/plugin/stats/KNNRemoteIndexBuildValue.java

@@ -0,0 +1,89 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.plugin.stats;
+
+import lombok.Getter;
+
+import java.util.concurrent.atomic.LongAdder;
+
+public enum KNNRemoteIndexBuildValue {
+
+    // Repository Accumulating Stats
+    WRITE_SUCCESS_COUNT("write_success_count"),
+    WRITE_FAILURE_COUNT("write_failure_count"),
+    WRITE_TIME("successful_write_time_in_millis"),
+    READ_SUCCESS_COUNT("read_success_count"),
+    READ_FAILURE_COUNT("read_failure_count"),
+    READ_TIME("successful_read_time_in_millis"),
+
+    // Remote Index Build Stats
+    REMOTE_INDEX_BUILD_CURRENT_OPERATIONS("remote_index_build_current_operations"),
+    REMOTE_INDEX_BUILD_CURRENT_SIZE("remote_index_build_current_size"),
+    REMOTE_INDEX_BUILD_TIME("remote_index_build_time_in_millis"),
+
+    // Client Stats
+    BUILD_REQUEST_SUCCESS_COUNT("build_request_success_count"),
+    BUILD_REQUEST_FAILURE_COUNT("build_request_failure_count"),
+    STATUS_REQUEST_SUCCESS_COUNT("status_request_success_count"),
+    STATUS_REQUEST_FAILURE_COUNT("status_request_failure_count"),
+    INDEX_BUILD_SUCCESS_COUNT("index_build_success_count"),
+    INDEX_BUILD_FAILURE_COUNT("index_build_failure_count"),
+    WAITING_TIME("waiting_time_in_ms");
+
+    @Getter
+    private final String name;
+    private final LongAdder value;
+
+    /**
+     * Constructor
+     *
+     * @param name name of the value
+     */
+    KNNRemoteIndexBuildValue(String name) {
+        this.name = name;
+        this.value = new LongAdder();
+    }
+
+    /**
+     * Get the value
+     * @return value
+     */
+    public Long getValue() {
+        return value.longValue();
+    }
+
+    /**
+     * Increment the value
+     */
+    public void increment() {
+        value.increment();
+    }
+
+    /**
+     * Decrement the value
+     */
+    public void decrement() {
+        value.decrement();
+    }
+
+    /**
+     * Increment the value by a specified amount
+     *
+     * @param delta The amount to increment
+     */
+    public void incrementBy(long delta) {
+        value.add(delta);
+    }
+
+    /**
+     * Decrement the value by a specified amount
+     *
+     * @param delta The amount to decrement
+     */
+    public void decrementBy(long delta) {
+        value.add(delta * -1);
+    }
+}