This is an automated email from the ASF dual-hosted git repository.
uros-b 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 81726013ac15 [SPARK-59784][SQL][TEST] Fix racy waits for SQL execution
metrics in tests
81726013ac15 is described below
commit 81726013ac158d6e05a80a0ac08a9d10c1ed29dc
Author: Sepuri Sai Krishna <[email protected]>
AuthorDate: Tue Oct 6 12:42:21 2026 +0000
[SPARK-59784][SQL][TEST] Fix racy waits for SQL execution metrics in tests
### What changes were proposed in this pull request?
This PR fixes a racy wait for SQL execution metrics that is shared by
several tests, including the flaky `driver side SQL metrics` test in
`SQLAppStatusListenerSuite`.
It adds a helper `SharedSparkSession.runAndWaitForExecution(func)`, which
drains the listener bus, records the id of the latest execution in the status
store, runs `func`, then waits (with `eventually`) until an execution with a
larger id has its metrics computed, and returns that id. It is used in:
- `SharedSparkSession.runAndFetchMetrics` (used by
`InMemoryTableMetricSuite` and `KeyGroupedPartitioningSuite`)
- `SQLAppStatusListenerSuite`: `driver side SQL metrics`, `SPARK-34338:
Report metrics from Datasource v2 scan`, `SPARK-39635: Report driver metrics
from Datasource v2 scan`, `SPARK-36030: Report metrics from Datasource v2 write`
- `PythonDataSourceSuite`: `SPARK-46424: Support Python metrics`
- `PythonUDFSuite`: `SPARK-34265: Instrument Python UDF execution using SQL
Metrics` and `SPARK-55046: pythonProcessingTime metric is available for Python
UDFs`
It also removes an unused execution count from `InMemoryTableMetricSuite`.
`SPARK-37578: Update output metrics from Datasource v2` in
`SQLAppStatusListenerSuite` and `BasicWriteJobStatsTrackerMetricSuite` still
wait on execution counts. This is on purpose: the first reads task output
metrics from its own listener, and the second uses its own local session where
it expects exactly two executions, so neither reads the metrics of the latest
execution by id.
### Why are the changes needed?
`driver side SQL metrics` fails intermittently in CI, e.g. in the scheduled
master and branch-4.x builds:
```
- driver side SQL metrics *** FAILED *** (21 milliseconds)
Map(3105033 -> "2", 3105032 -> "1", 3105031 -> "total (min, med, max
(stageId: taskId))
0 ms (0 ms, 0 ms, 0 ms (stage 5.0: task 10))") did not contain key
3105125 (SQLAppStatusListenerSuite.scala:607)
```
The tests above wait for the metrics of the execution they just ran like
this:
1. Wait for `executionsCount() < oldCount` to become false (or for
`executionsCount() >= oldCount` to hold). This is always satisfied immediately,
since the count never drops below the count taken before the query.
2. Wait until `executionsList().last.metricValues != null`, then read the
metrics of `executionsList().last`.
`SQLAppStatusListener` processes events asynchronously. If it has not
processed the new execution yet, `executionsList().last` is an older execution,
e.g. from a previous test, which already has metric values. The wait then ends
immediately and the test reads the wrong execution's metrics. The map in the
failure above is such a previous execution's metrics. The two `PythonUDFSuite`
tests read `executionsList().last` without waiting at all.
Waiting for a larger execution id rather than a larger execution count is
needed because the status store evicts old executions once
`spark.sql.ui.retainedExecutions` is reached, after which the count no longer
grows. Execution ids always grow.
In addition, several of the old waits were unbounded `while` /
`Thread.sleep` loops, so a real hang would never fail the test.
### Does this PR introduce _any_ user-facing change?
No. Test-only change.
### How was this patch tested?
- Reproduced the race deterministically with a temporary (not included)
suite that stalls the status listener queue for 2 seconds while the query runs:
the old wait logic (both the `SQLAppStatusListenerSuite` loops and the old
`runAndFetchMetrics`) fails every time with the same symptom as CI (it reads
the previous execution's metrics), and `runAndWaitForExecution` passes.
- Reproduced the baseline race from review deterministically with a
temporary (not included) suite that keeps an earlier execution's events queued
while the baseline is read: without draining the listener bus first,
`runAndWaitForExecution` returns that earlier execution; with it, it returns
the execution `func` ran.
- Verified with a temporary (not included) suite that
`runAndWaitForExecution` works when the status store is at its retention limit
(`spark.sql.ui.retainedExecutions=2`).
- Ran all affected suites locally, all passed:
`SQLAppStatusListenerWithInMemoryStoreSuite`,
`SQLAppStatusListenerWithRocksDBBackendSuite`, `InMemoryTableMetricSuite`,
`KeyGroupedPartitioningSuite`, `PythonDataSourceSuite` and `PythonUDFSuite`
(the Python suites with pandas and pyarrow installed, so no test was skipped).
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
Closes #59043 from SEPURI-SAI-KRISHNA/flaky-driver-side-sql-metrics.
Authored-by: Sepuri Sai Krishna <[email protected]>
Signed-off-by: Uros Bojanic <[email protected]>
(cherry picked from commit 3bdbac46beb9d5f865e49aa5e418cd3ed83c40e4)
Signed-off-by: Uros Bojanic <[email protected]>
---
.../datasources/InMemoryTableMetricSuite.scala | 3 -
.../execution/python/PythonDataSourceSuite.scala | 16 +----
.../sql/execution/python/PythonUDFSuite.scala | 6 +-
.../execution/ui/SQLAppStatusListenerSuite.scala | 76 +++++-----------------
.../apache/spark/sql/test/SharedSparkSession.scala | 33 ++++++----
5 files changed, 38 insertions(+), 96 deletions(-)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/InMemoryTableMetricSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/InMemoryTableMetricSuite.scala
index 09666f66572d..70c102ea3458 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/InMemoryTableMetricSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/InMemoryTableMetricSuite.scala
@@ -41,9 +41,6 @@ class InMemoryTableMetricSuite extends SharedSparkSession
with BeforeAndAfter {
private def testMetricOnDSv2(func: String => Unit, checker: Map[String,
String] => Unit): Unit = {
withTable("testcat.table_name") {
- val statusStore = spark.sharedState.statusStore
- val oldCount = statusStore.executionsList().size
-
val testCatalog =
spark.sessionState.catalogManager.catalog("testcat").asTableCatalog
testCatalog.createTable(
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala
index de96f23e519b..cbe44a22a301 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala
@@ -928,27 +928,13 @@ class PythonDataSourceSuite extends
PythonDataSourceSuiteBase {
val df = spark.read.format(dataSourceName).load()
val statusStore = spark.sharedState.statusStore
- val oldCount = statusStore.executionsList().size
-
- df.collect()
-
- // Wait until the new execution is started and being tracked.
- while (statusStore.executionsCount() < oldCount) {
- Thread.sleep(100)
- }
-
- // Wait for listener to finish computing the metrics for the execution.
- while (statusStore.executionsList().isEmpty ||
- statusStore.executionsList().last.metricValues == null) {
- Thread.sleep(100)
- }
+ val execId = runAndWaitForExecution(df.collect())
val executedPlan = df.queryExecution.executedPlan.collectFirst {
case p: BatchScanExec => p
}
assert(executedPlan.isDefined)
- val execId = statusStore.executionsList().last.executionId
val metrics = statusStore.executionMetrics(execId)
val pythonDataSent = executedPlan.get.metrics("pythonDataSent")
val pythonDataReceived = executedPlan.get.metrics("pythonDataReceived")
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
index 96fb33055057..9236ee9318d9 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
@@ -101,10 +101,9 @@ class PythonUDFSuite extends SharedSparkSession {
val df = base.groupBy(pythonTestUDF(base("a") + 1))
.agg(pythonTestUDF(pythonTestUDF(base("a") + 1)))
- df.count()
+ val lastExecId = runAndWaitForExecution(df.count())
val statusStore = spark.sharedState.statusStore
- val lastExecId = statusStore.executionsList().last.executionId
val executionMetrics =
statusStore.execution(lastExecId).get.metrics.mkString
for (metric <- pythonSQLMetrics) {
assert(executionMetrics.contains(metric))
@@ -185,10 +184,9 @@ class PythonUDFSuite extends SharedSparkSession {
val df = base.groupBy(pythonTestUDF(base("a") + 1))
.agg(pythonTestUDF(pythonTestUDF(base("a") + 1)))
- df.count()
+ val lastExecId = runAndWaitForExecution(df.count())
val statusStore = spark.sharedState.statusStore
- val lastExecId = statusStore.executionsList().last.executionId
val executionMetrics =
statusStore.execution(lastExecId).get.metrics.mkString
for (metric <- pythonSQLMetrics) {
assert(executionMetrics.contains(metric),
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListenerSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListenerSuite.scala
index 7d9c79ab9236..dda5c6e6baae 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListenerSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListenerSuite.scala
@@ -570,7 +570,6 @@ abstract class SQLAppStatusListenerSuite extends
SharedSparkSession with JsonTes
test("driver side SQL metrics") {
val statusStore = spark.sharedState.statusStore
- val oldCount = statusStore.executionsList().size
val expectedAccumValue = 12345L
val expectedAccumValue2 = 54321L
@@ -580,22 +579,12 @@ abstract class SQLAppStatusListenerSuite extends
SharedSparkSession with JsonTes
override lazy val executedPlan = physicalPlan
}
- SQLExecution.withNewExecutionId(dummyQueryExecution) {
- physicalPlan.execute().collect()
- }
-
- // Wait until the new execution is started and being tracked.
- while (statusStore.executionsCount() < oldCount) {
- Thread.sleep(100)
- }
-
- // Wait for listener to finish computing the metrics for the execution.
- while (statusStore.executionsList().isEmpty ||
- statusStore.executionsList().last.metricValues == null) {
- Thread.sleep(100)
+ val execId = runAndWaitForExecution {
+ SQLExecution.withNewExecutionId(dummyQueryExecution) {
+ physicalPlan.execute().collect()
+ }
}
- val execId = statusStore.executionsList().last.executionId
val metrics = statusStore.executionMetrics(execId)
val driverMetric = physicalPlan.metrics("dummy")
val driverMetric2 = physicalPlan.metrics("dummy2")
@@ -845,7 +834,6 @@ abstract class SQLAppStatusListenerSuite extends
SharedSparkSession with JsonTes
test("SPARK-34338: Report metrics from Datasource v2 scan") {
val statusStore = spark.sharedState.statusStore
- val oldCount = statusStore.executionsList().size
val schema = new StructType().add("i", "int").add("j", "int")
val physicalPlan = BatchScanExec(toAttributes(schema), new
CustomMetricScanBuilder(), Seq.empty,
@@ -855,22 +843,12 @@ abstract class SQLAppStatusListenerSuite extends
SharedSparkSession with JsonTes
override lazy val executedPlan = physicalPlan
}
- SQLExecution.withNewExecutionId(dummyQueryExecution) {
- physicalPlan.execute().collect()
- }
-
- // Wait until the new execution is started and being tracked.
- while (statusStore.executionsCount() < oldCount) {
- Thread.sleep(100)
- }
-
- // Wait for listener to finish computing the metrics for the execution.
- while (statusStore.executionsList().isEmpty ||
- statusStore.executionsList().last.metricValues == null) {
- Thread.sleep(100)
+ val execId = runAndWaitForExecution {
+ SQLExecution.withNewExecutionId(dummyQueryExecution) {
+ physicalPlan.execute().collect()
+ }
}
- val execId = statusStore.executionsList().last.executionId
val metrics = statusStore.executionMetrics(execId)
val expectedMetric = physicalPlan.metrics("custom_metric")
val expectedValue = "custom_metric: 12345, 12345"
@@ -884,7 +862,6 @@ abstract class SQLAppStatusListenerSuite extends
SharedSparkSession with JsonTes
test("SPARK-39635: Report driver metrics from Datasource v2 scan") {
val statusStore = spark.sharedState.statusStore
- val oldCount = statusStore.executionsList().size
val schema = new StructType().add("i", "int").add("j", "int")
val physicalPlan = BatchScanExec(toAttributes(schema), new
CustomDriverMetricScanBuilder(),
@@ -894,22 +871,12 @@ abstract class SQLAppStatusListenerSuite extends
SharedSparkSession with JsonTes
override lazy val executedPlan = physicalPlan
}
- SQLExecution.withNewExecutionId(dummyQueryExecution) {
- physicalPlan.execute().collect()
- }
-
- // Wait until the new execution is started and being tracked.
- while (statusStore.executionsCount() < oldCount) {
- Thread.sleep(100)
- }
-
- // Wait for listener to finish computing the metrics for the execution.
- while (statusStore.executionsList().isEmpty ||
- statusStore.executionsList().last.metricValues == null) {
- Thread.sleep(100)
+ val execId = runAndWaitForExecution {
+ SQLExecution.withNewExecutionId(dummyQueryExecution) {
+ physicalPlan.execute().collect()
+ }
}
- val execId = statusStore.executionsList().last.executionId
val metrics = statusStore.executionMetrics(execId)
val expectedMetric =
physicalPlan.metrics("custom_driver_metric_partition_count")
val expectedValue = "2"
@@ -920,25 +887,14 @@ abstract class SQLAppStatusListenerSuite extends
SharedSparkSession with JsonTes
test("SPARK-36030: Report metrics from Datasource v2 write") {
withTempDir { dir =>
val statusStore = spark.sharedState.statusStore
- val oldCount = statusStore.executionsList().size
val cls = classOf[CustomMetricsDataSource].getName
- spark.range(10).select($"id" as Symbol("i"), -$"id" as Symbol("j"))
- .write.format(cls)
- .option("path", dir.getCanonicalPath).mode("append").save()
-
- // Wait until the new execution is started and being tracked.
- eventually(timeout(10.seconds), interval(10.milliseconds)) {
- assert(statusStore.executionsCount() >= oldCount)
- }
-
- // Wait for listener to finish computing the metrics for the execution.
- eventually(timeout(10.seconds), interval(10.milliseconds)) {
- assert(statusStore.executionsList().nonEmpty &&
- statusStore.executionsList().last.metricValues != null)
+ val execId = runAndWaitForExecution {
+ spark.range(10).select($"id" as Symbol("i"), -$"id" as Symbol("j"))
+ .write.format(cls)
+ .option("path", dir.getCanonicalPath).mode("append").save()
}
- val execId = statusStore.executionsList().last.executionId
val metrics = statusStore.executionMetrics(execId)
val customMetric = metrics.find(_._2 == "custom_metric: 12345, 12345")
val innerMetric = metrics.find(_._2 == "inner_metric: 54321, 54321")
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
index 47e143339f40..d5cf8e2bf281 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala
@@ -26,27 +26,32 @@ import org.apache.spark.sql.classic
trait SharedSparkSession extends QueryTest with classic.SparkSessionBinder {
- // Runs func (which must trigger exactly one SQL execution) and returns the
SQL metrics of that
- // execution as a map keyed by (planNodeId, planNodeName, metricName) ->
metricValue.
- def runAndFetchMetrics(func: => Unit): Map[(Long, String, String), String] =
{
+ // Runs func (which must trigger exactly one SQL execution), waits until the
status store has
+ // computed the SQL metrics of that execution and returns its execution id.
+ def runAndWaitForExecution(func: => Unit): Long = {
val statusStore = spark.sharedState.statusStore
- val oldCount = statusStore.executionsList().size
+ // The listener updates the status store asynchronously, so the latest
execution in the store
+ // may still be one from before func. Drain the listener bus first so that
every execution
+ // started before func is in the store. Execution ids only grow, so then
wait for an execution
+ // with a larger id than any of them. Execution counts cannot be used for
this, as old
+ // executions are evicted once spark.sql.ui.retainedExecutions is reached.
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+ val lastExecId =
statusStore.executionsList().lastOption.map(_.executionId).getOrElse(-1L)
func
- // Wait until the new execution is started and being tracked.
- eventually(timeout(10.seconds), interval(10.milliseconds)) {
- assert(statusStore.executionsCount() >= oldCount)
- }
-
- // Wait for listener to finish computing the metrics for the execution.
eventually(timeout(10.seconds), interval(10.milliseconds)) {
- assert(statusStore.executionsList().nonEmpty &&
- statusStore.executionsList().last.metricValues != null)
+ val exec = statusStore.executionsList().lastOption
+ assert(exec.exists(e => e.executionId > lastExecId && e.metricValues !=
null))
+ exec.get.executionId
}
+ }
- val exec = statusStore.executionsList().last
- val execId = exec.executionId
+ // Runs func (which must trigger exactly one SQL execution) and returns the
SQL metrics of that
+ // execution as a map keyed by (planNodeId, planNodeName, metricName) ->
metricValue.
+ def runAndFetchMetrics(func: => Unit): Map[(Long, String, String), String] =
{
+ val execId = runAndWaitForExecution(func)
+ val statusStore = spark.sharedState.statusStore
val sqlMetrics = statusStore.planGraph(execId).allNodes
.flatMap(n => n.metrics.map(m => (m.accumulatorId, (n.id, n.name,
m.name))))
.toMap
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
