viirya opened a new pull request, #6813: URL: https://github.com/apache/datafusion-comet/pull/6813
## Which issue does this PR close? Closes #6812. ## Rationale for this change Spark's `ExecutionMemoryPool` registers a task's `memoryForTask` entry before its wait loop and removes it when the task's balance reaches zero. A request that waits below the task's minimum share then wakes up to `NoSuchElementException: key not found: <taskAttemptId>` ([SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444)). #6403 guarded the JVM shuffle allocator, but Comet's native memory pools go through `CometTaskMemoryManager.acquireMemory`, which had no such guard. A native `try_grow` waiting in Spark failed the task when another native plan or a JVM consumer of the same task released the task's last bytes in the meantime. The Spark fix, [apache/spark#58747](https://github.com/apache/spark/pull/58747), is still open against master. Even once it merges, the Spark versions Comet supports will not have it. On a Spark that has the fix, the exception never occurs and the retry is never taken. ## What changes are included in this PR? - `MissingTaskEntryRetry` holds the retry and the check for Spark's missing-entry exception that #6403 added to `CometUnifiedShuffleMemoryAllocator`. The shuffle allocator now uses it and behaves as before: it throws `SparkOutOfMemoryError` once the retries run out. - `CometTaskMemoryManager.acquireMemory` retries through the same helper. When the retries run out it returns a zero grant instead of throwing, so `try_grow` reports `ResourcesExhausted` and the operator spills, and `grow` carries the request as overcommit. The failed call was granted nothing, so `used` only counts what a successful attempt acquired. - `memory_management.md` now says the native pools are guarded too. - The test fixtures from `CometUnifiedShuffleMemoryAllocatorSuite` move to a shared `TaskMemoryTestUtils` trait. ## How are these changes tested? New tests in `CometTaskMemoryManagerSuite`: - Two tests reproduce the race against a real `UnifiedMemoryManager`. A native acquire waits in Spark while the task's last bytes are released, once by another native plan and once by a JVM consumer. Both failed with `NoSuchElementException: key not found: 0` before this change and are granted in full after it. - Using a task memory manager that fails on demand: a request that keeps losing the entry is refused with 0 after three attempts, one that loses it twice is granted on the third attempt, and any other `NoSuchElementException` is rethrown. `CometUnifiedShuffleMemoryAllocatorSuite` passes unchanged on the shared helper. This pull request and its description were written by Isaac. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
