[Remote Index Client] Implement remote client build awaiting functionality, validate encoder support

[owenhalpert]

Description

This PR is a followup to #2560. This PR adds a RemoteIndexWaiter interface to await the completion of the remote vector build. For the HTTP client, this is implemented as a simple polling mechanism in RemoteIndexPoller. While there is a fair amount of files changed due to refactoring, the crux of this PR is in the RemoteIndexPoller, RemoteIndexHTTPClient, and RemoteIndexBuildStrategy classes.

This PR also adds a check on the encoder type, enforcing that only HNSWFlat is supported by the remote build service.

This PR does not include:

• Metric collection

which will go in the next PR.

Related Issues

#2560 PR 1
#2518 LLD
#2391 Meta issue

Check List

• New functionality includes testing.
  - [ ] New functionality has been documented.
  - [ ] API changes companion pull request created.
• Commits are signed per the DCO using --signoff.
  - [ ] Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[jed326]

Thanks @owenhalpert, the overall approach looks good to me, just needs some cleanup in a few places.

[kotwanikunal]

On a high level - looking at the changes - I am not sure if we should add in all the client/HTTP constructs directly into the plugin. The job status constants, client configuration seems like a better fit for a client package.
This adds in a direct dependency between the build service and the kNN plugin - which I feel can be abstracted out.

[kotwanikunal]

assserts do not work in prod. Can we do null checks instead?

[jed326]

All the logic of getting the blob container comes from OpenSearch core so the assertion is just a sanity check on that

[navneet1v]

@jed326 if assertion fails will we fallback to CPU based index builds? and will there be a useful exception which can tell what happened? If no, please add these information

[jed326]

My point is there is not any way for blobContainer to be null, and none of the methods in the calling path to retrieve blobContainer throw any checked exceptions. The only thing that can throw an exception is getRepository, and creating a blobContainer on top of a repository reference is more or less just setting the file path to be written/read to. This assertion just helps make sure in tests if we do try to test/mock this method we are properly configuring things.

[navneet1v]

On a high level - looking at the changes - I am not sure if we should add in all the client/HTTP constructs directly into the plugin. The job status constants, client configuration seems like a better fit for a client package. This adds in a direct dependency between the build service and the kNN plugin - which I feel can be abstracted out.

@kotwanikunal , since remote index build service is not an hosted service like S3 which vends out a HTTP client. Also, we don't want to maintain another repo as the maintainence overhead will be high. Hence the client code is better suited to be in k-NN repo only. If you see right now, the code is abstracted in a different java package. Now if you think we should move the java package to a separate module then that is something which can be easily achieved in upcoming iterations. Let us know your thoughts.

[jed326]

I think it was mentioned before, but can we track in a GitHub issue somewhere to improve this parsing? I think that's a generic problem in k-NN and not specific to remote index build too.

[navneet1v]

+1

[owenhalpert]

Will get this out tomorrow

[owenhalpert]

#2586

[navneet1v]

@jed326 if assertion fails will we fallback to CPU based index builds? and will there be a useful exception which can tell what happened? If no, please add these information

[navneet1v]

+1

[navneet1v]

should this validation be part of fromXContent() ? I feel we should move this validation out of this function and put it at the place where we are validating the status and other attributes.

[owenhalpert]

Moving the actual content checks to the poller (while the status response object will just be checking the validity of the JSON).

[navneet1v]

when I mentioned that Polling function should not be part of client I mean the Poller class should call the awaitVectorBuild api from outside. What is being done here is exactly reversed. This is the reason why I still think the api name should be getIndexBuildStatus rather than awaitVectorBuild. Because polling is one way to wait for the IndexBuild to complete.

I think RemoteIndexPoller -> getIndexBuildStatus() rather than awaitVectorBuild() -> RemoteIndexPoller()

[owenhalpert]

Got it. Thanks. The new flow will be that the build strategy gets the Waiter implementation from a factory method. In this case it's a Poller. The Poller uses client.getBuildStatus and does the polling all on its own.

The client is specifically responsible for getting and sending requests — the Waiter will use the client to do so when needed.

[jed326]

LGTM, thanks @owenhalpert

[navneet1v]

Lets avoid IOException in method signature, it is just making the whole interfaces ugly. Lets catch the exception in the functions and then see if we want throw a proper runtime exception or we want to return some defaults.

[owenhalpert]

Good point. I think we can catch and return null for the parsing methods on IOExceptions so the supportsRemoteIndexBuild just returns false.

[navneet1v]

@jmazanec15 and @naveentatikonda can you please help in validating this logic to ensure that it is full proof, with all the different combinations that can be possible for the encoders

[jmazanec15]

Hmmm thinking about this a bit more. Do we only need the encoder on write? If so, I think itd be safest to retrieve it from the KnnVectorFieldType via mapper service (https://github.com/opensearch-project/k-NN/blob/main/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java). We should expose it as a method in https://github.com/opensearch-project/k-NN/blob/main/src/main/java/org/opensearch/knn/index/engine/KNNLibraryIndexingContext.java which will be available via the mapped field type.

If we need it on read, then we have to parse it out of this. But, from my understanding, we shouldnt need it on read.

[navneet1v]

Suggested change

	public RemoteBuildStatusRequest(RemoteBuildResponse remoteBuildResponse) {
	public static RemoteBuildStatusRequest build(RemoteBuildResponse remoteBuildResponse) {

[Vikasht34]

Why Status Build Request work based on RemoteBuildResponse rather than it has to work with minimum JobId which is enough for get status.?

[owenhalpert]

I think this keeps the design more extensible if the request/response were to contain something in addition to or instead of the job ID (e.g. tenantId).

[Vikasht34]

My Question is why constructor is taking RemoteBuildResponse , and then derive JOBID from here ? this is not flexibility , this is tight coupling with un related Response from some otheer API ...

[Vikasht34]

Why RemoteBuildStatusRequest is tied to RemoteBuildResponse , both are totally unrelated !! , either it should be tied to JOBID , TanentId and what ever it comes...

[owenhalpert]

That's a good point. I can make this change:

// in NativeIndexBuildStrategy
String jobId = remoteBuildResponse.getJobId();
RemoteBuildStatusRequest status = RemoteBuildStatusRequest.builder().jobId(jobId).build();

Now, instead of the two being related, the build strategy will process the remoteBuildResponse however it needs (in this case, just getting the job ID) and then separately build a status request with all necessary parameters (in this case, just the job ID).

In the future, we can draw more from the response, and easily add more parameters to the builder. Does that look better?

cc @navneet1v

[Vikasht34]

This is perfect !!

[navneet1v]

@owenhalpert please check why CIs are failing. Also, I think code is in pretty good shape and close to be getting shipped.

[Vikasht34]

Where we have defined Seervice API contract and how clinet can assume these parsing!!

[owenhalpert]

https://github.com/opensearch-project/remote-vector-index-builder/blob/main/API.md

[Vikasht34]

This is perfect !!

[Vikasht34]

I rewrote the entier class to make polling more intellizent with Jitter , distributed safe with exponential backoff but at the same time it respects user setting values . We can adustst MIN_POLL_INTERVAL_MS and max ones ..Let me know how it looks @navneet1v

import java.io.IOException;
import java.time.Duration;
import java.util.Objects;
import java.util.Random;

class RemoteIndexPoller implements RemoteIndexWaiter {

    private static final int INITIAL_DELAY_FACTOR = 3;
    private static final int MIN_POLL_INTERVAL_MS = 500;
    private static final int MAX_POLL_INTERVAL_MS = 30_000;
    private static final int STATUS_CHANGE_RESET = 2; // Reset interval if status changes within this threshold
    private static final double JITTER_LOWER = 0.8;
    private static final double JITTER_UPPER = 1.2;

    private final RemoteIndexClient client;
    private final Random random = new Random();

    RemoteIndexPoller(RemoteIndexClient client) {
        this.client = Objects.requireNonNull(client, "RemoteIndexClient cannot be null");
    }

    @SuppressWarnings("BusyWait")
    public RemoteBuildStatusResponse awaitVectorBuild(RemoteBuildStatusRequest request) throws InterruptedException, IOException {
        long startTime = System.nanoTime();
        long timeout = KNNSettings.getRemoteBuildClientTimeout().getNanos();
        long basePollInterval = KNNSettings.getRemoteBuildClientPollInterval().getMillis();
        long pollInterval = basePollInterval * INITIAL_DELAY_FACTOR;

        Thread.sleep(applyJitter(pollInterval)); // Initial delay

        int consecutiveSameStatus = 0;
        String lastStatus = "";

        for (int attempt = 1; System.nanoTime() - startTime < timeout; attempt++) {
            RemoteBuildStatusResponse response = client.getBuildStatus(request);
            String taskStatus = Objects.requireNonNullElse(response.getTaskStatus(), "").trim();

            logger.debug("Polling attempt {}: Task status = {}", attempt, taskStatus);

            switch (taskStatus) {
                case COMPLETED_INDEX_BUILD -> {
                    String fileName = Objects.requireNonNullElse(response.getFileName(), "").trim();
                    if (fileName.isEmpty()) {
                        throw new IOException("Invalid response: Missing file name for completed status.");
                    }
                    logger.info("Remote index build completed successfully in {} seconds.",
                        Duration.ofNanos(System.nanoTime() - startTime).toSeconds());
                    return response;
                }
                case FAILED_INDEX_BUILD -> throw new InterruptedException(
                    "Remote index build failed after " + Duration.ofNanos(System.nanoTime() - startTime).toMinutes() +
                        " minutes. Error: " + response.getErrorMessage());

                case RUNNING_INDEX_BUILD -> {
                    // Adjust polling intelligently
                    if (taskStatus.equals(lastStatus)) {
                        consecutiveSameStatus++;
                        if (consecutiveSameStatus >= STATUS_CHANGE_RESET) {
                            pollInterval = Math.min(pollInterval * 2, MAX_POLL_INTERVAL_MS); // Increase polling delay
                        }
                    } else {
                        consecutiveSameStatus = 0; // Reset when status changes
                        pollInterval = Math.max(pollInterval / 2, MIN_POLL_INTERVAL_MS); // Reduce polling delay
                    }
                    lastStatus = taskStatus;

                    // Prevent polling too close to timeout
                    long remainingTime = timeout - (System.nanoTime() - startTime);
                    if (remainingTime < pollInterval * 2) {
                        pollInterval = MIN_POLL_INTERVAL_MS;
                    }

                    long sleepTime = applyJitter(pollInterval);
                    logger.debug("Task still running. Waiting {} ms before next poll", sleepTime);
                    Thread.sleep(sleepTime);
                }
                default -> throw new IOException("Invalid task status: " + taskStatus);
            }
        }

        throw new InterruptedException("Remote index build timed out after " +
            Duration.ofNanos(System.nanoTime() - startTime).toMinutes() + " minutes.");
    }

    private long applyJitter(long baseInterval) {
        return (long) (baseInterval * (JITTER_LOWER + (random.nextDouble() * (JITTER_UPPER - JITTER_LOWER))));
    }
}

[jed326]

Thanks @Vikasht34. For the jitter aspect I can agree that would be generically good and I think would be fine to pull into this.

For exponential backoff I don't think we should use that logic for the poller for now and I think this is something we need to revisit based on benchmarking. We want to strike a balance between calling the remote index build service too often and waiting too long after the graph build has completed to get the status. For example, let's take a 1 minute poll interval. After polling 2 times (so we've waited 2 minutes), the build job is still just as likely to finish in 2min + 1ms as 2min + 2min, and I think it's better to over-index against the second case.

Based on benchmarking we should be able to establish p50/90/99 values on how long the remote index build takes for a given segment size and then we can set the polling interval as some function of that value. If the p50/p90/p99 is very narrow then exponential backoff probably wouldn't be necessary, but if it's a very wide range we can revisit this.