forked from apache/incubator-gluten
-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[GLUTEN-5438] feat: Dynamically sizing off-heap memory (apache#5439)
## What changes were proposed in this pull request? Today, in Spark we specify the on-heap and off-heap memory sizes as a configuration value read at the beginning of executing a job. With this change, we are exposing a new feature that is enabled with a new spark.gluten.memory.dynamic.offHeap.sizing.enabled setting. When this setting is configured to true, the offheap setting will be ignored in Gluten and we will size the offheap as the same size as the spark.executor.memory setting. We will then proceed to enforcing a total memory quota, calculated by the sum of what memory is committed and in use in the Java heap (calculated with Runtime.getRuntime().totalMemory() - Runtime.GetRuntime().freeMemory()) plus the tracked off-heap memory in TreeMemoryConsumer. When there is an allocation that would tide us over this total amount of committed memory, we will fail the allocation and trigger an OOM. Note that with this change, we perform the "quota check" when an allocation in the native engine is informed to Gluten. In practice, this means that it is possible that the Java codebase can oversubscribe memory as it allocates, which is under the on-heap quota, although there is enough off-heap usage where we should fail the allocation. A test exercising this setting is part of this change. Fixes: apache#5438 ## How was this patch tested? Manual testing with Spark and included test
- Loading branch information
1 parent
022c208
commit db8496b
Showing
8 changed files
with
291 additions
and
17 deletions.
There are no files selected for viewing
60 changes: 60 additions & 0 deletions
60
backends-velox/src/test/scala/org/apache/gluten/execution/DynamicOffHeapSizingTest.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,60 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one or more | ||
| * contributor license agreements. See the NOTICE file distributed with | ||
| * this work for additional information regarding copyright ownership. | ||
| * The ASF licenses this file to You under the Apache License, Version 2.0 | ||
| * (the "License"); you may not use this file except in compliance with | ||
| * the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package org.apache.gluten.execution | ||
|
|
||
| import org.apache.gluten.benchmarks.RandomParquetDataGenerator | ||
| import org.apache.gluten.tags.SkipTestTags | ||
|
|
||
| import org.apache.spark.SparkConf | ||
|
|
||
| @SkipTestTags | ||
| class DynamicOffHeapSizingTest extends VeloxWholeStageTransformerSuite { | ||
| override protected val resourcePath: String = "/tpch-data-parquet-velox" | ||
| override protected val fileFormat: String = "parquet" | ||
|
|
||
| private val dataGenerator = RandomParquetDataGenerator(System.currentTimeMillis()) | ||
| private val outputPath = getClass.getResource("/").getPath + "dynamicoffheapsizing_output.parquet" | ||
| private val AGG_SQL = | ||
| """select f_1, count(DISTINCT f_1) | ||
| |from tbl group | ||
| |group by 1""".stripMargin | ||
|
|
||
| override def beforeAll(): Unit = { | ||
| super.beforeAll() | ||
| } | ||
| override protected def sparkConf: SparkConf = { | ||
| super.sparkConf | ||
| .set("spark.shuffle.manager", "org.apache.spark.shuffle.sort.ColumnarShuffleManager") | ||
| .set("spark.executor.memory", "6GB") | ||
| .set("spark.gluten.memory.dynamic.offHeap.sizing.memory.fraction", "0.8") | ||
| .set("spark.gluten.memory.dynamic.offHeap.sizing.enabled", "true") | ||
| } | ||
|
|
||
| def getRootCause(e: Throwable): Throwable = { | ||
| if (e.getCause == null) { | ||
| return e | ||
| } | ||
| getRootCause(e.getCause) | ||
| } | ||
|
|
||
| test("Dynamic Off-Heap Sizing") { | ||
| System.gc() | ||
| dataGenerator.generateRandomData(spark, Some(outputPath)) | ||
| spark.read.format("parquet").load(outputPath).createOrReplaceTempView("tbl") | ||
| spark.sql(AGG_SQL) | ||
| } | ||
| } |
95 changes: 95 additions & 0 deletions
95
...re/src/main/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTarget.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,95 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one or more | ||
| * contributor license agreements. See the NOTICE file distributed with | ||
| * this work for additional information regarding copyright ownership. | ||
| * The ASF licenses this file to You under the Apache License, Version 2.0 | ||
| * (the "License"); you may not use this file except in compliance with | ||
| * the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package org.apache.gluten.memory.memtarget; | ||
|
|
||
| import org.apache.gluten.GlutenConfig; | ||
|
|
||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| import java.util.concurrent.atomic.AtomicLong; | ||
|
|
||
| public class DynamicOffHeapSizingMemoryTarget implements MemoryTarget { | ||
| private static final Logger LOG = LoggerFactory.getLogger(DynamicOffHeapSizingMemoryTarget.class); | ||
| private final MemoryTarget delegated; | ||
| // When dynamic off-heap sizing is enabled, the off-heap should be sized for the total usable | ||
| // memory, so we can use it as the max memory we will use. | ||
| private static final long MAX_MEMORY_IN_BYTES = GlutenConfig.getConf().offHeapMemorySize(); | ||
| private static final AtomicLong USED_OFFHEAP_BYTES = new AtomicLong(); | ||
|
|
||
| public DynamicOffHeapSizingMemoryTarget(MemoryTarget delegated) { | ||
| this.delegated = delegated; | ||
| } | ||
|
|
||
| @Override | ||
| public long borrow(long size) { | ||
| if (size == 0) { | ||
| return 0; | ||
| } | ||
|
|
||
| long totalMemory = Runtime.getRuntime().totalMemory(); | ||
| long freeMemory = Runtime.getRuntime().freeMemory(); | ||
| long usedOnHeapBytes = (totalMemory - freeMemory); | ||
| long usedOffHeapBytesNow = USED_OFFHEAP_BYTES.get(); | ||
|
|
||
| if (size + usedOffHeapBytesNow + usedOnHeapBytes > MAX_MEMORY_IN_BYTES) { | ||
| LOG.warn( | ||
| String.format( | ||
| "Failing allocation as unified memory is OOM. " | ||
| + "Used Off-heap: %d, Used On-Heap: %d, " | ||
| + "Free On-heap: %d, Total On-heap: %d, " | ||
| + "Max On-heap: %d, Allocation: %d.", | ||
| usedOffHeapBytesNow, | ||
| usedOnHeapBytes, | ||
| freeMemory, | ||
| totalMemory, | ||
| MAX_MEMORY_IN_BYTES, | ||
| size)); | ||
|
|
||
| return 0; | ||
| } | ||
|
|
||
| long reserved = delegated.borrow(size); | ||
|
|
||
| USED_OFFHEAP_BYTES.addAndGet(reserved); | ||
|
|
||
| return reserved; | ||
| } | ||
|
|
||
| @Override | ||
| public long repay(long size) { | ||
| long unreserved = delegated.repay(size); | ||
|
|
||
| USED_OFFHEAP_BYTES.addAndGet(-unreserved); | ||
|
|
||
| return unreserved; | ||
| } | ||
|
|
||
| @Override | ||
| public long usedBytes() { | ||
| return delegated.usedBytes(); | ||
| } | ||
|
|
||
| @Override | ||
| public <T> T accept(MemoryTargetVisitor<T> visitor) { | ||
| return visitor.visit(this); | ||
| } | ||
|
|
||
| public MemoryTarget delegated() { | ||
| return delegated; | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters