[TransferEngine] feature: introduce USE_NVMEOF to enable NVMe-oF separately. (kvcache-ai#106)

Signed-off-by: doujiang24 <doujiang24@gmail.com>

Author: doujiang24

CMakeLists.txt

@@ -27,6 +27,7 @@ add_definitions(-DCONFIG_ERDMA)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 option(USE_CUDA "option for using gpu direct" OFF)
+option(USE_NVMEOF "option for using NVMe over Fabric" OFF)
 option(USE_CXL "option for using cxl protocol" OFF)
 option(USE_ETCD "option for enable etcd as metadata server" ON)
 option(USE_REDIS "option for enable redis as metadata server" OFF)
@@ -38,7 +39,16 @@ option(WITH_RUST_EXAMPLE "build the Rust interface and sample code for the trans
 if (USE_CUDA)
   add_compile_definitions(USE_CUDA)
   message(STATUS "CUDA support is enabled")
+
+  if (USE_NVMEOF)
+    add_compile_definitions(USE_NVMEOF)
+    message(STATUS "NVMe-oF support is enabled")
+  endif()
+
   include_directories(/usr/local/cuda/include)
+  link_directories(/usr/local/cuda/lib)
+elseif(USE_NVMEOF)
+  message(FATAL_ERROR "Cannot enable USE_NVMEOF without USE_CUDA")
 endif()
 
 if (USE_REDIS)

mooncake-transfer-engine/example/transfer_engine_bench.cpp

@@ -30,7 +30,10 @@
 #ifdef USE_CUDA
 #include <bits/stdint-uintn.h>
 #include <cuda_runtime.h>
+
+#ifdef USE_NVMEOF
 #include <cufile.h>
+#endif
 
 #include <cassert>
 
@@ -41,7 +44,6 @@ static void checkCudaError(cudaError_t result, const char *message) {
         exit(EXIT_FAILURE);
     }
 }
-
 #endif
 
 const static int NR_SOCKETS =
@@ -258,7 +260,7 @@ std::string loadNicPriorityMatrix() {
            " \"cpu:1\": [[" +
            device_names +
            "], []], "
-           " \"gpu:0\": [[" +
+           " \"cuda:0\": [[" +
            device_names + "], []]}";
 }
 
@@ -294,7 +296,7 @@ int initiator() {
     if (FLAGS_use_vram) LOG(INFO) << "VRAM is used";
     for (int i = 0; i < buffer_num; ++i) {
         addr[i] = allocateMemoryPool(FLAGS_buffer_size, i, FLAGS_use_vram);
-        std::string name_prefix = FLAGS_use_vram ? "gpu:" : "cpu:";
+        std::string name_prefix = FLAGS_use_vram ? "cuda:" : "cpu:";
         int rc = engine->registerLocalMemory(addr[i], FLAGS_buffer_size,
                                              name_prefix + std::to_string(i));
         LOG_ASSERT(!rc);

mooncake-transfer-engine/src/CMakeLists.txt

@@ -27,5 +27,8 @@ target_link_libraries(transfer_engine PUBLIC transport rdma_transport ibverbs gl
 
 if (USE_CUDA)
   target_include_directories(transfer_engine PRIVATE /usr/local/cuda/include)
-  target_link_libraries(transfer_engine PUBLIC nvmeof_transport cuda cufile cudart rt)
+  target_link_libraries(transfer_engine PUBLIC cuda cudart rt)
+  if (USE_NVMEOF)
+    target_link_libraries(transfer_engine PUBLIC nvmeof_transport cufile)
+  endif()
 endif()

mooncake-transfer-engine/src/multi_transport.cpp

@@ -17,7 +17,7 @@
 #include "transport/rdma_transport/rdma_transport.h"
 #include "transport/tcp_transport/tcp_transport.h"
 #include "transport/transport.h"
-#ifdef USE_CUDA
+#ifdef USE_NVMEOF
 #include "transport/nvmeof_transport/nvmeof_transport.h"
 #endif
 
@@ -109,8 +109,7 @@ int MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id,
     status.transferred_bytes = task.transferred_bytes;
     uint64_t success_slice_count = task.success_slice_count;
     uint64_t failed_slice_count = task.failed_slice_count;
-    if (success_slice_count + failed_slice_count ==
-        task.slice_count) {
+    if (success_slice_count + failed_slice_count == task.slice_count) {
         if (failed_slice_count) {
             status.s = Transport::TransferStatusEnum::FAILED;
         } else {
@@ -131,7 +130,7 @@ Transport *MultiTransport::installTransport(const std::string &proto,
     } else if (std::string(proto) == "tcp") {
         transport = new TcpTransport();
     }
-#ifdef USE_CUDA
+#ifdef USE_NVMEOF
     else if (std::string(proto) == "nvmeof") {
         transport = new NVMeoFTransport();
     }

mooncake-transfer-engine/src/transport/CMakeLists.txt

@@ -6,7 +6,7 @@ add_library(transport OBJECT ${XPORT_SOURCES} $<TARGET_OBJECTS:rdma_transport>)
 add_subdirectory(tcp_transport)
 target_sources(transport PUBLIC $<TARGET_OBJECTS:tcp_transport>)
 
-if (USE_CUDA)
+if (USE_NVMEOF)
   add_subdirectory(nvmeof_transport)
   target_sources(transport PUBLIC $<TARGET_OBJECTS:nvmeof_transport>)
 endif()

mooncake-transfer-engine/src/transport/nvmeof_transport/cufile_context.cpp

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifdef USE_CUDA
+#ifdef USE_NVMEOF
 
 // TBD
 

mooncake-transfer-engine/tests/rdma_transport_test.cpp

@@ -36,7 +36,10 @@
 #ifdef USE_CUDA
 #include <bits/stdint-uintn.h>
 #include <cuda_runtime.h>
+
+#ifdef USE_NVMEOF
 #include <cufile.h>
+#endif
 
 #include <cassert>
 
@@ -47,7 +50,6 @@ static void checkCudaError(cudaError_t result, const char *message) {
         exit(EXIT_FAILURE);
     }
 }
-
 #endif
 
 #define NR_SOCKETS (1)
@@ -234,7 +236,7 @@ std::string loadNicPriorityMatrix() {
            " \"cpu:1\": [[" +
            device_names +
            "], []], "
-           " \"gpu:0\": [[" +
+           " \"cuda:0\": [[" +
            device_names + "], []]}";
 }
 
@@ -268,7 +270,7 @@ int initiator() {
 #ifdef USE_CUDA
     addr = allocateMemoryPool(ram_buffer_size, 0, FLAGS_use_vram);
     int rc = engine->registerLocalMemory(addr, ram_buffer_size,
-                                         FLAGS_use_vram ? "gpu:0" : "cpu:0");
+                                         FLAGS_use_vram ? "cuda:0" : "cpu:0");
     LOG_ASSERT(!rc);
 #else
     addr = allocateMemoryPool(ram_buffer_size, 0, false);

mooncake-transfer-engine/tests/rdma_transport_test2.cpp

@@ -90,7 +90,7 @@ std::string loadNicPriorityMatrix() {
            " \"cpu:1\": [[" +
            device_names +
            "], []], "
-           " \"gpu:0\": [[" +
+           " \"cuda:0\": [[" +
            device_names + "], []]}";
 }
 

mooncake-transfer-engine/tests/tcp_transport_test.cpp

@@ -25,7 +25,10 @@
 #ifdef USE_CUDA
 #include <bits/stdint-uintn.h>
 #include <cuda_runtime.h>
+
+#ifdef USE_NVMEOF
 #include <cufile.h>
+#endif
 
 #include <cassert>
 
@@ -36,7 +39,6 @@ static void checkCudaError(cudaError_t result, const char *message) {
         exit(EXIT_FAILURE);
     }
 }
-
 #endif
 
 #include "transfer_engine.h"