zhengruifeng commented on PR #58846:
URL: https://github.com/apache/spark/pull/58846#issuecomment-5699915279
### Driver peak-memory benchmark
I ran a synthetic benchmark of the post-aggregation driver step changed by
this PR. This is
not an end-to-end Word2Vec training benchmark: it isolates materializing
`synAgg` and copying
its vectors into `syn0Global` and `syn1Global`.
#### Setup
- Java 17, Spark `local-cluster[2,1,1024]`, 512 MiB per executor, and the
Spark core test
runner's 4 GiB driver heap
- 262,144 aggregated vectors, each containing 384 floats, across 64
partitions
- 384 MiB aggregated-vector payload and 384 MiB resident global model arrays
- Separate forked driver JVM for each mode and three repetitions per mode
- Driver used heap sampled every 5 ms through `MemoryMXBean`
- Full GC before measurement and after the action; the collected result is
kept reachable
through the post-action GC
#### Results
| Mode | Peak heap in three runs | Average peak | Average increase over
baseline | Average live heap after GC | Average action time |
|---|---:|---:|---:|---:|---:|
| `collect()` | 1,117.1 / 1,180.9 / 1,205.9 MiB | 1,168.0 MiB | 677.1 MiB |
894.8 MiB | 1.166 s |
| `runJob` result handler | 760.7 / 905.3 / 873.5 MiB | 846.5 MiB | 355.6
MiB | 492.7 MiB | 1.031 s |
The `runJob` version reduced average peak driver heap by 321.5 MiB (27.5%).
Relative to the
approximately 491 MiB baseline, peak growth fell by 47.5%. After a full GC,
the difference
was 402.1 MiB, showing that `collect()` retained approximately the complete
result payload
while partition results from `runJob` were reclaimable. Both modes produced
checksum
`524286.0` in every run.
#### How to reproduce
Save the following as `/tmp/Word2VecDriverMemoryBenchmark.scala` in a Spark
checkout:
<details>
<summary>Benchmark source</summary>
```scala
import java.lang.management.ManagementFactory
import java.lang.ref.Reference
import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong}
import org.apache.spark.{SparkConf, SparkContext}
object Word2VecDriverMemoryBenchmark {
private def usedHeap: Long =
ManagementFactory.getMemoryMXBean.getHeapMemoryUsage.getUsed
private def forceGc(): Unit = {
System.gc()
System.runFinalization()
Thread.sleep(500)
System.gc()
Thread.sleep(500)
}
private final class PeakHeapSampler extends Thread("peak-heap-sampler") {
private val running = new AtomicBoolean(true)
private val peak = new AtomicLong(usedHeap)
setDaemon(true)
override def run(): Unit = {
while (running.get()) {
peak.accumulateAndGet(usedHeap, Math.max)
Thread.sleep(5)
}
}
def finish(): Long = {
running.set(false)
join()
peak.get()
}
}
private def mib(bytes: Long): Double = bytes.toDouble / 1024 / 1024
def main(args: Array[String]): Unit = {
require(args.length == 4, "mode vectorCount vectorSize partitions")
val mode = args(0)
val vectorCount = args(1).toInt
val vectorSize = args(2).toInt
val partitions = args(3).toInt
require(vectorCount % 2 == 0)
require(mode == "collect" || mode == "runJob")
val conf = new SparkConf(false)
.setAppName(s"Word2VecDriverMemoryBenchmark-$mode")
.setMaster("local-cluster[2,1,1024]")
.set("spark.driver.host", "127.0.0.1")
.set("spark.driver.bindAddress", "127.0.0.1")
.set("spark.executor.memory", "512m")
.set("spark.ui.enabled", "false")
.set("spark.driver.maxResultSize", "0")
val sc = new SparkContext(conf)
sc.setLogLevel("ERROR")
try {
sc.parallelize(Seq(1), 1).count()
val vocabSize = vectorCount / 2
val syn0Global = new Array[Float](vocabSize * vectorSize)
val syn1Global = new Array[Float](vocabSize * vectorSize)
val synAgg = sc.range(0L, vectorCount.toLong, 1L, partitions).map { id
=>
val vec = new Array[Float](vectorSize)
vec(0) = id.toFloat
(id.toInt, vec)
}
def copyPartition(partitionSyn: Array[(Int, Array[Float])]): Unit = {
var i = 0
while (i < partitionSyn.length) {
val (index, vec) = partitionSyn(i)
if (index < vocabSize) {
Array.copy(vec, 0, syn0Global, index * vectorSize, vectorSize)
} else {
Array.copy(vec, 0, syn1Global, (index - vocabSize) * vectorSize,
vectorSize)
}
i += 1
}
}
forceGc()
val baseline = usedHeap
val sampler = new PeakHeapSampler
sampler.start()
val startNanos = System.nanoTime()
val retainedResult: AnyRef = mode match {
case "collect" =>
val result = synAgg.collect()
copyPartition(result)
result
case "runJob" =>
val updateSyn = (_: Int, result: Array[(Int, Array[Float])]) => {
copyPartition(result)
}
sc.runJob(
synAgg,
(iter: Iterator[(Int, Array[Float])]) => iter.toArray,
updateSyn)
null
}
val elapsedNanos = System.nanoTime() - startNanos
val heapAfterAction = usedHeap
forceGc()
val liveHeapAfterGc = usedHeap
val peakHeap = sampler.finish()
val checksum = syn0Global(0) + syn0Global((vocabSize - 1) *
vectorSize) +
syn1Global(0) + syn1Global((vocabSize - 1) * vectorSize)
Reference.reachabilityFence(retainedResult)
val payloadBytes = vectorCount.toLong * vectorSize *
java.lang.Float.BYTES
println(f"BENCH_RESULT mode=$mode payloadMiB=${mib(payloadBytes)}%.1f
" +
f"baselineMiB=${mib(baseline)}%.1f peakMiB=${mib(peakHeap)}%.1f " +
f"peakIncreaseMiB=${mib(peakHeap - baseline)}%.1f " +
f"afterActionMiB=${mib(heapAfterAction)}%.1f " +
f"liveAfterGcMiB=${mib(liveHeapAfterGc)}%.1f " +
f"elapsedSec=${elapsedNanos.toDouble / 1e9}%.3f
checksum=$checksum%.1f")
} finally {
sc.stop()
}
}
}
```
</details>
Run each mode from the repository root:
```bash
JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 build/sbt \
'project core' \
'set Test / unmanagedSources +=
file("/tmp/Word2VecDriverMemoryBenchmark.scala")' \
'Test/runMain Word2VecDriverMemoryBenchmark collect 262144 384 64' \
'Test/runMain Word2VecDriverMemoryBenchmark runJob 262144 384 64'
```
Repeat the command three times to compare peak ranges and averages. The
temporary source is
compiled into the core test classpath and does not modify tracked repository
files.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
