Celeborn Reborn in Rust

Most of what a distributed engine does is local. Shuffle is the exception: every mapper has to reach every reducer, and the way an engine moves that data decides whether a job can survive losing a worker and how much cluster it has to hold open while it runs.

Sail 0.7 made blocking shuffle a first-class execution mode, writing shuffle data to object storage so mapper and reducer stages no longer have to run at the same time. Sail 0.7.1 adds a second backend for that mode: Apache Celeborn, the remote shuffle service many teams already run alongside Spark.

The client lives in a new sail-celeborn crate, a Rust reimplementation of Celeborn’s lifecycle manager and shuffle client that speaks the same wire protocol as an existing JVM-based Celeborn cluster.

Shuffle Services for Distributed Compute

A shuffle operation involves M mappers and R reducers exchanging data across M × R logical streams.

Sail supports both pipelined shuffle and blocking shuffle. With pipelined shuffle, mapper and reducer tasks launch concurrently, and data is exchanged through Arrow Flight without durable materialization. Blocking shuffle takes a different approach: shuffle data is first written to storage and later read by a subsequent stage. Blocking shuffle improves job resilience and reduces concurrent compute resource requirements because mapper and reducer stages do not need to run simultaneously. The tradeoff is data materialization.

The central challenge in blocking shuffle is avoiding an M × R file layout, which would substantially amplify I/O costs. Sail’s storage-based blocking shuffle addresses this by inserting a distributed merge operation in the mapper stage. This operation gathers data for the same partition from every mapper in a streaming fashion, before writing the partition for each reducer to consume.

More generally, distributed systems often solve this problem through a dedicated shuffle service. Such a service receives streams from mappers, consolidates them, and redistributes partitions in a form readers can consume directly. It serves essentially the same purpose as Sail’s merge operation, while being independently maintainable and scalable.

Introducing Celeborn

Apache Celeborn is one such shuffle service. A Celeborn cluster consists of a master and workers. The master, or multiple masters in a high-availability deployment, coordinates application registration and shuffle slot allocation. Workers manage shuffle storage and serve the corresponding read and write requests.

Celeborn supports various compute engines via its client-side integration, which comprises two main components. The lifecycle manager registers the application and coordinates shuffle registration and slot allocation with the Celeborn master. The shuffle client interacts with Celeborn workers to read and write shuffle data. In Celeborn’s integrations for Apache Spark, the lifecycle manager runs in the Spark driver while the shuffle client runs in Spark executors.

Integrating Celeborn with Sail

As a unified distributed compute engine for the composable data stack, Sail is designed to integrate with established infrastructure. Starting with Sail 0.7.1, Sail supports Celeborn as another blocking shuffle backend.

The Celeborn integration reaches deeply into both Sail’s control plane and data plane. While the original Celeborn client is implemented in Java and Scala, Sail reimplements this logic in Rust so that it can run within the Sail engine without a JVM. Celeborn’s master and worker protocols involve TCP connections with Protocol Buffers message encoding and an additional binary framing layer. Sail implements the necessary encoding and decoding logic, allowing the Rust-based lifecycle manager and shuffle client to communicate directly with an existing JVM-based Celeborn cluster. Besides the transport layer, the lifecycle manager also has heavy business logic for interacting with the master during shuffle registration and failure recovery. We replicated the core logic in Rust and validated it with container-based integration tests.

The Celeborn client, split across the runtime boundary. The lifecycle manager runs in the Sail driver and a shuffle client runs in each Sail worker, reaching each other over Sail’s existing gRPC channel. Both also reach into the Celeborn cluster, but for different work: the lifecycle manager reserves slots and commits files, while the shuffle client moves the bytes.

Writing and Reading Shuffle Data

Sail implements both Celeborn client components using its actor model. The Sail driver actor owns the lifecycle manager actor, while each worker actor creates a shuffle client actor. Each actor processes mailbox messages sequentially and mutates its internal state without locks. This makes the implementation extensible and easy to reason about while porting the client logic from the upstream project.

During a shuffle write, the shuffle client contacts the lifecycle manager in the driver to obtain a shuffle ID. If necessary, the lifecycle manager registers the shuffle with the Celeborn master. The shuffle client then communicates directly with Celeborn workers to push data.

When a shuffle write task finishes, its shuffle client sends a mapper-end notification to the lifecycle manager. After the lifecycle manager receives a completion notification for every mapper, it commits the shuffle data on the relevant Celeborn workers.

After the write stage completes, Sail can schedule the read stage. The shuffle reader uses the shuffle client to obtain shuffle partition location information from the lifecycle manager, then contacts workers directly to fetch the data. Celeborn has already merged data from all mapper tasks for the relevant reducer partition, so each reducer partition is read as a single logical stream rather than M separate mapper streams.

The lifecycle manager and shuffle client communicate through Sail’s existing driver-worker service, which exposes Sail’s internal Celeborn-specific endpoints over the multiplexed gRPC channel. This is an improvement over the Java- and Scala-based Spark integration, where the lifecycle manager runs as a service in the Spark driver and requires an additional TCP port.

Celeborn does not prescribe the format of the actual shuffle data being managed. Sail serializes record batches as Arrow IPC and uses Celeborn’s configured push-data compression. Each reducer can therefore read a single stream of Arrow record batches gathered from the corresponding partitions of all mappers.

Sail’s Celeborn integration also supports failure recovery. For retryable push failures, the shuffle client asks the lifecycle manager to “revive” the affected partition. The lifecycle manager obtains a replacement partition location with a higher “epoch” from the master, and the client updates its routing information and retries when needed. Readers use the location history, including epochs, to locate shuffle data.

Celeborn Shuffle Backend in Action

Getting started with Celeborn-based blocking shuffle is easy. Define the following environment variables for the Sail server to specify celeborn as the shuffle backend and the master endpoint for your Celeborn cluster.

SAIL_CLUSTER__SHUFFLE_BACKEND__TYPE=celeborn
SAIL_CLUSTER__SHUFFLE_BACKEND__CELEBORN__MASTER_ENDPOINTS='["celeborn-master:12097"]'

The Celeborn backend is useful when storage-based blocking shuffle does not fit an existing data stack. For example, remote storage may not provide the throughput or latency required by a compute deployment, while an organization may already operate an Apache Celeborn cluster for Spark workloads. This separates scalable shuffle-data exchange from compute execution, allowing Sail to use an existing Celeborn deployment.

Getting Started with Sail

The Celeborn backend ships in Sail 0.7.1. Install or upgrade with pip install pysail==0.7.1, or see the installation guide for standalone binary and Docker options. The shuffle guide covers backend configuration in detail.

Join the Community

Sail wouldn’t have its current shape without its community, and we welcome contributions of all kinds: code, feature ideas, and especially bug reports with reproducible examples. Stay tuned for more feature announcements. You can follow along on GitHub or by joining our Slack Community.

Want to run Sail without operating the infrastructure? The LakeSail Platform offers fully managed Sail with built-in governance, observability, and BYOC deployment. Get started with a 14-day free trial.

Apache®, Apache Celeborn™, Apache Spark™, and Apache Arrow™ are trademarks (or registered trademarks) of The Apache Software Foundation. LakeSail, Inc. is an independent company and is not affiliated with, endorsed by, or sponsored by The Apache Software Foundation.