peterxcli commented on code in PR #5493:
URL: https://github.com/apache/datafusion-comet/pull/5493#discussion_r3910377417
##########
spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java:
##########
@@ -123,6 +143,114 @@ 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 only use this after spilling their own buffered data. The
wait fails fast when it
+ * can never succeed: when the request does not fit next to the memory this
thread itself still
+ * retains (e.g. the sorter's pointer array), or when all allocated memory
is retained by threads
+ * that are themselves blocked here 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. Interrupting the task (e.g. task
kill) 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();
+ long waitStart = 0;
+ long lastLog = 0;
+ try {
+ while (true) {
+ try {
+ return allocateMemoryBlock(size);
+ } catch (SparkOutOfMemoryError e) {
+ if (waitingThreads.put(self, size) == null) {
+ // Wake existing waiters so they re-evaluate the deadlock check
against the enlarged
+ // waiting set.
+ notifyAll();
+ }
+ // This thread cannot free what it retains while it waits, so a
request that does not
+ // fit next to its own retained memory can never be satisfied.
+ if (size > totalMemory - retainedMemory.getOrDefault(self, 0L)) {
+ throw e;
+ }
+ // The allocation just failed, so the request does not fit in the
unallocated pool.
+ // Waiting can only succeed while some thread can still free memory:
either a thread
+ // outside the waiting set retains pool memory, or another waiter's
request fits in the
+ // free pool, in which case that waiter can proceed and eventually
free what it retains.
+ if (allocatedMemory <= retainedByWaitingThreads() &&
!anyWaiterCanProceed()) {
+ throw e;
+ }
+ // The holders this wait depends on may themselves be blocked on
resources outside
+ // this pool (Spark execution memory, locks, I/O) that only a task
waiting here can
+ // release - a cycle this allocator cannot observe. Bound the wait
so such cycles
+ // unwind with the managed allocation error instead of hanging the
executor; Spark's
+ // task retry can then recover.
+ long now = System.currentTimeMillis();
+ if (waitStart == 0) {
+ waitStart = now;
+ lastLog = now;
+ logger.warn(
+ "Waiting for other tasks to free up {} bytes of Comet shuffle
pool memory", size);
+ } else if (now - waitStart >= memoryWaitTimeoutMs) {
+ logger.warn(
+ "Giving up after waiting {} ms for {} bytes of Comet shuffle
pool memory "
+ + "(see {})",
+ now - waitStart,
+ size,
+
CometConf$.MODULE$.COMET_SHUFFLE_JVM_MEMORY_WAIT_TIMEOUT().key());
+ throw e;
+ } else if (now - lastLog >= WAIT_LOG_INTERVAL_MS) {
+ lastLog = now;
+ logger.warn(
+ "Still waiting ({} ms so far) for {} bytes of Comet shuffle
pool memory; "
+ + "{} bytes free, {} thread(s) waiting",
+ now - waitStart,
+ size,
+ totalMemory - allocatedMemory,
+ waitingThreads.size());
+ }
+ try {
+ wait(
+ Math.max(
+ 1L, Math.min(WAIT_LOG_INTERVAL_MS, memoryWaitTimeoutMs -
(now - waitStart))));
Review Comment:
Addressed in 59b7cf2b2: allocateBlocking now checks Spark TaskContext
cancellation on each retry, polls at most once per second, and notifies
remaining waiters when it exits. The regression marks TaskContext interrupted
without interrupting the Java thread and requires prompt TaskKilledException;
CometDiskBlockWriterSuite passes 9/9.
##########
spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java:
##########
@@ -75,6 +98,8 @@ public final class CometBoundedShuffleMemoryAllocator extends
CometShuffleMemory
this.pageSize = pageSize;
this.totalMemory =
CometSparkSessionExtensions$.MODULE$.getCometShuffleMemorySize(conf,
SQLConf.get());
+ this.memoryWaitTimeoutMs =
+ (long)
CometConf$.MODULE$.COMET_SHUFFLE_JVM_MEMORY_WAIT_TIMEOUT().get();
Review Comment:
Addressed in acb71a1bd: allocateBlocking now resolves
spark.comet.shuffle.jvm.memoryWaitTimeout for each call. The shared-pool
regression initializes the singleton under a 100 ms task setting, then waits
under a 5 s task setting while the holder releases at 500 ms, so it
discriminates against retaining the first task setting.
--
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]
