[SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager #51708
Conversation
|
Besides Structured Streaming, RocksDB is also used in other areas like the Live UI. Don’t they require similar handling? |
|
@LuciferYang Sorry, I think I'm missing something. Could you elaborate on what the suggestion is here? |
|
@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. |
Sure, sounds good |
| "Setting this to 0 disables unmanaged memory polling.") | ||
| .version("4.1.0") | ||
| .timeConf(TimeUnit.MILLISECONDS) | ||
| .createWithDefaultString("1s") |
There was a problem hiding this comment.
To be safe and avoid a regression, shall we start with 0 by default, @ericm-db , @anishshri-db, @gatorsmile ?
| * @param unmanagedMemoryConsumer The consumer to register for memory tracking | ||
| */ | ||
| def registerUnmanagedMemoryConsumer( | ||
| unmanagedMemoryConsumer: UnmanagedMemoryConsumer): Unit = { |
| case class UnmanagedMemoryConsumerId( | ||
| componentType: String, | ||
| instanceKey: String | ||
| ) |
| * - Native libraries with custom memory allocation | ||
| * - Off-heap caches managed outside of Spark | ||
| */ | ||
| trait UnmanagedMemoryConsumer { |
There was a problem hiding this comment.
Shall we move this into a separate file, UnmanagedMemoryConsumer.scala?
| * @return Total memory usage in bytes across all tracked components | ||
| */ | ||
| def getMemoryUsage: Long = { | ||
|
|
There was a problem hiding this comment.
nit. Let's remove this redundant empty line.
| require(db != null && !db.isClosed, "RocksDB must be open to get memory usage") | ||
| RocksDB.mainMemorySources.map { memorySource => | ||
| getDBProperty(memorySource) | ||
| }.sum |
There was a problem hiding this comment.
Can we have a one-liner? Maybe, the following style?
- RocksDB.mainMemorySources.map { memorySource =>
- getDBProperty(memorySource)
- }.sum
+ RocksDB.mainMemorySources.map(getDBProperty).sum| * 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 = { |
There was a problem hiding this comment.
This looks like being invoked frequently in several places. What is the overload of this method?
There was a problem hiding this comment.
It's minimal, on the order of ns
| */ | ||
| object RocksDBMemoryManager extends Logging { | ||
| object RocksDBMemoryManager extends Logging with UnmanagedMemoryConsumer{ |
There was a problem hiding this comment.
I'm a little surprised because Scala linter didn't catch this, with UnmanagedMemoryConsumer{. Could you add a space like with UnmanagedMemoryConsumer {?
| boundedMemoryEnabled.toString)) { | ||
|
|
||
| import org.apache.spark.memory.UnifiedMemoryManager | ||
| import org.apache.spark.sql.streaming.Trigger |
There was a problem hiding this comment.
Do you have special reasons why we have these import statements in the test code body? Otherwise, please move this to the file header.
|
|
||
| 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 |
There was a problem hiding this comment.
This looks a little risky. Can we use eventually instead of Thread.sleep?
| 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 |
| } | ||
|
|
||
| // Verify memory tracking remains stable during continued operation | ||
| Thread.sleep(2000) // Let stream continue running |
|
@dongjoon-hyun thank you for all the feedback, I will address this |
|
@dongjoon-hyun Can you PTAL at this PR: #51778 |
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