rapids shuffle manager V2 phase 1: writer use as much memory as allowed and pipelined write

[binmahone]

This PR is actually a superset of #13479
Based on #13479 we further added SpillablePartialFileHandle, which is:

SpillablePartialFileHandle

A new spillable storage mechanism for partial shuffle files that optimizes I/O by preferring memory over disk.

Key Features

• Memory-first storage: Stores data in host memory buffer when memory usage is below threshold
• Automatic spill: Spills to disk when memory becomes scarce or buffer capacity is exceeded
• Dynamic buffer expansion: Buffer grows dynamically (doubling) up to configured max limit

Storage Modes

Mode	Description
FILE_ONLY	Always write to disk (used when memory is tight)
MEMORY_WITH_SPILL	Use memory buffer with automatic spill capability

Benefits

• Reduces disk I/O for shuffle operations when sufficient host memory is available
• Graceful fallback to disk storage under memory pressure
• Unified read/write interface regardless of storage location

[binmahone]

build

[greptile-apps]

_{6 files reviewed, 6 comments}

_{Edit Code Review Agent Settings | Greptile}

[binmahone]

Mostly nits. However there are still a lot of greptile comments that need to be looked at and addressed in some way.

Yes, @revans2 , greptile keeps popping new comments, which after double check seems to be over-engineering or just defensive programming. You comments are just resolved, please double check

[binmahone]

Just merged origin/main, so that 250923_optimize_nds_more is on top of latest main.

[binmahone]

build

[greptile-apps]

_{6 files reviewed, 6 comments}

_{Edit Code Review Agent Settings | Greptile}

[binmahone]

build

[greptile-apps]

_{20 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[abellina]

last comment, thank you for the updates so far. This is close.

[abellina]

Thank you for adding the comments. It is clearer to me what is going on here.

There are two things, one that has to be addressed in this PR is that this wait just means we are going to busy loop. The pattern to wait/notify is classically the following:

synchronized {
  while(condition==false) {
     wait( .. sure with optional timeout.. )
  }
}

synchronized {
  condition = true
  notify(); // or notifyAll() if multiple waiters.
}

Because we don't have a condition we are checking, this is just going to wait 1ms, sometimes, and just keep going, busy looping and consuming CPU.

I would like us to check a variable in the waiter and loop until the condition says "we should loop around and do work" in this PR.

As a subsequent improvement, we can look at doing this with a Future. As the get function in the Future should handle this waiter pattern.

[binmahone]

You're right that wait(1) with a 1ms timeout causes frequent spurious wakeups, which is inefficient. I've considered the alternatives:

Option 1: Classic condition flag pattern

var hasWork = false
while (!hasWork) {
  mergerCondition.wait()
}
hasWork = false

This requires additional synchronization to keep the hasWork flag consistent with the actual data state (partitionFutures, maxPartitionIdQueued, etc.), which adds complexity and potential for bugs.

Option 2: Increase timeout to wait(100)

Simple change - just increase the timeout from 1ms to 100ms.

I believe Option 2 is the most pragmatic choice:

• Minimal code change, lowest risk
• 100ms timeout eliminates busy-loop concerns
• Does NOT affect responsiveness: if notifyAll() is called, the thread wakes up immediately (not after 100ms)
• The timeout is purely defensive - in normal operation, the thread is always woken by notifyAll()

The current code already has the correct notifyAll() signals in place (after queuing tasks, after batch completion). The 100ms timeout just serves as a safety net in case a signal is missed, rather than spinning at 1ms intervals.

Would this approach address your concern?

[abellina]

Why do we want to poll? why don't we use the classic flag pattern? The flag tells us right away what it is that we are waiting on, it's a clear signal on what we need to proceed, instead of being notified somewhere else in the code and we need to go read why that notification happened. Imagine trying to debug a race here. If we had a clear flag that flipped, we can log it or detect it. I would really really like us to go Option 1 here.

[binmahone]

updated , please check again @abellina

[greptile-apps]

_{20 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{20 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[binmahone]

build

[binmahone]

build

[greptile-apps]

_{20 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

logic: infinite loop when oldCapacity is 0

if oldCapacity is 0, then newCapacity = 0 and 0 * 2 = 0 forever, causing the while loop to never terminate

while the default is 128MB, users can configure spark.rapids.memory.host.partialFileBufferInitialSize to 0, which causes this infinite loop

Suggested change

	var newCapacity = oldCapacity
	while (newCapacity < requiredCapacity && newCapacity < maxBufferSize) {
	newCapacity = math.min(newCapacity * 2, maxBufferSize)
	}
	// Calculate new capacity: keep doubling until >= requiredCapacity
	var newCapacity = if (oldCapacity == 0) 1 else oldCapacity
	while (newCapacity < requiredCapacity && newCapacity < maxBufferSize) {
	newCapacity = math.min(newCapacity * 2, maxBufferSize)
	}

[greptile-apps]

logic: missing validation allows 0 which causes infinite loop in SpillablePartialFileHandle.expandBuffer()

add .checkValue(v => v > 0, "must be > 0") before createWithDefault to prevent the infinite loop at SpillablePartialFileHandle.scala:134

Suggested change

	val PARTIAL_FILE_BUFFER_INITIAL_SIZE =
	conf("spark.rapids.memory.host.partialFileBufferInitialSize")
	.doc("The initial size in bytes for a host memory buffer used by " +
	"SpillablePartialFileHandle during shuffle write. This buffer allows shuffle " +
	"data to be kept in memory instead of writing to disk immediately, reducing " +
	"I/O overhead. The buffer can expand dynamically up to partialFileBufferMaxSize. " +
	"A smaller initial size reduces upfront memory allocation but may require more " +
	"expansions.")
	.startupOnly()
	.internal()
	.bytesConf(ByteUnit.BYTE)
	.createWithDefault(128L * 1024 * 1024) // 128MB default
	val PARTIAL_FILE_BUFFER_INITIAL_SIZE =
	conf("spark.rapids.memory.host.partialFileBufferInitialSize")
	.doc("The initial size in bytes for a host memory buffer used by " +
	"SpillablePartialFileHandle during shuffle write. This buffer allows shuffle " +
	"data to be kept in memory instead of writing to disk immediately, reducing " +
	"I/O overhead. The buffer can expand dynamically up to partialFileBufferMaxSize. " +
	"A smaller initial size reduces upfront memory allocation but may require more " +
	"expansions.")
	.startupOnly()
	.internal()
	.bytesConf(ByteUnit.BYTE)
	.checkValue(v => v > 0, "must be > 0")
	.createWithDefault(128L * 1024 * 1024) // 128MB default