Skip to content

feat: complete Celeborn map-side shuffle push lifecycle (6/n) - #5513

Open
pingzh wants to merge 1 commit into
apache:mainfrom
pingzh:pingzh-celeborn-map-side-push-lifecycle-part-6
Open

feat: complete Celeborn map-side shuffle push lifecycle (6/n)#5513
pingzh wants to merge 1 commit into
apache:mainfrom
pingzh:pingzh-celeborn-map-side-push-lifecycle-part-6

Conversation

@pingzh

@pingzh pingzh commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Part of #5352. This is the sixth foundational PR and does not close the issue.

Previous PRs:

Rationale for this change

The existing Celeborn integration provides the shuffle manager and partition-pusher foundation but does not yet complete the native map-side shuffle lifecycle.

Production use requires bounded frames, executor-wide backpressure, asynchronous completion tracking, cancellation handling, speculative-attempt coordination, retry-safe commits, and support for nested Arrow data.

What changes are included in this PR?

  • Wire native Comet shuffle writers into Celeborn without creating local shuffle files.
  • Add executor-wide byte admission covering native encoding, JNI copies, and Celeborn transport.
  • Track asynchronous push completion and preserve admission until requests actually complete.
  • Handle task cancellation, mapper completion, cleanup, and asynchronous transport failures.
  • Coordinate shuffle generations, speculative attempts, retries, and Spark output commits.
  • Handle stale commit ownership on Spark 3.4.
  • Support Celeborn 0.6 and 0.7 mapper-completion APIs.
  • Preserve existing CRC accounting, encrypted payload handling, and generic JNI callback compatibility.
  • Split oversized batches into complete row-aligned frames and compact nested Arrow dictionaries, lists, maps, and structs.
  • Add configurable limits:
    • spark.comet.shuffle.rss.maxFrameBytes
    • spark.comet.shuffle.rss.maxInFlightBytes
  • Register the new native shuffle writer suite in Linux and macOS CI.

How are these changes tested?

  • 101 JVM tests across:
    • CelebornShufflePartitionPusherSuite
    • CometCelebornShuffleManagerSuite
    • CometCelebornNativeShuffleWriterSuite
    • CometConfSuite
    • CometNativeShuffleSuite
  • 67 Rust shuffle tests, including bounded frames, nested Arrow compaction, backpressure, failed writers, and legacy callbacks.
  • 17 native planner tests covering local and remote shuffle destinations.
  • 3 JNI bridge tests.
  • Scala formatting, Rust formatting, Scala style checks, and CI suite-registration checks all pass.

@ziting-openai ziting-openai left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Native-first Celeborn initialization bypasses configured Spark shuffle encryption, and fixed maximum-frame admission serializes encoding under the default executor limits.

Comment thread spark/src/main/scala/org/apache/comet/shuffle/CelebornShufflePusherFactory.scala Outdated
Comment thread native/shuffle/src/writers/rss/rss_partition_writer.rs Outdated
@pingzh
pingzh force-pushed the pingzh-celeborn-map-side-push-lifecycle-part-6 branch 3 times, most recently from 16d7167 to 520ca35 Compare August 27, 2026 23:08
@pingzh
pingzh force-pushed the pingzh-celeborn-map-side-push-lifecycle-part-6 branch from 520ca35 to be5f4af Compare August 28, 2026 01:14
Comment thread native/shuffle/src/writers/rss/rss_partition_writer.rs Outdated
Comment thread spark/src/main/java/org/apache/comet/shuffle/CelebornShufflePartitionPusher.java Outdated
@pingzh
pingzh force-pushed the pingzh-celeborn-map-side-push-lifecycle-part-6 branch from be5f4af to 93b8336 Compare August 28, 2026 02:36

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Follow-up review at 93b8336ed2f61dfc82d177b6e1bb994eb7e341aa against base eabb5d4773091b983d8fce713f0e34b1cf93f877: verified the fixes for both prior admission P2s and rechecked the earlier fixes. No remaining P1/P2 found.

Validation: native core; 102 Rust tests; 128 JVM tests each on Spark 3.4, 3.5, and 4.0, and 76 each on 4.1 and 4.2; 96 allocation, 74 actual-JNI, and 48 stock Celeborn controls with old/base causal comparisons.

Tests used local/controlled transport, not a deployed Celeborn cluster. Encrypted native RSS is explicitly rejected; ordinary Spark crypto handling is retained. CI is still running.

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.

3 participants