[BUG] Retry UCX shuffle sends after GPU OOM by binmahone · Pull Request #15982 · NVIDIA/cudf-spark

binmahone · 2026-09-14T03:52:55Z

Description

The UCX shuffle server can encounter an OutOfMemoryError when SpillableDeviceBufferHandle.materialize() restores a spilled shuffle buffer while preparing a send. The observed production exception was emitted by an RMM allocation through JNI, but Java exposes the generic OutOfMemoryError type. The existing BufferSendState handler catches only Exception, so this error escaped the shuffle-server background thread. In addition, the server terminated when every send in a server batch failed preparation, even when each failure came from this materialization boundary.

From the user's perspective, under sustained memory pressure an affected UCX send can now retry for up to spark.network.timeout before failing instead of failing immediately; the existing setting therefore also bounds sender-side recovery. This change introduces no new configuration.

This change:

catches OutOfMemoryError only around SpillableDeviceBufferHandle.materialize(), rather than around the entire send-preparation path, and wraps it in RapidsShuffleSendPrepareException;
includes the original OutOfMemoryError text in the diagnostic because this Java type can represent either the observed native/RMM failure or a JVM-heap failure during disk-tier decompression;
gives each BufferSendState an independent continuous materialization-OOM episode: its first OOM starts the retry window, and only that same state successfully preparing a send resets the window;
evaluates retry and expiry per send state, independent of which unrelated requests happen to share a server batch;
closes only expired send states while requeueing co-batched states that remain inside their own windows;
applies exponential backoff from 10 ms to a 100 ms cap when no send can be prepared, avoiding a hot loop that repeatedly pays materialization spill sweeps and device synchronization;
bounds each retry window using a duration derived from spark.network.timeout, without claiming that the sender's deadline is synchronized with the receiver's polling deadline;
avoids throwing out of the single BSS worker when a materialization-OOM retry window expires;
preserves the existing fail-fast behavior when an entire zero-progress batch contains a non-memory preparation failure; and
adds deterministic unit coverage for initial retry, exact-boundary expiry cleanup, independent deadlines for co-batched requests, isolation from unrelated send progress, exponential-backoff bounds, and real BufferSendState reset behavior.

The server batch is only a scheduling unit. It is not a Spark task, stage, or shuffle failure boundary. Consequently, one expired transfer request must not close a fresh request or have its own retry window reset by another peer's progress.

The retry is implemented at the shuffle-server state-machine boundary rather than by mechanically wrapping the allocation in RmmRapidsRetryIterator.withRetryNoSplit. BufferReceiveState can register its known local task IDs with RmmSpark.shuffleThreadWorkingOnTasks, but the sender's BSS worker is serving remote waiting tasks and does not have their task identities. It therefore cannot participate correctly in the task-scoped RMM retry protocol.

An earlier prototype of the catch-and-requeue path was validated on an eight-node B200 Spark 3.5.3 workload. It completed 25 out of 25 Q21 attempts, executed the sender OOM retry branch 29 times, and produced the same full-result SHA-256 in every retained result: 614b224bd54f3248fcdd0762d4511ab1d352f2a034bf54ec1e7dc3358d5d49d3. The prototype fix is binmahone/spark-rapids@429b3bf3fd43f39744d7134dbb2a973115cefeda; the retained build/source revision is binmahone/spark-rapids@a3392fb3e935be3f86f47c5c476fa3bb80ad7c56.

That production evidence is intentionally limited: the prototype used a broader catch, fixed 10 ms delay, and no timeout, per-state expiry, or reset behavior. The 25/25 result validates production reachability of catch-and-requeue and the retained outputs, not the current PR's timeout, expiry, reset, exception-boundary, or backoff behavior.

The current behavior was validated locally at commit 90f4b4dd5dca3907c611bdb749927bbe302141e5 with Spark 3.5.3 and CUDA 13 on ARM64:

mvn -B -Dmaven.repo.local=<isolated-repository> -Dbuildver=353 \
  -Dcuda.version=cuda13 -Parm64 -DskipTests -Dmaven.scaladoc.skip=true \
  -Drat.skip=true install -pl tests -am
BUILD SUCCESS

mvn -B -Dmaven.repo.local=<isolated-repository> -Dbuildver=353 \
  -Dcuda.version=cuda13 -Parm64 -Dmaven.scaladoc.skip=true \
  -DwildcardSuites=com.nvidia.spark.rapids.shuffle.RapidsShuffleServerSuite \
  test -pl tests
Tests: succeeded 10, failed 0, canceled 0, ignored 0, pending 0
BUILD SUCCESS

The reset test uses a real BufferSendState and verifies the next retry episode receives a fresh window. Replacing resetOomRetryWindow with a no-op or removing its successful-preparation call site makes that test fail.

This change does not eliminate the underlying memory pressure, distinguish native/GPU and JVM-heap OOM with a typed exception, immediately propagate an exhausted retry window to the peer, or implement tier-aware sending. After local timeout cleanup, the peer can still discover the failed request through its own fetch timeout; immediate protocol-level failure propagation remains separate work.

No performance improvement is claimed.

This change was prepared with AI assistance and reviewed by the author.

Checklists

Documentation

Updated for new or modified user-facing features or behaviors
No user-facing change

Testing

Added or modified tests to cover new code paths
Covered by existing tests
Not required

Performance

Tests ran and results are added in the PR description
Issue filed with a link in the PR description
Not required

Signed-off-by: Hongbin Ma <mahongbin@apache.org>

greptile-apps · 2026-09-14T03:56:03Z

Confidence Score: 5/5

The PR appears safe to merge; no actionable new failures remain, and all previous findings were resolved or withdrawn.

Summary

This PR adds bounded, per-send recovery for shuffle-buffer materialization OOMs and updates the implementation since the previous review so unrelated co-batched progress no longer resets or expires another request’s retry window.

Wraps materialization-time OutOfMemoryError at the narrow allocation boundary.
Tracks retry windows and attempts independently per BufferSendState.
Uses bounded exponential retry backoff and closes only expired states.
Adds deterministic coverage for timeout cleanup, same-state reset, and mixed-batch isolation.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Prepare next send] --> B{Preparation result}
  B -->|Success| C[Reset this state's OOM window]
  C --> D[Submit transport send]
  B -->|Materialization OOM| E[Record this state's retry attempt]
  E --> F{Window expired?}
  F -->|No| G[Back off and requeue state]
  F -->|Yes| H[Close only expired state and log failures]
  B -->|Other failure| I{Any send prepared?}
  I -->|No| J[Preserve fail-fast behavior]
  I -->|Yes| K[Preserve existing mixed-batch requeue]

_{Reviews (3) · Last reviewed commit: "Make UCX send OOM retries state-local"}

Copilot

🟡 Changes recommended

Critical server retry and interruption-handling issues remain unresolved.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds bounded GPU OOM retry handling for UCX shuffle send preparation, with exponential backoff and deadline enforcement.

Changes:

Wraps preparation OOMs for retry handling.
Requeues recoverable OOM failures with bounded backoff.
Adds deterministic retry and deadline tests.

File summaries

File	Summary
`tests/src/test/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleServerSuite.scala`	Tests initial requeue and persistent-OOM deadline behavior.
`sql-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleServer.scala`	Implements retry and timeout handling; critical issues remain with interruption handling and per-state deadline enforcement.
`sql-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/BufferSendState.scala`	Wraps preparation OOMs and resets retry state; normal-path bookkeeping has a performance nit.

Review details

Suppressed comments (3)

sql-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/BufferSendState.scala:276

This reset executes for every successful send preparation, so the change does add work to the normal transfer path even though retry logging and sleeping are failure-only. The Performance: Not required checklist item is therefore not justified without a benchmark/profile or a documented measured basis for ruling out a runtime impact.

    resetOomRetryWindow()

sql-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleServer.scala:405

When this deadline path throws, the exception escapes the long-running bssExec loop (which calls doHandleTransferRequest without a catch), so the single-threaded worker exits. Subsequent transfer requests stay in pendingTransfersQueue and are never processed; this abandons the whole server rather than just releasing the timed-out send state. Handle deadline exhaustion without killing the loop, or explicitly restart/fail the server as a whole.

            val ise = new IllegalStateException(
              s"Unable to prepare any sends. $reason The sends will not be retried.")
            supressedErrors.foreach(ise.addSuppressed)
            throw ise

sql-plugin/src/main/scala/com/nvidia/spark/rapids/shuffle/RapidsShuffleServer.scala:393

This sleep runs on the single-threaded bssExecutor loop (UCXShuffleTransport.scala:299-305). If the OOM persists until spark.network.timeout (120 seconds by default), the worker cannot drain unrelated pending or continuing transfers for the entire recovery window, so one request can stall all shuffle sends. Use a delayed/non-blocking requeue or a separate scheduler rather than blocking this executor.

            waitBeforeOomRetry(backoffMillis)

Files reviewed: 3/3 changed files
Comments generated: 2
Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

binmahone · 2026-09-14T06:38:30Z

+  private[shuffle] def waitBeforeOomRetry(backoffMillis: Long): Unit =
+    Thread.sleep(backoffMillis)


The revised delay is a fixed 10 ms, matching the cluster-validated prototype, rather than a growing delay. Each attempt then returns to the outer queue-draining loop; it does not sleep continuously for spark.network.timeout. The worker already uses interruptible Object.wait(100) in that same loop, so interruption is existing executor-shutdown behavior rather than a new 120-second blocking mode. More importantly, timeout exhaustion no longer throws from doHandleTransferRequest: it closes the affected states, logs the suppressed causes, and returns, so persistent OOM cannot terminate the sole BSS worker. The unit test verifies that the timeout call returns and performs cleanup without another requeue.

binmahone · 2026-09-14T06:38:26Z

+        // Requeue sends that could not acquire enough memory. Non-memory failures
+        // reach this point only when another send made progress.
        addToContinueQueue(toTryAgain.toSeq)


This is intentionally a continuous zero-progress timeout rather than a per-state wall-clock deadline. If another BufferSendState prepares successfully, the BSS worker has made forward progress; the revised code resets the failed states' OOM retry accounting and requeues them. If a later batch makes no progress, a new timeout window starts. I added a unit test that verifies the reset and requeue behavior. This avoids expiring a request merely because it previously shared a batch with an OOM while still preventing an all-OOM zero-progress loop from running forever.

Signed-off-by: Hongbin Ma <mahongbin@apache.org>

binmahone · 2026-09-14T06:44:46Z

build

binmahone · 2026-09-14T06:59:31Z

The failing Maven install/verify matrix is blocked before compiling this PR by an integration regression currently reproducible on main at b23ab02e986767dfd26fde2e221d58658d9a2179.

After #15944 and #15957 were both merged, test_root_safe_provider_selection.py still creates a fake iceberg_runtime.py containing only coordinates(), while package-parallel-worlds.py now also reads system_runtime_path. The representative jobs all fail in that test with the same KeyError: 'system_runtime_path'; the remaining matrix jobs are then canceled by fail-fast. This PR does not modify either file.

The isolated fixture fix and reproduction are in #15985. After that base fix lands, I will refresh this PR's Maven checks.

Follow-up to #15944 and #15957. ### Description After both Iceberg packaging changes reached `main`, `package-parallel-worlds.py` began reading `system_runtime_path` from the loaded `iceberg_runtime.py` module. `RootSafeProviderSelectionTest` creates a minimal fake version of that module, but the fixture still defined only `coordinates`. As a result, the standard-assembler subtest failed with `KeyError: 'system_runtime_path'` before Maven compilation, causing install and verify matrix jobs on subsequent pull requests to fail or be canceled by fail-fast. This change adds the missing `system_runtime_path` function to the fake module. It returns `None`, matching the test fixture's existing behavior when no explicit system Iceberg runtime is configured. Production code is unchanged. Validation on current `main` (`b23ab02e986767dfd26fde2e221d58658d9a2179`): ```text $ python3 dist/scripts/tests/test_root_safe_provider_selection.py . ---------------------------------------------------------------------- Ran 1 test in 0.251s OK ``` The failure was also reproduced before the fixture change with the same `KeyError` reported by the Maven matrix in #15982. This change was prepared with AI assistance and reviewed by the author. ### Checklists Documentation - [ ] Updated for new or modified user-facing features or behaviors - [x] No user-facing change Testing - [x] Added or modified tests to cover new code paths - [ ] Covered by existing tests - [ ] Not required Performance - [ ] Tests ran and results are added in the PR description - [ ] Issue filed with a link in the PR description - [x] Not required Signed-off-by: Hongbin Ma <mahongbin@apache.org>

amahussein

Thanks for the fix. I agree that having stopOomRetries return rather than throw is the right shape.

Requesting changes

The one code ask is inline: the zero-progress branch uses forall, so one expired window closes every co-batched send, including a peer's that just started. Base closed the whole batch too, so this is an incomplete improvement, but the description claims the opposite.
Two asks about the evidence rather than the code. The prototype at 429b3bf3 has no retry window and a broader catch, so the 25/25 run supports catch-and-requeue but says nothing about the timeout, expiry or reset paths. And #13098 is not the gap: BufferReceiveState already retries on a non-task thread via RmmSpark.shuffleThreadWorkingOnTasks. The real reason is that the server has no task identity, since the waiting tasks are remote. That is a better argument.

amahussein · 2026-09-14T16:16:52Z

+            val buff = try {
+              spillable.materialize()
+            } catch {
+              case oom: OutOfMemoryError =>


Agreed there is no GPU subtype to catch: in cudf-spark-jni, GpuOOM and OffHeapOOM both extend RuntimeException. But on the disk tier materialize() decompresses via RapidsSerializerManager.wrapStream, and spark.shuffle.spill.compress defaults true, so an lz4 heap allocation lands inside this catch and is retried for the full window as "GPU memory exhausted".
Suggest: Appending oom.toString to the message would distinguish them.

Addressed in 90f4b4d. The wrapper message is now neutral (Memory exhausted), includes oom.toString, and preserves the original cause. The PR and issue now explicitly state that generic OutOfMemoryError at this boundary can also represent JVM-heap failure during disk-tier decompression; this change does not claim typed GPU/heap classification.

amahussein · 2026-09-14T16:18:01Z

+            // This allocation runs on the shuffle server's non-task thread. Yield to task threads
+            // so they can spill and release GPU memory before this retry.


The spill has already happened here: DeviceMemoryEventHandler.onAllocFailure runs synchronously inside the failing allocation, calls store.spill, and device-synchronizes up to oomMaxRetries (default 2) before RMM throws. So each retry re-pays a spill sweep plus up to two device-wide syncs, up to ~12k times over a 120s window. Worth rewording the comment and reconsidering the backoff #15981 specified.

Addressed in 90f4b4d. Exponential backoff is restored at 10/20/40/80/100 ms with a 100 ms cap, matching #15981. The comment now explains that the delay avoids a hot loop because each failed materialization can repeat spill sweeps and device synchronization; it no longer claims the sleep gives task threads their first opportunity to spill. Backoff bounds are covered by the suite.

amahussein · 2026-09-14T16:19:01Z

+    }
+  }
+
+  test("OOM retry accounting resets after shuffle send progress") {


Test 1 does pin the boundary: it drives a real BufferSendState. This one mocks both, so it verifies the call, not the effect. Replacing the body of resetOomRetryWindow with a no-op still passed all 8 tests, while deleting the call site here is caught. One case on a real state would close that gap.

Addressed in 90f4b4d. The mock interaction test was replaced with a real BufferSendState behavior test: it opens an OOM window, successfully prepares a send, and verifies that an OOM at the old deadline starts a fresh window. A no-op reset body or removal of the successful-preparation reset call makes the assertion return None instead of Some(1).

Confirmed by mutation: a no-op resetOomRetryWindow body now fails a test, where it passed all 8 before. That was the gap, so this is closed.

Signed-off-by: Hongbin Ma <mahongbin@apache.org>

binmahone · 2026-09-15T07:15:56Z

Addressed the requested changes in 90f4b4d. In addition to the inline fixes, I refreshed the PR and #15981 descriptions so the evidence and implementation contracts match: the 25/25 prototype validates only catch-and-requeue reachability/output correctness, not the current timeout/expiry/reset/backoff paths; and the RMM-retry rationale now identifies the sender server lacking the remote waiting task identities, while acknowledging that BufferReceiveState already registers known task IDs on its non-task thread. Local Spark 3.5.3/CUDA 13 ARM64 validation passed the full -pl tests -am install and RapidsShuffleServerSuite 10/10.

binmahone · 2026-09-15T07:31:49Z

build

binmahone · 2026-09-15T08:18:45Z

build

amahussein

Nothing blocking, and two small asks, neither of which needs a respin:

The Performance box ticks "Tests ran and results are added in the PR description", but the description has no timings, and says "No performance improvement is claimed". Either paste the wall-clock numbers from the 25/25 Q21 run, or tick "Not required" instead.
The template's Description section asks what the user experience will be. Under sustained GPU pressure a UCX send now fails after up to spark.network.timeout instead of immediately, and that setting now also bounds sender-side retry. One sentence saying so, and then whichever Documentation box matches.

For a follow-up rather than this PR: the backoff is skipped whenever any co-batched send makes progress, so a failing send re-materializes at loop speed until its window expires. That gating predates this PR and the window now bounds it, so it is not urgent, but a per-state earliest-retry timestamp would express the per-state contract this commit introduces. Happy to file it if you would rather not carry it.

amahussein · 2026-09-15T13:39:59Z

+        }
+
+        if (bssBuffers.isEmpty && retryableOom.nonEmpty) {
+          val retryAttempt = retryableOom.flatMap(_._3).max


Non-blocking, no need to respin for it. Changing this .max to .min passes all ten tests, so the aggregate is unpinned. max is the right choice: min would hold the delay at 10 ms whenever a fresh state joins the batch, which is the hot loop this commit removes. Worth one assertion next time you touch this file.

Agreed. max is intentional so a fresh co-batched state cannot hold the aggregate delay at 10 ms, but the existing tests do not distinguish max from min. I will keep the per-state earliest-retry timestamp and its mutation-resistant assertion as follow-up work rather than respinning this approved patch.

abellina · 2026-09-15T14:11:15Z

            val spillable = blockRange.block.bufferHandle.spillable
-            val buff = spillable.materialize()
+            val buff = try {
+              spillable.materialize()


A retry is OK to unblock, but it is not what I was expecting.

We are packing a bounce buffer here which is already allocated. If the source data could not be materialized to GPU, especially if it's in host mem currently, this retry should be "copy from the tier where the buffer is currently". The next issue then is ensuring we have host buffers for a disk->host->bounce buffer trampoline.

Currently getting a host bounce buffer is optional. I think we could make that required, in which case we would always be able to satisfy the send, guaranteeing that OOM errors go away. We still need to handle transport errors but this class of errors that could cause nightmare spills while we insist to materialize go away.

Agreed that tier-aware sending is the stronger architectural endpoint: if the source is already in host or disk tier, repeatedly rematerializing it on device can cause unnecessary spill churn. This PR remains the bounded recovery path that unblocks the current failure. Making a host bounce buffer reliably available and copying from the source tier into it needs a separate design covering reservation/capacity, lifecycle, and device/host/disk source paths; I will keep that scope out of this patch.

binmahone · 2026-09-15T14:27:34Z

Thanks, Ahmed. I updated the PR description to state the user-visible behavior under sustained memory pressure and that spark.network.timeout now also bounds sender-side recovery. I also changed the Performance checklist to Not required because this PR makes no performance claim.

binmahone · 2026-09-16T02:38:07Z

CI blocked by #16004

Retry UCX shuffle sends after GPU OOM

b2d3d5e

Signed-off-by: Hongbin Ma <mahongbin@apache.org>

Copilot AI lite review requested due to automatic review settings September 14, 2026 03:52

Copilot started reviewing on behalf of binmahone September 14, 2026 03:53 View session

greptile-apps Bot reviewed Sep 14, 2026

View reviewed changes

Copilot AI reviewed Sep 14, 2026

View reviewed changes

Address review feedback for UCX OOM recovery

de5166a

Signed-off-by: Hongbin Ma <mahongbin@apache.org>

binmahone mentioned this pull request Sep 14, 2026

Fix root-safe provider test runtime stub [reduced-it] #15985

Merged

8 tasks

sameerz requested review from abellina, amahussein and zpuller September 14, 2026 14:07

amahussein requested changes Sep 14, 2026

View reviewed changes

Make UCX send OOM retries state-local

90f4b4d

Signed-off-by: Hongbin Ma <mahongbin@apache.org>

binmahone mentioned this pull request Sep 15, 2026

[BUG] UCX shuffle sender can terminate after GPU OOM while materializing spilled buffers #15981

Open

amahussein approved these changes Sep 15, 2026

View reviewed changes

abellina reviewed Sep 15, 2026

View reviewed changes

		private[shuffle] def waitBeforeOomRetry(backoffMillis: Long): Unit =
		Thread.sleep(backoffMillis)

		// This allocation runs on the shuffle server's non-task thread. Yield to task threads
		// so they can spill and release GPU memory before this retry.

Uh oh!

Conversation

binmahone commented Sep 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Checklists

Uh oh!

greptile-apps Bot commented Sep 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (3)

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

binmahone commented Sep 14, 2026

Uh oh!

binmahone commented Sep 14, 2026

Uh oh!

amahussein left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

binmahone commented Sep 15, 2026

Uh oh!

binmahone commented Sep 15, 2026

Uh oh!

binmahone commented Sep 15, 2026

Uh oh!

amahussein left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

binmahone commented Sep 15, 2026

Uh oh!

binmahone commented Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

binmahone commented Sep 14, 2026 •

edited

Loading

greptile-apps Bot commented Sep 14, 2026 •

edited

Loading

binmahone commented Sep 16, 2026 •

edited

Loading