afterincomparableyum commented on code in PR #3677:
URL: https://github.com/apache/celeborn/pull/3677#discussion_r3213894560
##########
rust/celeborn-client/src/lib.rs:
##########
@@ -0,0 +1,180 @@
+//! Rust-friendly wrapper around `celeborn-client-sys`.
+
+use celeborn_client_sys::ffi;
+use cxx::UniquePtr;
+
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+ #[error("celeborn ffi error: {0}")]
+ Ffi(#[from] cxx::Exception),
+ #[error("invalid argument: {0}")]
+ InvalidArg(&'static str),
+}
+
+pub type Result<T> = std::result::Result<T, Error>;
+
+/// Configuration for connecting to a Celeborn LifecycleManager.
+pub struct Config {
+ pub app_id: String,
+ /// Max push buffer size in bytes. 0 means use cpp default (64kB).
+ pub push_buffer_max_size: i32,
+ /// Compression codec: "NONE", "LZ4", or "ZSTD".
+ pub shuffle_compression_codec: String,
+}
+
+impl Config {
+ pub fn new(app_id: String) -> Self {
+ Self {
+ app_id,
+ push_buffer_max_size: 0,
+ shuffle_compression_codec: "NONE".to_string(),
+ }
+ }
+}
+
+/// A Rust-friendly Celeborn shuffle client backed by the C++ implementation.
+pub struct ShuffleClient {
+ inner: UniquePtr<ffi::ShuffleClientHandle>,
+}
+
+impl ShuffleClient {
+ /// Connect to a running LifecycleManager at `lm_host:lm_port`.
+ pub fn connect(config: Config, lm_host: &str, lm_port: i32) ->
Result<Self> {
+ if config.app_id.is_empty() {
+ return Err(Error::InvalidArg("app_id is empty"));
+ }
+ if lm_port <= 0 {
+ return Err(Error::InvalidArg("lm_port must be > 0"));
+ }
+ let valid_codecs = ["NONE", "LZ4", "ZSTD"];
+ if !valid_codecs.contains(&config.shuffle_compression_codec.as_str()) {
+ return Err(Error::InvalidArg(
+ "shuffle_compression_codec must be NONE, LZ4, or ZSTD",
+ ));
+ }
+
+ cxx::let_cxx_string!(app_id_cxx = &config.app_id);
+ cxx::let_cxx_string!(codec_cxx = &config.shuffle_compression_codec);
+ let mut handle =
+ ffi::create_client(&app_id_cxx, config.push_buffer_max_size,
&codec_cxx)?;
+
+ cxx::let_cxx_string!(host_cxx = lm_host);
+ ffi::setup_lifecycle_manager(handle.pin_mut(), &host_cxx, lm_port)?;
+
+ Ok(Self { inner: handle })
+ }
+
+ /// Push data for a specific partition.
+ pub fn push_data(
+ &mut self,
+ shuffle_id: i32,
+ map_id: i32,
+ attempt_id: i32,
+ partition_id: i32,
+ data: &[u8],
+ num_mappers: i32,
+ num_partitions: i32,
+ ) -> Result<()> {
+ ffi::push_data(
+ self.inner.pin_mut(),
+ shuffle_id,
+ map_id,
+ attempt_id,
+ partition_id,
+ data,
+ num_mappers,
+ num_partitions,
+ )?;
+ Ok(())
+ }
+
+ /// Signal that a mapper has finished writing all its partitions.
+ pub fn mapper_end(
+ &mut self,
+ shuffle_id: i32,
+ map_id: i32,
+ attempt_id: i32,
+ num_mappers: i32,
+ ) -> Result<()> {
+ ffi::mapper_end(self.inner.pin_mut(), shuffle_id, map_id, attempt_id,
num_mappers)?;
+ Ok(())
+ }
+
+ /// Update reducer file group metadata for a given shuffle.
+ pub fn update_reducer_file_group(&mut self, shuffle_id: i32) -> Result<()>
{
+ ffi::update_reducer_file_group(self.inner.pin_mut(), shuffle_id)?;
+ Ok(())
+ }
+
+ /// Read all data for a partition with full control over parameters.
+ pub fn read_partition(
+ &mut self,
+ shuffle_id: i32,
+ partition_id: i32,
+ attempt_number: i32,
+ start_map_index: i32,
+ end_map_index: i32,
+ ) -> Result<Vec<u8>> {
+ let data = ffi::read_partition_full(
+ self.inner.pin_mut(),
+ shuffle_id,
+ partition_id,
+ attempt_number,
+ start_map_index,
+ end_map_index,
+ )?;
+ Ok(data)
+ }
+
+ /// Convenience: read all map outputs for a partition.
+ #[inline]
+ pub fn read_partition_all(
+ &mut self,
+ shuffle_id: i32,
+ partition_id: i32,
+ num_mappers: i32,
+ ) -> Result<Vec<u8>> {
+ self.read_partition(shuffle_id, partition_id, 0, 0, num_mappers)
+ }
+
+ /// Explicitly shut down the client. Preferred over relying on Drop.
+ ///
+ /// After calling `ffi::shutdown`, the underlying C++ handle is
intentionally
+ /// leaked to avoid a SIGSEGV caused by folly's `EventBase` destruction
race:
+ /// `TransportClient` destructor posts a callback to an `EventBase` that
may
+ /// already be torn down by `IOThreadPoolExecutor::join()`.
+ pub fn shutdown(mut self) -> Result<()> {
+ if let Some(pinned) = self.inner.as_mut() {
+ ffi::shutdown(pinned)?;
Review Comment:
If `ffi::shutdown` returns Err, the ? propagates and self is dropped
normally. `Drop::drop` then calls `ffi::shutdown(pinned)` a second time.
The comment on shutdown says the handle is intentionally leaked to avoid a
folly EventBase use-after-free during destruction, so calling shutdown twice
could re-trigger that teardown path.
I propose either ensuring the handle is leaked even when `ffi::shutdown`
returns an error (so Drop cannot call `ffi::shutdown` a second time), or adding
a “shutdown attempted” flag so Drop skips the shutdown call and only leaks the
handle.
##########
service/src/main/scala/org/apache/celeborn/service/deploy/lifecyclemanager/LifecycleManagerDaemon.scala:
##########
@@ -0,0 +1,120 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.service.deploy.lifecyclemanager
+
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.celeborn.client.LifecycleManager
+import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.internal.Logging
+import org.apache.celeborn.common.util.{SignalUtils, Utils}
+
+private[deploy] object LifecycleManagerDaemon extends Logging {
+
+ private[lifecyclemanager] val shutdownLatch: CountDownLatch = new
CountDownLatch(1)
+
+ private[lifecyclemanager] val currentInstance:
AtomicReference[LifecycleManager] =
+ new AtomicReference[LifecycleManager]()
+
+ private[lifecyclemanager] var exitFn: Int => Unit = (code: Int) =>
System.exit(code)
+
+ def main(args: Array[String]): Unit = {
+ SignalUtils.registerLogger(log)
+ val daemonArgs = LifecycleManagerDaemonArguments.parse(args)
+ val conf = new CelebornConf()
+
+ Utils.loadDefaultCelebornProperties(conf, daemonArgs.propertiesFile.orNull)
+ applyArgsToConf(daemonArgs, conf)
+
+ if (conf.authEnabledOnClient) {
+ logError(
+ "Standalone LifecycleManager does not support auth " +
+ "(cpp/Rust client lacks SASL); set celeborn.auth.enabled=false")
+ exitFn(1)
+ return
+ }
+
+ if (daemonArgs.port < 1024) {
Review Comment:
LifecycleManagerDaemonArguments already sys.exit(1) on this condition. Dead
code.
##########
rust/celeborn-client-sys/src/wrapper.cc:
##########
@@ -0,0 +1,141 @@
+#include "wrapper.h"
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace celeborn_ffi {
+
+std::unique_ptr<ShuffleClientHandle> create_client(
+ const std::string& app_id,
+ int32_t push_buffer_max_size,
+ const std::string& shuffle_compression_codec) {
+ try {
+ auto handle = std::make_unique<ShuffleClientHandle>();
+ handle->app_id = app_id;
+ handle->conf = std::make_shared<celeborn::conf::CelebornConf>();
+
+ if (push_buffer_max_size > 0) {
+ handle->conf->registerProperty(
+ celeborn::conf::CelebornConf::kClientPushBufferMaxSize,
+ std::to_string(push_buffer_max_size) + "b");
+ }
+
+ if (!shuffle_compression_codec.empty()) {
+ handle->conf->registerProperty(
+ celeborn::conf::CelebornConf::kShuffleCompressionCodec,
+ shuffle_compression_codec);
+ }
+
+ handle->endpoint =
std::make_shared<celeborn::client::ShuffleClientEndpoint>(
+ handle->conf);
+ handle->client = celeborn::client::ShuffleClientImpl::create(
+ app_id, handle->conf, *(handle->endpoint));
+
+ return handle;
+ } catch (const std::exception& e) {
+ throw std::runtime_error(std::string("celeborn-ffi: ") + e.what());
+ } catch (...) {
+ throw std::runtime_error("celeborn-ffi: unknown C++ exception");
+ }
+}
+
+void setup_lifecycle_manager(
+ ShuffleClientHandle& handle, const std::string& host, int32_t port) {
+ try {
+ handle.lifecycle_manager_host = host;
+ handle.client->setupLifecycleManagerRef(handle.lifecycle_manager_host,
port);
+ } catch (const std::exception& e) {
+ throw std::runtime_error(std::string("celeborn-ffi: ") + e.what());
+ } catch (...) {
+ throw std::runtime_error("celeborn-ffi: unknown C++ exception");
+ }
+}
+
+void shutdown(ShuffleClientHandle& handle) {
+ try {
+ handle.client->shutdown();
+ } catch (const std::exception& e) {
+ throw std::runtime_error(std::string("celeborn-ffi: ") + e.what());
+ } catch (...) {
+ throw std::runtime_error("celeborn-ffi: unknown C++ exception");
+ }
+}
+
+void push_data(ShuffleClientHandle& handle,
+ int32_t shuffle_id, int32_t map_id, int32_t attempt_id,
+ int32_t partition_id,
+ rust::Slice<const uint8_t> data,
+ int32_t num_mappers, int32_t num_partitions) {
+ try {
+ handle.client->pushData(
+ shuffle_id, map_id, attempt_id, partition_id,
+ data.data(), 0,
+ static_cast<int>(data.size()),
+ num_mappers, num_partitions);
+ } catch (const std::exception& e) {
+ throw std::runtime_error(std::string("celeborn-ffi: ") + e.what());
+ } catch (...) {
+ throw std::runtime_error("celeborn-ffi: unknown C++ exception");
+ }
+}
+
+void mapper_end(ShuffleClientHandle& handle,
+ int32_t shuffle_id, int32_t map_id,
+ int32_t attempt_id, int32_t num_mappers) {
+ try {
+ handle.client->mapperEnd(shuffle_id, map_id, attempt_id, num_mappers);
+ } catch (const std::exception& e) {
+ throw std::runtime_error(std::string("celeborn-ffi: ") + e.what());
+ } catch (...) {
+ throw std::runtime_error("celeborn-ffi: unknown C++ exception");
+ }
+}
+
+void update_reducer_file_group(ShuffleClientHandle& handle, int32_t
shuffle_id) {
+ try {
+ handle.client->updateReducerFileGroup(shuffle_id);
+ } catch (const std::exception& e) {
+ throw std::runtime_error(std::string("celeborn-ffi: ") + e.what());
+ } catch (...) {
+ throw std::runtime_error("celeborn-ffi: unknown C++ exception");
+ }
+}
+
+rust::Vec<uint8_t> read_partition_full(
+ ShuffleClientHandle& handle,
+ int32_t shuffle_id,
+ int32_t partition_id,
+ int32_t attempt_number,
+ int32_t start_map_index,
+ int32_t end_map_index) {
+ try {
+ auto stream = handle.client->readPartition(
+ shuffle_id, partition_id, attempt_number, start_map_index,
end_map_index);
+
+ rust::Vec<uint8_t> out;
+ out.reserve(64 * 1024);
+ std::vector<uint8_t> buf(64 * 1024);
+
+ while (true) {
+ int n = stream->read(buf.data(), 0, buf.size());
+ if (n == -1) {
+ break;
+ }
+ if (n <= 0) {
+ throw std::runtime_error(
+ "celeborn-ffi: CelebornInputStream::read returned unexpected
non-positive " +
+ std::to_string(n));
+ }
+ for (int i = 0; i < n; ++i) {
+ out.push_back(buf[i]);
Review Comment:
Per-byte push_back into rust::Vec<uint8_t> can introduce noticeable overhead
for large partition reads, and the initial reserve(64 * 1024) only avoids
reallocations for the first buffer chunk.
Consider accumulating into a std::vector<uint8_t> (or appending in larger
chunks) and copying once at the end, or resizing the destination per chunk and
using memcpy instead of pushing one byte at a time.
This doesn’t need to be addressed in the current PR, but it may be worth
leaving a TODO here to revisit for performance improvements.
##########
service/src/main/scala/org/apache/celeborn/service/deploy/lifecyclemanager/LifecycleManagerDaemon.scala:
##########
@@ -0,0 +1,120 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.service.deploy.lifecyclemanager
+
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.atomic.AtomicReference
+
+import org.apache.celeborn.client.LifecycleManager
+import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.internal.Logging
+import org.apache.celeborn.common.util.{SignalUtils, Utils}
+
+private[deploy] object LifecycleManagerDaemon extends Logging {
+
+ private[lifecyclemanager] val shutdownLatch: CountDownLatch = new
CountDownLatch(1)
+
+ private[lifecyclemanager] val currentInstance:
AtomicReference[LifecycleManager] =
+ new AtomicReference[LifecycleManager]()
+
+ private[lifecyclemanager] var exitFn: Int => Unit = (code: Int) =>
System.exit(code)
+
+ def main(args: Array[String]): Unit = {
+ SignalUtils.registerLogger(log)
+ val daemonArgs = LifecycleManagerDaemonArguments.parse(args)
+ val conf = new CelebornConf()
+
+ Utils.loadDefaultCelebornProperties(conf, daemonArgs.propertiesFile.orNull)
+ applyArgsToConf(daemonArgs, conf)
+
+ if (conf.authEnabledOnClient) {
+ logError(
+ "Standalone LifecycleManager does not support auth " +
+ "(cpp/Rust client lacks SASL); set celeborn.auth.enabled=false")
+ exitFn(1)
+ return
+ }
+
+ if (daemonArgs.port < 1024) {
+ logError(s"Port must be >= 1024, got ${daemonArgs.port}")
+ exitFn(1)
+ return
+ }
+
+ try {
+ val lifecycleManager = new LifecycleManager(daemonArgs.appId, conf)
+ currentInstance.set(lifecycleManager)
+
+ val shutdownHookThread = new
Thread("LifecycleManagerDaemon-ShutdownHook") {
+ override def run(): Unit = {
+ val watchdogTimeoutMs = conf.appHeartbeatTimeoutMs / 2
+ val watchdog = new Thread("LifecycleManagerDaemon-Watchdog") {
+ setDaemon(true)
+ override def run(): Unit = {
+ try {
+ Thread.sleep(watchdogTimeoutMs)
+ Runtime.getRuntime.halt(2)
+ } catch {
+ case _: InterruptedException => // normal exit
+ }
+ }
+ }
+ watchdog.start()
+
+ try {
+ val instance = currentInstance.get()
+ if (instance != null) {
+ instance.stop()
+ }
+ } catch {
+ case e: Exception =>
+ logError("Error stopping LifecycleManager during shutdown", e)
+ } finally {
+ watchdog.interrupt()
+ shutdownLatch.countDown()
+ }
+ }
+ }
+ Runtime.getRuntime.addShutdownHook(shutdownHookThread)
+
+ // scalastyle:off println
+ println(s"LifecycleManager bound at
${lifecycleManager.getHost}:${lifecycleManager.getPort}")
+ // scalastyle:on println
+
+ shutdownLatch.await()
+ exitFn(0)
+ } catch {
+ case e: Throwable =>
+ logError("Initialize LifecycleManager failed.", e)
+ e.printStackTrace()
+ exitFn(1)
+ }
+ }
+
+ private[lifecyclemanager] def runUntilStopped(lifecycleManager:
LifecycleManager): Unit = {
+ currentInstance.set(lifecycleManager)
+ shutdownLatch.await()
+ }
+
+ private[lifecyclemanager] def applyArgsToConf(
Review Comment:
The --host argument is basically ignored. LifecycleManagerDaemonArguments
parses --host into host: Option[String], but applyArgsToConf only writes
MASTER_ENDPOINTS and CLIENT_SHUFFLE_MANAGER_PORT. The host is never propagated.
LifecycleManager binds to lifecycleHost = Utils.localHostName(conf)
(LifecycleManager.scala:81), and Utils.localHostName only looks at the
CELEBORN_LOCAL_HOSTNAME env or the auto resolved hostname. There's no conf key
path for it.
So someone running start-lifecycle-manager.sh --host 10.0.0.5 gets the
default hostname with no warning. Either drop the option, set
CELEBORN_LOCAL_HOSTNAME from the script, or call Utils.setCustomHostname(host)
before constructing LifecycleManager.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
