andygrove commented on code in PR #6601:
URL: https://github.com/apache/datafusion-comet/pull/6601#discussion_r4209718015
##########
spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala:
##########
@@ -3671,6 +3717,65 @@ class CometIcebergWriteActionSuite
}
}
+ test("native acceleration: scheduler probe records handoff and accepted
files") {
Review Comment:
This test runs in mode `smoke`, so `beforeNative()` returns `None` and CI
never drives the native gate. That's why nothing caught the #6241 change. If
the native gate stays, could we add a local test that uses an `executor-loss`
probe with `release-all` touched up front, writes a few thousand rows at batch
size 1000, and asserts that `native-progress.json` lists files with bytes on
disk? I tried one locally and it runs in about a second.
##########
spark/src/test/scala/org/apache/comet/CometIcebergSchedulerFailureSuite.scala:
##########
@@ -0,0 +1,584 @@
+/*
+ * 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.comet
+
+import java.nio.charset.StandardCharsets.UTF_8
+import java.nio.file.{Files, Path, Paths}
+import java.util.UUID
+import java.util.concurrent.TimeUnit
+
+import scala.concurrent.{Await, Future}
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.duration.DurationLong
+import scala.jdk.CollectionConverters._
+
+import org.json4s.jackson.JsonMethods.parse
+import org.scalatest.DoNotDiscover
+
+import org.apache.spark.{CometListenerBusUtils, SparkConf}
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.comet.{CometIcebergWriteExec, IcebergCommitExec,
IcebergSchedulerTestProbe}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+
+/**
+ * Manual multi-host integration suite. Excluded from the ordinary single-host
CI matrix. Requires
+ * COMET_SCHEDULER_MASTER, SHARED_ROOT, EXECUTOR_CLASSPATH and KILL_COMMAND
(all with the
+ * COMET_SCHEDULER_ prefix). Missing infrastructure is a failure, never a
skipped pass. See
+ * docs/source/contributor-guide/iceberg-scheduler-failures.md for topology
and commands.
+ */
+@DoNotDiscover
+class CometIcebergSchedulerFailureSuite
+ extends CometTestBase
+ with AdaptiveSparkPlanHelper
+ with CometIcebergTestBase {
+ import IcebergSchedulerFiles._
+
+ private def required(name: String): String =
+ sys.env.getOrElse(
+ "COMET_SCHEDULER_" + name,
+ throw new IllegalArgumentException(s"COMET_SCHEDULER_$name is required"))
+
+ private var classpathDirectory: Option[Path] = None
+
+ private def timeoutMillis: Long =
+ sys.env.getOrElse("COMET_SCHEDULER_TIMEOUT_SECONDS", "180").toLong * 1000L
+
+ override protected def sparkConf: SparkConf = {
+ val master = required("MASTER")
+ require(!master.startsWith("local"), "A genuinely multi-host Spark cluster
is required")
+ super.sparkConf
+ .set("spark.shuffle.manager", "sort")
+ .setMaster(master)
+ .set("spark.executor.extraClassPath", executorClasspath())
+ .set("spark.executor.cores", "1")
+ .set("spark.cores.max", "4")
+ .set("spark.dynamicAllocation.enabled", "false")
+ .set("spark.shuffle.service.enabled", "false")
+ .set("spark.shuffle.push.enabled", "false")
+ .set("spark.decommission.enabled", "false")
+ .set("spark.task.maxFailures", "4")
+ .set("spark.speculation",
sys.env.getOrElse("COMET_SCHEDULER_SPECULATION_ENABLED", "true"))
+ .set("spark.speculation.interval", "100ms")
+ .set("spark.speculation.quantile", "0.5")
+ .set("spark.speculation.multiplier", "1.0")
+ .set("spark.speculation.minTaskRuntime", "500ms")
+ .set("spark.speculation.efficiency.enabled", "false")
+ .set("spark.sql.adaptive.enabled", "false")
+ .set("spark.sql.adaptive.coalescePartitions.enabled", "false")
+ .set("spark.sql.files.maxPartitionBytes", "1048576")
+ .set("spark.sql.files.openCostInBytes", "1048576")
+ .set(CometConf.COMET_BATCH_SIZE.key, "1000")
+ .set(CometConf.COMET_SHUFFLE_ENABLED.key, "false")
+ .set(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key, "true")
+ .set(CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.key, "true")
+ .set(
+ "spark.sql.extensions",
+ "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
+ }
+
+ // Use a snapshot of the Maven test JVM's unshaded classpath on all hosts.
This includes
+ // the actual reactor outputs, test helpers, Iceberg runtime and native
resource. Packaging
+ // only the shaded Comet jar would give the unshaded task closures
incompatible classes.
+ private def executorClasspath(): String = {
+ sys.env.get("COMET_SCHEDULER_EXECUTOR_CLASSPATH").getOrElse {
+ val shared =
Paths.get(required("SHARED_ROOT")).toAbsolutePath.normalize()
+ require(Files.isDirectory(shared), s"Shared root does not exist:
$shared")
+ val destination = Files.createDirectory(
+ shared.resolve("comet-scheduler-classpath-" +
UUID.randomUUID().toString))
+ classpathDirectory = Some(destination)
+ val sources = (System
+ .getProperty("java.class.path")
+ .split(java.io.File.pathSeparator)
+ .map(Paths.get(_))
+ .toSeq ++ Seq(
+
Paths.get(getClass.getProtectionDomain.getCodeSource.getLocation.toURI),
+ Paths.get(
+
classOf[CometIcebergWriteExec].getProtectionDomain.getCodeSource.getLocation.toURI)))
+ .map(_.toAbsolutePath.normalize())
+ .distinct
+ require(sources.forall(Files.exists(_)), s"Missing classpath entries:
$sources")
+ sources.zipWithIndex
+ .map { case (source, index) =>
+ val target = destination.resolve(index.toString + "-" +
source.getFileName.toString)
+ if (Files.isDirectory(source)) {
+ val tree = Files.walk(source)
+ try
+ tree.iterator().asScala.foreach { entry =>
+ val copy = target.resolve(source.relativize(entry))
+ if (Files.isDirectory(entry)) Files.createDirectories(copy)
+ else Files.copy(entry, copy)
+ }
+ finally tree.close()
+ } else Files.copy(source, target)
+ target.toString
+ }
+ .mkString(java.io.File.pathSeparator)
+ }
+ }
+
+ override protected def afterAll(): Unit = {
+ try super.afterAll()
+ finally classpathDirectory.foreach(dir => deleteRecursively(dir.toFile))
+ }
+
+ test("speculative native attempts: original wins") {
+ runScenario("speculation", speculativeWins = false)
+ }
+
+ test("speculative native attempts: speculative wins") {
+ runScenario("speculation", speculativeWins = true)
+ }
+
+ test("executor process loss during a concurrent native write: task
replacement") {
+ runScenario("executor-loss", speculativeWins = false)
+ }
+
+ test("executor process loss with shuffle output loss: stage re-execution") {
+ runScenario("stage-loss", speculativeWins = false)
+ }
+
+ private def preflight(root: Path): Unit = {
+ val directory = root.resolve("preflight")
+ Files.createDirectories(directory)
+ val challenge = UUID.randomUUID().toString
+ Files.write(directory.resolve("challenge"), challenge.getBytes(UTF_8))
+ val shared = directory.toString
+ val timeout = timeoutMillis
+ val job = Future {
+ spark.sparkContext
+ .parallelize(0 until 16, 16)
+ .mapPartitions { input =>
+ val dir = Paths.get(shared)
+ require(
+ new String(Files.readAllBytes(dir.resolve("challenge")), UTF_8) ==
challenge,
+ "Executors must see the driver's shared filesystem")
+ val a = context()
+ write(dir.resolve(s"worker-${a.attempt}.json"), a)
+ await("preflight release",
timeout)(Files.exists(dir.resolve("release")))(
+ paths(dir).mkString("\n"))
+ input
+ }
+ .collect()
+ }
+ try {
+ await("two executor processes on distinct scheduler hosts",
timeoutMillis)({
+ job.value.foreach(_.get)
+ val workers = paths(directory)
+ .filter(p =>
+ p.getFileName.toString.startsWith("worker-") &&
+ p.getFileName.toString.endsWith(".json"))
+ .map(read)
+ workers.map(_.executor).distinct.size >= 2 &&
workers.map(_.host).distinct.size >= 2
+ })(paths(directory).mkString("\n"))
+ } finally {
+ touch(directory.resolve("release"))
+ Await.ready(job, timeoutMillis.millis)
+ }
+ assert(Await.result(job, timeoutMillis.millis).toSet == (0 until 16).toSet)
+ }
+
+ private def runScenario(mode: String, speculativeWins: Boolean): Unit = {
+ assert(icebergAvailable, "A compatible Iceberg runtime must be on driver
and executors")
+ val shared = Paths.get(required("SHARED_ROOT")).toAbsolutePath.normalize()
+ require(Files.isDirectory(shared), s"Shared POSIX root does not exist:
$shared")
+ val runId = UUID.randomUUID().toString.replace("-", "")
+ val run = Files.createDirectory(shared.resolve(s"comet-scheduler-$runId"))
+ val gates = Files.createDirectory(run.resolve("gates"))
+ val warehouse = Files.createDirectory(run.resolve("warehouse"))
+ val source = run.resolve("source")
+ val catalog = "scheduler" + runId
+ val table = s"$catalog.db.target"
+ val group = "iceberg-scheduler-" + runId
+ val listener = new IcebergSchedulerEvents(group)
+ val probe = SharedIcebergSchedulerProbe(gates.toString, mode,
timeoutMillis)
+ var writeJob: Option[Future[Seq[SparkPlan]]] = None
+ var passed = false
+ var dataRoot: Option[Path] = None
+ var tableCreated = false
+ var speculativeWinner: Option[Long] = None
+ spark.sparkContext.addSparkListener(listener)
+ try {
+ preflight(run)
+ withSQLConf(
+ s"spark.sql.catalog.$catalog" ->
"org.apache.iceberg.spark.SparkCatalog",
+ s"spark.sql.catalog.$catalog.type" -> "hadoop",
+ s"spark.sql.catalog.$catalog.warehouse" -> warehouse.toString) {
+ spark.sql(s"CREATE NAMESPACE $catalog.db")
+ spark.sql(
+ s"CREATE TABLE $table (id INT) USING iceberg " +
+ "TBLPROPERTIES ('write.distribution-mode'='none', " +
+ "'write.target-file-size-bytes'='1')")
Review Comment:
I reproduced the executor-loss setup in local mode on the current head
merged onto main, and it can't get past the check before the kill. Since #6241
landed on October 5, the native writer holds back each partition's first data
page, 20,000 rows by default, before it opens any file. This table has no
`write.parquet.page-row-limit`, and each task writes 12,000 rows. So all four
native gates fired at `rows=2000` with an empty file list, and
`assert(target.rows >= 1000L && target.files.nonEmpty)` fails. With
`'write.parquet.page-row-limit'='1000'`, the same change #6241 made to the
mid-write retry test, each gate saw two files and the first was already on
disk. Could you add that, update the "Executor process loss" section of the doc
page to match, and re-run both loss scenarios on your cluster? Your earlier
runs predate #6241.
##########
spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala:
##########
@@ -111,6 +114,7 @@ case class IcebergCommitExec(
try {
messages.foreach(batchWrite.onDataWriterCommit)
IcebergWriteSummaryShim.commit(batchWrite, messages, child)
+ schedulerProbe.foreach(_.committed())
Review Comment:
`committed()` runs inside the `try` whose `catch` aborts the write. If the
probe throws after a successful commit, we call `batchWrite.abort` on messages
that are already in a snapshot. I tried a probe whose `committed()` throws. On
Spark 4.1 the INSERT failed even though the snapshot committed. On Spark 3.4
with Iceberg 1.5.2, where `SparkWrite.cleanupOnAbort` starts out `true`, the
abort also deleted the committed data file, so the table pointed at a missing
file. Your probe's `Files.createFile` throws on a duplicate marker by design,
so the suite can reach this. Could the call move below the `try`/`catch`, once
the commit has succeeded?
##########
native/core/src/execution/operators/iceberg_write.rs:
##########
@@ -1027,16 +1027,37 @@ async fn run_write_task(
}
};
+ let probe_locations = common
+ .test_probe
+ .as_ref()
+ .map(|_| abort_guard.generator.clone());
let outcome = async move {
+ let mut rows_written = 0u64;
+ let mut probe_reached = false;
while let Some(batch) = input.try_next().await? {
let mut decorated = decorate_batch_with_field_ids(batch,
&target_schema)?;
if nests_floats {
decorated = drop_unwritten_values(decorated)?;
}
+ let rows = decorated.num_rows() as u64;
let timer = write_time.timer();
writer.write(decorated, &properties).await?;
timer.done();
reservation.try_resize(open_files.bytes() +
writer.pending_bytes())?;
+ if let Some(probe) = common.test_probe.as_ref() {
Review Comment:
I'm wary of adding a test-only message to the wire format, and a gate that
polls the filesystem inside the native write loop, for a suite that CI doesn't
run. The page-row-limit problem shows how quietly it can drift. The native
writer pulls its JVM input one batch at a time through the Arrow C stream from
`CometArrowStream.inputObjects`, so holding that iterator after N rows parks
the writer at the same point as this gate. Could the mid-write pause live in
`IcebergSchedulerProbeRDD` instead? The attempt's files can be found by the
attempt-unique prefix from `file_name_prefix`, which `iceberg-writes.md`
already documents. That would remove the proto field and the native changes
entirely.
--
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]
