[SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager #51708

ericm-db · 2025-07-29T18:28:04Z

What changes were proposed in this pull request?

Currently, RocksDB memory is untracked and not included in memory decisions in Spark. We want to factor the RocksDB memory usage into memory allocations so we don't hit OOMs. This change introduces a background memory polling thread from the MemoryManager that queries RocksDB memory every X seconds (configurable via SQLConf).

Why are the changes needed?

This helps us avoid OOMs when RocksDB is used as the StateStoreProvider by taking other Spark allocations into account.

Does this PR introduce any user-facing change?

No

How was this patch tested?

Unit tests

Was this patch authored or co-authored using generative AI tooling?

No

…Manager

LuciferYang · 2025-07-29T19:25:49Z

Besides Structured Streaming, RocksDB is also used in other areas like the Live UI. Don’t they require similar handling?

ericm-db · 2025-07-29T19:35:11Z

@LuciferYang Sorry, I think I'm missing something. Could you elaborate on what the suggestion is here?

anishshri-db · 2025-07-30T18:34:13Z

@LuciferYang - which components are you referring to ?

LuciferYang · 2025-07-31T14:21:56Z

@LuciferYang - which components are you referring to ?

For instance, Spark Live UI can also utilize RocksDB as its storage backend, and I'm not sure if it encounters similar issues as well.

LuciferYang · 2025-07-31T14:25:36Z

@LuciferYang - which components are you referring to ?

For instance, Spark Live UI can also utilize RocksDB as its storage backend, and I'm not sure if it encounters similar issues as well.

I'm just quite curious about this. Even if such issues exist, we can still fix them in a separate PR.

core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala

ericm-db · 2025-07-31T15:52:18Z

I'm just quite curious about this. Even if such issues exist, we can still fix them in a separate PR.

Sure, sounds good

dongjoon-hyun · 2025-08-01T14:42:58Z

core/src/main/scala/org/apache/spark/internal/config/package.scala

+        "Setting this to 0 disables unmanaged memory polling.")
+      .version("4.1.0")
+      .timeConf(TimeUnit.MILLISECONDS)
+      .createWithDefaultString("1s")


To be safe and avoid a regression, shall we start with 0 by default, @ericm-db , @anishshri-db, @gatorsmile ?

dongjoon-hyun · 2025-08-01T14:46:48Z

core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala

+   * @param unmanagedMemoryConsumer The consumer to register for memory tracking
+   */
+  def registerUnmanagedMemoryConsumer(
+                                       unmanagedMemoryConsumer: UnmanagedMemoryConsumer): Unit = {


Indentation, @ericm-db ?

dongjoon-hyun · 2025-08-01T14:50:38Z

core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala

+case class UnmanagedMemoryConsumerId(
+                                      componentType: String,
+                                      instanceKey: String
+                                    )


Indentation, @ericm-db ?

dongjoon-hyun · 2025-08-01T14:52:36Z

core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala

+ * - Native libraries with custom memory allocation
+ * - Off-heap caches managed outside of Spark
+ */
+trait UnmanagedMemoryConsumer {


Shall we move this into a separate file, UnmanagedMemoryConsumer.scala?

dongjoon-hyun · 2025-08-01T14:53:52Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala

+   * @return Total memory usage in bytes across all tracked components
+   */
+  def getMemoryUsage: Long = {
+


nit. Let's remove this redundant empty line.

dongjoon-hyun · 2025-08-01T14:55:14Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala

+    require(db != null && !db.isClosed, "RocksDB must be open to get memory usage")
+    RocksDB.mainMemorySources.map { memorySource =>
+      getDBProperty(memorySource)
+    }.sum


Can we have a one-liner? Maybe, the following style?

- RocksDB.mainMemorySources.map { memorySource => - getDBProperty(memorySource) - }.sum + RocksDB.mainMemorySources.map(getDBProperty).sum

dongjoon-hyun · 2025-08-01T14:56:23Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala

+   * Updates the cached memory usage if enough time has passed.
+   * This is called from task thread operations, so it's already thread-safe.
+   */
+  def updateMemoryUsageIfNeeded(): Unit = {


This looks like being invoked frequently in several places. What is the overload of this method?

It's minimal, on the order of ns

dongjoon-hyun · 2025-08-01T14:58:22Z

...ore/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBMemoryManager.scala

+ * memory allocation decisions.
 */
-object RocksDBMemoryManager extends Logging {
+object RocksDBMemoryManager extends Logging with UnmanagedMemoryConsumer{


I'm a little surprised because Scala linter didn't catch this, with UnmanagedMemoryConsumer{. Could you add a space like with UnmanagedMemoryConsumer {?

dongjoon-hyun · 2025-08-01T15:00:28Z

...scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala

+            boundedMemoryEnabled.toString)) {
+
+          import org.apache.spark.memory.UnifiedMemoryManager
+          import org.apache.spark.sql.streaming.Trigger


Do you have special reasons why we have these import statements in the test code body? Otherwise, please move this to the file header.

dongjoon-hyun · 2025-08-01T15:01:55Z

...scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala

+
+          try {
+            // Let the stream run to establish RocksDB instances and generate state operations
+            Thread.sleep(2000) // 2 seconds should be enough for several processing cycles


This looks a little risky. Can we use eventually instead of Thread.sleep?

dongjoon-hyun · 2025-08-01T15:02:02Z

...scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala

+            val maxAttempts = 15 // 15 attempts with 1-second intervals = 15 seconds max
+
+            while (rocksDBMemory <= 0L && attempts < maxAttempts) {
+              Thread.sleep(1000) // Wait between checks to allow memory updates


dongjoon-hyun · 2025-08-01T15:02:09Z

...scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala

+            }
+
+            // Verify memory tracking remains stable during continued operation
+            Thread.sleep(2000) // Let stream continue running


ericm-db · 2025-08-01T16:12:06Z

@dongjoon-hyun thank you for all the feedback, I will address this

ericm-db · 2025-08-01T16:27:05Z

@dongjoon-hyun Can you PTAL at this PR: #51778

### What changes were proposed in this pull request? This PR aims to document newly added `core` module configurations as a part of Apache Spark 4.1.0 preparation. ### Why are the changes needed? To help the users use new features easily. - #47856 - #51130 - #51163 - #51604 - #51630 - #51708 - #51885 - #52091 - #52382 ### Does this PR introduce _any_ user-facing change? No behavior change because this is a documentation update. ### How was this patch tested? Manual review. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #52626 from dongjoon-hyun/SPARK-53926. Authored-by: Dongjoon Hyun <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>

…efore execution memory OOM ### What changes were proposed in this pull request? We have a log before OOM for off-heap memory allocation. Before the change, the log is: > 25/08/05 16:44:32 INFO TaskMemoryManager: 100 bytes of memory are used for execution and 100 bytes of memory are used for storage After: > 25/08/05 16:44:32 INFO TaskMemoryManager: 100 bytes of memory are used for execution and 100 bytes of memory are used for storage and 500 bytes of memory are used but unmanaged ### Why are the changes needed? Following #51708, to allow user to know the reason if the unmanaged memory causes OOM. ### Does this PR introduce _any_ user-facing change? Only changes a log message. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #51848 from zhztheplayer/wip-53128. Authored-by: Hongze Zhang <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>

…efore execution memory OOM ### What changes were proposed in this pull request? We have a log before OOM for off-heap memory allocation. Before the change, the log is: > 25/08/05 16:44:32 INFO TaskMemoryManager: 100 bytes of memory are used for execution and 100 bytes of memory are used for storage After: > 25/08/05 16:44:32 INFO TaskMemoryManager: 100 bytes of memory are used for execution and 100 bytes of memory are used for storage and 500 bytes of memory are used but unmanaged ### Why are the changes needed? Following #51708, to allow user to know the reason if the unmanaged memory causes OOM. ### Does this PR introduce _any_ user-facing change? Only changes a log message. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #51848 from zhztheplayer/wip-53128. Authored-by: Hongze Zhang <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]> (cherry picked from commit c4ad381) Signed-off-by: Dongjoon Hyun <[email protected]>

### What changes were proposed in this pull request? This PR aims to document newly added `core` module configurations as a part of Apache Spark 4.1.0 preparation. ### Why are the changes needed? To help the users use new features easily. - apache#47856 - apache#51130 - apache#51163 - apache#51604 - apache#51630 - apache#51708 - apache#51885 - apache#52091 - apache#52382 ### Does this PR introduce _any_ user-facing change? No behavior change because this is a documentation update. ### How was this patch tested? Manual review. ### Was this patch authored or co-authored using generative AI tooling? No. Closes apache#52626 from dongjoon-hyun/SPARK-53926. Authored-by: Dongjoon Hyun <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>

…efore execution memory OOM ### What changes were proposed in this pull request? We have a log before OOM for off-heap memory allocation. Before the change, the log is: > 25/08/05 16:44:32 INFO TaskMemoryManager: 100 bytes of memory are used for execution and 100 bytes of memory are used for storage After: > 25/08/05 16:44:32 INFO TaskMemoryManager: 100 bytes of memory are used for execution and 100 bytes of memory are used for storage and 500 bytes of memory are used but unmanaged ### Why are the changes needed? Following apache#51708, to allow user to know the reason if the unmanaged memory causes OOM. ### Does this PR introduce _any_ user-facing change? Only changes a log message. ### How was this patch tested? Existing tests. ### Was this patch authored or co-authored using generative AI tooling? No. Closes apache#51848 from zhztheplayer/wip-53128. Authored-by: Hongze Zhang <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>

[SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory …

5e855df

…Manager

github-actions bot added SQL STRUCTURED STREAMING CORE labels Jul 29, 2025

compilation errors

730c854

trigger

cdb5649

zhztheplayer reviewed Jul 31, 2025

View reviewed changes

core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala Show resolved Hide resolved

anishshri-db approved these changes Aug 1, 2025

View reviewed changes

anishshri-db closed this in a75e1aa Aug 1, 2025

dongjoon-hyun reviewed Aug 1, 2025

View reviewed changes

ericm-db mentioned this pull request Aug 1, 2025

[SPARK-53001][CORE][SQL][FOLLOW-UP] Disable spark.memory.unmanagedMemoryPollingInterval by default #51778

Closed

zhztheplayer mentioned this pull request Aug 5, 2025

[SPARK-53128][CORE] Include unmanaged memory bytes in the usage log before execution memory OOM #51848

Closed

dongjoon-hyun mentioned this pull request Oct 15, 2025

[SPARK-53926][DOCS] Document newly added core module configurations #52626

Closed

[SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager #51708

[SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager #51708

Uh oh!

Conversation

ericm-db commented Jul 29, 2025 • edited by gatorsmile Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

LuciferYang commented Jul 29, 2025

Uh oh!

ericm-db commented Jul 29, 2025

Uh oh!

anishshri-db commented Jul 30, 2025

Uh oh!

LuciferYang commented Jul 31, 2025

Uh oh!

LuciferYang commented Jul 31, 2025

Uh oh!

Uh oh!

ericm-db commented Jul 31, 2025

Uh oh!

dongjoon-hyun Aug 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

dongjoon-hyun Aug 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

ericm-db commented Aug 1, 2025

Uh oh!

ericm-db commented Aug 1, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ericm-db commented Jul 29, 2025 •

edited by gatorsmile

Loading

dongjoon-hyun Aug 1, 2025 •

edited

Loading

dongjoon-hyun Aug 1, 2025 •

edited

Loading