This is an automated email from the ASF dual-hosted git repository.
gaogaotiantian pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 2092cf6cbeae [SPARK-59773][PYTHON] Report Python worker metrics as JSON
2092cf6cbeae is described below
commit 2092cf6cbeaef96363edfbab7927598d3aba1ea7
Author: Spenser Sun <[email protected]>
AuthorDate: Wed Sep 30 16:37:03 2026 -0700
[SPARK-59773][PYTHON] Report Python worker metrics as JSON
### What changes were proposed in this pull request?
- Reuse `-3` as `METRICS_DATA` and send existing worker timestamps,
execution duration, and spill totals in one length-prefixed UTF-8 JSON object.
- Parse and validate the six fields inside `readWorkerMetrics`, returning
`WorkerMetrics` to the shared JVM handler.
- Preserve SQL metric keys, including `pythonProcessingTime` for
`pythonExecutionDurationMs`, and the existing initialization accounting for
reused workers.
### Why are the changes needed?
Named fields make reports extensible for
[SPARK-59772](https://issues.apache.org/jira/browse/SPARK-59772). Additional
measurements and collection wrappers remain separate work.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- Added JSON-reader checks to `BasePythonRunnerSuite` and driver-side spill
assertions to `pyspark.tests.test_worker`; existing UDF tests exercise worker
communication and SQL metrics.
- Worker and Arrow tests, focused Scala/SQL tests, Ruff, mypy, and
Scalastyle passed. The rebased tree passed the protocol and UDF metrics tests,
including reused-worker initialization; local Scala tests used LZ4 1.11.3
because 1.12.0 was unavailable.
- All 24 `bench_eval_type.ArrowBatchedUDFTimeBench` parameter combinations
passed an ASV `--quick` compatibility run locally.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-6, GPT-5)
Closes #59019 from Spenserrrr/udf-metrics-protocol.
Authored-by: Spenser Sun <[email protected]>
Signed-off-by: Tian Gao <[email protected]>
(cherry picked from commit e85102e700708123923b8fa1309190a42f3634bb)
Signed-off-by: Tian Gao <[email protected]>
---
.../org/apache/spark/api/python/PythonRunner.scala | 76 ++++++++++---
.../spark/api/python/BasePythonRunnerSuite.scala | 122 +++++++++++++++++++++
python/pyspark/serializers.py | 2 +-
python/pyspark/sql/pandas/serializers.py | 2 +-
python/pyspark/tests/test_worker.py | 20 ++++
python/pyspark/worker.py | 43 +++++---
.../sql/execution/python/PythonArrowOutput.scala | 8 +-
.../sql/execution/python/PythonUDFRunner.scala | 4 +-
8 files changed, 242 insertions(+), 35 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
index 4b9f5b3d7f1b..0f58098d79b3 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -29,6 +29,9 @@ import scala.jdk.CollectionConverters._
import scala.util.{Success, Try}
import scala.util.control.NonFatal
+import com.fasterxml.jackson.core.{JsonParser, JsonProcessingException}
+import com.fasterxml.jackson.databind.{DeserializationFeature, ObjectMapper}
+
import org.apache.spark._
import org.apache.spark.api.python.PythonFunction.PythonAccumulator
import org.apache.spark.internal.{Logging, MessageWithContext}
@@ -210,6 +213,53 @@ private[spark] object BasePythonRunner extends Logging {
}
}
+ // Values used to update the existing Python SQL metrics and task spill
metrics.
+ private[python] case class WorkerMetrics(
+ bootTimestampMs: Long,
+ initTimestampMs: Long,
+ finishTimestampMs: Long,
+ pythonExecutionDurationMs: Long,
+ memoryBytesSpilled: Long,
+ diskBytesSpilled: Long)
+
+ private lazy val workerMetricsMapper = new ObjectMapper()
+ .enable(JsonParser.Feature.STRICT_DUPLICATE_DETECTION)
+ .enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS)
+
+ /** Read and validate worker metrics after METRICS_DATA; ignore additional
fields. */
+ private[python] def readWorkerMetrics(stream: DataInputStream):
WorkerMetrics = {
+ val length = stream.readInt()
+ if (length <= 0) {
+ throw new SparkException(s"Invalid Python worker report length: $length")
+ }
+ val json = PythonWorkerUtils.readUTF(length, stream)
+ val report = try {
+ workerMetricsMapper.readTree(json)
+ } catch {
+ case e: JsonProcessingException =>
+ throw new SparkException("Malformed Python worker report JSON", e)
+ }
+ if (report == null || !report.isObject) {
+ throw new SparkException("Expected a Python worker JSON object")
+ }
+
+ def metricValue(name: String): Long = {
+ val value = report.get(name)
+ if (value == null || !value.isIntegralNumber || !value.canConvertToLong)
{
+ throw new SparkException(s"Missing or invalid Python worker metric:
$name")
+ }
+ value.longValue()
+ }
+
+ WorkerMetrics(
+ metricValue("bootTimestampMs"),
+ metricValue("initTimestampMs"),
+ metricValue("finishTimestampMs"),
+ metricValue("pythonExecutionDurationMs"),
+ metricValue("memoryBytesSpilled"),
+ metricValue("diskBytesSpilled"))
+ }
+
/**
* Creates a task identifier string for logging following Spark's standard
format.
* Format: "task <partition>.<attempt> in stage <stageId> (TID
<taskAttemptId>)"
@@ -857,12 +907,12 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
*/
protected def read(): OUT
- protected def handleTimingData(): Unit = {
- // Timing data from worker
- val bootTime = stream.readLong()
- val initTime = stream.readLong()
- val finishTime = stream.readLong()
- val processingTimeMs = stream.readLong()
+ protected def handleMetricsData(): Unit = {
+ val workerMetrics = BasePythonRunner.readWorkerMetrics(stream)
+ val bootTime = workerMetrics.bootTimestampMs
+ val initTime = workerMetrics.initTimestampMs
+ val finishTime = workerMetrics.finishTimestampMs
+ val pythonExecutionDurationMs = workerMetrics.pythonExecutionDurationMs
// A reused Python worker records bootTime before waiting for this task,
so it can precede
// startTime. Use the later timestamp to exclude the worker's idle time
from initialization.
val pythonWorkerInitializationStartTime = math.max(startTime, bootTime)
@@ -888,11 +938,9 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
metrics.get("pythonBootTime").foreach(_.add(boot))
metrics.get("pythonInitTime").foreach(_.add(init))
metrics.get("pythonTotalTime").foreach(_.add(total))
- metrics.get("pythonProcessingTime").foreach(_.add(processingTimeMs))
- val memoryBytesSpilled = stream.readLong()
- val diskBytesSpilled = stream.readLong()
- context.taskMetrics().incMemoryBytesSpilled(memoryBytesSpilled)
- context.taskMetrics().incDiskBytesSpilled(diskBytesSpilled)
+
metrics.get("pythonProcessingTime").foreach(_.add(pythonExecutionDurationMs))
+
context.taskMetrics().incMemoryBytesSpilled(workerMetrics.memoryBytesSpilled)
+ context.taskMetrics().incDiskBytesSpilled(workerMetrics.diskBytesSpilled)
}
protected def handlePythonException(): PythonException = {
@@ -1398,8 +1446,8 @@ private[spark] class PythonRunner(
batchesProcessed += 1
totalDataReceived += length
data
- case SpecialLengths.TIMING_DATA =>
- handleTimingData()
+ case SpecialLengths.METRICS_DATA =>
+ handleMetricsData()
read()
case SpecialLengths.PYTHON_EXCEPTION_THROWN =>
throw handlePythonException()
@@ -1422,7 +1470,7 @@ class PythonWorkerException(msg: String, cause: Throwable)
private[spark] object SpecialLengths {
val END_OF_DATA_SECTION = -1
val PYTHON_EXCEPTION_THROWN = -2
- val TIMING_DATA = -3
+ val METRICS_DATA = -3
val END_OF_STREAM = -4
val NULL = -5
val START_ARROW_STREAM = -6
diff --git
a/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala
b/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala
index 5477e8f7488c..f7675880e01f 100644
---
a/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala
+++
b/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala
@@ -17,6 +17,8 @@
package org.apache.spark.api.python
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream,
DataOutputStream, EOFException}
+import java.nio.charset.StandardCharsets
import java.util.concurrent.{CountDownLatch, TimeUnit}
import org.apache.spark.{SparkException, SparkFunSuite}
@@ -89,4 +91,124 @@ class BasePythonRunnerSuite extends SparkFunSuite {
BasePythonRunner.getWorkerMemoryMb(Some(1L), 2)
}
}
+
+ private val metricsReport =
+ "{\"bootTimestampMs\":1250,\"initTimestampMs\":2500," +
+ "\"finishTimestampMs\":3750,\"pythonExecutionDurationMs\":42," +
+ "\"memoryBytesSpilled\":7,\"diskBytesSpilled\":9}"
+ private val expectedMetrics = BasePythonRunner.WorkerMetrics(1250L, 2500L,
3750L, 42L, 7L, 9L)
+
+ private def framedMetricsStream(json: String, declaredLength: Option[Int] =
None,
+ endMarkers: Boolean = true): DataInputStream = {
+ val bytes = json.getBytes(StandardCharsets.UTF_8)
+ val buffer = new ByteArrayOutputStream()
+ val output = new DataOutputStream(buffer)
+ output.writeInt(declaredLength.getOrElse(bytes.length))
+ output.write(bytes)
+ if (endMarkers) {
+ output.writeInt(SpecialLengths.END_OF_DATA_SECTION)
+ output.writeInt(0) // No accumulator updates.
+ output.writeInt(SpecialLengths.END_OF_STREAM)
+ }
+ new DataInputStream(new ByteArrayInputStream(buffer.toByteArray))
+ }
+
+ private def checkEndMarkers(stream: DataInputStream): Unit = {
+ assert(stream.readInt() == SpecialLengths.END_OF_DATA_SECTION)
+ assert(stream.readInt() == 0)
+ assert(stream.readInt() == SpecialLengths.END_OF_STREAM)
+ assert(stream.read() == -1)
+ }
+
+ test("SPARK-59773: read worker metrics and leave the following sections
aligned") {
+ val withExtra = metricsReport.dropRight(1) +
",\"futureMetric\":{\"value\":[1,2]}}"
+ val stream = framedMetricsStream(withExtra)
+ assert(BasePythonRunner.readWorkerMetrics(stream) == expectedMetrics)
+ checkEndMarkers(stream)
+ }
+
+ test("SPARK-59773: worker metrics require all current fields") {
+ val fields = Seq(
+ "bootTimestampMs" -> "\"bootTimestampMs\":1250,",
+ "initTimestampMs" -> "\"initTimestampMs\":2500,",
+ "finishTimestampMs" -> "\"finishTimestampMs\":3750,",
+ "pythonExecutionDurationMs" -> "\"pythonExecutionDurationMs\":42,",
+ "memoryBytesSpilled" -> "\"memoryBytesSpilled\":7,",
+ "diskBytesSpilled" -> ",\"diskBytesSpilled\":9")
+ fields.foreach { case (name, field) =>
+ val stream = framedMetricsStream(metricsReport.replace(field, ""))
+ val error = intercept[SparkException] {
+ BasePythonRunner.readWorkerMetrics(stream)
+ }
+ assert(error.getMessage.contains(name))
+ checkEndMarkers(stream)
+ }
+ }
+
+ test("SPARK-59773: worker metrics reject non-int64 values") {
+ val invalidValues = Seq("null", "true", "\"42\"", "42.0", "[]", "{}",
+ "9223372036854775808", "-9223372036854775809")
+ val fields = Seq("pythonExecutionDurationMs" -> 42,
+ "memoryBytesSpilled" -> 7, "diskBytesSpilled" -> 9)
+ for {
+ (name, originalValue) <- fields
+ value <- invalidValues
+ } {
+ val field = "\"" + name + "\":"
+ val json = metricsReport.replace(field + originalValue, field + value)
+ val stream = framedMetricsStream(json)
+ val error = intercept[SparkException] {
+ BasePythonRunner.readWorkerMetrics(stream)
+ }
+ assert(error.getMessage.contains(name))
+ checkEndMarkers(stream)
+ }
+ }
+
+ test("SPARK-59773: worker metrics preserve int64 values exactly") {
+ Seq(Long.MinValue, 9007199254740993L, Long.MaxValue).foreach { value =>
+ val json = metricsReport.replace(
+ "\"pythonExecutionDurationMs\":42", "\"pythonExecutionDurationMs\":" +
value)
+ .replace("\"memoryBytesSpilled\":7", "\"memoryBytesSpilled\":" + value)
+ .replace("\"diskBytesSpilled\":9", "\"diskBytesSpilled\":" + value)
+ val stream = framedMetricsStream(json)
+ assert(BasePythonRunner.readWorkerMetrics(stream) ==
+ expectedMetrics.copy(pythonExecutionDurationMs = value,
+ memoryBytesSpilled = value, diskBytesSpilled = value))
+ checkEndMarkers(stream)
+ }
+ }
+
+ test("SPARK-59773: worker metrics reject malformed, duplicate, and
non-object JSON") {
+ val invalid = Seq("[]", "null", "42", metricsReport.dropRight(1),
metricsReport + "{}",
+ metricsReport.dropRight(1) + ",\"bootTimestampMs\":9}",
"{\"future\":1,\"future\":2}")
+ invalid.foreach { json =>
+ val stream = framedMetricsStream(json)
+ intercept[SparkException] {
+ BasePythonRunner.readWorkerMetrics(stream)
+ }
+ checkEndMarkers(stream)
+ }
+ }
+
+ test("SPARK-59773: worker metrics reject invalid lengths and truncated
frames") {
+ Seq(-1, 0).foreach { length =>
+ intercept[SparkException] {
+ BasePythonRunner.readWorkerMetrics(framedMetricsStream(metricsReport,
Some(length)))
+ }
+ }
+
+ val truncated = framedMetricsStream(metricsReport,
+ Some(metricsReport.getBytes(StandardCharsets.UTF_8).length + 1),
endMarkers = false)
+ intercept[EOFException] {
+ BasePythonRunner.readWorkerMetrics(truncated)
+ }
+ }
+
+ test("SPARK-59773: worker metrics accept a larger report with an additional
field") {
+ val json = metricsReport.dropRight(1) + ",\"future\":\"" + ("x" * (70 *
1024)) + "\"}"
+ val stream = framedMetricsStream(json)
+ assert(BasePythonRunner.readWorkerMetrics(stream) == expectedMetrics)
+ checkEndMarkers(stream)
+ }
}
diff --git a/python/pyspark/serializers.py b/python/pyspark/serializers.py
index 4266020ee999..ca6aecbd706d 100644
--- a/python/pyspark/serializers.py
+++ b/python/pyspark/serializers.py
@@ -88,7 +88,7 @@ __all__ = [
class SpecialLengths:
END_OF_DATA_SECTION = -1
PYTHON_EXCEPTION_THROWN = -2
- TIMING_DATA = -3
+ METRICS_DATA = -3
END_OF_STREAM = -4
NULL = -5
START_ARROW_STREAM = -6
diff --git a/python/pyspark/sql/pandas/serializers.py
b/python/pyspark/sql/pandas/serializers.py
index 9eed884dc7b7..6e70e69fa2f4 100644
--- a/python/pyspark/sql/pandas/serializers.py
+++ b/python/pyspark/sql/pandas/serializers.py
@@ -36,7 +36,7 @@ if TYPE_CHECKING:
class SpecialLengths:
END_OF_DATA_SECTION = -1
PYTHON_EXCEPTION_THROWN = -2
- TIMING_DATA = -3
+ METRICS_DATA = -3
END_OF_STREAM = -4
NULL = -5
START_ARROW_STREAM = -6
diff --git a/python/pyspark/tests/test_worker.py
b/python/pyspark/tests/test_worker.py
index 79df02daf9ea..402821a75288 100644
--- a/python/pyspark/tests/test_worker.py
+++ b/python/pyspark/tests/test_worker.py
@@ -145,6 +145,26 @@ class WorkerTests(ReusedPySparkTestCase):
self.assertEqual(sum(range(100)), acc2.value)
self.assertEqual(sum(range(100)), acc1.value)
+ def test_worker_metrics_include_spills(self):
+ accumulator = self.sc.accumulator(0)
+
+ def increment(value):
+ from pyspark import TaskContext, shuffle
+
+ # Supply known spill totals to check their transport into JVM task
metrics.
+ shuffle.MemoryBytesSpilled += 7
+ shuffle.DiskBytesSpilled += 9
+ accumulator.add(1)
+ return TaskContext.get().stageId(), value + 1
+
+ result = self.sc.parallelize([1, 2], 1).map(increment).collect()
+ self.assertEqual([value for _, value in result], [2, 3])
+ self.assertEqual(accumulator.value, 2)
+ self.sc._jsc.sc().listenerBus().waitUntilEmpty(10000)
+ stage = self.sc._jsc.sc().statusStore().lastStageAttempt(result[0][0])
+ self.assertEqual(stage.memoryBytesSpilled(), 14)
+ self.assertEqual(stage.diskBytesSpilled(), 18)
+
def test_reuse_worker_after_take(self):
rdd = self.sc.parallelize(range(100000), 1)
self.assertEqual(0, rdd.first())
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 80ac5ab95e24..7345f45e1ceb 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -73,7 +73,7 @@ from pyspark.serializers import (
CPickleSerializer,
SpecialLengths,
write_int,
- write_long,
+ write_with_length,
)
from pyspark.sql.conversion import (
ArrowBatchTransformer,
@@ -128,12 +128,23 @@ from pyspark.worker_util import (
)
-def report_times(outfile, boot, init, finish, processing_time_ms):
- write_int(SpecialLengths.TIMING_DATA, outfile)
- write_long(int(1000 * boot), outfile)
- write_long(int(1000 * init), outfile)
- write_long(int(1000 * finish), outfile)
- write_long(processing_time_ms, outfile)
+def report_metrics(
+ outfile, boot, init, finish, execution_duration_ms, memory_bytes_spilled,
disk_bytes_spilled
+):
+ """Write the worker metrics as a length-prefixed JSON report after
METRICS_DATA."""
+ payload = json.dumps(
+ {
+ "bootTimestampMs": int(1000 * boot),
+ "initTimestampMs": int(1000 * init),
+ "finishTimestampMs": int(1000 * finish),
+ "pythonExecutionDurationMs": execution_duration_ms,
+ "memoryBytesSpilled": memory_bytes_spilled,
+ "diskBytesSpilled": disk_bytes_spilled,
+ },
+ separators=(",", ":"),
+ ).encode("utf-8")
+ write_int(SpecialLengths.METRICS_DATA, outfile)
+ write_with_length(payload, outfile)
def chain(f, g):
@@ -4542,13 +4553,13 @@ def invoke_udf(message_receiver: SparkMessageReceiver,
outfile: BinaryIO):
serializer._flush_per_batch = True
run_process = pipelined_process if is_pipelined else process
- processing_start_time = time.time()
+ execution_start_time = time.time()
with capture_outputs():
if profiler:
profiler.profile(run_process)
else:
run_process()
- processing_time_ms = int(1000 * (time.time() - processing_start_time))
+ execution_duration_ms = int(1000 * (time.time() -
execution_start_time))
# Cleanup
# Reset task context to None. This is a guard code to avoid residual
context when worker
@@ -4559,15 +4570,21 @@ def invoke_udf(message_receiver: SparkMessageReceiver,
outfile: BinaryIO):
handle_worker_exception(e, outfile)
sys.exit(-1)
finish_time = time.time()
- report_times(outfile, boot_time, init_time, finish_time,
processing_time_ms)
- write_long(shuffle.MemoryBytesSpilled, outfile)
- write_long(shuffle.DiskBytesSpilled, outfile)
+ report_metrics(
+ outfile,
+ boot_time,
+ init_time,
+ finish_time,
+ execution_duration_ms,
+ shuffle.MemoryBytesSpilled,
+ shuffle.DiskBytesSpilled,
+ )
# Mark the beginning of the accumulators section of the output
write_int(SpecialLengths.END_OF_DATA_SECTION, outfile)
send_accumulator_updates(outfile)
- # Check end of stream — raises if the finish signal is not received
correctly.
+ # Check end of stream - raises if the finish signal is not received
correctly.
# Note: this call might fail due to other reasons (e.g. channel broke)
# which will terminate the worker process.
try:
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
index c66e16f2f04c..1cf2664f67b3 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
@@ -91,12 +91,12 @@ private[python] trait PythonArrowOutput[OUT <: AnyRef] {
self: BasePythonRunner[
super.handleEndOfDataSection()
}
- protected override def handleTimingData(): Unit = {
+ protected override def handleMetricsData(): Unit = {
// Get data size from pythonMetrics which is already being tracked
totalDataReceived = pythonMetrics.get("pythonDataReceived")
.map(_.value)
.getOrElse(0L)
- super.handleTimingData()
+ super.handleMetricsData()
}
protected override def read(): OUT = {
@@ -135,8 +135,8 @@ private[python] trait PythonArrowOutput[OUT <: AnyRef] {
self: BasePythonRunner[
}
read()
- case SpecialLengths.TIMING_DATA =>
- handleTimingData()
+ case SpecialLengths.METRICS_DATA =>
+ handleMetricsData()
read()
case SpecialLengths.PYTHON_EXCEPTION_THROWN =>
throw handlePythonException()
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonUDFRunner.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonUDFRunner.scala
index 654fec065820..b66057a50ecc 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonUDFRunner.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonUDFRunner.scala
@@ -129,8 +129,8 @@ abstract class BasePythonUDFRunner(
batchesProcessed += 1
totalDataReceived += length
obj
- case SpecialLengths.TIMING_DATA =>
- handleTimingData()
+ case SpecialLengths.METRICS_DATA =>
+ handleMetricsData()
read()
case SpecialLengths.PYTHON_EXCEPTION_THROWN =>
throw handlePythonException()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
