From 34fdce029352e224f27382da20f20e5c470835bc Mon Sep 17 00:00:00 2001 From: David Mollitor Date: Thu, 17 Sep 2026 15:41:23 +0000 Subject: [PATCH] [SPARK-59621][CORE] Serialize StatusUpdate manually to avoid Enumeration/BigDecimal overhead `CoarseGrainedClusterMessages.StatusUpdate` is sent from executor to driver on every task state change (`RUNNING` at start, `FINISHED`/`FAILED`/`KILLED` at end) -- roughly two or more per task, on the driver's RPC intake path. An empty-payload `StatusUpdate` Java-serializes to 1713 bytes, and two fields account for ~75% of it: `state` (a Scala `Enumeration` value, 642 bytes) and `taskCpus` (a `BigDecimal`, 638 bytes). A fresh `ObjectOutputStream` per RPC means no class-descriptor caching, so every message pays this in full. A Scala `Enumeration.Value` serializes a reference to its enclosing `Enumeration` object, dragging in the whole `TaskState` object; `scala.math.BigDecimal` drags in `java.math.BigDecimal` + `BigInteger` + `MathContext` + `RoundingMode` descriptors. The actual payload is ~31 bytes. Make `StatusUpdate` `Externalizable` with a compact manual encoding, mirroring `UpdateBlockInfo` in the same message family and reusing `TaskDescription`'s exact wire form for the fractional-CPU `BigDecimal` (`CpuAmount.toDisplayString` / `CpuAmount.normalize`): `state` as one byte, `taskCpus` as its normalized decimal string, `data` via `SerializableBuffer`'s existing channel-based serialization (no extra copy for large results), and `resources` as a size-prefixed nested map. This shrinks the empty-payload message from 1713 bytes to 191 bytes (~9x smaller), cutting steady serialization/GC/bandwidth on a per-task-frequency control message. It is behavior-preserving (all fields round-trip; fractional CPUs exactly) and is not a throughput claim. Verified with a new CoarseGrainedClusterMessagesSuite (round-trips all fields, every TaskState, fractional taskCpus exactly, empty payload, and a size guard). Co-authored-by: Isaac --- .../cluster/CoarseGrainedClusterMessage.scala | 59 +++++++++++-- .../CoarseGrainedClusterMessagesSuite.scala | 87 +++++++++++++++++++ 2 files changed, 137 insertions(+), 9 deletions(-) create mode 100644 core/src/test/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessagesSuite.scala diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala index e1be2af9294b3..02b253681c231 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala @@ -17,13 +17,14 @@ package org.apache.spark.scheduler.cluster +import java.io.{Externalizable, ObjectInput, ObjectOutput} import java.nio.ByteBuffer import org.apache.spark.TaskState.TaskState -import org.apache.spark.resource.{ResourceInformation, ResourceProfile} +import org.apache.spark.resource.{CpuAmount, ResourceInformation, ResourceProfile} import org.apache.spark.rpc.RpcEndpointRef import org.apache.spark.scheduler.{ExecutorLossReason, MiscellaneousProcessDetails} -import org.apache.spark.util.SerializableBuffer +import org.apache.spark.util.{SerializableBuffer, Utils} private[spark] sealed trait CoarseGrainedClusterMessage extends Serializable @@ -78,14 +79,54 @@ private[spark] object CoarseGrainedClusterMessages { case class LaunchedExecutor(executorId: String) extends CoarseGrainedClusterMessage + // Serialized manually (Externalizable): default Java serialization of `state` (a Scala + // Enumeration value) and `taskCpus` (a BigDecimal) each drag in a large class-descriptor graph + // -- together ~1.3KB per message on this hot, per-task-state-change path. case class StatusUpdate( - executorId: String, - taskId: Long, - state: TaskState, - data: SerializableBuffer, - taskCpus: BigDecimal, - resources: Map[String, Map[String, Long]] = Map.empty) - extends CoarseGrainedClusterMessage + var executorId: String, + var taskId: Long, + var state: TaskState, + var data: SerializableBuffer, + var taskCpus: BigDecimal, + var resources: Map[String, Map[String, Long]] = Map.empty) + extends CoarseGrainedClusterMessage with Externalizable { + + def this() = this(null, 0L, null, null, null, Map.empty) // For deserialization only + + override def writeExternal(out: ObjectOutput): Unit = Utils.tryOrIOException { + out.writeUTF(executorId) + out.writeLong(taskId) + out.writeByte(state.id) + // taskCpus as a normalized decimal string (like TaskDescription); round-trips exactly. + out.writeUTF(CpuAmount.toDisplayString(taskCpus)) + // Reuse SerializableBuffer's channel-based write (no extra copy for large results). + out.writeObject(data) + out.writeInt(resources.size) + resources.foreach { case (rName, addressAmounts) => + out.writeUTF(rName) + out.writeInt(addressAmounts.size) + addressAmounts.foreach { case (address, amount) => + out.writeUTF(address) + out.writeLong(amount) + } + } + } + + override def readExternal(in: ObjectInput): Unit = Utils.tryOrIOException { + executorId = in.readUTF() + taskId = in.readLong() + state = org.apache.spark.TaskState(in.readByte().toInt) + taskCpus = CpuAmount.normalize(BigDecimal(in.readUTF())) + data = in.readObject().asInstanceOf[SerializableBuffer] + val numResources = in.readInt() + resources = Iterator.fill(numResources) { + val rName = in.readUTF() + val numAddresses = in.readInt() + val addressAmounts = Iterator.fill(numAddresses)(in.readUTF() -> in.readLong()).toMap + rName -> addressAmounts + }.toMap + } + } object StatusUpdate { /** Alternate factory method that takes a ByteBuffer directly for the data field */ diff --git a/core/src/test/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessagesSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessagesSuite.scala new file mode 100644 index 0000000000000..ac40692362ea6 --- /dev/null +++ b/core/src/test/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessagesSuite.scala @@ -0,0 +1,87 @@ +/* + * 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.spark.scheduler.cluster + +import java.nio.ByteBuffer + +import org.apache.spark.{SparkConf, SparkFunSuite, TaskState} +import org.apache.spark.resource.CpuAmount +import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.StatusUpdate +import org.apache.spark.serializer.JavaSerializer + +class CoarseGrainedClusterMessagesSuite extends SparkFunSuite { + + private val ser = new JavaSerializer(new SparkConf(false)).newInstance() + + private def roundTrip(su: StatusUpdate): StatusUpdate = + ser.deserialize[StatusUpdate](ser.serialize(su)) + + private def statusUpdate( + state: TaskState.TaskState = TaskState.RUNNING, + data: ByteBuffer = ByteBuffer.allocate(0), + taskCpus: BigDecimal = CpuAmount.normalize(BigDecimal(1)), + resources: Map[String, Map[String, Long]] = Map.empty): StatusUpdate = + StatusUpdate("exec-17", 1234567L, state, data, taskCpus, resources) + + test("StatusUpdate round-trips all fields") { + val data = ByteBuffer.wrap(Array[Byte](1, 2, 3, 4, 5)) + val resources = Map( + "gpu" -> Map("0" -> 1L, "1" -> 2L), + "fpga" -> Map("addr-a" -> 7L)) + val cpus = CpuAmount.normalize(BigDecimal("0.5")) + val su = statusUpdate(TaskState.FINISHED, data, cpus, resources) + val rt = roundTrip(su) + + assert(rt.executorId === su.executorId) + assert(rt.taskId === su.taskId) + assert(rt.state === su.state) + assert(rt.taskCpus === cpus) + assert(rt.taskCpus.scale === cpus.scale) + assert(rt.data.value === su.data.value) + assert(rt.resources === resources) + } + + test("StatusUpdate round-trips every TaskState") { + TaskState.values.foreach { s => + assert(roundTrip(statusUpdate(state = s)).state === s) + } + } + + test("SPARK-58192: fractional taskCpus round-trips exactly") { + Seq("1", "0.5", "0.333333333", "2.25").foreach { amount => + val cpus = CpuAmount.normalize(BigDecimal(amount)) + val rt = roundTrip(statusUpdate(taskCpus = cpus)) + assert(rt.taskCpus === cpus, s"value for $amount") + assert(rt.taskCpus.scale === cpus.scale, s"scale for $amount") + } + } + + test("StatusUpdate round-trips with empty data and resources") { + val rt = roundTrip(statusUpdate()) + assert(rt.data.value.remaining() === 0) + assert(rt.resources.isEmpty) + } + + test("StatusUpdate is compact after manual serialization") { + // With default Java serialization, `state` (a Scala Enumeration value) and `taskCpus` (a + // BigDecimal) alone were ~1.3KB and an empty-payload message was ~1.7KB. Manual + // (Externalizable) encoding brings it to ~191 bytes; assert an ample upper bound. + val size = ser.serialize(statusUpdate()).remaining() + assert(size < 512, s"StatusUpdate serialized to $size bytes") + } +}