[andygrove]

Which issue does this PR close?

N/A. There is no tracking issue.

Rationale for this change

An application has to set spark.shuffle.manager to CometShuffleManager itself. If it does not, Comet disables itself and logs a warning, so the plugin alone gives a user nothing.

On Spark 3.4 and 3.5 the plugin cannot fix this, because SparkEnv.create builds the shuffle manager before any plugin runs. From Spark 4.0 (SPARK-45762) SparkContext creates it right after it initializes the plugins (SparkContext.scala: PluginContainer(...), then _env.initializeShuffleManager()), from the same SparkConf. Executors take their conf from the driver's sc.conf.getAll when they register, which happens after the plugin has run, and their own initializeShuffleManager() also follows their plugin initialization. So the driver plugin can choose the manager for the whole application.

What changes are included in this PR?

• CometDriverPlugin.init calls a new maybeSetShuffleManager, which sets spark.shuffle.manager to CometShuffleManager on Spark 4.0 and later. It does so only when spark.comet.enabled and Comet shuffle are on and the application did not set spark.shuffle.manager. An explicit value is always kept, including sort and the Celeborn manager, and Comet then disables itself as before. It runs before the cache serializer check, which only installs Comet's serializer with one of Comet's shuffle managers.
• Spark 3.4 and 3.5 are unchanged: the application still sets the manager itself.
• It also leaves the manager unset where Comet disables itself whatever the manager: on a big-endian host, with spark.sql.parquet.int96TimestampConversion=true, or when the native library fails to load and spark.comet.nativeLoadRequired is false. Comet would not run there, and Spark's exchanges would copy every row for want of a SortShuffleManager. CometSparkSessionExtensions.unsupportedEnvironmentReason holds these checks, and isCometLoaded and the plugin both call it. If the library fails to load and nativeLoadRequired is true, the plugin still sets the manager and leaves the error to isCometLoaded, as before.
• Docs: the spark.comet.shuffle.enabled description, the shuffle tuning guide, the plugin overview and the FAQ, and a new 1.3.0 entry in the upgrade guide. The entry covers the one behavior change: an application on Spark 4.0 or later that registers the plugin without a shuffle manager used to run without Comet and now runs with it. It needs no spark.comet.legacy.* key, and it points to Configuring Comet Memory, since these applications may have left spark.executor.memoryOverhead at Spark's default.

How are these changes tested?

• CometDriverPluginShuffleManagerSuite tests the decision: the version gate, a user-set value being kept (sort, tungsten-sort, Celeborn, a custom class), and Comet or Comet shuffle disabled, including through the deprecated spark.comet.exec.shuffle.enabled key, and an environment where Comet disables itself (spark.sql.parquet.int96TimestampConversion=true). The suite is in the plugin suite lists of pr_build_linux.yml and pr_build_macos.yml.
• CometPluginsUnifiedModeSuite names no shuffle manager and registers the plugin. A new test checks that on 4.0 and later SparkEnv holds CometShuffleManager and the plan has Comet shuffle, and that on 3.5 Comet stays disabled.
• I ran CometDriverPluginShuffleManagerSuite, CometPluginsUnifiedModeSuite, CometPluginsSuite, CometPluginsSparkShuffleManagerSuite and CometPluginsExtensionOnlySuite on the 4.0 profile, and the first two plus CometSparkSessionExtensionsSuite on 3.5. Executors only run in the local-mode tests, so the executor side rests on reading the Spark 4.0.0 source. It has not been run on a standalone or Kubernetes cluster.

The shuffle manager's handling of shuffles that Comet does not replace is tracked in #7488.

[kazuyukitanimura]

I read the executor side of the Spark source the same way you did. CoarseGrainedSchedulerBackend.sparkProperties is a lazy sc.conf.getAll, and Executor initializes plugins before initializeShuffleManager(). Every test runs in local mode, though, where executors share the driver's SparkEnv, so nothing exercises that path. If the reasoning were wrong, the driver would plan Comet shuffles that SortShuffleManager on the executors can't write. Would you be able to do one manual run on a standalone or Kubernetes cluster with Spark 4.0, the plugin, and no spark.shuffle.manager? Running sc.parallelize(1 to 100, 10).map(_ => org.apache.spark.SparkEnv.get.shuffleManager.getClass.getName).distinct.collect plus a GROUP BY query would do, with the result noted in the description.

docs/source/contributor-guide/plugin_overview.md:43 and docs/source/about/faq.md:130

▎ Two pages still describe the old behavior. plugin_overview.md lists each step CometDriverPlugin takes, and the new shuffle-manager step isn't there. Its Shuffle section also says the manager must be registered. The FAQ bullet "Comet's shuffle manager is not set" says Comet disables itself whenever spark.shuffle.manager does not name Comet's
▎ manager. That's no longer true on Spark 4.0 with CometPlugin. Could you update both in this PR?

docs/source/user-guide/latest/migration-guide.md:66

▎ The applications this entry affects are the ones that never ran Comet, so they're also the most likely to have left spark.executor.memoryOverhead at Spark's default. Comet's native memory outside the off-heap pool comes out of that overhead, and YARN or Kubernetes can kill an executor that runs out of it. The plugin warns at startup, but that
▎ warning is easy to miss after an upgrade. Could the entry link to Configuring Comet Memory?

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala:283

▎ With this PR, CometShuffleManager runs every shuffle in more applications, including RDD shuffles and Spark exchanges that Comet doesn't replace. I found three places where it handles those ordinary shuffles differently from SortShuffleManager. They predate this PR, so could you file a tracking issue and link it here?
▎ - Spark's ShuffleExchangeExec copies every row, because the manager isn't a SortShuffleManager.
▎ - getWriter delegates to the inner sortShuffleManager, which records each map task in its own taskIdMapsForShuffle. unregisterShuffle never calls sortShuffleManager.unregisterShuffle, so those entries stay on the executor heap until the application ends.
▎ - The resolver is built with a fresh map at line 107 instead of taskIdMapsForShuffle. IndexShuffleBlockResolver records blocks migrated in during decommissioning in that map, so unregisterShuffle never deletes them.

[kazuyukitanimura]

Thanks for working this out from the Spark 4.0 source. I checked the same ordering in v4.0.0 and it holds on the driver and the executors. One case concerns me. maybeSetShuffleManager checks spark.comet.enabled and Comet shuffle, but isCometLoaded can still disable Comet afterwards. That happens on a big-endian host, when spark.sql.parquet.int96TimestampConversion=true, or when the native library fails to load and spark.comet.nativeLoadRequired is left at false. In those applications this PR swaps SortShuffleManager for CometShuffleManager, and Comet never runs. Spark's ShuffleExchangeExec.needToCopyObjectsBeforeShuffle returns true for any manager that is not a SortShuffleManager, so every Spark exchange then copies each row before the shuffle. On 1.2 the same application ran on Spark's manager with no copy. Could the plugin skip setting the manager when it can see one of those conditions at startup? A small helper shared with isCometLoaded would keep the two from drifting. A test with one of those settings would lock it in.

[andygrove]

Done in bfc14c7. CometSparkSessionExtensions.unsupportedEnvironmentReason now holds the big-endian, INT96 and native-library checks. isCometLoaded and maybeSetShuffleManager both call it, so the two cannot drift. One difference: if the library fails to load and spark.comet.nativeLoadRequired is true, isCometLoaded throws, but the plugin treats that as supported and still sets the manager, so the failure shows up where it did before rather than at startup. CometDriverPluginShuffleManagerSuite has a test with spark.sql.parquet.int96TimestampConversion=true. The other two conditions can't be set in a unit test. The check only sees the startup value of the INT96 setting, since a session can still change it later.

[kazuyukitanimura]

CometDriverPluginShuffleManagerSuite is a new top-level suite. CI runs the plugin suites from an explicit list passed as -DwildcardSuites in .github/workflows/pr_build_linux.yml (around line 582) and pr_build_macos.yml (around line 248). Nothing in either list matches this name, so the version-gate tests and the keep-the-user's-value tests don't run in CI. Could you add it to both lists? Once it's registered, the run-all-spark-profiles label would cover 3.4 and 4.2, which the local runs in the description didn't include.

[andygrove]

Added to both pr_build_linux.yml and pr_build_macos.yml in bfc14c7.

[comphead]

Thanks @andygrove. I read the ordering again in v4.0.1 and v4.1.1 and it holds. sc.conf is the _conf passed to the driver SparkEnv, executors read sparkProperties lazily after the driver plugin ran, and registerWithExternalShuffleServer takes the class name from the conf. I have not run it on a cluster.

[comphead]

The guide's introduction says every behavior change below names a spark.comet.legacy.* key, and the 1.2.0 section says up front when none is needed. This entry names no key. Running an application in Comet instead of Spark is the "which operators run natively" kind of change that the introduction excludes, so could the 1.3.0 section say it needs no legacy key, as 1.2.0 does?

[andygrove]

The 1.3.0 section now says up front that it needs no spark.comet.legacy.* key, as 1.2.0 does (bfc14c7).

[andygrove]

cc @snmvaughan

[kazuyukitanimura]

Thanks for turning these around so quickly. I noticed nothing has run on bfc14c7 yet. CI (fork) says no run was found on your fork for that commit, and Required Checks here passes without a build. The commit changes isCometLoaded, which every Comet session goes through, so could you start the fork run? I'll add run-all-spark-profiles too, so 3.4,
3.5, 4.0 and 4.2 get covered along with your fork's 4.1 run.

Following up on the cluster run. @comphead and I both read the executor side the same way you did, but none of us has run it. Every test runs in local mode, where executors share the driver's SparkEnv, so nothing in CI covers this. If the reasoning were wrong, every Comet shuffle would fail on the executors. Could you do one run on a standalone or Kubernetes cluster with Spark 4.0, the plugin, and no spark.shuffle.manager before this merges? sc.parallelize(1 to 100, 10).map(_ => org.apache.spark.SparkEnv.get.shuffleManager.getClass.getName).distinct.collect plus a GROUP BY would be enough.