From 40c70b2fbc40f267f9cdeee67c3aca746a39a855 Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Wed, 2 Sep 2026 18:26:45 -0700 Subject: [PATCH 1/4] fix(core): account for PySpark memory in resource requests Signed-off-by: Partho Sarthi --- .../main/resources/bootstrap/tuningTable.yaml | 16 ++ .../spark/rapids/tool/tuning/AutoTuner.scala | 77 +++++++ .../config/PySparkMemoryTuningPolicy.scala | 2 + .../tuning/ProfilingAutoTunerSuiteV2.scala | 201 ++++++++++++++++++ .../tuning/QualificationAutoTunerSuite.scala | 2 + 5 files changed, 298 insertions(+) diff --git a/core/src/main/resources/bootstrap/tuningTable.yaml b/core/src/main/resources/bootstrap/tuningTable.yaml index 185c67c45..bad7edc16 100644 --- a/core/src/main/resources/bootstrap/tuningTable.yaml +++ b/core/src/main/resources/bootstrap/tuningTable.yaml @@ -174,6 +174,22 @@ tuningDefinitions: confType: name: byte defaultUnit: MiB + - label: spark.yarn.isPython + description: >- + Marks the application as using Python so YARN includes PySpark memory in executor resource requests. + enabled: false + level: cluster + category: tuning + confType: + name: boolean + - label: spark.kubernetes.resource.type + description: >- + Marks the application resource type as Python so Kubernetes includes PySpark memory in executor pod requests. + enabled: false + level: cluster + category: tuning + confType: + name: string - label: spark.executor.processTreeMetrics.enabled description: Enables executor process-tree metrics collection. enabled: true diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala index fbeaa5a5f..22c580d2a 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala @@ -303,6 +303,23 @@ abstract class AutoTuner( /** Factory method to create the config provider - must be implemented by subclasses */ protected def createConfigProvider(config: Option[TuningConfiguration]): ConfigProviderType + private def detachedTuningDefinition(key: String): TuningEntryDefinition = { + TuningEntryDefinition.getEntryDefinition(key).map { definition => + new TuningEntryDefinition( + definition.label, + definition.description, + definition.enabled, + definition.level, + definition.category, + definition.bootstrapEntry, + definition.defaultSpark, + definition.modifiedBy, + definition.confType, + definition.specialValues, + definition.comments) + }.getOrElse(TuningEntryDefinition(key, enabled = false)) + } + private def createPluginManager(sortAcrossPlugins: Boolean): TuningPluginManager = { TuningPluginManager.builder .withTunerInst(this) @@ -389,6 +406,13 @@ abstract class AutoTuner( tuningDefn.markAsEnable() } + // Keep detached definitions available without resolving sparkMaster during initialization. + // Only the selected cluster manager's entry is enabled after effective properties are ready. + Seq(PySparkMemoryTuningPolicy.YARN_IS_PYTHON_KEY, + PySparkMemoryTuningPolicy.KUBERNETES_RESOURCE_TYPE_KEY).foreach { key => + baseMap.getOrElseUpdate(key, detachedTuningDefinition(key)) + } + // Exclude properties specified in the skip list (Tool specific or // user specified using `exclude` section in target cluster) skippedRecommendations.foreach(baseMap.remove) @@ -791,6 +815,51 @@ abstract class AutoTuner( configProvider.getEntry(PySparkMemoryTuningPolicy.METRICS_POLLING_INTERVAL).getDefault) } + private def pySparkResourceAccountingRecommendation: Option[(String, String)] = { + sparkMaster.collect { + case Yarn => PySparkMemoryTuningPolicy.YARN_IS_PYTHON_KEY -> "true" + case Kubernetes => PySparkMemoryTuningPolicy.KUBERNETES_RESOURCE_TYPE_KEY -> "python" + } + } + + private def hasPositivePySparkMemory: Boolean = { + platform.getPySparkMemoryMB(getPropertyValue).exists(_ > 0L) + } + + private def enablePySparkResourceAccounting(): Unit = { + if (hasPositivePySparkMemory) { + pySparkResourceAccountingRecommendation.foreach { case (key, _) => + finalTuningTable.get(key).foreach(_.markAsEnable()) + } + } + } + + private def isPySparkResourceAccountingOutputEligible: Boolean = { + pySparkResourceAccountingRecommendation.forall { case (key, value) => + if (skippedRecommendations.contains(key)) { + false + } else if (ignoreRecommendation(key)) { + getPropertyValue(key).contains(value) + } else { + getPropertyValue(key).contains(value) || finalTuningTable.get(key).exists { definition => + val prospectiveEntry = TuningEntry.build( + key, getPropertyValue(key), None, Some(definition)) + prospectiveEntry.setRecommendedValue(value) + shouldIncludeInFinalRecommendations(prospectiveEntry) + } + } + } + } + + private def recommendPySparkResourceAccounting(): Unit = { + if (hasPositivePySparkMemory) { + pySparkResourceAccountingRecommendation.foreach { case (key, value) => + finalTuningTable.get(key).foreach(_.markAsEnable()) + appendRecommendation(key, value) + } + } + } + private def ceilToGiBInMB(value: BigDecimal): Option[Long] = { val gibibytes = (value / BigDecimal(1024)).setScale(0, BigDecimal.RoundingMode.CEILING) if (gibibytes.isValidLong) { @@ -906,6 +975,10 @@ abstract class AutoTuner( " Increase the selected source capacity or choose a larger executor layout." case "source-capability" => " executor overhead is supported only for YARN and Kubernetes targets." + case "resource-accounting" => + pySparkResourceAccountingRecommendation.map { case (key, value) => + s" Allow $key=$value so the cluster manager reserves PySpark memory." + }.getOrElse("") case "output-eligibility" => " Remove the selected source and PySpark memory from exclusion or limited-logic " + "lists to allow a full transfer." @@ -936,6 +1009,8 @@ abstract class AutoTuner( PySparkMemoryRebalanceSource.Overhead && !sparkMaster.contains(Yarn) && !sparkMaster.contains(Kubernetes)) { conflict("source-capability") + } else if (!isPySparkResourceAccountingOutputEligible) { + conflict("resource-accounting") } else if (delta > availableDelta) { conflict("capacity") } else { @@ -1443,6 +1518,7 @@ abstract class AutoTuner( } def calculateClusterLevelRecommendations(): Unit = { + enablePySparkResourceAccounting() pySparkMemoryAdjustment.filter(adjustment => adjustment.needsTelemetryRetry && pySparkMemoryTuningPolicy.recommendTelemetryConfigs) .foreach(_ => recommendPySparkTelemetrySettings()) @@ -1563,6 +1639,7 @@ abstract class AutoTuner( configProvider.getEntry("BATCH_SIZE_BYTES").getDefault) appendRecommendation("spark.locality.wait", configProvider.getEntry("LOCALITY_WAIT").getDefault) + recommendPySparkResourceAccounting() } def calculateJobLevelRecommendations(): Unit = { diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/config/PySparkMemoryTuningPolicy.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/config/PySparkMemoryTuningPolicy.scala index d7c484f00..933f145c8 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/config/PySparkMemoryTuningPolicy.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/config/PySparkMemoryTuningPolicy.scala @@ -51,6 +51,8 @@ object PySparkMemoryTuningPolicy { val RECOMMEND_TELEMETRY_CONFIGS = "PYSPARK_MEMORY_RECOMMEND_TELEMETRY_CONFIGS" val REBALANCE_SOURCE = "PYSPARK_MEMORY_REBALANCE_SOURCE" val PYSPARK_MEMORY_KEY = "spark.executor.pyspark.memory" + val YARN_IS_PYTHON_KEY = "spark.yarn.isPython" + val KUBERNETES_RESOURCE_TYPE_KEY = "spark.kubernetes.resource.type" val PROCESS_TREE_METRICS_KEY = "spark.executor.processTreeMetrics.enabled" val STAGE_EXECUTOR_METRICS_KEY = "spark.eventLog.logStageExecutorMetrics" val METRICS_POLLING_INTERVAL_KEY = "spark.executor.metrics.pollingInterval" diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala index 68ea2a230..d6934ad2f 100644 --- a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala @@ -2664,10 +2664,123 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { assert(!values.contains("spark.executor.processTreeMetrics.enabled")) assert(!values.contains("spark.eventLog.logStageExecutorMetrics")) assert(!values.contains("spark.executor.metrics.pollingInterval")) + assert(!values.contains("spark.yarn.isPython")) + assert(!values.contains("spark.kubernetes.resource.type")) val guidance = comments.mkString("\n") assert(!guidance.contains("PySpark memory autotuning needs a telemetry-enabled retry")) } + forAll(Table( + ("sparkMaster", "accountingKey", "accountingValue", "otherKey", "otherValue"), + (Yarn, "spark.yarn.isPython", "true", "spark.kubernetes.resource.type", "python"), + (Kubernetes, "spark.kubernetes.resource.type", "python", "spark.yarn.isPython", "true") + )) { (sparkMaster, accountingKey, accountingValue, otherKey, otherValue) => + test(s"positive PySpark memory recommends $accountingKey without Connect detection") { + val sourceProps = mutable.LinkedHashMap[String, String]( + "spark.executor.cores" -> "8", + "spark.executor.instances" -> "2", + "spark.executor.memory" -> "32g", + "spark.executor.pyspark.memory" -> "4g", + "spark.executor.resource.gpu.amount" -> "1", + "spark.plugins" -> "com.nvidia.spark.SQLPlugin", + otherKey -> otherValue) + val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), sourceProps, + Some(reliableProcessTreeMetricsSparkVersion)) + val platform = PlatformFactory.createInstance(PlatformNames.ONPREM) + configureEventLogClusterInfoForTest(platform, numCores = 8, numWorkers = 2, + sparkProperties = sourceProps.toMap) + + val autoTuner = buildAutoTunerForTests(infoProvider, platform, Some(sparkMaster)) + val (properties, _) = + autoTuner.getRecommendedProperties(showOnlyUpdatedProps = false) + val values = properties.map(property => property.name -> property.getTuneValue()).toMap + + assert(values("spark.executor.pyspark.memory") == "4g") + assert(values(accountingKey) == accountingValue) + assert(!values.contains(otherKey)) + } + } + + forAll(Table( + ("masterName", "sparkMaster"), + ("Standalone", Some(Standalone)), + ("Local", Some(Local)), + ("unspecified", None) + )) { (masterName, sparkMaster) => + test(s"positive PySpark memory emits no cluster-manager marker for $masterName") { + val sourceProps = mutable.LinkedHashMap[String, String]( + "spark.executor.cores" -> "8", + "spark.executor.instances" -> "2", + "spark.executor.memory" -> "32g", + "spark.executor.pyspark.memory" -> "4g", + "spark.executor.resource.gpu.amount" -> "1", + "spark.plugins" -> "com.nvidia.spark.SQLPlugin") + val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), sourceProps, + Some(reliableProcessTreeMetricsSparkVersion)) + val platform = PlatformFactory.createInstance(PlatformNames.ONPREM) + configureEventLogClusterInfoForTest(platform, numCores = 8, numWorkers = 2, + sparkProperties = sourceProps.toMap) + + val autoTuner = buildAutoTunerForTests(infoProvider, platform, sparkMaster) + val (properties, _) = + autoTuner.getRecommendedProperties(showOnlyUpdatedProps = false) + val values = properties.map(property => property.name -> property.getTuneValue()).toMap + + assert(values("spark.executor.pyspark.memory") == "4g") + assert(!values.contains("spark.yarn.isPython")) + assert(!values.contains("spark.kubernetes.resource.type")) + } + } + + test("target-enforced positive PySpark memory enables YARN resource accounting") { + val sourceProps = mutable.LinkedHashMap[String, String]( + "spark.executor.cores" -> "8", + "spark.executor.instances" -> "2", + "spark.executor.memory" -> "32g", + "spark.executor.resource.gpu.amount" -> "1", + "spark.plugins" -> "com.nvidia.spark.SQLPlugin") + val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), sourceProps, + Some(reliableProcessTreeMetricsSparkVersion)) + val targetClusterInfo = ToolTestUtils.buildTargetClusterInfo( + enforcedSparkProperties = Map("spark.executor.pyspark.memory" -> "4g")) + val platform = PlatformFactory.createInstance(PlatformNames.ONPREM, + Some(targetClusterInfo)) + configureEventLogClusterInfoForTest(platform, numCores = 8, numWorkers = 2, + sparkProperties = sourceProps.toMap) + + val autoTuner = buildAutoTunerForTests(infoProvider, platform, Some(Yarn)) + val (properties, _) = autoTuner.getRecommendedProperties(showOnlyUpdatedProps = false) + val values = properties.map(property => property.name -> property.getTuneValue()).toMap + + assert(values("spark.executor.pyspark.memory") == "4g") + assert(values("spark.yarn.isPython") == "true") + assert(!values.contains("spark.kubernetes.resource.type")) + } + + test("non-positive PySpark memory does not surface resource accounting properties") { + val sourceProps = mutable.LinkedHashMap[String, String]( + "spark.executor.cores" -> "8", + "spark.executor.instances" -> "2", + "spark.executor.memory" -> "32g", + "spark.executor.pyspark.memory" -> "0", + "spark.executor.resource.gpu.amount" -> "1", + "spark.plugins" -> "com.nvidia.spark.SQLPlugin", + "spark.yarn.isPython" -> "true", + "spark.kubernetes.resource.type" -> "python") + val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), sourceProps, + Some(reliableProcessTreeMetricsSparkVersion)) + val platform = PlatformFactory.createInstance(PlatformNames.ONPREM) + configureEventLogClusterInfoForTest(platform, numCores = 8, numWorkers = 2, + sparkProperties = sourceProps.toMap) + + val autoTuner = buildAutoTunerForTests(infoProvider, platform, Some(Yarn)) + val (properties, _) = autoTuner.getRecommendedProperties(showOnlyUpdatedProps = false) + val values = properties.map(property => property.name -> property.getTuneValue()).toMap + + assert(!values.contains("spark.yarn.isPython")) + assert(!values.contains("spark.kubernetes.resource.type")) + } + test("OVERHEAD PySpark rebalance on standalone emits no partial recommendation") { val sourceProps = mutable.LinkedHashMap[String, String]( "spark.executor.cores" -> "8", @@ -2810,6 +2923,8 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { assert(values("spark.executor.memory") == "29g") assert(values("spark.executor.pyspark.memory") == "7g") + assert(values("spark.kubernetes.resource.type") == "python") + assert(!values.contains("spark.yarn.isPython")) val coordinatedTotalMB = Seq("spark.executor.memory", "spark.executor.pyspark.memory") .map(key => StringUtils.convertToMB(values(key), Some(ByteUnit.BYTE))).sum assert(coordinatedTotalMB == 36L * 1024L) @@ -2847,6 +2962,8 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { assert(values("spark.executor.memory") == "29g") assert(values("spark.executor.pyspark.memory") == "7g") + assert(values("spark.yarn.isPython") == "true") + assert(!values.contains("spark.kubernetes.resource.type")) val coordinatedTotalMB = Seq("spark.executor.memory", "spark.executor.pyspark.memory") .map(key => StringUtils.convertToMB(values(key), Some(ByteUnit.BYTE))).sum assert(coordinatedTotalMB == 36L * 1024L) @@ -2921,6 +3038,88 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { assert(retryComments.exists(_.contains("telemetry-enabled retry"))) } + forAll(Table( + ("sparkMaster", "accountingKey", "enforcedValue", "isCompatible"), + (Yarn, "spark.yarn.isPython", "false", false), + (Yarn, "spark.yarn.isPython", "true", true), + (Kubernetes, "spark.kubernetes.resource.type", "java", false), + (Kubernetes, "spark.kubernetes.resource.type", "python", true) + )) { (sparkMaster, accountingKey, enforcedValue, isCompatible) => + test(s"PySpark rebalance handles enforced $accountingKey=$enforcedValue") { + val sourceProps = mutable.LinkedHashMap[String, String]( + "spark.executor.cores" -> "8", + "spark.executor.instances" -> "2", + "spark.executor.memory" -> "32g", + "spark.executor.pyspark.memory" -> "4g", + "spark.executor.resource.gpu.amount" -> "1", + "spark.plugins" -> "com.nvidia.spark.SQLPlugin") + val peakBytes = (BigDecimal("5.5") * BigDecimal(1024L * 1024L * 1024L)).toLong + val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), sourceProps, + Some(reliableProcessTreeMetricsSparkVersion), + pySparkMemoryEvidence = Seq(PySparkMemoryEvidence(1, 0, Seq(peakBytes)))) + val targetClusterInfo = ToolTestUtils.buildTargetClusterInfo( + cpuCores = Some(8), memoryGB = Some(128), gpuCount = Some(1), + gpuDevice = Some(GpuTypes.L4.toString), + enforcedSparkProperties = Map(accountingKey -> enforcedValue), + preserveSparkProperties = List("spark.executor.memory")) + val platform = PlatformFactory.createInstance(PlatformNames.ONPREM, + Some(targetClusterInfo)) + configureEventLogClusterInfoForTest(platform, numCores = 8, numWorkers = 2, + sparkProperties = sourceProps.toMap) + + val autoTuner = buildAutoTunerForTests(infoProvider, platform, Some(sparkMaster)) + val (properties, comments) = + autoTuner.getRecommendedProperties(showOnlyUpdatedProps = false) + val values = properties.map(property => property.name -> property.getTuneValue()).toMap + + assert(values("spark.executor.memory") == (if (isCompatible) "29g" else "32g")) + assert(values("spark.executor.pyspark.memory") == (if (isCompatible) "7g" else "4g")) + assert(values(accountingKey) == enforcedValue) + assert(comments.exists(_.comment.contains("constraint=resource-accounting")) == + !isCompatible, comments.mkString("\n")) + } + } + + forAll(Table( + ("sparkMaster", "accountingKey"), + (Yarn, "spark.yarn.isPython"), + (Kubernetes, "spark.kubernetes.resource.type") + )) { (sparkMaster, accountingKey) => + test(s"excluded $accountingKey blocks a partial PySpark rebalance") { + val sourceProps = mutable.LinkedHashMap[String, String]( + "spark.executor.cores" -> "8", + "spark.executor.instances" -> "2", + "spark.executor.memory" -> "32g", + "spark.executor.pyspark.memory" -> "4g", + "spark.executor.resource.gpu.amount" -> "1", + "spark.plugins" -> "com.nvidia.spark.SQLPlugin") + val peakBytes = (BigDecimal("5.5") * BigDecimal(1024L * 1024L * 1024L)).toLong + val infoProvider = getMockInfoProvider(0, Seq(0), Seq(0.0), sourceProps, + Some(reliableProcessTreeMetricsSparkVersion), + pySparkMemoryEvidence = Seq(PySparkMemoryEvidence(1, 0, Seq(peakBytes)))) + val targetClusterInfo = ToolTestUtils.buildTargetClusterInfo( + cpuCores = Some(8), memoryGB = Some(128), gpuCount = Some(1), + gpuDevice = Some(GpuTypes.L4.toString), + preserveSparkProperties = List("spark.executor.memory"), + excludeSparkProperties = List(accountingKey)) + val platform = PlatformFactory.createInstance(PlatformNames.ONPREM, + Some(targetClusterInfo)) + configureEventLogClusterInfoForTest(platform, numCores = 8, numWorkers = 2, + sparkProperties = sourceProps.toMap) + + val autoTuner = buildAutoTunerForTests(infoProvider, platform, Some(sparkMaster)) + val (properties, comments) = + autoTuner.getRecommendedProperties(showOnlyUpdatedProps = false) + val values = properties.map(property => property.name -> property.getTuneValue()).toMap + + assert(values("spark.executor.memory") == "32g") + assert(values("spark.executor.pyspark.memory") == "4g") + assert(!values.contains(accountingKey)) + assert(comments.count(_.comment.contains("constraint=resource-accounting")) == 1, + comments.mkString("\n")) + } + } + test("PySpark conflict comments identify capacity and source capability") { def run(heap: String, master: SparkMaster, overhead: Boolean): (Map[String, String], Seq[String]) = { @@ -3178,6 +3377,8 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { assert(values("spark.executor.metrics.pollingInterval") == "5000") assert(comments.exists(_.comment.contains("telemetry-enabled retry"))) assert(values("spark.executor.pyspark.memory") == "4g") + assert(values("spark.kubernetes.resource.type") == "python") + assert(!values.contains("spark.yarn.isPython")) } } diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerSuite.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerSuite.scala index ab990306a..475b4e16b 100644 --- a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerSuite.scala +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/QualificationAutoTunerSuite.scala @@ -2396,6 +2396,8 @@ class QualificationAutoTunerSuite extends BaseAutoTunerSuite { assert(values("spark.executor.memory") == "29g") assert(values("spark.executor.pyspark.memory") == "7g") + assert(values("spark.kubernetes.resource.type") == "python") + assert(!values.contains("spark.yarn.isPython")) assert(properties.find(_.name == "spark.executor.pyspark.memory").exists(_.isTuned())) assert(!comments.exists(_.comment.contains("constraint=")), comments.mkString("\n")) } From 1cc851934aef97d6af40527ed02e8d427e4ac3db Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Wed, 2 Sep 2026 18:34:27 -0700 Subject: [PATCH 2/4] docs(core): explain PySpark resource accounting helpers Signed-off-by: Partho Sarthi --- .../com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala index 22c580d2a..303d5a84e 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala @@ -303,6 +303,7 @@ abstract class AutoTuner( /** Factory method to create the config provider - must be implemented by subclasses */ protected def createConfigProvider(config: Option[TuningConfiguration]): ConfigProviderType + /** Copy a shared definition because enabling it mutates the entry. */ private def detachedTuningDefinition(key: String): TuningEntryDefinition = { TuningEntryDefinition.getEntryDefinition(key).map { definition => new TuningEntryDefinition( @@ -406,8 +407,6 @@ abstract class AutoTuner( tuningDefn.markAsEnable() } - // Keep detached definitions available without resolving sparkMaster during initialization. - // Only the selected cluster manager's entry is enabled after effective properties are ready. Seq(PySparkMemoryTuningPolicy.YARN_IS_PYTHON_KEY, PySparkMemoryTuningPolicy.KUBERNETES_RESOURCE_TYPE_KEY).foreach { key => baseMap.getOrElseUpdate(key, detachedTuningDefinition(key)) @@ -815,6 +814,10 @@ abstract class AutoTuner( configProvider.getEntry(PySparkMemoryTuningPolicy.METRICS_POLLING_INTERVAL).getDefault) } + /** + * Return the setting that makes the cluster manager include PySpark memory in the executor + * resource request. + */ private def pySparkResourceAccountingRecommendation: Option[(String, String)] = { sparkMaster.collect { case Yarn => PySparkMemoryTuningPolicy.YARN_IS_PYTHON_KEY -> "true" From 2559305d50112ab37060ca1590cc5f6c1d480951 Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Wed, 2 Sep 2026 18:47:27 -0700 Subject: [PATCH 3/4] docs(core): clarify PySpark accounting safeguards Signed-off-by: Partho Sarthi --- .../nvidia/spark/rapids/tool/tuning/AutoTuner.scala | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala index 303d5a84e..e39dbaf30 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala @@ -303,7 +303,11 @@ abstract class AutoTuner( /** Factory method to create the config provider - must be implemented by subclasses */ protected def createConfigProvider(config: Option[TuningConfiguration]): ConfigProviderType - /** Copy a shared definition because enabling it mutates the entry. */ + /** + * Return a per-tuner copy of a shared disabled definition. Resource-accounting entries are + * enabled only after the effective Spark master and properties are known; mutating the shared + * definition would leak that state into other AutoTuner instances. + */ private def detachedTuningDefinition(key: String): TuningEntryDefinition = { TuningEntryDefinition.getEntryDefinition(key).map { definition => new TuningEntryDefinition( @@ -837,6 +841,10 @@ abstract class AutoTuner( } } + /** + * Return whether the required accounting value is already effective or can be emitted. + * Rebalancing must not move memory into PySpark unless the cluster manager will reserve it. + */ private def isPySparkResourceAccountingOutputEligible: Boolean = { pySparkResourceAccountingRecommendation.forall { case (key, value) => if (skippedRecommendations.contains(key)) { From 52bff195628e00843bc35c5c3f1098ec7126ba95 Mon Sep 17 00:00:00 2001 From: Partho Sarthi Date: Wed, 2 Sep 2026 18:51:00 -0700 Subject: [PATCH 4/4] refactor(core): clarify PySpark memory reservation Signed-off-by: Partho Sarthi --- .../spark/rapids/tool/tuning/AutoTuner.scala | 35 +++++++++--------- .../tuning/ProfilingAutoTunerSuiteV2.scala | 36 +++++++++---------- 2 files changed, 36 insertions(+), 35 deletions(-) diff --git a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala index e39dbaf30..05aa95fd2 100644 --- a/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala +++ b/core/src/main/scala/com/nvidia/spark/rapids/tool/tuning/AutoTuner.scala @@ -304,9 +304,10 @@ abstract class AutoTuner( protected def createConfigProvider(config: Option[TuningConfiguration]): ConfigProviderType /** - * Return a per-tuner copy of a shared disabled definition. Resource-accounting entries are - * enabled only after the effective Spark master and properties are known; mutating the shared - * definition would leak that state into other AutoTuner instances. + * Return a private copy of a disabled PySpark memory-reservation definition, such as + * `spark.yarn.isPython` or `spark.kubernetes.resource.type`. Definitions loaded from the tuning + * table are shared, so enabling one in place would also enable it for later AutoTuner instances + * in the same JVM. */ private def detachedTuningDefinition(key: String): TuningEntryDefinition = { TuningEntryDefinition.getEntryDefinition(key).map { definition => @@ -822,7 +823,7 @@ abstract class AutoTuner( * Return the setting that makes the cluster manager include PySpark memory in the executor * resource request. */ - private def pySparkResourceAccountingRecommendation: Option[(String, String)] = { + private def pySparkMemoryReservationConfig: Option[(String, String)] = { sparkMaster.collect { case Yarn => PySparkMemoryTuningPolicy.YARN_IS_PYTHON_KEY -> "true" case Kubernetes => PySparkMemoryTuningPolicy.KUBERNETES_RESOURCE_TYPE_KEY -> "python" @@ -833,20 +834,20 @@ abstract class AutoTuner( platform.getPySparkMemoryMB(getPropertyValue).exists(_ > 0L) } - private def enablePySparkResourceAccounting(): Unit = { + private def enablePySparkMemoryReservationConfig(): Unit = { if (hasPositivePySparkMemory) { - pySparkResourceAccountingRecommendation.foreach { case (key, _) => + pySparkMemoryReservationConfig.foreach { case (key, _) => finalTuningTable.get(key).foreach(_.markAsEnable()) } } } /** - * Return whether the required accounting value is already effective or can be emitted. + * Return whether the required reservation value is already effective or can be emitted. * Rebalancing must not move memory into PySpark unless the cluster manager will reserve it. */ - private def isPySparkResourceAccountingOutputEligible: Boolean = { - pySparkResourceAccountingRecommendation.forall { case (key, value) => + private def isPySparkMemoryReservationConfigOutputEligible: Boolean = { + pySparkMemoryReservationConfig.forall { case (key, value) => if (skippedRecommendations.contains(key)) { false } else if (ignoreRecommendation(key)) { @@ -862,9 +863,9 @@ abstract class AutoTuner( } } - private def recommendPySparkResourceAccounting(): Unit = { + private def recommendPySparkMemoryReservationConfig(): Unit = { if (hasPositivePySparkMemory) { - pySparkResourceAccountingRecommendation.foreach { case (key, value) => + pySparkMemoryReservationConfig.foreach { case (key, value) => finalTuningTable.get(key).foreach(_.markAsEnable()) appendRecommendation(key, value) } @@ -986,8 +987,8 @@ abstract class AutoTuner( " Increase the selected source capacity or choose a larger executor layout." case "source-capability" => " executor overhead is supported only for YARN and Kubernetes targets." - case "resource-accounting" => - pySparkResourceAccountingRecommendation.map { case (key, value) => + case "memory-reservation" => + pySparkMemoryReservationConfig.map { case (key, value) => s" Allow $key=$value so the cluster manager reserves PySpark memory." }.getOrElse("") case "output-eligibility" => @@ -1020,8 +1021,8 @@ abstract class AutoTuner( PySparkMemoryRebalanceSource.Overhead && !sparkMaster.contains(Yarn) && !sparkMaster.contains(Kubernetes)) { conflict("source-capability") - } else if (!isPySparkResourceAccountingOutputEligible) { - conflict("resource-accounting") + } else if (!isPySparkMemoryReservationConfigOutputEligible) { + conflict("memory-reservation") } else if (delta > availableDelta) { conflict("capacity") } else { @@ -1529,7 +1530,7 @@ abstract class AutoTuner( } def calculateClusterLevelRecommendations(): Unit = { - enablePySparkResourceAccounting() + enablePySparkMemoryReservationConfig() pySparkMemoryAdjustment.filter(adjustment => adjustment.needsTelemetryRetry && pySparkMemoryTuningPolicy.recommendTelemetryConfigs) .foreach(_ => recommendPySparkTelemetrySettings()) @@ -1650,7 +1651,7 @@ abstract class AutoTuner( configProvider.getEntry("BATCH_SIZE_BYTES").getDefault) appendRecommendation("spark.locality.wait", configProvider.getEntry("LOCALITY_WAIT").getDefault) - recommendPySparkResourceAccounting() + recommendPySparkMemoryReservationConfig() } def calculateJobLevelRecommendations(): Unit = { diff --git a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala index d6934ad2f..48a745ef1 100644 --- a/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala +++ b/core/src/test/scala/com/nvidia/spark/rapids/tool/tuning/ProfilingAutoTunerSuiteV2.scala @@ -2671,11 +2671,11 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { } forAll(Table( - ("sparkMaster", "accountingKey", "accountingValue", "otherKey", "otherValue"), + ("sparkMaster", "reservationKey", "reservationValue", "otherKey", "otherValue"), (Yarn, "spark.yarn.isPython", "true", "spark.kubernetes.resource.type", "python"), (Kubernetes, "spark.kubernetes.resource.type", "python", "spark.yarn.isPython", "true") - )) { (sparkMaster, accountingKey, accountingValue, otherKey, otherValue) => - test(s"positive PySpark memory recommends $accountingKey without Connect detection") { + )) { (sparkMaster, reservationKey, reservationValue, otherKey, otherValue) => + test(s"positive PySpark memory recommends $reservationKey") { val sourceProps = mutable.LinkedHashMap[String, String]( "spark.executor.cores" -> "8", "spark.executor.instances" -> "2", @@ -2696,7 +2696,7 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { val values = properties.map(property => property.name -> property.getTuneValue()).toMap assert(values("spark.executor.pyspark.memory") == "4g") - assert(values(accountingKey) == accountingValue) + assert(values(reservationKey) == reservationValue) assert(!values.contains(otherKey)) } } @@ -2732,7 +2732,7 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { } } - test("target-enforced positive PySpark memory enables YARN resource accounting") { + test("target-enforced positive PySpark memory enables YARN memory reservation") { val sourceProps = mutable.LinkedHashMap[String, String]( "spark.executor.cores" -> "8", "spark.executor.instances" -> "2", @@ -2757,7 +2757,7 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { assert(!values.contains("spark.kubernetes.resource.type")) } - test("non-positive PySpark memory does not surface resource accounting properties") { + test("non-positive PySpark memory does not surface memory-reservation properties") { val sourceProps = mutable.LinkedHashMap[String, String]( "spark.executor.cores" -> "8", "spark.executor.instances" -> "2", @@ -3039,13 +3039,13 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { } forAll(Table( - ("sparkMaster", "accountingKey", "enforcedValue", "isCompatible"), + ("sparkMaster", "reservationKey", "enforcedValue", "isCompatible"), (Yarn, "spark.yarn.isPython", "false", false), (Yarn, "spark.yarn.isPython", "true", true), (Kubernetes, "spark.kubernetes.resource.type", "java", false), (Kubernetes, "spark.kubernetes.resource.type", "python", true) - )) { (sparkMaster, accountingKey, enforcedValue, isCompatible) => - test(s"PySpark rebalance handles enforced $accountingKey=$enforcedValue") { + )) { (sparkMaster, reservationKey, enforcedValue, isCompatible) => + test(s"PySpark rebalance handles enforced $reservationKey=$enforcedValue") { val sourceProps = mutable.LinkedHashMap[String, String]( "spark.executor.cores" -> "8", "spark.executor.instances" -> "2", @@ -3060,7 +3060,7 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { val targetClusterInfo = ToolTestUtils.buildTargetClusterInfo( cpuCores = Some(8), memoryGB = Some(128), gpuCount = Some(1), gpuDevice = Some(GpuTypes.L4.toString), - enforcedSparkProperties = Map(accountingKey -> enforcedValue), + enforcedSparkProperties = Map(reservationKey -> enforcedValue), preserveSparkProperties = List("spark.executor.memory")) val platform = PlatformFactory.createInstance(PlatformNames.ONPREM, Some(targetClusterInfo)) @@ -3074,18 +3074,18 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { assert(values("spark.executor.memory") == (if (isCompatible) "29g" else "32g")) assert(values("spark.executor.pyspark.memory") == (if (isCompatible) "7g" else "4g")) - assert(values(accountingKey) == enforcedValue) - assert(comments.exists(_.comment.contains("constraint=resource-accounting")) == + assert(values(reservationKey) == enforcedValue) + assert(comments.exists(_.comment.contains("constraint=memory-reservation")) == !isCompatible, comments.mkString("\n")) } } forAll(Table( - ("sparkMaster", "accountingKey"), + ("sparkMaster", "reservationKey"), (Yarn, "spark.yarn.isPython"), (Kubernetes, "spark.kubernetes.resource.type") - )) { (sparkMaster, accountingKey) => - test(s"excluded $accountingKey blocks a partial PySpark rebalance") { + )) { (sparkMaster, reservationKey) => + test(s"excluded $reservationKey blocks a partial PySpark rebalance") { val sourceProps = mutable.LinkedHashMap[String, String]( "spark.executor.cores" -> "8", "spark.executor.instances" -> "2", @@ -3101,7 +3101,7 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { cpuCores = Some(8), memoryGB = Some(128), gpuCount = Some(1), gpuDevice = Some(GpuTypes.L4.toString), preserveSparkProperties = List("spark.executor.memory"), - excludeSparkProperties = List(accountingKey)) + excludeSparkProperties = List(reservationKey)) val platform = PlatformFactory.createInstance(PlatformNames.ONPREM, Some(targetClusterInfo)) configureEventLogClusterInfoForTest(platform, numCores = 8, numWorkers = 2, @@ -3114,8 +3114,8 @@ class ProfilingAutoTunerSuiteV2 extends ProfilingAutoTunerSuiteBase { assert(values("spark.executor.memory") == "32g") assert(values("spark.executor.pyspark.memory") == "4g") - assert(!values.contains(accountingKey)) - assert(comments.count(_.comment.contains("constraint=resource-accounting")) == 1, + assert(!values.contains(reservationKey)) + assert(comments.count(_.comment.contains("constraint=memory-reservation")) == 1, comments.mkString("\n")) } }