Skip to content

[SPARK-57638][4.1][SQL] Avoid busy-waiting in Declarative Pipelines flow resolution #57571

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
LuciferYang wants to merge 1 commit into
base: branch-4.1
Choose a base branch
from
+194 −127
Open

[SPARK-57638][4.1][SQL] Avoid busy-waiting in Declarative Pipelines flow resolution #57571

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
263 changes: 136 additions & 127 deletions ...elines/src/main/scala/org/apache/spark/sql/pipelines/graph/DataflowGraphTransformer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,10 @@ import java.util.concurrent.{
ConcurrentLinkedDeque,
ConcurrentLinkedQueue,
ExecutionException,
ExecutorCompletionService,
Future
}

import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
import scala.util.control.NoStackTrace

Expand Down Expand Up @@ -134,159 +134,168 @@ class DataflowGraphTransformer(graph: DataflowGraph) extends AutoCloseable {
val failedFlowsQueue = new ConcurrentLinkedQueue[ResolutionFailedFlow]()
val failedDependentFlows = new ConcurrentHashMap[TableIdentifier, Seq[ResolutionFailedFlow]]()

var futures = ArrayBuffer[Future[Unit]]()
val completionService = new ExecutorCompletionService[Unit](executor)
var outstanding = 0
val toBeResolvedFlows = new ConcurrentLinkedDeque[Flow]()
toBeResolvedFlows.addAll(flows.asJava)

while (futures.nonEmpty || toBeResolvedFlows.peekFirst() != null) {
val (done, notDone) = futures.partition(_.isDone)
// Explicitly call future.get() to propagate exceptions one by one if any
// Waits on a finished resolution task and propagates its exception, if any.
def reap(finished: Future[Unit]): Unit = {
try {
done.foreach(_.get())
finished.get()
} catch {
case exn: ExecutionException =>
// Computation threw the exception that is the cause of exn
throw exn.getCause
}
futures = notDone
val flowOpt = {
// We only schedule [[batchSize]] number of flows in parallel.
if (futures.size < batchSize) {
Option(toBeResolvedFlows.pollFirst())
} else {
None
}
outstanding -= 1
}

while (outstanding > 0 || toBeResolvedFlows.peekFirst() != null) {
// Reap every resolution task that has already finished, without blocking.
var finished = completionService.poll()
while (finished != null) {
reap(finished)
finished = completionService.poll()
}
flowOpt.foreach { flow =>
futures.append(
executor.submit(
() =>
// We only schedule [[batchSize]] number of flows in parallel.
if (outstanding < batchSize && toBeResolvedFlows.peekFirst() != null) {
val flow = toBeResolvedFlows.pollFirst()
outstanding += 1
completionService.submit(
() =>
try {
try {
try {
// Note: Flow don't need their inputs passed, so for now we send empty Seq.
val result = transformer(flow, Seq.empty)
require(
result.forall(_.isInstanceOf[ResolvedFlow]),
"transformer must return a Seq[Flow]"
)
// Note: Flow don't need their inputs passed, so for now we send empty Seq.
val result = transformer(flow, Seq.empty)
require(
result.forall(_.isInstanceOf[ResolvedFlow]),
"transformer must return a Seq[Flow]"
)

val transformedFlows = result.map(_.asInstanceOf[ResolvedFlow])
resolvedFlowsMap.put(flow.identifier, transformedFlows)
resolvedFlows.addAll(transformedFlows.asJava)
} catch {
case e: TransformNodeRetryableException =>
val datasetIdentifier = e.datasetIdentifier
failedDependentFlows.compute(
datasetIdentifier,
(_, flows) => {
// Don't add the input flow back but the failed flow object
// back which has relevant failure information.
val failedFlow = e.failedNode
if (flows == null) {
Seq(failedFlow)
} else {
flows :+ failedFlow
}
val transformedFlows = result.map(_.asInstanceOf[ResolvedFlow])
resolvedFlowsMap.put(flow.identifier, transformedFlows)
resolvedFlows.addAll(transformedFlows.asJava)
} catch {
case e: TransformNodeRetryableException =>
val datasetIdentifier = e.datasetIdentifier
failedDependentFlows.compute(
datasetIdentifier,
(_, flows) => {
// Don't add the input flow back but the failed flow object
// back which has relevant failure information.
val failedFlow = e.failedNode
if (flows == null) {
Seq(failedFlow)
} else {
flows :+ failedFlow
}
)
// Between the time the flow started and finished resolving, perhaps the
// dependent dataset was resolved
resolvedFlowDestinationsMap.computeIfPresent(
datasetIdentifier,
(_, resolved) => {
if (resolved) {
// Check if the dataset that the flow is dependent on has been resolved
// and if so, remove all dependent flows from the failedDependentFlows and
// add them to the toBeResolvedFlows queue for retry.
failedDependentFlows.computeIfPresent(
datasetIdentifier,
(_, toRetryFlows) => {
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
null
}
)
}
resolved
}
)
// Between the time the flow started and finished resolving, perhaps the
// dependent dataset was resolved
resolvedFlowDestinationsMap.computeIfPresent(
datasetIdentifier,
(_, resolved) => {
if (resolved) {
// Check if the dataset that the flow is dependent on has been resolved
// and if so, remove all dependent flows from the failedDependentFlows and
// add them to the toBeResolvedFlows queue for retry.
failedDependentFlows.computeIfPresent(
datasetIdentifier,
(_, toRetryFlows) => {
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
null
}
)
}
resolved
}
)
case other: Throwable => throw other
}
// If all flows to this particular destination are resolved, move to the destination
// node transformer
if (flowsTo(flow.destinationIdentifier).forall({ flowToDestination =>
resolvedFlowsMap.containsKey(flowToDestination.identifier)
})) {
// If multiple flows completed in parallel, ensure we resolve the destination only
// once by electing a leader via computeIfAbsent
var isCurrentThreadLeader = false
resolvedFlowDestinationsMap.computeIfAbsent(flow.destinationIdentifier, _ => {
isCurrentThreadLeader = true
// Set initial value as false as flow destination is not resolved yet.
false
})
if (isCurrentThreadLeader) {
if (tableMap.contains(flow.destinationIdentifier)) {
val transformed =
transformer(
tableMap(flow.destinationIdentifier),
flowsTo(flow.destinationIdentifier)
)
resolvedTables.addAll(
transformed.collect { case t: Table => t }.asJava
)
case other: Throwable => throw other
}
// If all flows to this particular destination are resolved, move to the destination
// node transformer
if (flowsTo(flow.destinationIdentifier).forall({ flowToDestination =>
resolvedFlowsMap.containsKey(flowToDestination.identifier)
})) {
// If multiple flows completed in parallel, ensure we resolve the destination only
// once by electing a leader via computeIfAbsent
var isCurrentThreadLeader = false
resolvedFlowDestinationsMap.computeIfAbsent(flow.destinationIdentifier, _ => {
isCurrentThreadLeader = true
// Set initial value as false as flow destination is not resolved yet.
false
})
if (isCurrentThreadLeader) {
if (tableMap.contains(flow.destinationIdentifier)) {
resolvedFlows.addAll(
transformed.collect { case f: ResolvedFlow => f }.asJava
)
} else if (viewMap.contains(flow.destinationIdentifier)) {
resolvedViews.addAll {
val transformed =
transformer(
tableMap(flow.destinationIdentifier),
viewMap(flow.destinationIdentifier),
flowsTo(flow.destinationIdentifier)
)
resolvedTables.addAll(
transformed.collect { case t: Table => t }.asJava
)
resolvedFlows.addAll(
transformed.collect { case f: ResolvedFlow => f }.asJava
)
} else if (viewMap.contains(flow.destinationIdentifier)) {
resolvedViews.addAll {
val transformed =
transformer(
viewMap(flow.destinationIdentifier),
flowsTo(flow.destinationIdentifier)
)
transformed.map(_.asInstanceOf[View]).asJava
}
} else if (sinkMap.contains(flow.destinationIdentifier)) {
resolvedSinks.addAll {
val transformed =
transformer(
sinkMap(flow.destinationIdentifier), flowsTo(flow.destinationIdentifier)
)
require(
transformed.forall(_.isInstanceOf[Sink]),
"transformer must return a Seq[Sink]"
transformed.map(_.asInstanceOf[View]).asJava
}
} else if (sinkMap.contains(flow.destinationIdentifier)) {
resolvedSinks.addAll {
val transformed =
transformer(
sinkMap(flow.destinationIdentifier), flowsTo(flow.destinationIdentifier)
)
transformed.map(_.asInstanceOf[Sink]).asJava
}
} else {
throw new IllegalArgumentException(
s"Unsupported destination ${flow.destinationIdentifier.unquotedString}" +
s" in flow: ${flow.displayName} at transformDownNodes"
require(
transformed.forall(_.isInstanceOf[Sink]),
"transformer must return a Seq[Sink]"
)
transformed.map(_.asInstanceOf[Sink]).asJava
}
// Set flow destination as resolved now.
resolvedFlowDestinationsMap.computeIfPresent(
flow.destinationIdentifier,
(_, _) => {
// If there are any other node failures dependent on this destination, retry
// them
failedDependentFlows.computeIfPresent(
flow.destinationIdentifier,
(_, toRetryFlows) => {
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
null
}
)
true
}
} else {
throw new IllegalArgumentException(
s"Unsupported destination ${flow.destinationIdentifier.unquotedString}" +
s" in flow: ${flow.displayName} at transformDownNodes"
)
}
// Set flow destination as resolved now.
resolvedFlowDestinationsMap.computeIfPresent(
flow.destinationIdentifier,
(_, _) => {
// If there are any other node failures dependent on this destination, retry
// them
failedDependentFlows.computeIfPresent(
flow.destinationIdentifier,
(_, toRetryFlows) => {
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
null
}
)
true
}
)
}
} catch {
case ex: TransformNodeFailedException => failedFlowsQueue.add(ex.failedNode)
}
)
} catch {
case ex: TransformNodeFailedException => failedFlowsQueue.add(ex.failedNode)
}
)
} else if (outstanding > 0) {
// Nothing could be scheduled (slots full, or the queue is drained) but tasks are still
// running: block until the next finishes instead of busy-spinning on Future.isDone. The
// outstanding > 0 guard is required, not redundant: the poll() drain above can take
// outstanding to 0 with an empty queue, and then there is nothing to wait for - the loop
// should just exit rather than block forever in take().
reap(completionService.take())
}
}

Expand Down
58 changes: 58 additions & 0 deletions ...lines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectValidPipelineSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -512,6 +512,64 @@ class ConnectValidPipelineSuite extends PipelineTest with SharedSparkSession {
assert(g.flow(TableIdentifier("sink_flow")).isInstanceOf[StreamingFlow])
}

test("resolution terminates and resolves all flows when flow count exceeds parallelism") {
val session = spark
import session.implicits._

// DataflowGraphTransformer caps in-flight resolutions at `parallelism` (10). With more
// independent flows than that, the slots fill and the scheduler blocks on a finished task
// (the `take()` branch) once `parallelism` tasks are outstanding - the path the busy-wait
// rewrite changes. Small graphs in the other suites never reach this regime. This asserts the
// outcome only (all flows resolve and the call returns), so it is deterministic and has no
// timing dependence; a regression that deadlocked would hang here until the suite times out.
val numFlows = 25
class P extends TestGraphRegistrationContext(spark) {
(0 until numFlows).foreach { i =>
registerPersistedView(s"v$i", query = dfFlowFunc(Seq(i).toDF("x")))
}
}
val p = new P().resolveToDataflowGraph()

assert(p.resolved, "all flows should resolve when their count exceeds parallelism")
(0 until numFlows).foreach { i =>
assert(
p.resolvedFlow.contains(fullyQualifiedIdentifier(s"v$i")),
s"flow v$i was not resolved")
}
}

test("resolution re-queues retryable flows under load when consumers exceed parallelism") {
val session = spark
import session.implicits._

// A wide fan-out: many consumers reading from one source view, with the consumer count above
// `parallelism` (10), so slots fill and the loop blocks on take(). The consumers are
// registered (and therefore scheduled) before `src`, so the first batch resolves consumers
// whose `src` input is not yet available: each throws TransformNodeRetryableException and is
// parked as a dependent of `src`; once `src` resolves they are re-queued onto the deque and
// retried, re-driving the loop until every consumer resolves. This deterministically exercises
// the retryable re-queue path together with the blocking branch. Asserts only that everything
// resolves and the call returns - no timing assertions.
val numConsumers = 20
class P extends TestGraphRegistrationContext(spark) {
(0 until numConsumers).foreach { i =>
registerPersistedView(s"c$i", query = sqlFlowFunc(spark, "SELECT x FROM src"))
}
registerPersistedView("src", query = dfFlowFunc(Seq(1, 2, 3).toDF("x")))
}
val p = new P().resolveToDataflowGraph()

assert(p.resolved, "source and all consumers should resolve under load")
assert(
p.resolvedFlow.contains(fullyQualifiedIdentifier("src")),
"source flow was not resolved")
(0 until numConsumers).foreach { i =>
assert(
p.resolvedFlow.contains(fullyQualifiedIdentifier(s"c$i")),
s"consumer flow c$i was not resolved")
}
}

/** Verifies the [[DataflowGraph]] has the specified [[Flow]] with the specified schema. */
private def verifyFlowSchema(
pipeline: DataflowGraph,
Expand Down