[GPU] Handle segments too big for MSAI segment access#141872
Merged
ldematte merged 10 commits intoelastic:mainfrom Feb 10, 2026
Merged
[GPU] Handle segments too big for MSAI segment access#141872ldematte merged 10 commits intoelastic:mainfrom
ldematte merged 10 commits intoelastic:mainfrom
Conversation
ldematte
added
>bug
auto-backport
Collaborator
|
Hi @ldematte, I've created a changelog YAML for you. |
elasticsearchmachine
added
the
Team:Search Relevance
Collaborator
|
Pinging @elastic/es-search-relevance (Team:Search Relevance) |
Contributor
Author
|
@mayya-sharipova @ChrisHegarty I added some unit tests, but I think it would be nice to have some IT tests as well. Do you think it's worth it? I'd need your help here, I don't know how I can add cases with a different max chunk size (I could do that manually maybe, but I don't think it's the right way). |
libs/gpu-codec/src/main/java/org/elasticsearch/gpu/codec/DatasetUtilsImpl.java
Outdated
Show resolved
Hide resolved
| /** Returns a Dataset over an input slice */ | ||
| CuVSMatrix fromSlice(MemorySegmentAccessInput input, long pos, long len, int numVectors, int dims, CuVSMatrix.DataType dataType) | ||
| throws IOException; | ||
| CuVSMatrix fromInput(MemorySegment input, int numVectors, int dims, CuVSMatrix.DataType dataType); |
Contributor
There was a problem hiding this comment.
May be not related to this PR, do we need DatasetUtils class at all? Looks like we only have a single DatasetUtilsImpl now.
libs/gpu-codec/src/main/java/org/elasticsearch/gpu/codec/MemorySegmentUtils.java
Outdated
Show resolved
Hide resolved
libs/gpu-codec/src/main/java/org/elasticsearch/gpu/codec/MemorySegmentUtils.java
Show resolved
Hide resolved
libs/gpu-codec/src/main/java/org/elasticsearch/gpu/codec/MemorySegmentUtils.java
Show resolved
Hide resolved
libs/gpu-codec/src/main/java/org/elasticsearch/gpu/codec/MemorySegmentUtils.java
Show resolved
Hide resolved
mayya-sharipova
approved these changes
Feb 9, 2026
Contributor
mayya-sharipova
left a comment
mayya-sharipova
left a comment
There was a problem hiding this comment.
@ldematte Thanks Lorenzo for fixing the bug.
I also like how you restructured the code.
I will think about the testing. I left some comments, and after addressing them, the PR is good to be merged.
Collaborator
💚 Backport successful
|
ldematte
added a commit
to ldematte/elasticsearch
that referenced
this pull request
Feb 10, 2026
MSAI (MemorySegmentAccessInput) provides access to an index input via file memory mapping, returning a MemorySegment. However, MemorySegmentAccessInput may return null under some circumstances; in our usage, where we try to get a memory mapping for the whole file, it would fail if the file is too big (where "too big" is currently set at 16GB). This PR adds a fallback for this situation; if we fail to get a MemorySegment from MemorySegmentAccessInput, we assume the input as a whole is too big, and we fall back to using a temp file. Data is copied over to the temp file, and we use the FileChannel Java API directly to map the whole file. Using a temp file is necessary; if the input as a whole is too big (> 16G) we cannot reliably use directly on or off-heap memory, we need "something" to back that memory in case the OS needs to page it out (e.g. in the likely event we exceed the physical memory available). Closes elastic#141746
ldematte
added a commit
that referenced
this pull request
Feb 10, 2026
) MSAI (MemorySegmentAccessInput) provides access to an index input via file memory mapping, returning a MemorySegment. However, MemorySegmentAccessInput may return null under some circumstances; in our usage, where we try to get a memory mapping for the whole file, it would fail if the file is too big (where "too big" is currently set at 16GB). This PR adds a fallback for this situation; if we fail to get a MemorySegment from MemorySegmentAccessInput, we assume the input as a whole is too big, and we fall back to using a temp file. Data is copied over to the temp file, and we use the FileChannel Java API directly to map the whole file. Using a temp file is necessary; if the input as a whole is too big (> 16G) we cannot reliably use directly on or off-heap memory, we need "something" to back that memory in case the OS needs to page it out (e.g. in the likely event we exceed the physical memory available). Closes #141746
elasticsearchmachine
pushed a commit
that referenced
this pull request
Mar 10, 2026
MemorySegmentUtils directly cast the Directory to FSDirectory, but Elasticsearch wraps directories in Store$StoreDirectory (which extends FilterDirectory, not FSDirectory). When vector data exceeds MMapDirectory's max chunk size, the fallback path hit this cast and threw a ClassCastException, failing all shard merges. ClassCastException: Store$StoreDirectory cannot be cast to FSDirectory (Store$StoreDirectory is in module org.elasticsearch.server; FSDirectory is in module org.apache.lucene.core) Use FilterDirectory.unwrap() to peel through wrapper layers before casting. Also fix log message formatting for segment size values. Relates to #141872
mayya-sharipova
added a commit
to mayya-sharipova/elasticsearch
that referenced
this pull request
Mar 10, 2026
…3531) MemorySegmentUtils directly cast the Directory to FSDirectory, but Elasticsearch wraps directories in Store$StoreDirectory (which extends FilterDirectory, not FSDirectory). When vector data exceeds MMapDirectory's max chunk size, the fallback path hit this cast and threw a ClassCastException, failing all shard merges. ClassCastException: Store$StoreDirectory cannot be cast to FSDirectory (Store$StoreDirectory is in module org.elasticsearch.server; FSDirectory is in module org.apache.lucene.core) Use FilterDirectory.unwrap() to peel through wrapper layers before casting. Also fix log message formatting for segment size values. Relates to elastic#141872
elasticsearchmachine
pushed a commit
that referenced
this pull request
Mar 10, 2026
) (#143924) * Fix GPU merge ClassCastException with wrapped directories (#143531) MemorySegmentUtils directly cast the Directory to FSDirectory, but Elasticsearch wraps directories in Store$StoreDirectory (which extends FilterDirectory, not FSDirectory). When vector data exceeds MMapDirectory's max chunk size, the fallback path hit this cast and threw a ClassCastException, failing all shard merges. ClassCastException: Store$StoreDirectory cannot be cast to FSDirectory (Store$StoreDirectory is in module org.elasticsearch.server; FSDirectory is in module org.apache.lucene.core) Use FilterDirectory.unwrap() to peel through wrapper layers before casting. Also fix log message formatting for segment size values. Relates to #141872 * Fix compilation error: update CuVSGPUSupport to GPUSupport in GPUMergeFallbackIT
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
MSAI (
MemorySegmentAccessInput) provides access to an index input via file memory mapping, returning a MemorySegment.However,
MemorySegmentAccessInputmay return null under some circumstances; in our usage, where we try to get a memory mapping for the whole file, it would fail if the file is too big (where "too big" is currently set at 16GB).This PR adds a fallback for this situation; if we fail to get a MemorySegment from
MemorySegmentAccessInput, we assume the input as a whole is too big, and we fall back to using a temp file. Data is copied over to the temp file, and we use the FileChannel Java API directly to map the whole file.Using a temp file is necessary; if the input as a whole is too big (> 16G) we cannot reliably use directly on or off-heap memory, we need "something" to back that memory in case the OS needs to page it out (e.g. in the likely event we exceed the physical memory available).
Closes #141746