Yicong-Huang commented on code in PR #59111:
URL: https://github.com/apache/spark/pull/59111#discussion_r4212507857
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala:
##########
@@ -99,11 +96,96 @@ trait ExternalUDFExec extends UnaryExecNode {
// failure has already surfaced through the result iterator.
//
// For these UDF sessions, exhausting the data iterator covers execution
and
- // surfaces its errors. close() only cleans up protocol state and releases
or
- // invalidates the worker handle, so its return value is intentionally
ignored.
+ // surfaces its errors. close() cleans up protocol state and releases or
+ // invalidates the worker handle; a clean terminal response also carries
the
+ // execution's final or partial metrics.
//
- taskContext.addTaskCompletionListener[Unit](_ => session.close())
+ taskContext.addTaskCompletionListener[Unit] { _ =>
+ recordTerminalMetrics(session.close())
+ }
f(session)
}
+
+ protected def createUDFWorkerSession(
+ securityScope: Option[WorkerSecurityScope]): WorkerSession = {
+
SparkEnv.get.getExternalUDFDispatcher(workerSpec).createSession(securityScope)
+ }
+
+ private def recordTerminalMetrics(termination: Termination): Unit = {
+ val reported = termination match {
+ case Termination.Finished(response) if response.hasExecutionMetrics =>
+ Some(response.getExecutionMetrics)
+ case Termination.Cancelled(response) if response.hasExecutionMetrics =>
+ Some(response.getExecutionMetrics)
+ case _ => None
+ }
+ reported.foreach(ExternalUDFMetrics.update(metrics, _))
+ }
+}
+
+private[externalUDF] object ExternalUDFMetrics {
+ private val sizeMetrics = Map(
+ "bytesIn" -> "data received by the external UDF worker",
+ "bytesOut" -> "data returned by the external UDF worker")
+
+ private val countMetrics = Map(
+ "rowsIn" -> "rows received by the external UDF worker",
+ "rowsOut" -> "rows returned by the external UDF worker",
+ "batchesIn" -> "batches received by the external UDF worker",
+ "batchesOut" -> "batches returned by the external UDF worker")
+
+ private val timingMetrics = Map(
+ "initWallNanos" -> "external UDF worker initialization time",
+ "processingWallNanos" -> "external UDF worker processing time",
+ "receiveWallNanos" -> "time the external UDF worker waited for requests",
+ "sendWallNanos" -> "time the external UDF worker was blocked sending
responses",
+ "workWallNanos" -> "external UDF worker execution time",
+ "workCpuNanos" -> "external UDF worker CPU time",
+ "finishWallNanos" -> "external UDF worker finish time")
+
+ // Per-stream maxima the worker reports rather than totals. Like peakMemory,
+ // these are timing metrics so the figure to read is the per-task max in the
+ // UI's min/med/max breakdown; the aggregated total (a sum of per-task
maxima)
+ // is not meaningful on its own.
+ private val peakMetrics = Map(
+ "maxWorkWallNanos" -> "longest data unit executed by the external UDF
worker",
+ "maxWorkCpuNanos" -> "most CPU-intensive data unit executed by the
external UDF worker")
+
+ def create(sc: SparkContext): Map[String, SQLMetric] = {
+ sizeMetrics.map { case (name, description) =>
+ name -> SQLMetrics.createSizeMetric(sc, description)
+ } ++ countMetrics.map { case (name, description) =>
+ name -> SQLMetrics.createMetric(sc, description)
+ } ++ (timingMetrics ++ peakMetrics).map { case (name, description) =>
+ name -> SQLMetrics.createNanoTimingMetric(sc, description)
+ }
+ }
+
+ def update(target: Map[String, SQLMetric], reported: ExecutionMetrics): Unit
= {
+ def updateIfPresent(name: String, present: Boolean, value: => Long): Unit
= {
+ if (present) {
+ target(name) += value
+ }
+ }
+
+ updateIfPresent("bytesIn", reported.hasBytesIn, reported.getBytesIn)
+ updateIfPresent("bytesOut", reported.hasBytesOut, reported.getBytesOut)
+ updateIfPresent("rowsIn", reported.hasRowsIn, reported.getRowsIn)
+ updateIfPresent("rowsOut", reported.hasRowsOut, reported.getRowsOut)
+ updateIfPresent("batchesIn", reported.hasBatchesIn, reported.getBatchesIn)
+ updateIfPresent("batchesOut", reported.hasBatchesOut,
reported.getBatchesOut)
+ updateIfPresent("initWallNanos", reported.hasInitWallNanos,
reported.getInitWallNanos)
+ updateIfPresent(
+ "processingWallNanos", reported.hasProcessingWallNanos,
reported.getProcessingWallNanos)
+ updateIfPresent("receiveWallNanos", reported.hasReceiveWallNanos,
reported.getReceiveWallNanos)
+ updateIfPresent("sendWallNanos", reported.hasSendWallNanos,
reported.getSendWallNanos)
+ updateIfPresent("workWallNanos", reported.hasWorkWallNanos,
reported.getWorkWallNanos)
+ updateIfPresent("workCpuNanos", reported.hasWorkCpuNanos,
reported.getWorkCpuNanos)
+ updateIfPresent("finishWallNanos", reported.hasFinishWallNanos,
reported.getFinishWallNanos)
+ updateIfPresent(
+ "maxWorkWallNanos", reported.hasMaxWorkWallNanos,
reported.getMaxWorkWallNanos)
Review Comment:
**Nit (P3):** These two are per-stream maxima, but `update` merges them with
`+=`. When one task runs several sessions (e.g.
`df.select(udf(x)).coalesce(1)`, where `CoalescedRDD` computes every parent
partition in one task), the per-task value becomes the sum of the per-stream
maxima, not the max the comment above says to read. Could these keep the
per-task max instead, like `pythonPeakPickledBatchBytes` in
`BatchEvalPythonExec` (`if (m.isZero || v > m.value) m.set(v)`)?
**Verification:**
- **Regression:** Two terminal snapshots recorded into the same metrics in
one task produce the max of their max_work_* values and the sum of their
cumulative values.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala:
##########
@@ -99,11 +96,96 @@ trait ExternalUDFExec extends UnaryExecNode {
// failure has already surfaced through the result iterator.
//
// For these UDF sessions, exhausting the data iterator covers execution
and
- // surfaces its errors. close() only cleans up protocol state and releases
or
- // invalidates the worker handle, so its return value is intentionally
ignored.
+ // surfaces its errors. close() cleans up protocol state and releases or
+ // invalidates the worker handle; a clean terminal response also carries
the
+ // execution's final or partial metrics.
//
- taskContext.addTaskCompletionListener[Unit](_ => session.close())
+ taskContext.addTaskCompletionListener[Unit] { _ =>
+ recordTerminalMetrics(session.close())
+ }
f(session)
}
+
+ protected def createUDFWorkerSession(
+ securityScope: Option[WorkerSecurityScope]): WorkerSession = {
+
SparkEnv.get.getExternalUDFDispatcher(workerSpec).createSession(securityScope)
+ }
+
+ private def recordTerminalMetrics(termination: Termination): Unit = {
+ val reported = termination match {
+ case Termination.Finished(response) if response.hasExecutionMetrics =>
+ Some(response.getExecutionMetrics)
+ case Termination.Cancelled(response) if response.hasExecutionMetrics =>
+ Some(response.getExecutionMetrics)
+ case _ => None
+ }
+ reported.foreach(ExternalUDFMetrics.update(metrics, _))
+ }
+}
+
+private[externalUDF] object ExternalUDFMetrics {
+ private val sizeMetrics = Map(
+ "bytesIn" -> "data received by the external UDF worker",
+ "bytesOut" -> "data returned by the external UDF worker")
+
+ private val countMetrics = Map(
+ "rowsIn" -> "rows received by the external UDF worker",
+ "rowsOut" -> "rows returned by the external UDF worker",
+ "batchesIn" -> "batches received by the external UDF worker",
+ "batchesOut" -> "batches returned by the external UDF worker")
+
+ private val timingMetrics = Map(
+ "initWallNanos" -> "external UDF worker initialization time",
+ "processingWallNanos" -> "external UDF worker processing time",
+ "receiveWallNanos" -> "time the external UDF worker waited for requests",
+ "sendWallNanos" -> "time the external UDF worker was blocked sending
responses",
+ "workWallNanos" -> "external UDF worker execution time",
+ "workCpuNanos" -> "external UDF worker CPU time",
+ "finishWallNanos" -> "external UDF worker finish time")
+
+ // Per-stream maxima the worker reports rather than totals. Like peakMemory,
Review Comment:
**Nit (P3):** nit: `peakMemory` is a size metric (`createSizeMetric` in
`SortExec`/`HashAggregateExec`), so "Like peakMemory, these are timing metrics"
reads wrong. Maybe: "Like peakMemory, the figure to read is the per-task max in
the UI's min/med/max breakdown; they use nano timing metrics for formatting."
**Verification:**
- **Inspection:** The peak metric comment describes peakMemory only as an
analogue for reading the per-task max.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExecSuite.scala:
##########
@@ -396,6 +434,69 @@ class ExecuteExternalUDFExecSuite extends QueryTest with
SharedSparkSession {
parameters = Map.empty)
}
+ gridTest("external UDF records typed terminal metrics from")(
+ Seq("finish", "cancel")) { terminal =>
+ val reported = ExecutionMetrics.newBuilder()
+ .setBytesIn(1)
+ .setBytesOut(2)
+ .setRowsIn(3)
+ .setRowsOut(4)
+ .setBatchesIn(5)
+ .setBatchesOut(6)
+ .setInitWallNanos(7)
+ .setProcessingWallNanos(8)
+ .setReceiveWallNanos(9)
+ .setSendWallNanos(10)
+ .setWorkWallNanos(11)
+ .setWorkCpuNanos(12)
+ .setFinishWallNanos(13)
+ .setMaxWorkWallNanos(14)
+ .setMaxWorkCpuNanos(15)
+ .build()
+ val execution = testExecution(
+ EchoResponses,
+ rowCount = 5L,
+ terminalMetrics = Some(reported),
+ cancelOnClose = terminal == "cancel")
+ val plan = execution.plan.asInstanceOf[TestExecuteExternalUDFExec]
+
+ terminal match {
+ case "finish" => plan.executeCollect()
+ case "cancel" => plan.executeTake(1)
+ case other => fail(s"unexpected terminal: $other")
+ }
+
+ Seq(
+ "bytesIn", "bytesOut", "rowsIn", "rowsOut", "batchesIn", "batchesOut",
+ "initWallNanos", "processingWallNanos", "receiveWallNanos",
"sendWallNanos",
+ "workWallNanos", "workCpuNanos", "finishWallNanos", "maxWorkWallNanos",
+ "maxWorkCpuNanos"
+ ).zipWithIndex.foreach { case (name, index) =>
+ assert(plan.metrics(name).value === index + 1L)
+ }
+ }
+
+ test("external UDF preserves missing and measured zero terminal metrics") {
+ val reported = ExecutionMetrics.newBuilder()
+ .setBytesIn(0)
+ .build()
+ val execution = testExecution(
+ EchoResponses,
+ rowCount = 1L,
+ terminalMetrics = Some(reported))
+ val plan = execution.plan.asInstanceOf[TestExecuteExternalUDFExec]
+
+ plan.executeCollect()
+
+ assert(plan.metrics("bytesIn").value === 0L)
+ assert(!plan.metrics("bytesIn").isZero)
+ assert(plan.metrics("bytesOut").isZero)
Review Comment:
**Nit (P3):** This only checks `bytesOut` for absence. A miswired `hasX`
guard on any timing/peak field (e.g. `sendWallNanos` using
`hasReceiveWallNanos`) would still pass this and the grid test. Could we assert
`isZero` for every non-count metric other than `bytesIn`?
**Verification:**
- **Behavior:** With only bytes_in reported, every other size, timing and
peak metric stays unreported while counts report raw zero.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala:
##########
@@ -99,11 +96,84 @@ trait ExternalUDFExec extends UnaryExecNode {
// failure has already surfaced through the result iterator.
//
// For these UDF sessions, exhausting the data iterator covers execution
and
- // surfaces its errors. close() only cleans up protocol state and releases
or
- // invalidates the worker handle, so its return value is intentionally
ignored.
+ // surfaces its errors. close() cleans up protocol state and releases or
+ // invalidates the worker handle; a clean terminal response also carries
the
+ // execution's final or partial metrics.
//
- taskContext.addTaskCompletionListener[Unit](_ => session.close())
+ taskContext.addTaskCompletionListener[Unit] { _ =>
+ recordTerminalMetrics(session.close())
+ }
f(session)
}
+
+ protected def createUDFWorkerSession(
+ securityScope: Option[WorkerSecurityScope]): WorkerSession = {
+
SparkEnv.get.getExternalUDFDispatcher(workerSpec).createSession(securityScope)
+ }
+
+ private def recordTerminalMetrics(termination: Termination): Unit = {
+ val reported = termination match {
+ case Termination.Finished(response) if response.hasExecutionMetrics =>
+ Some(response.getExecutionMetrics)
+ case Termination.Cancelled(response) if response.hasExecutionMetrics =>
+ Some(response.getExecutionMetrics)
+ case _ => None
+ }
+ reported.foreach(ExternalUDFMetrics.update(metrics, _))
+ }
+}
+
+private[externalUDF] object ExternalUDFMetrics {
+ private val sizeMetrics = Map(
+ "bytesIn" -> "data received by the external UDF worker",
+ "bytesOut" -> "data returned by the external UDF worker")
+
+ private val countMetrics = Map(
+ "rowsIn" -> "rows received by the external UDF worker",
+ "rowsOut" -> "rows returned by the external UDF worker",
+ "batchesIn" -> "batches received by the external UDF worker",
+ "batchesOut" -> "batches returned by the external UDF worker")
+
+ private val timingMetrics = Map(
+ "initWallNanos" -> "external UDF worker initialization time",
+ "processingWallNanos" -> "external UDF worker processing time",
+ "receiveWallNanos" -> "time the external UDF worker waited for requests",
+ "sendWallNanos" -> "time the external UDF worker was blocked sending
responses",
+ "workWallNanos" -> "external UDF worker execution time",
+ "workCpuNanos" -> "external UDF worker CPU time",
+ "finishWallNanos" -> "external UDF worker finish time")
+
+ def create(sc: SparkContext): Map[String, SQLMetric] = {
+ sizeMetrics.map { case (name, description) =>
+ name -> SQLMetrics.createSizeMetric(sc, description)
+ } ++ countMetrics.map { case (name, description) =>
+ name -> SQLMetrics.createMetric(sc, description)
+ } ++ timingMetrics.map { case (name, description) =>
+ name -> SQLMetrics.createNanoTimingMetric(sc, description)
+ }
+ }
+
+ def update(target: Map[String, SQLMetric], reported: ExecutionMetrics): Unit
= {
+ def updateIfPresent(name: String, present: Boolean, value: => Long): Unit
= {
+ if (present) {
+ target(name) += value
+ }
Review Comment:
Thanks, the two max fields look good. One catch: `update` merges them with
`+=`, so when one task runs several sessions (e.g. `coalesce(1)` after the UDF,
where CoalescedRDD computes all parent partitions in one task), the per-task
value becomes the sum of the per-stream maxima rather than the max. Could these
two keep the per-task max instead, like `pythonPeakPickledBatchBytes` does (`if
(m.isZero || v > m.value) m.set(v)`)?
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4149439049","thread_id":"inline:4149439049","verdict_sha256":"87477ce65a50a26d03c88ebb991591a37555f280ae823d9e3898c044afa9c70c"}
-->
--
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]