diff --git a/presto-native-execution/CMakeLists.txt b/presto-native-execution/CMakeLists.txt index 55e09fdc20519..9464e27e49835 100644 --- a/presto-native-execution/CMakeLists.txt +++ b/presto-native-execution/CMakeLists.txt @@ -121,6 +121,7 @@ if(PRESTO_ENABLE_CUDF) set(VELOX_ENABLE_PARSE ON) set(VELOX_ENABLE_DUCKDB ON) add_compile_definitions(PRESTO_ENABLE_CUDF) + add_compile_definitions(VELOX_ENABLE_CUDF) enable_language(CUDA) # Determine CUDA_ARCHITECTURES automatically. cmake_policy(SET CMP0104 NEW) diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp index f19081ee50872..183188e70242c 100644 --- a/presto-native-execution/presto_cpp/main/PrestoServer.cpp +++ b/presto-native-execution/presto_cpp/main/PrestoServer.cpp @@ -70,6 +70,7 @@ #ifdef PRESTO_ENABLE_CUDF #include "velox/experimental/cudf/exec/ToCudf.h" +#include "velox/velox/experimental/cudf-exchange/Communicator.h" #endif #ifdef PRESTO_ENABLE_REMOTE_FUNCTIONS @@ -155,19 +156,45 @@ bool isSharedLibrary(const fs::path& path) { return pathExt == kLinuxSharedLibExt || pathExt == kMacOSSharedLibExt; } -void registerVeloxCudf() { +std::shared_ptr registerVeloxCudf() { + std::shared_ptr serverThread = nullptr; #ifdef PRESTO_ENABLE_CUDF facebook::velox::cudf_velox::CudfOptions::getInstance().setPrefix( SystemConfig::instance()->prestoDefaultNamespacePrefix()); facebook::velox::cudf_velox::registerCudf(); + + auto coordURL = SystemConfig::instance()->discoveryUri(); + + if (! coordURL.hasValue()) { + PRESTO_STARTUP_LOG(ERROR) << "No coordinator Uri, can't create CudfExchange"; + } + else { + std::string coordStr = coordURL.value(); + PRESTO_STARTUP_LOG(INFO) << "In registerVeloxCudf " << coordStr; + auto server = facebook::velox::cudf_exchange::Communicator::initAndGet( + SystemConfig::instance()->cudfServerPort(), coordStr); + if (server) { + serverThread = std::make_shared( + &facebook::velox::cudf_exchange::Communicator::run, server.get()); + } + } + PRESTO_STARTUP_LOG(INFO) << "cuDF is registered."; #endif + return serverThread; } -void unregisterVeloxCudf() { +void unregisterVeloxCudf(std::shared_ptr serverThread) { #ifdef PRESTO_ENABLE_CUDF facebook::velox::cudf_velox::unregisterCudf(); PRESTO_SHUTDOWN_LOG(INFO) << "cuDF is unregistered."; + if (serverThread) { + auto server = facebook::velox::cudf_exchange::Communicator::getInstance(); + server->stop(); + server.reset(); + PRESTO_SHUTDOWN_LOG(INFO) << "Joining UCX Communicator thread."; + serverThread->join(); + } #endif } @@ -398,7 +425,7 @@ void PrestoServer::run() { }); } } - registerVeloxCudf(); + auto communicatorThread = registerVeloxCudf(); registerFunctions(); registerRemoteFunctions(); registerVectorSerdes(); @@ -649,7 +676,7 @@ void PrestoServer::run() { unregisterFileReadersAndWriters(); unregisterFileSystems(); unregisterConnectors(); - unregisterVeloxCudf(); + unregisterVeloxCudf(communicatorThread); PRESTO_SHUTDOWN_LOG(INFO) << "Joining Driver CPU Executor '" << driverCpuExecutor_->getName() diff --git a/presto-native-execution/presto_cpp/main/common/Configs.cpp b/presto-native-execution/presto_cpp/main/common/Configs.cpp index 5d8a86af6b596..4762c2353b01d 100644 --- a/presto-native-execution/presto_cpp/main/common/Configs.cpp +++ b/presto-native-execution/presto_cpp/main/common/Configs.cpp @@ -157,6 +157,7 @@ SystemConfig::SystemConfig() { NONE_PROP(kHttpsCertPath), NONE_PROP(kHttpsKeyPath), NONE_PROP(kHttpsClientCertAndKeyPath), + NONE_PROP(kCudfServerPort), NUM_PROP(kExchangeHttpClientNumIoThreadsHwMultiplier, 1.0), NUM_PROP(kExchangeHttpClientNumCpuThreadsHwMultiplier, 1.0), NUM_PROP(kConnectorNumCpuThreadsHwMultiplier, 0.0), @@ -274,6 +275,10 @@ int SystemConfig::httpServerHttpPort() const { return requiredProperty(kHttpServerHttpPort); } +int SystemConfig::cudfServerPort() const { + return requiredProperty(kCudfServerPort); +} + bool SystemConfig::httpServerReusePort() const { return optionalProperty(kHttpServerReusePort).value(); } diff --git a/presto-native-execution/presto_cpp/main/common/Configs.h b/presto-native-execution/presto_cpp/main/common/Configs.h index baaa9dbf06d4c..47de476348266 100644 --- a/presto-native-execution/presto_cpp/main/common/Configs.h +++ b/presto-native-execution/presto_cpp/main/common/Configs.h @@ -628,6 +628,8 @@ class SystemConfig : public ConfigBase { static constexpr std::string_view kHeartbeatFrequencyMs{ "heartbeat-frequency-ms"}; + static constexpr std::string_view kCudfServerPort{"exchange.cudf.server.port"}; + static constexpr std::string_view kExchangeMaxErrorDuration{ "exchange.max-error-duration"}; @@ -771,6 +773,8 @@ class SystemConfig : public ConfigBase { int httpServerHttpPort() const; + int cudfServerPort() const; + bool httpServerReusePort() const; bool httpServerBindToNodeInternalAddressOnlyEnabled() const; diff --git a/presto-native-execution/presto_cpp/main/connectors/CMakeLists.txt b/presto-native-execution/presto_cpp/main/connectors/CMakeLists.txt index 45cba0ced7c7c..48d720a6cecfe 100644 --- a/presto-native-execution/presto_cpp/main/connectors/CMakeLists.txt +++ b/presto-native-execution/presto_cpp/main/connectors/CMakeLists.txt @@ -21,7 +21,7 @@ endif() if(PRESTO_ENABLE_CUDF) target_link_libraries(presto_connectors velox_cudf_parquet_connector - cudf::cudf) + ucxx::ucxx cudf::cudf) endif() target_link_libraries(presto_connectors presto_velox_expr_conversion diff --git a/presto-native-execution/presto_cpp/main/connectors/PrestoToVeloxConnector.cpp b/presto-native-execution/presto_cpp/main/connectors/PrestoToVeloxConnector.cpp index bbd8763138d13..f9c4f29174a58 100644 --- a/presto-native-execution/presto_cpp/main/connectors/PrestoToVeloxConnector.cpp +++ b/presto-native-execution/presto_cpp/main/connectors/PrestoToVeloxConnector.cpp @@ -28,7 +28,6 @@ #include "velox/connectors/hive/iceberg/IcebergSplit.h" #include "velox/connectors/tpch/TpchConnector.h" #include "velox/connectors/tpch/TpchConnectorSplit.h" -#include "velox/experimental/cudf/exec/ToCudf.h" #ifdef PRESTO_ENABLE_CUDF #include "velox/experimental/cudf/connectors/parquet/ParquetConnector.h" diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp index e484dab310148..d421687fc1d45 100644 --- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp +++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp @@ -1938,7 +1938,9 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan( const std::shared_ptr& tableWriteInfo, const protocol::TaskId& taskId) { core::PlanFragment planFragment; + planFragment.prestoId = std::stoi(fragment.id); + std::cout << "!!! PLAN ID: " << fragment.id << std::endl; // Convert the fragment info first. const auto& descriptor = fragment.stageExecutionDescriptor; planFragment.executionStrategy = @@ -1956,7 +1958,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan( if (auto output = std::dynamic_pointer_cast( fragment.root)) { - planFragment.planNode = toVeloxQueryPlan(output, tableWriteInfo, taskId); + planFragment.planNode = toVeloxQueryPlan(output, tableWriteInfo, taskId, (planFragment.prestoId == 0)); return planFragment; } @@ -2007,7 +2009,8 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan( partitionedOutputNodeId, outputType, toVeloxSerdeKind((partitioningScheme.encoding)), - sourceNode); + sourceNode, + (planFragment.prestoId == 0)); return planFragment; case protocol::SystemPartitioning::FIXED: { switch (systemPartitioningHandle->function) { @@ -2021,7 +2024,8 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan( partitionedOutputNodeId, outputType, toVeloxSerdeKind((partitioningScheme.encoding)), - sourceNode); + sourceNode, + (planFragment.prestoId == 0)); return planFragment; } planFragment.planNode = @@ -2034,7 +2038,8 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan( std::make_shared(), outputType, toVeloxSerdeKind((partitioningScheme.encoding)), - sourceNode); + sourceNode, + (planFragment.prestoId == 0)); return planFragment; } case protocol::SystemPartitionFunction::HASH: { @@ -2047,7 +2052,8 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan( partitionedOutputNodeId, outputType, toVeloxSerdeKind((partitioningScheme.encoding)), - sourceNode); + sourceNode, + (planFragment.prestoId == 0)); return planFragment; } planFragment.planNode = @@ -2061,7 +2067,8 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan( inputType, keyChannels, constValues), outputType, toVeloxSerdeKind((partitioningScheme.encoding)), - sourceNode); + sourceNode, + (planFragment.prestoId == 0)); return planFragment; } case protocol::SystemPartitionFunction::BROADCAST: { @@ -2070,7 +2077,8 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan( 1, outputType, toVeloxSerdeKind((partitioningScheme.encoding)), - sourceNode); + sourceNode, + (planFragment.prestoId == 0)); return planFragment; } default: @@ -2089,7 +2097,8 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan( partitionedOutputNodeId, std::move(outputType), toVeloxSerdeKind((partitioningScheme.encoding)), - std::move(sourceNode)); + std::move(sourceNode), + (planFragment.prestoId == 0)); return planFragment; } default: @@ -2108,7 +2117,8 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan( partitionedOutputNodeId, outputType, toVeloxSerdeKind((partitioningScheme.encoding)), - sourceNode); + sourceNode, + (planFragment.prestoId == 0)); return planFragment; } @@ -2124,19 +2134,22 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan( std::shared_ptr(std::move(spec)), toRowType(partitioningScheme.outputLayout, typeParser_), toVeloxSerdeKind((partitioningScheme.encoding)), - sourceNode); + sourceNode, + (planFragment.prestoId == 0)); return planFragment; } core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan( const std::shared_ptr& node, const std::shared_ptr& tableWriteInfo, - const protocol::TaskId& taskId) { + const protocol::TaskId& taskId, + const bool isRootFragment) { return core::PartitionedOutputNode::single( node->id, toRowType(node->outputVariables, typeParser_), VectorSerde::Kind::kPresto, - toVeloxQueryPlan(node->source, tableWriteInfo, taskId)); + toVeloxQueryPlan(node->source, tableWriteInfo, taskId), + isRootFragment); } core::PlanNodePtr VeloxInteractiveQueryPlanConverter::toVeloxQueryPlan( diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h index 6a728a0b4c44c..8d56d957c390d 100644 --- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h +++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h @@ -58,7 +58,8 @@ class VeloxQueryPlanConverterBase { velox::core::PlanNodePtr toVeloxQueryPlan( const std::shared_ptr& node, const std::shared_ptr& tableWriteInfo, - const protocol::TaskId& taskId); + const protocol::TaskId& taskId, + const bool isRootFragment = false); velox::core::PlanNodePtr toVeloxQueryPlan( const std::shared_ptr& node,