fix: size the native runtime for standalone executors without spark.executor.cores

docs/source/contributor-guide/development.md

@@ -117,8 +117,10 @@ is torn down on plugin shutdown (via `release_runtime`) so that the tokio worker
 and the JVM can shut down cleanly:
 
 - **Worker threads:** one per executor core by default (`spark.executor.cores`, or the thread
-  count of a `local[N]` or `local[*]` master, and one when `spark.executor.cores` is not set
-  outside local mode), configurable via `COMET_WORKER_THREADS`
+  count of a `local[N]` or `local[*]` master, the cores per worker of a `local-cluster` master,
+  the JVM's available processors on a standalone cluster, and one on YARN and Kubernetes, when
+  `spark.executor.cores` is not set), configurable via `COMET_WORKER_THREADS`. See
+  `CometExecIterator.numDriverOrExecutorCores`.
 - **Max blocking threads:** 512 by default, configurable via `COMET_MAX_BLOCKING_THREADS`
 - All async I/O (S3, HTTP, Parquet reads) runs on worker threads as non-blocking futures
 

docs/source/user-guide/latest/tuning.md

@@ -53,11 +53,24 @@ larger, so queries over such data spill more. For these workloads:
 
 ## Configuring Tokio Runtime
 
-Comet uses a global tokio runtime per executor process. By default it starts one worker thread per executor core
-(`spark.executor.cores`, or the thread count of `local[N]` and `local[*]` masters) and allows up to 512 blocking
-threads, which is tokio's default. If `spark.executor.cores` is not set outside local mode, Comet starts a single
-worker thread. These values can be overridden using the environment variables `COMET_WORKER_THREADS` and
-`COMET_MAX_BLOCKING_THREADS`.
+Comet uses a global tokio runtime per executor process. By default it starts one worker thread per core that the
+executor runs tasks on, and allows up to 512 blocking threads, which is tokio's default. Comet takes the number of
+cores from the first of these that applies:
+
+- the thread count of a `local`, `local[N]` or `local[*]` master, in local mode
+- `spark.executor.cores`, when it is set
+- the cores per worker `C` of a `local-cluster[N, C, M]` master
+- the number of processors available to the executor's JVM on a standalone (`spark://`) cluster, since a standalone
+  executor without `spark.executor.cores` takes every core its worker offers, and a worker offers all of its
+  machine's cores unless it is started with a different number (`SPARK_WORKER_CORES`)
+- one on YARN and Kubernetes, which is their default for `spark.executor.cores`
+
+On any other cluster manager, such as Mesos, Comet cannot tell how many cores the executor has when
+`spark.executor.cores` is not set. It starts one worker thread for each processor available to the executor's JVM, and
+logs a warning unless `COMET_WORKER_THREADS` is set. Too few worker threads cost more than a few idle ones: native plans
+that read no input from the JVM, such as a native Parquet scan feeding a sort, run entirely on the worker threads, so
+tasks share them when an executor runs more tasks at once than it has worker threads. The worker and blocking thread
+counts can be overridden using the environment variables `COMET_WORKER_THREADS` and `COMET_MAX_BLOCKING_THREADS`.
 
 ## Metrics Overhead
 

spark/src/main/scala/org/apache/comet/CometExecIterator.scala

@@ -24,6 +24,7 @@ import java.util.Locale
 import java.util.concurrent.{ConcurrentHashMap, Executors, ThreadFactory, TimeUnit}
 import java.util.concurrent.atomic.AtomicBoolean
 
+import scala.util.Try
 import scala.util.control.NonFatal
 
 import org.apache.arrow.c.ArrowArrayStream
@@ -553,7 +554,14 @@ object CometExecIterator extends Logging {
    * memory limit and never reads those settings. Comet does not know how any other cluster
    * manager sizes its executors.
    */
-  def isContainerSizedFromOverhead(master: String): Boolean =
+  def isContainerSizedFromOverhead(master: String): Boolean = isYarnOrKubernetes(master)
+
+  /**
+   * Whether `master` is YARN or Kubernetes, the cluster managers that give an executor a
+   * container and a default of one core, which Comet resolves its memory overhead and core count
+   * from.
+   */
+  private def isYarnOrKubernetes(master: String): Boolean =
     master == "yarn" || master.startsWith("k8s://")
 
   /**
@@ -922,26 +930,73 @@ object CometExecIterator extends Logging {
     }
   }
 
+  private val unknownCoresWarned = new AtomicBoolean(false)
+
+  /**
+   * The number of cores this JVM runs tasks on, which sizes the native runtime's worker threads.
+   * If the master does not tell, falls back to the processors available to the JVM and logs a
+   * warning once, unless `COMET_WORKER_THREADS` sets the worker threads regardless.
+   */
   private def numDriverOrExecutorCores(conf: SparkConf): Int = {
-    def convertToInt(threads: String): Int = {
-      if (threads == "*") Runtime.getRuntime.availableProcessors() else threads.toInt
-    }
+    val master = conf.get("spark.master")
+    val availableProcessors = Runtime.getRuntime.availableProcessors()
+    numDriverOrExecutorCores(master, conf.getOption("spark.executor.cores"), availableProcessors)
+      .getOrElse {
+        // The native side takes COMET_WORKER_THREADS over this count, when it parses as a number.
+        val workerThreadsOverridden = sys.env
+          .get("COMET_WORKER_THREADS")
+          .exists(threads => Try(threads.toLong).filter(_ >= 0).isSuccess)
+        if (!workerThreadsOverridden && unknownCoresWarned.compareAndSet(false, true)) {
+          logWarning(
+            s"Comet cannot tell how many cores this executor has from spark.master=$master when " +
+              "spark.executor.cores is not set, so its native runtime starts one worker thread " +
+              s"for each of the $availableProcessors processors available to the JVM. Native " +
+              "plans that read no input from the JVM run on the worker threads, so concurrent " +
+              "tasks share them. Set spark.executor.cores, or the COMET_WORKER_THREADS " +
+              s"environment variable, to the executor's core count. ${CometConf.TUNING_GUIDE}.")
+        }
+        availableProcessors
+      }
+  }
 
-    // If running in local mode, get number of threads from the spark.master setting.
-    // See https://spark.apache.org/docs/latest/submitting-applications.html#master-urls
-    // for supported formats
+  // Master URL formats, as SparkContext parses them. See
+  // https://spark.apache.org/docs/latest/submitting-applications.html#master-urls
+  // `local[*]` means using all available cores and `local[2]` means using 2 cores.
+  private val LOCAL_N_REGEX = """local\[([0-9]+|\*)\]""".r
+  // `local[num-worker-threads, max-failures]`
+  private val LOCAL_N_FAILURES_REGEX = """local\[([0-9]+|\*)\s*,\s*([0-9]+)\]""".r
+  // `local-cluster[num-workers, cores-per-worker, memory-per-worker-mib]`
+  private val LOCAL_CLUSTER_REGEX =
+    """local-cluster\[\s*([0-9]+)\s*,\s*([0-9]+)\s*,\s*([0-9]+)\s*]""".r
 
-    // `local[*]` means using all available cores and `local[2]` means using 2 cores.
-    val LOCAL_N_REGEX = """local\[([0-9]+|\*)\]""".r
-    // Also handle format `local[num-worker-threads, max-failures]
-    val LOCAL_N_FAILURES_REGEX = """local\[([0-9]+|\*)\s*,\s*([0-9]+)\]""".r
+  /**
+   * The number of cores the driver (in local mode) or an executor runs tasks on, given the master
+   * URL, `spark.executor.cores` and the number of processors available to the JVM. Every task
+   * takes `spark.task.cpus` of these cores, so this is at least the number of tasks that run at
+   * once. None if it cannot be determined.
+   */
+  def numDriverOrExecutorCores(
+      master: String,
+      executorCores: Option[String],
+      availableProcessors: Int): Option[Int] = {
+    def localThreads(threads: String): Int =
+      if (threads == "*") availableProcessors else threads.toInt
 
-    val master = conf.get("spark.master")
     master match {
-      case "local" => 1
-      case LOCAL_N_REGEX(threads) => convertToInt(threads)
-      case LOCAL_N_FAILURES_REGEX(threads, _) => convertToInt(threads)
-      case _ => conf.get("spark.executor.cores", "1").toInt
+      // Local mode runs tasks on the master's threads and ignores spark.executor.cores.
+      case "local" => Some(1)
+      case LOCAL_N_REGEX(threads) => Some(localThreads(threads))
+      case LOCAL_N_FAILURES_REGEX(threads, _) => Some(localThreads(threads))
+      case _ if executorCores.isDefined => executorCores.map(_.toInt)
+      // Without spark.executor.cores, a standalone executor takes every core its worker offers,
+      // but Spark does not set spark.executor.cores on the executor. A local-cluster worker
+      // offers the master's cores per worker, and a standalone worker offers all of the
+      // machine's processors unless it is started with a different number.
+      case LOCAL_CLUSTER_REGEX(_, coresPerWorker, _) => Some(coresPerWorker.toInt)
+      case _ if master.startsWith("spark://") => Some(availableProcessors)
+      // The default of spark.executor.cores on YARN and Kubernetes.
+      case _ if isYarnOrKubernetes(master) => Some(1)
+      case _ => None
     }
   }
 

spark/src/test/scala/org/apache/spark/CometExecIteratorLifecycleSuite.scala

@@ -826,6 +826,39 @@ class CometExecIteratorLifecycleSuite extends CometTestBase {
     assert(memoryUsageLogInterval(Some("10 seconds please")) == 0L)
   }
 
+  test("the native runtime gets a worker thread for every core the executor runs tasks on") {
+    import CometExecIterator.numDriverOrExecutorCores
+    def cores(master: String, executorCores: Option[String] = None): Option[Int] =
+      numDriverOrExecutorCores(master, executorCores, availableProcessors = 12)
+
+    // Local mode runs tasks on the master's threads and ignores spark.executor.cores.
+    assert(cores("local") == Some(1))
+    assert(cores("local[4]") == Some(4))
+    assert(cores("local[*]") == Some(12))
+    assert(cores("local[4, 3]") == Some(4))
+    assert(cores("local[*,3]") == Some(12))
+    assert(cores("local[4]", executorCores = Some("2")) == Some(4))
+
+    // Elsewhere spark.executor.cores is the executor's core count when it is set.
+    assert(cores("spark://host:7077", executorCores = Some("3")) == Some(3))
+    assert(cores("local-cluster[2, 3, 1024]", executorCores = Some("1")) == Some(1))
+    assert(cores("yarn", executorCores = Some("8")) == Some(8))
+    assert(cores("k8s://https://host:6443", executorCores = Some("8")) == Some(8))
+    assert(cores("mesos://host:5050", executorCores = Some("8")) == Some(8))
+
+    // Without it, a standalone executor takes every core its worker offers: the master's cores
+    // per worker in local-cluster mode, and by default all of the machine's processors.
+    assert(cores("local-cluster[2, 3, 1024]") == Some(3))
+    assert(cores("local-cluster[2,3,1024]") == Some(3))
+    assert(cores("spark://host:7077") == Some(12))
+    assert(cores("spark://host1:7077,host2:7077") == Some(12))
+    // YARN and Kubernetes executors default to one core.
+    assert(cores("yarn") == Some(1))
+    assert(cores("k8s://https://host:6443") == Some(1))
+    // Other cluster managers do not tell.
+    assert(cores("mesos://host:5050").isEmpty)
+  }
+
   /**
    * Exports an empty input stream the way a native plan's JVM input is exported, over a reader
    * that counts how many times it is closed. Only the stream's release callback closes it.