Skip to content

feat: support Apache Celeborn remote shuffle service - #5467

Closed
pingzh wants to merge 12 commits into
apache:mainfrom
pingzh:pingzh-support-celeborn-rss
Closed

feat: support Apache Celeborn remote shuffle service#5467
pingzh wants to merge 12 commits into
apache:mainfrom
pingzh:pingzh-support-celeborn-rss

Conversation

@pingzh

@pingzh pingzh commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #5352.

Rationale for this change

Remote shuffle services such as Apache Celeborn are widely used in cloud-native Spark deployments, but Comet’s native shuffle currently cannot integrate with them. This prevents users from combining Comet’s native execution with existing Celeborn infrastructure.

This change enables native Comet shuffle data to be written directly to Celeborn and read back through a native-aware shuffle reader while preserving existing fallback behavior.

What changes are included in this PR?

  • Add a native RSS partition writer that encodes complete Arrow IPC shuffle frames and pushes them through task-scoped JNI callbacks.
  • Introduce destination-aware native shuffle planning and execution.
  • Add a Celeborn partition pusher with bounded frame sizes, executor-wide backpressure, asynchronous completion tracking, cancellation, retries, and cleanup.
  • Implement CometCelebornShuffleManager, native shuffle writers, and a Celeborn-backed shuffle reader.
  • Coordinate shuffle generations, speculative attempts, stage retries, and Spark commit ownership.
  • Support unmodified Apache Celeborn 0.7.0+ clients, including their public mapper-completion and partition-reader APIs.
  • Enable Celeborn-backed native shuffle only for supported plans when native shuffle is explicitly configured.
  • Preserve existing Spark/Celeborn fallback paths for unsupported plans, disabled configurations, and Spark I/O encryption.
  • Add configuration options for Celeborn integration, maximum shuffle frame size, and executor-wide in-flight shuffle bytes.
  • Document configuration requirements and limitations.

Example configuration:

spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManager
spark.celeborn.master.endpoints=celeborn-master:9097
spark.comet.enabled=true
spark.comet.exec.enabled=true
spark.comet.shuffle.enabled=true
spark.comet.shuffle.mode=native

How are these changes tested?

The full Maven reactor passes 136 tests across five suites:

  • CelebornShufflePartitionPusherSuite: 64 tests.
  • CometCelebornNativeShuffleWriterSuite: 11 tests.
  • CometCelebornShuffleManagerSuite: 20 tests.
  • CometCelebornShuffleReaderSuite: 25 tests.
  • CometSparkSessionExtensionsSuite: 16 tests.

Native Rust coverage includes 36 passing tests for shuffle plan compatibility, JNI callbacks, RSS partition writing, shuffle destinations, and native planner integration.

Coverage includes public Celeborn API compatibility, nested Arrow batches, backpressure, asynchronous push completion, cancellation, task failures, speculative execution, shuffle retries, reducer reads, metrics, unsupported-plan fallback, and encrypted shuffle fallback.

The native library builds successfully, and Rust formatting, ScalaStyle, and Spotless checks all pass.

Int.box(abandoned.stageAttempt),
Int.box(abandoned.mapId),
Int.box(abandoned.taskAttempt),
UnknownReason)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Release commit ownership compatibly with Spark 3.4.3. The supported spark-3.4 profile pins Spark 3.4.3, whose OutputCommitCoordinator.taskCompleted does not clear authorizedCommitters when the current owner fails; it calls dagScheduler.stageFailed instead. Consequently invoking it here with UnknownReason leaves the abandoned attempt holding the commit lock, so the original or genuine retry is permanently denied and production shuffle stages can fail. This is already reproducible on the current head: the Spark 3.4 shuffle CI job fails all three ownership/speculation tests (early map ownership, speculation arriving first, and speculative replacement owner). Please implement a Spark-3.4.3-compatible owner-release path and retain regression coverage. AI-assisted review (Codex).

@pingzh

pingzh commented Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

close this , I will break it down to smaller PRs

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support Apache Celeborn remote shuffle service

2 participants