[TransferEngine] Support Status return value (#125)

Co-authored-by: luoyuan.luo <[email protected]>

Author: yuan-luo

CMakeLists.txt

@@ -47,7 +47,7 @@ execute_process(
 string(STRIP ${PYTHON_SYS_PATH} PYTHON_SYS_PATH)
 
 set(PYBIND11_FINDPYTHON ON)
-find_package(pybind11 CONFIG REQUIRED) 
+find_package(pybind11 CONFIG REQUIRED)
 
 if (USE_CUDA)
   add_compile_definitions(USE_CUDA)

mooncake-integration/vllm/vllm_adaptor.cpp

@@ -201,13 +201,13 @@ int VLLMAdaptor::transferSync(const char *target_hostname, uintptr_t buffer,
     entry.target_id = handle;
     entry.target_offset = peer_buffer_address;
 
-    int ret = engine_->submitTransfer(batch_id, {entry});
-    if (ret < 0) return -1;
+    Status s = engine_->submitTransfer(batch_id, {entry});
+    if (!s.ok()) return -1;
 
     TransferStatus status;
     while (true) {
-        int ret = engine_->getTransferStatus(batch_id, 0, status);
-        LOG_ASSERT(!ret);
+        Status s = engine_->getTransferStatus(batch_id, 0, status);
+        LOG_ASSERT(s.ok());
         if (status.s == TransferStatusEnum::COMPLETED) {
             engine_->freeBatchID(batch_id);
             return 0;

mooncake-integration/vllm/vllm_adaptor.h

@@ -24,6 +24,7 @@
 #include <stack>
 #include <vector>
 
+#include "common/base/status.h"
 #include "transfer_engine.h"
 #include "transport/rdma_transport/rdma_transport.h"
 #include "transport/transport.h"

mooncake-p2p-store/src/p2pstore/transfer_engine.go

@@ -140,7 +140,7 @@ func (engine *TransferEngine) submitTransfer(batchID BatchID, requests []Transfe
 	}
 
 	ret := C.submitTransfer(engine.engine, C.batch_id_t(batchID), &requestSlice[0], C.size_t(len(requests)))
-	if ret < 0 {
+	if ret != 0 {
 		return ErrTransferEngine
 	}
 	return nil
@@ -149,15 +149,15 @@ func (engine *TransferEngine) submitTransfer(batchID BatchID, requests []Transfe
 func (engine *TransferEngine) getTransferStatus(batchID BatchID, taskID int) (int, uint64, error) {
 	var status C.transfer_status_t
 	ret := C.getTransferStatus(engine.engine, C.batch_id_t(batchID), C.size_t(taskID), &status)
-	if ret < 0 {
+	if ret != 0 {
 		return -1, 0, ErrTransferEngine
 	}
 	return int(status.status), uint64(status.transferred_bytes), nil
 }
 
 func (engine *TransferEngine) freeBatchID(batchID BatchID) error {
 	ret := C.freeBatchID(engine.engine, C.batch_id_t(batchID))
-	if ret < 0 {
+	if ret != 0 {
 		return ErrTransferEngine
 	}
 	return nil

mooncake-store/src/client.cpp

@@ -391,10 +391,10 @@ ErrorCode Client::TransferData(
         return ErrorCode::TRANSFER_FAIL;
     }
 
-    int error_code = transfer_engine_->submitTransfer(batch_id, transfer_tasks);
-    if (error_code != 0) {
+    Status s = transfer_engine_->submitTransfer(batch_id, transfer_tasks);
+    if (!s.ok()) {
         LOG(ERROR) << "Failed to submit all transfers, error code is "
-                   << error_code;
+                   << s.code();
         transfer_engine_->freeBatchID(batch_id);
         return ErrorCode::TRANSFER_FAIL;
     }
@@ -415,11 +415,10 @@ ErrorCode Client::TransferData(
         }
         for (size_t i = 0; i < batch_size; ++i) {
             TransferStatus status;
-            error_code =
-                transfer_engine_->getTransferStatus(batch_id, i, status);
-            if (error_code != 0) {
+            s = transfer_engine_->getTransferStatus(batch_id, i, status);
+            if (!s.ok()) {
                 LOG(ERROR) << "Transfer " << i
-                           << " error, error_code=" << error_code;
+                           << " error, error_code=" << s.code();
                 transfer_engine_->freeBatchID(batch_id);
                 return ErrorCode::TRANSFER_FAIL;
             }

mooncake-transfer-engine/example/transfer_engine_bench.cpp

@@ -24,6 +24,7 @@
 #include <sstream>
 #include <unordered_map>
 
+#include "common/base/status.h"
 #include "transfer_engine.h"
 #include "transport/transport.h"
 
@@ -160,7 +161,7 @@ static inline std::string calculateRate(uint64_t data_bytes, double duration) {
 volatile bool running = true;
 std::atomic<size_t> total_batch_count(0);
 
-int initiatorWorker(TransferEngine *engine, SegmentID segment_id, int thread_id,
+Status initiatorWorker(TransferEngine *engine, SegmentID segment_id, int thread_id,
                     void *addr) {
     bindToSocket(thread_id % NR_SOCKETS);
     TransferRequest::OpCode opcode;
@@ -184,7 +185,7 @@ int initiatorWorker(TransferEngine *engine, SegmentID segment_id, int thread_id,
     size_t batch_count = 0;
     while (running) {
         auto batch_id = engine->allocateBatchID(FLAGS_batch_size);
-        int ret = 0;
+        Status s;
         std::vector<TransferRequest> requests;
         for (int i = 0; i < FLAGS_batch_size; ++i) {
             TransferRequest entry;
@@ -199,14 +200,14 @@ int initiatorWorker(TransferEngine *engine, SegmentID segment_id, int thread_id,
             requests.emplace_back(entry);
         }
 
-        ret = engine->submitTransfer(batch_id, requests);
-        LOG_ASSERT(!ret);
+        s = engine->submitTransfer(batch_id, requests);
+        LOG_ASSERT(s.ok());
         for (int task_id = 0; task_id < FLAGS_batch_size; ++task_id) {
             bool completed = false;
             TransferStatus status;
             while (!completed) {
-                int ret = engine->getTransferStatus(batch_id, task_id, status);
-                LOG_ASSERT(!ret);
+                Status s = engine->getTransferStatus(batch_id, task_id, status);
+                LOG_ASSERT(s.ok());
                 if (status.s == TransferStatusEnum::COMPLETED)
                     completed = true;
                 else if (status.s == TransferStatusEnum::FAILED) {
@@ -217,13 +218,13 @@ int initiatorWorker(TransferEngine *engine, SegmentID segment_id, int thread_id,
             }
         }
 
-        ret = engine->freeBatchID(batch_id);
-        LOG_ASSERT(!ret);
+        s = engine->freeBatchID(batch_id);
+        LOG_ASSERT(s.ok());
         batch_count++;
     }
     LOG(INFO) << "Worker " << thread_id << " stopped!";
     total_batch_count.fetch_add(batch_count);
-    return 0;
+    return Status::OK();
 }
 
 std::string formatDeviceNames(const std::string &device_names) {