Add a debug option to check if memory allocation is covered by retry framework

[thirtiseven]

Contributes to #13672
Depends on NVIDIA/spark-rapids-jni#4061

Description

This PR add a debug option to check if memory allocation is covered by retry framework, so we can find if we have code that will do a host/device memory allocation is covered by retry framework, which is potentially cause OOM.

The option is default off by default and should be only useful for developers. We can run it from time to time to find uncovered cases, or set up a CI to run periodically.

Here is the python script to convert the csv to a flamegraph: https://gist.github.com/thirtiseven/f72e14c1a860dc40a75f1351c3616a68

csv result file generated by this pr:

uncovered_allocations.csv

with following command:

SPARK_RAPIDS_RETRY_COVERAGE_TRACKING=true RANDOM_SELECT=0.05 ./integration_tests/run_pyspark_from_build.sh

Checklists

• This PR has added documentation for new or modified features or behaviors.
• This PR has added new tests or modified existing tests to cover new code paths.
  (Please explain in the PR description how the new code paths are tested, such as names of the new/existing tests that cover them.)
• Performance testing has been performed and its results are added in the PR description. Or, an issue has been filed with a link in the PR description.

[greptile-apps]

Greptile Summary

This PR adds a debug feature to detect memory allocations that are not protected by the retry framework (withRetry/withRetryNoSplit). The implementation introduces:

• AllocationRetryCoverageTracker: A new tracking utility that captures stack traces of uncovered allocations and logs them to a CSV file. Enabled via the SPARK_RAPIDS_RETRY_COVERAGE_TRACKING environment variable (disabled by default).

• RetryStateTracker enhancements: Added isInRetryBlock() method with a thread-local depth counter to properly track nested retry blocks. The depth tracking ensures correct behavior when retry blocks are nested.

• Integration points:

  • Device allocations are tracked via RMM's onAllocated() callback (requires RMM debug mode)
  • Host allocations are tracked in HostAlloc after successful allocation
  • The shell script propagates the environment variable to executors

Key Design Decisions:

• Off by default to avoid performance overhead in production
• Uses thread-local depth counters to handle nested retry blocks correctly
• Deduplicates logged stacks using ConcurrentHashMap to reduce noise
• Conditionally enables RMM debug mode only when tracking is enabled
• CSV output appends across multiple JVM processes for integration test runs

The feature helps developers identify OOM-prone code paths that should be wrapped in retry logic, addressing issue #13672.

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The implementation is well-designed with proper synchronization, gated behind an environment variable, and disabled by default. The changes are additive and don't modify existing behavior when the feature is disabled. Thread-local depth tracking correctly handles nested retry blocks. No logic or syntax issues found.
• No files require special attention

Important Files Changed

Filename	Overview
sql-plugin/src/main/scala/com/nvidia/spark/rapids/AllocationRetryCoverageTracker.scala	New file adds debug tracking for memory allocations outside retry blocks; well-structured with proper synchronization and CSV output
sql-plugin/src/main/scala/com/nvidia/spark/rapids/RmmRapidsRetryIterator.scala	Adds RetryStateTracker with thread-local depth counter to track nested retry blocks; properly integrates with iterator lifecycle
sql-plugin/src/main/scala/com/nvidia/spark/rapids/DeviceMemoryEventHandler.scala	Adds onAllocated callback to check device allocation coverage when RMM debug mode is enabled

Sequence Diagram

sequenceDiagram
    participant App as Application Code
    participant RRI as RmmRapidsRetryIterator
    participant RST as RetryStateTracker
    participant Alloc as Allocator (Host/Device)
    participant ARTC as AllocationRetryCoverageTracker
    participant RMM as RMM (Debug Mode)
    participant CSV as CSV Output File

    Note over App,CSV: Enable tracking: export SPARK_RAPIDS_RETRY_COVERAGE_TRACKING=true

    App->>RRI: withRetryNoSplit(data)(fn)
    activate RRI
    RRI->>RST: enterRetryBlock()
    activate RST
    RST->>RST: increment localRetryBlockDepth
    deactivate RST
    
    RRI->>App: execute fn(data)
    activate App
    App->>Alloc: allocate memory
    activate Alloc
    
    alt Device Memory Allocation
        Alloc->>RMM: allocation occurs
        RMM->>ARTC: onAllocated() callback
    else Host Memory Allocation
        Alloc->>ARTC: checkHostAllocation()
    end
    
    activate ARTC
    ARTC->>RST: isInRetryBlock()
    RST-->>ARTC: return depth > 0 (true)
    Note over ARTC: Allocation is COVERED - no logging
    deactivate ARTC
    
    Alloc-->>App: memory allocated
    deactivate Alloc
    App-->>RRI: result
    deactivate App
    
    RRI->>RST: exitRetryBlock()
    activate RST
    RST->>RST: decrement localRetryBlockDepth
    deactivate RST
    deactivate RRI

    Note over App,CSV: Scenario: Uncovered Allocation

    App->>Alloc: direct allocation (no retry wrapper)
    activate Alloc
    
    alt Device Memory
        Alloc->>RMM: allocation occurs
        RMM->>ARTC: onAllocated() callback
    else Host Memory
        Alloc->>ARTC: checkHostAllocation()
    end
    
    activate ARTC
    ARTC->>RST: isInRetryBlock()
    RST-->>ARTC: return depth == 0 (false)
    ARTC->>ARTC: capture stack trace
    ARTC->>ARTC: filter to Rapids classes
    ARTC->>ARTC: deduplicate by stackKey
    ARTC->>CSV: write uncovered allocation record
    Note over CSV: kind,call_stack
    deactivate ARTC
    
    Alloc-->>App: memory allocated
    deactivate Alloc

Loading

[greptile-apps]

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[thirtiseven]

Resulting files generated by the python script which can be opend by https://www.speedscope.app/

uncovered_allocations_DEVICE.speedscope.json

uncovered_allocations_HOST.speedscope.json

[greptile-apps]

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (2)

1. sql-plugin/src/main/scala/com/nvidia/spark/rapids/AllocationRetryCoverageTracker.scala, line 122-127 (link)

   logic: the detection logic will incorrectly flag allocations inside OOMRetryState class methods as covered, even though OOMRetryState isn't a retry framework method - it's just a helper class that happens to have "Retry" in its name

2. sql-plugin/src/main/scala/com/nvidia/spark/rapids/AllocationRetryCoverageTracker.scala, line 146-153 (link)

   logic: check-then-act race condition - between checking !loggedStacks.contains(stackKey) on line 146 and calling loggedStacks.add(stackKey) on line 147, another thread could add the same key, causing both threads to log duplicate entries to the file. is duplicate logging acceptable in this debugging tool, or should we ensure strictly unique entries?

_{5 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[thirtiseven]

Additional Comments (2)

1. sql-plugin/src/main/scala/com/nvidia/spark/rapids/AllocationRetryCoverageTracker.scala, line 122-127 (link)
   logic: the detection logic will incorrectly flag allocations inside OOMRetryState class methods as covered, even though OOMRetryState isn't a retry framework method - it's just a helper class that happens to have "Retry" in its name
2. sql-plugin/src/main/scala/com/nvidia/spark/rapids/AllocationRetryCoverageTracker.scala, line 146-153 (link)

Should be ok.

logic: check-then-act race condition - between checking !loggedStacks.contains(stackKey) on line 146 and calling loggedStacks.add(stackKey) on line 147, another thread could add the same key, causing both threads to log duplicate entries to the file. is duplicate logging acceptable in this debugging tool, or should we ensure strictly unique entries?

5 files reviewed, 2 comments

Edit Code Review Agent Settings | Greptile

It's acceptable.

[abellina]

Note we need to upgrade copyrights (2026)

I like the idea, had one comment.

[abellina]

should we, instead of pattern matching, use the RetryStateTracker? It presently tracks if the current thread is retrying, but could it also try when it is in a retry block? We could also gate that on whether we have enabled the feature.

[thirtiseven]

Yes that's a better solution, thanks. Updated.

[Copilot]

Pull request overview

This PR introduces a debugging feature to track memory allocations that are not covered by the retry framework (withRetry/withRetryNoSplit blocks), helping developers identify code paths that could potentially cause OOM errors without proper retry handling.

Key changes:

• Adds AllocationRetryCoverageTracker to detect and log uncovered host and device memory allocations
• Enhances RetryStateTracker with depth-based tracking to detect if code is executing inside a retry block
• Integrates tracking hooks into HostAlloc for host memory and DeviceMemoryEventHandler for device memory

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
sql-plugin/src/main/scala/com/nvidia/spark/rapids/AllocationRetryCoverageTracker.scala	New utility class that tracks memory allocations not covered by retry framework, logs findings to CSV file
sql-plugin/src/main/scala/com/nvidia/spark/rapids/RmmRapidsRetryIterator.scala	Adds retry block depth tracking to RetryStateTracker and wraps iterator next() with enter/exit calls
sql-plugin/src/main/scala/com/nvidia/spark/rapids/HostAlloc.scala	Integrates coverage tracker check after successful host memory allocation
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuDeviceManager.scala	Enables RMM debug mode when retry coverage tracking is enabled to get allocation callbacks
sql-plugin/src/main/scala/com/nvidia/spark/rapids/DeviceMemoryEventHandler.scala	Implements onAllocated callback to check device memory allocation coverage
integration_tests/run_pyspark_from_build.sh	Propagates SPARK_RAPIDS_RETRY_COVERAGE_TRACKING environment variable to test executors

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[greptile-apps]

Additional Comments (1)

1. sql-plugin/src/main/scala/com/nvidia/spark/rapids/RmmRapidsRetryIterator.scala, line 585-587 (link)

   logic: clearCurThreadRetrying() called here may conflict with the same call at line 755 in parent class RmmRapidsRetryIterator.next(). The parent's next() already handles clearing retry state after completing all retry attempts. Calling it here in the finally block means it will execute even on successful paths, potentially clearing state prematurely or redundantly.

_{6 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

1. sql-plugin/src/main/scala/com/nvidia/spark/rapids/RmmRapidsRetryIterator.scala, line 586-587 (link)

   logic: clearCurThreadRetrying() should be called before exitRetryBlock() only on exception path. On success path, the parent's next() already manages setCurThreadRetrying state. Calling clear here may incorrectly reset retry state for nested retry blocks.

_{6 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[firestarman]

Overall it looks good to me.
Can you add some example lines in the csv file got from you local tests into the description ?

[firestarman]

Better have more reviews from others

[thirtiseven]

Overall it looks good to me. Can you add some example lines in the csv file got from you local tests into the description ?

Thanks. I added the results CSV file to the description.

[abellina]

I ran this tool locally and I think this is a must. Thanks for doing this.