You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "ericm-db (via GitHub)" <gi...@apache.org> on 2025/07/29 18:28:06 UTC
[PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
ericm-db opened a new pull request, #51708:
URL: https://github.com/apache/spark/pull/51708
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'common/utils/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
Currently, RocksDB memory is untracked and not included in memory decisions in Spark (particularly when Photon is enabled). 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?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
Avoid OOMs
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as new features, bug fixes, or other behavior changes. Documentation-only updates are not considered user-facing changes.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
No
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
Unit tests
### Was this patch authored or co-authored using generative AI tooling?
<!--
If generative AI tooling has been used in the process of authoring this patch, please include the
phrase: 'Generated-by: ' followed by the name of the tool and its version.
If no, write 'No'.
Please refer to the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248186361
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -242,3 +481,57 @@ object UnifiedMemoryManager {
(usableMemory * memoryFraction).toLong
}
}
+
+/**
+ * Identifier for an unmanaged memory consumer.
+ *
+ * @param componentType The type of component (e.g., "RocksDB", "NativeLibrary")
+ * @param instanceKey A unique key to identify this specific instance of the component.
+ * For shared memory consumers, this should be a common key across
+ * all instances to avoid double counting.
+ */
+case class UnmanagedMemoryConsumerId(
+ componentType: String,
+ instanceKey: String
+ )
+
+/**
+ * Interface for components that consume memory outside of Spark's unified memory management.
+ *
+ * Components implementing this trait can register themselves with the memory manager
+ * to have their memory usage tracked and factored into memory allocation decisions.
+ * This helps prevent OOM errors when unmanaged components use significant memory.
+ *
+ * Examples of unmanaged memory consumers:
+ * - RocksDB state stores in structured streaming
+ * - Native libraries with custom memory allocation
+ * - Off-heap caches managed outside of Spark
+ */
+trait UnmanagedMemoryConsumer {
Review Comment:
Shall we move this into a separate file, `UnmanagedMemoryConsumer.scala`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248188981
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala:
##########
@@ -1339,6 +1382,53 @@ class RocksDB(
private def getDBProperty(property: String): Long = db.getProperty(property).toLong
+ /**
+ * Returns the current memory usage of this RocksDB instance in bytes.
+ * WARNING: This method should only be called from the task thread when
+ * RocksDB is in a safe state.
+ *
+ * This includes memory from all major RocksDB components:
+ * - Table readers (indexes and filters in memory)
+ * - Memtables (write buffers)
+ * - Block cache (cached data blocks)
+ * - Block cache pinned usage (blocks pinned in cache)
+ *
+ * @return Total memory usage in bytes across all tracked components
+ */
+ def getMemoryUsage: Long = {
+
Review Comment:
nit. Let's remove this redundant empty line.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248166153
##########
core/src/main/scala/org/apache/spark/internal/config/package.scala:
##########
@@ -500,6 +500,16 @@ package object config {
.doubleConf
.createWithDefault(0.6)
+ private[spark] val UNMANAGED_MEMORY_POLLING_INTERVAL =
+ ConfigBuilder("spark.memory.unmanagedMemoryPollingInterval")
+ .doc("Interval for polling unmanaged memory users to track their memory usage. " +
+ "Unmanaged memory users are components that manage their own memory outside of " +
+ "Spark's core memory management, such as RocksDB for Streaming State Store. " +
+ "Setting this to 0 disables unmanaged memory polling.")
+ .version("4.1.0")
+ .timeConf(TimeUnit.MILLISECONDS)
+ .createWithDefaultString("1s")
Review Comment:
To be safe and avoid a regression, shall we start with `0` by default, @ericm-db , @anishshri-db, @gatorsmile ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248365535
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala:
##########
@@ -1339,6 +1382,53 @@ class RocksDB(
private def getDBProperty(property: String): Long = db.getProperty(property).toLong
+ /**
+ * Returns the current memory usage of this RocksDB instance in bytes.
+ * WARNING: This method should only be called from the task thread when
+ * RocksDB is in a safe state.
+ *
+ * This includes memory from all major RocksDB components:
+ * - Table readers (indexes and filters in memory)
+ * - Memtables (write buffers)
+ * - Block cache (cached data blocks)
+ * - Block cache pinned usage (blocks pinned in cache)
+ *
+ * @return Total memory usage in bytes across all tracked components
+ */
+ def getMemoryUsage: Long = {
+
+ require(db != null && !db.isClosed, "RocksDB must be open to get memory usage")
+ RocksDB.mainMemorySources.map { memorySource =>
+ getDBProperty(memorySource)
+ }.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 = {
Review Comment:
It's minimal, on the order of ns
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "zhztheplayer (via GitHub)" <gi...@apache.org>.
zhztheplayer commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2245595044
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -56,7 +64,47 @@ private[spark] class UnifiedMemoryManager(
conf,
numCores,
onHeapStorageRegionSize,
- maxHeapMemory - onHeapStorageRegionSize) {
+ maxHeapMemory - onHeapStorageRegionSize) with Logging {
+
+ /**
+ * Unmanaged memory tracking infrastructure.
+ *
+ * Unmanaged memory refers to memory consumed by components that manage their own memory
+ * outside of Spark's unified memory management system. Examples include:
+ * - RocksDB state stores used in structured streaming
+ * - Native libraries with their own memory management
+ * - Off-heap caches managed by unmanaged systems
Review Comment:
Just curious, why don't directly register to storage memory?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248206169
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala:
##########
@@ -314,4 +314,80 @@ class RocksDBStateStoreIntegrationSuite extends StreamTest
assert(changelogVersionsPresent(dirForPartition0) == List(3L, 4L))
assert(snapshotVersionsPresent(dirForPartition0).contains(5L))
}
+
+ // Test with both bounded memory enabled and disabled
+ Seq(true, false).foreach { boundedMemoryEnabled =>
+ test(s"RocksDB memory tracking integration with UnifiedMemoryManager" +
+ s" with boundedMemory=$boundedMemoryEnabled") {
+ withTempDir { dir =>
+ withSQLConf(
+ (SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName),
+ (SQLConf.CHECKPOINT_LOCATION.key -> dir.getCanonicalPath),
+ (SQLConf.SHUFFLE_PARTITIONS.key -> "5"),
+ (SQLConf.STREAMING_MAINTENANCE_INTERVAL.key -> (5 * 60 * 1000).toString),
+ ("spark.memory.unmanagedMemoryPollingInterval" -> "100ms"),
+ ("spark.sql.streaming.stateStore.rocksdb.boundedMemoryUsage" ->
+ boundedMemoryEnabled.toString)) {
+
+ import org.apache.spark.memory.UnifiedMemoryManager
+ import org.apache.spark.sql.streaming.Trigger
+
+ // Use rate stream to ensure continuous state operations that trigger memory updates
+ val query = spark.readStream
+ .format("rate")
+ .option("rowsPerSecond", "10") // Continuous but not overwhelming
+ .load()
+ .selectExpr("value % 100 as key", "value")
+ .groupBy("key")
+ .agg(count("*").as("count"), max("value").as("max_value"))
+ .writeStream
+ .format("console")
+ .outputMode("update")
+ .trigger(Trigger.ProcessingTime(200)) // Regular triggers to ensure state operations
+ .start()
+
+ 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
Review Comment:
This looks a little risky. Can we use `eventually` instead of `Thread.sleep`?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala:
##########
@@ -314,4 +314,80 @@ class RocksDBStateStoreIntegrationSuite extends StreamTest
assert(changelogVersionsPresent(dirForPartition0) == List(3L, 4L))
assert(snapshotVersionsPresent(dirForPartition0).contains(5L))
}
+
+ // Test with both bounded memory enabled and disabled
+ Seq(true, false).foreach { boundedMemoryEnabled =>
+ test(s"RocksDB memory tracking integration with UnifiedMemoryManager" +
+ s" with boundedMemory=$boundedMemoryEnabled") {
+ withTempDir { dir =>
+ withSQLConf(
+ (SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName),
+ (SQLConf.CHECKPOINT_LOCATION.key -> dir.getCanonicalPath),
+ (SQLConf.SHUFFLE_PARTITIONS.key -> "5"),
+ (SQLConf.STREAMING_MAINTENANCE_INTERVAL.key -> (5 * 60 * 1000).toString),
+ ("spark.memory.unmanagedMemoryPollingInterval" -> "100ms"),
+ ("spark.sql.streaming.stateStore.rocksdb.boundedMemoryUsage" ->
+ boundedMemoryEnabled.toString)) {
+
+ import org.apache.spark.memory.UnifiedMemoryManager
+ import org.apache.spark.sql.streaming.Trigger
+
+ // Use rate stream to ensure continuous state operations that trigger memory updates
+ val query = spark.readStream
+ .format("rate")
+ .option("rowsPerSecond", "10") // Continuous but not overwhelming
+ .load()
+ .selectExpr("value % 100 as key", "value")
+ .groupBy("key")
+ .agg(count("*").as("count"), max("value").as("max_value"))
+ .writeStream
+ .format("console")
+ .outputMode("update")
+ .trigger(Trigger.ProcessingTime(200)) // Regular triggers to ensure state operations
+ .start()
+
+ 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
+
+ // Now check for memory tracking - the continuous stream should trigger memory updates
+ var rocksDBMemory = 0L
+ var attempts = 0
+ 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
Review Comment:
ditto.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248173918
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -191,14 +255,189 @@ private[spark] class UnifiedMemoryManager(
}
}
-object UnifiedMemoryManager {
+object UnifiedMemoryManager extends Logging {
// Set aside a fixed amount of memory for non-storage, non-execution purposes.
// This serves a function similar to `spark.memory.fraction`, but guarantees that we reserve
// sufficient memory for the system even for small heaps. E.g. if we have a 1GB JVM, then
// the memory used for execution and storage will be (1024 - 300) * 0.6 = 434MB by default.
private val RESERVED_SYSTEM_MEMORY_BYTES = 300 * 1024 * 1024
+ private val unmanagedMemoryConsumers =
+ new ConcurrentHashMap[UnmanagedMemoryConsumerId, UnmanagedMemoryConsumer]
+
+ // Cached unmanaged memory usage values updated by polling
+ private val unmanagedOnHeapUsed = new AtomicLong(0L)
+ private val unmanagedOffHeapUsed = new AtomicLong(0L)
+
+ // Atomic flag to ensure polling is only started once per JVM
+ private val pollingStarted = new AtomicBoolean(false)
+
+ /**
+ * Register an unmanaged memory consumer to track its memory usage.
+ *
+ * Unmanaged memory consumers are components that manage their own memory outside
+ * of Spark's unified memory management system. By registering, their memory usage
+ * will be periodically polled and factored into Spark's memory allocation decisions.
+ *
+ * @param unmanagedMemoryConsumer The consumer to register for memory tracking
+ */
+ def registerUnmanagedMemoryConsumer(
+ unmanagedMemoryConsumer: UnmanagedMemoryConsumer): Unit = {
Review Comment:
Indentation, @ericm-db ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on PR #51708:
URL: https://github.com/apache/spark/pull/51708#issuecomment-3145086062
@dongjoon-hyun thank you for all the feedback, I will address this
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on PR #51708:
URL: https://github.com/apache/spark/pull/51708#issuecomment-3145124644
@dongjoon-hyun Can you PTAL at this PR: https://github.com/apache/spark/pull/51778
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #51708:
URL: https://github.com/apache/spark/pull/51708#issuecomment-3140194866
> > @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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #51708:
URL: https://github.com/apache/spark/pull/51708#issuecomment-3140181352
> @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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2246616627
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -56,7 +64,47 @@ private[spark] class UnifiedMemoryManager(
conf,
numCores,
onHeapStorageRegionSize,
- maxHeapMemory - onHeapStorageRegionSize) {
+ maxHeapMemory - onHeapStorageRegionSize) with Logging {
+
+ /**
+ * Unmanaged memory tracking infrastructure.
+ *
+ * Unmanaged memory refers to memory consumed by components that manage their own memory
+ * outside of Spark's unified memory management system. Examples include:
+ * - RocksDB state stores used in structured streaming
+ * - Native libraries with their own memory management
+ * - Off-heap caches managed by unmanaged systems
Review Comment:
Yea +1 - probably better to account against UMM in this case
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248191897
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala:
##########
@@ -1339,6 +1382,53 @@ class RocksDB(
private def getDBProperty(property: String): Long = db.getProperty(property).toLong
+ /**
+ * Returns the current memory usage of this RocksDB instance in bytes.
+ * WARNING: This method should only be called from the task thread when
+ * RocksDB is in a safe state.
+ *
+ * This includes memory from all major RocksDB components:
+ * - Table readers (indexes and filters in memory)
+ * - Memtables (write buffers)
+ * - Block cache (cached data blocks)
+ * - Block cache pinned usage (blocks pinned in cache)
+ *
+ * @return Total memory usage in bytes across all tracked components
+ */
+ def getMemoryUsage: Long = {
+
+ require(db != null && !db.isClosed, "RocksDB must be open to get memory usage")
+ RocksDB.mainMemorySources.map { memorySource =>
+ getDBProperty(memorySource)
+ }.sum
Review Comment:
Can we have a one-liner? Maybe, the following style?
```scala
- RocksDB.mainMemorySources.map { memorySource =>
- getDBProperty(memorySource)
- }.sum
+ RocksDB.mainMemorySources.map(getDBProperty).sum
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248166153
##########
core/src/main/scala/org/apache/spark/internal/config/package.scala:
##########
@@ -500,6 +500,16 @@ package object config {
.doubleConf
.createWithDefault(0.6)
+ private[spark] val UNMANAGED_MEMORY_POLLING_INTERVAL =
+ ConfigBuilder("spark.memory.unmanagedMemoryPollingInterval")
+ .doc("Interval for polling unmanaged memory users to track their memory usage. " +
+ "Unmanaged memory users are components that manage their own memory outside of " +
+ "Spark's core memory management, such as RocksDB for Streaming State Store. " +
+ "Setting this to 0 disables unmanaged memory polling.")
+ .version("4.1.0")
+ .timeConf(TimeUnit.MILLISECONDS)
+ .createWithDefaultString("1s")
Review Comment:
To be safe and avoid a regression, shall we start with `0` by default, @ericm-db , @anishshri-db, @gatorsmile ?
For non-rocksdb users, `1s` causes an unnecessary overhead, doesn't it?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248182158
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -242,3 +481,57 @@ object UnifiedMemoryManager {
(usableMemory * memoryFraction).toLong
}
}
+
+/**
+ * Identifier for an unmanaged memory consumer.
+ *
+ * @param componentType The type of component (e.g., "RocksDB", "NativeLibrary")
+ * @param instanceKey A unique key to identify this specific instance of the component.
+ * For shared memory consumers, this should be a common key across
+ * all instances to avoid double counting.
+ */
+case class UnmanagedMemoryConsumerId(
+ componentType: String,
+ instanceKey: String
+ )
Review Comment:
Indentation, @ericm-db ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248203198
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala:
##########
@@ -314,4 +314,80 @@ class RocksDBStateStoreIntegrationSuite extends StreamTest
assert(changelogVersionsPresent(dirForPartition0) == List(3L, 4L))
assert(snapshotVersionsPresent(dirForPartition0).contains(5L))
}
+
+ // Test with both bounded memory enabled and disabled
+ Seq(true, false).foreach { boundedMemoryEnabled =>
+ test(s"RocksDB memory tracking integration with UnifiedMemoryManager" +
+ s" with boundedMemory=$boundedMemoryEnabled") {
+ withTempDir { dir =>
+ withSQLConf(
+ (SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName),
+ (SQLConf.CHECKPOINT_LOCATION.key -> dir.getCanonicalPath),
+ (SQLConf.SHUFFLE_PARTITIONS.key -> "5"),
+ (SQLConf.STREAMING_MAINTENANCE_INTERVAL.key -> (5 * 60 * 1000).toString),
+ ("spark.memory.unmanagedMemoryPollingInterval" -> "100ms"),
+ ("spark.sql.streaming.stateStore.rocksdb.boundedMemoryUsage" ->
+ boundedMemoryEnabled.toString)) {
+
+ import org.apache.spark.memory.UnifiedMemoryManager
+ import org.apache.spark.sql.streaming.Trigger
Review 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "zhztheplayer (via GitHub)" <gi...@apache.org>.
zhztheplayer commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2245595044
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -56,7 +64,47 @@ private[spark] class UnifiedMemoryManager(
conf,
numCores,
onHeapStorageRegionSize,
- maxHeapMemory - onHeapStorageRegionSize) {
+ maxHeapMemory - onHeapStorageRegionSize) with Logging {
+
+ /**
+ * Unmanaged memory tracking infrastructure.
+ *
+ * Unmanaged memory refers to memory consumed by components that manage their own memory
+ * outside of Spark's unified memory management system. Examples include:
+ * - RocksDB state stores used in structured streaming
+ * - Native libraries with their own memory management
+ * - Off-heap caches managed by unmanaged systems
Review Comment:
Just curious, why not directly register to storage memory?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #51708:
URL: https://github.com/apache/spark/pull/51708#issuecomment-3133768313
Besides Structured Streaming, RocksDB is also used in other areas like the Live UI. Don’t they require similar handling?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on PR #51708:
URL: https://github.com/apache/spark/pull/51708#issuecomment-3133790390
@LuciferYang Sorry, I think I'm missing something. Could you elaborate on what the suggestion is here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248194594
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala:
##########
@@ -1339,6 +1382,53 @@ class RocksDB(
private def getDBProperty(property: String): Long = db.getProperty(property).toLong
+ /**
+ * Returns the current memory usage of this RocksDB instance in bytes.
+ * WARNING: This method should only be called from the task thread when
+ * RocksDB is in a safe state.
+ *
+ * This includes memory from all major RocksDB components:
+ * - Table readers (indexes and filters in memory)
+ * - Memtables (write buffers)
+ * - Block cache (cached data blocks)
+ * - Block cache pinned usage (blocks pinned in cache)
+ *
+ * @return Total memory usage in bytes across all tracked components
+ */
+ def getMemoryUsage: Long = {
+
+ require(db != null && !db.isClosed, "RocksDB must be open to get memory usage")
+ RocksDB.mainMemorySources.map { memorySource =>
+ getDBProperty(memorySource)
+ }.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 = {
Review Comment:
This looks like being invoked very frequently in several places. What is the overload of this method?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db closed pull request #51708: [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager
URL: https://github.com/apache/spark/pull/51708
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248194594
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala:
##########
@@ -1339,6 +1382,53 @@ class RocksDB(
private def getDBProperty(property: String): Long = db.getProperty(property).toLong
+ /**
+ * Returns the current memory usage of this RocksDB instance in bytes.
+ * WARNING: This method should only be called from the task thread when
+ * RocksDB is in a safe state.
+ *
+ * This includes memory from all major RocksDB components:
+ * - Table readers (indexes and filters in memory)
+ * - Memtables (write buffers)
+ * - Block cache (cached data blocks)
+ * - Block cache pinned usage (blocks pinned in cache)
+ *
+ * @return Total memory usage in bytes across all tracked components
+ */
+ def getMemoryUsage: Long = {
+
+ require(db != null && !db.isClosed, "RocksDB must be open to get memory usage")
+ RocksDB.mainMemorySources.map { memorySource =>
+ getDBProperty(memorySource)
+ }.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 = {
Review Comment:
This looks like being invoked frequently in several places. What is the overload of this method?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on PR #51708:
URL: https://github.com/apache/spark/pull/51708#issuecomment-3137431607
@LuciferYang - which components are you referring to ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2245784785
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -56,7 +64,47 @@ private[spark] class UnifiedMemoryManager(
conf,
numCores,
onHeapStorageRegionSize,
- maxHeapMemory - onHeapStorageRegionSize) {
+ maxHeapMemory - onHeapStorageRegionSize) with Logging {
+
+ /**
+ * Unmanaged memory tracking infrastructure.
+ *
+ * Unmanaged memory refers to memory consumed by components that manage their own memory
+ * outside of Spark's unified memory management system. Examples include:
+ * - RocksDB state stores used in structured streaming
+ * - Native libraries with their own memory management
+ * - Off-heap caches managed by unmanaged systems
Review Comment:
@zhztheplayer Good question.
Looking at the class structures, I think it makes sense for UMM to be the integration point, as this approach is more extensible.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248166153
##########
core/src/main/scala/org/apache/spark/internal/config/package.scala:
##########
@@ -500,6 +500,16 @@ package object config {
.doubleConf
.createWithDefault(0.6)
+ private[spark] val UNMANAGED_MEMORY_POLLING_INTERVAL =
+ ConfigBuilder("spark.memory.unmanagedMemoryPollingInterval")
+ .doc("Interval for polling unmanaged memory users to track their memory usage. " +
+ "Unmanaged memory users are components that manage their own memory outside of " +
+ "Spark's core memory management, such as RocksDB for Streaming State Store. " +
+ "Setting this to 0 disables unmanaged memory polling.")
+ .version("4.1.0")
+ .timeConf(TimeUnit.MILLISECONDS)
+ .createWithDefaultString("1s")
Review Comment:
To be safe, shall we start with `0` by default, @ericm-db , @anishshri-db, @gatorsmile ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on PR #51708:
URL: https://github.com/apache/spark/pull/51708#issuecomment-3140484283
> I'm just quite curious about this. Even if such issues exist, we can still fix them in a separate PR.
Sure, sounds good
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "ericm-db (via GitHub)" <gi...@apache.org>.
ericm-db commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2246597584
##########
core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala:
##########
@@ -56,7 +64,47 @@ private[spark] class UnifiedMemoryManager(
conf,
numCores,
onHeapStorageRegionSize,
- maxHeapMemory - onHeapStorageRegionSize) {
+ maxHeapMemory - onHeapStorageRegionSize) with Logging {
+
+ /**
+ * Unmanaged memory tracking infrastructure.
+ *
+ * Unmanaged memory refers to memory consumed by components that manage their own memory
+ * outside of Spark's unified memory management system. Examples include:
+ * - RocksDB state stores used in structured streaming
+ * - Native libraries with their own memory management
+ * - Off-heap caches managed by unmanaged systems
Review Comment:
Also, it looks like UMM accounts for both execution and storage memory, which is what we want.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248206684
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreIntegrationSuite.scala:
##########
@@ -314,4 +314,80 @@ class RocksDBStateStoreIntegrationSuite extends StreamTest
assert(changelogVersionsPresent(dirForPartition0) == List(3L, 4L))
assert(snapshotVersionsPresent(dirForPartition0).contains(5L))
}
+
+ // Test with both bounded memory enabled and disabled
+ Seq(true, false).foreach { boundedMemoryEnabled =>
+ test(s"RocksDB memory tracking integration with UnifiedMemoryManager" +
+ s" with boundedMemory=$boundedMemoryEnabled") {
+ withTempDir { dir =>
+ withSQLConf(
+ (SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName),
+ (SQLConf.CHECKPOINT_LOCATION.key -> dir.getCanonicalPath),
+ (SQLConf.SHUFFLE_PARTITIONS.key -> "5"),
+ (SQLConf.STREAMING_MAINTENANCE_INTERVAL.key -> (5 * 60 * 1000).toString),
+ ("spark.memory.unmanagedMemoryPollingInterval" -> "100ms"),
+ ("spark.sql.streaming.stateStore.rocksdb.boundedMemoryUsage" ->
+ boundedMemoryEnabled.toString)) {
+
+ import org.apache.spark.memory.UnifiedMemoryManager
+ import org.apache.spark.sql.streaming.Trigger
+
+ // Use rate stream to ensure continuous state operations that trigger memory updates
+ val query = spark.readStream
+ .format("rate")
+ .option("rowsPerSecond", "10") // Continuous but not overwhelming
+ .load()
+ .selectExpr("value % 100 as key", "value")
+ .groupBy("key")
+ .agg(count("*").as("count"), max("value").as("max_value"))
+ .writeStream
+ .format("console")
+ .outputMode("update")
+ .trigger(Trigger.ProcessingTime(200)) // Regular triggers to ensure state operations
+ .start()
+
+ 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
+
+ // Now check for memory tracking - the continuous stream should trigger memory updates
+ var rocksDBMemory = 0L
+ var attempts = 0
+ 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
+ rocksDBMemory = UnifiedMemoryManager.getMemoryByComponentType("RocksDB")
+ attempts += 1
+
+ if (rocksDBMemory > 0L) {
+ logInfo(s"RocksDB memory detected: $rocksDBMemory bytes " +
+ s"after $attempts attempts with boundedMemory=$boundedMemoryEnabled")
+ }
+ }
+
+ // Verify memory tracking remains stable during continued operation
+ Thread.sleep(2000) // Let stream continue running
Review Comment:
ditto.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53001] Integrate RocksDB Memory Usage with the Unified Memory Manager [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51708:
URL: https://github.com/apache/spark/pull/51708#discussion_r2248198840
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBMemoryManager.scala:
##########
@@ -17,22 +17,93 @@
package org.apache.spark.sql.execution.streaming.state
+import java.util.concurrent.ConcurrentHashMap
+
+import scala.jdk.CollectionConverters._
+
import org.rocksdb._
+import org.apache.spark.SparkEnv
import org.apache.spark.internal.{Logging, MDC}
import org.apache.spark.internal.LogKeys._
+import org.apache.spark.memory.{MemoryMode, UnifiedMemoryManager, UnmanagedMemoryConsumer, UnmanagedMemoryConsumerId}
/**
* Singleton responsible for managing cache and write buffer manager associated with all RocksDB
* state store instances running on a single executor if boundedMemoryUsage is enabled for RocksDB.
* If boundedMemoryUsage is disabled, a new cache object is returned.
+ * This also implements UnmanagedMemoryConsumer to report RocksDB memory usage to Spark's
+ * UnifiedMemoryManager, allowing Spark to account for RocksDB memory when making
+ * memory allocation decisions.
*/
-object RocksDBMemoryManager extends Logging {
+object RocksDBMemoryManager extends Logging with UnmanagedMemoryConsumer{
Review Comment:
I'm a little surprised because Scala linter didn't catch this, `with UnmanagedMemoryConsumer{`. Could you add a space like `with UnmanagedMemoryConsumer {`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org