sunchao commented on code in PR #5493:
URL: https://github.com/apache/datafusion-comet/pull/5493#discussion_r3926724407
##########
spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java:
##########
@@ -123,6 +147,122 @@ public synchronized MemoryBlock allocate(long required) {
return allocateMemoryBlock(size);
}
+ /**
+ * Like {@link #allocate(long)}, but waits for other tasks of this shared
pool to free memory,
+ * mirroring how Spark's unified memory manager blocks a task until memory
becomes available.
+ * Callers must first spill buffered data they can cheaply release; memory
this thread still
+ * retains (e.g. the sorter's pointer array or sibling writers' pages) is
included in the liveness
+ * checks below. The wait fails fast when it can never succeed: when the
request does not fit next
+ * to the requester's retained memory, or when all allocated memory is
retained by blocked threads
+ * and none of their requests fits in the free pool. Because the holders it
depends on may in turn
+ * be blocked on resources outside this pool that only a task waiting here
can release, the wait
+ * is also bounded by `spark.comet.shuffle.jvm.memoryWaitTimeout`, after
which the managed
+ * allocation error is thrown and Spark's task retry can recover. Task
cancellation or Java
+ * interruption aborts the wait.
+ */
+ @Override
+ public synchronized MemoryBlock allocateBlocking(long required) {
+ long memoryWaitTimeoutMs =
+ (long)
CometConf$.MODULE$.COMET_SHUFFLE_JVM_MEMORY_WAIT_TIMEOUT().get();
Review Comment:
Could this lookup move into the first actual wait? `allocateBlocking` can
succeed immediately after the caller spills, so that path does not need to
resolve the timeout.
Across two JVM forks, the current successful allocation/free path measured
329–336 ns and 432 B of Java heap per pair. The prototype with a capacity check
and deferred lookup measured 110–126 ns and 112 B. The capacity-check-only
control retained roughly the current cost, which points to the eager
configuration lookup as the avoidable work on this path. These are component
measurements and the absolute cost is small per page.
Read the setting once per waiting request, from that task's configuration,
so this preserves the per-task timeout fix already made in this PR.
##########
spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java:
##########
@@ -123,6 +147,122 @@ public synchronized MemoryBlock allocate(long required) {
return allocateMemoryBlock(size);
}
+ /**
+ * Like {@link #allocate(long)}, but waits for other tasks of this shared
pool to free memory,
+ * mirroring how Spark's unified memory manager blocks a task until memory
becomes available.
+ * Callers must first spill buffered data they can cheaply release; memory
this thread still
+ * retains (e.g. the sorter's pointer array or sibling writers' pages) is
included in the liveness
+ * checks below. The wait fails fast when it can never succeed: when the
request does not fit next
+ * to the requester's retained memory, or when all allocated memory is
retained by blocked threads
+ * and none of their requests fits in the free pool. Because the holders it
depends on may in turn
+ * be blocked on resources outside this pool that only a task waiting here
can release, the wait
+ * is also bounded by `spark.comet.shuffle.jvm.memoryWaitTimeout`, after
which the managed
+ * allocation error is thrown and Spark's task retry can recover. Task
cancellation or Java
+ * interruption aborts the wait.
+ */
+ @Override
+ public synchronized MemoryBlock allocateBlocking(long required) {
+ long memoryWaitTimeoutMs =
+ (long)
CometConf$.MODULE$.COMET_SHUFFLE_JVM_MEMORY_WAIT_TIMEOUT().get();
+ long size = Math.max(pageSize, required);
+ Thread self = Thread.currentThread();
+ TaskContext taskContext = TaskContext.get();
+ long waitStart = 0;
+ long lastLog = 0;
+ try {
+ while (true) {
+ if (taskContext != null) {
+ taskContext.killTaskIfInterrupted();
+ }
+ try {
+ return allocateMemoryBlock(size);
+ } catch (SparkOutOfMemoryError e) {
Review Comment:
Could the retry path use a private allocation helper that reports
pool-capacity shortage without constructing an exception? Each small free can
wake every waiter even when their requests still do not fit, so these retries
construct and format a new `SparkOutOfMemoryError`, including its stack trace,
under the shared allocator monitor.
In the controlled allocator experiment described in the review body, a
capacity-check prototype reduced waiter Java heap allocation over 1.2 seconds
from 29.5 MiB to 0.22 MiB with 8 waiters, and from 54.2 MiB to 0.49 MiB with 32
waiters.
A private nonthrowing helper could serve both APIs: immediate allocation
converts a capacity miss to the managed error, while blocking allocation
creates that error only when it must fail. Keep the existing ownership,
liveness, cancellation, timeout and notification behavior, and preserve
oversized-request, page-table and native-allocation errors.
--
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]
