Skip to content

[Bug] Warm-up does not pick up .veb file for MOS.#3228

Draft
0ctopus13prime wants to merge 2 commits intoopensearch-project:mainfrom
0ctopus13prime:faiss-sq-warmup
Draft

[Bug] Warm-up does not pick up .veb file for MOS.#3228
0ctopus13prime wants to merge 2 commits intoopensearch-project:mainfrom
0ctopus13prime:faiss-sq-warmup

Conversation

@0ctopus13prime
Copy link
Copy Markdown
Collaborator

Description

Previously, warmup was handled externally by MemoryOptimizedSearchWarmup — it manually located engine files on disk, opened them via the Directory, and page-faulted bytes by seeking every 4096 bytes. Full-precision vectors were loaded separately by iterating FloatVectorValues through the LeafReader. The codec readers themselves had no warmup awareness; they threw a WarmupInitializationException when they received a null target vector during warmup, which the caller silently caught.

This PR pushes warmup responsibility down into the codec readers and the VectorSearcher implementations where the actual index structures live.

The key changes:

  1. Introduced WarmUpCollector, a sentinel KnnCollector singleton. Instead of passing null as the target vector and catching exceptions, warmup is now signaled cleanly by passing WarmUpCollector.INSTANCE through the existing search() API. Codec readers detect it via WarmUpCollector.isWarmUpRequest() and branch into warmup logic.

  2. Added warmUp() to the VectorSearcher interface. FaissMemoryOptimizedSearcher.warmUp() clones the IndexInput and reads every byte of the FAISS graph, then warms up flat vectors based on encoding type (float or byte). This replaces the old external byte-skipping approach with a complete sequential read that actually touches every page.

  3. Codec readers now own the warmup orchestration. NativeEngines990KnnVectorsReader delegates to VectorSearcher.warmUp() and additionally warms up .vec files for qframe fields via WarmupUtil.readAll(). Faiss1040ScalarQuantizedKnnVectorsReader delegates to VectorSearcher.warmUp() for the .faiss file, then explicitly iterates full-precision vectors to warm up .veb and .vec.

  4. MemoryOptimizedSearchWarmup was simplified — it no longer touches the Directory or reads files directly. It just identifies memory-optimized fields and calls search() with the WarmUpCollector sentinel, letting each codec reader handle the rest.

  5. Removed WarmupInitializationException and the null-target-vector hack from FaissMemoryOptimizedSearcher.

  6. Added WarmupUtil as a shared utility for reading IndexInput, FloatVectorValues, and ByteVectorValues into the page cache, with an optimization that reads directly from the backing IndexInput slice when the values implement HasIndexSlice.

Related Issues

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

Check List

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

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.

Signed-off-by: Dooyong Kim <kdooyong@amazon.com>
@codecov
Copy link
Copy Markdown

codecov bot commented Mar 28, 2026

Codecov Report

❌ Patch coverage is 93.93939% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 83.15%. Comparing base (683d0a3) to head (09c4020).

Files with missing lines Patch % Lines
...java/org/opensearch/knn/index/util/WarmupUtil.java 84.21% 0 Missing and 3 partials ⚠️
...yoptsearch/faiss/FaissMemoryOptimizedSearcher.java 85.71% 0 Missing and 1 partial ⚠️
Additional details and impacted files
@@             Coverage Diff              @@
##               main    #3228      +/-   ##
============================================
- Coverage     83.18%   83.15%   -0.04%     
- Complexity     4215     4232      +17     
============================================
  Files           453      455       +2     
  Lines         15455    15468      +13     
  Branches       1972     1976       +4     
============================================
+ Hits          12857    12863       +6     
- Misses         1805     1808       +3     
- Partials        793      797       +4     

☔ 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.

@shatejas
Copy link
Copy Markdown
Collaborator

Hi @0ctopus13prime, I took a look at this and I believe there is a better way to revamp warmup API.

While warmup can fundamentally be achieved with search the low level reality for this specific case is different, there is completely separate logic for warmup which calls for separation of concerns and not break the single-responsibility principle. The intent of moving warmup to reader seems spot on here. At the same time I want to propose a more extensible way to implement this and make it generic

  1. Introduce an interface for reader
interface WarmableReader {
   void warmUp();
}
  1. Implement it in all codecs
public void Native99EngineReader implements WarmableReader {

void warmUp() {
// NO-op if the field is not MOS
// warmup graph files, think through if you want to wamup flat vectors
}

}
  1. Check field if it has a vector value to warmup. The intent here is to simplify and push use case specific warmup logic to readers. This might help get rid of use case specific checks in warmup API. For instance you might be able to get away with checking if MOS is present or if there is quantization or not
ArrayList<FieldInfo> fields = new ArrayList<>();
        for (FieldInfo field : leafReader.getFieldInfos()) {
            if (field.hasVectorValues()) {
                fields.add(field);
            }
        }
        return fields;
  1. Modify warmUpField(FieldInfo field, SegmentReader segmentReader, Directory directory) to warmup per field readers
          KnnVectorsReader perFieldReader = segmentReader.getVectorReader();
            assert perFieldReader instanceof PerFieldKnnVectorsFormat.FieldsReader;
            KnnVectorsReader vectorsReader = ((PerFieldKnnVectorsFormat.FieldsReader) perFieldReader).getFieldReader(field.getName());
            if (vectorsReader instanceof WarmableReader warmableReader) {
                log.debug("Warming up reader for field: {}", field.getName());
                warmableReader.warmup();
            }

This might be missing some edge cases, but should cover most of the cases. Let me know your thoughts on this

@0ctopus13prime
Copy link
Copy Markdown
Collaborator Author

Hi @shatejas
Thank you for the suggestion, yeah, it looks good on me. I believe what you suggest will make code structure better.
Will update it shortly.

Signed-off-by: Dooyong Kim <kdooyong@amazon.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants