fix: make CometDiskBlockWriter spill registry per-task instead of executor-global - #5493
fix: make CometDiskBlockWriter spill registry per-task instead of executor-global#5493peterxcli wants to merge 13 commits into
Conversation
…cutor-global The static currentWriters list in CometDiskBlockWriter was shared by all shuffle tasks in the executor, so a task under memory pressure could spill other tasks' writers, count their memory as freed for itself, and retain writers of failed tasks forever. It also allowed an ABBA deadlock between the writer monitor and the registry lock when two tasks spilled concurrently, and unsynchronized cross-thread reads could throw ConcurrentModificationException. Replace the static registry with a task-owned list created by CometBypassMergeSortShuffleWriter and passed to its partition writers, and delete the dead spillingWriters path (nothing ever added to it). All access to a writer is now confined to its own task's thread.
61a10fb to
ccbf5b7
Compare
In on-heap mode all tasks share one CometBoundedShuffleMemoryAllocator, so with per-task spill registries a task whose own writers hold nothing spillable failed with SparkOutOfMemoryError when other tasks held the pool. Preserve allocation progress by letting the post-spill retry in SpillWriter.initialCurrentPage wait (bounded by a timeout) for other tasks to free pool memory; free() notifies waiters. This point is only reached after the task has spilled everything it owns, so waiting tasks hold no pool memory and the tasks holding it can always progress and free it. The off-heap allocator keeps its non-waiting behavior since Spark's memory manager already arbitrates between tasks.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
[P2] Rechecked 3cee056cd1ded0240892ef53ebfef515fd6c1e4e: the existing shared-pool finding is partially fixed. An independently released 1.5-second holder now succeeds, but a real Spark 3.5.9 Comet JVM shuffle with a 1-MiB pool and 256-KiB pages still raises SparkOutOfMemoryError at the 60-second allocation deadline when another task independently pauses for 65 seconds while holding four pages. The exact base Java-class overlay and current-head 2-MiB control each return all 752 correct rows. This remains limited to opt-in on-heap compatibility; the probe allows one task attempt, so normal Spark retries may recover.
Any finite deadline fails some workload: a task holding the shared on-heap pool longer than the cap still made waiters fail with SparkOutOfMemoryError. Wait until memory is freed instead, mirroring how Spark's unified memory manager blocks a task until memory becomes available. Liveness is preserved because a waiting task has already spilled everything it owns, so pool memory is only held by tasks that can progress and free it; interrupting the task (e.g. task kill) still aborts the wait. Log once when a task starts waiting so long waits are diagnosable.
…non-fatal The unbounded wait could hang when it can never succeed: the sort-based path retains its pointer array through an empty spill(), so a waiting task may itself hold pool memory, and two such tasks can block each other forever; a single request larger than the pool waited on an empty pool. Track the pool memory retained per thread and give up before waiting when the request cannot fit next to the requester's own retained memory, or when all allocated memory is retained by threads that are themselves blocked waiting (nothing can be freed anymore). Also stop rethrowing SparkOutOfMemoryError when the wait is interrupted: it extends OutOfMemoryError, so Spark's killed-task handler does not match it and an intentional task kill was reported as ExceptionFailure. Throw a RuntimeException wrapping the InterruptedException instead so kills are classified as TaskKilled.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
[P2] Follow-up to the allocation-progress thread: the new check at CometBoundedShuffleMemoryAllocator.java:163 in b8fadf58 can reject a satisfiable schedule. Notified threads remain in waitingThreads until allocation returns. With a 1-MiB pool and 256-KiB pages, after a healthy holder finishes, two unsafe-shuffle tasks each retain 32,768 pointer bytes and request 999,448 / 262,144 bytes. If the large waiter resumes first, it throws with 983,040 bytes free, although the smaller waiter can finish and release its array for it. In fresh Spark 3.5.9 opt-in on-heap tests, small-first arrival failed 10/10; exact 4c12797 Java overlays and current reversed arrival each returned three correct rows 10/10, and the matching 2-MiB control passed 3/3. Please let a satisfiable notified request proceed before declaring deadlock. This is an incremental regression against 4c12797, reproduced with local[3,1].
The deadlock check counted notified-but-not-yet-resumed waiters as permanently blocked: after a holder freed the pool, a large waiter that resumed first could declare deadlock even though a smaller waiter, still registered as waiting, could satisfy its request from the free pool, finish, and release memory for the large one. Track each waiter's pending request size and only declare deadlock when, in addition to all allocated memory being retained by blocked waiters, none of their requests fits in the free pool. A waiter whose request fits will proceed once it reacquires the monitor and eventually free what it retains, so the wait remains live.
The core of this is clearly right. A static registry that lets task A spill task B's buffers, then reads B's My concerns are all about the second half of the PR. The description does not mention The "What changes are included in this PR?" section only describes the registry change. But this PR also adds about 80 lines of new blocking-allocation logic to I do understand the connection: once a task can no longer take memory from other tasks, it needs somewhere to go when the pool is full. But that is an argument that deserves to be made explicitly, and it may well be an argument for a second PR rather than a second half.
From an operator's point of view that is a hung task with one line of explanation. Would you add a bounded wait, either
Entries are removed from Interrupt is rethrown as The comment says throwing a plain Test flakiness
|
A task blocked waiting for the shared on-heap pool logged a single warning and then went silent, so a long stall (e.g. a holder stuck on I/O) was invisible to operators. Use a timed wait and re-log every 30 seconds with the elapsed time, requested size, free bytes, and number of waiting threads. The wait itself remains unbounded (a fixed deadline was shown in review to fail workloads whose holder legitimately outlives the cap); the existing fail-fast checks still cover the provably-unsatisfiable cases. Also harden the suite's thread-state polls to accept TIMED_WAITING now that the allocator uses a timed wait.
|
Thanks for the review @andygrove. Point-by-point: PR description — agreed, that was a real gap. The description now has a dedicated section explaining
Interrupt classification — the exception type is what drives it. Spark's killed-task clause in Test flakiness — the suite runs in ~8 s locally with 60 s |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed cab20e38 with five independent reviewers. Timed reevaluation addresses the previous missed-notification stall in source. The distinct P2 below is source-verified; no current local runtime reproduction was performed.
…fails SpillSorter's constructor allocates the one-entry ShuffleInMemorySorter array and then the real pointer array. When the second allocation failed, the first was orphaned: the writer is never handed to Spark, so no cleanup path could ever free it. Pre-existing, the orphan was only a tiny pool leak; with blocking allocation it also broke the wait's progress condition, since the orphaned bytes look like memory a live holder could still free, turning a formerly prompt SparkOutOfMemoryError for a later full-pool request into an unbounded wait. Free the partially-constructed sorter's allocations before rethrowing, so every allocated byte again has an owner with a cleanup path - the invariant the blocking wait relies on.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
[P2] The constructor leak in the existing orphan-allocation thread is fixed at a21f3f19, but bypass failures can still orphan pages. With opt-in on-heap JVM shuffle, a 1-MiB pool, 256-KiB pages and bypass threshold 3, a three-partition task can spill its largest writer and retain 524,288 sibling bytes. Before its 512,024-byte retry, an independent four-partition unsafe task can retain a 32,768-byte pointer array and wait for 999,448 bytes. Only 491,520 bytes remain free, so the all-waiter check throws SparkOutOfMemoryError in the bypass task.
Both bypass write() and Spark's ShuffleWriteProcessor catch only Exception; stop(false) is skipped, and task-memory cleanup does not own these private pages. The unrelated unsafe job can then wait indefinitely, including through every 30-second recheck. Please reclaim the bypass pages on this error path. The cleanup gap is pre-existing; the introduced regression is unbounded waiting. This is verified against current source and Spark 3.5.9 contracts, not an executed reproduction.
…rror Spark's ShuffleWriteProcessor only calls stop(false) when the shuffle writer throws an Exception, so a fatal error such as SparkOutOfMemoryError (e.g. from the blocking allocator's fail-fast checks) skipped cleanup and orphaned the task's buffered pages. Those pages live in the executor-shared bounded pool that Spark's task-memory cleanup does not own, so on-heap they starved other tasks' allocations - and with blocking allocation, an unrelated waiter could wait on them indefinitely. Catch Throwable in CometBypassMergeSortShuffleWriter.write() and free all partition writers' memory before rethrowing (freeMemory is a no-op when stop(false) runs afterwards on Exception paths). CometUnsafeShuffleWriter already handles generic throwables via its success/finally pattern.
|
Addressed the bypass orphan-pages finding in a1a4a84. You're right that |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed a1a4a84. The earlier bypass-write failure and failed-constructor cleanup are addressed. [P2] One case remains under the existing orphan-allocation discussion, in Spark 3.5.9's pre-write path.
The unsafe writer allocates its pointer array before Spark evaluates rdd.iterator. Comet's eager round-robin input sort can then throw a managed SparkOutOfMemoryError. The writer's write() cleanup has not started, the processor's Exception handler skips stop(false), and Spark task-memory cleanup does not own the 32-KiB Comet allocation.
With on-heap JVM shuffle and a 1-MiB Comet pool/page size, a bypass task on another thread can then wait indefinitely on that abandoned allocation. Controlled runs with the actual 3.5.9 processor, writers and shared Spark memory manager reproduced head waiting versus base failing promptly; explicit orphan cleanup resumed native output. These used constrained managed memory, manual task contexts and a fixed CI native binary, not a full scheduler run.
Please give allocations made before write() task-completion cleanup, or defer them until the writer's cleanup scope owns them.
…ask completion The unsafe writer's sorter allocates its pointer array at construction, before Spark has evaluated the shuffle input iterator and called write(). A fatal error during iterator evaluation (e.g. SparkOutOfMemoryError from an eager input sort) therefore skipped every existing cleanup path: the writer's success/finally scope never starts, ShuffleWriteProcessor only calls stop(false) for exceptions, and Spark's task-memory cleanup does not own Comet's pool. The orphaned allocation could then starve other tasks' blocking waits indefinitely. Register a TaskCompletionListener at writer construction that runs sorter.cleanupResources() if the sorter is still alive; listeners run in Task.run's finally regardless of failure type, and cleanupResources() is idempotent so this is a no-op whenever write()/stop() already cleaned up.
|
Addressed the pre-write orphan finding in 844b308, taking the "give allocations task-completion cleanup" option: |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
[P2] The orphan-allocation issue remains partially addressed at 844b308. SpillSorter still calls serializeSchema outside the cleanup block after adopting its pointer array. A nested integer field with parquet.field.id=2147483648L passes the ordinary Spark 3.5 schema/encoder path, then throws in the Parquet ID conversion. The outer sorter is never assigned, so the new completion listener cannot release the 32 KiB array; in the opt-in on-heap JVM path that orphan can keep another job's first-page allocation waiting indefinitely. Please cover later constructor failures as well. This is verified from source; the passing pre-write test does not exercise it, and no current runtime reproduction is claimed.
The constructor cleanup only wrapped the pointer-array allocation, but serializeSchema (and the Native library load) run after the array is adopted and can still throw - e.g. schema serialization rejects a nested field whose parquet.field.id exceeds the 32-bit range. The enclosing sorter field is never assigned in that case, so not even the unsafe writer's task-completion listener could reclaim the adopted 32 KiB array. Widen the cleanup to the rest of the constructor and track whether the sorter adopted the array, so exactly one owner frees it: before adoption the array is freed directly (the sorter still owns only its initial one-entry array), after adoption inMemSorter.free() releases it, avoiding a double free that would trip TaskMemoryManager assertions in off-heap mode. Also skip null entries in the bypass writer's failure cleanup, in case a mid-loop writer construction failure left the tail of the array unassigned.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
[P2] Account for live owners blocked in Spark's execution-memory pool
At 4b592b9, the reported orphan allocations are reclaimed, but the allocation-progress condition still permits a live cross-pool cycle.
In opt-in on-heap JVM shuffle on Spark 3.5.9, with round-robin pre-sorting, bypass disabled and 4096-entry pointer buffers, each unsafe writer holds 32 KiB before its input iterator is built. B can finish sorting while it alone owns more than three quarters of the fully growable Spark execution pool; A then blocks below its minimum share in Spark while retaining its Comet array. For B's first nonfinal 976-KiB binary payload, its approximately 999448-byte Comet request cannot fit in a 1-MiB pool with both arrays (983040 bytes free), yet fits beside its own array (1015808 bytes). A is outside waitingThreads, so B waits too. Each task holds memory the other needs; neither can reach cleanup, and timed retries see unchanged state.
BASE fails and unwinds at this retry instead. Please avoid retaining pre-write Comet allocations across blocking input construction, or otherwise break this cross-pool cycle before waiting indefinitely. This is verified from current and matching public 3.5.9 source, not an executed reproduction.
Spark evaluates the shuffle input iterator between constructing the shuffle writer and calling write(), and that evaluation can block on Spark's execution-memory pool (e.g. Comet's eager round-robin input sort). Because the writer's constructor allocated the sorter's pointer array up front, a task blocked in that window retained Comet pool memory while holding no place in the Comet allocator's waiting set, so a second task blocked in the Comet pool saw it as a live holder and waited forever: task A waits on Spark memory held by B, task B waits on Comet memory held by A, and neither can reach cleanup. Move open() from the constructor into write(), so no Comet pool memory is retained across blocking input construction and the cross-pool cycle cannot form. Everything that uses the sorter already runs inside write(), and stop(), peak-memory accounting, and the task-completion listener all null-guard it. The listener stays as a backstop for fatal errors during write() itself.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
[P2] Following up on the cross-pool finding: moving open() fixes eager input construction, but it still reserves the Comet pointer array before records.hasNext(). On Spark 3.5 with on-heap JVM shuffle, bypass disabled, constrained pools and large rows, ordinary lazy Sort or partial HashAggregate input can then block in Spark execution memory. A second task can retain its nonfinal Spark input pages while waiting for a Comet page that fits beside its own pointer array but not both arrays. The first task is absent from waitingThreads, so the checks keep allowing a stable two-pool wait and neither task reaches cleanup. BASE instead raised a managed allocation error and unwound. Please also prevent retention across this lazy input boundary, or make the retry unwind when that dependency cannot progress. This is a source-derived reachable schedule, not an executed reproduction.
Deferring the sorter allocation to write() still reserved the pointer array before the first records.hasNext(), where lazy input (a sort or partial aggregate materializing its child) can block on Spark execution memory - and mid-write hasNext() calls can block the same way, so no placement of the allocation removes the cross-pool cycle entirely: a task blocked in Spark's pool retains Comet memory, while a Comet waiter holds Spark-side input pages the first task needs. Two changes: allocate the sorter only after the input has produced its first record, so the common lazy-materialization window retains nothing; and make a Comet waiter unwind (after at least one full wait interval) when the memory it waits for is retained by threads parked in an untimed wait inside Spark's execution-memory pool, detected from their stack frames (ExecutionMemoryPool / UnifiedMemoryManager). That is precise for the cycle's dependency: sleeping, latch-parked, or computing holders do not match and keep the wait alive, matching the previously-reviewed healthy-holder scenarios. Unwinding raises the managed SparkOutOfMemoryError, as base did, so the task releases its memory and tasks in the other pool can progress.
|
Addressed the lazy-input cross-pool follow-up in ab5c033 with both halves of your suggestion, since neither alone closes it: deferral cannot help once rows are buffered (mid-write
Added a regression test that builds the real two-pool state: one task owns the whole Spark execution pool (a real |
3 participants
Which issue does this PR close?
Part of #5212 (positions 6, 7, 9, and 10). It does not close the epic.
Rationale for this change
CometDiskBlockWriter.currentWritersis declaredstatic final, so despite its comment ("of same shuffle task") it is an executor-global registry shared by every shuffle task in the JVM. This causes several problems, all reproduced with a deterministic two-task reproducer against currentmain:spill()sorts the global list and force-spills the largest writers regardless of owner. In the reproducer, task A force-flushed task B's 3,145,728 buffered bytes (2,751 rows written to B's file from A's thread, disk-spill metrics charged to B), counted all of it toward its owntotalFreed, and spilled none of its own data. With-ea(as CI runs), task A then crashes withAssertionErrorininitialCurrentPagebecause its own active page was never freed.insertRowholds the writer monitor and then takes the registry lock insidespill(), while a concurrently spilling task holds the registry lock and requests the victim writer's monitor. A benchmark with two concurrent tasks under a shared memory cap deadlocked in 3 of 3 rounds (confirmed viaThreadMXBean.findDeadlockedThreads()).ConcurrentModificationException(position 7): the spill comparator reads other tasks'allocatedPageslists unsynchronized while their owners mutate them; observed live in the same concurrent benchmark.spillingandtotalWrittenare mutated from other tasks' threads without proper synchronization.stop(false)frees memory but never removes writers from the static list, so they are retained for the executor lifetime (12/12 writers still strongly reachable after GC across three simulated failed tasks). The retained writers also keep row addresses that point into freed pages, so a later task spilling them would hand freed addresses to native code.spillingWritersis never populated, so its loop infreeMemory()is unreachable.What changes are included in this PR?
1. Task-owned spill registry. Replace the static registry with a task-owned
LinkedListcreated inCometBypassMergeSortShuffleWriter.write()and passed to each partition'sCometDiskBlockWriter. Spilling now only ever touches the requesting task's writers, matching Spark's own shuffle semantics, and every writer is confined to its task's thread (which also removes the deadlock, theConcurrentModificationException, and the cross-thread mutation of positions 7/9 structurally). Writers of a failed task become unreachable together with the task's shuffle writer, removing the leak. The deadspillingWritersfield and its loop infreeMemory()are deleted.2. Blocking allocation for the shared on-heap pool (
CometBoundedShuffleMemoryAllocator.allocateBlocking). The registry change removes a task's only recourse when the shared pool is full: previously it force-spilled other tasks' buffers. In the opt-in on-heap compatibility mode all tasks share one bounded allocator, so without a replacement a task whose own writers hold nothing spillable would fail withSparkOutOfMemoryErrorwhenever other tasks held the pool (reproduced in review with a real Spark 3.5.9 shuffle). The replacement mirrors Spark's unified memory manager: after a task has spilled everything it owns (the retry inSpillWriter.initialCurrentPage, the only caller), it waits for other tasks to free pool memory instead of failing, andfree()notifies waiters.The waiting logic keeps per-thread accounting (which thread allocated each page, how much each thread retains, and which threads are blocked waiting with what request size) in order to fail fast, before or during a wait, exactly when the wait provably cannot succeed:
spill()), which also covers a single request larger than the whole pool; orSparkOutOfMemoryErrorso its task releases memory and the other pool can progress; sleeping, latch-parked, or computing holders deliberately do not match.To keep such cycles from forming in the first place, the unsafe writer also allocates its sorter lazily:
open()moved from the constructor intowrite(), after the firstrecords.hasNext(), so no Comet pool memory is retained while Spark evaluates or lazily materializes the shuffle input (which can block on Spark execution memory). Writer-held allocations that fatal errors would otherwise orphan (Spark only callsstop(false)for exceptions) are reclaimed on every path: the bypass writer frees its partition writers' pages on anyThrowableout ofwrite(),SpillSorterfrees everything it allocated when its constructor fails part-way, and a task-completion listener backstops the unsafe writer.A wait that can still succeed is not bounded by any deadline (a fixed deadline was shown in review to fail workloads where a healthy holder simply takes longer than the cap), but the allocator logs a warning when a wait starts and every 30 seconds while it lasts, so a stalled task is visible in executor logs. Interrupting a waiting task surfaces a NonFatal
RuntimeException(not the fatalSparkOutOfMemoryError) so that an intentional task kill is classified by Spark's killed-task handler (case _: InterruptedException | NonFatal(_) if task.reasonIfKilled.isDefinedinExecutor.TaskRunner) asTaskKilledrather thanExceptionFailure. The off-heap (CometUnifiedShuffleMemoryAllocator) path is unchanged: it delegates to Spark's memory manager, which already arbitrates memory between tasks.Single-task spill behavior is unchanged: in an A/B benchmark the spilled byte counts were identical in every run, and wall-time differences were within noise. The two-concurrent-task benchmark that deadlocks on
maincompletes in all repetitions with this change.How are these changes tested?
New
CometDiskBlockWriterSuite(runs in ~7 s) covers both halves:TaskMemoryManagers sharing oneTestMemoryManager; one task is driven into memory pressure and must resolve it by spilling only its own writer while the other task's buffered rows, spill metrics, and file remain untouched. Fails onmain(the victim task's buffer is force-flushed and the requesting task trips theinitialCurrentPageassertion), passes with this change.InterruptedExceptioncause; and a notified waiter is not declared deadlocked while another waiter's smaller request can proceed (both waiters complete regardless of wake order).🤖 Generated with Claude Code