Feature/offloadrpc

mooncake-integration/store/mooncake_perf_points.def

@@ -108,6 +108,16 @@ PERF_KEY_DEF(GET_SSD_OFFLOAD_RPC,            "real_client.cpp::batch_get_into_of
 PERF_KEY_DEF(GET_SSD_TRANSFER_DATA,          "real_client.cpp::batch_get_into_offload_object_internal",  "TransferData")
 PERF_KEY_DEF(GET_SSD_RELEASE_BUFFER,         "real_client.cpp::batch_get_into_offload_object_internal",  "ReleaseBuffer")
 
+// === Owner-side offload sub-steps (pull & push) ===
+PERF_KEY_DEF(GET_SSD_OWNER_READ,             "file_storage.cpp::BatchGet",                               "OwnerSsdRead")
+PERF_KEY_DEF(GET_SSD_OWNER_ALLOC,            "file_storage.cpp::BatchGet",                               "OwnerAllocBuffer")
+PERF_KEY_DEF(GET_SSD_OWNER_LOAD,             "file_storage.cpp::BatchGet",                               "OwnerDiskLoad")
+PERF_KEY_DEF(GET_SSD_OWNER_LOAD_PLAN,        "storage_backend.cpp::BatchLoad",                           "OwnerLoadPlan")
+PERF_KEY_DEF(GET_SSD_OWNER_LOAD_URING,       "storage_backend.cpp::BatchLoad",                           "OwnerLoadUring")
+PERF_KEY_DEF(GET_SSD_OWNER_LOAD_POSIX,       "storage_backend.cpp::BatchLoad",                           "OwnerLoadPosix")
+PERF_KEY_DEF(GET_SSD_OWNER_PUSH_WRITE,       "real_client.cpp::batch_get_offload_object_push",           "OwnerPushWrite")
+PERF_KEY_DEF(GET_SSD_OWNER_RELEASE,          "file_storage.cpp::ReleaseBuffer",                          "OwnerReleaseBuffer")
+
 // === UB / URMA endpoint first connection path ===
 PERF_KEY_DEF(UB_HANDSHAKE_ENCODE,            "transfer_metadata.cpp::TransferHandshakeUtil::encode",     "Encode")
 PERF_KEY_DEF(UB_HANDSHAKE_DECODE,            "transfer_metadata.cpp::TransferHandshakeUtil::decode",     "Decode")

mooncake-store/include/client_service.h

@@ -497,6 +497,21 @@ class Client {
         const std::unordered_map<std::string, std::vector<Slice>>&
             batch_slices);
 
+    /**
+     * @brief Push counterpart of BatchGetOffloadObject, run on the data owner.
+     * WRITEs each key's staged ClientBuffer blob (src_pointers[i]) directly
+     * into the requester's destination slices over the transfer engine.
+     * @param requester_te_addr The requester's Transfer Engine endpoint.
+     * @param keys List of keys (one per src pointer / dst slice group).
+     * @param src_pointers Owner-local ClientBuffer addresses, one per key.
+     * @param dst_slices Per-key destination slices on the requester.
+     */
+    tl::expected<void, ErrorCode> BatchPushOffloadObject(
+        const std::string& requester_te_addr,
+        const std::vector<std::string>& keys,
+        const std::vector<uint64_t>& src_pointers,
+        const std::vector<std::vector<OffloadDstSlice>>& dst_slices);
+
     /**
      * @brief Notifies the master that offloading of specified objects has
      * succeeded.

mooncake-store/include/pyclient.h

@@ -153,6 +153,18 @@ class ClientRequester {
                              const std::vector<std::string> &keys,
                              const std::vector<int64_t> &sizes);
 
+    /**
+     * @brief Push-mode offload: invoke the owner's push handler, which WRITEs
+     * the data straight into our destination slices. A single RPC replaces the
+     * pull path's get-pointers + READ + release_offload_buffer sequence.
+     * @param client_addr Network address of the remote FileStorage owner.
+     * @param req keys/sizes plus our TE endpoint and destination slices.
+     */
+    tl::expected<BatchGetOffloadObjectPushResponse, ErrorCode>
+    batch_get_offload_object_push(
+        const std::string &client_addr,
+        const BatchGetOffloadObjectPushRequest &req);
+
     /**
      * @brief Notifies remote FileStorage to release buffer after transfer
      * completion. This is a fire-and-forget call - errors are logged but not

mooncake-store/include/real_client.h

@@ -682,6 +682,20 @@ class RealClient : public PyClient {
     batch_get_offload_object(const std::vector<std::string> &keys,
                              const std::vector<int64_t> &sizes);
 
+    /**
+     * @brief Push-mode offload handler, run on the data owner. Reads the
+     * requested keys from SSD into the local ClientBuffer, WRITEs them
+     * directly into the requester's destination slices over the transfer
+     * engine, then releases the ClientBuffer locally. The requester therefore
+     * needs no follow-up READ nor a separate release_offload_buffer RPC.
+     * @param req keys/sizes plus the requester's TE endpoint and dst slices.
+     * @return per-batch result; data already landed in requester memory.
+     */
+    async_simple::coro::Lazy<
+        tl::expected<BatchGetOffloadObjectPushResponse, ErrorCode>>
+    batch_get_offload_object_push(
+        const BatchGetOffloadObjectPushRequest &req);
+
     /**
      * @brief Releases buffer associated with a specific batch_id.
      * Called by remote client after transfer completion.

mooncake-store/include/rpc_types.h

@@ -273,4 +273,42 @@ struct BatchGetOffloadObjectResponse {
 YLT_REFL(BatchGetOffloadObjectResponse, batch_id, pointers,
          transfer_engine_addr, gc_ttl_ms);
 
+// One destination slice on the requester side. In push mode the data owner
+// writes (one-sided) the on-disk blob of a key directly into these addresses.
+struct OffloadDstSlice {
+    uint64_t addr;
+    uint64_t size;
+
+    OffloadDstSlice() : addr(0), size(0) {}
+    OffloadDstSlice(uint64_t addr_param, uint64_t size_param)
+        : addr(addr_param), size(size_param) {}
+};
+YLT_REFL(OffloadDstSlice, addr, size);
+
+// Push-mode offload request. Unlike the pull path (where the requester gets
+// back ClientBuffer pointers and issues the RDMA READ itself), here the
+// requester hands the owner its own transfer engine endpoint and destination
+// slice addresses. The owner reads SSD into its registered ClientBuffer and
+// then WRITEs straight into the requester's memory, so the requester needs no
+// follow-up READ nor a separate release_offload_buffer RPC.
+struct BatchGetOffloadObjectPushRequest {
+    std::vector<std::string> keys;  // tenant-scoped storage keys
+    std::vector<int64_t> sizes;     // total bytes per key (for FileStorage)
+    std::string requester_te_addr;  // requester's transfer engine endpoint
+    std::vector<std::vector<OffloadDstSlice>> dst_slices;  // per-key dst slices
+
+    BatchGetOffloadObjectPushRequest() = default;
+};
+YLT_REFL(BatchGetOffloadObjectPushRequest, keys, sizes, requester_te_addr,
+         dst_slices);
+
+struct BatchGetOffloadObjectPushResponse {
+    ErrorCode error_code;  // overall result; data is already in requester memory
+
+    BatchGetOffloadObjectPushResponse() : error_code(ErrorCode::OK) {}
+    explicit BatchGetOffloadObjectPushResponse(ErrorCode error_code_param)
+        : error_code(error_code_param) {}
+};
+YLT_REFL(BatchGetOffloadObjectPushResponse, error_code);
+
 }  // namespace mooncake

mooncake-store/include/transfer_task.h

@@ -18,6 +18,7 @@
 #include "transfer_engine.h"
 #include "types.h"
 #include "replica.h"
+#include "rpc_types.h"
 #include "storage_backend.h"
 #include "client_metric.h"
 #ifdef USE_NOF
@@ -585,6 +586,16 @@ class TransferSubmitter {
         const std::unordered_map<std::string, std::vector<Slice>>&
             batched_slices);
 
+    // Push counterpart of submit_batch_get_offload_object, run on the data
+    // owner. WRITEs each key's on-disk blob (staged in the owner's local
+    // ClientBuffer at src_pointers[i]) directly into the requester's
+    // destination slices, opening the requester's segment once.
+    std::optional<TransferFuture> submit_batch_push_offload_object(
+        const std::string& requester_te_addr,
+        const std::vector<std::string>& keys,
+        const std::vector<uint64_t>& src_pointers,
+        const std::vector<std::vector<OffloadDstSlice>>& dst_slices);
+
     /**
      * @brief Pure comparison helper: returns true iff both endpoints are
      * non-empty and identical. Exposed for unit testing of the locality

mooncake-store/src/client_service.cpp

@@ -3223,6 +3223,26 @@ tl::expected<void, ErrorCode> Client::BatchGetOffloadObject(
     return {};
 }
 
+tl::expected<void, ErrorCode> Client::BatchPushOffloadObject(
+    const std::string& requester_te_addr,
+    const std::vector<std::string>& keys,
+    const std::vector<uint64_t>& src_pointers,
+    const std::vector<std::vector<OffloadDstSlice>>& dst_slices) {
+    auto future = transfer_submitter_->submit_batch_push_offload_object(
+        requester_te_addr, keys, src_pointers, dst_slices);
+    if (!future) {
+        MC_LOG(ERROR) << "Failed to submit push transfer operation";
+        return tl::make_unexpected(ErrorCode::TRANSFER_FAIL);
+    }
+    MC_VLOG(1) << "Using transfer strategy: " << future->strategy();
+    auto result = future->get();
+    if (result != ErrorCode::OK) {
+        MC_LOG(ERROR) << "Push transfer failed, error code is " << result;
+        return tl::make_unexpected(result);
+    }
+    return {};
+}
+
 tl::expected<void, ErrorCode> Client::NotifyOffloadSuccess(
     const std::vector<std::string>& keys,
     const std::vector<StorageObjectMetadata>& metadatas) {

mooncake-store/src/file_storage.cpp

@@ -12,6 +12,14 @@
 #include "file_interface.h"
 #endif
 
+// ubdiag perf points for the offload owner-side SSD path. UBDIAG_PROGRAM_NAME
+// must match the other TUs (store_py/real_client/client_service); each TU's
+// static initializer writes the shared detail::AutoProgramName(), and leaving
+// it nullptr here could clobber the program name depending on init order.
+#define UBDIAG_PERF_DEF_FILE "mooncake_perf_points.def"
+#define UBDIAG_PROGRAM_NAME "mooncake_store"
+#include "ubdiag/auto_perf.h"
+
 namespace mooncake {
 
 using gpu_staging::CopyDeviceToHost;
@@ -323,13 +331,28 @@ tl::expected<void, ErrorCode> FileStorage::Init() {
 tl::expected<FileStorage::BatchGetResult, ErrorCode> FileStorage::BatchGet(
     const std::vector<std::string>& keys, const std::vector<int64_t>& sizes) {
     auto start_time = std::chrono::steady_clock::now();
+    // Owner-side SSD read total (OwnerSsdRead); broken down into buffer
+    // allocation (OwnerAllocBuffer) and disk load (OwnerDiskLoad). End() is on
+    // the success path only — the error early-returns let the dtor Abandon the
+    // unfinished total sample.
+    UbDiag::PerfPoint pt_read(PerfKey::GET_SSD_OWNER_READ,
+                              UbDiag::PerfLevel::MODULE);
+    pt_read.Start();
+    UbDiag::PerfPoint pt_alloc(PerfKey::GET_SSD_OWNER_ALLOC,
+                               UbDiag::PerfLevel::MODULE);
+    pt_alloc.Start();
     auto allocate_res = AllocateBatch(keys, sizes);
+    pt_alloc.End(allocate_res ? 0 : -1);
     if (!allocate_res) {
         LOG(ERROR) << "Failed to allocate batch objects";
         return tl::make_unexpected(allocate_res.error());
     }
     auto allocated_batch = allocate_res.value();
+    UbDiag::PerfPoint pt_load(PerfKey::GET_SSD_OWNER_LOAD,
+                              UbDiag::PerfLevel::MODULE);
+    pt_load.Start();
     auto result = BatchLoad(allocated_batch->slices);
+    pt_load.End(result ? 0 : -1);
     if (!result) {
         LOG(ERROR) << "Batch load object failed,err_code = " << result.error();
         return tl::make_unexpected(result.error());
@@ -358,6 +381,7 @@ tl::expected<FileStorage::BatchGetResult, ErrorCode> FileStorage::BatchGet(
                             .count();
     VLOG(1) << "Time taken for FileStorage::BatchGet: " << elapsed_time
             << "us, key size: " << keys.size() << ", batch_id: " << batch_id;
+    pt_read.End(0);
     return batch_result;
 }
 
@@ -1010,16 +1034,23 @@ void FileStorage::ClientBufferGCThreadFunc() {
 }
 
 bool FileStorage::ReleaseBuffer(uint64_t batch_id) {
+    // Owner-side ClientBuffer release (OwnerReleaseBuffer); shared by the pull
+    // path (release_offload_buffer RPC) and the push path (inline after WRITE).
+    UbDiag::PerfPoint pt_release(PerfKey::GET_SSD_OWNER_RELEASE,
+                                 UbDiag::PerfLevel::MODULE);
+    pt_release.Start();
     MutexLocker locker(&client_buffer_mutex_);
     auto it = client_buffer_allocated_batches_.find(batch_id);
     if (it != client_buffer_allocated_batches_.end()) {
         VLOG(1) << "Releasing buffer for batch_id: " << batch_id
                 << " (transfer completed)";
         client_buffer_allocated_batches_.erase(it);
+        pt_release.End(0);
         return true;
     }
     VLOG(1) << "batch_id " << batch_id
             << " not found (may have been GC'd already)";
+    pt_release.End(-1);
     return false;
 }