Implement Optimized embedding generation in text embedding processor

[will-hwang]

Description

This PR merges the previously approved PRs (#1191 , #1217 ) from feature branch onto main.

See below for Benchmark results

Setup

Version: OS 3.0 alpha
Cluster: 4 r6g.2xlarge nodes (1 coordinator node 3 data nodes)

Benchmark Results

Dataset: Trec-covid (https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/trec-covid.zip)
Ingest Pipeline

{
	"description": "An NLP ingest pipeline",
	"processors": [
		{
			"text_embedding": {
				"model_id": "-l5Dr5UB17VvHPspuB6C",
				"field_map": {
					"passage_text": "passage_embedding"
				},
				"skip_existing": true/false
			}
		}
	]
}

Index Mapping

{
  "settings": {
    "index.knn": true,
    "number_of_shards": 3,
    "default_pipeline": "nlp-ingest-pipeline"
  },
  "mappings": {
    "properties": {
      "id": {
        "type": "text"
      },
      "passage_embedding": {
        "type": "knn_vector",
        "dimension": 768,
        "method": {
          "engine": "lucene",
          "space_type": "l2",
          "name": "hnsw",
          "parameters": {}
        }
      },
      "text": {
          "type": "text"
      }
    }
  }
}

Ingest Latency

The following table presents latency measurements of initial ingest operation (in milliseconds) with skip_existing feature enabled and disabled. The Percent difference columns show the relative performance impact between the two.

Operation	Doc Size	Batch Size	skip_existing_off latency (ms)	skip_existing_on latency (ms)	Percent Difference
Single Ingest	1000	1	459896.33	463579.12	0.80%
Single Ingest	2000	1	932563.3	938318.45	0.62%
Single Ingest	3000	1	1400710.43	1401216.8	0.04%
Batch Ingest	171332	200	2247191.38	2192883.54	-2.42%
Batch Ingest	171332	500	2065514.4	2011408.73	-2.62%

Update Latency

The following table presents the latency measurements of update operation after identical ingest operation (in milliseconds) with skip_existing feature enabled and disabled. The Percent difference columns show the relative performance impact between the two.

Operation	Doc Size	Batch Size	skip_existing_off latency (ms)	skip_existing_on latency (ms)	Percent Difference
Single Update	1000	1	459896.33	100957.91	-78.05%
Single Update	2000	1	932563.3	201715.16	-78.37%
Single Update	3000	1	1400710.43	292020.02	-79.15%
Batch Update	171332	200	2247191.38	352767.2	-84.30%
Batch Update	171332	500	2065514.4	293164.93	-85.81%

Related Issues

Resolves #[Issue number to be closed when this PR is merged]

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.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 0.00%. Comparing base (fdd4c8c) to head (3dc9d7c).
Report is 1 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #1238       +/-   ##
============================================
- Coverage     81.90%       0   -81.91%     
============================================
  Files           208       0      -208     
  Lines          9594       0     -9594     
  Branches       1632       0     -1632     
============================================
- Hits           7858       0     -7858     
+ Misses         1105       0     -1105     
+ Partials        631       0      -631     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[junqiu-lei]

Thanks for sharing the benchmark results, LGTM.

[heemin32]

Could you share the index mapping for knn field?

[heemin32]

Can we test with shard number 3?

{
  "settings": {
    "index.knn": true,
    "number_of_shards": 3,
    "default_pipeline": "nlp-ingest-pipeline"
  },
  "mappings": {
    "properties": {
      "id": {
        "type": "text"
      },
      "passage_embedding": {
        "type": "knn_vector",
        "dimension": 768,
        "method": {
          "engine": "lucene",
          "space_type": "l2",
          "name": "hnsw",
          "parameters": {}
        }
      },
      "text": {
          "type": "text"
      }
    }
  }
}

[bzhangam]

nit: Should we simply return ingestDocumentWrappers rather than construct a new empty list?

[will-hwang]

sure, either should work

[bzhangam]

I think here we do updateWithExceptions(ingestDocumentWrappers, e) in multiple catch blocks. But we only consider the case that IngestDocumentWrapper might already run into exception and not sent for inference here which does not make sense. It's possible the other catch block can override the exception of IngestDocumentWrapper incorrectly. e.g. The IngestDocumentWrapper can have exception A and not sent for inference but now we can potentially override its exception related to inference.

I think we should either propagate the exception to one catch block and handle it with the consideration that IngestDocumentWrapper can already run into the exception. Or we should modify updateWithExceptions function to account for that.

[will-hwang]

this was added to ensure that if any of the lines while handling the exceptions in for loop throws exceptions, it will be handled properly. I understand that this may not be needed, but was brought up as a possible risk in previous pr

[heemin32]

When would this situation occur? The IngestDocumentWrapper might have already encountered an exception and was not sent for inference. It would be helpful if you could point out the code where the exception is set before inference.

I agree that there's an inconsistency in how we handle this—checking if an exception is already set only in this specific place but not elsewhere.

[will-hwang]

see link

[will-hwang]

maybe we can propagate the exception to top level and set the exceptions there?

[heemin32]

I don't think we can propagate the exception. One option is to remove this check and override the exception for all documents. When an exception occurs here, retaining the original exception may not make a significant difference.

[will-hwang]

another option is to keep the check for all exception cases, which will only store the first exception encountered for all documents.

[will-hwang]

this way, we'll know which ingest document failed with which exact exception

[heemin32]

that would work as well.

[bzhangam]

I wonder do we really check if the doc already run into the exception and skip the embedding generation for it? If that's the case it makes sense to keep the actual exception fail the ingest of the doc which can be helpful for debugging.

[bzhangam]

Wonder why we want to clearly try catch the exception here? If we don't catch it here won't it be auto caught by the exception handler?

[will-hwang]

exception handler is catching the exception thrown only in doBatchExecute method here. we do not catch exceptions if any of the lines between 221-236 throws exceptions.

[bzhangam]

Why we need this try catch? If we can throw the exception when we update the exception in ingestDocumentWrappers then in the catch block we will throw exception again and will not invoke handler.accept. If we can throw the exception when we do handler.accept shouldn't the exception handler of the handler to handle it rather than handle it here?

[heemin32]

Good catch. handler.accept() should go to finally block.

[bzhangam]

Why we need this catch block? If we are worry about the exception when we do sortByLengthAndReturnOriginalOrder can we simply try catch it? It's not easy to understand each try catch block if we nest them a lot.

[will-hwang]

i'm okay with simplifying the try/catch, or even removing if we know for sure the underlying methods will not throw exceptions. But this was brought up as a possible concern in the previous PR, so decided to add them

[heemin32]

I think @bzhangam is talking about

try {
  sortByLengthAndReturnOriginalOrder()
} catch {
}

where we only wrap try-catch around sortByLengthAndReturnOriginalOrder. but not removing it.

I think having try-catch wrapping the entire method might be simpler and safer.
If the code is hard to read, we can try to introduce a method to return lambda.

protected void doSubBatchExecute(
        List<IngestDocumentWrapper> ingestDocumentWrappers,
        List<String> inferenceList,
        List<DataForInference> dataForInferences,
        Consumer<List<IngestDocumentWrapper>> handler
    ) {
   ...
   doBatchExecute(inferenceList, batchExecuteHandler());
   ...
}


private Consumer<InferenceResult> batchExecuteHandler() {
 ...
}

[bzhangam]

Yeah. I'm more thinking can we make the try catch cleaner. I think when we do try catch we want to do some special logic to handle the exception. But here we do the same logic with multiple nested try catch which is kind complicated the code.

[heemin32]

Agree. I think keep passing handler to downstream method is not a good practice as that will make us to keep adding all those try-catch for every method. In general, we should pass handler only when we make async call.

@bzhangam If you have more specific recommendation on how we should refactor the code here, please provide one.

[will-hwang]

i've refactored the code to separate out the try catch to a separate method as heemin suggested. Do you have other suggestions on how to make try/catch cleaner?

[bzhangam]

By using wrap if we have an exception in the onResponse we will auto invoke the onFailure to handle the exception which is helpful in our use case since what we want to do in the try catch is simply the same thing we want to do in the onFailure.

    static <Response> ActionListener<Response> wrap(final CheckedConsumer<Response, ? extends Exception> onResponse, final Consumer<Exception> onFailure) {
        return new ActionListener<Response>() {
            public void onResponse(Response response) {
                try {
                    onResponse.accept(response);
                } catch (Exception e) {
                    this.onFailure(e);
                }

            }

            public void onFailure(Exception e) {
                onFailure.accept(e);
            }
        };
    }

With this I think we can remove some try catch in the onResponse part.

[bzhangam]

Do we need to get all doc? What if some docs don't have a value for inference?

[will-hwang]

we need to get all docs first to check if docs have or don't have value for inference

[bzhangam]

If the doc in the request doesn't have a value for the neural field which means we should not generate embedding for it since there is no value there. In that I think we can skip pulling the existing doc since we know there is no neural field we need to inference.

[bzhangam]

Same as above do we really want to update the exception for all docs? Do we really need to have multiple try catch here? It seems like we are doing the same thing in the catch block. Can we pop up the exception and simply handle it at the top layer?

[will-hwang]

ActionListener.wrap(response -> {}) spawns a new thread, and needs to be handled separately

[bzhangam]

Wonder why we need to sort the inference list here? Seems like we sort it and then restore it without doing anything special.

[will-hwang]

sort is done before inference call is made, and results is the retrieved embedding from the inference call. I think it's safe to leave as is for now

[bzhangam]

I understand the sort is done before the inference call but I don't know why we sort it. Because we simply send the sorted data to the inference call without any logic e.g. group or batch the data. If that's the case this seems like a redundant logic.

[bzhangam]

I think we can create an issue to track this no need to address it in this PR. Probably it's better to check with the author to see why we add this logic here.

[will-hwang]

Can we test with shard number 3?

{
  "settings": {
    "index.knn": true,
    "number_of_shards": 3,
    "default_pipeline": "nlp-ingest-pipeline"
  },
  "mappings": {
    "properties": {
      "id": {
        "type": "text"
      },
      "passage_embedding": {
        "type": "knn_vector",
        "dimension": 768,
        "method": {
          "engine": "lucene",
          "space_type": "l2",
          "name": "hnsw",
          "parameters": {}
        }
      },
      "text": {
          "type": "text"
      }
    }
  }
}

@heemin32 updated in description

[bzhangam]

I think we should not set it for all since some docs may not have the data for inference and we are not getting the existing doc for them.

[bzhangam]

If we don't plan to do something special here I think we can rely on the caller to handle the exception.

[will-hwang]

removed to handle it at caller

[bzhangam]

I think this is not needed since the exception handler will catch this exception. They are wrapped by an ActionListener wrap in doBatchExecute.

[will-hwang]

removed

[bzhangam]

I think we should log the error here to help debug since we cannot add the exception to the doc in this case.

[will-hwang]

added

[heemin32]

LGTM. Thanks.